Add DualFastShutter (#1795)

* Add DualFastShutter

* Add doc strings to shutters

* Add tests for dual_fast_shutter

* Added tests for read

* Use asycio.gather in test

* Improve I09 beamline device layout

* Added doc string to SourceSelector and get_obj_from_selected_source

* Added test_selectable_source

* Added additional tests for selectable_source

* minor test update

* Added read only signals for shutter_device_name

* Fixed tests

* Fix test flakeyness

* Fix controller arg

* Fixed flakey test

* Update doc string

* Remove pyproject.toml changes as now added to ophyd-async main

* Imrpove doc strings

* Setup new test names

* Revert remaining changes to pyproject.toml

* Added dual_fast_shutter_fixture

* Added code coverage

* Restored pyproject.toml

Author: oliwenmandiamond

src/dodal/beamlines/i09.py

@@ -1,3 +1,5 @@
+from ophyd_async.core import InOut
+
 from dodal.common.beamlines.beamline_utils import (
     device_factory,
 )
@@ -7,10 +9,14 @@
     PitchAndRollCrystal,
     StationaryCrystal,
 )
-from dodal.devices.electron_analyser.base import DualEnergySource
+from dodal.devices.electron_analyser.base import (
+    DualEnergySource,
+)
 from dodal.devices.electron_analyser.vgscienta import VGScientaDetector
+from dodal.devices.fast_shutter import DualFastShutter, GenericFastShutter
 from dodal.devices.i09 import Grating, LensMode, PassEnergy, PsuMode
 from dodal.devices.pgm import PlaneGratingMonochromator
+from dodal.devices.selectable_source import SourceSelector
 from dodal.devices.synchrotron import Synchrotron
 from dodal.devices.temperture_controller import (
     Lakeshore336,
@@ -19,7 +25,8 @@
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
 BL = get_beamline_name("i09")
-PREFIX = BeamlinePrefix(BL)
+I_PREFIX = BeamlinePrefix(BL, suffix="I")
+J_PREFIX = BeamlinePrefix(BL, suffix="J")
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
@@ -29,36 +36,65 @@ def synchrotron() -> Synchrotron:
     return Synchrotron()
 
 
+@device_factory()
+def source_selector() -> SourceSelector:
+    return SourceSelector()
+
+
 @device_factory()
 def pgm() -> PlaneGratingMonochromator:
     return PlaneGratingMonochromator(
-        prefix=f"{BeamlinePrefix(BL, suffix='J').beamline_prefix}-MO-PGM-01:",
-        grating=Grating,
+        prefix=f"{J_PREFIX.beamline_prefix}-MO-PGM-01:", grating=Grating
     )
 
 
 @device_factory()
 def dcm() -> DoubleCrystalMonochromatorWithDSpacing:
     return DoubleCrystalMonochromatorWithDSpacing(
-        f"{PREFIX.beamline_prefix}-MO-DCM-01:", PitchAndRollCrystal, StationaryCrystal
+        f"{I_PREFIX.beamline_prefix}-MO-DCM-01:", PitchAndRollCrystal, StationaryCrystal
+    )
+
+
+@device_factory()
+def dual_energy_source() -> DualEnergySource:
+    return DualEnergySource(
+        dcm().energy_in_eV,
+        pgm().energy.user_readback,
+        source_selector().selected_source,
+    )
+
+
+@device_factory()
+def fsi1() -> GenericFastShutter[InOut]:
+    return GenericFastShutter[InOut](
+        f"{I_PREFIX.beamline_prefix}-EA-FSHTR-01:CTRL", InOut.OUT, InOut.IN
+    )
+
+
+@device_factory()
+def fsj1() -> GenericFastShutter[InOut]:
+    return GenericFastShutter[InOut](
+        f"{J_PREFIX.beamline_prefix}-EA-FSHTR-01:CTRL", InOut.OUT, InOut.IN
     )
 
 
 @device_factory()
-def energy_source() -> DualEnergySource:
-    return DualEnergySource(dcm().energy_in_eV, pgm().energy.user_readback)
+def dual_fast_shutter() -> DualFastShutter[InOut]:
+    return DualFastShutter[InOut](fsi1(), fsj1(), source_selector().selected_source)
 
 
 # Connect will work again after this work completed
 # https://jira.diamond.ac.uk/browse/I09-651
 @device_factory()
 def ew4000() -> VGScientaDetector[LensMode, PsuMode, PassEnergy]:
     return VGScientaDetector[LensMode, PsuMode, PassEnergy](
-        prefix=f"{PREFIX.beamline_prefix}-EA-DET-01:CAM:",
+        prefix=f"{I_PREFIX.beamline_prefix}-EA-DET-01:CAM:",
         lens_mode_type=LensMode,
         psu_mode_type=PsuMode,
         pass_energy_type=PassEnergy,
-        energy_source=energy_source(),
+        energy_source=dual_energy_source(),
+        shutter=dual_fast_shutter(),
+        source_selector=source_selector(),
     )
 
 

src/dodal/beamlines/p60.py

@@ -11,6 +11,7 @@
     PassEnergy,
     PsuMode,
 )
