Merge branch 'main' of github.com:GlacioHack/xdem into coreg-allow-ra…

…ster-classes

Added warnings for overridden 'transform' argument.
Improved argument validation for 'Coreg'

pyproject.toml

@@ -0,0 +1,6 @@
+[tool.pytest.ini_options]
+addopts = "--doctest-modules"
+testpaths = [
+	"tests",
+	"xdem"
+]

tests/test_coreg.py

@@ -131,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
 
@@ -377,7 +377,7 @@ 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, self.ref.transform)
@@ -408,7 +408,7 @@ 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)
+        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
@@ -418,7 +418,7 @@ 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, self.ref.transform)
@@ -532,8 +532,85 @@ def test_coreg_raster_and_ndarray_args(_) -> None:
         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")
+        dem_arr2 = dem_arr + 1
+        transform = rio.transform.from_origin(0, 5, 1, 1)
 
+        dem_arr2_fixed = coreg.BiasCorr().fit(dem_arr, dem_arr2, transform=transform).apply(dem_arr2, transform=transform)
 
+        assert np.array_equal(dem_arr, dem_arr2_fixed)
 
 
 

xdem/coreg.py

@@ -447,13 +447,13 @@ 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, reference_dem: np.ndarray | np.ma.masked_array | RasterType,
+    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,
             subsample: Union[float, int] = 1.0,
-            verbose: bool = False):
+            verbose: bool = False) -> CoregType:
         """
         Estimate the coregistration transform on the given DEMs.
 
@@ -469,20 +469,36 @@ def fit(self, reference_dem: np.ndarray | np.ma.masked_array | RasterType,
         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 'reference_dem' is not a Raster, just extract the data (shape validation comes later)
-        elif isinstance(dem_to_be_aligned, gu.Raster):
-            dem_to_be_aligned = dem_to_be_aligned.data
-
-        # If 'reference_dem' is a Raster, use its transform (if not given) and then extract its data.
-        if isinstance(reference_dem, gu.Raster):
-
-            if transform is None:
-                transform = reference_dem.transform
-            reference_dem = reference_dem.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)
@@ -526,6 +542,8 @@ def fit(self, reference_dem: np.ndarray | np.ma.masked_array | RasterType,
         # Flag that the fitting function has been called.
         self._fit_called = True
 
+        return self
+
     @overload
     def apply(self, dem: RasterType, transform: rio.transform.Affine | None) -> RasterType: ...
 
@@ -537,12 +555,12 @@ def apply(self, dem: np.ma.masked_array, transform: rio.transform.Affine | None)
 
 
     def apply(self, dem: np.ndarray | np.ma.masked_array | RasterType,
-              transform: rio.transform.Affine | None = None) -> Union[np.ndarray, np.ma.masked_array]:
+              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.
         """
@@ -552,9 +570,11 @@ def apply(self, dem: np.ndarray | np.ma.masked_array | RasterType,
         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")
+                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)

xdem/spstats.py

@@ -68,7 +68,7 @@ def interp_nd_binning(df: pd.DataFrame, list_var_names: Union[str,list[str]], st
 
     # Extrapolated linearly outside the 2D frame.
     >>> fun((-1, 1))
-    array(-5.)
+    array(-1.)
     """
     # if list of variable input is simply a string
     if isinstance(list_var_names,str):

xdem/terrain.py

@@ -260,14 +260,14 @@ def get_terrain_attribute(
                [1, 1, 1],
                [2, 2, 2]])
         >>> slope, aspect = get_terrain_attribute(dem, ["slope", "aspect"], resolution=1)
-        >>> slope
-        array([[45., 45., 45.],
-               [45., 45., 45.],
-               [45., 45., 45.]], dtype=float32)
+        >>> slope  # Note the flattening edge effect; see 'get_quadric_coefficients()' for more.
+        array([[26.56505118, 26.56505118, 26.56505118],
+               [45.        , 45.        , 45.        ],
+               [26.56505118, 26.56505118, 26.56505118]])
         >>> aspect
         array([[0., 0., 0.],
                [0., 0., 0.],
-               [0., 0., 0.]], dtype=float32)
+               [0., 0., 0.]])
 
     :returns: One or multiple arrays of the requested attribute(s)
     """
@@ -450,15 +450,15 @@ def aspect(dem: np.ndarray | np.ma.masked_array, degrees: bool = True) -> np.nda
         >>> aspect(dem, degrees=True)
         array([[0., 0., 0.],
                [0., 0., 0.],
-               [0., 0., 0.]], dtype=float32)
+               [0., 0., 0.]])
         >>> dem.T
         array([[0, 1, 2],
                [0, 1, 2],
                [0, 1, 2]])
         >>> aspect(dem.T, degrees=True)
         array([[270., 270., 270.],
                [270., 270., 270.],
-               [270., 270., 270.]], dtype=float32)
+               [270., 270., 270.]])
 
     """
     return get_terrain_attribute(dem, attribute="aspect", resolution=1.0, degrees=degrees)