From ef07916ce8463c9a926386747715bc0d2aa71dfc Mon Sep 17 00:00:00 2001
From: Leonardo Uieda <leouieda@gmail.com>
Date: Mon, 24 Oct 2022 14:11:44 +0100
Subject: [PATCH] Add option to return 1D arrays in grid_coordinates

The meshgrid=True option enables and disables calling numpy.meshgrid on
the output arrays. Can be useful when building xarray grids, where we
end up making the 2D arrays and then extracting the 1D arrays.
---
 verde/coordinates.py            | 30 +++++++++++++++++++++++++++++-
 verde/tests/test_coordinates.py |  8 ++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/verde/coordinates.py b/verde/coordinates.py
index 067b1fc73..b46b504b8 100644
--- a/verde/coordinates.py
+++ b/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)
diff --git a/verde/tests/test_coordinates.py b/verde/tests/test_coordinates.py
index 38cfe7198..b61f29ed0 100644
--- a/verde/tests/test_coordinates.py
+++ b/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
+        )