Allow Rasters as inputs for Coreg.fit() and Coreg.apply()

docs/source/code/coregistration.py

@@ -53,10 +53,10 @@
 
 bias_corr = coreg.BiasCorr()
 # Note that the transform argument is not needed, since it is a simple vertical correction.
-bias_corr.fit(ref_data, tba_data, inlier_mask=inlier_mask)
+bias_corr.fit(ref_data, tba_data, inlier_mask=inlier_mask, transform=reference_dem.transform)
 
 # Apply the bias to a DEM
-corrected_dem = bias_corr.apply(tba_data, transform=None)  # The transform does not need to be given for bias
+corrected_dem = bias_corr.apply(tba_data, transform=dem_to_be_aligned.transform)
 
 # Use median bias instead
 bias_median = coreg.BiasCorr(bias_func=np.median)

tests/test_coreg.py

@@ -18,6 +18,7 @@
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
     from xdem import coreg, examples, spatial_tools, misc
+    import xdem
 
 
 def load_examples() -> tuple[gu.georaster.Raster, gu.georaster.Raster, gu.geovector.Vector]:
@@ -130,7 +131,7 @@ def test_bias(self):
         # Check that this is indeed a new object
         assert biascorr is not biascorr2
         # Fit the corrected DEM to see if the bias will be close to or at zero
-        biascorr2.fit(reference_dem=self.ref.data, dem_to_be_aligned=tba_unbiased, inlier_mask=self.inlier_mask)
+        biascorr2.fit(reference_dem=self.ref.data, dem_to_be_aligned=tba_unbiased, transform=self.ref.transform, inlier_mask=self.inlier_mask)
         # Test the bias
         assert abs(biascorr2._meta.get("bias")) < 0.01
 
@@ -376,10 +377,10 @@ def test_z_scale_corr(self):
         scaled_dem = self.ref.data * factor
 
         # Fit the correction
-        zcorr.fit(self.ref.data, scaled_dem)
+        zcorr.fit(self.ref.data, scaled_dem, transform=self.ref.transform)
 
         # Apply the correction
-        unscaled_dem = zcorr.apply(scaled_dem, None)
+        unscaled_dem = zcorr.apply(scaled_dem, self.ref.transform)
 
         # Make sure the difference is now minimal
         diff = (self.ref.data - unscaled_dem).filled(np.nan)
@@ -407,8 +408,8 @@ def test_z_scale_corr(self):
         dem_with_nans += error_field * 3
 
         # Try the fit now with the messed up DEM as reference.
-        zcorr.fit(dem_with_nans, scaled_dem)
-        unscaled_dem = zcorr.apply(scaled_dem, None)
+        zcorr.fit(dem_with_nans, scaled_dem, transform=self.ref.transform)
+        unscaled_dem = zcorr.apply(scaled_dem, self.ref.transform)
         diff = (dem_with_nans - unscaled_dem).filled(np.nan)
         assert np.abs(np.nanmedian(diff)) < 0.05
 
@@ -417,10 +418,10 @@ def test_z_scale_corr(self):
 
         # Try to correct using a nonlinear correction.
         zcorr_nonlinear = coreg.ZScaleCorr(degree=2)
-        zcorr_nonlinear.fit(dem_with_nans, scaled_dem)
+        zcorr_nonlinear.fit(dem_with_nans, scaled_dem, transform=self.ref.transform)
 
         # Make sure the difference is minimal
-        unscaled_dem = zcorr_nonlinear.apply(scaled_dem, None)
+        unscaled_dem = zcorr_nonlinear.apply(scaled_dem, self.ref.transform)
         diff = (dem_with_nans - unscaled_dem).filled(np.nan)
         assert np.abs(np.nanmedian(diff)) < 0.05
 
@@ -484,6 +485,123 @@ def test_blockwise_coreg(self, pipeline, subdivision):
         # Check that offsets were actually calculated.
         assert np.sum(np.abs(np.linalg.norm(stats[["x_off", "y_off", "z_off"]], axis=0))) > 0
 
