mx-bluesky 0.3.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

mx_bluesky/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from mx_bluesky.hyperion.external_interaction.callbacks.plan_reactive_callback import (
+    PlanReactiveCallback,
+)
+from mx_bluesky.hyperion.external_interaction.nexus.nexus_utils import (
+    create_beam_and_attenuator_parameters,
+    vds_type_based_on_bit_depth,
+)
+from mx_bluesky.hyperion.external_interaction.nexus.write_nexus import NexusWriter
+from mx_bluesky.hyperion.log import NEXUS_LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.parameters.gridscan import ThreeDGridScan
+
+if TYPE_CHECKING:
+    from event_model.documents import Event, EventDescriptor, RunStart
+
+
+class GridscanNexusFileCallback(PlanReactiveCallback):
+    """Callback class to handle the creation of Nexus files based on experiment \
+    parameters. Initialises on recieving a 'start' document for the \
+    'run_gridscan_move_and_tidy' sub plan, which must also contain the run parameters, \
+    as metadata under the 'hyperion_internal_parameters' key. Actually writes the \
+    nexus files on updates the timestamps on recieving the 'ispyb_reading_hardware' event \
+    document, and finalises the files on getting a 'stop' document for the whole run.
+
+    To use, subscribe the Bluesky RunEngine to an instance of this class.
+    E.g.:
+        nexus_file_handler_callback = NexusFileCallback(parameters)
+        RE.subscribe(nexus_file_handler_callback)
+    Or decorate a plan using bluesky.preprocessors.subs_decorator.
+
+    See: https://blueskyproject.io/bluesky/callbacks.html#ways-to-invoke-callbacks
+    """
+
+    def __init__(self) -> None:
+        super().__init__(NEXUS_LOGGER)
+        self.run_start_uid: str | None = None
+        self.nexus_writer_1: NexusWriter | None = None
+        self.nexus_writer_2: NexusWriter | None = None
+        self.descriptors: dict[str, EventDescriptor] = {}
+        self.log = NEXUS_LOGGER
+
+    def activity_gated_start(self, doc: RunStart):
+        if doc.get("subplan_name") == CONST.PLAN.GRIDSCAN_OUTER:
+            json_params = doc.get("hyperion_parameters")
+            NEXUS_LOGGER.info(
+                f"Nexus writer received start document with experiment parameters {json_params}"
+            )
+            parameters = ThreeDGridScan.from_json(json_params)
+            d_size = parameters.detector_params.detector_size_constants.det_size_pixels
+            grid_n_img_1 = parameters.scan_indices[1]
+            grid_n_img_2 = parameters.num_images - grid_n_img_1
+            data_shape_1 = (grid_n_img_1, d_size.width, d_size.height)
+            data_shape_2 = (grid_n_img_2, d_size.width, d_size.height)
+            run_number_2 = parameters.detector_params.run_number + 1
+            self.nexus_writer_1 = NexusWriter(
+                parameters, data_shape_1, parameters.scan_points_first_grid
+            )
+            self.nexus_writer_2 = NexusWriter(
+                parameters,
+                data_shape_2,
+                parameters.scan_points_second_grid,
+                run_number=run_number_2,
+                vds_start_index=parameters.scan_indices[1],
+                omega_start_deg=90,
+            )
+            self.run_start_uid = doc.get("uid")
+
+    def activity_gated_descriptor(self, doc: EventDescriptor):
+        self.descriptors[doc["uid"]] = doc
+
+    def activity_gated_event(self, doc: Event) -> Event | None:
+        assert (event_descriptor := self.descriptors.get(doc["descriptor"])) is not None
+        if event_descriptor.get("name") == CONST.DESCRIPTORS.HARDWARE_READ_DURING:
+            data = doc["data"]
+            for nexus_writer in [self.nexus_writer_1, self.nexus_writer_2]:
+                assert nexus_writer, "Nexus callback did not receive start doc"
+                (
+                    nexus_writer.beam,
+                    nexus_writer.attenuator,
+                ) = create_beam_and_attenuator_parameters(
+                    data["dcm-energy_in_kev"],
+                    data["flux_flux_reading"],
+                    data["attenuator-actual_transmission"],
+                )
+                vds_data_type = vds_type_based_on_bit_depth(
+                    doc["data"]["eiger_bit_depth"]
+                )
+                nexus_writer.create_nexus_file(vds_data_type)
+                NEXUS_LOGGER.info(f"Nexus file created at {nexus_writer.data_filename}")
+
+        return super().activity_gated_event(doc)

