Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Correctly pass parameters through all of Hyperion's UDC parameters #746

Merged
merged 12 commits into from
Jan 17, 2025
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from collections.abc import Callable, Sequence
from time import time
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, TypeVar

import numpy as np
from bluesky import preprocessors as bpp
Expand Down Expand Up @@ -64,6 +64,9 @@ def ispyb_activation_wrapper(plan_generator: MsgGenerator, parameters):
)


T = TypeVar("T", bound="GridCommon")


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
Expand All @@ -81,12 +84,14 @@ class GridscanISPyBCallback(BaseISPyBCallback):

def __init__(
self,
param_type: type[T],
*,
emit: Callable[..., Any] | None = None,
) -> None:
super().__init__(emit=emit)
self.ispyb: StoreInIspyb
self.ispyb_ids: IspybIds = IspybIds()
self.param_type = param_type
self._start_of_fgs_uid: str | None = None
self._processing_start_time: float | None = None

Expand All @@ -101,7 +106,7 @@ def activity_gated_start(self, doc: RunStart):
)
mx_bluesky_parameters = doc.get("mx_bluesky_parameters")
assert isinstance(mx_bluesky_parameters, str)
self.params = GridCommon.model_validate_json(mx_bluesky_parameters)
self.params = self.param_type.model_validate_json(mx_bluesky_parameters)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because GridCommon is now abstract, we need to use an implemented version of this model here. Since the type is bound to GridCommon, I don't think this will affect generalisability.

Similar change done in nexus_callback.py

self.ispyb = StoreInIspyb(self.ispyb_config)
data_collection_group_info = populate_data_collection_group(self.params)

Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, TypeVar

from mx_bluesky.common.external_interaction.callbacks.common.plan_reactive_callback import (
PlanReactiveCallback,
Expand All @@ -11,12 +11,16 @@
)
from mx_bluesky.common.external_interaction.nexus.write_nexus import NexusWriter
from mx_bluesky.common.parameters.constants import DocDescriptorNames, PlanNameConstants
from mx_bluesky.common.parameters.gridscan import ThreeDGridScan
from mx_bluesky.common.parameters.gridscan import (
SpecifiedThreeDGridScan,
)
from mx_bluesky.common.utils.log import NEXUS_LOGGER

if TYPE_CHECKING:
from event_model.documents import Event, EventDescriptor, RunStart

T = TypeVar("T", bound="SpecifiedThreeDGridScan")


class GridscanNexusFileCallback(PlanReactiveCallback):
"""Callback class to handle the creation of Nexus files based on experiment \
Expand All @@ -35,8 +39,9 @@ class GridscanNexusFileCallback(PlanReactiveCallback):
See: https://blueskyproject.io/bluesky/callbacks.html#ways-to-invoke-callbacks
"""

def __init__(self) -> None:
def __init__(self, param_type: type[T]) -> None:
super().__init__(NEXUS_LOGGER)
self.param_type = param_type
self.run_start_uid: str | None = None
self.nexus_writer_1: NexusWriter | None = None
self.nexus_writer_2: NexusWriter | None = None
Expand All @@ -50,7 +55,7 @@ def activity_gated_start(self, doc: RunStart):
NEXUS_LOGGER.info(
f"Nexus writer received start document with experiment parameters {mx_bluesky_parameters}"
)
parameters = ThreeDGridScan.model_validate_json(mx_bluesky_parameters)
parameters = self.param_type.model_validate_json(mx_bluesky_parameters)
d_size = parameters.detector_params.detector_size_constants.det_size_pixels
grid_n_img_1 = parameters.scan_indices[1]
grid_n_img_2 = parameters.num_images - grid_n_img_1
Expand Down
3 changes: 2 additions & 1 deletion src/mx_bluesky/common/parameters/components.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
GridscanParamConstants,
)

PARAMETER_VERSION = Version.parse("5.2.0")
PARAMETER_VERSION = Version.parse("5.3.0")


class RotationAxis(StrEnum):
Expand Down Expand Up @@ -151,6 +151,7 @@ class DiffractionExperiment(
transmission_frac: float = Field(default=0.1)
ispyb_experiment_type: IspybExperimentType
storage_directory: str
use_roi_mode: bool = Field(default=GridscanParamConstants.USE_ROI)

@model_validator(mode="before")
@classmethod
Expand Down
64 changes: 5 additions & 59 deletions src/mx_bluesky/common/parameters/gridscan.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,6 @@
from __future__ import annotations

import os

from dodal.devices.aperturescatterguard import ApertureValue
from dodal.devices.detector import (
DetectorParams,
)
from dodal.devices.fast_grid_scan import (
ZebraGridScanParams,
)
Expand All @@ -23,86 +18,37 @@
XyzStarts,
)
from mx_bluesky.common.parameters.constants import (
DetectorParamConstants,
GridscanParamConstants,
HardwareConstants,
)
from mx_bluesky.common.parameters.robot_load import RobotLoadAndEnergyChange

"""Parameter models in this file are abstract. They should be inherited by a top-level model"""


class GridCommon(
DiffractionExperimentWithSample,
OptionalGonioAngleStarts,
):
"""Parameters used in every MX diffraction experiment using grids. This model should be used by plans which have no knowledge of the grid specifications - i.e before automatic grid detection has completed"""

