Merge pull request #321 from OpenCOMPES/calibration_creation_date

Calibration creation date

.pre-commit-config.yaml

@@ -15,10 +15,10 @@ repos:
     # Ruff version.
     rev: v0.1.7
     hooks:
-      # Run the linter.
-      - id: ruff
       # Run the formatter.
       - id: ruff-format
+      # Run the linter.
+      - id: ruff
 -   repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.7.1
     hooks:

sed/loader/mpes/loader.py

@@ -840,7 +840,7 @@ def gather_metadata(
         ].keys()
 
         lens_volts = np.array(
-            [metadata["file"][f"KTOF:Lens:{lens}:V"] for lens in lens_list],
+            [metadata["file"].get(f"KTOF:Lens:{lens}:V", np.NaN) for lens in lens_list],
         )
         for mode, value in self._config["metadata"]["lens_mode_config"].items():
             lens_volts_config = np.array([value[k] for k in lens_list])

tests/calibrator/test_delay.py

@@ -96,7 +96,12 @@ def test_delay_parameters_from_delay_range_mm() -> None:
     dc = DelayCalibrator(config=config)
     with pytest.raises(NotImplementedError):
         dc.append_delay_axis(df, delay_range_mm=(1, 15))
-    df, metadata = dc.append_delay_axis(df, delay_range_mm=(1, 15), time0=1)
+    df, metadata = dc.append_delay_axis(
+        df,
+        delay_range_mm=(1, 15),
+        time0=1,
+        adc_range=config["delay"]["adc_range"],
+    )
     assert "delay" in df.columns
     assert "delay_range" in metadata["calibration"]
     assert "adc_range" in metadata["calibration"]
@@ -171,7 +176,7 @@ def test_add_offset_from_args(df=test_dataframe) -> None:
     cfg_ = cfg.copy()
     cfg_.pop("delay")
     config = parse_config(
-        config=cfg,
+        config=cfg_,
         folder_config={},
         user_config={},
         system_config={},
@@ -182,11 +187,33 @@ def test_add_offset_from_args(df=test_dataframe) -> None:
         constant=1,
         flip_delay_axis=True,
         columns="bam",
-        weights=0.001,
     )
     assert "delay" in df.columns
     assert "bam" in dc.offsets.keys()
     expected = -np.array(
-        delay_stage_vals + bam_vals * 0.001 + 1,
+        delay_stage_vals + bam_vals * 1 + 1,
+    )
+    np.testing.assert_allclose(expected, df["delay"])
+
+
+def test_add_offset_from_dict(df=test_dataframe) -> None:
+    """test that the timing offset is corrected for correctly from config"""
+    cfg_ = cfg.copy()
+    offsets = cfg["delay"]["offsets"]  # type:ignore
+    offsets["bam"].pop("weight")
+    offsets["flip_delay_axis"] = False
+    cfg_.pop("delay")
+    config = parse_config(
+        config=cfg_,
+        folder_config={},
+        user_config={},
+        system_config={},
     )
+
+    expected = np.asarray(delay_stage_vals + bam_vals * 1 + 1)
+
+    dc = DelayCalibrator(config=config)
+    df, _ = dc.add_offsets(df.copy(), offsets=offsets)
+    assert "delay" in df.columns
+    assert "bam" in dc.offsets.keys()
     np.testing.assert_allclose(expected, df["delay"])

tests/calibrator/test_energy.py

@@ -8,6 +8,7 @@
 from importlib.util import find_spec
 from typing import Any
 from typing import Dict
+from typing import Literal
 
 import dask.dataframe
 import numpy as np
@@ -565,13 +566,18 @@ def test_apply_energy_correction_raises(correction_type: str) -> None:
         assert config["dataframe"]["corrected_tof_column"] in df.columns
 
 
-def test_add_offsets_functionality() -> None:
+@pytest.mark.parametrize(
+    "energy_scale",
+    ["kinetic", "binding"],
+)
+def test_add_offsets_functionality(energy_scale: str) -> None:
     """test the add_offsets function"""
