From 90728276a0da8f92604dd6a95c66dddf56e6de5d Mon Sep 17 00:00:00 2001 From: Tom Willemsen Date: Tue, 7 Nov 2023 17:16:57 +0000 Subject: [PATCH] Merge tidy-ups --- deploy/deploy_hyperion.py | 3 --- setup.cfg | 2 +- .../vmxm_flyscan_xray_centre_plan.py | 4 +-- .../callbacks/ispyb_callback_base.py | 6 ++--- .../callbacks/rotation/callback_collection.py | 9 +++++-- .../callbacks/rotation/nexus_callback.py | 10 +++++-- .../xray_centre/callback_collection.py | 16 +++++++++--- .../callbacks/xray_centre/nexus_callback.py | 26 ++++++++++++++----- .../xray_centre/tests/test_nexus_handler.py | 9 ++++--- .../external_interaction/nexus/nexus_utils.py | 14 +--------- .../external_interaction/nexus/write_nexus.py | 18 ++++++++++--- .../unit_tests/test_write_nexus.py | 26 ++++++++++++++++--- 12 files changed, 96 insertions(+), 47 deletions(-) diff --git a/deploy/deploy_hyperion.py b/deploy/deploy_hyperion.py index aaae3d5c5..7daf91c0b 100644 --- a/deploy/deploy_hyperion.py +++ b/deploy/deploy_hyperion.py @@ -67,13 +67,10 @@ def get_hyperion_release_dir_from_args(repo: repo) -> str: if args.beamline == "dev": print("Running as dev") return "/tmp/hyperion_release_test/bluesky" -<<<<<<< HEAD elif args.beamline == "i03": return f"/dls_sw/{args.beamline}/software/bluesky" elif args.beamline == "i02-1": return f"/dls_sw/{args.beamline}/software/bluesky" -======= ->>>>>>> origin/main else: return f"/dls_sw/{args.beamline}/software/bluesky" diff --git a/setup.cfg b/setup.cfg index ada021446..fb801130c 100644 --- a/setup.cfg +++ b/setup.cfg @@ -36,7 +36,7 @@ install_requires = xarray doct databroker - dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@9adc9ee996ff70a1e3f340def3f85e10a318af82 + dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@93923f92ead4cfae30d8f2257b8c8f55815ff207 pydantic<2.0 # See https://github.com/DiamondLightSource/hyperion/issues/774 scipy diff --git a/src/hyperion/experiment_plans/vmxm_flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/vmxm_flyscan_xray_centre_plan.py index ed58e2f72..e0000e437 100755 --- a/src/hyperion/experiment_plans/vmxm_flyscan_xray_centre_plan.py +++ b/src/hyperion/experiment_plans/vmxm_flyscan_xray_centre_plan.py @@ -30,7 +30,7 @@ VmxmFastGridScanCallbackCollection, ) from hyperion.parameters import external_parameters -from hyperion.parameters.constants import ISPYB_PLAN_NAME, SIM_BEAMLINE +from hyperion.parameters.constants import ISPYB_HARDWARE_READ_PLAN, SIM_BEAMLINE from hyperion.tracing import TRACER from hyperion.utils.context import device_composite_from_context, setup_context @@ -90,7 +90,7 @@ def run_gridscan( ): fgs_motors = fgs_composite.fast_grid_scan - yield from bps.create(name=ISPYB_PLAN_NAME) + yield from bps.create(name=ISPYB_HARDWARE_READ_PLAN) yield from bps.read(fgs_composite.synchrotron.machine_status.synchrotron_mode) yield from bps.save() diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index 134247f75..114dc2d11 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -78,9 +78,9 @@ def event(self, doc: dict): "flux_flux_reading", 0.0 ) - LOGGER.info("Creating ispyb entry.") - self.ispyb_ids = self.ispyb.begin_deposition() - LOGGER.info(f"Recieved ISPYB IDs: {self.ispyb_ids}") + LOGGER.info("Creating ispyb entry.") + self.ispyb_ids = self.ispyb.begin_deposition() + 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 diff --git a/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py b/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py index d23be80de..1a8b3a713 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py @@ -15,6 +15,7 @@ from hyperion.external_interaction.callbacks.rotation.zocalo_callback import ( RotationZocaloCallback, ) +from hyperion.external_interaction.nexus.nexus_utils import create_i03_goniometer_axes if TYPE_CHECKING: from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( @@ -25,7 +26,9 @@ @dataclass(frozen=True, order=True) class RotationCallbackCollection(AbstractPlanCallbackCollection): """Groups the callbacks for external interactions for a rotation scan. - Cast to a list to pass it to Bluesky.preprocessors.subs_decorator().""" + Cast to a list to pass it to Bluesky.preprocessors.subs_decorator(). + + Note: specific to i03.""" nexus_handler: RotationNexusFileCallback ispyb_handler: RotationISPyBCallback @@ -33,7 +36,9 @@ class RotationCallbackCollection(AbstractPlanCallbackCollection): @classmethod def from_params(cls, parameters: RotationInternalParameters): - nexus_handler = RotationNexusFileCallback() + nexus_handler = RotationNexusFileCallback( + create_goniometer_axes=create_i03_goniometer_axes + ) ispyb_handler = RotationISPyBCallback(parameters) zocalo_handler = RotationZocaloCallback( parameters.hyperion_params.zocalo_environment, ispyb_handler diff --git a/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py b/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py index 3135da361..21f29ba08 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py @@ -2,7 +2,11 @@ from bluesky.callbacks import CallbackBase -from hyperion.external_interaction.nexus.write_nexus import NexusWriter +from hyperion.external_interaction.nexus.nexus_utils import create_i03_goniometer_axes +from hyperion.external_interaction.nexus.write_nexus import ( + CreateGoniometerProtocol, + NexusWriter, +) from hyperion.log import LOGGER from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, @@ -24,10 +28,11 @@ class RotationNexusFileCallback(CallbackBase): Usually used as part of a RotationCallbackCollection. """ - def __init__(self): + def __init__(self, create_goniometer_axes: CreateGoniometerProtocol): self.run_uid: str | None = None self.parameters: RotationInternalParameters | None = None self.writer: NexusWriter | None = None + self.create_goniometer_axes = create_goniometer_axes def start(self, doc: dict): if doc.get("subplan_name") == "rotation_scan_with_cleanup": @@ -42,5 +47,6 @@ def start(self, doc: dict): self.parameters, self.parameters.get_scan_points(), self.parameters.get_data_shape(), + create_goniometer_func=self.create_goniometer_axes, ) self.writer.create_nexus_file() diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py b/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py index b08f7b7e8..a92eb38eb 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py @@ -16,6 +16,10 @@ from hyperion.external_interaction.callbacks.xray_centre.zocalo_callback import ( XrayCentreZocaloCallback, ) +from hyperion.external_interaction.nexus.nexus_utils import ( + create_i03_goniometer_axes, + create_vmxm_goniometer_axes, +) if TYPE_CHECKING: from hyperion.parameters.internal_parameters import InternalParameters @@ -25,7 +29,9 @@ class XrayCentreCallbackCollection(AbstractPlanCallbackCollection): """Groups the callbacks for external interactions in the fast grid scan, and connects the Zocalo and ISPyB handlers. Cast to a list to pass it to - Bluesky.preprocessors.subs_decorator().""" + Bluesky.preprocessors.subs_decorator(). + + Note: currently specific to i03 as it creates i03's goniometer axes.""" nexus_handler: GridscanNexusFileCallback ispyb_handler: GridscanISPyBCallback @@ -33,7 +39,7 @@ class XrayCentreCallbackCollection(AbstractPlanCallbackCollection): @classmethod def from_params(cls, parameters: InternalParameters): - nexus_handler = GridscanNexusFileCallback() + nexus_handler = GridscanNexusFileCallback(create_i03_goniometer_axes) ispyb_handler = GridscanISPyBCallback(parameters) zocalo_handler = XrayCentreZocaloCallback(parameters, ispyb_handler) callback_collection = cls( @@ -46,14 +52,16 @@ def from_params(cls, parameters: InternalParameters): @dataclass(frozen=True, order=True) class VmxmFastGridScanCallbackCollection(AbstractPlanCallbackCollection): - """Like XRayCentreCallbackCollection, but without zocalo.""" + """Like XRayCentreCallbackCollection, but without zocalo. + + Note: currently specific to VMXm as it creates VMXm's goniometer axes.""" nexus_handler: Gridscan2DNexusFileCallback ispyb_handler: GridscanISPyBCallback @classmethod def from_params(cls, parameters: InternalParameters): - nexus_handler = Gridscan2DNexusFileCallback() + nexus_handler = Gridscan2DNexusFileCallback(create_vmxm_goniometer_axes) ispyb_handler = GridscanISPyBCallback(parameters) callback_collection = cls( nexus_handler=nexus_handler, diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py index facd4a4b6..8d49c97a1 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py @@ -2,7 +2,10 @@ from bluesky.callbacks import CallbackBase -from hyperion.external_interaction.nexus.write_nexus import NexusWriter +from hyperion.external_interaction.nexus.write_nexus import ( + CreateGoniometerProtocol, + NexusWriter, +) from hyperion.log import LOGGER from hyperion.parameters.constants import ISPYB_HARDWARE_READ_PLAN from hyperion.parameters.plan_specific.gridscan_internal_params import ( @@ -29,11 +32,12 @@ class GridscanNexusFileCallback(CallbackBase): Usually used as part of an FGSCallbackCollection. """ - def __init__(self) -> None: + def __init__(self, create_goniometer_func: CreateGoniometerProtocol) -> None: self.parameters: GridscanInternalParameters | None = None self.run_start_uid: str | None = None self.nexus_writer_1: NexusWriter | None = None self.nexus_writer_2: NexusWriter | None = None + self.create_goniometer_func = create_goniometer_func def start(self, doc: dict): if doc.get("subplan_name") == "run_gridscan_move_and_tidy": @@ -57,9 +61,14 @@ def descriptor(self, doc): nexus_data_1 = self.parameters.get_nexus_info(1) LOGGER.info(f"Nexus data 1: {nexus_data_1}") nexus_data_2 = self.parameters.get_nexus_info(2) - self.nexus_writer_1 = NexusWriter(self.parameters, **nexus_data_1) + self.nexus_writer_1 = NexusWriter( + self.parameters, + create_goniometer_func=self.create_goniometer_func, + **nexus_data_1, + ) self.nexus_writer_2 = NexusWriter( self.parameters, + create_goniometer_func=self.create_goniometer_func, **nexus_data_2, vds_start_index=nexus_data_1["data_shape"][0], ) @@ -70,10 +79,11 @@ def descriptor(self, doc): class Gridscan2DNexusFileCallback(CallbackBase): """Similar to above, but for a 2D gridscan""" - def __init__(self) -> None: + def __init__(self, create_goniometer_func: CreateGoniometerProtocol) -> None: self.parameters: GridscanInternalParameters | None = None self.run_start_uid: str | None = None self.nexus_writer: NexusWriter | None = None + self.create_goniometer_func = create_goniometer_func def start(self, doc: dict): if doc.get("subplan_name") == "run_gridscan_move_and_tidy": @@ -85,7 +95,7 @@ def start(self, doc: dict): self.run_start_uid = doc.get("uid") def descriptor(self, doc): - if doc.get("name") == ISPYB_PLAN_NAME: + if doc.get("name") == ISPYB_HARDWARE_READ_PLAN: assert ( self.parameters is not None ), "Nexus callback did not receive parameters before being asked to write!" @@ -96,5 +106,9 @@ def descriptor(self, doc): LOGGER.info("Initialising nexus writer") nexus_data = self.parameters.get_nexus_info(1) LOGGER.info(f"Nexus data: {nexus_data}") - self.nexus_writer = NexusWriter(self.parameters, **nexus_data) + self.nexus_writer = NexusWriter( + self.parameters, + create_goniometer_func=self.create_goniometer_func, + **nexus_data, + ) self.nexus_writer.create_nexus_file() diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py index 376ee401d..4f6380dc1 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py @@ -5,6 +5,7 @@ from hyperion.external_interaction.callbacks.xray_centre.nexus_callback import ( GridscanNexusFileCallback, ) +from hyperion.external_interaction.nexus.nexus_utils import create_i03_goniometer_axes from hyperion.parameters.constants import ISPYB_HARDWARE_READ_PLAN from hyperion.parameters.external_parameters import from_file as default_raw_params from hyperion.parameters.plan_specific.gridscan_internal_params import ( @@ -36,7 +37,7 @@ def test_writers_not_setup_on_plan_start_doc( nexus_writer: MagicMock, dummy_params: GridscanInternalParameters, ): - nexus_handler = GridscanNexusFileCallback() + nexus_handler = GridscanNexusFileCallback(create_i03_goniometer_axes) nexus_writer.assert_not_called() nexus_handler.start( { @@ -51,7 +52,7 @@ def test_writers_dont_create_on_init_but_do_on_ispyb_event( nexus_writer: MagicMock, dummy_params: GridscanInternalParameters, ): - nexus_handler = GridscanNexusFileCallback() + nexus_handler = GridscanNexusFileCallback(create_i03_goniometer_axes) assert nexus_handler.nexus_writer_1 is None assert nexus_handler.nexus_writer_2 is None @@ -85,7 +86,7 @@ def test_writers_do_create_one_file_each_on_start_doc_for_run_gridscan( ): nexus_writer.side_effect = [MagicMock(), MagicMock()] - nexus_handler = GridscanNexusFileCallback() + nexus_handler = GridscanNexusFileCallback(create_i03_goniometer_axes) nexus_handler.start( { "subplan_name": "run_gridscan_move_and_tidy", @@ -113,7 +114,7 @@ def test_writers_do_create_one_file_each_on_start_doc_for_run_gridscan( def test_sensible_error_if_writing_triggered_before_params_received( nexus_writer: MagicMock, dummy_params ): - nexus_handler = GridscanNexusFileCallback() + nexus_handler = GridscanNexusFileCallback(create_i03_goniometer_axes) with pytest.raises(AssertionError) as excinfo: nexus_handler.descriptor( { diff --git a/src/hyperion/external_interaction/nexus/nexus_utils.py b/src/hyperion/external_interaction/nexus/nexus_utils.py index 278cda946..5d450265a 100644 --- a/src/hyperion/external_interaction/nexus/nexus_utils.py +++ b/src/hyperion/external_interaction/nexus/nexus_utils.py @@ -10,8 +10,6 @@ from hyperion.external_interaction.ispyb.ispyb_dataclass import IspybParams -BL = get_beamline_name(None) - def create_i03_goniometer_axes( omega_start: float, @@ -72,7 +70,7 @@ def create_vmxm_goniometer_axes( omega_start: float, scan_points: dict | None, x_y_z_increments: tuple[float, float, float] = (0.0, 0.0, 0.0), -): +) -> Goniometer: """Returns a Nexgen 'Goniometer' object with the dependency chain of I03's Smargon goniometer. If scan points is provided these values will be used in preference to those from the params object. @@ -117,16 +115,6 @@ def create_vmxm_goniometer_axes( return Goniometer(gonio_axes, scan_points) -if BL == "i03": - create_goniometer_axes = create_i03_goniometer_axes -elif BL == "i02-1": - create_goniometer_axes = create_vmxm_goniometer_axes -else: - - def create_goniometer_axes(*a, **k): - raise ValueError("Cannot create goniometer axes on unknown beamline") - - def get_start_and_predicted_end_time(time_expected: float) -> tuple[str, str]: time_format = r"%Y-%m-%dT%H:%M:%SZ" start = datetime.utcfromtimestamp(time.time()) diff --git a/src/hyperion/external_interaction/nexus/write_nexus.py b/src/hyperion/external_interaction/nexus/write_nexus.py index 8c029aced..7c330bcfe 100644 --- a/src/hyperion/external_interaction/nexus/write_nexus.py +++ b/src/hyperion/external_interaction/nexus/write_nexus.py @@ -6,6 +6,7 @@ import math from pathlib import Path +from typing import Protocol from nexgen.nxs_utils import Detector, Goniometer, Source from nexgen.nxs_write.NXmxWriter import NXmxFileWriter @@ -13,18 +14,30 @@ from hyperion.external_interaction.nexus.nexus_utils import ( create_beam_and_attenuator_parameters, create_detector_parameters, - create_goniometer_axes, get_start_and_predicted_end_time, ) from hyperion.parameters.internal_parameters import InternalParameters +class CreateGoniometerProtocol(Protocol): + @staticmethod + def __call__( + omega_start: float, + scan_points: dict | None, + x_y_z_increments: tuple[float, float, float] = (0.0, 0.0, 0.0), + chi: float = 0.0, + phi: float = 0.0, + ) -> Goniometer: + ... + + class NexusWriter: def __init__( self, parameters: InternalParameters, scan_points: dict, data_shape: tuple[int, int, int], + create_goniometer_func: CreateGoniometerProtocol, omega_start: float | None = None, run_number: int | None = None, vds_start_index: int = 0, @@ -66,12 +79,11 @@ def __init__( self.master_file: Path = ( self.directory / f"{self.filename}_{self.run_number}_master.h5" ) - try: chi = parameters.experiment_params.chi_start except Exception: chi = 0.0 - self.goniometer: Goniometer = create_goniometer_axes( + self.goniometer: Goniometer = create_goniometer_func( self.omega_start, self.scan_points, chi=chi ) diff --git a/src/hyperion/external_interaction/unit_tests/test_write_nexus.py b/src/hyperion/external_interaction/unit_tests/test_write_nexus.py index ffc3d22b9..50c3eb31b 100644 --- a/src/hyperion/external_interaction/unit_tests/test_write_nexus.py +++ b/src/hyperion/external_interaction/unit_tests/test_write_nexus.py @@ -7,6 +7,7 @@ import pytest from dodal.devices.fast_grid_scan_common import GridAxis, GridScanParams +from hyperion.external_interaction.nexus.nexus_utils import create_i03_goniometer_axes from hyperion.external_interaction.nexus.write_nexus import NexusWriter from hyperion.parameters.plan_specific.gridscan_internal_params import ( GridscanInternalParameters, @@ -26,10 +27,15 @@ def assert_end_data_correct(nexus_writer: NexusWriter): @pytest.fixture def dummy_nexus_writers(test_fgs_params: GridscanInternalParameters): nexus_info_1 = test_fgs_params.get_nexus_info(1) - nexus_writer_1 = NexusWriter(test_fgs_params, **nexus_info_1) + nexus_writer_1 = NexusWriter( + test_fgs_params, + create_goniometer_func=create_i03_goniometer_axes, + **nexus_info_1, + ) nexus_info_2 = test_fgs_params.get_nexus_info(2) nexus_writer_2 = NexusWriter( test_fgs_params, + create_goniometer_func=create_i03_goniometer_axes, **nexus_info_2, vds_start_index=nexus_info_1["data_shape"][0], ) @@ -49,8 +55,16 @@ def create_nexus_writers_with_many_images(parameters: GridscanInternalParameters parameters.experiment_params.y_steps = y parameters.experiment_params.z_steps = z parameters.hyperion_params.detector_params.num_triggers = x * y + x * z - nexus_writer_1 = NexusWriter(parameters, **parameters.get_nexus_info(1)) - nexus_writer_2 = NexusWriter(parameters, **parameters.get_nexus_info(2)) + nexus_writer_1 = NexusWriter( + parameters, + create_goniometer_func=create_i03_goniometer_axes, + **parameters.get_nexus_info(1), + ) + nexus_writer_2 = NexusWriter( + parameters, + create_goniometer_func=create_i03_goniometer_axes, + **parameters.get_nexus_info(2), + ) yield nexus_writer_1, nexus_writer_2 @@ -71,7 +85,11 @@ def dummy_nexus_writers_with_more_images(test_fgs_params: GridscanInternalParame @pytest.fixture def single_dummy_file(test_fgs_params: GridscanInternalParameters): - nexus_writer = NexusWriter(test_fgs_params, **test_fgs_params.get_nexus_info(1)) + nexus_writer = NexusWriter( + test_fgs_params, + create_goniometer_func=create_i03_goniometer_axes, + **test_fgs_params.get_nexus_info(1), + ) yield nexus_writer for file in [nexus_writer.nexus_file, nexus_writer.master_file]: if os.path.isfile(file):