Skip to content

Commit

Permalink
470 Full experiment plan to replace execute request (#538)
Browse files Browse the repository at this point in the history
* Implementation of grand unified experiment plan

* Add load_centre_collect to the experiment registry

* Additional system test for full experiment plan

* Make ruff happpy

* Changes following PR review comments

* Make pyright happy
rtuck99 authored Oct 7, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 7b9efa7 commit 22b1ea6
Showing 22 changed files with 1,388 additions and 558 deletions.
24 changes: 16 additions & 8 deletions docs/developer/hyperion/reference/param_hierarchy.puml
Original file line number Diff line number Diff line change
@@ -8,7 +8,9 @@ package Mixins {
class WithSample
class WithScan
class WithOavCentring
class WithOptionalEnergyChange
class WithSnapshot
class WithVisit
class OptionalXyzStarts
class XyzStarts
class OptionalGonioAngleStarts
@@ -22,49 +24,51 @@ package Experiments {
class DiffractionExperimentWithSample
class GridCommon
class GridScanWithEdgeDetect
class LoadCentreCollect
class PinTipCentreThenXrayCentre
class RotationScan
class MultiRotationScan
class RobotLoadAndEnergyChange
class RobotLoadThenCentre
class SpecifiedGridScan
class ThreeDGridScan
}

class HyperionParameters
note bottom: Base class for all experiment parameter models

class TemporaryIspybExtras
note bottom: To be removed

note top: Base class for all experiment parameter models

BaseModel <|-- HyperionParameters
BaseModel <|-- SplitScan
BaseModel <|-- OptionalGonioAngleStarts
BaseModel <|-- OptionalXyzStarts
BaseModel <|-- TemporaryIspybExtras
BaseModel <|-- WithOavCentring
BaseModel <|-- WithOptionalEnergyChange
BaseModel <|-- WithSnapshot
BaseModel <|-- WithSample
BaseModel <|-- WithScan
BaseModel <|-- WithVisit
BaseModel <|-- XyzStarts

RotationScan *-- TemporaryIspybExtras
MultiRotationScan *-- TemporaryIspybExtras
OptionalGonioAngleStarts <|-- RotationScanPerSweep
OptionalXyzStarts <|-- RotationScanPerSweep
DiffractionExperimentWithSample <|-- RotationExperiment
HyperionParameters <|-- DiffractionExperiment
WithSnapshot <|-- DiffractionExperiment
WithOptionalEnergyChange <|-- DiffractionExperiment
WithVisit <|-- DiffractionExperiment
DiffractionExperiment <|-- DiffractionExperimentWithSample
WithSample <|-- DiffractionExperimentWithSample
DiffractionExperimentWithSample <|-- GridCommon
GridCommon <|-- GridScanWithEdgeDetect
GridCommon <|-- PinTipCentreThenXrayCentre
GridCommon <|-- RobotLoadThenCentre
RobotLoadThenCentre *-- RobotLoadAndEnergyChange
RobotLoadThenCentre *-- PinTipCentreThenXrayCentre
GridCommon <|-- SpecifiedGridScan
WithScan <|-- SpecifiedGridScan
SpecifiedGridScan <|-- ThreeDGridScan
SplitScan <|-- ThreeDGridScan
WithOptionalEnergyChange <|-- ThreeDGridScan
WithOavCentring <|-- GridCommon
WithScan <|-- RotationScan
RotationScanPerSweep <|-- RotationScan
@@ -75,4 +79,8 @@ SplitScan <|-- MultiRotationScan
XyzStarts <|-- SpecifiedGridScan
OptionalGonioAngleStarts <|-- GridCommon
OptionalGonioAngleStarts <|-- RotationScan
HyperionParameters <|-- RobotLoadAndEnergyChange
WithSample <|-- RobotLoadAndEnergyChange
WithSnapshot <|-- RobotLoadAndEnergyChange
WithOptionalEnergyChange <|-- RobotLoadAndEnergyChange
@enduml
Original file line number Diff line number Diff line change
@@ -7,12 +7,14 @@
import mx_bluesky.hyperion.experiment_plans.rotation_scan_plan as rotation_scan_plan
from mx_bluesky.hyperion.experiment_plans import (
grid_detect_then_xray_centre_plan,
load_centre_collect_full_plan,
pin_centre_then_xray_centre_plan,
robot_load_then_centre_plan,
)
from mx_bluesky.hyperion.external_interaction.callbacks.common.callback_util import (
CallbacksFactory,
create_gridscan_callbacks,
create_load_centre_collect_callbacks,
create_robot_load_and_centre_callbacks,
create_rotation_callbacks,
)
@@ -22,6 +24,7 @@
RobotLoadThenCentre,
ThreeDGridScan,
)
from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
from mx_bluesky.hyperion.parameters.rotation import MultiRotationScan, RotationScan


@@ -42,6 +45,7 @@ class ExperimentRegistryEntry(TypedDict):
| MultiRotationScan
| PinTipCentreThenXrayCentre
| RobotLoadThenCentre
| LoadCentreCollect
]
callbacks_factory: CallbacksFactory

