From 33b1ee2b41e544228861394385b76e917b2067b2 Mon Sep 17 00:00:00 2001
From: Romain Hugonnet <romain.hugonnet@gmail.com>
Date: Mon, 27 Nov 2023 20:58:14 -0900
Subject: [PATCH] Fix bug in `CoregPipeline.fit()` (#450)

* Fix bug and remove _fit_func

* Re-add the fit_pts func

* Add test to check pipeline output consistency

* Linting

* Fix random states and increase tolerance

* Increase tolerance of tilt test
---
 tests/test_coreg/test_affine.py |  2 +-
 tests/test_coreg/test_base.py   | 38 +++++++++++++++++++++++++++++++++
 xdem/coreg/base.py              |  2 +-
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tests/test_coreg/test_affine.py b/tests/test_coreg/test_affine.py
index c1c1fdd2..baf5e54e 100644
--- a/tests/test_coreg/test_affine.py
+++ b/tests/test_coreg/test_affine.py
@@ -288,7 +288,7 @@ def test_tilt(self) -> None:
         assert np.abs(np.mean(periglacial_offset)) < np.abs(np.mean(pre_offset))
 
         # Check that the mean periglacial offset is low
-        assert np.abs(np.mean(periglacial_offset)) < 0.01
+        assert np.abs(np.mean(periglacial_offset)) < 0.02
 
     def test_icp_opencv(self) -> None:
         warnings.simplefilter("error")
diff --git a/tests/test_coreg/test_base.py b/tests/test_coreg/test_base.py
index 2e0bfe73..d7ff64ef 100644
--- a/tests/test_coreg/test_base.py
+++ b/tests/test_coreg/test_base.py
@@ -676,6 +676,44 @@ def test_coreg_add(self) -> None:
         vshift5 = vshift3 + vshift3
         assert vshift5.to_matrix()[2, 3] == vshift * 4
 
+    def test_pipeline_consistency(self) -> None:
+        """Check that pipelines properties are respected: reflectivity, fusion of same coreg"""
+
+        # Test 1: Fusion of same coreg
+        # Many vertical shifts
+        many_vshifts = coreg.VerticalShift() + coreg.VerticalShift() + coreg.VerticalShift()
+        many_vshifts.fit(**self.fit_params, random_state=42)
+        aligned_dem, _ = many_vshifts.apply(self.tba.data, transform=self.ref.transform, crs=self.ref.crs)
+
+        # The last steps should have shifts of EXACTLY zero
+        assert many_vshifts.pipeline[1]._meta["vshift"] == pytest.approx(0, abs=10e-5)
+        assert many_vshifts.pipeline[2]._meta["vshift"] == pytest.approx(0, abs=10e-5)
+
+        # Many horizontal + vertical shifts
+        many_nks = coreg.NuthKaab() + coreg.NuthKaab() + coreg.NuthKaab()
+        many_nks.fit(**self.fit_params, random_state=42)
+        aligned_dem, _ = many_nks.apply(self.tba.data, transform=self.ref.transform, crs=self.ref.crs)
+
+        # The last steps should have shifts of NEARLY zero
+        assert many_nks.pipeline[1]._meta["vshift"] == pytest.approx(0, abs=0.02)
+        assert many_nks.pipeline[1]._meta["offset_east_px"] == pytest.approx(0, abs=0.02)
+        assert many_nks.pipeline[1]._meta["offset_north_px"] == pytest.approx(0, abs=0.02)
+        assert many_nks.pipeline[2]._meta["vshift"] == pytest.approx(0, abs=0.02)
+        assert many_nks.pipeline[2]._meta["offset_east_px"] == pytest.approx(0, abs=0.02)
+        assert many_nks.pipeline[2]._meta["offset_north_px"] == pytest.approx(0, abs=0.02)
+
+        # Test 2: Reflectivity
+        # Those two pipelines should give almost the same result
+        nk_vshift = coreg.NuthKaab() + coreg.VerticalShift()
+        vshift_nk = coreg.VerticalShift() + coreg.NuthKaab()
+
+        nk_vshift.fit(**self.fit_params, random_state=42)
+        aligned_dem, _ = nk_vshift.apply(self.tba.data, transform=self.ref.transform, crs=self.ref.crs)
+        vshift_nk.fit(**self.fit_params, random_state=42)
+        aligned_dem, _ = vshift_nk.apply(self.tba.data, transform=self.ref.transform, crs=self.ref.crs)
+
+        assert np.allclose(nk_vshift.to_matrix(), vshift_nk.to_matrix(), atol=10e-1)
+
 
 class TestBlockwiseCoreg:
     ref, tba, outlines = load_examples()  # Load example reference, to-be-aligned and mask.
diff --git a/xdem/coreg/base.py b/xdem/coreg/base.py
index 191c0e71..0ce059c1 100644
--- a/xdem/coreg/base.py
+++ b/xdem/coreg/base.py
@@ -1516,7 +1516,7 @@ def fit(
 
             main_args_fit = {
                 "reference_dem": ref_dem,
-                "dem_to_be_aligned": tba_dem,
+                "dem_to_be_aligned": tba_dem_mod,
                 "inlier_mask": inlier_mask,
                 "transform": out_transform,
                 "crs": crs,