mx-bluesky 0.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

mx_bluesky/hyperion/external_interaction/ispyb/ispyb_store.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import asdict
+from typing import TYPE_CHECKING
+
+import ispyb
+import ispyb.sqlalchemy
+from ispyb.connector.mysqlsp.main import ISPyBMySQLSPConnector as Connector
+from ispyb.sp.mxacquisition import MXAcquisition
+from ispyb.strictordereddict import StrictOrderedDict
+from pydantic import BaseModel
+
+from mx_bluesky.hyperion.external_interaction.ispyb.data_model import (
+    DataCollectionGridInfo,
+    DataCollectionGroupInfo,
+    DataCollectionInfo,
+    ScanDataInfo,
+)
+from mx_bluesky.hyperion.external_interaction.ispyb.ispyb_utils import (
+    get_current_time_string,
+    get_session_id_from_visit,
+)
+from mx_bluesky.hyperion.log import ISPYB_LOGGER
+from mx_bluesky.hyperion.tracing import TRACER
+
+if TYPE_CHECKING:
+    pass
+
+I03_EIGER_DETECTOR = 78
+EIGER_FILE_SUFFIX = "h5"
+
+
+class IspybIds(BaseModel):
+    data_collection_ids: tuple[int, ...] = ()
+    data_collection_group_id: int | None = None
+    grid_ids: tuple[int, ...] = ()
+
+
+class StoreInIspyb:
+    def __init__(self, ispyb_config: str) -> None:
+        self.ISPYB_CONFIG_PATH: str = ispyb_config
+        self._data_collection_group_id: int | None
+
+    def begin_deposition(
+        self,
+        data_collection_group_info: DataCollectionGroupInfo,
+        scan_data_infos: Sequence[ScanDataInfo],
+    ) -> IspybIds:
+        ispyb_ids = IspybIds()
+        if scan_data_infos[0].data_collection_info:
+            ispyb_ids.data_collection_group_id = scan_data_infos[
+                0
+            ].data_collection_info.parent_id
+
+        return self._begin_or_update_deposition(
+            ispyb_ids, data_collection_group_info, scan_data_infos
+        )
+
+    def update_deposition(
+        self,
+        ispyb_ids,
+        scan_data_infos: Sequence[ScanDataInfo],
+    ) -> IspybIds:
+        assert (
+            ispyb_ids.data_collection_group_id
+        ), "Attempted to store scan data without a collection group"
+        assert (
+            ispyb_ids.data_collection_ids
+        ), "Attempted to store scan data without a collection"
+        return self._begin_or_update_deposition(ispyb_ids, None, scan_data_infos)
+
+    def _begin_or_update_deposition(
+        self,
+        ispyb_ids,
+        data_collection_group_info: DataCollectionGroupInfo | None,
+        scan_data_infos,
+    ) -> IspybIds:
+        with ispyb.open(self.ISPYB_CONFIG_PATH) as conn:
+            assert conn is not None, "Failed to connect to ISPyB"
+            if data_collection_group_info:
+                ispyb_ids.data_collection_group_id = (
+                    self._store_data_collection_group_table(
+                        conn,
+                        data_collection_group_info,
+                        ispyb_ids.data_collection_group_id,
+                    )
+                )
+            else:
+                assert ispyb_ids.data_collection_group_id, "Attempt to update data collection without a data collection group ID"
+
+            grid_ids = list(ispyb_ids.grid_ids)
+            data_collection_ids_out = list(ispyb_ids.data_collection_ids)
+            for scan_data_info in scan_data_infos:
+                data_collection_id = scan_data_info.data_collection_id
+                if (
+                    scan_data_info.data_collection_info
+                    and not scan_data_info.data_collection_info.parent_id
+                ):
+                    scan_data_info.data_collection_info.parent_id = (
+                        ispyb_ids.data_collection_group_id
+                    )
+
+                new_data_collection_id, grid_id = self._store_single_scan_data(
+                    conn, scan_data_info, data_collection_id
+                )
+                if not data_collection_id:
+                    data_collection_ids_out.append(new_data_collection_id)
+                if grid_id:
+                    grid_ids.append(grid_id)
+            ispyb_ids = IspybIds(
+                data_collection_ids=tuple(data_collection_ids_out),
+                grid_ids=tuple(grid_ids),
+                data_collection_group_id=ispyb_ids.data_collection_group_id,
+            )
+        return ispyb_ids
+
+    def end_deposition(self, ispyb_ids: IspybIds, success: str, reason: str):
+        assert (
+            ispyb_ids.data_collection_ids
+        ), "Can't end ISPyB deposition, data_collection IDs are missing"
+        assert (
+            ispyb_ids.data_collection_group_id is not None
+        ), "Cannot end ISPyB deposition without data collection group ID"
+
+        for id_ in ispyb_ids.data_collection_ids:
+            ISPYB_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 = get_current_time_string()
+            self._update_scan_with_end_time_and_status(
+                current_time,
+                run_status,
+                reason,
+                id_,
+                ispyb_ids.data_collection_group_id,
+            )
+
+    def append_to_comment(
+        self, data_collection_id: int, comment: str, delimiter: str = " "
+    ) -> None:
+        with ispyb.open(self.ISPYB_CONFIG_PATH) as conn:
+            assert conn is not None, "Failed to connect to ISPyB!"
+            mx_acquisition: MXAcquisition = conn.mx_acquisition
+            mx_acquisition.update_data_collection_append_comments(
+                data_collection_id, comment, delimiter
+            )
+
+    def _update_scan_with_end_time_and_status(
+        self,
+        end_time: str,
+        run_status: str,
+        reason: str,
+        data_collection_id: int,
+        data_collection_group_id: int,
+    ) -> None:
+        if reason is not None and reason != "":
+            self.append_to_comment(data_collection_id, f"{run_status} reason: {reason}")
+
+        with ispyb.open(self.ISPYB_CONFIG_PATH) as conn:
+            assert conn is not None, "Failed to connect to ISPyB!"
+
+            mx_acquisition: MXAcquisition = conn.mx_acquisition
+
+            params = mx_acquisition.get_data_collection_params()
+            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 _store_position_table(
+        self, conn: Connector, dc_pos_info, data_collection_id
+    ) -> int:
+        mx_acquisition: MXAcquisition = conn.mx_acquisition
+
+        params = mx_acquisition.get_dc_position_params()
+        params["id"] = data_collection_id
+        params |= asdict(dc_pos_info)
+
+        return mx_acquisition.update_dc_position(list(params.values()))
+
+    def _store_data_collection_group_table(
+        self,
+        conn: Connector,
+        dcg_info: DataCollectionGroupInfo,
+        data_collection_group_id: int | None = None,
+    ) -> int:
+        mx_acquisition: MXAcquisition = conn.mx_acquisition
+
+        params = mx_acquisition.get_data_collection_group_params()
+        if data_collection_group_id:
+            params["id"] = data_collection_group_id
+        params["parent_id"] = get_session_id_from_visit(conn, dcg_info.visit_string)
+        params |= {k: v for k, v in asdict(dcg_info).items() if k != "visit_string"}
+
+        return self._upsert_data_collection_group(conn, params)
+
+    def _store_data_collection_table(
+        self, conn, data_collection_id, data_collection_info
+    ):
+        params = self._fill_common_data_collection_params(
+            conn, data_collection_id, data_collection_info
+        )
+        return self._upsert_data_collection(conn, params)
+
+    def _store_single_scan_data(
+        self, conn, scan_data_info, data_collection_id=None
+    ) -> tuple[int, int | None]:
+        data_collection_id = self._store_data_collection_table(
+            conn, data_collection_id, scan_data_info.data_collection_info
+        )
+
+        if scan_data_info.data_collection_position_info:
+            self._store_position_table(
+                conn,
+                scan_data_info.data_collection_position_info,
+                data_collection_id,
+            )
+
+        grid_id = None
+        if scan_data_info.data_collection_grid_info:
+            grid_id = self._store_grid_info_table(
+                conn,
+                data_collection_id,
+                scan_data_info.data_collection_grid_info,
+            )
+        return data_collection_id, grid_id
+
+    def _store_grid_info_table(
+        self,
+        conn: Connector,
+        ispyb_data_collection_id: int,
+        dc_grid_info: DataCollectionGridInfo,
+    ) -> int:
+        mx_acquisition: MXAcquisition = conn.mx_acquisition
+        params = mx_acquisition.get_dc_grid_params()
+        params |= dc_grid_info.as_dict()
+        params["parentid"] = ispyb_data_collection_id
+        return mx_acquisition.upsert_dc_grid(list(params.values()))
+
+    def _fill_common_data_collection_params(
+        self, conn, data_collection_id, data_collection_info: DataCollectionInfo
+    ) -> StrictOrderedDict:
+        mx_acquisition: MXAcquisition = conn.mx_acquisition
+        params = mx_acquisition.get_data_collection_params()
+
+        if data_collection_id:
+            params["id"] = data_collection_id
+        if data_collection_info.visit_string:
+            # This is only needed for populating the DataCollectionGroup
+            params["visit_id"] = get_session_id_from_visit(
+                conn, data_collection_info.visit_string
+            )
+        params |= {
+            k: v for k, v in asdict(data_collection_info).items() if k != "visit_string"
+        }
+
+        return params
+
+    @staticmethod
+    @TRACER.start_as_current_span("_upsert_data_collection_group")
+    def _upsert_data_collection_group(
+        conn: Connector, params: StrictOrderedDict
+    ) -> int:
+        return conn.mx_acquisition.upsert_data_collection_group(list(params.values()))
+
+    @staticmethod
+    @TRACER.start_as_current_span("_upsert_data_collection")
+    def _upsert_data_collection(conn: Connector, params: StrictOrderedDict) -> int:
+        return conn.mx_acquisition.upsert_data_collection(list(params.values()))

