Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Calibration creation date #321

Merged
merged 13 commits into from
Feb 21, 2024
Merged
Show file tree
Hide file tree
Changes from 12 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
265 changes: 160 additions & 105 deletions sed/calibrator/delay.py

Large diffs are not rendered by default.

300 changes: 173 additions & 127 deletions sed/calibrator/energy.py

Large diffs are not rendered by default.

300 changes: 188 additions & 112 deletions sed/calibrator/momentum.py

Large diffs are not rendered by default.

832 changes: 545 additions & 287 deletions sed/core/processor.py

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion sed/loader/mpes/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
35 changes: 31 additions & 4 deletions tests/calibrator/test_delay.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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={},
Expand All @@ -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"])
52 changes: 33 additions & 19 deletions tests/calibrator/test_energy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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):
Expand Down
18 changes: 17 additions & 1 deletion tests/calibrator/test_momentum.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
)
Loading