mx-bluesky 1.5.11__py3-none-any.whl → 1.5.14__py3-none-any.whl

mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py

@@ -17,6 +17,7 @@ from dodal.devices.fast_grid_scan import (
     set_fast_grid_scan_params,
 )
 from dodal.devices.flux import Flux
+from dodal.devices.i04.beamsize import Beamsize
 from dodal.devices.i04.transfocator import Transfocator
 from dodal.devices.mx_phase1.beamstop import Beamstop
 from dodal.devices.oav.oav_detector import OAV
@@ -25,7 +26,7 @@ from dodal.devices.robot import BartRobot
 from dodal.devices.s4_slit_gaps import S4SlitGaps
 from dodal.devices.smargon import Smargon
 from dodal.devices.synchrotron import Synchrotron
-from dodal.devices.undulator import Undulator
+from dodal.devices.undulator import UndulatorInKeV
 from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra.zebra import Zebra
 from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
@@ -33,7 +34,11 @@ from dodal.devices.zocalo import ZocaloResults
 from dodal.plans.preprocessors.verify_undulator_gap import (
     verify_undulator_gap_before_run_decorator,
 )
+from pydantic import BaseModel
 
+from mx_bluesky.beamlines.i04.external_interaction.config_server import (
+    get_i04_config_client,
+)
 from mx_bluesky.common.device_setup_plans.setup_zebra_and_shutter import (
     setup_zebra_for_gridscan,
     tidy_up_zebra_after_gridscan,
@@ -58,6 +63,7 @@ from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
     GridscanNexusFileCallback,
 )