+    def test_coreg_raster_and_ndarray_args(_) -> None:
+
+        # Create a small sample-DEM
+        dem1 = xdem.DEM.from_array(
+            np.arange(25, dtype="int32").reshape(5, 5),
+            transform=rio.transform.from_origin(0, 5, 1, 1),
+            crs=4326,
+            nodata=-9999
+        )
+        # Assign a funny value to one particular pixel. This is to validate that reprojection works perfectly.
+        dem1.data[0, 1, 1] = 100
+
+        # Translate the DEM 1 "meter" right and add a bias
+        dem2 = dem1.reproject(dst_bounds=rio.coords.BoundingBox(1, 0, 6, 5))
+        dem2 += 1
+
+        # Create a biascorr for Rasters ("_r") and for arrays ("_a")
+        biascorr_r = coreg.BiasCorr()
+        biascorr_a = biascorr_r.copy()
+
+        # Fit the data
+        biascorr_r.fit(
+            reference_dem=dem1,
+            dem_to_be_aligned=dem2
+        )
+        biascorr_a.fit(
+            reference_dem=dem1.data,
+            dem_to_be_aligned=dem2.reproject(dem1).data,
+            transform=dem1.transform
+        )
+
+        # Validate that they ended up giving the same result.
+        assert biascorr_r._meta["bias"] == biascorr_a._meta["bias"]
+
+        # De-shift dem2
+        dem2_r = biascorr_r.apply(dem2)
+        dem2_a = biascorr_a.apply(dem2.data, dem2.transform)
+
+        # Validate that the return formats were the expected ones, and that they are equal.
+        assert isinstance(dem2_r, xdem.DEM)
+        assert isinstance(dem2_a, np.ma.masked_array)
+        assert np.array_equal(dem2_r, dem2_r)
+
+        # If apply on a masked_array was given without a transform, it should fail.
+        with pytest.raises(ValueError, match="'transform' must be given"):
+            biascorr_a.apply(dem2.data)
+
+        with pytest.warns(UserWarning, match="DEM .* overrides the given 'transform'"):
+            biascorr_a.apply(dem2, transform=dem2.transform)
+
+
+    @pytest.mark.parametrize("combination", [
+        ("dem1", "dem2", "None", "fit", "passes", ""),
+        ("dem1", "dem2", "None", "apply", "passes", ""),
+        ("dem1.data", "dem2.data", "dem1.transform", "fit", "passes", ""),
+        ("dem1.data", "dem2.data", "dem1.transform", "apply", "passes", ""),
+        ("dem1", "dem2.data", "dem1.transform", "fit", "warns", "'reference_dem' .* overrides the given 'transform'"),
+        ("dem1.data", "dem2", "dem1.transform", "fit", "warns", "'dem_to_be_aligned' .* overrides .*"),
+        ("dem1.data", "dem2.data", "None", "fit", "error", "'transform' must be given if both DEMs are array-like."),
+        ("dem1", "dem2.data", "None", "apply", "error", "'transform' must be given if DEM is array-like."),
+        ("dem1", "dem2", "dem2.transform", "apply", "warns", "DEM .* overrides the given 'transform'"),
+        ("None", "None", "None", "fit", "error", "Both DEMs need to be array-like"),
+        ("dem1 + np.nan", "dem2", "None", "fit", "error", "'reference_dem' had only NaNs"),
+        ("dem1", "dem2 + np.nan", "None", "fit", "error", "'dem_to_be_aligned' had only NaNs"),
+    ])
+    def test_coreg_raises(_, combination: tuple[str, str, str, str, str, str]) -> None:
+        """
+        Assert that the expected warnings/errors are triggered under different circumstances.
+
+        The 'combination' param contains this in order:
+            1. The reference_dem (will be eval'd)
+            2. The dem to be aligned (will be eval'd)
+            3. The transform to use (will be eval'd)
+            4. Which coreg method to assess
+            5. The expected outcome of the test.
+            6. The error/warning message (if applicable)
+        """
+        warnings.simplefilter("error")
+
+        ref_dem, tba_dem, transform, testing_step, result, text = combination
+        # Create a small sample-DEM
+        dem1 = xdem.DEM.from_array(
+            np.arange(25, dtype="int32").reshape(5, 5),
+            transform=rio.transform.from_origin(0, 5, 1, 1),
+            crs=4326,
+            nodata=-9999
+        )
+        dem2 = dem1.copy()
+
+        # Evaluate the parametrization (e.g. 'dem2.transform')
+        ref_dem, tba_dem, transform = map(eval, (ref_dem, tba_dem, transform))
+
+        # Use BiasCorr as a representative example.
+        biascorr = xdem.coreg.BiasCorr()
+
+        fit_func = lambda: biascorr.fit(ref_dem, tba_dem, transform=transform)
+        apply_func = lambda: biascorr.apply(tba_dem, transform=transform)
+
+        # Try running the methods in order and validate the result.
+        for method, method_call in [("fit", fit_func), ("apply", apply_func)]:
+            with warnings.catch_warnings():
+                if method != testing_step:  # E.g. skip warnings for 'fit' if 'apply' is being tested.
+                    warnings.simplefilter("ignore")
+
+                if result == "warns" and testing_step == method:
+                    with pytest.warns(UserWarning, match=text):
+                        method_call()
+                elif result == "error" and testing_step == method:
+                    with pytest.raises(ValueError, match=text):
+                        method_call()
+                else:
+                    method_call()
+
+                if testing_step == "fit":   # If we're testing 'fit', 'apply' does not have to be run.
+                    return
+
+
     def test_coreg_oneliner(_) -> None:
         """Test that a DEM can be coregistered in one line by chaining calls."""
         dem_arr = np.ones((5, 5), dtype="int32")