grid_width_um: float = Field(default=GridscanParamConstants.WIDTH_UM)
exposure_time_s: float = Field(default=GridscanParamConstants.EXPOSURE_TIME_S)
use_roi_mode: bool = Field(default=GridscanParamConstants.USE_ROI)

ispyb_experiment_type: IspybExperimentType = Field(
default=IspybExperimentType.GRIDSCAN_3D
)
selected_aperture: ApertureValue | None = Field(default=ApertureValue.SMALL)

@property
def detector_params(self):
self.det_dist_to_beam_converter_path = (
self.det_dist_to_beam_converter_path
or DetectorParamConstants.BEAM_XY_LUT_PATH
)
optional_args = {}
if self.run_number:
optional_args["run_number"] = self.run_number
assert self.detector_distance_mm is not None, (
"Detector distance must be filled before generating DetectorParams"
)
os.makedirs(self.storage_directory, exist_ok=True)
return DetectorParams(
detector_size_constants=DetectorParamConstants.DETECTOR,
expected_energy_ev=self.demand_energy_ev,
exposure_time=self.exposure_time_s,
directory=self.storage_directory,
prefix=self.file_name,
detector_distance=self.detector_distance_mm,
omega_start=self.omega_start_deg or 0,
omega_increment=0,
num_images_per_trigger=1,
num_triggers=self.num_images,
use_roi_mode=self.use_roi_mode,
det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
trigger_mode=self.trigger_mode,
**optional_args,
)


class RobotLoadThenCentre(GridCommon):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed with @rtuck99 - we're so far off any non-hyperion beamlines doing anything other than a grid scan, it's probably good to get rid of some of the common parameters like this one

thawing_time: float = Field(default=HardwareConstants.THAWING_TIME)
tip_offset_um: float = Field(default=HardwareConstants.TIP_OFFSET_UM)

def robot_load_params(self):
my_params = self.model_dump()
return RobotLoadAndEnergyChange(**my_params)

def pin_centre_then_xray_centre_params(self):
my_params = self.model_dump()
del my_params["thawing_time"]
return PinTipCentreThenXrayCentre(**my_params)


class GridScanWithEdgeDetect(GridCommon):
box_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)


class PinTipCentreThenXrayCentre(GridCommon):
tip_offset_um: float = 0


class SpecifiedGrid(XyzStarts, WithScan):
"""A specified grid is one which has defined values for the start position,
grid and box sizes, etc., as opposed to parameters for a plan which will create
those parameters at some point (e.g. through optical pin detection)."""


class ThreeDGridScan(
class SpecifiedThreeDGridScan(
GridCommon,
SpecifiedGrid,
SplitScan,
Expand Down
16 changes: 0 additions & 16 deletions src/mx_bluesky/common/parameters/robot_load.py

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,14 @@
from mx_bluesky.common.utils.tracing import TRACER
from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_x_y_z
from mx_bluesky.hyperion.experiment_plans.common.xrc_result import XRayCentreResult
from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
from mx_bluesky.hyperion.parameters.gridscan import HyperionSpecifiedThreeDGridScan


def change_aperture_then_move_to_xtal(
best_hit: XRayCentreResult,
smargon: Smargon,
aperture_scatterguard: ApertureScatterguard,
parameters: HyperionThreeDGridScan | None = None,
parameters: HyperionSpecifiedThreeDGridScan | None = None,
):
"""For the given x-ray centring result,
* Change the aperture so that the beam size is comparable to the crystal size
Expand Down
18 changes: 9 additions & 9 deletions src/mx_bluesky/hyperion/experiment_plans/experiment_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,6 @@

import mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan as flyscan_xray_centre_plan
import mx_bluesky.hyperion.experiment_plans.rotation_scan_plan as rotation_scan_plan
from mx_bluesky.common.parameters.gridscan import (
GridScanWithEdgeDetect,
PinTipCentreThenXrayCentre,
RobotLoadThenCentre,
)
from mx_bluesky.hyperion.experiment_plans import (
grid_detect_then_xray_centre_plan,
load_centre_collect_full_plan,
Expand All @@ -23,8 +18,13 @@
create_robot_load_and_centre_callbacks,
create_rotation_callbacks,
)
from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
from mx_bluesky.hyperion.parameters.gridscan import (
GridScanWithEdgeDetect,
HyperionSpecifiedThreeDGridScan,
PinTipCentreThenXrayCentre,
)
from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
from mx_bluesky.hyperion.parameters.robot_load import RobotLoadThenCentre
from mx_bluesky.hyperion.parameters.rotation import MultiRotationScan, RotationScan


Expand All @@ -39,21 +39,21 @@ def do_nothing():
class ExperimentRegistryEntry(TypedDict):
setup: Callable
param_type: type[
HyperionThreeDGridScan
HyperionSpecifiedThreeDGridScan
| GridScanWithEdgeDetect
| RotationScan
| MultiRotationScan
| PinTipCentreThenXrayCentre
| RobotLoadThenCentre
| LoadCentreCollect
| RobotLoadThenCentre
]
callbacks_factory: CallbacksFactory


PLAN_REGISTRY: dict[str, ExperimentRegistryEntry] = {
"flyscan_xray_centre": {
"setup": flyscan_xray_centre_plan.create_devices,
"param_type": HyperionThreeDGridScan,
"param_type": HyperionSpecifiedThreeDGridScan,
"callbacks_factory": create_gridscan_callbacks,
},
"grid_detect_then_xray_centre": {
Expand Down
Loading
Loading