diff --git a/src/dodal/beamlines/i09.py b/src/dodal/beamlines/i09.py index b7c989c4a18..b7c9c0c9d1c 100644 --- a/src/dodal/beamlines/i09.py +++ b/src/dodal/beamlines/i09.py @@ -18,6 +18,7 @@ from dodal.devices.hutch_shutter import EXP_SHUTTER_2_INFIX, HutchShutter from dodal.devices.motors import XYZAzimuthPolarStage from dodal.devices.pgm import PlaneGratingMonochromator +from dodal.devices.scaler import ScalerController, SimpleChannelScaler from dodal.devices.selectable_source import SourceSelector from dodal.devices.synchrotron import Synchrotron from dodal.devices.temperture_controller import Lakeshore336 @@ -27,6 +28,7 @@ BL = get_beamline_name("i09") I_PREFIX = BeamlinePrefix(BL, suffix="I") J_PREFIX = BeamlinePrefix(BL, suffix="J") +L_PREFIX = BeamlinePrefix(BL, suffix="L") set_log_beamline(BL) set_utils_beamline(BL) @@ -139,3 +141,53 @@ def intensity_protection() -> SignalRW[IntensityProtection]: return epics_signal_rw( IntensityProtection, f"{I_PREFIX.beamline_prefix}-DI-EAN-01:PROT:ILK" ) + + +@devices.factory +def scaler1() -> ScalerController: + return ScalerController(f"{I_PREFIX.beamline_prefix}-EA-SCLR-01") + + +@devices.factory +def hm3amp20_1(scaler1: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler1, channel=2) + + +@devices.factory +def sm5amp8_1(scaler1: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler1, channel=3) + + +@devices.factory +def smpmamp39_1(scaler1: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler1, channel=4) + + +@devices.factory +def rfdamp10_1(scaler1: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler1, channel=5) + + +@devices.factory +def scaler2() -> ScalerController: + return ScalerController(f"{L_PREFIX}-VA-SCLR-01") + + +@devices.factory +def hm3amp20(scaler2: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler2, channel=2) + + +@devices.factory +def sm5amp8(scaler2: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler2, channel=3) + + +@devices.factory +def smpmamp39(scaler2: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler2, channel=4) + + +@devices.factory +def rfdamp10(scaler2: ScalerController) -> SimpleChannelScaler: + return SimpleChannelScaler(scaler2, channel=5) diff --git a/src/dodal/devices/scaler.py b/src/dodal/devices/scaler.py new file mode 100644 index 00000000000..9e7aa63d69b --- /dev/null +++ b/src/dodal/devices/scaler.py @@ -0,0 +1,94 @@ +import asyncio + +from bluesky.protocols import Reading, Triggerable +from ophyd_async.core import ( + AsyncStatus, + DeviceMock, + Reference, + StandardReadable, + StandardReadableFormat, + default_mock_class, + set_and_wait_for_value, + set_mock_value, +) +from ophyd_async.epics.core import epics_signal_rw, wait_for_good_state + + +class MockScalerController(DeviceMock["ScalerController"]): + async def connect(self, device: "ScalerController"): + set_mock_value(device.counting, False) + + async def _complete(): + await asyncio.sleep(0.2) + set_mock_value(device.counting, False) + + def _on_value(value: dict[str, Reading[bool]]): + if value[device.counting.name]["value"] is True: + asyncio.create_task(_complete()) + + # Can't use callback_on_mock_put as this is called before the mock put, we need + # to simulate after mock put update. + device.counting.subscribe_reading(_on_value) + + +@default_mock_class(MockScalerController) +class ScalerController(StandardReadable, Triggerable): + """Scaler controller that is triggerable. It will set the counting signal to True + and then waits for it to be False. + """ + + def __init__(self, prefix: str, name: str = ""): + self.counting = epics_signal_rw(bool, prefix + ".CNT") + with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL): + self.count_time = epics_signal_rw(float, prefix + ".TP") + + self._acquire_status: AsyncStatus | None = None + # Store the prefix so that the SimpleChannelScaler can reuse. + self.prefix = prefix + + super().__init__(name) + + @AsyncStatus.wrap + async def trigger(self): + self._acquire_status = await set_and_wait_for_value( + self.counting, True, wait_for_set_completion=True + ) + await self._acquire_status + await wait_for_good_state(self.counting, {False}) + + +class SimpleChannelScaler(StandardReadable, Triggerable): + """Create individual channel for a scaler. A ScalerController is used for the + Trigger logic. It will also add this instance signals as readables to the + ScannableController and also add the controllers count_period signal to this + classes read configuration. + """ + + def __init__( + self, + scalar_controller: ScalerController, + channel: int, + name: str = "", + ): + self._scaler_controller_ref = Reference(scalar_controller) + + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.count = epics_signal_rw( + float, f"{scalar_controller.prefix}.S{channel}" + ) + + super().__init__(name) + + scalar_controller.add_readables([self]) + # Avoid circular read configuration by specifying individual signal + self.add_readables( + [scalar_controller.count_time], StandardReadableFormat.CONFIG_SIGNAL + ) + + @AsyncStatus.wrap + async def set(self, value: float): + await self.count.set(value) + + @AsyncStatus.wrap + async def trigger(self): + await self._scaler_controller_ref().trigger() diff --git a/tests/devices/test_scaler.py b/tests/devices/test_scaler.py new file mode 100644 index 00000000000..4539c81ff52 --- /dev/null +++ b/tests/devices/test_scaler.py @@ -0,0 +1,104 @@ +import time +from unittest.mock import AsyncMock, call + +import pytest +from ophyd_async.core import get_mock_put, init_devices +from ophyd_async.testing import assert_configuration, assert_reading, partial_reading + +from dodal.devices.scaler import ScalerController, SimpleChannelScaler + + +@pytest.fixture +def scaler1() -> ScalerController: + with init_devices(mock=True): + scaler1 = ScalerController("TEST:") + return scaler1 + + +@pytest.fixture +def hm3amp20_1(scaler1: ScalerController) -> SimpleChannelScaler: + with init_devices(mock=True): + hm3amp20_1 = SimpleChannelScaler(scaler1, 1) + return hm3amp20_1 + + +@pytest.fixture +def sm5amp8_1(scaler1: ScalerController) -> SimpleChannelScaler: + with init_devices(mock=True): + sm5amp8_1 = SimpleChannelScaler(scaler1, 2) + return sm5amp8_1 + + +async def test_scaler_controller_read(scaler1: ScalerController) -> None: + await assert_reading(scaler1, {}) + + +async def test_scaler_controller_read_with_multiple_channels( + scaler1: ScalerController, + hm3amp20_1: SimpleChannelScaler, + sm5amp8_1: SimpleChannelScaler, +) -> None: + await assert_reading( + scaler1, + { + "hm3amp20_1-count": partial_reading(0), + "sm5amp8_1-count": partial_reading(0), + }, + ) + + +async def test_scaler_controller_read_configuration(scaler1: ScalerController) -> None: + await assert_configuration(scaler1, {"scaler1-count_time": partial_reading(0.0)}) + + +async def test_scaler_controller_trigger_waits_for_counting_to_finish( + scaler1: ScalerController, +) -> None: + start = time.monotonic() + await scaler1.trigger() + elapsed = time.monotonic() - start + assert elapsed >= 0.2 + assert await scaler1.counting.get_value() is False + + +async def test_scaler_controller_trigger_sets_counting_true_then_false( + scaler1: ScalerController, +) -> None: + values = [] + scaler1.counting.subscribe(values.append) + await scaler1.trigger() + + states = [] + expected_states = [False, True, False] + for v in values: + states.append(v[scaler1.counting.name]["value"]) + assert states == expected_states + + +async def test_simple_channel_scaler_read(hm3amp20_1: SimpleChannelScaler) -> None: + await assert_reading(hm3amp20_1, {"hm3amp20_1-count": partial_reading(0)}) + + +async def test_simple_channel_scaler_read_configuration( + hm3amp20_1: SimpleChannelScaler, +) -> None: + await assert_configuration(hm3amp20_1, {"scaler1-count_time": partial_reading(0)}) + + +async def test_simple_channel_scaler_set(hm3amp20_1: SimpleChannelScaler) -> None: + value = 10 + await hm3amp20_1.set(value) + get_mock_put(hm3amp20_1.count).assert_awaited_once_with(value) + + +async def test_simple_channel_scaler_trigger( + scaler1: ScalerController, hm3amp20_1: SimpleChannelScaler, sm5amp8_1 +) -> None: + mock_trigger = AsyncMock() + scaler1.trigger = mock_trigger + + await hm3amp20_1.trigger() + mock_trigger.assert_awaited_once() + + await sm5amp8_1.trigger() + mock_trigger.assert_has_calls([call(), call()])