diff --git a/mx_bec/devices/beam_profile.py b/mx_bec/devices/beam_profile.py new file mode 100644 index 0000000..2157376 --- /dev/null +++ b/mx_bec/devices/beam_profile.py @@ -0,0 +1,117 @@ +"""A device for determining beam shape and location from sample camera image analysis.""" + +from abc import ABC, abstractmethod +from typing import Literal + +import numpy as np +from ophyd import Component as Cpt +from ophyd import EpicsSignal, EpicsSignalRO, Kind +from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase + + +class BeamProfile(PSIDeviceBase, ABC): + """Common interface for different methods of getting beam location and size""" + + # Readback + x_pos_px = Cpt[EpicsSignalRO] + y_pos_px = Cpt[EpicsSignalRO] + x_sig_px = Cpt[EpicsSignalRO] + y_sig_px = Cpt[EpicsSignalRO] + + @abstractmethod + def prepare_plugin(self): + """Set parameters, etc., which only need to be set up once before using + the profile results (e.g. threshold value...)""" + + @abstractmethod + def enable_computation(self): + """Ensure plugins or whatever else is needed is configured to provide + analysis results""" + + @abstractmethod + def disable_computation(self): + """Save resources by turning computations off again""" + + +class GaussianBeamProfile(BeamProfile): + """Use image analysis of the scintillator to determine the beam centre and width on the sample + camera image, by fitting a Gaussian to the image profile. + The analysis is provided by AD plugins in EPICS, we merely configure it and read the results.""" + + # Config + cbs_enabled = Cpt(EpicsSignal, name="cbs_enabled", suffix="EnableCallbacks", kind=Kind.config) + x_compute = Cpt(EpicsSignal, name="x_compute", suffix="X:Compute", kind=Kind.config) + y_compute = Cpt(EpicsSignal, name="y_compute", suffix="X:Compute", kind=Kind.config) + + # Readback + x_pos_px = Cpt(EpicsSignalRO, name="x_pos_px", suffix="X:Mu_RBV", kind=Kind.normal) + y_pos_px = Cpt(EpicsSignalRO, name="y_pos_px", suffix="Y:Mu_RBV", kind=Kind.normal) + x_sig_px = Cpt(EpicsSignalRO, name="x_sig_px", suffix="X:Sigma_RBV", kind=Kind.normal) + y_sig_px = Cpt(EpicsSignalRO, name="x_sig_px", suffix="Y:Sigma_RBV", kind=Kind.normal) + + def prepare_plugin(self): ... + + def enable_computation(self): + """Ensure all the configuration parameters are set.""" + # TODO: add making sure the correct plugins and callbacks are wired together. + # TODO: make sure we are at max zoom + self._set_computation(1).wait() + + def disable_computation(self): + """Ensure all the image analysis computations are disabled.""" + self._set_computation(0).wait() + + def _set_computation(self, enabled: Literal[0, 1]): + st1 = self.cbs_enabled.set(enabled) + st2 = self.x_compute.set(enabled) + st3 = self.y_compute.set(enabled) + return st1 and st2 and st3 + + +class CentroidBeamProfile(BeamProfile): + """Use image analysis of the scintillator to determine the beam centre and width on the sample + camera image, by finding the pixels over a threshold in a histogram of each direction. + The analysis is provided by AD plugins in EPICS, we merely configure it and read the results.""" + + # Config + histo_enabled = Cpt( + EpicsSignal, name="histo_enabled", suffix="ComputeHistogram", kind=Kind.config + ) + centroid_enabled = Cpt( + EpicsSignal, name="centroid_enabled", suffix="ComputeCentroid", kind=Kind.config + ) + centroid_thresh = Cpt( + EpicsSignal, name="centroid_thresh", suffix="CentroidThreshold", kind=Kind.config + ) + + # Data for setup + histo_arr = Cpt(EpicsSignal, name="histo_arr", suffix="Histogram_RBV", kind=Kind.omitted) + + def prepare_plugin(self): + self._update_threshold() + + def enable_computation(self): + """Ensure all the configuration parameters are set.""" + # TODO: add making sure the correct plugins and callbacks are wired together. + # TODO: make sure we are at max zoom + self._set_computation(1).wait() + + def disable_computation(self): + """Ensure all the image analysis computations are disabled.""" + self._set_computation(0).wait() + + def _update_threshold(self): + self.histo_enabled.set(1).wait() + histogram: np.typing.NDArray[np.int64] = self.histo_arr.get() # type: ignore + cs = np.cumsum(histogram) + cs_thresh = np.sum(histogram) * 0.97 + exceeds_thresh_indices = np.where(cs > cs_thresh) + self.centroid_thresh.set(exceeds_thresh_indices[0][0]).wait() + + def _set_computation(self, enabled: Literal[0, 1]): + return self.centroid_enabled.set(enabled) + + x_pos_px = Cpt(EpicsSignalRO, name="x_pos_px", suffix="CentroidX_RBV", kind=Kind.normal) + y_pos_px = Cpt(EpicsSignalRO, name="y_pos_px", suffix="CentroidY_RBV", kind=Kind.normal) + x_sig_px = Cpt(EpicsSignalRO, name="x_sig_px", suffix="SigmaX_RBV", kind=Kind.normal) + y_sig_px = Cpt(EpicsSignalRO, name="x_sig_px", suffix="SigmaY_RBV", kind=Kind.normal) diff --git a/mx_bec/devices/beam_steering.py b/mx_bec/devices/beam_steering.py new file mode 100644 index 0000000..fd997be --- /dev/null +++ b/mx_bec/devices/beam_steering.py @@ -0,0 +1,144 @@ +"""A device for stepping the beam towards the centre""" + +from __future__ import annotations + +import threading +from typing import Literal, cast + +from bec_lib.logger import bec_logger +from ophyd import Component as Cpt +from ophyd import DeviceStatus, Kind, Signal +from ophyd_devices import EpicsMotorEC +from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase + +from mx_bec.devices.beam_profile import BeamProfile + +logger = bec_logger.logger + + +class BeamSteerer(PSIDeviceBase): + """Based on the beam profile, step towards the centre by moving the focussing mirrors. + Assumes everything is done at max zoom.""" + + def __init__( + self, + *, + name: str, + prefix: str = "", + scan_info: ScanInfo | None = None, + device_manager: DeviceManagerBase | None = None, + **kwargs, + ): + super().__init__( + name=name, prefix=prefix, scan_info=scan_info, device_manager=device_manager, **kwargs + ) + self._busy = threading.Lock() + self._abort = threading.Event() + + # Config + hfm_step_per_x_px = Cpt(Signal, name="hfm_step_per_x_px", kind=Kind.config) + vfm_step_per_y_px = Cpt(Signal, name="hfm_step_per_y_px", kind=Kind.config) + sample_loc_x_px = Cpt(Signal, name="sample_loc_x_px", kind=Kind.config) + sample_loc_y_px = Cpt(Signal, name="sample_loc_y_px", kind=Kind.config) + px_tolerance = Cpt(Signal, name="px_tolerance", kind=Kind.config) + step_fraction = Cpt(Signal, name="step_fraction", kind=Kind.config) + + def _px_from_centre(self, bp: BeamProfile, axis: Literal["x", "y"]): + sample_loc = self.sample_loc_x_px if axis == "x" else self.sample_loc_y_px + curr_loc = bp.x_pos_px if axis == "x" else bp.y_pos_px + return float(sample_loc.get() - curr_loc.get()) # type: ignore + + def step_towards_centre(self): + """Take a step towards the beam centre by adjusting the HFM and VFM. Uses BeamProfile to + determine the x and y location on the samcam in px, the step_per_[x,y]_px parameters to + determine how far to move the mirrors, and multiplies this by the step fraction.""" + if self.device_manager is None: + raise RuntimeError( + "This can only be called in a connected BEC session with access to other devices." + ) + + hfm_motor = cast(EpicsMotorEC | None, self.device_manager.devices.get("hfm_yr")) + vfm_motor = cast(EpicsMotorEC | None, self.device_manager.devices.get("vfm_yw")) + profile_device = self.device_manager.devices.get("beam_steering").user_parameter.get( + "profile_provider" + ) + bp = cast(BeamProfile | None, self.device_manager.devices.get(profile_device)) + if hfm_motor is None or vfm_motor is None or bp is None: + raise RuntimeError( + "This device needs the focussing mirrors and the beam profile, " + "supply them in 'needs' in the BEC device config." + ) + logger.info( + f"Beam steering stepping towards target {self.sample_loc_x_px.get()}, {self.sample_loc_y_px.get()} (x, y px on samcam image)..." + ) + return self._step_towards_centre(hfm_motor, vfm_motor, bp) + + def _step_towards_centre( + self, hfm_motor: EpicsMotorEC, vfm_motor: EpicsMotorEC, beam_profile: BeamProfile + ): + x_dist = self._px_from_centre(beam_profile, "x") + y_dist = self._px_from_centre(beam_profile, "y") + tol = float(self.px_tolerance.get()) # type: ignore + damp_factor = float(self.step_fraction.get()) # type: ignore + hfm_step_per_x_px = float(self.hfm_step_per_x_px.get()) # type: ignore + vfm_step_per_y_px = float(self.vfm_step_per_y_px.get()) # type: ignore + logger.info(f"Beam steering - samcam distance: {x_dist=} px, {y_dist=} px") + + if abs(x_dist) > tol: + x_full_step = x_dist * hfm_step_per_x_px + x_dampened_step = x_full_step * damp_factor + new_x_pos = hfm_motor.user_setpoint.get() + x_dampened_step # type: ignore + logger.info(f"Moving HFM_YR {x_dampened_step} units to {new_x_pos}") + x_status = hfm_motor.move(new_x_pos) + else: + x_status = None + + if abs(y_dist) > tol: + y_full_step = y_dist * vfm_step_per_y_px + y_dampened_step = y_full_step * damp_factor + new_y_pos = vfm_motor.user_setpoint.get() + y_dampened_step # type: ignore + logger.info(f"Moving VFM_YW {y_dampened_step} units to {new_y_pos}") + y_status = vfm_motor.move(new_y_pos) + else: + y_status = None + + if x_status is not None and y_status is not None: + return x_status and y_status + if x_status is not None: + return x_status + if y_status is not None: + return y_status + return None + + def _iterate_steps(self, status): + logger.info("Running beam steering...") + step_count = 1 + try: + while (_step_status := self.step_towards_centre()) is not None: + logger.info(f"Beam steering iteration: {step_count}") + step_count += 1 + if self._abort.is_set(): + raise RuntimeError(f"{self.name}: aborted") + _step_status.wait() + except Exception as exc: + status.set_exception(exc) + else: + status.set_finished() + finally: + self._busy.release() + + def trigger(self): + """External interface for to iterate stepping towards the centre until finished""" + + if not self._busy.acquire(blocking=False): + raise RuntimeError(f"{self.name} is still busy") + self._abort.clear() + status = DeviceStatus(self, timeout=60, settle_time=0.0) + threading.Thread( + target=self._iterate_steps, args=(status,), daemon=True, name=f"{self.name}_trigger" + ).start() + return status + + def stop(self, *, success=False): + self._abort.set() + super().stop(success=success)