From 9044b57f554f8ad6b4326f2f3b10b988a268b736 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Thu, 19 Feb 2026 11:05:12 +0000 Subject: [PATCH 1/8] wip --- .../beamlines/i02_1/i02_1_gridscan_plan.py | 29 +++++++++++++ .../callbacks/common/ispyb_callback_base.py | 42 ++++++++++++------- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index fe19a767c5..7cf26ab848 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -1,4 +1,5 @@ from functools import partial +from pathlib import Path import bluesky.preprocessors as bpp import pydantic @@ -14,6 +15,9 @@ from dodal.devices.s4_slit_gaps import S4SlitGaps from dodal.devices.undulator import BaseUndulator from dodal.devices.zebra.zebra import Zebra +from pydantic import BaseModel, PrivateAttr +from pydantic_extra_types.semantic_version import SemanticVersion +from semver import Version from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import ( setup_zebra_for_gridscan, @@ -122,6 +126,31 @@ def _tidy_plan( yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra) +PARAMETER_VERSION = Version.parse("1.0.0") + + +def get_internal_param_version() -> SemanticVersion: + return SemanticVersion.validate_from_str(str(PARAMETER_VERSION)) + + +class ExternalGridScanParams(BaseModel): + gda_parameter_version: SemanticVersion + visit: str + file_name: str + storage_directory: Path + exposure_time_s: float + snapshot_directory: Path + x_start_um: float + y_start_um: float + z_start_um: float + x_steps: int + y_steps: int + sample_id: int | None = None + + # Internal parameter version compatible with this external model + _internal_param_version: str = PrivateAttr(default="6.0.0") + + def i02_1_gridscan_plan( parameters: SpecifiedTwoDGridScan, composite: FlyScanXRayCentreComposite = inject(""), diff --git a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py index 028e630a88..19a2b4ac35 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py @@ -129,25 +129,33 @@ def activity_gated_event(self, doc: Event) -> Event: return self.tag_doc(doc) def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]: + _data = doc["data"] + assert self.params, "Event handled before activity_gated_start received params" ISPYB_ZOCALO_CALLBACK_LOGGER.info( "ISPyB handler received event from read hardware" ) - synchrotron_mode = doc["data"]["synchrotron-synchrotron_mode"] + synchrotron_mode = _data["synchrotron-synchrotron_mode"] assert isinstance(synchrotron_mode, SynchrotronMode) hwscan_data_collection_info = DataCollectionInfo( - undulator_gap1=doc["data"]["undulator-current_gap"], + undulator_gap1=_data["undulator-current_gap"], synchrotron_mode=synchrotron_mode.value, - slitgap_horizontal=doc["data"]["s4_slit_gaps-xgap"], - slitgap_vertical=doc["data"]["s4_slit_gaps-ygap"], + slitgap_horizontal=_data["s4_slit_gaps-xgap"], + slitgap_vertical=_data["s4_slit_gaps-ygap"], ) hwscan_data_collection_info = _update_based_on_energy( doc, self.params.detector_params, hwscan_data_collection_info ) + + # VMXm doesn't have a gonio-y position, allow None + pos_y = _data.get("gonio-y", None) + if pos_y: + pos_y = float(pos_y) + hwscan_position_info = DataCollectionPositionInfo( - pos_x=float(doc["data"]["gonio-x"]), - pos_y=float(doc["data"]["gonio-y"]), - pos_z=float(doc["data"]["gonio-z"]), + pos_x=float(_data["gonio-x"]), + pos_y=pos_y, + pos_z=float(_data["gonio-z"]), ) scan_data_infos = self.populate_info_for_update( hwscan_data_collection_info, hwscan_position_info, self.params @@ -160,18 +168,24 @@ def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]: def _handle_ispyb_transmission_flux_read( self, doc: Event ) -> Sequence[ScanDataInfo]: + _data = doc["data"] + assert self.params - aperture = doc["data"]["aperture_scatterguard-selected_aperture"] - beamsize_x_mm = doc["data"]["beamsize-x_um"] / 1000 - beamsize_y_mm = doc["data"]["beamsize-y_um"] / 1000 + aperture = _data.get( + "aperture_scatterguard-selected_aperture", "Not implemented" + ) + beamsize_x_mm = ( + _data["beamsize-x_um"] / 1000 + ) # todo beamsize for vmxm is complicated. Do hacky thing where beamsize is read from fake device until we make real one + beamsize_y_mm = _data["beamsize-y_um"] / 1000 hwscan_data_collection_info = DataCollectionInfo( beamsize_at_samplex=beamsize_x_mm, beamsize_at_sampley=beamsize_y_mm, - flux=doc["data"]["flux-flux_reading"], - detector_mode="ROI" if doc["data"]["eiger_cam_roi_mode"] else "FULL", - ispyb_detector_id=doc["data"]["eiger-ispyb_detector_id"], + flux=_data["flux-flux_reading"], + detector_mode="ROI" if _data["eiger_cam_roi_mode"] else "FULL", + ispyb_detector_id=_data["eiger-ispyb_detector_id"], ) - if transmission := doc["data"]["attenuator-actual_transmission"]: + if transmission := _data["attenuator-actual_transmission"]: # Ispyb wants the transmission in a percentage, we use fractions hwscan_data_collection_info.transmission = transmission * 100 hwscan_data_collection_info = _update_based_on_energy( From d40ef7962bf7be968be335424b5547d9ba68ec3d Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Mon, 23 Feb 2026 17:10:58 +0000 Subject: [PATCH 2/8] wip --- .../i02_1/external_interaction/__init__.py | 0 .../callbacks/__init__.py | 0 .../callbacks/gridscan/__init__.py | 0 .../callbacks/gridscan/ispyb_callback.py | 315 ++++++++++++++++++ .../beamlines/i02_1/i02_1_gridscan_plan.py | 15 +- .../callbacks/common/ispyb_callback_base.py | 29 +- 6 files changed, 348 insertions(+), 11 deletions(-) create mode 100644 src/mx_bluesky/beamlines/i02_1/external_interaction/__init__.py create mode 100644 src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/__init__.py create mode 100644 src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/__init__.py create mode 100644 src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/__init__.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/__init__.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/__init__.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py new file mode 100644 index 0000000000..9859ba0ef1 --- /dev/null +++ b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from time import time +from typing import TYPE_CHECKING, Any, TypeVar + +from bluesky import preprocessors as bpp +from bluesky.utils import MsgGenerator, make_decorator + +from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import I02_1FgsParams +from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import ( + BaseISPyBCallback, + D, +) +from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import ( + populate_data_collection_group, + populate_remaining_data_collection_info, +) +from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( + GridscanPlane, + _smargon_omega_to_xyxz_plane, +) +from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping import ( + construct_comment_for_gridscan, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionGridInfo, + DataCollectionGroupInfo, + DataCollectionInfo, + DataCollectionPositionInfo, + Orientation, + ScanDataInfo, +) +from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( + IspybIds, + StoreInIspyb, +) +from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample +from mx_bluesky.common.parameters.constants import PlanNameConstants +from mx_bluesky.common.utils.exceptions import ( + ISPyBDepositionNotMadeError, + SampleError, +) +from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER, set_dcgid_tag + +if TYPE_CHECKING: + from event_model import RunStart, RunStop + +T = TypeVar("T", bound="I02_1FgsParams") +ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" + + +def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): + return bpp.set_run_key_wrapper( + bpp.run_wrapper( + plan_generator, + md={ + "activate_callbacks": ["GridscanISPyBCallback"], + "subplan_name": PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, + "mx_bluesky_parameters": parameters.model_dump_json(), + }, + ), + PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, + ) + + +ispyb_activation_decorator = make_decorator(ispyb_activation_wrapper) + + +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 + receiving an 'event' document for the 'ispyb_reading_hardware' event, and updates the + deposition on receiving its final 'stop' document. + + To use, subscribe the Bluesky RunEngine to an instance of this class. + E.g.: + ispyb_handler_callback = FGSISPyBCallback(parameters) + run_engine.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, + param_type: type[T], + *, + emit: Callable[..., Any] | None = None, + ) -> None: + super().__init__(emit=emit) + self.ispyb: StoreInIspyb + self.param_type = param_type + self._start_of_fgs_uid: str | None = None + self._processing_start_time: float | None = None + self._grid_plane_to_id_map: dict[GridscanPlane, int] = {} + self._grid_plane_to_width_map: dict[GridscanPlane, int] = {} + self.data_collection_group_info: DataCollectionGroupInfo | None + + def activity_gated_start(self, doc: RunStart): + if doc.get("subplan_name") == PlanNameConstants.DO_FGS: + self._start_of_fgs_uid = doc.get("uid") + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "ISPyB callback received start document with experiment parameters and " + f"uid: {self._start_of_fgs_uid}" + ) + mx_bluesky_parameters = doc.get("mx_bluesky_parameters") + assert isinstance(mx_bluesky_parameters, str) + self.params = self.param_type.model_validate_json(mx_bluesky_parameters) + assert isinstance(self.params, I02_1FgsParams) + self.ispyb = StoreInIspyb(self.ispyb_config) + self.data_collection_group_info = populate_data_collection_group( + self.params + ) + + # todo fix this: define scan_data_infos here and then overwrite later + scan_data_infos = [] + assert self.params.num_grids > 0 + for grid in range(self.params.num_grids): + scan_data_infos.append( + ScanDataInfo( + data_collection_info=populate_remaining_data_collection_info( + f"MX-Bluesky: Xray centring {grid + 1}/{self.params.num_grids} -", + None, + DataCollectionInfo(), + self.params, + ) + ) + ) + + # todo make a function which populates all of this. "fill deposition with grid info" + + self.ispyb_ids = self.ispyb.begin_deposition( + self.data_collection_group_info, scan_data_infos + ) + # Use grid information given by GDA to complete ispyb info + scan_data_infos = self._get_scan_infos(doc) + self.ispyb_ids = self.ispyb.update_deposition( + self.ispyb_ids, scan_data_infos + ) + self.ispyb.update_data_collection_group_table( + self.data_collection_group_info, self.ispyb_ids.data_collection_group_id + ) + + set_dcgid_tag(self.ispyb_ids.data_collection_group_id) + return super().activity_gated_start(doc) + + def _add_processing_time_to_comment(self, processing_start_time: float): + assert self.data_collection_group_info, ASSERT_START_BEFORE_EVENT_DOC_MESSAGE + proc_time = time() - processing_start_time + crystal_summary = f"Zocalo processing took {proc_time:.2f} s." + + self.data_collection_group_info.comments = ( + self.data_collection_group_info.comments or "" + ) + crystal_summary + + self.ispyb.append_to_comment( + self.ispyb_ids.data_collection_ids[0], crystal_summary + ) + + def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: + """ + so grid information is available immediately after the plan is triggered. + In contrast, i03 and i04 use the OAV to automatically detect their grid. + """ + assert isinstance(self.params, I02_1FgsParams) + assert self.ispyb_ids.data_collection_ids, "No current data collection" + assert self.data_collection_group_info, "No data collection group" + data = doc["data"] + scan_data_infos = [] + + for grid_num in range(self.params.num_grids): + omega = data.get("gonio-omega", self.params.omega_starts_deg[grid_num]) + + # Don't need to do deal with the grid plane here since vmxm only do + # one plane, but leave it in so it's easier to standardise in the future + grid_plane = _smargon_omega_to_xyxz_plane(omega) + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + f"Generating dc info for gridplane {grid_plane}, omega {omega}" + ) + data_collection_number = self.data_collection_number_from_gridplane( + grid_plane + ) + file_template = f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5" + # Snapshots have already been taken in GDA + data_collection_info = DataCollectionInfo( + xtal_snapshot1=str(self.params.path_to_xtal_snapshot), + xtal_snapshot2=str(self.params.path_to_xtal_snapshot), + xtal_snapshot3=str(self.params.path_to_xtal_snapshot), + n_images=self.params.num_images, + data_collection_number=data_collection_number, + file_template=file_template, + ) + data_collection_grid_info = DataCollectionGridInfo( + dx_in_mm=self.params.x_step_size_um * 1000, + dy_in_mm=self.params.y_step_sizes_um[grid_num] * 1000, + steps_x=self.params.x_steps, + steps_y=self.params.y_steps[grid_num], + microns_per_pixel_x=self.params.microns_per_pixel_x, + microns_per_pixel_y=self.params.microns_per_pixel_y, + snapshot_offset_x_pixel=self.params.upper_left_x, + snapshot_offset_y_pixel=self.params.upper_left_y, + orientation=Orientation.HORIZONTAL, + snaked=True, + ) + data_collection_info.comments = construct_comment_for_gridscan( + data_collection_grid_info + ) + + # Grid plane logic isn't needed for VMX, but keep it for now anyway + data_collection_id = self.ispyb_ids.data_collection_ids[ + 0 if grid_plane == GridscanPlane.OMEGA_XY else 1 + ] + self._grid_plane_to_id_map[grid_plane] = data_collection_id + self._grid_plane_to_width_map[grid_plane] = ( + data_collection_grid_info.steps_y + ) + + y_steps = self._grid_plane_to_width_map.get(GridscanPlane.OMEGA_XY, "_") + self.data_collection_group_info.comments = ( + f"Diffraction grid scan of {data_collection_grid_info.steps_x} by " + f"{y_steps}." + ) + + self._populate_axis_info(data_collection_info, doc["data"]) + + # todo do all this stuff as soon as possible after plan starts + + scan_data_info = ScanDataInfo( + data_collection_info=data_collection_info, + data_collection_id=data_collection_id, + data_collection_grid_info=data_collection_grid_info, + ) + + scan_data_infos.append(scan_data_info) + + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "Updating ispyb data collection after loading grid params" + ) + + return scan_data_infos + + def _populate_axis_info(self, data_collection_info: DataCollectionInfo, doc: dict): + if (omega_start := doc.get("gonio-omega")) 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 + if (chi_start := doc.get("gonio-chi")) is not None: + data_collection_info.chi_start = chi_start + + 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" + ) + assert isinstance(self.params, I02_1FgsParams) + scan_data_infos = [] + for grid_num in range(self.params.num_grids): + scan_data_info = ScanDataInfo( + data_collection_info=event_sourced_data_collection_info, + data_collection_id=self.ispyb_ids.data_collection_ids[grid_num], + ) + scan_data_infos.append(scan_data_info) + return scan_data_infos + + def activity_gated_stop(self, doc: RunStop) -> RunStop: + assert self.data_collection_group_info, ( + f"No data collection group info - stop document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" + ) + if doc.get("run_start") == self._start_of_fgs_uid: + self._processing_start_time = time() + if doc.get("run_start") == self._start_of_fgs_uid: + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "ISPyB callback received stop document corresponding to start document " + f"with uid: {self._start_of_fgs_uid}." + ) + if self.ispyb_ids == IspybIds(): + raise ISPyBDepositionNotMadeError( + "ispyb was not initialised at run start" + ) + exception_type, message = SampleError.type_and_message_from_reason( + doc.get("reason", "") + ) + if exception_type: + doc["reason"] = message + self.data_collection_group_info.comments = message + elif self._processing_start_time: + self._add_processing_time_to_comment(self._processing_start_time) + self.ispyb.update_data_collection_group_table( + self.data_collection_group_info, + self.ispyb_ids.data_collection_group_id, + ) + self.data_collection_group_info = None + self._grid_plane_to_id_map.clear() + self._grid_plane_to_width_map.clear() + return super().activity_gated_stop(doc) + return self.tag_doc(doc) + + def tag_doc(self, doc: D) -> D: + doc = super().tag_doc(doc) + assert isinstance(doc, dict) + if self._grid_plane_to_id_map: + doc["grid_plane_to_id_map"] = self._grid_plane_to_id_map + return doc # type: ignore + + def data_collection_number_from_gridplane(self, plane) -> int: + assert self.params + base_number = self.params.detector_params.run_number + return base_number if plane == GridscanPlane.OMEGA_XY else base_number + 1 diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index 7cf26ab848..fec811f46c 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -151,8 +151,21 @@ class ExternalGridScanParams(BaseModel): _internal_param_version: str = PrivateAttr(default="6.0.0") +class I02_1FgsParams(SpecifiedTwoDGridScan): # noqa: N801 + """For VMXm gridscans, GDA currently takes the snapshots and provides bluesky with a path, and + sends over the grid parameters""" + + path_to_xtal_snapshot: Path + beam_size_x: float + beam_size_y: float + microns_per_pixel_x: float + microns_per_pixel_y: float + upper_left_x: int # position of X,Y for the top left of the grid, in pixels + upper_left_y: int + + def i02_1_gridscan_plan( - parameters: SpecifiedTwoDGridScan, + parameters: I02_1FgsParams, composite: FlyScanXRayCentreComposite = inject(""), ) -> MsgGenerator: """BlueAPI entry point for i02-1 grid scans""" diff --git a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py index 19a2b4ac35..a6ef1989d5 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py @@ -9,6 +9,7 @@ from dodal.devices.detector.det_resolution import resolution from dodal.devices.synchrotron import SynchrotronMode +from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import I02_1FgsParams from mx_bluesky.common.external_interaction.callbacks.common.plan_reactive_callback import ( PlanReactiveCallback, ) @@ -147,14 +148,9 @@ def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]: doc, self.params.detector_params, hwscan_data_collection_info ) - # VMXm doesn't have a gonio-y position, allow None - pos_y = _data.get("gonio-y", None) - if pos_y: - pos_y = float(pos_y) - hwscan_position_info = DataCollectionPositionInfo( pos_x=float(_data["gonio-x"]), - pos_y=pos_y, + pos_y=float(_data.get("gonio-y")), pos_z=float(_data["gonio-z"]), ) scan_data_infos = self.populate_info_for_update( @@ -174,10 +170,23 @@ def _handle_ispyb_transmission_flux_read( aperture = _data.get( "aperture_scatterguard-selected_aperture", "Not implemented" ) - beamsize_x_mm = ( - _data["beamsize-x_um"] / 1000 - ) # todo beamsize for vmxm is complicated. Do hacky thing where beamsize is read from fake device until we make real one - beamsize_y_mm = _data["beamsize-y_um"] / 1000 + beamsize_x_mm = _data.get("beamsize-x_um", None) + if beamsize_x_mm: + beamsize_x_mm = beamsize_x_mm / 1000 + beamsize_y_mm = _data.get("beamsize-y_um", None) + if beamsize_y_mm: + beamsize_y_mm = beamsize_y_mm / 1000 + if not (beamsize_x_mm and beamsize_y_mm): + # todo write issue about vmxm beamsize and link here + try: + assert isinstance(self.params, I02_1FgsParams) + beamsize_x_mm = self.params.beam_size_x + beamsize_y_mm = self.params.beam_size_y + except AssertionError: + ISPYB_ZOCALO_CALLBACK_LOGGER.warning( + "ISPyB callbacks couldn't get beamsize" + ) + hwscan_data_collection_info = DataCollectionInfo( beamsize_at_samplex=beamsize_x_mm, beamsize_at_sampley=beamsize_y_mm, From 3c7c55e07c1c020501dcfd72528945d75376dbfe Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Wed, 25 Feb 2026 17:07:06 +0000 Subject: [PATCH 3/8] wip --- src/mx_bluesky/beamlines/i02_1/composites.py | 16 + .../callbacks/gridscan/ispyb_callback.py | 315 ----------------- .../beamlines/i02_1/i02_1_gridscan_plan.py | 53 +-- .../i04_grid_detect_then_xray_centre_plan.py | 10 +- ...ommon_grid_detect_then_xray_centre_plan.py | 2 +- .../callbacks/{xray_centre => }/__init__.py | 0 .../callbacks/common/ispyb_callback_base.py | 8 +- .../callbacks/common/ispyb_mapping.py | 2 +- .../callbacks/grid/__init__.py | 0 .../grid/grid_detect_and_scan/__init__.py | 0 .../grid_detect_and_scan}/ispyb_callback.py | 9 +- .../grid_detect_and_scan}/ispyb_mapping.py | 0 .../grid_detect_and_scan}/nexus_callback.py | 0 .../callbacks/grid/gridscan/__init__.py | 0 .../callbacks/grid/gridscan/ispyb_callback.py | 320 ++++++++++++++++++ .../pin_centre_then_gridscan_plan.py | 2 +- .../callbacks/__main__.py | 16 +- tests/conftest.py | 2 +- .../callbacks/test_external_callbacks.py | 2 +- .../test_ispyb_dev_connection.py | 10 +- .../test_load_centre_collect_full_plan.py | 32 +- .../test_zocalo_system.py | 2 +- .../inner_plans/test_do_fgs.py | 2 +- .../test_common_flyscan_xray_centre_plan.py | 20 +- ...ommon_grid_detect_then_xray_centre_plan.py | 6 +- .../test_grid_detection_plan.py | 8 +- .../callbacks/test_zocalo_handler.py | 8 +- .../xray_centre/test_ispyb_callback.py | 24 +- .../xray_centre/test_ispyb_handler.py | 18 +- .../xray_centre/test_ispyb_mapping.py | 4 +- .../xray_centre/test_nexus_handler.py | 8 +- tests/unit_tests/conftest.py | 22 +- .../hyperion/experiment_plans/conftest.py | 6 +- .../test_hyperion_flyscan_xray_centre_plan.py | 12 +- .../test_pin_centre_then_xray_centre_plan.py | 2 +- .../nexus/test_write_nexus.py | 2 +- 36 files changed, 493 insertions(+), 450 deletions(-) create mode 100644 src/mx_bluesky/beamlines/i02_1/composites.py rename src/mx_bluesky/common/external_interaction/callbacks/{xray_centre => }/__init__.py (100%) create mode 100644 src/mx_bluesky/common/external_interaction/callbacks/grid/__init__.py create mode 100644 src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/__init__.py rename src/mx_bluesky/common/external_interaction/callbacks/{xray_centre => grid/grid_detect_and_scan}/ispyb_callback.py (97%) rename src/mx_bluesky/common/external_interaction/callbacks/{xray_centre => grid/grid_detect_and_scan}/ispyb_mapping.py (100%) rename src/mx_bluesky/common/external_interaction/callbacks/{xray_centre => grid/grid_detect_and_scan}/nexus_callback.py (100%) create mode 100644 src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/__init__.py create mode 100644 src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py diff --git a/src/mx_bluesky/beamlines/i02_1/composites.py b/src/mx_bluesky/beamlines/i02_1/composites.py new file mode 100644 index 0000000000..795757bb55 --- /dev/null +++ b/src/mx_bluesky/beamlines/i02_1/composites.py @@ -0,0 +1,16 @@ +from pathlib import Path + +from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan + + +class I02_1FgsParams(SpecifiedTwoDGridScan): # noqa: N801 + """For VMXm gridscans, GDA currently takes the snapshots and provides bluesky with a path, and + sends over the grid parameters""" + + path_to_xtal_snapshot: Path + beam_size_x: float + beam_size_y: float + microns_per_pixel_x: float + microns_per_pixel_y: float + upper_left_x: int # position of X,Y for the top left of the grid, in pixels + upper_left_y: int diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py index 9859ba0ef1..e69de29bb2 100644 --- a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py @@ -1,315 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Sequence -from time import time -from typing import TYPE_CHECKING, Any, TypeVar - -from bluesky import preprocessors as bpp -from bluesky.utils import MsgGenerator, make_decorator - -from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import I02_1FgsParams -from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import ( - BaseISPyBCallback, - D, -) -from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import ( - populate_data_collection_group, - populate_remaining_data_collection_info, -) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanPlane, - _smargon_omega_to_xyxz_plane, -) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping import ( - construct_comment_for_gridscan, -) -from mx_bluesky.common.external_interaction.ispyb.data_model import ( - DataCollectionGridInfo, - DataCollectionGroupInfo, - DataCollectionInfo, - DataCollectionPositionInfo, - Orientation, - ScanDataInfo, -) -from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( - IspybIds, - StoreInIspyb, -) -from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample -from mx_bluesky.common.parameters.constants import PlanNameConstants -from mx_bluesky.common.utils.exceptions import ( - ISPyBDepositionNotMadeError, - SampleError, -) -from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER, set_dcgid_tag - -if TYPE_CHECKING: - from event_model import RunStart, RunStop - -T = TypeVar("T", bound="I02_1FgsParams") -ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" - - -def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): - return bpp.set_run_key_wrapper( - bpp.run_wrapper( - plan_generator, - md={ - "activate_callbacks": ["GridscanISPyBCallback"], - "subplan_name": PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, - "mx_bluesky_parameters": parameters.model_dump_json(), - }, - ), - PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, - ) - - -ispyb_activation_decorator = make_decorator(ispyb_activation_wrapper) - - -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 - receiving an 'event' document for the 'ispyb_reading_hardware' event, and updates the - deposition on receiving its final 'stop' document. - - To use, subscribe the Bluesky RunEngine to an instance of this class. - E.g.: - ispyb_handler_callback = FGSISPyBCallback(parameters) - run_engine.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, - param_type: type[T], - *, - emit: Callable[..., Any] | None = None, - ) -> None: - super().__init__(emit=emit) - self.ispyb: StoreInIspyb - self.param_type = param_type - self._start_of_fgs_uid: str | None = None - self._processing_start_time: float | None = None - self._grid_plane_to_id_map: dict[GridscanPlane, int] = {} - self._grid_plane_to_width_map: dict[GridscanPlane, int] = {} - self.data_collection_group_info: DataCollectionGroupInfo | None - - def activity_gated_start(self, doc: RunStart): - if doc.get("subplan_name") == PlanNameConstants.DO_FGS: - self._start_of_fgs_uid = doc.get("uid") - ISPYB_ZOCALO_CALLBACK_LOGGER.info( - "ISPyB callback received start document with experiment parameters and " - f"uid: {self._start_of_fgs_uid}" - ) - mx_bluesky_parameters = doc.get("mx_bluesky_parameters") - assert isinstance(mx_bluesky_parameters, str) - self.params = self.param_type.model_validate_json(mx_bluesky_parameters) - assert isinstance(self.params, I02_1FgsParams) - self.ispyb = StoreInIspyb(self.ispyb_config) - self.data_collection_group_info = populate_data_collection_group( - self.params - ) - - # todo fix this: define scan_data_infos here and then overwrite later - scan_data_infos = [] - assert self.params.num_grids > 0 - for grid in range(self.params.num_grids): - scan_data_infos.append( - ScanDataInfo( - data_collection_info=populate_remaining_data_collection_info( - f"MX-Bluesky: Xray centring {grid + 1}/{self.params.num_grids} -", - None, - DataCollectionInfo(), - self.params, - ) - ) - ) - - # todo make a function which populates all of this. "fill deposition with grid info" - - self.ispyb_ids = self.ispyb.begin_deposition( - self.data_collection_group_info, scan_data_infos - ) - # Use grid information given by GDA to complete ispyb info - scan_data_infos = self._get_scan_infos(doc) - self.ispyb_ids = self.ispyb.update_deposition( - self.ispyb_ids, scan_data_infos - ) - self.ispyb.update_data_collection_group_table( - self.data_collection_group_info, self.ispyb_ids.data_collection_group_id - ) - - set_dcgid_tag(self.ispyb_ids.data_collection_group_id) - return super().activity_gated_start(doc) - - def _add_processing_time_to_comment(self, processing_start_time: float): - assert self.data_collection_group_info, ASSERT_START_BEFORE_EVENT_DOC_MESSAGE - proc_time = time() - processing_start_time - crystal_summary = f"Zocalo processing took {proc_time:.2f} s." - - self.data_collection_group_info.comments = ( - self.data_collection_group_info.comments or "" - ) + crystal_summary - - self.ispyb.append_to_comment( - self.ispyb_ids.data_collection_ids[0], crystal_summary - ) - - def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: - """ - so grid information is available immediately after the plan is triggered. - In contrast, i03 and i04 use the OAV to automatically detect their grid. - """ - assert isinstance(self.params, I02_1FgsParams) - assert self.ispyb_ids.data_collection_ids, "No current data collection" - assert self.data_collection_group_info, "No data collection group" - data = doc["data"] - scan_data_infos = [] - - for grid_num in range(self.params.num_grids): - omega = data.get("gonio-omega", self.params.omega_starts_deg[grid_num]) - - # Don't need to do deal with the grid plane here since vmxm only do - # one plane, but leave it in so it's easier to standardise in the future - grid_plane = _smargon_omega_to_xyxz_plane(omega) - ISPYB_ZOCALO_CALLBACK_LOGGER.info( - f"Generating dc info for gridplane {grid_plane}, omega {omega}" - ) - data_collection_number = self.data_collection_number_from_gridplane( - grid_plane - ) - file_template = f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5" - # Snapshots have already been taken in GDA - data_collection_info = DataCollectionInfo( - xtal_snapshot1=str(self.params.path_to_xtal_snapshot), - xtal_snapshot2=str(self.params.path_to_xtal_snapshot), - xtal_snapshot3=str(self.params.path_to_xtal_snapshot), - n_images=self.params.num_images, - data_collection_number=data_collection_number, - file_template=file_template, - ) - data_collection_grid_info = DataCollectionGridInfo( - dx_in_mm=self.params.x_step_size_um * 1000, - dy_in_mm=self.params.y_step_sizes_um[grid_num] * 1000, - steps_x=self.params.x_steps, - steps_y=self.params.y_steps[grid_num], - microns_per_pixel_x=self.params.microns_per_pixel_x, - microns_per_pixel_y=self.params.microns_per_pixel_y, - snapshot_offset_x_pixel=self.params.upper_left_x, - snapshot_offset_y_pixel=self.params.upper_left_y, - orientation=Orientation.HORIZONTAL, - snaked=True, - ) - data_collection_info.comments = construct_comment_for_gridscan( - data_collection_grid_info - ) - - # Grid plane logic isn't needed for VMX, but keep it for now anyway - data_collection_id = self.ispyb_ids.data_collection_ids[ - 0 if grid_plane == GridscanPlane.OMEGA_XY else 1 - ] - self._grid_plane_to_id_map[grid_plane] = data_collection_id - self._grid_plane_to_width_map[grid_plane] = ( - data_collection_grid_info.steps_y - ) - - y_steps = self._grid_plane_to_width_map.get(GridscanPlane.OMEGA_XY, "_") - self.data_collection_group_info.comments = ( - f"Diffraction grid scan of {data_collection_grid_info.steps_x} by " - f"{y_steps}." - ) - - self._populate_axis_info(data_collection_info, doc["data"]) - - # todo do all this stuff as soon as possible after plan starts - - scan_data_info = ScanDataInfo( - data_collection_info=data_collection_info, - data_collection_id=data_collection_id, - data_collection_grid_info=data_collection_grid_info, - ) - - scan_data_infos.append(scan_data_info) - - ISPYB_ZOCALO_CALLBACK_LOGGER.info( - "Updating ispyb data collection after loading grid params" - ) - - return scan_data_infos - - def _populate_axis_info(self, data_collection_info: DataCollectionInfo, doc: dict): - if (omega_start := doc.get("gonio-omega")) 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 - if (chi_start := doc.get("gonio-chi")) is not None: - data_collection_info.chi_start = chi_start - - 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" - ) - assert isinstance(self.params, I02_1FgsParams) - scan_data_infos = [] - for grid_num in range(self.params.num_grids): - scan_data_info = ScanDataInfo( - data_collection_info=event_sourced_data_collection_info, - data_collection_id=self.ispyb_ids.data_collection_ids[grid_num], - ) - scan_data_infos.append(scan_data_info) - return scan_data_infos - - def activity_gated_stop(self, doc: RunStop) -> RunStop: - assert self.data_collection_group_info, ( - f"No data collection group info - stop document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" - ) - if doc.get("run_start") == self._start_of_fgs_uid: - self._processing_start_time = time() - if doc.get("run_start") == self._start_of_fgs_uid: - ISPYB_ZOCALO_CALLBACK_LOGGER.info( - "ISPyB callback received stop document corresponding to start document " - f"with uid: {self._start_of_fgs_uid}." - ) - if self.ispyb_ids == IspybIds(): - raise ISPyBDepositionNotMadeError( - "ispyb was not initialised at run start" - ) - exception_type, message = SampleError.type_and_message_from_reason( - doc.get("reason", "") - ) - if exception_type: - doc["reason"] = message - self.data_collection_group_info.comments = message - elif self._processing_start_time: - self._add_processing_time_to_comment(self._processing_start_time) - self.ispyb.update_data_collection_group_table( - self.data_collection_group_info, - self.ispyb_ids.data_collection_group_id, - ) - self.data_collection_group_info = None - self._grid_plane_to_id_map.clear() - self._grid_plane_to_width_map.clear() - return super().activity_gated_stop(doc) - return self.tag_doc(doc) - - def tag_doc(self, doc: D) -> D: - doc = super().tag_doc(doc) - assert isinstance(doc, dict) - if self._grid_plane_to_id_map: - doc["grid_plane_to_id_map"] = self._grid_plane_to_id_map - return doc # type: ignore - - def data_collection_number_from_gridplane(self, plane) -> int: - assert self.params - base_number = self.params.detector_params.run_number - return base_number if plane == GridscanPlane.OMEGA_XY else base_number + 1 diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index fec811f46c..42645ffd53 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -19,6 +19,7 @@ from pydantic_extra_types.semantic_version import SemanticVersion from semver import Version +from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import ( setup_zebra_for_gridscan, tidy_up_zebra_after_gridscan, @@ -31,14 +32,15 @@ ) from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloCallback, + ZocaloInfoGenerator, + ZocaloStartInfo, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, - generate_start_info_from_omega_map, -) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) +from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( + GridscanISPyBCallback, +) from mx_bluesky.common.parameters.constants import ( EnvironmentConstants, PlanNameConstants, @@ -47,8 +49,30 @@ FlyScanEssentialDevices, GonioWithOmegaType, ) -from mx_bluesky.common.parameters.gridscan import GenericGrid from mx_bluesky.common.utils.log import LOGGER +from mx_bluesky.common.utils.utils import number_of_frames_from_scan_spec + + +# todo this all needs chanigng +def generate_start_info_from_omega_map() -> ZocaloInfoGenerator: + """ + Generate the zocalo trigger info from bluesky runs where the frame number is + computed using metadata added to the document by the ISPyB callback and the + run start which together can be used to determine the correct frame numbering. + """ + doc = yield [] + omega_to_scan_spec = doc["omega_to_scan_spec"] + start_frame = 0 + infos = [] + for i, omega in enumerate([0, 90]): + frames = number_of_frames_from_scan_spec(omega_to_scan_spec[omega]) + infos.append( + ZocaloStartInfo( + doc["grid_plane_to_id_map"][omega], None, start_frame, frames, i + ) + ) + start_frame += frames + yield infos def create_gridscan_callbacks() -> tuple[ @@ -57,11 +81,11 @@ def create_gridscan_callbacks() -> tuple[ return ( GridscanNexusFileCallback(param_type=SpecifiedTwoDGridScan), GridscanISPyBCallback( - param_type=GenericGrid, + param_type=I02_1FgsParams, emit=ZocaloCallback( PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV, - generate_start_info_from_omega_map, + generate_start_info_from_omega_map, # todo dont need this ), ), ) @@ -151,19 +175,6 @@ class ExternalGridScanParams(BaseModel): _internal_param_version: str = PrivateAttr(default="6.0.0") -class I02_1FgsParams(SpecifiedTwoDGridScan): # noqa: N801 - """For VMXm gridscans, GDA currently takes the snapshots and provides bluesky with a path, and - sends over the grid parameters""" - - path_to_xtal_snapshot: Path - beam_size_x: float - beam_size_y: float - microns_per_pixel_x: float - microns_per_pixel_y: float - upper_left_x: int # position of X,Y for the top left of the grid, in pixels - upper_left_y: int - - def i02_1_gridscan_plan( parameters: I02_1FgsParams, composite: FlyScanXRayCentreComposite = inject(""), diff --git a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py index 68a8450a49..305fac78d2 100644 --- a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py +++ b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py @@ -62,11 +62,11 @@ from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, generate_start_info_from_omega_map, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) from mx_bluesky.common.parameters.components import get_param_version @@ -270,11 +270,11 @@ def get_ready_for_oav_and_close_shutter( def create_gridscan_callbacks() -> tuple[ - GridscanNexusFileCallback, GridscanISPyBCallback + GridscanNexusFileCallback, GridDetectAndScanISPyBCallback ]: return ( GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan), - GridscanISPyBCallback( + GridDetectAndScanISPyBCallback( param_type=GenericGrid, emit=ZocaloCallback( PlanNameConstants.DO_FGS, diff --git a/src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py b/src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py index 2bf71cdf0c..3d9a94dc3e 100644 --- a/src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py +++ b/src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py @@ -31,7 +31,7 @@ GridDetectionCallback, GridParamUpdate, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( ispyb_activation_wrapper, ) from mx_bluesky.common.parameters.constants import ( diff --git a/src/mx_bluesky/common/external_interaction/callbacks/xray_centre/__init__.py b/src/mx_bluesky/common/external_interaction/callbacks/__init__.py similarity index 100% rename from src/mx_bluesky/common/external_interaction/callbacks/xray_centre/__init__.py rename to src/mx_bluesky/common/external_interaction/callbacks/__init__.py diff --git a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py index a6ef1989d5..5de9ea845b 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py @@ -9,7 +9,6 @@ from dodal.devices.detector.det_resolution import resolution from dodal.devices.synchrotron import SynchrotronMode -from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import I02_1FgsParams from mx_bluesky.common.external_interaction.callbacks.common.plan_reactive_callback import ( PlanReactiveCallback, ) @@ -179,10 +178,9 @@ def _handle_ispyb_transmission_flux_read( if not (beamsize_x_mm and beamsize_y_mm): # todo write issue about vmxm beamsize and link here try: - assert isinstance(self.params, I02_1FgsParams) - beamsize_x_mm = self.params.beam_size_x - beamsize_y_mm = self.params.beam_size_y - except AssertionError: + beamsize_x_mm = self.params.beam_size_x # type: ignore + beamsize_y_mm = self.params.beam_size_y # type: ignore + except Exception: ISPYB_ZOCALO_CALLBACK_LOGGER.warning( "ISPyB callbacks couldn't get beamsize" ) diff --git a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_mapping.py b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_mapping.py index 7df7c2a176..33cbeea8a5 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_mapping.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_mapping.py @@ -48,7 +48,7 @@ def populate_remaining_data_collection_info( data_collection_info.ybeam = beam_position[1] data_collection_info.start_time = get_current_time_string() if data_collection_info.data_collection_number is not None: - # Do not write the file template if we don't have sufficient information - for gridscans we may not + # Do not write the file template if we don't have sufficient information - for gridscans we may not # know the data collection number until later data_collection_info.file_template = f"{params.detector_params.prefix}_{data_collection_info.data_collection_number}_master.h5" return data_collection_info diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/__init__.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/__init__.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py similarity index 97% rename from src/mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_callback.py rename to src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py index ebfcd270b0..d0c4025ec3 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py @@ -21,7 +21,7 @@ from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloInfoGenerator, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( construct_comment_for_gridscan, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( @@ -66,7 +66,7 @@ def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): bpp.run_wrapper( plan_generator, md={ - "activate_callbacks": ["GridscanISPyBCallback"], + "activate_callbacks": ["GridDetectAndScanISPyBCallback"], "subplan_name": PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, "mx_bluesky_parameters": parameters.model_dump_json(), }, @@ -78,12 +78,15 @@ def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): ispyb_activation_decorator = make_decorator(ispyb_activation_wrapper) -class GridscanISPyBCallback(BaseISPyBCallback): +class GridDetectAndScanISPyBCallback(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 receiving an 'event' document for the 'ispyb_reading_hardware' event, and updates the deposition on receiving its final 'stop' document. + This callback is specifically for detecting and scanning two grids. In the future + this callback should be made compatible to a generic number of grids + To use, subscribe the Bluesky RunEngine to an instance of this class. E.g.: ispyb_handler_callback = FGSISPyBCallback(parameters) diff --git a/src/mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_mapping.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_mapping.py similarity index 100% rename from src/mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_mapping.py rename to src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_mapping.py diff --git a/src/mx_bluesky/common/external_interaction/callbacks/xray_centre/nexus_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/nexus_callback.py similarity index 100% rename from src/mx_bluesky/common/external_interaction/callbacks/xray_centre/nexus_callback.py rename to src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/nexus_callback.py diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/__init__.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py new file mode 100644 index 0000000000..b7e6d6a893 --- /dev/null +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py @@ -0,0 +1,320 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from time import time +from typing import TYPE_CHECKING, Any, TypeVar + +from bluesky import preprocessors as bpp +from bluesky.utils import MsgGenerator, make_decorator + +from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams +from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import ( + BaseISPyBCallback, + D, +) +from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import ( + populate_data_collection_group, + populate_remaining_data_collection_info, +) +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridscanPlane, + _smargon_omega_to_xyxz_plane, +) +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( + construct_comment_for_gridscan, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionGridInfo, + DataCollectionGroupInfo, + DataCollectionInfo, + DataCollectionPositionInfo, + Orientation, + ScanDataInfo, +) +from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( + IspybIds, + StoreInIspyb, +) +from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample +from mx_bluesky.common.parameters.constants import PlanNameConstants +from mx_bluesky.common.parameters.gridscan import SpecifiedGrids +from mx_bluesky.common.utils.exceptions import ( + ISPyBDepositionNotMadeError, + SampleError, +) +from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER, set_dcgid_tag + +if TYPE_CHECKING: + from event_model import RunStart, RunStop + +T = TypeVar("T", bound="I02_1FgsParams") +ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" + + +def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): + return bpp.set_run_key_wrapper( + bpp.run_wrapper( + plan_generator, + md={ + "activate_callbacks": ["GridDetectAndScanISPyBCallback"], + "subplan_name": PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, + "mx_bluesky_parameters": parameters.model_dump_json(), + }, + ), + PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, + ) + + +ispyb_activation_decorator = make_decorator(ispyb_activation_wrapper) + + +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 + receiving an 'event' document for the 'ispyb_reading_hardware' event, and updates the + deposition on receiving its final 'stop' document. + + To use, subscribe the Bluesky RunEngine to an instance of this class. + E.g.: + ispyb_handler_callback = FGSISPyBCallback(parameters) + run_engine.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, + param_type: type[T], + *, + emit: Callable[..., Any] | None = None, + ) -> None: + super().__init__(emit=emit) + self.ispyb: StoreInIspyb + self.param_type = param_type + self._start_of_fgs_uid: str | None = None + self._processing_start_time: float | None = None + self._grid_plane_to_id_map: dict[GridscanPlane, int] = {} + self._grid_plane_to_width_map: dict[GridscanPlane, int] = {} + self.data_collection_group_info: DataCollectionGroupInfo | None + + def activity_gated_start(self, doc: RunStart): + if doc.get("subplan_name") == PlanNameConstants.DO_FGS: + self._start_of_fgs_uid = doc.get("uid") + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "ISPyB callback received start document with experiment parameters and " + f"uid: {self._start_of_fgs_uid}" + ) + mx_bluesky_parameters = doc.get("mx_bluesky_parameters") + assert isinstance(mx_bluesky_parameters, str) + self.params = self.param_type.model_validate_json(mx_bluesky_parameters) + + # Fill ispyb deposition with all relevant info, including grid info + self.fill_gridscan_deposition_and_store(lambda: self._get_scan_infos(doc)) + + set_dcgid_tag(self.ispyb_ids.data_collection_group_id) + return super().activity_gated_start(doc) + + def _add_processing_time_to_comment(self, processing_start_time: float): + assert self.data_collection_group_info, ASSERT_START_BEFORE_EVENT_DOC_MESSAGE + proc_time = time() - processing_start_time + crystal_summary = f"Zocalo processing took {proc_time:.2f} s." + + self.data_collection_group_info.comments = ( + self.data_collection_group_info.comments or "" + ) + crystal_summary + + self.ispyb.append_to_comment( + self.ispyb_ids.data_collection_ids[0], crystal_summary + ) + + def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: + """ + so grid information is available immediately after the plan is triggered. + In contrast, i03 and i04 use the OAV to automatically detect their grid. + """ + assert isinstance(self.params, I02_1FgsParams) + assert self.ispyb_ids.data_collection_ids, "No current data collection" + assert self.data_collection_group_info, "No data collection group" + data = doc["data"] + scan_data_infos = [] + + for grid_num in range(self.params.num_grids): + omega = data.get("gonio-omega", self.params.omega_starts_deg[grid_num]) + + # Don't need to do deal with the grid plane here since vmxm only do + # one plane, but leave it in so it's easier to standardise in the future + grid_plane = _smargon_omega_to_xyxz_plane(omega) + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + f"Generating dc info for gridplane {grid_plane}, omega {omega}" + ) + data_collection_number = self.data_collection_number_from_gridplane( + grid_plane + ) + file_template = f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5" + # Snapshots have already been taken in GDA + + # todo the _get_scan_infos here should be vmxm specific since they require us asserting that params includes path_to_xtal_snapshot + # and do this by asserting params are vmxm params + + data_collection_info = DataCollectionInfo( + xtal_snapshot1=str(self.params.path_to_xtal_snapshot), + xtal_snapshot2=str(self.params.path_to_xtal_snapshot), + xtal_snapshot3=str(self.params.path_to_xtal_snapshot), + n_images=self.params.num_images, + data_collection_number=data_collection_number, + file_template=file_template, + ) + data_collection_grid_info = DataCollectionGridInfo( + dx_in_mm=self.params.x_step_size_um * 1000, + dy_in_mm=self.params.y_step_sizes_um[grid_num] * 1000, + steps_x=self.params.x_steps, + steps_y=self.params.y_steps[grid_num], + microns_per_pixel_x=self.params.microns_per_pixel_x, + microns_per_pixel_y=self.params.microns_per_pixel_y, + snapshot_offset_x_pixel=self.params.upper_left_x, + snapshot_offset_y_pixel=self.params.upper_left_y, + orientation=Orientation.HORIZONTAL, + snaked=True, + ) + data_collection_info.comments = construct_comment_for_gridscan( + data_collection_grid_info + ) + + # Grid plane logic isn't needed for VMX, but keep it for now anyway + data_collection_id = self.ispyb_ids.data_collection_ids[ + 0 if grid_plane == GridscanPlane.OMEGA_XY else 1 + ] + self._grid_plane_to_id_map[grid_plane] = data_collection_id + self._grid_plane_to_width_map[grid_plane] = ( + data_collection_grid_info.steps_y + ) + + y_steps = self._grid_plane_to_width_map.get(GridscanPlane.OMEGA_XY, "_") + self.data_collection_group_info.comments = ( + f"Diffraction grid scan of {data_collection_grid_info.steps_x} by " + f"{y_steps}." + ) + + self._populate_axis_info(data_collection_info, doc["data"]) + + # todo do all this stuff as soon as possible after plan starts + + scan_data_info = ScanDataInfo( + data_collection_info=data_collection_info, + data_collection_id=data_collection_id, + data_collection_grid_info=data_collection_grid_info, + ) + + scan_data_infos.append(scan_data_info) + + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "Updating ispyb data collection after loading grid params" + ) + + return scan_data_infos + + def _populate_axis_info(self, data_collection_info: DataCollectionInfo, doc: dict): + if (omega_start := doc.get("gonio-omega")) 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 + if (chi_start := doc.get("gonio-chi")) is not None: + data_collection_info.chi_start = chi_start + + 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" + ) + assert isinstance(self.params, I02_1FgsParams) + scan_data_infos = [] + for grid_num in range(self.params.num_grids): + scan_data_info = ScanDataInfo( + data_collection_info=event_sourced_data_collection_info, + data_collection_id=self.ispyb_ids.data_collection_ids[grid_num], + ) + scan_data_infos.append(scan_data_info) + return scan_data_infos + + def activity_gated_stop(self, doc: RunStop) -> RunStop: + assert self.data_collection_group_info, ( + f"No data collection group info - stop document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" + ) + if doc.get("run_start") == self._start_of_fgs_uid: + self._processing_start_time = time() + if doc.get("run_start") == self._start_of_fgs_uid: + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "ISPyB callback received stop document corresponding to start document " + f"with uid: {self._start_of_fgs_uid}." + ) + if self.ispyb_ids == IspybIds(): + raise ISPyBDepositionNotMadeError( + "ispyb was not initialised at run start" + ) + exception_type, message = SampleError.type_and_message_from_reason( + doc.get("reason", "") + ) + if exception_type: + doc["reason"] = message + self.data_collection_group_info.comments = message + elif self._processing_start_time: + self._add_processing_time_to_comment(self._processing_start_time) + self.ispyb.update_data_collection_group_table( + self.data_collection_group_info, + self.ispyb_ids.data_collection_group_id, + ) + self.data_collection_group_info = None + self._grid_plane_to_id_map.clear() + self._grid_plane_to_width_map.clear() + return super().activity_gated_stop(doc) + return self.tag_doc(doc) + + def tag_doc(self, doc: D) -> D: + doc = super().tag_doc(doc) + assert isinstance(doc, dict) + if self._grid_plane_to_id_map: + doc["grid_plane_to_id_map"] = self._grid_plane_to_id_map + return doc # type: ignore + + def data_collection_number_from_gridplane(self, plane) -> int: + assert self.params + base_number = self.params.detector_params.run_number + return base_number if plane == GridscanPlane.OMEGA_XY else base_number + 1 + + def fill_gridscan_deposition_and_store( + self, make_scan_infos_with_grid_info: Callable[..., Sequence[ScanDataInfo]] + ): + assert isinstance(self.params, SpecifiedGrids) + + # Do initial deposition using all info except grid info + self.ispyb = StoreInIspyb(self.ispyb_config) + self.data_collection_group_info = populate_data_collection_group(self.params) + scan_data_infos = [] + assert self.params.num_grids > 0 + for grid in range(self.params.num_grids): + scan_data_infos.append( + ScanDataInfo( + data_collection_info=populate_remaining_data_collection_info( + f"MX-Bluesky: Xray centring {grid + 1}/{self.params.num_grids} -", + None, + DataCollectionInfo(), + self.params, + ) + ) + ) + self.ispyb_ids = self.ispyb.begin_deposition( + self.data_collection_group_info, scan_data_infos + ) + # Now use grid information to complete deposition + scan_data_infos: list[ScanDataInfo] = list(make_scan_infos_with_grid_info()) + self.ispyb_ids = self.ispyb.update_deposition(self.ispyb_ids, scan_data_infos) + self.ispyb.update_data_collection_group_table( + self.data_collection_group_info, self.ispyb_ids.data_collection_group_id + ) diff --git a/src/mx_bluesky/hyperion/experiment_plans/pin_centre_then_gridscan_plan.py b/src/mx_bluesky/hyperion/experiment_plans/pin_centre_then_gridscan_plan.py index 8d407b4827..43133e6d7a 100644 --- a/src/mx_bluesky/hyperion/experiment_plans/pin_centre_then_gridscan_plan.py +++ b/src/mx_bluesky/hyperion/experiment_plans/pin_centre_then_gridscan_plan.py @@ -19,7 +19,7 @@ PinTipCentringComposite, pin_tip_centre_plan, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( ispyb_activation_wrapper, ) from mx_bluesky.common.parameters.constants import OavConstants, PlanNameConstants diff --git a/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py b/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py index 2f6696466e..cad6482812 100644 --- a/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py +++ b/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py @@ -25,16 +25,16 @@ from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloCallback, ) -from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import ( - SampleHandlingCallback, -) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, generate_start_info_from_omega_map, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) +from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import ( + SampleHandlingCallback, +) from mx_bluesky.common.utils.log import ( ISPYB_ZOCALO_CALLBACK_LOGGER, NEXUS_LOGGER, @@ -75,11 +75,11 @@ def create_gridscan_callbacks() -> tuple[ - GridscanNexusFileCallback, GridscanISPyBCallback + GridscanNexusFileCallback, GridDetectAndScanISPyBCallback ]: return ( GridscanNexusFileCallback(param_type=HyperionSpecifiedThreeDGridScan), - GridscanISPyBCallback( + GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams, emit=ZocaloCallback( CONST.PLAN.DO_FGS, CONST.ZOCALO_ENV, generate_start_info_from_omega_map diff --git a/tests/conftest.py b/tests/conftest.py index f91f9447e8..de0b1b9dec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -90,7 +90,7 @@ from mx_bluesky.beamlines.i04.external_interaction.config_server import ( get_i04_config_client, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridscanPlane, ) from mx_bluesky.common.parameters.constants import ( diff --git a/tests/system_tests/hyperion/external_interaction/callbacks/test_external_callbacks.py b/tests/system_tests/hyperion/external_interaction/callbacks/test_external_callbacks.py index 1a8cf6a5f0..f7a28afbdd 100644 --- a/tests/system_tests/hyperion/external_interaction/callbacks/test_external_callbacks.py +++ b/tests/system_tests/hyperion/external_interaction/callbacks/test_external_callbacks.py @@ -38,7 +38,7 @@ fetch_xrc_results_from_zocalo, zocalo_stage_decorator, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( ispyb_activation_decorator, ) from mx_bluesky.common.parameters.components import WithSnapshot diff --git a/tests/system_tests/hyperion/external_interaction/test_ispyb_dev_connection.py b/tests/system_tests/hyperion/external_interaction/test_ispyb_dev_connection.py index aafcf3932c..abd93a289b 100644 --- a/tests/system_tests/hyperion/external_interaction/test_ispyb_dev_connection.py +++ b/tests/system_tests/hyperion/external_interaction/test_ispyb_dev_connection.py @@ -21,10 +21,10 @@ populate_data_collection_group, populate_remaining_data_collection_info, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( construct_comment_for_gridscan, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( @@ -416,7 +416,9 @@ def test_ispyb_deposition_in_gridscan( set_mock_value( grid_detect_then_xray_centre_composite.s4_slit_gaps.ygap.user_readback, 0.1 ) - ispyb_callback = GridscanISPyBCallback(GenericGridWithHyperionDetectorParams) + ispyb_callback = GridDetectAndScanISPyBCallback( + GenericGridWithHyperionDetectorParams + ) run_engine.subscribe(ispyb_callback) run_engine( grid_detect_then_xray_centre( diff --git a/tests/system_tests/hyperion/external_interaction/test_load_centre_collect_full_plan.py b/tests/system_tests/hyperion/external_interaction/test_load_centre_collect_full_plan.py index dc7152ebd1..2d626df06d 100644 --- a/tests/system_tests/hyperion/external_interaction/test_load_centre_collect_full_plan.py +++ b/tests/system_tests/hyperion/external_interaction/test_load_centre_collect_full_plan.py @@ -27,12 +27,12 @@ from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import ( get_proposal_and_session_from_visit_string, ) +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, +) from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import ( SampleHandlingCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, -) from mx_bluesky.common.utils.exceptions import ( CrystalNotFoundError, WarningError, @@ -294,7 +294,7 @@ def test_execute_load_centre_collect_full( tmp_path, robot_load_cb: RobotLoadISPyBCallback, ): - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -473,7 +473,7 @@ def move_to_initial_omega(): yield from bps.mv(load_centre_collect_composite.gonio.omega, initial_omega) run_engine(move_to_initial_omega()) - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -565,7 +565,7 @@ def test_load_centre_collect_updates_bl_sample_status_pin_tip_detection_fail( fetch_blsample: Callable[..., Any], ): robot_load_cb = RobotLoadISPyBCallback() - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) sample_handling_cb = SampleHandlingCallback() @@ -599,7 +599,7 @@ def test_load_centre_collect_updates_bl_sample_status_grid_detection_fail_tip_no fetch_blsample: Callable[..., Any], ): robot_load_cb = RobotLoadISPyBCallback() - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) sample_handling_cb = SampleHandlingCallback() @@ -651,7 +651,7 @@ def test_load_centre_collect_updates_bl_sample_status_gridscan_no_diffraction( fetch_blsample: Callable[..., Any], ): robot_load_cb = RobotLoadISPyBCallback() - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) sample_handling_cb = SampleHandlingCallback() @@ -683,7 +683,7 @@ def test_load_centre_collect_updates_bl_sample_status_rotation_failure( fetch_blsample: Callable[..., Any], ): robot_load_cb = RobotLoadISPyBCallback() - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) sample_handling_cb = SampleHandlingCallback() @@ -741,7 +741,7 @@ def test_load_centre_collect_gridscan_result_at_edge_of_grid( load_centre_collect_composite.zocalo.my_zocalo_result = _with_sample_ids( zocalo_result, [SimConstants.ST_SAMPLE_ID] ) - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -775,7 +775,7 @@ def test_execute_load_centre_collect_capture_rotation_snapshots( ): load_centre_collect_params.multi_rotation_scan.snapshot_directory = tmp_path - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -857,7 +857,7 @@ def test_load_centre_collect_multisample_pin_reports_correct_sample_ids_in_ispyb fetch_datacollection_attribute: Callable[..., Any], ): load_centre_collect_composite.zocalo.my_zocalo_result = zocalo_result - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -910,7 +910,7 @@ def test_load_centre_collect_multisample_pin_reports_correct_sample_ids_in_ispyb fetch_datacollection_ids_for_group_id: Callable[..., Any], ): load_centre_collect_composite.zocalo.my_zocalo_result = zocalo_result - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -976,7 +976,7 @@ def test_load_centre_collect_multisample_pin_reports_correct_sample_ids_robot_lo robot_load_cb: RobotLoadISPyBCallback, ): load_centre_collect_composite.zocalo.my_zocalo_result = zocalo_result - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -1033,7 +1033,7 @@ def test_load_centre_collect_multisample_pin_updates_sample_status_for_parent_sa fetch_blsample: Callable[..., Any], ): load_centre_collect_composite.zocalo.my_zocalo_result = zocalo_result - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() @@ -1167,7 +1167,7 @@ def test_load_centre_collect_generate_rotation_snapshots( SNAPSHOT_GENERATION_ZOCALO_RESULT ) - ispyb_gridscan_cb = GridscanISPyBCallback( + ispyb_gridscan_cb = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_rotation_cb = RotationISPyBCallback() diff --git a/tests/system_tests/hyperion/external_interaction/test_zocalo_system.py b/tests/system_tests/hyperion/external_interaction/test_zocalo_system.py index 29f70abc86..9b5972e807 100644 --- a/tests/system_tests/hyperion/external_interaction/test_zocalo_system.py +++ b/tests/system_tests/hyperion/external_interaction/test_zocalo_system.py @@ -13,7 +13,7 @@ from mx_bluesky.common.experiment_plans.inner_plans.read_hardware import ( read_hardware_for_zocalo, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridscanPlane, ispyb_activation_wrapper, ) diff --git a/tests/unit_tests/common/experiment_plans/inner_plans/test_do_fgs.py b/tests/unit_tests/common/experiment_plans/inner_plans/test_do_fgs.py index 921ed85540..23f75acd3a 100644 --- a/tests/unit_tests/common/experiment_plans/inner_plans/test_do_fgs.py +++ b/tests/unit_tests/common/experiment_plans/inner_plans/test_do_fgs.py @@ -18,7 +18,7 @@ from mx_bluesky.common.experiment_plans.inner_plans.do_fgs import ( kickoff_and_complete_gridscan, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridscanPlane, ) from mx_bluesky.common.parameters.constants import ( diff --git a/tests/unit_tests/common/experiment_plans/test_common_flyscan_xray_centre_plan.py b/tests/unit_tests/common/experiment_plans/test_common_flyscan_xray_centre_plan.py index dcb4b8f8d6..ea516e365e 100644 --- a/tests/unit_tests/common/experiment_plans/test_common_flyscan_xray_centre_plan.py +++ b/tests/unit_tests/common/experiment_plans/test_common_flyscan_xray_centre_plan.py @@ -37,11 +37,11 @@ from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, ispyb_activation_wrapper, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( @@ -78,7 +78,7 @@ def mock_plan(): def run_engine_with_subs_snapshots_already_taken(run_engine_with_subs, test_event_data): run_engine, subscriptions = run_engine_with_subs ispyb_gridscan_callback = [ - sub for sub in subscriptions if isinstance(sub, GridscanISPyBCallback) + sub for sub in subscriptions if isinstance(sub, GridDetectAndScanISPyBCallback) ][0] ispyb_gridscan_callback.active = True ispyb_gridscan_callback.start( @@ -94,7 +94,7 @@ def run_engine_with_subs_snapshots_already_taken(run_engine_with_subs, test_even @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb", modified_store_grid_scan_mock, ) class TestFlyscanXrayCentrePlan: @@ -122,7 +122,9 @@ def test_when_run_gridscan_called_ispyb_deposition_made_and_records_errors( test_three_d_grid_params: SpecifiedThreeDGridScan, beamline_specific: BeamlineSpecificFGSFeatures, ): - ispyb_callback = GridscanISPyBCallback(param_type=SpecifiedThreeDGridScan) + ispyb_callback = GridDetectAndScanISPyBCallback( + param_type=SpecifiedThreeDGridScan + ) run_engine.subscribe(ispyb_callback) error = None @@ -319,7 +321,7 @@ def test_when_grid_scan_ran_then_eiger_disarmed_before_zocalo_end( test_three_d_grid_params: SpecifiedThreeDGridScan, run_engine_with_subs_snapshots_already_taken: tuple[ RunEngine, - tuple[GridscanNexusFileCallback, GridscanISPyBCallback], + tuple[GridscanNexusFileCallback, GridDetectAndScanISPyBCallback], ], beamline_specific: BeamlineSpecificFGSFeatures, ): @@ -342,7 +344,7 @@ def test_when_grid_scan_ran_then_eiger_disarmed_before_zocalo_end( ) with patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter.create_nexus_file", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback.NexusWriter.create_nexus_file", autospec=True, ): [run_engine.subscribe(cb) for cb in (nexus_cb, ispyb_cb)] @@ -576,7 +578,7 @@ def test_when_gridscan_succeeds_and_results_fetched_ispyb_comment_appended_to( run_gridscan: MagicMock, run_engine_with_subs: tuple[ RunEngine, - tuple[GridscanNexusFileCallback, GridscanISPyBCallback], + tuple[GridscanNexusFileCallback, GridDetectAndScanISPyBCallback], ], test_three_d_grid_params: SpecifiedThreeDGridScan, fake_fgs_composite: FlyScanEssentialDevices, diff --git a/tests/unit_tests/common/experiment_plans/test_common_grid_detect_then_xray_centre_plan.py b/tests/unit_tests/common/experiment_plans/test_common_grid_detect_then_xray_centre_plan.py index 38cba4d938..83390b533b 100644 --- a/tests/unit_tests/common/experiment_plans/test_common_grid_detect_then_xray_centre_plan.py +++ b/tests/unit_tests/common/experiment_plans/test_common_grid_detect_then_xray_centre_plan.py @@ -24,7 +24,7 @@ from mx_bluesky.common.experiment_plans.inner_plans.xrc_results_utils import ( _fire_xray_centre_result_event, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( ispyb_activation_wrapper, ) from mx_bluesky.common.parameters.constants import ( @@ -267,7 +267,7 @@ def test_detect_grid_and_do_gridscan_does_not_activate_ispyb_callback( msg for msg in msgs if msg.command == "open_run" - and "GridscanISPyBCallback" in msg.kwargs["activate_callbacks"] + and "GridDetectAndScanISPyBCallback" in msg.kwargs["activate_callbacks"] ] assert not activations @@ -340,7 +340,7 @@ def test_grid_detect_then_xray_centre_activates_ispyb_callback( assert_message_and_return_remaining( msgs_from_simulated_grid_detect_then_xray_centre, lambda msg: msg.command == "open_run" - and "GridscanISPyBCallback" in msg.kwargs["activate_callbacks"], + and "GridDetectAndScanISPyBCallback" in msg.kwargs["activate_callbacks"], ) diff --git a/tests/unit_tests/common/experiment_plans/test_grid_detection_plan.py b/tests/unit_tests/common/experiment_plans/test_grid_detection_plan.py index 3f3b526763..e6c181e27b 100644 --- a/tests/unit_tests/common/experiment_plans/test_grid_detection_plan.py +++ b/tests/unit_tests/common/experiment_plans/test_grid_detection_plan.py @@ -26,8 +26,8 @@ from mx_bluesky.common.external_interaction.callbacks.common.grid_detection_callback import ( GridDetectionCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, ispyb_activation_wrapper, ) from mx_bluesky.common.parameters.gridscan import GenericGrid, SpecifiedThreeDGridScan @@ -218,7 +218,7 @@ async def test_when_grid_detection_plan_run_then_ispyb_callback_gets_correct_val ): params = OAVParameters("loopCentring", test_config_files["oav_config_json"]) composite, _ = fake_devices - cb = GridscanISPyBCallback(param_type=GenericGrid) + cb = GridDetectAndScanISPyBCallback(param_type=GenericGrid) cb.data_collection_group_info = dummy_rotation_data_collection_group_info run_engine.subscribe(cb) @@ -232,7 +232,7 @@ async def test_when_grid_detection_plan_run_then_ispyb_callback_gets_correct_val assert_event( cb.activity_gated_start.mock_calls[0], # pyright:ignore - {"activate_callbacks": ["GridscanISPyBCallback"]}, + {"activate_callbacks": ["GridDetectAndScanISPyBCallback"]}, ) assert_event( cb.activity_gated_event.mock_calls[0], # pyright: ignore diff --git a/tests/unit_tests/common/external_interaction/callbacks/test_zocalo_handler.py b/tests/unit_tests/common/external_interaction/callbacks/test_zocalo_handler.py index aca1a0eb62..705e95ae9d 100644 --- a/tests/unit_tests/common/external_interaction/callbacks/test_zocalo_handler.py +++ b/tests/unit_tests/common/external_interaction/callbacks/test_zocalo_handler.py @@ -86,10 +86,10 @@ def test_handler_stores_collection_ispyb_ids_come_in_as_subplan( autospec=True, ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback.NexusWriter", ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb", ) def test_execution_of_do_fgs_triggers_zocalo_calls( self, @@ -156,10 +156,10 @@ def test_execution_of_do_fgs_triggers_zocalo_calls( autospec=True, ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback.NexusWriter", ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb", ) def test_do_fgs_triggers_zocalo_calls_when_snapshots_in_reverse_order( self, diff --git a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_callback.py b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_callback.py index b1c9d67e69..0a3c62112d 100644 --- a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_callback.py +++ b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_callback.py @@ -10,8 +10,8 @@ from mx_bluesky.common.experiment_plans.inner_plans.read_hardware import ( read_hardware_plan, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, GridscanPlane, _smargon_omega_to_xyxz_plane, ) @@ -73,7 +73,7 @@ ) class TestXrayCentreISPyBCallback: def test_activity_gated_start_3d(self, mock_ispyb_conn, test_event_data, tmp_path): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.activity_gated_start( @@ -116,7 +116,7 @@ def test_activity_gated_start_3d(self, mock_ispyb_conn, test_event_data, tmp_pat def test_reason_provided_if_crystal_not_found_error( self, mock_update_data_collection_group_table, mock_ispyb_conn, test_event_data ): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.activity_gated_start( @@ -137,7 +137,7 @@ def test_reason_provided_if_crystal_not_found_error( ) def test_hardware_read_event_3d(self, mock_ispyb_conn, test_event_data): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.activity_gated_start( @@ -166,7 +166,7 @@ def test_hardware_read_event_3d(self, mock_ispyb_conn, test_event_data): assert update_dc_requests[1].body == expected_upsert def test_flux_read_events_3d(self, mock_ispyb_conn, test_event_data): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.activity_gated_start( @@ -231,7 +231,7 @@ def test_activity_gated_event_oav_snapshot_triggered( snapshot_events: list[str], first_comment: str, ): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.activity_gated_start( @@ -334,7 +334,7 @@ def test_activity_gated_event_oav_snapshot_triggered( async def test_ispyb_callback_handles_read_hardware_in_run_engine( self, run_engine, mock_ispyb_conn, dummy_rotation_data_collection_group_info ): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback._handle_ispyb_hardware_read = MagicMock() @@ -349,7 +349,7 @@ async def test_ispyb_callback_handles_read_hardware_in_run_engine( @subs_decorator(callback) @run_decorator( md={ - "activate_callbacks": ["GridscanISPyBCallback"], + "activate_callbacks": ["GridDetectAndScanISPyBCallback"], }, ) def test_plan(): @@ -366,7 +366,7 @@ def test_plan(): callback._handle_ispyb_transmission_flux_read.assert_called_once() @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.GridscanISPyBCallback._handle_oav_grid_snapshot_triggered", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.GridDetectAndScanISPyBCallback._handle_oav_grid_snapshot_triggered", ) @patch( "mx_bluesky.common.external_interaction.ispyb.ispyb_store.StoreInIspyb.update_deposition", @@ -381,7 +381,7 @@ def test_given_event_doc_before_start_doc_received_then_exception_raised( mock__handle_oav_grid_snapshot_triggered, test_event_data, ): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.activity_gated_descriptor( @@ -400,7 +400,7 @@ def test_given_event_doc_before_start_doc_received_then_exception_raised( def test_ispyb_callback_clears_state_after_run_stop( self, test_event_data, mock_ispyb_conn ): - callback = GridscanISPyBCallback( + callback = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) callback.active = True diff --git a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py index 1abe4d71a5..d967e716f5 100644 --- a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py +++ b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py @@ -3,8 +3,8 @@ import pytest from graypy import GELFTCPHandler -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, ) from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( IspybIds, @@ -48,12 +48,12 @@ def mock_store_in_ispyb(config, *args, **kwargs) -> StoreInIspyb: MagicMock(return_value=td.DUMMY_TIME_STRING), ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb", mock_store_in_ispyb, ) class TestXrayCentreIspybHandler: def test_fgs_failing_results_in_bad_run_status_in_ispyb(self, test_event_data): - ispyb_handler = GridscanISPyBCallback( + ispyb_handler = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_handler.activity_gated_start( @@ -88,7 +88,7 @@ def test_fgs_failing_results_in_bad_run_status_in_ispyb(self, test_event_data): def test_fgs_raising_no_exception_results_in_good_run_status_in_ispyb( self, test_event_data ): - ispyb_handler = GridscanISPyBCallback( + ispyb_handler = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_handler.activity_gated_start( @@ -134,7 +134,7 @@ def test_given_ispyb_callback_started_writing_to_ispyb_when_messages_logged_then ) gelf_handler.emit = MagicMock() - ispyb_handler = GridscanISPyBCallback( + ispyb_handler = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_handler.activity_gated_start( @@ -170,7 +170,7 @@ def test_given_ispyb_callback_finished_writing_to_ispyb_when_messages_logged_the ) gelf_handler.emit = MagicMock() - ispyb_handler = GridscanISPyBCallback( + ispyb_handler = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams ) ispyb_handler.activity_gated_start( @@ -197,13 +197,13 @@ def test_given_ispyb_callback_finished_writing_to_ispyb_when_messages_logged_the assert not hasattr(latest_record, "dc_group_id") @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.time", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.time", side_effect=[2, 100], ) def test_given_fgs_plan_finished_when_zocalo_results_event_then_expected_comment_deposited( self, mock_time, dummy_rotation_data_collection_group_info, test_event_data ): - ispyb_handler = GridscanISPyBCallback( + ispyb_handler = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams, ) diff --git a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_mapping.py b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_mapping.py index 93d5fc510c..1ab24436f2 100644 --- a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_mapping.py +++ b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_mapping.py @@ -2,7 +2,7 @@ import pytest -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( construct_comment_for_gridscan, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( @@ -64,7 +64,7 @@ def test_ispyb_deposition_rounds_position_to_int( ], ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping.oav_utils.bottom_right_from_top_left", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping.oav_utils.bottom_right_from_top_left", autospec=True, ) def test_ispyb_deposition_rounds_box_size_int( diff --git a/tests/unit_tests/common/external_interaction/xray_centre/test_nexus_handler.py b/tests/unit_tests/common/external_interaction/xray_centre/test_nexus_handler.py index c5f970bb05..a48bf30fd6 100644 --- a/tests/unit_tests/common/external_interaction/xray_centre/test_nexus_handler.py +++ b/tests/unit_tests/common/external_interaction/xray_centre/test_nexus_handler.py @@ -5,7 +5,7 @@ import pytest from numpy.typing import DTypeLike -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) from mx_bluesky.hyperion.parameters.gridscan import HyperionSpecifiedThreeDGridScan @@ -34,7 +34,7 @@ def test_writers_not_called_on_plan_start_doc( @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter" + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback.NexusWriter" ) def test_writers_dont_create_on_init_but_do_on_during_collection_read_event( mock_nexus_writer: MagicMock, @@ -73,7 +73,7 @@ def test_writers_dont_create_on_init_but_do_on_during_collection_read_event( ], ) @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter" + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback.NexusWriter" ) def test_given_different_bit_depths_then_writers_created_wth_correct_virtual_dataset_size( mock_nexus_writer: MagicMock, @@ -108,7 +108,7 @@ def test_given_different_bit_depths_then_writers_created_wth_correct_virtual_dat @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter" + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback.NexusWriter" ) def test_beam_and_attenuator_set_on_ispyb_transmission_event( mock_nexus_writer: MagicMock, diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index e8a66b1dbb..f373841b3a 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -49,11 +49,11 @@ from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, generate_start_info_from_omega_map, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( @@ -170,11 +170,11 @@ def assert_event(mock_call, expected): def create_gridscan_callbacks() -> tuple[ - GridscanNexusFileCallback, GridscanISPyBCallback + GridscanNexusFileCallback, GridDetectAndScanISPyBCallback ]: return ( GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan), - GridscanISPyBCallback( + GridDetectAndScanISPyBCallback( param_type=SpecifiedThreeDGridScan, emit=ZocaloCallback( PlanNameConstants.DO_FGS, @@ -210,7 +210,7 @@ def mock_subscriptions(test_three_d_grid_params): "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger", ), patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb" + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb" ) as mock_store_in_ispyb, ): mock_store_in_ispyb.return_value.begin_deposition.return_value = IspybIds( @@ -227,13 +227,17 @@ def mock_subscriptions(test_three_d_grid_params): yield (nexus_callback, ispyb_callback) -ReWithSubs = tuple[RunEngine, tuple[GridscanNexusFileCallback | GridscanISPyBCallback]] +ReWithSubs = tuple[ + RunEngine, tuple[GridscanNexusFileCallback | GridDetectAndScanISPyBCallback] +] @pytest.fixture def run_engine_with_subs( run_engine: RunEngine, - mock_subscriptions: tuple[GridscanNexusFileCallback | GridscanISPyBCallback], + mock_subscriptions: tuple[ + GridscanNexusFileCallback | GridDetectAndScanISPyBCallback + ], ) -> Generator[ReWithSubs, Any, None]: for cb in list(mock_subscriptions): run_engine.subscribe(cb) @@ -271,7 +275,7 @@ def make_event_doc(data, descriptor="abc123") -> Event: def run_generic_ispyb_handler_setup( - ispyb_handler: GridscanISPyBCallback, + ispyb_handler: GridDetectAndScanISPyBCallback, params: SpecifiedThreeDGridScan, ): """This is useful when testing 'run_gridscan_and_move(...)' because this stuff diff --git a/tests/unit_tests/hyperion/experiment_plans/conftest.py b/tests/unit_tests/hyperion/experiment_plans/conftest.py index b4d98e96af..30a98d21b9 100644 --- a/tests/unit_tests/hyperion/experiment_plans/conftest.py +++ b/tests/unit_tests/hyperion/experiment_plans/conftest.py @@ -110,10 +110,10 @@ def mock_subscriptions(): autospec=True, ), patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb.append_to_comment" + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb.append_to_comment" ), patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb.begin_deposition", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb.begin_deposition", new=MagicMock( return_value=IspybIds( data_collection_ids=(0, 0), data_collection_group_id=0 @@ -121,7 +121,7 @@ def mock_subscriptions(): ), ), patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb.update_deposition", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb.update_deposition", new=MagicMock( return_value=IspybIds( data_collection_ids=(0, 0), diff --git a/tests/unit_tests/hyperion/experiment_plans/test_hyperion_flyscan_xray_centre_plan.py b/tests/unit_tests/hyperion/experiment_plans/test_hyperion_flyscan_xray_centre_plan.py index d2c2f6aae1..286cfc7d37 100644 --- a/tests/unit_tests/hyperion/experiment_plans/test_hyperion_flyscan_xray_centre_plan.py +++ b/tests/unit_tests/hyperion/experiment_plans/test_hyperion_flyscan_xray_centre_plan.py @@ -18,10 +18,10 @@ BeamlineSpecificFGSFeatures, common_flyscan_xray_centre, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( - GridscanISPyBCallback, +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( + GridDetectAndScanISPyBCallback, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) from mx_bluesky.common.parameters.constants import ( @@ -47,7 +47,9 @@ modified_store_grid_scan_mock, ) -ReWithSubs = tuple[RunEngine, tuple[GridscanNexusFileCallback, GridscanISPyBCallback]] +ReWithSubs = tuple[ + RunEngine, tuple[GridscanNexusFileCallback, GridDetectAndScanISPyBCallback] +] class CompleteError(Exception): @@ -70,7 +72,7 @@ def fgs_composite_with_panda_pcap( @patch( - "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb", + "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.StoreInIspyb", modified_store_grid_scan_mock, ) class TestFlyscanXrayCentrePlan: diff --git a/tests/unit_tests/hyperion/experiment_plans/test_pin_centre_then_xray_centre_plan.py b/tests/unit_tests/hyperion/experiment_plans/test_pin_centre_then_xray_centre_plan.py index 369f1c52d8..d873531075 100644 --- a/tests/unit_tests/hyperion/experiment_plans/test_pin_centre_then_xray_centre_plan.py +++ b/tests/unit_tests/hyperion/experiment_plans/test_pin_centre_then_xray_centre_plan.py @@ -149,7 +149,7 @@ def test_pin_centre_then_gridscan_plan_activates_ispyb_callback_before_pin_tip_c msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "open_run" - and "GridscanISPyBCallback" in msg.kwargs["activate_callbacks"], + and "GridDetectAndScanISPyBCallback" in msg.kwargs["activate_callbacks"], ) msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "pin_tip_centre_plan" diff --git a/tests/unit_tests/hyperion/external_interaction/nexus/test_write_nexus.py b/tests/unit_tests/hyperion/external_interaction/nexus/test_write_nexus.py index 694f248fc3..5bb9f813b8 100644 --- a/tests/unit_tests/hyperion/external_interaction/nexus/test_write_nexus.py +++ b/tests/unit_tests/hyperion/external_interaction/nexus/test_write_nexus.py @@ -16,7 +16,7 @@ ZebraGridScanParamsThreeD, ) -from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( _create_writers_from_params, ) from mx_bluesky.common.external_interaction.nexus.nexus_utils import ( From 453d802b8de9caf924382126a010ca982be21b3f Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Fri, 27 Feb 2026 14:48:32 +0000 Subject: [PATCH 4/8] Make common ispyb callback for gridscans with no grid detect --- .../callbacks/gridscan/ispyb_callback.py | 85 ++++++++++ .../beamlines/i02_1/i02_1_gridscan_plan.py | 42 ++--- .../i04_grid_detect_then_xray_centre_plan.py | 23 ++- .../callbacks/common/ispyb_callback_base.py | 2 +- .../grid_detect_and_scan/ispyb_callback.py | 57 ++----- .../callbacks/grid/gridscan/ispyb_callback.py | 156 ++++-------------- .../callbacks/grid/utils.py | 73 ++++++++ .../common/parameters/components.py | 2 +- src/mx_bluesky/common/parameters/constants.py | 1 + src/mx_bluesky/common/parameters/gridscan.py | 2 +- .../callbacks/__main__.py | 14 +- 11 files changed, 236 insertions(+), 221 deletions(-) create mode 100644 src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py index e69de29bb2..0c6830078b 100644 --- a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py @@ -0,0 +1,85 @@ +from collections.abc import Sequence + +from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( + construct_comment_for_gridscan, +) +from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( + GridscanISPyBCallback, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionGridInfo, + DataCollectionInfo, + Orientation, + ScanDataInfo, +) +from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER + + +class I021GridscanISPyBCallback(GridscanISPyBCallback): + def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: + """ + For VMXm, grid information is available immediately after the plan is triggered. + """ + assert isinstance(self.params, I02_1FgsParams) + assert self.ispyb_ids.data_collection_ids, "No current data collection" + assert self.data_collection_group_info, "No data collection group" + data = doc["data"] + scan_data_infos = [] + + for grid_num in range(self.params.num_grids): + omega = data.get("gonio-omega", self.params.omega_starts_deg[grid_num]) + + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + f"Generating dc info for gridplane XY, omega {omega}" + ) + data_collection_number = self.params.detector_params.run_number + file_template = f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5" + # Snapshots have already been taken in GDA + + data_collection_info = DataCollectionInfo( + xtal_snapshot1=str(self.params.path_to_xtal_snapshot), + xtal_snapshot2=str(self.params.path_to_xtal_snapshot), + xtal_snapshot3=str(self.params.path_to_xtal_snapshot), + n_images=self.params.num_images, + data_collection_number=data_collection_number, + file_template=file_template, + ) + data_collection_grid_info = DataCollectionGridInfo( + dx_in_mm=self.params.x_step_size_um * 1000, + dy_in_mm=self.params.y_step_sizes_um[grid_num] * 1000, + steps_x=self.params.x_steps, + steps_y=self.params.y_steps[grid_num], + microns_per_pixel_x=self.params.microns_per_pixel_x, + microns_per_pixel_y=self.params.microns_per_pixel_y, + snapshot_offset_x_pixel=self.params.upper_left_x, + snapshot_offset_y_pixel=self.params.upper_left_y, + orientation=Orientation.HORIZONTAL, + snaked=True, + ) + data_collection_info.comments = construct_comment_for_gridscan( + data_collection_grid_info + ) + + data_collection_id = self.ispyb_ids.data_collection_ids[0] + + self.data_collection_group_info.comments = ( + f"Diffraction grid scan of {data_collection_grid_info.steps_x} by " + f"{self.params.y_steps}." + ) + + self._populate_axis_info(data_collection_info, doc["data"]) + + scan_data_info = ScanDataInfo( + data_collection_info=data_collection_info, + data_collection_id=data_collection_id, + data_collection_grid_info=data_collection_grid_info, + ) + + scan_data_infos.append(scan_data_info) + + ISPYB_ZOCALO_CALLBACK_LOGGER.info( + "Updating ispyb data collection after loading grid params" + ) + + return scan_data_infos diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index 42645ffd53..3afa851898 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -32,8 +32,6 @@ ) from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( ZocaloCallback, - ZocaloInfoGenerator, - ZocaloStartInfo, ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, @@ -41,6 +39,9 @@ from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( GridscanISPyBCallback, ) +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + generate_start_info_from_omega_map, +) from mx_bluesky.common.parameters.constants import ( EnvironmentConstants, PlanNameConstants, @@ -50,34 +51,11 @@ GonioWithOmegaType, ) from mx_bluesky.common.utils.log import LOGGER -from mx_bluesky.common.utils.utils import number_of_frames_from_scan_spec - - -# todo this all needs chanigng -def generate_start_info_from_omega_map() -> ZocaloInfoGenerator: - """ - Generate the zocalo trigger info from bluesky runs where the frame number is - computed using metadata added to the document by the ISPyB callback and the - run start which together can be used to determine the correct frame numbering. - """ - doc = yield [] - omega_to_scan_spec = doc["omega_to_scan_spec"] - start_frame = 0 - infos = [] - for i, omega in enumerate([0, 90]): - frames = number_of_frames_from_scan_spec(omega_to_scan_spec[omega]) - infos.append( - ZocaloStartInfo( - doc["grid_plane_to_id_map"][omega], None, start_frame, frames, i - ) - ) - start_frame += frames - yield infos - - -def create_gridscan_callbacks() -> tuple[ - GridscanNexusFileCallback, GridscanISPyBCallback -]: + + +def create_gridscan_callbacks( + params: I02_1FgsParams, +) -> tuple[GridscanNexusFileCallback, GridscanISPyBCallback]: return ( GridscanNexusFileCallback(param_type=SpecifiedTwoDGridScan), GridscanISPyBCallback( @@ -85,7 +63,7 @@ def create_gridscan_callbacks() -> tuple[ emit=ZocaloCallback( PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV, - generate_start_info_from_omega_map, # todo dont need this + lambda: generate_start_info_from_omega_map(params.omega_starts_deg), ), ), ) @@ -182,7 +160,7 @@ def i02_1_gridscan_plan( """BlueAPI entry point for i02-1 grid scans""" beamline_specific = construct_i02_1_specific_features(composite, parameters) - callbacks = create_gridscan_callbacks() + callbacks = create_gridscan_callbacks(parameters) @bpp.subs_decorator(callbacks) def decorated_flyscan_plan(): diff --git a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py index 305fac78d2..2fb0518a2f 100644 --- a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py +++ b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py @@ -64,14 +64,17 @@ ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridDetectAndScanISPyBCallback, - generate_start_info_from_omega_map, ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + generate_start_info_from_omega_map, +) from mx_bluesky.common.parameters.components import get_param_version from mx_bluesky.common.parameters.constants import ( EnvironmentConstants, + GridscanParamConstants, OavConstants, PlanGroupCheckpointConstants, PlanNameConstants, @@ -81,6 +84,7 @@ ) from mx_bluesky.common.parameters.gridscan import ( GenericGrid, + SpecifiedGrids, SpecifiedThreeDGridScan, ) from mx_bluesky.common.preprocessors.preprocessors import ( @@ -208,7 +212,10 @@ def _inner_grid_detect_then_xrc(): # Hyperion handles its callbacks differently to BlueAPI-managed plans, see # https://github.com/DiamondLightSource/mx-bluesky/issues/1117 flyscan_event_handler = XRayCentreEventHandler() - callbacks = *create_gridscan_callbacks(), flyscan_event_handler + callbacks = ( + *create_gridscan_callbacks(grid_common_params.specified_grid_params), + flyscan_event_handler, + ) @bpp.subs_decorator(callbacks) @verify_undulator_gap_before_run_decorator(composite) @@ -269,17 +276,19 @@ def get_ready_for_oav_and_close_shutter( yield from bps.wait(group) -def create_gridscan_callbacks() -> tuple[ - GridscanNexusFileCallback, GridDetectAndScanISPyBCallback -]: +def create_gridscan_callbacks( + params: SpecifiedGrids | None, +) -> tuple[GridscanNexusFileCallback, GridDetectAndScanISPyBCallback]: return ( GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan), GridDetectAndScanISPyBCallback( - param_type=GenericGrid, + param_type=SpecifiedThreeDGridScan, emit=ZocaloCallback( PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV, - generate_start_info_from_omega_map, + lambda: generate_start_info_from_omega_map( + [GridscanParamConstants.OMEGA_1, GridscanParamConstants.OMEGA_2] + ), ), ), ) diff --git a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py index 5de9ea845b..2860891d8a 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py @@ -176,7 +176,7 @@ def _handle_ispyb_transmission_flux_read( if beamsize_y_mm: beamsize_y_mm = beamsize_y_mm / 1000 if not (beamsize_x_mm and beamsize_y_mm): - # todo write issue about vmxm beamsize and link here + # VMXm don't have a beamsize device in dodal yet, they get beamsize sent in from GDA try: beamsize_x_mm = self.params.beam_size_x # type: ignore beamsize_y_mm = self.params.beam_size_y # type: ignore diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py index d0c4025ec3..8dca22d8bd 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py @@ -8,7 +8,6 @@ from bluesky import preprocessors as bpp from bluesky.utils import MsgGenerator, make_decorator -from dodal.devices.zocalo import ZocaloStartInfo from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import ( BaseISPyBCallback, @@ -18,12 +17,14 @@ populate_data_collection_group, populate_remaining_data_collection_info, ) -from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( - ZocaloInfoGenerator, -) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( construct_comment_for_gridscan, ) +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + ASSERT_START_BEFORE_EVENT_DOC_MESSAGE, + add_processing_time_to_comment, + common_populate_axis_info, +) from mx_bluesky.common.external_interaction.ispyb.data_model import ( DataCollectionGridInfo, DataCollectionGroupInfo, @@ -38,13 +39,12 @@ ) from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample from mx_bluesky.common.parameters.constants import DocDescriptorNames, PlanNameConstants -from mx_bluesky.common.parameters.gridscan import GenericGrid +from mx_bluesky.common.parameters.gridscan import SpecifiedGrids from mx_bluesky.common.utils.exceptions import ( ISPyBDepositionNotMadeError, SampleError, ) from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER, set_dcgid_tag -from mx_bluesky.common.utils.utils import number_of_frames_from_scan_spec OMEGA_TOLERANCE = 1 @@ -57,8 +57,7 @@ class GridscanPlane(StrEnum): if TYPE_CHECKING: from event_model import Event, RunStart, RunStop -T = TypeVar("T", bound="GenericGrid") -ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" +T = TypeVar("T", bound="SpecifiedGrids") def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): @@ -173,16 +172,8 @@ def activity_gated_event(self, doc: Event): return doc def _add_processing_time_to_comment(self, processing_start_time: float): - assert self.data_collection_group_info, ASSERT_START_BEFORE_EVENT_DOC_MESSAGE - proc_time = time() - processing_start_time - crystal_summary = f"Zocalo processing took {proc_time:.2f} s." - - self.data_collection_group_info.comments = ( - self.data_collection_group_info.comments or "" - ) + crystal_summary - - self.ispyb.append_to_comment( - self.ispyb_ids.data_collection_ids[0], crystal_summary + add_processing_time_to_comment( + self, processing_start_time, self.data_collection_group_info ) def _handle_oav_grid_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]: @@ -258,14 +249,7 @@ def _handle_oav_grid_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]: return [scan_data_info] def _populate_axis_info(self, data_collection_info: DataCollectionInfo, doc: dict): - if (omega_start := doc.get("gonio-omega")) 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 - if (chi_start := doc.get("gonio-chi")) is not None: - data_collection_info.chi_start = chi_start + common_populate_axis_info(data_collection_info, doc) def populate_info_for_update( self, @@ -340,27 +324,6 @@ def data_collection_number_from_gridplane(self, plane) -> int: return base_number if plane == GridscanPlane.OMEGA_XY else base_number + 1 -def generate_start_info_from_omega_map() -> ZocaloInfoGenerator: - """ - Generate the zocalo trigger info from bluesky runs where the frame number is - computed using metadata added to the document by the ISPyB callback and the - run start which together can be used to determine the correct frame numbering. - """ - doc = yield [] - omega_to_scan_spec = doc["omega_to_scan_spec"] - start_frame = 0 - infos = [] - for i, omega in enumerate([GridscanPlane.OMEGA_XY, GridscanPlane.OMEGA_XZ]): - frames = number_of_frames_from_scan_spec(omega_to_scan_spec[omega]) - infos.append( - ZocaloStartInfo( - doc["grid_plane_to_id_map"][omega], None, start_frame, frames, i - ) - ) - start_frame += frames - yield infos - - def _smargon_omega_to_xyxz_plane(smargon_omega: float) -> GridscanPlane: modulo_180 = abs(smargon_omega) % 180 is_xy = isclose(modulo_180, 0, abs_tol=OMEGA_TOLERANCE) diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py index b7e6d6a893..05d70296d2 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py @@ -1,5 +1,6 @@ from __future__ import annotations +from abc import abstractmethod from collections.abc import Callable, Sequence from time import time from typing import TYPE_CHECKING, Any, TypeVar @@ -7,10 +8,8 @@ from bluesky import preprocessors as bpp from bluesky.utils import MsgGenerator, make_decorator -from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import ( BaseISPyBCallback, - D, ) from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import ( populate_data_collection_group, @@ -18,17 +17,15 @@ ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridscanPlane, - _smargon_omega_to_xyxz_plane, ) -from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( - construct_comment_for_gridscan, +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + add_processing_time_to_comment, + common_populate_axis_info, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( - DataCollectionGridInfo, DataCollectionGroupInfo, DataCollectionInfo, DataCollectionPositionInfo, - Orientation, ScanDataInfo, ) from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( @@ -47,8 +44,7 @@ if TYPE_CHECKING: from event_model import RunStart, RunStop -T = TypeVar("T", bound="I02_1FgsParams") -ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" +T = TypeVar("T", bound="SpecifiedGrids") def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): @@ -56,12 +52,12 @@ def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): bpp.run_wrapper( plan_generator, md={ - "activate_callbacks": ["GridDetectAndScanISPyBCallback"], - "subplan_name": PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, + "activate_callbacks": ["GridscanISPyBCallback"], + "subplan_name": PlanNameConstants.TRIGGER_GRIDSCAN_ISPYB_CALLBACK, "mx_bluesky_parameters": parameters.model_dump_json(), }, ), - PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN, + PlanNameConstants.TRIGGER_GRIDSCAN_ISPYB_CALLBACK, ) @@ -74,6 +70,9 @@ class GridscanISPyBCallback(BaseISPyBCallback): receiving an 'event' document for the 'ispyb_reading_hardware' event, and updates the deposition on receiving its final 'stop' document. + This callback should be used when grid parameters have been sent in to BlueAPI as part + of entry parameters. + To use, subscribe the Bluesky RunEngine to an instance of this class. E.g.: ispyb_handler_callback = FGSISPyBCallback(parameters) @@ -116,113 +115,20 @@ def activity_gated_start(self, doc: RunStart): return super().activity_gated_start(doc) def _add_processing_time_to_comment(self, processing_start_time: float): - assert self.data_collection_group_info, ASSERT_START_BEFORE_EVENT_DOC_MESSAGE - proc_time = time() - processing_start_time - crystal_summary = f"Zocalo processing took {proc_time:.2f} s." - - self.data_collection_group_info.comments = ( - self.data_collection_group_info.comments or "" - ) + crystal_summary - - self.ispyb.append_to_comment( - self.ispyb_ids.data_collection_ids[0], crystal_summary + add_processing_time_to_comment( + self, processing_start_time, self.data_collection_group_info ) - def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: - """ - so grid information is available immediately after the plan is triggered. - In contrast, i03 and i04 use the OAV to automatically detect their grid. - """ - assert isinstance(self.params, I02_1FgsParams) - assert self.ispyb_ids.data_collection_ids, "No current data collection" - assert self.data_collection_group_info, "No data collection group" - data = doc["data"] - scan_data_infos = [] - - for grid_num in range(self.params.num_grids): - omega = data.get("gonio-omega", self.params.omega_starts_deg[grid_num]) - - # Don't need to do deal with the grid plane here since vmxm only do - # one plane, but leave it in so it's easier to standardise in the future - grid_plane = _smargon_omega_to_xyxz_plane(omega) - ISPYB_ZOCALO_CALLBACK_LOGGER.info( - f"Generating dc info for gridplane {grid_plane}, omega {omega}" - ) - data_collection_number = self.data_collection_number_from_gridplane( - grid_plane - ) - file_template = f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5" - # Snapshots have already been taken in GDA - - # todo the _get_scan_infos here should be vmxm specific since they require us asserting that params includes path_to_xtal_snapshot - # and do this by asserting params are vmxm params + @abstractmethod + def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: ... - data_collection_info = DataCollectionInfo( - xtal_snapshot1=str(self.params.path_to_xtal_snapshot), - xtal_snapshot2=str(self.params.path_to_xtal_snapshot), - xtal_snapshot3=str(self.params.path_to_xtal_snapshot), - n_images=self.params.num_images, - data_collection_number=data_collection_number, - file_template=file_template, - ) - data_collection_grid_info = DataCollectionGridInfo( - dx_in_mm=self.params.x_step_size_um * 1000, - dy_in_mm=self.params.y_step_sizes_um[grid_num] * 1000, - steps_x=self.params.x_steps, - steps_y=self.params.y_steps[grid_num], - microns_per_pixel_x=self.params.microns_per_pixel_x, - microns_per_pixel_y=self.params.microns_per_pixel_y, - snapshot_offset_x_pixel=self.params.upper_left_x, - snapshot_offset_y_pixel=self.params.upper_left_y, - orientation=Orientation.HORIZONTAL, - snaked=True, - ) - data_collection_info.comments = construct_comment_for_gridscan( - data_collection_grid_info - ) - - # Grid plane logic isn't needed for VMX, but keep it for now anyway - data_collection_id = self.ispyb_ids.data_collection_ids[ - 0 if grid_plane == GridscanPlane.OMEGA_XY else 1 - ] - self._grid_plane_to_id_map[grid_plane] = data_collection_id - self._grid_plane_to_width_map[grid_plane] = ( - data_collection_grid_info.steps_y - ) - - y_steps = self._grid_plane_to_width_map.get(GridscanPlane.OMEGA_XY, "_") - self.data_collection_group_info.comments = ( - f"Diffraction grid scan of {data_collection_grid_info.steps_x} by " - f"{y_steps}." - ) - - self._populate_axis_info(data_collection_info, doc["data"]) - - # todo do all this stuff as soon as possible after plan starts - - scan_data_info = ScanDataInfo( - data_collection_info=data_collection_info, - data_collection_id=data_collection_id, - data_collection_grid_info=data_collection_grid_info, - ) - - scan_data_infos.append(scan_data_info) - - ISPYB_ZOCALO_CALLBACK_LOGGER.info( - "Updating ispyb data collection after loading grid params" - ) - - return scan_data_infos + """ + Use grid parameters to create a sequence of ScanDataInfos. See + i02-1's gridscan ispyb callback for example implementation. + """ def _populate_axis_info(self, data_collection_info: DataCollectionInfo, doc: dict): - if (omega_start := doc.get("gonio-omega")) 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 - if (chi_start := doc.get("gonio-chi")) is not None: - data_collection_info.chi_start = chi_start + common_populate_axis_info(data_collection_info, doc) def populate_info_for_update( self, @@ -233,7 +139,7 @@ def populate_info_for_update( assert self.ispyb_ids.data_collection_ids, ( "Expect at least one valid data collection to record scan data" ) - assert isinstance(self.params, I02_1FgsParams) + assert isinstance(self.params, SpecifiedGrids) scan_data_infos = [] for grid_num in range(self.params.num_grids): scan_data_info = ScanDataInfo( @@ -271,22 +177,12 @@ def activity_gated_stop(self, doc: RunStop) -> RunStop: self.ispyb_ids.data_collection_group_id, ) self.data_collection_group_info = None - self._grid_plane_to_id_map.clear() - self._grid_plane_to_width_map.clear() return super().activity_gated_stop(doc) return self.tag_doc(doc) - def tag_doc(self, doc: D) -> D: - doc = super().tag_doc(doc) - assert isinstance(doc, dict) - if self._grid_plane_to_id_map: - doc["grid_plane_to_id_map"] = self._grid_plane_to_id_map - return doc # type: ignore - def data_collection_number_from_gridplane(self, plane) -> int: assert self.params - base_number = self.params.detector_params.run_number - return base_number if plane == GridscanPlane.OMEGA_XY else base_number + 1 + return self.params.detector_params.run_number def fill_gridscan_deposition_and_store( self, make_scan_infos_with_grid_info: Callable[..., Sequence[ScanDataInfo]] @@ -313,8 +209,12 @@ def fill_gridscan_deposition_and_store( self.data_collection_group_info, scan_data_infos ) # Now use grid information to complete deposition - scan_data_infos: list[ScanDataInfo] = list(make_scan_infos_with_grid_info()) - self.ispyb_ids = self.ispyb.update_deposition(self.ispyb_ids, scan_data_infos) + scan_data_infos_list: list[ScanDataInfo] = list( + make_scan_infos_with_grid_info() + ) + self.ispyb_ids = self.ispyb.update_deposition( + self.ispyb_ids, scan_data_infos_list + ) self.ispyb.update_data_collection_group_table( self.data_collection_group_info, self.ispyb_ids.data_collection_group_id ) diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py new file mode 100644 index 0000000000..45f98cb71d --- /dev/null +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from time import time + +from dodal.devices.zocalo import ZocaloStartInfo + +from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import ( + BaseISPyBCallback, +) +from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( + ZocaloInfoGenerator, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionGroupInfo, + DataCollectionInfo, +) +from mx_bluesky.common.parameters.constants import PlanNameConstants +from mx_bluesky.common.utils.utils import number_of_frames_from_scan_spec + +ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" + + +def generate_start_info_from_omega_map( + omega_positions: list[int], +) -> ZocaloInfoGenerator: + """ + Generate the zocalo trigger info from bluesky runs where the frame number is + computed using metadata added to the document by the ISPyB callback and the + run start which together can be used to determine the correct frame numbering. + """ + doc = yield [] + omega_to_scan_spec = doc["omega_to_scan_spec"] + start_frame = 0 + infos = [] + omegas_str = [str(omega) for omega in omega_positions] + for i, omega in enumerate(omegas_str): + frames = number_of_frames_from_scan_spec(omega_to_scan_spec[omega]) + infos.append( + ZocaloStartInfo( + doc["grid_plane_to_id_map"][omega], None, start_frame, frames, i + ) + ) + start_frame += frames + yield infos + + +def common_populate_axis_info(data_collection_info: DataCollectionInfo, doc: dict): + if (omega_start := doc.get("gonio-omega")) 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 + if (chi_start := doc.get("gonio-chi")) is not None: + data_collection_info.chi_start = chi_start + + +def add_processing_time_to_comment( + callback: BaseISPyBCallback, + processing_start_time: float, + data_collection_group_info: DataCollectionGroupInfo | None, +): + assert data_collection_group_info, ASSERT_START_BEFORE_EVENT_DOC_MESSAGE + proc_time = time() - processing_start_time + crystal_summary = f"Zocalo processing took {proc_time:.2f} s." + + data_collection_group_info.comments = ( + data_collection_group_info.comments or "" + ) + crystal_summary + + callback.ispyb.append_to_comment( + callback.ispyb_ids.data_collection_ids[0], crystal_summary + ) diff --git a/src/mx_bluesky/common/parameters/components.py b/src/mx_bluesky/common/parameters/components.py index ca7aebf68e..1c54bf9a00 100644 --- a/src/mx_bluesky/common/parameters/components.py +++ b/src/mx_bluesky/common/parameters/components.py @@ -271,7 +271,7 @@ def _start_for_axis(self, axis: XyzAxis, grid: int) -> float: class OptionalGonioAngleStarts(BaseModel): # Gridscans have different omega starts - omega_starts_deg: list[float] = [0, 90] + omega_starts_deg: list[int] = [0, 90] phi_start_deg: float | None = None chi_start_deg: float | None = None diff --git a/src/mx_bluesky/common/parameters/constants.py b/src/mx_bluesky/common/parameters/constants.py index 158dbe28af..a7961b1c6e 100644 --- a/src/mx_bluesky/common/parameters/constants.py +++ b/src/mx_bluesky/common/parameters/constants.py @@ -71,6 +71,7 @@ class PlanNameConstants: GRIDSCAN_OUTER = "run_gridscan_move_and_tidy" DO_FGS = "do_fgs" FLYSCAN_RESULTS = "xray_centre_results" + TRIGGER_GRIDSCAN_ISPYB_CALLBACK = "trigger gridscan ispyb callback" # Rotation scan ROTATION_MULTI = "multi_rotation_wrapper" ROTATION_MULTI_OUTER = "multi_rotation_outer" diff --git a/src/mx_bluesky/common/parameters/gridscan.py b/src/mx_bluesky/common/parameters/gridscan.py index 71469e56eb..e6bebfa63e 100644 --- a/src/mx_bluesky/common/parameters/gridscan.py +++ b/src/mx_bluesky/common/parameters/gridscan.py @@ -124,7 +124,7 @@ class SpecifiedGrids(GenericGrid, XyzStarts, WithScan, Generic[GridScanParamType # See https://github.com/DiamondLightSource/mx-bluesky/issues/1634 for a better structure for this # class - omega_starts_deg: list[float] = Field( + omega_starts_deg: list[int] = Field( default=[GridscanParamConstants.OMEGA_1, GridscanParamConstants.OMEGA_2] ) x_step_size_um: PositiveFloat = Field( diff --git a/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py b/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py index cad6482812..297446fe56 100644 --- a/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py +++ b/src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py @@ -27,14 +27,17 @@ ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridDetectAndScanISPyBCallback, - generate_start_info_from_omega_map, ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + generate_start_info_from_omega_map, +) from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import ( SampleHandlingCallback, ) +from mx_bluesky.common.parameters.constants import GridscanParamConstants from mx_bluesky.common.utils.log import ( ISPYB_ZOCALO_CALLBACK_LOGGER, NEXUS_LOGGER, @@ -63,7 +66,6 @@ from mx_bluesky.hyperion.parameters.cli import CallbackArgs, parse_callback_args from mx_bluesky.hyperion.parameters.constants import CONST from mx_bluesky.hyperion.parameters.gridscan import ( - GenericGridWithHyperionDetectorParams, HyperionSpecifiedThreeDGridScan, ) @@ -80,9 +82,13 @@ def create_gridscan_callbacks() -> tuple[ return ( GridscanNexusFileCallback(param_type=HyperionSpecifiedThreeDGridScan), GridDetectAndScanISPyBCallback( - param_type=GenericGridWithHyperionDetectorParams, + param_type=HyperionSpecifiedThreeDGridScan, emit=ZocaloCallback( - CONST.PLAN.DO_FGS, CONST.ZOCALO_ENV, generate_start_info_from_omega_map + CONST.PLAN.DO_FGS, + CONST.ZOCALO_ENV, + lambda: generate_start_info_from_omega_map( + [GridscanParamConstants.OMEGA_1, GridscanParamConstants.OMEGA_2] + ), ), ), ) From 6d2dcd82599335eb5540a9f3e23d588add9bdd53 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Fri, 27 Feb 2026 16:11:43 +0000 Subject: [PATCH 5/8] Fixes --- pyproject.toml | 2 +- .../beamlines/i02_1/i02_1_gridscan_plan.py | 4 +-- .../i04_grid_detect_then_xray_centre_plan.py | 9 +++--- .../grid_detect_and_scan/ispyb_callback.py | 4 +-- .../callbacks/grid/gridscan/ispyb_callback.py | 18 ++++++++---- .../callbacks/grid/utils.py | 28 +++++++++++++++++++ .../xray_centre/test_ispyb_handler.py | 8 ++++-- tests/unit_tests/conftest.py | 9 ++++-- 8 files changed, 62 insertions(+), 20 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9b8fedcbe5..07bc7310d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,7 @@ dev = [ "tox-uv", "types-mock", "types-requests", - "sphinxcontrib.mermaid", # To build dodal docs + "sphinxcontrib.mermaid", # To build dodal docs ] [project.scripts] diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index 3afa851898..702f16de97 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -40,7 +40,7 @@ GridscanISPyBCallback, ) from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( - generate_start_info_from_omega_map, + generate_start_info_from_num_grids, ) from mx_bluesky.common.parameters.constants import ( EnvironmentConstants, @@ -63,7 +63,7 @@ def create_gridscan_callbacks( emit=ZocaloCallback( PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV, - lambda: generate_start_info_from_omega_map(params.omega_starts_deg), + lambda: generate_start_info_from_num_grids(params), ), ), ) diff --git a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py index 2fb0518a2f..ae11e4c232 100644 --- a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py +++ b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py @@ -84,7 +84,6 @@ ) from mx_bluesky.common.parameters.gridscan import ( GenericGrid, - SpecifiedGrids, SpecifiedThreeDGridScan, ) from mx_bluesky.common.preprocessors.preprocessors import ( @@ -213,7 +212,7 @@ def _inner_grid_detect_then_xrc(): # https://github.com/DiamondLightSource/mx-bluesky/issues/1117 flyscan_event_handler = XRayCentreEventHandler() callbacks = ( - *create_gridscan_callbacks(grid_common_params.specified_grid_params), + *create_gridscan_callbacks(), flyscan_event_handler, ) @@ -276,9 +275,9 @@ def get_ready_for_oav_and_close_shutter( yield from bps.wait(group) -def create_gridscan_callbacks( - params: SpecifiedGrids | None, -) -> tuple[GridscanNexusFileCallback, GridDetectAndScanISPyBCallback]: +def create_gridscan_callbacks() -> tuple[ + GridscanNexusFileCallback, GridDetectAndScanISPyBCallback +]: return ( GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan), GridDetectAndScanISPyBCallback( diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py index 8dca22d8bd..c118a51bfd 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py @@ -39,7 +39,7 @@ ) from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample from mx_bluesky.common.parameters.constants import DocDescriptorNames, PlanNameConstants -from mx_bluesky.common.parameters.gridscan import SpecifiedGrids +from mx_bluesky.common.parameters.gridscan import GenericGrid from mx_bluesky.common.utils.exceptions import ( ISPyBDepositionNotMadeError, SampleError, @@ -57,7 +57,7 @@ class GridscanPlane(StrEnum): if TYPE_CHECKING: from event_model import Event, RunStart, RunStop -T = TypeVar("T", bound="SpecifiedGrids") +T = TypeVar("T", bound="GenericGrid") def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py index 05d70296d2..75962b48eb 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py @@ -15,9 +15,6 @@ populate_data_collection_group, populate_remaining_data_collection_info, ) -from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( - GridscanPlane, -) from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( add_processing_time_to_comment, common_populate_axis_info, @@ -45,6 +42,7 @@ from event_model import RunStart, RunStop T = TypeVar("T", bound="SpecifiedGrids") +D = TypeVar("D") def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters): @@ -93,8 +91,7 @@ def __init__( self.param_type = param_type self._start_of_fgs_uid: str | None = None self._processing_start_time: float | None = None - self._grid_plane_to_id_map: dict[GridscanPlane, int] = {} - self._grid_plane_to_width_map: dict[GridscanPlane, int] = {} + self._grid_num_to_id_map: dict[int, int] = {} self.data_collection_group_info: DataCollectionGroupInfo | None def activity_gated_start(self, doc: RunStart): @@ -142,9 +139,11 @@ def populate_info_for_update( assert isinstance(self.params, SpecifiedGrids) scan_data_infos = [] for grid_num in range(self.params.num_grids): + id = self.ispyb_ids.data_collection_ids[grid_num] + self._grid_num_to_id_map[grid_num] = id scan_data_info = ScanDataInfo( data_collection_info=event_sourced_data_collection_info, - data_collection_id=self.ispyb_ids.data_collection_ids[grid_num], + data_collection_id=id, ) scan_data_infos.append(scan_data_info) return scan_data_infos @@ -180,6 +179,13 @@ def activity_gated_stop(self, doc: RunStop) -> RunStop: return super().activity_gated_stop(doc) return self.tag_doc(doc) + def tag_doc(self, doc: D) -> D: + doc = super().tag_doc(doc) + assert isinstance(doc, dict) + if self._grid_num_to_id_map: + doc["_grid_num_to_id_map"] = self._grid_num_to_id_map + return doc # type: ignore + def data_collection_number_from_gridplane(self, plane) -> int: assert self.params return self.params.detector_params.run_number diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py index 45f98cb71d..18ac853878 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py @@ -15,6 +15,7 @@ DataCollectionInfo, ) from mx_bluesky.common.parameters.constants import PlanNameConstants +from mx_bluesky.common.parameters.gridscan import SpecifiedGrids from mx_bluesky.common.utils.utils import number_of_frames_from_scan_spec ASSERT_START_BEFORE_EVENT_DOC_MESSAGE = f"No data collection group info - event document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" @@ -44,6 +45,33 @@ def generate_start_info_from_omega_map( yield infos +def generate_start_info_from_num_grids( + params: SpecifiedGrids, +) -> ZocaloInfoGenerator: + """ + Generate the zocalo trigger info from bluesky runs where the grid specs + are immediately known from entry parameters. Metadata added to the document + by the ispyb callback maps the data collection id to the grid number. + """ + + doc = yield [] + start_frame = 0 + infos = [] + for grid_num in range(params.num_grids): + frames = len(params.scan_points[grid_num]) + infos.append( + ZocaloStartInfo( + doc["_grid_num_to_id_map"][grid_num], + None, + start_frame, + frames, + grid_num, + ) + ) + start_frame += frames + yield infos + + def common_populate_axis_info(data_collection_info: DataCollectionInfo, doc: dict): if (omega_start := doc.get("gonio-omega")) is not None: omega_in_gda_space = -omega_start diff --git a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py index d967e716f5..9e6359da6b 100644 --- a/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py +++ b/tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_handler.py @@ -196,12 +196,16 @@ def test_given_ispyb_callback_finished_writing_to_ispyb_when_messages_logged_the latest_record = gelf_handler.emit.call_args.args[-1] assert not hasattr(latest_record, "dc_group_id") + @patch( + "mx_bluesky.common.external_interaction.callbacks.grid.utils.time", + new=MagicMock(side_effect=[100]), + ) @patch( "mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback.time", - side_effect=[2, 100], + new=MagicMock(side_effect=[2]), ) def test_given_fgs_plan_finished_when_zocalo_results_event_then_expected_comment_deposited( - self, mock_time, dummy_rotation_data_collection_group_info, test_event_data + self, dummy_rotation_data_collection_group_info, test_event_data ): ispyb_handler = GridDetectAndScanISPyBCallback( param_type=GenericGridWithHyperionDetectorParams, diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index f373841b3a..2cd2120ac8 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -51,11 +51,13 @@ ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import ( GridDetectAndScanISPyBCallback, - generate_start_info_from_omega_map, ) from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import ( GridscanNexusFileCallback, ) +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + generate_start_info_from_omega_map, +) from mx_bluesky.common.external_interaction.ispyb.data_model import ( DataCollectionGroupInfo, ) @@ -66,6 +68,7 @@ from mx_bluesky.common.parameters.constants import ( DocDescriptorNames, EnvironmentConstants, + GridscanParamConstants, PlanNameConstants, ) from mx_bluesky.common.parameters.device_composites import ( @@ -179,7 +182,9 @@ def create_gridscan_callbacks() -> tuple[ emit=ZocaloCallback( PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV, - generate_start_info_from_omega_map, + lambda: generate_start_info_from_omega_map( + [GridscanParamConstants.OMEGA_1, GridscanParamConstants.OMEGA_2] + ), ), ), ) From 2e49666e44dd2dd1825446548b20953c79718abc Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Fri, 27 Feb 2026 16:38:51 +0000 Subject: [PATCH 6/8] Add test stubs --- .../beamlines/i02_1/i02_1_gridscan_plan.py | 2 ++ .../i04_grid_detect_then_xray_centre_plan.py | 2 +- .../grid_detect_and_scan/ispyb_callback.py | 4 ++-- .../callbacks/grid/gridscan/ispyb_callback.py | 6 ++--- .../callbacks/grid/utils.py | 2 +- .../i02_1/test_gridscan_ispyb_callback.py | 1 + .../i02_1/test_i02_1_gridscan_plan.py | 23 +++++++++++++++---- .../common/test_ispyb_callback_base.py | 5 ++++ .../callbacks/grid/__init__.py | 0 .../grid/test_ispyb_gridscan_callback.py | 1 + .../callbacks/grid/test_utils.py | 1 + 11 files changed, 35 insertions(+), 12 deletions(-) create mode 100644 tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py create mode 100644 tests/unit_tests/common/external_interaction/callbacks/grid/__init__.py create mode 100644 tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py create mode 100644 tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index 702f16de97..282de76d6e 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -38,6 +38,7 @@ ) from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( GridscanISPyBCallback, + ispyb_activation_decorator, ) from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( generate_start_info_from_num_grids, @@ -162,6 +163,7 @@ def i02_1_gridscan_plan( beamline_specific = construct_i02_1_specific_features(composite, parameters) callbacks = create_gridscan_callbacks(parameters) + @ispyb_activation_decorator @bpp.subs_decorator(callbacks) def decorated_flyscan_plan(): yield from common_flyscan_xray_centre(composite, parameters, beamline_specific) diff --git a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py index ae11e4c232..d205ccff53 100644 --- a/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py +++ b/src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py @@ -281,7 +281,7 @@ def create_gridscan_callbacks() -> tuple[ return ( GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan), GridDetectAndScanISPyBCallback( - param_type=SpecifiedThreeDGridScan, + param_type=GenericGrid, emit=ZocaloCallback( PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV, diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py index c118a51bfd..39c914fdfd 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/grid_detect_and_scan/ispyb_callback.py @@ -22,7 +22,7 @@ ) from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( ASSERT_START_BEFORE_EVENT_DOC_MESSAGE, - add_processing_time_to_comment, + common_add_processing_time_to_comment, common_populate_axis_info, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( @@ -172,7 +172,7 @@ def activity_gated_event(self, doc: Event): return doc def _add_processing_time_to_comment(self, processing_start_time: float): - add_processing_time_to_comment( + common_add_processing_time_to_comment( self, processing_start_time, self.data_collection_group_info ) diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py index 75962b48eb..ab0b523cbe 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py @@ -16,7 +16,7 @@ populate_remaining_data_collection_info, ) from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( - add_processing_time_to_comment, + common_add_processing_time_to_comment, common_populate_axis_info, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( @@ -112,7 +112,7 @@ def activity_gated_start(self, doc: RunStart): return super().activity_gated_start(doc) def _add_processing_time_to_comment(self, processing_start_time: float): - add_processing_time_to_comment( + common_add_processing_time_to_comment( self, processing_start_time, self.data_collection_group_info ) @@ -150,7 +150,7 @@ def populate_info_for_update( def activity_gated_stop(self, doc: RunStop) -> RunStop: assert self.data_collection_group_info, ( - f"No data collection group info - stop document has been emitted before a {PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN} start document" + f"No data collection group info - stop document has been emitted before a {PlanNameConstants.DO_FGS} start document" ) if doc.get("run_start") == self._start_of_fgs_uid: self._processing_start_time = time() diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py index 18ac853878..2729ef9f90 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py @@ -83,7 +83,7 @@ def common_populate_axis_info(data_collection_info: DataCollectionInfo, doc: dic data_collection_info.chi_start = chi_start -def add_processing_time_to_comment( +def common_add_processing_time_to_comment( callback: BaseISPyBCallback, processing_start_time: float, data_collection_group_info: DataCollectionGroupInfo | None, diff --git a/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py b/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py new file mode 100644 index 0000000000..5ed1f1320a --- /dev/null +++ b/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py @@ -0,0 +1 @@ +def test_get_scan_infos_gives_expected_output(): ... diff --git a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py index 49f5cb5513..68a8de4fee 100644 --- a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py +++ b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py @@ -14,6 +14,7 @@ from dodal.devices.zebra.zebra import Zebra from pydantic import ValidationError +from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import ( FlyScanXRayCentreComposite, construct_i02_1_specific_features, @@ -27,8 +28,8 @@ @pytest.fixture -def fgs_params_two_d(tmp_path) -> SpecifiedTwoDGridScan: - return SpecifiedTwoDGridScan( +def fgs_params_two_d(tmp_path) -> I02_1FgsParams: + return I02_1FgsParams( x_start_um=0, y_starts_um=[0], z_starts_um=[0], @@ -41,6 +42,13 @@ def fgs_params_two_d(tmp_path) -> SpecifiedTwoDGridScan: storage_directory=str(tmp_path), x_steps=5, y_steps=[3], + path_to_xtal_snapshot=tmp_path, + beam_size_x=0, + beam_size_y=0, + microns_per_pixel_x=1, + microns_per_pixel_y=1, + upper_left_x=0, + upper_left_y=0, ) @@ -96,7 +104,7 @@ def _check_lengths_are_same(self): # type: ignore def test_two_d_grid_scan_validation( y_starts_um: list[float], z_starts_um: list[float], - omega_starts_deg: list[float], + omega_starts_deg: list[int], y_step_sizes_um: list[float], y_steps: list[int], should_raise: bool, @@ -135,7 +143,7 @@ def test_i02_1_flyscan_xray_centre_in_re( mock_common_scan: MagicMock, mock_create_features: MagicMock, run_engine: RunEngine, - fgs_params_two_d: SpecifiedTwoDGridScan, + fgs_params_two_d: I02_1FgsParams, fgs_composite: FlyScanXRayCentreComposite, ): expected_features = construct_i02_1_specific_features( @@ -143,7 +151,12 @@ def test_i02_1_flyscan_xray_centre_in_re( ) mock_create_features.return_value = expected_features - run_engine(i02_1_gridscan_plan(fgs_params_two_d, fgs_composite)) + run_engine( + i02_1_gridscan_plan(fgs_params_two_d, fgs_composite) + ) # todo fix this typing mock_common_scan.assert_called_once_with( fgs_composite, fgs_params_two_d, expected_features ) + + +def test_ispyb_callbacks_activated_when_expected(): ... diff --git a/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py b/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py index b09ad0c957..a6c66f06b3 100644 --- a/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py +++ b/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py @@ -58,3 +58,8 @@ def test_plan(): with pytest.raises(ValueError): run_engine(test_plan()) + + +def test_handle_ispyb_transmission_flux_read_if_no_beamsize(): + # Test exception and check if params exist then it still works + ... diff --git a/tests/unit_tests/common/external_interaction/callbacks/grid/__init__.py b/tests/unit_tests/common/external_interaction/callbacks/grid/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py b/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py new file mode 100644 index 0000000000..c567398461 --- /dev/null +++ b/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py @@ -0,0 +1 @@ +# todo use tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_callback.py as a template to test this file diff --git a/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py b/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py new file mode 100644 index 0000000000..00e89fc4a0 --- /dev/null +++ b/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py @@ -0,0 +1 @@ +def test_generate_start_info_from_num_grids(): ... From 7c691b774ea9e145e9eafa5cb922e60eec4ce610 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Mon, 2 Mar 2026 15:23:47 +0000 Subject: [PATCH 7/8] Add fixes and tests --- .../callbacks/gridscan/ispyb_callback.py | 13 +- .../beamlines/i02_1/i02_1_gridscan_plan.py | 6 +- .../experiment_plans/inner_plans/do_fgs.py | 4 +- .../callbacks/grid/gridscan/ispyb_callback.py | 6 +- .../callbacks/grid/utils.py | 3 +- .../common/parameters/components.py | 2 + tests/unit_tests/beamlines/i02_1/conftest.py | 30 +++++ .../i02_1/test_gridscan_ispyb_callback.py | 77 +++++++++++- .../i02_1/test_i02_1_gridscan_plan.py | 84 ++++++++----- .../common/test_ispyb_callback_base.py | 51 +++++++- .../grid/test_ispyb_gridscan_callback.py | 113 +++++++++++++++++- .../callbacks/grid/test_utils.py | 29 ++++- .../common/parameters/test_gridscan.py | 4 +- 13 files changed, 369 insertions(+), 53 deletions(-) create mode 100644 tests/unit_tests/beamlines/i02_1/conftest.py diff --git a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py index 0c6830078b..3e9affe91f 100644 --- a/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py @@ -5,7 +5,7 @@ construct_comment_for_gridscan, ) from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( - GridscanISPyBCallback, + GridscanISPyBCallback as CommonGridscanISPyBCallback, ) from mx_bluesky.common.external_interaction.ispyb.data_model import ( DataCollectionGridInfo, @@ -16,7 +16,11 @@ from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER -class I021GridscanISPyBCallback(GridscanISPyBCallback): +def _make_comment(x_steps: int, y_steps: int) -> str: + return f"Diffraction grid scan of {x_steps} by {y_steps}." + + +class GridscanISPyBCallback(CommonGridscanISPyBCallback): def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: """ For VMXm, grid information is available immediately after the plan is triggered. @@ -63,9 +67,8 @@ def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: data_collection_id = self.ispyb_ids.data_collection_ids[0] - self.data_collection_group_info.comments = ( - f"Diffraction grid scan of {data_collection_grid_info.steps_x} by " - f"{self.params.y_steps}." + self.data_collection_group_info.comments = _make_comment( + self.params.x_steps, self.params.y_steps[0] ) self._populate_axis_info(data_collection_info, doc["data"]) diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index 282de76d6e..a348dbbb4e 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -24,6 +24,9 @@ setup_zebra_for_gridscan, tidy_up_zebra_after_gridscan, ) +from mx_bluesky.beamlines.i02_1.external_interaction.callbacks.gridscan.ispyb_callback import ( + GridscanISPyBCallback, +) from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import ( BeamlineSpecificFGSFeatures, @@ -37,7 +40,6 @@ GridscanNexusFileCallback, ) from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( - GridscanISPyBCallback, ispyb_activation_decorator, ) from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( @@ -163,8 +165,8 @@ def i02_1_gridscan_plan( beamline_specific = construct_i02_1_specific_features(composite, parameters) callbacks = create_gridscan_callbacks(parameters) - @ispyb_activation_decorator @bpp.subs_decorator(callbacks) + @ispyb_activation_decorator(parameters) def decorated_flyscan_plan(): yield from common_flyscan_xray_centre(composite, parameters, beamline_specific) diff --git a/src/mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py b/src/mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py index d7e973a926..452cd61f45 100644 --- a/src/mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py +++ b/src/mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Sequence from time import time import bluesky.plan_stubs as bps @@ -65,7 +65,7 @@ def kickoff_and_complete_gridscan( detector: EigerDetector, # Once Eiger inherits from StandardDetector, use that type instead synchrotron: Synchrotron, scan_points: list[AxesPoints[Axis]], - omega_starts_deg: list[float], + omega_starts_deg: Sequence[float], plan_during_collection: Callable[[], MsgGenerator] | None = None, ): """Triggers a grid scan motion program and waits for completion, accounting for synchrotron topup. diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py index ab0b523cbe..82ed95e03b 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/gridscan/ispyb_callback.py @@ -95,7 +95,7 @@ def __init__( self.data_collection_group_info: DataCollectionGroupInfo | None def activity_gated_start(self, doc: RunStart): - if doc.get("subplan_name") == PlanNameConstants.DO_FGS: + if doc.get("subplan_name") == PlanNameConstants.TRIGGER_GRIDSCAN_ISPYB_CALLBACK: self._start_of_fgs_uid = doc.get("uid") ISPYB_ZOCALO_CALLBACK_LOGGER.info( "ISPyB callback received start document with experiment parameters and " @@ -186,10 +186,6 @@ def tag_doc(self, doc: D) -> D: doc["_grid_num_to_id_map"] = self._grid_num_to_id_map return doc # type: ignore - def data_collection_number_from_gridplane(self, plane) -> int: - assert self.params - return self.params.detector_params.run_number - def fill_gridscan_deposition_and_store( self, make_scan_infos_with_grid_info: Callable[..., Sequence[ScanDataInfo]] ): diff --git a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py index 2729ef9f90..e115d27299 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/grid/utils.py @@ -55,9 +55,10 @@ def generate_start_info_from_num_grids( """ doc = yield [] - start_frame = 0 + infos = [] for grid_num in range(params.num_grids): + start_frame = params.scan_indices[grid_num] frames = len(params.scan_points[grid_num]) infos.append( ZocaloStartInfo( diff --git a/src/mx_bluesky/common/parameters/components.py b/src/mx_bluesky/common/parameters/components.py index 1c54bf9a00..8fb57e2b03 100644 --- a/src/mx_bluesky/common/parameters/components.py +++ b/src/mx_bluesky/common/parameters/components.py @@ -271,6 +271,8 @@ def _start_for_axis(self, axis: XyzAxis, grid: int) -> float: class OptionalGonioAngleStarts(BaseModel): # Gridscans have different omega starts + # See https://github.com/DiamondLightSource/mx-bluesky/issues/1631 for why + # we use int omega_starts_deg: list[int] = [0, 90] phi_start_deg: float | None = None diff --git a/tests/unit_tests/beamlines/i02_1/conftest.py b/tests/unit_tests/beamlines/i02_1/conftest.py new file mode 100644 index 0000000000..56e1de725b --- /dev/null +++ b/tests/unit_tests/beamlines/i02_1/conftest.py @@ -0,0 +1,30 @@ +import pytest + +from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams +from mx_bluesky.common.parameters.components import get_param_version + + +@pytest.fixture +def fgs_params_two_d(tmp_path) -> I02_1FgsParams: + return I02_1FgsParams( + x_start_um=0, + y_starts_um=[0], + z_starts_um=[0], + y_step_sizes_um=[10], + omega_starts_deg=[0], + parameter_model_version=get_param_version(), + sample_id=0, + visit="cm0000-0", + file_name="test_file", + storage_directory=str(tmp_path), + x_steps=5, + y_steps=[3], + path_to_xtal_snapshot=tmp_path, + beam_size_x=0, + beam_size_y=0, + microns_per_pixel_x=1, + microns_per_pixel_y=1, + upper_left_x=0, + upper_left_y=0, + detector_distance_mm=100, + ) diff --git a/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py b/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py index 5ed1f1320a..c6280d8915 100644 --- a/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py +++ b/tests/unit_tests/beamlines/i02_1/test_gridscan_ispyb_callback.py @@ -1 +1,76 @@ -def test_get_scan_infos_gives_expected_output(): ... +from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams +from mx_bluesky.beamlines.i02_1.external_interaction.callbacks.gridscan.ispyb_callback import ( + GridscanISPyBCallback, + _make_comment, +) +from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import ( + construct_comment_for_gridscan, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionGridInfo, + DataCollectionGroupInfo, + DataCollectionInfo, + Orientation, + ScanDataInfo, +) +from mx_bluesky.common.external_interaction.ispyb.ispyb_store import IspybIds + + +def _get_expected_scan_info(params: I02_1FgsParams, dcid: int): + dc_grid_info = DataCollectionGridInfo( + params.x_step_size_um * 1000, + params.y_step_sizes_um[0] * 1000, + params.x_steps, + params.y_steps[0], + params.microns_per_pixel_x, + params.microns_per_pixel_y, + params.upper_left_x, + params.upper_left_y, + Orientation.HORIZONTAL, + True, + ) + xtal = str(params.path_to_xtal_snapshot) + dc_info = DataCollectionInfo( + omega_start=0, + data_collection_number=1, + xtal_snapshot1=xtal, + xtal_snapshot2=xtal, + xtal_snapshot3=xtal, + n_images=params.x_steps * params.y_steps[0], + axis_end=0, + axis_range=0, + axis_start=0, + file_template=f"{params.file_name}_1_master.h5", + comments=construct_comment_for_gridscan(dc_grid_info), + ) + return [ + ScanDataInfo( + data_collection_info=dc_info, + data_collection_id=dcid, + data_collection_grid_info=dc_grid_info, + ) + ] + + +def test_get_scan_infos_gives_expected_output( + fgs_params_two_d: I02_1FgsParams, +): + callback = GridscanISPyBCallback(param_type=I02_1FgsParams) + callback.params = fgs_params_two_d + doc = {} + doc["data"] = { + "gonio-omega": 0, + } + callback.ispyb_ids = IspybIds() + callback.ispyb_ids.data_collection_group_id = 0 + callback.ispyb_ids.data_collection_ids = ((0),) + callback.data_collection_group_info = DataCollectionGroupInfo( + "0", + "SAD", + None, + comments=_make_comment(fgs_params_two_d.x_steps, fgs_params_two_d.y_steps[0]), + ) + scan_info = callback._get_scan_infos(doc) + assert scan_info[0].data_collection_grid_info + + assert scan_info == _get_expected_scan_info(fgs_params_two_d, 0) diff --git a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py index 68a8de4fee..2ae688a210 100644 --- a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py +++ b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py @@ -21,37 +21,20 @@ i02_1_gridscan_plan, ) from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan +from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import ( + populate_data_collection_group, + populate_remaining_data_collection_info, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionInfo, + ScanDataInfo, +) from mx_bluesky.common.parameters.components import get_param_version from mx_bluesky.common.parameters.device_composites import ( GonioWithOmega, ) -@pytest.fixture -def fgs_params_two_d(tmp_path) -> I02_1FgsParams: - return I02_1FgsParams( - x_start_um=0, - y_starts_um=[0], - z_starts_um=[0], - y_step_sizes_um=[10], - omega_starts_deg=[0], - parameter_model_version=get_param_version(), - sample_id=0, - visit="visit", - file_name="test_file", - storage_directory=str(tmp_path), - x_steps=5, - y_steps=[3], - path_to_xtal_snapshot=tmp_path, - beam_size_x=0, - beam_size_y=0, - microns_per_pixel_x=1, - microns_per_pixel_y=1, - upper_left_x=0, - upper_left_y=0, - ) - - @pytest.fixture def zebra_fgs_two_d() -> ZebraFastGridScanTwoD: device = i02_1.zebra_fast_grid_scan.build(connect_immediately=True, mock=True) @@ -151,12 +134,55 @@ def test_i02_1_flyscan_xray_centre_in_re( ) mock_create_features.return_value = expected_features - run_engine( - i02_1_gridscan_plan(fgs_params_two_d, fgs_composite) - ) # todo fix this typing + run_engine(i02_1_gridscan_plan(fgs_params_two_d, fgs_composite)) mock_common_scan.assert_called_once_with( fgs_composite, fgs_params_two_d, expected_features ) -def test_ispyb_callbacks_activated_when_expected(): ... +@patch( + "mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan.construct_i02_1_specific_features", +) +@patch( + "mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan.common_flyscan_xray_centre", + new=MagicMock(), +) +@patch( + "mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback.StoreInIspyb" +) +def test_ispyb_activated_correct_params( + mock_store_ispyb: MagicMock, + mock_create_features: MagicMock, + run_engine: RunEngine, + fgs_params_two_d: I02_1FgsParams, + fgs_composite: FlyScanXRayCentreComposite, +): + mock_ispyb = MagicMock() + + mock_store_ispyb.return_value = mock_ispyb + expected_features = construct_i02_1_specific_features( + fgs_composite, fgs_params_two_d + ) + run_engine.md["data"] = {} + + mock_create_features.return_value = expected_features + # uid = "test uid" + + # def do_run_start_and_run_stop(): + # yield from RunStart(uid=uid, time=0) + # yield from RunStop(uid=uid, ) + + run_engine(i02_1_gridscan_plan(fgs_params_two_d, fgs_composite)) + initial_group_info = populate_data_collection_group(fgs_params_two_d) + initial_group_info.comments = f"Diffraction grid scan of {fgs_params_two_d.x_steps} by {fgs_params_two_d.y_steps[0]}.Zocalo processing took 0.00 s." + initial_scan_info = ScanDataInfo( + data_collection_info=populate_remaining_data_collection_info( + "MX-Bluesky: Xray centring 1/1 -", + None, + DataCollectionInfo(), + fgs_params_two_d, + ) + ) + mock_ispyb.begin_deposition.assert_called_once_with( + initial_group_info, [initial_scan_info] + ) diff --git a/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py b/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py index a6c66f06b3..b58e01755e 100644 --- a/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py +++ b/tests/unit_tests/common/external_interaction/callbacks/common/test_ispyb_callback_base.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call, patch import bluesky.preprocessors as bpp import pytest @@ -60,6 +60,49 @@ def test_plan(): run_engine(test_plan()) -def test_handle_ispyb_transmission_flux_read_if_no_beamsize(): - # Test exception and check if params exist then it still works - ... +def _get_working_doc(): + return { + "data": { + "flux-flux_reading": 0, + "eiger-ispyb_detector_id": 0, + "eiger_cam_roi_mode": None, + "attenuator-actual_transmission": None, + } + } + + +@patch( + "mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base.ISPYB_ZOCALO_CALLBACK_LOGGER" +) +def test_handle_ispyb_transmission_flux_read_if_no_beamsize_warning( + mock_logger: MagicMock, + test_three_d_grid_params: SpecifiedThreeDGridScan, +): + callback = BaseISPyBCallback() + callback.params = test_three_d_grid_params + doc = _get_working_doc() + callback._handle_ispyb_transmission_flux_read(doc) # type: ignore + mock_logger.warning.assert_has_calls( + [call("ISPyB callbacks couldn't get beamsize")] + ) + + +@patch( + "mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base.ISPYB_ZOCALO_CALLBACK_LOGGER" +) +def test_handle_ispyb_transmission_flux_read_if_params_specify_beamsize( + mock_logger: MagicMock, + test_three_d_grid_params: SpecifiedThreeDGridScan, +): + test_three_d_grid_params.beam_size_x = 0 # type: ignore + test_three_d_grid_params.beam_size_y = 1 # type: ignore + + callback = BaseISPyBCallback() + callback.params = test_three_d_grid_params + doc = _get_working_doc() + callback._handle_ispyb_transmission_flux_read(doc) # type: ignore + + assert ( + call("ISPyB callbacks couldn't get beamsize") + not in mock_logger.warning.call_args_list + ) diff --git a/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py b/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py index c567398461..841d69867b 100644 --- a/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py +++ b/tests/unit_tests/common/external_interaction/callbacks/grid/test_ispyb_gridscan_callback.py @@ -1 +1,112 @@ -# todo use tests/unit_tests/common/external_interaction/xray_centre/test_ispyb_callback.py as a template to test this file +from collections.abc import Sequence +from unittest.mock import MagicMock, patch + +import pytest +from event_model import RunStop + +from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import ( + GridscanISPyBCallback, +) +from mx_bluesky.common.external_interaction.ispyb.data_model import ( + DataCollectionGroupInfo, + DataCollectionInfo, + ScanDataInfo, +) +from mx_bluesky.common.external_interaction.ispyb.ispyb_store import ( + IspybIds, + StoreInIspyb, +) +from mx_bluesky.common.parameters.constants import PlanNameConstants +from mx_bluesky.common.parameters.gridscan import ( + SpecifiedThreeDGridScan, +) +from mx_bluesky.common.utils.exceptions import ISPyBDepositionNotMadeError + + +class Callback(GridscanISPyBCallback): + def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]: + return [ScanDataInfo(data_collection_info=DataCollectionInfo())] + + +@patch( + "mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback.BaseISPyBCallback.activity_gated_start" +) +def test_gridscan_callback_start_calls_correct_funcs( + mock_start: MagicMock, test_three_d_grid_params: SpecifiedThreeDGridScan +): + cb = Callback(SpecifiedThreeDGridScan) + cb.fill_gridscan_deposition_and_store = MagicMock() + doc = { + "subplan_name": PlanNameConstants.DO_FGS, + "mx_bluesky_parameters": test_three_d_grid_params.model_dump_json(), + } + cb.activity_gated_start(doc) # type: ignore + cb.fill_gridscan_deposition_and_store.assert_called_once() + mock_start.assert_called_once() + + +def test_populate_info_for_update(test_three_d_grid_params: SpecifiedThreeDGridScan): + cb = Callback(SpecifiedThreeDGridScan) + cb.params = test_three_d_grid_params + cb.ispyb_ids = IspybIds(data_collection_ids=(0, 1)) + cb._grid_num_to_id_map = {0: 0, 1: 1} + es_dcid = DataCollectionInfo() + infos = cb.populate_info_for_update(es_dcid, None, test_three_d_grid_params) + assert infos == [ + ScanDataInfo(data_collection_id=0, data_collection_info=DataCollectionInfo()), + ScanDataInfo(data_collection_id=1, data_collection_info=DataCollectionInfo()), + ] + + +def test_stop_errors_if_empty_ispyb_id(): + cb = Callback(SpecifiedThreeDGridScan) + cb.ispyb_ids = IspybIds() + cb.data_collection_group_info = DataCollectionGroupInfo("", "", None) + doc: RunStop = { + "time": 0, + "uid": "0", + "exit_status": "success", + "run_start": None, # type: ignore + } + with pytest.raises(ISPyBDepositionNotMadeError): + cb.activity_gated_stop(doc) + + +def test_exception_added_onto_comments(): + cb = Callback(SpecifiedThreeDGridScan) + cb.ispyb = StoreInIspyb("") + cb.ispyb.update_data_collection_group_table = MagicMock() + cb.ispyb_ids = IspybIds(data_collection_ids=(0,)) + cb.data_collection_group_info = DataCollectionGroupInfo("", "", None) + reason = "test reason" + doc: RunStop = { + "time": 0, + "uid": "0", + "exit_status": "success", + "run_start": None, # type: ignore + "reason": f"[test]: {reason}", + } + cb.activity_gated_stop(doc) + cb.ispyb.update_data_collection_group_table.assert_called_once_with( + DataCollectionGroupInfo("", "", None, comments=reason), None + ) + + +@patch( + "mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback.StoreInIspyb" +) +def test_fill_gridscan_deposition_and_store( + mock_store: MagicMock, + test_three_d_grid_params: SpecifiedThreeDGridScan, +): + cb = Callback(SpecifiedThreeDGridScan) + cb.params = test_three_d_grid_params + ispyb = StoreInIspyb("") + ispyb.begin_deposition = MagicMock() + ispyb.update_deposition = MagicMock() + ispyb.update_data_collection_group_table = MagicMock() + mock_store.return_value = ispyb + cb.fill_gridscan_deposition_and_store(MagicMock()) + ispyb.begin_deposition.assert_called_once() + ispyb.update_deposition.assert_called_once() + ispyb.update_data_collection_group_table.assert_called_once() diff --git a/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py b/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py index 00e89fc4a0..9cb8718504 100644 --- a/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py +++ b/tests/unit_tests/common/external_interaction/callbacks/grid/test_utils.py @@ -1 +1,28 @@ -def test_generate_start_info_from_num_grids(): ... +from bluesky.run_engine import RunEngine +from dodal.devices.zocalo import ZocaloStartInfo + +from mx_bluesky.common.external_interaction.callbacks.grid.utils import ( + generate_start_info_from_num_grids, +) +from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan + + +def test_generate_start_info_from_num_grids( + test_three_d_grid_params: SpecifiedThreeDGridScan, run_engine: RunEngine +): + zocalo_info_gen = generate_start_info_from_num_grids(test_three_d_grid_params) + next(zocalo_info_gen) + infos = zocalo_info_gen.send({"_grid_num_to_id_map": {0: 0, 1: 1, 2: 2}}) + + expected_infos = [ + ZocaloStartInfo( + ispyb_dcid=num, + filename=None, + start_frame_index=test_three_d_grid_params.scan_indices[num], + number_of_frames=len(test_three_d_grid_params.scan_points[num]), + message_index=num, + ) + for num in range(test_three_d_grid_params.num_grids) + ] + + assert infos == expected_infos diff --git a/tests/unit_tests/common/parameters/test_gridscan.py b/tests/unit_tests/common/parameters/test_gridscan.py index 60e592e15d..8eb612f4a1 100644 --- a/tests/unit_tests/common/parameters/test_gridscan.py +++ b/tests/unit_tests/common/parameters/test_gridscan.py @@ -34,7 +34,7 @@ def fast_gridscan_params(): ... # type: ignore def test_specified_grids_validation_error( y_starts_um: list[float], z_starts_um: list[float], - omega_starts_deg: list[float], + omega_starts_deg: list[int], y_step_sizes_um: list[float], y_steps: list[int], should_raise: bool, @@ -82,7 +82,7 @@ def _check_lengths_are_same(self): # type: ignore def test_three_d_grid_scan_validation( y_starts_um: list[float], z_starts_um: list[float], - omega_starts_deg: list[float], + omega_starts_deg: list[int], y_step_sizes_um: list[float], y_steps: list[int], should_raise: bool, From 78dac2891095bf61491c4da7451068ac72b8660a Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Mon, 2 Mar 2026 15:25:04 +0000 Subject: [PATCH 8/8] Remove comment --- tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py index 2ae688a210..723cec4e91 100644 --- a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py +++ b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py @@ -166,11 +166,6 @@ def test_ispyb_activated_correct_params( run_engine.md["data"] = {} mock_create_features.return_value = expected_features - # uid = "test uid" - - # def do_run_start_and_run_stop(): - # yield from RunStart(uid=uid, time=0) - # yield from RunStop(uid=uid, ) run_engine(i02_1_gridscan_plan(fgs_params_two_d, fgs_composite)) initial_group_info = populate_data_collection_group(fgs_params_two_d)