Make attenuator with enum filter values

src/dodal/beamlines/i02_1.py

@@ -0,0 +1,38 @@
+"""Beamline i02-1 is also known as VMXm, or I02J"""
+
+from dodal.common.beamlines.beamline_utils import (
+    device_factory,
+)
+from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
+from dodal.devices.attenuator.attenuator import EnumFilterAttenuator
+from dodal.devices.attenuator.filter_selections import (
+    I02_1FilterFourSelections,
+    I02_1FilterOneSelections,
+    I02_1FilterThreeSelections,
+    I02_1FilterTwoSelections,
+)
+from dodal.log import set_beamline as set_log_beamline
+from dodal.utils import BeamlinePrefix, get_beamline_name
+
+BL = get_beamline_name("i02-1")
+PREFIX = BeamlinePrefix(BL)
+set_log_beamline(BL)
+set_utils_beamline(BL)
+
+
+@device_factory()
+def attenuator() -> EnumFilterAttenuator:
+    """Get the i02-1 attenuator device, instantiate it if it hasn't already been.
+    If this is called when already instantiated in i02-1, it will return the existing object.
+    """
+
+    return EnumFilterAttenuator(
+        num_filters=4,
+        filter_selection=[
+            I02_1FilterOneSelections,
+            I02_1FilterTwoSelections,
+            I02_1FilterThreeSelections,
+            I02_1FilterFourSelections,
+        ],
+        prefix="BL02J-OP-ATTN-01:",
+    )

src/dodal/beamlines/i03.py

@@ -13,7 +13,7 @@
     ApertureScatterguard,
     load_positions_from_beamline_parameters,
 )
-from dodal.devices.attenuator import Attenuator
+from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.backlight import Backlight
 from dodal.devices.cryostream import CryoStream
 from dodal.devices.dcm import DCM
@@ -80,12 +80,12 @@ def aperture_scatterguard(
 
 def attenuator(
     wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
-) -> Attenuator:
+) -> BinaryFilterAttenuator:
     """Get the i03 attenuator device, instantiate it if it hasn't already been.
     If this is called when already instantiated in i03, it will return the existing object.
     """
     return device_instantiation(
-        Attenuator,
+        BinaryFilterAttenuator,
         "attenuator",
         "-EA-ATTN-01:",
         wait_for_connection,

src/dodal/beamlines/i04.py

@@ -6,7 +6,7 @@
     ApertureScatterguard,
     load_positions_from_beamline_parameters,
 )
-from dodal.devices.attenuator import Attenuator
+from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.backlight import Backlight
 from dodal.devices.dcm import DCM
 from dodal.devices.detector import DetectorParams
@@ -138,12 +138,12 @@ def sample_shutter(
 
 def attenuator(
     wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
-) -> Attenuator:
+) -> BinaryFilterAttenuator:
     """Get the i04 attenuator device, instantiate it if it hasn't already been.
     If this is called when already instantiated in i04, it will return the existing object.
     """
     return device_instantiation(
-        Attenuator,
+        BinaryFilterAttenuator,
         "attenuator",
         "-EA-ATTN-01:",
         wait_for_connection,

src/dodal/beamlines/i24.py

@@ -1,6 +1,6 @@
 from dodal.common.beamlines.beamline_utils import BL, device_instantiation
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
-from dodal.devices.attenuator import ReadOnlyAttenuator
+from dodal.devices.attenuator.attenuator import ReadOnlyAttenuator
 from dodal.devices.detector import DetectorParams
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.hutch_shutter import HutchShutter

src/dodal/beamlines/p99.py

@@ -1,6 +1,8 @@
 from dodal.common.beamlines.beamline_utils import device_factory, set_beamline
+from dodal.devices.attenuator.filter import FilterMotor
+from dodal.devices.attenuator.filter_selections import p99StageSelections
 from dodal.devices.motors import XYZPositioner
-from dodal.devices.p99.sample_stage import FilterMotor, SampleAngleStage
+from dodal.devices.p99.sample_stage import SampleAngleStage
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
@@ -17,7 +19,9 @@ def angle_stage() -> SampleAngleStage:
 
 @device_factory()
 def filter() -> FilterMotor:
-    return FilterMotor(f"{PREFIX.beamline_prefix}-MO-STAGE-02:MP:SELECT")
+    return FilterMotor(
+        p99StageSelections, f"{PREFIX.beamline_prefix}-MO-STAGE-02:MP:SELECT"
+    )
 
 
 @device_factory()

src/dodal/devices/attenuator.py → src/dodal/devices/attenuator/attenuator.py

@@ -11,6 +11,8 @@
 )
 from ophyd_async.epics.core import epics_signal_r, epics_signal_rw, epics_signal_x
 
