Gridders apply projections using only the first two coordinates

verde/base/base_classes.py

@@ -300,7 +300,9 @@ def grid(
         if projection is None:
             data = check_data(self.predict(coordinates))
         else:
-            data = check_data(self.predict(projection(*coordinates)))
+            data = check_data(
+                self.predict(project_coordinates(coordinates, projection))
+            )
         data_names = get_data_names(data, data_names)
         coords = {dims[1]: coordinates[0][0, :], dims[0]: coordinates[1][:, 0]}
         attrs = {"metadata": "Generated by {}".format(repr(self))}
@@ -378,7 +380,9 @@ def scatter(
         if projection is None:
             data = check_data(self.predict(coordinates))
         else:
-            data = check_data(self.predict(projection(*coordinates)))
+            data = check_data(
+                self.predict(project_coordinates(coordinates, projection))
+            )
         data_names = get_data_names(data, data_names)
         columns = [(dims[0], coordinates[1]), (dims[1], coordinates[0])]
         columns.extend(zip(data_names, data))
@@ -480,14 +484,14 @@ def profile(
         # geographic coordinates since we don't do actual distances on a
         # sphere).
         if projection is not None:
-            point1 = projection(*point1)
-            point2 = projection(*point2)
+            point1 = project_coordinates(point1, projection)
+            point2 = project_coordinates(point2, projection)
         coordinates, distances = profile_coordinates(point1, point2, size, **kwargs)
         data = check_data(self.predict(coordinates))
         # Project the coordinates back to have geographic coordinates for the
         # profile but Cartesian distances.
         if projection is not None:
-            coordinates = projection(*coordinates, inverse=True)
+            coordinates = project_coordinates(coordinates, projection, inverse=True)
         data_names = get_data_names(data, data_names)
         columns = [
             (dims[0], coordinates[1]),
@@ -506,6 +510,42 @@ def _get_dims(self, dims):
         return self.dims
 
 
+def project_coordinates(coordinates, projection, **kwargs):
+    """
+    Apply projection to given coordiantes
+
+    Allows to apply projections to any number of coordinates, assuming
+    that the first ones are ``longitude`` and ``latitude``.
+
+    Examples
+    --------
+
+    >>> # Define a custom projection function
+    >>> def projection(lon, lat, inverse=False):
+    ...     "Simple projection of geographic coordinates"
+    ...     radius = 1000
+    ...     if inverse:
+    ...         return (lon / radius, lat / radius)
+    ...     return (lon * radius, lat * radius)
+
+    >>> # Apply the projection to a set of coordinates containing:
+    >>> # longitude, latitude and height
+    >>> coordinates = (1., -2., 3.)
+    >>> project_coordinates(coordinates, projection)
+    (1000.0, -2000.0, 3.0)
+
+    >>> # Apply the inverse projection
+    >>> coordinates = (-500.0, 1500.0, -19.0)
+    >>> project_coordinates(coordinates, projection, inverse=True)
+    (-0.5, 1.5, -19.0)
+
+    """
+    proj_coordinates = projection(*coordinates[:2], **kwargs)
+    if len(coordinates) > 2:
+        proj_coordinates += tuple(coordinates[2:])
+    return proj_coordinates
+
+
 def get_data_names(data, data_names):
     """
     Get default names for data fields if none are given based on the data.

verde/tests/test_base.py

@@ -201,6 +201,70 @@ def proj(lon, lat, inverse=False):
     npt.assert_allclose(prof.distance, distance_true)
 
 
+def test_basegridder_projection_multiple_coordinates():
+    "Test BaseGridder when passing in a projection with multiple coordinates"
+
+    # Lets say the projection is doubling the coordinates
+    def proj(lon, lat, inverse=False):
+        "Project from the new coordinates to the original"
+        if inverse:
+            return (lon / 2, lat / 2)
+        return (lon * 2, lat * 2)
+
+    # Values in "geographic" coordinates
+    region = (0, 10, -10, -5)
+    shape = (51, 31)
+    angular, linear = 2, 100
+    coordinates = scatter_points(region, 1000, random_state=0, extra_coords=(1, 2))
+    data = angular * coordinates[0] + linear
+    # Project before passing to our Cartesian gridder
+    proj_coordinates = proj(coordinates[0], coordinates[1]) + coordinates[2:]
+    grd = PolyGridder().fit(proj_coordinates, data)
+
+    # Check the estimated coefficients
+    # The grid is estimated in projected coordinates (which are twice as large)
+    # so the rate of change (angular) should be half to get the same values.
+    npt.assert_allclose(grd.coefs_, [linear, angular / 2])
+
+    # The actual values for a grid
+    coordinates_true = grid_coordinates(region, shape, extra_coords=(13, 17))
+    data_true = angular * coordinates_true[0] + linear
+
+    # Check the scatter
+    scat = grd.scatter(
+        region, 1000, random_state=0, projection=proj, extra_coords=(13, 17)
+    )
+    npt.assert_allclose(scat.scalars, data)
+    npt.assert_allclose(scat.easting, coordinates[0])
+    npt.assert_allclose(scat.northing, coordinates[1])
+
+    # Check the grid
+    grid = grd.grid(region, shape, projection=proj, extra_coords=(13, 17))
+    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])
+
+    # Check the profile
+    prof = grd.profile(
+        (region[0], region[-1]),
+        (region[1], region[-1]),
+        shape[1],
+        projection=proj,
+        extra_coords=(13, 17),
+    )
+    npt.assert_allclose(prof.scalars, data_true[-1, :])
+    # Coordinates should still be evenly spaced since the projection is a
+    # multiplication.
+    npt.assert_allclose(prof.easting, coordinates_true[0][0, :])
+    npt.assert_allclose(prof.northing, coordinates_true[1][-1, :])
+    # Distance should still be in the projected coordinates. If the projection
+    # is from geographic, we shouldn't be returning distances in degrees but in
+    # projected meters. The distances will be evenly spaced in unprojected
+    # coordinates.
+    distance_true = np.linspace(region[0] * 2, region[1] * 2, shape[1])
+    npt.assert_allclose(prof.distance, distance_true)
+
+
 def test_check_fit_input():
     "Make sure no exceptions are raised for standard cases"
     size = 20