Repository: fatiando/verde

Issue:
Support n-dimensional coordinates in BlockReduce
**Description of the desired feature**

The `verde.BlockReduce` class implements the useful `filter` method, which takes `coordinates` and `data` and reduce the number of observation points based on a `reduction` function (like `numpy.median` for example). It returns two variables: `blocked_coords` and `blocked_data`.

It works very nice when working with 2D data, but it has a few drawbacks when working with 3D data (like harmonic functions measured at an observation point given by three coordinates: `easting`, `northing`, `vertical`). 

On the first place, `filter()` only takes into account the first two elements of `coordinates`, while ignoring the rest.

For example, the following code produces synthetic data on 3D coordinates and performs a reduction:
```python
import numpy as np
import verde as vd

coordinates = vd.scatter_points(
    region=[-100, 100, -100, 100], size=1000, random_state=1, extra_coords=30
)
data = coordinates[0] ** 2 + coordinates[1] ** 2

reducer = vd.BlockReduce(reduction=np.median, spacing=30)
blocked_coords, blocked_data = reducer.filter(coordinates=coordinates, data=data)

print(blocked_coords)
>>> (array([-86.92565699, -60.28602231, -25.32019348,  -4.57180259,
            32.46048571,  60.19388801,  80.80329182, -86.43861313,
           -57.08922825, -28.58767822,  -4.39373841,  24.38387358,
            54.43560591,  83.28117061, -91.03093308, -60.19683937,
           -30.88785459,   2.97782241,  26.18629589,  59.72078359,
            83.7203556 , -85.72747096, -62.27263981, -28.49776653,
             6.73066899,  31.24586634,  57.05920564,  86.37222333,
           -87.2323049 , -57.17757264, -30.85717616,   0.55945368,
            28.54618533,  52.93460056,  85.69458141, -93.88008954,
           -55.94320244, -32.18752505,  -2.4177569 ,  33.28991613,
            53.49497377,  88.16117837, -81.17297821, -60.16279652,
           -28.8589658 ,   4.29077631,  29.20026531,  58.5711356 ,
            84.06017411]), array([-88.86031746, -89.50180723, -86.27246784, -83.98780455,
           -82.65943472, -84.53167597, -86.08382994, -59.01557084,
           -59.53812904, -53.70741289, -55.88118363, -51.96555856,
           -59.85450239, -57.87436432, -32.15916302, -31.3737262 ,
           -26.94510776, -30.52611593, -28.19867331, -28.67106966,
           -25.66908466,  -0.21052559,   3.65631607,  -1.51519052,
             1.42388471,  -1.77800774,   2.4122055 ,  -1.79829417,
            28.07622516,  26.46597493,  31.65570674,  30.14057418,
            22.29381973,  27.53528751,  23.2476737 ,  57.30252866,
            61.34936215,  53.12962358,  60.42677294,  60.80440403,
            55.66017801,  62.55215476,  87.11733217,  85.40857097,
            86.12887843,  86.60298467,  86.83554537,  82.74792497,
            86.82501974]))

print(blocked_data)
>>> [14921.99174523 11683.43129776  8153.8981112   7180.46986919
      7816.19690596 11092.80064568 14452.44779479 11619.70650905
      6640.19910392  4033.28000418  3195.33368894  3759.37649066
      6137.97625454 10753.62083976  9371.23765705  4617.49869354
      1815.99051967  1052.94798326  1815.61428988  4672.96671468
      7750.9894085   7383.27255981  3884.07463202   914.12579464
       134.56754578   985.78279509  3313.47562402  7461.68635671
      8247.88174796  3866.0306672   1938.65793764  1019.45739813
      1507.97681279  3593.67734767  7923.85177954 11679.50716615
      6712.13290729  4220.50927868  3738.56502148  5045.05923082
      6390.75792236 11697.77175237 14373.29412264 10905.69778448
      8132.98001843  7598.11106611  8443.15892087 10253.55144211
     15203.06797102]
```
So after the reduction, the vertical dimension is completely lost. One way to patch this is to add this vertical coordinate as part of the `data` argument, but that's not ideal: it's not specifically data and it's problematic when trying to process this kind of data with `verde.Chain`.

I would propose that `vd.BlockReduce.filter` use only the first two coordinates to build the blocks and compute the reduction, but also reduce all the subsequent coordinates. So `block_coordinates` should have the same number of elements as `coordinates`.

This issue is related to #138 and #189 .

**Are you willing to help implement and maintain this feature?** Yes


Comments:

Patch:
diff --git a/verde/blockreduce.py b/verde/blockreduce.py
index 2ececfe..c92ddd6 100644
--- a/verde/blockreduce.py
+++ b/verde/blockreduce.py
@@ -76,6 +76,10 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         If True, then the returned coordinates correspond to the center of each
         block. Otherwise, the coordinates are calculated by applying the same
         reduction operation to the input coordinates.
+    drop_coords : bool
+        If True, only the reduced ``easting`` and ``northing`` coordinates are returned,
+        dropping any other ones. If False, all coordinates are reduced and returned.
+        Default True.
 
     See also
     --------
