import uuid
from collections.abc import Hashable
from typing import Any
import bluesky.plan_stubs as bps
from bluesky.plan_stubs import abs_set
from bluesky.protocols import Movable
from bluesky.utils import MsgGenerator, plan
from dodal.devices.slits import Slits
from ophyd_async.core import SignalRW, StandardReadable
from ophyd_async.epics.motor import Motor
from pydantic import RootModel
from sm_bluesky.log import LOGGER
[docs]
class MotorTable(RootModel):
"""RootModel for motor tables"""
root: dict[str, float]
class HighLowLimitsDevice(StandardReadable):
low_limit_travel: SignalRW[float]
high_limit_travel: SignalRW[float]
[docs]
@plan
def move_motor_with_look_up(
slit: Movable[float],
size: float,
motor_table: dict[str, float],
use_motor_position: bool = False,
wait: bool = True,
group: Hashable | None = None,
) -> MsgGenerator[None]:
"""Perform a step scan with the range and starting motor position
given/calculated by using a look up table(dictionary).
Move to the peak position after the scan and update the lookup table.
Parameters
----------
motor: Motor
Motor devices that is being centre.
size: float
The motor position or name in the motor_table.
motor_table: dict[str, float],
Look up table for motor position,
use_motor_position: bool = False,
If Ture it will take motor position as size.
wait: bool = True,
If Ture, it will wait until position is reached.
group: Hashable | None = None,
Bluesky group identifier used by ‘wait’.
"""
MotorTable.model_validate(motor_table)
if use_motor_position:
yield from abs_set(slit, size, wait=wait, group=group)
elif str(int(size)) in motor_table:
yield from abs_set(slit, motor_table[str(int(size))], wait=wait, group=group)
else:
raise ValueError(
f"No slit with size={size}. Available slit size: {motor_table}"
)
[docs]
@plan
def set_slit_size(
xy_slit: Slits,
x_size: float,
y_size: float | None = None,
wait: bool = True,
group: Hashable | None = None,
) -> MsgGenerator[None]:
"""Set opening of x-y slit.
Parameters
----------
xy_slit: Slits
A slits device.
x_size: float
The x opening size.
y_size: float
The y opening size.
wait: bool
If this is True it will wait for all motions to finish.
group (optional): Hashable
Bluesky group identifier used by ‘wait’.
"""
if wait and group is None:
group = f"{xy_slit.name}_wait"
if y_size is None:
y_size = x_size
LOGGER.info(f"Setting {xy_slit.name} to x = {x_size}, y = {y_size}.")
yield from bps.abs_set(xy_slit.x_gap, x_size, group=group)
yield from bps.abs_set(xy_slit.y_gap, y_size, group=group)
if wait:
LOGGER.info(f"Waiting for {xy_slit.name} to finish move.")
yield from bps.wait(group=group)
[docs]
@plan
def check_within_limit(
values: list[float], device: HighLowLimitsDevice | Motor
) -> MsgGenerator[None]:
"""Check if the given values are within the limits of the device.
Parameters
----------
values : List[float]
The values to check.
device : HighLowLimitsDevice
The device to check the limits of.
Raises
------
ValueError
If any value is outside the device's limits.
"""
LOGGER.info(f"Check {device.name} limits.")
lower_limit = yield from bps.rd(device.low_limit_travel)
high_limit = yield from bps.rd(device.high_limit_travel)
for value in values:
if not lower_limit < value < high_limit:
raise ValueError(
f"{device.name} move request of {value} is beyond limits:"
f"{lower_limit} < {high_limit}"
)
[docs]
@plan
def get_motor_positions(*arg: Motor) -> MsgGenerator:
"""
Get the motor positions of the given motors and store them in a list.
Parameters
----------
arg : Motor
The motors to get the positions of.
Returns
-------
Iterator[Tuple[str, float]]
An iterator of tuples containing the motor name and its position.
"""
motor_position = []
for motor in arg:
motor_position.append(motor)
position = yield from bps.rd(motor)
motor_position.append(position)
LOGGER.info(f"Stored motor, position = {motor_position}.")
return motor_position
[docs]
@plan
def get_velocity_and_step_size(
scan_motor: Motor, ideal_velocity: float, ideal_step_size: float
) -> MsgGenerator[Any]:
"""
Adjust the step size if the required velocity is higher than the max value.
Parameters
----------
scan_motor : Motor
The motor which will move continuously.
ideal_velocity : float
The desired velocity.
ideal_step_size : float
The non-scanning motor step size.
Returns
-------
Iterator[Tuple[float, float]]
An iterator containing the adjusted velocity and step size.
"""
if ideal_velocity <= 0.0:
raise ValueError(f"{scan_motor.name} speed: {ideal_velocity} <= 0")
max_velocity = yield from bps.rd(scan_motor.max_velocity) # type: ignore
# if motor does not move fast enough increase step_motor step size
if ideal_velocity > max_velocity:
ideal_step_size = ideal_step_size * (ideal_velocity / max_velocity)
ideal_velocity = round(max_velocity, 3)
return ideal_velocity, ideal_step_size
[docs]
@plan
def cache_speed(
motor_and_speeds: list[Motor],
) -> MsgGenerator[dict[Motor, float]]:
"""Cache the current velocity of each motor.
Parameters
----------
motor_and_speeds : list[Motor]
List of motor devices whose velocity should be cached.
Returns
-------
dict[Motor, float]
Mapping of each motor to its current velocity.
"""
speeds = {}
for axis in motor_and_speeds:
speed = yield from bps.rd(axis.velocity)
speeds[axis] = speed
return speeds
[docs]
@plan
def restore_speed(
motor_and_speeds: dict[Motor, float],
group: str | None = None,
wait_for_all: bool = True,
) -> MsgGenerator:
"""Restore cached velocities for motors.
Parameters
----------
motor_and_speeds : dict[Motor, float]
Mapping of motor devices to the velocity values to restore.
group : str | None, optional
Optional Bluesky group identifier used during the restore moves.
If omitted, a unique reset group name is generated.
wait_for_all : bool, optional
If True, wait for all velocity restore operations to complete.
"""
reset_group = f"reset-{group if group else str(uuid.uuid4())[:6]}"
for device, speed in motor_and_speeds.items():
yield from bps.abs_set(device.velocity, speed, group=reset_group)
if wait_for_all:
yield from bps.wait(reset_group)
