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

mx_bluesky/beamlines/i24/serial/setup_beamline/setup_zebra_plans.py

@@ -9,8 +9,6 @@ the edm screen, while on the schematics they are 0 indexed. Thus, `Soft In 1` fr
 schematics corresponds to soft_in_2 in the code.
 """
 
-import logging
-
 import bluesky.plan_stubs as bps
 from dodal.devices.zebra import (
     AND3,
@@ -35,6 +33,8 @@ from dodal.devices.zebra import (
     Zebra,
 )
 
+from mx_bluesky.beamlines.i24.serial.log import SSX_LOGGER
+
 # Detector specific outs
 TTL_EIGER = 1
 TTL_PILATUS = 2
@@ -48,8 +48,6 @@ SHUTTER_MODE = {
 GATE_START = 1.0
 SHUTTER_OPEN_TIME = 0.05  # For pp with long delays
 
-logger = logging.getLogger("I24ssx.setup_zebra")
-
 
 def get_zebra_settings_for_extruder(
     exp_time: float,
@@ -70,29 +68,29 @@ def get_zebra_settings_for_extruder(
 
 
 def arm_zebra(zebra: Zebra):
-    yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True)
-    logger.info("Zebra armed.")
+    yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    SSX_LOGGER.info("Zebra armed.")
 
 
 def disarm_zebra(zebra: Zebra):
-    yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True)
-    logger.info("Zebra disarmed.")
+    yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    SSX_LOGGER.info("Zebra disarmed.")
 
 
 def open_fast_shutter(zebra: Zebra):
     yield from bps.abs_set(zebra.inputs.soft_in_2, SoftInState.YES, wait=True)
-    logger.info("Fast shutter open.")
+    SSX_LOGGER.info("Fast shutter open.")
 
 
 def close_fast_shutter(zebra: Zebra):
     yield from bps.abs_set(zebra.inputs.soft_in_2, SoftInState.NO, wait=True)
-    logger.info("Fast shutter closed.")
+    SSX_LOGGER.info("Fast shutter closed.")
 
 
 def set_shutter_mode(zebra: Zebra, mode: str):
     # SOFT_IN:B0 has to be disabled for manual mode
     yield from bps.abs_set(zebra.inputs.soft_in_1, SHUTTER_MODE[mode], wait=True)
-    logger.info(f"Shutter mode set to {mode}.")
+    SSX_LOGGER.info(f"Shutter mode set to {mode}.")
 
 
 def setup_pc_sources(
@@ -124,12 +122,12 @@ def setup_zebra_for_quickshot_plan(
         exp_time (float): Collection exposure time, in s.
         num_images (float): Number of images to be collected.
     """
-    logger.info("Setup ZEBRA for quickshot collection.")
+    SSX_LOGGER.info("Setup ZEBRA for quickshot collection.")
     yield from bps.abs_set(zebra.pc.arm_source, ArmSource.SOFT, group=group)
     yield from setup_pc_sources(zebra, TrigSource.TIME, TrigSource.EXTERNAL)
 
     gate_width = exp_time * num_images + 0.5
-    logger.info(f"Gate start set to {GATE_START}, with width {gate_width}.")
+    SSX_LOGGER.info(f"Gate start set to {GATE_START}, with width {gate_width}.")
     yield from bps.abs_set(zebra.pc.gate_start, GATE_START, group=group)
     yield from bps.abs_set(zebra.pc.gate_width, gate_width, group=group)
 
