Merge pull request #575 from OpenCOMPES/fix_pose_adjustment

Fix pose adjustment

Author: rettigl

src/sed/calibrator/momentum.py

@@ -874,6 +874,10 @@ def reset_deformation(self, **kwds):
             stackaxis=0,
         ).astype("float64")
 
+        # reset image
+        if self.slice_corrected is not None:
+            self.slice_corrected = self.slice
+
         self.rdeform_field = coordmat[1, ...]
         self.cdeform_field = coordmat[0, ...]
 

src/sed/config/default.yaml

@@ -37,7 +37,7 @@ energy:
   # Number of bins to use for energy calibration traces
   bins: 1000
   # Bin ranges to use for energy calibration curves (for tof_binning=1)
-  ranges: [100000, 150000]
+  ranges: [256000, 276000]
   # Option to normalize energy calibration traces
   normalize: True
   # Pixel range for smoothing
@@ -54,7 +54,7 @@ energy:
   energy_scale: "kinetic"
   # Approximate position of the high-energy-cutoff in tof_column bins,
   # used for displaying a graph to choose the energy correction function parameters.
-  tof_fermi: 132250
+  tof_fermi: 264500
   # TOF range to visualize for the correction tool around tof_fermi
   tof_width: [-600, 1000]
   # x-integration range for the correction tool around the center pixel
@@ -72,7 +72,7 @@ momentum:
   # Bin numbers used for the respective axes
   bins: [512, 512, 300]
   # bin ranges to use (in unbinned detector coordinates)
-  ranges: [[-256, 1792], [-256, 1792], [132000, 138000]]
+  ranges: [[-256, 1792], [-256, 1792], [252000, 280000]]
   # The x/y pixel ranges of the detector
   detector_ranges: [[0, 2048], [0, 2048]]
   # The center pixel of the detector in the binned x/y coordinates
@@ -110,4 +110,4 @@ histogram:
   # Axes names starting with "@" refer to keys in the "dataframe" section
   axes: ["@x", "@y", "@tof"]
   # default ranges to use for histogram visualization (in unbinned detector coordinates)
-  ranges: [[0, 1800], [0, 1800], [0, 150000]]
+  ranges: [[0, 1800], [0, 1800], [0, 300000]]

src/sed/core/processor.py

@@ -762,6 +762,7 @@ def pose_adjustment(
 
         if not use_correction:
             self.mc.reset_deformation()
+            reset = False
 
         if self.mc.cdeform_field is None or self.mc.rdeform_field is None:
             # Generate distortion correction from config values

tests/calibrator/test_energy.py

@@ -455,7 +455,12 @@ def test_adjust_energy_correction_raises(correction_type: str) -> None:
     Args:
         correction_type (str): type of correction to test
     """
-    config = parse_config(config={}, folder_config={}, user_config={}, system_config={})
+    config = parse_config(
+        config={"dataframe": {"tof_binning": 2}},
+        folder_config={},
+        user_config={},
+        system_config={},
+    )
     ec = EnergyCalibrator(
         config=config,
         loader=get_loader("mpes", config=config),
@@ -499,7 +504,12 @@ def test_energy_correction_from_dict_kwds(correction_type: str, correction_kwd:
         correction_type (str): type of correction to test
         correction_kwd (dict): parameters to pass to the function
     """
-    config = parse_config(config={}, folder_config={}, user_config={}, system_config={})
+    config = parse_config(
+        config={"dataframe": {"tof_binning": 2}},
+        folder_config={},
+        user_config={},
+        system_config={},
+    )
     sample_df = pd.DataFrame(sample, columns=columns)
     ec = EnergyCalibrator(
         config=config,

tests/test_processor.py

@@ -287,7 +287,7 @@ def test_copy_tool() -> None:
 )
 def test_momentum_correction_workflow(features: np.ndarray) -> None:
     """Test for the momentum correction workflow"""
-    config = {"core": {"loader": "mpes"}}
+    config = {"core": {"loader": "mpes"}, "dataframe": {"tof_binning": 2}}
     processor = SedProcessor(
         folder=df_folder,
         config=config,
@@ -296,7 +296,7 @@ def test_momentum_correction_workflow(features: np.ndarray) -> None:
         system_config={},
         verbose=True,
     )
-    processor.bin_and_load_momentum_calibration(apply=True)
+    processor.bin_and_load_momentum_calibration(plane=200, width=20, apply=True)
     assert processor.mc.slice is not None
     assert processor.mc.pouter is None
     if len(features) == 5 or len(features) == 7:
@@ -339,7 +339,7 @@ def test_momentum_correction_workflow(features: np.ndarray) -> None:
 
 def test_pose_adjustment() -> None:
     """Test for the pose correction and application of momentum correction workflow"""
-    config = {"core": {"loader": "mpes"}}
+    config = {"core": {"loader": "mpes"}, "dataframe": {"tof_binning": 2}}
     processor = SedProcessor(
         folder=df_folder,
         config=config,
@@ -351,7 +351,7 @@ def test_pose_adjustment() -> None:
     # pose adjustment w/o loaded image
     processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)
 
-    processor.bin_and_load_momentum_calibration(apply=True)
+    processor.bin_and_load_momentum_calibration(plane=200, width=20, apply=True)
     # test pose adjustment
     processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)
 
@@ -365,7 +365,7 @@ def test_pose_adjustment() -> None:
     )
     with pytest.raises(ValueError):
         processor.apply_momentum_correction()
