diff --git a/verde/base/base_classes.py b/verde/base/base_classes.py index c0e944f57..40e15fbf4 100644 --- a/verde/base/base_classes.py +++ b/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. diff --git a/verde/tests/test_base.py b/verde/tests/test_base.py index d637bd4dc..ad8ea59c8 100644 --- a/verde/tests/test_base.py +++ b/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