Repository: fatiando/verde
Pull request #70 (2018-05-28 13:23:41)

Refactor distance_mask to be more useful
Invert the mask to make it compatible with xarray.Dataset.where without
having to negate it (~).
Allow the function to take coordinates or parameters to generate a
regular grid automatically. This way we don't need to meshgrid the grid
coordinates and just give region and spacing.


Patch:
diff --git a/examples/distance_mask.py b/examples/distance_mask.py
index e67ca0c..3121170 100644
--- a/examples/distance_mask.py
+++ b/examples/distance_mask.py
@@ -24,13 +24,15 @@ coordinates = vd.grid_coordinates(region, spacing=spacing)
 dummy_data = np.ones_like(coordinates[0])
 
 # Generate a mask for points that are more than 2 grid spacings away from any
-# data point. The mask is True for points that are too far.
-mask = vd.distance_mask(coordinates, (data.longitude, data.latitude),
-                        maxdist=spacing*2)
+# data point. The mask is True for points that are within the maximum distance.
+# Here, we'll provide the grid coordinates to the function but we could also
+# give it a region and spacing instead if we hadn't generated the coordinates.
+mask = vd.distance_mask((data.longitude, data.latitude), maxdist=spacing*2,
+                        coordinates=coordinates)
 print(mask)
 
 # Turn points that are too far into NaNs so they won't show up in our plot
-dummy_data[mask] = np.nan
+dummy_data[~mask] = np.nan
 
 # Make a plot of the masked data and the data locations.
 crs = ccrs.PlateCarree()
diff --git a/examples/spline.py b/examples/spline.py
index f88682a..832bc08 100644
--- a/examples/spline.py
+++ b/examples/spline.py
@@ -56,11 +56,11 @@ print("\nScore: {:.3f}".format(score))
 # Create a grid of the vertical velocity and mask it to only show points close
 # to the actual data.
 region = vd.get_region(coordinates)
+mask = vd.distance_mask((data.longitude, data.latitude), maxdist=0.5,
+                        region=region, spacing=spacing)
 grid = chain.grid(region=region, spacing=spacing, projection=projection,
                   dims=['latitude', 'longitude'], data_names=['velocity'])
-mask = vd.distance_mask(np.meshgrid(grid.longitude, grid.latitude),
-                        (data.longitude, data.latitude), maxdist=0.5)
-grid = grid.where(~mask)
+grid = grid.where(mask)
 
 
 def setup_map(ax):
diff --git a/examples/vector2d.py b/examples/vector2d.py
index 8d862f3..c51c3ec 100644
--- a/examples/vector2d.py
+++ b/examples/vector2d.py
@@ -50,14 +50,13 @@ chain.fit(*train)
 score = chain.score(*test)
 print("Cross-validation R^2 score: {:.2f}".format(score))
 
-# Interpolate our horizontal GPS velocities onto a regular geographic grid
+# Interpolate our horizontal GPS velocities onto a regular geographic grid and
+# mask the data that are far from the observation points
 grid = chain.grid(region, spacing=spacing, projection=projection,
                   dims=['latitude', 'longitude'])
-# Mask the data that are far from the observation points
-mask = vd.distance_mask(np.meshgrid(grid.longitude, grid.latitude),
-                        (data.longitude, data.latitude),
-                        maxdist=0.5)
-grid = grid.where(~mask)
+mask = vd.distance_mask((data.longitude, data.latitude), maxdist=0.5,
+                        region=region, spacing=spacing)
+grid = grid.where(mask)
 
 # Calculate residuals between the predictions and the original input data.
 # Even though we aren't using regularization or regularly distributed forces,
diff --git a/verde/grid_math.py b/verde/grid_math.py
index c8d504e..312f3b3 100644
--- a/verde/grid_math.py
+++ b/verde/grid_math.py
@@ -4,25 +4,44 @@ Operations on spatial data: block operations, derivatives, etc.
 import numpy as np
 from scipy.spatial import cKDTree  # pylint: disable=no-name-in-module
 
+from .coordinates import grid_coordinates
 