-    processor.bin_and_load_momentum_calibration(apply=True)
+    processor.bin_and_load_momentum_calibration(plane=200, width=20, apply=True)
     processor.define_features(
         features=feature7,
         rotation_symmetry=6,
@@ -383,10 +383,72 @@ def test_pose_adjustment() -> None:
         processor.pose_adjustment(**adjust_params, apply=True, illegal_kwd=True)
 
 
+def test_pose_adjustment_use_correction() -> None:
+    """Test for the use of momentum correction in the pose adjustment workflow"""
+    config = {"core": {"loader": "mpes"}, "dataframe": {"tof_binning": 2}}
+    processor = SedProcessor(
+        folder=df_folder,
+        config=config,
+        folder_config={},
+        user_config={},
+        system_config={},
+        verbose=True,
+    )
+    processor.bin_and_load_momentum_calibration(plane=200, width=20, apply=True)
+    # uncorrected reference
+    processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)
+    uncorrected_slice = processor.mc.slice_transformed.copy()
+    uncorrected_cdeform_field = processor.mc.cdeform_field.copy()
+    uncorrected_rdeform_field = processor.mc.rdeform_field.copy()
+    # apply correction, and use pose adjustment without correction
+    processor.define_features(
+        features=feature7,
+        rotation_symmetry=6,
+        include_center=True,
+        apply=True,
+    )
+    processor.generate_splinewarp(use_center=True)
+    processor.save_splinewarp(filename="sed_config_pose_adjustments_use_corrections.yaml")
+    processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)
+    np.testing.assert_allclose(processor.mc.slice_transformed, uncorrected_slice)
+    np.testing.assert_allclose(processor.mc.cdeform_field, uncorrected_cdeform_field)
+    np.testing.assert_allclose(processor.mc.rdeform_field, uncorrected_rdeform_field)
+    # apply correction, and use pose adjustment with correction
+    processor.pose_adjustment(**adjust_params, use_correction=True, apply=True)
+    assert not np.array_equal(processor.mc.slice_transformed, uncorrected_slice)
+    assert not np.array_equal(processor.mc.cdeform_field, uncorrected_cdeform_field)
+    assert not np.array_equal(processor.mc.rdeform_field, uncorrected_rdeform_field)
+    corrected_cdeform_field = processor.mc.cdeform_field.copy()
+    corrected_rdeform_field = processor.mc.rdeform_field.copy()
+
+    processor = SedProcessor(
+        folder=df_folder,
+        config=config,
+        folder_config="sed_config_pose_adjustments_use_corrections.yaml",
+        user_config={},
+        system_config={},
+        verbose=True,
+    )
+    # apply pose correction without loaded image
+    # uncorrected reference
+    processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)
+    np.testing.assert_allclose(processor.mc.cdeform_field, uncorrected_cdeform_field)
+    np.testing.assert_allclose(processor.mc.rdeform_field, uncorrected_rdeform_field)
+    # apply correction, and use pose adjustment with correction
+    processor.generate_splinewarp(use_center=True)
+    processor.pose_adjustment(**adjust_params, use_correction=True, apply=True)
+    np.testing.assert_allclose(processor.mc.cdeform_field, corrected_cdeform_field)
+    np.testing.assert_allclose(processor.mc.rdeform_field, corrected_rdeform_field)
+    # apply correction, and use pose adjustment without correction
+    processor.pose_adjustment(**adjust_params, use_correction=False, apply=True)
+    np.testing.assert_allclose(processor.mc.cdeform_field, uncorrected_cdeform_field)
+    np.testing.assert_allclose(processor.mc.rdeform_field, uncorrected_rdeform_field)
+
+
 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"}}
+    config = {"core": {"loader": "mpes"}, "dataframe": {"tof_binning": 2}}
     # pose adjustment w/ loaded image
     processor = SedProcessor(
         folder=df_folder,
@@ -396,7 +458,7 @@ def test_pose_adjustment_save_load() -> None:
         system_config={},
         verbose=True,
     )
-    processor.bin_and_load_momentum_calibration(apply=True)
+    processor.bin_and_load_momentum_calibration(plane=200, width=20, apply=True)
     processor.define_features(
         features=feature7,
         rotation_symmetry=6,
@@ -458,7 +520,7 @@ def test_pose_adjustment_save_load() -> None:
 
 def test_momentum_calibration_workflow() -> None:
     """Test the calibration of the momentum axes"""
-    config = {"core": {"loader": "mpes"}}
+    config = {"core": {"loader": "mpes"}, "dataframe": {"tof_binning": 2}}
     processor = SedProcessor(
         folder=df_folder,
         config=config,
@@ -471,7 +533,7 @@ def test_momentum_calibration_workflow() -> None:
         processor.apply_momentum_calibration()
     with pytest.raises(ValueError):
         processor.calibrate_momentum_axes(apply=True)
-    processor.bin_and_load_momentum_calibration(apply=True)
+    processor.bin_and_load_momentum_calibration(plane=200, width=20, apply=True)
     with pytest.raises(AssertionError):
         processor.calibrate_momentum_axes(point_a=point_a, apply=True)
     processor.calibrate_momentum_axes(point_a=point_a, k_distance=k_distance, apply=True)
@@ -508,7 +570,7 @@ def test_momentum_calibration_workflow() -> None:
 
 def test_energy_correction() -> None:
     """Test energy correction workflow."""
-    config = {"core": {"loader": "mpes"}}
+    config = {"core": {"loader": "mpes"}, "dataframe": {"tof_binning": 2}}
     processor = SedProcessor(
         folder=df_folder,
         config=config,