mx_bluesky/hyperion/external_interaction/ispyb/ispyb_utils.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import datetime
+import os
+
+from ispyb import NoResult
+from ispyb.connector.mysqlsp.main import ISPyBMySQLSPConnector as Connector
+from ispyb.sp.core import Core
+
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})(/?$)"
+
+
+def get_ispyb_config():
+    return os.environ.get("ISPYB_CONFIG_PATH", CONST.SIM.ISPYB_CONFIG)
+
+
+def get_session_id_from_visit(conn: Connector, visit: str):
+    try:
+        core: Core = conn.core
+        return core.retrieve_visit_id(visit)
+    except NoResult as e:
+        raise NoResult(f"No session ID found in ispyb for visit {visit}") from e
+
+
+def get_current_time_string():
+    now = datetime.datetime.now()
+    return now.strftime("%Y-%m-%d %H:%M:%S")

mx_bluesky/hyperion/external_interaction/nexus/nexus_utils.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import time
+from datetime import datetime, timedelta
+
+import numpy as np
+from dodal.devices.detector import DetectorParams
+from dodal.devices.zebra import RotationDirection
+from nexgen.nxs_utils import Attenuator, Axis, Beam, Detector, EigerDetector, Goniometer
+from nexgen.nxs_utils.axes import TransformationType
+from numpy.typing import DTypeLike
+
+from mx_bluesky.hyperion.log import NEXUS_LOGGER
+from mx_bluesky.hyperion.utils.utils import convert_eV_to_angstrom
+
+
+def vds_type_based_on_bit_depth(detector_bit_depth: int) -> DTypeLike:
+    """Works out the datatype for the VDS, based on the bit depth from the detector."""
+    if detector_bit_depth == 8:
+        return np.uint8
+    elif detector_bit_depth == 16:
+        return np.uint16
+    elif detector_bit_depth == 32:
+        return np.uint32
+    else:
+        NEXUS_LOGGER.error(
+            f"Unknown detector bit depth {detector_bit_depth}, assuming 16-bit"
+        )
+        return np.uint16
+
+
+def create_goniometer_axes(
+    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,
+    rotation_direction: RotationDirection = RotationDirection.NEGATIVE,
+):
+    """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.
+
+    Args:
+        omega_start (float): the starting position of omega, the only extra value that
+                             needs to be specified except for the scan points.
+        scan_points (dict):  a dictionary of points in the scan for each axis. Obtained
+                             by calculating the scan path with scanspec and calling
+                             consume() on it.
+        x_y_z_increments:    optionally, specify the increments between each image for
+                             the x, y, and z axes. Will be ignored if scan_points
+                             is provided.
+    """
+    gonio_axes = [
+        Axis(
+            "omega",
+            ".",
+            TransformationType.ROTATION,
+            (1.0 * rotation_direction.multiplier, 0.0, 0.0),
+            omega_start,
+        ),
+        Axis(
+            name="sam_z",
+            depends="omega",
+            transformation_type=TransformationType.TRANSLATION,
+            vector=(0.0, 0.0, 1.0),
+            start_pos=0.0,
+            increment=x_y_z_increments[2],
+        ),
+        Axis(
+            name="sam_y",
+            depends="sam_z",
+            transformation_type=TransformationType.TRANSLATION,
+            vector=(0.0, 1.0, 0.0),
+            start_pos=0.0,
+            increment=x_y_z_increments[1],
+        ),
+        Axis(
+            name="sam_x",
+            depends="sam_y",
+            transformation_type=TransformationType.TRANSLATION,
+            vector=(1.0, 0.0, 0.0),
+            start_pos=0.0,
+            increment=x_y_z_increments[0],
+        ),
+        Axis(
+            "chi", "sam_x", TransformationType.ROTATION, (0.006, -0.0264, 0.9996), chi
+        ),
+        Axis("phi", "chi", TransformationType.ROTATION, (-1, -0.0025, -0.0056), phi),
+    ]
+    return Goniometer(gonio_axes, scan_points)
+
+
+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())
+    end_est = start + timedelta(seconds=time_expected)
+    return start.strftime(time_format), end_est.strftime(time_format)
+
+
+def create_detector_parameters(detector_params: DetectorParams) -> Detector:
+    """Returns the detector information in a format that nexgen wants.
+
+    Args:
+        detector_params (DetectorParams): The detector params as Hyperion stores them.
+
+    Returns:
+        Detector: Detector description for nexgen.
+    """
+    detector_pixels = detector_params.get_detector_size_pizels()
+
+    eiger_params = EigerDetector(
+        "Eiger 16M", (detector_pixels.height, detector_pixels.width), "Si", 46051, 0
+    )
+
+    detector_axes = [
+        Axis(
+            "det_z",
+            ".",
+            TransformationType.TRANSLATION,
+            (0.0, 0.0, 1.0),
+            detector_params.detector_distance,
+        )
+    ]
+    # Eiger parameters, axes, beam_center, exp_time, [fast, slow]
+    return Detector(
+        eiger_params,
+        detector_axes,
+        list(
+            detector_params.get_beam_position_pixels(detector_params.detector_distance)
+        ),
+        detector_params.exposure_time,
+        [(-1.0, 0.0, 0.0), (0.0, -1.0, 0.0)],
+    )
+
+
+def create_beam_and_attenuator_parameters(
+    energy_kev: float, flux: float, transmission_fraction: float
+) -> tuple[Beam, Attenuator]:
+    """Create beam and attenuator objects that nexgen can understands
+
+    Returns:
+        tuple[Beam, Attenuator]: Descriptions of the beam and attenuator for nexgen.
+    """
+    return (
+        Beam(convert_eV_to_angstrom(energy_kev * 1000), flux),  # pyright: ignore
+        Attenuator(transmission_fraction),  # pyright: ignore
+    )