+from mx_bluesky.common.parameters.components import PARAMETER_VERSION
 from mx_bluesky.common.parameters.constants import (
     EnvironmentConstants,
     OavConstants,
@@ -67,13 +73,28 @@ from mx_bluesky.common.parameters.constants import (
 from mx_bluesky.common.parameters.device_composites import (
     GridDetectThenXRayCentreComposite,
 )
-from mx_bluesky.common.parameters.gridscan import GridCommon, SpecifiedThreeDGridScan
+from mx_bluesky.common.parameters.gridscan import (
+    GridCommon,
+    SpecifiedThreeDGridScan,
+)
 from mx_bluesky.common.preprocessors.preprocessors import (
     transmission_and_xbpm_feedback_for_collection_decorator,
 )
+from mx_bluesky.common.utils.exceptions import CrystalNotFoundError
 from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.common.utils.utils import (
+    fix_transmission_and_exposure_time_for_current_wavelength,
+)
 
-DEFAULT_BEAMSIZE_MICRONS = 20
+DEFAULT_XRC_BEAMSIZE_MICRONS = 20
+
+
+class I04AutoXrcParams(BaseModel):
+    sample_id: int
+    file_name: str
+    visit: str
+    detector_distance_mm: float
+    storage_directory: str
 
 
 def _change_beamsize(
@@ -91,19 +112,20 @@ def _change_beamsize(
 
 
 # See https://github.com/DiamondLightSource/blueapi/issues/506 for using device composites
-def i04_grid_detect_then_xray_centre(
-    parameters: GridCommon,
+def i04_default_grid_detect_and_xray_centre(
+    parameters: I04AutoXrcParams,
     aperture_scatterguard: ApertureScatterguard = inject("aperture_scatterguard"),
     attenuator: BinaryFilterAttenuator = inject("attenuator"),
     backlight: Backlight = inject("backlight"),
     beamstop: Beamstop = inject("beamstop"),
+    beamsize: Beamsize = inject("beamsize"),
     dcm: DoubleCrystalMonochromator = inject("dcm"),
     zebra_fast_grid_scan: ZebraFastGridScanThreeD = inject("zebra_fast_grid_scan"),
     flux: Flux = inject("flux"),
     oav: OAV = inject("oav"),
     pin_tip_detection: PinTipDetection = inject("pin_tip_detection"),
     s4_slit_gaps: S4SlitGaps = inject("s4_slit_gaps"),
-    undulator: Undulator = inject("undulator"),
+    undulator: UndulatorInKeV = inject("undulator"),
     xbpm_feedback: XBPMFeedback = inject("xbpm_feedback"),
     zebra: Zebra = inject("zebra"),
     robot: BartRobot = inject("robot"),
@@ -124,7 +146,6 @@ def i04_grid_detect_then_xray_centre(
     - Changes the aperture to match the beam size to the crystal size
     - Moves the sample to the crystal centre of mass
 
-
     i04's implementation of this plan is very similar to Hyperion. However, since i04
     isn't running in a continuous Bluesky UDC loop, we take additional steps in beamline
     tidy-up.
@@ -139,6 +160,7 @@ def i04_grid_detect_then_xray_centre(
         attenuator,
         backlight,
         beamstop,
+        beamsize,
         dcm,
         detector_motion,
         zebra_fast_grid_scan,
@@ -152,9 +174,18 @@ def i04_grid_detect_then_xray_centre(
         robot,
         sample_shutter,
     )
-    initial_beamsize = yield from bps.rd(transfocator.beamsize_set_microns)
+    initial_beamsize = yield from bps.rd(transfocator.current_vertical_size_rbv)
+
+    initial_x = yield from bps.rd(smargon.x.user_readback)
+    initial_y = yield from bps.rd(smargon.y.user_readback)
+    initial_z = yield from bps.rd(smargon.z.user_readback)
+
+    _current_wavelength_a = yield from bps.rd(composite.dcm.wavelength_in_a)
+    grid_common_params = _get_grid_common_params(_current_wavelength_a, parameters)
 
     def tidy_beamline():
+        yield from bps.mv(transfocator, initial_beamsize)
+
         if not udc:
             yield from get_ready_for_oav_and_close_shutter(
                 composite.smargon,
@@ -162,7 +193,6 @@ def i04_grid_detect_then_xray_centre(
                 composite.aperture_scatterguard,
                 composite.detector_motion,
             )
-        yield from bps.mv(transfocator, initial_beamsize)
 
     @bpp.finalize_decorator(tidy_beamline)
     def _inner_grid_detect_then_xrc():
@@ -174,20 +204,30 @@ def i04_grid_detect_then_xray_centre(
         @bpp.subs_decorator(callbacks)
         @verify_undulator_gap_before_run_decorator(composite)
         @transmission_and_xbpm_feedback_for_collection_decorator(
-            composite, parameters.transmission_frac, PlanNameConstants.GRIDSCAN_OUTER
+            composite,
+            grid_common_params.transmission_frac,
+            PlanNameConstants.GRIDSCAN_OUTER,
         )
         def grid_detect_then_xray_centre_with_callbacks():
             yield from grid_detect_then_xray_centre(
                 composite=composite,
-                parameters=parameters,
+                parameters=grid_common_params,
                 xrc_params_type=SpecifiedThreeDGridScan,
                 construct_beamline_specific=construct_i04_specific_features,
                 oav_config=oav_config,
             )
 
-        yield from grid_detect_then_xray_centre_with_callbacks()
+        try:
+            yield from grid_detect_then_xray_centre_with_callbacks()
+        except CrystalNotFoundError:
+            yield from bps.mv(
+                smargon.x, initial_x, smargon.y, initial_y, smargon.z, initial_z
+            )
+            raise
 
-    yield from _change_beamsize(transfocator, DEFAULT_BEAMSIZE_MICRONS, parameters)
+    yield from _change_beamsize(
+        transfocator, DEFAULT_XRC_BEAMSIZE_MICRONS, grid_common_params
+    )
     yield from _inner_grid_detect_then_xrc()
 
 
@@ -252,6 +292,7 @@ def construct_i04_specific_features(
         xrc_composite.flux.flux_reading,
         xrc_composite.dcm.energy_in_keV,
         xrc_composite.eiger.bit_depth,
+        xrc_composite.beamsize,
     ]
 
     tidy_plan = partial(
@@ -278,3 +319,37 @@ def construct_i04_specific_features(
         signals_to_read_during_collection,
         get_xrc_results_from_zocalo=True,
     )
+
+
+def _get_grid_common_params(
+    _current_wavelength_a: float, parameters: I04AutoXrcParams
+) -> GridCommon:
+    """Calculate scaled transmission and exposure by comparing current beamline energy to default energy"""
+    _assumed_wavelength_a = (
+        get_i04_config_client().get_feature_flags().ASSUMED_WAVELENGTH_IN_A
+    )
+    _unscaled_transmission = (
+        get_i04_config_client().get_feature_flags().XRC_UNSCALED_TRANSMISSION_FRAC
+    )
+    _unscaled_exposure_time_s = (
+        get_i04_config_client().get_feature_flags().XRC_UNSCALED_EXPOSURE_TIME_S
+    )
+    transmission_frac, exposure_time_s = (
+        fix_transmission_and_exposure_time_for_current_wavelength(
+            _current_wavelength_a,
+            _assumed_wavelength_a,
+            _unscaled_transmission,
+            _unscaled_exposure_time_s,
+        )
+    )
+
+    return GridCommon(
+        sample_id=parameters.sample_id,
+        file_name=parameters.file_name,
+        visit=parameters.visit,
+        detector_distance_mm=parameters.detector_distance_mm,
+        storage_directory=parameters.storage_directory,
+        transmission_frac=transmission_frac,
+        exposure_time_s=exposure_time_s,
+        parameter_model_version=PARAMETER_VERSION,
+    )

mx_bluesky/beamlines/i04/external_interaction/config_server.py

@@ -0,0 +1,15 @@
+from functools import cache
+
+from mx_bluesky.beamlines.i04.parameters.constants import (
+    I04FeatureSettings,
+    I04FeatureSettingsSources,
+)
+from mx_bluesky.common.external_interaction.config_server import MXConfigClient
+
+
+@cache
+def get_i04_config_client() -> MXConfigClient[I04FeatureSettings]:
+    return MXConfigClient(
+        feature_sources=I04FeatureSettingsSources,
+        feature_dc=I04FeatureSettings,
+    )

mx_bluesky/beamlines/i04/oav_centering_plans/oav_imaging.py

@@ -0,0 +1,115 @@
+import os
+import time
+
+import bluesky.plan_stubs as bps
+from bluesky.utils import MsgGenerator
+from dodal.common import inject
+from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
+from dodal.devices.backlight import Backlight
+from dodal.devices.mx_phase1.beamstop import Beamstop, BeamstopPositions
+from dodal.devices.oav.oav_detector import OAV
+from dodal.devices.robot import BartRobot, PinMounted
+from dodal.devices.scintillator import InOut, Scintillator
+from dodal.devices.xbpm_feedback import XBPMFeedback
+from dodal.devices.zebra.zebra_controlled_shutter import (
+    ZebraShutter,
+    ZebraShutterControl,
+    ZebraShutterState,
+)
+from ophyd_async.core import InOut as core_INOUT
+
+from mx_bluesky.common.utils.exceptions import BeamlineStateError
+
+initial_wait_group = "Wait for scint to move in"
+
+
+def take_oav_image_with_scintillator_in(
+    image_name: str | None = None,
+    image_path: str = "dls_sw/i04/software/bluesky/scratch",
+    transmission: float = 1,
+    attenuator: BinaryFilterAttenuator = inject("attenuator"),
+    shutter: ZebraShutter = inject("sample_shutter"),
+    oav: OAV = inject("oav"),
+    robot: BartRobot = inject("robot"),
+    beamstop: Beamstop = inject("beamstop"),
+    backlight: Backlight = inject("backlight"),
+    scintillator: Scintillator = inject("scintillator"),
+    xbpm_feedback: XBPMFeedback = inject("xbpm_feedback"),
+) -> MsgGenerator:
+    """
+    Takes an OAV image at specified transmission after necessary checks and preparation steps.
+
+    Args:
+        image_name: Name of the OAV image to be saved
+        image_path: Path where the image should be saved
+        transmission: Transmission of the beam, takes a value from 0 to 1 where
+                    1 lets all the beam through and 0 lets none of the beam through.
+        devices: These are the specific ophyd-devices used for the plan, the
+                    defaults are always correct.
+    """
+
+    yield from _prepare_beamline_for_scintillator_images(
+        robot, beamstop, backlight, scintillator, xbpm_feedback, initial_wait_group
+    )
+
+    yield from bps.abs_set(attenuator, transmission, group=initial_wait_group)
+
+    if image_name is None:
+        image_name = f"{time.time_ns()}ATT{transmission * 100}"
+
+    yield from bps.wait(initial_wait_group)
+
+    yield from bps.abs_set(shutter.control_mode, ZebraShutterControl.MANUAL, wait=True)
+    yield from bps.abs_set(shutter, ZebraShutterState.OPEN, wait=True)
+
+    take_and_save_oav_image(file_path=image_path, file_name=image_name, oav=oav)
+
+
+def _prepare_beamline_for_scintillator_images(
+    robot: BartRobot,
+    beamstop: Beamstop,
+    backlight: Backlight,
+    scintillator: Scintillator,
+    xbpm_feedback: XBPMFeedback,
+    group: str,
+) -> MsgGenerator:
+    """
+    Prepares the beamline for oav image by making sure the pin is NOT mounted and
+    the beam is on (feedback check). Finally, the scintillator is moved in.
+    """
+    pin_mounted = yield from bps.rd(robot.gonio_pin_sensor)
+    if pin_mounted == PinMounted.PIN_MOUNTED:
+        raise BeamlineStateError("Pin should not be mounted!")
+
+    yield from bps.trigger(xbpm_feedback, group=group)
+
+    yield from bps.abs_set(
+        beamstop.selected_pos, BeamstopPositions.DATA_COLLECTION, group=group
+    )
+
+    yield from bps.abs_set(backlight, core_INOUT.OUT, group=group)
+
+    yield from bps.abs_set(scintillator.selected_pos, InOut.IN, group=group)
+
+
+def take_and_save_oav_image(
+    file_name: str,
+    file_path: str,
+    oav: OAV,
+) -> MsgGenerator:
+    """
+    Plan which takes and saves an OAV image to the specified path.
+     Args:
+        file_name: Filename specifying the name of the image,
+        file_path: Path as a string specifying where the image should be saved,
+        oav: The OAV to take the image with
+    """
+    group = "oav image path setting"
+    full_file_path = file_path + "/" + file_name
+    if not os.path.exists(full_file_path):
+        yield from bps.abs_set(oav.snapshot.filename, file_name, group=group)
+        yield from bps.abs_set(oav.snapshot.directory, file_path, group=group)
+        yield from bps.wait(group)
+        yield from bps.trigger(oav.snapshot, wait=True)
+    else:
+        raise FileExistsError("OAV image file path already exists")

mx_bluesky/beamlines/i04/parameters/constants.py

@@ -0,0 +1,21 @@
+from pydantic.dataclasses import dataclass
+
+from mx_bluesky.common.parameters.constants import (
+    FeatureSettings,
+    FeatureSettingSources,
+)
+
+
+# These currently exist in GDA domain.properties
+class I04FeatureSettingsSources(FeatureSettingSources):
+    ASSUMED_WAVELENGTH_IN_A = "gda.px.expttable.default.wavelength"
+    XRC_UNSCALED_TRANSMISSION_FRAC = "gda.mx.bluesky.i04.xrc.unscaled_transmission_frac"
+    XRC_UNSCALED_EXPOSURE_TIME_S = "gda.mx.bluesky.i04.xrc.unscaled_exposure_time_s"
+
+
+# Use these defaults if we can't read from the config server
+@dataclass
+class I04FeatureSettings(FeatureSettings):
+    ASSUMED_WAVELENGTH_IN_A: float = 0.95373
+    XRC_UNSCALED_TRANSMISSION_FRAC: int = 1
+    XRC_UNSCALED_EXPOSURE_TIME_S: float = 0.007

mx_bluesky/beamlines/i04/redis_to_murko_forwarder.py

@@ -1,21 +1,27 @@
 import io
 import json
+import logging
 import pickle
 from datetime import timedelta
+from logging import StreamHandler
 from typing import TypedDict
 
 import numpy as np
 import zmq
 from dodal.devices.i04.constants import RedisConstants
-from dodal.devices.i04.murko_results import MurkoResult
+from dodal.devices.i04.murko_results import RESULTS_COMPLETE_MESSAGE, MurkoResult
 from numpy.typing import NDArray
 from PIL import Image
 from redis import StrictRedis
 
+from mx_bluesky.beamlines.i04.callbacks.murko_callback import (
+    FORWARDING_COMPLETE_MESSAGE,
+)
 from mx_bluesky.common.utils.log import LOGGER
 
 MURKO_ADDRESS = "tcp://i04-murko-prod.diamond.ac.uk:8008"
 
+
 FullMurkoResults = dict[str, list[MurkoResult]]
 
 
@@ -112,6 +118,12 @@ class BatchMurkoForwarder:
             self.redis_client.expire(redis_key, timedelta(days=7))
         self.redis_client.publish("murko-results", pickle.dumps(results))
 
+    def send_stop_message_to_redis(self):
+        LOGGER.info(f"Publishing results complete message: {RESULTS_COMPLETE_MESSAGE}")
+        self.redis_client.publish(
+            "murko-results", pickle.dumps(RESULTS_COMPLETE_MESSAGE)
+        )
+
     def add(self, sample_id: str, uuid: str, image: NDArray):
         """Add an image to the batch to send to murko."""
         image_size = get_image_size(image)
@@ -159,6 +171,13 @@ class RedisListener:
         if message and message["type"] == "message":
             data = json.loads(message["data"])
             LOGGER.info(f"Received from redis: {data}")
+            if data == FORWARDING_COMPLETE_MESSAGE:
+                LOGGER.info(
+                    f"Received forwarding complete message: {FORWARDING_COMPLETE_MESSAGE}"
+                )
+                self.forwarder.flush()
+                self.forwarder.send_stop_message_to_redis()
+                return
             uuid = data["uuid"]
             sample_id = data["sample_id"]
 
@@ -188,6 +207,10 @@ class RedisListener:
 
 
 def main():
+    stream_handler = StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    LOGGER.addHandler(stream_handler)
+
     client = RedisListener()
     client.listen_for_image_data_forever()