From 3841ed975df12739b424c65682cb2266aa64fe47 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Tue, 28 Apr 2026 16:19:28 -0400 Subject: [PATCH 01/33] added osx-arm64 workspace. this branch should parallel main, but with osx support --- .gitignore | 3 +++ pyproject.toml | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index bfb34fe..959bdb7 100644 --- a/.gitignore +++ b/.gitignore @@ -164,3 +164,6 @@ Thumbs.db # IDEs .vscode/ .cursor/ + +# Other +notes.* \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index f1e8460..70a46ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -164,7 +164,7 @@ isort.required-imports = ["from __future__ import annotations"] [tool.pixi.project] channels = ["conda-forge"] -platforms = ["linux-64"] +platforms = ["linux-64", "osx-arm64"] [tool.pixi.pypi-dependencies] cditools = { path = ".", editable = true } From 3aa49bbb83657eed684cfe9366538e83496d4e30 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 13:52:06 -0400 Subject: [PATCH 02/33] finished most of the logic for the attenuator bank started writing tests; ignored eiger_async tests for now so that running tests is somewhat useful --- pyproject.toml | 15 ++-- src/cditools/attenuator.py | 172 +++++++++++++++++++++++++++++++++++++ tests/test_attenuator.py | 27 ++++++ 3 files changed, 207 insertions(+), 7 deletions(-) create mode 100644 src/cditools/attenuator.py create mode 100644 tests/test_attenuator.py diff --git a/pyproject.toml b/pyproject.toml index 70a46ef..3a0f051 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ test = [ "tiled[minimal-client]", "tiled[minimal-server]", "ophyd >=v1.10.6", + "pytest-watcher" ] dev = [ "caproto[standard] >=0.4.2rc1,!=1.2.0", @@ -92,15 +93,12 @@ dev-dependencies = [ [tool.pytest.ini_options] minversion = "6.0" -addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"] +# TODO - fix the eiger_async module and tests +addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config", "--ignore=tests/test_eiger_async.py"] xfail_strict = true -filterwarnings = [ - "ignore", -] +filterwarnings = "ignore" log_cli_level = "INFO" -testpaths = [ - "tests", -] +testpaths = "tests" [tool.coverage] run.source = ["cditools"] @@ -174,3 +172,6 @@ default = { solve-group = "default" } dev = { features = ["dev"], solve-group = "default" } docs = { features = ["docs"], solve-group = "default" } test = { features = ["test"], solve-group = "default" } + +[tool.pixi.dependencies] +xraylib = ">=4.2.1,<5" diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py new file mode 100644 index 0000000..6eeec35 --- /dev/null +++ b/src/cditools/attenuator.py @@ -0,0 +1,172 @@ +from dataclasses import dataclass + +import xraylib +import numpy as np +import math +from ophyd_async.core import DeviceVector, StandardReadable +from ophyd_async.epics.core import EpicsDevice + + +@dataclass +class AttenuatorCombination: + attenuation: float + attenuators: list[int] + +# The available attenuations can be calculated with the utility +# methods below, but they do not change often, +# so we hardcode them here +# TODO - there will eventually be eight filters +AVAILABLE_ATTENUATIONS = [ + AttenuatorCombination(attenuation=0.08, attenuators=[0, 1, 2, 3]), + AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), + AttenuatorCombination(attenuation=0.104, attenuators=[0, 2, 3]), + AttenuatorCombination(attenuation=0.124, attenuators=[2, 3]), + AttenuatorCombination(attenuation=0.165, attenuators=[0, 1, 3]), + AttenuatorCombination(attenuation=0.196, attenuators=[1, 3]), + AttenuatorCombination(attenuation=0.214, attenuators=[0, 3]), + AttenuatorCombination(attenuation=0.256, attenuators=[3]), + AttenuatorCombination(attenuation=0.312, attenuators=[0, 1, 2]), + AttenuatorCombination(attenuation=0.372, attenuators=[1, 2]), + AttenuatorCombination(attenuation=0.406, attenuators=[0, 2]), + AttenuatorCombination(attenuation=0.484, attenuators=[2]), + AttenuatorCombination(attenuation=0.644, attenuators=[0, 1]), + AttenuatorCombination(attenuation=0.768, attenuators=[1]), + AttenuatorCombination(attenuation=0.839, attenuators=[0]), + AttenuatorCombination(attenuation=1.0, attenuators=[]) +] + +class Attenuator(EpicsDevice): + + filter_material = "Al" + filter_density = xraylib.ElementDensity(13) # g/cm³ + filter_material_z = 13 + + def __init__(self, prefix: str, num: int, thickness: int): + """ + pv - the PV for this filter + """ + self.prefix = prefix + self.num = num + self.pv = f"{self.prefix}:D0{self.num+1}-Cmd" + self.thickness = thickness # microns + super().__init__(prefix=self.prefix) + + def __repr__(self): + return str(self.thickness) + + @property + def thickness_cm(self): + # Thickness is in microns, so convert to cm + return self.thickness * 1e-4 + + @property + def linear_atten_coefficient(self) -> float: + """ + Calculates µ, the linear attenuation coefficient of this material, + at this thickness, and this beam energy. + + photon energy in KeV + xraylib.CS_Total in cm²/g + linear_atten_coefficient in cm⁻¹ + """ + photon_energy = 8.6 # KeV TODO - get the right number; this is taken from bmm + mass_atten_cross_section = xraylib.CS_Total(self.filter_material_z, photon_energy) # + return mass_atten_cross_section * self.filter_density + + @property + def attenuation(self): + """ + Attenuation is the fraction of remaining beam + """ + return np.exp(-self.linear_atten_coefficient * self.thickness_cm ) + + +class AttenuatorBank(StandardReadable, EpicsDevice): + """ + The ioc for the iologik1 lives on xf09id1-inst-ioc1 + """ + prefix = 'XF:09ID1-ES{IOLOGIK1:E1212}:' + thicknesses = [16, 24, 66, 124] # microns + available_attenuations = AVAILABLE_ATTENUATIONS + + # TODO - create filter devices with ophyd-async + def __init__(self): + with self.add_children_as_readables(): + self.attenuators = DeviceVector( + { + i: Attenuator(self.prefix, i, self.thicknesses[i]) + for i in range(0,len(self.thicknesses)) + } + ) + super().__init__(prefix=self.prefix) + + + def set_attenuation(self, target_attenuation: float): + pass + + def find_closest_attenuation(self, target_attenuation: float) -> AttenuatorCombination: + """ + This could be faster if we implemented binary search, + but that seems like overkill for our use case. The search space + is small, so we start in the middle, and work up or down. + """ + best_idx = len(self.available_attenuations) // 2 + atten = self.available_attenuations[best_idx].attenuation + diff = float('inf') + new_diff = abs(target_attenuation - atten) + + while new_diff < diff: + diff = new_diff + # TODO - the (in|de)crement can surely be combined into a clever one liner + if target_attenuation > atten: + atten = self.available_attenuations[best_idx+1].attenuation + new_diff = abs(target_attenuation - atten) + if new_diff < diff: + best_idx += 1 + else: + atten = self.available_attenuations[best_idx-1].attenuation + new_diff = abs(target_attenuation - atten) + if new_diff < diff: + best_idx -= 1 + # Break if we have reached the end of the list + if best_idx >= len(self.available_attenuations) or best_idx < 0: + break + # TODO - should return just the found attentuation? or also the + # requested attenuation and/or the difference? + return self.available_attenuations[best_idx] + + """ + These are utility methods that should not be called during production. + They are used to calculate the available attenuations from all + combinations of attenuators. The result is then used as the + AttenuationBank()._available_attenuations attribute. + """ + def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: + """ + It is more efficient to precompute all possible total + attenuations, and simply look up the closest one. + """ + available_attenuations = list() + for combination in self._powerset(): + attens = [self.attenuators[a] for a in self.attenuators if a in combination] + total_atten = self._calculate_total_attenuation(*attens) + available_attenuations.append(AttenuatorCombination(total_atten, combination)) + # We want the available attenuations sorted so we can efficiently search through them + available_attenuations.sort(key=lambda a: a.attenuation) # type: ignore + return available_attenuations + + def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float : + return round(float(math.prod([a.attenuation for a in attenuators])), 3) + + def _powerset(self) -> list[list[int]] : + """ + This is a famously n*O(2^n) problem. + """ + powerset = [] + for i in range(1 << len(self.attenuators)): + combination = [] + for j in range(len(self.attenuators)): + if i & (1 << j): + combination.append(j) + powerset.append(combination) + return powerset diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py new file mode 100644 index 0000000..6feca52 --- /dev/null +++ b/tests/test_attenuator.py @@ -0,0 +1,27 @@ +from ophyd_async.core import init_devices +import pytest_asyncio +from ophyd_async.testing import * + +from cditools.attenuator import AttenuatorBank + + +pytest_plugins = ('pytest_asyncio',) + +@pytest_asyncio.fixture +async def mock_attenuator_bank(): + async with init_devices(mock=True): + mock_attenuator_bank = AttenuatorBank() + yield mock_attenuator_bank + +def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): + nearest = mock_attenuator_bank.find_closest_attenuation(0.7) + assert nearest.attenuation == 0.644 + + nearest2 = mock_attenuator_bank.find_closest_attenuation(0.2) + assert nearest2.attenuation == 0.196 + + nearest3 = mock_attenuator_bank.find_closest_attenuation(0.02) + assert nearest3.attenuation == 0.08 + + nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) + assert nearest4.attenuation == 1 \ No newline at end of file From c44bd46760b9c7bbbf510e218f3ad137b9a86dc3 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 14:38:03 -0400 Subject: [PATCH 03/33] small cleanups --- src/cditools/attenuator.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 6eeec35..4498bd0 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -38,8 +38,8 @@ class AttenuatorCombination: class Attenuator(EpicsDevice): filter_material = "Al" - filter_density = xraylib.ElementDensity(13) # g/cm³ filter_material_z = 13 + filter_density = xraylib.ElementDensity(filter_material_z) # g/cm³ def __init__(self, prefix: str, num: int, thickness: int): """ @@ -70,7 +70,7 @@ def linear_atten_coefficient(self) -> float: linear_atten_coefficient in cm⁻¹ """ photon_energy = 8.6 # KeV TODO - get the right number; this is taken from bmm - mass_atten_cross_section = xraylib.CS_Total(self.filter_material_z, photon_energy) # + mass_atten_cross_section = xraylib.CS_Total(self.filter_material_z, photon_energy) return mass_atten_cross_section * self.filter_density @property @@ -78,7 +78,7 @@ def attenuation(self): """ Attenuation is the fraction of remaining beam """ - return np.exp(-self.linear_atten_coefficient * self.thickness_cm ) + return np.exp(-self.linear_atten_coefficient * self.thickness_cm) class AttenuatorBank(StandardReadable, EpicsDevice): @@ -89,7 +89,6 @@ class AttenuatorBank(StandardReadable, EpicsDevice): thicknesses = [16, 24, 66, 124] # microns available_attenuations = AVAILABLE_ATTENUATIONS - # TODO - create filter devices with ophyd-async def __init__(self): with self.add_children_as_readables(): self.attenuators = DeviceVector( @@ -160,7 +159,7 @@ def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float : def _powerset(self) -> list[list[int]] : """ - This is a famously n*O(2^n) problem. + This is a famously O(n*2^n) problem. """ powerset = [] for i in range(1 << len(self.attenuators)): From 010e3ecf28854f9119706c43d86d2deff5d75fe4 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 14:54:50 -0400 Subject: [PATCH 04/33] fixed formatting issues --- .gitignore | 2 +- src/cditools/attenuator.py | 54 ++++++++++++++++++++++---------------- tests/test_attenuator.py | 10 ++++--- 3 files changed, 39 insertions(+), 27 deletions(-) diff --git a/.gitignore b/.gitignore index 959bdb7..b34d1dd 100644 --- a/.gitignore +++ b/.gitignore @@ -166,4 +166,4 @@ Thumbs.db .cursor/ # Other -notes.* \ No newline at end of file +notes.* diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 4498bd0..0d508dd 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -1,8 +1,10 @@ +from __future__ import annotations + +import math from dataclasses import dataclass -import xraylib import numpy as np -import math +import xraylib from ophyd_async.core import DeviceVector, StandardReadable from ophyd_async.epics.core import EpicsDevice @@ -12,6 +14,7 @@ class AttenuatorCombination: attenuation: float attenuators: list[int] + # The available attenuations can be calculated with the utility # methods below, but they do not change often, # so we hardcode them here @@ -32,14 +35,14 @@ class AttenuatorCombination: AttenuatorCombination(attenuation=0.644, attenuators=[0, 1]), AttenuatorCombination(attenuation=0.768, attenuators=[1]), AttenuatorCombination(attenuation=0.839, attenuators=[0]), - AttenuatorCombination(attenuation=1.0, attenuators=[]) + AttenuatorCombination(attenuation=1.0, attenuators=[]), ] -class Attenuator(EpicsDevice): +class Attenuator(EpicsDevice): filter_material = "Al" filter_material_z = 13 - filter_density = xraylib.ElementDensity(filter_material_z) # g/cm³ + filter_density = xraylib.ElementDensity(filter_material_z) # g/cm³ def __init__(self, prefix: str, num: int, thickness: int): """ @@ -47,8 +50,8 @@ def __init__(self, prefix: str, num: int, thickness: int): """ self.prefix = prefix self.num = num - self.pv = f"{self.prefix}:D0{self.num+1}-Cmd" - self.thickness = thickness # microns + self.pv = f"{self.prefix}:D0{self.num + 1}-Cmd" + self.thickness = thickness # microns super().__init__(prefix=self.prefix) def __repr__(self): @@ -69,8 +72,10 @@ def linear_atten_coefficient(self) -> float: xraylib.CS_Total in cm²/g linear_atten_coefficient in cm⁻¹ """ - photon_energy = 8.6 # KeV TODO - get the right number; this is taken from bmm - mass_atten_cross_section = xraylib.CS_Total(self.filter_material_z, photon_energy) + photon_energy = 8.6 # KeV TODO - get the right number; this is taken from bmm + mass_atten_cross_section = xraylib.CS_Total( + self.filter_material_z, photon_energy + ) return mass_atten_cross_section * self.filter_density @property @@ -85,8 +90,9 @@ class AttenuatorBank(StandardReadable, EpicsDevice): """ The ioc for the iologik1 lives on xf09id1-inst-ioc1 """ - prefix = 'XF:09ID1-ES{IOLOGIK1:E1212}:' - thicknesses = [16, 24, 66, 124] # microns + + prefix = "XF:09ID1-ES{IOLOGIK1:E1212}:" + thicknesses = (16, 24, 66, 124) # microns available_attenuations = AVAILABLE_ATTENUATIONS def __init__(self): @@ -94,16 +100,17 @@ def __init__(self): self.attenuators = DeviceVector( { i: Attenuator(self.prefix, i, self.thicknesses[i]) - for i in range(0,len(self.thicknesses)) + for i in range(len(self.thicknesses)) } ) super().__init__(prefix=self.prefix) - def set_attenuation(self, target_attenuation: float): pass - def find_closest_attenuation(self, target_attenuation: float) -> AttenuatorCombination: + def find_closest_attenuation( + self, target_attenuation: float + ) -> AttenuatorCombination: """ This could be faster if we implemented binary search, but that seems like overkill for our use case. The search space @@ -111,19 +118,19 @@ def find_closest_attenuation(self, target_attenuation: float) -> AttenuatorCombi """ best_idx = len(self.available_attenuations) // 2 atten = self.available_attenuations[best_idx].attenuation - diff = float('inf') + diff = float("inf") new_diff = abs(target_attenuation - atten) while new_diff < diff: diff = new_diff # TODO - the (in|de)crement can surely be combined into a clever one liner if target_attenuation > atten: - atten = self.available_attenuations[best_idx+1].attenuation + atten = self.available_attenuations[best_idx + 1].attenuation new_diff = abs(target_attenuation - atten) if new_diff < diff: best_idx += 1 else: - atten = self.available_attenuations[best_idx-1].attenuation + atten = self.available_attenuations[best_idx - 1].attenuation new_diff = abs(target_attenuation - atten) if new_diff < diff: best_idx -= 1 @@ -140,24 +147,27 @@ def find_closest_attenuation(self, target_attenuation: float) -> AttenuatorCombi combinations of attenuators. The result is then used as the AttenuationBank()._available_attenuations attribute. """ + def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: """ It is more efficient to precompute all possible total attenuations, and simply look up the closest one. """ - available_attenuations = list() + available_attenuations = [] for combination in self._powerset(): attens = [self.attenuators[a] for a in self.attenuators if a in combination] total_atten = self._calculate_total_attenuation(*attens) - available_attenuations.append(AttenuatorCombination(total_atten, combination)) + available_attenuations.append( + AttenuatorCombination(total_atten, combination) + ) # We want the available attenuations sorted so we can efficiently search through them - available_attenuations.sort(key=lambda a: a.attenuation) # type: ignore + available_attenuations.sort(key=lambda a: a.attenuation) # type: ignore[attr-defined] return available_attenuations - def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float : + def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float: return round(float(math.prod([a.attenuation for a in attenuators])), 3) - def _powerset(self) -> list[list[int]] : + def _powerset(self) -> list[list[int]]: """ This is a famously O(n*2^n) problem. """ diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 6feca52..5274e62 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -1,11 +1,12 @@ -from ophyd_async.core import init_devices +from __future__ import annotations + import pytest_asyncio -from ophyd_async.testing import * +from ophyd_async.core import init_devices from cditools.attenuator import AttenuatorBank +pytest_plugins = ("pytest_asyncio",) -pytest_plugins = ('pytest_asyncio',) @pytest_asyncio.fixture async def mock_attenuator_bank(): @@ -13,6 +14,7 @@ async def mock_attenuator_bank(): mock_attenuator_bank = AttenuatorBank() yield mock_attenuator_bank + def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest = mock_attenuator_bank.find_closest_attenuation(0.7) assert nearest.attenuation == 0.644 @@ -24,4 +26,4 @@ def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): assert nearest3.attenuation == 0.08 nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) - assert nearest4.attenuation == 1 \ No newline at end of file + assert nearest4.attenuation == 1 From f8e34736e1440c147567dd4718bb0bf445f20d87 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 15:02:01 -0400 Subject: [PATCH 05/33] moved xraylib dependency --- pyproject.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3a0f051..32d14bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "ophyd", "ophyd-async[ca] >=0.19", "h5py", + "xraylib >=4.2.1,<5" ] [project.optional-dependencies] @@ -173,5 +174,5 @@ dev = { features = ["dev"], solve-group = "default" } docs = { features = ["docs"], solve-group = "default" } test = { features = ["test"], solve-group = "default" } -[tool.pixi.dependencies] -xraylib = ">=4.2.1,<5" +# [tool.pixi.dependencies] +# xraylib = ">=4.2.1,<5" From f35622a732c21b268ae206712062d8ba5b15aac4 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 16:32:02 -0400 Subject: [PATCH 06/33] added open, close, and get_status methods with corresponding signals --- pyproject.toml | 3 --- src/cditools/attenuator.py | 24 +++++++++++++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 32d14bf..681e587 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -173,6 +173,3 @@ default = { solve-group = "default" } dev = { features = ["dev"], solve-group = "default" } docs = { features = ["docs"], solve-group = "default" } test = { features = ["test"], solve-group = "default" } - -# [tool.pixi.dependencies] -# xraylib = ">=4.2.1,<5" diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 0d508dd..294fd9f 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -5,8 +5,8 @@ import numpy as np import xraylib -from ophyd_async.core import DeviceVector, StandardReadable -from ophyd_async.epics.core import EpicsDevice +from ophyd_async.core import DeviceVector, SignalW, StandardReadable, SignalR, StrictEnum +from ophyd_async.epics.core import EpicsDevice, epics_signal_r, epics_signal_w @dataclass @@ -38,6 +38,10 @@ class AttenuatorCombination: AttenuatorCombination(attenuation=1.0, attenuators=[]), ] +class AttenuatorEnum(StrictEnum): + LOW = "Low" + HIGH = "High" + class Attenuator(EpicsDevice): filter_material = "Al" @@ -50,13 +54,23 @@ def __init__(self, prefix: str, num: int, thickness: int): """ self.prefix = prefix self.num = num - self.pv = f"{self.prefix}:D0{self.num + 1}-Cmd" + self.cmd = epics_signal_w(AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Cmd") + self.status = epics_signal_r(AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Sts") self.thickness = thickness # microns super().__init__(prefix=self.prefix) def __repr__(self): return str(self.thickness) + async def open(self): + await self.cmd.set(AttenuatorEnum.HIGH) + + async def close(self): + await self.cmd.set(AttenuatorEnum.LOW) + + async def get_status(self): + await self.status.get_value() + @property def thickness_cm(self): # Thickness is in microns, so convert to cm @@ -88,10 +102,10 @@ def attenuation(self): class AttenuatorBank(StandardReadable, EpicsDevice): """ - The ioc for the iologik1 lives on xf09id1-inst-ioc1 + The ioc for the iologik1 lives on xf09id1-inst-ioc1.nsls2.bnl.gov """ - prefix = "XF:09ID1-ES{IOLOGIK1:E1212}:" + prefix = "XF:09ID1-ES{IOLOGIK1:E1212}" thicknesses = (16, 24, 66, 124) # microns available_attenuations = AVAILABLE_ATTENUATIONS From bc3547da5ddf097a5e339deb0f9e5a6273bf1593 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 22:25:30 -0400 Subject: [PATCH 07/33] added test, added test env dependency --- pyproject.toml | 3 +++ tests/test_attenuator.py | 9 +++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 681e587..e46c384 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,6 +168,9 @@ platforms = ["linux-64", "osx-arm64"] [tool.pixi.pypi-dependencies] cditools = { path = ".", editable = true } +[tool.pixi.feature.test.pypi-dependencies] +ophyd-async = {version = ">=0.10.0a4", extras = ["ca"]} + [tool.pixi.environments] default = { solve-group = "default" } dev = { features = ["dev"], solve-group = "default" } diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 5274e62..adccffe 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -1,17 +1,19 @@ from __future__ import annotations import pytest_asyncio -from ophyd_async.core import init_devices +from ophyd_async.core import init_devices, callback_on_mock_put, set_mock_value -from cditools.attenuator import AttenuatorBank +from cditools.attenuator import AttenuatorBank, AVAILABLE_ATTENUATIONS pytest_plugins = ("pytest_asyncio",) @pytest_asyncio.fixture async def mock_attenuator_bank(): + async with init_devices(mock=True): mock_attenuator_bank = AttenuatorBank() + yield mock_attenuator_bank @@ -27,3 +29,6 @@ def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) assert nearest4.attenuation == 1 + +def test_up_to_date_available_attenuations(mock_attenuator_bank: AttenuatorBank): + assert mock_attenuator_bank._calculate_available_attentuations() == AVAILABLE_ATTENUATIONS # type: ignore \ No newline at end of file From 542b6795d0f6a5f1a4b6adfc8d0a581c04027eb7 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 8 May 2026 22:30:11 -0400 Subject: [PATCH 08/33] fixed dependency structure, formatted --- pyproject.toml | 6 ++---- src/cditools/attenuator.py | 11 +++++++++-- tests/test_attenuator.py | 11 +++++++---- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e46c384..4a19076 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,8 @@ test = [ "tiled[minimal-client]", "tiled[minimal-server]", "ophyd >=v1.10.6", - "pytest-watcher" + "pytest-watcher", + "ophyd-async[ca] >=0.10.0a4" ] dev = [ "caproto[standard] >=0.4.2rc1,!=1.2.0", @@ -168,9 +169,6 @@ platforms = ["linux-64", "osx-arm64"] [tool.pixi.pypi-dependencies] cditools = { path = ".", editable = true } -[tool.pixi.feature.test.pypi-dependencies] -ophyd-async = {version = ">=0.10.0a4", extras = ["ca"]} - [tool.pixi.environments] default = { solve-group = "default" } dev = { features = ["dev"], solve-group = "default" } diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 294fd9f..ccd5b41 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -5,7 +5,11 @@ import numpy as np import xraylib -from ophyd_async.core import DeviceVector, SignalW, StandardReadable, SignalR, StrictEnum +from ophyd_async.core import ( + DeviceVector, + StandardReadable, + StrictEnum, +) from ophyd_async.epics.core import EpicsDevice, epics_signal_r, epics_signal_w @@ -38,6 +42,7 @@ class AttenuatorCombination: AttenuatorCombination(attenuation=1.0, attenuators=[]), ] + class AttenuatorEnum(StrictEnum): LOW = "Low" HIGH = "High" @@ -55,7 +60,9 @@ def __init__(self, prefix: str, num: int, thickness: int): self.prefix = prefix self.num = num self.cmd = epics_signal_w(AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Cmd") - self.status = epics_signal_r(AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Sts") + self.status = epics_signal_r( + AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Sts" + ) self.thickness = thickness # microns super().__init__(prefix=self.prefix) diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index adccffe..e9cc7a0 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -1,16 +1,15 @@ from __future__ import annotations import pytest_asyncio -from ophyd_async.core import init_devices, callback_on_mock_put, set_mock_value +from ophyd_async.core import init_devices -from cditools.attenuator import AttenuatorBank, AVAILABLE_ATTENUATIONS +from cditools.attenuator import AVAILABLE_ATTENUATIONS, AttenuatorBank pytest_plugins = ("pytest_asyncio",) @pytest_asyncio.fixture async def mock_attenuator_bank(): - async with init_devices(mock=True): mock_attenuator_bank = AttenuatorBank() @@ -30,5 +29,9 @@ def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) assert nearest4.attenuation == 1 + def test_up_to_date_available_attenuations(mock_attenuator_bank: AttenuatorBank): - assert mock_attenuator_bank._calculate_available_attentuations() == AVAILABLE_ATTENUATIONS # type: ignore \ No newline at end of file + assert ( + mock_attenuator_bank._calculate_available_attentuations() # type: ignore[reportPrivateUsage] + == AVAILABLE_ATTENUATIONS + ) From 1881d814ab2c9f1d2a8cbc13d875854ddc2b3fed Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 11 May 2026 10:59:42 -0400 Subject: [PATCH 09/33] simplified finding closest attenuator logic --- src/cditools/attenuator.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index ccd5b41..9e57673 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -141,23 +141,17 @@ def find_closest_attenuation( atten = self.available_attenuations[best_idx].attenuation diff = float("inf") new_diff = abs(target_attenuation - atten) + inc = 1 if target_attenuation > atten else -1 while new_diff < diff: diff = new_diff - # TODO - the (in|de)crement can surely be combined into a clever one liner - if target_attenuation > atten: - atten = self.available_attenuations[best_idx + 1].attenuation - new_diff = abs(target_attenuation - atten) - if new_diff < diff: - best_idx += 1 - else: - atten = self.available_attenuations[best_idx - 1].attenuation - new_diff = abs(target_attenuation - atten) - if new_diff < diff: - best_idx -= 1 - # Break if we have reached the end of the list - if best_idx >= len(self.available_attenuations) or best_idx < 0: + # break if we are about to check oustide the list + if best_idx + inc >= len(self.available_attenuations) or best_idx + inc < 0: break + atten = self.available_attenuations[best_idx + inc].attenuation + new_diff = abs(target_attenuation - atten) + if new_diff < diff: + best_idx += inc # TODO - should return just the found attentuation? or also the # requested attenuation and/or the difference? return self.available_attenuations[best_idx] From 3ef4648f1c8837dda7a824f4baf1e22148f8ca5a Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 11 May 2026 18:04:50 -0400 Subject: [PATCH 10/33] added set_attenuation method --- src/cditools/attenuator.py | 34 ++++++++++++++++++++-------- tests/test_attenuator.py | 45 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 68 insertions(+), 11 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 9e57673..b666dbe 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import math from dataclasses import dataclass @@ -44,8 +45,8 @@ class AttenuatorCombination: class AttenuatorEnum(StrictEnum): - LOW = "Low" - HIGH = "High" + LOW = "Low" # off + HIGH = "High" # on class Attenuator(EpicsDevice): @@ -55,7 +56,11 @@ class Attenuator(EpicsDevice): def __init__(self, prefix: str, num: int, thickness: int): """ - pv - the PV for this filter + prefix - the common prefix for the attenuator bank + num - an integer denoting which attenuator within the bank this is + cmd - the write PV to open and close the attenuator + status - the read PV for the status (high or low) + thickness - the thickness of the attenuator in microns """ self.prefix = prefix self.num = num @@ -67,16 +72,17 @@ def __init__(self, prefix: str, num: int, thickness: int): super().__init__(prefix=self.prefix) def __repr__(self): - return str(self.thickness) + return f"{str(self.thickness)} microns, {self.filter_material}" async def open(self): - await self.cmd.set(AttenuatorEnum.HIGH) + await self.cmd.set(AttenuatorEnum.LOW) async def close(self): - await self.cmd.set(AttenuatorEnum.LOW) + await self.cmd.set(AttenuatorEnum.HIGH) async def get_status(self): - await self.status.get_value() + status = await self.status.get_value() + return status @property def thickness_cm(self): @@ -126,8 +132,18 @@ def __init__(self): ) super().__init__(prefix=self.prefix) - def set_attenuation(self, target_attenuation: float): - pass + async def get_status(self): + status = await asyncio.gather(*(a.get_status() for _, a in self.attenuators.items())) + return status + + async def set_attenuation(self, target_attenuation: float): + attenuation_combination = self.find_closest_attenuation(target_attenuation) + # use gather()? + for num, atten, in self.attenuators.items(): + if num in attenuation_combination.attenuators: + await atten.close() + else: + await atten.open() def find_closest_attenuation( self, target_attenuation: float diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index e9cc7a0..7ef6207 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -1,9 +1,11 @@ from __future__ import annotations +from ophyd_async.testing import assert_reading +import pytest import pytest_asyncio -from ophyd_async.core import init_devices +from ophyd_async.core import init_devices, get_mock_put, set_mock_value -from cditools.attenuator import AVAILABLE_ATTENUATIONS, AttenuatorBank +from cditools.attenuator import AVAILABLE_ATTENUATIONS, AttenuatorBank, AttenuatorEnum pytest_plugins = ("pytest_asyncio",) @@ -16,6 +18,45 @@ async def mock_attenuator_bank(): yield mock_attenuator_bank +@pytest.mark.asyncio +async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): + atten_mock0 = get_mock_put(mock_attenuator_bank.attenuators[0].cmd) + atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].cmd) + atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].cmd) + atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].cmd) + + # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), + combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 + await mock_attenuator_bank.set_attenuation(combo0.attenuation) + atten_mock0.assert_called_with(AttenuatorEnum.LOW) + atten_mock1.assert_called_with(AttenuatorEnum.HIGH) + atten_mock2.assert_called_with(AttenuatorEnum.HIGH) + atten_mock3.assert_called_with(AttenuatorEnum.HIGH) + + # AttenuatorCombination(attenuation=0.768, attenuators=[1]), + combo1 = AVAILABLE_ATTENUATIONS[-3] + await mock_attenuator_bank.set_attenuation(combo1.attenuation) + atten_mock0.assert_called_with(AttenuatorEnum.LOW) + atten_mock1.assert_called_with(AttenuatorEnum.HIGH) + atten_mock2.assert_called_with(AttenuatorEnum.LOW) + atten_mock3.assert_called_with(AttenuatorEnum.LOW) + +@pytest.mark.asyncio +async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): + set_mock_value(mock_attenuator_bank.attenuators[0].status, AttenuatorEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[1].status, AttenuatorEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[2].status, AttenuatorEnum.HIGH) + set_mock_value(mock_attenuator_bank.attenuators[3].status, AttenuatorEnum.LOW) + + await assert_reading(mock_attenuator_bank, { + "mock_attenuator_bank": { + "attenuator0": { "status": "LOW" }, + "attenuator1": { "status": "LOW" }, + "attenuator2": { "status": "HIGH" }, + "attenuator3": { "status": "LOW" }, + } + }) + def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest = mock_attenuator_bank.find_closest_attenuation(0.7) assert nearest.attenuation == 0.644 From 734a0260a4a9c9858831f81e420f574c68886fff Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 11 May 2026 23:29:08 -0400 Subject: [PATCH 11/33] minor fixes --- src/cditools/attenuator.py | 14 +++++++++++--- tests/test_attenuator.py | 18 +++++++----------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index b666dbe..276e4a4 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -65,9 +65,14 @@ def __init__(self, prefix: str, num: int, thickness: int): self.prefix = prefix self.num = num self.cmd = epics_signal_w(AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Cmd") + self.mode = epics_signal_r(bool, f"{self.prefix}:DIO{self.num + 1}-Mode") self.status = epics_signal_r( AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Sts" ) + self.in_status = epics_signal_r( + AttenuatorEnum, f"{self.prefix}:DI{self.num + 1}-Sts" + ) + self.thickness = thickness # microns super().__init__(prefix=self.prefix) @@ -138,12 +143,13 @@ async def get_status(self): async def set_attenuation(self, target_attenuation: float): attenuation_combination = self.find_closest_attenuation(target_attenuation) - # use gather()? + coros = [] for num, atten, in self.attenuators.items(): if num in attenuation_combination.attenuators: - await atten.close() + coros.append(atten.close()) else: - await atten.open() + coros.append(atten.open()) + await asyncio.gather(*coros) def find_closest_attenuation( self, target_attenuation: float @@ -168,6 +174,8 @@ def find_closest_attenuation( new_diff = abs(target_attenuation - atten) if new_diff < diff: best_idx += inc + else: # if diff did not change, then we have found the best option + break # TODO - should return just the found attentuation? or also the # requested attenuation and/or the difference? return self.available_attenuations[best_idx] diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 7ef6207..7b23f78 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -1,6 +1,5 @@ from __future__ import annotations -from ophyd_async.testing import assert_reading import pytest import pytest_asyncio from ophyd_async.core import init_devices, get_mock_put, set_mock_value @@ -14,10 +13,8 @@ async def mock_attenuator_bank(): async with init_devices(mock=True): mock_attenuator_bank = AttenuatorBank() - yield mock_attenuator_bank - @pytest.mark.asyncio async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): atten_mock0 = get_mock_put(mock_attenuator_bank.attenuators[0].cmd) @@ -48,14 +45,13 @@ async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): set_mock_value(mock_attenuator_bank.attenuators[2].status, AttenuatorEnum.HIGH) set_mock_value(mock_attenuator_bank.attenuators[3].status, AttenuatorEnum.LOW) - await assert_reading(mock_attenuator_bank, { - "mock_attenuator_bank": { - "attenuator0": { "status": "LOW" }, - "attenuator1": { "status": "LOW" }, - "attenuator2": { "status": "HIGH" }, - "attenuator3": { "status": "LOW" }, - } - }) + assert await mock_attenuator_bank.get_status() == [ + AttenuatorEnum.LOW, + AttenuatorEnum.LOW, + AttenuatorEnum.HIGH, + AttenuatorEnum.LOW + ] + def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest = mock_attenuator_bank.find_closest_attenuation(0.7) From d2b1b0ba515e9f749bb4c61bfd0caf9fc0e548e1 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 11 May 2026 23:35:26 -0400 Subject: [PATCH 12/33] fixed style things --- src/cditools/attenuator.py | 21 ++++++++++++--------- tests/test_attenuator.py | 14 ++++++++------ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 276e4a4..de228f3 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -45,8 +45,8 @@ class AttenuatorCombination: class AttenuatorEnum(StrictEnum): - LOW = "Low" # off - HIGH = "High" # on + LOW = "Low" # off + HIGH = "High" # on class Attenuator(EpicsDevice): @@ -77,7 +77,7 @@ def __init__(self, prefix: str, num: int, thickness: int): super().__init__(prefix=self.prefix) def __repr__(self): - return f"{str(self.thickness)} microns, {self.filter_material}" + return f"{self.thickness!s} microns, {self.filter_material}" async def open(self): await self.cmd.set(AttenuatorEnum.LOW) @@ -86,8 +86,7 @@ async def close(self): await self.cmd.set(AttenuatorEnum.HIGH) async def get_status(self): - status = await self.status.get_value() - return status + return await self.status.get_value() @property def thickness_cm(self): @@ -138,13 +137,17 @@ def __init__(self): super().__init__(prefix=self.prefix) async def get_status(self): - status = await asyncio.gather(*(a.get_status() for _, a in self.attenuators.items())) - return status + return await asyncio.gather( + *(a.get_status() for _, a in self.attenuators.items()) + ) async def set_attenuation(self, target_attenuation: float): attenuation_combination = self.find_closest_attenuation(target_attenuation) coros = [] - for num, atten, in self.attenuators.items(): + for ( + num, + atten, + ) in self.attenuators.items(): if num in attenuation_combination.attenuators: coros.append(atten.close()) else: @@ -174,7 +177,7 @@ def find_closest_attenuation( new_diff = abs(target_attenuation - atten) if new_diff < diff: best_idx += inc - else: # if diff did not change, then we have found the best option + else: # if diff did not change, then we have found the best option break # TODO - should return just the found attentuation? or also the # requested attenuation and/or the difference? diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 7b23f78..408ccd6 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -2,7 +2,7 @@ import pytest import pytest_asyncio -from ophyd_async.core import init_devices, get_mock_put, set_mock_value +from ophyd_async.core import get_mock_put, init_devices, set_mock_value from cditools.attenuator import AVAILABLE_ATTENUATIONS, AttenuatorBank, AttenuatorEnum @@ -15,6 +15,7 @@ async def mock_attenuator_bank(): mock_attenuator_bank = AttenuatorBank() yield mock_attenuator_bank + @pytest.mark.asyncio async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): atten_mock0 = get_mock_put(mock_attenuator_bank.attenuators[0].cmd) @@ -38,6 +39,7 @@ async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): atten_mock2.assert_called_with(AttenuatorEnum.LOW) atten_mock3.assert_called_with(AttenuatorEnum.LOW) + @pytest.mark.asyncio async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): set_mock_value(mock_attenuator_bank.attenuators[0].status, AttenuatorEnum.LOW) @@ -46,11 +48,11 @@ async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): set_mock_value(mock_attenuator_bank.attenuators[3].status, AttenuatorEnum.LOW) assert await mock_attenuator_bank.get_status() == [ - AttenuatorEnum.LOW, - AttenuatorEnum.LOW, - AttenuatorEnum.HIGH, - AttenuatorEnum.LOW - ] + AttenuatorEnum.LOW, + AttenuatorEnum.LOW, + AttenuatorEnum.HIGH, + AttenuatorEnum.LOW, + ] def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): From 0dc3e900d1cda4cec5226a72cdf9d57b7dc26fea Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Tue, 12 May 2026 15:42:02 -0400 Subject: [PATCH 13/33] add movable protocol to both Attenuator classes, wrap set method in AsyncStatus --- pyproject.toml | 1 - src/cditools/attenuator.py | 40 ++++++++++++++++++++++------------ tests/test_attenuator.py | 44 ++++++++++++++++++++++---------------- 3 files changed, 51 insertions(+), 34 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4a19076..e376a4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,6 @@ test = [ "tiled[minimal-server]", "ophyd >=v1.10.6", "pytest-watcher", - "ophyd-async[ca] >=0.10.0a4" ] dev = [ "caproto[standard] >=0.4.2rc1,!=1.2.0", diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index de228f3..c590a13 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -6,12 +6,15 @@ import numpy as np import xraylib +from bluesky.protocols import Movable from ophyd_async.core import ( + AsyncStatus, DeviceVector, StandardReadable, StrictEnum, + # AsyncMovable, ) -from ophyd_async.epics.core import EpicsDevice, epics_signal_r, epics_signal_w +from ophyd_async.epics.core import EpicsDevice, epics_signal_r, epics_signal_rw @dataclass @@ -44,12 +47,12 @@ class AttenuatorCombination: ] -class AttenuatorEnum(StrictEnum): +class AttenuatorStatusEnum(StrictEnum): LOW = "Low" # off HIGH = "High" # on -class Attenuator(EpicsDevice): +class Attenuator(EpicsDevice, Movable[AttenuatorStatusEnum]): filter_material = "Al" filter_material_z = 13 filter_density = xraylib.ElementDensity(filter_material_z) # g/cm³ @@ -64,26 +67,34 @@ def __init__(self, prefix: str, num: int, thickness: int): """ self.prefix = prefix self.num = num - self.cmd = epics_signal_w(AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Cmd") + self.thickness = thickness # microns + + self.cmd = epics_signal_rw( + AttenuatorStatusEnum, f"{self.prefix}:DO{self.num + 1}-Cmd" + ) self.mode = epics_signal_r(bool, f"{self.prefix}:DIO{self.num + 1}-Mode") - self.status = epics_signal_r( - AttenuatorEnum, f"{self.prefix}:DO{self.num + 1}-Sts" + self.status = epics_signal_rw( + AttenuatorStatusEnum, f"{self.prefix}:DO{self.num + 1}-Sts" ) - self.in_status = epics_signal_r( - AttenuatorEnum, f"{self.prefix}:DI{self.num + 1}-Sts" + self.in_status = epics_signal_rw( + AttenuatorStatusEnum, f"{self.prefix}:DI{self.num + 1}-Sts" ) - self.thickness = thickness # microns super().__init__(prefix=self.prefix) def __repr__(self): return f"{self.thickness!s} microns, {self.filter_material}" + @AsyncStatus.wrap + async def set(self, value: AttenuatorStatusEnum): + await self.cmd.set(value) + + # TODO - replace these with `yield from bps.mv()` async def open(self): - await self.cmd.set(AttenuatorEnum.LOW) + await self.cmd.set(AttenuatorStatusEnum.LOW) async def close(self): - await self.cmd.set(AttenuatorEnum.HIGH) + await self.cmd.set(AttenuatorStatusEnum.HIGH) async def get_status(self): return await self.status.get_value() @@ -117,7 +128,7 @@ def attenuation(self): return np.exp(-self.linear_atten_coefficient * self.thickness_cm) -class AttenuatorBank(StandardReadable, EpicsDevice): +class AttenuatorBank(StandardReadable, EpicsDevice, Movable[float]): """ The ioc for the iologik1 lives on xf09id1-inst-ioc1.nsls2.bnl.gov """ @@ -141,8 +152,9 @@ async def get_status(self): *(a.get_status() for _, a in self.attenuators.items()) ) - async def set_attenuation(self, target_attenuation: float): - attenuation_combination = self.find_closest_attenuation(target_attenuation) + @AsyncStatus.wrap + async def set(self, value: float): + attenuation_combination = self.find_closest_attenuation(value) coros = [] for ( num, diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 408ccd6..570a42a 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -4,7 +4,11 @@ import pytest_asyncio from ophyd_async.core import get_mock_put, init_devices, set_mock_value -from cditools.attenuator import AVAILABLE_ATTENUATIONS, AttenuatorBank, AttenuatorEnum +from cditools.attenuator import ( + AVAILABLE_ATTENUATIONS, + AttenuatorBank, + AttenuatorStatusEnum, +) pytest_plugins = ("pytest_asyncio",) @@ -25,33 +29,35 @@ async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 - await mock_attenuator_bank.set_attenuation(combo0.attenuation) - atten_mock0.assert_called_with(AttenuatorEnum.LOW) - atten_mock1.assert_called_with(AttenuatorEnum.HIGH) - atten_mock2.assert_called_with(AttenuatorEnum.HIGH) - atten_mock3.assert_called_with(AttenuatorEnum.HIGH) + await mock_attenuator_bank.set(combo0.attenuation) + atten_mock0.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock1.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock3.assert_called_with(AttenuatorStatusEnum.HIGH) # AttenuatorCombination(attenuation=0.768, attenuators=[1]), combo1 = AVAILABLE_ATTENUATIONS[-3] - await mock_attenuator_bank.set_attenuation(combo1.attenuation) - atten_mock0.assert_called_with(AttenuatorEnum.LOW) - atten_mock1.assert_called_with(AttenuatorEnum.HIGH) - atten_mock2.assert_called_with(AttenuatorEnum.LOW) - atten_mock3.assert_called_with(AttenuatorEnum.LOW) + await mock_attenuator_bank.set(combo1.attenuation) + atten_mock0.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock1.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock2.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock3.assert_called_with(AttenuatorStatusEnum.LOW) @pytest.mark.asyncio async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): - set_mock_value(mock_attenuator_bank.attenuators[0].status, AttenuatorEnum.LOW) - set_mock_value(mock_attenuator_bank.attenuators[1].status, AttenuatorEnum.LOW) - set_mock_value(mock_attenuator_bank.attenuators[2].status, AttenuatorEnum.HIGH) - set_mock_value(mock_attenuator_bank.attenuators[3].status, AttenuatorEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[0].status, AttenuatorStatusEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[1].status, AttenuatorStatusEnum.LOW) + set_mock_value( + mock_attenuator_bank.attenuators[2].status, AttenuatorStatusEnum.HIGH + ) + set_mock_value(mock_attenuator_bank.attenuators[3].status, AttenuatorStatusEnum.LOW) assert await mock_attenuator_bank.get_status() == [ - AttenuatorEnum.LOW, - AttenuatorEnum.LOW, - AttenuatorEnum.HIGH, - AttenuatorEnum.LOW, + AttenuatorStatusEnum.LOW, + AttenuatorStatusEnum.LOW, + AttenuatorStatusEnum.HIGH, + AttenuatorStatusEnum.LOW, ] From e7066c9dc19999b3b348392b46af163ee82eae5a Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Tue, 12 May 2026 16:41:02 -0400 Subject: [PATCH 14/33] fixed read, write pvs for position --- src/cditools/attenuator.py | 22 ++++++++-------------- tests/test_attenuator.py | 16 ++++++++-------- 2 files changed, 16 insertions(+), 22 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index c590a13..c812287 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -69,14 +69,11 @@ def __init__(self, prefix: str, num: int, thickness: int): self.num = num self.thickness = thickness # microns - self.cmd = epics_signal_rw( - AttenuatorStatusEnum, f"{self.prefix}:DO{self.num + 1}-Cmd" + self.position = epics_signal_rw( + AttenuatorStatusEnum, f"{self.prefix}:DO{self.num + 1}-Sts", write_pv=f"{self.prefix}:DO{self.num + 1}-Cmd", ) - self.mode = epics_signal_r(bool, f"{self.prefix}:DIO{self.num + 1}-Mode") - self.status = epics_signal_rw( - AttenuatorStatusEnum, f"{self.prefix}:DO{self.num + 1}-Sts" - ) - self.in_status = epics_signal_rw( + self.mode = epics_signal_rw(bool, f"{self.prefix}:DIO{self.num + 1}-Mode") + self.in_status = epics_signal_r( AttenuatorStatusEnum, f"{self.prefix}:DI{self.num + 1}-Sts" ) @@ -87,17 +84,14 @@ def __repr__(self): @AsyncStatus.wrap async def set(self, value: AttenuatorStatusEnum): - await self.cmd.set(value) + await self.position.set(value) # TODO - replace these with `yield from bps.mv()` async def open(self): - await self.cmd.set(AttenuatorStatusEnum.LOW) + await self.position.set(AttenuatorStatusEnum.LOW) async def close(self): - await self.cmd.set(AttenuatorStatusEnum.HIGH) - - async def get_status(self): - return await self.status.get_value() + await self.position.set(AttenuatorStatusEnum.HIGH) @property def thickness_cm(self): @@ -149,7 +143,7 @@ def __init__(self): async def get_status(self): return await asyncio.gather( - *(a.get_status() for _, a in self.attenuators.items()) + *(a.position.get_value() for _, a in self.attenuators.items()) ) @AsyncStatus.wrap diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 570a42a..9757def 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -22,10 +22,10 @@ async def mock_attenuator_bank(): @pytest.mark.asyncio async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): - atten_mock0 = get_mock_put(mock_attenuator_bank.attenuators[0].cmd) - atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].cmd) - atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].cmd) - atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].cmd) + atten_mock0 = get_mock_put(mock_attenuator_bank.attenuators[0].position) + atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].position) + atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].position) + atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].position) # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 @@ -46,12 +46,12 @@ async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): @pytest.mark.asyncio async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): - set_mock_value(mock_attenuator_bank.attenuators[0].status, AttenuatorStatusEnum.LOW) - set_mock_value(mock_attenuator_bank.attenuators[1].status, AttenuatorStatusEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[0].position, AttenuatorStatusEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW) set_mock_value( - mock_attenuator_bank.attenuators[2].status, AttenuatorStatusEnum.HIGH + mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.HIGH ) - set_mock_value(mock_attenuator_bank.attenuators[3].status, AttenuatorStatusEnum.LOW) + set_mock_value(mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.LOW) assert await mock_attenuator_bank.get_status() == [ AttenuatorStatusEnum.LOW, From 706af8d25ea33baff308b0ddf1dfb13813b0f45d Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Tue, 12 May 2026 17:11:18 -0400 Subject: [PATCH 15/33] more comments, use asyncmovable protocol, changed cmd to position --- src/cditools/attenuator.py | 18 ++++++++++-------- tests/test_attenuator.py | 12 +++++++++--- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index c812287..b235489 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -6,13 +6,12 @@ import numpy as np import xraylib -from bluesky.protocols import Movable from ophyd_async.core import ( + AsyncMovable, AsyncStatus, DeviceVector, StandardReadable, StrictEnum, - # AsyncMovable, ) from ophyd_async.epics.core import EpicsDevice, epics_signal_r, epics_signal_rw @@ -48,11 +47,11 @@ class AttenuatorCombination: class AttenuatorStatusEnum(StrictEnum): - LOW = "Low" # off - HIGH = "High" # on + LOW = "Low" # off / not obstructing + HIGH = "High" # on / obstructing -class Attenuator(EpicsDevice, Movable[AttenuatorStatusEnum]): +class Attenuator(EpicsDevice, AsyncMovable[AttenuatorStatusEnum]): filter_material = "Al" filter_material_z = 13 filter_density = xraylib.ElementDensity(filter_material_z) # g/cm³ @@ -70,7 +69,9 @@ def __init__(self, prefix: str, num: int, thickness: int): self.thickness = thickness # microns self.position = epics_signal_rw( - AttenuatorStatusEnum, f"{self.prefix}:DO{self.num + 1}-Sts", write_pv=f"{self.prefix}:DO{self.num + 1}-Cmd", + AttenuatorStatusEnum, + f"{self.prefix}:DO{self.num + 1}-Sts", + write_pv=f"{self.prefix}:DO{self.num + 1}-Cmd", ) self.mode = epics_signal_rw(bool, f"{self.prefix}:DIO{self.num + 1}-Mode") self.in_status = epics_signal_r( @@ -86,11 +87,12 @@ def __repr__(self): async def set(self, value: AttenuatorStatusEnum): await self.position.set(value) - # TODO - replace these with `yield from bps.mv()` async def open(self): + """Open means open to allowing the beam to pass unobstructed""" await self.position.set(AttenuatorStatusEnum.LOW) async def close(self): + """Closed means obstructing the beam""" await self.position.set(AttenuatorStatusEnum.HIGH) @property @@ -122,7 +124,7 @@ def attenuation(self): return np.exp(-self.linear_atten_coefficient * self.thickness_cm) -class AttenuatorBank(StandardReadable, EpicsDevice, Movable[float]): +class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): """ The ioc for the iologik1 lives on xf09id1-inst-ioc1.nsls2.bnl.gov """ diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 9757def..f18d356 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -46,12 +46,18 @@ async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): @pytest.mark.asyncio async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): - set_mock_value(mock_attenuator_bank.attenuators[0].position, AttenuatorStatusEnum.LOW) - set_mock_value(mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW) + set_mock_value( + mock_attenuator_bank.attenuators[0].position, AttenuatorStatusEnum.LOW + ) + set_mock_value( + mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW + ) set_mock_value( mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.HIGH ) - set_mock_value(mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.LOW) + set_mock_value( + mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.LOW + ) assert await mock_attenuator_bank.get_status() == [ AttenuatorStatusEnum.LOW, From 56803ef883862b5889e2f7ff89715a230717c1c5 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Wed, 13 May 2026 11:09:58 -0400 Subject: [PATCH 16/33] realized that DeviceVector can have arbitrary integer keys, so does not have to be zero-indexed --- src/cditools/attenuator.py | 48 +++++++++++++++++++------------------- tests/test_attenuator.py | 45 +++++++++++++++++++++++++++-------- 2 files changed, 59 insertions(+), 34 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index b235489..fce14b4 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -27,21 +27,21 @@ class AttenuatorCombination: # so we hardcode them here # TODO - there will eventually be eight filters AVAILABLE_ATTENUATIONS = [ - AttenuatorCombination(attenuation=0.08, attenuators=[0, 1, 2, 3]), - AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), - AttenuatorCombination(attenuation=0.104, attenuators=[0, 2, 3]), - AttenuatorCombination(attenuation=0.124, attenuators=[2, 3]), - AttenuatorCombination(attenuation=0.165, attenuators=[0, 1, 3]), - AttenuatorCombination(attenuation=0.196, attenuators=[1, 3]), - AttenuatorCombination(attenuation=0.214, attenuators=[0, 3]), - AttenuatorCombination(attenuation=0.256, attenuators=[3]), - AttenuatorCombination(attenuation=0.312, attenuators=[0, 1, 2]), - AttenuatorCombination(attenuation=0.372, attenuators=[1, 2]), - AttenuatorCombination(attenuation=0.406, attenuators=[0, 2]), - AttenuatorCombination(attenuation=0.484, attenuators=[2]), - AttenuatorCombination(attenuation=0.644, attenuators=[0, 1]), - AttenuatorCombination(attenuation=0.768, attenuators=[1]), - AttenuatorCombination(attenuation=0.839, attenuators=[0]), + AttenuatorCombination(attenuation=0.08, attenuators=[1, 2, 3, 4]), + AttenuatorCombination(attenuation=0.095, attenuators=[2, 3, 4]), + AttenuatorCombination(attenuation=0.104, attenuators=[1, 3, 4]), + AttenuatorCombination(attenuation=0.124, attenuators=[3, 4]), + AttenuatorCombination(attenuation=0.165, attenuators=[1, 2, 4]), + AttenuatorCombination(attenuation=0.196, attenuators=[2, 4]), + AttenuatorCombination(attenuation=0.214, attenuators=[1, 4]), + AttenuatorCombination(attenuation=0.256, attenuators=[4]), + AttenuatorCombination(attenuation=0.312, attenuators=[1, 2, 3]), + AttenuatorCombination(attenuation=0.372, attenuators=[2, 3]), + AttenuatorCombination(attenuation=0.406, attenuators=[1, 3]), + AttenuatorCombination(attenuation=0.484, attenuators=[3]), + AttenuatorCombination(attenuation=0.644, attenuators=[1, 2]), + AttenuatorCombination(attenuation=0.768, attenuators=[2]), + AttenuatorCombination(attenuation=0.839, attenuators=[1]), AttenuatorCombination(attenuation=1.0, attenuators=[]), ] @@ -60,9 +60,9 @@ def __init__(self, prefix: str, num: int, thickness: int): """ prefix - the common prefix for the attenuator bank num - an integer denoting which attenuator within the bank this is - cmd - the write PV to open and close the attenuator - status - the read PV for the status (high or low) thickness - the thickness of the attenuator in microns + + position - the read / write PV to open and close the attenuator """ self.prefix = prefix self.num = num @@ -70,12 +70,12 @@ def __init__(self, prefix: str, num: int, thickness: int): self.position = epics_signal_rw( AttenuatorStatusEnum, - f"{self.prefix}:DO{self.num + 1}-Sts", - write_pv=f"{self.prefix}:DO{self.num + 1}-Cmd", + f"{self.prefix}:DO{self.num}-Sts", + write_pv=f"{self.prefix}:DO{self.num}-Cmd", ) - self.mode = epics_signal_rw(bool, f"{self.prefix}:DIO{self.num + 1}-Mode") + self.mode = epics_signal_rw(bool, f"{self.prefix}:DIO{self.num}-Mode") self.in_status = epics_signal_r( - AttenuatorStatusEnum, f"{self.prefix}:DI{self.num + 1}-Sts" + AttenuatorStatusEnum, f"{self.prefix}:DI{self.num}-Sts" ) super().__init__(prefix=self.prefix) @@ -137,8 +137,8 @@ def __init__(self): with self.add_children_as_readables(): self.attenuators = DeviceVector( { - i: Attenuator(self.prefix, i, self.thicknesses[i]) - for i in range(len(self.thicknesses)) + i: Attenuator(self.prefix, i, self.thicknesses[i - 1]) + for i in range(1, len(self.thicknesses) + 1) } ) super().__init__(prefix=self.prefix) @@ -226,6 +226,6 @@ def _powerset(self) -> list[list[int]]: combination = [] for j in range(len(self.attenuators)): if i & (1 << j): - combination.append(j) + combination.append(j + 1) # +1 because attenuators are 1-indexed powerset.append(combination) return powerset diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index f18d356..36b7fd6 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -20,43 +20,68 @@ async def mock_attenuator_bank(): yield mock_attenuator_bank +@pytest.mark.asyncio +async def test_attenuators_indexed_at_1(mock_attenuator_bank: AttenuatorBank): + with pytest.raises(KeyError): + mock_attenuator_bank.attenuators[0] + + atten1 = mock_attenuator_bank.attenuators[1] + assert atten1.num == 1 + assert atten1.thickness == 16 + assert atten1.position.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DO1-Sts" + assert atten1.mode.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DIO1-Mode" + assert atten1.in_status.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DI1-Sts" + + atten2 = mock_attenuator_bank.attenuators[2] + assert atten2.num == 2 + assert atten2.thickness == 24 + + atten3 = mock_attenuator_bank.attenuators[3] + assert atten3.num == 3 + assert atten3.thickness == 66 + + atten4 = mock_attenuator_bank.attenuators[4] + assert atten4.num == 4 + assert atten4.thickness == 124 + + @pytest.mark.asyncio async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): - atten_mock0 = get_mock_put(mock_attenuator_bank.attenuators[0].position) atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].position) atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].position) atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].position) + atten_mock4 = get_mock_put(mock_attenuator_bank.attenuators[4].position) # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 await mock_attenuator_bank.set(combo0.attenuation) - atten_mock0.assert_called_with(AttenuatorStatusEnum.LOW) - atten_mock1.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) atten_mock3.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock4.assert_called_with(AttenuatorStatusEnum.HIGH) # AttenuatorCombination(attenuation=0.768, attenuators=[1]), combo1 = AVAILABLE_ATTENUATIONS[-3] await mock_attenuator_bank.set(combo1.attenuation) - atten_mock0.assert_called_with(AttenuatorStatusEnum.LOW) - atten_mock1.assert_called_with(AttenuatorStatusEnum.HIGH) - atten_mock2.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) atten_mock3.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock4.assert_called_with(AttenuatorStatusEnum.LOW) @pytest.mark.asyncio async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): set_mock_value( - mock_attenuator_bank.attenuators[0].position, AttenuatorStatusEnum.LOW + mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW ) set_mock_value( - mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW + mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.LOW ) set_mock_value( - mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.HIGH + mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH ) set_mock_value( - mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.LOW + mock_attenuator_bank.attenuators[4].position, AttenuatorStatusEnum.LOW ) assert await mock_attenuator_bank.get_status() == [ From d27711ee4e0473888b4a455efda8a51c659295e5 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Wed, 13 May 2026 11:19:03 -0400 Subject: [PATCH 17/33] made thicknesses a constant --- src/cditools/attenuator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index fce14b4..9da37b5 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -22,6 +22,8 @@ class AttenuatorCombination: attenuators: list[int] +THICKNESSES = (16, 24, 66, 124) # microns + # The available attenuations can be calculated with the utility # methods below, but they do not change often, # so we hardcode them here @@ -130,7 +132,7 @@ class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): """ prefix = "XF:09ID1-ES{IOLOGIK1:E1212}" - thicknesses = (16, 24, 66, 124) # microns + thicknesses = THICKNESSES available_attenuations = AVAILABLE_ATTENUATIONS def __init__(self): From 0745eb56cf141b9d674fa9e28fcecc136d70b483 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 14 May 2026 14:31:29 -0400 Subject: [PATCH 18/33] corrected terms transmission and attenuation, added some tests --- src/cditools/attenuator.py | 62 +++++++++++++++++++++++--------------- 1 file changed, 37 insertions(+), 25 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 9da37b5..bacbbb4 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -18,7 +18,7 @@ @dataclass class AttenuatorCombination: - attenuation: float + transmission: float attenuators: list[int] @@ -29,22 +29,22 @@ class AttenuatorCombination: # so we hardcode them here # TODO - there will eventually be eight filters AVAILABLE_ATTENUATIONS = [ - AttenuatorCombination(attenuation=0.08, attenuators=[1, 2, 3, 4]), - AttenuatorCombination(attenuation=0.095, attenuators=[2, 3, 4]), - AttenuatorCombination(attenuation=0.104, attenuators=[1, 3, 4]), - AttenuatorCombination(attenuation=0.124, attenuators=[3, 4]), - AttenuatorCombination(attenuation=0.165, attenuators=[1, 2, 4]), - AttenuatorCombination(attenuation=0.196, attenuators=[2, 4]), - AttenuatorCombination(attenuation=0.214, attenuators=[1, 4]), - AttenuatorCombination(attenuation=0.256, attenuators=[4]), - AttenuatorCombination(attenuation=0.312, attenuators=[1, 2, 3]), - AttenuatorCombination(attenuation=0.372, attenuators=[2, 3]), - AttenuatorCombination(attenuation=0.406, attenuators=[1, 3]), - AttenuatorCombination(attenuation=0.484, attenuators=[3]), - AttenuatorCombination(attenuation=0.644, attenuators=[1, 2]), - AttenuatorCombination(attenuation=0.768, attenuators=[2]), - AttenuatorCombination(attenuation=0.839, attenuators=[1]), - AttenuatorCombination(attenuation=1.0, attenuators=[]), + AttenuatorCombination(transmission=0.08, attenuators=[1, 2, 3, 4]), + AttenuatorCombination(transmission=0.095, attenuators=[2, 3, 4]), + AttenuatorCombination(transmission=0.104, attenuators=[1, 3, 4]), + AttenuatorCombination(transmission=0.124, attenuators=[3, 4]), + AttenuatorCombination(transmission=0.165, attenuators=[1, 2, 4]), + AttenuatorCombination(transmission=0.196, attenuators=[2, 4]), + AttenuatorCombination(transmission=0.214, attenuators=[1, 4]), + AttenuatorCombination(transmission=0.256, attenuators=[4]), + AttenuatorCombination(transmission=0.312, attenuators=[1, 2, 3]), + AttenuatorCombination(transmission=0.372, attenuators=[2, 3]), + AttenuatorCombination(transmission=0.406, attenuators=[1, 3]), + AttenuatorCombination(transmission=0.484, attenuators=[3]), + AttenuatorCombination(transmission=0.644, attenuators=[1, 2]), + AttenuatorCombination(transmission=0.768, attenuators=[2]), + AttenuatorCombination(transmission=0.839, attenuators=[1]), + AttenuatorCombination(transmission=1.0, attenuators=[]), ] @@ -119,12 +119,15 @@ def linear_atten_coefficient(self) -> float: return mass_atten_cross_section * self.filter_density @property - def attenuation(self): - """ - Attenuation is the fraction of remaining beam - """ + def transmission(self): + """Transmission is the fraction of remaining beam""" return np.exp(-self.linear_atten_coefficient * self.thickness_cm) + @property + def attenuation(self): + """Attenuation is the fraction of the beam removed""" + return 1 - self.transmission + class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): """ @@ -173,7 +176,7 @@ def find_closest_attenuation( is small, so we start in the middle, and work up or down. """ best_idx = len(self.available_attenuations) // 2 - atten = self.available_attenuations[best_idx].attenuation + atten = self.available_attenuations[best_idx].transmission diff = float("inf") new_diff = abs(target_attenuation - atten) inc = 1 if target_attenuation > atten else -1 @@ -183,7 +186,7 @@ def find_closest_attenuation( # break if we are about to check oustide the list if best_idx + inc >= len(self.available_attenuations) or best_idx + inc < 0: break - atten = self.available_attenuations[best_idx + inc].attenuation + atten = self.available_attenuations[best_idx + inc].transmission new_diff = abs(target_attenuation - atten) if new_diff < diff: best_idx += inc @@ -213,11 +216,11 @@ def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: AttenuatorCombination(total_atten, combination) ) # We want the available attenuations sorted so we can efficiently search through them - available_attenuations.sort(key=lambda a: a.attenuation) # type: ignore[attr-defined] + available_attenuations.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] return available_attenuations def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float: - return round(float(math.prod([a.attenuation for a in attenuators])), 3) + return round(float(math.prod([a.transmission for a in attenuators])), 3) def _powerset(self) -> list[list[int]]: """ @@ -231,3 +234,12 @@ def _powerset(self) -> list[list[int]]: combination.append(j + 1) # +1 because attenuators are 1-indexed powerset.append(combination) return powerset + + +""" +from cditools.attenuator import AttenuatorBank + +bank = AttenuatorBank() +atten = bank.attenuators[1] + +""" From dee2024d56eb843e014b6bb8a45e3586a24b6640 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 14 May 2026 14:31:57 -0400 Subject: [PATCH 19/33] corrected terms transmission and attenuation, added some tests --- pyproject.toml | 6 +++++- tests/test_attenuator.py | 28 ++++++++++++++++++++++------ 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e376a4e..2aba173 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,8 @@ dependencies = [ "ophyd", "ophyd-async[ca] >=0.19", "h5py", - "xraylib >=4.2.1,<5" + "xraylib >=4.2.1,<5", + "xrayutilities>=1.7.12,<2" ] [project.optional-dependencies] @@ -168,6 +169,9 @@ platforms = ["linux-64", "osx-arm64"] [tool.pixi.pypi-dependencies] cditools = { path = ".", editable = true } +[tool.pixi.dependencies] +xrayutilities = ">=1.7.12,<2" + [tool.pixi.environments] default = { solve-group = "default" } dev = { features = ["dev"], solve-group = "default" } diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 36b7fd6..462753b 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -6,6 +6,7 @@ from cditools.attenuator import ( AVAILABLE_ATTENUATIONS, + Attenuator, AttenuatorBank, AttenuatorStatusEnum, ) @@ -20,6 +21,21 @@ async def mock_attenuator_bank(): yield mock_attenuator_bank +@pytest_asyncio.fixture +async def mock_attenuator(mock_attenuator_bank: AttenuatorBank): + async with init_devices(mock=True): + mock_attenuator = Attenuator(mock_attenuator_bank.prefix, 1, 16) + yield mock_attenuator + + +def test_transmission(mock_attenuator: Attenuator): + assert mock_attenuator.transmission == pytest.approx(0.84, abs=0.01) + + +def test_attenuation(mock_attenuator: Attenuator): + assert mock_attenuator.attenuation == pytest.approx(0.16, abs=0.01) + + @pytest.mark.asyncio async def test_attenuators_indexed_at_1(mock_attenuator_bank: AttenuatorBank): with pytest.raises(KeyError): @@ -54,7 +70,7 @@ async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 - await mock_attenuator_bank.set(combo0.attenuation) + await mock_attenuator_bank.set(combo0.transmission) atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) atten_mock3.assert_called_with(AttenuatorStatusEnum.HIGH) @@ -62,7 +78,7 @@ async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): # AttenuatorCombination(attenuation=0.768, attenuators=[1]), combo1 = AVAILABLE_ATTENUATIONS[-3] - await mock_attenuator_bank.set(combo1.attenuation) + await mock_attenuator_bank.set(combo1.transmission) atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) atten_mock3.assert_called_with(AttenuatorStatusEnum.LOW) @@ -94,16 +110,16 @@ async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest = mock_attenuator_bank.find_closest_attenuation(0.7) - assert nearest.attenuation == 0.644 + assert nearest.transmission == 0.644 nearest2 = mock_attenuator_bank.find_closest_attenuation(0.2) - assert nearest2.attenuation == 0.196 + assert nearest2.transmission == 0.196 nearest3 = mock_attenuator_bank.find_closest_attenuation(0.02) - assert nearest3.attenuation == 0.08 + assert nearest3.transmission == 0.08 nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) - assert nearest4.attenuation == 1 + assert nearest4.transmission == 1 def test_up_to_date_available_attenuations(mock_attenuator_bank: AttenuatorBank): From aee0a9dc2438b27fbb83f42aa1859d2c4e12dd92 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 14 May 2026 16:09:52 -0400 Subject: [PATCH 20/33] made photon_energy and units parameters to pass into calculations, used xrayutilities to calculate the transmission --- src/cditools/attenuator.py | 106 +++++++++++++++++++------------------ tests/test_attenuator.py | 31 ++++++++--- 2 files changed, 77 insertions(+), 60 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index bacbbb4..57e1e14 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -6,6 +6,7 @@ import numpy as np import xraylib +import xrayutilities as xu from ophyd_async.core import ( AsyncMovable, AsyncStatus, @@ -28,22 +29,23 @@ class AttenuatorCombination: # methods below, but they do not change often, # so we hardcode them here # TODO - there will eventually be eight filters + AVAILABLE_ATTENUATIONS = [ - AttenuatorCombination(transmission=0.08, attenuators=[1, 2, 3, 4]), - AttenuatorCombination(transmission=0.095, attenuators=[2, 3, 4]), - AttenuatorCombination(transmission=0.104, attenuators=[1, 3, 4]), - AttenuatorCombination(transmission=0.124, attenuators=[3, 4]), - AttenuatorCombination(transmission=0.165, attenuators=[1, 2, 4]), - AttenuatorCombination(transmission=0.196, attenuators=[2, 4]), - AttenuatorCombination(transmission=0.214, attenuators=[1, 4]), - AttenuatorCombination(transmission=0.256, attenuators=[4]), - AttenuatorCombination(transmission=0.312, attenuators=[1, 2, 3]), - AttenuatorCombination(transmission=0.372, attenuators=[2, 3]), - AttenuatorCombination(transmission=0.406, attenuators=[1, 3]), - AttenuatorCombination(transmission=0.484, attenuators=[3]), - AttenuatorCombination(transmission=0.644, attenuators=[1, 2]), - AttenuatorCombination(transmission=0.768, attenuators=[2]), - AttenuatorCombination(transmission=0.839, attenuators=[1]), + AttenuatorCombination(transmission=0.084, attenuators=[1, 2, 3, 4]), + AttenuatorCombination(transmission=0.1, attenuators=[2, 3, 4]), + AttenuatorCombination(transmission=0.109, attenuators=[1, 3, 4]), + AttenuatorCombination(transmission=0.129, attenuators=[3, 4]), + AttenuatorCombination(transmission=0.171, attenuators=[1, 2, 4]), + AttenuatorCombination(transmission=0.203, attenuators=[2, 4]), + AttenuatorCombination(transmission=0.222, attenuators=[1, 4]), + AttenuatorCombination(transmission=0.263, attenuators=[4]), + AttenuatorCombination(transmission=0.32, attenuators=[1, 2, 3]), + AttenuatorCombination(transmission=0.38, attenuators=[2, 3]), + AttenuatorCombination(transmission=0.414, attenuators=[1, 3]), + AttenuatorCombination(transmission=0.492, attenuators=[3]), + AttenuatorCombination(transmission=0.65, attenuators=[1, 2]), + AttenuatorCombination(transmission=0.772, attenuators=[2]), + AttenuatorCombination(transmission=0.842, attenuators=[1]), AttenuatorCombination(transmission=1.0, attenuators=[]), ] @@ -97,36 +99,30 @@ async def close(self): """Closed means obstructing the beam""" await self.position.set(AttenuatorStatusEnum.HIGH) - @property - def thickness_cm(self): - # Thickness is in microns, so convert to cm - return self.thickness * 1e-4 + def transmission(self, photon_energy: float, units: str = "KeV"): + """Transmission is the fraction of remaining beam""" + # return np.exp(-self.linear_atten_coefficient(photon_energy) * self.thickness_cm) + abs_len = self._absorption_length(photon_energy, units=units) + return np.exp(-self.thickness / abs_len) - @property - def linear_atten_coefficient(self) -> float: - """ - Calculates µ, the linear attenuation coefficient of this material, - at this thickness, and this beam energy. + def attenuation(self, photon_energy: float, units: str = "KeV"): + """Attenuation is the fraction of the beam removed""" + return 1 - self.transmission(photon_energy, units=units) - photon energy in KeV - xraylib.CS_Total in cm²/g - linear_atten_coefficient in cm⁻¹ + def _absorption_length(self, photon_energy: float, units: str = "KeV") -> float: """ - photon_energy = 8.6 # KeV TODO - get the right number; this is taken from bmm - mass_atten_cross_section = xraylib.CS_Total( - self.filter_material_z, photon_energy - ) - return mass_atten_cross_section * self.filter_density - - @property - def transmission(self): - """Transmission is the fraction of remaining beam""" - return np.exp(-self.linear_atten_coefficient * self.thickness_cm) + Calculates L, the characteristic absorption length of this material, + at this beam energy. - @property - def attenuation(self): - """Attenuation is the fraction of the beam removed""" - return 1 - self.transmission + photon energy in KeV or eV + absorption length in microns + """ + if units == "KeV": + photon_energy = photon_energy * 1e3 + elif units != "eV": + msg = "Photon energy units must be eV or KeV" + raise RuntimeError(msg) + return xu.materials.Al.absorption_length(photon_energy) # type: ignore[reportArgumentType] class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): @@ -196,14 +192,9 @@ def find_closest_attenuation( # requested attenuation and/or the difference? return self.available_attenuations[best_idx] - """ - These are utility methods that should not be called during production. - They are used to calculate the available attenuations from all - combinations of attenuators. The result is then used as the - AttenuationBank()._available_attenuations attribute. - """ - - def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: + def _calculate_available_attentuations( + self, photon_energy: float, units: str = "KeV" + ) -> list[AttenuatorCombination]: """ It is more efficient to precompute all possible total attenuations, and simply look up the closest one. @@ -211,7 +202,9 @@ def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: available_attenuations = [] for combination in self._powerset(): attens = [self.attenuators[a] for a in self.attenuators if a in combination] - total_atten = self._calculate_total_attenuation(*attens) + total_atten = self._calculate_total_attenuation( + *attens, photon_energy=photon_energy, units=units + ) available_attenuations.append( AttenuatorCombination(total_atten, combination) ) @@ -219,8 +212,17 @@ def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: available_attenuations.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] return available_attenuations - def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float: - return round(float(math.prod([a.transmission for a in attenuators])), 3) + def _calculate_total_attenuation( + self, *attenuators: Attenuator, photon_energy: float, units: str = "KeV" + ) -> float: + return round( + float( + math.prod( + [a.transmission(photon_energy, units=units) for a in attenuators] + ) + ), + 3, + ) def _powerset(self) -> list[list[int]]: """ diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 462753b..cc6a7ab 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -12,6 +12,7 @@ ) pytest_plugins = ("pytest_asyncio",) +photon_energy = 8.6 # KeV @pytest_asyncio.fixture @@ -28,12 +29,26 @@ async def mock_attenuator(mock_attenuator_bank: AttenuatorBank): yield mock_attenuator -def test_transmission(mock_attenuator: Attenuator): - assert mock_attenuator.transmission == pytest.approx(0.84, abs=0.01) +def test_transmission_kev(mock_attenuator: Attenuator): + assert mock_attenuator.transmission(photon_energy) == pytest.approx(0.84, abs=0.01) -def test_attenuation(mock_attenuator: Attenuator): - assert mock_attenuator.attenuation == pytest.approx(0.16, abs=0.01) +def test_transmission_ev(mock_attenuator: Attenuator): + photon_energy = 8600 # eV + assert mock_attenuator.transmission(photon_energy, units="eV") == pytest.approx( + 0.84, abs=0.01 + ) + + +def test_attenuation_kev(mock_attenuator: Attenuator): + assert mock_attenuator.attenuation(photon_energy) == pytest.approx(0.16, abs=0.01) + + +def test_attenuation_ev(mock_attenuator: Attenuator): + photon_energy = 8600 # eV + assert mock_attenuator.attenuation(photon_energy, units="eV") == pytest.approx( + 0.16, abs=0.01 + ) @pytest.mark.asyncio @@ -110,13 +125,13 @@ async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): nearest = mock_attenuator_bank.find_closest_attenuation(0.7) - assert nearest.transmission == 0.644 + assert nearest.transmission == 0.65 nearest2 = mock_attenuator_bank.find_closest_attenuation(0.2) - assert nearest2.transmission == 0.196 + assert nearest2.transmission == 0.203 nearest3 = mock_attenuator_bank.find_closest_attenuation(0.02) - assert nearest3.transmission == 0.08 + assert nearest3.transmission == 0.084 nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) assert nearest4.transmission == 1 @@ -124,6 +139,6 @@ def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): def test_up_to_date_available_attenuations(mock_attenuator_bank: AttenuatorBank): assert ( - mock_attenuator_bank._calculate_available_attentuations() # type: ignore[reportPrivateUsage] + mock_attenuator_bank._calculate_available_attentuations(photon_energy) # type: ignore[reportPrivateUsage] == AVAILABLE_ATTENUATIONS ) From 8a21bdaed3ea86bdeab0cd2222205645ca3020a5 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 14 May 2026 16:25:28 -0400 Subject: [PATCH 21/33] fully removed xraylib; minor refactoring --- pyproject.toml | 1 - src/cditools/attenuator.py | 8 ++------ 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2aba173..aecced6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,6 @@ dependencies = [ "ophyd", "ophyd-async[ca] >=0.19", "h5py", - "xraylib >=4.2.1,<5", "xrayutilities>=1.7.12,<2" ] diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 57e1e14..230c73e 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -5,7 +5,6 @@ from dataclasses import dataclass import numpy as np -import xraylib import xrayutilities as xu from ophyd_async.core import ( AsyncMovable, @@ -29,7 +28,6 @@ class AttenuatorCombination: # methods below, but they do not change often, # so we hardcode them here # TODO - there will eventually be eight filters - AVAILABLE_ATTENUATIONS = [ AttenuatorCombination(transmission=0.084, attenuators=[1, 2, 3, 4]), AttenuatorCombination(transmission=0.1, attenuators=[2, 3, 4]), @@ -56,9 +54,7 @@ class AttenuatorStatusEnum(StrictEnum): class Attenuator(EpicsDevice, AsyncMovable[AttenuatorStatusEnum]): - filter_material = "Al" - filter_material_z = 13 - filter_density = xraylib.ElementDensity(filter_material_z) # g/cm³ + filter_material = xu.materials.Al def __init__(self, prefix: str, num: int, thickness: int): """ @@ -122,7 +118,7 @@ def _absorption_length(self, photon_energy: float, units: str = "KeV") -> float: elif units != "eV": msg = "Photon energy units must be eV or KeV" raise RuntimeError(msg) - return xu.materials.Al.absorption_length(photon_energy) # type: ignore[reportArgumentType] + return self.filter_material.absorption_length(photon_energy) # type: ignore[reportArgumentType] class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): From 8f3d086d2f2be5df2fcf76465f6c97fdc87cb3cf Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 14 May 2026 16:43:42 -0400 Subject: [PATCH 22/33] added tests --- tests/test_attenuator.py | 242 +++++++++++++++++++++------------------ 1 file changed, 129 insertions(+), 113 deletions(-) diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index cc6a7ab..0b5cef8 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -3,6 +3,7 @@ import pytest import pytest_asyncio from ophyd_async.core import get_mock_put, init_devices, set_mock_value +from ophyd_async.testing import assert_value from cditools.attenuator import ( AVAILABLE_ATTENUATIONS, @@ -29,116 +30,131 @@ async def mock_attenuator(mock_attenuator_bank: AttenuatorBank): yield mock_attenuator -def test_transmission_kev(mock_attenuator: Attenuator): - assert mock_attenuator.transmission(photon_energy) == pytest.approx(0.84, abs=0.01) - - -def test_transmission_ev(mock_attenuator: Attenuator): - photon_energy = 8600 # eV - assert mock_attenuator.transmission(photon_energy, units="eV") == pytest.approx( - 0.84, abs=0.01 - ) - - -def test_attenuation_kev(mock_attenuator: Attenuator): - assert mock_attenuator.attenuation(photon_energy) == pytest.approx(0.16, abs=0.01) - - -def test_attenuation_ev(mock_attenuator: Attenuator): - photon_energy = 8600 # eV - assert mock_attenuator.attenuation(photon_energy, units="eV") == pytest.approx( - 0.16, abs=0.01 - ) - - -@pytest.mark.asyncio -async def test_attenuators_indexed_at_1(mock_attenuator_bank: AttenuatorBank): - with pytest.raises(KeyError): - mock_attenuator_bank.attenuators[0] - - atten1 = mock_attenuator_bank.attenuators[1] - assert atten1.num == 1 - assert atten1.thickness == 16 - assert atten1.position.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DO1-Sts" - assert atten1.mode.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DIO1-Mode" - assert atten1.in_status.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DI1-Sts" - - atten2 = mock_attenuator_bank.attenuators[2] - assert atten2.num == 2 - assert atten2.thickness == 24 - - atten3 = mock_attenuator_bank.attenuators[3] - assert atten3.num == 3 - assert atten3.thickness == 66 - - atten4 = mock_attenuator_bank.attenuators[4] - assert atten4.num == 4 - assert atten4.thickness == 124 - - -@pytest.mark.asyncio -async def test_set_attenuators(mock_attenuator_bank: AttenuatorBank): - atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].position) - atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].position) - atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].position) - atten_mock4 = get_mock_put(mock_attenuator_bank.attenuators[4].position) - - # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), - combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 - await mock_attenuator_bank.set(combo0.transmission) - atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) - atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) - atten_mock3.assert_called_with(AttenuatorStatusEnum.HIGH) - atten_mock4.assert_called_with(AttenuatorStatusEnum.HIGH) - - # AttenuatorCombination(attenuation=0.768, attenuators=[1]), - combo1 = AVAILABLE_ATTENUATIONS[-3] - await mock_attenuator_bank.set(combo1.transmission) - atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) - atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) - atten_mock3.assert_called_with(AttenuatorStatusEnum.LOW) - atten_mock4.assert_called_with(AttenuatorStatusEnum.LOW) - - -@pytest.mark.asyncio -async def test_get_bank_status(mock_attenuator_bank: AttenuatorBank): - set_mock_value( - mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW - ) - set_mock_value( - mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.LOW - ) - set_mock_value( - mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH - ) - set_mock_value( - mock_attenuator_bank.attenuators[4].position, AttenuatorStatusEnum.LOW - ) - - assert await mock_attenuator_bank.get_status() == [ - AttenuatorStatusEnum.LOW, - AttenuatorStatusEnum.LOW, - AttenuatorStatusEnum.HIGH, - AttenuatorStatusEnum.LOW, - ] - - -def test_find_closest_attenuation(mock_attenuator_bank: AttenuatorBank): - nearest = mock_attenuator_bank.find_closest_attenuation(0.7) - assert nearest.transmission == 0.65 - - nearest2 = mock_attenuator_bank.find_closest_attenuation(0.2) - assert nearest2.transmission == 0.203 - - nearest3 = mock_attenuator_bank.find_closest_attenuation(0.02) - assert nearest3.transmission == 0.084 - - nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) - assert nearest4.transmission == 1 - - -def test_up_to_date_available_attenuations(mock_attenuator_bank: AttenuatorBank): - assert ( - mock_attenuator_bank._calculate_available_attentuations(photon_energy) # type: ignore[reportPrivateUsage] - == AVAILABLE_ATTENUATIONS - ) +class TestAttenuator: + @pytest.mark.asyncio + async def test_open(self, mock_attenuator: Attenuator): + set_mock_value(mock_attenuator.position, AttenuatorStatusEnum.HIGH) + await mock_attenuator.open() + await assert_value(mock_attenuator.position, AttenuatorStatusEnum.LOW) + + @pytest.mark.asyncio + async def test_close(self, mock_attenuator: Attenuator): + set_mock_value(mock_attenuator.position, AttenuatorStatusEnum.LOW) + await mock_attenuator.close() + await assert_value(mock_attenuator.position, AttenuatorStatusEnum.HIGH) + + def test_transmission_kev(self, mock_attenuator: Attenuator): + assert mock_attenuator.transmission(photon_energy) == pytest.approx( + 0.84, abs=0.01 + ) + + def test_transmission_ev(self, mock_attenuator: Attenuator): + photon_energy = 8600 # eV + assert mock_attenuator.transmission(photon_energy, units="eV") == pytest.approx( + 0.84, abs=0.01 + ) + + def test_attenuation_kev(self, mock_attenuator: Attenuator): + assert mock_attenuator.attenuation(photon_energy) == pytest.approx( + 0.16, abs=0.01 + ) + + def test_attenuation_ev(self, mock_attenuator: Attenuator): + photon_energy = 8600 # eV + assert mock_attenuator.attenuation(photon_energy, units="eV") == pytest.approx( + 0.16, abs=0.01 + ) + + +class TestAttenuatorBank: + @pytest.mark.asyncio + async def test_attenuators_indexed_at_1(self, mock_attenuator_bank: AttenuatorBank): + with pytest.raises(KeyError): + mock_attenuator_bank.attenuators[0] + + atten1 = mock_attenuator_bank.attenuators[1] + assert atten1.num == 1 + assert atten1.thickness == 16 + assert atten1.position.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DO1-Sts" + assert atten1.mode.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DIO1-Mode" + assert ( + atten1.in_status.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DI1-Sts" + ) + + atten2 = mock_attenuator_bank.attenuators[2] + assert atten2.num == 2 + assert atten2.thickness == 24 + + atten3 = mock_attenuator_bank.attenuators[3] + assert atten3.num == 3 + assert atten3.thickness == 66 + + atten4 = mock_attenuator_bank.attenuators[4] + assert atten4.num == 4 + assert atten4.thickness == 124 + + @pytest.mark.asyncio + async def test_set_attenuators(self, mock_attenuator_bank: AttenuatorBank): + atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].position) + atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].position) + atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].position) + atten_mock4 = get_mock_put(mock_attenuator_bank.attenuators[4].position) + + # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), + combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 + await mock_attenuator_bank.set(combo0.transmission) + atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock3.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock4.assert_called_with(AttenuatorStatusEnum.HIGH) + + # AttenuatorCombination(attenuation=0.768, attenuators=[1]), + combo1 = AVAILABLE_ATTENUATIONS[-3] + await mock_attenuator_bank.set(combo1.transmission) + atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) + atten_mock3.assert_called_with(AttenuatorStatusEnum.LOW) + atten_mock4.assert_called_with(AttenuatorStatusEnum.LOW) + + @pytest.mark.asyncio + async def test_get_bank_status(self, mock_attenuator_bank: AttenuatorBank): + set_mock_value( + mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW + ) + set_mock_value( + mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.LOW + ) + set_mock_value( + mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH + ) + set_mock_value( + mock_attenuator_bank.attenuators[4].position, AttenuatorStatusEnum.LOW + ) + + assert await mock_attenuator_bank.get_status() == [ + AttenuatorStatusEnum.LOW, + AttenuatorStatusEnum.LOW, + AttenuatorStatusEnum.HIGH, + AttenuatorStatusEnum.LOW, + ] + + def test_find_closest_attenuation(self, mock_attenuator_bank: AttenuatorBank): + nearest = mock_attenuator_bank.find_closest_attenuation(0.7) + assert nearest.transmission == 0.65 + + nearest2 = mock_attenuator_bank.find_closest_attenuation(0.2) + assert nearest2.transmission == 0.203 + + nearest3 = mock_attenuator_bank.find_closest_attenuation(0.02) + assert nearest3.transmission == 0.084 + + nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) + assert nearest4.transmission == 1 + + def test_up_to_date_available_attenuations( + self, mock_attenuator_bank: AttenuatorBank + ): + assert ( + mock_attenuator_bank._calculate_available_attentuations(photon_energy) # type: ignore[reportPrivateUsage] + == AVAILABLE_ATTENUATIONS + ) From be8f3b49b12c15beab750306e703a3ae51142e20 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 14 May 2026 19:15:43 -0400 Subject: [PATCH 23/33] pass energy object into AttenuatorBank on creation, so it can poll for photon energy --- src/cditools/attenuator.py | 103 ++++++++++++++----------------------- tests/test_attenuator.py | 63 ++++++++++++++++++----- 2 files changed, 88 insertions(+), 78 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 230c73e..490a2d7 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -15,37 +15,20 @@ ) from ophyd_async.epics.core import EpicsDevice, epics_signal_r, epics_signal_rw +from cditools.motors import Energy + @dataclass class AttenuatorCombination: transmission: float attenuators: list[int] + @property + def attenuation(self): + return 1 - self.transmission -THICKNESSES = (16, 24, 66, 124) # microns -# The available attenuations can be calculated with the utility -# methods below, but they do not change often, -# so we hardcode them here -# TODO - there will eventually be eight filters -AVAILABLE_ATTENUATIONS = [ - AttenuatorCombination(transmission=0.084, attenuators=[1, 2, 3, 4]), - AttenuatorCombination(transmission=0.1, attenuators=[2, 3, 4]), - AttenuatorCombination(transmission=0.109, attenuators=[1, 3, 4]), - AttenuatorCombination(transmission=0.129, attenuators=[3, 4]), - AttenuatorCombination(transmission=0.171, attenuators=[1, 2, 4]), - AttenuatorCombination(transmission=0.203, attenuators=[2, 4]), - AttenuatorCombination(transmission=0.222, attenuators=[1, 4]), - AttenuatorCombination(transmission=0.263, attenuators=[4]), - AttenuatorCombination(transmission=0.32, attenuators=[1, 2, 3]), - AttenuatorCombination(transmission=0.38, attenuators=[2, 3]), - AttenuatorCombination(transmission=0.414, attenuators=[1, 3]), - AttenuatorCombination(transmission=0.492, attenuators=[3]), - AttenuatorCombination(transmission=0.65, attenuators=[1, 2]), - AttenuatorCombination(transmission=0.772, attenuators=[2]), - AttenuatorCombination(transmission=0.842, attenuators=[1]), - AttenuatorCombination(transmission=1.0, attenuators=[]), -] +THICKNESSES = (16, 24, 66, 124) # microns class AttenuatorStatusEnum(StrictEnum): @@ -95,17 +78,16 @@ async def close(self): """Closed means obstructing the beam""" await self.position.set(AttenuatorStatusEnum.HIGH) - def transmission(self, photon_energy: float, units: str = "KeV"): + def transmission(self, photon_energy: float, egu: str = "KeV"): """Transmission is the fraction of remaining beam""" - # return np.exp(-self.linear_atten_coefficient(photon_energy) * self.thickness_cm) - abs_len = self._absorption_length(photon_energy, units=units) + abs_len = self._absorption_length(photon_energy, egu=egu) return np.exp(-self.thickness / abs_len) - def attenuation(self, photon_energy: float, units: str = "KeV"): + def attenuation(self, photon_energy: float, egu: str = "KeV"): """Attenuation is the fraction of the beam removed""" - return 1 - self.transmission(photon_energy, units=units) + return 1 - self.transmission(photon_energy, egu=egu) - def _absorption_length(self, photon_energy: float, units: str = "KeV") -> float: + def _absorption_length(self, photon_energy: float, egu: str = "KeV") -> float: """ Calculates L, the characteristic absorption length of this material, at this beam energy. @@ -113,9 +95,9 @@ def _absorption_length(self, photon_energy: float, units: str = "KeV") -> float: photon energy in KeV or eV absorption length in microns """ - if units == "KeV": + if egu == "KeV": photon_energy = photon_energy * 1e3 - elif units != "eV": + elif egu != "eV": msg = "Photon energy units must be eV or KeV" raise RuntimeError(msg) return self.filter_material.absorption_length(photon_energy) # type: ignore[reportArgumentType] @@ -128,9 +110,10 @@ class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): prefix = "XF:09ID1-ES{IOLOGIK1:E1212}" thicknesses = THICKNESSES - available_attenuations = AVAILABLE_ATTENUATIONS - def __init__(self): + def __init__(self, energy: Energy): + self.energy = energy + with self.add_children_as_readables(): self.attenuators = DeviceVector( { @@ -140,6 +123,14 @@ def __init__(self): ) super().__init__(prefix=self.prefix) + @property + def photon_energy(self): + return self.energy.energy.readback.get() + + @property + def egu(self): + return self.energy.egu + async def get_status(self): return await asyncio.gather( *(a.position.get_value() for _, a in self.attenuators.items()) @@ -167,18 +158,19 @@ def find_closest_attenuation( but that seems like overkill for our use case. The search space is small, so we start in the middle, and work up or down. """ - best_idx = len(self.available_attenuations) // 2 - atten = self.available_attenuations[best_idx].transmission + available_attenuations = self._calculate_available_attentuations() + best_idx = len(available_attenuations) // 2 + atten = available_attenuations[best_idx].transmission diff = float("inf") new_diff = abs(target_attenuation - atten) inc = 1 if target_attenuation > atten else -1 while new_diff < diff: diff = new_diff - # break if we are about to check oustide the list - if best_idx + inc >= len(self.available_attenuations) or best_idx + inc < 0: + # break if we are about to check outside the list + if best_idx + inc >= len(available_attenuations) or best_idx + inc < 0: break - atten = self.available_attenuations[best_idx + inc].transmission + atten = available_attenuations[best_idx + inc].transmission new_diff = abs(target_attenuation - atten) if new_diff < diff: best_idx += inc @@ -186,11 +178,9 @@ def find_closest_attenuation( break # TODO - should return just the found attentuation? or also the # requested attenuation and/or the difference? - return self.available_attenuations[best_idx] + return available_attenuations[best_idx] - def _calculate_available_attentuations( - self, photon_energy: float, units: str = "KeV" - ) -> list[AttenuatorCombination]: + def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: """ It is more efficient to precompute all possible total attenuations, and simply look up the closest one. @@ -198,9 +188,7 @@ def _calculate_available_attentuations( available_attenuations = [] for combination in self._powerset(): attens = [self.attenuators[a] for a in self.attenuators if a in combination] - total_atten = self._calculate_total_attenuation( - *attens, photon_energy=photon_energy, units=units - ) + total_atten = self._calculate_total_attenuation(*attens) available_attenuations.append( AttenuatorCombination(total_atten, combination) ) @@ -208,17 +196,11 @@ def _calculate_available_attentuations( available_attenuations.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] return available_attenuations - def _calculate_total_attenuation( - self, *attenuators: Attenuator, photon_energy: float, units: str = "KeV" - ) -> float: - return round( - float( - math.prod( - [a.transmission(photon_energy, units=units) for a in attenuators] - ) - ), - 3, - ) + def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float: + transmissions = [ + a.transmission(self.photon_energy, self.egu) for a in attenuators + ] + return round(float(math.prod(transmissions)), 3) def _powerset(self) -> list[list[int]]: """ @@ -232,12 +214,3 @@ def _powerset(self) -> list[list[int]]: combination.append(j + 1) # +1 because attenuators are 1-indexed powerset.append(combination) return powerset - - -""" -from cditools.attenuator import AttenuatorBank - -bank = AttenuatorBank() -atten = bank.attenuators[1] - -""" diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 0b5cef8..13e3218 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -1,25 +1,51 @@ from __future__ import annotations +from unittest.mock import MagicMock + import pytest import pytest_asyncio from ophyd_async.core import get_mock_put, init_devices, set_mock_value from ophyd_async.testing import assert_value from cditools.attenuator import ( - AVAILABLE_ATTENUATIONS, Attenuator, AttenuatorBank, + AttenuatorCombination, AttenuatorStatusEnum, ) +from cditools.motors import Energy pytest_plugins = ("pytest_asyncio",) photon_energy = 8.6 # KeV +# These are the attenuations at photon_energy = 8.6 KeV +TEST_ATTENUATIONS = [ + AttenuatorCombination(transmission=0.084, attenuators=[1, 2, 3, 4]), + AttenuatorCombination(transmission=0.1, attenuators=[2, 3, 4]), + AttenuatorCombination(transmission=0.109, attenuators=[1, 3, 4]), + AttenuatorCombination(transmission=0.129, attenuators=[3, 4]), + AttenuatorCombination(transmission=0.171, attenuators=[1, 2, 4]), + AttenuatorCombination(transmission=0.203, attenuators=[2, 4]), + AttenuatorCombination(transmission=0.222, attenuators=[1, 4]), + AttenuatorCombination(transmission=0.263, attenuators=[4]), + AttenuatorCombination(transmission=0.32, attenuators=[1, 2, 3]), + AttenuatorCombination(transmission=0.38, attenuators=[2, 3]), + AttenuatorCombination(transmission=0.414, attenuators=[1, 3]), + AttenuatorCombination(transmission=0.492, attenuators=[3]), + AttenuatorCombination(transmission=0.65, attenuators=[1, 2]), + AttenuatorCombination(transmission=0.772, attenuators=[2]), + AttenuatorCombination(transmission=0.842, attenuators=[1]), + AttenuatorCombination(transmission=1.0, attenuators=[]), +] + @pytest_asyncio.fixture async def mock_attenuator_bank(): async with init_devices(mock=True): - mock_attenuator_bank = AttenuatorBank() + mock_energy = MagicMock(spec=Energy) + mock_energy.energy.readback.get.return_value = photon_energy + mock_energy.egu = "KeV" + mock_attenuator_bank = AttenuatorBank(mock_energy) yield mock_attenuator_bank @@ -50,7 +76,7 @@ def test_transmission_kev(self, mock_attenuator: Attenuator): def test_transmission_ev(self, mock_attenuator: Attenuator): photon_energy = 8600 # eV - assert mock_attenuator.transmission(photon_energy, units="eV") == pytest.approx( + assert mock_attenuator.transmission(photon_energy, egu="eV") == pytest.approx( 0.84, abs=0.01 ) @@ -61,12 +87,25 @@ def test_attenuation_kev(self, mock_attenuator: Attenuator): def test_attenuation_ev(self, mock_attenuator: Attenuator): photon_energy = 8600 # eV - assert mock_attenuator.attenuation(photon_energy, units="eV") == pytest.approx( + assert mock_attenuator.attenuation(photon_energy, egu="eV") == pytest.approx( 0.16, abs=0.01 ) class TestAttenuatorBank: + @pytest.mark.asyncio + async def test_attenuation_bank_creation( + self, mock_attenuator_bank: AttenuatorBank + ): + assert mock_attenuator_bank.energy.energy.readback.get() == 8.6 + assert mock_attenuator_bank.photon_energy == 8.6 + + second_energy = MagicMock(spec=Energy) + second_energy.energy.readback.get.return_value = 6 + second_bank = AttenuatorBank(second_energy) + assert second_bank.energy.energy.readback.get() == 6 + assert second_bank.photon_energy == 6 + @pytest.mark.asyncio async def test_attenuators_indexed_at_1(self, mock_attenuator_bank: AttenuatorBank): with pytest.raises(KeyError): @@ -100,16 +139,14 @@ async def test_set_attenuators(self, mock_attenuator_bank: AttenuatorBank): atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].position) atten_mock4 = get_mock_put(mock_attenuator_bank.attenuators[4].position) - # AttenuatorCombination(attenuation=0.095, attenuators=[1, 2, 3]), - combo0 = AVAILABLE_ATTENUATIONS[1] # attenuators 1,2,3 + combo0 = TEST_ATTENUATIONS[1] # attenuators 2, 3, 4 await mock_attenuator_bank.set(combo0.transmission) atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) atten_mock3.assert_called_with(AttenuatorStatusEnum.HIGH) atten_mock4.assert_called_with(AttenuatorStatusEnum.HIGH) - # AttenuatorCombination(attenuation=0.768, attenuators=[1]), - combo1 = AVAILABLE_ATTENUATIONS[-3] + combo1 = TEST_ATTENUATIONS[-3] # attenuator 2 await mock_attenuator_bank.set(combo1.transmission) atten_mock1.assert_called_with(AttenuatorStatusEnum.LOW) atten_mock2.assert_called_with(AttenuatorStatusEnum.HIGH) @@ -151,10 +188,10 @@ def test_find_closest_attenuation(self, mock_attenuator_bank: AttenuatorBank): nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) assert nearest4.transmission == 1 - def test_up_to_date_available_attenuations( + def test_find_closest_attenuation_with_alt_energies( self, mock_attenuator_bank: AttenuatorBank ): - assert ( - mock_attenuator_bank._calculate_available_attentuations(photon_energy) # type: ignore[reportPrivateUsage] - == AVAILABLE_ATTENUATIONS - ) + nearest = mock_attenuator_bank.find_closest_attenuation(0.7) + assert nearest == AttenuatorCombination(transmission=0.65, attenuators=[1, 2]) + + # third_photon_energy = 5 From 0c4b80186637d2744f4ed4d4d5a7bd12af277b10 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 15 May 2026 12:50:36 -0400 Subject: [PATCH 24/33] more comprehensive status for attenuator bank --- src/cditools/attenuator.py | 38 +++++++++++++++++++--- tests/test_attenuator.py | 65 ++++++++++++++++++++++++++------------ 2 files changed, 77 insertions(+), 26 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 490a2d7..cf57cf0 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -66,6 +66,10 @@ def __init__(self, prefix: str, num: int, thickness: int): def __repr__(self): return f"{self.thickness!s} microns, {self.filter_material}" + @property + def name(self): + return f"attenuator_{self.num}" + @AsyncStatus.wrap async def set(self, value: AttenuatorStatusEnum): await self.position.set(value) @@ -79,7 +83,7 @@ async def close(self): await self.position.set(AttenuatorStatusEnum.HIGH) def transmission(self, photon_energy: float, egu: str = "KeV"): - """Transmission is the fraction of remaining beam""" + """Transmission is the fraction of beam remaining""" abs_len = self._absorption_length(photon_energy, egu=egu) return np.exp(-self.thickness / abs_len) @@ -92,8 +96,10 @@ def _absorption_length(self, photon_energy: float, egu: str = "KeV") -> float: Calculates L, the characteristic absorption length of this material, at this beam energy. - photon energy in KeV or eV - absorption length in microns + photon energy: the beam energy + egu: the engineering units of the beam energy (KeV or eV) + absorption length: the characteristic absorption length of the + filter material (microns) """ if egu == "KeV": photon_energy = photon_energy * 1e3 @@ -131,10 +137,31 @@ def photon_energy(self): def egu(self): return self.energy.egu - async def get_status(self): - return await asyncio.gather( + async def get_status(self): # type: ignore[reportUnknownParameterType] + """ + Status polls the bluesky energy object for the current beam energy, and + returns that energy, each filter position, each transmission, and + the total transmission. + """ + status = {} + active_attens = [] + en = self.photon_energy + egu = self.egu + positions = await asyncio.gather( *(a.position.get_value() for _, a in self.attenuators.items()) ) + for i, pos in zip(self.attenuators, positions): + atten = self.attenuators[i] + is_active = pos == AttenuatorStatusEnum.HIGH + if is_active: + active_attens.append(atten) + transmission = atten.transmission(en, egu) if is_active else 0 + status[atten.name] = {"active": is_active, "transmission": transmission} + status["active_attenuators"] = [a.num for a in active_attens] + status["photon_energy"] = en + status["egu"] = egu + status["total_transmission"] = self._calculate_total_attenuation(*active_attens) + return status @AsyncStatus.wrap async def set(self, value: float): @@ -150,6 +177,7 @@ async def set(self, value: float): coros.append(atten.open()) await asyncio.gather(*coros) + # TODO - fix the language here (transmission / attenuation) def find_closest_attenuation( self, target_attenuation: float ) -> AttenuatorCombination: diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 13e3218..91fe1f7 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -119,6 +119,7 @@ async def test_attenuators_indexed_at_1(self, mock_attenuator_bank: AttenuatorBa assert ( atten1.in_status.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DI1-Sts" ) + assert atten1.name == "attenuator_1" atten2 = mock_attenuator_bank.attenuators[2] assert atten2.num == 2 @@ -133,7 +134,7 @@ async def test_attenuators_indexed_at_1(self, mock_attenuator_bank: AttenuatorBa assert atten4.thickness == 124 @pytest.mark.asyncio - async def test_set_attenuators(self, mock_attenuator_bank: AttenuatorBank): + async def test_set_attenuation(self, mock_attenuator_bank: AttenuatorBank): atten_mock1 = get_mock_put(mock_attenuator_bank.attenuators[1].position) atten_mock2 = get_mock_put(mock_attenuator_bank.attenuators[2].position) atten_mock3 = get_mock_put(mock_attenuator_bank.attenuators[3].position) @@ -154,26 +155,48 @@ async def test_set_attenuators(self, mock_attenuator_bank: AttenuatorBank): atten_mock4.assert_called_with(AttenuatorStatusEnum.LOW) @pytest.mark.asyncio - async def test_get_bank_status(self, mock_attenuator_bank: AttenuatorBank): - set_mock_value( - mock_attenuator_bank.attenuators[1].position, AttenuatorStatusEnum.LOW - ) - set_mock_value( - mock_attenuator_bank.attenuators[2].position, AttenuatorStatusEnum.LOW - ) - set_mock_value( - mock_attenuator_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH - ) - set_mock_value( - mock_attenuator_bank.attenuators[4].position, AttenuatorStatusEnum.LOW - ) - - assert await mock_attenuator_bank.get_status() == [ - AttenuatorStatusEnum.LOW, - AttenuatorStatusEnum.LOW, - AttenuatorStatusEnum.HIGH, - AttenuatorStatusEnum.LOW, - ] + async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): + mock_attenuator_bank.set(1) + status = await mock_attenuator_bank.get_status() + assert status == { + "active_attenuators": [], + "photon_energy": 8.6, + "egu": "KeV", + "total_transmission": 1, + "attenuator_1": {"active": False, "transmission": 0}, + "attenuator_2": {"active": False, "transmission": 0}, + "attenuator_3": {"active": False, "transmission": 0}, + "attenuator_4": {"active": False, "transmission": 0}, + } + + # Test with different energy and attenuations + async with init_devices(mock=True): + second_energy = MagicMock(spec=Energy) + second_energy.energy.readback.get.return_value = 12 + second_energy.egu = "KeV" + second_bank = AttenuatorBank(second_energy) + set_mock_value(second_bank.attenuators[1].position, AttenuatorStatusEnum.LOW) + set_mock_value(second_bank.attenuators[2].position, AttenuatorStatusEnum.HIGH) + set_mock_value(second_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH) + set_mock_value(second_bank.attenuators[4].position, AttenuatorStatusEnum.LOW) + + status = await second_bank.get_status() + assert status == { + "active_attenuators": [2, 3], + "photon_energy": 12, + "egu": "KeV", + "total_transmission": pytest.approx(0.699), + "attenuator_1": {"active": False, "transmission": 0}, + "attenuator_2": { + "active": True, + "transmission": pytest.approx(0.909, rel=0.001), + }, + "attenuator_3": { + "active": True, + "transmission": pytest.approx(0.769, rel=0.001), + }, + "attenuator_4": {"active": False, "transmission": 0}, + } def test_find_closest_attenuation(self, mock_attenuator_bank: AttenuatorBank): nearest = mock_attenuator_bank.find_closest_attenuation(0.7) From cd07820a075d6a6341eee77f026c241bbe58f4d5 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 15 May 2026 14:07:39 -0400 Subject: [PATCH 25/33] fixed terminology so it is clear that the value set on the attenuator bank is the transmission --- src/cditools/attenuator.py | 40 ++++++++++++++++++++------------------ tests/test_attenuator.py | 10 +++++----- 2 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index cf57cf0..4e4e621 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -160,12 +160,15 @@ async def get_status(self): # type: ignore[reportUnknownParameterType] status["active_attenuators"] = [a.num for a in active_attens] status["photon_energy"] = en status["egu"] = egu - status["total_transmission"] = self._calculate_total_attenuation(*active_attens) + status["total_transmission"] = self._calculate_total_transmission( + *active_attens + ) return status @AsyncStatus.wrap async def set(self, value: float): - attenuation_combination = self.find_closest_attenuation(value) + """Set the transmission for the attenuator bank""" + attenuation_combination = self.find_closest_transmission(value) coros = [] for ( num, @@ -177,21 +180,20 @@ async def set(self, value: float): coros.append(atten.open()) await asyncio.gather(*coros) - # TODO - fix the language here (transmission / attenuation) - def find_closest_attenuation( - self, target_attenuation: float + def find_closest_transmission( + self, target_transmission: float ) -> AttenuatorCombination: """ This could be faster if we implemented binary search, but that seems like overkill for our use case. The search space is small, so we start in the middle, and work up or down. """ - available_attenuations = self._calculate_available_attentuations() + available_attenuations = self._calculate_available_transmissions() best_idx = len(available_attenuations) // 2 atten = available_attenuations[best_idx].transmission diff = float("inf") - new_diff = abs(target_attenuation - atten) - inc = 1 if target_attenuation > atten else -1 + new_diff = abs(target_transmission - atten) + inc = 1 if target_transmission > atten else -1 while new_diff < diff: diff = new_diff @@ -199,7 +201,7 @@ def find_closest_attenuation( if best_idx + inc >= len(available_attenuations) or best_idx + inc < 0: break atten = available_attenuations[best_idx + inc].transmission - new_diff = abs(target_attenuation - atten) + new_diff = abs(target_transmission - atten) if new_diff < diff: best_idx += inc else: # if diff did not change, then we have found the best option @@ -208,23 +210,23 @@ def find_closest_attenuation( # requested attenuation and/or the difference? return available_attenuations[best_idx] - def _calculate_available_attentuations(self) -> list[AttenuatorCombination]: + def _calculate_available_transmissions(self) -> list[AttenuatorCombination]: """ - It is more efficient to precompute all possible total - attenuations, and simply look up the closest one. + Calculates all possible transmissions for the attenuator bank, using + the powerset of the available attenuators. """ - available_attenuations = [] + available_transmissions = [] for combination in self._powerset(): attens = [self.attenuators[a] for a in self.attenuators if a in combination] - total_atten = self._calculate_total_attenuation(*attens) - available_attenuations.append( - AttenuatorCombination(total_atten, combination) + total_transmission = self._calculate_total_transmission(*attens) + available_transmissions.append( + AttenuatorCombination(total_transmission, combination) ) # We want the available attenuations sorted so we can efficiently search through them - available_attenuations.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] - return available_attenuations + available_transmissions.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] + return available_transmissions - def _calculate_total_attenuation(self, *attenuators: Attenuator) -> float: + def _calculate_total_transmission(self, *attenuators: Attenuator) -> float: transmissions = [ a.transmission(self.photon_energy, self.egu) for a in attenuators ] diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 91fe1f7..164c1f9 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -199,22 +199,22 @@ async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): } def test_find_closest_attenuation(self, mock_attenuator_bank: AttenuatorBank): - nearest = mock_attenuator_bank.find_closest_attenuation(0.7) + nearest = mock_attenuator_bank.find_closest_transmission(0.7) assert nearest.transmission == 0.65 - nearest2 = mock_attenuator_bank.find_closest_attenuation(0.2) + nearest2 = mock_attenuator_bank.find_closest_transmission(0.2) assert nearest2.transmission == 0.203 - nearest3 = mock_attenuator_bank.find_closest_attenuation(0.02) + nearest3 = mock_attenuator_bank.find_closest_transmission(0.02) assert nearest3.transmission == 0.084 - nearest4 = mock_attenuator_bank.find_closest_attenuation(0.98) + nearest4 = mock_attenuator_bank.find_closest_transmission(0.98) assert nearest4.transmission == 1 def test_find_closest_attenuation_with_alt_energies( self, mock_attenuator_bank: AttenuatorBank ): - nearest = mock_attenuator_bank.find_closest_attenuation(0.7) + nearest = mock_attenuator_bank.find_closest_transmission(0.7) assert nearest == AttenuatorCombination(transmission=0.65, attenuators=[1, 2]) # third_photon_energy = 5 From 1d6a948db454d6ef28b9a18402966cc239d6d223 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 15 May 2026 14:54:42 -0400 Subject: [PATCH 26/33] made prefix an argument for bank creation --- src/cditools/attenuator.py | 4 ++-- tests/test_attenuator.py | 15 +++++++-------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 4e4e621..ca84ec0 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -114,10 +114,10 @@ class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): The ioc for the iologik1 lives on xf09id1-inst-ioc1.nsls2.bnl.gov """ - prefix = "XF:09ID1-ES{IOLOGIK1:E1212}" thicknesses = THICKNESSES - def __init__(self, energy: Energy): + def __init__(self, prefix: str, energy: Energy): + self.prefix = prefix self.energy = energy with self.add_children_as_readables(): diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 164c1f9..c867a74 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -17,6 +17,7 @@ pytest_plugins = ("pytest_asyncio",) photon_energy = 8.6 # KeV +prefix = "test-prefix" # These are the attenuations at photon_energy = 8.6 KeV TEST_ATTENUATIONS = [ @@ -45,7 +46,7 @@ async def mock_attenuator_bank(): mock_energy = MagicMock(spec=Energy) mock_energy.energy.readback.get.return_value = photon_energy mock_energy.egu = "KeV" - mock_attenuator_bank = AttenuatorBank(mock_energy) + mock_attenuator_bank = AttenuatorBank(prefix, mock_energy) yield mock_attenuator_bank @@ -102,7 +103,7 @@ async def test_attenuation_bank_creation( second_energy = MagicMock(spec=Energy) second_energy.energy.readback.get.return_value = 6 - second_bank = AttenuatorBank(second_energy) + second_bank = AttenuatorBank(prefix, second_energy) assert second_bank.energy.energy.readback.get() == 6 assert second_bank.photon_energy == 6 @@ -114,11 +115,9 @@ async def test_attenuators_indexed_at_1(self, mock_attenuator_bank: AttenuatorBa atten1 = mock_attenuator_bank.attenuators[1] assert atten1.num == 1 assert atten1.thickness == 16 - assert atten1.position.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DO1-Sts" - assert atten1.mode.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DIO1-Mode" - assert ( - atten1.in_status.source == "mock+ca://XF:09ID1-ES{IOLOGIK1:E1212}:DI1-Sts" - ) + assert atten1.position.source == "mock+ca://test-prefix:DO1-Sts" + assert atten1.mode.source == "mock+ca://test-prefix:DIO1-Mode" + assert atten1.in_status.source == "mock+ca://test-prefix:DI1-Sts" assert atten1.name == "attenuator_1" atten2 = mock_attenuator_bank.attenuators[2] @@ -174,7 +173,7 @@ async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): second_energy = MagicMock(spec=Energy) second_energy.energy.readback.get.return_value = 12 second_energy.egu = "KeV" - second_bank = AttenuatorBank(second_energy) + second_bank = AttenuatorBank(prefix, second_energy) set_mock_value(second_bank.attenuators[1].position, AttenuatorStatusEnum.LOW) set_mock_value(second_bank.attenuators[2].position, AttenuatorStatusEnum.HIGH) set_mock_value(second_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH) From 09138c1941dc59592b7cba101883e0fe68360bfc Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Fri, 15 May 2026 15:41:08 -0400 Subject: [PATCH 27/33] made material for attenuators a parameter --- src/cditools/attenuator.py | 15 +++++++++------ tests/test_attenuator.py | 11 +++++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index ca84ec0..615cad3 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -37,9 +37,7 @@ class AttenuatorStatusEnum(StrictEnum): class Attenuator(EpicsDevice, AsyncMovable[AttenuatorStatusEnum]): - filter_material = xu.materials.Al - - def __init__(self, prefix: str, num: int, thickness: int): + def __init__(self, prefix: str, num: int, material: str, thickness: int): """ prefix - the common prefix for the attenuator bank num - an integer denoting which attenuator within the bank this is @@ -49,6 +47,7 @@ def __init__(self, prefix: str, num: int, thickness: int): """ self.prefix = prefix self.num = num + self.filter_material = getattr(xu.materials, material) self.thickness = thickness # microns self.position = epics_signal_rw( @@ -116,15 +115,19 @@ class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): thicknesses = THICKNESSES - def __init__(self, prefix: str, energy: Energy): + def __init__( + self, prefix: str, atten_configs: list[tuple[str, int]], energy: Energy + ): self.prefix = prefix self.energy = energy with self.add_children_as_readables(): self.attenuators = DeviceVector( { - i: Attenuator(self.prefix, i, self.thicknesses[i - 1]) - for i in range(1, len(self.thicknesses) + 1) + i: Attenuator( + self.prefix, i, atten_configs[i - 1][0], atten_configs[i - 1][1] + ) + for i in range(1, len(atten_configs) + 1) } ) super().__init__(prefix=self.prefix) diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index c867a74..9ad9dbe 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -18,6 +18,7 @@ pytest_plugins = ("pytest_asyncio",) photon_energy = 8.6 # KeV prefix = "test-prefix" +attenuator_configs = [("Al", 16), ("Al", 24), ("Al", 66), ("Al", 124)] # These are the attenuations at photon_energy = 8.6 KeV TEST_ATTENUATIONS = [ @@ -46,14 +47,14 @@ async def mock_attenuator_bank(): mock_energy = MagicMock(spec=Energy) mock_energy.energy.readback.get.return_value = photon_energy mock_energy.egu = "KeV" - mock_attenuator_bank = AttenuatorBank(prefix, mock_energy) + mock_attenuator_bank = AttenuatorBank(prefix, attenuator_configs, mock_energy) yield mock_attenuator_bank @pytest_asyncio.fixture async def mock_attenuator(mock_attenuator_bank: AttenuatorBank): async with init_devices(mock=True): - mock_attenuator = Attenuator(mock_attenuator_bank.prefix, 1, 16) + mock_attenuator = Attenuator(mock_attenuator_bank.prefix, 1, "Al", 16) yield mock_attenuator @@ -103,7 +104,7 @@ async def test_attenuation_bank_creation( second_energy = MagicMock(spec=Energy) second_energy.energy.readback.get.return_value = 6 - second_bank = AttenuatorBank(prefix, second_energy) + second_bank = AttenuatorBank(prefix, attenuator_configs, second_energy) assert second_bank.energy.energy.readback.get() == 6 assert second_bank.photon_energy == 6 @@ -173,7 +174,7 @@ async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): second_energy = MagicMock(spec=Energy) second_energy.energy.readback.get.return_value = 12 second_energy.egu = "KeV" - second_bank = AttenuatorBank(prefix, second_energy) + second_bank = AttenuatorBank(prefix, attenuator_configs, second_energy) set_mock_value(second_bank.attenuators[1].position, AttenuatorStatusEnum.LOW) set_mock_value(second_bank.attenuators[2].position, AttenuatorStatusEnum.HIGH) set_mock_value(second_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH) @@ -215,5 +216,3 @@ def test_find_closest_attenuation_with_alt_energies( ): nearest = mock_attenuator_bank.find_closest_transmission(0.7) assert nearest == AttenuatorCombination(transmission=0.65, attenuators=[1, 2]) - - # third_photon_energy = 5 From 22fa679617a77b1d65b88da292f538702918afc3 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 18 May 2026 11:40:38 -0400 Subject: [PATCH 28/33] MNT: minor cleanups added sha back to github action; properly formatted docstrings; pass in material and thickness to constructor; added pixi format task; pass energy into utility methods; remove property reading energy --- pyproject.toml | 3 ++ src/cditools/attenuator.py | 89 +++++++++++++++++++++++--------------- tests/test_attenuator.py | 18 ++++---- 3 files changed, 66 insertions(+), 44 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aecced6..7ca1b9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,6 +171,9 @@ cditools = { path = ".", editable = true } [tool.pixi.dependencies] xrayutilities = ">=1.7.12,<2" +[tool.pixi.feature.dev.tasks] +format = "ruff format ." + [tool.pixi.environments] default = { solve-group = "default" } dev = { features = ["dev"], solve-group = "default" } diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 615cad3..038940f 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -28,22 +28,30 @@ def attenuation(self): return 1 - self.transmission -THICKNESSES = (16, 24, 66, 124) # microns - - class AttenuatorStatusEnum(StrictEnum): LOW = "Low" # off / not obstructing HIGH = "High" # on / obstructing class Attenuator(EpicsDevice, AsyncMovable[AttenuatorStatusEnum]): - def __init__(self, prefix: str, num: int, material: str, thickness: int): + def __init__(self, prefix: str, num: int, material: str, thickness: float): """ - prefix - the common prefix for the attenuator bank - num - an integer denoting which attenuator within the bank this is - thickness - the thickness of the attenuator in microns - - position - the read / write PV to open and close the attenuator + Parameters + ---------- + prefix : str + The common prefix for the attenuator bank + num : int + An integer denoting which attenuator within the bank this is + thickness : float + The thickness of the attenuator in microns + + Attributes + ---------- + position : SignalRW[AttenuatorStatusEnum] + The read / write PV to open and close the attenuator and get + the current state of the attenuator + mode : SignalRW[bool] + in_status : SignalR[AttenuatorStatusEnum] """ self.prefix = prefix self.num = num @@ -95,10 +103,17 @@ def _absorption_length(self, photon_energy: float, egu: str = "KeV") -> float: Calculates L, the characteristic absorption length of this material, at this beam energy. - photon energy: the beam energy - egu: the engineering units of the beam energy (KeV or eV) - absorption length: the characteristic absorption length of the - filter material (microns) + Parameters + ---------- + photon energy : float + The beam energy + egu : {'KeV', 'eV'} + The engineering units of the beam energy + + Returns + ------- + float + The characteristic absorption length of the filter material (microns) """ if egu == "KeV": photon_energy = photon_energy * 1e3 @@ -113,10 +128,8 @@ class AttenuatorBank(StandardReadable, EpicsDevice, AsyncMovable[float]): The ioc for the iologik1 lives on xf09id1-inst-ioc1.nsls2.bnl.gov """ - thicknesses = THICKNESSES - def __init__( - self, prefix: str, atten_configs: list[tuple[str, int]], energy: Energy + self, prefix: str, atten_configs: list[tuple[str, float]], energy: Energy ): self.prefix = prefix self.energy = energy @@ -124,18 +137,17 @@ def __init__( with self.add_children_as_readables(): self.attenuators = DeviceVector( { - i: Attenuator( - self.prefix, i, atten_configs[i - 1][0], atten_configs[i - 1][1] - ) - for i in range(1, len(atten_configs) + 1) + i: Attenuator(self.prefix, i, material, thickness) + for i, (material, thickness) in enumerate(atten_configs, start=1) } ) super().__init__(prefix=self.prefix) - @property - def photon_energy(self): - return self.energy.energy.readback.get() + # @property + # def photon_energy(self): + # return self.energy.energy.readback.get() + # TODO - make this not a property @property def egu(self): return self.energy.egu @@ -148,7 +160,7 @@ async def get_status(self): # type: ignore[reportUnknownParameterType] """ status = {} active_attens = [] - en = self.photon_energy + energy = self.energy.energy.readback.get() egu = self.egu positions = await asyncio.gather( *(a.position.get_value() for _, a in self.attenuators.items()) @@ -158,20 +170,21 @@ async def get_status(self): # type: ignore[reportUnknownParameterType] is_active = pos == AttenuatorStatusEnum.HIGH if is_active: active_attens.append(atten) - transmission = atten.transmission(en, egu) if is_active else 0 + transmission = atten.transmission(energy, egu) if is_active else 0 status[atten.name] = {"active": is_active, "transmission": transmission} status["active_attenuators"] = [a.num for a in active_attens] - status["photon_energy"] = en + status["photon_energy"] = energy status["egu"] = egu status["total_transmission"] = self._calculate_total_transmission( - *active_attens + energy, *active_attens ) return status @AsyncStatus.wrap async def set(self, value: float): """Set the transmission for the attenuator bank""" - attenuation_combination = self.find_closest_transmission(value) + photon_energy = self.energy.energy.readback.get() + attenuation_combination = self.find_closest_transmission(photon_energy, value) coros = [] for ( num, @@ -184,14 +197,14 @@ async def set(self, value: float): await asyncio.gather(*coros) def find_closest_transmission( - self, target_transmission: float + self, photon_energy: float, target_transmission: float ) -> AttenuatorCombination: """ This could be faster if we implemented binary search, but that seems like overkill for our use case. The search space is small, so we start in the middle, and work up or down. """ - available_attenuations = self._calculate_available_transmissions() + available_attenuations = self._calculate_available_transmissions(photon_energy) best_idx = len(available_attenuations) // 2 atten = available_attenuations[best_idx].transmission diff = float("inf") @@ -213,7 +226,9 @@ def find_closest_transmission( # requested attenuation and/or the difference? return available_attenuations[best_idx] - def _calculate_available_transmissions(self) -> list[AttenuatorCombination]: + def _calculate_available_transmissions( + self, photon_energy: float + ) -> list[AttenuatorCombination]: """ Calculates all possible transmissions for the attenuator bank, using the powerset of the available attenuators. @@ -221,7 +236,9 @@ def _calculate_available_transmissions(self) -> list[AttenuatorCombination]: available_transmissions = [] for combination in self._powerset(): attens = [self.attenuators[a] for a in self.attenuators if a in combination] - total_transmission = self._calculate_total_transmission(*attens) + total_transmission = self._calculate_total_transmission( + photon_energy, *attens + ) available_transmissions.append( AttenuatorCombination(total_transmission, combination) ) @@ -229,10 +246,10 @@ def _calculate_available_transmissions(self) -> list[AttenuatorCombination]: available_transmissions.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] return available_transmissions - def _calculate_total_transmission(self, *attenuators: Attenuator) -> float: - transmissions = [ - a.transmission(self.photon_energy, self.egu) for a in attenuators - ] + def _calculate_total_transmission( + self, photon_energy: float, *attenuators: Attenuator + ) -> float: + transmissions = [a.transmission(photon_energy, self.egu) for a in attenuators] return round(float(math.prod(transmissions)), 3) def _powerset(self) -> list[list[int]]: diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 9ad9dbe..85e0257 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -18,7 +18,7 @@ pytest_plugins = ("pytest_asyncio",) photon_energy = 8.6 # KeV prefix = "test-prefix" -attenuator_configs = [("Al", 16), ("Al", 24), ("Al", 66), ("Al", 124)] +attenuator_configs = [("Al", 16.0), ("Al", 24.0), ("Al", 66.0), ("Al", 124.0)] # These are the attenuations at photon_energy = 8.6 KeV TEST_ATTENUATIONS = [ @@ -100,13 +100,13 @@ async def test_attenuation_bank_creation( self, mock_attenuator_bank: AttenuatorBank ): assert mock_attenuator_bank.energy.energy.readback.get() == 8.6 - assert mock_attenuator_bank.photon_energy == 8.6 + # assert mock_attenuator_bank.photon_energy == 8.6 second_energy = MagicMock(spec=Energy) second_energy.energy.readback.get.return_value = 6 second_bank = AttenuatorBank(prefix, attenuator_configs, second_energy) assert second_bank.energy.energy.readback.get() == 6 - assert second_bank.photon_energy == 6 + # assert second_bank.photon_energy == 6 @pytest.mark.asyncio async def test_attenuators_indexed_at_1(self, mock_attenuator_bank: AttenuatorBank): @@ -199,20 +199,22 @@ async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): } def test_find_closest_attenuation(self, mock_attenuator_bank: AttenuatorBank): - nearest = mock_attenuator_bank.find_closest_transmission(0.7) + en = mock_attenuator_bank.energy.energy.readback.get() + nearest = mock_attenuator_bank.find_closest_transmission(en, 0.7) assert nearest.transmission == 0.65 - nearest2 = mock_attenuator_bank.find_closest_transmission(0.2) + nearest2 = mock_attenuator_bank.find_closest_transmission(en, 0.2) assert nearest2.transmission == 0.203 - nearest3 = mock_attenuator_bank.find_closest_transmission(0.02) + nearest3 = mock_attenuator_bank.find_closest_transmission(en, 0.02) assert nearest3.transmission == 0.084 - nearest4 = mock_attenuator_bank.find_closest_transmission(0.98) + nearest4 = mock_attenuator_bank.find_closest_transmission(en, 0.98) assert nearest4.transmission == 1 def test_find_closest_attenuation_with_alt_energies( self, mock_attenuator_bank: AttenuatorBank ): - nearest = mock_attenuator_bank.find_closest_transmission(0.7) + en = mock_attenuator_bank.energy.energy.readback.get() + nearest = mock_attenuator_bank.find_closest_transmission(en, 0.7) assert nearest == AttenuatorCombination(transmission=0.65, attenuators=[1, 2]) From fec48280ad3d71b3eb3620b3e234e925edf5d724 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 18 May 2026 13:32:55 -0400 Subject: [PATCH 29/33] PRF: minor speed up removed sorting of available transmissions because we can find the best one in linear time; also it is easier to read --- src/cditools/attenuator.py | 45 ++++++++++---------------------------- tests/test_attenuator.py | 6 ++--- 2 files changed, 15 insertions(+), 36 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 038940f..681c4cd 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -89,15 +89,15 @@ async def close(self): """Closed means obstructing the beam""" await self.position.set(AttenuatorStatusEnum.HIGH) + def attenuation(self, photon_energy: float, egu: str = "KeV"): + """Attenuation is the fraction of the beam removed""" + return 1 - self.transmission(photon_energy, egu=egu) + def transmission(self, photon_energy: float, egu: str = "KeV"): """Transmission is the fraction of beam remaining""" abs_len = self._absorption_length(photon_energy, egu=egu) return np.exp(-self.thickness / abs_len) - def attenuation(self, photon_energy: float, egu: str = "KeV"): - """Attenuation is the fraction of the beam removed""" - return 1 - self.transmission(photon_energy, egu=egu) - def _absorption_length(self, photon_energy: float, egu: str = "KeV") -> float: """ Calculates L, the characteristic absorption length of this material, @@ -152,9 +152,9 @@ def __init__( def egu(self): return self.energy.egu - async def get_status(self): # type: ignore[reportUnknownParameterType] + async def read(self): # type: ignore[reportUnknownParameterType] """ - Status polls the bluesky energy object for the current beam energy, and + Polls the bluesky energy object for the current beam energy, and returns that energy, each filter position, each transmission, and the total transmission. """ @@ -199,31 +199,10 @@ async def set(self, value: float): def find_closest_transmission( self, photon_energy: float, target_transmission: float ) -> AttenuatorCombination: - """ - This could be faster if we implemented binary search, - but that seems like overkill for our use case. The search space - is small, so we start in the middle, and work up or down. - """ available_attenuations = self._calculate_available_transmissions(photon_energy) - best_idx = len(available_attenuations) // 2 - atten = available_attenuations[best_idx].transmission - diff = float("inf") - new_diff = abs(target_transmission - atten) - inc = 1 if target_transmission > atten else -1 - - while new_diff < diff: - diff = new_diff - # break if we are about to check outside the list - if best_idx + inc >= len(available_attenuations) or best_idx + inc < 0: - break - atten = available_attenuations[best_idx + inc].transmission - new_diff = abs(target_transmission - atten) - if new_diff < diff: - best_idx += inc - else: # if diff did not change, then we have found the best option - break - # TODO - should return just the found attentuation? or also the - # requested attenuation and/or the difference? + best_idx = np.argmin( + [abs(target_transmission - _.transmission) for _ in available_attenuations] + ) return available_attenuations[best_idx] def _calculate_available_transmissions( @@ -231,7 +210,9 @@ def _calculate_available_transmissions( ) -> list[AttenuatorCombination]: """ Calculates all possible transmissions for the attenuator bank, using - the powerset of the available attenuators. + the powerset of the available attenuators. The result is not sorted, + as it is more efficient to scan linearly each time for the closest + match. """ available_transmissions = [] for combination in self._powerset(): @@ -242,8 +223,6 @@ def _calculate_available_transmissions( available_transmissions.append( AttenuatorCombination(total_transmission, combination) ) - # We want the available attenuations sorted so we can efficiently search through them - available_transmissions.sort(key=lambda a: a.transmission) # type: ignore[attr-defined] return available_transmissions def _calculate_total_transmission( diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index 85e0257..efbe304 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -155,9 +155,9 @@ async def test_set_attenuation(self, mock_attenuator_bank: AttenuatorBank): atten_mock4.assert_called_with(AttenuatorStatusEnum.LOW) @pytest.mark.asyncio - async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): + async def test_read(self, mock_attenuator_bank: AttenuatorBank): mock_attenuator_bank.set(1) - status = await mock_attenuator_bank.get_status() + status = await mock_attenuator_bank.read() assert status == { "active_attenuators": [], "photon_energy": 8.6, @@ -180,7 +180,7 @@ async def test_get_status(self, mock_attenuator_bank: AttenuatorBank): set_mock_value(second_bank.attenuators[3].position, AttenuatorStatusEnum.HIGH) set_mock_value(second_bank.attenuators[4].position, AttenuatorStatusEnum.LOW) - status = await second_bank.get_status() + status = await second_bank.read() assert status == { "active_attenuators": [2, 3], "photon_energy": 12, From 80917caef49c699a9883f1a6c36c3b411b6c8e21 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Mon, 18 May 2026 13:39:37 -0400 Subject: [PATCH 30/33] made photon_energy and egu get methods instead of properties --- src/cditools/attenuator.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 681c4cd..6ca0e14 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -143,13 +143,10 @@ def __init__( ) super().__init__(prefix=self.prefix) - # @property - # def photon_energy(self): - # return self.energy.energy.readback.get() + def get_photon_energy(self): + return self.energy.energy.readback.get() - # TODO - make this not a property - @property - def egu(self): + def get_egu(self): return self.energy.egu async def read(self): # type: ignore[reportUnknownParameterType] @@ -160,8 +157,8 @@ async def read(self): # type: ignore[reportUnknownParameterType] """ status = {} active_attens = [] - energy = self.energy.energy.readback.get() - egu = self.egu + energy = self.get_photon_energy() + egu = self.get_egu() positions = await asyncio.gather( *(a.position.get_value() for _, a in self.attenuators.items()) ) @@ -183,8 +180,9 @@ async def read(self): # type: ignore[reportUnknownParameterType] @AsyncStatus.wrap async def set(self, value: float): """Set the transmission for the attenuator bank""" - photon_energy = self.energy.energy.readback.get() - attenuation_combination = self.find_closest_transmission(photon_energy, value) + attenuation_combination = self.find_closest_transmission( + self.get_photon_energy(), value + ) coros = [] for ( num, @@ -228,7 +226,9 @@ def _calculate_available_transmissions( def _calculate_total_transmission( self, photon_energy: float, *attenuators: Attenuator ) -> float: - transmissions = [a.transmission(photon_energy, self.egu) for a in attenuators] + transmissions = [ + a.transmission(photon_energy, self.get_egu()) for a in attenuators + ] return round(float(math.prod(transmissions)), 3) def _powerset(self) -> list[list[int]]: From 6be318224d0f4cab7c250bfb34e20fca1cb00be9 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 18 Jun 2026 10:04:52 -0400 Subject: [PATCH 31/33] Update src/cditools/attenuator.py Co-authored-by: Thomas A Caswell --- src/cditools/attenuator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 6ca0e14..011b073 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -118,7 +118,7 @@ def _absorption_length(self, photon_energy: float, egu: str = "KeV") -> float: if egu == "KeV": photon_energy = photon_energy * 1e3 elif egu != "eV": - msg = "Photon energy units must be eV or KeV" + msg = f"Photon energy units must be eV or KeV (not {egu=})" raise RuntimeError(msg) return self.filter_material.absorption_length(photon_energy) # type: ignore[reportArgumentType] From 525f87f12c44b69b280c95e43e88e1d1b650fc1d Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 18 Jun 2026 13:54:11 -0400 Subject: [PATCH 32/33] remove active_attenuators from status --- src/cditools/attenuator.py | 1 - tests/test_attenuator.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index 011b073..ed958af 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -169,7 +169,6 @@ async def read(self): # type: ignore[reportUnknownParameterType] active_attens.append(atten) transmission = atten.transmission(energy, egu) if is_active else 0 status[atten.name] = {"active": is_active, "transmission": transmission} - status["active_attenuators"] = [a.num for a in active_attens] status["photon_energy"] = energy status["egu"] = egu status["total_transmission"] = self._calculate_total_transmission( diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index efbe304..fea6f93 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -159,7 +159,6 @@ async def test_read(self, mock_attenuator_bank: AttenuatorBank): mock_attenuator_bank.set(1) status = await mock_attenuator_bank.read() assert status == { - "active_attenuators": [], "photon_energy": 8.6, "egu": "KeV", "total_transmission": 1, @@ -182,7 +181,6 @@ async def test_read(self, mock_attenuator_bank: AttenuatorBank): status = await second_bank.read() assert status == { - "active_attenuators": [2, 3], "photon_energy": 12, "egu": "KeV", "total_transmission": pytest.approx(0.699), From 4385f2caa349d205c7985f291abe70eb1bc6ccd2 Mon Sep 17 00:00:00 2001 From: Dan Henriksen Date: Thu, 18 Jun 2026 15:59:45 -0400 Subject: [PATCH 33/33] added describe method to match read --- src/cditools/attenuator.py | 38 +++++++++++++++++++++++++++++++- tests/test_attenuator.py | 44 +++++++++++++++++++++++++++++++++++--- 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/src/cditools/attenuator.py b/src/cditools/attenuator.py index ed958af..f466ae7 100644 --- a/src/cditools/attenuator.py +++ b/src/cditools/attenuator.py @@ -2,10 +2,12 @@ import asyncio import math +from collections import OrderedDict from dataclasses import dataclass import numpy as np import xrayutilities as xu +from event_model import DataKey # type: ignore[import-untyped] from ophyd_async.core import ( AsyncMovable, AsyncStatus, @@ -155,7 +157,7 @@ async def read(self): # type: ignore[reportUnknownParameterType] returns that energy, each filter position, each transmission, and the total transmission. """ - status = {} + status = OrderedDict() active_attens = [] energy = self.get_photon_energy() egu = self.get_egu() @@ -176,6 +178,40 @@ async def read(self): # type: ignore[reportUnknownParameterType] ) return status + async def describe(self) -> OrderedDict[str, DataKey]: + """Describe the structure of values returned by read().""" + + description = OrderedDict() + + for atten in self.attenuators.values(): + description[atten.name] = DataKey( + source=atten.position.source, + dtype="string", + shape=[], + ) + energy_source = getattr( + self.energy.energy.readback, + "source", + f"ca://{self.prefix}:photon_energy", + ) + description["photon_energy"] = DataKey( + source=energy_source, + dtype="number", + shape=[], + ) + description["egu"] = DataKey( + source=f"ca://{self.prefix}:egu", + dtype="string", + shape=[], + ) + description["total_transmission"] = DataKey( + source=f"ca://{self.prefix}:total_transmission", + dtype="number", + shape=[], + ) + + return description + @AsyncStatus.wrap async def set(self, value: float): """Set the transmission for the attenuator bank""" diff --git a/tests/test_attenuator.py b/tests/test_attenuator.py index fea6f93..63ea0ca 100644 --- a/tests/test_attenuator.py +++ b/tests/test_attenuator.py @@ -181,9 +181,6 @@ async def test_read(self, mock_attenuator_bank: AttenuatorBank): status = await second_bank.read() assert status == { - "photon_energy": 12, - "egu": "KeV", - "total_transmission": pytest.approx(0.699), "attenuator_1": {"active": False, "transmission": 0}, "attenuator_2": { "active": True, @@ -194,6 +191,47 @@ async def test_read(self, mock_attenuator_bank: AttenuatorBank): "transmission": pytest.approx(0.769, rel=0.001), }, "attenuator_4": {"active": False, "transmission": 0}, + "photon_energy": 12, + "egu": "KeV", + "total_transmission": pytest.approx(0.699), + } + + @pytest.mark.asyncio + async def test_describe(self, mock_attenuator_bank: AttenuatorBank): + description = await mock_attenuator_bank.describe() + + expected_keys = { + "attenuator_1", + "attenuator_2", + "attenuator_3", + "attenuator_4", + "photon_energy", + "egu", + "total_transmission", + } + assert set(description.keys()) == expected_keys + + for i in range(1, 5): + assert description[f"attenuator_{i}"] == { + "source": mock_attenuator_bank.attenuators[i].position.source, + "dtype": "string", + "shape": [], + } + + assert description["photon_energy"] == { + "source": mock_attenuator_bank.energy.energy.readback.source, + "dtype": "number", + "shape": [], + } + assert description["egu"] == { + "source": f"ca://{mock_attenuator_bank.prefix}:egu", + "dtype": "string", + "shape": [], + } + assert description["total_transmission"] == { + "source": f"ca://{mock_attenuator_bank.prefix}:total_transmission", + "dtype": "number", + "shape": [], } def test_find_closest_attenuation(self, mock_attenuator_bank: AttenuatorBank):