mx-bluesky 1.4.1a0__py3-none-any.whl → 1.4.3__py3-none-any.whl

mx_bluesky/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.4.1a0'
-__version_tuple__ = version_tuple = (1, 4, 1)
+__version__ = version = '1.4.3'
+__version_tuple__ = version_tuple = (1, 4, 3)

mx_bluesky/beamlines/i04/redis_to_murko_forwarder.py

@@ -0,0 +1,178 @@
+import io
+import json
+import pickle
+from typing import TypedDict
+
+import numpy as np
+import zmq
+from dodal.beamlines.i04 import MURKO_REDIS_DB, REDIS_HOST, REDIS_PASSWORD
+from numpy.typing import NDArray
+from PIL import Image
+from redis import StrictRedis
+
+from mx_bluesky.common.utils.log import LOGGER
+
+MURKO_ADDRESS = "tcp://i04-murko-prod.diamond.ac.uk:8008"
+
+
+class MurkoRequest(TypedDict):
+    to_predict: NDArray
+    model_img_size: tuple[int, int]
+    save: bool
+    min_size: int
+    description: list
+    prefix: list[str]
+
+
+def get_image_size(image: NDArray) -> tuple[int, int]:
+    """Returns the width and height of a numpy image"""
+    return image.shape[1], image.shape[0]
+
+
+def send_to_murko_and_get_results(request: MurkoRequest) -> dict:
+    LOGGER.info(f"Sending {request['prefix']} to murko")
+    context = zmq.Context()
+    socket = context.socket(zmq.REQ)
+    socket.connect(MURKO_ADDRESS)
+    socket.send(pickle.dumps(request))
+    raw_results = socket.recv()
+    results = pickle.loads(raw_results)
+    LOGGER.info(f"Got {len(results['descriptions'])} results")
+    return results
+
+
+def correlate_results_to_uuids(request: MurkoRequest, murko_results: dict) -> list:
+    results = []
+    uuids = request["prefix"]
+
+    width, height = get_image_size(request["to_predict"][0])
+
+    for uuid, prediction in zip(uuids, murko_results["descriptions"], strict=False):
+        coords = prediction["most_likely_click"]
+        y_coord = coords[0] * height
+        x_coord = coords[1] * width
+        results.append(
+            {"uuid": uuid, "x_pixel_coord": x_coord, "y_pixel_coord": y_coord}
+        )
+    return results
+
+
+class BatchMurkoForwarder:
+    def __init__(self, redis_client: StrictRedis, batch_size: int):
+        """
+        Holds image data streamed from redis and forwards it to murko when:
+            * A set number have been received
+            * The shape of the images changes
+            * When `flush` is called
+
+        Once data has been forwarded this will then wait on the results and put them
+        back in redis.
+
+        Args:
+            redis_client: The client to send murko results back to redis.
+            batch_size: How many results to accumulate until they are flushed to redis.
+        """
+        self.redis_client = redis_client
+        self.batch_size = batch_size
+        self._uuids_and_images: dict[str, NDArray] = {}
+        self._last_image_size: tuple[int, int] | None = None
+        self._last_sample_id = ""
+
+    def _handle_batch_of_images(self, sample_id, images, uuids):
+        request_arguments: MurkoRequest = {
+            "model_img_size": (256, 320),
+            "to_predict": np.array(images),
+            "save": False,
+            "min_size": 64,
+            "description": [
+                "foreground",
+                "crystal",
+                "loop_inside",
+                "loop",
+                ["crystal", "loop"],
+                ["crystal", "loop", "stem"],
+            ],
+            "prefix": uuids,
+        }
+        predictions = send_to_murko_and_get_results(request_arguments)
+        results = correlate_results_to_uuids(request_arguments, predictions)
+        self._send_murko_results_to_redis(sample_id, results)
+
+    def _send_murko_results_to_redis(self, sample_id: str, results: list):
+        for result in results:
+            self.redis_client.hset(
+                f"murko:{sample_id}:results", result["uuid"], json.dumps(result)
+            )
+        self.redis_client.publish("murko-results", json.dumps(results))
+
+    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)
+        self._last_sample_id = sample_id
+        if self._last_image_size and self._last_image_size != image_size:
+            self.flush()
+        self._uuids_and_images[uuid] = image
+        self._last_image_size = image_size
+        if len(self._uuids_and_images.keys()) >= self.batch_size:
+            self.flush()
+
+    def flush(self):
+        """Flush the batch to murko."""
+        if self._uuids_and_images:
+            self._handle_batch_of_images(
+                self._last_sample_id,
+                list(self._uuids_and_images.values()),
+                list(self._uuids_and_images.keys()),
+            )
+        self._uuids_and_images = {}
+        self._last_image_size = None
+
+
+class RedisListener:
+    TIMEOUT_S = 2
+
+    def __init__(
+        self,
+        redis_host=REDIS_HOST,
+        redis_password=REDIS_PASSWORD,
+        db=MURKO_REDIS_DB,
+        redis_channel="murko",
+    ):
+        self.redis_client = StrictRedis(
+            host=redis_host,
+            password=redis_password,
+            db=db,
+        )
+        self.pubsub = self.redis_client.pubsub()
+        self.channel = redis_channel
+        self.forwarder = BatchMurkoForwarder(self.redis_client, 10)
+
+    def _get_and_handle_message(self):
+        message = self.pubsub.get_message(timeout=self.TIMEOUT_S)
+        if message and message["type"] == "message":
+            data = json.loads(message["data"])
+            LOGGER.info(f"Received from redis: {data}")
+            uuid = data["uuid"]
+            sample_id = data["sample_id"]
+
+            # Images are put in redis as raw jpeg bytes, murko needs numpy arrays
+            raw_image = self.redis_client.hget(f"murko:{sample_id}:raw", uuid)
+            assert isinstance(raw_image, bytes)
+            image = Image.open(io.BytesIO(raw_image))
+            image = np.asarray(image)
+
+            self.forwarder.add(sample_id, uuid, image)
+
+        elif not message:
+            self.forwarder.flush()
+
+    def listen_for_image_data_forever(self):
+        self.pubsub.subscribe(self.channel)
+
+        while True:
+            self._get_and_handle_message()
+
+
+if __name__ == "__main__":
+    client = RedisListener()
+    client.listen_for_image_data_forever()