+from dodal.devices.selectable_source import SourceSelector
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
@@ -20,6 +21,11 @@
 set_utils_beamline(BL)
 
 
+@device_factory()
+def source_selector() -> SourceSelector:
+    return SourceSelector()
+
+
 @device_factory()
 def al_kalpha_source() -> LabXraySourceReadable:
     return LabXraySourceReadable(LabXraySource.AL_KALPHA)
@@ -32,7 +38,11 @@ def mg_kalpha_source() -> LabXraySourceReadable:
 
 @device_factory()
 def energy_source() -> DualEnergySource:
-    return DualEnergySource(al_kalpha_source().energy_ev, mg_kalpha_source().energy_ev)
+    return DualEnergySource(
+        al_kalpha_source().energy_ev,
+        mg_kalpha_source().energy_ev,
+        source_selector().selected_source,
+    )
 
 
 # Connect will work again after this work completed

src/dodal/devices/electron_analyser/base/__init__.py

@@ -15,7 +15,7 @@
     GenericAnalyserDriverIO,
     TAbstractAnalyserDriverIO,
 )
-from .base_enums import EnergyMode, SelectedSource
+from .base_enums import EnergyMode
 from .base_region import (
     AbstractBaseRegion,
     AbstractBaseSequence,
@@ -27,7 +27,7 @@
     TLensMode,
 )
 from .base_util import to_binding_energy, to_kinetic_energy
-from .energy_sources import DualEnergySource, EnergySource
+from .energy_sources import AbstractEnergySource, DualEnergySource, EnergySource
 
 __all__ = [
     "ElectronAnalyserController",
@@ -42,7 +42,6 @@
     "GenericAnalyserDriverIO",
     "TAbstractAnalyserDriverIO",
     "EnergyMode",
-    "SelectedSource",
     "AbstractBaseRegion",
     "AbstractBaseSequence",
     "GenericRegion",
@@ -53,6 +52,7 @@
     "TLensMode",
     "to_binding_energy",
     "to_kinetic_energy",
+    "AbstractEnergySource",
     "DualEnergySource",
     "EnergySource",
 ]

src/dodal/devices/electron_analyser/base/base_controller.py

@@ -12,10 +12,9 @@
     GenericRegion,
     TAbstractBaseRegion,
 )
