Merge pull request #513 from OpenCOMPES/fix-apply-saved-transformation

fix application of saved transformation, and add tests for this
OpenCOMPES · Nov 11, 2024 · 16fd8e8 · 16fd8e8
2 parents 4c4fada + 7673020
commit 16fd8e8
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/sed/calibrator/momentum.py b/sed/calibrator/momentum.py
@@ -1732,7 +1732,7 @@ def apply_corrections(
                         self.spline_warp_estimate(verbose=verbose)
                     if self.transformations:
                         # Apply config pose adjustments
-                        self.pose_adjustment()
+                        self.pose_adjustment(apply=True)
                 else:
                     raise ValueError("No corrections or transformations defined!")
 

diff --git a/sed/core/processor.py b/sed/core/processor.py
@@ -711,6 +711,7 @@ def pose_adjustment(
         if self.mc.slice_corrected is None:
             if self.mc.slice is None:
                 self.mc.slice = np.zeros(self._config["momentum"]["bins"][0:2])
+                self.mc.bin_ranges = self._config["momentum"]["ranges"]
             self.mc.slice_corrected = self.mc.slice
 
         if not use_correction:

diff --git a/tests/test_processor.py b/tests/test_processor.py
@@ -371,6 +371,7 @@ def test_pose_adjustment_save_load() -> None:
     """Test for the saving and loading of pose correction and application of momentum correction
     workflow"""
     config = {"core": {"loader": "mpes"}}
+    # pose adjustment w/ loaded image
     processor = SedProcessor(
         folder=df_folder,
         config=config,
@@ -379,7 +380,6 @@ def test_pose_adjustment_save_load() -> None:
         system_config={},
         verbose=True,
     )
-    # pose adjustment w/o loaded image
     processor.bin_and_load_momentum_calibration(apply=True)
     processor.define_features(
         features=feature7,
@@ -391,6 +391,33 @@ def test_pose_adjustment_save_load() -> None:
     processor.save_splinewarp(filename="sed_config_pose_adjustments.yaml")
     processor.pose_adjustment(**adjust_params, apply=True)  # type: ignore[arg-type]
     processor.save_transformations(filename="sed_config_pose_adjustments.yaml")
+    processor.apply_momentum_correction()
+    cdeform_field = processor.mc.cdeform_field.copy()
+    rdeform_field = processor.mc.rdeform_field.copy()
+    # pose adjustment w/o loaded image
+    processor = SedProcessor(
+        folder=df_folder,
+        config=config,
+        folder_config="sed_config_pose_adjustments.yaml",
+        user_config={},
+        system_config={},
+        verbose=True,
+    )
+    processor.pose_adjustment(**adjust_params, apply=True)  # type: ignore[arg-type]
+    processor.apply_momentum_correction()
+    # pose adjustment w/o loaded image and w/o correction
+    processor = SedProcessor(
+        folder=df_folder,
+        config=config,
+        folder_config="sed_config_pose_adjustments.yaml",
+        user_config={},
+        system_config={},
+        verbose=True,
+    )
+    processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)  # type: ignore[arg-type]
+    assert np.all(processor.mc.cdeform_field != cdeform_field)
+    assert np.all(processor.mc.rdeform_field != rdeform_field)
+    # from config
     processor = SedProcessor(
         folder=df_folder,
         config=config,
@@ -403,6 +430,8 @@ def test_pose_adjustment_save_load() -> None:
     assert "Xm" in processor.dataframe.columns
     assert "Ym" in processor.dataframe.columns
     assert "momentum_correction" in processor.attributes.metadata
+    np.testing.assert_allclose(processor.mc.cdeform_field, cdeform_field)
+    np.testing.assert_allclose(processor.mc.rdeform_field, rdeform_field)
     os.remove("sed_config_pose_adjustments.yaml")