mx_bluesky/hyperion/external_interaction/nexus/write_nexus.py

@@ -0,0 +1,114 @@
+"""
+Define beamline parameters for I03, Eiger detector and give an example of writing a
+gridscan.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+from dodal.devices.zebra import RotationDirection
+from dodal.utils import get_beamline_name
+from nexgen.nxs_utils import Attenuator, Beam, Detector, Goniometer, Source
+from nexgen.nxs_write.nxmx_writer import NXmxFileWriter
+from numpy.typing import DTypeLike
+from scanspec.core import AxesPoints
+
+from mx_bluesky.hyperion.external_interaction.nexus.nexus_utils import (
+    create_detector_parameters,
+    create_goniometer_axes,
+    get_start_and_predicted_end_time,
+)
+from mx_bluesky.hyperion.parameters.components import DiffractionExperimentWithSample
+
+
+class NexusWriter:
+    def __init__(
+        self,
+        parameters: DiffractionExperimentWithSample,
+        data_shape: tuple[int, int, int],
+        scan_points: AxesPoints,
+        *,
+        run_number: int | None = None,
+        omega_start_deg: float = 0,
+        chi_start_deg: float = 0,
+        phi_start_deg: float = 0,
+        vds_start_index: int = 0,
+        # override default values when there is more than one collection per
+        # detector arming event:
+        full_num_of_images: int | None = None,
+        meta_data_run_number: int | None = None,
+        rotation_direction: RotationDirection = RotationDirection.NEGATIVE,
+    ) -> None:
+        self.beam: Beam | None = None
+        self.attenuator: Attenuator | None = None
+        self.scan_points: dict = scan_points
+        self.data_shape: tuple[int, int, int] = data_shape
+        self.run_number: int = (
+            run_number if run_number else parameters.detector_params.run_number
+        )
+        self.detector: Detector = create_detector_parameters(parameters.detector_params)
+        self.source: Source = Source(get_beamline_name("S03"))
+        self.directory: Path = Path(parameters.storage_directory)
+        self.start_index: int = vds_start_index
+        self.full_num_of_images: int = full_num_of_images or parameters.num_images
+        self.data_filename: str = (
+            f"{parameters.file_name}_{meta_data_run_number}"
+            if meta_data_run_number
+            else parameters.detector_params.full_filename
+        )
+        self.nexus_file: Path = (
+            self.directory / f"{parameters.file_name}_{self.run_number}.nxs"
+        )
+        self.master_file: Path = (
+            self.directory / f"{parameters.file_name}_{self.run_number}_master.h5"
+        )
+        self.goniometer: Goniometer = create_goniometer_axes(
+            omega_start_deg,
+            self.scan_points,
+            chi=chi_start_deg,
+            phi=phi_start_deg,
+            rotation_direction=rotation_direction,
+        )
+
+    def create_nexus_file(self, bit_depth: DTypeLike):
+        """
+        Creates a nexus file based on the parameters supplied when this object was
+        initialised.
+        """
+        start_time, est_end_time = get_start_and_predicted_end_time(
+            self.detector.exp_time * self.full_num_of_images
+        )
+
+        assert self.beam is not None
+        assert self.attenuator is not None
+
+        vds_shape = self.data_shape
+
+        for filename in [self.nexus_file, self.master_file]:
+            NXmx_Writer = NXmxFileWriter(
+                filename,
+                self.goniometer,
+                self.detector,
+                self.source,
+                self.beam,
+                self.attenuator,
+                self.full_num_of_images,
+            )
+            NXmx_Writer.write(
+                image_filename=f"{self.data_filename}",
+                start_time=start_time,
+                est_end_time=est_end_time,
+            )
+            NXmx_Writer.write_vds(
+                vds_offset=self.start_index, vds_shape=vds_shape, vds_dtype=bit_depth
+            )
+
+    def get_image_datafiles(self, max_images_per_file=1000):
+        return [
+            self.directory / f"{self.data_filename}_{h5_num + 1:06}.h5"
+            for h5_num in range(
+                math.ceil(self.full_num_of_images / max_images_per_file)
+            )
+        ]

mx_bluesky/hyperion/log.py

@@ -0,0 +1,99 @@
+import logging
+from logging.handlers import TimedRotatingFileHandler
+from os import environ
+from pathlib import Path
+
+from dodal.log import (
+    ERROR_LOG_BUFFER_LINES,
+    CircularMemoryHandler,
+    DodalLogHandlers,
+    integrate_bluesky_and_ophyd_logging,
+    set_up_all_logging_handlers,
+)
+from dodal.log import LOGGER as dodal_logger
+
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+LOGGER = logging.getLogger("Hyperion")
+LOGGER.setLevel("DEBUG")
+LOGGER.parent = dodal_logger
+__logger_handlers: DodalLogHandlers | None = None
+
+ISPYB_LOGGER = logging.getLogger("Hyperion ISPyB and Zocalo callbacks")
+ISPYB_LOGGER.setLevel(logging.DEBUG)
+
+NEXUS_LOGGER = logging.getLogger("Hyperion NeXus callbacks")
+NEXUS_LOGGER.setLevel(logging.DEBUG)
+
+ALL_LOGGERS = [LOGGER, ISPYB_LOGGER, NEXUS_LOGGER]
+
+
+class ExperimentMetadataTagFilter(logging.Filter):
+    dc_group_id: str | None = None
+    run_uid: str | None = None
+
+    def filter(self, record):
+        if self.dc_group_id:
+            record.dc_group_id = self.dc_group_id
+        if self.run_uid:
+            record.run_uid = self.run_uid
+        return True
+
+
+tag_filter = ExperimentMetadataTagFilter()
+
+
+def set_dcgid_tag(dcgid):
+    """Set the datacollection group id as a tag on all subsequent log messages.
+    Setting to None will remove the tag."""
+    tag_filter.dc_group_id = dcgid
+
+
+def set_uid_tag(uid):
+    tag_filter.run_uid = uid
+
+
+def do_default_logging_setup(dev_mode=False):
+    handlers = set_up_all_logging_handlers(
+        dodal_logger,
+        _get_logging_dir(),
+        "hyperion.log",
+        dev_mode,
+        ERROR_LOG_BUFFER_LINES,
+        CONST.GRAYLOG_PORT,
+    )
+    integrate_bluesky_and_ophyd_logging(dodal_logger)
+    handlers["graylog_handler"].addFilter(tag_filter)
+
+    global __logger_handlers
+    __logger_handlers = handlers
+
+
+def _get_debug_handler() -> CircularMemoryHandler:
+    assert (
+        __logger_handlers is not None
+    ), "You can only use this after running the default logging setup"
+    return __logger_handlers["debug_memory_handler"]
+
+
+def flush_debug_handler() -> str:
+    """Writes the contents of the circular debug log buffer to disk and returns the written filename"""
+    handler = _get_debug_handler()
+    assert isinstance(
+        handler.target, TimedRotatingFileHandler
+    ), "Circular memory handler doesn't have an appropriate fileHandler target"
+    handler.flush()
+    return handler.target.baseFilename
+
+
+def _get_logging_dir() -> Path:
+    """Get the path to write the hyperion log files to.
+
+    If the HYPERION_LOG_DIR environment variable exists then logs will be put in here.
+    If no environment variable is found it will default it to the ./tmp/dev directory.
+
+    Returns:
+        logging_path (Path): Path to the log file for the file handler to write to.
+    """
+    logging_path = Path(environ.get("HYPERION_LOG_DIR") or "./tmp/dev/")
+    return logging_path

mx_bluesky/hyperion/parameters/__init__.py

@@ -0,0 +1,2 @@
+"""This module handles the translation between externally supplied parameters and the
+internal parameter model."""