Add option to pass coordinates to the grid method (#326)

Allow `BaseGridder.grid` to take the coordinates of the preexisting grid either
as 2d arrays (meshgrid) or as 1d arrays. Raise a `FutureWarning` if the
`spacing`, `shape` or `region` arguments are passed, since they will be removed
in v2.0.0. Write a new `meshgrid_from_1d` function that creates a meshgrid out
of 1d horizontal arrays, checks if they have a single dimension, and supports
for passing extra coordinates as 2d-arrays.

verde/base/base_classes.py

@@ -9,13 +9,19 @@
 """
 from abc import ABCMeta, abstractmethod
 
+import warnings
 import pandas as pd
 from sklearn.base import BaseEstimator
 from sklearn.model_selection import BaseCrossValidator
 
 from ..coordinates import grid_coordinates, profile_coordinates, scatter_points
 from .utils import check_data, check_data_names, score_estimator
-from ..utils import make_xarray_grid
+from ..utils import (
+    make_xarray_grid,
+    meshgrid_from_1d,
+    get_ndim_horizontal_coords,
+    check_meshgrid,
+)
 
 
 # Pylint doesn't like X, y scikit-learn argument names.
@@ -359,19 +365,24 @@ def grid(
         dims=None,
         data_names=None,
         projection=None,
-        **kwargs
+        coordinates=None,
+        **kwargs,
     ):  # pylint: disable=too-many-locals
         """
         Interpolate the data onto a regular grid.
 
-        The grid can be specified by either the number of points in each
-        dimension (the *shape*) or by the grid node spacing. See
-        :func:`verde.grid_coordinates` for details. Other arguments for
-        :func:`verde.grid_coordinates` can be passed as extra keyword arguments
-        (``kwargs``) to this method.
+        The grid can be specified by two methods:
 
-        If the interpolator collected the input data region, then it will be
-        used if ``region=None``. Otherwise, you must specify the grid region.
+        - Pass the actual *coordinates* of the grid points, as generated by
+          :func:`verde.grid_coordinates` or from an existing
+          :class:`xarray.Dataset` grid.
+        - Let the method define a new grid by either passing the number of
+          points in each dimension (the *shape*) or by the grid node *spacing*.
+          If the interpolator collected the input data region, then it will be
+          used if ``region=None``. Otherwise, you must specify the grid region.
+          See :func:`verde.grid_coordinates` for details. Other arguments for
+          :func:`verde.grid_coordinates` can be passed as extra keyword
+          arguments (``kwargs``) to this method.
 
         Use the *dims* and *data_names* arguments to set custom names for the
         dimensions and the data field(s) in the output :class:`xarray.Dataset`.
@@ -381,12 +392,15 @@ def grid(
         ----------
         region : list = [W, E, S, N]
             The west, east, south, and north boundaries of a given region.
+            Use only if ``coordinates`` is None.
         shape : tuple = (n_north, n_east) or None
             The number of points in the South-North and West-East directions,
             respectively.
+            Use only if ``coordinates`` is None.
         spacing : tuple = (s_north, s_east) or None
             The grid spacing in the South-North and West-East directions,
             respectively.
+            Use only if ``coordinates`` is None.
         dims : list or None
             The names of the northing and easting data dimensions,
             respectively, in the output grid. Default is determined from the
@@ -408,6 +422,13 @@ def grid(
             will be used to project the generated grid coordinates before
             passing them into ``predict``. For example, you can use this to
             generate a geographic grid from a Cartesian gridder.
+        coordinates : tuple of arrays
+            Tuple of arrays containing the coordinates of the grid in the
+            following order: (easting, northing, vertical, ...).
+            The easting and northing arrays could be 1d or 2d arrays, if
+            they are 2d they must be part of a meshgrid.
+            If coordinates are passed, ``region``, ``shape``, and ``spacing``
+            are ignored.
 
         Returns
         -------
@@ -420,16 +441,48 @@ def grid(
         verde.grid_coordinates : Generate the coordinate values for the grid.
 
         """
-        dims = self._get_dims(dims)
-        region = get_instance_region(self, region)
-        coordinates = grid_coordinates(region, shape=shape, spacing=spacing, **kwargs)
+        if coordinates is not None and (spacing is not None or shape is not None):
+            raise ValueError(
+                "Both coordinates and spacing or shape were provided. "
+                + "Please pass only coordinates or either the spacing or the shape."
+            )
+        if coordinates is not None and region is not None:
+            raise ValueError(
+                "Both coordinates and region were provided. "
+                + "Please pass region only if spacing or shape is specified."
+            )
+        # Raise deprecation warning for the region, shape and spacing arguments
+        if spacing is not None or shape is not None or region is not None:
+            warnings.warn(
+                "The 'spacing', 'shape' and 'region' arguments will be removed "
+                + "in Verde v2.0.0. "
+                + "Please use the 'verde.grid_coordinates' function to define "
+                + "grid coordinates and pass them as the 'coordinates' argument.",
+                FutureWarning,
+            )
+        # Get grid coordinates from coordinates parameter
+        if coordinates is not None:
+            ndim = get_ndim_horizontal_coords(*coordinates[:2])
+            if ndim == 1:
+                # Build a meshgrid if easting and northing are 1d
+                coordinates = meshgrid_from_1d(coordinates)
+            else:
+                check_meshgrid(coordinates)
+        # Build the grid coordinates through vd.grid_coordinates
+        else:
+            region = get_instance_region(self, region)
+            coordinates = grid_coordinates(
+                region, shape=shape, spacing=spacing, **kwargs
+            )
+        # Predict on the grid points (project the coordinates if needed)
         if projection is None:
             data = check_data(self.predict(coordinates))
         else:
             data = check_data(
                 self.predict(project_coordinates(coordinates, projection))
             )
-        # Get names for data and any extra coordinates
+        # Get names for dims, data and any extra coordinates
+        dims = self._get_dims(dims)
         data_names = self._get_data_names(data, data_names)
         extra_coords_names = self._get_extra_coords_names(coordinates)
         # Create xarray.Dataset
@@ -455,7 +508,7 @@ def scatter(
         dims=None,
         data_names=None,
         projection=None,
-        **kwargs
+        **kwargs,
     ):
         """
         Interpolate values onto a random scatter of points.
@@ -534,7 +587,7 @@ def profile(
         dims=None,
         data_names=None,
         projection=None,
-        **kwargs
+        **kwargs,
     ):
         """
         Interpolate data along a profile between two points.

verde/tests/test_base.py

@@ -8,6 +8,7 @@
 """
 Test the base classes and their utility functions.
 """
+import warnings
 import numpy as np
 import numpy.testing as npt
 import pytest
@@ -134,16 +135,28 @@ def test_basegridder():
         # A region should be given because it hasn't been assigned
         grd.grid()
 
+    npt.assert_allclose(grd.coefs_, [linear, angular])
+
+    # Check predictions by comparing against expected values
     coordinates_true = grid_coordinates(region, shape=shape)
     data_true = angular * coordinates_true[0] + linear
-    grid = grd.grid(region=region, shape=shape)
+    # Grid passing region and shape
+    grids = []
+    grids.append(grd.grid(region=region, shape=shape))
+    # Grid passing grid coordinates
+    grids.append(grd.grid(coordinates=coordinates_true))
+    # Grid passing grid coordinates as 1d arrays
+    grids.append(grd.grid(coordinates=tuple(np.unique(c) for c in coordinates_true)))
+    # Grid on profile
     prof = grd.profile((0, -10), (10, -10), 30)
+    # Grid on scatter
+    scat = grd.scatter(region=region, size=1000, random_state=0)
 
-    npt.assert_allclose(grd.coefs_, [linear, angular])
-    npt.assert_allclose(grid.scalars.values, data_true)
-    npt.assert_allclose(grid.easting.values, coordinates_true[0][0, :])
-    npt.assert_allclose(grid.northing.values, coordinates_true[1][:, 0])
-    npt.assert_allclose(grd.scatter(region, 1000, random_state=0).scalars, data)
+    for grid in grids:
+        npt.assert_allclose(grid.scalars.values, data_true)
+        npt.assert_allclose(grid.easting.values, coordinates_true[0][0, :])
+        npt.assert_allclose(grid.northing.values, coordinates_true[1][:, 0])
+    npt.assert_allclose(scat.scalars, data)
     npt.assert_allclose(
         prof.scalars,
         angular * coordinates_true[0][0, :] + linear,
@@ -356,6 +369,36 @@ def proj(lon, lat, inverse=False):
     npt.assert_allclose(prof.distance, distance_true)
 
 
+def test_basegridder_grid_invalid_arguments():
+    """
+    Test if errors and warnings are raised on invalid arguments to grid method
+    """
+    region = (0, 10, -10, -5)
+    angular, linear = 2, 100
+    coordinates = scatter_points(region, 1000, random_state=0, extra_coords=(1, 2))
+    data = angular * coordinates[0] + linear
+    grd = PolyGridder().fit(coordinates, data)
+    # Check error is raised if coordinates and shape are passed
+    grid_coords = (np.linspace(*region[:2], 11), np.linspace(*region[2:], 7))
+    with pytest.raises(ValueError):
+        grd.grid(coordinates=grid_coords, shape=(30, 30))
+    # Check error is raised if coordinates and spacing are passed
+    with pytest.raises(ValueError):
+        grd.grid(coordinates=grid_coords, spacing=10)
+    # Check error is raised if both coordinates and region are passed
+    with pytest.raises(ValueError):
+        grd.grid(coordinates=grid_coords, region=region)
+    # Check if FutureWarning is raised after passing region, spacing or shape
+    with warnings.catch_warnings(record=True) as warns:
+        grd.grid(region=region, shape=(4, 4))
+        assert len(warns) == 1
+        assert issubclass(warns[0].category, FutureWarning)
+    with warnings.catch_warnings(record=True) as warns:
+        grd.grid(region=region, spacing=1)
+        assert len(warns) == 1
+        assert issubclass(warns[0].category, FutureWarning)
+
+
 def test_check_fit_input():
     "Make sure no exceptions are raised for standard cases"
     size = 20

verde/tests/test_utils.py

@@ -15,7 +15,7 @@
 from scipy.spatial import cKDTree  # pylint: disable=no-name-in-module
 import pytest
 
-from ..coordinates import grid_coordinates
+from ..coordinates import grid_coordinates, scatter_points
 from ..utils import (
     parse_engine,
     dummy_jit,
@@ -24,6 +24,8 @@
     partition_by_sum,
     make_xarray_grid,
     meshgrid_to_1d,
+    meshgrid_from_1d,
+    get_ndim_horizontal_coords,
 )
 from .. import utils
 
@@ -306,3 +308,30 @@ def test_meshgrid_to_1d_invalid():
     northing = np.linspace(-4, -4, 9)
     with pytest.raises(ValueError):
         meshgrid_to_1d((easting, northing))
+
+
+def test_meshgrid_from_1d_invalid():
+    """
+    Check if error is raised after non 1d arrays passed to meshgrid_from_1d
+    """
+    coordinates = grid_coordinates(region=(0, 10, -5, 5), shape=(11, 11))
+    with pytest.raises(ValueError):
+        meshgrid_from_1d(coordinates)
+
+
+def test_check_ndim_easting_northing():
+    """
+    Test if check_ndim_easting_northing works as expected
+    """
+    # Easting and northing as 1d arrays
+    # pylint: disable=unbalanced-tuple-unpacking
+    easting, northing = scatter_points((-5, 5, 0, 4), 50, random_state=42)
+    assert get_ndim_horizontal_coords(easting, northing) == 1
+    # Easting and northing as 2d arrays
+    easting, northing = grid_coordinates((-5, 5, 0, 4), spacing=1)
+    assert get_ndim_horizontal_coords(easting, northing) == 2
+    # Check if error is raised after easting and northing with different ndims
+    easting = np.linspace(0, 5, 6)
+    northing = np.linspace(-5, 5, 16).reshape(4, 4)
+    with pytest.raises(ValueError):
+        get_ndim_horizontal_coords(easting, northing)

verde/utils.py

@@ -341,13 +341,7 @@ def make_xarray_grid(
 
     """
     # 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."
-        )
+    ndim = get_ndim_horizontal_coords(*coordinates[:2])
     # Convert 2d horizontal coordinates to 1d arrays if needed
     if ndim == 2:
         coordinates = meshgrid_to_1d(coordinates)
@@ -418,6 +412,92 @@ def meshgrid_to_1d(coordinates):
     return coordinates
 
 
+def meshgrid_from_1d(coordinates):
+    """
+    Convert horizontal coordinates of 2d grids from 1d-arrays to 2d-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. The horizontal coordinates (easting and northing) should
+        be 1d arrays, any extra coordinate should be an array with a shape of
+        ``(northing.size, easting.size)``.
+
+    Returns
+    -------
+    coordinates : tuple of arrays
+        Arrays with coordinates of each point in the grid.
+        The horizontal coordinates have been converted to 2d-arrays with the
+        same shape, forming a meshgrid. All extra coordinates have not been
+        modified.
+
+    Examples
+    --------
+
+    >>> easting = np.linspace(0, 4, 5)
+    >>> northing = np.linspace(-3, 3, 7)
+    >>> height = 2 * np.ones((7, 5))
+    >>> coordinates = (easting, northing, height)
+    >>> easting, northing, height = meshgrid_from_1d(coordinates)
+    >>> print(easting)
+    [[0. 1. 2. 3. 4.]
+     [0. 1. 2. 3. 4.]
+     [0. 1. 2. 3. 4.]
+     [0. 1. 2. 3. 4.]
+     [0. 1. 2. 3. 4.]
+     [0. 1. 2. 3. 4.]
+     [0. 1. 2. 3. 4.]]
+    >>> print(northing)
+    [[-3. -3. -3. -3. -3.]
+     [-2. -2. -2. -2. -2.]
+     [-1. -1. -1. -1. -1.]
+     [ 0.  0.  0.  0.  0.]
+     [ 1.  1.  1.  1.  1.]
+     [ 2.  2.  2.  2.  2.]
+     [ 3.  3.  3.  3.  3.]]
+
+    """
+    ndim = get_ndim_horizontal_coords(*coordinates[:2])
+    if ndim != 1:
+        raise ValueError(
+            "Horizontal coordinates must be 1d-arrays. " + f"{ndim}d-arrays provided."
+        )
+    easting, northing = np.meshgrid(coordinates[0], coordinates[1])
+    coordinates = (easting, northing, *coordinates[2:])
+    coordinates = check_coordinates(coordinates)
+    return coordinates
+
+
+def get_ndim_horizontal_coords(easting, northing):
+    """
+    Return the number of dimensions of the horizontal coordinates arrays
+
+    Also check if the two horizontal coordinates arrays same dimensions.
+
+    Parameters
+    ----------
+    easting : nd-array
+        Array for the easting coordinates
+    northing : nd-array
+        Array for the northing coordinates
+
+    Returns
+    -------
+    ndim : int
+        Number of dimensions of the ``easting`` and ``northing`` arrays.
+    """
+    ndim = np.ndim(easting)
+    if ndim != np.ndim(northing):
+        raise ValueError(
+            "Horizontal coordinates dimensions mismatch. "
+            + f"The easting coordinate array has {easting.ndim} dimensions "
+            + f"while the northing has {northing.ndim}."
+        )
+    return ndim
+
+
 def check_meshgrid(coordinates):
     """
     Check if the given horizontal coordinates define a meshgrid