-from dodal.devices.electron_analyser.base.energy_sources import (
-    AbstractEnergySource,
-    DualEnergySource,
-)
+from dodal.devices.electron_analyser.base.energy_sources import AbstractEnergySource
+from dodal.devices.fast_shutter import FastShutter
+from dodal.devices.selectable_source import SourceSelector
 
 
 class ElectronAnalyserController(
@@ -32,7 +31,9 @@ def __init__(
         self,
         driver: TAbstractAnalyserDriverIO,
         energy_source: AbstractEnergySource,
-        deadtime: float,
+        shutter: FastShutter | None = None,
+        source_selector: SourceSelector | None = None,
+        deadtime: float = 0,
         image_mode: ADImageMode = ADImageMode.SINGLE,
     ):
         """
@@ -45,13 +46,19 @@ def __init__(
             image_mode: The image mode to configure the driver with before measuring.
         """
         self.energy_source = energy_source
+        self.shutter = shutter
+        self.source_selector = source_selector
         super().__init__(driver, deadtime, image_mode)
 
-    async def setup_with_region(self, region: TAbstractBaseRegion):
+    async def setup_with_region(self, region: TAbstractBaseRegion) -> None:
         """Logic to set the driver with a region."""
+        if self.source_selector is not None:
+            await self.source_selector.set(region.excitation_energy_source)
+
+        # Should this be moved to a VGScientController only?
+        if self.shutter is not None:
+            await self.shutter.set(self.shutter.close_state)
 
-        if isinstance(self.energy_source, DualEnergySource):
-            self.energy_source.selected_source.set(region.excitation_energy_source)
         excitation_energy = await self.energy_source.energy.get_value()
         epics_region = region.prepare_for_epics(excitation_energy)
         await self.driver.set(epics_region)
@@ -62,6 +69,10 @@ async def prepare(self, trigger_info: TriggerInfo) -> None:
         # axis calculation.
         excitation_energy = await self.energy_source.energy.get_value()
         await self.driver.cached_excitation_energy.set(excitation_energy)
+
+        if self.shutter is not None:
+            await self.shutter.set(self.shutter.open_state)
+
         await super().prepare(trigger_info)
 
 

src/dodal/devices/electron_analyser/base/base_enums.py

@@ -4,8 +4,3 @@
 class EnergyMode(StrictEnum):
     KINETIC = "Kinetic"
     BINDING = "Binding"
-
-
-class SelectedSource(StrictEnum):
-    SOURCE1 = "source1"
-    SOURCE2 = "source2"

src/dodal/devices/electron_analyser/base/base_region.py

@@ -6,11 +6,12 @@
 from ophyd_async.core import StrictEnum, SupersetEnum
 from pydantic import BaseModel, Field, model_validator
 
-from dodal.devices.electron_analyser.base.base_enums import EnergyMode, SelectedSource
+from dodal.devices.electron_analyser.base.base_enums import EnergyMode
 from dodal.devices.electron_analyser.base.base_util import (
     to_binding_energy,
     to_kinetic_energy,
 )
+from dodal.devices.selectable_source import SelectedSource
 
 AnyAcqMode: TypeAlias = StrictEnum
 AnyLensMode: TypeAlias = SupersetEnum | StrictEnum

src/dodal/devices/electron_analyser/base/energy_sources.py

@@ -3,14 +3,14 @@
 from ophyd_async.core import (
     Reference,
     SignalR,
+    SignalRW,
     StandardReadable,
     StandardReadableFormat,
     derived_signal_r,
     soft_signal_r_and_setter,
-    soft_signal_rw,
 )
 
-from dodal.devices.electron_analyser.base.base_enums import SelectedSource
+from dodal.devices.selectable_source import SelectedSource, get_obj_from_selected_source
 
 
 class AbstractEnergySource(StandardReadable):
@@ -51,51 +51,52 @@ def energy(self) -> SignalR[float]:
         return self._source_ref()
 
 
+def get_float_from_selected_source(
+    selected: SelectedSource, s1: float, s2: float
+) -> float:
+    """Wrapper function to provide type hints for derived signal."""
+    return get_obj_from_selected_source(selected, s1, s2)
+
+
 class DualEnergySource(AbstractEnergySource):
     """
     Holds two EnergySource devices and provides a signal to read energy depending on
-    which source is selected. This is controlled by a selected_source signal which can
-    switch source using SelectedSource enum. Both sources energy is recorded in the
-    read, the energy signal is used as a helper signal to know which source is being
-    used.
+    which source is selected. The energy is the one that corrosponds to the
+    selected_source signal. For example, selected_source is source1 if selected_source
+    is at SelectedSource.SOURCE1 and vise versa for source2 and SelectedSource.SOURCE2.
     """
 
     def __init__(
-        self, source1: SignalR[float], source2: SignalR[float], name: str = ""
+        self,
+        source1: SignalR[float],
+        source2: SignalR[float],
+        selected_source: SignalRW[SelectedSource],
+        name: str = "",
     ):
         """
         Args:
-            source1: Default energy signal to select.
-            source2: Secondary energy signal to select.
-            name: name of this device.
+            source1: Energy source that corrosponds to SelectedSource.SOURCE1.
+            source2: Energy source that corrosponds to SelectedSource.SOURCE2.
+            selected_source: Signal that decides the active energy source.
+            name: Name of this device.
         """
 
+        self.selected_source_ref = Reference(selected_source)
         with self.add_children_as_readables():
-            self.selected_source = soft_signal_rw(
-                SelectedSource, initial_value=SelectedSource.SOURCE1
-            )
             self.source1 = EnergySource(source1)
             self.source2 = EnergySource(source2)
 
         self._selected_energy = derived_signal_r(
-            self._get_excitation_energy,
+            get_float_from_selected_source,
             "eV",
-            selected_source=self.selected_source,
-            source1=self.source1.energy,
-            source2=self.source2.energy,
+            selected=self.selected_source_ref(),
+            s1=self.source1.energy,
+            s2=self.source2.energy,
         )
+        self.add_readables([selected_source])
 
         super().__init__(name)
 
-    def _get_excitation_energy(
-        self, selected_source: SelectedSource, source1: float, source2: float
-    ) -> float:
-        match selected_source:
-            case SelectedSource.SOURCE1:
-                return source1
-            case SelectedSource.SOURCE2:
-                return source2
-
     @property
     def energy(self) -> SignalR[float]:
         return self._selected_energy

src/dodal/devices/electron_analyser/specs/specs_detector.py

@@ -5,15 +5,14 @@
 )
 from dodal.devices.electron_analyser.base.base_detector import ElectronAnalyserDetector
 from dodal.devices.electron_analyser.base.base_region import TLensMode, TPsuMode