mx_bluesky/hyperion/external_interaction/callbacks/zocalo_callback.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from bluesky.callbacks import CallbackBase
+from dodal.devices.zocalo import ZocaloStartInfo, ZocaloTrigger
+
+from mx_bluesky.hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
+from mx_bluesky.hyperion.log import ISPYB_LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.utils.utils import number_of_frames_from_scan_spec
+
+if TYPE_CHECKING:
+    from event_model.documents import Event, EventDescriptor, RunStart, RunStop
+
+
+class ZocaloCallback(CallbackBase):
+    """Callback class to handle the triggering of Zocalo processing.
+    Sends zocalo a run_start signal on receiving a start document for the specified
+    sub-plan, and sends a run_end signal on receiving a stop document for the same plan.
+
+    The metadata of the sub-plan this starts on must include a zocalo_environment.
+
+    Shouldn't be subscribed directly to the RunEngine, instead should be passed to the
+    `emit` argument of an ISPyB callback which appends DCIDs to the relevant start doc.
+    """
+
+    def _reset_state(self):
+        self.run_uid: str | None = None
+        self.triggering_plan: str | None = None
+        self.zocalo_interactor: ZocaloTrigger | None = None
+        self.zocalo_info: list[ZocaloStartInfo] = []
+        self.descriptors: dict[str, EventDescriptor] = {}
+
+    def __init__(
+        self,
+    ):
+        super().__init__()
+        self._reset_state()
+
+    def start(self, doc: RunStart):
+        ISPYB_LOGGER.info("Zocalo handler received start document.")
+        if triggering_plan := doc.get(CONST.TRIGGER.ZOCALO):
+            self.triggering_plan = triggering_plan
+            assert isinstance(zocalo_environment := doc.get("zocalo_environment"), str)
+            ISPYB_LOGGER.info(f"Zocalo environment set to {zocalo_environment}.")
+            self.zocalo_interactor = ZocaloTrigger(zocalo_environment)
+
+        if self.triggering_plan and doc.get("subplan_name") == self.triggering_plan:
+            self.run_uid = doc.get("uid")
+            assert isinstance(scan_points := doc.get("scan_points"), list)
+            if (
+                isinstance(ispyb_ids := doc.get("ispyb_dcids"), tuple)
+                and len(ispyb_ids) > 0
+            ):
+                ids_and_shape = list(zip(ispyb_ids, scan_points, strict=False))
+                start_frame = 0
+                self.zocalo_info = []
+                for idx, id_and_shape in enumerate(ids_and_shape):
+                    id, shape = id_and_shape
+                    num_frames = number_of_frames_from_scan_spec(shape)
+                    self.zocalo_info.append(
+                        ZocaloStartInfo(id, None, start_frame, num_frames, idx)
+                    )
+                    start_frame += num_frames
+            else:
+                raise ISPyBDepositionNotMade(
+                    f"No ISPyB IDs received by the start of {self.triggering_plan=}"
+                )
+
+    def descriptor(self, doc: EventDescriptor):
+        self.descriptors[doc["uid"]] = doc
+
+    def event(self, doc: Event) -> Event:
+        event_descriptor = self.descriptors[doc["descriptor"]]
+        if event_descriptor.get("name") == CONST.DESCRIPTORS.ZOCALO_HW_READ:
+            filename = doc["data"]["eiger_odin_file_writer_id"]
+            for start_info in self.zocalo_info:
+                start_info.filename = filename
+                assert self.zocalo_interactor is not None
+                self.zocalo_interactor.run_start(start_info)
+        return doc
+
+    def stop(self, doc: RunStop):
+        if doc.get("run_start") == self.run_uid:
+            ISPYB_LOGGER.info(
+                f"Zocalo handler received stop document, for run {doc.get('run_start')}."
+            )
+            assert self.zocalo_interactor is not None
+            for info in self.zocalo_info:
+                self.zocalo_interactor.run_end(info.ispyb_dcid)
+            self._reset_state()

mx_bluesky/hyperion/external_interaction/config_server.py

