feat: add EnergyOptimizer class with transition step calculation and tests

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)

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