[WIP] Allow apply_matrix() with multiprocessing

tests/test_coreg/test_base.py

@@ -6,6 +6,7 @@
 import re
 import sys
 import warnings
+from pathlib import Path
 from typing import Any, Callable, Iterable, Mapping
 
 import geopandas as gpd
@@ -16,6 +17,7 @@
 import rasterio as rio
 from geoutils import Raster, Vector
 from geoutils.raster import RasterType
+from geoutils.raster.distributed_computing import MultiprocConfig
 from scipy.ndimage import binary_dilation
 
 import xdem
@@ -1068,3 +1070,33 @@ def test_apply_matrix__raster_realdata(self) -> None:
         diff_it_gd = z_points_gd[valids] - z_points_it[valids]
         assert np.percentile(np.abs(diff_it_gd), 99) < 1.2  # 99% of values are within a 1.20 meter (instead of 90%)
         assert np.percentile(np.abs(diff_it_gd), 50) < 0.03  # 10 times more precise than above
+
+    @pytest.mark.parametrize("chunk_size", [5, 30])
+    @pytest.mark.parametrize("matrix", list_matrices)
+    @pytest.mark.parametrize("invert", [True, False])
+    @pytest.mark.parametrize("resampling", [None, "nearest", "linear", "cubic", "quintic"])
+    def test_apply_matrix_multi(
+        self, matrix: NDArrayf, chunk_size: int, invert: bool, resampling: str, tmp_path: Path
+    ) -> None:
+
+        # Create a synthetic raster and convert to point cloud
+        dem_arr = np.linspace(0, 2, 25).reshape(5, 5)
+        transform = rio.transform.from_origin(0, 5, 1, 1)
+        dem = gu.Raster.from_array(dem_arr, transform=transform, crs=4326, nodata=100)
+        epc = dem.to_pointcloud(data_column_name="z").ds
+
+        # If a centroid was not given, default to the center of the DEM (at Z=0).
+        centroid = (np.mean(epc.geometry.x.values), np.mean(epc.geometry.y.values), 0.0)
+        resample = resampling is not None
+
+        # Apply affine transformation to both datasets
+        ref_dem = apply_matrix(dem, matrix, invert=invert, centroid=centroid, resample=resample)
+
+        mp_config = MultiprocConfig(chunk_size=5, outfile=tmp_path / "test.tif")
+        apply_matrix(dem, matrix, multiproc_config=mp_config, invert=invert, centroid=centroid, resample=resample)
+        res_dem = xdem.DEM(tmp_path / "test.tif")
+
+        assert ref_dem.nodata == res_dem.nodata
+        assert ref_dem.crs == res_dem.crs
+        assert ref_dem.transform == res_dem.transform
+        assert ref_dem.get_stats("mean") == res_dem.get_stats("mean")

xdem/coreg/base.py

@@ -54,6 +54,7 @@
 from geoutils.raster import Raster, RasterType, raster
 from geoutils.raster._geotransformations import _resampling_method_from_str
 from geoutils.raster.array import get_array_and_mask
+from geoutils.raster.distributed_computing import map_overlap_multiproc_save
 from geoutils.raster.georeferencing import _cast_pixel_interpretation, _coords
 from geoutils.raster.geotransformations import _translate
 
