Allow make_xarray_grid to take horizontal coordinates as 1d arrays (#300)

Now make_xarray_grid could take 1d arrays for horizontal coordinates, such as
how xr.DataArray stores the single dimension coordinates. Add a new
non-public meshgrid_to_1d function to convert 2d meshgrids (like the ones
generated by vd.grid_coordinates) to 1d arrays containing a single coordinate
of the corresponding axis. This function checks if all elements inside the 2d
array are equal along the right axis, i.e. a vaild meshgrid.

Author: santisoler

verde/__init__.py

@@ -21,7 +21,12 @@
     longitude_continuity,
 )
 from .mask import distance_mask, convexhull_mask
-from .utils import variance_to_weights, maxabs, grid_to_table, make_xarray_grid
+from .utils import (
+    variance_to_weights,
+    maxabs,
+    grid_to_table,
+    make_xarray_grid,
+)
 from .io import load_surfer
 from .distances import median_distance
 from .blockreduce import BlockReduce, BlockMean

verde/tests/test_utils.py

@@ -23,6 +23,7 @@
     grid_to_table,
     partition_by_sum,
     make_xarray_grid,
+    meshgrid_to_1d,
 )
 from .. import utils
 
@@ -218,3 +219,75 @@ def test_make_xarray_grid_invalid_extra_coords():
         make_xarray_grid(
             coordinates, data, data_names="dummy", extra_coords_names=["upward", "time"]
         )
+
+
+def test_make_xarray_grid_invalid_2d_coordinates():
+    """
+    Check if error is raised if invaild 2d coordinates array are passed
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    easting, northing = grid_coordinates(region, spacing=spacing)
+    # Change only one element of the easting array
+    easting[2, 2] = -1000
+    data = np.ones_like(easting)
+    with pytest.raises(ValueError):
+        make_xarray_grid((easting, northing), data, data_names="dummy")
+
+
+def test_make_xarray_grid_coordinates_as_1d_arrays():
+    """
+    Check if it can handle coordinates as 1d-arrays
+    """
+    region = (-10, -5, 6, 10)
+    easting = np.linspace(*region[:2], 6, dtype=float)
+    northing = np.linspace(*region[2:], 5, dtype=float)
+    data = np.ones((northing.size, easting.size))
+    grid = make_xarray_grid((easting, northing), data, data_names="dummy")
+    npt.assert_allclose(grid.easting, [-10, -9, -8, -7, -6, -5])
+    npt.assert_allclose(grid.northing, [6, 7, 8, 9, 10])
+    npt.assert_allclose(grid.dummy, 1)
+    assert grid.dummy.shape == (5, 6)
+
+
+def test_make_xarray_grid_invalid_mixed_coordinates():
+    """
+    Check if error is raised when horizontal coordinates have mixed dimensions
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    easting, northing = grid_coordinates(region, spacing=spacing)
+    data = np.ones_like(easting)
+    # easting is 1d, but northing is 2d
+    with pytest.raises(ValueError):
+        make_xarray_grid((easting[0, :], northing), data, data_names="dummy")
+    # northing is 1d, but easting is 2d
+    with pytest.raises(ValueError):
+        make_xarray_grid((easting, northing[:, 0]), data, data_names="dummy")
+
+
+def test_meshgrid_to_1d_invalid():
+    """
+    Check if error is raised after invalid meshgrid
+    """
+    region = (-10, -5, 6, 10)
+    # Modify one element of easting
+    easting, northing = grid_coordinates(region=region, spacing=1)
+    easting[2, 2] = -9999
+    with pytest.raises(ValueError):
+        meshgrid_to_1d((easting, northing))
+    # Modify one element of northing
+    easting, northing = grid_coordinates(region=region, spacing=1)
+    northing[2, 3] = -9999
+    with pytest.raises(ValueError):
+        meshgrid_to_1d((easting, northing))
+    # Pass invalid shapes
+    easting = np.arange(16).reshape(4, 4)
+    northing = np.arange(9).reshape(3, 3)
+    with pytest.raises(ValueError):
+        meshgrid_to_1d((easting, northing))
+    # Pass 1d arrays
+    easting = np.linspace(0, 10, 11)
+    northing = np.linspace(-4, -4, 9)
+    with pytest.raises(ValueError):
+        meshgrid_to_1d((easting, northing))

verde/utils.py

