Fix last test issues

tests/test_coreg/test_affine.py

@@ -168,7 +168,7 @@ def test_gradientdescending(self, subsample: int = 10000, inlier_mask: bool = Tr
         )
         assert gds._meta["offset_east_px"] == pytest.approx(-0.496000, rel=1e-1, abs=0.1)
         assert gds._meta["offset_north_px"] == pytest.approx(-0.1875, rel=1e-1, abs=0.1)
-        assert gds._meta["vshift"] == pytest.approx(-1.8730, rel=1e-1)
+        assert gds._meta["vshift"] == pytest.approx(-2.39, rel=1e-1)
 
     @pytest.mark.parametrize("shift_px", [(1, 1), (2, 2)])  # type: ignore
     @pytest.mark.parametrize("coreg_class", [coreg.NuthKaab, coreg.GradientDescending, coreg.ICP])  # type: ignore

tests/test_coreg/test_base.py

@@ -20,7 +20,7 @@
     import xdem
     from xdem import coreg, examples, misc, spatialstats
     from xdem._typing import NDArrayf
-    from xdem.coreg.base import Coreg, apply_matrix_rst
+    from xdem.coreg.base import Coreg, _apply_matrix_rst
 
 
 def load_examples() -> tuple[RasterType, RasterType, Vector]:
@@ -831,7 +831,7 @@ def test_apply_matrix() -> None:
     vshift = 5
     matrix = np.diag(np.ones(4, float))
     matrix[2, 3] = vshift
-    transformed_dem = apply_matrix_rst(ref_arr, ref.transform, matrix)
+    transformed_dem = _apply_matrix_rst(ref_arr, ref.transform, matrix)
     reverted_dem = transformed_dem - vshift
 
     # Check that the reverted DEM has the exact same values as the initial one
@@ -850,7 +850,7 @@ def test_apply_matrix() -> None:
     matrix[0, 3] = pixel_shift * tba.res[0]
     matrix[2, 3] = -vshift
 
-    transformed_dem = apply_matrix_rst(shifted_dem, ref.transform, matrix, resampling="bilinear")
+    transformed_dem = _apply_matrix_rst(shifted_dem, ref.transform, matrix, resampling="bilinear")
     diff = np.asarray(ref_arr - transformed_dem)
 
     # Check that the median is very close to zero
@@ -876,14 +876,14 @@ def rotation_matrix(rotation: float = 30) -> NDArrayf:
         np.mean([ref.bounds.top, ref.bounds.bottom]),
         ref.data.mean(),
     )
-    rotated_dem = apply_matrix_rst(ref.data.squeeze(), ref.transform, rotation_matrix(rotation), centroid=centroid)
+    rotated_dem = _apply_matrix_rst(ref.data.squeeze(), ref.transform, rotation_matrix(rotation), centroid=centroid)
     # Make sure that the rotated DEM is way off, but is centered around the same approximate point.
     assert np.abs(np.nanmedian(rotated_dem - ref.data.data)) < 1
     assert spatialstats.nmad(rotated_dem - ref.data.data) > 500
 
     # Apply a rotation in the opposite direction
     unrotated_dem = (
-            apply_matrix_rst(rotated_dem, ref.transform, rotation_matrix(-rotation * 0.99), centroid=centroid) + 4.0
+            _apply_matrix_rst(rotated_dem, ref.transform, rotation_matrix(-rotation * 0.99), centroid=centroid) + 4.0
     )  # TODO: Check why the 0.99 rotation and +4 vertical shift were introduced.
 
     diff = np.asarray(ref.data.squeeze() - unrotated_dem)

xdem/coreg/__init__.py

@@ -10,7 +10,7 @@
     Tilt,
     VerticalShift,
 )
