diff options
| -rw-r--r-- | numpy/lib/function_base.py | 22 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 39 |
2 files changed, 33 insertions, 28 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 01c44eec5..63b191b07 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -3125,11 +3125,14 @@ def add_newdoc(place, obj, doc): # Based on scitools meshgrid def meshgrid(*xi, **kwargs): """ - Return coordinate matrices from two or more coordinate vectors. + Return coordinate matrices from coordinate vectors. Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector fields over N-D grids, given one-dimensional coordinate arrays x1, x2,..., xn. + + .. versionchanged:: 1.9 + 1-D and 0-D cases are allowed. Parameters ---------- @@ -3178,6 +3181,8 @@ def meshgrid(*xi, **kwargs): for i in range(nx): for j in range(ny): # treat xv[j,i], yv[j,i] + + In the 1-D and 0-D case, the indexing and sparse keywords have no effect. See Also -------- @@ -3214,28 +3219,23 @@ def meshgrid(*xi, **kwargs): >>> h = plt.contourf(x,y,z) """ - if len(xi) < 2: - raise ValueError( - 'meshgrid() takes 2 or more arguments ' - '(%d given)' % int(len(xi) > 0)) - - args = np.atleast_1d(*xi) - ndim = len(args) + ndim = len(xi) copy_ = kwargs.get('copy', True) sparse = kwargs.get('sparse', False) indexing = kwargs.get('indexing', 'xy') + if not indexing in ['xy', 'ij']: raise ValueError( "Valid values for `indexing` are 'xy' and 'ij'.") s0 = (1,) * ndim - output = [x.reshape(s0[:i] + (-1,) + s0[i + 1::]) - for i, x in enumerate(args)] + output = [np.asanyarray(x).reshape(s0[:i] + (-1,) + s0[i + 1::]) + for i, x in enumerate(xi)] shape = [x.size for x in output] - if indexing == 'xy': + if indexing == 'xy' and ndim > 1: # switch first and second axis output[0].shape = (1, -1) + (1,)*(ndim - 2) output[1].shape = (-1, 1) + (1,)*(ndim - 2) diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index d1069cf92..2255c0b36 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -1311,28 +1311,33 @@ class TestMsort(TestCase): class TestMeshgrid(TestCase): def test_simple(self): [X, Y] = meshgrid([1, 2, 3], [4, 5, 6, 7]) - assert_(np.all(X == np.array([[1, 2, 3], - [1, 2, 3], - [1, 2, 3], - [1, 2, 3]]))) - assert_(np.all(Y == np.array([[4, 4, 4], - [5, 5, 5], - [6, 6, 6], - [7, 7, 7]]))) + assert_array_equal(X, np.array([[1, 2, 3], + [1, 2, 3], + [1, 2, 3], + [1, 2, 3]])) + assert_array_equal(Y, np.array([[4, 4, 4], + [5, 5, 5], + [6, 6, 6], + [7, 7, 7]])) def test_single_input(self): - assert_raises(ValueError, meshgrid, np.arange(5)) + [X] = meshgrid([1, 2, 3, 4]) + assert_array_equal(X, np.array([1, 2, 3, 4])) + + def test_no_input(self): + args = [] + assert_array_equal([], meshgrid(*args)) def test_indexing(self): x = [1, 2, 3] y = [4, 5, 6, 7] [X, Y] = meshgrid(x, y, indexing='ij') - assert_(np.all(X == np.array([[1, 1, 1, 1], - [2, 2, 2, 2], - [3, 3, 3, 3]]))) - assert_(np.all(Y == np.array([[4, 5, 6, 7], - [4, 5, 6, 7], - [4, 5, 6, 7]]))) + assert_array_equal(X, np.array([[1, 1, 1, 1], + [2, 2, 2, 2], + [3, 3, 3, 3]])) + assert_array_equal(Y, np.array([[4, 5, 6, 7], + [4, 5, 6, 7], + [4, 5, 6, 7]])) # Test expected shapes: z = [8, 9] @@ -1345,8 +1350,8 @@ class TestMeshgrid(TestCase): def test_sparse(self): [X, Y] = meshgrid([1, 2, 3], [4, 5, 6, 7], sparse=True) - assert_(np.all(X == np.array([[1, 2, 3]]))) - assert_(np.all(Y == np.array([[4], [5], [6], [7]]))) + assert_array_equal(X, np.array([[1, 2, 3]])) + assert_array_equal(Y, np.array([[4], [5], [6], [7]])) class TestPiecewise(TestCase): |