-from dodal.devices.electron_analyser.base.energy_sources import (
-    DualEnergySource,
-    EnergySource,
-)
+from dodal.devices.electron_analyser.base.energy_sources import AbstractEnergySource
 from dodal.devices.electron_analyser.specs.specs_driver_io import SpecsAnalyserDriverIO
 from dodal.devices.electron_analyser.specs.specs_region import (
     SpecsRegion,
     SpecsSequence,
 )
+from dodal.devices.fast_shutter import FastShutter
+from dodal.devices.selectable_source import SourceSelector
 
 
 class SpecsDetector(
@@ -29,7 +28,9 @@ def __init__(
         prefix: str,
         lens_mode_type: type[TLensMode],
         psu_mode_type: type[TPsuMode],
-        energy_source: DualEnergySource | EnergySource,
+        energy_source: AbstractEnergySource,
+        shutter: FastShutter | None = None,
+        source_selector: SourceSelector | None = None,
         name: str = "",
     ):
         # Save to class so takes part with connect()
@@ -39,7 +40,7 @@ def __init__(
 
         controller = ElectronAnalyserController[
             SpecsAnalyserDriverIO[TLensMode, TPsuMode], SpecsRegion[TLensMode, TPsuMode]
-        ](self.driver, energy_source, 0)
+        ](self.driver, energy_source, shutter, source_selector)
 
         sequence_class = SpecsSequence[lens_mode_type, psu_mode_type]