@@ -77,6 +81,11 @@ class ExperimentRegistryEntry(TypedDict):
"param_type": MultiRotationScan,
"callbacks_factory": create_rotation_callbacks,
},
"load_centre_collect_full_plan": {
"setup": load_centre_collect_full_plan.create_devices,
"param_type": LoadCentreCollect,
"callbacks_factory": create_load_centre_collect_callbacks,
},
}


Original file line number Diff line number Diff line change
@@ -77,6 +77,12 @@ class SmargonSpeedException(Exception):
pass


class CrystalNotFoundException(WarningException):
"""Raised if grid detection completed normally but no crystal was found."""

pass


@dataclasses.dataclass
class FlyScanXRayCentreComposite:
"""All devices which are directly or indirectly required by this plan"""
@@ -190,47 +196,50 @@ def run_gridscan_and_move(

LOGGER.info("Grid scan finished, getting results.")

with TRACER.start_span("wait_for_zocalo"):
yield from bps.trigger_and_read(
[fgs_composite.zocalo], name=ZOCALO_READING_PLAN_NAME
)
LOGGER.info("Zocalo triggered and read, interpreting results.")
xray_centre, bbox_size = yield from get_processing_result(fgs_composite.zocalo)
LOGGER.info(f"Got xray centre: {xray_centre}, bbox size: {bbox_size}")
if xray_centre is not None:
xray_centre = parameters.FGS_params.grid_position_to_motor_position(
xray_centre
try:
with TRACER.start_span("wait_for_zocalo"):
yield from bps.trigger_and_read(
[fgs_composite.zocalo], name=ZOCALO_READING_PLAN_NAME
)
else:
xray_centre = initial_xyz
LOGGER.warning("No X-ray centre received")
if bbox_size is not None:
with TRACER.start_span("change_aperture"):
yield from set_aperture_for_bbox_size(
fgs_composite.aperture_scatterguard, bbox_size
LOGGER.info("Zocalo triggered and read, interpreting results.")
xray_centre, bbox_size = yield from get_processing_result(
fgs_composite.zocalo
)
LOGGER.info(f"Got xray centre: {xray_centre}, bbox size: {bbox_size}")
if xray_centre is not None:
xray_centre = parameters.FGS_params.grid_position_to_motor_position(
xray_centre
)
else:
LOGGER.warning("No bounding box size received")

# once we have the results, go to the appropriate position
LOGGER.info("Moving to centre of mass.")
with TRACER.start_span("move_to_result"):
x, y, z = xray_centre
yield from move_x_y_z(fgs_composite.sample_motors, x, y, z, wait=True)

if parameters.FGS_params.set_stub_offsets:
LOGGER.info("Recentring smargon co-ordinate system to this point.")
yield from bps.mv(
fgs_composite.sample_motors.stub_offsets, StubPosition.CURRENT_AS_CENTER
)

# Turn off dev/shm streaming to avoid filling disk, see https://github.com/DiamondLightSource/hyperion/issues/1395
LOGGER.info("Turning off Eiger dev/shm streaming")
yield from bps.abs_set(fgs_composite.eiger.odin.fan.dev_shm_enable, 0)
else:
LOGGER.warning("No X-ray centre received")
raise CrystalNotFoundException()
if bbox_size is not None:
with TRACER.start_span("change_aperture"):
yield from set_aperture_for_bbox_size(
fgs_composite.aperture_scatterguard, bbox_size
)
else:
LOGGER.warning("No bounding box size received")

# once we have the results, go to the appropriate position
LOGGER.info("Moving to centre of mass.")
with TRACER.start_span("move_to_result"):
x, y, z = xray_centre
yield from move_x_y_z(fgs_composite.sample_motors, x, y, z, wait=True)

if parameters.FGS_params.set_stub_offsets:
LOGGER.info("Recentring smargon co-ordinate system to this point.")
yield from bps.mv(
fgs_composite.sample_motors.stub_offsets, StubPosition.CURRENT_AS_CENTER
)
finally:
# Turn off dev/shm streaming to avoid filling disk, see https://github.com/DiamondLightSource/hyperion/issues/1395
LOGGER.info("Turning off Eiger dev/shm streaming")
yield from bps.abs_set(fgs_composite.eiger.odin.fan.dev_shm_enable, 0)

# Wait on everything before returning to GDA (particularly apertures), can be removed
# when we do not return to GDA here
yield from bps.wait()
# Wait on everything before returning to GDA (particularly apertures), can be removed
# when we do not return to GDA here
yield from bps.wait()


@bpp.set_run_key_decorator(CONST.PLAN.GRIDSCAN_MAIN)
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import dataclasses

from blueapi.core import BlueskyContext
from dodal.devices.oav.oav_parameters import OAVParameters

from mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan import (
RobotLoadThenCentreComposite,
robot_load_then_centre,
)
from mx_bluesky.hyperion.experiment_plans.rotation_scan_plan import (
RotationScanComposite,
multi_rotation_scan,
)
from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
from mx_bluesky.hyperion.utils.context import device_composite_from_context


@dataclasses.dataclass
class LoadCentreCollectComposite(RobotLoadThenCentreComposite, RotationScanComposite):
"""Composite that provides access to the required devices."""

pass


def create_devices(context: BlueskyContext) -> LoadCentreCollectComposite:
"""Create the necessary devices for the plan."""
return device_composite_from_context(context, LoadCentreCollectComposite)


def load_centre_collect_full_plan(
composite: LoadCentreCollectComposite,
params: LoadCentreCollect,
oav_params: OAVParameters | None = None,
):
"""Attempt a complete data collection experiment, consisting of the following:
* Load the sample if necessary
* Move to the specified goniometer start angles
* Perform optical centring, then X-ray centring
* If X-ray centring finds a diffracting centre then move to that centre and
* do a collection with the specified parameters.
"""
if not oav_params:
oav_params = OAVParameters(context="xrayCentring")
yield from robot_load_then_centre(composite, params.robot_load_then_centre)

yield from multi_rotation_scan(composite, params.multi_rotation_scan, oav_params)
Original file line number Diff line number Diff line change
@@ -44,3 +44,21 @@ def create_rotation_callbacks() -> (
tuple[RotationNexusFileCallback, RotationISPyBCallback]
):
return (RotationNexusFileCallback(), RotationISPyBCallback(emit=ZocaloCallback()))


def create_load_centre_collect_callbacks() -> (
tuple[
GridscanNexusFileCallback,
GridscanISPyBCallback,
RobotLoadISPyBCallback,
RotationNexusFileCallback,
RotationISPyBCallback,
]
):
return (
GridscanNexusFileCallback(),
GridscanISPyBCallback(emit=ZocaloCallback()),
RobotLoadISPyBCallback(),
RotationNexusFileCallback(),
RotationISPyBCallback(emit=ZocaloCallback()),
)
Original file line number Diff line number Diff line change
@@ -110,9 +110,9 @@ def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]:
slitgap_vertical=doc["data"]["s4_slit_gaps_ygap"],
)
hwscan_position_info = DataCollectionPositionInfo(
pos_x=doc["data"]["smargon-x"],
pos_y=doc["data"]["smargon-y"],
pos_z=doc["data"]["smargon-z"],
pos_x=float(doc["data"]["smargon-x"]),
pos_y=float(doc["data"]["smargon-y"]),
pos_z=float(doc["data"]["smargon-z"]),
)
scan_data_infos = self.populate_info_for_update(
hwscan_data_collection_info, hwscan_position_info, self.params
18 changes: 9 additions & 9 deletions src/mx_bluesky/hyperion/parameters/components.py
Original file line number Diff line number Diff line change
@@ -147,6 +147,15 @@ class WithOptionalEnergyChange(BaseModel):

class WithVisit(BaseModel):
visit: str = Field(min_length=1)
zocalo_environment: str = Field(default=CONST.ZOCALO_ENV)
beamline: str = Field(default=CONST.I03.BEAMLINE, pattern=r"BL\d{2}[BIJS]")
det_dist_to_beam_converter_path: str = Field(
default=CONST.PARAM.DETECTOR.BEAM_XY_LUT_PATH
)
insertion_prefix: str = Field(
default=CONST.I03.INSERTION_PREFIX, pattern=r"SR\d{2}[BIJS]"
)
detector_distance_mm: float | None = Field(default=None, gt=0)


class DiffractionExperiment(
@@ -157,16 +166,7 @@ class DiffractionExperiment(
file_name: str
exposure_time_s: float = Field(gt=0)
comment: str = Field(default="")
beamline: str = Field(default=CONST.I03.BEAMLINE, pattern=r"BL\d{2}[BIJS]")
insertion_prefix: str = Field(
default=CONST.I03.INSERTION_PREFIX, pattern=r"SR\d{2}[BIJS]"
)
det_dist_to_beam_converter_path: str = Field(
default=CONST.PARAM.DETECTOR.BEAM_XY_LUT_PATH
)
zocalo_environment: str = Field(default=CONST.ZOCALO_ENV)
trigger_mode: TriggerMode = Field(default=TriggerMode.FREE_RUN)
detector_distance_mm: float | None = Field(default=None, gt=0)
run_number: int | None = Field(default=None, ge=0)
selected_aperture: ApertureValue | None = Field(default=None)
transmission_frac: float = Field(default=0.1)
50 changes: 50 additions & 0 deletions src/mx_bluesky/hyperion/parameters/load_centre_collect.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
from typing import TypeVar

from pydantic import BaseModel, model_validator

from mx_bluesky.hyperion.parameters.components import (
HyperionParameters,
WithSample,
WithVisit,
)
from mx_bluesky.hyperion.parameters.gridscan import (
RobotLoadThenCentre,
)
from mx_bluesky.hyperion.parameters.rotation import MultiRotationScan

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


def construct_from_values(parent_context: dict, key: str, t: type[T]) -> T:
values = dict(parent_context)
values |= values[key]
return t(**values)


class LoadCentreCollect(HyperionParameters, WithVisit, WithSample):
"""Experiment parameters to perform the combined robot load,
pin-tip centre and rotation scan operations."""

robot_load_then_centre: RobotLoadThenCentre
multi_rotation_scan: MultiRotationScan

@model_validator(mode="before")
@classmethod
def validate_model(cls, values):
allowed_keys = (
LoadCentreCollect.model_fields.keys()
| RobotLoadThenCentre.model_fields.keys()
| MultiRotationScan.model_fields.keys()
)
disallowed_keys = values.keys() - allowed_keys
assert (
disallowed_keys == set()
), f"Unexpected fields found in LoadCentreCollect {disallowed_keys}"

values["robot_load_then_centre"] = construct_from_values(
values, "robot_load_then_centre", RobotLoadThenCentre
)
values["multi_rotation_scan"] = construct_from_values(
values, "multi_rotation_scan", MultiRotationScan
)
return values
Loading

0 comments on commit 22b1ea6

Please sign in to comment.