mx_bluesky/beamlines/i24/serial/__init__.py

@@ -18,9 +18,6 @@ from .fixed_target.i24ssx_Chip_Manager_py3v1 import (
     moveto,
     moveto_preset,
     pumpprobe_calc,
-    save_screen_map,
-    upload_parameters,
-    write_parameter_file,
 )
 from .log import clean_up_log_config_at_end, setup_collection_logs
 from .setup_beamline.setup_detector import setup_detector_stage
@@ -44,9 +41,6 @@ __all__ = [
     "load_lite_map",
     "load_stock_map",
     "pumpprobe_calc",
-    "save_screen_map",
-    "upload_parameters",
-    "write_parameter_file",
     "setup_collection_logs",
     "clean_up_log_config_at_end",
 ]

mx_bluesky/beamlines/i24/serial/dcid.py

@@ -2,29 +2,26 @@ import datetime
 import json
 import math
 import os
-import re
 import subprocess
-import warnings
 from functools import lru_cache
 
+import bluesky.plan_stubs as bps
 import requests
+from dodal.beamlines import i24
+from dodal.devices.i24.beam_center import DetectorBeamCenter
+from dodal.devices.i24.dcm import DCM
+from dodal.devices.i24.focus_mirrors import FocusMirrorsMode
+from dodal.devices.i24.pilatus_metadata import PilatusMetadata
 
+from mx_bluesky.beamlines.i24.serial.fixed_target.ft_utils import PumpProbeSetting
 from mx_bluesky.beamlines.i24.serial.log import SSX_LOGGER
