mx-bluesky 1.2.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

mx_bluesky/__init__.py

@@ -1,6 +1,11 @@
-from importlib.metadata import version
+"""Top level API.
 
-__version__ = version("mx_bluesky")
-del version
+.. data:: __version__
+    :type: str
+
+    Version number as calculated by https://github.com/pypa/setuptools_scm
+"""
+
+from ._version import __version__
 
 __all__ = ["__version__"]

mx_bluesky/__main__.py

@@ -1,19 +1,24 @@
-from argparse import ArgumentParser
+"""Interface for ``python -m mx_bluesky``."""
 
-from mx_bluesky.example import run_plan
+from argparse import ArgumentParser
+from collections.abc import Sequence
 
 from . import __version__
 
 __all__ = ["main"]
 
 
-def main(args=None):
+def main(args: Sequence[str] | None = None) -> None:
+    """Argument parser for the CLI."""
     parser = ArgumentParser()
-    parser.add_argument("-v", "--version", action="version", version=__version__)
-    args = parser.parse_args(args)
-    run_plan()
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        version=__version__,
+    )
+    parser.parse_args(args)
 
 
-# test with: python -m mx_bluesky
 if __name__ == "__main__":
     main()

mx_bluesky/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.2.0'
-__version_tuple__ = version_tuple = (1, 2, 0)
+__version__ = version = '1.4.1'
+__version_tuple__ = version_tuple = (1, 4, 1)

mx_bluesky/beamlines/i04/callbacks/murko_callback.py

@@ -1,13 +1,16 @@
 import copy
 import json
+from datetime import timedelta
 
 from bluesky.callbacks import CallbackBase
 from dodal.log import LOGGER
-from event_model.documents import Event, RunStart
+from event_model.documents import Event, RunStart, RunStop
 from redis import StrictRedis
 
 
 class MurkoCallback(CallbackBase):