@@ -138,7 +136,7 @@ def setup_zebra_for_quickshot_plan(
 
     if wait:
         yield from bps.wait(group)
-    logger.info("Finished setting up zebra.")
+    SSX_LOGGER.info("Finished setting up zebra.")
 
 
 def set_logic_gates_for_porto_triggering(
@@ -204,7 +202,7 @@ def setup_zebra_for_extruder_with_pump_probe_plan(
         pulse1_delay (float, optional): Delay to start pulse1 (the laser control) after \
             gate start. Defaults to 0.0.
     """
-    logger.info("Setup ZEBRA for pump probe extruder collection.")
+    SSX_LOGGER.info("Setup ZEBRA for pump probe extruder collection.")
 
     yield from set_shutter_mode(zebra, "manual")
 
@@ -226,7 +224,7 @@ def setup_zebra_for_extruder_with_pump_probe_plan(
     gate_width, gate_step = get_zebra_settings_for_extruder(
         exp_time, pump_exp, pump_delay
     )
-    logger.info(
+    SSX_LOGGER.info(
         f"""
         Gate start set to {GATE_START}, with calculated width {gate_width}
         and step {gate_step}.
@@ -241,13 +239,13 @@ def setup_zebra_for_extruder_with_pump_probe_plan(
     # Settings for extruder pump probe:
     # PULSE1_DLY is the start (0 usually), PULSE1_WID is the laser dwell set on edm
     # PULSE2_DLY is the laser delay set on edm, PULSE2_WID is the exposure time
-    logger.info(
+    SSX_LOGGER.info(
         f"Pulse1 starting at {pulse1_delay} with width set to laser dwell {pump_exp}."
     )
     yield from bps.abs_set(zebra.output.pulse_1.input, PC_GATE, group=group)
     yield from bps.abs_set(zebra.output.pulse_1.delay, pulse1_delay, group=group)
     yield from bps.abs_set(zebra.output.pulse_1.width, pump_exp, group=group)
-    logger.info(
+    SSX_LOGGER.info(
         f"""
         Pulse2 starting at laser delay {pump_delay} with width set to \
         exposure time {exp_time}.
@@ -259,7 +257,7 @@ def setup_zebra_for_extruder_with_pump_probe_plan(
 
     if wait:
         yield from bps.wait(group)
-    logger.info("Finished setting up zebra.")
+    SSX_LOGGER.info("Finished setting up zebra.")
 
 
 def setup_zebra_for_fastchip_plan(
@@ -302,7 +300,7 @@ def setup_zebra_for_fastchip_plan(
         start_time_offset (float): Delay on the start of the position compare. \
             Defaults to 0.0 (standard chip collection).
     """
-    logger.info("Setup ZEBRA for a fixed target collection.")
+    SSX_LOGGER.info("Setup ZEBRA for a fixed target collection.")
 
     yield from set_shutter_mode(zebra, "manual")
 
@@ -341,7 +339,7 @@ def setup_zebra_for_fastchip_plan(
 
     if wait:
         yield from bps.wait(group)
-    logger.info("Finished setting up zebra.")
+    SSX_LOGGER.info("Finished setting up zebra.")
 
 
 def open_fast_shutter_at_each_position_plan(
@@ -371,10 +369,10 @@ def open_fast_shutter_at_each_position_plan(
         num_exposures (int): Number of times data is collected in each aperture.
         exposure_time (float): Exposure time for each shot.
     """
-    logger.info(
+    SSX_LOGGER.info(
         "ZEBRA setup for fastchip collection with long delays between exposures."
     )
-    logger.debug("Controlling the fast shutter on PULSE2.")
+    SSX_LOGGER.debug("Controlling the fast shutter on PULSE2.")
     # Output panel pulse_2 settings
     yield from bps.abs_set(zebra.output.pulse_2.input, PC_GATE, group=group)
     yield from bps.abs_set(zebra.output.pulse_2.delay, 0.0, group=group)
@@ -386,7 +384,7 @@ def open_fast_shutter_at_each_position_plan(
 
     if wait:
         yield from bps.wait(group=group)
-    logger.debug("Finished setting up for long delays.")
+    SSX_LOGGER.debug("Finished setting up for long delays.")
 
 
 def reset_pc_gate_and_pulse(zebra: Zebra, group: str = "reset_pc"):
@@ -444,16 +442,16 @@ def zebra_return_to_normal_plan(
 
     if wait:
         yield from bps.wait(group)
-    logger.info("Zebra settings back to normal.")
+    SSX_LOGGER.info("Zebra settings back to normal.")
 
 
 def reset_zebra_when_collection_done_plan(zebra: Zebra):
     """
     End of collection zebra operations: close fast shutter, disarm and reset settings.
     """
-    logger.debug("Close the fast shutter.")
+    SSX_LOGGER.debug("Close the fast shutter.")
     yield from close_fast_shutter(zebra)
-    logger.debug("Disarm the zebra.")
+    SSX_LOGGER.debug("Disarm the zebra.")
     yield from disarm_zebra(zebra)
-    logger.debug("Set zebra back to normal.")
+    SSX_LOGGER.debug("Set zebra back to normal.")
     yield from zebra_return_to_normal_plan(zebra, wait=True)

mx_bluesky/beamlines/i24/serial/write_nexus.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import pathlib
 import pprint
@@ -9,14 +8,13 @@ from typing import Literal
 import requests
 
 from mx_bluesky.beamlines.i24.serial.fixed_target.ft_utils import ChipType, MappingType
+from mx_bluesky.beamlines.i24.serial.log import SSX_LOGGER
 from mx_bluesky.beamlines.i24.serial.parameters import (
     ExtruderParameters,
     FixedTargetParameters,
 )
 from mx_bluesky.beamlines.i24.serial.setup_beamline import Eiger, caget, cagetstring
 
-logger = logging.getLogger("I24ssx.nexus_writer")
-
 
 def call_nexgen(
     chip_prog_dict: dict | None,
@@ -54,23 +52,23 @@ def call_nexgen(
     )
     t0 = time.time()
     max_wait = 60  # seconds
-    logger.info(f"Watching for {meta_h5}")
+    SSX_LOGGER.info(f"Watching for {meta_h5}")
     while time.time() - t0 < max_wait:
         if meta_h5.exists():
-            logger.info(f"Found {meta_h5} after {time.time() - t0:.1f} seconds")
+            SSX_LOGGER.info(f"Found {meta_h5} after {time.time() - t0:.1f} seconds")
             time.sleep(5)
             break
-        logger.debug(f"Waiting for {meta_h5}")
+        SSX_LOGGER.debug(f"Waiting for {meta_h5}")
         time.sleep(1)
     if not meta_h5.exists():
-        logger.warning(f"Giving up waiting for {meta_h5} after {max_wait} seconds")
+        SSX_LOGGER.warning(f"Giving up waiting for {meta_h5} after {max_wait} seconds")
         return False
 
     transmission = (float(caget(Eiger.pv.transmission)),)
 
     if det_type == Eiger.name:
         bit_depth = int(caget(Eiger.pv.bit_depth))
-        logger.debug(
+        SSX_LOGGER.debug(
             f"Call to nexgen server with the following chip definition: \n{chip_prog_dict}"
         )
 
@@ -96,10 +94,12 @@ def call_nexgen(
             "wavelength": wavelength,
             "bit_depth": bit_depth,
         }
-        logger.info(f"Sending POST request to {url} with payload:")
-        logger.info(pprint.pformat(payload))
+        SSX_LOGGER.info(f"Sending POST request to {url} with payload:")
+        SSX_LOGGER.info(pprint.pformat(payload))
         response = requests.post(url, headers=headers, json=payload)
-        logger.info(f"Response: {response.text} (status code: {response.status_code})")
+        SSX_LOGGER.info(
+            f"Response: {response.text} (status code: {response.status_code})"
+        )
         # the following will raise an error if the request was unsuccessful
         return response.status_code == requests.codes.ok
     return False

mx_bluesky/common/device_setup_plans/read_hardware_for_setup.py

@@ -0,0 +1,14 @@
+import bluesky.plan_stubs as bps
+from dodal.devices.eiger import EigerDetector
+
+from mx_bluesky.common.parameters.constants import DocDescriptorNames
+
+
+def read_hardware_for_zocalo(detector: EigerDetector):
+    """ "
+    If the RunEngine is subscribed to the ZocaloCallback, this plan will also trigger zocalo.
+    A bluesky run must be open to use this plan
+    """
+    yield from bps.create(name=DocDescriptorNames.ZOCALO_HW_READ)
+    yield from bps.read(detector.odin.file_writer.id)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from bps.save()

mx_bluesky/common/external_interaction/config_server.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from functools import cache
+
+from daq_config_server.client import ConfigServer
+from pydantic import BaseModel, Field, model_validator
+
+
+class FeatureFlags(BaseModel, ABC):
+    """Abstract class to use ConfigServer to toggle features for an experiment
+
+    A module wanting to use FeatureFlags should inherit this class, add boolean features
+    as attributes, and implement a get_config_server method, which returns a cached creation of
+    ConfigServer. See HyperionFeatureFlags for an example
+
+    Values supplied upon class instantiation will always take priority over the config server. If connection to the server cannot
+    be made AND values were not supplied, attributes will use their default values
+    """
+
+    # Feature values supplied at construction will override values from the config server
+    overriden_features: dict = Field(default_factory=dict, exclude=True)
+
+    @staticmethod
+    @cache
+    @abstractmethod
+    def get_config_server() -> ConfigServer: ...
+
+    @model_validator(mode="before")
+    @classmethod
+    def mark_overridden_features(cls, values):
+        assert isinstance(values, dict)
+        values["overriden_features"] = values.copy()
+        return values
+
+    def _get_flags(self):
+        flags = type(self).get_config_server().best_effort_get_all_feature_flags()
+        return {f: flags[f] for f in flags if f in self.model_fields.keys()}
+
+    def update_self_from_server(self):
+        """Used to update the feature flags from the server during a plan. Where there are flags which were explicitly set from externally supplied parameters, these values will be used instead."""
+        for flag, value in self._get_flags().items():
+            updated_value = (
+                value
+                if flag not in self.overriden_features.keys()
+                else self.overriden_features[flag]
+            )
+            setattr(self, flag, updated_value)

mx_bluesky/common/parameters/components.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import os
+from abc import abstractmethod
+from collections.abc import Sequence
+from enum import StrEnum
+from pathlib import Path
+from typing import Literal, SupportsInt, cast
+
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import (
+    DetectorParams,
+    TriggerMode,
+)
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+)
+from pydantic_extra_types.semantic_version import SemanticVersion
+from scanspec.core import AxesPoints
+from semver import Version
+
+from mx_bluesky.common.parameters.constants import (
+    TEST_MODE,
+    DetectorParamConstants,
+    GridscanParamConstants,
+)
+
+PARAMETER_VERSION = Version.parse("5.2.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 MxBlueskyParameters(BaseModel):
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    def __hash__(self) -> int:
+        return self.model_dump_json().__hash__()
+
+    parameter_model_version: SemanticVersion
+
+    @field_validator("parameter_model_version")
+    @classmethod
+    def _validate_version(cls, version: Version):
+        assert (
+            version >= Version(major=PARAMETER_VERSION.major)
+        ), f"Parameter version too old! This version of hyperion uses {PARAMETER_VERSION}"
+        assert (
+            version <= Version(major=PARAMETER_VERSION.major + 1)
+        ), f"Parameter version too new! This version of hyperion uses {PARAMETER_VERSION}"
+        return version
+
+
+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):
+    beamline: str = Field(default="BL03I", pattern=r"BL\d{2}[BIJS]")
+    visit: str = Field(min_length=1)
+    det_dist_to_beam_converter_path: str = Field(
+        default=DetectorParamConstants.BEAM_XY_LUT_PATH
+    )
+    insertion_prefix: str = "SR03S" if TEST_MODE else "SR03I"
+    detector_distance_mm: float | None = Field(default=None, gt=0)
+
+
+class DiffractionExperiment(
+    MxBlueskyParameters, 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_directories(cls, values):
+        os.makedirs(values["storage_directory"], exist_ok=True)
+
+        values["snapshot_directory"] = values.get(
+            "snapshot_directory",
+            Path(values["storage_directory"], "snapshots").as_posix(),
+        )
+        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 WithPandaGridScan(BaseModel):
+    """For experiments which use a PandA for constant-motion grid scans"""
+
+    panda_runup_distance_mm: float = Field(
+        default=GridscanParamConstants.PANDA_RUN_UP_DISTANCE_MM
+    )
+
+
+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 MultiXtalSelection(BaseModel):
+    name: str
+
+
+class TopNByMaxCountSelection(MultiXtalSelection):
+    name: Literal["TopNByMaxCount"] = "TopNByMaxCount"  #  pyright: ignore [reportIncompatibleVariableOverride]
+    n: int
+
+
+class WithCentreSelection(BaseModel):
+    select_centres: TopNByMaxCountSelection = Field(
+        discriminator="name", default=TopNByMaxCountSelection(n=1)
+    )
+
+    @property
+    def selection_params(self) -> MultiXtalSelection:
+        """A helper property because pydantic does not allow polymorphism with base classes
+        # only type unions"""
+        cast1 = cast(MultiXtalSelection, self.select_centres)
+        return cast1
+
+
+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