xdem/coreg.py

@@ -9,10 +9,11 @@
 import tempfile
 import warnings
 from enum import Enum
-from typing import Any, Callable, Optional, Union, Sequence, TypeVar
+from typing import Any, Callable, Optional, overload, Union, Sequence, TypeVar
 
 import fiona
 import geoutils as gu
+from geoutils.georaster import RasterType
 import numpy as np
 import rasterio as rio
 import rasterio.warp  # pylint: disable=unused-import
@@ -446,8 +447,8 @@ def __init__(self, meta: Optional[dict[str, Any]] = None, matrix: Optional[np.nd
                 valid_matrix = pytransform3d.transformations.check_transform(matrix)
             self._meta["matrix"] = valid_matrix
 
-    def fit(self: CoregType, reference_dem: Union[np.ndarray, np.ma.masked_array],
-            dem_to_be_aligned: Union[np.ndarray, np.ma.masked_array],
+    def fit(self: CoregType, reference_dem: np.ndarray | np.ma.masked_array | RasterType,
+            dem_to_be_aligned: np.ndarray | np.ma.masked_array | RasterType,
             inlier_mask: Optional[np.ndarray] = None,
             transform: Optional[rio.transform.Affine] = None,
             weights: Optional[np.ndarray] = None,
@@ -468,6 +469,36 @@ def fit(self: CoregType, reference_dem: Union[np.ndarray, np.ma.masked_array],
         if weights is not None:
             raise NotImplementedError("Weights have not yet been implemented")
 
+        # Validate that both inputs are valid array-like (or Raster) types.
+        if not all(hasattr(dem, "__array_interface__") for dem in (reference_dem, dem_to_be_aligned)):
+            raise ValueError(
+                "Both DEMs need to be array-like (implement a numpy array interface)."
+                f"'reference_dem': {reference_dem}, 'dem_to_be_aligned': {dem_to_be_aligned}"
+            )
+
+        # If both DEMs are Rasters, validate that 'dem_to_be_aligned' is in the right grid. Then extract its data.
+        if isinstance(dem_to_be_aligned, gu.Raster) and isinstance(reference_dem, gu.Raster):
+            dem_to_be_aligned = dem_to_be_aligned.reproject(reference_dem, silent=True).data
+
+        # If any input is a Raster, use its transform if 'transform is None'.
+        # If 'transform' was given and any input is a Raster, trigger a warning.
+        # Finally, extract only the data of the raster.
+        for name, dem in [("reference_dem", reference_dem), ("dem_to_be_aligned", dem_to_be_aligned)]:
+            if hasattr(dem, "transform"):
+                if transform is None:
+                    transform = getattr(dem, "transform")
+                elif transform is not None:
+                    warnings.warn(f"'{name}' of type {type(dem)} overrides the given 'transform'")
+
+                """
+                if name == "reference_dem":
+                    reference_dem = dem.data
+                else:
+                    dem_to_be_aligned = dem.data
+                """
+
+        if transform is None:
+            raise ValueError("'transform' must be given if both DEMs are array-like.")
 
         ref_dem, ref_mask = xdem.spatial_tools.get_array_and_mask(reference_dem)
         tba_dem, tba_mask = xdem.spatial_tools.get_array_and_mask(dem_to_be_aligned)
@@ -513,19 +544,38 @@ def fit(self: CoregType, reference_dem: Union[np.ndarray, np.ma.masked_array],
 
         return self
 
-    def apply(self, dem: Union[np.ndarray, np.ma.masked_array],
-              transform: rio.transform.Affine) -> Union[np.ndarray, np.ma.masked_array]:
+    @overload
+    def apply(self, dem: RasterType, transform: rio.transform.Affine | None) -> RasterType: ...
+
+    @overload
+    def apply(self, dem: np.ndarray, transform: rio.transform.Affine | None) -> np.ndarray: ...
+
+    @overload
+    def apply(self, dem: np.ma.masked_array, transform: rio.transform.Affine | None) -> np.ma.masked_array: ...
+
+
+    def apply(self, dem: np.ndarray | np.ma.masked_array | RasterType,
+              transform: rio.transform.Affine | None = None) -> RasterType | np.ndarray | np.ma.masked_array:
         """
         Apply the estimated transform to a DEM.
 
-        :param dem: A DEM to apply the transform on.
-        :param transform: The transform object of the DEM. TODO: Remove??
+        :param dem: A DEM array or Raster to apply the transform on.
+        :param transform: The transform object of the DEM. Required if 'dem' is an array and not a Raster.
 
         :returns: The transformed DEM.
         """
         if not self._fit_called and self._meta.get("matrix") is None:
             raise AssertionError(".fit() does not seem to have been called yet")
 
+        if isinstance(dem, gu.Raster):
+            if transform is None:
+                transform = dem.transform
+            else:
+                warnings.warn(f"DEM of type {type(dem)} overrides the given 'transform'")
+        else:
+            if transform is None:
+                raise ValueError("'transform' must be given if DEM is array-like.")
+
         # The array to provide the functions will be an ndarray with NaNs for masked out areas.
         dem_array, dem_mask = xdem.spatial_tools.get_array_and_mask(dem)
         if np.all(dem_mask):
@@ -549,8 +599,18 @@ def apply(self, dem: Union[np.ndarray, np.ma.masked_array],
             else:
                 raise ValueError("Coreg method is non-rigid but has no implemented _apply_func")
 
-        # Return the array in the same format as it was given (ndarray or masked_array)
-        return np.ma.masked_array(applied_dem, mask=dem.mask) if isinstance(dem, np.ma.masked_array) else applied_dem
+        # If the DEM was a masked_array, copy the mask to the new DEM
+        if hasattr(dem, "mask"):
+            applied_dem = np.ma.masked_array(applied_dem, mask=dem.mask)  # type: ignore
+        # If the DEM was a Raster with a mask, copy the mask to the new DEM
+        elif hasattr(dem, "data") and hasattr(dem.data, "mask"):
+            applied_dem = np.ma.masked_array(applied_dem, mask=dem.data.mask)  # type: ignore
+
+        # If the input was a Raster, return a Raster as well.
+        if isinstance(dem, gu.Raster):
+            return dem.from_array(applied_dem, transform, dem.crs, nodata=dem.nodata)
+
+        return applied_dem
 
     def apply_pts(self, coords: np.ndarray) -> np.ndarray:
         """