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

mx_bluesky/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, cast
+
+from mx_bluesky.hyperion.external_interaction.callbacks.common.ispyb_mapping import (
+    populate_data_collection_group,
+    populate_remaining_data_collection_info,
+)
+from mx_bluesky.hyperion.external_interaction.callbacks.ispyb_callback_base import (
+    BaseISPyBCallback,
+)
+from mx_bluesky.hyperion.external_interaction.callbacks.rotation.ispyb_mapping import (
+    populate_data_collection_info_for_rotation,
+)
+from mx_bluesky.hyperion.external_interaction.ispyb.data_model import (
+    DataCollectionInfo,
+    DataCollectionPositionInfo,
+    ScanDataInfo,
+)
+from mx_bluesky.hyperion.external_interaction.ispyb.ispyb_store import (
+    IspybIds,
+    StoreInIspyb,
+)
+from mx_bluesky.hyperion.log import ISPYB_LOGGER, set_dcgid_tag
+from mx_bluesky.hyperion.parameters.components import IspybExperimentType
+from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.parameters.rotation import RotationScan
+
+if TYPE_CHECKING:
+    from event_model.documents import Event, RunStart, RunStop
+
+
+class RotationISPyBCallback(BaseISPyBCallback):
+    """Callback class to handle the deposition of experiment parameters into the ISPyB
+    database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on
+    recieving an 'event' document for the 'ispyb_reading_hardware' event, and updates the
+    deposition on recieving its final 'stop' document.
+
+    To use, subscribe the Bluesky RunEngine to an instance of this class.
+    E.g.:
+        ispyb_handler_callback = RotationISPyBCallback(parameters)
+        RE.subscribe(ispyb_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,
+        *,
+        emit: Callable[..., Any] | None = None,
+    ) -> None:
+        super().__init__(emit=emit)
+        self.last_sample_id: int | None = None
+        self.ispyb_ids: IspybIds = IspybIds()
+
+    def activity_gated_start(self, doc: RunStart):
+        if doc.get("subplan_name") == CONST.PLAN.ROTATION_OUTER:
+            ISPYB_LOGGER.info(
+                "ISPyB callback received start document with experiment parameters."
+            )
+            self.params = RotationScan.from_json(doc.get("hyperion_parameters"))
+            dcgid = (
+                self.ispyb_ids.data_collection_group_id
+                if (self.params.sample_id == self.last_sample_id)
+                else None
+            )
+            if (
+                self.params.ispyb_experiment_type
+                == IspybExperimentType.CHARACTERIZATION
+            ):
+                ISPYB_LOGGER.info("Screening collection - using new DCG")
+                dcgid = None
+                self.last_sample_id = None
+            else:
+                ISPYB_LOGGER.info(
+                    f"Collection is {self.params.ispyb_experiment_type} - storing sampleID to bundle images"
+                )
+                self.last_sample_id = self.params.sample_id
+            self.ispyb = StoreInIspyb(self.ispyb_config)
+            ISPYB_LOGGER.info("Beginning ispyb deposition")
+            data_collection_group_info = populate_data_collection_group(self.params)
+            data_collection_info = populate_data_collection_info_for_rotation(
+                cast(RotationScan, self.params)
+            )
+            data_collection_info = populate_remaining_data_collection_info(
+                self.params.comment,
+                dcgid,
+                data_collection_info,
+                self.params,
+            )
+            data_collection_info.parent_id = dcgid
+            scan_data_info = ScanDataInfo(
+                data_collection_info=data_collection_info,
+            )
+            self.ispyb_ids = self.ispyb.begin_deposition(
+                data_collection_group_info, [scan_data_info]
+            )
+        ISPYB_LOGGER.info("ISPYB handler received start document.")
+        if doc.get("subplan_name") == CONST.PLAN.ROTATION_MAIN:
+            self.uid_to_finalize_on = doc.get("uid")
+        return super().activity_gated_start(doc)
+
+    def populate_info_for_update(
+        self,
+        event_sourced_data_collection_info: DataCollectionInfo,
+        event_sourced_position_info: DataCollectionPositionInfo | None,
+        params,
+    ) -> Sequence[ScanDataInfo]:
+        assert (
+            self.ispyb_ids.data_collection_ids
+        ), "Expect an existing DataCollection to update"
+
+        return [
+            ScanDataInfo(
+                data_collection_info=event_sourced_data_collection_info,
+                data_collection_id=self.ispyb_ids.data_collection_ids[0],
+                data_collection_position_info=event_sourced_position_info,
+            )
+        ]
+
+    def _handle_ispyb_hardware_read(self, doc: Event):
+        """Use the hardware read values to create the ispyb comment"""
+        scan_data_infos = super()._handle_ispyb_hardware_read(doc)
+        motor_positions_mm = [
+            doc["data"]["smargon-x"],
+            doc["data"]["smargon-y"],
+            doc["data"]["smargon-z"],
+        ]
+        assert (
+            self.params
+        ), "handle_ispyb_hardware_read triggered before activity_gated_start"
+        motor_positions_um = [position * 1000 for position in motor_positions_mm]
+        comment = f"Sample position (µm): ({motor_positions_um[0]:.0f}, {motor_positions_um[1]:.0f}, {motor_positions_um[2]:.0f}) {self.params.comment} "
+        scan_data_infos[0].data_collection_info.comments = comment
+        return scan_data_infos
+
+    def activity_gated_event(self, doc: Event):
+        doc = super().activity_gated_event(doc)
+        set_dcgid_tag(self.ispyb_ids.data_collection_group_id)
+
+        descriptor_name = self.descriptors[doc["descriptor"]].get("name")
+        if descriptor_name == CONST.DESCRIPTORS.OAV_ROTATION_SNAPSHOT_TRIGGERED:
+            scan_data_infos = self._handle_oav_rotation_snapshot_triggered(doc)
+            self.ispyb_ids = self.ispyb.update_deposition(
+                self.ispyb_ids, scan_data_infos
+            )
+
+        return doc
+
+    def _handle_oav_rotation_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
+        assert self.ispyb_ids.data_collection_ids, "No current data collection"
+        assert self.params, "ISPyB handler didn't receive parameters!"
+        data = doc["data"]
+        self._oav_snapshot_event_idx += 1
+        data_collection_info = DataCollectionInfo(
+            **{
+                f"xtal_snapshot{self._oav_snapshot_event_idx}": data.get(
+                    "oav_snapshot_last_saved_path"
+                )
+            }
+        )
+        scan_data_info = ScanDataInfo(
+            data_collection_id=self.ispyb_ids.data_collection_ids[-1],
+            data_collection_info=data_collection_info,
+        )
+        return [scan_data_info]
+
+    def activity_gated_stop(self, doc: RunStop) -> RunStop:
+        if doc.get("run_start") == self.uid_to_finalize_on:
+            self.uid_to_finalize_on = None
+            return super().activity_gated_stop(doc)
+        return self._tag_doc(doc)

mx_bluesky/hyperion/external_interaction/callbacks/rotation/ispyb_mapping.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from mx_bluesky.hyperion.external_interaction.ispyb.data_model import DataCollectionInfo
+from mx_bluesky.hyperion.parameters.rotation import RotationScan
+
+
+def populate_data_collection_info_for_rotation(params: RotationScan):
+    info = DataCollectionInfo(
+        omega_start=params.omega_start_deg,
+        data_collection_number=params.detector_params.run_number,  # type:ignore # the validator always makes this int
+        n_images=params.num_images,
+        axis_range=params.rotation_increment_deg,
+        axis_start=params.omega_start_deg,
+        axis_end=(params.omega_start_deg + params.scan_width_deg),
+        kappa_start=params.kappa_start_deg,
+    )
+    return info

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

@@ -0,0 +1,102 @@
+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.rotation import RotationScan
+
+from ..logging_callback import format_doc_for_log
+
+if TYPE_CHECKING:
+    from event_model.documents import Event, EventDescriptor, RunStart
+
+
+class RotationNexusFileCallback(PlanReactiveCallback):
+    """Callback class to handle the creation of Nexus files based on experiment
+    parameters for rotation scans
+
+    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_uid: str | None = None
+        self.writer: NexusWriter | None = None
+        self.descriptors: dict[str, EventDescriptor] = {}
+        # used when multiple collections are made in one detector arming event:
+        self.full_num_of_images: int | None = None
+        self.meta_data_run_number: int | None = None
+
+    def activity_gated_descriptor(self, doc: EventDescriptor):
+        self.descriptors[doc["uid"]] = doc
+
+    def activity_gated_event(self, doc: Event):
+        event_descriptor = self.descriptors.get(doc["descriptor"])
+        if event_descriptor is None:
+            NEXUS_LOGGER.warning(
+                f"Rotation Nexus handler {self} received event doc {format_doc_for_log(doc)} and "
+                "has no corresponding descriptor record"
+            )
+            return doc
+        if event_descriptor.get("name") == CONST.DESCRIPTORS.HARDWARE_READ_DURING:
+            NEXUS_LOGGER.info(
+                f"Nexus handler received event from read hardware {format_doc_for_log(doc)}"
+            )
+            data = doc["data"]
+            assert self.writer, "Nexus writer not initialised"
+            (
+                self.writer.beam,
+                self.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"])
+            self.writer.create_nexus_file(vds_data_type)
+            NEXUS_LOGGER.info(f"Nexus file created at {self.writer.data_filename}")
+        return doc
+
+    def activity_gated_start(self, doc: RunStart):
+        if doc.get("subplan_name") == CONST.PLAN.ROTATION_MULTI:
+            self.full_num_of_images = doc.get("full_num_of_images")
+            self.meta_data_run_number = doc.get("meta_data_run_number")
+        if doc.get("subplan_name") == CONST.PLAN.ROTATION_OUTER:
+            self.run_uid = doc.get("uid")
+            json_params = doc.get("hyperion_parameters")
+            NEXUS_LOGGER.info(
+                f"Nexus writer received start document with experiment parameters {json_params}"
+            )
+            parameters = RotationScan.from_json(json_params)
+            NEXUS_LOGGER.info("Setting up nexus file...")
+            det_size = (
+                parameters.detector_params.detector_size_constants.det_size_pixels
+            )
+            shape = (parameters.num_images, det_size.width, det_size.height)
+            self.writer = NexusWriter(
+                parameters,
+                shape,
+                parameters.scan_points,
+                omega_start_deg=parameters.omega_start_deg,
+                chi_start_deg=parameters.chi_start_deg or 0,
+                phi_start_deg=parameters.phi_start_deg or 0,
+                vds_start_index=parameters.nexus_vds_start_img,
+                full_num_of_images=self.full_num_of_images,
+                meta_data_run_number=self.meta_data_run_number,
+                rotation_direction=parameters.rotation_direction,
+            )

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

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from time import time
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from blueapi.core import MsgGenerator
+from bluesky import preprocessors as bpp
+from dodal.devices.zocalo.zocalo_results import ZOCALO_READING_PLAN_NAME
+
+from mx_bluesky.hyperion.external_interaction.callbacks.common.ispyb_mapping import (
+    populate_data_collection_group,
+    populate_remaining_data_collection_info,
+)
+from mx_bluesky.hyperion.external_interaction.callbacks.ispyb_callback_base import (
+    BaseISPyBCallback,
+)
+from mx_bluesky.hyperion.external_interaction.callbacks.logging_callback import (
+    format_doc_for_log,
+)
+from mx_bluesky.hyperion.external_interaction.callbacks.xray_centre.ispyb_mapping import (
+    construct_comment_for_gridscan,
+    populate_xy_data_collection_info,
+    populate_xz_data_collection_info,
+)
+from mx_bluesky.hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
+from mx_bluesky.hyperion.external_interaction.ispyb.data_model import (
+    DataCollectionGridInfo,
+    DataCollectionInfo,
+    DataCollectionPositionInfo,
+    Orientation,
+    ScanDataInfo,
+)
+from mx_bluesky.hyperion.external_interaction.ispyb.ispyb_store import (
+    IspybIds,
+    StoreInIspyb,
+)
+from mx_bluesky.hyperion.log import ISPYB_LOGGER, set_dcgid_tag
+from mx_bluesky.hyperion.parameters.components import DiffractionExperimentWithSample
+from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.parameters.gridscan import (
+    GridCommon,
+)
+
+if TYPE_CHECKING:
+    from event_model import Event, RunStart, RunStop
+
+
+def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters):
+    return bpp.run_wrapper(
+        plan_generator,
+        md={
+            "activate_callbacks": ["GridscanISPyBCallback"],
+            "subplan_name": CONST.PLAN.GRID_DETECT_AND_DO_GRIDSCAN,
+            "hyperion_parameters": parameters.model_dump_json(),
+        },
+    )
+
+
+class GridscanISPyBCallback(BaseISPyBCallback):
+    """Callback class to handle the deposition of experiment parameters into the ISPyB
+    database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on
+    recieving an 'event' document for the 'ispyb_reading_hardware' event, and updates the
+    deposition on recieving its final 'stop' document.
+
+    To use, subscribe the Bluesky RunEngine to an instance of this class.
+    E.g.:
+        ispyb_handler_callback = FGSISPyBCallback(parameters)
+        RE.subscribe(ispyb_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,
+        *,
+        emit: Callable[..., Any] | None = None,
+    ) -> None:
+        super().__init__(emit=emit)
+        self.ispyb: StoreInIspyb
+        self.ispyb_ids: IspybIds = IspybIds()
+        self._start_of_fgs_uid: str | None = None
+        self._processing_start_time: float | None = None
+
+    def activity_gated_start(self, doc: RunStart):
+        if doc.get("subplan_name") == CONST.PLAN.DO_FGS:
+            self._start_of_fgs_uid = doc.get("uid")
+        if doc.get("subplan_name") == CONST.PLAN.GRID_DETECT_AND_DO_GRIDSCAN:
+            self.uid_to_finalize_on = doc.get("uid")
+            ISPYB_LOGGER.info(
+                "ISPyB callback received start document with experiment parameters and "
+                f"uid: {self.uid_to_finalize_on}"
+            )
+            self.params = GridCommon.from_json(doc.get("hyperion_parameters"))
+            self.ispyb = StoreInIspyb(self.ispyb_config)
+            data_collection_group_info = populate_data_collection_group(self.params)
+
+            scan_data_infos = [
+                ScanDataInfo(
+                    data_collection_info=populate_remaining_data_collection_info(
+                        None,
+                        None,
+                        populate_xy_data_collection_info(
+                            self.params.detector_params,
+                        ),
+                        self.params,
+                    ),
+                ),
+                ScanDataInfo(
+                    data_collection_info=populate_remaining_data_collection_info(
+                        None,
+                        None,
+                        populate_xz_data_collection_info(self.params.detector_params),
+                        self.params,
+                    )
+                ),
+            ]
+
+            self.ispyb_ids = self.ispyb.begin_deposition(
+                data_collection_group_info, scan_data_infos
+            )
+            set_dcgid_tag(self.ispyb_ids.data_collection_group_id)
+        return super().activity_gated_start(doc)
+
+    def activity_gated_event(self, doc: Event):
+        doc = super().activity_gated_event(doc)
+
+        descriptor_name = self.descriptors[doc["descriptor"]].get("name")
+        if descriptor_name == ZOCALO_READING_PLAN_NAME:
+            self._handle_zocalo_read_event(doc)
+        elif descriptor_name == CONST.DESCRIPTORS.OAV_GRID_SNAPSHOT_TRIGGERED:
+            scan_data_infos = self._handle_oav_grid_snapshot_triggered(doc)
+            self.ispyb_ids = self.ispyb.update_deposition(
+                self.ispyb_ids, scan_data_infos
+            )
+
+        return doc
+
+    def _handle_zocalo_read_event(self, doc):
+        crystal_summary = ""
+        if self._processing_start_time is not None:
+            proc_time = time() - self._processing_start_time
+            crystal_summary = f"Zocalo processing took {proc_time:.2f} s. "
+        bboxes: list[np.ndarray] = []
+        ISPYB_LOGGER.info(
+            f"Amending comment based on Zocalo reading doc: {format_doc_for_log(doc)}"
+        )
+        raw_results = doc["data"]["zocalo-results"]
+        if len(raw_results) > 0:
+            for n, res in enumerate(raw_results):
+                bb = res["bounding_box"]
+                diff = np.array(bb[1]) - np.array(bb[0])
+                bboxes.append(diff)
+
+                nicely_formatted_com = [
+                    f"{np.round(com, 2)}" for com in res["centre_of_mass"]
+                ]
+                crystal_summary += (
+                    f"Crystal {n + 1}: "
+                    f"Strength {res['total_count']}; "
+                    f"Position (grid boxes) {nicely_formatted_com}; "
+                    f"Size (grid boxes) {bboxes[n]}; "
+                )
+        else:
+            crystal_summary += "Zocalo found no crystals in this gridscan."
+        assert (
+            self.ispyb_ids.data_collection_ids
+        ), "No data collection to add results to"
+        self.ispyb.append_to_comment(
+            self.ispyb_ids.data_collection_ids[0], crystal_summary
+        )
+
+    def _handle_oav_grid_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
+        assert self.ispyb_ids.data_collection_ids, "No current data collection"
+        assert self.params, "ISPyB handler didn't receive parameters!"
+        data = doc["data"]
+        data_collection_id = None
+        data_collection_info = DataCollectionInfo(
+            xtal_snapshot1=data.get("oav_grid_snapshot_last_path_full_overlay"),
+            xtal_snapshot2=data.get("oav_grid_snapshot_last_path_outer"),
+            xtal_snapshot3=data.get("oav_grid_snapshot_last_saved_path"),
+            n_images=(
+                data["oav_grid_snapshot_num_boxes_x"]
+                * data["oav_grid_snapshot_num_boxes_y"]
+            ),
+        )
+        microns_per_pixel_x = data["oav_grid_snapshot_microns_per_pixel_x"]
+        microns_per_pixel_y = data["oav_grid_snapshot_microns_per_pixel_y"]
+        data_collection_grid_info = DataCollectionGridInfo(
+            dx_in_mm=data["oav_grid_snapshot_box_width"] * microns_per_pixel_x / 1000,
+            dy_in_mm=data["oav_grid_snapshot_box_width"] * microns_per_pixel_y / 1000,
+            steps_x=data["oav_grid_snapshot_num_boxes_x"],
+            steps_y=data["oav_grid_snapshot_num_boxes_y"],
+            microns_per_pixel_x=microns_per_pixel_x,
+            microns_per_pixel_y=microns_per_pixel_y,
+            snapshot_offset_x_pixel=int(data["oav_grid_snapshot_top_left_x"]),
+            snapshot_offset_y_pixel=int(data["oav_grid_snapshot_top_left_y"]),
+            orientation=Orientation.HORIZONTAL,
+            snaked=True,
+        )
+        data_collection_info.comments = construct_comment_for_gridscan(
+            data_collection_grid_info
+        )
+        if len(self.ispyb_ids.data_collection_ids) > self._oav_snapshot_event_idx:
+            data_collection_id = self.ispyb_ids.data_collection_ids[
+                self._oav_snapshot_event_idx
+            ]
+        self._populate_axis_info(data_collection_info, doc["data"]["smargon-omega"])
+
+        scan_data_info = ScanDataInfo(
+            data_collection_info=data_collection_info,
+            data_collection_id=data_collection_id,
+            data_collection_grid_info=data_collection_grid_info,
+        )
+        ISPYB_LOGGER.info("Updating ispyb data collection after oav snapshot.")
+        self._oav_snapshot_event_idx += 1
+        return [scan_data_info]
+
+    def _populate_axis_info(
+        self, data_collection_info: DataCollectionInfo, omega_start: float | None
+    ):
+        if omega_start is not None:
+            omega_in_gda_space = -omega_start
+            data_collection_info.omega_start = omega_in_gda_space
+            data_collection_info.axis_start = omega_in_gda_space
+            data_collection_info.axis_end = omega_in_gda_space
+            data_collection_info.axis_range = 0
+
+    def populate_info_for_update(
+        self,
+        event_sourced_data_collection_info: DataCollectionInfo,
+        event_sourced_position_info: DataCollectionPositionInfo | None,
+        params: DiffractionExperimentWithSample,
+    ) -> Sequence[ScanDataInfo]:
+        assert (
+            self.ispyb_ids.data_collection_ids
+        ), "Expect at least one valid data collection to record scan data"
+        xy_scan_data_info = ScanDataInfo(
+            data_collection_info=event_sourced_data_collection_info,
+            data_collection_id=self.ispyb_ids.data_collection_ids[0],
+        )
+        scan_data_infos = [xy_scan_data_info]
+
+        data_collection_id = (
+            self.ispyb_ids.data_collection_ids[1]
+            if len(self.ispyb_ids.data_collection_ids) > 1
+            else None
+        )
+        xz_scan_data_info = ScanDataInfo(
+            data_collection_info=event_sourced_data_collection_info,
+            data_collection_id=data_collection_id,
+        )
+        scan_data_infos.append(xz_scan_data_info)
+        return scan_data_infos
+
+    def activity_gated_stop(self, doc: RunStop) -> RunStop:
+        if doc.get("run_start") == self._start_of_fgs_uid:
+            self._processing_start_time = time()
+        if doc.get("run_start") == self.uid_to_finalize_on:
+            ISPYB_LOGGER.info(
+                "ISPyB callback received stop document corresponding to start document "
+                f"with uid: {self.uid_to_finalize_on}."
+            )
+            if self.ispyb_ids == IspybIds():
+                raise ISPyBDepositionNotMade("ispyb was not initialised at run start")
+            return super().activity_gated_stop(doc)
+        return self._tag_doc(doc)

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

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import numpy
+from dodal.devices.detector import DetectorParams
+from dodal.devices.oav import utils as oav_utils
+
+from mx_bluesky.hyperion.external_interaction.ispyb.data_model import (
+    DataCollectionGridInfo,
+    DataCollectionInfo,
+)
+
+
+def populate_xz_data_collection_info(detector_params: DetectorParams):
+    assert (
+        detector_params.omega_start is not None
+        and detector_params.run_number is not None
+    ), "StoreGridscanInIspyb failed to get parameters"
+    run_number = detector_params.run_number + 1
+    info = DataCollectionInfo(
+        data_collection_number=run_number,
+    )
+    return info
+
+
+def populate_xy_data_collection_info(detector_params: DetectorParams):
+    return DataCollectionInfo(
+        data_collection_number=detector_params.run_number,
+    )
+
+
+def construct_comment_for_gridscan(grid_info: DataCollectionGridInfo) -> str:
+    assert grid_info is not None, "StoreGridScanInIspyb failed to get parameters"
+
+    bottom_right = oav_utils.bottom_right_from_top_left(
+        numpy.array(
+            [grid_info.snapshot_offset_x_pixel, grid_info.snapshot_offset_y_pixel]
+        ),  # type: ignore
+        grid_info.steps_x,
+        grid_info.steps_y,
+        grid_info.dx_in_mm,
+        grid_info.dy_in_mm,
+        grid_info.microns_per_pixel_x,
+        grid_info.microns_per_pixel_y,
+    )
+    return (
+        "Hyperion: Xray centring - Diffraction grid scan of "
+        f"{grid_info.steps_x} by "
+        f"{grid_info.steps_y} images in "
+        f"{(grid_info.dx_in_mm * 1e3):.1f} um by "
+        f"{(grid_info.dy_in_mm * 1e3):.1f} um steps. "
+        f"Top left (px): [{int(grid_info.snapshot_offset_x_pixel)},{int(grid_info.snapshot_offset_y_pixel)}], "
+        f"bottom right (px): [{bottom_right[0]},{bottom_right[1]}]."
+    )