-from mx_bluesky.beamlines.i24.serial.parameters import SSXType
-from mx_bluesky.beamlines.i24.serial.setup_beamline import (
-    Detector,
-    Eiger,
-    Pilatus,
-    caget,
-    cagetstring,
-    pv,
+from mx_bluesky.beamlines.i24.serial.parameters import (
+    BeamSettings,
+    DetectorName,
+    ExtruderParameters,
+    FixedTargetParameters,
 )
-
-try:
-    from typing import Literal
-except ImportError:
-    pass
-
+from mx_bluesky.beamlines.i24.serial.setup_beamline import Detector, Eiger, Pilatus
 
 # Collection start/end script to kick off analysis
 COLLECTION_START_SCRIPT = "/dls_sw/i24/scripts/RunAtStartOfCollect-i24-ssx.sh"
@@ -49,19 +46,54 @@ def get_auth_header() -> dict:
     return {"Authorization": "Bearer " + token}
 
 
-class DCID:
+def read_beam_info_from_hardware(
+    dcm: DCM,
+    mirrors: FocusMirrorsMode,
+    beam_center: DetectorBeamCenter,
+    detector_name: DetectorName,
+):
+    """ Read the beam information from hardware.
+
+    Args:
+        dcm (DCM): The decm device.
+        mirrors (FocusMirrorMode): The device describing the focus mirror mode settings.
+        beam_center (DetectorBeamCenter): A device to set and read the beam center on \
+            the detector.
+        detector_name (DetectorName): The detector currently in use.
+
+    Returns:
+        BeamSettings parameter model.
     """
-    Interfaces with ISPyB to allow ssx DCID/synchweb interaction.
+    wavelength = yield from bps.rd(dcm.wavelength_in_a)
+    beamsize_x = yield from bps.rd(mirrors.beam_size_x)
+    beamsize_y = yield from bps.rd(mirrors.beam_size_y)
+    pixel_size = (
+        Eiger().pixel_size_mm if detector_name == "eiger" else Pilatus().pixel_size_mm
+    )
+    beam_center_x = yield from bps.rd(beam_center.beam_x)
+    beam_center_y = yield from bps.rd(beam_center.beam_y)
+    return BeamSettings(
+        wavelength_in_a=wavelength,
+        beam_size_in_um=(beamsize_x, beamsize_y),
+        beam_center_in_mm=(
+            beam_center_x * pixel_size[0],
+            beam_center_y * pixel_size[1],
+        ),
+    )
+
+
+class DCID:
+    """ Interfaces with ISPyB to allow ssx DCID/synchweb interaction.
 
     Args:
-        server: The URL for the bridge server, if not the default.
-        emit_errors:
-            If False, errors while interacting with the DCID server will
-            not be propagated to the caller. This decides if you want to
-            stop collection if you can't get a DCID
-        timeout: Length of time to wait for the DB server before giving up
-        ssx_type: The type of SSX experiment this is for
-        detector: The detector in use for current collection.
+        server (str, optional): The URL for the bridge server, if not the default.
+        emit_errors (bool, optional): If False, errors while interacting with the DCID \
+            server will not be propagated to the caller. This decides if you want to \
+            stop collection if you can't get a DCID. Defaults to True.
+        timeout (float, optional): Length of time in s to wait for the DB server before \
+            giving up. Defaults to 10 s.
+        expt_parameters (ExtruderParameters | FixedTargetParameters): Collection \
+            parameters input by user.
 
 
     Attributes:
@@ -75,47 +107,44 @@ class DCID:
         server: str | None = None,
         emit_errors: bool = True,
         timeout: float = 10,
-        ssx_type: SSXType = SSXType.FIXED,
-        detector: Detector | Literal["eiger", "pilatus"] | None = None,
+        expt_params: ExtruderParameters | FixedTargetParameters,
     ):
+        self.parameters = expt_params
         self.detector: Detector
         # Handle case of string literal
-        if detector == "eiger":
-            self.detector = Eiger()
-        elif detector == "pilatus":
-            self.detector = Pilatus()
-        elif detector is None:
-            self.detector = Pilatus()
-            warnings.warn(
-                "Please pass detector= to DCID. Pilatus assumed, this will be removed in the future.",
-                UserWarning,
-                stacklevel=5,
-            )
+        match expt_params.detector_name:
+            case "eiger":
+                self.detector = Eiger()
+            case "pilatus":
+                self.detector = Pilatus()
 
         self.server = server or DEFAULT_ISPYB_SERVER
         self.emit_errors = emit_errors
         self.error = False
         self.timeout = timeout
-        self.ssx_type = SSXType(ssx_type)
         self.dcid = None
 
     def generate_dcid(
         self,
-        visit: str,
+        beam_settings: BeamSettings,
         image_dir: str,
+        file_template: str,
         num_images: int,
-        exposure_time: float,
-        start_time: datetime.datetime | None = None,
         shots_per_position: int = 1,
-        pump_exposure_time: float | None = None,
-        pump_delay: float = 0,
-        pump_status: int = 0,
+        start_time: datetime.datetime | None = None,
+        pump_probe: bool = False,
     ):
         """Generate an ispyb DCID.
 
         Args:
-            visit: The name of the visit e.g. "mx12345-4"
-            image_dir: The location the images will be written
+            beam_settings (BeamSettings): Information about the beam read from hardware.
+            image_dir (str): The location the images will be written to.
+            num_images (int): Total number of images to be collected.
+            shots_per_position (int, optional): Number of exposures per position in a \
+                chip. Defaults to 1, which works for extruder.
+            start_time(datetime, optional): Collection start time. Defaults to None.
+            pump_probe (bool, optional): If True, a pump probe collection is running. \
+                Defaults to False.
         """
         try:
             if not start_time:
@@ -123,22 +152,18 @@ class DCID:
             elif not start_time.timetz:
                 start_time = start_time.astimezone()
 
-            # Gather data from the beamline
-            detector_distance = float(caget(self.detector.pv.detector_distance))
-            wavelength = float(caget(self.detector.pv.wavelength))
-            resolution = get_resolution(self.detector, detector_distance, wavelength)
-            beamsize_x, beamsize_y = get_beamsize()
-            transmission = float(caget(self.detector.pv.transmission)) * 100
-            xbeam, ybeam = get_beam_center(self.detector)
+            resolution = get_resolution(
+                self.detector,
+                self.parameters.detector_distance_mm,
+                beam_settings.wavelength_in_a,
+            )
+            beamsize_x, beamsize_y = beam_settings.beam_size_in_um
+            transmission = self.parameters.transmission * 100
+            xbeam, ybeam = beam_settings.beam_center_in_mm
 
             if isinstance(self.detector, Pilatus):
-                # Mirror the construction that the PPU does
-                fileTemplate = get_pilatus_filename_template_from_pvs()
                 startImageNumber = 0
             elif isinstance(self.detector, Eiger):
-                # Eiger base filename is directly written to the PV
-                # Nexgen then uses this to write the .nxs file
-                fileTemplate = str(cagetstring(self.detector.pv.file_name)) + ".nxs"
                 startImageNumber = 1
             else:
                 raise ValueError("Unknown detector:", self.detector)
@@ -147,48 +172,48 @@ class DCID:
                 {
                     "name": "Xray probe",
                     "offset": 0,
-                    "duration": exposure_time,
-                    "period": exposure_time,
+                    "duration": self.parameters.exposure_time_s,
+                    "period": self.parameters.exposure_time_s,
                     "repetition": shots_per_position,
                     "eventType": "XrayDetection",
                 }
             ]
-            if pump_status > 0:
-                # https://confluence.diamond.ac.uk/pages/viewpage.action?pageId=131238829
-                # https://confluence.diamond.ac.uk/display/MXTech/Dynamics+and+fixed+targets
-                # pump_status = 0: no pump probe
-                # pump_status = 1: pump then probe
-                # pump_status = 2: pump within probe
-                # pump_status = 3-7: different EAVA modes (i.e. also pump then probe)
-                if pump_status != 2 and self.ssx_type is SSXType.FIXED:
-                    # Pump status could be 1 for extruder but not have this.
-                    # pump then probe - pump_delay corresponds to time *before* first image
-                    pump_delay = -pump_delay
+            if pump_probe:
+                match self.parameters:
+                    case FixedTargetParameters():
+                        # pump then probe - pump_delay corresponds to time *before* first image
+                        pump_delay = (
+                            -self.parameters.laser_delay_s
+                            if self.parameters.pump_repeat
+                            is not PumpProbeSetting.Short2
+                            else self.parameters.laser_delay_s
+                        )
+                    case ExtruderParameters():
+                        pump_delay = self.parameters.laser_delay_s
                 events.append(
                     {
                         "name": "Laser probe",
                         "offset": pump_delay,
-                        "duration": pump_exposure_time,
-                        # "period": None,
+                        "duration": self.parameters.laser_dwell_s,
                         "repetition": 1,
                         "eventType": "LaserExcitation",
                     },
                 )
 
             data = {
-                "detectorDistance": float(detector_distance),
+                "detectorDistance": self.parameters.detector_distance_mm,
                 "detectorId": self.detector.id,
-                "exposureTime": float(exposure_time),
-                "fileTemplate": fileTemplate,
-                "imageDirectory": str(image_dir),
-                "numberOfImages": int(num_images),
-                "resolution": float(resolution),
+                "exposureTime": self.parameters.exposure_time_s,
+                "fileTemplate": file_template,
+                "imageDirectory": image_dir,
+                "numberOfImages": num_images,
+                "resolution": resolution,
                 "startImageNumber": startImageNumber,
                 "startTime": start_time.isoformat(),
-                "transmission": float(transmission),
-                "visit": visit,
-                "wavelength": float(wavelength),
-                "group": {"experimentType": self.ssx_type.value},
+                "transmission": transmission,
+                "visit": self.parameters.visit.name,
+                "wavelength": beam_settings.wavelength_in_a,
+                "group": {"experimentType": self.parameters.ispyb_experiment_type},
                 "xBeam": xbeam,
                 "yBeam": ybeam,
                 "ssx": {
@@ -334,84 +359,33 @@ class DCID:
             SSX_LOGGER.warning("Error completing DCID: %s (%s)", e, resp_str)
 
 
-def get_pilatus_filename_template_from_pvs() -> str:
-    """
-    Get the template file path by querying the detector PVs.
-
-    Returns: A template string, with the image numbers replaced with '#'
-    """
-
-    filename = cagetstring(pv.pilat_filename)
-    filename_template = cagetstring(pv.pilat_filetemplate)
-    file_number = int(caget(pv.pilat_filenumber))
-    # Exploit fact that passing negative numbers will put the - before the 0's
-    expected_filename = str(filename_template % (filename, f"{file_number:05d}_", -9))
-    # Now, find the -09 part of this
-    numberpart = re.search(r"(-0+9)", expected_filename)
-    # Make sure this was the only one
-    if numberpart is not None:
-        assert re.search(r"(-0+9)", expected_filename[numberpart.end() :]) is None
-        template_fill = "#" * len(numberpart.group(0))
-        return (
-            expected_filename[: numberpart.start()]
-            + template_fill
-            + expected_filename[numberpart.end() :]
-        )
-    else:
-        raise ValueError(f"{filename=} did not contain the numbers for templating")
-
-
-def get_beamsize() -> tuple[float | None, float | None]:
+def get_pilatus_filename_template_from_device():
     """
-    Read the PVs to get the current beamsize.
+    Get the template file path by querying the detector PVs, mirror the construction \
+    that the PPU does.
 
     Returns:
-        A tuple (x, y) of beam size (in µm). These values can be 'None'
-        if the focus mode was unrecognised.
+        A template string, with the image numbers replaced with '#'
     """
-    # These I24 modes are from GDA
-    focus_modes = {
-        "focus10": ("7x7", 7, 7),
-        "focus20d": ("20x20", 20, 20),
-        "focus30d": ("30x30", 30, 30),
-        "focus50d": ("50x50", 50, 50),
-        "focus1050d": ("10x50", 10, 50),
-        "focus5010d": ("50x10", 50, 10),
-        "focus3010d": ("30x10", 30, 10),
-    }
-    v_mode = caget("BL24I-OP-MFM-01:G0:TARGETAPPLY")
-    h_mode = caget("BL24I-OP-MFM-01:G1:TARGETAPPLY")
-    # Validate these and note an error otherwise
-    if not v_mode.startswith("VMFM") or v_mode[4:] not in focus_modes:
-        SSX_LOGGER.error("Unrecognised vertical beam mode %s", v_mode)
-    if not h_mode.startswith("HMFM") or h_mode[4:] not in focus_modes:
-        SSX_LOGGER.error("Unrecognised horizontal beam mode %s", h_mode)
-    _, h, _ = focus_modes.get(h_mode[4:], (None, None, None))
-    _, _, v = focus_modes.get(v_mode[4:], (None, None, None))
-
-    return (h, v)
+    pilatus_metadata: PilatusMetadata = i24.pilatus_metadata()
+
+    filename_template = yield from bps.rd(pilatus_metadata.filename_template)
+    return filename_template
 
 
 def get_resolution(detector: Detector, distance: float, wavelength: float) -> float:
-    """
-    Calculate the inscribed resolution for detector.
+    """ Calculate the inscribed resolution for detector.
 
-    This assumes perfectly centered beam as I don't know where to
-    extract the beam position parameters yet.
+    This assumes perfectly centered beam as I don't know where to extract the beam \
+    position parameters yet.
 
     Args:
-        distance: Distance to detector (mm)
-        wavelength: Beam wavelength (Å)
+        detector (Detector): Detector instance, Eiger() or Pilatus().
+        distance (float): Distance to detector, in mm.
+        wavelength (float): Beam wavelength, in Å.
 
     Returns:
-        Maximum resolution (Å)
+        Maximum resolution, in Å.
     """
     width = detector.image_size_mm[0]
     return round(wavelength / (2 * math.sin(math.atan(width / (2 * distance)) / 2)), 2)
-
-
-def get_beam_center(detector: Detector) -> tuple[float, float]:
-    """Get the detector beam center, in mm"""
-    beamX = float(caget(detector.pv.beamx)) * detector.pixel_size_mm[0]
-    beamY = float(caget(detector.pv.beamy)) * detector.pixel_size_mm[1]
-    return (beamX, beamY)

mx_bluesky/beamlines/i24/serial/extruder/EX-gui-edm/DiamondExtruder-I24-py3v1.edl

@@ -370,7 +370,7 @@ font "arial-medium-r-24.0"
 buttonLabel "Start"
 numCmds 1
 command {
-  0 "blueapi -c CONFIG_LOCATION controller run run_extruder_plan '\{\"zebra\":\"zebra\",\"aperture\":\"aperture\",\"backlight\":\"backlight\",\"beamstop\":\"beamstop\",\"detector_stage\":\"detector_motion\",\"shutter\":\"shutter\",\"dcm\":\"dcm\"\}'"
+  0 "blueapi -c CONFIG_LOCATION controller run run_extruder_plan '\{\"zebra\":\"zebra\",\"aperture\":\"aperture\",\"backlight\":\"backlight\",\"beamstop\":\"beamstop\",\"detector_stage\":\"detector_motion\",\"shutter\":\"shutter\",\"dcm\":\"dcm\",\"mirrors\":\"focus_mirrors\",\"attenuator\":\"attenuator\"\}'"
 }
 endObjectProperties