From 73b0137b166e38d88e05db00464b60472e5df1cf Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 08:13:43 +0100 Subject: [PATCH 01/42] make gridscan artemisparam and ispybparam classes --- .../ispyb/ispyb_dataclass.py | 65 +++++++++++++++++-- .../external_interaction/nexus/write_nexus.py | 6 +- src/artemis/parameters/internal_parameters.py | 36 ++-------- .../plan_specific/fgs_internal_params.py | 52 +++++++++++++-- .../grid_scan_with_edge_detect_params.py | 34 +++++++++- .../rotation_scan_internal_params.py | 50 +++++++++++++- .../tests/test_data/artemis_parameters.json | 65 +++++++++++++++++++ .../tests/test_internal_parameters.py | 34 ++++++---- .../parameters/tests/test_pydantic_classes.py | 35 ++++++++++ .../tests/test_schema_validation.py | 2 +- 10 files changed, 320 insertions(+), 59 deletions(-) create mode 100644 src/artemis/parameters/tests/test_data/artemis_parameters.json create mode 100644 src/artemis/parameters/tests/test_pydantic_classes.py diff --git a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py index 05a2a016a..8b53598bd 100644 --- a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py @@ -4,7 +4,7 @@ import numpy as np from pydantic import BaseModel, validator -ISPYB_PARAM_DEFAULTS = { +GRIDSCAN_ISPYB_PARAM_DEFAULTS = { "sample_id": None, "sample_barcode": None, "visit_path": "", @@ -33,6 +33,9 @@ class IspybParams(BaseModel): visit_path: str + + +class RotationIspybParams(IspybParams): microns_per_pixel_x: float microns_per_pixel_y: float upper_left: np.ndarray @@ -87,11 +90,61 @@ def _parse_position( xtal_snapshots_omega_start: Optional[List[str]] = None xtal_snapshots_omega_end: Optional[List[str]] = None - def __eq__(self, other) -> bool: - if not isinstance(other, IspybParams): - return NotImplemented - else: - return self.json() == other.json() + +class GridscanIspybParams(IspybParams): + microns_per_pixel_x: float + microns_per_pixel_y: float + upper_left: np.ndarray + position: np.ndarray + + class Config: + arbitrary_types_allowed = True + json_encoders = {np.ndarray: lambda a: a.tolist()} + + def dict(self, **kwargs): + as_dict = super().dict(**kwargs) + as_dict["upper_left"] = as_dict["upper_left"].tolist() + as_dict["position"] = as_dict["position"].tolist() + return as_dict + + @validator("upper_left", pre=True) + def _parse_upper_left( + cls, upper_left: list[int | float] | np.ndarray, values: Dict[str, Any] + ) -> np.ndarray: + assert len(upper_left) == 3 + if isinstance(upper_left, np.ndarray): + return upper_left + return np.array(upper_left) + + @validator("position", pre=True) + def _parse_position( + cls, position: list[int | float] | np.ndarray, values: Dict[str, Any] + ) -> np.ndarray: + assert len(position) == 3 + if isinstance(position, np.ndarray): + return position + return np.array(position) + + transmission: float + flux: float + wavelength: float + beam_size_x: float + beam_size_y: float + focal_spot_size_x: float + focal_spot_size_y: float + comment: str + resolution: float + + sample_id: Optional[int] = None + sample_barcode: Optional[str] = None + + # Optional from GDA as populated by Ophyd + undulator_gap: Optional[float] = None + synchrotron_mode: Optional[str] = None + slit_gap_size_x: Optional[float] = None + slit_gap_size_y: Optional[float] = None + xtal_snapshots_omega_start: Optional[List[str]] = None + xtal_snapshots_omega_end: Optional[List[str]] = None class Orientation(Enum): diff --git a/src/artemis/external_interaction/nexus/write_nexus.py b/src/artemis/external_interaction/nexus/write_nexus.py index 8e4a19974..590b9d44d 100644 --- a/src/artemis/external_interaction/nexus/write_nexus.py +++ b/src/artemis/external_interaction/nexus/write_nexus.py @@ -30,7 +30,7 @@ from scanspec.core import Path as ScanPath from scanspec.specs import Line -from artemis.external_interaction.ispyb.ispyb_dataclass import IspybParams +from artemis.external_interaction.ispyb.ispyb_dataclass import GridscanIspybParams from artemis.parameters.internal_parameters import InternalParameters @@ -171,12 +171,12 @@ def create_detector_parameters(detector_params: DetectorParams) -> Detector: def create_beam_and_attenuator_parameters( - ispyb_params: IspybParams, + ispyb_params: GridscanIspybParams, ) -> Tuple[Beam, Attenuator]: """Create beam and attenuator dictionaries that nexgen can understand. Args: - ispyb_params (IspybParams): An IspybParams object holding all required data. + ispyb_params (GridscanIspybParams): An GridscanIspybParams object holding all required data. Returns: Tuple[Beam, Attenuator]: Descriptions of the beam and attenuator for nexgen. diff --git a/src/artemis/parameters/internal_parameters.py b/src/artemis/parameters/internal_parameters.py index 8c80c7c3e..b3bda1ff4 100644 --- a/src/artemis/parameters/internal_parameters.py +++ b/src/artemis/parameters/internal_parameters.py @@ -5,10 +5,6 @@ from pydantic import BaseModel, root_validator from semver import Version -from artemis.external_interaction.ispyb.ispyb_dataclass import ( - ISPYB_PARAM_DEFAULTS, - IspybParams, -) from artemis.parameters.constants import ( DEFAULT_EXPERIMENT_TYPE, DETECTOR_PARAM_DEFAULTS, @@ -41,14 +37,6 @@ class ArtemisParameters(BaseModel): insertion_prefix: str = SIM_INSERTION_PREFIX experiment_type: str = DEFAULT_EXPERIMENT_TYPE detector_params: DetectorParams = DetectorParams(**DETECTOR_PARAM_DEFAULTS) - ispyb_params: IspybParams = IspybParams(**ISPYB_PARAM_DEFAULTS) - - class Config: - arbitrary_types_allowed = True - json_encoders = { - **DetectorParams.Config.json_encoders, - **IspybParams.Config.json_encoders, - } def flatten_dict(d: dict, parent_items: dict = {}) -> dict: @@ -72,21 +60,6 @@ def flatten_dict(d: dict, parent_items: dict = {}) -> dict: return items -def artemis_param_key_definitions(): - artemis_param_field_keys = [ - "zocalo_environment", - "beamline", - "insertion_prefix", - "experiment_type", - ] - detector_field_keys = list(DetectorParams.__annotations__.keys()) - # not an annotation but specified as field encoder in DetectorParams: - detector_field_keys.append("detector") - ispyb_field_keys = list(IspybParams.__annotations__.keys()) - - return artemis_param_field_keys, detector_field_keys, ispyb_field_keys - - def fetch_subdict_from_bucket( list_of_keys: list[str], bucket: dict[str, Any] ) -> dict[str, Any]: @@ -116,6 +89,7 @@ def get_extracted_experiment_and_flat_artemis_params( def extract_artemis_params_from_flat_dict( external_params: dict[str, Any], + artemis_param_key_definitions: tuple[list[str], list[str], list[str]], ) -> dict[str, Any]: all_params_bucket = flatten_dict(external_params) @@ -123,7 +97,7 @@ def extract_artemis_params_from_flat_dict( artemis_param_field_keys, detector_field_keys, ispyb_field_keys, - ) = artemis_param_key_definitions() + ) = artemis_param_key_definitions artemis_params_args: dict[str, Any] = fetch_subdict_from_bucket( artemis_param_field_keys, all_params_bucket @@ -147,7 +121,6 @@ class Config: use_enum_values = True arbitrary_types_allowed = True json_encoders = { - **ArtemisParameters.Config.json_encoders, ParameterVersion: lambda pv: str(pv), } @@ -160,6 +133,11 @@ def _preprocess_all(cls, values): values["artemis_params"] = flatten_dict(values) return values + @staticmethod + @abstractmethod + def _artemis_param_key_definitions(): + ... + @abstractmethod def _preprocess_experiment_params( cls, diff --git a/src/artemis/parameters/plan_specific/fgs_internal_params.py b/src/artemis/parameters/plan_specific/fgs_internal_params.py index 5935bd2a1..9aed12f72 100644 --- a/src/artemis/parameters/plan_specific/fgs_internal_params.py +++ b/src/artemis/parameters/plan_specific/fgs_internal_params.py @@ -3,10 +3,22 @@ from typing import Any import numpy as np -from dodal.devices.detector import TriggerMode +from dodal.devices.detector import DetectorParams, TriggerMode from dodal.devices.fast_grid_scan import GridScanParams from pydantic import validator +from artemis.external_interaction.ispyb.ispyb_dataclass import ( + GRIDSCAN_ISPYB_PARAM_DEFAULTS, + GridscanIspybParams, + IspybParams, +) +from artemis.parameters.constants import ( + DEFAULT_EXPERIMENT_TYPE, + DETECTOR_PARAM_DEFAULTS, + SIM_BEAMLINE, + SIM_INSERTION_PREFIX, + SIM_ZOCALO_ENV, +) from artemis.parameters.internal_parameters import ( ArtemisParameters, InternalParameters, @@ -15,17 +27,47 @@ ) +class GridscanArtemisParameters(ArtemisParameters): + ispyb_params: GridscanIspybParams = GridscanIspybParams( + **GRIDSCAN_ISPYB_PARAM_DEFAULTS + ) + + class Config: + arbitrary_types_allowed = True + json_encoders = { + **DetectorParams.Config.json_encoders, + **GridscanIspybParams.Config.json_encoders, + } + + class FGSInternalParameters(InternalParameters): experiment_params: GridScanParams - artemis_params: ArtemisParameters + artemis_params: GridscanArtemisParameters class Config: arbitrary_types_allowed = True json_encoders = { **GridScanParams.Config.json_encoders, - **ArtemisParameters.Config.json_encoders, + **GridscanArtemisParameters.Config.json_encoders, } + @staticmethod + def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: + artemis_param_field_keys = [ + "zocalo_environment", + "beamline", + "insertion_prefix", + "experiment_type", + ] + detector_field_keys = list(DetectorParams.__annotations__.keys()) + # not an annotation but specified as field encoder in DetectorParams: + detector_field_keys.append("detector") + ispyb_field_keys = list(IspybParams.__annotations__.keys()) + list( + GridscanIspybParams.__annotations__.keys() + ) + + return artemis_param_field_keys, detector_field_keys, ispyb_field_keys + @validator("experiment_params", pre=True) def _preprocess_experiment_params( cls, @@ -49,5 +91,7 @@ def _preprocess_artemis_params( all_params["num_images_per_trigger"] = 1 all_params["trigger_mode"] = TriggerMode.FREE_RUN all_params["upper_left"] = np.array(all_params["upper_left"]) - artemis_param_dict = extract_artemis_params_from_flat_dict(all_params) + artemis_param_dict = extract_artemis_params_from_flat_dict( + all_params, cls._artemis_param_key_definitions() + ) return ArtemisParameters(**artemis_param_dict) diff --git a/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py b/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py index ab2bc7d7f..5dac01325 100644 --- a/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py +++ b/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py @@ -4,11 +4,15 @@ import numpy as np from dataclasses_json import DataClassJsonMixin -from dodal.devices.detector import TriggerMode +from dodal.devices.detector import DetectorParams, TriggerMode from dodal.parameters.experiment_parameter_base import AbstractExperimentParameterBase from pydantic import validator from pydantic.dataclasses import dataclass +from artemis.external_interaction.ispyb.ispyb_dataclass import ( + GridscanIspybParams, + IspybParams, +) from artemis.parameters.internal_parameters import ( ArtemisParameters, InternalParameters, @@ -17,6 +21,9 @@ flatten_dict, get_extracted_experiment_and_flat_artemis_params, ) +from artemis.parameters.plan_specific.fgs_internal_params import ( + GridscanArtemisParameters, +) @dataclass @@ -36,7 +43,7 @@ def get_num_images(self): class GridScanWithEdgeDetectInternalParameters(InternalParameters): experiment_params: GridScanWithEdgeDetectParams - artemis_params: ArtemisParameters + artemis_params: GridscanArtemisParameters def __init__(self, data): prepared_args = get_extracted_experiment_and_flat_artemis_params( @@ -44,6 +51,23 @@ def __init__(self, data): ) super().__init__(**prepared_args) + @staticmethod + def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: + artemis_param_field_keys = [ + "zocalo_environment", + "beamline", + "insertion_prefix", + "experiment_type", + ] + detector_field_keys = list(DetectorParams.__annotations__.keys()) + # not an annotation but specified as field encoder in DetectorParams: + detector_field_keys.append("detector") + ispyb_field_keys = list(IspybParams.__annotations__.keys()) + list( + GridscanIspybParams.__annotations__.keys() + ) + + return artemis_param_field_keys, detector_field_keys, ispyb_field_keys + @validator("experiment_params", pre=True) def _preprocess_experiment_params( cls, @@ -67,4 +91,8 @@ def _preprocess_artemis_params( all_params["num_images_per_trigger"] = 1 all_params["trigger_mode"] = TriggerMode.FREE_RUN all_params["upper_left"] = np.array([0, 0, 0]) - return ArtemisParameters(**extract_artemis_params_from_flat_dict(all_params)) + return ArtemisParameters( + **extract_artemis_params_from_flat_dict( + all_params, cls._artemis_param_key_definitions() + ) + ) diff --git a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py index 22f544cad..4757f8565 100644 --- a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py +++ b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py @@ -3,11 +3,25 @@ from typing import Any, Optional import numpy as np +from dodal.devices.detector import DetectorParams from dodal.devices.motors import XYZLimitBundle from dodal.devices.zebra import RotationDirection from dodal.parameters.experiment_parameter_base import AbstractExperimentParameterBase from pydantic import BaseModel, validator +from artemis.external_interaction.ispyb.ispyb_dataclass import ( + GRIDSCAN_ISPYB_PARAM_DEFAULTS, + GridscanIspybParams, + IspybParams, + RotationIspybParams, +) +from artemis.parameters.constants import ( + DEFAULT_EXPERIMENT_TYPE, + DETECTOR_PARAM_DEFAULTS, + SIM_BEAMLINE, + SIM_INSERTION_PREFIX, + SIM_ZOCALO_ENV, +) from artemis.parameters.internal_parameters import ( ArtemisParameters, InternalParameters, @@ -16,6 +30,17 @@ ) +class RotationArtemisParameters(ArtemisParameters): + ispyb_params: IspybParams = RotationIspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) + + class Config: + arbitrary_types_allowed = True + json_encoders = { + **DetectorParams.Config.json_encoders, + **GridscanIspybParams.Config.json_encoders, + } + + class RotationScanParams(BaseModel, AbstractExperimentParameterBase): """ Holder class for the parameters of a rotation data collection. @@ -61,7 +86,24 @@ def get_num_images(self): class RotationInternalParameters(InternalParameters): experiment_params: RotationScanParams - artemis_params: ArtemisParameters + artemis_params: RotationArtemisParameters + + @staticmethod + def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: + artemis_param_field_keys = [ + "zocalo_environment", + "beamline", + "insertion_prefix", + "experiment_type", + ] + detector_field_keys = list(DetectorParams.__annotations__.keys()) + # not an annotation but specified as field encoder in DetectorParams: + detector_field_keys.append("detector") + ispyb_field_keys = list(IspybParams.__annotations__.keys()) + list( + GridscanIspybParams.__annotations__.keys() + ) + + return artemis_param_field_keys, detector_field_keys, ispyb_field_keys @validator("experiment_params", pre=True) def _preprocess_experiment_params( @@ -88,4 +130,8 @@ def _preprocess_artemis_params( all_params["num_triggers"] = 1 all_params["num_images_per_trigger"] = all_params["num_images"] all_params["upper_left"] = np.array(all_params["upper_left"]) - return ArtemisParameters(**extract_artemis_params_from_flat_dict(all_params)) + return ArtemisParameters( + **extract_artemis_params_from_flat_dict( + all_params, cls._artemis_param_key_definitions() + ) + ) diff --git a/src/artemis/parameters/tests/test_data/artemis_parameters.json b/src/artemis/parameters/tests/test_data/artemis_parameters.json new file mode 100644 index 000000000..afa635a0e --- /dev/null +++ b/src/artemis/parameters/tests/test_data/artemis_parameters.json @@ -0,0 +1,65 @@ +{ + "zocalo_environment": "devrmq", + "beamline": "BL03S", + "insertion_prefix": "SR03S", + "experiment_type": "fast_grid_scan", + "detector_params": { + "current_energy": 100.0, + "exposure_time": 0.1, + "directory": "/tmp/", + "prefix": "file_name", + "run_number": 0, + "detector_distance": 100.0, + "omega_start": 0.0, + "omega_increment": 0.0, + "num_images_per_trigger": 1, + "num_triggers": 60, + "use_roi_mode": false, + "det_dist_to_beam_converter_path": "src/artemis/unit_tests/test_lookup_table.txt", + "trigger_mode": 2, + "detector_size_constants": "EIGER2_X_16M", + "beam_xy_converter": null, + "start_index": 0, + "nexus_file_run_number": 0 + }, + "ispyb_params": { + "visit_path": "/tmp/cm31105-4/", + "microns_per_pixel_x": 1.0, + "microns_per_pixel_y": 1.0, + "upper_left": [ + 10.0, + 20.0, + 30.0 + ], + "position": [ + 10.0, + 20.0, + 30.0 + ], + "transmission": 1.0, + "flux": 10.0, + "wavelength": 0.01, + "beam_size_x": 1.0, + "beam_size_y": 1.0, + "focal_spot_size_x": 1.0, + "focal_spot_size_y": 1.0, + "comment": "test", + "resolution": 1.0, + "sample_id": null, + "sample_barcode": null, + "undulator_gap": null, + "synchrotron_mode": null, + "slit_gap_size_x": 1.0, + "slit_gap_size_y": 1.0, + "xtal_snapshots_omega_start": [ + "test_1_y", + "test_2_y", + "test_3_y" + ], + "xtal_snapshots_omega_end": [ + "test_1_z", + "test_2_z", + "test_3_z" + ] + } +} \ No newline at end of file diff --git a/src/artemis/parameters/tests/test_internal_parameters.py b/src/artemis/parameters/tests/test_internal_parameters.py index 67811b7c2..583cbd0cc 100644 --- a/src/artemis/parameters/tests/test_internal_parameters.py +++ b/src/artemis/parameters/tests/test_internal_parameters.py @@ -8,20 +8,22 @@ from pydantic import ValidationError from artemis.external_interaction.ispyb.ispyb_dataclass import ( - ISPYB_PARAM_DEFAULTS, - IspybParams, + GRIDSCAN_ISPYB_PARAM_DEFAULTS, + GridscanIspybParams, ) from artemis.parameters import external_parameters from artemis.parameters.external_parameters import from_file from artemis.parameters.internal_parameters import ( - ArtemisParameters, InternalParameters, extract_artemis_params_from_flat_dict, fetch_subdict_from_bucket, flatten_dict, get_extracted_experiment_and_flat_artemis_params, ) -from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters +from artemis.parameters.plan_specific.fgs_internal_params import ( + FGSInternalParameters, + GridscanArtemisParameters, +) @pytest.fixture @@ -50,12 +52,12 @@ def test_cant_initialise_abstract_internalparams(): def test_ispyb_param_dict(): - ispyb_params = IspybParams(**ISPYB_PARAM_DEFAULTS) + ispyb_params = GridscanIspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) as_dict = ispyb_params.dict() assert isinstance(as_dict.get("position"), list) - modified_params = copy.deepcopy(ISPYB_PARAM_DEFAULTS) + modified_params = copy.deepcopy(GRIDSCAN_ISPYB_PARAM_DEFAULTS) modified_params["position"] = [123, 7777777, 3] - modified_ispyb_params = IspybParams(**modified_params) + modified_ispyb_params = GridscanIspybParams(**modified_params) assert ispyb_params != modified_ispyb_params assert isinstance(modified_ispyb_params.position, np.ndarray) modified_as_dict = modified_ispyb_params.dict() @@ -92,25 +94,35 @@ def test_flatten(): def test_artemis_params_needs_values_from_experiment(raw_params): extracted_artemis_param_dict = extract_artemis_params_from_flat_dict( - flatten_dict(raw_params) + flatten_dict(raw_params), FGSInternalParameters._artemis_param_key_definitions() ) with pytest.raises(ValidationError): - artemis_params = ArtemisParameters(**extracted_artemis_param_dict) + artemis_params = GridscanArtemisParameters(**extracted_artemis_param_dict) with pytest.raises(UnboundLocalError): assert artemis_params is not None +def test_artemis_parameters_only_from_file(): + with open("src/artemis/parameters/tests/test_data/artemis_parameters.json") as f: + artemis_param_dict = json.load(f) + artemis_params_deserialised = GridscanArtemisParameters(**artemis_param_dict) + ispyb = artemis_params_deserialised.ispyb_params + detector = artemis_params_deserialised.detector_params + assert isinstance(ispyb, GridscanIspybParams) + assert isinstance(detector, DetectorParams) + + def test_artemis_params_can_be_deserialised_from_internal_representation(raw_params): internal_params = FGSInternalParameters(**raw_params) artemis_param_json = internal_params.artemis_params.json() artemis_param_dict = json.loads(artemis_param_json) assert artemis_param_dict.get("ispyb_params") is not None assert artemis_param_dict.get("detector_params") is not None - artemis_params_deserialised = ArtemisParameters(**artemis_param_dict) + artemis_params_deserialised = GridscanArtemisParameters(**artemis_param_dict) assert internal_params.artemis_params == artemis_params_deserialised ispyb = artemis_params_deserialised.ispyb_params detector = artemis_params_deserialised.detector_params - assert isinstance(ispyb, IspybParams) + assert isinstance(ispyb, GridscanIspybParams) assert isinstance(detector, DetectorParams) diff --git a/src/artemis/parameters/tests/test_pydantic_classes.py b/src/artemis/parameters/tests/test_pydantic_classes.py new file mode 100644 index 000000000..2addea458 --- /dev/null +++ b/src/artemis/parameters/tests/test_pydantic_classes.py @@ -0,0 +1,35 @@ +from pydantic import BaseModel + +from artemis.external_interaction.ispyb.ispyb_dataclass import ( + GRIDSCAN_ISPYB_PARAM_DEFAULTS, + GridscanIspybParams, + IspybParams, +) + + +class FirstModel(BaseModel): + x: int + + +class SecondModel(FirstModel): + y: int + + +data = {"x": 3, "y": 4} + + +def test_first_model(): + first = FirstModel(**data) + + +def test_second_model(): + second = SecondModel(**data) + + +def test_first_isp_model(): + first = IspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) + + +def test_second_isp_model(): + second = GridscanIspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) + assert second.visit_path == "" diff --git a/src/artemis/parameters/tests/test_schema_validation.py b/src/artemis/parameters/tests/test_schema_validation.py index 5a3ce6308..2065af8cf 100644 --- a/src/artemis/parameters/tests/test_schema_validation.py +++ b/src/artemis/parameters/tests/test_schema_validation.py @@ -42,7 +42,7 @@ def test_good_params_artemisparams_validates(): jsonschema.validate(params["artemis_params"], artemis_schema, resolver=resolver) -def test_good_params_ispybparams_validates(): +def test_good_params_GridscanIspybParams_validates(): jsonschema.validate( params["artemis_params"]["ispyb_params"], ispyb_schema, resolver=resolver ) From 565dbc333b384b0e6c2e3f84bace24b5e73aaf82 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 08:42:43 +0100 Subject: [PATCH 02/42] make calling of subclasses consistent --- .../parameters/plan_specific/fgs_internal_params.py | 12 ++++-------- .../plan_specific/rotation_scan_internal_params.py | 13 ++++--------- 2 files changed, 8 insertions(+), 17 deletions(-) diff --git a/src/artemis/parameters/plan_specific/fgs_internal_params.py b/src/artemis/parameters/plan_specific/fgs_internal_params.py index 9aed12f72..07a1ee265 100644 --- a/src/artemis/parameters/plan_specific/fgs_internal_params.py +++ b/src/artemis/parameters/plan_specific/fgs_internal_params.py @@ -12,13 +12,6 @@ GridscanIspybParams, IspybParams, ) -from artemis.parameters.constants import ( - DEFAULT_EXPERIMENT_TYPE, - DETECTOR_PARAM_DEFAULTS, - SIM_BEAMLINE, - SIM_INSERTION_PREFIX, - SIM_ZOCALO_ENV, -) from artemis.parameters.internal_parameters import ( ArtemisParameters, InternalParameters, @@ -32,6 +25,9 @@ class GridscanArtemisParameters(ArtemisParameters): **GRIDSCAN_ISPYB_PARAM_DEFAULTS ) + def __init__(self, **kwargs): + super().__init__(**kwargs) # TODO REMOVE JUST FOR DEBUGGING + class Config: arbitrary_types_allowed = True json_encoders = { @@ -94,4 +90,4 @@ def _preprocess_artemis_params( artemis_param_dict = extract_artemis_params_from_flat_dict( all_params, cls._artemis_param_key_definitions() ) - return ArtemisParameters(**artemis_param_dict) + return GridscanArtemisParameters(**artemis_param_dict) diff --git a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py index 4757f8565..efe0827aa 100644 --- a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py +++ b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py @@ -15,13 +15,6 @@ IspybParams, RotationIspybParams, ) -from artemis.parameters.constants import ( - DEFAULT_EXPERIMENT_TYPE, - DETECTOR_PARAM_DEFAULTS, - SIM_BEAMLINE, - SIM_INSERTION_PREFIX, - SIM_ZOCALO_ENV, -) from artemis.parameters.internal_parameters import ( ArtemisParameters, InternalParameters, @@ -31,7 +24,9 @@ class RotationArtemisParameters(ArtemisParameters): - ispyb_params: IspybParams = RotationIspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) + ispyb_params: RotationIspybParams = RotationIspybParams( + **GRIDSCAN_ISPYB_PARAM_DEFAULTS + ) class Config: arbitrary_types_allowed = True @@ -130,7 +125,7 @@ def _preprocess_artemis_params( all_params["num_triggers"] = 1 all_params["num_images_per_trigger"] = all_params["num_images"] all_params["upper_left"] = np.array(all_params["upper_left"]) - return ArtemisParameters( + return RotationArtemisParameters( **extract_artemis_params_from_flat_dict( all_params, cls._artemis_param_key_definitions() ) From 169795725aae20e550c52756253ce82b22f44318 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 08:48:55 +0100 Subject: [PATCH 03/42] move ispybparams contents to base class --- .../ispyb/ispyb_dataclass.py | 56 ++----------------- .../test_rotation_internal_parameters.py | 4 +- 2 files changed, 7 insertions(+), 53 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py index 8b53598bd..bfbb6c889 100644 --- a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py @@ -33,12 +33,8 @@ class IspybParams(BaseModel): visit_path: str - - -class RotationIspybParams(IspybParams): microns_per_pixel_x: float microns_per_pixel_y: float - upper_left: np.ndarray position: np.ndarray class Config: @@ -47,19 +43,9 @@ class Config: def dict(self, **kwargs): as_dict = super().dict(**kwargs) - as_dict["upper_left"] = as_dict["upper_left"].tolist() as_dict["position"] = as_dict["position"].tolist() return as_dict - @validator("upper_left", pre=True) - def _parse_upper_left( - cls, upper_left: list[int | float] | np.ndarray, values: Dict[str, Any] - ) -> np.ndarray: - assert len(upper_left) == 3 - if isinstance(upper_left, np.ndarray): - return upper_left - return np.array(upper_left) - @validator("position", pre=True) def _parse_position( cls, position: list[int | float] | np.ndarray, values: Dict[str, Any] @@ -91,20 +77,16 @@ def _parse_position( xtal_snapshots_omega_end: Optional[List[str]] = None +class RotationIspybParams(IspybParams): + ... + + class GridscanIspybParams(IspybParams): - microns_per_pixel_x: float - microns_per_pixel_y: float upper_left: np.ndarray - position: np.ndarray - - class Config: - arbitrary_types_allowed = True - json_encoders = {np.ndarray: lambda a: a.tolist()} def dict(self, **kwargs): as_dict = super().dict(**kwargs) as_dict["upper_left"] = as_dict["upper_left"].tolist() - as_dict["position"] = as_dict["position"].tolist() return as_dict @validator("upper_left", pre=True) @@ -116,36 +98,6 @@ def _parse_upper_left( return upper_left return np.array(upper_left) - @validator("position", pre=True) - def _parse_position( - cls, position: list[int | float] | np.ndarray, values: Dict[str, Any] - ) -> np.ndarray: - assert len(position) == 3 - if isinstance(position, np.ndarray): - return position - return np.array(position) - - transmission: float - flux: float - wavelength: float - beam_size_x: float - beam_size_y: float - focal_spot_size_x: float - focal_spot_size_y: float - comment: str - resolution: float - - sample_id: Optional[int] = None - sample_barcode: Optional[str] = None - - # Optional from GDA as populated by Ophyd - undulator_gap: Optional[float] = None - synchrotron_mode: Optional[str] = None - slit_gap_size_x: Optional[float] = None - slit_gap_size_y: Optional[float] = None - xtal_snapshots_omega_start: Optional[List[str]] = None - xtal_snapshots_omega_end: Optional[List[str]] = None - class Orientation(Enum): HORIZONTAL = "horizontal" diff --git a/src/artemis/parameters/plan_specific/tests/test_rotation_internal_parameters.py b/src/artemis/parameters/plan_specific/tests/test_rotation_internal_parameters.py index e0af0a25a..4f027415a 100644 --- a/src/artemis/parameters/plan_specific/tests/test_rotation_internal_parameters.py +++ b/src/artemis/parameters/plan_specific/tests/test_rotation_internal_parameters.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock import numpy as np +import pytest from dodal.devices.det_dim_constants import EIGER2_X_16M_SIZE from dodal.devices.motors import XYZLimitBundle @@ -53,7 +54,8 @@ def test_rotation_parameters_load_from_file(): ispyb_params = internal_parameters.artemis_params.ispyb_params np.testing.assert_array_equal(ispyb_params.position, np.array([10, 20, 30])) - np.testing.assert_array_equal(ispyb_params.upper_left, np.array([10, 20, 30])) + with pytest.raises(AttributeError): + ispyb_params.upper_left detector_params = internal_parameters.artemis_params.detector_params From 1edec586bc15e252a62e1ebacc4f155f01c96ab0 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 08:59:07 +0100 Subject: [PATCH 04/42] remove test file --- .../parameters/tests/test_pydantic_classes.py | 35 ------------------- 1 file changed, 35 deletions(-) delete mode 100644 src/artemis/parameters/tests/test_pydantic_classes.py diff --git a/src/artemis/parameters/tests/test_pydantic_classes.py b/src/artemis/parameters/tests/test_pydantic_classes.py deleted file mode 100644 index 2addea458..000000000 --- a/src/artemis/parameters/tests/test_pydantic_classes.py +++ /dev/null @@ -1,35 +0,0 @@ -from pydantic import BaseModel - -from artemis.external_interaction.ispyb.ispyb_dataclass import ( - GRIDSCAN_ISPYB_PARAM_DEFAULTS, - GridscanIspybParams, - IspybParams, -) - - -class FirstModel(BaseModel): - x: int - - -class SecondModel(FirstModel): - y: int - - -data = {"x": 3, "y": 4} - - -def test_first_model(): - first = FirstModel(**data) - - -def test_second_model(): - second = SecondModel(**data) - - -def test_first_isp_model(): - first = IspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) - - -def test_second_isp_model(): - second = GridscanIspybParams(**GRIDSCAN_ISPYB_PARAM_DEFAULTS) - assert second.visit_path == "" From 159333e01f32d4b5930083e4d4af2543d9562f4a Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 10:39:59 +0100 Subject: [PATCH 05/42] fix ispyb handler and tests --- .../callbacks/fgs/ispyb_callback.py | 12 +- .../callbacks/fgs/tests/conftest.py | 10 +- .../ispyb/store_in_ispyb.py | 120 +++++++++++------- .../system_tests/conftest.py | 14 +- .../system_tests/test_ispyb_dev_connection.py | 20 +-- .../unit_tests/test_store_in_ispyb.py | 26 ++-- 6 files changed, 115 insertions(+), 87 deletions(-) diff --git a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py index 0fc002042..e90b7694d 100644 --- a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py @@ -7,9 +7,9 @@ from artemis.external_interaction.exceptions import ISPyBDepositionNotMade from artemis.external_interaction.ispyb.store_in_ispyb import ( - StoreInIspyb, - StoreInIspyb2D, - StoreInIspyb3D, + Store2DGridscanInIspyb, + Store3DGridscanInIspyb, + StoreGridscanInIspyb, ) from artemis.log import LOGGER, set_dcgid_tag from artemis.parameters.constants import ISPYB_PLAN_NAME, SIM_ISPYB_CONFIG @@ -42,10 +42,10 @@ def __init__(self, parameters: FGSInternalParameters): "Using dev ISPyB database. If you want to use the real database, please" " set the ISPYB_CONFIG_PATH environment variable." ) - self.ispyb: StoreInIspyb = ( - StoreInIspyb3D(ispyb_config, self.params) + self.ispyb: StoreGridscanInIspyb = ( + Store3DGridscanInIspyb(ispyb_config, self.params) if self.params.experiment_params.is_3d_grid_scan - else StoreInIspyb2D(ispyb_config, self.params) + else Store2DGridscanInIspyb(ispyb_config, self.params) ) self.ispyb_ids: tuple = (None, None, None) self.uid_to_finalize_on = None diff --git a/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py b/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py index cad9600dd..ab232e8ea 100644 --- a/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py +++ b/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py @@ -16,7 +16,7 @@ def nexus_writer(): @pytest.fixture def mock_ispyb_get_time(): with patch( - "artemis.external_interaction.callbacks.fgs.ispyb_callback.StoreInIspyb3D.get_current_time_string" + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.get_current_time_string" ) as p: yield p @@ -24,7 +24,7 @@ def mock_ispyb_get_time(): @pytest.fixture def mock_ispyb_store_grid_scan(): with patch( - "artemis.external_interaction.callbacks.fgs.ispyb_callback.StoreInIspyb3D.store_grid_scan" + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.store_grid_scan" ) as p: yield p @@ -32,7 +32,7 @@ def mock_ispyb_store_grid_scan(): @pytest.fixture def mock_ispyb_update_time_and_status(): with patch( - "artemis.external_interaction.callbacks.fgs.ispyb_callback.StoreInIspyb3D.update_grid_scan_with_end_time_and_status" + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.update_grid_scan_with_end_time_and_status" ) as p: yield p @@ -40,7 +40,7 @@ def mock_ispyb_update_time_and_status(): @pytest.fixture def mock_ispyb_begin_deposition(): with patch( - "artemis.external_interaction.callbacks.fgs.ispyb_callback.StoreInIspyb3D.begin_deposition" + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.begin_deposition" ) as p: yield p @@ -48,7 +48,7 @@ def mock_ispyb_begin_deposition(): @pytest.fixture def mock_ispyb_end_deposition(): with patch( - "artemis.external_interaction.callbacks.fgs.ispyb_callback.StoreInIspyb3D.end_deposition" + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.end_deposition" ) as p: yield p diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index feef623db..7e6fd2b30 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -8,10 +8,12 @@ import dodal.devices.oav.utils as oav_utils import ispyb import ispyb.sqlalchemy -import numpy as np +from dodal.devices.detector import DetectorParams +from ispyb.sp.core import Core +from ispyb.sp.mxacquisition import MXAcquisition from sqlalchemy.connectors import Connector -from artemis.external_interaction.ispyb.ispyb_dataclass import Orientation +from artemis.external_interaction.ispyb.ispyb_dataclass import IspybParams, Orientation from artemis.log import LOGGER from artemis.tracing import TRACER @@ -22,28 +24,52 @@ class StoreInIspyb(ABC): + ispyb_params: IspybParams | None = None + detector_params: DetectorParams | None = None + run_number: int | None = None + omega_start: int | None = None + experiment_type: str | None = None + xtal_snapshots: list[str] | None = None + conn: Connector = None + mx_acquisition: MXAcquisition | None = None + core: Core | None = None + datacollection_ids: tuple[int, ...] | None = None + datacollection_group_id: int | None = None + + def __init__(self, ispyb_config, parameters=None): + self.ISPYB_CONFIG_PATH: str = ispyb_config + self.full_params: InternalParameters = parameters + + @abstractmethod + def _store_scan_data(self): + pass + + @abstractmethod + def begin_deposition(self, success: str, reason: str): + pass + + @abstractmethod + def end_deposition(self, success: str, reason: str): + pass + + def append_to_comment( + self, datacollection_id: int, comment: str, delimiter: str = " " + ) -> None: + with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: + self.mx_acquisition = self.conn.mx_acquisition + self.mx_acquisition.update_data_collection_append_comments( + datacollection_id, comment, delimiter + ) + + +class StoreGridscanInIspyb(StoreInIspyb): VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})/" def __init__(self, ispyb_config, parameters=None): - self.ISPYB_CONFIG_PATH = ispyb_config - self.full_params = parameters - self.ispyb_params = None - self.detector_params = None - self.run_number = None - self.omega_start = None - self.experiment_type = None - self.xtal_snapshots = None + super().__init__(ispyb_config, parameters) self.upper_left = None self.y_steps = None self.y_step_size = None - - # writing to ispyb - self.conn: Connector = None - self.mx_acquisition = None - self.core = None - - self.datacollection_ids = None - self.datacollection_group_id = None self.grid_ids = None def begin_deposition(self): @@ -71,6 +97,10 @@ def end_deposition(self, success: str, reason: str): else: run_status = "DataCollection Successful" current_time = self.get_current_time_string() + assert ( + self.datacollection_ids is not None + ), "Can't end ISPyB deposition, datacollection IDs are missing" + assert self.datacollection_group_id is not None for id in self.datacollection_ids: self.update_grid_scan_with_end_time_and_status( current_time, run_status, reason, id, self.datacollection_group_id @@ -83,12 +113,10 @@ def store_grid_scan(self, full_params: InternalParameters): self.run_number = self.detector_params.run_number self.omega_start = self.detector_params.omega_start self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_start - self.upper_left = np.array( - [ - self.ispyb_params.upper_left[0], - self.ispyb_params.upper_left[1], - ] - ) + self.upper_left = [ + int(self.ispyb_params.upper_left[0]), + int(self.ispyb_params.upper_left[1]), + ] self.y_steps = full_params.experiment_params.y_steps self.y_step_size = full_params.experiment_params.y_step_size @@ -98,19 +126,6 @@ def store_grid_scan(self, full_params: InternalParameters): return self._store_scan_data() - @abstractmethod - def _store_scan_data(self): - pass - - def append_to_comment( - self, datacollection_id: int, comment: str, delimiter: str = " " - ) -> None: - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition = self.conn.mx_acquisition - self.mx_acquisition.update_data_collection_append_comments( - datacollection_id, comment, delimiter - ) - def update_grid_scan_with_end_time_and_status( self, end_time: str, @@ -119,11 +134,13 @@ def update_grid_scan_with_end_time_and_status( datacollection_id: int, datacollection_group_id: int, ) -> None: + assert self.ispyb_params is not None + assert self.detector_params is not None if reason is not None and reason != "": self.append_to_comment(datacollection_id, f"{run_status} reason: {reason}") with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition = self.conn.mx_acquisition + self.mx_acquisition: MXAcquisition = self.conn.mx_acquisition params = self.mx_acquisition.get_data_collection_params() params["id"] = datacollection_id @@ -134,6 +151,8 @@ def update_grid_scan_with_end_time_and_status( self.mx_acquisition.upsert_data_collection(list(params.values())) def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: + assert self.ispyb_params is not None + params = self.mx_acquisition.get_dc_grid_params() params["parentid"] = ispyb_data_collection_id @@ -152,6 +171,8 @@ def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: return self.mx_acquisition.upsert_dc_grid(list(params.values())) def _construct_comment(self) -> str: + assert self.ispyb_params is not None + bottom_right = oav_utils.bottom_right_from_top_left( self.upper_left, self.full_params.experiment_params.x_steps, @@ -173,6 +194,10 @@ def _construct_comment(self) -> str: @TRACER.start_as_current_span("store_ispyb_datacollection_table") def _store_data_collection_table(self, data_collection_group_id: int) -> int: + assert self.ispyb_params is not None + assert self.detector_params is not None + assert self.core is not None + assert self.xtal_snapshots is not None try: session_id = self.core.retrieve_visit_id(self.get_visit_string()) except ispyb.NoResult: @@ -236,6 +261,7 @@ def _store_data_collection_table(self, data_collection_group_id: int) -> int: return self.mx_acquisition.upsert_data_collection(list(params.values())) def _store_position_table(self, dc_id: int) -> int: + assert self.ispyb_params is not None params = self.mx_acquisition.get_dc_position_params() params["id"] = dc_id @@ -243,11 +269,13 @@ def _store_position_table(self, dc_id: int) -> int: params["pos_x"], params["pos_y"], params["pos_z"], - ) = self.ispyb_params.position + ) = self.ispyb_params.position.tolist() return self.mx_acquisition.update_dc_position(list(params.values())) def _store_data_collection_group_table(self) -> int: + assert self.core is not None + assert self.ispyb_params is not None try: session_id = self.core.retrieve_visit_id(self.get_visit_string()) except ispyb.NoResult: @@ -279,7 +307,7 @@ def get_visit_string_from_path(self, path): return match.group(1) if match else None -class StoreInIspyb3D(StoreInIspyb): +class Store3DGridscanInIspyb(StoreGridscanInIspyb): def __init__(self, ispyb_config, parameters=None): super().__init__(ispyb_config, parameters) self.experiment_type = "Mesh3D" @@ -315,17 +343,15 @@ def __prepare_second_scan_params(self): self.omega_start += 90 self.run_number += 1 self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_end - self.upper_left = np.array( - [ - self.ispyb_params.upper_left[0], - self.ispyb_params.upper_left[2], - ] - ) + self.upper_left = [ + int(self.ispyb_params.upper_left[0]), + int(self.ispyb_params.upper_left[2]), + ] self.y_steps = self.full_params.experiment_params.z_steps self.y_step_size = self.full_params.experiment_params.z_step_size -class StoreInIspyb2D(StoreInIspyb): +class Store2DGridscanInIspyb(StoreGridscanInIspyb): def __init__(self, ispyb_config, parameters=None): super().__init__(ispyb_config, parameters) self.experiment_type = "mesh" diff --git a/src/artemis/external_interaction/system_tests/conftest.py b/src/artemis/external_interaction/system_tests/conftest.py index 8807e5330..a7e283f6d 100644 --- a/src/artemis/external_interaction/system_tests/conftest.py +++ b/src/artemis/external_interaction/system_tests/conftest.py @@ -11,8 +11,8 @@ import artemis.external_interaction.zocalo.zocalo_interaction from artemis.external_interaction.ispyb.store_in_ispyb import ( - StoreInIspyb2D, - StoreInIspyb3D, + Store2DGridscanInIspyb, + Store3DGridscanInIspyb, ) from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters @@ -76,7 +76,7 @@ def fetch_comment() -> Callable: @pytest.fixture def dummy_params(): - dummy_params = FGSInternalParameters(default_raw_params()) + dummy_params = FGSInternalParameters(**default_raw_params()) dummy_params.artemis_params.ispyb_params.upper_left = np.array([100, 100, 50]) dummy_params.artemis_params.ispyb_params.microns_per_pixel_x = 0.8 dummy_params.artemis_params.ispyb_params.microns_per_pixel_y = 0.8 @@ -87,13 +87,13 @@ def dummy_params(): @pytest.fixture -def dummy_ispyb(dummy_params) -> StoreInIspyb2D: - return StoreInIspyb2D(ISPYB_CONFIG, dummy_params) +def dummy_ispyb(dummy_params) -> Store2DGridscanInIspyb: + return Store2DGridscanInIspyb(ISPYB_CONFIG, dummy_params) @pytest.fixture -def dummy_ispyb_3d(dummy_params) -> StoreInIspyb3D: - return StoreInIspyb3D(ISPYB_CONFIG, dummy_params) +def dummy_ispyb_3d(dummy_params) -> Store3DGridscanInIspyb: + return Store3DGridscanInIspyb(ISPYB_CONFIG, dummy_params) @pytest.fixture diff --git a/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py b/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py index bb2af06d7..693ab9b04 100644 --- a/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py +++ b/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py @@ -1,9 +1,9 @@ import pytest from artemis.external_interaction.ispyb.store_in_ispyb import ( - StoreInIspyb, - StoreInIspyb2D, - StoreInIspyb3D, + Store2DGridscanInIspyb, + Store3DGridscanInIspyb, + StoreGridscanInIspyb, ) from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters @@ -26,7 +26,7 @@ def test_ispyb_get_comment_from_collection_correctly(fetch_comment): @pytest.mark.s03 def test_ispyb_deposition_comment_correct_on_failure( - dummy_ispyb: StoreInIspyb2D, fetch_comment + dummy_ispyb: Store2DGridscanInIspyb, fetch_comment ): dcid = dummy_ispyb.begin_deposition() dummy_ispyb.end_deposition("fail", "could not connect to devices") @@ -38,7 +38,7 @@ def test_ispyb_deposition_comment_correct_on_failure( @pytest.mark.s03 def test_ispyb_deposition_comment_correct_for_3D_on_failure( - dummy_ispyb_3d: StoreInIspyb3D, fetch_comment + dummy_ispyb_3d: Store3DGridscanInIspyb, fetch_comment ): dcid = dummy_ispyb_3d.begin_deposition() dcid1 = dcid[0][0] @@ -58,10 +58,10 @@ def test_ispyb_deposition_comment_correct_for_3D_on_failure( @pytest.mark.parametrize( "StoreClass, exp_num_of_grids, success", [ - (StoreInIspyb2D, 1, False), - (StoreInIspyb2D, 1, True), - (StoreInIspyb3D, 2, False), - (StoreInIspyb3D, 2, True), + (Store2DGridscanInIspyb, 1, False), + (Store2DGridscanInIspyb, 1, True), + (Store3DGridscanInIspyb, 2, False), + (Store3DGridscanInIspyb, 2, True), ], ) def test_can_store_2D_ispyb_data_correctly_when_in_error( @@ -69,7 +69,7 @@ def test_can_store_2D_ispyb_data_correctly_when_in_error( ): test_params = FGSInternalParameters(**default_raw_params()) test_params.artemis_params.ispyb_params.visit_path = "/tmp/cm31105-4/" - ispyb: StoreInIspyb = StoreClass(ISPYB_CONFIG, test_params) + ispyb: StoreGridscanInIspyb = StoreClass(ISPYB_CONFIG, test_params) dc_ids, grid_ids, dcg_id = ispyb.begin_deposition() assert len(dc_ids) == exp_num_of_grids diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index ddf15ef2a..4edc9c7e6 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -7,8 +7,8 @@ from mockito import mock, when from artemis.external_interaction.ispyb.store_in_ispyb import ( - StoreInIspyb2D, - StoreInIspyb3D, + Store2DGridscanInIspyb, + Store3DGridscanInIspyb, ) from artemis.parameters.constants import SIM_ISPYB_CONFIG from artemis.parameters.external_parameters import from_file as default_raw_params @@ -34,7 +34,7 @@ def dummy_params(): @pytest.fixture def dummy_ispyb(dummy_params): - store_in_ispyb_2d = StoreInIspyb2D(SIM_ISPYB_CONFIG, dummy_params) + store_in_ispyb_2d = Store2DGridscanInIspyb(SIM_ISPYB_CONFIG, dummy_params) store_in_ispyb_2d.get_current_datacollection_comment = MagicMock() store_in_ispyb_2d.get_current_datacollection_comment.return_value = "" return store_in_ispyb_2d @@ -42,7 +42,7 @@ def dummy_ispyb(dummy_params): @pytest.fixture def dummy_ispyb_3d(dummy_params): - store_in_ispyb_3d = StoreInIspyb3D(SIM_ISPYB_CONFIG, dummy_params) + store_in_ispyb_3d = Store3DGridscanInIspyb(SIM_ISPYB_CONFIG, dummy_params) store_in_ispyb_3d.get_current_datacollection_comment = MagicMock() store_in_ispyb_3d.get_current_datacollection_comment.return_value = "" return store_in_ispyb_3d @@ -97,7 +97,9 @@ def test_store_grid_scan(ispyb_conn, dummy_ispyb, dummy_params): @patch("ispyb.open", new_callable=mock_open) def test_store_3d_grid_scan( - ispyb_conn, dummy_ispyb_3d: StoreInIspyb3D, dummy_params: FGSInternalParameters + ispyb_conn, + dummy_ispyb_3d: Store3DGridscanInIspyb, + dummy_params: FGSInternalParameters, ): ispyb_conn.return_value.mx_acquisition = mock() ispyb_conn.return_value.core = mock() @@ -225,7 +227,7 @@ def test_sample_id(default_params, actual): @patch("ispyb.open") def test_given_real_sampleid_when_grid_scan_stored_then_sample_id_set( - ispyb_conn, dummy_ispyb: StoreInIspyb2D, dummy_params: FGSInternalParameters + ispyb_conn, dummy_ispyb: Store2DGridscanInIspyb, dummy_params: FGSInternalParameters ): expected_sample_id = "0001" dummy_params.artemis_params.ispyb_params.sample_id = expected_sample_id @@ -242,7 +244,7 @@ def test_sample_id(default_params, actual): @patch("ispyb.open") def test_fail_result_run_results_in_bad_run_status( mock_ispyb_conn: MagicMock, - dummy_ispyb: StoreInIspyb2D, + dummy_ispyb: Store2DGridscanInIspyb, ): setup_mock_return_values(mock_ispyb_conn) mock_mx_aquisition = ( @@ -265,7 +267,7 @@ def test_fail_result_run_results_in_bad_run_status( @patch("ispyb.open") def test_no_exception_during_run_results_in_good_run_status( mock_ispyb_conn: MagicMock, - dummy_ispyb: StoreInIspyb2D, + dummy_ispyb: Store2DGridscanInIspyb, ): setup_mock_return_values(mock_ispyb_conn) mock_mx_aquisition = ( @@ -286,7 +288,7 @@ def test_no_exception_during_run_results_in_good_run_status( @patch("ispyb.open") def test_ispyb_deposition_comment_correct( mock_ispyb_conn: MagicMock, - dummy_ispyb: StoreInIspyb2D, + dummy_ispyb: Store2DGridscanInIspyb, ): setup_mock_return_values(mock_ispyb_conn) mock_mx_aquisition = ( @@ -306,7 +308,7 @@ def test_ispyb_deposition_comment_correct( @patch("ispyb.open") def test_ispyb_deposition_rounds_to_int( mock_ispyb_conn: MagicMock, - dummy_ispyb: StoreInIspyb2D, + dummy_ispyb: Store2DGridscanInIspyb, ): setup_mock_return_values(mock_ispyb_conn) mock_mx_aquisition = ( @@ -329,7 +331,7 @@ def test_ispyb_deposition_rounds_to_int( @patch("ispyb.open") def test_ispyb_deposition_comment_for_3D_correct( mock_ispyb_conn: MagicMock, - dummy_ispyb_3d: StoreInIspyb3D, + dummy_ispyb_3d: Store3DGridscanInIspyb, ): setup_mock_return_values(mock_ispyb_conn) mock_mx_aquisition = ( @@ -351,7 +353,7 @@ def test_ispyb_deposition_comment_for_3D_correct( @patch("ispyb.open") def test_given_x_and_y_steps_different_from_total_images_when_grid_scan_stored_then_num_images_correct( - ispyb_conn, dummy_ispyb: StoreInIspyb2D, dummy_params: FGSInternalParameters + ispyb_conn, dummy_ispyb: Store2DGridscanInIspyb, dummy_params: FGSInternalParameters ): expected_number_of_steps = 200 * 3 dummy_params.experiment_params.x_steps = 200 From 2950597ac43f1348cb10898ccb84fd1334f8b97a Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 11:08:51 +0100 Subject: [PATCH 06/42] tidy up a bit --- .../external_interaction/callbacks/fgs/zocalo_callback.py | 6 ++---- .../external_interaction/system_tests/test_zocalo_system.py | 2 +- src/artemis/system_tests/test_fgs_plan.py | 6 +++--- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/artemis/external_interaction/callbacks/fgs/zocalo_callback.py b/src/artemis/external_interaction/callbacks/fgs/zocalo_callback.py index 3090a0467..f67b0242f 100644 --- a/src/artemis/external_interaction/callbacks/fgs/zocalo_callback.py +++ b/src/artemis/external_interaction/callbacks/fgs/zocalo_callback.py @@ -117,12 +117,10 @@ def wait_for_results(self, fallback_xyz: ndarray) -> tuple[ndarray, Optional[lis self.ispyb.append_to_comment(crystal_summary) raw_centre = np.array([*(raw_results[0]["centre_of_mass"])]) + adjusted_centre = raw_centre - np.array([0.5, 0.5, 0.5]) # _wait_for_result returns the centre of the grid box, but we want the corner - results = np.array( - [raw_centre[0] - 0.5, raw_centre[1] - 0.5, raw_centre[2] - 0.5] - ) - xray_centre = self.grid_position_to_motor_position(results) + xray_centre = self.grid_position_to_motor_position(adjusted_centre) bbox_size: list[int] | None = bboxes[0] diff --git a/src/artemis/external_interaction/system_tests/test_zocalo_system.py b/src/artemis/external_interaction/system_tests/test_zocalo_system.py index 842210320..a24415272 100644 --- a/src/artemis/external_interaction/system_tests/test_zocalo_system.py +++ b/src/artemis/external_interaction/system_tests/test_zocalo_system.py @@ -51,7 +51,7 @@ def test_given_a_result_with_no_diffraction_when_zocalo_called_then_move_to_fall ): fallback = np.array([1, 2, 3]) zc, centre = run_zocalo_with_dev_ispyb("NO_DIFF", fallback) - assert centre == fallback + assert np.allclose(centre, fallback) @pytest.mark.s03 diff --git a/src/artemis/system_tests/test_fgs_plan.py b/src/artemis/system_tests/test_fgs_plan.py index 9a0e66105..8a2cde968 100644 --- a/src/artemis/system_tests/test_fgs_plan.py +++ b/src/artemis/system_tests/test_fgs_plan.py @@ -211,6 +211,6 @@ def test_WHEN_plan_run_THEN_move_to_centre_returned_from_zocalo_expected_centre( RE(get_plan(params, callbacks)) # The following numbers are derived from the centre returned in fake_zocalo - assert fgs_composite.sample_motors.x.user_readback.get() == pytest.approx(0.05) - assert fgs_composite.sample_motors.y.user_readback.get() == pytest.approx(0.15) - assert fgs_composite.sample_motors.z.user_readback.get() == pytest.approx(0.25) + assert fgs_composite.sample_motors.x.user_readback.get() == pytest.approx(-0.05) + assert fgs_composite.sample_motors.y.user_readback.get() == pytest.approx(0.05) + assert fgs_composite.sample_motors.z.user_readback.get() == pytest.approx(0.15) From 19790afe99da3e7da1db27423428035b200f03cd Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 13:00:30 +0100 Subject: [PATCH 07/42] add correct typing to everything --- .../ispyb/store_in_ispyb.py | 65 +++++++++++++++---- 1 file changed, 54 insertions(+), 11 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 7e6fd2b30..97a6b572c 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -13,12 +13,23 @@ from ispyb.sp.mxacquisition import MXAcquisition from sqlalchemy.connectors import Connector -from artemis.external_interaction.ispyb.ispyb_dataclass import IspybParams, Orientation +from artemis.external_interaction.ispyb.ispyb_dataclass import ( + GridscanIspybParams, + IspybParams, + Orientation, + RotationIspybParams, +) from artemis.log import LOGGER from artemis.tracing import TRACER if TYPE_CHECKING: - from artemis.parameters.internal_parameters import InternalParameters + from artemis.parameters.plan_specific.fgs_internal_params import ( + FGSInternalParameters, + ) + from artemis.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, + ) + I03_EIGER_DETECTOR = 78 EIGER_FILE_SUFFIX = "h5" @@ -36,9 +47,8 @@ class StoreInIspyb(ABC): datacollection_ids: tuple[int, ...] | None = None datacollection_group_id: int | None = None - def __init__(self, ispyb_config, parameters=None): + def __init__(self, ispyb_config, parameters=None) -> None: self.ISPYB_CONFIG_PATH: str = ispyb_config - self.full_params: InternalParameters = parameters @abstractmethod def _store_scan_data(self): @@ -62,15 +72,42 @@ def append_to_comment( ) +class StoreRotationInIspyb(StoreInIspyb): + ispyb_params: RotationIspybParams | None = None + + def __init__(self, ispyb_config, parameters=None) -> None: + self.full_params: RotationInternalParameters | None = parameters + super().__init__(ispyb_config, parameters) + + def store_rotation_scan(self): + with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: + self.mx_acquisition = self.conn.mx_acquisition + self.core = self.conn.core + return self._store_scan_data() + + def _store_scan_data(self): + pass + + @abstractmethod + def begin_deposition(self, success: str, reason: str): + pass + + @abstractmethod + def end_deposition(self, success: str, reason: str): + pass + + class StoreGridscanInIspyb(StoreInIspyb): VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})/" - - def __init__(self, ispyb_config, parameters=None): + ispyb_params: GridscanIspybParams | None = None + supper_left: list[int] | None = None + y_steps: int | None = None + y_step_size: int | None = None + grid_ids: tuple[int, ...] | None = None + + def __init__(self, ispyb_config, parameters=None) -> None: + self.full_params: FGSInternalParameters | None = parameters super().__init__(ispyb_config, parameters) - self.upper_left = None - self.y_steps = None - self.y_step_size = None - self.grid_ids = None def begin_deposition(self): ( @@ -106,7 +143,7 @@ def end_deposition(self, success: str, reason: str): current_time, run_status, reason, id, self.datacollection_group_id ) - def store_grid_scan(self, full_params: InternalParameters): + def store_grid_scan(self, full_params: FGSInternalParameters): self.full_params = full_params self.ispyb_params = full_params.artemis_params.ispyb_params self.detector_params = full_params.artemis_params.detector_params @@ -152,6 +189,8 @@ def update_grid_scan_with_end_time_and_status( def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: assert self.ispyb_params is not None + assert self.full_params is not None + assert self.upper_left is not None params = self.mx_acquisition.get_dc_grid_params() @@ -172,6 +211,9 @@ def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: def _construct_comment(self) -> str: assert self.ispyb_params is not None + assert self.full_params is not None + assert self.upper_left is not None + assert self.y_step_size is not None bottom_right = oav_utils.bottom_right_from_top_left( self.upper_left, @@ -197,6 +239,7 @@ def _store_data_collection_table(self, data_collection_group_id: int) -> int: assert self.ispyb_params is not None assert self.detector_params is not None assert self.core is not None + assert self.full_params is not None assert self.xtal_snapshots is not None try: session_id = self.core.retrieve_visit_id(self.get_visit_string()) From d1d6c0194e628cf9e96f7207d1f9e549f92f7065 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 14:00:08 +0100 Subject: [PATCH 08/42] move som emore functions to base class --- .../ispyb/store_in_ispyb.py | 73 ++++++++++--------- 1 file changed, 38 insertions(+), 35 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 97a6b572c..4f69dbbfa 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -71,12 +71,47 @@ def append_to_comment( datacollection_id, comment, delimiter ) + def get_current_time_string(self): + now = datetime.datetime.now() + return now.strftime("%Y-%m-%d %H:%M:%S") + + def get_visit_string(self): + visit_path_match = self.get_visit_string_from_path(self.ispyb_params.visit_path) + if visit_path_match: + return visit_path_match + else: + return self.get_visit_string_from_path(self.detector_params.directory) + + def get_visit_string_from_path(self, path): + match = re.search(self.VISIT_PATH_REGEX, path) if path else None + return match.group(1) if match else None + + def _store_data_collection_group_table(self) -> int: + assert self.core is not None + assert self.ispyb_params is not None + assert self.mx_acquisition is not None + try: + session_id = self.core.retrieve_visit_id(self.get_visit_string()) + except ispyb.NoResult: + raise Exception( + f"Not found - session ID for visit {self.get_visit_string()} where self.ispyb_params.visit_path is {self.ispyb_params.visit_path}" + ) + + params = self.mx_acquisition.get_data_collection_group_params() + params["parentid"] = session_id + params["experimenttype"] = self.experiment_type + params["sampleid"] = self.ispyb_params.sample_id + params["sample_barcode"] = self.ispyb_params.sample_barcode + + return self.mx_acquisition.upsert_data_collection_group(list(params.values())) + class StoreRotationInIspyb(StoreInIspyb): ispyb_params: RotationIspybParams | None = None def __init__(self, ispyb_config, parameters=None) -> None: self.full_params: RotationInternalParameters | None = parameters + self.experiment_type = "SAD" super().__init__(ispyb_config, parameters) def store_rotation_scan(self): @@ -88,14 +123,15 @@ def store_rotation_scan(self): def _store_scan_data(self): pass - @abstractmethod def begin_deposition(self, success: str, reason: str): pass - @abstractmethod def end_deposition(self, success: str, reason: str): pass + def _construct_comment(self) -> str: + return "Hyperion rotation scan" + class StoreGridscanInIspyb(StoreInIspyb): VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})/" @@ -316,39 +352,6 @@ def _store_position_table(self, dc_id: int) -> int: return self.mx_acquisition.update_dc_position(list(params.values())) - def _store_data_collection_group_table(self) -> int: - assert self.core is not None - assert self.ispyb_params is not None - try: - session_id = self.core.retrieve_visit_id(self.get_visit_string()) - except ispyb.NoResult: - raise Exception( - f"Not found - session ID for visit {self.get_visit_string()} where self.ispyb_params.visit_path is {self.ispyb_params.visit_path}" - ) - - params = self.mx_acquisition.get_data_collection_group_params() - params["parentid"] = session_id - params["experimenttype"] = self.experiment_type - params["sampleid"] = self.ispyb_params.sample_id - params["sample_barcode"] = self.ispyb_params.sample_barcode - - return self.mx_acquisition.upsert_data_collection_group(list(params.values())) - - def get_current_time_string(self): - now = datetime.datetime.now() - return now.strftime("%Y-%m-%d %H:%M:%S") - - def get_visit_string(self): - visit_path_match = self.get_visit_string_from_path(self.ispyb_params.visit_path) - if visit_path_match: - return visit_path_match - else: - return self.get_visit_string_from_path(self.detector_params.directory) - - def get_visit_string_from_path(self, path): - match = re.search(self.VISIT_PATH_REGEX, path) if path else None - return match.group(1) if match else None - class Store3DGridscanInIspyb(StoreGridscanInIspyb): def __init__(self, ispyb_config, parameters=None): From f00eeb62be620cfe44a793436b517dfce2b90284 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 14:12:13 +0100 Subject: [PATCH 09/42] make most of store datacollection table generic --- .../ispyb/store_in_ispyb.py | 172 ++++++++++-------- 1 file changed, 101 insertions(+), 71 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 4f69dbbfa..eaad2b70e 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -3,7 +3,7 @@ import datetime import re from abc import ABC, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import dodal.devices.oav.utils as oav_utils import ispyb @@ -54,6 +54,16 @@ def __init__(self, ispyb_config, parameters=None) -> None: def _store_scan_data(self): pass + @abstractmethod + def _construct_comment(self) -> str: + pass + + @abstractmethod + def _mutate_datacollection_params_for_experiment( + self, params: dict[str, Any] + ) -> dict[str, Any]: + pass + @abstractmethod def begin_deposition(self, success: str, reason: str): pass @@ -66,7 +76,7 @@ def append_to_comment( self, datacollection_id: int, comment: str, delimiter: str = " " ) -> None: with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition = self.conn.mx_acquisition + self.mx_acquisition: MXAcquisition = self.conn.mx_acquisition self.mx_acquisition.update_data_collection_append_comments( datacollection_id, comment, delimiter ) @@ -105,6 +115,77 @@ def _store_data_collection_group_table(self) -> int: return self.mx_acquisition.upsert_data_collection_group(list(params.values())) + @TRACER.start_as_current_span("store_ispyb_datacollection_table") + def _store_data_collection_table(self, data_collection_group_id: int) -> int: + assert self.ispyb_params is not None + assert self.detector_params is not None + assert self.core is not None + assert self.xtal_snapshots is not None + assert self.mx_acquisition is not None + try: + session_id = self.core.retrieve_visit_id(self.get_visit_string()) + except ispyb.NoResult: + raise Exception( + f"Not found - session ID for visit {self.get_visit_string()}" + ) + + params = self.mx_acquisition.get_data_collection_params() + + params = self._mutate_datacollection_params_for_experiment(params) + + params["visitid"] = session_id + params["parentid"] = data_collection_group_id + params["sampleid"] = self.ispyb_params.sample_id + params["detectorid"] = I03_EIGER_DETECTOR + params["axis_start"] = self.omega_start + + params["axis_range"] = 0 + params["focal_spot_size_at_samplex"] = self.ispyb_params.focal_spot_size_x + params["focal_spot_size_at_sampley"] = self.ispyb_params.focal_spot_size_y + params["slitgap_vertical"] = self.ispyb_params.slit_gap_size_y + params["slitgap_horizontal"] = self.ispyb_params.slit_gap_size_x + params["beamsize_at_samplex"] = self.ispyb_params.beam_size_x + params["beamsize_at_sampley"] = self.ispyb_params.beam_size_y + params["transmission"] = self.ispyb_params.transmission + params["comments"] = self._construct_comment() + params["datacollection_number"] = self.run_number + params["detector_distance"] = self.detector_params.detector_distance + params["exp_time"] = self.detector_params.exposure_time + params["imgdir"] = self.detector_params.directory + params["imgprefix"] = self.detector_params.prefix + params["imgsuffix"] = EIGER_FILE_SUFFIX + + # Both overlap and n_passes included for backwards compatibility, + # planned to be removed later + params["n_passes"] = 1 + params["overlap"] = 0 + + params["flux"] = self.ispyb_params.flux + params["omegastart"] = self.omega_start + params["start_image_number"] = 1 + params["resolution"] = self.ispyb_params.resolution + params["wavelength"] = self.ispyb_params.wavelength + beam_position = self.detector_params.get_beam_position_mm( + self.detector_params.detector_distance + ) + params["xbeam"], params["ybeam"] = beam_position + ( + params["xtal_snapshot1"], + params["xtal_snapshot2"], + params["xtal_snapshot3"], + ) = self.xtal_snapshots + params["synchrotron_mode"] = self.ispyb_params.synchrotron_mode + params["undulator_gap1"] = self.ispyb_params.undulator_gap + params["starttime"] = self.get_current_time_string() + + # temporary file template until nxs filewriting is integrated and we can use + # that file name + params[ + "file_template" + ] = f"{self.detector_params.prefix}_{self.run_number}_master.h5" + + return self.mx_acquisition.upsert_data_collection(list(params.values())) + class StoreRotationInIspyb(StoreInIspyb): ispyb_params: RotationIspybParams | None = None @@ -114,6 +195,16 @@ def __init__(self, ispyb_config, parameters=None) -> None: self.experiment_type = "SAD" super().__init__(ispyb_config, parameters) + def _mutate_datacollection_params_for_experiment( + self, params: dict[str, Any] + ) -> dict[str, Any]: + assert self.full_params is not None + params["axis_end"] = ( + self.omega_start + self.full_params.experiment_params.rotation_angle + ) + params["n_images"] = self.full_params.experiment_params.get_num_images() + return params + def store_rotation_scan(self): with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: self.mx_acquisition = self.conn.mx_acquisition @@ -199,6 +290,14 @@ def store_grid_scan(self, full_params: FGSInternalParameters): return self._store_scan_data() + def _mutate_datacollection_params_for_experiment( + self, params: dict[str, Any] + ) -> dict[str, Any]: + assert self.full_params is not None + params["axis_end"] = self.omega_start + params["n_images"] = self.full_params.experiment_params.x_steps * self.y_steps + return params + def update_grid_scan_with_end_time_and_status( self, end_time: str, @@ -270,75 +369,6 @@ def _construct_comment(self) -> str: f"bottom right (px): [{bottom_right[0]},{bottom_right[1]}]." ) - @TRACER.start_as_current_span("store_ispyb_datacollection_table") - def _store_data_collection_table(self, data_collection_group_id: int) -> int: - assert self.ispyb_params is not None - assert self.detector_params is not None - assert self.core is not None - assert self.full_params is not None - assert self.xtal_snapshots is not None - try: - session_id = self.core.retrieve_visit_id(self.get_visit_string()) - except ispyb.NoResult: - raise Exception( - f"Not found - session ID for visit {self.get_visit_string()}" - ) - - params = self.mx_acquisition.get_data_collection_params() - params["visitid"] = session_id - params["parentid"] = data_collection_group_id - params["sampleid"] = self.ispyb_params.sample_id - params["detectorid"] = I03_EIGER_DETECTOR - params["axis_start"] = self.omega_start - params["axis_end"] = self.omega_start - params["axis_range"] = 0 - params["focal_spot_size_at_samplex"] = self.ispyb_params.focal_spot_size_x - params["focal_spot_size_at_sampley"] = self.ispyb_params.focal_spot_size_y - params["slitgap_vertical"] = self.ispyb_params.slit_gap_size_y - params["slitgap_horizontal"] = self.ispyb_params.slit_gap_size_x - params["beamsize_at_samplex"] = self.ispyb_params.beam_size_x - params["beamsize_at_sampley"] = self.ispyb_params.beam_size_y - params["transmission"] = self.ispyb_params.transmission - params["comments"] = self._construct_comment() - params["datacollection_number"] = self.run_number - params["detector_distance"] = self.detector_params.detector_distance - params["exp_time"] = self.detector_params.exposure_time - params["imgdir"] = self.detector_params.directory - params["imgprefix"] = self.detector_params.prefix - params["imgsuffix"] = EIGER_FILE_SUFFIX - params["n_images"] = self.full_params.experiment_params.x_steps * self.y_steps - - # Both overlap and n_passes included for backwards compatibility, - # planned to be removed later - params["n_passes"] = 1 - params["overlap"] = 0 - - params["flux"] = self.ispyb_params.flux - params["omegastart"] = self.omega_start - params["start_image_number"] = 1 - params["resolution"] = self.ispyb_params.resolution - params["wavelength"] = self.ispyb_params.wavelength - beam_position = self.detector_params.get_beam_position_mm( - self.detector_params.detector_distance - ) - params["xbeam"], params["ybeam"] = beam_position - ( - params["xtal_snapshot1"], - params["xtal_snapshot2"], - params["xtal_snapshot3"], - ) = self.xtal_snapshots - params["synchrotron_mode"] = self.ispyb_params.synchrotron_mode - params["undulator_gap1"] = self.ispyb_params.undulator_gap - params["starttime"] = self.get_current_time_string() - - # temporary file template until nxs filewriting is integrated and we can use - # that file name - params[ - "file_template" - ] = f"{self.detector_params.prefix}_{self.run_number}_master.h5" - - return self.mx_acquisition.upsert_data_collection(list(params.values())) - def _store_position_table(self, dc_id: int) -> int: assert self.ispyb_params is not None params = self.mx_acquisition.get_dc_position_params() From 8b2f090db84611281c9d8ee3f7889b8647a73f25 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 14:20:04 +0100 Subject: [PATCH 10/42] add _store_scan_data for rotation --- .../ispyb/store_in_ispyb.py | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index eaad2b70e..f6353b847 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -32,6 +32,7 @@ I03_EIGER_DETECTOR = 78 EIGER_FILE_SUFFIX = "h5" +VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})/" class StoreInIspyb(ABC): @@ -93,9 +94,24 @@ def get_visit_string(self): return self.get_visit_string_from_path(self.detector_params.directory) def get_visit_string_from_path(self, path): - match = re.search(self.VISIT_PATH_REGEX, path) if path else None + match = re.search(VISIT_PATH_REGEX, path) if path else None return match.group(1) if match else None + def _store_position_table(self, dc_id: int) -> int: + assert self.ispyb_params is not None + assert self.mx_acquisition is not None + + params = self.mx_acquisition.get_dc_position_params() + + params["id"] = dc_id + ( + params["pos_x"], + params["pos_y"], + params["pos_z"], + ) = self.ispyb_params.position.tolist() + + return self.mx_acquisition.update_dc_position(list(params.values())) + def _store_data_collection_group_table(self) -> int: assert self.core is not None assert self.ispyb_params is not None @@ -212,7 +228,12 @@ def store_rotation_scan(self): return self._store_scan_data() def _store_scan_data(self): - pass + data_collection_group_id = self._store_data_collection_group_table() + data_collection_id = self._store_data_collection_table(data_collection_group_id) + + self._store_position_table(data_collection_id) + + return data_collection_id, data_collection_group_id def begin_deposition(self, success: str, reason: str): pass @@ -225,7 +246,6 @@ def _construct_comment(self) -> str: class StoreGridscanInIspyb(StoreInIspyb): - VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})/" ispyb_params: GridscanIspybParams | None = None supper_left: list[int] | None = None y_steps: int | None = None @@ -369,19 +389,6 @@ def _construct_comment(self) -> str: f"bottom right (px): [{bottom_right[0]},{bottom_right[1]}]." ) - def _store_position_table(self, dc_id: int) -> int: - assert self.ispyb_params is not None - params = self.mx_acquisition.get_dc_position_params() - - params["id"] = dc_id - ( - params["pos_x"], - params["pos_y"], - params["pos_z"], - ) = self.ispyb_params.position.tolist() - - return self.mx_acquisition.update_dc_position(list(params.values())) - class Store3DGridscanInIspyb(StoreGridscanInIspyb): def __init__(self, ispyb_config, parameters=None): From b613bd47e5b248dc7de91f4c6205241bef74299d Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 14:39:40 +0100 Subject: [PATCH 11/42] add begin and end depostion for rotation --- .../callbacks/fgs/tests/conftest.py | 2 +- .../ispyb/store_in_ispyb.py | 94 ++++++++++++------- 2 files changed, 59 insertions(+), 37 deletions(-) diff --git a/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py b/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py index ab232e8ea..63ca0bfed 100644 --- a/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py +++ b/src/artemis/external_interaction/callbacks/fgs/tests/conftest.py @@ -32,7 +32,7 @@ def mock_ispyb_store_grid_scan(): @pytest.fixture def mock_ispyb_update_time_and_status(): with patch( - "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.update_grid_scan_with_end_time_and_status" + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb.update_scan_with_end_time_and_status" ) as p: yield p diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index f6353b847..4dd6778c1 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -45,7 +45,6 @@ class StoreInIspyb(ABC): conn: Connector = None mx_acquisition: MXAcquisition | None = None core: Core | None = None - datacollection_ids: tuple[int, ...] | None = None datacollection_group_id: int | None = None def __init__(self, ispyb_config, parameters=None) -> None: @@ -97,6 +96,30 @@ def get_visit_string_from_path(self, path): match = re.search(VISIT_PATH_REGEX, path) if path else None return match.group(1) if match else None + def update_scan_with_end_time_and_status( + self, + end_time: str, + run_status: str, + reason: str, + datacollection_id: int, + datacollection_group_id: int, + ) -> None: + assert self.ispyb_params is not None + assert self.detector_params is not None + if reason is not None and reason != "": + self.append_to_comment(datacollection_id, f"{run_status} reason: {reason}") + + with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: + self.mx_acquisition: MXAcquisition = self.conn.mx_acquisition + + params = self.mx_acquisition.get_data_collection_params() + params["id"] = datacollection_id + params["parentid"] = datacollection_group_id + params["endtime"] = end_time + params["run_status"] = run_status + + self.mx_acquisition.upsert_data_collection(list(params.values())) + def _store_position_table(self, dc_id: int) -> int: assert self.ispyb_params is not None assert self.mx_acquisition is not None @@ -205,6 +228,7 @@ def _store_data_collection_table(self, data_collection_group_id: int) -> int: class StoreRotationInIspyb(StoreInIspyb): ispyb_params: RotationIspybParams | None = None + datacollection_id: int | None = None def __init__(self, ispyb_config, parameters=None) -> None: self.full_params: RotationInternalParameters | None = parameters @@ -221,12 +245,6 @@ def _mutate_datacollection_params_for_experiment( params["n_images"] = self.full_params.experiment_params.get_num_images() return params - def store_rotation_scan(self): - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition = self.conn.mx_acquisition - self.core = self.conn.core - return self._store_scan_data() - def _store_scan_data(self): data_collection_group_id = self._store_data_collection_group_table() data_collection_id = self._store_data_collection_table(data_collection_group_id) @@ -235,11 +253,37 @@ def _store_scan_data(self): return data_collection_id, data_collection_group_id - def begin_deposition(self, success: str, reason: str): - pass + def begin_deposition(self): + with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: + self.mx_acquisition = self.conn.mx_acquisition + self.core = self.conn.core + return self._store_scan_data() def end_deposition(self, success: str, reason: str): - pass + """Write the end of datacollection data. + Args: + success (str): The success of the run, could be fail or abort + reason (str): If the run failed, the reason why + """ + LOGGER.info( + f"End ispyb deposition with status '{success}' and reason '{reason}'." + ) + if success == "fail" or success == "abort": + run_status = "DataCollection Unsuccessful" + else: + run_status = "DataCollection Successful" + current_time = self.get_current_time_string() + assert ( + self.datacollection_id is not None + ), "Can't end ISPyB deposition, datacollection IDs are missing" + assert self.datacollection_group_id is not None + self.update_scan_with_end_time_and_status( + current_time, + run_status, + reason, + self.datacollection_id, + self.datacollection_group_id, + ) def _construct_comment(self) -> str: return "Hyperion rotation scan" @@ -247,7 +291,8 @@ def _construct_comment(self) -> str: class StoreGridscanInIspyb(StoreInIspyb): ispyb_params: GridscanIspybParams | None = None - supper_left: list[int] | None = None + datacollection_ids: tuple[int, ...] | None = None + upper_left: list[int] | None = None y_steps: int | None = None y_step_size: int | None = None grid_ids: tuple[int, ...] | None = None @@ -286,7 +331,7 @@ def end_deposition(self, success: str, reason: str): ), "Can't end ISPyB deposition, datacollection IDs are missing" assert self.datacollection_group_id is not None for id in self.datacollection_ids: - self.update_grid_scan_with_end_time_and_status( + self.update_scan_with_end_time_and_status( current_time, run_status, reason, id, self.datacollection_group_id ) @@ -318,34 +363,11 @@ def _mutate_datacollection_params_for_experiment( params["n_images"] = self.full_params.experiment_params.x_steps * self.y_steps return params - def update_grid_scan_with_end_time_and_status( - self, - end_time: str, - run_status: str, - reason: str, - datacollection_id: int, - datacollection_group_id: int, - ) -> None: - assert self.ispyb_params is not None - assert self.detector_params is not None - if reason is not None and reason != "": - self.append_to_comment(datacollection_id, f"{run_status} reason: {reason}") - - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition: MXAcquisition = self.conn.mx_acquisition - - params = self.mx_acquisition.get_data_collection_params() - params["id"] = datacollection_id - params["parentid"] = datacollection_group_id - params["endtime"] = end_time - params["run_status"] = run_status - - self.mx_acquisition.upsert_data_collection(list(params.values())) - def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: assert self.ispyb_params is not None assert self.full_params is not None assert self.upper_left is not None + assert self.mx_acquisition is not None params = self.mx_acquisition.get_dc_grid_params() From 05a4a04882ff4a14f28ca7995b75e13bf5b9e047 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 15:14:55 +0100 Subject: [PATCH 12/42] create rotation ispyb callback --- .../callbacks/rotation/ispyb_callback.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py diff --git a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py new file mode 100644 index 000000000..3fe9bb898 --- /dev/null +++ b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import os +from typing import Dict + +from bluesky.callbacks import CallbackBase + +from artemis.external_interaction.exceptions import ISPyBDepositionNotMade +from artemis.external_interaction.ispyb.store_in_ispyb import StoreRotationInIspyb +from artemis.log import LOGGER, set_dcgid_tag +from artemis.parameters.constants import ISPYB_PLAN_NAME, SIM_ISPYB_CONFIG +from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters + + +class RotationISPyBHandlerCallback(CallbackBase): + """Callback class to handle the deposition of experiment parameters into the ISPyB + database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on + recieving an 'event' document for the 'ispyb_readings' event, and updates the + deposition on recieving it's final 'stop' document. + + To use, subscribe the Bluesky RunEngine to an instance of this class. + E.g.: + nexus_file_handler_callback = NexusFileHandlerCallback(parameters) + RE.subscribe(nexus_file_handler_callback) + Or decorate a plan using bluesky.preprocessors.subs_decorator. + + See: https://blueskyproject.io/bluesky/callbacks.html#ways-to-invoke-callbacks + + Usually used as part of an FGSCallbackCollection. + """ + + def __init__(self, parameters: FGSInternalParameters): + self.params = parameters + self.descriptors: Dict[str, dict] = {} + ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) + if ispyb_config == SIM_ISPYB_CONFIG: + LOGGER.warning( + "Using dev ISPyB database. If you want to use the real database, please" + " set the ISPYB_CONFIG_PATH environment variable." + ) + self.ispyb: StoreRotationInIspyb = StoreRotationInIspyb( + ispyb_config, self.params + ) + + self.ispyb_id: int | None = None + self.uid_to_finalize_on: str | None = None + + def append_to_comment(self, comment: str): + try: + self.ispyb.append_to_comment(self.ispyb_id, comment) + except TypeError: + LOGGER.warning("ISPyB deposition not initialised, can't update comment.") + + def descriptor(self, doc: dict): + self.descriptors[doc["uid"]] = doc + + def start(self, doc: dict): + if self.uid_to_finalize_on is None: + self.uid_to_finalize_on = doc.get("uid") + + def event(self, doc: dict): + LOGGER.debug("ISPyB handler received event document.") + event_descriptor = self.descriptors[doc["descriptor"]] + + if event_descriptor.get("name") == ISPYB_PLAN_NAME: + self.params.artemis_params.ispyb_params.undulator_gap = doc["data"][ + "undulator_gap" + ] + self.params.artemis_params.ispyb_params.synchrotron_mode = doc["data"][ + "synchrotron_machine_status_synchrotron_mode" + ] + self.params.artemis_params.ispyb_params.slit_gap_size_x = doc["data"][ + "s4_slit_gaps_xgap" + ] + self.params.artemis_params.ispyb_params.slit_gap_size_y = doc["data"][ + "s4_slit_gaps_ygap" + ] + + LOGGER.info("Creating ispyb entry.") + self.ispyb_ids = self.ispyb.begin_deposition() + set_dcgid_tag(self.ispyb_id) + + def stop(self, doc: dict): + if doc.get("run_start") == self.uid_to_finalize_on: + LOGGER.debug("ISPyB handler received stop document.") + exit_status = doc.get("exit_status") + reason = doc.get("reason") + if self.ispyb_ids == (None, None, None): + raise ISPyBDepositionNotMade("ispyb was not initialised at run start") + self.ispyb.end_deposition(exit_status, reason) + set_dcgid_tag(None) From a509718373d6bc4715002ff3ee1bd380caf02de9 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 15:22:29 +0100 Subject: [PATCH 13/42] add ispyb reading stub to rotation scan plan --- .../read_hardware_for_setup.py | 25 +++++++++++++++++++ .../experiment_plans/fast_grid_scan_plan.py | 21 ++-------------- .../experiment_plans/rotation_scan_plan.py | 12 ++++++++- 3 files changed, 38 insertions(+), 20 deletions(-) create mode 100644 src/artemis/device_setup_plans/read_hardware_for_setup.py diff --git a/src/artemis/device_setup_plans/read_hardware_for_setup.py b/src/artemis/device_setup_plans/read_hardware_for_setup.py new file mode 100644 index 000000000..0a94846b9 --- /dev/null +++ b/src/artemis/device_setup_plans/read_hardware_for_setup.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import bluesky.plan_stubs as bps +from dodal.beamlines.i03 import S4SlitGaps, Synchrotron, Undulator + +import artemis.log +from artemis.parameters.constants import ISPYB_PLAN_NAME + + +def read_hardware_for_ispyb( + undulator: Undulator, + synchrotron: Synchrotron, + s4_slit_gaps: S4SlitGaps, +): + artemis.log.LOGGER.info( + "Reading status of beamline parameters for ispyb deposition." + ) + yield from bps.create( + name=ISPYB_PLAN_NAME + ) # gives name to event *descriptor* document + yield from bps.read(undulator.gap) + yield from bps.read(synchrotron.machine_status.synchrotron_mode) + yield from bps.read(s4_slit_gaps.xgap) + yield from bps.read(s4_slit_gaps.ygap) + yield from bps.save() diff --git a/src/artemis/experiment_plans/fast_grid_scan_plan.py b/src/artemis/experiment_plans/fast_grid_scan_plan.py index 9e632e263..1f0bfc5b6 100755 --- a/src/artemis/experiment_plans/fast_grid_scan_plan.py +++ b/src/artemis/experiment_plans/fast_grid_scan_plan.py @@ -25,6 +25,7 @@ from dodal.devices.fast_grid_scan import set_fast_grid_scan_params import artemis.log +from artemis.device_setup_plans.read_hardware_for_setup import read_hardware_for_ispyb from artemis.device_setup_plans.setup_zebra import ( set_zebra_shutter_to_manual, setup_zebra_for_fgs, @@ -35,7 +36,7 @@ get_beamline_parameters, get_beamline_prefixes, ) -from artemis.parameters.constants import ISPYB_PLAN_NAME, SIM_BEAMLINE +from artemis.parameters.constants import SIM_BEAMLINE from artemis.tracing import TRACER if TYPE_CHECKING: @@ -124,24 +125,6 @@ def set_aperture(): yield from set_aperture() -def read_hardware_for_ispyb( - undulator: Undulator, - synchrotron: Synchrotron, - s4_slit_gaps: S4SlitGaps, -): - artemis.log.LOGGER.info( - "Reading status of beamline parameters for ispyb deposition." - ) - yield from bps.create( - name=ISPYB_PLAN_NAME - ) # gives name to event *descriptor* document - yield from bps.read(undulator.gap) - yield from bps.read(synchrotron.machine_status.synchrotron_mode) - yield from bps.read(s4_slit_gaps.xgap) - yield from bps.read(s4_slit_gaps.ygap) - yield from bps.save() - - @bpp.set_run_key_decorator("move_xyz") @bpp.run_decorator(md={"subplan_name": "move_xyz"}) def move_xyz( diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index 83c63e579..3d1f61a07 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -17,6 +17,7 @@ from dodal.devices.zebra import RotationDirection, Zebra from ophyd.epics_motor import EpicsMotor +from artemis.device_setup_plans.read_hardware_for_setup import read_hardware_for_ispyb from artemis.device_setup_plans.setup_zebra import ( arm_zebra, disarm_zebra, @@ -118,7 +119,14 @@ def rotation_scan_plan( yield from setup_sample_environment(detector_motion, backlight) LOGGER.info(f"moving omega to beginning, start_angle={start_angle}") yield from move_to_start_w_buffer(smargon.omega, start_angle) - LOGGER.info("wait for any previous moves...") + + # get some information for the ispyb deposition and trigger the callback + yield from read_hardware_for_ispyb( + i03.undulator(), + i03.synchrotron(), + i03.s4_slit_gaps(), + ) + LOGGER.info( f"setting up zebra w: start_angle={start_angle}, scan_width={scan_width}" ) @@ -133,6 +141,8 @@ def rotation_scan_plan( ), group="setup_zebra", ) + + LOGGER.info("wait for any previous moves...") # wait for all the setup tasks at once yield from bps.wait("setup_senv") yield from bps.wait("move_to_start") From bd12a83676e6e4772a75c7d8de149711efa6cc44 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 16:13:58 +0100 Subject: [PATCH 14/42] tidy rotation plan and add new devices to test --- .../experiment_plans/rotation_scan_plan.py | 19 ++++++------ .../tests/test_rotation_scan_plan.py | 29 ++++++++++++++----- 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index 3d1f61a07..e995cf867 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -3,12 +3,7 @@ from typing import TYPE_CHECKING import bluesky.plan_stubs as bps -from bluesky.preprocessors import ( - finalize_decorator, - finalize_wrapper, - stage_decorator, - subs_decorator, -) +import bluesky.preprocessors as bpp from dodal.beamlines import i03 from dodal.devices.backlight import Backlight from dodal.devices.detector_motion import DetectorMotion @@ -96,6 +91,8 @@ def set_speed(axis: EpicsMotor, image_width, exposure_time, wait=True): ) +@bpp.set_run_key_decorator("rotation_scan_main") +@bpp.run_decorator(md={"subplan_name": "rotation_scan_main"}) def rotation_scan_plan( params: RotationInternalParameters, eiger: EigerDetector, @@ -161,7 +158,7 @@ def rotation_scan_plan( def cleanup_plan(eiger, zebra, smargon, detector_motion, backlight): yield from cleanup_sample_environment(zebra, detector_motion) - yield from finalize_wrapper(disarm_zebra(zebra), bps.wait("cleanup_senv")) + yield from bpp.finalize_wrapper(disarm_zebra(zebra), bps.wait("cleanup_senv")) def get_plan( @@ -180,14 +177,16 @@ def get_plan( "backlight": backlight, } - @subs_decorator(list(subscriptions)) + @bpp.subs_decorator(list(subscriptions)) def rotation_scan_plan_with_stage_and_cleanup( params: RotationInternalParameters, ): eiger.set_detector_parameters(params.artemis_params.detector_params) - @stage_decorator([eiger]) - @finalize_decorator(lambda: cleanup_plan(**devices)) + @bpp.stage_decorator([eiger]) + @bpp.set_run_key_decorator("rotation_scan_with_cleanup") + @bpp.run_decorator(md={"subplan_name": "rotation_scan_with_cleanup"}) + @bpp.finalize_decorator(lambda: cleanup_plan(**devices)) def rotation_with_cleanup_and_stage(params): yield from rotation_scan_plan(params, **devices) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index fc649458f..4c47b274a 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -5,6 +5,7 @@ import pytest from bluesky.utils import Msg +from dodal.beamlines import i03 from ophyd.status import Status from artemis.experiment_plans.rotation_scan_plan import ( @@ -110,15 +111,29 @@ def test_rotation_plan( smargon.omega.velocity.set = mock_omega_sets smargon.omega.set = mock_omega_sets - with patch( - "bluesky.preprocessors.__read_and_stash_a_motor", - __fake_read, + undulator = i03.undulator(fake_with_ophyd_sim=True) + synchrotron = i03.synchrotron(fake_with_ophyd_sim=True) + slit_gaps = i03.s4_slit_gaps(fake_with_ophyd_sim=True) + + with ( + patch("dodal.beamlines.i03.undulator", return_value=undulator), + patch("dodal.beamlines.i03.synchrotron", return_value=synchrotron), + patch("dodal.beamlines.i03.s4_slit_gaps", return_value=slit_gaps), ): - RE( - rotation_scan_plan( - test_rotation_params, eiger, smargon, zebra, backlight, detector_motion + with patch( + "bluesky.preprocessors.__read_and_stash_a_motor", + __fake_read, + ): + RE( + rotation_scan_plan( + test_rotation_params, + eiger, + smargon, + zebra, + backlight, + detector_motion, + ) ) - ) # once for each velocity set and once for each position set for a total of 4 calls assert mock_omega_sets.call_count == 4 From 2f049298ea2a05f09fe8bca8d9a560855663a8af Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 6 Jun 2023 17:07:26 +0100 Subject: [PATCH 15/42] add system test for ispyb deposition in plan --- .../experiment_plans/rotation_scan_plan.py | 35 +++++++---------- .../tests/test_rotation_scan_plan.py | 39 +++++++++++++++++++ .../rotation/rotation_callback_collection.py | 7 +++- .../ispyb/store_in_ispyb.py | 7 ++-- 4 files changed, 64 insertions(+), 24 deletions(-) diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index e995cf867..d8fd3dc01 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -10,6 +10,7 @@ from dodal.devices.eiger import DetectorParams, EigerDetector from dodal.devices.smargon import Smargon from dodal.devices.zebra import RotationDirection, Zebra +from ophyd.device import Device from ophyd.epics_motor import EpicsMotor from artemis.device_setup_plans.read_hardware_for_setup import read_hardware_for_ispyb @@ -32,12 +33,17 @@ ) -def create_devices(): - i03.eiger(wait_for_connection=False) - i03.smargon() - i03.zebra() - i03.detector_motion() - i03.backlight() +def create_devices() -> dict[str, Device]: + """Ensures necessary devices have been instantiated and returns a dict with + references to them""" + devices = { + "eiger": i03.eiger(wait_for_connection=False), + "smargon": i03.smargon(), + "zebra": i03.zebra(), + "detector_motion": i03.detector_motion(), + "backlight": i03.backlight(), + } + return devices DIRECTION = RotationDirection.NEGATIVE @@ -164,26 +170,15 @@ def cleanup_plan(eiger, zebra, smargon, detector_motion, backlight): def get_plan( params: RotationInternalParameters, subscriptions: RotationCallbackCollection ): - eiger = i03.eiger(wait_for_connection=False) - smargon = i03.smargon() - zebra = i03.zebra() - detector_motion = i03.detector_motion() - backlight = i03.backlight() - devices = { - "eiger": eiger, - "smargon": smargon, - "zebra": zebra, - "detector_motion": detector_motion, - "backlight": backlight, - } + devices = create_devices() @bpp.subs_decorator(list(subscriptions)) def rotation_scan_plan_with_stage_and_cleanup( params: RotationInternalParameters, ): - eiger.set_detector_parameters(params.artemis_params.detector_params) + devices["eiger"].set_detector_parameters(params.artemis_params.detector_params) - @bpp.stage_decorator([eiger]) + @bpp.stage_decorator([devices["eiger"]]) @bpp.set_run_key_decorator("rotation_scan_with_cleanup") @bpp.run_decorator(md={"subplan_name": "rotation_scan_with_cleanup"}) @bpp.finalize_decorator(lambda: cleanup_plan(**devices)) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 4c47b274a..de3e54871 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -182,3 +182,42 @@ def test_cleanup_happens( ) ) cleanup_plan.assert_called_once() + + +@pytest.mark.s03() +@patch("bluesky.plan_stubs.wait") +def test_ispyb_deposition_in_plan( + bps_wait: MagicMock, + RE, + test_rotation_params, +): + def fake_create_devices(): + devices = { + "eiger": i03.eiger(wait_for_connection=False, fake_with_ophyd_sim=True), + "smargon": i03.smargon(), + "zebra": i03.zebra(), + "detector_motion": i03.detector_motion(fake_with_ophyd_sim=True), + "backlight": i03.backlight(fake_with_ophyd_sim=True), + } + return devices + + undulator = i03.undulator(fake_with_ophyd_sim=True) + synchrotron = i03.synchrotron(fake_with_ophyd_sim=True) + slit_gaps = i03.s4_slit_gaps(fake_with_ophyd_sim=True) + + with ( + patch( + "artemis.experiment_plans.rotation_scan_plan.create_devices", + fake_create_devices, + ), + patch("dodal.beamlines.i03.undulator", return_value=undulator), + patch("dodal.beamlines.i03.synchrotron", return_value=synchrotron), + patch("dodal.beamlines.i03.s4_slit_gaps", return_value=slit_gaps), + patch("dodal.devices.eiger.EigerDetector.stage"), + ): + RE( + get_plan( + test_rotation_params, + RotationCallbackCollection.from_params(test_rotation_params), + ) + ) diff --git a/src/artemis/external_interaction/callbacks/rotation/rotation_callback_collection.py b/src/artemis/external_interaction/callbacks/rotation/rotation_callback_collection.py index 3ad4b26f6..59a36a35f 100644 --- a/src/artemis/external_interaction/callbacks/rotation/rotation_callback_collection.py +++ b/src/artemis/external_interaction/callbacks/rotation/rotation_callback_collection.py @@ -6,6 +6,9 @@ from artemis.external_interaction.callbacks.abstract_plan_callback_collection import ( AbstractPlanCallbackCollection, ) +from artemis.external_interaction.callbacks.rotation.ispyb_callback import ( + RotationISPyBHandlerCallback, +) from artemis.external_interaction.callbacks.rotation.nexus_callback import ( RotationNexusFileHandlerCallback, ) @@ -21,11 +24,13 @@ class RotationCallbackCollection(AbstractPlanCallbackCollection): Bluesky.preprocessors.subs_decorator().""" nexus_handler: RotationNexusFileHandlerCallback + ispyb_handler: RotationISPyBHandlerCallback @classmethod def from_params(cls, parameters: InternalParameters): nexus_handler = RotationNexusFileHandlerCallback(parameters) + ispyb_handler = RotationISPyBHandlerCallback(parameters) callback_collection = cls( - nexus_handler=nexus_handler, + nexus_handler=nexus_handler, ispyb_handler=ispyb_handler ) return callback_collection diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 4dd6778c1..50657ab00 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -227,11 +227,12 @@ def _store_data_collection_table(self, data_collection_group_id: int) -> int: class StoreRotationInIspyb(StoreInIspyb): - ispyb_params: RotationIspybParams | None = None + ispyb_params: RotationIspybParams datacollection_id: int | None = None - def __init__(self, ispyb_config, parameters=None) -> None: - self.full_params: RotationInternalParameters | None = parameters + def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None: + self.full_params: RotationInternalParameters = parameters + self.ispyb_params = self.full_params.artemis_params.ispyb_params self.experiment_type = "SAD" super().__init__(ispyb_config, parameters) From 6a0fc58239e876d7d4ef40e80edf2c4b989d56fe Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 7 Jun 2023 09:00:46 +0100 Subject: [PATCH 16/42] remove self.conn/acq/core and instead pass references to active connection --- .../tests/test_rotation_scan_plan.py | 9 +- .../ispyb/store_in_ispyb.py | 111 +++++++++--------- 2 files changed, 61 insertions(+), 59 deletions(-) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index de3e54871..68cdce94d 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -6,7 +6,7 @@ import pytest from bluesky.utils import Msg from dodal.beamlines import i03 -from ophyd.status import Status +from ophyd.status import Status, SubscriptionStatus from artemis.experiment_plans.rotation_scan_plan import ( DIRECTION, @@ -192,8 +192,13 @@ def test_ispyb_deposition_in_plan( test_rotation_params, ): def fake_create_devices(): + with patch( + "dodal.devices.eiger_odin.EigerOdin.create_finished_status", + return_value=lambda: True, + ): + eiger = i03.eiger(wait_for_connection=False, fake_with_ophyd_sim=True) devices = { - "eiger": i03.eiger(wait_for_connection=False, fake_with_ophyd_sim=True), + "eiger": eiger, "smargon": i03.smargon(), "zebra": i03.zebra(), "detector_motion": i03.detector_motion(fake_with_ophyd_sim=True), diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 50657ab00..9508bf836 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -9,9 +9,9 @@ import ispyb import ispyb.sqlalchemy from dodal.devices.detector import DetectorParams +from ispyb.connector.mysqlsp.main import ISPyBMySQLSPConnector as Connector from ispyb.sp.core import Core from ispyb.sp.mxacquisition import MXAcquisition -from sqlalchemy.connectors import Connector from artemis.external_interaction.ispyb.ispyb_dataclass import ( GridscanIspybParams, @@ -42,16 +42,13 @@ class StoreInIspyb(ABC): omega_start: int | None = None experiment_type: str | None = None xtal_snapshots: list[str] | None = None - conn: Connector = None - mx_acquisition: MXAcquisition | None = None - core: Core | None = None datacollection_group_id: int | None = None def __init__(self, ispyb_config, parameters=None) -> None: self.ISPYB_CONFIG_PATH: str = ispyb_config @abstractmethod - def _store_scan_data(self): + def _store_scan_data(self, conn: Connector): pass @abstractmethod @@ -75,9 +72,9 @@ def end_deposition(self, success: str, reason: str): def append_to_comment( self, datacollection_id: int, comment: str, delimiter: str = " " ) -> None: - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition: MXAcquisition = self.conn.mx_acquisition - self.mx_acquisition.update_data_collection_append_comments( + with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + mx_acquisition: MXAcquisition = conn.mx_acquisition + mx_acquisition.update_data_collection_append_comments( datacollection_id, comment, delimiter ) @@ -109,22 +106,21 @@ def update_scan_with_end_time_and_status( if reason is not None and reason != "": self.append_to_comment(datacollection_id, f"{run_status} reason: {reason}") - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition: MXAcquisition = self.conn.mx_acquisition + with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + mx_acquisition: MXAcquisition = conn.mx_acquisition - params = self.mx_acquisition.get_data_collection_params() + params = mx_acquisition.get_data_collection_params() params["id"] = datacollection_id params["parentid"] = datacollection_group_id params["endtime"] = end_time params["run_status"] = run_status - self.mx_acquisition.upsert_data_collection(list(params.values())) + mx_acquisition.upsert_data_collection(list(params.values())) - def _store_position_table(self, dc_id: int) -> int: + def _store_position_table(self, conn: Connector, dc_id: int) -> int: assert self.ispyb_params is not None - assert self.mx_acquisition is not None - - params = self.mx_acquisition.get_dc_position_params() + mx_acquisition: MXAcquisition = conn.mx_acquisition + params = mx_acquisition.get_dc_position_params() params["id"] = dc_id ( @@ -133,43 +129,46 @@ def _store_position_table(self, dc_id: int) -> int: params["pos_z"], ) = self.ispyb_params.position.tolist() - return self.mx_acquisition.update_dc_position(list(params.values())) + return mx_acquisition.update_dc_position(list(params.values())) - def _store_data_collection_group_table(self) -> int: - assert self.core is not None + def _store_data_collection_group_table(self, conn: Connector) -> int: assert self.ispyb_params is not None - assert self.mx_acquisition is not None + core: Core = conn.core + mx_acquisition: MXAcquisition = conn.mx_acquisition + try: - session_id = self.core.retrieve_visit_id(self.get_visit_string()) + session_id = core.retrieve_visit_id(self.get_visit_string()) except ispyb.NoResult: raise Exception( f"Not found - session ID for visit {self.get_visit_string()} where self.ispyb_params.visit_path is {self.ispyb_params.visit_path}" ) - params = self.mx_acquisition.get_data_collection_group_params() + params = mx_acquisition.get_data_collection_group_params() params["parentid"] = session_id params["experimenttype"] = self.experiment_type params["sampleid"] = self.ispyb_params.sample_id params["sample_barcode"] = self.ispyb_params.sample_barcode - return self.mx_acquisition.upsert_data_collection_group(list(params.values())) + return mx_acquisition.upsert_data_collection_group(list(params.values())) @TRACER.start_as_current_span("store_ispyb_datacollection_table") - def _store_data_collection_table(self, data_collection_group_id: int) -> int: + def _store_data_collection_table( + self, conn: Connector, data_collection_group_id: int + ) -> int: assert self.ispyb_params is not None assert self.detector_params is not None - assert self.core is not None assert self.xtal_snapshots is not None - assert self.mx_acquisition is not None + + core: Core = conn.core + mx_acquisition: MXAcquisition = conn.mx_acquisition try: - session_id = self.core.retrieve_visit_id(self.get_visit_string()) + session_id = core.retrieve_visit_id(self.get_visit_string()) except ispyb.NoResult: raise Exception( f"Not found - session ID for visit {self.get_visit_string()}" ) - params = self.mx_acquisition.get_data_collection_params() - + params = mx_acquisition.get_data_collection_params() params = self._mutate_datacollection_params_for_experiment(params) params["visitid"] = session_id @@ -177,7 +176,6 @@ def _store_data_collection_table(self, data_collection_group_id: int) -> int: params["sampleid"] = self.ispyb_params.sample_id params["detectorid"] = I03_EIGER_DETECTOR params["axis_start"] = self.omega_start - params["axis_range"] = 0 params["focal_spot_size_at_samplex"] = self.ispyb_params.focal_spot_size_x params["focal_spot_size_at_sampley"] = self.ispyb_params.focal_spot_size_y @@ -223,7 +221,7 @@ def _store_data_collection_table(self, data_collection_group_id: int) -> int: "file_template" ] = f"{self.detector_params.prefix}_{self.run_number}_master.h5" - return self.mx_acquisition.upsert_data_collection(list(params.values())) + return mx_acquisition.upsert_data_collection(list(params.values())) class StoreRotationInIspyb(StoreInIspyb): @@ -255,10 +253,8 @@ def _store_scan_data(self): return data_collection_id, data_collection_group_id def begin_deposition(self): - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition = self.conn.mx_acquisition - self.core = self.conn.core - return self._store_scan_data() + with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + return self._store_scan_data(conn) def end_deposition(self, success: str, reason: str): """Write the end of datacollection data. @@ -350,11 +346,8 @@ def store_grid_scan(self, full_params: FGSInternalParameters): self.y_steps = full_params.experiment_params.y_steps self.y_step_size = full_params.experiment_params.y_step_size - with ispyb.open(self.ISPYB_CONFIG_PATH) as self.conn: - self.mx_acquisition = self.conn.mx_acquisition - self.core = self.conn.core - - return self._store_scan_data() + with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + return self._store_scan_data(conn) def _mutate_datacollection_params_for_experiment( self, params: dict[str, Any] @@ -364,13 +357,15 @@ def _mutate_datacollection_params_for_experiment( params["n_images"] = self.full_params.experiment_params.x_steps * self.y_steps return params - def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: + def _store_grid_info_table( + self, conn: Connector, ispyb_data_collection_id: int + ) -> int: assert self.ispyb_params is not None assert self.full_params is not None assert self.upper_left is not None - assert self.mx_acquisition is not None - params = self.mx_acquisition.get_dc_grid_params() + mx_acquisition: MXAcquisition = conn.mx_acquisition + params = mx_acquisition.get_dc_grid_params() params["parentid"] = ispyb_data_collection_id params["dxInMm"] = self.full_params.experiment_params.x_step_size @@ -385,7 +380,7 @@ def _store_grid_info_table(self, ispyb_data_collection_id: int) -> int: params["orientation"] = Orientation.HORIZONTAL.value params["snaked"] = True - return self.mx_acquisition.upsert_dc_grid(list(params.values())) + return mx_acquisition.upsert_dc_grid(list(params.values())) def _construct_comment(self) -> str: assert self.ispyb_params is not None @@ -418,26 +413,26 @@ def __init__(self, ispyb_config, parameters=None): super().__init__(ispyb_config, parameters) self.experiment_type = "Mesh3D" - def _store_scan_data(self): - data_collection_group_id = self._store_data_collection_group_table() + def _store_scan_data(self, conn: Connector): + data_collection_group_id = self._store_data_collection_group_table(conn) data_collection_id_1 = self._store_data_collection_table( - data_collection_group_id + conn, data_collection_group_id ) - self._store_position_table(data_collection_id_1) + self._store_position_table(conn, data_collection_id_1) - grid_id_1 = self._store_grid_info_table(data_collection_id_1) + grid_id_1 = self._store_grid_info_table(conn, data_collection_id_1) self.__prepare_second_scan_params() data_collection_id_2 = self._store_data_collection_table( - data_collection_group_id + conn, data_collection_group_id ) - self._store_position_table(data_collection_id_2) + self._store_position_table(conn, data_collection_id_2) - grid_id_2 = self._store_grid_info_table(data_collection_id_2) + grid_id_2 = self._store_grid_info_table(conn, data_collection_id_2) return ( [data_collection_id_1, data_collection_id_2], @@ -462,13 +457,15 @@ def __init__(self, ispyb_config, parameters=None): super().__init__(ispyb_config, parameters) self.experiment_type = "mesh" - def _store_scan_data(self): - data_collection_group_id = self._store_data_collection_group_table() + def _store_scan_data(self, conn: Connector): + data_collection_group_id = self._store_data_collection_group_table(conn) - data_collection_id = self._store_data_collection_table(data_collection_group_id) + data_collection_id = self._store_data_collection_table( + conn, data_collection_group_id + ) - self._store_position_table(data_collection_id) + self._store_position_table(conn, data_collection_id) - grid_id = self._store_grid_info_table(data_collection_id) + grid_id = self._store_grid_info_table(conn, data_collection_id) return [data_collection_id], [grid_id], data_collection_group_id From 20960dfeb2d74560ffec941041e45d3c4d792a52 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 7 Jun 2023 10:14:34 +0100 Subject: [PATCH 17/42] add test for making ispyb deposition in plan --- .../experiment_plans/rotation_scan_plan.py | 3 +- .../tests/test_rotation_scan_plan.py | 68 +++++++++++++------ .../callbacks/rotation/ispyb_callback.py | 12 ++-- .../ispyb/store_in_ispyb.py | 46 ++++++++----- 4 files changed, 85 insertions(+), 44 deletions(-) diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index d8fd3dc01..ed77d1e83 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -92,8 +92,9 @@ def move_to_end_w_buffer(axis: EpicsMotor, scan_width: float, wait: float = True def set_speed(axis: EpicsMotor, image_width, exposure_time, wait=True): + speed_for_rotation = image_width / exposure_time yield from bps.abs_set( - axis.velocity, image_width / exposure_time, group="set_speed", wait=True + axis.velocity, speed_for_rotation, group="set_speed", wait=wait ) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 68cdce94d..eb6555704 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -6,7 +6,7 @@ import pytest from bluesky.utils import Msg from dodal.beamlines import i03 -from ophyd.status import Status, SubscriptionStatus +from ophyd.status import Status from artemis.experiment_plans.rotation_scan_plan import ( DIRECTION, @@ -19,6 +19,9 @@ from artemis.external_interaction.callbacks.rotation.rotation_callback_collection import ( RotationCallbackCollection, ) +from artemis.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, +) if TYPE_CHECKING: from dodal.devices.backlight import Backlight @@ -184,45 +187,68 @@ def test_cleanup_happens( cleanup_plan.assert_called_once() +@pytest.fixture() +def fake_create_devices( + eiger: EigerDetector, + smargon: Smargon, + zebra: Zebra, +): + eiger.stage = MagicMock() + eiger.unstage = MagicMock() + mock_omega_sets = MagicMock(return_value=Status(done=True, success=True)) + mock_arm_disarm = MagicMock( + side_effect=zebra.pc.armed.set, return_value=Status(done=True, success=True) + ) + zebra.pc.arm_demand.set = mock_arm_disarm + smargon.omega.velocity.set = mock_omega_sets + smargon.omega.set = mock_omega_sets + zebra.pc.arm_demand.set = mock_arm_disarm + + devices = { + "eiger": i03.eiger(wait_for_connection=False, fake_with_ophyd_sim=True), + "smargon": smargon, + "zebra": zebra, + "detector_motion": i03.detector_motion(fake_with_ophyd_sim=True), + "backlight": i03.backlight(fake_with_ophyd_sim=True), + } + return devices + + @pytest.mark.s03() @patch("bluesky.plan_stubs.wait") def test_ispyb_deposition_in_plan( - bps_wait: MagicMock, + bps_wait, + fake_create_devices, RE, - test_rotation_params, + test_rotation_params: RotationInternalParameters, ): - def fake_create_devices(): - with patch( - "dodal.devices.eiger_odin.EigerOdin.create_finished_status", - return_value=lambda: True, - ): - eiger = i03.eiger(wait_for_connection=False, fake_with_ophyd_sim=True) - devices = { - "eiger": eiger, - "smargon": i03.smargon(), - "zebra": i03.zebra(), - "detector_motion": i03.detector_motion(fake_with_ophyd_sim=True), - "backlight": i03.backlight(fake_with_ophyd_sim=True), - } - return devices - undulator = i03.undulator(fake_with_ophyd_sim=True) synchrotron = i03.synchrotron(fake_with_ophyd_sim=True) slit_gaps = i03.s4_slit_gaps(fake_with_ophyd_sim=True) + test_rotation_params.experiment_params.image_width = 0.27 + test_rotation_params.artemis_params.ispyb_params.beam_size_x = 23 + test_rotation_params.artemis_params.ispyb_params.beam_size_y = 47 + test_rotation_params.artemis_params.detector_params.exposure_time = 0.023 + test_rotation_params.artemis_params.ispyb_params.wavelength = 0.71 + callbacks = RotationCallbackCollection.from_params(test_rotation_params) + with ( patch( "artemis.experiment_plans.rotation_scan_plan.create_devices", - fake_create_devices, + lambda: fake_create_devices, ), patch("dodal.beamlines.i03.undulator", return_value=undulator), patch("dodal.beamlines.i03.synchrotron", return_value=synchrotron), patch("dodal.beamlines.i03.s4_slit_gaps", return_value=slit_gaps), - patch("dodal.devices.eiger.EigerDetector.stage"), + patch( + "bluesky.preprocessors.__read_and_stash_a_motor", + __fake_read, + ), ): RE( get_plan( test_rotation_params, - RotationCallbackCollection.from_params(test_rotation_params), + callbacks, ) ) diff --git a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py index 3fe9bb898..dc693765d 100644 --- a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py @@ -42,12 +42,12 @@ def __init__(self, parameters: FGSInternalParameters): ispyb_config, self.params ) - self.ispyb_id: int | None = None + self.ispyb_ids: tuple[int, int] | tuple[None, None] = (None, None) self.uid_to_finalize_on: str | None = None def append_to_comment(self, comment: str): try: - self.ispyb.append_to_comment(self.ispyb_id, comment) + self.ispyb.append_to_comment(self.ispyb_ids[0], comment) except TypeError: LOGGER.warning("ISPyB deposition not initialised, can't update comment.") @@ -78,14 +78,16 @@ def event(self, doc: dict): LOGGER.info("Creating ispyb entry.") self.ispyb_ids = self.ispyb.begin_deposition() - set_dcgid_tag(self.ispyb_id) + set_dcgid_tag(self.ispyb_ids[1]) def stop(self, doc: dict): if doc.get("run_start") == self.uid_to_finalize_on: LOGGER.debug("ISPyB handler received stop document.") exit_status = doc.get("exit_status") reason = doc.get("reason") - if self.ispyb_ids == (None, None, None): - raise ISPyBDepositionNotMade("ispyb was not initialised at run start") + if self.ispyb_ids == (None, None): + raise ISPyBDepositionNotMade( + "ISPyB deposition was not initialised at run start!" + ) self.ispyb.end_deposition(exit_status, reason) set_dcgid_tag(None) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 9508bf836..26dfe56fc 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -42,7 +42,7 @@ class StoreInIspyb(ABC): omega_start: int | None = None experiment_type: str | None = None xtal_snapshots: list[str] | None = None - datacollection_group_id: int | None = None + data_collection_group_id: int | None = None def __init__(self, ispyb_config, parameters=None) -> None: self.ISPYB_CONFIG_PATH: str = ispyb_config @@ -226,11 +226,19 @@ def _store_data_collection_table( class StoreRotationInIspyb(StoreInIspyb): ispyb_params: RotationIspybParams - datacollection_id: int | None = None + data_collection_id: int | None = None def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None: self.full_params: RotationInternalParameters = parameters - self.ispyb_params = self.full_params.artemis_params.ispyb_params + self.ispyb_params: RotationInternalParameters = ( + parameters.artemis_params.ispyb_params + ) + self.detector_params = parameters.artemis_params.detector_params + if self.ispyb_params.xtal_snapshots_omega_start: + self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_start + else: + self.xtal_snapshots = [] + LOGGER.warning("No xtal snapshot paths sent to ISPyB!") self.experiment_type = "SAD" super().__init__(ispyb_config, parameters) @@ -239,16 +247,20 @@ def _mutate_datacollection_params_for_experiment( ) -> dict[str, Any]: assert self.full_params is not None params["axis_end"] = ( - self.omega_start + self.full_params.experiment_params.rotation_angle + self.full_params.experiment_params.omega_start + + self.full_params.experiment_params.rotation_angle ) params["n_images"] = self.full_params.experiment_params.get_num_images() return params - def _store_scan_data(self): - data_collection_group_id = self._store_data_collection_group_table() - data_collection_id = self._store_data_collection_table(data_collection_group_id) - - self._store_position_table(data_collection_id) + def _store_scan_data(self, conn: Connector): + data_collection_group_id = self._store_data_collection_group_table(conn) + self.data_collection_group_id = data_collection_group_id + data_collection_id = self._store_data_collection_table( + conn, data_collection_group_id + ) + self.data_collection_id = data_collection_id + self._store_position_table(conn, data_collection_id) return data_collection_id, data_collection_group_id @@ -271,15 +283,15 @@ def end_deposition(self, success: str, reason: str): run_status = "DataCollection Successful" current_time = self.get_current_time_string() assert ( - self.datacollection_id is not None + self.data_collection_id is not None ), "Can't end ISPyB deposition, datacollection IDs are missing" - assert self.datacollection_group_id is not None + assert self.data_collection_group_id is not None self.update_scan_with_end_time_and_status( current_time, run_status, reason, - self.datacollection_id, - self.datacollection_group_id, + self.data_collection_id, + self.data_collection_group_id, ) def _construct_comment(self) -> str: @@ -302,9 +314,9 @@ def begin_deposition(self): ( self.datacollection_ids, self.grid_ids, - self.datacollection_group_id, + self.data_collection_group_id, ) = self.store_grid_scan(self.full_params) - return self.datacollection_ids, self.grid_ids, self.datacollection_group_id + return self.datacollection_ids, self.grid_ids, self.data_collection_group_id def end_deposition(self, success: str, reason: str): """Write the end of datacollection data. @@ -326,10 +338,10 @@ def end_deposition(self, success: str, reason: str): assert ( self.datacollection_ids is not None ), "Can't end ISPyB deposition, datacollection IDs are missing" - assert self.datacollection_group_id is not None + assert self.data_collection_group_id is not None for id in self.datacollection_ids: self.update_scan_with_end_time_and_status( - current_time, run_status, reason, id, self.datacollection_group_id + current_time, run_status, reason, id, self.data_collection_group_id ) def store_grid_scan(self, full_params: FGSInternalParameters): From 997d302c8d531b4d3d44d2cdbfd4652100130f74 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 7 Jun 2023 10:58:04 +0100 Subject: [PATCH 18/42] test other deposition attributes --- .../tests/test_rotation_scan_plan.py | 35 ++++++++++++++++--- .../ispyb/store_in_ispyb.py | 5 ++- .../system_tests/conftest.py | 26 ++++++++++++++ 3 files changed, 60 insertions(+), 6 deletions(-) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index eb6555704..e507ba424 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -19,6 +19,10 @@ from artemis.external_interaction.callbacks.rotation.rotation_callback_collection import ( RotationCallbackCollection, ) +from artemis.external_interaction.system_tests.conftest import ( # noqa + fetch_comment, + fetch_datacollection_attribute, +) from artemis.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, ) @@ -221,16 +225,24 @@ def test_ispyb_deposition_in_plan( fake_create_devices, RE, test_rotation_params: RotationInternalParameters, + fetch_comment, + fetch_datacollection_attribute, ): undulator = i03.undulator(fake_with_ophyd_sim=True) synchrotron = i03.synchrotron(fake_with_ophyd_sim=True) slit_gaps = i03.s4_slit_gaps(fake_with_ophyd_sim=True) - test_rotation_params.experiment_params.image_width = 0.27 - test_rotation_params.artemis_params.ispyb_params.beam_size_x = 23 - test_rotation_params.artemis_params.ispyb_params.beam_size_y = 47 - test_rotation_params.artemis_params.detector_params.exposure_time = 0.023 - test_rotation_params.artemis_params.ispyb_params.wavelength = 0.71 + test_wl = 0.71 + test_bs_x = 0.023 + test_bs_y = 0.047 + test_exp_time = 0.023 + test_img_wid = 0.27 + + test_rotation_params.experiment_params.image_width = test_img_wid + test_rotation_params.artemis_params.ispyb_params.beam_size_x = test_bs_x + test_rotation_params.artemis_params.ispyb_params.beam_size_y = test_bs_y + test_rotation_params.artemis_params.detector_params.exposure_time = test_exp_time + test_rotation_params.artemis_params.ispyb_params.wavelength = test_wl callbacks = RotationCallbackCollection.from_params(test_rotation_params) with ( @@ -252,3 +264,16 @@ def test_ispyb_deposition_in_plan( callbacks, ) ) + + dcid = callbacks.ispyb_handler.ispyb_ids[0] + comment = fetch_comment(dcid) + assert comment == "Hyperion rotation scan" + wavelength = fetch_datacollection_attribute(dcid, "wavelength") + beamsize_x = fetch_datacollection_attribute(dcid, "beamSizeAtSampleX") + beamsize_y = fetch_datacollection_attribute(dcid, "beamSizeAtSampleY") + exposure = fetch_datacollection_attribute(dcid, "exposureTime") + + assert wavelength == test_wl + assert beamsize_x == test_bs_x + assert beamsize_y == test_bs_y + assert exposure == test_exp_time diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index 26dfe56fc..c427cfa71 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -176,7 +176,6 @@ def _store_data_collection_table( params["sampleid"] = self.ispyb_params.sample_id params["detectorid"] = I03_EIGER_DETECTOR params["axis_start"] = self.omega_start - params["axis_range"] = 0 params["focal_spot_size_at_samplex"] = self.ispyb_params.focal_spot_size_x params["focal_spot_size_at_sampley"] = self.ispyb_params.focal_spot_size_y params["slitgap_vertical"] = self.ispyb_params.slit_gap_size_y @@ -234,6 +233,8 @@ def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None parameters.artemis_params.ispyb_params ) self.detector_params = parameters.artemis_params.detector_params + self.omega_start = self.detector_params.omega_start + if self.ispyb_params.xtal_snapshots_omega_start: self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_start else: @@ -246,6 +247,7 @@ def _mutate_datacollection_params_for_experiment( self, params: dict[str, Any] ) -> dict[str, Any]: assert self.full_params is not None + params["axis_range"] = self.full_params.experiment_params.rotation_angle params["axis_end"] = ( self.full_params.experiment_params.omega_start + self.full_params.experiment_params.rotation_angle @@ -365,6 +367,7 @@ def _mutate_datacollection_params_for_experiment( self, params: dict[str, Any] ) -> dict[str, Any]: assert self.full_params is not None + params["axis_range"] = 0 params["axis_end"] = self.omega_start params["n_images"] = self.full_params.experiment_params.x_steps * self.y_steps return params diff --git a/src/artemis/external_interaction/system_tests/conftest.py b/src/artemis/external_interaction/system_tests/conftest.py index a7e283f6d..27b0052fb 100644 --- a/src/artemis/external_interaction/system_tests/conftest.py +++ b/src/artemis/external_interaction/system_tests/conftest.py @@ -66,6 +66,24 @@ def get_current_datacollection_comment(Session: Callable, dcid: int) -> str: return current_comment +def get_current_datacollection_attribute( + Session: Callable, dcid: int, attr: str +) -> str: + """Read the specified field 'attr' from the given datacollection id's ISPyB entry. + Returns an empty string if the attribute is not found. + """ + try: + with Session() as session: + query = session.query(DataCollection).filter( + DataCollection.dataCollectionId == dcid + ) + first_result = query.first() + data: str = getattr(first_result, attr) + except Exception: + data = "" + return data + + @pytest.fixture def fetch_comment() -> Callable: url = ispyb.sqlalchemy.url(ISPYB_CONFIG) @@ -74,6 +92,14 @@ def fetch_comment() -> Callable: return partial(get_current_datacollection_comment, Session) +@pytest.fixture +def fetch_datacollection_attribute() -> Callable: + url = ispyb.sqlalchemy.url(ISPYB_CONFIG) + engine = create_engine(url, connect_args={"use_pure": True}) + Session = sessionmaker(engine) + return partial(get_current_datacollection_attribute, Session) + + @pytest.fixture def dummy_params(): dummy_params = FGSInternalParameters(**default_raw_params()) From 8421bf32b1bb85ce65a99dcd9a93dfec02ea2753 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 7 Jun 2023 11:13:06 +0100 Subject: [PATCH 19/42] fix ispyb unit tests --- .../unit_tests/test_store_in_ispyb.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index 4edc9c7e6..2fbc90987 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -73,18 +73,18 @@ def test_store_grid_scan(ispyb_conn, dummy_ispyb, dummy_params): ispyb_conn.return_value.mx_acquisition = mock() ispyb_conn.return_value.core = mock() - when(dummy_ispyb)._store_position_table(TEST_DATA_COLLECTION_IDS[0]).thenReturn( - TEST_POSITION_ID - ) - when(dummy_ispyb)._store_data_collection_group_table().thenReturn( + when(dummy_ispyb)._store_position_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[0] + ).thenReturn(TEST_POSITION_ID) + when(dummy_ispyb)._store_data_collection_group_table(ispyb_conn()).thenReturn( TEST_DATA_COLLECTION_GROUP_ID ) when(dummy_ispyb)._store_data_collection_table( - TEST_DATA_COLLECTION_GROUP_ID + ispyb_conn(), TEST_DATA_COLLECTION_GROUP_ID ).thenReturn(TEST_DATA_COLLECTION_IDS[0]) - when(dummy_ispyb)._store_grid_info_table(TEST_DATA_COLLECTION_IDS[0]).thenReturn( - TEST_GRID_INFO_ID - ) + when(dummy_ispyb)._store_grid_info_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[0] + ).thenReturn(TEST_GRID_INFO_ID) assert dummy_ispyb.experiment_type == "mesh" @@ -104,22 +104,22 @@ def test_store_3d_grid_scan( ispyb_conn.return_value.mx_acquisition = mock() ispyb_conn.return_value.core = mock() - when(dummy_ispyb_3d)._store_position_table(TEST_DATA_COLLECTION_IDS[0]).thenReturn( - TEST_POSITION_ID - ) - when(dummy_ispyb_3d)._store_position_table(TEST_DATA_COLLECTION_IDS[1]).thenReturn( - TEST_POSITION_ID - ) - when(dummy_ispyb_3d)._store_data_collection_group_table().thenReturn( + when(dummy_ispyb_3d)._store_position_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[0] + ).thenReturn(TEST_POSITION_ID) + when(dummy_ispyb_3d)._store_position_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[1] + ).thenReturn(TEST_POSITION_ID) + when(dummy_ispyb_3d)._store_data_collection_group_table(ispyb_conn()).thenReturn( TEST_DATA_COLLECTION_GROUP_ID ) - when(dummy_ispyb_3d)._store_grid_info_table(TEST_DATA_COLLECTION_IDS[0]).thenReturn( - TEST_GRID_INFO_ID - ) - when(dummy_ispyb_3d)._store_grid_info_table(TEST_DATA_COLLECTION_IDS[1]).thenReturn( - TEST_GRID_INFO_ID - ) + when(dummy_ispyb_3d)._store_grid_info_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[0] + ).thenReturn(TEST_GRID_INFO_ID) + when(dummy_ispyb_3d)._store_grid_info_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[1] + ).thenReturn(TEST_GRID_INFO_ID) dummy_ispyb_3d._store_data_collection_table = Mock( side_effect=TEST_DATA_COLLECTION_IDS From fc58c6db98e0c992c77c7665e4d64884ec236084 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 7 Jun 2023 11:32:45 +0100 Subject: [PATCH 20/42] add unit test --- .../unit_tests/test_store_in_ispyb.py | 51 +++++++++++++++++-- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index 2fbc90987..f7b6b718c 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -9,10 +9,14 @@ from artemis.external_interaction.ispyb.store_in_ispyb import ( Store2DGridscanInIspyb, Store3DGridscanInIspyb, + StoreRotationInIspyb, ) from artemis.parameters.constants import SIM_ISPYB_CONFIG from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters +from artemis.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, +) TEST_DATA_COLLECTION_IDS = [12, 13] TEST_DATA_COLLECTION_GROUP_ID = 34 @@ -32,19 +36,31 @@ def dummy_params(): return dummy_params +@pytest.fixture +def dummy_rotation_params(): + dummy_params = RotationInternalParameters( + **default_raw_params( + "src/artemis/parameters/tests/test_data/good_test_rotation_scan_parameters.json" + ) + ) + return dummy_params + + @pytest.fixture def dummy_ispyb(dummy_params): store_in_ispyb_2d = Store2DGridscanInIspyb(SIM_ISPYB_CONFIG, dummy_params) - store_in_ispyb_2d.get_current_datacollection_comment = MagicMock() - store_in_ispyb_2d.get_current_datacollection_comment.return_value = "" return store_in_ispyb_2d +@pytest.fixture +def dummy_rotation_ispyb(dummy_rotation_params): + store_in_ispyb = StoreRotationInIspyb(SIM_ISPYB_CONFIG, dummy_rotation_params) + return store_in_ispyb + + @pytest.fixture def dummy_ispyb_3d(dummy_params): store_in_ispyb_3d = Store3DGridscanInIspyb(SIM_ISPYB_CONFIG, dummy_params) - store_in_ispyb_3d.get_current_datacollection_comment = MagicMock() - store_in_ispyb_3d.get_current_datacollection_comment.return_value = "" return store_in_ispyb_3d @@ -68,6 +84,33 @@ def test_regex_string(dummy_ispyb, visit_path: str, expected_match: str): assert dummy_ispyb.get_visit_string_from_path(visit_path) == expected_match +@patch("ispyb.open", new_callable=mock_open) +def test_store_rotation_scan( + ispyb_conn, dummy_rotation_ispyb: StoreRotationInIspyb, dummy_rotation_params +): + ispyb_conn.return_value.mx_acquisition = mock() + ispyb_conn.return_value.core = mock() + + when(dummy_rotation_ispyb)._store_position_table( + ispyb_conn(), TEST_DATA_COLLECTION_IDS[0] + ).thenReturn(TEST_POSITION_ID) + + when(dummy_rotation_ispyb)._store_data_collection_group_table( + ispyb_conn() + ).thenReturn(TEST_DATA_COLLECTION_GROUP_ID) + + when(dummy_rotation_ispyb)._store_data_collection_table( + ispyb_conn(), TEST_DATA_COLLECTION_GROUP_ID + ).thenReturn(TEST_DATA_COLLECTION_IDS[0]) + + assert dummy_rotation_ispyb.experiment_type == "SAD" + + assert dummy_rotation_ispyb._store_scan_data(ispyb_conn()) == ( + TEST_DATA_COLLECTION_IDS[0], + TEST_DATA_COLLECTION_GROUP_ID, + ) + + @patch("ispyb.open", new_callable=mock_open) def test_store_grid_scan(ispyb_conn, dummy_ispyb, dummy_params): ispyb_conn.return_value.mx_acquisition = mock() From cff2c163f9d20a4452ef94c4e724ae0f93b27aa4 Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 15 Jun 2023 18:55:27 +0100 Subject: [PATCH 21/42] make ispybconfig constant --- .../external_interaction/system_tests/conftest.py | 11 +++++------ .../system_tests/test_ispyb_dev_connection.py | 5 ++--- src/artemis/parameters/constants.py | 3 +++ 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/artemis/external_interaction/system_tests/conftest.py b/src/artemis/external_interaction/system_tests/conftest.py index 27b0052fb..73a09c817 100644 --- a/src/artemis/external_interaction/system_tests/conftest.py +++ b/src/artemis/external_interaction/system_tests/conftest.py @@ -14,11 +14,10 @@ Store2DGridscanInIspyb, Store3DGridscanInIspyb, ) +from artemis.parameters.constants import DEF_ISPYB_DATABASE_CFG from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters -ISPYB_CONFIG = "/dls_sw/dasc/mariadb/credentials/ispyb-dev.cfg" - TEST_RESULT_LARGE = [ { "centre_of_mass": [1, 2, 3], @@ -86,7 +85,7 @@ def get_current_datacollection_attribute( @pytest.fixture def fetch_comment() -> Callable: - url = ispyb.sqlalchemy.url(ISPYB_CONFIG) + url = ispyb.sqlalchemy.url(DEF_ISPYB_DATABASE_CFG) engine = create_engine(url, connect_args={"use_pure": True}) Session = sessionmaker(engine) return partial(get_current_datacollection_comment, Session) @@ -94,7 +93,7 @@ def fetch_comment() -> Callable: @pytest.fixture def fetch_datacollection_attribute() -> Callable: - url = ispyb.sqlalchemy.url(ISPYB_CONFIG) + url = ispyb.sqlalchemy.url(DEF_ISPYB_DATABASE_CFG) engine = create_engine(url, connect_args={"use_pure": True}) Session = sessionmaker(engine) return partial(get_current_datacollection_attribute, Session) @@ -114,12 +113,12 @@ def dummy_params(): @pytest.fixture def dummy_ispyb(dummy_params) -> Store2DGridscanInIspyb: - return Store2DGridscanInIspyb(ISPYB_CONFIG, dummy_params) + return Store2DGridscanInIspyb(DEF_ISPYB_DATABASE_CFG, dummy_params) @pytest.fixture def dummy_ispyb_3d(dummy_params) -> Store3DGridscanInIspyb: - return Store3DGridscanInIspyb(ISPYB_CONFIG, dummy_params) + return Store3DGridscanInIspyb(DEF_ISPYB_DATABASE_CFG, dummy_params) @pytest.fixture diff --git a/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py b/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py index 693ab9b04..e43469323 100644 --- a/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py +++ b/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py @@ -5,11 +5,10 @@ Store3DGridscanInIspyb, StoreGridscanInIspyb, ) +from artemis.parameters.constants import DEF_ISPYB_DATABASE_CFG from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters -ISPYB_CONFIG = "/dls_sw/dasc/mariadb/credentials/ispyb-dev.cfg" - @pytest.mark.s03 def test_ispyb_get_comment_from_collection_correctly(fetch_comment): @@ -69,7 +68,7 @@ def test_can_store_2D_ispyb_data_correctly_when_in_error( ): test_params = FGSInternalParameters(**default_raw_params()) test_params.artemis_params.ispyb_params.visit_path = "/tmp/cm31105-4/" - ispyb: StoreGridscanInIspyb = StoreClass(ISPYB_CONFIG, test_params) + ispyb: StoreGridscanInIspyb = StoreClass(DEF_ISPYB_DATABASE_CFG, test_params) dc_ids, grid_ids, dcg_id = ispyb.begin_deposition() assert len(dc_ids) == exp_num_of_grids diff --git a/src/artemis/parameters/constants.py b/src/artemis/parameters/constants.py index a94718e84..8d58942f1 100644 --- a/src/artemis/parameters/constants.py +++ b/src/artemis/parameters/constants.py @@ -11,7 +11,10 @@ } PARAMETER_VERSION = 0.2 +# this one is for reading SIM_ISPYB_CONFIG = "src/artemis/external_interaction/unit_tests/test_config.cfg" +# this one is for making depositions: +DEF_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-dev.cfg" PARAMETER_SCHEMA_DIRECTORY = "src/artemis/parameters/schemas/" DETECTOR_PARAM_DEFAULTS = { From 74502328b065bd7326fc53779240cf271f25d639 Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 15 Jun 2023 19:02:20 +0100 Subject: [PATCH 22/42] fix stuff from merge --- src/artemis/experiment_plans/rotation_scan_plan.py | 4 +--- src/artemis/system_tests/test_fgs_plan.py | 7 +++---- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index a25c0bbfa..30f70e3a8 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -168,9 +168,7 @@ def cleanup_plan(eiger, zebra, smargon, detector_motion, backlight): yield from bpp.finalize_wrapper(disarm_zebra(zebra), bps.wait("cleanup_senv")) -def get_plan( - params: RotationInternalParameters, subscriptions: RotationCallbackCollection -): +def get_plan(params: RotationInternalParameters): devices = create_devices() subscriptions = RotationCallbackCollection.from_params(params) diff --git a/src/artemis/system_tests/test_fgs_plan.py b/src/artemis/system_tests/test_fgs_plan.py index 91166d3ae..43010bcfa 100644 --- a/src/artemis/system_tests/test_fgs_plan.py +++ b/src/artemis/system_tests/test_fgs_plan.py @@ -22,11 +22,10 @@ fetch_comment, zocalo_env, ) -from artemis.external_interaction.system_tests.test_ispyb_dev_connection import ( - ISPYB_CONFIG, -) from artemis.parameters.beamline_parameters import GDABeamlineParameters -from artemis.parameters.constants import BEAMLINE_PARAMETER_PATHS, SIM_BEAMLINE +from artemis.parameters.constants import BEAMLINE_PARAMETER_PATHS +from artemis.parameters.constants import DEF_ISPYB_DATABASE_CFG as ISPYB_CONFIG +from artemis.parameters.constants import SIM_BEAMLINE from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters From 7fce67eb93807c8b60a543e95b8cc2f71a4bb7a4 Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 15 Jun 2023 19:26:11 +0100 Subject: [PATCH 23/42] fix tests --- .../experiment_plans/tests/test_rotation_scan_plan.py | 5 ++--- src/artemis/external_interaction/ispyb/ispyb_dataclass.py | 6 ++++++ src/artemis/parameters/internal_parameters.py | 3 --- .../plan_specific/grid_scan_with_edge_detect_params.py | 5 ++++- 4 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 3b939b8ce..d56f84c22 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -77,7 +77,6 @@ def test_get_plan( eiger: EigerDetector, detector_motion: DetectorMotion, backlight: Backlight, - mock_subscriptions, ): eiger.stage = MagicMock() eiger.unstage = MagicMock() @@ -92,8 +91,8 @@ def test_get_plan( return_value=detector_motion, ), patch( - "artemis.experiment_plans.fast_grid_scan_plan.FGSCallbackCollection.from_params", - lambda _: mock_subscriptions, + "artemis.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", + lambda _: MagicMock(), ), ): RE(get_plan(test_rotation_params)) diff --git a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py index bfbb6c889..287d14e26 100644 --- a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py @@ -37,6 +37,9 @@ class IspybParams(BaseModel): microns_per_pixel_y: float position: np.ndarray + def __init__(self, **kwargs): + super().__init__(**kwargs) # TODO REMOVE JUST FOR DEBUGGING + class Config: arbitrary_types_allowed = True json_encoders = {np.ndarray: lambda a: a.tolist()} @@ -84,6 +87,9 @@ class RotationIspybParams(IspybParams): class GridscanIspybParams(IspybParams): upper_left: np.ndarray + def __init__(self, **kwargs): + super().__init__(**kwargs) # TODO REMOVE JUST FOR DEBUGGING + def dict(self, **kwargs): as_dict = super().dict(**kwargs) as_dict["upper_left"] = as_dict["upper_left"].tolist() diff --git a/src/artemis/parameters/internal_parameters.py b/src/artemis/parameters/internal_parameters.py index c17661e2c..b3bda1ff4 100644 --- a/src/artemis/parameters/internal_parameters.py +++ b/src/artemis/parameters/internal_parameters.py @@ -2,7 +2,6 @@ from typing import Any from dodal.devices.eiger import DetectorParams -from dodal.parameters.experiment_parameter_base import AbstractExperimentParameterBase from pydantic import BaseModel, root_validator from semver import Version @@ -117,8 +116,6 @@ def extract_artemis_params_from_flat_dict( class InternalParameters(BaseModel): params_version: ParameterVersion - experiment_params: AbstractExperimentParameterBase - artemis_params: ArtemisParameters class Config: use_enum_values = True diff --git a/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py b/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py index 36415f045..fc558fa64 100644 --- a/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py +++ b/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py @@ -49,6 +49,9 @@ class Config: **ArtemisParameters.Config.json_encoders, } + def __init__(self, **args): + super().__init__(**args) + @staticmethod def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: artemis_param_field_keys = [ @@ -89,7 +92,7 @@ def _preprocess_artemis_params( all_params["num_images_per_trigger"] = 1 all_params["trigger_mode"] = TriggerMode.FREE_RUN all_params["upper_left"] = np.array([0, 0, 0]) - return ArtemisParameters( + return GridscanArtemisParameters( **extract_artemis_params_from_flat_dict( all_params, cls._artemis_param_key_definitions() ) From 0402d036e05b6cdc2f9a772e76767e62df457e7d Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 16 Jun 2023 08:15:40 +0100 Subject: [PATCH 24/42] pass experiment type to superclass --- .../ispyb/store_in_ispyb.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index c427cfa71..e033d37d0 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -44,8 +44,9 @@ class StoreInIspyb(ABC): xtal_snapshots: list[str] | None = None data_collection_group_id: int | None = None - def __init__(self, ispyb_config, parameters=None) -> None: + def __init__(self, ispyb_config: str, experiment_type: str) -> None: self.ISPYB_CONFIG_PATH: str = ispyb_config + self.experiment_type = experiment_type @abstractmethod def _store_scan_data(self, conn: Connector): @@ -240,8 +241,7 @@ def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None else: self.xtal_snapshots = [] LOGGER.warning("No xtal snapshot paths sent to ISPyB!") - self.experiment_type = "SAD" - super().__init__(ispyb_config, parameters) + super().__init__(ispyb_config, "SAD") def _mutate_datacollection_params_for_experiment( self, params: dict[str, Any] @@ -308,9 +308,14 @@ class StoreGridscanInIspyb(StoreInIspyb): y_step_size: int | None = None grid_ids: tuple[int, ...] | None = None - def __init__(self, ispyb_config, parameters=None) -> None: + def __init__( + self, + ispyb_config: str, + experiment_type: str, + parameters: FGSInternalParameters = None, + ) -> None: self.full_params: FGSInternalParameters | None = parameters - super().__init__(ispyb_config, parameters) + super().__init__(ispyb_config, experiment_type) def begin_deposition(self): ( @@ -425,8 +430,7 @@ def _construct_comment(self) -> str: class Store3DGridscanInIspyb(StoreGridscanInIspyb): def __init__(self, ispyb_config, parameters=None): - super().__init__(ispyb_config, parameters) - self.experiment_type = "Mesh3D" + super().__init__(ispyb_config, "Mesh3D", parameters) def _store_scan_data(self, conn: Connector): data_collection_group_id = self._store_data_collection_group_table(conn) @@ -469,8 +473,7 @@ def __prepare_second_scan_params(self): class Store2DGridscanInIspyb(StoreGridscanInIspyb): def __init__(self, ispyb_config, parameters=None): - super().__init__(ispyb_config, parameters) - self.experiment_type = "mesh" + super().__init__(ispyb_config, "mesh", parameters) def _store_scan_data(self, conn: Connector): data_collection_group_id = self._store_data_collection_group_table(conn) From e5cceb42a92bc7a4bf932ddf627f7f8a49ea9512 Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 16 Jun 2023 08:36:15 +0100 Subject: [PATCH 25/42] move end deposition logic to parent class --- .../ispyb/store_in_ispyb.py | 106 ++++++++---------- 1 file changed, 45 insertions(+), 61 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index e033d37d0..a829c97e3 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -57,7 +57,7 @@ def _construct_comment(self) -> str: pass @abstractmethod - def _mutate_datacollection_params_for_experiment( + def _mutate_data_collection_params_for_experiment( self, params: dict[str, Any] ) -> dict[str, Any]: pass @@ -71,12 +71,12 @@ def end_deposition(self, success: str, reason: str): pass def append_to_comment( - self, datacollection_id: int, comment: str, delimiter: str = " " + self, data_collection_id: int, comment: str, delimiter: str = " " ) -> None: with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: mx_acquisition: MXAcquisition = conn.mx_acquisition mx_acquisition.update_data_collection_append_comments( - datacollection_id, comment, delimiter + data_collection_id, comment, delimiter ) def get_current_time_string(self): @@ -99,25 +99,48 @@ def update_scan_with_end_time_and_status( end_time: str, run_status: str, reason: str, - datacollection_id: int, - datacollection_group_id: int, + data_collection_id: int, + data_collection_group_id: int, ) -> None: assert self.ispyb_params is not None assert self.detector_params is not None if reason is not None and reason != "": - self.append_to_comment(datacollection_id, f"{run_status} reason: {reason}") + self.append_to_comment(data_collection_id, f"{run_status} reason: {reason}") with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: mx_acquisition: MXAcquisition = conn.mx_acquisition params = mx_acquisition.get_data_collection_params() - params["id"] = datacollection_id - params["parentid"] = datacollection_group_id + params["id"] = data_collection_id + params["parentid"] = data_collection_group_id params["endtime"] = end_time params["run_status"] = run_status mx_acquisition.upsert_data_collection(list(params.values())) + def _end_deposition(self, dcid: int, success: str, reason: str): + """Write the end of data_collection data. + Args: + success (str): The success of the run, could be fail or abort + reason (str): If the run failed, the reason why + """ + LOGGER.info( + f"End ispyb deposition with status '{success}' and reason '{reason}'." + ) + if success == "fail" or success == "abort": + run_status = "DataCollection Unsuccessful" + else: + run_status = "DataCollection Successful" + current_time = self.get_current_time_string() + assert self.data_collection_group_id is not None + self.update_scan_with_end_time_and_status( + current_time, + run_status, + reason, + dcid, + self.data_collection_group_id, + ) + def _store_position_table(self, conn: Connector, dc_id: int) -> int: assert self.ispyb_params is not None mx_acquisition: MXAcquisition = conn.mx_acquisition @@ -152,7 +175,7 @@ def _store_data_collection_group_table(self, conn: Connector) -> int: return mx_acquisition.upsert_data_collection_group(list(params.values())) - @TRACER.start_as_current_span("store_ispyb_datacollection_table") + @TRACER.start_as_current_span("store_ispyb_data_collection_table") def _store_data_collection_table( self, conn: Connector, data_collection_group_id: int ) -> int: @@ -170,7 +193,7 @@ def _store_data_collection_table( ) params = mx_acquisition.get_data_collection_params() - params = self._mutate_datacollection_params_for_experiment(params) + params = self._mutate_data_collection_params_for_experiment(params) params["visitid"] = session_id params["parentid"] = data_collection_group_id @@ -185,7 +208,7 @@ def _store_data_collection_table( params["beamsize_at_sampley"] = self.ispyb_params.beam_size_y params["transmission"] = self.ispyb_params.transmission params["comments"] = self._construct_comment() - params["datacollection_number"] = self.run_number + params["data_collection_number"] = self.run_number params["detector_distance"] = self.detector_params.detector_distance params["exp_time"] = self.detector_params.exposure_time params["imgdir"] = self.detector_params.directory @@ -243,7 +266,7 @@ def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None LOGGER.warning("No xtal snapshot paths sent to ISPyB!") super().__init__(ispyb_config, "SAD") - def _mutate_datacollection_params_for_experiment( + def _mutate_data_collection_params_for_experiment( self, params: dict[str, Any] ) -> dict[str, Any]: assert self.full_params is not None @@ -271,30 +294,10 @@ def begin_deposition(self): return self._store_scan_data(conn) def end_deposition(self, success: str, reason: str): - """Write the end of datacollection data. - Args: - success (str): The success of the run, could be fail or abort - reason (str): If the run failed, the reason why - """ - LOGGER.info( - f"End ispyb deposition with status '{success}' and reason '{reason}'." - ) - if success == "fail" or success == "abort": - run_status = "DataCollection Unsuccessful" - else: - run_status = "DataCollection Successful" - current_time = self.get_current_time_string() assert ( self.data_collection_id is not None - ), "Can't end ISPyB deposition, datacollection IDs are missing" - assert self.data_collection_group_id is not None - self.update_scan_with_end_time_and_status( - current_time, - run_status, - reason, - self.data_collection_id, - self.data_collection_group_id, - ) + ), "Can't end ISPyB deposition, data_collection IDs is missing" + self._end_deposition(self.data_collection_id, success, reason) def _construct_comment(self) -> str: return "Hyperion rotation scan" @@ -302,7 +305,7 @@ def _construct_comment(self) -> str: class StoreGridscanInIspyb(StoreInIspyb): ispyb_params: GridscanIspybParams | None = None - datacollection_ids: tuple[int, ...] | None = None + data_collection_ids: tuple[int, ...] | None = None upper_left: list[int] | None = None y_steps: int | None = None y_step_size: int | None = None @@ -319,37 +322,18 @@ def __init__( def begin_deposition(self): ( - self.datacollection_ids, + self.data_collection_ids, self.grid_ids, self.data_collection_group_id, ) = self.store_grid_scan(self.full_params) - return self.datacollection_ids, self.grid_ids, self.data_collection_group_id + return self.data_collection_ids, self.grid_ids, self.data_collection_group_id def end_deposition(self, success: str, reason: str): - """Write the end of datacollection data. - - Args: - success (str): The success of the run, could be fail or abort - reason (str):If the run failed, the reason why - """ - LOGGER.info( - f"End ispyb deposition with status '{success}' and reason '{reason}'." - ) - if success == "fail": - run_status = "DataCollection Unsuccessful" - elif success == "abort": - run_status = "DataCollection Unsuccessful" - else: - run_status = "DataCollection Successful" - current_time = self.get_current_time_string() assert ( - self.datacollection_ids is not None - ), "Can't end ISPyB deposition, datacollection IDs are missing" - assert self.data_collection_group_id is not None - for id in self.datacollection_ids: - self.update_scan_with_end_time_and_status( - current_time, run_status, reason, id, self.data_collection_group_id - ) + self.data_collection_ids is not None + ), "Can't end ISPyB deposition, data_collection IDs are missing" + for id in self.data_collection_ids: + self._end_deposition(id, success, reason) def store_grid_scan(self, full_params: FGSInternalParameters): self.full_params = full_params @@ -368,7 +352,7 @@ def store_grid_scan(self, full_params: FGSInternalParameters): with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: return self._store_scan_data(conn) - def _mutate_datacollection_params_for_experiment( + def _mutate_data_collection_params_for_experiment( self, params: dict[str, Any] ) -> dict[str, Any]: assert self.full_params is not None From 0fe777ec4d2bcd9893d2964f3a54dd261e43d1d9 Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 16 Jun 2023 08:40:00 +0100 Subject: [PATCH 26/42] set dev config in test and fix typo --- .../experiment_plans/tests/test_rotation_scan_plan.py | 2 ++ .../external_interaction/system_tests/conftest.py | 10 +++++----- .../system_tests/test_ispyb_dev_connection.py | 4 ++-- src/artemis/parameters/constants.py | 2 +- src/artemis/system_tests/test_fgs_plan.py | 2 +- 5 files changed, 11 insertions(+), 9 deletions(-) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index d56f84c22..a3531e418 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -23,6 +23,7 @@ fetch_comment, fetch_datacollection_attribute, ) +from artemis.parameters.constants import DEV_ISPYB_DATABASE_CFG from artemis.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, ) @@ -248,6 +249,7 @@ def test_ispyb_deposition_in_plan( test_rotation_params.artemis_params.detector_params.exposure_time = test_exp_time test_rotation_params.artemis_params.ispyb_params.wavelength = test_wl callbacks = RotationCallbackCollection.from_params(test_rotation_params) + callbacks.ispyb_handler.ispyb.ISPYB_CONFIG_PATH = DEV_ISPYB_DATABASE_CFG with ( patch( diff --git a/src/artemis/external_interaction/system_tests/conftest.py b/src/artemis/external_interaction/system_tests/conftest.py index 73a09c817..7e60b7d40 100644 --- a/src/artemis/external_interaction/system_tests/conftest.py +++ b/src/artemis/external_interaction/system_tests/conftest.py @@ -14,7 +14,7 @@ Store2DGridscanInIspyb, Store3DGridscanInIspyb, ) -from artemis.parameters.constants import DEF_ISPYB_DATABASE_CFG +from artemis.parameters.constants import DEV_ISPYB_DATABASE_CFG from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters @@ -85,7 +85,7 @@ def get_current_datacollection_attribute( @pytest.fixture def fetch_comment() -> Callable: - url = ispyb.sqlalchemy.url(DEF_ISPYB_DATABASE_CFG) + url = ispyb.sqlalchemy.url(DEV_ISPYB_DATABASE_CFG) engine = create_engine(url, connect_args={"use_pure": True}) Session = sessionmaker(engine) return partial(get_current_datacollection_comment, Session) @@ -93,7 +93,7 @@ def fetch_comment() -> Callable: @pytest.fixture def fetch_datacollection_attribute() -> Callable: - url = ispyb.sqlalchemy.url(DEF_ISPYB_DATABASE_CFG) + url = ispyb.sqlalchemy.url(DEV_ISPYB_DATABASE_CFG) engine = create_engine(url, connect_args={"use_pure": True}) Session = sessionmaker(engine) return partial(get_current_datacollection_attribute, Session) @@ -113,12 +113,12 @@ def dummy_params(): @pytest.fixture def dummy_ispyb(dummy_params) -> Store2DGridscanInIspyb: - return Store2DGridscanInIspyb(DEF_ISPYB_DATABASE_CFG, dummy_params) + return Store2DGridscanInIspyb(DEV_ISPYB_DATABASE_CFG, dummy_params) @pytest.fixture def dummy_ispyb_3d(dummy_params) -> Store3DGridscanInIspyb: - return Store3DGridscanInIspyb(DEF_ISPYB_DATABASE_CFG, dummy_params) + return Store3DGridscanInIspyb(DEV_ISPYB_DATABASE_CFG, dummy_params) @pytest.fixture diff --git a/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py b/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py index e43469323..44e4c031a 100644 --- a/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py +++ b/src/artemis/external_interaction/system_tests/test_ispyb_dev_connection.py @@ -5,7 +5,7 @@ Store3DGridscanInIspyb, StoreGridscanInIspyb, ) -from artemis.parameters.constants import DEF_ISPYB_DATABASE_CFG +from artemis.parameters.constants import DEV_ISPYB_DATABASE_CFG from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters @@ -68,7 +68,7 @@ def test_can_store_2D_ispyb_data_correctly_when_in_error( ): test_params = FGSInternalParameters(**default_raw_params()) test_params.artemis_params.ispyb_params.visit_path = "/tmp/cm31105-4/" - ispyb: StoreGridscanInIspyb = StoreClass(DEF_ISPYB_DATABASE_CFG, test_params) + ispyb: StoreGridscanInIspyb = StoreClass(DEV_ISPYB_DATABASE_CFG, test_params) dc_ids, grid_ids, dcg_id = ispyb.begin_deposition() assert len(dc_ids) == exp_num_of_grids diff --git a/src/artemis/parameters/constants.py b/src/artemis/parameters/constants.py index 8d58942f1..44740da4f 100644 --- a/src/artemis/parameters/constants.py +++ b/src/artemis/parameters/constants.py @@ -14,7 +14,7 @@ # this one is for reading SIM_ISPYB_CONFIG = "src/artemis/external_interaction/unit_tests/test_config.cfg" # this one is for making depositions: -DEF_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-dev.cfg" +DEV_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-dev.cfg" PARAMETER_SCHEMA_DIRECTORY = "src/artemis/parameters/schemas/" DETECTOR_PARAM_DEFAULTS = { diff --git a/src/artemis/system_tests/test_fgs_plan.py b/src/artemis/system_tests/test_fgs_plan.py index 43010bcfa..b1f4247b7 100644 --- a/src/artemis/system_tests/test_fgs_plan.py +++ b/src/artemis/system_tests/test_fgs_plan.py @@ -24,7 +24,7 @@ ) from artemis.parameters.beamline_parameters import GDABeamlineParameters from artemis.parameters.constants import BEAMLINE_PARAMETER_PATHS -from artemis.parameters.constants import DEF_ISPYB_DATABASE_CFG as ISPYB_CONFIG +from artemis.parameters.constants import DEV_ISPYB_DATABASE_CFG as ISPYB_CONFIG from artemis.parameters.constants import SIM_BEAMLINE from artemis.parameters.external_parameters import from_file as default_raw_params from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters From a723870ea0a95ed2b0e6b2c9d202c175c6934797 Mon Sep 17 00:00:00 2001 From: David Perl Date: Mon, 17 Jul 2023 16:15:18 +0100 Subject: [PATCH 27/42] tidy --- src/artemis/experiment_plans/tests/test_grid_detection_plan.py | 2 ++ src/artemis/external_interaction/ispyb/ispyb_dataclass.py | 3 --- src/artemis/parameters/plan_specific/fgs_internal_params.py | 3 --- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/src/artemis/experiment_plans/tests/test_grid_detection_plan.py b/src/artemis/experiment_plans/tests/test_grid_detection_plan.py index 2cdd2f825..033b5ff48 100644 --- a/src/artemis/experiment_plans/tests/test_grid_detection_plan.py +++ b/src/artemis/experiment_plans/tests/test_grid_detection_plan.py @@ -71,6 +71,8 @@ def test_grid_detection_plan_runs_and_triggers_snapshots( oav.mxsc.pin_tip.tip_x.sim_put(100) oav.mxsc.pin_tip.tip_y.sim_put(100) + sleep(0.1) + params = OAVParameters(context="loopCentring", **test_config_files) gridscan_params = GridScanParams() diff --git a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py index 7596d775e..7002dfe1b 100644 --- a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py @@ -37,9 +37,6 @@ class IspybParams(BaseModel): microns_per_pixel_y: float position: np.ndarray - def __init__(self, **kwargs): - super().__init__(**kwargs) # TODO REMOVE JUST FOR DEBUGGING - class Config: arbitrary_types_allowed = True json_encoders = {np.ndarray: lambda a: a.tolist()} diff --git a/src/artemis/parameters/plan_specific/fgs_internal_params.py b/src/artemis/parameters/plan_specific/fgs_internal_params.py index 300bf3e49..c866dbc33 100644 --- a/src/artemis/parameters/plan_specific/fgs_internal_params.py +++ b/src/artemis/parameters/plan_specific/fgs_internal_params.py @@ -27,9 +27,6 @@ class GridscanArtemisParameters(ArtemisParameters): **GRIDSCAN_ISPYB_PARAM_DEFAULTS ) - def __init__(self, **kwargs): - super().__init__(**kwargs) # TODO REMOVE JUST FOR DEBUGGING - class Config: arbitrary_types_allowed = True json_encoders = { From 6c10d67c2178917fc70566914cdc85b335ae0c35 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 18 Jul 2023 09:40:24 +0100 Subject: [PATCH 28/42] fix more merge issues --- .../experiment_plans/rotation_scan_plan.py | 4 +- .../experiment_plans/tests/conftest.py | 20 ++++++++ .../tests/test_grid_detection_plan.py | 2 +- .../tests/test_rotation_scan_plan.py | 46 ++++++++++++------- .../rotation/tests/test_rotation_callbacks.py | 8 ---- .../rotation_scan_internal_params.py | 6 +++ .../tests/test_data/artemis_parameters.json | 2 +- 7 files changed, 59 insertions(+), 29 deletions(-) diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index 227d98423..fa587163a 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -130,9 +130,7 @@ def rotation_scan_plan( # get some information for the ispyb deposition and trigger the callback yield from read_hardware_for_ispyb( - i03.undulator(), - i03.synchrotron(), - i03.s4_slit_gaps(), + i03.undulator(), i03.synchrotron(), i03.s4_slit_gaps(), i03.flux() ) LOGGER.info( diff --git a/src/artemis/experiment_plans/tests/conftest.py b/src/artemis/experiment_plans/tests/conftest.py index 0495848b3..ce80dfcbf 100644 --- a/src/artemis/experiment_plans/tests/conftest.py +++ b/src/artemis/experiment_plans/tests/conftest.py @@ -69,6 +69,26 @@ def detector_motion(): return i03.detector_motion(fake_with_ophyd_sim=True) +@pytest.fixture +def undulator(): + return i03.undulator(fake_with_ophyd_sim=True) + + +@pytest.fixture +def s4_slit_gaps(): + return i03.s4_slit_gaps(fake_with_ophyd_sim=True) + + +@pytest.fixture +def synchrotron(): + return i03.synchrotron(fake_with_ophyd_sim=True) + + +@pytest.fixture +def flux(): + return i03.flux(fake_with_ophyd_sim=True) + + @pytest.fixture def aperture_scatterguard(): return i03.aperture_scatterguard( diff --git a/src/artemis/experiment_plans/tests/test_grid_detection_plan.py b/src/artemis/experiment_plans/tests/test_grid_detection_plan.py index 033b5ff48..fbd990f11 100644 --- a/src/artemis/experiment_plans/tests/test_grid_detection_plan.py +++ b/src/artemis/experiment_plans/tests/test_grid_detection_plan.py @@ -68,8 +68,8 @@ def test_grid_detection_plan_runs_and_triggers_snapshots( sleep(0.1) - oav.mxsc.pin_tip.tip_x.sim_put(100) oav.mxsc.pin_tip.tip_y.sim_put(100) + oav.mxsc.pin_tip.tip_x.sim_put(100) sleep(0.1) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 2602461f4..27f61145f 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -114,6 +114,10 @@ def test_rotation_plan( detector_motion: DetectorMotion, backlight: Backlight, mock_rotation_subscriptions: RotationCallbackCollection, + synchrotron, + s4_slit_gaps, + undulator, + flux, ): mock_omega_sets = MagicMock(return_value=Status(done=True, success=True)) @@ -125,12 +129,19 @@ def test_rotation_plan( smargon.omega.velocity.set = mock_omega_sets smargon.omega.set = mock_omega_sets - with patch( - "bluesky.preprocessors.__read_and_stash_a_motor", - __fake_read, - ), patch( - "artemis.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", - lambda _: mock_rotation_subscriptions, + with ( + patch( + "bluesky.preprocessors.__read_and_stash_a_motor", + __fake_read, + ), + patch( + "artemis.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", + lambda _: mock_rotation_subscriptions, + ), + patch("dodal.beamlines.i03.undulator", lambda: undulator), + patch("dodal.beamlines.i03.synchrotron", lambda: synchrotron), + patch("dodal.beamlines.i03.s4_slit_gaps", lambda: s4_slit_gaps), + patch("dodal.beamlines.i03.flux", lambda: flux), ): RE( rotation_scan_plan( @@ -171,11 +182,13 @@ def test_cleanup_happens( cleanup_plan.assert_not_called() # check that failure is handled in composite plan with ( - patch("dodal.beamlines.i03.smargon", return_value=smargon), - patch("dodal.beamlines.i03.eiger", return_value=eiger), - patch("dodal.beamlines.i03.zebra", return_value=zebra), - patch("dodal.beamlines.i03.backlight", return_value=backlight), - patch("dodal.beamlines.i03.detector_motion", return_value=detector_motion), + patch("dodal.beamlines.i03.smargon", return_value=lambda: smargon), + patch("dodal.beamlines.i03.eiger", return_value=lambda: eiger), + patch("dodal.beamlines.i03.zebra", return_value=lambda: zebra), + patch("dodal.beamlines.i03.backlight", return_value=lambda: backlight), + patch( + "dodal.beamlines.i03.detector_motion", return_value=lambda: detector_motion + ), patch( "artemis.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", lambda _: mock_rotation_subscriptions, @@ -227,11 +240,11 @@ def test_ispyb_deposition_in_plan( test_rotation_params: RotationInternalParameters, fetch_comment, fetch_datacollection_attribute, + undulator, + synchrotron, + s4_slit_gaps, + flux, ): - undulator = i03.undulator(fake_with_ophyd_sim=True) - synchrotron = i03.synchrotron(fake_with_ophyd_sim=True) - slit_gaps = i03.s4_slit_gaps(fake_with_ophyd_sim=True) - test_wl = 0.71 test_bs_x = 0.023 test_bs_y = 0.047 @@ -253,7 +266,8 @@ def test_ispyb_deposition_in_plan( ), patch("dodal.beamlines.i03.undulator", return_value=undulator), patch("dodal.beamlines.i03.synchrotron", return_value=synchrotron), - patch("dodal.beamlines.i03.s4_slit_gaps", return_value=slit_gaps), + patch("dodal.beamlines.i03.s4_slit_gaps", return_value=s4_slit_gaps), + patch("dodal.beamlines.i03.flux", return_value=flux), patch( "bluesky.preprocessors.__read_and_stash_a_motor", __fake_read, diff --git a/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index e51bb1b95..7ddd312f4 100644 --- a/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -6,9 +6,6 @@ import pytest from bluesky.run_engine import RunEngine -from artemis.external_interaction.callbacks.rotation.nexus_callback import ( - RotationNexusFileHandlerCallback, -) from artemis.external_interaction.callbacks.rotation.rotation_callback_collection import ( RotationCallbackCollection, ) @@ -27,11 +24,6 @@ def params(): ) -def test_callback_collection_init(params): - callbacks = RotationCallbackCollection() - assert isinstance(callbacks.nexus_handler, RotationNexusFileHandlerCallback) - - def fake_get_plan( parameters: RotationInternalParameters, subscriptions: RotationCallbackCollection ): diff --git a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py index 7db1b7fa3..13d5a7d2e 100644 --- a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py +++ b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py @@ -88,6 +88,12 @@ class RotationInternalParameters(InternalParameters): experiment_params: RotationScanParams artemis_params: RotationArtemisParameters + class Config: + arbitrary_types_allowed = True + json_encoders = { + **RotationArtemisParameters.Config.json_encoders, + } + @staticmethod def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: artemis_param_field_keys = [ diff --git a/src/artemis/parameters/tests/test_data/artemis_parameters.json b/src/artemis/parameters/tests/test_data/artemis_parameters.json index afa635a0e..4ba84ae17 100644 --- a/src/artemis/parameters/tests/test_data/artemis_parameters.json +++ b/src/artemis/parameters/tests/test_data/artemis_parameters.json @@ -4,7 +4,7 @@ "insertion_prefix": "SR03S", "experiment_type": "fast_grid_scan", "detector_params": { - "current_energy": 100.0, + "current_energy_ev": 100.0, "exposure_time": 0.1, "directory": "/tmp/", "prefix": "file_name", From 5a7111ddf707cacf6f5f7aba0f040b16d3e5e02b Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 18 Jul 2023 14:32:36 +0100 Subject: [PATCH 29/42] mock ispyb callback in nexus writing tests --- src/artemis/experiment_plans/tests/conftest.py | 2 ++ .../tests/test_rotation_scan_plan.py | 12 +++++------- .../rotation/tests/test_rotation_callbacks.py | 5 +++++ .../system_tests/test_write_rotation_nexus.py | 5 ++++- .../plan_specific/rotation_scan_internal_params.py | 1 - 5 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/artemis/experiment_plans/tests/conftest.py b/src/artemis/experiment_plans/tests/conftest.py index ce80dfcbf..889f8a5d5 100644 --- a/src/artemis/experiment_plans/tests/conftest.py +++ b/src/artemis/experiment_plans/tests/conftest.py @@ -173,6 +173,8 @@ def mock_subscriptions(test_fgs_params): def mock_rotation_subscriptions(test_rotation_params): with patch( "artemis.external_interaction.callbacks.rotation.rotation_callback_collection.RotationNexusFileHandlerCallback" + ), patch( + "artemis.external_interaction.callbacks.rotation.rotation_callback_collection.RotationISPyBHandlerCallback" ): subscriptions = RotationCallbackCollection.from_params(test_rotation_params) return subscriptions diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 27f61145f..970ed2f9f 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -182,13 +182,11 @@ def test_cleanup_happens( cleanup_plan.assert_not_called() # check that failure is handled in composite plan with ( - patch("dodal.beamlines.i03.smargon", return_value=lambda: smargon), - patch("dodal.beamlines.i03.eiger", return_value=lambda: eiger), - patch("dodal.beamlines.i03.zebra", return_value=lambda: zebra), - patch("dodal.beamlines.i03.backlight", return_value=lambda: backlight), - patch( - "dodal.beamlines.i03.detector_motion", return_value=lambda: detector_motion - ), + patch("dodal.beamlines.i03.smargon", return_value=smargon), + patch("dodal.beamlines.i03.eiger", return_value=eiger), + patch("dodal.beamlines.i03.zebra", return_value=zebra), + patch("dodal.beamlines.i03.backlight", return_value=backlight), + patch("dodal.beamlines.i03.detector_motion", return_value=detector_motion), patch( "artemis.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", lambda _: mock_rotation_subscriptions, diff --git a/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index 7ddd312f4..8946a0f91 100644 --- a/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/artemis/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -47,6 +47,8 @@ def test_nexus_handler_gets_documents_in_mock_plan(params: RotationInternalParam cb = RotationCallbackCollection.from_params(params) cb.nexus_handler.start = MagicMock() cb.nexus_handler.stop = MagicMock() + cb.ispyb_handler.start = MagicMock() + cb.ispyb_handler.stop = MagicMock() RE(fake_get_plan(params, cb)) @@ -68,6 +70,9 @@ def test_nexus_handler_triggers_write_file_when_told( cb = RotationCallbackCollection.from_params(params) + cb.ispyb_handler.start = MagicMock() + cb.ispyb_handler.stop = MagicMock() + RE(fake_get_plan(params, cb)) assert os.path.isfile("/tmp/file_name_0.nxs") diff --git a/src/artemis/external_interaction/system_tests/test_write_rotation_nexus.py b/src/artemis/external_interaction/system_tests/test_write_rotation_nexus.py index 9c1ceec3b..27dc2742a 100644 --- a/src/artemis/external_interaction/system_tests/test_write_rotation_nexus.py +++ b/src/artemis/external_interaction/system_tests/test_write_rotation_nexus.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp @@ -75,6 +75,9 @@ def test_rotation_scan_nexus_output_compared_to_existing_file( RE = RunEngine({}) cb = RotationCallbackCollection.from_params(test_params) + cb.ispyb_handler.start = MagicMock() + cb.ispyb_handler.stop = MagicMock() + cb.ispyb_handler.event = MagicMock() with patch( "artemis.external_interaction.nexus.write_nexus.get_current_time", return_value="test_time", diff --git a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py index 13d5a7d2e..5d63253f1 100644 --- a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py +++ b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py @@ -138,7 +138,6 @@ def _preprocess_artemis_params( all_params["omega_increment"] = 0 all_params["num_triggers"] = 1 all_params["num_images_per_trigger"] = all_params["num_images"] - all_params["upper_left"] = np.array(all_params["upper_left"]) return RotationArtemisParameters( **extract_artemis_params_from_flat_dict( all_params, cls._artemis_param_key_definitions() From 5ea651c064829c0831e98218cd79f4a796520115 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 18 Jul 2023 14:53:27 +0100 Subject: [PATCH 30/42] remove debugging line --- src/artemis/external_interaction/ispyb/ispyb_dataclass.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py index 7002dfe1b..2fb8faded 100644 --- a/src/artemis/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/artemis/external_interaction/ispyb/ispyb_dataclass.py @@ -84,9 +84,6 @@ class RotationIspybParams(IspybParams): class GridscanIspybParams(IspybParams): upper_left: np.ndarray - def __init__(self, **kwargs): - super().__init__(**kwargs) # TODO REMOVE JUST FOR DEBUGGING - def dict(self, **kwargs): as_dict = super().dict(**kwargs) as_dict["upper_left"] = as_dict["upper_left"].tolist() From 46704bf5042dbdb11d340039d30576bc69e9f810 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 25 Jul 2023 13:56:27 +0100 Subject: [PATCH 31/42] tidy bits and comments --- src/artemis/experiment_plans/fast_grid_scan_plan.py | 4 +--- src/artemis/experiment_plans/rotation_scan_plan.py | 2 +- .../experiment_plans/tests/test_rotation_scan_plan.py | 2 +- .../external_interaction/callbacks/fgs/ispyb_callback.py | 2 +- .../callbacks/rotation/ispyb_callback.py | 8 ++++---- 5 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/artemis/experiment_plans/fast_grid_scan_plan.py b/src/artemis/experiment_plans/fast_grid_scan_plan.py index 6660144bd..0a04a5c88 100755 --- a/src/artemis/experiment_plans/fast_grid_scan_plan.py +++ b/src/artemis/experiment_plans/fast_grid_scan_plan.py @@ -75,9 +75,7 @@ def __init__( fake_with_ophyd_sim=fake, aperture_positions=aperture_positions ) self.backlight = i03.backlight(fake_with_ophyd_sim=fake) - self.eiger = i03.eiger( - wait_for_connection=False, fake_with_ophyd_sim=fake, params=detector_params - ) + self.eiger = i03.eiger(fake_with_ophyd_sim=fake, params=detector_params) self.fast_grid_scan = i03.fast_grid_scan(fake_with_ophyd_sim=fake) self.flux = i03.flux(fake_with_ophyd_sim=fake) self.s4_slit_gaps = i03.s4_slit_gaps(fake_with_ophyd_sim=fake) diff --git a/src/artemis/experiment_plans/rotation_scan_plan.py b/src/artemis/experiment_plans/rotation_scan_plan.py index 45240ab1e..f4412df76 100644 --- a/src/artemis/experiment_plans/rotation_scan_plan.py +++ b/src/artemis/experiment_plans/rotation_scan_plan.py @@ -39,7 +39,7 @@ def create_devices() -> dict[str, Device]: """Ensures necessary devices have been instantiated and returns a dict with references to them""" return { - "eiger": i03.eiger(wait_for_connection=False), + "eiger": i03.eiger(), "smargon": i03.smargon(), "zebra": i03.zebra(), "detector_motion": i03.detector_motion(), diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 9ca3305d1..6444509b9 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -217,7 +217,7 @@ def fake_create_devices( zebra.pc.arm_demand.set = mock_arm_disarm devices = { - "eiger": i03.eiger(wait_for_connection=False, fake_with_ophyd_sim=True), + "eiger": i03.eiger(fake_with_ophyd_sim=True), "smargon": smargon, "zebra": zebra, "detector_motion": i03.detector_motion(fake_with_ophyd_sim=True), diff --git a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py index a85276e10..a2cd0dc44 100644 --- a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py @@ -20,7 +20,7 @@ class FGSISPyBHandlerCallback(CallbackBase): """Callback class to handle the deposition of experiment parameters into the ISPyB database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on recieving an 'event' document for the 'ispyb_readings' event, and updates the - deposition on recieving it's final 'stop' document. + deposition on recieving its final 'stop' document. To use, subscribe the Bluesky RunEngine to an instance of this class. E.g.: diff --git a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py index dc693765d..e690957f8 100644 --- a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py @@ -16,17 +16,17 @@ class RotationISPyBHandlerCallback(CallbackBase): """Callback class to handle the deposition of experiment parameters into the ISPyB database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on recieving an 'event' document for the 'ispyb_readings' event, and updates the - deposition on recieving it's final 'stop' document. + deposition on recieving its final 'stop' document. To use, subscribe the Bluesky RunEngine to an instance of this class. E.g.: - nexus_file_handler_callback = NexusFileHandlerCallback(parameters) - RE.subscribe(nexus_file_handler_callback) + ispyb_handler_callback = RotationISPyBHandlerCallback(parameters) + RE.subscribe(ispyb_handler_callback) Or decorate a plan using bluesky.preprocessors.subs_decorator. See: https://blueskyproject.io/bluesky/callbacks.html#ways-to-invoke-callbacks - Usually used as part of an FGSCallbackCollection. + Usually used as part of a RotationCallbackCollection. """ def __init__(self, parameters: FGSInternalParameters): From 31bc2e7e71ffe9817fed0abf1ea6e50abd41b506 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 25 Jul 2023 14:22:24 +0100 Subject: [PATCH 32/42] make ispyb callback base class --- .../callbacks/fgs/ispyb_callback.py | 72 +++------------- .../callbacks/ispyb_callback_base.py | 83 +++++++++++++++++++ .../callbacks/rotation/ispyb_callback.py | 66 +++------------ 3 files changed, 107 insertions(+), 114 deletions(-) create mode 100644 src/artemis/external_interaction/callbacks/ispyb_callback_base.py diff --git a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py index a2cd0dc44..5a5bd6773 100644 --- a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py @@ -1,22 +1,19 @@ from __future__ import annotations -import os -from typing import Dict - -from bluesky.callbacks import CallbackBase - +from artemis.external_interaction.callbacks.ispyb_callback_base import ( + BaseISPyBHandlerCallback, +) from artemis.external_interaction.exceptions import ISPyBDepositionNotMade from artemis.external_interaction.ispyb.store_in_ispyb import ( Store2DGridscanInIspyb, Store3DGridscanInIspyb, StoreGridscanInIspyb, ) -from artemis.log import LOGGER, set_dcgid_tag -from artemis.parameters.constants import ISPYB_PLAN_NAME, SIM_ISPYB_CONFIG +from artemis.log import set_dcgid_tag from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters -class FGSISPyBHandlerCallback(CallbackBase): +class FGSISPyBHandlerCallback(BaseISPyBHandlerCallback): """Callback class to handle the deposition of experiment parameters into the ISPyB database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on recieving an 'event' document for the 'ispyb_readings' event, and updates the @@ -34,67 +31,24 @@ class FGSISPyBHandlerCallback(CallbackBase): """ def __init__(self, parameters: FGSInternalParameters): - self.params = parameters - self.descriptors: Dict[str, dict] = {} - ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) - if ispyb_config == SIM_ISPYB_CONFIG: - LOGGER.warning( - "Using dev ISPyB database. If you want to use the real database, please" - " set the ISPYB_CONFIG_PATH environment variable." - ) + super.__init__(parameters) self.ispyb: StoreGridscanInIspyb = ( - Store3DGridscanInIspyb(ispyb_config, self.params) + Store3DGridscanInIspyb(self.ispyb_config, self.params) if self.params.experiment_params.is_3d_grid_scan - else Store2DGridscanInIspyb(ispyb_config, self.params) + else Store2DGridscanInIspyb(self.ispyb_config, self.params) ) self.ispyb_ids: tuple = (None, None, None) - self.uid_to_finalize_on = None def append_to_comment(self, comment: str): - try: - for id in self.ispyb_ids[0]: - self.ispyb.append_to_comment(id, comment) - except TypeError: - LOGGER.warning("ISPyB deposition not initialised, can't update comment.") - - def descriptor(self, doc: dict): - self.descriptors[doc["uid"]] = doc - - def start(self, doc: dict): - if self.uid_to_finalize_on is None: - self.uid_to_finalize_on = doc.get("uid") + for id in self.ispyb_ids[0]: + self._append_to_comment(id, comment) def event(self, doc: dict): - LOGGER.debug("ISPyB handler received event document.") - event_descriptor = self.descriptors[doc["descriptor"]] - - if event_descriptor.get("name") == ISPYB_PLAN_NAME: - self.params.artemis_params.ispyb_params.undulator_gap = doc["data"][ - "undulator_gap" - ] - self.params.artemis_params.ispyb_params.synchrotron_mode = doc["data"][ - "synchrotron_machine_status_synchrotron_mode" - ] - self.params.artemis_params.ispyb_params.slit_gap_size_x = doc["data"][ - "s4_slit_gaps_xgap" - ] - self.params.artemis_params.ispyb_params.slit_gap_size_y = doc["data"][ - "s4_slit_gaps_ygap" - ] - self.params.artemis_params.ispyb_params.transmission = doc["data"][ - "attenuator_actual_transmission" - ] - - LOGGER.info("Creating ispyb entry.") - self.ispyb_ids = self.ispyb.begin_deposition() - set_dcgid_tag(self.ispyb_ids[2]) + super().event(doc) + set_dcgid_tag(self.ispyb_ids[2]) def stop(self, doc: dict): if doc.get("run_start") == self.uid_to_finalize_on: - LOGGER.debug("ISPyB handler received stop document.") - exit_status = doc.get("exit_status") - reason = doc.get("reason") if self.ispyb_ids == (None, None, None): raise ISPyBDepositionNotMade("ispyb was not initialised at run start") - self.ispyb.end_deposition(exit_status, reason) - set_dcgid_tag(None) + super().stop(doc) diff --git a/src/artemis/external_interaction/callbacks/ispyb_callback_base.py b/src/artemis/external_interaction/callbacks/ispyb_callback_base.py new file mode 100644 index 000000000..8ee2bafba --- /dev/null +++ b/src/artemis/external_interaction/callbacks/ispyb_callback_base.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import os +from typing import Dict + +from bluesky.callbacks import CallbackBase + +from artemis.external_interaction.ispyb.store_in_ispyb import StoreInIspyb +from artemis.log import LOGGER, set_dcgid_tag +from artemis.parameters.constants import ISPYB_PLAN_NAME, SIM_ISPYB_CONFIG +from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters + + +class BaseISPyBHandlerCallback(CallbackBase): + def __init__(self, parameters: FGSInternalParameters): + """Subclasses should run super().__init__() with parameters, then set + self.ispyb to the type of ispyb relevant to the experiment and define the type + for self.ispyb_ids.""" + self.params = parameters + self.descriptors: Dict[str, dict] = {} + self.ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) + if self.ispyb_config == SIM_ISPYB_CONFIG: + LOGGER.warning( + "Using dev ISPyB database. If you want to use the real database, please" + " set the ISPYB_CONFIG_PATH environment variable." + ) + self.uid_to_finalize_on = None + + def _append_to_comment(self, id: int, comment: str): + assert isinstance(self.ispyb, StoreInIspyb) + try: + self.ispyb.append_to_comment(id, comment) + except TypeError: + LOGGER.warning("ISPyB deposition not initialised, can't update comment.") + + def descriptor(self, doc: dict): + self.descriptors[doc["uid"]] = doc + + def start(self, doc: dict): + if self.uid_to_finalize_on is None: + self.uid_to_finalize_on = doc.get("uid") + + def event(self, doc: dict): + """Subclasses should extend this to add a call to set_dcig_tag from + artemis.log""" + + LOGGER.debug("ISPyB handler received event document.") + assert isinstance( + self.ispyb, StoreInIspyb + ), "ISPyB deposition can't be initialised!" + event_descriptor = self.descriptors[doc["descriptor"]] + + if event_descriptor.get("name") == ISPYB_PLAN_NAME: + self.params.artemis_params.ispyb_params.undulator_gap = doc["data"][ + "undulator_gap" + ] + self.params.artemis_params.ispyb_params.synchrotron_mode = doc["data"][ + "synchrotron_machine_status_synchrotron_mode" + ] + self.params.artemis_params.ispyb_params.slit_gap_size_x = doc["data"][ + "s4_slit_gaps_xgap" + ] + self.params.artemis_params.ispyb_params.slit_gap_size_y = doc["data"][ + "s4_slit_gaps_ygap" + ] + self.params.artemis_params.ispyb_params.transmission = doc["data"][ + "attenuator_actual_transmission" + ] + + LOGGER.info("Creating ispyb entry.") + self.ispyb_ids = self.ispyb.begin_deposition() + + def stop(self, doc: dict): + """Subclasses must check that they are recieving a stop document for the correct + uid to use this method!""" + assert isinstance( + self.ispyb, StoreInIspyb + ), "ISPyB handler recieved stop document, but deposition object doesn't exist!" + LOGGER.debug("ISPyB handler received stop document.") + exit_status = doc.get("exit_status") + reason = doc.get("reason") + self.ispyb.end_deposition(exit_status, reason) + set_dcgid_tag(None) diff --git a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py index e690957f8..9eebeaf1e 100644 --- a/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/rotation/ispyb_callback.py @@ -1,18 +1,15 @@ from __future__ import annotations -import os -from typing import Dict - -from bluesky.callbacks import CallbackBase - +from artemis.external_interaction.callbacks.ispyb_callback_base import ( + BaseISPyBHandlerCallback, +) from artemis.external_interaction.exceptions import ISPyBDepositionNotMade from artemis.external_interaction.ispyb.store_in_ispyb import StoreRotationInIspyb -from artemis.log import LOGGER, set_dcgid_tag -from artemis.parameters.constants import ISPYB_PLAN_NAME, SIM_ISPYB_CONFIG +from artemis.log import set_dcgid_tag from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters -class RotationISPyBHandlerCallback(CallbackBase): +class RotationISPyBHandlerCallback(BaseISPyBHandlerCallback): """Callback class to handle the deposition of experiment parameters into the ISPyB database. Listens for 'event' and 'descriptor' documents. Creates the ISpyB entry on recieving an 'event' document for the 'ispyb_readings' event, and updates the @@ -30,64 +27,23 @@ class RotationISPyBHandlerCallback(CallbackBase): """ def __init__(self, parameters: FGSInternalParameters): - self.params = parameters - self.descriptors: Dict[str, dict] = {} - ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) - if ispyb_config == SIM_ISPYB_CONFIG: - LOGGER.warning( - "Using dev ISPyB database. If you want to use the real database, please" - " set the ISPYB_CONFIG_PATH environment variable." - ) + super().__init__(parameters) self.ispyb: StoreRotationInIspyb = StoreRotationInIspyb( - ispyb_config, self.params + self.ispyb_config, self.params ) - self.ispyb_ids: tuple[int, int] | tuple[None, None] = (None, None) - self.uid_to_finalize_on: str | None = None def append_to_comment(self, comment: str): - try: - self.ispyb.append_to_comment(self.ispyb_ids[0], comment) - except TypeError: - LOGGER.warning("ISPyB deposition not initialised, can't update comment.") - - def descriptor(self, doc: dict): - self.descriptors[doc["uid"]] = doc - - def start(self, doc: dict): - if self.uid_to_finalize_on is None: - self.uid_to_finalize_on = doc.get("uid") + self._append_to_comment(self.ispyb_ids[0], comment) def event(self, doc: dict): - LOGGER.debug("ISPyB handler received event document.") - event_descriptor = self.descriptors[doc["descriptor"]] - - if event_descriptor.get("name") == ISPYB_PLAN_NAME: - self.params.artemis_params.ispyb_params.undulator_gap = doc["data"][ - "undulator_gap" - ] - self.params.artemis_params.ispyb_params.synchrotron_mode = doc["data"][ - "synchrotron_machine_status_synchrotron_mode" - ] - self.params.artemis_params.ispyb_params.slit_gap_size_x = doc["data"][ - "s4_slit_gaps_xgap" - ] - self.params.artemis_params.ispyb_params.slit_gap_size_y = doc["data"][ - "s4_slit_gaps_ygap" - ] - - LOGGER.info("Creating ispyb entry.") - self.ispyb_ids = self.ispyb.begin_deposition() - set_dcgid_tag(self.ispyb_ids[1]) + super().event(doc) + set_dcgid_tag(self.ispyb_ids[1]) def stop(self, doc: dict): if doc.get("run_start") == self.uid_to_finalize_on: - LOGGER.debug("ISPyB handler received stop document.") - exit_status = doc.get("exit_status") - reason = doc.get("reason") if self.ispyb_ids == (None, None): raise ISPyBDepositionNotMade( "ISPyB deposition was not initialised at run start!" ) - self.ispyb.end_deposition(exit_status, reason) - set_dcgid_tag(None) + super().stop(doc) From 2e4ca3cb019f6b3f0b0f0b3c8950f7a19bf9b8b4 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 25 Jul 2023 14:28:43 +0100 Subject: [PATCH 33/42] get rid of class variables in storeinispyb --- .../callbacks/fgs/ispyb_callback.py | 4 +-- .../ispyb/store_in_ispyb.py | 36 ++++++++----------- 2 files changed, 17 insertions(+), 23 deletions(-) diff --git a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py index 5a5bd6773..2d48f7639 100644 --- a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py @@ -21,8 +21,8 @@ class FGSISPyBHandlerCallback(BaseISPyBHandlerCallback): To use, subscribe the Bluesky RunEngine to an instance of this class. E.g.: - nexus_file_handler_callback = NexusFileHandlerCallback(parameters) - RE.subscribe(nexus_file_handler_callback) + ispyb_handler_callback = FGSISPyBHandlerCallback(parameters) + RE.subscribe(ispyb_handler_callback) Or decorate a plan using bluesky.preprocessors.subs_decorator. See: https://blueskyproject.io/bluesky/callbacks.html#ways-to-invoke-callbacks diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index a829c97e3..a4caa7848 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -36,15 +36,14 @@ class StoreInIspyb(ABC): - ispyb_params: IspybParams | None = None - detector_params: DetectorParams | None = None - run_number: int | None = None - omega_start: int | None = None - experiment_type: str | None = None - xtal_snapshots: list[str] | None = None - data_collection_group_id: int | None = None - def __init__(self, ispyb_config: str, experiment_type: str) -> None: + self.ispyb_params: IspybParams | None = None + self.detector_params: DetectorParams | None = None + self.run_number: int | None = None + self.omega_start: int | None = None + self.experiment_type: str | None = None + self.xtal_snapshots: list[str] | None = None + self.data_collection_group_id: int | None = None self.ISPYB_CONFIG_PATH: str = ispyb_config self.experiment_type = experiment_type @@ -248,16 +247,12 @@ def _store_data_collection_table( class StoreRotationInIspyb(StoreInIspyb): - ispyb_params: RotationIspybParams - data_collection_id: int | None = None - def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None: self.full_params: RotationInternalParameters = parameters - self.ispyb_params: RotationInternalParameters = ( - parameters.artemis_params.ispyb_params - ) + self.ispyb_params: RotationIspybParams = parameters.artemis_params.ispyb_params self.detector_params = parameters.artemis_params.detector_params self.omega_start = self.detector_params.omega_start + self.data_collection_id: int | None = None if self.ispyb_params.xtal_snapshots_omega_start: self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_start @@ -304,13 +299,6 @@ def _construct_comment(self) -> str: class StoreGridscanInIspyb(StoreInIspyb): - ispyb_params: GridscanIspybParams | None = None - data_collection_ids: tuple[int, ...] | None = None - upper_left: list[int] | None = None - y_steps: int | None = None - y_step_size: int | None = None - grid_ids: tuple[int, ...] | None = None - def __init__( self, ispyb_config: str, @@ -318,6 +306,12 @@ def __init__( parameters: FGSInternalParameters = None, ) -> None: self.full_params: FGSInternalParameters | None = parameters + self.ispyb_params: GridscanIspybParams | None = None + self.data_collection_ids: tuple[int, ...] | None = None + self.upper_left: list[int] | None = None + self.y_steps: int | None = None + self.y_step_size: int | None = None + self.grid_ids: tuple[int, ...] | None = None super().__init__(ispyb_config, experiment_type) def begin_deposition(self): From 2fa6ebc583e622fbcf7ac9e49be36fe37bd6504b Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 25 Jul 2023 14:49:08 +0100 Subject: [PATCH 34/42] more artemis param key definitions to superclass --- src/artemis/parameters/internal_parameters.py | 15 ++++++++--- .../plan_specific/fgs_internal_params.py | 22 ++++++---------- .../rotation_scan_internal_params.py | 25 ++++++------------- 3 files changed, 27 insertions(+), 35 deletions(-) diff --git a/src/artemis/parameters/internal_parameters.py b/src/artemis/parameters/internal_parameters.py index 3ee228d09..d9be75f0e 100644 --- a/src/artemis/parameters/internal_parameters.py +++ b/src/artemis/parameters/internal_parameters.py @@ -135,10 +135,19 @@ def _preprocess_all(cls, values): values["artemis_params"] = flatten_dict(values) return values - @staticmethod @abstractmethod - def _artemis_param_key_definitions(): - ... + def _artemis_param_key_definitions(self): + artemis_param_field_keys = [ + "zocalo_environment", + "beamline", + "insertion_prefix", + "experiment_type", + ] + detector_field_keys = list(DetectorParams.__annotations__.keys()) + # not an annotation but specified as field encoder in DetectorParams: + detector_field_keys.append("detector") + ispyb_field_keys = list(IspybParams.__annotations__.keys()) + return artemis_param_field_keys, detector_field_keys, ispyb_field_keys @abstractmethod def _preprocess_experiment_params( diff --git a/src/artemis/parameters/plan_specific/fgs_internal_params.py b/src/artemis/parameters/plan_specific/fgs_internal_params.py index c866dbc33..04b68ace8 100644 --- a/src/artemis/parameters/plan_specific/fgs_internal_params.py +++ b/src/artemis/parameters/plan_specific/fgs_internal_params.py @@ -45,21 +45,13 @@ class Config: **GridscanArtemisParameters.Config.json_encoders, } - @staticmethod - def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: - artemis_param_field_keys = [ - "zocalo_environment", - "beamline", - "insertion_prefix", - "experiment_type", - ] - detector_field_keys = list(DetectorParams.__annotations__.keys()) - # not an annotation but specified as field encoder in DetectorParams: - detector_field_keys.append("detector") - ispyb_field_keys = list(IspybParams.__annotations__.keys()) + list( - GridscanIspybParams.__annotations__.keys() - ) - + def _artemis_param_key_definitions(self) -> tuple[list[str], list[str], list[str]]: + ( + artemis_param_field_keys, + detector_field_keys, + ispyb_field_keys, + ) = super()._artemis_param_key_definitions() + ispyb_field_keys += list(GridscanIspybParams.__annotations__.keys()) return artemis_param_field_keys, detector_field_keys, ispyb_field_keys @validator("experiment_params", pre=True) diff --git a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py index 5d63253f1..faa4fb015 100644 --- a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py +++ b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py @@ -13,8 +13,6 @@ from artemis.external_interaction.ispyb.ispyb_dataclass import ( GRIDSCAN_ISPYB_PARAM_DEFAULTS, - GridscanIspybParams, - IspybParams, RotationIspybParams, ) from artemis.parameters.internal_parameters import ( @@ -34,7 +32,7 @@ class Config: arbitrary_types_allowed = True json_encoders = { **DetectorParams.Config.json_encoders, - **GridscanIspybParams.Config.json_encoders, + **RotationIspybParams.Config.json_encoders, } @@ -94,20 +92,13 @@ class Config: **RotationArtemisParameters.Config.json_encoders, } - @staticmethod - def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: - artemis_param_field_keys = [ - "zocalo_environment", - "beamline", - "insertion_prefix", - "experiment_type", - ] - detector_field_keys = list(DetectorParams.__annotations__.keys()) - # not an annotation but specified as field encoder in DetectorParams: - detector_field_keys.append("detector") - ispyb_field_keys = list(IspybParams.__annotations__.keys()) + list( - GridscanIspybParams.__annotations__.keys() - ) + def _artemis_param_key_definitions(self) -> tuple[list[str], list[str], list[str]]: + ( + artemis_param_field_keys, + detector_field_keys, + ispyb_field_keys, + ) = super()._artemis_param_key_definitions() + ispyb_field_keys += list(RotationIspybParams.__annotations__.keys()) return artemis_param_field_keys, detector_field_keys, ispyb_field_keys From e266a37df721df1a8f0ca02a07308355f73cbcdc Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 10:04:43 +0100 Subject: [PATCH 35/42] fixes and test fixes --- .../experiment_plans/tests/conftest.py | 3 +- .../tests/test_fast_grid_scan_plan.py | 50 ++++++++++++++++++- .../tests/test_rotation_scan_plan.py | 15 ++++-- .../callbacks/fgs/ispyb_callback.py | 2 +- .../fgs/tests/test_zocalo_handler.py | 20 ++++++-- .../callbacks/ispyb_callback_base.py | 4 ++ .../ispyb/store_in_ispyb.py | 6 +-- .../unit_tests/test_store_in_ispyb.py | 7 ++- src/artemis/parameters/internal_parameters.py | 3 +- .../plan_specific/fgs_internal_params.py | 6 +-- .../grid_scan_with_edge_detect_params.py | 24 +++------ .../rotation_scan_internal_params.py | 5 +- src/artemis/system_tests/test_main_system.py | 2 +- test_parameter_defaults.json | 2 +- 14 files changed, 107 insertions(+), 42 deletions(-) diff --git a/src/artemis/experiment_plans/tests/conftest.py b/src/artemis/experiment_plans/tests/conftest.py index 242d1f7db..6c48a339c 100644 --- a/src/artemis/experiment_plans/tests/conftest.py +++ b/src/artemis/experiment_plans/tests/conftest.py @@ -14,6 +14,7 @@ from artemis.external_interaction.callbacks.rotation.rotation_callback_collection import ( RotationCallbackCollection, ) +from artemis.external_interaction.ispyb.store_in_ispyb import Store3DGridscanInIspyb from artemis.external_interaction.system_tests.conftest import TEST_RESULT_LARGE from artemis.parameters.external_parameters import from_file as raw_params_from_file from artemis.parameters.internal_parameters import InternalParameters @@ -175,7 +176,7 @@ def mock_subscriptions(test_fgs_params): subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( TEST_RESULT_LARGE ) - subscriptions.ispyb_handler.ispyb = MagicMock() + subscriptions.ispyb_handler.ispyb = MagicMock(spec=Store3DGridscanInIspyb) subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: [[0, 0], 0, 0] return subscriptions diff --git a/src/artemis/experiment_plans/tests/test_fast_grid_scan_plan.py b/src/artemis/experiment_plans/tests/test_fast_grid_scan_plan.py index a5da658a9..278317f0e 100644 --- a/src/artemis/experiment_plans/tests/test_fast_grid_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_fast_grid_scan_plan.py @@ -32,6 +32,7 @@ from artemis.external_interaction.callbacks.logging_callback import ( VerbosePlanExecutionLoggingCallback, ) +from artemis.external_interaction.ispyb.store_in_ispyb import Store3DGridscanInIspyb from artemis.external_interaction.system_tests.conftest import ( TEST_RESULT_LARGE, TEST_RESULT_MEDIUM, @@ -39,6 +40,7 @@ ) from artemis.log import set_up_logging_handlers from artemis.parameters import external_parameters +from artemis.parameters.constants import ISPYB_PLAN_NAME from artemis.parameters.plan_specific.fgs_internal_params import FGSInternalParameters @@ -91,7 +93,7 @@ def test_read_hardware_for_ispyb_updates_from_ophyd_devices( fake_fgs_composite.flux.flux_reading.sim_put(flux_test_value) test_ispyb_callback = FGSISPyBHandlerCallback(test_fgs_params) - test_ispyb_callback.ispyb = MagicMock() + test_ispyb_callback.ispyb = MagicMock(spec=Store3DGridscanInIspyb) RE.subscribe(test_ispyb_callback) def standalone_read_hardware_for_ispyb(und, syn, slits, attn, fl): @@ -135,6 +137,22 @@ def test_results_adjusted_and_passed_to_move_xyz( set_up_logging_handlers(logging_level="INFO", dev_mode=True) RE.subscribe(VerbosePlanExecutionLoggingCallback()) + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "123abc", "name": ISPYB_PLAN_NAME} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "123abc", + "data": { + "undulator_gap": 0, + "synchrotron_machine_status_synchrotron_mode": 0, + "s4_slit_gaps_xgap": 0, + "s4_slit_gaps_ygap": 0, + "attenuator_actual_transmission": 0, + }, + } + ) + mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( TEST_RESULT_LARGE ) @@ -214,6 +232,21 @@ def test_individual_plans_triggered_once_and_only_once_in_composite_run( test_fgs_params: FGSInternalParameters, RE: RunEngine, ): + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "123abc", "name": ISPYB_PLAN_NAME} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "123abc", + "data": { + "undulator_gap": 0, + "synchrotron_machine_status_synchrotron_mode": 0, + "s4_slit_gaps_xgap": 0, + "s4_slit_gaps_ygap": 0, + "attenuator_actual_transmission": 0, + }, + } + ) set_up_logging_handlers(logging_level="INFO", dev_mode=True) RE.subscribe(VerbosePlanExecutionLoggingCallback()) @@ -248,6 +281,21 @@ def test_logging_within_plan( test_fgs_params: FGSInternalParameters, RE: RunEngine, ): + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "123abc", "name": ISPYB_PLAN_NAME} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "123abc", + "data": { + "undulator_gap": 0, + "synchrotron_machine_status_synchrotron_mode": 0, + "s4_slit_gaps_xgap": 0, + "s4_slit_gaps_ygap": 0, + "attenuator_actual_transmission": 0, + }, + } + ) set_up_logging_handlers(logging_level="INFO", dev_mode=True) RE.subscribe(VerbosePlanExecutionLoggingCallback()) diff --git a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py index 02163725b..829ac339c 100644 --- a/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/artemis/experiment_plans/tests/test_rotation_scan_plan.py @@ -208,13 +208,13 @@ def fake_create_devices( eiger.stage = MagicMock() eiger.unstage = MagicMock() mock_omega_sets = MagicMock(return_value=Status(done=True, success=True)) + mock_arm_disarm = MagicMock( - side_effect=zebra.pc.armed.set, return_value=Status(done=True, success=True) + side_effect=zebra.pc.arm.armed.set, return_value=Status(done=True, success=True) ) - zebra.pc.arm_demand.set = mock_arm_disarm + zebra.pc.arm.set = mock_arm_disarm smargon.omega.velocity.set = mock_omega_sets smargon.omega.set = mock_omega_sets - zebra.pc.arm_demand.set = mock_arm_disarm devices = { "eiger": i03.eiger(fake_with_ophyd_sim=True), @@ -226,10 +226,12 @@ def fake_create_devices( return devices -@pytest.mark.s03() +@pytest.mark.s03 @patch("bluesky.plan_stubs.wait") +@patch("artemis.external_interaction.nexus.write_nexus.NexusWriter") def test_ispyb_deposition_in_plan( bps_wait, + nexus_writer, fake_create_devices, RE, test_rotation_params: RotationInternalParameters, @@ -267,11 +269,14 @@ def test_ispyb_deposition_in_plan( "bluesky.preprocessors.__read_and_stash_a_motor", __fake_read, ), + patch( + "artemis.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", + lambda _: callbacks, + ), ): RE( get_plan( test_rotation_params, - callbacks, ) ) diff --git a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py index 2d48f7639..9627d37c6 100644 --- a/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py +++ b/src/artemis/external_interaction/callbacks/fgs/ispyb_callback.py @@ -31,7 +31,7 @@ class FGSISPyBHandlerCallback(BaseISPyBHandlerCallback): """ def __init__(self, parameters: FGSInternalParameters): - super.__init__(parameters) + super().__init__(parameters) self.ispyb: StoreGridscanInIspyb = ( Store3DGridscanInIspyb(self.ispyb_config, self.params) if self.params.experiment_params.is_3d_grid_scan diff --git a/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py b/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py index 7b318de92..5fb0d71fb 100644 --- a/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py +++ b/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock, call +from unittest.mock import MagicMock, call, patch import numpy as np import pytest @@ -75,12 +75,21 @@ def test_execution_of_run_gridscan_triggers_zocalo_calls( callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_not_called() +@patch( + "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + autospec=True, +) def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_called( + store_3d_grid_scan, dummy_params: FGSInternalParameters, ): callbacks = FGSCallbackCollection.from_params(dummy_params) + callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) + callbacks.ispyb_handler.descriptor(td.test_descriptor_document) + callbacks.ispyb_handler.event(td.test_event_document) + mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.ispyb_ids = (0, 0, 100) + callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) expected_centre_grid_coords = np.array([1, 2, 3]) single_crystal_result = [ { @@ -93,8 +102,9 @@ def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_cal callbacks.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( single_crystal_result ) + results = callbacks.zocalo_handler.wait_for_results(np.array([0, 0, 0])) - found_centre = callbacks.zocalo_handler.wait_for_results(np.array([0, 0, 0]))[0] + found_centre = results[0] callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( 100 ) @@ -111,7 +121,7 @@ def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_ ): callbacks = FGSCallbackCollection.from_params(dummy_params) mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.ispyb_ids = (0, 0, 100) + callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) callbacks.zocalo_handler.zocalo_interactor.wait_for_result.side_effect = ( NoDiffractionFound() ) @@ -140,7 +150,7 @@ def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_b ): callbacks = FGSCallbackCollection.from_params(dummy_params) mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.ispyb_ids = (0, 0, 100) + callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) expected_centre_grid_coords = np.array([4, 6, 2]) multi_crystal_result = [ { diff --git a/src/artemis/external_interaction/callbacks/ispyb_callback_base.py b/src/artemis/external_interaction/callbacks/ispyb_callback_base.py index 8ee2bafba..52225e74d 100644 --- a/src/artemis/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/artemis/external_interaction/callbacks/ispyb_callback_base.py @@ -69,6 +69,10 @@ def event(self, doc: dict): LOGGER.info("Creating ispyb entry.") self.ispyb_ids = self.ispyb.begin_deposition() + LOGGER.info( + f"Recieved ISPYB dcids: {self.ispyb_ids[0]}, grid ids: " + f"{self.ispyb_ids[1]}, dc group id: {self.ispyb_ids[2]}" + ) def stop(self, doc: dict): """Subclasses must check that they are recieving a stop document for the correct diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index a4caa7848..bdd870227 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -32,7 +32,7 @@ I03_EIGER_DETECTOR = 78 EIGER_FILE_SUFFIX = "h5" -VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})/" +VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})(/?$)" class StoreInIspyb(ABC): @@ -248,6 +248,7 @@ def _store_data_collection_table( class StoreRotationInIspyb(StoreInIspyb): def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None: + super().__init__(ispyb_config, "SAD") self.full_params: RotationInternalParameters = parameters self.ispyb_params: RotationIspybParams = parameters.artemis_params.ispyb_params self.detector_params = parameters.artemis_params.detector_params @@ -259,7 +260,6 @@ def __init__(self, ispyb_config, parameters: RotationInternalParameters) -> None else: self.xtal_snapshots = [] LOGGER.warning("No xtal snapshot paths sent to ISPyB!") - super().__init__(ispyb_config, "SAD") def _mutate_data_collection_params_for_experiment( self, params: dict[str, Any] @@ -305,6 +305,7 @@ def __init__( experiment_type: str, parameters: FGSInternalParameters = None, ) -> None: + super().__init__(ispyb_config, experiment_type) self.full_params: FGSInternalParameters | None = parameters self.ispyb_params: GridscanIspybParams | None = None self.data_collection_ids: tuple[int, ...] | None = None @@ -312,7 +313,6 @@ def __init__( self.y_steps: int | None = None self.y_step_size: int | None = None self.grid_ids: tuple[int, ...] | None = None - super().__init__(ispyb_config, experiment_type) def begin_deposition(self): ( diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index 6b75e8d46..4d3bc12b4 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -75,13 +75,18 @@ def test_get_current_time_string(dummy_ispyb): "visit_path, expected_match", [ ("/dls/i03/data/2022/cm6477-45/", "cm6477-45"), + ("/dls/i03/data/2022/cm6477-45", "cm6477-45"), ("/dls/i03/data/2022/mx54663-1/", "mx54663-1"), + ("/dls/i03/data/2022/mx54663-1", "mx54663-1"), ("/dls/i03/data/2022/mx53-1/", None), + ("/dls/i03/data/2022/mx53-1", None), ("/dls/i03/data/2022/mx5563-1565/", None), + ("/dls/i03/data/2022/mx5563-1565", None), ], ) def test_regex_string(dummy_ispyb, visit_path: str, expected_match: str): - assert dummy_ispyb.get_visit_string_from_path(visit_path) == expected_match + test_visit_path = dummy_ispyb.get_visit_string_from_path(visit_path) + assert test_visit_path == expected_match @patch("ispyb.open", new_callable=mock_open) diff --git a/src/artemis/parameters/internal_parameters.py b/src/artemis/parameters/internal_parameters.py index d9be75f0e..e5565d9c6 100644 --- a/src/artemis/parameters/internal_parameters.py +++ b/src/artemis/parameters/internal_parameters.py @@ -135,8 +135,9 @@ def _preprocess_all(cls, values): values["artemis_params"] = flatten_dict(values) return values + @staticmethod @abstractmethod - def _artemis_param_key_definitions(self): + def _artemis_param_key_definitions(): artemis_param_field_keys = [ "zocalo_environment", "beamline", diff --git a/src/artemis/parameters/plan_specific/fgs_internal_params.py b/src/artemis/parameters/plan_specific/fgs_internal_params.py index 04b68ace8..0daffcadb 100644 --- a/src/artemis/parameters/plan_specific/fgs_internal_params.py +++ b/src/artemis/parameters/plan_specific/fgs_internal_params.py @@ -12,7 +12,6 @@ from artemis.external_interaction.ispyb.ispyb_dataclass import ( GRIDSCAN_ISPYB_PARAM_DEFAULTS, GridscanIspybParams, - IspybParams, ) from artemis.parameters.internal_parameters import ( ArtemisParameters, @@ -45,12 +44,13 @@ class Config: **GridscanArtemisParameters.Config.json_encoders, } - def _artemis_param_key_definitions(self) -> tuple[list[str], list[str], list[str]]: + @staticmethod + def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: ( artemis_param_field_keys, detector_field_keys, ispyb_field_keys, - ) = super()._artemis_param_key_definitions() + ) = InternalParameters._artemis_param_key_definitions() ispyb_field_keys += list(GridscanIspybParams.__annotations__.keys()) return artemis_param_field_keys, detector_field_keys, ispyb_field_keys diff --git a/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py b/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py index 61e2effd9..d896e4646 100644 --- a/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py +++ b/src/artemis/parameters/plan_specific/grid_scan_with_edge_detect_params.py @@ -8,10 +8,7 @@ from pydantic import validator from pydantic.dataclasses import dataclass -from artemis.external_interaction.ispyb.ispyb_dataclass import ( - GridscanIspybParams, - IspybParams, -) +from artemis.external_interaction.ispyb.ispyb_dataclass import GridscanIspybParams from artemis.parameters.internal_parameters import ( ArtemisParameters, InternalParameters, @@ -53,19 +50,12 @@ def __init__(self, **args): @staticmethod def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: - artemis_param_field_keys = [ - "zocalo_environment", - "beamline", - "insertion_prefix", - "experiment_type", - ] - detector_field_keys = list(DetectorParams.__annotations__.keys()) - # not an annotation but specified as field encoder in DetectorParams: - detector_field_keys.append("detector") - ispyb_field_keys = list(IspybParams.__annotations__.keys()) + list( - GridscanIspybParams.__annotations__.keys() - ) - + ( + artemis_param_field_keys, + detector_field_keys, + ispyb_field_keys, + ) = InternalParameters._artemis_param_key_definitions() + ispyb_field_keys += list(GridscanIspybParams.__annotations__.keys()) return artemis_param_field_keys, detector_field_keys, ispyb_field_keys @validator("experiment_params", pre=True) diff --git a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py index faa4fb015..bd0461092 100644 --- a/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py +++ b/src/artemis/parameters/plan_specific/rotation_scan_internal_params.py @@ -92,12 +92,13 @@ class Config: **RotationArtemisParameters.Config.json_encoders, } - def _artemis_param_key_definitions(self) -> tuple[list[str], list[str], list[str]]: + @staticmethod + def _artemis_param_key_definitions() -> tuple[list[str], list[str], list[str]]: ( artemis_param_field_keys, detector_field_keys, ispyb_field_keys, - ) = super()._artemis_param_key_definitions() + ) = InternalParameters._artemis_param_key_definitions() ispyb_field_keys += list(RotationIspybParams.__annotations__.keys()) return artemis_param_field_keys, detector_field_keys, ispyb_field_keys diff --git a/src/artemis/system_tests/test_main_system.py b/src/artemis/system_tests/test_main_system.py index 6c2164355..a41638e5f 100644 --- a/src/artemis/system_tests/test_main_system.py +++ b/src/artemis/system_tests/test_main_system.py @@ -343,7 +343,7 @@ def test_when_blueskyrunner_initiated_then_plans_are_setup_and_devices_connected smargon.return_value.wait_for_connection.assert_called() s4_slits.return_value.wait_for_connection.assert_called() fast_grid_scan.return_value.wait_for_connection.assert_called() - eiger.return_value.wait_for_connection.assert_not_called() # can't wait on eiger + eiger.return_value.wait_for_connection.assert_called() backlight.return_value.wait_for_connection.assert_called() aperture_scatterguard.return_value.wait_for_connection.assert_called() oav.return_value.wait_for_connection.assert_called() diff --git a/test_parameter_defaults.json b/test_parameter_defaults.json index 186360380..479ed7b52 100644 --- a/test_parameter_defaults.json +++ b/test_parameter_defaults.json @@ -14,7 +14,7 @@ "det_dist_to_beam_converter_path": "src/artemis/unit_tests/test_lookup_table.txt" }, "ispyb_params": { - "visit_path": "", + "visit_path": "/tmp/cm31105-4", "microns_per_pixel_x": 0.0, "microns_per_pixel_y": 0.0, "upper_left": [ From 7eb6a6b55e93071726f719cad5f450c12f95f473 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 10:11:25 +0100 Subject: [PATCH 36/42] mock ispyb connection in zoc tests --- .../callbacks/fgs/tests/test_zocalo_handler.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py b/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py index 5fb0d71fb..ac295babc 100644 --- a/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py +++ b/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py @@ -116,7 +116,12 @@ def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_cal np.testing.assert_array_equal(found_centre, expected_centre_motor_coords) +@patch( + "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + autospec=True, +) def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_THEN_fallback_centre_used( + store_3d_grid_scan, dummy_params, ): callbacks = FGSCallbackCollection.from_params(dummy_params) @@ -135,7 +140,12 @@ def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_ np.testing.assert_array_equal(found_centre, fallback_position) +@patch( + "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + autospec=True, +) def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_exception( + store_3d_grid_scan, dummy_params, ): callbacks = FGSCallbackCollection.from_params(dummy_params) @@ -145,7 +155,12 @@ def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_excepti callbacks.zocalo_handler.start(td.test_do_fgs_start_document) +@patch( + "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + autospec=True, +) def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_bbox_from_first( + store_3d_grid_scan, dummy_params: FGSInternalParameters, ): callbacks = FGSCallbackCollection.from_params(dummy_params) From 822ebd324aa07cafa5c40832ef6729ca72423c54 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 10:45:06 +0100 Subject: [PATCH 37/42] fix mock --- .../callbacks/fgs/tests/test_zocalo_handler.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py b/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py index ac295babc..f8fe77f84 100644 --- a/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py +++ b/src/artemis/external_interaction/callbacks/fgs/tests/test_zocalo_handler.py @@ -76,7 +76,7 @@ def test_execution_of_run_gridscan_triggers_zocalo_calls( @patch( - "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb", autospec=True, ) def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_called( @@ -117,7 +117,7 @@ def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_cal @patch( - "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb", autospec=True, ) def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_THEN_fallback_centre_used( @@ -141,7 +141,7 @@ def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_ @patch( - "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb", autospec=True, ) def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_exception( @@ -156,7 +156,7 @@ def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_excepti @patch( - "artemis.external_interaction.ispyb.store_in_ispyb.Store3DGridscanInIspyb", + "artemis.external_interaction.callbacks.fgs.ispyb_callback.Store3DGridscanInIspyb", autospec=True, ) def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_bbox_from_first( From f6003d34281a30063790347e4bf83ebfa931df08 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 11:33:30 +0100 Subject: [PATCH 38/42] add some unhappy path tests --- .../unit_tests/test_store_in_ispyb.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index 4d3bc12b4..7c758a3d6 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -115,6 +115,35 @@ def test_store_rotation_scan( TEST_DATA_COLLECTION_GROUP_ID, ) + assert dummy_rotation_ispyb.begin_deposition() == ( + TEST_DATA_COLLECTION_IDS[0], + TEST_DATA_COLLECTION_GROUP_ID, + ) + + +@patch("ispyb.open", new_callable=mock_open) +def test_store_rotation_scan_failures( + ispyb_conn, + dummy_rotation_ispyb: StoreRotationInIspyb, + dummy_rotation_params: RotationInternalParameters, +): + ispyb_conn.return_value.mx_acquisition = mock() + ispyb_conn.return_value.core = mock() + + dummy_rotation_ispyb.data_collection_id = None + + with pytest.raises(AssertionError): + dummy_rotation_ispyb.end_deposition("", "") + + with patch("artemis.log.LOGGER.warning", autospec=True) as warning: + dummy_rotation_params.artemis_params.ispyb_params.xtal_snapshots_omega_start = ( + None + ) + ispyb_no_snapshots = StoreRotationInIspyb( # noqa + SIM_ISPYB_CONFIG, dummy_rotation_params + ) + warning.assert_called_once_with("No xtal snapshot paths sent to ISPyB!") + @patch("ispyb.open", new_callable=mock_open) def test_store_grid_scan(ispyb_conn, dummy_ispyb, dummy_params): From 530acdc75743d8838b2b41b59281821e53acbe95 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 13:11:33 +0100 Subject: [PATCH 39/42] add couple more tests --- .../ispyb/store_in_ispyb.py | 3 +- .../unit_tests/test_store_in_ispyb.py | 112 ++++++++++++++++++ 2 files changed, 114 insertions(+), 1 deletion(-) diff --git a/src/artemis/external_interaction/ispyb/store_in_ispyb.py b/src/artemis/external_interaction/ispyb/store_in_ispyb.py index bdd870227..c6bf11dfe 100755 --- a/src/artemis/external_interaction/ispyb/store_in_ispyb.py +++ b/src/artemis/external_interaction/ispyb/store_in_ispyb.py @@ -192,7 +192,6 @@ def _store_data_collection_table( ) params = mx_acquisition.get_data_collection_params() - params = self._mutate_data_collection_params_for_experiment(params) params["visitid"] = session_id params["parentid"] = data_collection_group_id @@ -243,6 +242,8 @@ def _store_data_collection_table( "file_template" ] = f"{self.detector_params.prefix}_{self.run_number}_master.h5" + params = self._mutate_data_collection_params_for_experiment(params) + return mx_acquisition.upsert_data_collection(list(params.values())) diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index 7c758a3d6..e62a1b2db 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -1,4 +1,6 @@ +import json import re +from copy import deepcopy from unittest.mock import MagicMock, Mock, mock_open, patch import numpy as np @@ -26,6 +28,87 @@ TIME_FORMAT_REGEX = r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}" +EMPTY_DATA_COLLECTION_PARAMS = { + "id": None, + "parentid": None, + "visitid": None, + "sampleid": None, + "detectorid": None, + "positionid": None, + "apertureid": None, + "datacollectionnumber": None, + "starttime": None, + "endtime": None, + "runstatus": None, + "axisstart": None, + "axisend": None, + "axisrange": None, + "overlap": None, + "nimages": None, + "startimagenumber": None, + "npasses": None, + "exptime": None, + "imgdir": None, + "imgprefix": None, + "imgsuffix": None, + "imgcontainersubpath": None, + "filetemplate": None, + "wavelength": None, + "resolution": None, + "detectordistance": None, + "xbeam": None, + "ybeam": None, + "comments": None, + "slitgapvertical": None, + "slitgaphorizontal": None, + "transmission": None, + "synchrotronmode": None, + "xtalsnapshot1": None, + "xtalsnapshot2": None, + "xtalsnapshot3": None, + "xtalsnapshot4": None, + "rotationaxis": None, + "phistart": None, + "kappastart": None, + "omegastart": None, + "resolutionatcorner": None, + "detector2theta": None, + "undulatorgap1": None, + "undulatorgap2": None, + "undulatorgap3": None, + "beamsizeatsamplex": None, + "beamsizeatsampley": None, + "avgtemperature": None, + "actualcenteringposition": None, + "beamshape": None, + "focalspotsizeatsamplex": None, + "focalspotsizeatsampley": None, + "polarisation": None, + "flux": None, + "processeddatafile": None, + "datfile": None, + "magnification": None, + "totalabsorbeddose": None, + "binning": None, + "particlediameter": None, + "boxsizectf": None, + "minresolution": None, + "mindefocus": None, + "maxdefocus": None, + "defocusstepsize": None, + "amountastigmatism": None, + "extractsize": None, + "bgradius": None, + "voltage": None, + "objaperture": None, + "c1aperture": None, + "c2aperture": None, + "c3aperture": None, + "c1lens": None, + "c2lens": None, + "c3lens": None, +} + @pytest.fixture def dummy_params(): @@ -89,6 +172,35 @@ def test_regex_string(dummy_ispyb, visit_path: str, expected_match: str): assert test_visit_path == expected_match +@patch("ispyb.open", new_callable=mock_open) +def test_mutate_params( + ispyb_conn, + dummy_rotation_ispyb: StoreRotationInIspyb, + dummy_ispyb_3d: Store3DGridscanInIspyb, + dummy_params: FGSInternalParameters, + dummy_rotation_params: RotationInternalParameters, +): + rotation_dict = deepcopy(EMPTY_DATA_COLLECTION_PARAMS) + fgs_dict = deepcopy(EMPTY_DATA_COLLECTION_PARAMS) + + dummy_ispyb_3d.y_steps = 5 + + rotation_transformed = ( + dummy_rotation_ispyb._mutate_data_collection_params_for_experiment( + rotation_dict + ) + ) + assert rotation_transformed["axis_range"] == 180.0 + assert rotation_transformed["axis_end"] == 180.0 + assert rotation_transformed["n_images"] == 1800 + + fgs_transformed = dummy_ispyb_3d._mutate_data_collection_params_for_experiment( + fgs_dict + ) + assert fgs_transformed["axis_range"] == 0 + assert fgs_transformed["n_images"] == 200 + + @patch("ispyb.open", new_callable=mock_open) def test_store_rotation_scan( ispyb_conn, dummy_rotation_ispyb: StoreRotationInIspyb, dummy_rotation_params From 8cb1c17af6d94ec0bc8d47f6f302988e0a1dce79 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 13:15:19 +0100 Subject: [PATCH 40/42] fix lint --- .../external_interaction/unit_tests/test_store_in_ispyb.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py index e62a1b2db..0f194b597 100644 --- a/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/artemis/external_interaction/unit_tests/test_store_in_ispyb.py @@ -1,4 +1,3 @@ -import json import re from copy import deepcopy from unittest.mock import MagicMock, Mock, mock_open, patch From da82a0cbe0dfa0919586267ce424c80c040eb8c1 Mon Sep 17 00:00:00 2001 From: David Perl Date: Wed, 26 Jul 2023 13:53:10 +0100 Subject: [PATCH 41/42] add test for matching params --- .../tests/test_internal_parameters.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/src/artemis/parameters/tests/test_internal_parameters.py b/src/artemis/parameters/tests/test_internal_parameters.py index bd317ae62..a6de6ffd4 100644 --- a/src/artemis/parameters/tests/test_internal_parameters.py +++ b/src/artemis/parameters/tests/test_internal_parameters.py @@ -24,6 +24,9 @@ FGSInternalParameters, GridscanArtemisParameters, ) +from artemis.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, +) @pytest.fixture @@ -33,6 +36,23 @@ def raw_params(): ) +@pytest.fixture +def rotation_raw_params(): + return external_parameters.from_file( + "src/artemis/parameters/tests/test_data/good_test_rotation_scan_parameters.json" + ) + + +@pytest.fixture +def gridscan_params(raw_params): + return FGSInternalParameters(**raw_params) + + +@pytest.fixture +def rotation_params(rotation_raw_params): + return RotationInternalParameters(**rotation_raw_params) + + TEST_PARAM_DICT = { "layer_1": { "a": 23, @@ -175,6 +195,38 @@ def test_fetch_subdict(raw_params): assert len(subdict_2) == 3 +def test_param_fields_match_components_they_should_use( + gridscan_params: FGSInternalParameters, + rotation_params: RotationInternalParameters, +): + r_params = rotation_params.artemis_params.ispyb_params + g_params = gridscan_params.artemis_params.ispyb_params + + r_calculated_ispyb_param_keys = list( + rotation_params._artemis_param_key_definitions()[2] + ) + g_calculated_ispyb_param_keys = list( + gridscan_params._artemis_param_key_definitions()[2] + ) + + from artemis.external_interaction.ispyb.ispyb_dataclass import IspybParams + + base_ispyb_annotation_keys = list(IspybParams.__annotations__.keys()) + r_ispyb_annotation_keys = list(r_params.__class__.__annotations__.keys()) + g_ispyb_annotation_keys = list(g_params.__class__.__annotations__.keys()) + + assert ( + r_calculated_ispyb_param_keys + == base_ispyb_annotation_keys + r_ispyb_annotation_keys + ) + assert ( + g_calculated_ispyb_param_keys + == base_ispyb_annotation_keys + g_ispyb_annotation_keys + ) + assert "upper_left" in g_ispyb_annotation_keys + assert "upper_left" not in r_ispyb_annotation_keys + + def test_internal_params_eq(): params = external_parameters.from_file( "src/artemis/parameters/tests/test_data/good_test_parameters.json" From 2f3a95f9ae2f16a262afba495d576232f263a5ba Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Thu, 27 Jul 2023 11:06:07 +0100 Subject: [PATCH 42/42] (#368) Fix logging in BaseISPyBHandlerCallback --- .../external_interaction/callbacks/ispyb_callback_base.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/artemis/external_interaction/callbacks/ispyb_callback_base.py b/src/artemis/external_interaction/callbacks/ispyb_callback_base.py index 52225e74d..834b7a053 100644 --- a/src/artemis/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/artemis/external_interaction/callbacks/ispyb_callback_base.py @@ -69,10 +69,7 @@ def event(self, doc: dict): LOGGER.info("Creating ispyb entry.") self.ispyb_ids = self.ispyb.begin_deposition() - LOGGER.info( - f"Recieved ISPYB dcids: {self.ispyb_ids[0]}, grid ids: " - f"{self.ispyb_ids[1]}, dc group id: {self.ispyb_ids[2]}" - ) + LOGGER.info(f"Recieved ISPYB IDs: {self.ispyb_ids}") def stop(self, doc: dict): """Subclasses must check that they are recieving a stop document for the correct