+    DATA_EXPIRY_DAYS = 7
+
     def __init__(self, redis_host: str, redis_password: str, redis_db: int = 0):
         self.redis_client = StrictRedis(
             host=redis_host, password=redis_password, db=redis_db
@@ -15,15 +18,17 @@ class MurkoCallback(CallbackBase):
         self.last_uuid = None
 
     def start(self, doc: RunStart) -> RunStart | None:
+        self.sample_id = doc.get("sample_id")
         self.murko_metadata = {
             "zoom_percentage": doc.get("zoom_percentage"),
             "microns_per_x_pixel": doc.get("microns_per_x_pixel"),
             "microns_per_y_pixel": doc.get("microns_per_y_pixel"),
             "beam_centre_i": doc.get("beam_centre_i"),
             "beam_centre_j": doc.get("beam_centre_j"),
-            "sample_id": doc.get("sample_id"),
+            "sample_id": self.sample_id,
         }
         self.last_uuid = None
+        LOGGER.info(f"Starting to stream metadata to murko under {self.sample_id}")
         return doc
 
     def event(self, doc: Event) -> Event:
@@ -40,6 +45,11 @@ class MurkoCallback(CallbackBase):
         metadata["uuid"] = uuid
 
         # Send metadata to REDIS and trigger murko
-        self.redis_client.hset("test-metadata", uuid, json.dumps(metadata))
+        redis_key = f"murko:{metadata['sample_id']}:metadata"
+        self.redis_client.hset(redis_key, uuid, json.dumps(metadata))
+        self.redis_client.expire(redis_key, timedelta(days=self.DATA_EXPIRY_DAYS))
         self.redis_client.publish("murko", json.dumps(metadata))
-        LOGGER.info("Metadata sent to redis")
+
+    def stop(self, doc: RunStop) -> RunStop | None:
+        LOGGER.info(f"Finished streaming {self.sample_id} to murko")
+        return doc

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/i04/thawing_plan.py

@@ -1,11 +1,13 @@
+from collections.abc import Callable
+
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
 from bluesky.preprocessors import run_decorator, subs_decorator
-from dls_bluesky_core.core import MsgGenerator
+from bluesky.utils import MsgGenerator
 from dodal.beamlines.i04 import MURKO_REDIS_DB, REDIS_HOST, REDIS_PASSWORD
 from dodal.common import inject
 from dodal.devices.oav.oav_detector import OAV
-from dodal.devices.oav.oav_to_redis_forwarder import OAVToRedisForwarder
+from dodal.devices.oav.oav_to_redis_forwarder import OAVToRedisForwarder, Source
 from dodal.devices.robot import BartRobot
 from dodal.devices.smargon import Smargon
 from dodal.devices.thawer import Thawer, ThawerStates
@@ -22,31 +24,62 @@ def thaw_and_stream_to_redis(
     oav: OAV = inject("oav"),
     oav_to_redis_forwarder: OAVToRedisForwarder = inject("oav_to_redis_forwarder"),
 ) -> MsgGenerator:
-    zoom_percentage = yield from bps.rd(oav.zoom_controller.percentage)
+    zoom_percentage = yield from bps.rd(oav.zoom_controller.percentage)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     sample_id = yield from bps.rd(robot.sample_id)
 
-    yield from bps.abs_set(oav.zoom_controller.level, "1.0x", wait=True)
-    yield from bps.abs_set(oav_to_redis_forwarder.sample_id, sample_id)
+    sample_id = int(sample_id)
+    zoom_level_before_thawing = yield from bps.rd(oav.zoom_controller.level)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+
+    yield from bps.mv(oav.zoom_controller.level, "1.0x")  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+
+    def switch_forwarder_to_ROI() -> MsgGenerator:
+        yield from bps.complete(oav_to_redis_forwarder, wait=True)
+        yield from bps.mv(
+            # See: https://github.com/bluesky/bluesky/issues/1809
+            oav_to_redis_forwarder.selected_source,  # type: ignore
+            Source.ROI.value,  # type: ignore
+        )
+        yield from bps.kickoff(oav_to_redis_forwarder, wait=True)
+
+    microns_per_pixel_x = yield from bps.rd(oav.microns_per_pixel_x)
+    microns_per_pixel_y = yield from bps.rd(oav.microns_per_pixel_y)
+    beam_centre_i = yield from bps.rd(oav.beam_centre_i)
+    beam_centre_j = yield from bps.rd(oav.beam_centre_j)
 
     @subs_decorator(MurkoCallback(REDIS_HOST, REDIS_PASSWORD, MURKO_REDIS_DB))
     @run_decorator(
         md={
-            "microns_per_x_pixel": oav.parameters.micronsPerXPixel,
-            "microns_per_y_pixel": oav.parameters.micronsPerYPixel,
-            "beam_centre_i": oav.parameters.beam_centre_i,
-            "beam_centre_j": oav.parameters.beam_centre_j,
+            "microns_per_x_pixel": microns_per_pixel_x,
+            "microns_per_y_pixel": microns_per_pixel_y,
+            "beam_centre_i": beam_centre_i,
+            "beam_centre_j": beam_centre_j,
             "zoom_percentage": zoom_percentage,
             "sample_id": sample_id,
         }
     )
     def _thaw_and_stream_to_redis():
+        yield from bps.mv(
+            oav_to_redis_forwarder.sample_id,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+            sample_id,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+            oav_to_redis_forwarder.selected_source,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+            Source.FULL_SCREEN.value,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        )
+
         yield from bps.kickoff(oav_to_redis_forwarder, wait=True)
         yield from bps.monitor(smargon.omega.user_readback, name="smargon")
         yield from bps.monitor(oav_to_redis_forwarder.uuid, name="oav")
-        yield from thaw(time_to_thaw, rotation, thawer, smargon)
+        yield from thaw(
+            time_to_thaw, rotation, thawer, smargon, switch_forwarder_to_ROI
+        )
         yield from bps.complete(oav_to_redis_forwarder)
 
-    yield from _thaw_and_stream_to_redis()
+    def cleanup():
+        yield from bps.mv(oav.zoom_controller.level, zoom_level_before_thawing)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+
+    yield from bpp.contingency_wrapper(
+        _thaw_and_stream_to_redis(),
+        final_plan=cleanup,
+    )
 
 
 def thaw(
@@ -54,6 +87,7 @@ def thaw(
     rotation: float = 360,
     thawer: Thawer = inject("thawer"),
     smargon: Smargon = inject("smargon"),
+    plan_between_rotations: Callable[[], MsgGenerator] | None = None,
 ) -> MsgGenerator:
     """Rotates the sample and thaws it at the same time.
 
@@ -64,6 +98,8 @@ def thaw(
         thawer (Thawer, optional): The thawing device. Defaults to inject("thawer").
         smargon (Smargon, optional): The smargon used to rotate.
                                      Defaults to inject("smargon")
+        plan_between_rotations (MsgGenerator, optional): A plan to run between rotations
+                                    of the smargon. Defaults to no plan.
     """
     inital_velocity = yield from bps.rd(smargon.omega.velocity)
     new_velocity = abs(rotation / time_to_thaw) * 2.0
@@ -72,6 +108,8 @@ def thaw(
         yield from bps.abs_set(smargon.omega.velocity, new_velocity, wait=True)
         yield from bps.abs_set(thawer.control, ThawerStates.ON, wait=True)
         yield from bps.rel_set(smargon.omega, rotation, wait=True)
+        if plan_between_rotations:
+            yield from plan_between_rotations()
         yield from bps.rel_set(smargon.omega, -rotation, wait=True)
 
     def cleanup():

mx_bluesky/beamlines/i24/serial/__init__.py

@@ -22,6 +22,7 @@ from .fixed_target.i24ssx_Chip_Manager_py3v1 import (
     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
 
 __all__ = [
@@ -46,4 +47,6 @@ __all__ = [
     "save_screen_map",
     "upload_parameters",
     "write_parameter_file",
+    "setup_collection_logs",
+    "clean_up_log_config_at_end",
 ]