@@ -0,0 +1,35 @@
+from daq_config_server.client import ConfigServer
+from pydantic import BaseModel
+
+from mx_bluesky.hyperion.log import LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+_CONFIG_SERVER: ConfigServer | None = None
+
+
+def config_server() -> ConfigServer:
+    global _CONFIG_SERVER
+    if _CONFIG_SERVER is None:
+        _CONFIG_SERVER = ConfigServer(CONST.CONFIG_SERVER_URL, LOGGER)
+    return _CONFIG_SERVER
+
+
+class FeatureFlags(BaseModel):
+    # The default value will be used as the fallback when doing a best-effort fetch
+    # from the service
+    use_panda_for_gridscan: bool = False
+    use_gpu_for_gridscan: bool = False
+    set_stub_offsets: bool = False
+
+    @classmethod
+    def _get_flags(cls):
+        flags = config_server().best_effort_get_all_feature_flags()
+        return {f: flags[f] for f in flags if f in cls.__fields__.keys()}
+
+    @classmethod
+    def best_effort(cls):
+        return cls(**cls._get_flags())
+
+    def update_self_from_server(self):
+        for flag, value in self._get_flags().items():
+            setattr(self, flag, value)

mx_bluesky/hyperion/external_interaction/exceptions.py

@@ -0,0 +1,13 @@
+from mx_bluesky.hyperion.exceptions import WarningException
+
+
+class ISPyBDepositionNotMade(Exception):
+    """Raised when the ISPyB or Zocalo callbacks can't access ISPyB deposition numbers."""
+
+    pass
+
+
+class NoCentreFoundException(WarningException):
+    """Error for if zocalo is unable to find the centre during a gridscan."""
+
+    pass

mx_bluesky/hyperion/external_interaction/ispyb/data_model.py

@@ -0,0 +1,95 @@
+from dataclasses import asdict, dataclass
+from enum import Enum
+
+
+class Orientation(Enum):
+    HORIZONTAL = "horizontal"
+    VERTICAL = "vertical"
+
+
+@dataclass()
+class DataCollectionGroupInfo:
+    visit_string: str
+    experiment_type: str
+    sample_id: int | None
+    sample_barcode: str | None = None
+
+
+@dataclass(kw_only=True)
+class DataCollectionInfo:
+    omega_start: float | None = None
+    data_collection_number: int | None = None
+    xtal_snapshot1: str | None = None
+    xtal_snapshot2: str | None = None
+    xtal_snapshot3: str | None = None
+    xtal_snapshot4: str | None = None
+
+    n_images: int | None = None
+    axis_range: float | None = None
+    axis_end: float | None = None
+    kappa_start: float | None = None
+
+    parent_id: int | None = None
+    visit_string: str | None = None
+    sample_id: int | None = None
+    detector_id: int | None = None
+    axis_start: float | None = None
+    focal_spot_size_at_samplex: float | None = None
+    focal_spot_size_at_sampley: float | None = None
+    slitgap_vertical: float | None = None
+    slitgap_horizontal: float | None = None
+    beamsize_at_samplex: float | None = None
+    beamsize_at_sampley: float | None = None
+    transmission: float | None = None
+    comments: str | None = None
+    detector_distance: float | None = None
+    exp_time: float | None = None
+    imgdir: str | None = None
+    file_template: str | None = None
+    imgprefix: str | None = None
+    imgsuffix: str | None = None
+    n_passes: int | None = None
+    overlap: int | None = None
+    flux: float | None = None
+    start_image_number: int | None = None
+    resolution: float | None = None
+    wavelength: float | None = None
+    xbeam: float | None = None
+    ybeam: float | None = None
+    synchrotron_mode: str | None = None
+    undulator_gap1: float | None = None
+    start_time: str | None = None
+
+
+@dataclass
+class DataCollectionPositionInfo:
+    pos_x: float
+    pos_y: float
+    pos_z: float
+
+
+@dataclass
+class DataCollectionGridInfo:
+    dx_in_mm: float
+    dy_in_mm: float
+    steps_x: int
+    steps_y: int
+    microns_per_pixel_x: float
+    microns_per_pixel_y: float
+    snapshot_offset_x_pixel: int
+    snapshot_offset_y_pixel: int
+    orientation: Orientation
+    snaked: bool
+
+    def as_dict(self):
+        d = asdict(self)
+        d["orientation"] = self.orientation.value
+        return d
+
+
+@dataclass(kw_only=True)
+class ScanDataInfo:
+    data_collection_info: DataCollectionInfo
+    data_collection_id: int | None = None
+    data_collection_position_info: DataCollectionPositionInfo | None = None
+    data_collection_grid_info: DataCollectionGridInfo | None = None

