Add support for n-dimensional coordinates on BlockReduce

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,13 +97,15 @@ def __init__(
         adjust="spacing",
         center_coordinates=False,
         shape=None,
+        drop_coords=True,
     ):
         self.reduction = reduction
         self.shape = shape
         self.spacing = spacing
         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 @@ def filter(self, coordinates, data, weights=None):
         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 @@ def filter(self, coordinates, data, weights=None):
         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 @@ def _block_coordinates(self, coordinates, block_coordinates, labels):
         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 @@ def _block_coordinates(self, coordinates, block_coordinates, labels):
         ----------
         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 @@ def _block_coordinates(self, coordinates, block_coordinates, labels):
         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 @@ def __init__(
         center_coordinates=False,
         uncertainty=False,
         shape=None,
+        drop_coords=True,
     ):
         super().__init__(
             reduction=np.average,
@@ -319,6 +342,7 @@ def __init__(
             region=region,
             adjust=adjust,
             center_coordinates=center_coordinates,
+            drop_coords=drop_coords,
         )
         self.uncertainty = uncertainty
 
@@ -334,8 +358,8 @@ def filter(self, coordinates, data, weights=None):
         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 @@ def filter(self, coordinates, data, weights=None):
         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

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)