+from dodal.devices.attenuator.filter import FilterMotor
+from dodal.devices.attenuator.filter_selections import FilterSelection
 from dodal.log import LOGGER
 
 
@@ -27,8 +29,9 @@ def __init__(self, prefix: str, name: str = "") -> None:
         super().__init__(name)
 
 
-class Attenuator(ReadOnlyAttenuator, Movable):
+class BinaryFilterAttenuator(ReadOnlyAttenuator, Movable):
     """The attenuator will insert filters into the beam to reduce its transmission.
+    In this attenuator, each filter can be in one of two states: IN or OUT
 
     This device should be set with:
         yield from bps.set(attenuator, desired_transmission)
@@ -83,3 +86,29 @@ async def set(self, value: float):
                 for i in range(16)
             ]
         )
+
+
+class EnumFilterAttenuator(ReadOnlyAttenuator):
+    """The attenuator will insert filters into the beam to reduce its transmission.
+
+    This device is currently working, but feature incomplete. See https://github.com/DiamondLightSource/dodal/issues/972
+
+    In this attenuator, the state of a filter corresponds to the selected material,
+    e.g Ag50, in contrast to either being 'IN' or 'OUT'; see BinaryFilterAttenuator.
+    """
+
+    def __init__(
+        self,
+        num_filters: int,
+        filter_selection: list[type[FilterSelection]],
+        prefix: str,
+        name: str = "",
+    ):
+        with self.add_children_as_readables():
+            self.filters: DeviceVector[FilterMotor] = DeviceVector(
+                {
+                    i: FilterMotor(filter_selection[i], f"{prefix}MP{i+1}:", name)
+                    for i in range(num_filters)
+                }
+            )
+        super().__init__(name=name, prefix=prefix)

src/dodal/devices/attenuator/filter.py

@@ -0,0 +1,13 @@
+from ophyd_async.core import StandardReadable
+from ophyd_async.epics.core import epics_signal_rw
+
+from dodal.devices.attenuator.filter_selections import FilterSelection
+
+
+class FilterMotor(StandardReadable):
+    def __init__(
+        self, filter_selections: type[FilterSelection], prefix: str, name: str = ""
+    ):
+        with self.add_children_as_readables():
+            self.user_setpoint = epics_signal_rw(filter_selections, f"{prefix}SELECT")
+        super().__init__(name=name)

src/dodal/devices/attenuator/filter_selections.py

@@ -0,0 +1,75 @@
+from ophyd_async.core import SubsetEnum
+
+
+class FilterSelection(SubsetEnum): ...
+
+
+class p99StageSelections(FilterSelection):
+    EMPTY = "Empty"
+    MN5UM = "Mn 5um"
+    FE = "Fe (empty)"
+    CO5UM = "Co 5um"
+    NI5UM = "Ni 5um"
+    CU5UM = "Cu 5um"
+    ZN5UM = "Zn 5um"
+    ZR = "Zr (empty)"
+    MO = "Mo (empty)"
+    RH = "Rh (empty)"
+    PD = "Pd (empty)"
+    AG = "Ag (empty)"
+    CD25UM = "Cd 25um"
+    W = "W (empty)"
+    PT = "Pt (empty)"
+    USER = "User"
+
+
+class I02_1FilterOneSelections(FilterSelection):
+    EMPTY = "Empty"
+    AL8 = "Al8"
+    AL15 = "Al15"
+    AL25 = "Al25"
+    AL1000 = "Al1000"
+    TI50 = "Ti50"
+    TI100 = "Ti100"
+    TI200 = "Ti200"
+    TI400 = "Ti400"
+    TWO_TIMES_TI500 = "2xTi500"
+
+
+class I02_1FilterTwoSelections(FilterSelection):
+    EMPTY = "Empty"
+    AL50 = "Al50"
+    AL100 = "Al100"
+    AL125 = "Al125"
+    AL250 = "Al250"
+    AL500 = "Al500"
+    AL1000 = "Al1000"
+    TI50 = "Ti50"
+    TI100 = "Ti100"
+    TWO_TIMES_TI500 = "2xTi500"
+
+
+class I02_1FilterThreeSelections(FilterSelection):
+    EMPTY = "Empty"
+    AL15 = "Al15"
+    AL25 = "Al25"
+    AL50 = "Al50"
+    AL100 = "Al100"
+    AL250 = "Al250"
+    AL1000 = "Al1000"
+    TI50 = "Ti50"
+    TI100 = "Ti100"
+    TI200 = "Ti200"
+
+
+class I02_1FilterFourSelections(FilterSelection):
+    EMPTY = "Empty"
+    AL15 = "Al15"
+    AL25 = "Al25"
+    AL50 = "Al50"
+    AL100 = "Al100"
+    AL250 = "Al250"
+    AL500 = "Al500"
+    TI300 = "Ti300"
+    TI400 = "Ti400"
+    TI500 = "Ti500"

src/dodal/devices/p99/sample_stage.py

@@ -1,5 +1,5 @@
-from ophyd_async.core import StandardReadable, SubsetEnum
-from ophyd_async.epics.core import epics_signal_rw, epics_signal_rw_rbv
+from ophyd_async.core import StandardReadable
+from ophyd_async.epics.core import epics_signal_rw_rbv
 
 
 class SampleAngleStage(StandardReadable):
@@ -9,29 +9,3 @@ def __init__(self, prefix: str, name: str = ""):
             self.roll = epics_signal_rw_rbv(float, prefix + "WRITEROLL", ":RBV")
             self.pitch = epics_signal_rw_rbv(float, prefix + "WRITEPITCH", ":RBV")
         super().__init__(name=name)
-
-
-class p99StageSelections(SubsetEnum):
-    EMPTY = "Empty"
-    MN5UM = "Mn 5um"
-    FE = "Fe (empty)"
-    CO5UM = "Co 5um"
-    NI5UM = "Ni 5um"
-    CU5UM = "Cu 5um"
-    ZN5UM = "Zn 5um"
-    ZR = "Zr (empty)"
-    MO = "Mo (empty)"
-    RH = "Rh (empty)"
-    PD = "Pd (empty)"
-    AG = "Ag (empty)"
-    CD25UM = "Cd 25um"
-    W = "W (empty)"
-    PT = "Pt (empty)"
-    USER = "User"
-
-
-class FilterMotor(StandardReadable):
-    def __init__(self, prefix: str, name: str = ""):
-        with self.add_children_as_readables():
-            self.user_setpoint = epics_signal_rw(p99StageSelections, prefix)
-        super().__init__(name=name)

tests/devices/unit_tests/p99/test_p99_stage.py

@@ -2,10 +2,10 @@
 from ophyd_async.core import DeviceCollector
 from ophyd_async.testing import set_mock_value
 
+from dodal.devices.attenuator.filter import FilterMotor
+from dodal.devices.attenuator.filter_selections import p99StageSelections
 from dodal.devices.p99.sample_stage import (
-    FilterMotor,
     SampleAngleStage,
-    p99StageSelections,
 )
 
 # Long enough for multiple asyncio event loop cycles to run so
@@ -26,7 +26,11 @@ async def sim_sampleAngleStage():
 @pytest.fixture
 async def sim_filter_wheel():
     async with DeviceCollector(mock=True):
-        sim_filter_wheel = FilterMotor("p99-MO-TABLE-01:", name="sim_filter_wheel")
+        sim_filter_wheel = FilterMotor(
+            filter_selections=p99StageSelections,
+            prefix="p99-MO-TABLE-01:",
+            name="sim_filter_wheel",
+        )
     yield sim_filter_wheel
 
 

tests/devices/unit_tests/test_attenuator.py

@@ -6,29 +6,33 @@
 from ophyd_async.core import DeviceCollector
 from ophyd_async.testing import callback_on_mock_put, set_mock_value
 
-from dodal.devices.attenuator import Attenuator
+from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 
 CALCULATED_VALUE = [True, False, True] * 6  # Some "random" values
 
 
 @pytest.fixture
 async def fake_attenuator():
     async with DeviceCollector(mock=True):
-        fake_attenuator: Attenuator = Attenuator("", "attenuator")
+        fake_attenuator: BinaryFilterAttenuator = BinaryFilterAttenuator(
+            "", "attenuator"
+        )
 
     return fake_attenuator
 
 
-async def test_set_transmission_success(fake_attenuator: Attenuator):
+async def test_set_transmission_success(fake_attenuator: BinaryFilterAttenuator):
     await fake_attenuator.set(1.0)
 
 
-def test_set_transmission_in_run_engine(fake_attenuator: Attenuator, RE: RunEngine):
+def test_set_transmission_in_run_engine(
+    fake_attenuator: BinaryFilterAttenuator, RE: RunEngine
+):
     RE(bps.abs_set(fake_attenuator, 1, wait=True))
 
 
 async def test_given_attenuator_sets_filters_to_expected_value_then_set_returns(
-    fake_attenuator: Attenuator,
+    fake_attenuator: BinaryFilterAttenuator,
 ):
     def mock_apply_values(*args, **kwargs):
         for i in range(16):
@@ -43,7 +47,7 @@ def mock_apply_values(*args, **kwargs):
 
 
 async def test_given_attenuator_fails_to_set_filters_then_set_timeout(
-    fake_attenuator: Attenuator,
+    fake_attenuator: BinaryFilterAttenuator,
 ):
     def mock_apply_values(*args, **kwargs):
         for i in range(16):