@@ -317,12 +317,22 @@ def make_xarray_grid(
         dummy     (northing, easting) float64 1.0 1.0 1.0 1.0 1.0 1.0
 
     """
-    coordinates = check_coordinates(coordinates)
+    # Check dimensions of the horizontal coordinates of the regular grid
+    ndim = np.ndim(coordinates[0])
+    if ndim != np.ndim(coordinates[1]):
+        raise ValueError(
+            "Incompatible dimensions between horizontal coordinates of the regular grid: "
+            + f"'{ndim}' and '{np.ndim(coordinates[1])}'. "
+            + "Both coordinates must have the same number of dimensions."
+        )
+    # Convert 2d horizontal coordinates to 1d arrays if needed
+    if ndim == 2:
+        coordinates = meshgrid_to_1d(coordinates)
     data = check_data(data)
     data_names = check_data_names(data, data_names)
     # dims is like shape with order (rows, cols) for the array
     # so the first element is northing and second is easting
-    coords = {dims[1]: coordinates[0][0, :], dims[0]: coordinates[1][:, 0]}
+    coords = {dims[1]: coordinates[0], dims[0]: coordinates[1]}
     # Extra coordinates are handled like 2D data arrays with
     # the same dims and the data.
     if coordinates[2:]:
@@ -333,6 +343,86 @@ def make_xarray_grid(
     return xr.Dataset(data_vars, coords)
 
 
+def meshgrid_to_1d(coordinates):
+    """
+    Convert horizontal coordinates of 2d grids into 1d-arrays
+
+    Parameters
+    ----------
+    coordinates : tuple of arrays
+        Arrays with coordinates of each point in the grid. Each array must
+        contain values for a dimension in the order: easting, northing,
+        vertical, etc. All arrays must be 2d and need to have the same
+        *shape*. The horizontal coordinates should be actual meshgrids.
+
+    Returns
+    -------
+    coordinates : tuple of arrays
+        Arrays with coordinates of each point in the grid. The horizontal
+        coordinates have been converted to 1d-arrays, having only a single
+        coordinate point per its corresponding axis.
+        All extra coordinates have not been modified.
+
+    Examples
+    --------
+
+    >>> import verde as vd
+    >>> coordinates = vd.grid_coordinates(
+    ...     region=(0, 4, -3, 3), spacing=1, extra_coords=2
+    ... )
+    >>> easting, northing, height = meshgrid_to_1d(coordinates)
+    >>> print(easting)
+    [0. 1. 2. 3. 4.]
+    >>> print(northing)
+    [-3. -2. -1.  0.  1.  2.  3.]
+    >>> print(height)
+    [[2. 2. 2. 2. 2.]
+     [2. 2. 2. 2. 2.]
+     [2. 2. 2. 2. 2.]
+     [2. 2. 2. 2. 2.]
+     [2. 2. 2. 2. 2.]
+     [2. 2. 2. 2. 2.]
+     [2. 2. 2. 2. 2.]]
+    """
+    check_coordinates(coordinates)
+    check_meshgrid(coordinates)
+    easting, northing = coordinates[0][0, :], coordinates[1][:, 0]
+    coordinates = (easting, northing, *coordinates[2:])
+    return coordinates
+
+
+def check_meshgrid(coordinates):
+    """
+    Check if the given horizontal coordinates define a meshgrid
+
+    Check if the rows of the easting 2d-array are identical. Check if the
+    columns of the northing 2d-array are identical. This function does not
+    check if the easting and northing coordinates are evenly spaced.
+
+    Parameters
+    ----------
+    coordinates : tuple of arrays
+        Arrays with coordinates of each point in the grid. Each array must
+        contain values for a dimension in the order: easting, northing,
+        vertical, etc. Only easting and northing will be checked, the other
+        ones will be ignored. All arrays must be 2d and need to have the same
+        *shape*.
+    """
+    # Get the two first arrays as easting and northing
+    easting, northing = coordinates[:2]
+    # Check if all elements of easting along the zeroth axis are equal
+    msg = (
+        "Invalid coordinate array. The arrays for the horizontal "
+        + "coordinates of a regular grid must be meshgrids."
+    )
+    if not np.allclose(easting[0, :], easting):
+        raise ValueError(msg)
+    # Check if all elements of northing along the first axis are equal
+    # (need to make northing[:, 0] a vertical array so numpy can compare)
+    if not np.allclose(northing[:, 0][:, None], northing):
+        raise ValueError(msg)
+
+
 def grid_to_table(grid):
     """
     Convert a grid to a table with the values and coordinates of each point.