-from xdem.coreg.base import BlockwiseCoreg, Coreg, CoregPipeline, apply_matrix_rst  # noqa
+from xdem.coreg.base import BlockwiseCoreg, Coreg, CoregPipeline, apply_matrix, invert_matrix  # noqa
 from xdem.coreg.biascorr import (  # noqa
     BiasCorr,
     BiasCorr1D,

xdem/coreg/base.py

@@ -707,7 +707,47 @@ def invert_matrix(matrix: NDArrayf) -> NDArrayf:
         return pytransform3d.transformations.invert_transform(checked_matrix)
 
 
-def apply_matrix_rst(
+def apply_matrix(
+    elev: gu.Raster | NDArrayf | gpd.GeoDataFrame,
+    matrix: NDArrayf,
+    invert: bool = False,
+    centroid: tuple[float, float, float] | None = None,
+    resampling: int | str = "bilinear",
+    transform: rio.transform.Affine = None,
+    z_name: str = "z",
+) -> NDArrayf | gu.Raster | gpd.GeoDataFrame:
+    """
+    Apply a 3D affine transformation matrix to a 3D elevation point cloud or 2.5D DEM.
+
+    :param elev: Elevation point cloud or DEM to transform, either a 2D array (requires transform) or
+        geodataframe (requires z_name).
+    :param matrix: Affine (4x4) transformation matrix to apply to the DEM.
+    :param invert: Whether to invert the transformation matrix.
+    :param centroid: The X/Y/Z transformation centroid. Irrelevant for pure translations. Defaults to the midpoint (Z=0).
+    :param resampling: The resampling method to use, only for DEM 2.5D transformation. Can be `nearest`, `bilinear`,
+        `cubic` or an integer from 0-5.
+    :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.
+    :return:
+    """
+
+    if isinstance(elev, gpd.GeoDataFrame):
+        return _apply_matrix_pts(epc=elev, matrix=matrix, invert=invert, centroid=centroid, z_name=z_name)
+    else:
+        if isinstance(elev, gu.Raster):
+            transform = elev.transform
+            dem = elev.data
+        else:
+            dem = elev
+
+        # TODO: Add exception for translation to update only geotransform, maybe directly in apply_matrix?
+        applied_dem = _apply_matrix_rst(dem=dem, transform=transform, matrix=matrix, invert=invert, centroid=centroid,
+                                        resampling=resampling)
+        if isinstance(elev, gu.Raster):
+            applied_dem = gu.Raster.from_array(applied_dem, transform, elev.crs, elev.nodata)
+        return applied_dem
+
+def _apply_matrix_rst(
     dem: NDArrayf,
     transform: rio.transform.Affine,
     matrix: NDArrayf,
@@ -845,7 +885,7 @@ def apply_matrix_rst(
 
     return transformed_dem
 
-def apply_matrix_pts(
+def _apply_matrix_pts(
     epc: gpd.GeoDataFrame,
     matrix: NDArrayf,
     invert: bool = False,
@@ -859,7 +899,7 @@ def apply_matrix_pts(
     :param matrix: Affine (4x4) transformation matrix to apply to the DEM.
     :param invert: Whether to invert the transformation matrix.
     :param centroid: The X/Y/Z transformation centroid. Irrelevant for pure translations. Defaults to the midpoint (Z=0).
-    :param z_name:
+    :param z_name: Column name to use as elevation, only for point elevation data passed as geodataframe.
 
     :return: Transformed elevation point cloud.
     """
@@ -1484,7 +1524,7 @@ def _apply_func(self, **kwargs: Any) -> tuple[np.ndarray | gpd.GeoDataFrame, aff
 
                     # Apply the matrix around the centroid (if defined, otherwise just from the center).
                     transform = kwargs.pop("transform")
-                    applied_elev = apply_matrix_rst(
+                    applied_elev = _apply_matrix_rst(
                         dem=kwargs.pop("elev"),
                         transform=transform,
                         matrix=self.to_matrix(),
@@ -1506,10 +1546,10 @@ def _apply_func(self, **kwargs: Any) -> tuple[np.ndarray | gpd.GeoDataFrame, aff
             except NotImplementedCoregApply:
                 if self.is_affine:  # This only works on it's rigid, however.
 
-                    applied_elev = apply_matrix_pts(epc=kwargs["elev"],
-                                                    matrix=self.to_matrix(),
-                                                    centroid=self._meta.get("centroid"),
-                                                    z_name=kwargs.pop("z_name"))
+                    applied_elev = _apply_matrix_pts(epc=kwargs["elev"],
+                                                     matrix=self.to_matrix(),
+                                                     centroid=self._meta.get("centroid"),
+                                                     z_name=kwargs.pop("z_name"))
 
                 else:
                     raise ValueError("Cannot transform, Coreg method is non-affine and has no implemented _apply_pts.")

xdem/coreg/biascorr.py

@@ -891,7 +891,7 @@ def _fit_rst_rst(  # type: ignore
     ) -> None:
 
         # If already passed by user, pass along
-        if self._meta["terrain_attribute"] in bias_vars:
+        if bias_vars is not None and self._meta["terrain_attribute"] in bias_vars:
             attr = bias_vars[self._meta["terrain_attribute"]]
 
         # If only declared during instantiation
@@ -933,7 +933,7 @@ def _fit_rst_pts(  # type: ignore
     ) -> None:
 
         # If already passed by user, pass along
-        if self._meta["terrain_attribute"] in bias_vars:
+        if bias_vars is not None and self._meta["terrain_attribute"] in bias_vars:
             attr = bias_vars[self._meta["terrain_attribute"]]
 
         # If only declared during instantiation