Skip to content

Commit

Permalink
Update to work with ophyd_async xspress3 (DiamondLightSource/hyperion…
Browse files Browse the repository at this point in the history
…#1434)

* Update optimise_attenuation plan to work with ophyd_async xspress3

* merged main into branch

---------

Co-authored-by: Raymond Fan <[email protected]>
Co-authored-by: Relm-Arrowny <[email protected]>
3 people authored Jun 19, 2024
1 parent fbdc42d commit 4e36438
Showing 3 changed files with 26 additions and 36 deletions.
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
@@ -46,7 +46,7 @@ install_requires =
bluesky >= 1.13.0a3
blueapi >= 0.4.3-rc1
# needed for CI while xpress is broken, remove on merging
dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@b6b22383cfb31cb5e4c5b88de0c785463cc1c8e6
dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git

[options.entry_points]
console_scripts =
18 changes: 7 additions & 11 deletions src/hyperion/experiment_plans/optimise_attenuation_plan.py
Original file line number Diff line number Diff line change
@@ -6,7 +6,7 @@
import numpy as np
from blueapi.core import BlueskyContext
from dodal.devices.attenuator import Attenuator
from dodal.devices.xspress3_mini.xspress3_mini import Xspress3Mini
from dodal.devices.xspress3.xspress3 import Xspress3
from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterState

from hyperion.log import LOGGER
@@ -28,7 +28,7 @@ class OptimizeAttenuationComposite:

attenuator: Attenuator
sample_shutter: ZebraShutter
xspress3mini: Xspress3Mini
xspress3mini: Xspress3


def create_devices(context: BlueskyContext) -> OptimizeAttenuationComposite:
@@ -75,11 +75,6 @@ def is_counts_within_target(total_count, lower_count_limit, upper_count_limit) -
return False


def arm_devices(xspress3mini: Xspress3Mini):
yield from bps.abs_set(xspress3mini.do_arm, 1, wait=True)
LOGGER.info("Arming Xspress3Mini complete")


def calculate_new_direction(direction: Direction, deadtime, deadtime_threshold):
if direction == Direction.POSITIVE:
if deadtime > deadtime_threshold:
@@ -154,7 +149,8 @@ def open_and_run():
yield from bps.abs_set(
composite.sample_shutter, ZebraShutterState.OPEN, wait=True
)
yield from bps.abs_set(composite.xspress3mini.do_arm, 1, wait=True)
yield from bps.stage(composite.xspress3mini, wait=True)
yield from bps.unstage(composite.xspress3mini, wait=True)

yield from open_and_run()

@@ -240,8 +236,8 @@ def deadtime_optimisation(
for cycle in range(0, max_cycles):
yield from do_device_optimise_iteration(composite, transmission)

total_time = float(composite.xspress3mini.channel_1.total_time.get())
reset_ticks = float(composite.xspress3mini.channel_1.reset_ticks.get())
total_time = yield from bps.rd(composite.xspress3mini.channels[1].total_time)
reset_ticks = yield from bps.rd(composite.xspress3mini.channels[1].reset_ticks)

LOGGER.info(f"Current total time = {total_time}")
LOGGER.info(f"Current reset ticks = {reset_ticks}")
@@ -354,7 +350,7 @@ def total_counts_optimisation(
yield from do_device_optimise_iteration(composite, transmission)

data = np.array(
(yield from bps.rd(composite.xspress3mini.dt_corrected_latest_mca))
(yield from bps.rd(composite.xspress3mini.dt_corrected_latest_mca[1]))
)
total_count = sum(data[int(low_roi) : int(high_roi)])
LOGGER.info(f"Total count is {total_count}")
42 changes: 18 additions & 24 deletions tests/unit_tests/experiment_plans/test_optimise_attenuation_plan.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
from bluesky.run_engine import RunEngine
from dodal.beamlines import i03
from ophyd.status import Status
from ophyd_async.core import set_mock_value

from hyperion.experiment_plans import optimise_attenuation_plan
from hyperion.experiment_plans.optimise_attenuation_plan import (
AttenuationOptimisationFailedException,
Direction,
OptimizeAttenuationComposite,
arm_devices,
calculate_new_direction,
check_parameters,
deadtime_calc_new_transmission,
@@ -57,8 +58,8 @@ def get_good_status():
@pytest.fixture
def fake_composite_mocked_sets(fake_composite: OptimizeAttenuationComposite):
with patch.object(
fake_composite.xspress3mini.do_arm,
"set",
fake_composite.xspress3mini,
"stage",
MagicMock(return_value=get_good_status()),
), patch.object(
fake_composite.sample_shutter,
@@ -93,8 +94,9 @@ def test_is_deadtime_is_optimised_logs_warning_when_upper_transmission_limit_is_
def test_total_counts_calc_new_transmission_raises_warning_on_high_transmission(
RE: RunEngine, mock_emit, fake_composite_mocked_sets: OptimizeAttenuationComposite
):
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca.sim_put( # type: ignore
[1, 1, 1, 1, 1, 1]
set_mock_value(
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca[1],
np.array([1, 1, 1, 1, 1, 1]),
)
RE(
total_counts_optimisation(
@@ -140,8 +142,8 @@ def test_deadtime_optimisation_calculates_deadtime_correctly(
RE: RunEngine,
fake_composite: OptimizeAttenuationComposite,
):
fake_composite.xspress3mini.channel_1.total_time.sim_put(100) # type: ignore
fake_composite.xspress3mini.channel_1.reset_ticks.sim_put(101) # type: ignore
set_mock_value(fake_composite.xspress3mini.channels[1].total_time, 100)
set_mock_value(fake_composite.xspress3mini.channels[1].reset_ticks, 101)

with patch(
"hyperion.experiment_plans.optimise_attenuation_plan.is_deadtime_optimised",
@@ -220,8 +222,9 @@ def test_total_count_exception_raised_after_max_cycles_reached(
RE: RunEngine, fake_composite_mocked_sets: OptimizeAttenuationComposite
):
optimise_attenuation_plan.is_counts_within_target = MagicMock(return_value=False)
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca.sim_put( # type: ignore
[1, 1, 1, 1, 1, 1]
set_mock_value(
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca[1],
np.array([1, 1, 1, 1, 1, 1]),
)
with pytest.raises(AttenuationOptimisationFailedException):
RE(
@@ -231,15 +234,6 @@ def test_total_count_exception_raised_after_max_cycles_reached(
)


def test_arm_devices_runs_correct_functions(
RE: RunEngine, fake_composite: OptimizeAttenuationComposite
):
fake_composite.xspress3mini.detector_state.sim_put("Acquire") # type: ignore
fake_composite.xspress3mini.arm = MagicMock(return_value=get_good_status())
RE(arm_devices(fake_composite.xspress3mini))
fake_composite.xspress3mini.arm.assert_called_once()


@pytest.mark.parametrize(
"direction, transmission, increment, upper_limit, lower_limit, new_transmission",
[
@@ -267,8 +261,9 @@ def test_deadtime_calc_new_transmission_raises_error_on_low_ransmission():
def test_total_count_calc_new_transmission_raises_error_on_low_ransmission(
RE: RunEngine, fake_composite_mocked_sets: OptimizeAttenuationComposite
):
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca.sim_put( # type: ignore
[1, 1, 1, 1, 1, 1]
set_mock_value(
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca[1],
np.array([1, 1, 1, 1, 1, 1]),
)
with pytest.raises(AttenuationOptimisationFailedException):
RE(
@@ -287,18 +282,17 @@ def test_total_count_calc_new_transmission_raises_error_on_low_ransmission(
)


@patch("hyperion.experiment_plans.optimise_attenuation_plan.arm_devices", autospec=True)
def test_total_counts_gets_within_target(
mock_arm_devices,
RE: RunEngine,
fake_composite_mocked_sets: OptimizeAttenuationComposite,
):
# For simplicity we just increase the data array each iteration. In reality it's the transmission value that affects the array
def update_data(_):
nonlocal iteration
iteration += 1
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca.sim_put( # type: ignore
([50, 50, 50, 50, 50]) * iteration
set_mock_value(
fake_composite_mocked_sets.xspress3mini.dt_corrected_latest_mca[1],
np.array(([50, 50, 50, 50, 50]) * iteration),
)
return get_good_status()

0 comments on commit 4e36438

Please sign in to comment.