@@ -93,6 +97,7 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         adjust="spacing",
         center_coordinates=False,
         shape=None,
+        drop_coords=True,
     ):
         self.reduction = reduction
         self.shape = shape
@@ -100,6 +105,7 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         self.region = region
         self.adjust = adjust
         self.center_coordinates = center_coordinates
+        self.drop_coords = drop_coords
 
     def filter(self, coordinates, data, weights=None):
         """
@@ -120,8 +126,8 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         coordinates : tuple of arrays
             Arrays with the coordinates of each data point. Should be in the
             following order: (easting, northing, vertical, ...). Only easting
-            and northing will be used, all subsequent coordinates will be
-            ignored.
+            and northing will be used to create the blocks. If ``drop_coords`` is
+            ``False``, all other coordinates will be reduced along with the data.
         data : array or tuple of arrays
             The data values at each point. If you want to reduce more than one
             data component, pass in multiple arrays as elements of a tuple. All
@@ -134,8 +140,10 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         Returns
         -------
         blocked_coordinates : tuple of arrays
-            (easting, northing) arrays with the coordinates of each block that
-            contains data.
+            Tuple containing arrays with the coordinates of each block that contains
+            data. If ``drop_coords`` is ``True``, the tuple will only contain
+            (``easting``, ``northing``). If ``drop_coords`` is ``False``, it will
+            contain (``easting``, ``northing``, ``vertical``, ...).
         blocked_data : array
             The block reduced data values.
 
@@ -175,6 +183,9 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         If self.center_coordinates, the coordinates will be the center of each
         block. Otherwise, will apply the reduction to the coordinates.
 
+        If self.drop_coords, only the easting and northing coordinates will be returned.
+        If False, all coordinates will be reduced.
+
         Blocks without any data will be omitted.
 
         *block_coordinates* and *labels* should be the outputs of
@@ -184,9 +195,7 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         ----------
         coordinates : tuple of arrays
             Arrays with the coordinates of each data point. Should be in the
-            following order: (easting, northing, vertical, ...). Only easting
-            and northing will be used, all subsequent coordinates will be
-            ignored.
+            following order: (easting, northing, vertical, ...).
         block_coordinates : tuple of arrays
             (easting, northing) arrays with the coordinates of the center of
             each block.
@@ -197,22 +206,31 @@ class BlockReduce(BaseEstimator):  # pylint: disable=too-few-public-methods
         Returns
         -------
         coordinates : tuple of arrays
-            (easting, northing) arrays with the coordinates assigned to each
-            non-empty block.
+            Tuple containing arrays with the coordinates of each block that contains
+            data. If ``drop_coords`` is ``True``, the tuple will only contain
+            (``easting``, ``northing``). If ``drop_coords`` is ``False``, it will
+            contain (``easting``, ``northing``, ``vertical``, ...).
 
         """
-        if self.center_coordinates:
-            unique = np.unique(labels)
-            return tuple(i[unique] for i in block_coordinates)
         # Doing the coordinates separately from the data because in case of
         # weights the reduction applied to then is different (no weights
         # ever)
-        easting, northing = coordinates[:2]
-        table = pd.DataFrame(
-            dict(easting=easting.ravel(), northing=northing.ravel(), block=labels)
-        )
+        if self.drop_coords:
+            coordinates = coordinates[:2]
+        coords = {
+            "coordinate{}".format(i): coord.ravel()
+            for i, coord in enumerate(coordinates)
+        }
+        coords["block"] = labels
+        table = pd.DataFrame(coords)
         grouped = table.groupby("block").aggregate(self.reduction)
-        return grouped.easting.values, grouped.northing.values
+        if self.center_coordinates:
+            unique = np.unique(labels)
+            for i, block_coord in enumerate(block_coordinates[:2]):
+                grouped["coordinate{}".format(i)] = block_coord[unique].ravel()
+        return tuple(
+            grouped["coordinate{}".format(i)].values for i in range(len(coordinates))
+        )
 
 
 class BlockMean(BlockReduce):  # pylint: disable=too-few-public-methods
@@ -287,6 +305,10 @@ class BlockMean(BlockReduce):  # pylint: disable=too-few-public-methods
         If True, then the returned coordinates correspond to the center of each
         block. Otherwise, the coordinates are calculated by applying the same
         reduction operation to the input coordinates.
+    drop_coords : bool
+        If True, only the reduced ``easting`` and ``northing`` coordinates are returned,
+        dropping any other ones. If False, all coordinates are reduced and returned.
+        Default True.
     uncertainty : bool
         If True, the blocked weights will be calculated by uncertainty
         propagation of the data uncertainties. If this is case, then the input
@@ -311,6 +333,7 @@ class BlockMean(BlockReduce):  # pylint: disable=too-few-public-methods
         center_coordinates=False,
         uncertainty=False,
         shape=None,
+        drop_coords=True,
     ):
         super().__init__(
             reduction=np.average,
@@ -319,6 +342,7 @@ class BlockMean(BlockReduce):  # pylint: disable=too-few-public-methods
             region=region,
             adjust=adjust,
             center_coordinates=center_coordinates,
+            drop_coords=drop_coords,
         )
         self.uncertainty = uncertainty
 
