mx-bluesky 1.2.0__py3-none-any.whl → 1.4.1a0__py3-none-any.whl

mx_bluesky/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py

@@ -5,10 +5,16 @@ from time import time
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
-from blueapi.core import MsgGenerator
 from bluesky import preprocessors as bpp
-from dodal.devices.zocalo.zocalo_results import ZOCALO_READING_PLAN_NAME
+from bluesky.utils import MsgGenerator
+from dodal.devices.zocalo.zocalo_results import (
+    ZOCALO_READING_PLAN_NAME,
+    get_processing_results_from_event,
+)
 
+from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample
+from mx_bluesky.common.parameters.constants import PlanNameConstants
+from mx_bluesky.common.utils.log import set_dcgid_tag
 from mx_bluesky.hyperion.external_interaction.callbacks.common.ispyb_mapping import (
     populate_data_collection_group,
     populate_remaining_data_collection_info,
@@ -36,8 +42,7 @@ from mx_bluesky.hyperion.external_interaction.ispyb.ispyb_store import (
     IspybIds,
     StoreInIspyb,
 )
-from mx_bluesky.hyperion.log import ISPYB_LOGGER, set_dcgid_tag
-from mx_bluesky.hyperion.parameters.components import DiffractionExperimentWithSample
+from mx_bluesky.hyperion.log import ISPYB_LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
 from mx_bluesky.hyperion.parameters.gridscan import (
     GridCommon,
@@ -85,7 +90,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         self._processing_start_time: float | None = None
 
     def activity_gated_start(self, doc: RunStart):
-        if doc.get("subplan_name") == CONST.PLAN.DO_FGS:
+        if doc.get("subplan_name") == PlanNameConstants.DO_FGS:
             self._start_of_fgs_uid = doc.get("uid")
         if doc.get("subplan_name") == CONST.PLAN.GRID_DETECT_AND_DO_GRIDSCAN:
             self.uid_to_finalize_on = doc.get("uid")
@@ -93,7 +98,9 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                 "ISPyB callback received start document with experiment parameters and "
                 f"uid: {self.uid_to_finalize_on}"
             )
-            self.params = GridCommon.from_json(doc.get("hyperion_parameters"))
+            hyperion_params = doc.get("hyperion_parameters")
+            assert isinstance(hyperion_params, str)
+            self.params = GridCommon.model_validate_json(hyperion_params)
             self.ispyb = StoreInIspyb(self.ispyb_config)
             data_collection_group_info = populate_data_collection_group(self.params)
 
@@ -147,7 +154,8 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         ISPYB_LOGGER.info(
             f"Amending comment based on Zocalo reading doc: {format_doc_for_log(doc)}"
         )
-        raw_results = doc["data"]["zocalo-results"]
+
+        raw_results = get_processing_results_from_event("zocalo", doc)
         if len(raw_results) > 0:
             for n, res in enumerate(raw_results):
                 bb = res["bounding_box"]
@@ -178,25 +186,25 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         data = doc["data"]
         data_collection_id = None
         data_collection_info = DataCollectionInfo(
-            xtal_snapshot1=data.get("oav_grid_snapshot_last_path_full_overlay"),
-            xtal_snapshot2=data.get("oav_grid_snapshot_last_path_outer"),
-            xtal_snapshot3=data.get("oav_grid_snapshot_last_saved_path"),
+            xtal_snapshot1=data.get("oav-grid_snapshot-last_path_full_overlay"),
+            xtal_snapshot2=data.get("oav-grid_snapshot-last_path_outer"),
+            xtal_snapshot3=data.get("oav-grid_snapshot-last_saved_path"),
             n_images=(
-                data["oav_grid_snapshot_num_boxes_x"]
-                * data["oav_grid_snapshot_num_boxes_y"]
+                data["oav-grid_snapshot-num_boxes_x"]
+                * data["oav-grid_snapshot-num_boxes_y"]
             ),
         )
-        microns_per_pixel_x = data["oav_grid_snapshot_microns_per_pixel_x"]
-        microns_per_pixel_y = data["oav_grid_snapshot_microns_per_pixel_y"]
+        microns_per_pixel_x = data["oav-microns_per_pixel_x"]
+        microns_per_pixel_y = data["oav-microns_per_pixel_y"]
         data_collection_grid_info = DataCollectionGridInfo(
-            dx_in_mm=data["oav_grid_snapshot_box_width"] * microns_per_pixel_x / 1000,
-            dy_in_mm=data["oav_grid_snapshot_box_width"] * microns_per_pixel_y / 1000,
-            steps_x=data["oav_grid_snapshot_num_boxes_x"],
-            steps_y=data["oav_grid_snapshot_num_boxes_y"],
+            dx_in_mm=data["oav-grid_snapshot-box_width"] * microns_per_pixel_x / 1000,
+            dy_in_mm=data["oav-grid_snapshot-box_width"] * microns_per_pixel_y / 1000,
+            steps_x=data["oav-grid_snapshot-num_boxes_x"],
+            steps_y=data["oav-grid_snapshot-num_boxes_y"],
             microns_per_pixel_x=microns_per_pixel_x,
             microns_per_pixel_y=microns_per_pixel_y,
-            snapshot_offset_x_pixel=int(data["oav_grid_snapshot_top_left_x"]),
-            snapshot_offset_y_pixel=int(data["oav_grid_snapshot_top_left_y"]),
+            snapshot_offset_x_pixel=int(data["oav-grid_snapshot-top_left_x"]),
+            snapshot_offset_y_pixel=int(data["oav-grid_snapshot-top_left_y"]),
             orientation=Orientation.HORIZONTAL,
             snaked=True,
         )

mx_bluesky/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py

@@ -12,7 +12,7 @@ from mx_bluesky.hyperion.external_interaction.nexus.nexus_utils import (
 from mx_bluesky.hyperion.external_interaction.nexus.write_nexus import NexusWriter
 from mx_bluesky.hyperion.log import NEXUS_LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.gridscan import ThreeDGridScan
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
 
 if TYPE_CHECKING:
     from event_model.documents import Event, EventDescriptor, RunStart
@@ -45,11 +45,12 @@ class GridscanNexusFileCallback(PlanReactiveCallback):
 
     def activity_gated_start(self, doc: RunStart):
         if doc.get("subplan_name") == CONST.PLAN.GRIDSCAN_OUTER:
-            json_params = doc.get("hyperion_parameters")
+            hyperion_params = doc.get("hyperion_parameters")
+            assert isinstance(hyperion_params, str)
             NEXUS_LOGGER.info(
-                f"Nexus writer received start document with experiment parameters {json_params}"
+                f"Nexus writer received start document with experiment parameters {hyperion_params}"
             )
-            parameters = ThreeDGridScan.from_json(json_params)
+            parameters = HyperionThreeDGridScan.model_validate_json(hyperion_params)
             d_size = parameters.detector_params.detector_size_constants.det_size_pixels
             grid_n_img_1 = parameters.scan_indices[1]
             grid_n_img_2 = parameters.num_images - grid_n_img_1

mx_bluesky/hyperion/external_interaction/config_server.py

@@ -1,35 +1,18 @@
+from functools import cache
+
 from daq_config_server.client import ConfigServer
-from pydantic import BaseModel
 
+from mx_bluesky.common.external_interaction.config_server import FeatureFlags
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
 
-_CONFIG_SERVER: ConfigServer | None = None
-
-
-def config_server() -> ConfigServer:
-    global _CONFIG_SERVER
-    if _CONFIG_SERVER is None:
-        _CONFIG_SERVER = ConfigServer(CONST.CONFIG_SERVER_URL, LOGGER)
-    return _CONFIG_SERVER
-
-
-class FeatureFlags(BaseModel):
-    # The default value will be used as the fallback when doing a best-effort fetch
-    # from the service
-    use_panda_for_gridscan: bool = False
-    use_gpu_for_gridscan: bool = False
-    set_stub_offsets: bool = False
-
-    @classmethod
-    def _get_flags(cls):
-        flags = config_server().best_effort_get_all_feature_flags()
-        return {f: flags[f] for f in flags if f in cls.__fields__.keys()}
 
-    @classmethod
-    def best_effort(cls):
-        return cls(**cls._get_flags())
+class HyperionFeatureFlags(FeatureFlags):
+    @staticmethod
+    @cache
+    def get_config_server() -> ConfigServer:
+        return ConfigServer(CONST.CONFIG_SERVER_URL, LOGGER)
 
-    def update_self_from_server(self):
-        for flag, value in self._get_flags().items():
-            setattr(self, flag, value)
+    use_panda_for_gridscan: bool = CONST.I03.USE_PANDA_FOR_GRIDSCAN
+    compare_cpu_and_gpu_zocalo: bool = CONST.I03.COMPARE_CPU_AND_GPU_ZOCALO
+    set_stub_offsets: bool = CONST.I03.SET_STUB_OFFSETS

mx_bluesky/hyperion/external_interaction/ispyb/exp_eye_store.py

@@ -35,7 +35,7 @@ class ExpeyeInteraction:
 
     def __init__(self) -> None:
         url, token = _get_base_url_and_token()
-        self.base_url = url + "/core"
+        self.base_url = url
         self.auth = BearerAuth(token)
 
     def _send_and_get_response(self, url, data, send_func) -> dict:

mx_bluesky/hyperion/external_interaction/ispyb/ispyb_store.py

@@ -11,6 +11,7 @@ from ispyb.sp.mxacquisition import MXAcquisition
 from ispyb.strictordereddict import StrictOrderedDict
 from pydantic import BaseModel
 
+from mx_bluesky.common.utils.tracing import TRACER
 from mx_bluesky.hyperion.external_interaction.ispyb.data_model import (
     DataCollectionGridInfo,
     DataCollectionGroupInfo,
@@ -22,7 +23,6 @@ from mx_bluesky.hyperion.external_interaction.ispyb.ispyb_utils import (
     get_session_id_from_visit,
 )
 from mx_bluesky.hyperion.log import ISPYB_LOGGER
-from mx_bluesky.hyperion.tracing import TRACER
 
 if TYPE_CHECKING:
     pass

mx_bluesky/hyperion/external_interaction/nexus/nexus_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import time
-from datetime import datetime, timedelta
+from datetime import UTC, datetime, timedelta
 
 import numpy as np
 from dodal.devices.detector import DetectorParams
@@ -93,7 +93,7 @@ def create_goniometer_axes(
 
 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())
+    start = datetime.fromtimestamp(time.time(), tz=UTC)
     end_est = start + timedelta(seconds=time_expected)
     return start.strftime(time_format), end_est.strftime(time_format)
 

mx_bluesky/hyperion/external_interaction/nexus/write_nexus.py

@@ -15,12 +15,12 @@ from nexgen.nxs_write.nxmx_writer import NXmxFileWriter
 from numpy.typing import DTypeLike
 from scanspec.core import AxesPoints
 
+from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample
 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:

mx_bluesky/hyperion/log.py

@@ -1,23 +1,10 @@
 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)
@@ -26,74 +13,3 @@ 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/components.py

@@ -1,254 +1,7 @@
-from __future__ import annotations
+from pydantic import BaseModel, Field
 
-import json
-from abc import abstractmethod
-from collections.abc import Sequence
-from enum import StrEnum
-from pathlib import Path
-from typing import SupportsInt, TypeVar
+from mx_bluesky.hyperion.external_interaction.config_server import HyperionFeatureFlags
 
-from dodal.devices.aperturescatterguard import ApertureValue
-from dodal.devices.detector import (
-    DetectorParams,
-    TriggerMode,
-)
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-    field_serializer,
-    field_validator,
-    model_validator,
-)
-from scanspec.core import AxesPoints
-from semver import Version
 
-from mx_bluesky.hyperion.external_interaction.config_server import FeatureFlags
-from mx_bluesky.hyperion.parameters.constants import CONST
-
-T = TypeVar("T")
-
-
-class ParameterVersion(Version):
-    @classmethod
-    def _parse(cls, version):
-        if isinstance(version, cls):
-            return version
-        return cls.parse(version)
-
-
-PARAMETER_VERSION = ParameterVersion.parse("5.1.0")
-
-
-class RotationAxis(StrEnum):
-    OMEGA = "omega"
-    PHI = "phi"
-    CHI = "chi"
-    KAPPA = "kappa"
-
-
-class XyzAxis(StrEnum):
-    X = "sam_x"
-    Y = "sam_y"
-    Z = "sam_z"
-
-
-class IspybExperimentType(StrEnum):
-    # Enum values from ispyb column data type
-    SAD = "SAD"  # at or slightly above the peak
-    SAD_INVERSE_BEAM = "SAD - Inverse Beam"
-    OSC = "OSC"  # "native" (in the absence of a heavy atom)
-    COLLECT_MULTIWEDGE = (
-        "Collect - Multiwedge"  # "poorly determined" ~ EDNA complex strategy???
-    )
-    MAD = "MAD"
-    HELICAL = "Helical"
-    MULTI_POSITIONAL = "Multi-positional"
-    MESH = "Mesh"
-    BURN = "Burn"
-    MAD_INVERSE_BEAM = "MAD - Inverse Beam"
-    CHARACTERIZATION = "Characterization"
-    DEHYDRATION = "Dehydration"
-    TOMO = "tomo"
-    EXPERIMENT = "experiment"
-    EM = "EM"
-    PDF = "PDF"
-    PDF_BRAGG = "PDF+Bragg"
-    BRAGG = "Bragg"
-    SINGLE_PARTICLE = "single particle"
-    SERIAL_FIXED = "Serial Fixed"
-    SERIAL_JET = "Serial Jet"
-    STANDARD = "Standard"  # Routine structure determination experiment
-    TIME_RESOLVED = "Time Resolved"  # Investigate the change of a system over time
-    DLS_ANVIL_HP = "Diamond Anvil High Pressure"  # HP sample environment pressure cell
-    CUSTOM = "Custom"  # Special or non-standard data collection
-    XRF_MAP = "XRF map"
-    ENERGY_SCAN = "Energy scan"
-    XRF_SPECTRUM = "XRF spectrum"
-    XRF_MAP_XAS = "XRF map xas"
-    MESH_3D = "Mesh3D"
-    SCREENING = "Screening"
-    STILL = "Still"
-    SSX_CHIP = "SSX-Chip"
-    SSX_JET = "SSX-Jet"
-
-    # Aliases for historic hyperion experiment type mapping
-    ROTATION = "SAD"
-    GRIDSCAN_2D = "mesh"
-    GRIDSCAN_3D = "Mesh3D"
-
-
-class HyperionParameters(BaseModel):
-    model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-        extra="allow",
-    )
-
-    def __hash__(self) -> int:
-        return self.json().__hash__()
-
-    features: FeatureFlags = Field(default=FeatureFlags())
-    parameter_model_version: ParameterVersion
-
-    @field_serializer("parameter_model_version")
-    def serialize_parameter_version(self, version: ParameterVersion):
-        return str(version)
-
-    @field_validator("parameter_model_version", mode="before")
-    @classmethod
-    def _validate_version(cls, version_str: str):
-        version = ParameterVersion.parse(version_str)
-        assert (
-            version >= ParameterVersion(major=PARAMETER_VERSION.major)
-        ), f"Parameter version too old! This version of hyperion uses {PARAMETER_VERSION}"
-        assert (
-            version <= ParameterVersion(major=PARAMETER_VERSION.major + 1)
-        ), f"Parameter version too new! This version of hyperion uses {PARAMETER_VERSION}"
-        return version
-
-    @classmethod
-    def from_json(cls, input: str | None):
-        assert input is not None
-        return cls(**json.loads(input))
-
-
-class WithSnapshot(BaseModel):
-    snapshot_directory: Path
-    snapshot_omegas_deg: list[float] | None = None
-
-    @property
-    def take_snapshots(self) -> bool:
-        return bool(self.snapshot_omegas_deg)
-
-
-class WithOptionalEnergyChange(BaseModel):
-    demand_energy_ev: float | None = Field(default=None, gt=0)
-
-
-class WithVisit(BaseModel):
-    visit: str = Field(min_length=1)
-    zocalo_environment: str = Field(default=CONST.ZOCALO_ENV)
-    beamline: str = Field(default=CONST.I03.BEAMLINE, pattern=r"BL\d{2}[BIJS]")
-    det_dist_to_beam_converter_path: str = Field(
-        default=CONST.PARAM.DETECTOR.BEAM_XY_LUT_PATH
-    )
-    insertion_prefix: str = Field(
-        default=CONST.I03.INSERTION_PREFIX, pattern=r"SR\d{2}[BIJS]"
-    )
-    detector_distance_mm: float | None = Field(default=None, gt=0)
-
-
-class DiffractionExperiment(
-    HyperionParameters, WithSnapshot, WithOptionalEnergyChange, WithVisit
-):
-    """For all experiments which use beam"""
-
-    file_name: str
-    exposure_time_s: float = Field(gt=0)
-    comment: str = Field(default="")
-    trigger_mode: TriggerMode = Field(default=TriggerMode.FREE_RUN)
-    run_number: int | None = Field(default=None, ge=0)
-    selected_aperture: ApertureValue | None = Field(default=None)
-    transmission_frac: float = Field(default=0.1)
-    ispyb_experiment_type: IspybExperimentType
-    storage_directory: str
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_snapshot_directory(cls, values):
-        snapshot_dir = values.get(
-            "snapshot_directory", Path(values["storage_directory"], "snapshots")
-        )
-        values["snapshot_directory"] = (
-            snapshot_dir if isinstance(snapshot_dir, Path) else Path(snapshot_dir)
-        )
-        return values
-
-    @property
-    def num_images(self) -> int:
-        return 0
-
-    @property
-    @abstractmethod
-    def detector_params(self) -> DetectorParams: ...
-
-
-class WithScan(BaseModel):
-    """For experiments where the scan is known"""
-
-    @property
-    @abstractmethod
-    def scan_points(self) -> AxesPoints: ...
-
-    @property
-    @abstractmethod
-    def num_images(self) -> int: ...
-
-
-class SplitScan(BaseModel):
-    @property
-    @abstractmethod
-    def scan_indices(self) -> Sequence[SupportsInt]:
-        """Should return the first index of each scan (i.e. for each nexus file)"""
-        ...
-
-
-class WithSample(BaseModel):
-    sample_id: int
-    sample_puck: int | None = None
-    sample_pin: int | None = None
-
-
-class DiffractionExperimentWithSample(DiffractionExperiment, WithSample): ...
-
-
-class WithOavCentring(BaseModel):
-    oav_centring_file: str = Field(default=CONST.I03.OAV_CENTRING_FILE)
-
-
-class OptionalXyzStarts(BaseModel):
-    x_start_um: float | None = None
-    y_start_um: float | None = None
-    z_start_um: float | None = None
-
-
-class XyzStarts(BaseModel):
-    x_start_um: float
-    y_start_um: float
-    z_start_um: float
-
-    def _start_for_axis(self, axis: XyzAxis) -> float:
-        match axis:
-            case XyzAxis.X:
-                return self.x_start_um
-            case XyzAxis.Y:
-                return self.y_start_um
-            case XyzAxis.Z:
-                return self.z_start_um
-
-
-class OptionalGonioAngleStarts(BaseModel):
-    omega_start_deg: float | None = None
-    phi_start_deg: float | None = None
-    chi_start_deg: float | None = None
-    kappa_start_deg: float | None = None
+class WithHyperionFeatures(BaseModel):
+    features: HyperionFeatureFlags = Field(default=HyperionFeatureFlags())