From 6cce80e93c2e3cac3d43d098279c1f122acfbeaa Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Wed, 24 Apr 2024 15:10:19 +0100 Subject: [PATCH] Port UndulatorDCM to ophyd-async and remove energy signal --- src/dodal/devices/undulator_dcm.py | 67 ++++++---- .../unit_tests/test_focusing_mirror.py | 8 +- .../devices/unit_tests/test_undulator_dcm.py | 124 ++++++++++-------- 3 files changed, 107 insertions(+), 92 deletions(-) diff --git a/src/dodal/devices/undulator_dcm.py b/src/dodal/devices/undulator_dcm.py index f0eee51016..cca1c8933a 100644 --- a/src/dodal/devices/undulator_dcm.py +++ b/src/dodal/devices/undulator_dcm.py @@ -1,14 +1,16 @@ +import asyncio + import numpy as np +from bluesky.protocols import Movable from numpy import argmin, loadtxt, ndarray -from ophyd.status import Status -from ophyd_async.core import SignalW, StandardReadable +from ophyd_async.core import AsyncStatus, StandardReadable from dodal.devices.dcm import DCM from dodal.devices.undulator import Undulator, UndulatorGapAccess from dodal.log import LOGGER -ENERGY_TIMEOUT_S = 30 -STATUS_TIMEOUT_S = 10 +ENERGY_TIMEOUT_S: float = 30.0 +STATUS_TIMEOUT_S: float = 10.0 # Enable to allow testing when the beamline is down, do not change in production! TEST_MODE = False @@ -19,6 +21,7 @@ class AccessError(Exception): def _get_energy_distance_table(lookup_table_path: str) -> ndarray: + # TODO: Make IO async return loadtxt(lookup_table_path, comments=["#", "Units"]) @@ -30,60 +33,66 @@ def _get_closest_gap_for_energy( return table[1][idx] -class UndulatorDCM(StandardReadable): +class UndulatorDCM(StandardReadable, Movable): """ Composite device to handle changing beamline energies """ - class EnergySignal(Signal): - parent: "UndulatorDCM" + _setpoint: float + + def __init__(self, undulator: Undulator, dcm: DCM, name: str = ""): + self.undulator = undulator + self.dcm = dcm + + super().__init__(name) - def set(self, value, *, timeout=None, settle_time=None, **kwargs) -> Status: + def set(self, value: float) -> AsyncStatus: + async def _set(value: float): + # TODO: Break up into smaller methods energy_kev = value - access_level = self.parent.undulator.gap_access.get(as_string=True) - if access_level == UndulatorGapAccess.DISABLED.value and not TEST_MODE: + access_level = await self.undulator.gap_access.get_value() + if access_level is UndulatorGapAccess.DISABLED and not TEST_MODE: raise AccessError( "Undulator gap access is disabled. Contact Control Room" ) # Get 2d np.array converting energies to undulator gap distance, from lookup table energy_to_distance_table = _get_energy_distance_table( - self.parent.undulator.lookup_table_path + self.undulator.lookup_table_path ) LOGGER.info(f"Setting DCM energy to {energy_kev:.2f} kev") - status = self.parent.dcm.energy_in_kev.move( - energy_kev, timeout=ENERGY_TIMEOUT_S - ) + statuses = [ + self.dcm.energy_in_kev.set( + energy_kev, + timeout=ENERGY_TIMEOUT_S, + ) + ] # Use the lookup table to get the undulator gap associated with this dcm energy gap_to_match_dcm_energy = _get_closest_gap_for_energy( - energy_kev * 1000, energy_to_distance_table + energy_kev * 1000, + energy_to_distance_table, ) # Check if undulator gap is close enough to the value from the DCM - current_gap = self.parent.undulator.current_gap.get() + current_gap = await self.undulator.current_gap.get_value() if ( abs(gap_to_match_dcm_energy - current_gap) - > self.parent.undulator.gap_discrepancy_tolerance_mm + > self.undulator.gap_discrepancy_tolerance_mm ): LOGGER.info( f"Undulator gap mismatch. {abs(gap_to_match_dcm_energy-current_gap):.3f}mm is outside tolerance.\ Moving gap to nominal value, {gap_to_match_dcm_energy:.3f}mm" ) if not TEST_MODE: - status &= self.parent.undulator.gap_motor.move( - gap_to_match_dcm_energy, timeout=STATUS_TIMEOUT_S + statuses.append( + self.undulator.gap_motor.set( + gap_to_match_dcm_energy, + timeout=STATUS_TIMEOUT_S, + ) ) + await asyncio.gather(*statuses) - return status - - energy_kev = Component(EnergySignal) - - def __init__(self, undulator: Undulator, dcm: DCM, *args, **kwargs): - self.undulator = undulator - self.dcm = dcm - self.energy_kev = SignalW() - - super().__init__(*args, **kwargs) + return AsyncStatus(_set(value)) diff --git a/tests/devices/unit_tests/test_focusing_mirror.py b/tests/devices/unit_tests/test_focusing_mirror.py index 5636fe94e0..32a6974073 100644 --- a/tests/devices/unit_tests/test_focusing_mirror.py +++ b/tests/devices/unit_tests/test_focusing_mirror.py @@ -63,9 +63,7 @@ def not_ok(_): def test_mirror_set_voltage_sets_and_waits_happy_path( vfm_mirror_voltages_with_set: VFMMirrorVoltages, ): - vfm_mirror_voltages_with_set._channel14_voltage_device._setpoint_v.set.return_value = ( - NullStatus() - ) + vfm_mirror_voltages_with_set._channel14_voltage_device._setpoint_v.set.return_value = NullStatus() vfm_mirror_voltages_with_set._channel14_voltage_device._demand_accepted.sim_put( MirrorVoltageDemand.OK ) @@ -103,9 +101,7 @@ def test_mirror_set_voltage_sets_and_waits_set_fail( def test_mirror_set_voltage_sets_and_waits_settle_timeout_expires( vfm_mirror_voltages_with_set_timing_out: VFMMirrorVoltages, ): - vfm_mirror_voltages_with_set_timing_out._channel14_voltage_device._setpoint_v.set.return_value = ( - NullStatus() - ) + vfm_mirror_voltages_with_set_timing_out._channel14_voltage_device._setpoint_v.set.return_value = NullStatus() status: StatusBase = vfm_mirror_voltages_with_set_timing_out.voltage_channels[ 0 diff --git a/tests/devices/unit_tests/test_undulator_dcm.py b/tests/devices/unit_tests/test_undulator_dcm.py index 424756355f..e27342d04b 100644 --- a/tests/devices/unit_tests/test_undulator_dcm.py +++ b/tests/devices/unit_tests/test_undulator_dcm.py @@ -1,9 +1,13 @@ +import asyncio from unittest.mock import MagicMock, patch import numpy as np import pytest -from ophyd.sim import make_fake_device -from ophyd.status import Status +from ophyd_async.core import ( + AsyncStatus, + DeviceCollector, + set_sim_value, +) from dodal.devices.dcm import DCM from dodal.devices.undulator import Undulator, UndulatorGapAccess @@ -18,26 +22,24 @@ @pytest.fixture -def fake_undulator_dcm() -> UndulatorDCM: - undulator: Undulator = make_fake_device(Undulator)( - name="undulator", - lookup_table_path="./tests/devices/unit_tests/test_beamline_undulator_to_gap_lookup_table.txt", - ) - dcm: DCM = make_fake_device(DCM)( - name="dcm", daq_configuration_path=MOCK_DAQ_CONFIG_PATH - ) - undulator_dcm: UndulatorDCM = make_fake_device(UndulatorDCM)( - undulator, dcm, name="undulator_dcm" - ) +async def fake_undulator_dcm() -> UndulatorDCM: + async with DeviceCollector(sim=True): + undulator = Undulator( + "UND-01", + name="undulator", + lookup_table_path="./tests/devices/unit_tests/test_beamline_undulator_to_gap_lookup_table.txt", + ) + dcm = DCM("DCM-01", name="dcm", daq_configuration_path=MOCK_DAQ_CONFIG_PATH) + undulator_dcm = UndulatorDCM(undulator, dcm, name="undulator_dcm") return undulator_dcm -def test_when_gap_access_is_disabled_set_energy_then_error_is_raised( +async def test_when_gap_access_is_disabled_set_energy_then_error_is_raised( fake_undulator_dcm: UndulatorDCM, ): - fake_undulator_dcm.undulator.gap_access.sim_put(UndulatorGapAccess.DISABLED.value) # type: ignore + set_sim_value(fake_undulator_dcm.undulator.gap_access, UndulatorGapAccess.DISABLED) with pytest.raises(AccessError): - fake_undulator_dcm.energy_kev.set(5) + await fake_undulator_dcm.set(5) def test_energy_to_distance_table_correct_format(fake_undulator_dcm: UndulatorDCM): @@ -60,80 +62,88 @@ def test_correct_closest_distance_to_energy_from_table(dcm_energy, expected_outp @patch("dodal.devices.undulator_dcm.loadtxt") @patch("dodal.devices.undulator_dcm.LOGGER") -def test_if_gap_is_wrong_then_logger_info_is_called_and_gap_is_set_correctly( +async def test_if_gap_is_wrong_then_logger_info_is_called_and_gap_is_set_correctly( mock_logger: MagicMock, mock_load: MagicMock, fake_undulator_dcm: UndulatorDCM ): - fake_undulator_dcm.undulator.current_gap.sim_put(5.3) # type: ignore - fake_undulator_dcm.undulator.gap_motor.move = MagicMock() - fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock() + set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.3) + set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.readback, 5.7) + + # fake_undulator_dcm.undulator.gap_motor.move = MagicMock() + # fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock() mock_load.return_value = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]]) - fake_undulator_dcm.dcm.energy_in_kev.user_readback.sim_put(5.7) # type: ignore - fake_undulator_dcm.energy_kev.set(6.9) + await fake_undulator_dcm.set(6.9) - fake_undulator_dcm.dcm.energy_in_kev.move.assert_called_once_with(6.9, timeout=30) - fake_undulator_dcm.undulator.gap_motor.move.assert_called_once_with( - 6.045, timeout=10 - ) + assert (await fake_undulator_dcm.dcm.energy_in_kev.setpoint.get_value()) == 6.9 + assert (await fake_undulator_dcm.undulator.gap_motor.setpoint.get_value()) == 6.045 mock_logger.info.assert_called() @patch("dodal.devices.undulator_dcm.loadtxt") @patch("dodal.devices.undulator_dcm.LOGGER") @patch("dodal.devices.undulator_dcm.TEST_MODE", True) -def test_when_gap_access_is_not_checked_if_test_mode_enabled( +async def test_when_gap_access_is_not_checked_if_test_mode_enabled( mock_logger: MagicMock, mock_load: MagicMock, fake_undulator_dcm: UndulatorDCM ): - fake_undulator_dcm.undulator.gap_access.sim_put(UndulatorGapAccess.DISABLED.value) # type: ignore - fake_undulator_dcm.undulator.current_gap.sim_put(5.3) # type: ignore - fake_undulator_dcm.undulator.gap_motor.move = MagicMock() - fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock() + set_sim_value(fake_undulator_dcm.undulator.gap_access, UndulatorGapAccess.DISABLED) + set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.3) + set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.readback, 5.7) + + set_sim_value(fake_undulator_dcm.undulator.gap_motor.setpoint, 0.0) + set_sim_value(fake_undulator_dcm.undulator.gap_motor.readback, 0.0) + + # fake_undulator_dcm.undulator.gap_motor.move = MagicMock() + # fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock() mock_load.return_value = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]]) - fake_undulator_dcm.dcm.energy_in_kev.user_readback.sim_put(5.7) # type: ignore - fake_undulator_dcm.energy_kev.set(6.9) + await fake_undulator_dcm.set(6.9) + + assert (await fake_undulator_dcm.dcm.energy_in_kev.setpoint.get_value()) == 6.9 + # Verify undulator has not been asked to move + assert (await fake_undulator_dcm.undulator.gap_motor.setpoint.get_value()) == 0.0 - fake_undulator_dcm.dcm.energy_in_kev.move.assert_called_once_with(6.9, timeout=30) - fake_undulator_dcm.undulator.gap_motor.move.assert_not_called() mock_logger.info.assert_called() @patch("dodal.devices.undulator_dcm.loadtxt") @patch("dodal.devices.undulator_dcm.LOGGER") -def test_if_gap_is_already_correct_then_dont_move_gap( +async def test_if_gap_is_already_correct_then_dont_move_gap( mock_logger: MagicMock, mock_load: MagicMock, fake_undulator_dcm: UndulatorDCM ): - fake_undulator_dcm.undulator.gap_motor.move = MagicMock() - fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock() + set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.setpoint, 0.0) + set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.readback, 0.0) + # fake_undulator_dcm.undulator.gap_motor.move = MagicMock() + # fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock() mock_load.return_value = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]]) - fake_undulator_dcm.undulator.current_gap.sim_put(5.4605) # type: ignore + set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.4605) - fake_undulator_dcm.energy_kev.set(5.8).wait(timeout=0.01) + status = fake_undulator_dcm.set(5.8) + await asyncio.wait_for(status, timeout=0.01) - fake_undulator_dcm.undulator.gap_motor.move.assert_not_called() + # Verify undulator has not been asked to move + assert (await fake_undulator_dcm.undulator.gap_motor.setpoint.get_value()) == 0.0 mock_logger.info.assert_called_once() -def test_energy_set_only_complete_when_all_statuses_are_finished( +async def test_energy_set_only_complete_when_all_statuses_are_finished( fake_undulator_dcm: UndulatorDCM, ): - dcm_energy_move_status = Status() - undulator_gap_move_status = Status() + set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.0) + + release_dcm = asyncio.Event() + release_undulator = asyncio.Event() - fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock( - return_value=dcm_energy_move_status + fake_undulator_dcm.dcm.energy_in_kev.set = MagicMock( + return_value=AsyncStatus(release_dcm.wait()) ) - fake_undulator_dcm.undulator.gap_motor.move = MagicMock( - return_value=undulator_gap_move_status + fake_undulator_dcm.undulator.gap_motor.set = MagicMock( + return_value=AsyncStatus(release_undulator.wait()) ) - _get_energy_distance_table = MagicMock() - _get_closest_gap_for_energy = MagicMock(return_value=10) - fake_undulator_dcm.undulator.current_gap.sim_put(5) # type: ignore - status: Status = fake_undulator_dcm.energy_kev.set(5.8) + status = fake_undulator_dcm.set(5.0) - assert not status.success - dcm_energy_move_status.set_finished() - assert not status.success - undulator_gap_move_status.set_finished() - status.wait(timeout=0.01) + assert not status.done + release_dcm.set() + assert not status.done + release_undulator.set() + await asyncio.wait_for(status, timeout=0.01)