mx_bluesky/hyperion/external_interaction/ispyb/exp_eye_store.py

@@ -0,0 +1,125 @@
+import configparser
+
+from requests import patch, post
+from requests.auth import AuthBase
+
+from mx_bluesky.hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
+from mx_bluesky.hyperion.external_interaction.ispyb.ispyb_utils import (
+    get_current_time_string,
+    get_ispyb_config,
+)
+
+RobotActionID = int
+
+
+class BearerAuth(AuthBase):
+    def __init__(self, token):
+        self.token = token
+
+    def __call__(self, r):
+        r.headers["authorization"] = "Bearer " + self.token
+        return r
+
+
+def _get_base_url_and_token() -> tuple[str, str]:
+    config = configparser.ConfigParser()
+    conf = get_ispyb_config()
+    config.read(conf)
+    expeye_config = config["expeye"]
+    return expeye_config["url"], expeye_config["token"]
+
+
+class ExpeyeInteraction:
+    CREATE_ROBOT_ACTION = "/proposals/{proposal}/sessions/{visit_number}/robot-actions"
+    UPDATE_ROBOT_ACTION = "/robot-actions/{action_id}"
+
+    def __init__(self) -> None:
+        url, token = _get_base_url_and_token()
+        self.base_url = url + "/core"
+        self.auth = BearerAuth(token)
+
+    def _send_and_get_response(self, url, data, send_func) -> dict:
+        response = send_func(url, auth=self.auth, json=data)
+        if not response.ok:
+            raise ISPyBDepositionNotMade(f"Could not write {data} to {url}: {response}")
+        return response.json()
+
+    def start_load(
+        self,
+        proposal_reference: str,
+        visit_number: int,
+        sample_id: int,
+        dewar_location: int,
+        container_location: int,
+    ) -> RobotActionID:
+        """Create a robot load entry in ispyb.
+
+        Args:
+            proposal_reference (str): The proposal of the experiment e.g. cm37235
+            visit_number (int): The visit number for the proposal, usually this can be
+                                found added to the end of the proposal e.g. the data for
+                                visit number 2 of proposal cm37235 is in cm37235-2
+            sample_id (int): The id of the sample in the database
+            dewar_location (int): Which puck in the dewar the sample is in
+            container_location (int): Which pin in that puck has the sample
+
+        Returns:
+            RobotActionID: The id of the robot load action that is created
+        """
+        url = self.base_url + self.CREATE_ROBOT_ACTION.format(
+            proposal=proposal_reference, visit_number=visit_number
+        )
+
+        data = {
+            "startTimestamp": get_current_time_string(),
+            "sampleId": sample_id,
+            "actionType": "LOAD",
+            "containerLocation": container_location,
+            "dewarLocation": dewar_location,
+        }
+        response = self._send_and_get_response(url, data, post)
+        return response["robotActionId"]
+
+    def update_barcode_and_snapshots(
+        self,
+        action_id: RobotActionID,
+        barcode: str,
+        snapshot_before_path: str,
+        snapshot_after_path: str,
+    ):
+        """Update the barcode and snapshots of an existing robot action.
+
+        Args:
+            action_id (RobotActionID): The id of the action to update
+            barcode (str): The barcode to give the action
+            snapshot_before_path (str): Path to the snapshot before robot load
+            snapshot_after_path (str): Path to the snapshot after robot load
+        """
+        url = self.base_url + self.UPDATE_ROBOT_ACTION.format(action_id=action_id)
+
+        data = {
+            "sampleBarcode": barcode,
+            "xtalSnapshotBefore": snapshot_before_path,
+            "xtalSnapshotAfter": snapshot_after_path,
+        }
+        self._send_and_get_response(url, data, patch)
+
+    def end_load(self, action_id: RobotActionID, status: str, reason: str):
+        """Finish an existing robot action, providing final information about how it went
+
+        Args:
+            action_id (RobotActionID): The action to finish.
+            status (str): The status of the action at the end, "success" for success,
+                          otherwise error
+            reason (str): If the status is in error than the reason for that error
+        """
+        url = self.base_url + self.UPDATE_ROBOT_ACTION.format(action_id=action_id)
+
+        run_status = "SUCCESS" if status == "success" else "ERROR"
+
+        data = {
+            "endTimestamp": get_current_time_string(),
+            "status": run_status,
+            "message": reason,
+        }
+        self._send_and_get_response(url, data, patch)