@@ -334,8 +358,8 @@ class BlockMean(BlockReduce):  # pylint: disable=too-few-public-methods
         coordinates : tuple of arrays
             Arrays with the coordinates of each data point. Should be in the
             following order: (easting, northing, vertical, ...). Only easting
-            and northing will be used, all subsequent coordinates will be
-            ignored.
+            and northing will be used to create the blocks. If ``drop_coords`` is
+            ``False``, all other coordinates will be reduced along with the data.
         data : array or tuple of arrays
             The data values at each point. If you want to reduce more than one
             data component, pass in multiple arrays as elements of a tuple. All
@@ -350,8 +374,10 @@ class BlockMean(BlockReduce):  # pylint: disable=too-few-public-methods
         Returns
         -------
         blocked_coordinates : tuple of arrays
-            (easting, northing) arrays with the coordinates of each block that
-            contains data.
+            Tuple containing arrays with the coordinates of each block that contains
+            data. If ``drop_coords`` is ``True``, the tuple will only contain
+            (``easting``, ``northing``). If ``drop_coords`` is ``False``, it will
+            contain (``easting``, ``northing``, ``vertical``, ...).
         blocked_mean : array or tuple of arrays
             The block averaged data values.
         blocked_weights : array or tuple of arrays


Test patch:
diff --git a/verde/tests/test_blockreduce.py b/verde/tests/test_blockreduce.py
index 66b6969..165bc07 100644
--- a/verde/tests/test_blockreduce.py
+++ b/verde/tests/test_blockreduce.py
@@ -76,6 +76,69 @@ def test_block_reduce_weights():
     npt.assert_allclose(block_data, 20)
 
 
+def test_block_reduce_drop_coords():
+    "Try reducing constant values in a regular grid dropping extra coordinates"
+    region = (-5, 0, 5, 10)
+    east, north, down, time = grid_coordinates(
+        region, spacing=0.1, pixel_register=True, extra_coords=[70, 1]
+    )
+    data = 20 * np.ones_like(east)
+    reducer = BlockReduce(np.mean, spacing=1, drop_coords=True)
+    block_coords, block_data = reducer.filter((east, north, down, time), data)
+    assert len(block_coords) == 2
+    assert len(block_coords[0]) == 25
+    assert len(block_coords[1]) == 25
+    assert len(block_data) == 25
+    npt.assert_allclose(block_data, 20)
+    npt.assert_allclose(block_coords[0][:5], np.linspace(-4.5, -0.5, 5))
+    npt.assert_allclose(block_coords[1][::5], np.linspace(5.5, 9.5, 5))
+
+
+def test_block_reduce_multiple_coordinates():
+    "Try reducing constant values in a regular grid with n-dimensional coordinates"
+    region = (-5, 0, 5, 10)
+    east, north, down, time = grid_coordinates(
+        region, spacing=0.1, pixel_register=True, extra_coords=[70, 1]
+    )
+    data = 20 * np.ones_like(east)
+    reducer = BlockReduce(np.mean, spacing=1, drop_coords=False)
+    block_coords, block_data = reducer.filter((east, north, down, time), data)
+    assert len(block_coords) == 4
+    assert len(block_coords[0]) == 25
+    assert len(block_coords[1]) == 25
+    assert len(block_coords[2]) == 25
+    assert len(block_coords[3]) == 25
+    assert len(block_data) == 25
+    npt.assert_allclose(block_data, 20)
+    npt.assert_allclose(block_coords[0][:5], np.linspace(-4.5, -0.5, 5))
+    npt.assert_allclose(block_coords[1][::5], np.linspace(5.5, 9.5, 5))
+    npt.assert_allclose(block_coords[2][::5], 70 * np.ones(5))
+    npt.assert_allclose(block_coords[3][::5], np.ones(5))
+
+
+def test_block_reduce_scatter_multiple_coordinates():
+    "Try reducing constant values in a dense enough scatter with n-dimensional coords"
+    region = (-5, 0, 5, 10)
+    coordinates = scatter_points(
+        region, size=10000, random_state=0, extra_coords=[70, 1]
+    )
+    data = 20 * np.ones_like(coordinates[0])
+    block_coords, block_data = BlockReduce(
+        np.mean, 1, region=region, center_coordinates=True, drop_coords=False
+    ).filter(coordinates, data)
+    assert len(block_coords) == 4
+    assert len(block_coords[0]) == 25
+    assert len(block_coords[1]) == 25
+    assert len(block_coords[2]) == 25
+    assert len(block_coords[3]) == 25
+    assert len(block_data) == 25
+    npt.assert_allclose(block_data, 20)
+    npt.assert_allclose(block_coords[0][:5], np.linspace(-4.5, -0.5, 5))
+    npt.assert_allclose(block_coords[1][::5], np.linspace(5.5, 9.5, 5))
+    npt.assert_allclose(block_coords[2][::5], 70 * np.ones(5))
+    npt.assert_allclose(block_coords[3][::5], np.ones(5))
+
+
 def test_block_reduce_multiple_components():
     "Try reducing multiple components in a regular grid"
     region = (-5, 0, 5, 10)