Add option to return 1D arrays in grid_coordinates

verde/coordinates.py

@@ -195,6 +195,7 @@ def grid_coordinates(
     adjust="spacing",
     pixel_register=False,
     extra_coords=None,
+    meshgrid=True,
 ):
     """
     Generate the coordinates for each point on a regular grid.
@@ -235,6 +236,14 @@ def grid_coordinates(
         contain a constant value. Will generate an extra array per value given
         in *extra_coords*. Use this to generate arrays of constant heights or
         times, for example, that might be needed to evaluate a gridder.
+    meshgrid : bool
+        If True, will call :func:`numpy.meshgrid` on the generated coordinates
+        and return 2D arrays (useful if you need the coordinate values for
+        every single point on a grid). Otherwise, will return 1D coordinate
+        arrays (useful if you're making a :class:`xarray.DataArray` or looping
+        over grid coordinates with two ``for`` loops). Passing False to
+        *meshgrid* is **incompatible with *extra_coords* and an exception will
+        be raised** if used together (:class:`ValueError`).
 
     Returns
     -------
@@ -284,6 +293,18 @@ def grid_coordinates(
      [ 7.5  7.5  7.5]
      [10.  10.  10. ]]
 
+    If you don't need the 2D arrays, then use ``meshgrid=False``:
+
+    >>> east, north = grid_coordinates(
+    ...     region=(0, 5, 0, 10), spacing=2.5, meshgrid=False,
+    ... )
+    >>> print(east.shape, north.shape)
+    (3,) (5,)
+    >>> print(east)
+    [0.  2.5 5. ]
+    >>> print(north)
+    [ 0.   2.5  5.   7.5 10. ]
+
     The spacing can be different for northing and easting, respectively:
 
     >>> east, north = grid_coordinates(region=(-5, 1, 0, 10), spacing=(2.5, 1))
@@ -444,8 +465,15 @@ def grid_coordinates(
     if pixel_register:
         east_lines = east_lines[:-1] + (east_lines[1] - east_lines[0]) / 2
         north_lines = north_lines[:-1] + (north_lines[1] - north_lines[0]) / 2
-    coordinates = list(np.meshgrid(east_lines, north_lines))
+    if meshgrid:
+        coordinates = list(np.meshgrid(east_lines, north_lines))
+    else:
+        coordinates = (east_lines, north_lines)
     if extra_coords is not None:
+        if not meshgrid:
+            raise ValueError(
+                "Using 'meshgrid=False' and 'extra_coords' is 'verde.grid_coordinates' is not allowed."
+            )
         for value in np.atleast_1d(extra_coords):
             coordinates.append(np.ones_like(coordinates[0]) * value)
     return tuple(coordinates)

verde/tests/test_coordinates.py

@@ -283,3 +283,11 @@ def test_invalid_geographic_coordinates():
     latitude[0] = 100
     with pytest.raises(ValueError):
         longitude_continuity([longitude, latitude], region)
+
+
+def test_meshgrid_extra_coords_error():
+    "Should raise an exception if meshgrid=False and extra_coords are used"
+    with pytest.raises(ValueError):
+        grid_coordinates(
+            region=(0, 1, 0, 3), spacing=0.1, meshgrid=False, extra_coords=10
+        )