-def distance_mask(coordinates, data_coordinates, maxdist):
+
+def distance_mask(data_coordinates, maxdist, coordinates=None, region=None,
+                  spacing=None, shape=None, **kwargs):
     """
     Create a mask for points that are too far from the given data points.
 
-    Produces a mask array that is True when a point is more than *maxdist* from
-    the closest data point.
+    Produces a mask array that is False when a point is more than *maxdist*
+    from the closest data point and True otherwise.
+
+    The points that will be masked can be specified by their coordinates or by
+    a region and shape/spacing, in which case a regular grid will be generated
+    for the mask. If specifying a region, extra keyword arguments to this
+    function will be passed along to :func:`verde.grid_coordinates`.
 
     Parameters
     ----------
-    coordinates : tuple of arrays
-        Arrays with the coordinates of each point that will be masked. Should
-        be in the following order: (easting, northing, vertical, ...). Only
-        easting and northing will be used, all subsequent coordinates will be
-        ignored.
     data_coordinates : tuple of arrays
         Same as *coordinates* but for the data points.
     maxdist : float
         The maximum distance that a point can be from the closest data point.
+    coordinates : None or tuple of arrays
+        Arrays with the coordinates of each point that will be masked. Should
+        be in the following order: (easting, northing, vertical, ...). Only
+        easting and northing will be used, all subsequent coordinates will be
+        ignored. If not given, the *region* and *spacing* or *shape* must be
+        provided.
+    region : list = [W, E, S, N] or None
+        The boundaries of a given region in Cartesian or geographic
+        coordinates.
+    shape : tuple = (n_north, n_east) or None
+        The number of points in the South-North and West-East directions,
+        respectively.
+    spacing : float, tuple = (s_north, s_east), or None
+        The grid spacing in the South-North and West-East directions,
+        respectively. A single value means that the spacing is equal in both
+        directions.
 
     Returns
     -------
@@ -33,17 +52,34 @@ def distance_mask(coordinates, data_coordinates, maxdist):
     --------
 
     >>> from verde import grid_coordinates
-    >>> coords = grid_coordinates((0, 5, -10, -5), spacing=1)
-    >>> mask = distance_mask(coords, (2.5, -7.5), maxdist=2)
+    >>> region = (0, 5, -10, -5)
+    >>> spacing = 1
+    >>> coords = grid_coordinates(region, spacing=spacing)
+    >>> mask = distance_mask((2.5, -7.5), maxdist=2, coordinates=coords)
+    >>> print(mask)
+    [[False False False False False False]
+     [False False  True  True False False]
+     [False  True  True  True  True False]
+     [False  True  True  True  True False]
+     [False False  True  True False False]
+     [False False False False False False]]
+    >>> mask = distance_mask((3.5, -7.5), maxdist=2, region=region,
+    ...                      spacing=spacing)
     >>> print(mask)
-    [[ True  True  True  True  True  True]
-     [ True  True False False  True  True]
-     [ True False False False False  True]
-     [ True False False False False  True]
-     [ True  True False False  True  True]
-     [ True  True  True  True  True  True]]
+    [[False False False False False False]
+     [False False False  True  True False]
+     [False False  True  True  True  True]
+     [False False  True  True  True  True]
+     [False False False  True  True False]
+     [False False False False False False]]
 
     """
+    if coordinates is None:
+        if region is None:
+            raise ValueError("Either coordinates or region and shape/spacing "
+                             "must be given to generate the mask.")
+        coordinates = grid_coordinates(region=region, spacing=spacing,
+                                       shape=shape, **kwargs)
     data_easting, data_northing = data_coordinates[:2]
     easting, northing = coordinates[:2]
     data_easting = np.atleast_1d(data_easting)
@@ -52,4 +88,4 @@ def distance_mask(coordinates, data_coordinates, maxdist):
     tree = cKDTree(data_points)
     points = np.transpose((easting.ravel(), northing.ravel()))
     distance = tree.query(points)[0].reshape(easting.shape)
-    return distance > maxdist
+    return distance <= maxdist


Test patch:
diff --git a/verde/tests/test_grid_math.py b/verde/tests/test_grid_math.py
new file mode 100644
index 0000000..ba454cf
--- /dev/null
+++ b/verde/tests/test_grid_math.py
@@ -0,0 +1,27 @@
+"""
+Test the functions in verde.grid_math
+"""
+import pytest
+
+from ..grid_math import distance_mask
+
+
+def test_distance_mask():
+    "Check that the mask works for basic input"
+    region = (0, 5, -10, -5)
+    spacing = 1
+    mask = distance_mask((2.5, -7.5), maxdist=2, region=region,
+                         spacing=spacing)
+    true = [[False, False, False, False, False, False],
+            [False, False, True, True, False, False],
+            [False, True, True, True, True, False],
+            [False, True, True, True, True, False],
+            [False, False, True, True, False, False],
+            [False, False, False, False, False, False]]
+    assert mask.tolist() == true
+
+
+def test_distance_mask_fails():
+    "Check that the function fails if no coordinates or region are given"
+    with pytest.raises(ValueError):
+        distance_mask((2.5, -7.5), maxdist=2)