@@ -1545,7 +1546,6 @@ def _apply_matrix_rst(
 
     :returns: Transformed DEM, Transform.
     """
-
     # Invert matrix if required
     if invert:
         matrix = invert_matrix(matrix)
@@ -1665,6 +1665,7 @@ def apply_matrix(
     resampling: Literal["nearest", "linear", "cubic", "quintic"] = "linear",
     transform: rio.transform.Affine = None,
     z_name: str = "z",
+    multiproc_config: gu.raster.MultiprocConfig | None = None,
     **kwargs: Any,
 ) -> tuple[NDArrayf, affine.Affine]: ...
 
@@ -1679,6 +1680,7 @@ def apply_matrix(
     resampling: Literal["nearest", "linear", "cubic", "quintic"] = "linear",
     transform: rio.transform.Affine = None,
     z_name: str = "z",
+    multiproc_config: gu.raster.MultiprocConfig | None = None,
     **kwargs: Any,
 ) -> gu.Raster | gpd.GeoDataFrame: ...
 
@@ -1692,6 +1694,7 @@ def apply_matrix(
     resampling: Literal["nearest", "linear", "cubic", "quintic"] = "linear",
     transform: rio.transform.Affine = None,
     z_name: str = "z",
+    multiproc_config: gu.raster.MultiprocConfig | None = None,
     **kwargs: Any,
 ) -> tuple[NDArrayf, affine.Affine] | gu.Raster | gpd.GeoDataFrame:
     """
@@ -1725,45 +1728,112 @@ def apply_matrix(
         information, see scipy.ndimage.map_coordinates and scipy.interpolate.interpn. Default is linear.
     :param transform: Geotransform of the DEM, only for DEM passed as 2D array.
     :param z_name: Column name to use as elevation, only for point elevation data passed as geodataframe.
+    :param multiproc_config: ...
     :param kwargs: Keywords passed to _apply_matrix_rst for testing.
 
     :return: Affine transformed elevation point cloud or DEM.
     """
 
+    mp_backend = multiproc_config is not None
+    # The check below can only run on Xarray
+    # dask_backend = da is not None and elev._chunks is not None
+    dask_backend = False
+
     # Apply matrix to elevation point cloud
     if isinstance(elev, gpd.GeoDataFrame):
+        if mp_backend or dask_backend:
+            raise ValueError(
+                "Cannot use Multiprocessing or Dask simultaneously to apply a 3D affine transformation matrix"
+                "to a 3D elevation point cloud."
+            )
         return _apply_matrix_pts(epc=elev, matrix=matrix, invert=invert, centroid=centroid, z_name=z_name)
+
     # Or apply matrix to raster (often requires re-gridding)
     else:
+        if mp_backend and dask_backend:
+            raise ValueError(
+                "Cannot use Multiprocessing and Dask simultaneously. To use Dask, remove mp_config parameter "
+                "from reproject(). To use Multiprocessing, open the file without chunks."
+            )
 
         # First, we apply the affine matrix for the array/transform
         if isinstance(elev, gu.Raster):
             transform = elev.transform
             dem = elev.data.filled(np.nan)
         else:
             dem = elev
-        applied_dem, out_transform = _apply_matrix_rst(
-            dem=dem,
-            transform=transform,
-            matrix=matrix,
-            invert=invert,
-            centroid=centroid,
-            resampling=resampling,
-            **kwargs,
-        )
 
-        # Then, if resample is True, we reproject the DEM from its out_transform onto the transform
-        if resample:
-            applied_dem = _reproject_horizontal_shift_samecrs(
-                applied_dem, src_transform=out_transform, dst_transform=transform, resampling=resampling
+        # If using Multiprocessing backend, process and return None (files written on disk)
+        if mp_backend:
+            # Get depth of overlap
+            depth = 10  # ath.ceil(depth)
+
+            # Block function to pass
+            def _apply_matrix_rst_block(
+                block: Raster,
+                matrix: NDArrayf,
+                invert: bool = False,
+                centroid: tuple[float, float, float] | None = None,
+                resample: bool = True,
+                resampling: Literal["nearest", "linear", "cubic", "quintic"] = "linear",
+                **kwargs: Any,
+            ) -> Raster:
+                """Block function for multiprocessing."""
+                filtered_block, new_transform = _apply_matrix_rst(
+                    block.data.filled(np.nan), block.transform, matrix, invert, centroid, resampling, **kwargs
+                )
+                if resample:
+                    print(resample, type(filtered_block))
+                    filtered_block = _reproject_horizontal_shift_samecrs(
+                        filtered_block,
+                        src_transform=new_transform,
+                        dst_transform=block.transform,
+                        resampling=resampling,
+                    )
+                    new_transform = block.transform
+                rast = gu.Raster.from_array(filtered_block, new_transform, block.crs, block.nodata)
+                return rast
+
+            dem_res = map_overlap_multiproc_save(
+                _apply_matrix_rst_block,
+                elev,
+                multiproc_config,
+                matrix,
+                invert,
+                centroid,
+                resample,
+                resampling,
+                depth=depth,
+                **kwargs,
             )
-            out_transform = transform
 
-        # We return a raster if input was a raster
-        if isinstance(elev, gu.Raster):
-            applied_dem = gu.Raster.from_array(applied_dem, out_transform, elev.crs, elev.nodata)
+            return dem_res
+        elif dask_backend:
+            raise ValueError("Cannot use Dask....")
+            return None
+        else:
+            applied_dem, out_transform = _apply_matrix_rst(
+                dem=dem,
+                transform=transform,
+                matrix=matrix,
+                invert=invert,
+                centroid=centroid,
+                resampling=resampling,
+                **kwargs,
+            )
+
+            # Then, if resample is True, we reproject the DEM from its out_transform onto the transform
+            if resample:
+                applied_dem = _reproject_horizontal_shift_samecrs(
+                    applied_dem, src_transform=out_transform, dst_transform=transform, resampling=resampling
+                )
+                out_transform = transform
+
+            # We return a raster if input was a raster
+            if isinstance(elev, gu.Raster):
+                applied_dem = gu.Raster.from_array(applied_dem, out_transform, elev.crs, elev.nodata)
+
             return applied_dem
-        return applied_dem, out_transform
 
 
 ###########################################