Merge pull request numpy#4112 from rhewett/meshgrid_single

seberg · seberg · commit 0856259fde53 · 2013-12-22T12:39:42.000-08:00
ENH: Allow meshgrid to take 1D and 0D inputs
diff --git 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
@@ -1311,28 +1311,33 @@ def test_simple(self):
 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 @@ def test_indexing(self):
 
     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):
