test: add a shift before coregistration test

GlacioHack · Nov 22, 2024 · d5e4602 · d5e4602
1 parent ff6e437
commit d5e4602
Showing 1 changed file with 53 additions and 4 deletions.
diff --git a/tests/test_coreg/test_workflows.py b/tests/test_coreg/test_workflows.py
@@ -256,9 +256,58 @@ def test_dem_coregistration(self) -> None:
         out_fig.close()
 
         # Testing different coreg method
-        dem_coreg, coreg_method, coreg_stats, inlier_mask = dem_coregistration(
+        dem_coreg2, coreg_method2, coreg_stats2, inlier_mask2 = dem_coregistration(
             tba_dem, ref_dem, coreg_method=xdem.coreg.Deramp()
         )
-        assert isinstance(coreg_method, xdem.coreg.Deramp)
-        assert abs(coreg_stats["med_orig"].values) > abs(coreg_stats["med_coreg"].values)
-        assert coreg_stats["nmad_orig"].values > coreg_stats["nmad_coreg"].values
+        assert isinstance(coreg_method2, xdem.coreg.Deramp)
+        assert abs(coreg_stats2["med_orig"].values) > abs(coreg_stats2["med_coreg"].values)
+        assert coreg_stats2["nmad_orig"].values > coreg_stats2["nmad_coreg"].values
+
+        # Testing with initial shift
+        test_shift_list = [10, 5]
+        tba_dem_shift = tba_dem.copy()
+        dem_coreg2, coreg_method2, coreg_stats2, inlier_mask2 = dem_coregistration(
+            tba_dem_shift, ref_dem, estimated_initial_shift=test_shift_list
+        )
+        tba_dem_shift_test = tba_dem.translate(test_shift_list[0] * tba_dem.res[0], test_shift_list[1] * tba_dem.res[1])
+        assert tba_dem_shift.raster_equal(tba_dem_shift_test)
+        assert isinstance(coreg_method2, xdem.coreg.CoregPipeline)
+        assert isinstance(coreg_method, xdem.coreg.CoregPipeline)
+        assert isinstance(coreg_method2.pipeline[0], xdem.coreg.AffineCoreg)
+        assert (
+            coreg_method2.pipeline[0].meta["outputs"]["affine"]["shift_x"]
+            == coreg_method.pipeline[0].meta["outputs"]["affine"]["shift_x"]
+        )
+        assert (
+            coreg_method2.pipeline[0].meta["outputs"]["affine"]["shift_y"]
+            == coreg_method.pipeline[0].meta["outputs"]["affine"]["shift_y"]
+        )
+
+        # Testing without coreg pipeline
+        test_shift_tuple = (-5, 2)  # tuple
+        coreg_simple = xdem.coreg.affine.DhMinimize()
+
+        tba_dem_shift = tba_dem.copy()
+        dem_coreg2, coreg_method2, coreg_stats2, inlier_mask2 = dem_coregistration(
+            tba_dem_shift, ref_dem, coreg_method=coreg_simple, estimated_initial_shift=test_shift_tuple
+        )
+        dem_coreg3, coreg_method3, coreg_stats3, inlier_mask3 = dem_coregistration(
+            tba_dem, ref_dem, coreg_method=coreg_simple
+        )
+        assert isinstance(coreg_method2, xdem.coreg.AffineCoreg)
+        assert isinstance(coreg_method3, xdem.coreg.AffineCoreg)
+        assert coreg_method2.meta["outputs"]["affine"]["shift_x"] == coreg_method3.meta["outputs"]["affine"]["shift_x"]
+        assert coreg_method2.meta["outputs"]["affine"]["shift_y"] == coreg_method3.meta["outputs"]["affine"]["shift_y"]
+
+        # Check if the appropriate exception is raised with an initial shift and without affine coreg
+        with pytest.raises(TypeError, match=r".*affine.*"):
+            dem_coregistration(
+                tba_dem_shift, ref_dem, coreg_method=xdem.coreg.Deramp(), estimated_initial_shift=test_shift_tuple
+            )
+        with pytest.raises(TypeError, match=r".*affine.*"):
+            dem_coregistration(
+                tba_dem_shift,
+                ref_dem,
+                coreg_method=xdem.coreg.Deramp() + xdem.coreg.TerrainBias(),
+                estimated_initial_shift=test_shift_tuple,
+            )