diff --git a/addams_bec/bec_ipython_client/plugins/energy_optimizer/addams_energy_optimizer.py b/addams_bec/bec_ipython_client/plugins/energy_optimizer/addams_energy_optimizer.py new file mode 100644 index 0000000..b3b360a --- /dev/null +++ b/addams_bec/bec_ipython_client/plugins/energy_optimizer/addams_energy_optimizer.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Tuple + +if TYPE_CHECKING: # pragma: no cover + from bec_lib.devicemanager import DeviceManagerBase as DeviceManager + + +class EnergyOptimizer: + def __init__(self, device_manager: DeviceManager): + self.device_manager = device_manager + self.strips = { + "Si": {"energy_range": [5, 10]}, + "Rh": {"energy_range": [8, 23]}, + "Pt": {"energy_range": [20, 40]}, + } + self.overlap_energyies = {"Si": {"Rh": 9}, "Rh": {"Si": 9, "Pt": 22}, "Pt": {"Rh": 22}} + + @staticmethod + def get_transition_steps(start_energy: float, target_energy: float) -> List[Tuple[float, str]]: + """ + Get the required steps to transition from one energy to another. + + Args: + start_energy: The starting energy in keV. + target_energy: The target energy in keV. + + Returns: + A list of tuples containing the energy and the strip name for each step. + """ + strips = {"Si": (5, 10), "Rh": (8, 23), "Pt": (20, 40)} + overlap_energyies = {"Si": {"Rh": 9}, "Rh": {"Si": 9, "Pt": 22}, "Pt": {"Rh": 22}} + + def get_strip(energy: float) -> str: + for strip, (low, high) in strips.items(): + if low <= energy <= high: + return strip + return "" + + def find_overlap(from_strip: str, to_strip: str) -> float: + return overlap_energyies[from_strip][to_strip] + + path = [] + + if get_strip(start_energy) == "": + raise ValueError("Start energy is out of range for available strips") + if not any(low <= target_energy <= high for low, high in strips.values()): + raise ValueError("End energy is out of range for available strips") + + current_energy = start_energy + + # TODO: this should be replaced with a readout from the PV + current_strip = get_strip(current_energy) + + # if the target energy is covered by the current strip, return the path + if strips[current_strip][0] <= target_energy <= strips[current_strip][1]: + return [(target_energy, current_strip)] + + target_strip = get_strip(target_energy) + + available_strips = list(strips.keys()) + current_index = available_strips.index(current_strip) + target_index = available_strips.index(target_strip) + step = 1 if target_index > current_index else -1 + + for i in range(current_index, target_index, step): + next_strip = available_strips[i + step] + overlap_energy = find_overlap(available_strips[i], next_strip) + path.append((overlap_energy, next_strip)) + current_strip = next_strip + + path.append((target_energy, target_strip)) + + return path + + +if __name__ == "__main__": # pragma: no cover + from unittest.mock import MagicMock + + device_manager = MagicMock() + optimizer = EnergyOptimizer(device_manager) + steps = optimizer.get_transition_steps(30, 20) + print(steps) diff --git a/tests/tests_bec_ipython_client/test_energy_optimizer.py b/tests/tests_bec_ipython_client/test_energy_optimizer.py new file mode 100644 index 0000000..eec8a80 --- /dev/null +++ b/tests/tests_bec_ipython_client/test_energy_optimizer.py @@ -0,0 +1,38 @@ +from unittest import mock + +import pytest + +from addams_bec.bec_ipython_client.plugins.energy_optimizer.addams_energy_optimizer import ( + EnergyOptimizer, +) + + +@pytest.fixture +def optimizer(): + dm = mock.MagicMock() + yield EnergyOptimizer(dm) + + +@pytest.mark.parametrize( + "start_energy, target_energy, expected", + [ + (5, 10, [(10, "Si")]), + (5, 20, [(9, "Rh"), (20, "Rh")]), + (5, 40, [(9, "Rh"), (22, "Pt"), (40, "Pt")]), + (5, 5, [(5, "Si")]), + (5, 4, ValueError), + (5, 41, ValueError), + (2, 8, ValueError), + (18, 40, [(22, "Pt"), (40, "Pt")]), + (18, 5, [(9, "Si"), (5, "Si")]), + (18, 10, [(10, "Rh")]), + (18, 20, [(20, "Rh")]), + (25, 7, [(22, "Rh"), (9, "Si"), (7, "Si")]), + ], +) +def test_get_transition_steps(optimizer, start_energy, target_energy, expected): + if expected == ValueError: + with pytest.raises(ValueError): + optimizer.get_transition_steps(start_energy, target_energy) + else: + assert optimizer.get_transition_steps(start_energy, target_energy) == expected