+    scale_sign: Literal[-1, 1] = -1 if energy_scale == "binding" else 1
     config = parse_config(
         config={
             "energy": {
                 "calibration": {
-                    "energy_scale": "kinetic",
+                    "energy_scale": energy_scale,
                 },
                 "offsets": {
                     "constant": 1,
@@ -610,26 +616,41 @@ def test_add_offsets_functionality() -> None:
         loader=get_loader("flash", config=config),
     )
     res, meta = ec.add_offsets(t_df)
-    exp_vals = df["energy"].copy() + 1
-    exp_vals += df["off1"] - df["off1"].mean()
-    exp_vals -= df["off2"]
-    exp_vals += df["off3"].mean()
+    exp_vals = df["energy"].copy() + 1 * scale_sign
+    exp_vals += (df["off1"] - df["off1"].mean()) * scale_sign
+    exp_vals -= df["off2"] * scale_sign
+    exp_vals += df["off3"].mean() * scale_sign
     np.testing.assert_allclose(res["energy"].values, exp_vals.values)
-    exp_meta = params.copy()
+    exp_meta: Dict[str, Any] = {}
     exp_meta["applied"] = True
+    exp_meta["offsets"] = ec.offsets
     assert meta == exp_meta
     # test with explicit params
     ec = EnergyCalibrator(
         config=config,
         loader=get_loader("flash", config=config),
     )
     t_df = dask.dataframe.from_pandas(df.copy(), npartitions=2)
-    res, meta = ec.add_offsets(t_df, **params)  # type: ignore[arg-type] # pylint disable=unexpected-keyword-arg
+    res, meta = ec.add_offsets(t_df, **params)  # type: ignore
     np.testing.assert_allclose(res["energy"].values, exp_vals.values)
-    params["applied"] = True
-    assert meta == params
-
-    # test with different energy scale
+    exp_meta = {}
+    exp_meta["applied"] = True
+    exp_meta["offsets"] = ec.offsets
+    assert meta == exp_meta
+    # test with minimal parameters
+    ec = EnergyCalibrator(
+        config=config,
+        loader=get_loader("flash", config=config),
+    )
+    t_df = dask.dataframe.from_pandas(df.copy(), npartitions=2)
+    res, meta = ec.add_offsets(t_df, weights=-1, columns="off1")
+    res, meta = ec.add_offsets(res, columns="off1")
+    exp_vals = df["energy"].copy()
+    np.testing.assert_allclose(res["energy"].values, exp_vals.values)
+    exp_meta = {}
+    exp_meta["applied"] = True
+    exp_meta["offsets"] = ec.offsets
+    assert meta == exp_meta
 
 
 def test_add_offset_raises() -> None:
@@ -657,13 +678,6 @@ def test_add_offset_raises() -> None:
         },
     )
     t_df = dask.dataframe.from_pandas(df.copy(), npartitions=2)
-    # no sign in config
-    with pytest.raises(KeyError):
-        cfg = deepcopy(cfg_dict)
-        cfg["energy"]["offsets"]["off1"].pop("weight")
-        config = parse_config(config=cfg, folder_config={}, user_config={}, system_config={})
-        ec = EnergyCalibrator(config=cfg, loader=get_loader("flash", config=config))
-        _ = ec.add_offsets(t_df)
 
     # no energy scale
     with pytest.raises(ValueError):

tests/calibrator/test_momentum.py

@@ -162,7 +162,7 @@ def test_apply_correction() -> None:
     assert "Xm" in df.columns
     assert "Ym" in df.columns
     assert metadata["correction"]["applied"] is True
-    np.testing.assert_equal(metadata["correction"]["prefs"], features)
+    np.testing.assert_equal(metadata["correction"]["reference_points"], features)
     assert metadata["correction"]["cdeform_field"].shape == momentum_map.shape
     assert metadata["correction"]["rdeform_field"].shape == momentum_map.shape
 
@@ -378,3 +378,19 @@ def test_momentum_calibration_two_points() -> None:
             metadata["calibration"][key],
             value,
         )
+    # Test with passing calibration parameters
+    calibration = mc.calibration.copy()
+    calibration.pop("creation_date")
+    df, _, _ = get_loader(loader_name="mpes", config=config).read_dataframe(
+        folders=df_folder,
+        collect_metadata=False,
+    )
+    mc = MomentumCorrector(config=config)
+    df, metadata = mc.append_k_axis(df, **calibration)
+    assert "kx" in df.columns
+    assert "ky" in df.columns
+    for key, value in mc.calibration.items():
+        np.testing.assert_equal(
+            metadata["calibration"][key],
+            value,
+        )