From caf13106c162a176b8a5023a181db5a1b5959770 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Thu, 27 Nov 2025 15:17:51 +1100 Subject: [PATCH 01/17] support partial udc --- .../HardwareObjects/ANSTO/PrefectWorkflow.py | 24 +- .../config/default_params_pins.yml | 9 + .../ANSTO/prefect_flows/partial_udc.py | 205 ++++++++++ .../prefect_flows/schemas/partial_udc.py | 351 ++++++++++++++++++ .../prefect_flows/schemas/prefect_workflow.py | 1 + .../prefect_flows/sync_prefect_client.py | 25 +- mxcubecore/configuration/ansto/config.py | 4 + .../ansto/mockup/prefect_flows.xml | 5 + .../configuration/ansto/prefect_flows.xml | 5 + 9 files changed, 620 insertions(+), 9 deletions(-) create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py diff --git a/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py b/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py index 47057ae911..7cc2e03d4f 100644 --- a/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py +++ b/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py @@ -11,7 +11,6 @@ import redis import yaml from gevent.event import Event -from mx3_beamline_library.devices.motors import md3 from mxcubecore import HardwareRepository as HWR from mxcubecore.BaseHardwareObjects import HardwareObject @@ -22,6 +21,7 @@ from .prefect_flows.full_dataset_collection_flow import FullDatasetFlow from .prefect_flows.grid_scan_flow import GridScanFlow from .prefect_flows.one_shot_flow import OneShotFlow +from .prefect_flows.partial_udc import PartialUDCFlow from .prefect_flows.schemas.prefect_workflow import PrefectFlows from .prefect_flows.screening_flow import ScreeningFlow from .redis_utils import get_redis_connection @@ -494,6 +494,24 @@ def start_prefect_workflow(self, sample_id: int | None) -> None: self.state.value = "ON" raise QueueExecutionException("dialog_box_parameters is empty", self) + elif self.workflow_name == PrefectFlows.partial_udc: + logging.getLogger("HWR").info(f"Starting workflow: {self.workflow_name}") + + self.partial_udc_flow = PartialUDCFlow( + state=self._state, + resolution=self.resolution, + sample_id=sample_id, + ) + dialog_box_parameters = self.open_dialog(self.partial_udc_flow.dialog_box()) + if dialog_box_parameters: + logging.getLogger("HWR").info( + f"Dialog box parameters: {dialog_box_parameters}" + ) + self.partial_udc_flow.run(dialog_box_parameters=dialog_box_parameters) + else: + self.state.value = "ON" + raise QueueExecutionException("dialog_box_parameters is empty", self) + else: logging.getLogger("HWR").error( f"Workflow {self.workflow_name} not supported" @@ -538,6 +556,10 @@ def _save_default_collection_params_to_redis(self) -> None: for key, value in default_params[collection_type].items(): self._save_params_to_redis(collection_type, key, value) + collection_type = "partial_udc" + for key, value in default_params[collection_type].items(): + self._save_params_to_redis(collection_type, key, value) + def _save_params_to_redis( self, collection_type: str, key: str, value: int | str | bool | None ) -> None: diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml index ac10965d74..62c040dc01 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml @@ -35,3 +35,12 @@ one_shot: photon_energy: 13.0 # keV resolution: 2.0 # Ångström (16M) transmission: 0.5 # Percentage + +partial_udc: + # TODO: add more steps later + md3_alignment_y_speed: 1.0 # mm/s + omega_range: 0.0 # Degrees + photon_energy: 13.0 # keV + resolution: 1.5 # Ångström (16M) + transmission: 1.0 # Percentage + detector_roi_mode: "4M" # 4M or disabled diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py new file mode 100644 index 0000000000..308ca79f88 --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -0,0 +1,205 @@ +import logging + +from mx3_beamline_library.devices.beam import energy_master + +from mxcubecore.configuration.ansto.config import settings +from mxcubecore.queue_entry.base_queue_entry import QueueExecutionException + +from ..redis_utils import get_redis_connection +from ..Resolution import Resolution +from .abstract_flow import AbstractPrefectWorkflow +from .schemas.partial_udc import ( + GridScanParams, + OpticalCenteringExtraConfig, + OpticalCenteringParams, + PartialUDCDialogBox, + SingleLoopDataCollectionConfig, +) +from .sync_prefect_client import MX3SyncPrefectClient + + +class PartialUDCFlow(AbstractPrefectWorkflow): + """Prefect Raster Workflow""" + + def __init__( + self, + state, + resolution: Resolution, + sample_id: int | None, + ) -> None: + super().__init__(state, resolution, sample_id=sample_id) + + self._collection_type = "partial_udc" + self.grid_step_map = { + "5x5": (5.0, 5.0), + "10x10": (10.0, 10.0), + "20x20": (20.0, 20.0), + } + + def run(self, dialog_box_parameters: dict) -> None: + """ + + + Parameters + ---------- + dialog_box_parameters : dict + A dictionary containing parameters from the dialog box + + Returns + ------- + None + """ + + self._state.value = "RUNNING" + + dialog_box_model = PartialUDCDialogBox.model_validate(dialog_box_parameters) + head_type = self.get_head_type() + + if not settings.ADD_DUMMY_PIN_TO_DB: + if self.sample_id is None: + logging.getLogger("HWR").info("Getting sample from the data layer...") + sample_id = self.get_sample_id_of_mounted_sample(dialog_box_model) + logging.getLogger("HWR").info(f"Mounted sample id: {sample_id}") + else: + logging.getLogger("HWR").info( + f"Hand-mount mode, sample id: {self.sample_id}" + ) + sample_id = self.sample_id + + else: + logging.getLogger("HWR").warning( + "SIM mode! The sample id will not be obtained from the data layer. " + "Setting sample id to 1. " + "Ensure that this sample id exists in the db before launching the flow" + ) + sample_id = 1 + + photon_energy = energy_master.get() + + with get_redis_connection() as redis_connection: + default_resolution = float(redis_connection.get("grid_scan:resolution")) + + detector_distance = self._resolution_to_distance( + default_resolution, + energy=photon_energy, + ) + logging.getLogger("HWR").info( + f"Detector distance corresponding to {default_resolution} A: {detector_distance} [m]" + ) + + # TODO: partial udc for plates not supported yet + if head_type == "Plate": + use_centring_table = False + msg = "Partial UDC for plates is not supported yet" + logging.getLogger("user_level_log").error(msg) + raise QueueExecutionException(msg, self) + else: + use_centring_table = True + + partial_udc_config = SingleLoopDataCollectionConfig( + optical_centering=OpticalCenteringParams( + beam_position=(612, 512), + extra_config=OpticalCenteringExtraConfig(grid_height_scale_factor=2), + grid_step=self.grid_step_map[dialog_box_model.grid_step], + calibrated_alignment_z=0.85, # TODO: maybe get from redis? + ), + grid_scan=GridScanParams( + omega_range=0, + md3_alignment_y_speed=dialog_box_model.md3_alignment_y_speed, + detector_distance=detector_distance, + photon_energy=photon_energy, + transmission=dialog_box_model.transmission / 100, + crystal_finder_threshold=1, # TODO: can user set this? + number_of_processes=settings.GRID_SCAN_NUMBER_OF_PROCESSES, + ), + screening=None, + full_dataset=None, + mount_pin_at_start_of_flow=False, + add_dummy_pin_to_db=settings.ADD_DUMMY_PIN_TO_DB, + ) + prefect_parameters = { + "sample_id": sample_id, + "pin": {"id": 1, "puck": 1}, + "prepick_pin": {"id": 2, "puck": 1}, + "config": partial_udc_config.model_dump(exclude_none=True), + } + + logging.getLogger("HWR").info( + f"Parameters sent to prefect flow: {prefect_parameters}" + ) + + # Remember the collection params for the next collection + self._save_dialog_box_params_to_redis(dialog_box_model) + + partial_udc_flow = MX3SyncPrefectClient( + name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters + ) + try: + partial_udc_flow.trigger_data_collection(sample_id, mode="partial_udc") + logging.getLogger("user_level_log").info( + "Partial UDC completed successfully." + ) + self._state.value = "ON" + self.mxcubecore_workflow_aborted = False + except Exception as ex: + raise QueueExecutionException(str(ex), self) from ex + + def dialog_box(self) -> dict: + """ + Workflow dialog box. Returns a dictionary that follows a JSON schema + + Returns + ------- + dialog : dict + A dictionary following the JSON schema. + """ + properties = { + "md3_alignment_y_speed": { + "title": "Alignment Y Speed [mm/s]", + "type": "number", + "minimum": 0.1, + "maximum": 14.8, + "default": float(self._get_dialog_box_param("md3_alignment_y_speed")), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, + "maximum": 100, + "default": float(self._get_dialog_box_param("transmission")), + "widget": "textarea", + }, + "detector_roi_mode": { + "title": "Detector ROI Mode", + "type": "string", + "enum": ["4M", "disabled"], + "default": str(self._get_dialog_box_param("detector_roi_mode")), + "widget": "select", + }, + "grid_step": { + "title": "Grid Step [um]", + "type": "string", + "enum": ["5x5", "10x10", "20x20"], + "default": str(self._get_dialog_box_param("grid_step")), + "widget": "select", + }, + } + tray_conditional: dict | None = None + if self.get_head_type() == "Plate": + tray_properties, tray_conditional = self.build_tray_dialog_schema() + properties.update(tray_properties) + + dialog = { + "properties": properties, + "required": [ + "md3_alignment_y_speed", + "transmission", + ], + "dialogName": "Grid Scan Parameters", + } + + if tray_conditional: + dialog.update(tray_conditional) + + return dialog diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py new file mode 100644 index 0000000000..b1f6634a83 --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py @@ -0,0 +1,351 @@ +from enum import StrEnum + +from pydantic import ( + BaseModel, + Field, + PositiveFloat, + PositiveInt, + model_validator, +) + +from .grid_scan import ( + GridScanDialogBox, + GridScanParams, +) + + +class PartialUDCDialogBox(GridScanDialogBox): + grid_step: str = Field( + description="This is set as a string in the UI and is latter mapped to a tuple" + ) + + +# Below are the schemas copied from bluesky_worker single loop data collection flow schema +# Only a subset of the full schema is included here +class OpticalCenteringExtraConfig(BaseModel): + grid_height_scale_factor: float = 2 # Cut the grid height by this amount + + +class OpticalCenteringParams(BaseModel): + beam_position: list[int] | tuple[int, int] = Field( + default=(612, 512), + description="Global default. These are pixel (x,y) coordinates", + ) + grid_step: list[float] | tuple[float, float] = Field( + default=(80, 80), + description="Global default. (x,y) grid step measured in micrometers", + ) + calibrated_alignment_z: float = Field( + default=0.480, + description="Calibration parameter. Should be updated regularly", + ) + extra_config: OpticalCenteringExtraConfig = OpticalCenteringExtraConfig() + + +class GridScanParams(BaseModel): + """Parameters for executing grid scans""" + + omega_range: float = Field(description="Global default. Measured in degrees.") + md3_alignment_y_speed: float = Field(description="Global default. Measured in mm/s") + detector_distance: float = Field(description="Global default. Measured in meters.") + photon_energy: float = Field(description="Global default. Measured in keV.") + transmission: float = Field(strict=True, ge=0, le=1, default=0.1) + crystal_finder_threshold: int = Field(default=1, description="Global default.") + + number_of_processes: int | None = None + + +class DataCollectionBase(BaseModel): + + start_omega: float = Field( + default=0, description="Output from PyBest. Measured in degrees." + ) + omega_range: PositiveFloat = Field( + description="Output from PyBest. Measured in degrees." + ) + number_of_passes: int = 1 + count_time: PositiveFloat | None = None + + detector_distance: PositiveFloat = Field( + default=0.3, description="Output from PyBest. Measured in m." + ) + photon_energy: PositiveFloat = Field( + default=13, description="Global default. Measured in keV." + ) + transmission: float = Field(strict=True, ge=0, le=1, default=0.1) + beam_size: tuple[PositiveInt, PositiveInt] | list[PositiveInt] = Field( + default=(80, 80), + description="Determined by the crystal finder. Not currently used. " + "Measured in um.", + ) + + # Any pair of the following parameters can be defined: + exposure_time: PositiveFloat | None = Field( + default=None, description="Measured in seconds." + ) + number_of_frames: PositiveInt | None = Field( + default=None, + ) + oscillation_speed: PositiveFloat | None = Field( + default=None, + description="Measured in degrees/s", + ) + frame_rate: PositiveFloat | None = Field( + default=None, + description="Measured in Hz", + ) + + @model_validator(mode="after") + def calculate_params(self): # noqa + if len(self.__dict__) == 0: + return self + + exposure_time = self.exposure_time + number_of_frames = self.number_of_frames + oscillation_speed = self.oscillation_speed + frame_rate = self.frame_rate + omega_range = self.omega_range + + # Ensure exactly two parameters are set + self.ensure_params_are_consistent( + exposure_time, number_of_frames, oscillation_speed, frame_rate, omega_range + ) + + # Calculate the missing parameters + if exposure_time is None: + exposure_time = self.calculate_exposure_time( + number_of_frames, oscillation_speed, frame_rate, omega_range + ) + self.exposure_time = exposure_time + if number_of_frames is None: + number_of_frames = self.calculate_number_of_frames( + exposure_time, oscillation_speed, frame_rate, omega_range + ) + self.number_of_frames = number_of_frames + if oscillation_speed is None: + oscillation_speed = self.calculate_oscillation_speed( + exposure_time, number_of_frames, frame_rate, omega_range + ) + self.oscillation_speed = oscillation_speed + if frame_rate is None: + frame_rate = self.calculate_frame_rate(exposure_time, number_of_frames) + self.frame_rate = frame_rate + + return self + + def ensure_params_are_consistent( + self, + exposure_time, + number_of_frames, + oscillation_speed, + frame_rate, + omega_range, + ): + set_params = [] + for param in [ + exposure_time, + number_of_frames, + oscillation_speed, + frame_rate, + ]: + if param is not None: + set_params.append(param) + + if len(set_params) == 4: + if round(exposure_time, 3) != round(number_of_frames / frame_rate, 3): + raise ConsistencyError( + "Exposure time, number of frames, and frame rate are inconsistent." + ) + if round(oscillation_speed, 3) != round( + omega_range * frame_rate / number_of_frames, 3 + ): + raise ConsistencyError( + "Oscillation speed, omega range, and number of frames are " + "inconsistent." + ) + elif len(set_params) == 3: + if exposure_time is None: + if round(oscillation_speed, 3) != round( + omega_range * frame_rate / number_of_frames, 3 + ): + raise ConsistencyError( + "Oscillation speed, omega range, and number of frames are " + "inconsistent." + ) + elif number_of_frames is None: + if round(exposure_time, 3) != round(omega_range / oscillation_speed, 3): + raise ConsistencyError( + "Exposure time, omega range, and oscillation speed are " + "inconsistent." + ) + elif oscillation_speed is None: + if round(exposure_time, 3) != round(number_of_frames / frame_rate, 3): + raise ConsistencyError( + "Exposure time, number of frames, and frame rate are inconsistent." + ) + elif frame_rate is None: + if round(exposure_time, 3) != round(omega_range / oscillation_speed, 3): + raise ConsistencyError( + "Exposure time, number of frames, and oscillation speed are " + "inconsistent." + ) + elif len(set_params) < 1: + raise ConsistencyError( + "Exactly 2 or more parameters of 'exposure_time', 'number_of_frames', " + "'oscillation_speed', 'frame_rate' must be set." + ) + + def calculate_exposure_time( + self, number_of_frames, oscillation_speed, frame_rate, omega_range + ): + if number_of_frames is None and frame_rate is not None: + number_of_frames = omega_range * frame_rate / oscillation_speed + elif number_of_frames and oscillation_speed: + frame_rate = number_of_frames * oscillation_speed / omega_range + return number_of_frames / frame_rate + + def calculate_number_of_frames( + self, exposure_time, oscillation_speed, frame_rate, omega_range + ): + if exposure_time and frame_rate: + return round(exposure_time * frame_rate) + elif exposure_time and oscillation_speed: + raise ValueError( + "Cannot calculate frame rate and number of frames " + "from exposure time and oscillation speed" + ) + + @staticmethod + def calculate_oscillation_speed( + exposure_time, number_of_frames, frame_rate, omega_range + ): + if number_of_frames is not None and exposure_time is not None: + frame_rate = number_of_frames / exposure_time + elif exposure_time is not None and frame_rate is not None: + number_of_frames = exposure_time * frame_rate + return omega_range * frame_rate / number_of_frames + + @staticmethod + def calculate_frame_rate(exposure_time, number_of_frames): + return number_of_frames / exposure_time + + +class ScreeningParams(DataCollectionBase): + """Parameters for collecting screening data""" + + +class FullDatasetParams(DataCollectionBase): + """Parameters for collecting full datasets""" + + rotation_axis_offset: int | None = Field( + default=None, + description="Output from RadDose. Not yet implemented. Measured in um.", + ) + + +class ProjectStrategy(BaseModel): + """Parameters for what to collect on, and when to stop""" + + use_screening: bool = Field(default=True, description="User defined") + use_dose: bool = Field( + default=False, description="User defined. Not yet implemented." + ) + use_rotation_axis_offset: bool = Field( + default=True, description="User defined. Not yet implemented." + ) + + # Full dataset targets + target_av_dose: float | None = Field( + default=None, description="User defined. Not yet implemented. Measured in MGy." + ) + target_completeness: float | None = Field( + default=None, + description="User defined. Measured in percent.", + ) + target_multiplicity: float | None = Field( + default=None, description="User defined. Measured in factor." + ) + target_resolution: float | None = Field( + default=None, description="User defined. Measured in Å." + ) + + # Screening targets + screening_target_resolution: float | None = Field( + default=None, + description="If None, it is automatically calculated as " + "target_resolution + 1 [Å] (if target_resolution is not None).", + ) + + # Project criteria + n_datasets_better_than: int | None = Field( + default=None, description="User defined." + ) + + @model_validator(mode="before") + @classmethod + def set_screening_target_resolution(cls, values: dict): # noqa + """if target_resolution is not None, sets the screening target + resolution to (target_resolution + 1) [Å] + """ + if values.get("screening_target_resolution") is None: + if values.get("target_resolution") is not None: + values["screening_target_resolution"] = values["target_resolution"] + 1 + return values + return values + + @model_validator(mode="before") + @classmethod + def validate_n_datasets_better_than(cls, values: dict): # noqa + """ + Checks that at least one target value has been given + if `n_datasets_better_than` is not None + """ + keys = [ + "target_av_dose", + "target_completeness", + "target_multiplicity", + "target_resolution", + ] + key_vals = [] + if values.get("n_datasets_better_than") is not None: + for key in keys: + key_vals.append(values.get(key)) + if not any(key_vals): + raise ValueError( + "n_datasets_better_than has been specified, but no target values " + "have been given. Update the ProjectStrategy schema." + ) + return values + return values + + +class ProcessingPipeline(StrEnum): + DIALS = "dials" + FAST_DP = "fast_dp" + DIALS_AND_FAST_DP = "dials_and_fast_dp" + DOZOR = "dozor" + FAST_DP_AND_XIA2 = "fast_dp_and_xia2" + XIA2 = "xia2" + + +class SingleLoopDataCollectionConfig(BaseModel): + """Input of the UDC and partial UDC prefect flow""" + + optical_centering: OpticalCenteringParams = OpticalCenteringParams() + grid_scan: GridScanParams + screening: ScreeningParams | None = None + full_dataset: FullDatasetParams | None = None + project_strategy: ProjectStrategy = ProjectStrategy() + processing_pipeline: ProcessingPipeline = ProcessingPipeline.FAST_DP + + mount_pin_at_start_of_flow: bool = Field( + default=True, description="Only must be changed for debugging" + ) + # Parameters used only for debugging purposes + add_dummy_pin_to_db: bool = Field( + default=False, description="Only must be True for debugging" + ) + + +class ConsistencyError(Exception): + pass diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/prefect_workflow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/prefect_workflow.py index 03f7f583dc..66ecbd15b1 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/prefect_workflow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/prefect_workflow.py @@ -6,3 +6,4 @@ class PrefectFlows(str, Enum): collect_dataset = "Collect Dataset" grid_scan = "Grid Scan" one_shot = "One Shot" + partial_udc = "Partial UDC" diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/sync_prefect_client.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/sync_prefect_client.py index e3f78aba43..ab71e121fe 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/sync_prefect_client.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/sync_prefect_client.py @@ -1,5 +1,6 @@ import logging from enum import StrEnum +from typing import Literal from uuid import UUID import redis @@ -100,7 +101,10 @@ def trigger_flow(self, wait=False, poll_interval=3) -> FlowRun: sleep(poll_interval) def trigger_data_collection( - self, sample_id: str, poll_interval: float = 3.0 + self, + sample_id: int, + poll_interval: float = 3.0, + mode: Literal["mxcube", "partial_udc"] = "mxcube", ) -> None: """ Triggers a prefect flow but only waits until data collection is finished @@ -108,16 +112,23 @@ def trigger_data_collection( Parameters ---------- + sample_id : int + The database sample id poll_interval : float, optional The poll interval in seconds, by default 3.0 seconds + mode : Literal["mxcube", "partial_udc"], optional + The mode to use for setting the redis state key, by default "mxcube" Returns ------- None """ - self.redis_connection.set( - f"mxcube_scan_state:{sample_id}", FlowState.RUNNING, ex=3600 - ) + if mode == "mxcube": + state_key = f"mxcube_scan_state:{sample_id}" + else: + state_key = f"state:{sample_id}" + + self.redis_connection.set(state_key, FlowState.RUNNING, ex=3600) with get_client(sync_client=True) as client: deployment_id = client.read_deployment_by_name(self.name).id @@ -126,16 +137,14 @@ def trigger_data_collection( ) self.flow_run_id = response.id - flow_status = self.redis_connection.get(f"mxcube_scan_state:{sample_id}") + flow_status = self.redis_connection.get(state_key) while flow_status not in [ FlowState.FAILED, FlowState.READY_TO_UNMOUNT, FlowState.COMPLETED, ]: sleep(poll_interval) - flow_status = self.redis_connection.get( - f"mxcube_scan_state:{sample_id}" - ) + flow_status = self.redis_connection.get(state_key) flow_run = client.read_flow_run(self.flow_run_id) flow_state = flow_run.state if flow_state and flow_state.type == StateType.FAILED: diff --git a/mxcubecore/configuration/ansto/config.py b/mxcubecore/configuration/ansto/config.py index 620fa177c1..87401e9969 100644 --- a/mxcubecore/configuration/ansto/config.py +++ b/mxcubecore/configuration/ansto/config.py @@ -8,6 +8,10 @@ class PrefectSettings(BaseSettings): PREFECT_URI: str = Field("http://localhost:4200", env="PREFECT_URI") + PARTIAL_UDC_DEPLOYMENT_NAME: str = Field( + "single-loop-data-collection-flow/single_loop_data_collection_flow", + env="PARTIAL_UDC_DEPLOYMENT_NAME", + ) TRAY_CALIBRATION_DEPLOYMENT_NAME: str = Field( "tray-calibration/tray_calibration", env="TRAY_CALIBRATION_DEPLOYMENT_NAME" ) diff --git a/mxcubecore/configuration/ansto/mockup/prefect_flows.xml b/mxcubecore/configuration/ansto/mockup/prefect_flows.xml index 1d629def14..49fc9c6a70 100644 --- a/mxcubecore/configuration/ansto/mockup/prefect_flows.xml +++ b/mxcubecore/configuration/ansto/mockup/prefect_flows.xml @@ -20,4 +20,9 @@ Grid Scan grid + + Partial UDC + Partial UDC + samplegrid + \ No newline at end of file diff --git a/mxcubecore/configuration/ansto/prefect_flows.xml b/mxcubecore/configuration/ansto/prefect_flows.xml index 1d629def14..49fc9c6a70 100644 --- a/mxcubecore/configuration/ansto/prefect_flows.xml +++ b/mxcubecore/configuration/ansto/prefect_flows.xml @@ -20,4 +20,9 @@ Grid Scan grid + + Partial UDC + Partial UDC + samplegrid + \ No newline at end of file From acbccbb461c90e406e4e7b7d34d5bbdb3fb95dfb Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Thu, 27 Nov 2025 16:08:36 +1100 Subject: [PATCH 02/17] use pin default values --- .../HardwareObjects/ANSTO/PrefectWorkflow.py | 4 ---- .../ANSTO/prefect_flows/abstract_flow.py | 13 +++++++++-- .../config/default_params_pins.yml | 10 +-------- .../ANSTO/prefect_flows/partial_udc.py | 22 +++++++++++++++---- 4 files changed, 30 insertions(+), 19 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py b/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py index 7cc2e03d4f..93a0af9b6a 100644 --- a/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py +++ b/mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py @@ -556,10 +556,6 @@ def _save_default_collection_params_to_redis(self) -> None: for key, value in default_params[collection_type].items(): self._save_params_to_redis(collection_type, key, value) - collection_type = "partial_udc" - for key, value in default_params[collection_type].items(): - self._save_params_to_redis(collection_type, key, value) - def _save_params_to_redis( self, collection_type: str, key: str, value: int | str | bool | None ) -> None: diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py index d4850fb158..e8bce5f921 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py @@ -28,6 +28,7 @@ from .schemas.full_dataset import FullDatasetDialogBox from .schemas.grid_scan import GridScanDialogBox from .schemas.one_shot import OneShotDialogBox +from .schemas.prefect_workflow import PrefectFlows from .schemas.screening import ScreeningDialogBox @@ -334,13 +335,14 @@ def _get_dialog_box_param( "lab_name", "project_name", ], + collection_type: str | None = None, ) -> str | int | float | None: """ Retrieve a parameter value from Redis. Parameters ---------- - Literal[ + parameter : Literal[ "exposure_time", "omega_range", "number_of_frames", @@ -355,17 +357,24 @@ def _get_dialog_box_param( "project_name", ] A parameter saved in redis + collection_type : str | None + The collection type. If None, use the current collection type Returns ------- str | int | float The last value set in redis for a given parameter """ + if collection_type is not None: + type = collection_type + else: + type = self._collection_type + with get_redis_connection() as redis_connection: if parameter in ["lab_name", "project_name", "auto_create_well"]: value = redis_connection.get(f"mxcube_common_params:{parameter}") else: - value = redis_connection.get(f"{self._collection_type}:{parameter}") + value = redis_connection.get(f"{type}:{parameter}") if parameter == "auto_create_well": if value is not None: diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml index 62c040dc01..cfde85b08e 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/config/default_params_pins.yml @@ -27,6 +27,7 @@ grid_scan: resolution: 1.5 # Ångström (16M) transmission: 1.0 # Percentage detector_roi_mode: "4M" # 4M or disabled + grid_step: "20x20" one_shot: exposure_time: 0.5 # Total exposure time, seconds @@ -35,12 +36,3 @@ one_shot: photon_energy: 13.0 # keV resolution: 2.0 # Ångström (16M) transmission: 0.5 # Percentage - -partial_udc: - # TODO: add more steps later - md3_alignment_y_speed: 1.0 # mm/s - omega_range: 0.0 # Degrees - photon_energy: 13.0 # keV - resolution: 1.5 # Ångström (16M) - transmission: 1.0 # Percentage - detector_roi_mode: "4M" # 4M or disabled diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 308ca79f88..1ada174be4 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -159,7 +159,11 @@ def dialog_box(self) -> dict: "type": "number", "minimum": 0.1, "maximum": 14.8, - "default": float(self._get_dialog_box_param("md3_alignment_y_speed")), + "default": float( + self._get_dialog_box_param( + "md3_alignment_y_speed", collection_type="grid_scan" + ) + ), "widget": "textarea", }, "transmission": { @@ -167,21 +171,31 @@ def dialog_box(self) -> dict: "type": "number", "minimum": 0, "maximum": 100, - "default": float(self._get_dialog_box_param("transmission")), + "default": float( + self._get_dialog_box_param( + "transmission", collection_type="grid_scan" + ) + ), "widget": "textarea", }, "detector_roi_mode": { "title": "Detector ROI Mode", "type": "string", "enum": ["4M", "disabled"], - "default": str(self._get_dialog_box_param("detector_roi_mode")), + "default": str( + self._get_dialog_box_param( + "detector_roi_mode", collection_type="grid_scan" + ) + ), "widget": "select", }, "grid_step": { "title": "Grid Step [um]", "type": "string", "enum": ["5x5", "10x10", "20x20"], - "default": str(self._get_dialog_box_param("grid_step")), + "default": str( + self._get_dialog_box_param("grid_step", collection_type="grid_scan") + ), "widget": "select", }, } From 033b22e47962ccd7b75e1f7e509a2a937961cc92 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Thu, 27 Nov 2025 16:27:19 +1100 Subject: [PATCH 03/17] add screening to partial udc --- .../ANSTO/prefect_flows/partial_udc.py | 157 ++++++++++++++++-- .../prefect_flows/schemas/partial_udc.py | 8 +- 2 files changed, 149 insertions(+), 16 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 1ada174be4..adc68f4427 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -13,6 +13,7 @@ OpticalCenteringExtraConfig, OpticalCenteringParams, PartialUDCDialogBox, + ScreeningParams, SingleLoopDataCollectionConfig, ) from .sync_prefect_client import MX3SyncPrefectClient @@ -55,10 +56,13 @@ def run(self, dialog_box_parameters: dict) -> None: dialog_box_model = PartialUDCDialogBox.model_validate(dialog_box_parameters) head_type = self.get_head_type() + gs_model = dialog_box_model.grid_scan + sc_model = dialog_box_model.screening + if not settings.ADD_DUMMY_PIN_TO_DB: if self.sample_id is None: logging.getLogger("HWR").info("Getting sample from the data layer...") - sample_id = self.get_sample_id_of_mounted_sample(dialog_box_model) + sample_id = self.get_sample_id_of_mounted_sample(gs_model) logging.getLogger("HWR").info(f"Mounted sample id: {sample_id}") else: logging.getLogger("HWR").info( @@ -100,19 +104,32 @@ def run(self, dialog_box_parameters: dict) -> None: optical_centering=OpticalCenteringParams( beam_position=(612, 512), extra_config=OpticalCenteringExtraConfig(grid_height_scale_factor=2), - grid_step=self.grid_step_map[dialog_box_model.grid_step], + grid_step=self.grid_step_map[gs_model.grid_step], calibrated_alignment_z=0.85, # TODO: maybe get from redis? ), grid_scan=GridScanParams( omega_range=0, - md3_alignment_y_speed=dialog_box_model.md3_alignment_y_speed, + md3_alignment_y_speed=gs_model.md3_alignment_y_speed, detector_distance=detector_distance, photon_energy=photon_energy, - transmission=dialog_box_model.transmission / 100, + transmission=gs_model.transmission / 100, crystal_finder_threshold=1, # TODO: can user set this? number_of_processes=settings.GRID_SCAN_NUMBER_OF_PROCESSES, ), - screening=None, + screening=ScreeningParams( + omega_range=sc_model.omega_range, + exposure_time=sc_model.exposure_time, + number_of_passes=1, + count_time=None, + number_of_frames=sc_model.number_of_frames, + detector_distance=self._resolution_to_distance( + sc_model.resolution, + energy=photon_energy, + ), + photon_energy=photon_energy, + transmission=sc_model.transmission / 100, + beam_size=(80, 80), + ), full_dataset=None, mount_pin_at_start_of_flow=False, add_dummy_pin_to_db=settings.ADD_DUMMY_PIN_TO_DB, @@ -129,7 +146,13 @@ def run(self, dialog_box_parameters: dict) -> None: ) # Remember the collection params for the next collection - self._save_dialog_box_params_to_redis(dialog_box_model) + # TODO: save dialog params can have the collection type as argument + original_collection_type = self._collection_type + self._collection_type = "grid_scan" + self._save_dialog_box_params_to_redis(gs_model) + self._collection_type = "screening" + self._save_dialog_box_params_to_redis(sc_model) + self._collection_type = original_collection_type partial_udc_flow = MX3SyncPrefectClient( name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters @@ -153,7 +176,8 @@ def dialog_box(self) -> dict: dialog : dict A dictionary following the JSON schema. """ - properties = { + # Grid Scan Properties + gs_properties = { "md3_alignment_y_speed": { "title": "Alignment Y Speed [mm/s]", "type": "number", @@ -199,21 +223,124 @@ def dialog_box(self) -> dict: "widget": "select", }, } + tray_conditional: dict | None = None if self.get_head_type() == "Plate": tray_properties, tray_conditional = self.build_tray_dialog_schema() - properties.update(tray_properties) + gs_properties.update(tray_properties) + + # Screening Properties + resolution_limits = self.resolution.get_limits() + sc_properties = { + "exposure_time": { + "title": "Total Exposure Time [s]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + self._get_dialog_box_param( + "exposure_time", collection_type="screening" + ) + ), + "widget": "textarea", + }, + "omega_range": { + "title": "Omega Range [degrees]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + self._get_dialog_box_param( + "omega_range", collection_type="screening" + ) + ), + "widget": "textarea", + }, + "number_of_frames": { + "title": "Number of Frames", + "type": "integer", + "minimum": 1, + "default": int( + self._get_dialog_box_param( + "number_of_frames", collection_type="screening" + ) + ), + "widget": "textarea", + }, + "resolution": { + "title": "Resolution [Å]", + "type": "number", + "minimum": resolution_limits[0], + "maximum": resolution_limits[1], + "default": float( + self._get_dialog_box_param( + "resolution", collection_type="screening" + ) + ), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, + "maximum": 100, + "default": float( + self._get_dialog_box_param( + "transmission", collection_type="screening" + ) + ), + "widget": "textarea", + }, + "crystal_counter": { + "title": "Crystal ID", + "type": "integer", + "minimum": 0, + "default": int( + self._get_dialog_box_param( + "crystal_counter", collection_type="screening" + ) + ), + "widget": "textarea", + }, + } + + if settings.ADD_DUMMY_PIN_TO_DB: + # Dev only + sc_properties["sample_id"] = { + "title": "Database sample id (dev only)", + "type": "integer", + "default": 1, + "widget": "textarea", + } dialog = { - "properties": properties, - "required": [ - "md3_alignment_y_speed", - "transmission", - ], - "dialogName": "Grid Scan Parameters", + "properties": { + "grid_scan": { + "type": "object", + "title": "Grid Scan Parameters", + "properties": gs_properties, + "required": [ + "md3_alignment_y_speed", + "transmission", + ], + }, + "screening": { + "type": "object", + "title": "Screening Parameters", + "properties": sc_properties, + "required": [ + "exposure_time", + "omega_range", + "number_of_frames", + "resolution", + "crystal_counter", + "transmission", + ], + }, + }, + "required": ["grid_scan", "screening"], + "dialogName": "Partial UDC Parameters", } if tray_conditional: - dialog.update(tray_conditional) + dialog["properties"]["grid_scan"].update(tray_conditional) return dialog diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py index b1f6634a83..818a28ba8f 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py @@ -12,14 +12,20 @@ GridScanDialogBox, GridScanParams, ) +from .screening import ScreeningDialogBox -class PartialUDCDialogBox(GridScanDialogBox): +class PartialUDCGridScanDialogBox(GridScanDialogBox): grid_step: str = Field( description="This is set as a string in the UI and is latter mapped to a tuple" ) +class PartialUDCDialogBox(BaseModel): + grid_scan: PartialUDCGridScanDialogBox + screening: ScreeningDialogBox + + # Below are the schemas copied from bluesky_worker single loop data collection flow schema # Only a subset of the full schema is included here class OpticalCenteringExtraConfig(BaseModel): From 44b7080dbc89337184d70f27fe3dc28c105ac14f Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Thu, 27 Nov 2025 16:36:47 +1100 Subject: [PATCH 04/17] add collection params arg --- .../ANSTO/prefect_flows/abstract_flow.py | 7 ++++++- .../ANSTO/prefect_flows/partial_udc.py | 14 ++++---------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py index e8bce5f921..df46552905 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py @@ -300,6 +300,7 @@ def _save_dialog_box_params_to_redis( | GridScanDialogBox | OneShotDialogBox ), + collection_type: str | None = None, ) -> None: """ Save the last set parameters from the dialog box to Redis. @@ -309,6 +310,10 @@ def _save_dialog_box_params_to_redis( dialog_box : ScreeningDialogBox | FullDatasetDialogBox | GridScanDialogBox | OneShotDialogBox A dialog box pydantic model """ + if collection_type is not None: + type = collection_type + else: + type = self._collection_type with get_redis_connection() as redis_connection: for key, value in dialog_box.model_dump(exclude_none=True).items(): if isinstance(value, bool): @@ -317,7 +322,7 @@ def _save_dialog_box_params_to_redis( if key in ["lab_name", "project_name", "auto_create_well"]: redis_connection.set(f"mxcube_common_params:{key}", value) else: - redis_connection.set(f"{self._collection_type}:{key}", value) + redis_connection.set(f"{type}:{key}", value) def _get_dialog_box_param( self, diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index adc68f4427..2c76cc507f 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -136,8 +136,8 @@ def run(self, dialog_box_parameters: dict) -> None: ) prefect_parameters = { "sample_id": sample_id, - "pin": {"id": 1, "puck": 1}, - "prepick_pin": {"id": 2, "puck": 1}, + "pin": None, + "prepick_pin": None, "config": partial_udc_config.model_dump(exclude_none=True), } @@ -146,14 +146,8 @@ def run(self, dialog_box_parameters: dict) -> None: ) # Remember the collection params for the next collection - # TODO: save dialog params can have the collection type as argument - original_collection_type = self._collection_type - self._collection_type = "grid_scan" - self._save_dialog_box_params_to_redis(gs_model) - self._collection_type = "screening" - self._save_dialog_box_params_to_redis(sc_model) - self._collection_type = original_collection_type - + self._save_dialog_box_params_to_redis(gs_model, collection_type="grid_scan") + self._save_dialog_box_params_to_redis(sc_model, collection_type="screening") partial_udc_flow = MX3SyncPrefectClient( name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters ) From 4d76a207550e79f881cb10aed448d7dc55c42379 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Thu, 27 Nov 2025 16:40:05 +1100 Subject: [PATCH 05/17] rename variables --- .../ANSTO/prefect_flows/partial_udc.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 2c76cc507f..863fe648eb 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -56,13 +56,13 @@ def run(self, dialog_box_parameters: dict) -> None: dialog_box_model = PartialUDCDialogBox.model_validate(dialog_box_parameters) head_type = self.get_head_type() - gs_model = dialog_box_model.grid_scan - sc_model = dialog_box_model.screening + grid_scan_model = dialog_box_model.grid_scan + screening_model = dialog_box_model.screening if not settings.ADD_DUMMY_PIN_TO_DB: if self.sample_id is None: logging.getLogger("HWR").info("Getting sample from the data layer...") - sample_id = self.get_sample_id_of_mounted_sample(gs_model) + sample_id = self.get_sample_id_of_mounted_sample(grid_scan_model) logging.getLogger("HWR").info(f"Mounted sample id: {sample_id}") else: logging.getLogger("HWR").info( @@ -104,30 +104,30 @@ def run(self, dialog_box_parameters: dict) -> None: optical_centering=OpticalCenteringParams( beam_position=(612, 512), extra_config=OpticalCenteringExtraConfig(grid_height_scale_factor=2), - grid_step=self.grid_step_map[gs_model.grid_step], + grid_step=self.grid_step_map[grid_scan_model.grid_step], calibrated_alignment_z=0.85, # TODO: maybe get from redis? ), grid_scan=GridScanParams( omega_range=0, - md3_alignment_y_speed=gs_model.md3_alignment_y_speed, + md3_alignment_y_speed=grid_scan_model.md3_alignment_y_speed, detector_distance=detector_distance, photon_energy=photon_energy, - transmission=gs_model.transmission / 100, + transmission=grid_scan_model.transmission / 100, crystal_finder_threshold=1, # TODO: can user set this? number_of_processes=settings.GRID_SCAN_NUMBER_OF_PROCESSES, ), screening=ScreeningParams( - omega_range=sc_model.omega_range, - exposure_time=sc_model.exposure_time, + omega_range=screening_model.omega_range, + exposure_time=screening_model.exposure_time, number_of_passes=1, count_time=None, - number_of_frames=sc_model.number_of_frames, + number_of_frames=screening_model.number_of_frames, detector_distance=self._resolution_to_distance( - sc_model.resolution, + screening_model.resolution, energy=photon_energy, ), photon_energy=photon_energy, - transmission=sc_model.transmission / 100, + transmission=screening_model.transmission / 100, beam_size=(80, 80), ), full_dataset=None, @@ -146,8 +146,8 @@ def run(self, dialog_box_parameters: dict) -> None: ) # Remember the collection params for the next collection - self._save_dialog_box_params_to_redis(gs_model, collection_type="grid_scan") - self._save_dialog_box_params_to_redis(sc_model, collection_type="screening") + self._save_dialog_box_params_to_redis(grid_scan_model, collection_type="grid_scan") + self._save_dialog_box_params_to_redis(screening_model, collection_type="screening") partial_udc_flow = MX3SyncPrefectClient( name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters ) From 68dabc0fb6d6a22e4976bd75a6c1b10b4a46b241 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 09:26:39 +1100 Subject: [PATCH 06/17] add screening props to dialog box folder --- .../ANSTO/prefect_flows/partial_udc.py | 91 ++----------------- .../schemas/dialog_boxes/screening.py | 63 +++++++++++++ .../schemas/dialog_boxes/utils.py | 62 +++++++++++++ .../ANSTO/prefect_flows/screening_flow.py | 57 +----------- 4 files changed, 136 insertions(+), 137 deletions(-) create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 863fe648eb..8200cd6f1b 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -8,6 +8,7 @@ from ..redis_utils import get_redis_connection from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.screening import get_screening_schema from .schemas.partial_udc import ( GridScanParams, OpticalCenteringExtraConfig, @@ -146,8 +147,12 @@ def run(self, dialog_box_parameters: dict) -> None: ) # Remember the collection params for the next collection - self._save_dialog_box_params_to_redis(grid_scan_model, collection_type="grid_scan") - self._save_dialog_box_params_to_redis(screening_model, collection_type="screening") + self._save_dialog_box_params_to_redis( + grid_scan_model, collection_type="grid_scan" + ) + self._save_dialog_box_params_to_redis( + screening_model, collection_type="screening" + ) partial_udc_flow = MX3SyncPrefectClient( name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters ) @@ -225,85 +230,7 @@ def dialog_box(self) -> dict: # Screening Properties resolution_limits = self.resolution.get_limits() - sc_properties = { - "exposure_time": { - "title": "Total Exposure Time [s]", - "type": "number", - "exclusiveMinimum": 0, - "default": float( - self._get_dialog_box_param( - "exposure_time", collection_type="screening" - ) - ), - "widget": "textarea", - }, - "omega_range": { - "title": "Omega Range [degrees]", - "type": "number", - "exclusiveMinimum": 0, - "default": float( - self._get_dialog_box_param( - "omega_range", collection_type="screening" - ) - ), - "widget": "textarea", - }, - "number_of_frames": { - "title": "Number of Frames", - "type": "integer", - "minimum": 1, - "default": int( - self._get_dialog_box_param( - "number_of_frames", collection_type="screening" - ) - ), - "widget": "textarea", - }, - "resolution": { - "title": "Resolution [Å]", - "type": "number", - "minimum": resolution_limits[0], - "maximum": resolution_limits[1], - "default": float( - self._get_dialog_box_param( - "resolution", collection_type="screening" - ) - ), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, - "maximum": 100, - "default": float( - self._get_dialog_box_param( - "transmission", collection_type="screening" - ) - ), - "widget": "textarea", - }, - "crystal_counter": { - "title": "Crystal ID", - "type": "integer", - "minimum": 0, - "default": int( - self._get_dialog_box_param( - "crystal_counter", collection_type="screening" - ) - ), - "widget": "textarea", - }, - } - - if settings.ADD_DUMMY_PIN_TO_DB: - # Dev only - sc_properties["sample_id"] = { - "title": "Database sample id (dev only)", - "type": "integer", - "default": 1, - "widget": "textarea", - } + screening_properties = get_screening_schema(resolution_limits) dialog = { "properties": { @@ -319,7 +246,7 @@ def dialog_box(self) -> dict: "screening": { "type": "object", "title": "Screening Parameters", - "properties": sc_properties, + "properties": screening_properties, "required": [ "exposure_time", "omega_range", diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py new file mode 100644 index 0000000000..a3283d85c8 --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py @@ -0,0 +1,63 @@ +from .utils import get_dialog_box_param + + +def get_screening_schema(resolution_limits: tuple[float, float]) -> dict: + properties = { + "exposure_time": { + "title": "Total Exposure Time [s]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + get_dialog_box_param("exposure_time", collection_type="screening") + ), + "widget": "textarea", + }, + "omega_range": { + "title": "Omega Range [degrees]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + get_dialog_box_param("omega_range", collection_type="screening") + ), + "widget": "textarea", + }, + "number_of_frames": { + "title": "Number of Frames", + "type": "integer", + "minimum": 1, + "default": int( + get_dialog_box_param("number_of_frames", collection_type="screening") + ), + "widget": "textarea", + }, + "resolution": { + "title": "Resolution [Å]", + "type": "number", + "minimum": resolution_limits[0], + "maximum": resolution_limits[1], + "default": float( + get_dialog_box_param("resolution", collection_type="screening") + ), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, + "maximum": 100, + "default": float( + get_dialog_box_param("transmission", collection_type="screening") + ), + "widget": "textarea", + }, + "crystal_counter": { + "title": "Crystal ID", + "type": "integer", + "minimum": 0, + "default": int( + get_dialog_box_param("crystal_counter", collection_type="screening") + ), + "widget": "textarea", + }, + } + return properties diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py new file mode 100644 index 0000000000..bc25dc9135 --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py @@ -0,0 +1,62 @@ +from typing import Literal + +from ....redis_utils import get_redis_connection + + +def get_dialog_box_param( + parameter: Literal[ + "exposure_time", + "omega_range", + "number_of_frames", + "processing_pipeline", + "crystal_counter", + "photon_energy", + "resolution", + "md3_alignment_y_speed", + "transmission", + "auto_create_well", + "lab_name", + "project_name", + ], + collection_type: Literal["screening", "full_dataset", "grid_scan", "one_shot"], +) -> str | int | float | None: + """ + Retrieve a parameter value from Redis. + + Parameters + ---------- + parameter : Literal[ + "exposure_time", + "omega_range", + "number_of_frames", + "processing_pipeline", + "crystal_counter", + "photon_energy", + "resolution", + "md3_alignment_y_speed", + "transmission", + "auto_create_well", + "lab_name", + "project_name", + ] + A parameter saved in redis + collection_type : str | None + The collection type. If None, use the current collection type + + Returns + ------- + str | int | float + The last value set in redis for a given parameter + """ + + with get_redis_connection() as redis_connection: + if parameter in ["lab_name", "project_name", "auto_create_well"]: + value = redis_connection.get(f"mxcube_common_params:{parameter}") + else: + value = redis_connection.get(f"{collection_type}:{parameter}") + + if parameter == "auto_create_well": + if value is not None: + return bool(int(value)) + else: + return value diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/screening_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/screening_flow.py index e04ca7508e..ff1340afbb 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/screening_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/screening_flow.py @@ -7,6 +7,7 @@ from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.screening import get_screening_schema from .schemas.screening import ( ScreeningDialogBox, ScreeningParams, @@ -126,61 +127,7 @@ def dialog_box(self) -> dict: """ resolution_limits = self.resolution.get_limits() - properties = { - "exposure_time": { - "title": "Total Exposure Time [s]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(self._get_dialog_box_param("exposure_time")), - "widget": "textarea", - }, - "omega_range": { - "title": "Omega Range [degrees]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(self._get_dialog_box_param("omega_range")), - "widget": "textarea", - }, - "number_of_frames": { - "title": "Number of Frames", - "type": "integer", - "minimum": 1, - "default": int(self._get_dialog_box_param("number_of_frames")), - "widget": "textarea", - }, - "resolution": { - "title": "Resolution [Å]", - "type": "number", - "minimum": resolution_limits[0], - "maximum": resolution_limits[1], - "default": float(self._get_dialog_box_param("resolution")), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, - "maximum": 100, - "default": float(self._get_dialog_box_param("transmission")), - "widget": "textarea", - }, - "crystal_counter": { - "title": "Crystal ID", - "type": "integer", - "minimum": 0, - "default": int(self._get_dialog_box_param("crystal_counter")), - "widget": "textarea", - }, - } - - if settings.ADD_DUMMY_PIN_TO_DB: - # Dev only - properties["sample_id"] = { - "title": "Database sample id (dev only)", - "type": "integer", - "default": 1, - "widget": "textarea", - } + properties = get_screening_schema(resolution_limits) tray_conditional: dict | None = None if self.get_head_type() == "Plate": From 8e04e3f8f10060c55dfd6fefc0da8ead2dd903e4 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 09:41:37 +1100 Subject: [PATCH 07/17] add common grid scan schema --- .../ANSTO/prefect_flows/grid_scan_flow.py | 27 +--------- .../ANSTO/prefect_flows/partial_udc.py | 52 ++----------------- .../schemas/dialog_boxes/grid_scan.py | 48 +++++++++++++++++ 3 files changed, 54 insertions(+), 73 deletions(-) create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/grid_scan_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/grid_scan_flow.py index 65fe210c34..08e0ea2482 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/grid_scan_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/grid_scan_flow.py @@ -16,6 +16,7 @@ from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.grid_scan import get_grid_scan_schema from .schemas.grid_scan import ( GridScanDialogBox, GridScanParams, @@ -334,31 +335,7 @@ def dialog_box(self) -> dict: dialog : dict A dictionary following the JSON schema. """ - properties = { - "md3_alignment_y_speed": { - "title": "Alignment Y Speed [mm/s]", - "type": "number", - "minimum": 0.1, - "maximum": 14.8, - "default": float(self._get_dialog_box_param("md3_alignment_y_speed")), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, - "maximum": 100, - "default": float(self._get_dialog_box_param("transmission")), - "widget": "textarea", - }, - "detector_roi_mode": { - "title": "Detector ROI Mode", - "type": "string", - "enum": ["4M", "disabled"], - "default": str(self._get_dialog_box_param("detector_roi_mode")), - "widget": "select", - }, - } + properties = get_grid_scan_schema(partial_udc=False) tray_conditional: dict | None = None if self.get_head_type() == "Plate": tray_properties, tray_conditional = self.build_tray_dialog_schema() diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 8200cd6f1b..a303b1543b 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -8,6 +8,7 @@ from ..redis_utils import get_redis_connection from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.grid_scan import get_grid_scan_schema from .schemas.dialog_boxes.screening import get_screening_schema from .schemas.partial_udc import ( GridScanParams, @@ -176,57 +177,12 @@ def dialog_box(self) -> dict: A dictionary following the JSON schema. """ # Grid Scan Properties - gs_properties = { - "md3_alignment_y_speed": { - "title": "Alignment Y Speed [mm/s]", - "type": "number", - "minimum": 0.1, - "maximum": 14.8, - "default": float( - self._get_dialog_box_param( - "md3_alignment_y_speed", collection_type="grid_scan" - ) - ), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, - "maximum": 100, - "default": float( - self._get_dialog_box_param( - "transmission", collection_type="grid_scan" - ) - ), - "widget": "textarea", - }, - "detector_roi_mode": { - "title": "Detector ROI Mode", - "type": "string", - "enum": ["4M", "disabled"], - "default": str( - self._get_dialog_box_param( - "detector_roi_mode", collection_type="grid_scan" - ) - ), - "widget": "select", - }, - "grid_step": { - "title": "Grid Step [um]", - "type": "string", - "enum": ["5x5", "10x10", "20x20"], - "default": str( - self._get_dialog_box_param("grid_step", collection_type="grid_scan") - ), - "widget": "select", - }, - } + grid_scan_properties = get_grid_scan_schema(partial_udc=True) tray_conditional: dict | None = None if self.get_head_type() == "Plate": tray_properties, tray_conditional = self.build_tray_dialog_schema() - gs_properties.update(tray_properties) + grid_scan_properties.update(tray_properties) # Screening Properties resolution_limits = self.resolution.get_limits() @@ -237,7 +193,7 @@ def dialog_box(self) -> dict: "grid_scan": { "type": "object", "title": "Grid Scan Parameters", - "properties": gs_properties, + "properties": grid_scan_properties, "required": [ "md3_alignment_y_speed", "transmission", diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py new file mode 100644 index 0000000000..edc83fc94f --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py @@ -0,0 +1,48 @@ +from .utils import get_dialog_box_param + + +def get_grid_scan_schema(partial_udc: bool = False) -> dict: + properties = { + "md3_alignment_y_speed": { + "title": "Alignment Y Speed [mm/s]", + "type": "number", + "minimum": 0.1, + "maximum": 14.8, + "default": float( + get_dialog_box_param( + "md3_alignment_y_speed", collection_type="grid_scan" + ) + ), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, + "maximum": 100, + "default": float( + get_dialog_box_param("transmission", collection_type="grid_scan") + ), + "widget": "textarea", + }, + "detector_roi_mode": { + "title": "Detector ROI Mode", + "type": "string", + "enum": ["4M", "disabled"], + "default": str( + get_dialog_box_param("detector_roi_mode", collection_type="grid_scan") + ), + "widget": "select", + }, + } + if partial_udc: + properties["grid_step"] = { + "title": "Grid Step [um]", + "type": "string", + "enum": ["5x5", "10x10", "20x20"], + "default": str( + get_dialog_box_param("grid_step", collection_type="grid_scan") + ), + "widget": "select", + } + return properties From 52e34f4c0321494a84d91b648adbfd369da4616a Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 09:50:55 +1100 Subject: [PATCH 08/17] add full dataset schema --- .../full_dataset_collection_flow.py | 67 +--------------- .../schemas/dialog_boxes/full_dataset.py | 77 +++++++++++++++++++ 2 files changed, 79 insertions(+), 65 deletions(-) create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/full_dataset_collection_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/full_dataset_collection_flow.py index 84d9e386a0..a382bc4362 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/full_dataset_collection_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/full_dataset_collection_flow.py @@ -7,6 +7,7 @@ from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.full_dataset import get_full_dataset_schema from .schemas.full_dataset import ( FullDatasetDialogBox, FullDatasetParams, @@ -125,72 +126,8 @@ def dialog_box(self) -> dict: """ resolution_limits = self.resolution.get_limits() - properties = { - "exposure_time": { - "title": "Total Exposure Time [s]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(self._get_dialog_box_param("exposure_time")), - "widget": "textarea", - }, - "omega_range": { - "title": "Omega Range [degrees]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(self._get_dialog_box_param("omega_range")), - "widget": "textarea", - }, - "number_of_frames": { - "title": "Number of Frames", - "type": "integer", - "minimum": 1, - "default": int(self._get_dialog_box_param("number_of_frames")), - "widget": "textarea", - }, - "resolution": { - "title": "Resolution [Å]", - "type": "number", - "minimum": resolution_limits[0], - "maximum": resolution_limits[1], - "default": float(self._get_dialog_box_param("resolution")), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, - "maximum": 100, - "default": float(self._get_dialog_box_param("transmission")), - "widget": "textarea", - }, - "processing_pipeline": { - "title": "Data Processing Pipeline", - "type": "string", - "enum": [ - "dials", - "fast_dp", - "dials_and_fast_dp", - "fast_dp_and_xia2", - "xia2", - ], - "default": self._get_dialog_box_param("processing_pipeline"), - }, - "crystal_counter": { - "title": "Crystal ID", - "type": "integer", - "minimum": 0, - "default": int(self._get_dialog_box_param("crystal_counter")), - "widget": "textarea", - }, - } + properties = get_full_dataset_schema(resolution_limits) - if settings.ADD_DUMMY_PIN_TO_DB: - properties["sample_id"] = { - "title": "Database sample id (dev only)", - "type": "integer", - "default": 1, - "widget": "textarea", - } tray_conditional: dict | None = None if self.get_head_type() == "Plate": tray_properties, tray_conditional = self.build_tray_dialog_schema() diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py new file mode 100644 index 0000000000..0e42ee9696 --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py @@ -0,0 +1,77 @@ +from .utils import get_dialog_box_param + + +def get_full_dataset_schema(resolution_limits: tuple[float, float]) -> dict: + properties = { + "exposure_time": { + "title": "Total Exposure Time [s]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + get_dialog_box_param("exposure_time", collection_type="full_dataset") + ), + "widget": "textarea", + }, + "omega_range": { + "title": "Omega Range [degrees]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + get_dialog_box_param("omega_range", collection_type="full_dataset") + ), + "widget": "textarea", + }, + "number_of_frames": { + "title": "Number of Frames", + "type": "integer", + "minimum": 1, + "default": int( + get_dialog_box_param("number_of_frames", collection_type="full_dataset") + ), + "widget": "textarea", + }, + "resolution": { + "title": "Resolution [Å]", + "type": "number", + "minimum": resolution_limits[0], + "maximum": resolution_limits[1], + "default": float( + get_dialog_box_param("resolution", collection_type="full_dataset") + ), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, + "maximum": 100, + "default": float( + get_dialog_box_param("transmission", collection_type="full_dataset") + ), + "widget": "textarea", + }, + "processing_pipeline": { + "title": "Data Processing Pipeline", + "type": "string", + "enum": [ + "dials", + "fast_dp", + "dials_and_fast_dp", + "fast_dp_and_xia2", + "xia2", + ], + "default": get_dialog_box_param( + "processing_pipeline", collection_type="full_dataset" + ), + }, + "crystal_counter": { + "title": "Crystal ID", + "type": "integer", + "minimum": 0, + "default": int( + get_dialog_box_param("crystal_counter", collection_type="full_dataset") + ), + "widget": "textarea", + }, + } + return properties From 877d78f30c328f24fe7cecab10620d6bb3073df0 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 10:14:35 +1100 Subject: [PATCH 09/17] add one shot schema --- .../ANSTO/prefect_flows/one_shot_flow.py | 43 +------------------ .../schemas/dialog_boxes/one_shot.py | 37 ++++++++++++++++ .../schemas/dialog_boxes/utils.py | 4 +- 3 files changed, 41 insertions(+), 43 deletions(-) create mode 100644 mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py index e12ef04e0a..503ed9da58 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py @@ -12,6 +12,7 @@ OneShotParams, ) from .sync_prefect_client import MX3SyncPrefectClient +from .schemas.dialog_boxes.one_shot import get_one_shot_schema class OneShotFlow(AbstractPrefectWorkflow): @@ -110,47 +111,7 @@ def dialog_box(self) -> dict: A dictionary following the JSON schema. """ resolution_limits = self.resolution.get_limits() - - properties = { - "exposure_time": { - "title": "Total Exposure Time [s]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(self._get_dialog_box_param("exposure_time")), - "widget": "textarea", - }, - "omega_range": { - "title": "Omega Range [degrees]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(self._get_dialog_box_param("omega_range")), - "widget": "textarea", - }, - "resolution": { - "title": "Resolution [Å]", - "type": "number", - "minimum": resolution_limits[0], - "maximum": resolution_limits[1], - "default": float(self._get_dialog_box_param("resolution")), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, # TODO: get limits from PV? - "maximum": 100, - "default": float(self._get_dialog_box_param("transmission")), - "widget": "textarea", - }, - } - - if settings.ADD_DUMMY_PIN_TO_DB: - properties["sample_id"] = { - "title": "Database sample id (dev only)", - "type": "integer", - "default": 1, - "widget": "textarea", - } + properties = get_one_shot_schema(resolution_limits) tray_conditional: dict | None = None if self.get_head_type() == "Plate": diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py new file mode 100644 index 0000000000..3b9793c335 --- /dev/null +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py @@ -0,0 +1,37 @@ +from .utils import get_dialog_box_param + + +def get_one_shot_schema(resolution_limits: tuple[float, float]) -> dict: + properties = { + "exposure_time": { + "title": "Total Exposure Time [s]", + "type": "number", + "exclusiveMinimum": 0, + "default": float(get_dialog_box_param("exposure_time", collection_type="one_shot")), + "widget": "textarea", + }, + "omega_range": { + "title": "Omega Range [degrees]", + "type": "number", + "exclusiveMinimum": 0, + "default": float(get_dialog_box_param("omega_range", collection_type="one_shot")), + "widget": "textarea", + }, + "resolution": { + "title": "Resolution [Å]", + "type": "number", + "minimum": resolution_limits[0], + "maximum": resolution_limits[1], + "default": float(get_dialog_box_param("resolution", collection_type="one_shot")), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, # TODO: get limits from PV? + "maximum": 100, + "default": float(get_dialog_box_param("transmission", collection_type="one_shot")), + "widget": "textarea", + }, + } + return properties \ No newline at end of file diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py index bc25dc9135..4481d9f555 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/utils.py @@ -40,8 +40,8 @@ def get_dialog_box_param( "project_name", ] A parameter saved in redis - collection_type : str | None - The collection type. If None, use the current collection type + collection_type : Literal["screening", "full_dataset", "grid_scan", "one_shot"] + The collection type Returns ------- From 04a4085fb1b16bdafba0af604aae6d5be94b8794 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 10:24:37 +1100 Subject: [PATCH 10/17] pre-commit --- .../ANSTO/prefect_flows/one_shot_flow.py | 2 +- .../schemas/dialog_boxes/one_shot.py | 74 ++++++++++--------- 2 files changed, 42 insertions(+), 34 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py index 503ed9da58..ad54dfcfc6 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py @@ -7,12 +7,12 @@ from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.one_shot import get_one_shot_schema from .schemas.one_shot import ( OneShotDialogBox, OneShotParams, ) from .sync_prefect_client import MX3SyncPrefectClient -from .schemas.dialog_boxes.one_shot import get_one_shot_schema class OneShotFlow(AbstractPrefectWorkflow): diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py index 3b9793c335..bd86c9d508 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/one_shot.py @@ -1,37 +1,45 @@ -from .utils import get_dialog_box_param +from .utils import get_dialog_box_param def get_one_shot_schema(resolution_limits: tuple[float, float]) -> dict: properties = { - "exposure_time": { - "title": "Total Exposure Time [s]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(get_dialog_box_param("exposure_time", collection_type="one_shot")), - "widget": "textarea", - }, - "omega_range": { - "title": "Omega Range [degrees]", - "type": "number", - "exclusiveMinimum": 0, - "default": float(get_dialog_box_param("omega_range", collection_type="one_shot")), - "widget": "textarea", - }, - "resolution": { - "title": "Resolution [Å]", - "type": "number", - "minimum": resolution_limits[0], - "maximum": resolution_limits[1], - "default": float(get_dialog_box_param("resolution", collection_type="one_shot")), - "widget": "textarea", - }, - "transmission": { - "title": "Transmission [%]", - "type": "number", - "minimum": 0, # TODO: get limits from PV? - "maximum": 100, - "default": float(get_dialog_box_param("transmission", collection_type="one_shot")), - "widget": "textarea", - }, - } - return properties \ No newline at end of file + "exposure_time": { + "title": "Total Exposure Time [s]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + get_dialog_box_param("exposure_time", collection_type="one_shot") + ), + "widget": "textarea", + }, + "omega_range": { + "title": "Omega Range [degrees]", + "type": "number", + "exclusiveMinimum": 0, + "default": float( + get_dialog_box_param("omega_range", collection_type="one_shot") + ), + "widget": "textarea", + }, + "resolution": { + "title": "Resolution [Å]", + "type": "number", + "minimum": resolution_limits[0], + "maximum": resolution_limits[1], + "default": float( + get_dialog_box_param("resolution", collection_type="one_shot") + ), + "widget": "textarea", + }, + "transmission": { + "title": "Transmission [%]", + "type": "number", + "minimum": 0, # TODO: get limits from PV? + "maximum": 100, + "default": float( + get_dialog_box_param("transmission", collection_type="one_shot") + ), + "widget": "textarea", + }, + } + return properties From bb22884da7e5e821a2158858f42cbfa507e8a887 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 10:25:10 +1100 Subject: [PATCH 11/17] add dataset to partial udc --- .../ANSTO/prefect_flows/partial_udc.py | 40 ++++++++++++++++++- .../prefect_flows/schemas/partial_udc.py | 2 + 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index a303b1543b..009e8a288c 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -8,9 +8,11 @@ from ..redis_utils import get_redis_connection from ..Resolution import Resolution from .abstract_flow import AbstractPrefectWorkflow +from .schemas.dialog_boxes.full_dataset import get_full_dataset_schema from .schemas.dialog_boxes.grid_scan import get_grid_scan_schema from .schemas.dialog_boxes.screening import get_screening_schema from .schemas.partial_udc import ( + FullDatasetParams, GridScanParams, OpticalCenteringExtraConfig, OpticalCenteringParams, @@ -60,6 +62,7 @@ def run(self, dialog_box_parameters: dict) -> None: grid_scan_model = dialog_box_model.grid_scan screening_model = dialog_box_model.screening + full_dataset_model = dialog_box_model.full_dataset if not settings.ADD_DUMMY_PIN_TO_DB: if self.sample_id is None: @@ -132,7 +135,20 @@ def run(self, dialog_box_parameters: dict) -> None: transmission=screening_model.transmission / 100, beam_size=(80, 80), ), - full_dataset=None, + full_dataset=FullDatasetParams( + omega_range=full_dataset_model.omega_range, + exposure_time=full_dataset_model.exposure_time, + number_of_passes=1, + count_time=None, + number_of_frames=full_dataset_model.number_of_frames, + detector_distance=self._resolution_to_distance( + full_dataset_model.resolution, + energy=photon_energy, + ), + photon_energy=photon_energy, + transmission=full_dataset_model.transmission / 100, + beam_size=(80, 80), + ), mount_pin_at_start_of_flow=False, add_dummy_pin_to_db=settings.ADD_DUMMY_PIN_TO_DB, ) @@ -154,6 +170,9 @@ def run(self, dialog_box_parameters: dict) -> None: self._save_dialog_box_params_to_redis( screening_model, collection_type="screening" ) + self._save_dialog_box_params_to_redis( + full_dataset_model, collection_type="full_dataset" + ) partial_udc_flow = MX3SyncPrefectClient( name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters ) @@ -188,6 +207,9 @@ def dialog_box(self) -> dict: resolution_limits = self.resolution.get_limits() screening_properties = get_screening_schema(resolution_limits) + # Full Dataset Properties + full_dataset_properties = get_full_dataset_schema(resolution_limits) + dialog = { "properties": { "grid_scan": { @@ -212,8 +234,22 @@ def dialog_box(self) -> dict: "transmission", ], }, + "full_dataset": { + "type": "object", + "title": "Dataset Parameters", + "properties": full_dataset_properties, + "required": [ + "exposure_time", + "omega_range", + "number_of_frames", + "resolution", + "crystal_counter", + "transmission", + "processing_pipeline", + ], + }, }, - "required": ["grid_scan", "screening"], + "required": ["grid_scan", "screening", "full_dataset"], "dialogName": "Partial UDC Parameters", } diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py index 818a28ba8f..72095ffc4f 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py @@ -8,6 +8,7 @@ model_validator, ) +from .full_dataset import FullDatasetDialogBox from .grid_scan import ( GridScanDialogBox, GridScanParams, @@ -24,6 +25,7 @@ class PartialUDCGridScanDialogBox(GridScanDialogBox): class PartialUDCDialogBox(BaseModel): grid_scan: PartialUDCGridScanDialogBox screening: ScreeningDialogBox + full_dataset: FullDatasetDialogBox # Below are the schemas copied from bluesky_worker single loop data collection flow schema From 364413e9dc812733f5fd8a8bc6a5dc787d877a9f Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 13:49:03 +1100 Subject: [PATCH 12/17] hide schema if not selected --- .../ANSTO/prefect_flows/partial_udc.py | 191 +++++++++++------- .../prefect_flows/schemas/partial_udc.py | 6 +- 2 files changed, 126 insertions(+), 71 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 009e8a288c..8e24bb84c1 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -17,6 +17,7 @@ OpticalCenteringExtraConfig, OpticalCenteringParams, PartialUDCDialogBox, + ProjectStrategy, ScreeningParams, SingleLoopDataCollectionConfig, ) @@ -56,9 +57,14 @@ def run(self, dialog_box_parameters: dict) -> None: """ self._state.value = "RUNNING" + head_type = self.get_head_type() + # TODO: partial udc for plates not supported yet + if head_type == "Plate": + msg = "Partial UDC for plates is not supported yet" + logging.getLogger("user_level_log").error(msg) + raise QueueExecutionException(msg, self) dialog_box_model = PartialUDCDialogBox.model_validate(dialog_box_parameters) - head_type = self.get_head_type() grid_scan_model = dialog_box_model.grid_scan screening_model = dialog_box_model.screening @@ -96,32 +102,8 @@ def run(self, dialog_box_parameters: dict) -> None: f"Detector distance corresponding to {default_resolution} A: {detector_distance} [m]" ) - # TODO: partial udc for plates not supported yet - if head_type == "Plate": - use_centring_table = False - msg = "Partial UDC for plates is not supported yet" - logging.getLogger("user_level_log").error(msg) - raise QueueExecutionException(msg, self) - else: - use_centring_table = True - - partial_udc_config = SingleLoopDataCollectionConfig( - optical_centering=OpticalCenteringParams( - beam_position=(612, 512), - extra_config=OpticalCenteringExtraConfig(grid_height_scale_factor=2), - grid_step=self.grid_step_map[grid_scan_model.grid_step], - calibrated_alignment_z=0.85, # TODO: maybe get from redis? - ), - grid_scan=GridScanParams( - omega_range=0, - md3_alignment_y_speed=grid_scan_model.md3_alignment_y_speed, - detector_distance=detector_distance, - photon_energy=photon_energy, - transmission=grid_scan_model.transmission / 100, - crystal_finder_threshold=1, # TODO: can user set this? - number_of_processes=settings.GRID_SCAN_NUMBER_OF_PROCESSES, - ), - screening=ScreeningParams( + if dialog_box_model.perform_screening: + screening_params = ScreeningParams( omega_range=screening_model.omega_range, exposure_time=screening_model.exposure_time, number_of_passes=1, @@ -134,8 +116,12 @@ def run(self, dialog_box_parameters: dict) -> None: photon_energy=photon_energy, transmission=screening_model.transmission / 100, beam_size=(80, 80), - ), - full_dataset=FullDatasetParams( + ) + else: + screening_params = None + + if dialog_box_model.perform_full_dataset: + full_dataset_params = FullDatasetParams( omega_range=full_dataset_model.omega_range, exposure_time=full_dataset_model.exposure_time, number_of_passes=1, @@ -148,9 +134,33 @@ def run(self, dialog_box_parameters: dict) -> None: photon_energy=photon_energy, transmission=full_dataset_model.transmission / 100, beam_size=(80, 80), + ) + else: + full_dataset_params = None + + partial_udc_config = SingleLoopDataCollectionConfig( + optical_centering=OpticalCenteringParams( + beam_position=(612, 512), + extra_config=OpticalCenteringExtraConfig(grid_height_scale_factor=2), + grid_step=self.grid_step_map[grid_scan_model.grid_step], + calibrated_alignment_z=0.85, # TODO: maybe get from redis? + ), + grid_scan=GridScanParams( + omega_range=0, + md3_alignment_y_speed=grid_scan_model.md3_alignment_y_speed, + detector_distance=detector_distance, + photon_energy=photon_energy, + transmission=grid_scan_model.transmission / 100, + crystal_finder_threshold=1, # TODO: can user set this? + number_of_processes=settings.GRID_SCAN_NUMBER_OF_PROCESSES, ), + screening=screening_params, + full_dataset=full_dataset_params, mount_pin_at_start_of_flow=False, add_dummy_pin_to_db=settings.ADD_DUMMY_PIN_TO_DB, + project_strategy=ProjectStrategy( + use_screening=dialog_box_model.perform_screening + ), ) prefect_parameters = { "sample_id": sample_id, @@ -167,12 +177,25 @@ def run(self, dialog_box_parameters: dict) -> None: self._save_dialog_box_params_to_redis( grid_scan_model, collection_type="grid_scan" ) - self._save_dialog_box_params_to_redis( - screening_model, collection_type="screening" - ) - self._save_dialog_box_params_to_redis( - full_dataset_model, collection_type="full_dataset" - ) + if dialog_box_model.perform_screening: + self._save_dialog_box_params_to_redis( + screening_model, collection_type="screening" + ) + if dialog_box_model.perform_full_dataset: + self._save_dialog_box_params_to_redis( + full_dataset_model, collection_type="full_dataset" + ) + + with get_redis_connection() as redis_connection: + redis_connection.set( + "partial_udc:perform_screening", + 1 if dialog_box_model.perform_screening else 0, + ) + redis_connection.set( + "partial_udc:perform_full_dataset", + 1 if dialog_box_model.perform_full_dataset else 0, + ) + partial_udc_flow = MX3SyncPrefectClient( name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters ) @@ -198,11 +221,6 @@ def dialog_box(self) -> dict: # Grid Scan Properties grid_scan_properties = get_grid_scan_schema(partial_udc=True) - tray_conditional: dict | None = None - if self.get_head_type() == "Plate": - tray_properties, tray_conditional = self.build_tray_dialog_schema() - grid_scan_properties.update(tray_properties) - # Screening Properties resolution_limits = self.resolution.get_limits() screening_properties = get_screening_schema(resolution_limits) @@ -210,6 +228,18 @@ def dialog_box(self) -> dict: # Full Dataset Properties full_dataset_properties = get_full_dataset_schema(resolution_limits) + with get_redis_connection() as redis_connection: + perform_screening = bool( + int( + redis_connection.get("partial_udc:perform_screening") or 0 + ) # False by default + ) + perform_full_dataset = bool( + int( + redis_connection.get("partial_udc:perform_full_dataset") or 0 + ) # False by default + ) + dialog = { "properties": { "grid_scan": { @@ -221,39 +251,62 @@ def dialog_box(self) -> dict: "transmission", ], }, - "screening": { - "type": "object", - "title": "Screening Parameters", - "properties": screening_properties, - "required": [ - "exposure_time", - "omega_range", - "number_of_frames", - "resolution", - "crystal_counter", - "transmission", - ], + "perform_screening": { + "type": "boolean", + "title": "Perform Screening", + "default": perform_screening, }, - "full_dataset": { - "type": "object", - "title": "Dataset Parameters", - "properties": full_dataset_properties, - "required": [ - "exposure_time", - "omega_range", - "number_of_frames", - "resolution", - "crystal_counter", - "transmission", - "processing_pipeline", - ], + "perform_full_dataset": { + "type": "boolean", + "title": "Perform Full Dataset", + "default": perform_full_dataset, }, }, - "required": ["grid_scan", "screening", "full_dataset"], + "required": ["grid_scan"], + "allOf": [ + { + "if": {"properties": {"perform_screening": {"const": True}}}, + "then": { + "properties": { + "screening": { + "type": "object", + "title": "Screening Parameters", + "properties": screening_properties, + "required": [ + "exposure_time", + "omega_range", + "number_of_frames", + "resolution", + "crystal_counter", + "transmission", + ], + } + } + }, + }, + { + "if": {"properties": {"perform_full_dataset": {"const": True}}}, + "then": { + "properties": { + "full_dataset": { + "type": "object", + "title": "Dataset Parameters", + "properties": full_dataset_properties, + "required": [ + "exposure_time", + "omega_range", + "number_of_frames", + "resolution", + "crystal_counter", + "transmission", + "processing_pipeline", + ], + } + } + }, + }, + ], "dialogName": "Partial UDC Parameters", } - if tray_conditional: - dialog["properties"]["grid_scan"].update(tray_conditional) - return dialog diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py index 72095ffc4f..a746c44051 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py @@ -24,8 +24,10 @@ class PartialUDCGridScanDialogBox(GridScanDialogBox): class PartialUDCDialogBox(BaseModel): grid_scan: PartialUDCGridScanDialogBox - screening: ScreeningDialogBox - full_dataset: FullDatasetDialogBox + perform_screening: bool = False + screening: ScreeningDialogBox | None = None + perform_full_dataset: bool = False + full_dataset: FullDatasetDialogBox | None = None # Below are the schemas copied from bluesky_worker single loop data collection flow schema From ebefaf444cf5f16af9c0b67dde95ab4404cfa275 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Fri, 28 Nov 2025 15:52:47 +1100 Subject: [PATCH 13/17] update partial udc schema --- .../ANSTO/prefect_flows/partial_udc.py | 82 ++++++++++++------- .../schemas/dialog_boxes/full_dataset.py | 11 ++- .../schemas/dialog_boxes/grid_scan.py | 4 +- .../schemas/dialog_boxes/screening.py | 11 ++- .../prefect_flows/schemas/partial_udc.py | 4 +- 5 files changed, 72 insertions(+), 40 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 8e24bb84c1..88459838d5 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -37,14 +37,39 @@ def __init__( self._collection_type = "partial_udc" self.grid_step_map = { + "2.5x2.5": (2.5, 2.5), "5x5": (5.0, 5.0), "10x10": (10.0, 10.0), "20x20": (20.0, 20.0), } + def _check_redis_params(self) -> None: + """Check that required parameters are in redis""" + with get_redis_connection() as redis_connection: + if ( + redis_connection.hget("top_camera_target_coords", "x_pixel_target") + is None + ): + msg = ( + "Top camera target coordinates are not set. " + + "Run optical centering calibration before running Partial UDC." + ) + logging.getLogger("user_level_log").error(msg) + raise QueueExecutionException(msg, self) + if ( + redis_connection.hget("top_camera_pixels_per_mm", "pixels_per_mm_x") + is None + ): + msg = ( + "Top camera pixels per mm are not set. " + + "Run optical centering calibration before running Partial UDC." + ) + logging.getLogger("user_level_log").error(msg) + raise QueueExecutionException(msg, self) + def run(self, dialog_box_parameters: dict) -> None: """ - + Runs the Partial UDC prefect flow Parameters ---------- @@ -55,6 +80,7 @@ def run(self, dialog_box_parameters: dict) -> None: ------- None """ + self._check_redis_params() self._state.value = "RUNNING" head_type = self.get_head_type() @@ -102,7 +128,7 @@ def run(self, dialog_box_parameters: dict) -> None: f"Detector distance corresponding to {default_resolution} A: {detector_distance} [m]" ) - if dialog_box_model.perform_screening: + if dialog_box_model.run_screening: screening_params = ScreeningParams( omega_range=screening_model.omega_range, exposure_time=screening_model.exposure_time, @@ -120,7 +146,7 @@ def run(self, dialog_box_parameters: dict) -> None: else: screening_params = None - if dialog_box_model.perform_full_dataset: + if dialog_box_model.run_full_dataset: full_dataset_params = FullDatasetParams( omega_range=full_dataset_model.omega_range, exposure_time=full_dataset_model.exposure_time, @@ -159,7 +185,7 @@ def run(self, dialog_box_parameters: dict) -> None: mount_pin_at_start_of_flow=False, add_dummy_pin_to_db=settings.ADD_DUMMY_PIN_TO_DB, project_strategy=ProjectStrategy( - use_screening=dialog_box_model.perform_screening + use_screening=dialog_box_model.run_screening ), ) prefect_parameters = { @@ -177,23 +203,23 @@ def run(self, dialog_box_parameters: dict) -> None: self._save_dialog_box_params_to_redis( grid_scan_model, collection_type="grid_scan" ) - if dialog_box_model.perform_screening: + if dialog_box_model.run_screening: self._save_dialog_box_params_to_redis( screening_model, collection_type="screening" ) - if dialog_box_model.perform_full_dataset: + if dialog_box_model.run_full_dataset: self._save_dialog_box_params_to_redis( full_dataset_model, collection_type="full_dataset" ) with get_redis_connection() as redis_connection: redis_connection.set( - "partial_udc:perform_screening", - 1 if dialog_box_model.perform_screening else 0, + "partial_udc:run_screening", + 1 if dialog_box_model.run_screening else 0, ) redis_connection.set( - "partial_udc:perform_full_dataset", - 1 if dialog_box_model.perform_full_dataset else 0, + "partial_udc:run_full_dataset", + 1 if dialog_box_model.run_full_dataset else 0, ) partial_udc_flow = MX3SyncPrefectClient( @@ -223,20 +249,22 @@ def dialog_box(self) -> dict: # Screening Properties resolution_limits = self.resolution.get_limits() - screening_properties = get_screening_schema(resolution_limits) + screening_properties = get_screening_schema(resolution_limits, partial_udc=True) # Full Dataset Properties - full_dataset_properties = get_full_dataset_schema(resolution_limits) + full_dataset_properties = get_full_dataset_schema( + resolution_limits, partial_udc=True + ) with get_redis_connection() as redis_connection: - perform_screening = bool( + run_screening = bool( int( - redis_connection.get("partial_udc:perform_screening") or 0 + redis_connection.get("partial_udc:run_screening") or 0 ) # False by default ) - perform_full_dataset = bool( + run_full_dataset = bool( int( - redis_connection.get("partial_udc:perform_full_dataset") or 0 + redis_connection.get("partial_udc:run_full_dataset") or 0 ) # False by default ) @@ -244,40 +272,39 @@ def dialog_box(self) -> dict: "properties": { "grid_scan": { "type": "object", - "title": "Grid Scan Parameters", + "title": "Optical and X-ray Centering", "properties": grid_scan_properties, "required": [ "md3_alignment_y_speed", "transmission", ], }, - "perform_screening": { + "run_screening": { "type": "boolean", - "title": "Perform Screening", - "default": perform_screening, + "title": "Screen", + "default": run_screening, }, - "perform_full_dataset": { + "run_full_dataset": { "type": "boolean", - "title": "Perform Full Dataset", - "default": perform_full_dataset, + "title": "Dataset", + "default": run_full_dataset, }, }, "required": ["grid_scan"], "allOf": [ { - "if": {"properties": {"perform_screening": {"const": True}}}, + "if": {"properties": {"run_screening": {"const": True}}}, "then": { "properties": { "screening": { "type": "object", - "title": "Screening Parameters", + "title": "Screen Parameters", "properties": screening_properties, "required": [ "exposure_time", "omega_range", "number_of_frames", "resolution", - "crystal_counter", "transmission", ], } @@ -285,7 +312,7 @@ def dialog_box(self) -> dict: }, }, { - "if": {"properties": {"perform_full_dataset": {"const": True}}}, + "if": {"properties": {"run_full_dataset": {"const": True}}}, "then": { "properties": { "full_dataset": { @@ -297,7 +324,6 @@ def dialog_box(self) -> dict: "omega_range", "number_of_frames", "resolution", - "crystal_counter", "transmission", "processing_pipeline", ], diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py index 0e42ee9696..9379567404 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/full_dataset.py @@ -1,7 +1,9 @@ from .utils import get_dialog_box_param -def get_full_dataset_schema(resolution_limits: tuple[float, float]) -> dict: +def get_full_dataset_schema( + resolution_limits: tuple[float, float], partial_udc: bool = False +) -> dict: properties = { "exposure_time": { "title": "Total Exposure Time [s]", @@ -64,7 +66,9 @@ def get_full_dataset_schema(resolution_limits: tuple[float, float]) -> dict: "processing_pipeline", collection_type="full_dataset" ), }, - "crystal_counter": { + } + if not partial_udc: + properties["crystal_counter"] = { "title": "Crystal ID", "type": "integer", "minimum": 0, @@ -72,6 +76,5 @@ def get_full_dataset_schema(resolution_limits: tuple[float, float]) -> dict: get_dialog_box_param("crystal_counter", collection_type="full_dataset") ), "widget": "textarea", - }, - } + } return properties diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py index edc83fc94f..8ac0c6baaa 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/grid_scan.py @@ -37,9 +37,9 @@ def get_grid_scan_schema(partial_udc: bool = False) -> dict: } if partial_udc: properties["grid_step"] = { - "title": "Grid Step [um]", + "title": "Grid Step Size [μm]", "type": "string", - "enum": ["5x5", "10x10", "20x20"], + "enum": ["2.5x2.5", "5x5", "10x10", "20x20"], "default": str( get_dialog_box_param("grid_step", collection_type="grid_scan") ), diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py index a3283d85c8..4018385bfe 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/dialog_boxes/screening.py @@ -1,7 +1,9 @@ from .utils import get_dialog_box_param -def get_screening_schema(resolution_limits: tuple[float, float]) -> dict: +def get_screening_schema( + resolution_limits: tuple[float, float], partial_udc: bool = False +) -> dict: properties = { "exposure_time": { "title": "Total Exposure Time [s]", @@ -50,7 +52,9 @@ def get_screening_schema(resolution_limits: tuple[float, float]) -> dict: ), "widget": "textarea", }, - "crystal_counter": { + } + if not partial_udc: + properties["crystal_counter"] = { "title": "Crystal ID", "type": "integer", "minimum": 0, @@ -58,6 +62,5 @@ def get_screening_schema(resolution_limits: tuple[float, float]) -> dict: get_dialog_box_param("crystal_counter", collection_type="screening") ), "widget": "textarea", - }, - } + } return properties diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py index a746c44051..7bbd786da0 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/partial_udc.py @@ -24,9 +24,9 @@ class PartialUDCGridScanDialogBox(GridScanDialogBox): class PartialUDCDialogBox(BaseModel): grid_scan: PartialUDCGridScanDialogBox - perform_screening: bool = False + run_screening: bool = False screening: ScreeningDialogBox | None = None - perform_full_dataset: bool = False + run_full_dataset: bool = False full_dataset: FullDatasetDialogBox | None = None From 26935ce8c95ae8037c2747b989eb8237ed343caf Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Mon, 1 Dec 2025 14:42:37 +1100 Subject: [PATCH 14/17] correctly use cal params --- .../ANSTO/prefect_flows/partial_udc.py | 56 +++++++++++++++++-- 1 file changed, 50 insertions(+), 6 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py index 88459838d5..af80a588fc 100644 --- a/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py +++ b/mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py @@ -25,7 +25,12 @@ class PartialUDCFlow(AbstractPrefectWorkflow): - """Prefect Raster Workflow""" + """Partial UDC Prefect Flow. Allows running: + - Optical and X-ray Centering + - Optical and X-ray Centering + Screening + - Optical and X-ray Centering + Full Dataset + - Optical and X-ray Centering + Screening + Full Dataset + """ def __init__( self, @@ -43,8 +48,20 @@ def __init__( "20x20": (20.0, 20.0), } - def _check_redis_params(self) -> None: - """Check that required parameters are in redis""" + def _check_optical_centering_params(self) -> None: + """Check that required optical centering parameters are in redis. + Currently checks por the following parameters: + - top_camera_target_coords:x_pixel_target + - top_camera_pixels_per_mm:pixels_per_mm_x + + Additionally gets the following optional parameters from redis: + - optical_centering:calibrated_alignment_z (default: 0.487 m) + - optical_centering:grid_height_scale_factor (default: 2) + + Returns + ------- + None + """ with get_redis_connection() as redis_connection: if ( redis_connection.hget("top_camera_target_coords", "x_pixel_target") @@ -67,6 +84,30 @@ def _check_redis_params(self) -> None: logging.getLogger("user_level_log").error(msg) raise QueueExecutionException(msg, self) + # The values below are saved in redis so that they can be changed on the fly + # if needed while we tune the optical centering params. + self.calibrated_alignment_z = redis_connection.get( + "optical_centering:calibrated_alignment_z" + ) + if self.calibrated_alignment_z is None: + self.calibrated_alignment_z = 0.487 + logging.getLogger("HWR").warning( + f"Calibrated alignment Z not found in redis, using default value: {self.calibrated_alignment_z} m" + ) + else: + self.calibrated_alignment_z = float(self.calibrated_alignment_z) + + self.grid_height_scale_factor = redis_connection.get( + "optical_centering:grid_height_scale_factor" + ) + if self.grid_height_scale_factor is None: + self.grid_height_scale_factor = 2 + logging.getLogger("HWR").warning( + f"Grid height scale factor not found in redis, using default value: {self.grid_height_scale_factor}" + ) + else: + self.grid_height_scale_factor = float(self.grid_height_scale_factor) + def run(self, dialog_box_parameters: dict) -> None: """ Runs the Partial UDC prefect flow @@ -80,7 +121,6 @@ def run(self, dialog_box_parameters: dict) -> None: ------- None """ - self._check_redis_params() self._state.value = "RUNNING" head_type = self.get_head_type() @@ -90,6 +130,8 @@ def run(self, dialog_box_parameters: dict) -> None: logging.getLogger("user_level_log").error(msg) raise QueueExecutionException(msg, self) + self._check_optical_centering_params() + dialog_box_model = PartialUDCDialogBox.model_validate(dialog_box_parameters) grid_scan_model = dialog_box_model.grid_scan @@ -167,9 +209,11 @@ def run(self, dialog_box_parameters: dict) -> None: partial_udc_config = SingleLoopDataCollectionConfig( optical_centering=OpticalCenteringParams( beam_position=(612, 512), - extra_config=OpticalCenteringExtraConfig(grid_height_scale_factor=2), + extra_config=OpticalCenteringExtraConfig( + grid_height_scale_factor=self.grid_height_scale_factor + ), grid_step=self.grid_step_map[grid_scan_model.grid_step], - calibrated_alignment_z=0.85, # TODO: maybe get from redis? + calibrated_alignment_z=self.calibrated_alignment_z, ), grid_scan=GridScanParams( omega_range=0, From 85a3f234f1febed423bd47c4c225824b458b3fc6 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Mon, 1 Dec 2025 15:01:22 +1100 Subject: [PATCH 15/17] update deployment name --- mxcubecore/configuration/ansto/config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mxcubecore/configuration/ansto/config.py b/mxcubecore/configuration/ansto/config.py index 87401e9969..5a62e23210 100644 --- a/mxcubecore/configuration/ansto/config.py +++ b/mxcubecore/configuration/ansto/config.py @@ -41,7 +41,8 @@ class PrefectSettings(BaseSettings): "mxcube-screening/plans", env="SCREENING_DEPLOYMENT_NAME" ) SAMPLE_CENTERING_PREFECT_DEPLOYMENT_NAME: str = Field( - "mxcube-sample-centering/plans", env="SAMPLE_CENTERING_PREFECT_DEPLOYMENT_NAME" + "optical-centering/optical_centering", + env="SAMPLE_CENTERING_PREFECT_DEPLOYMENT_NAME", ) USE_TOP_CAMERA: bool = Field( True, env="USE_TOP_CAMERA", description="False only for development" From 1b42f9844347616f7d3d7af809ea50e073135cd3 Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Mon, 1 Dec 2025 15:05:59 +1100 Subject: [PATCH 16/17] update optical centering payload --- mxcubecore/HardwareObjects/ANSTO/Diffractometer.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/mxcubecore/HardwareObjects/ANSTO/Diffractometer.py b/mxcubecore/HardwareObjects/ANSTO/Diffractometer.py index e91c9f1bb1..42f5fdc297 100644 --- a/mxcubecore/HardwareObjects/ANSTO/Diffractometer.py +++ b/mxcubecore/HardwareObjects/ANSTO/Diffractometer.py @@ -367,14 +367,8 @@ def start_automatic_centring( try: sample_centering = MX3SyncPrefectClient( name=settings.SAMPLE_CENTERING_PREFECT_DEPLOYMENT_NAME, - parameters={ - "sample_id": "test", - "beam_position": [self.beam_position[0], self.beam_position[1]], - "use_top_camera": settings.USE_TOP_CAMERA, - "calibrated_alignment_z": settings.CALIBRATED_ALIGNMENT_Z, - }, + parameters={}, ) - # NOTE: using asyncio.run() does not seem to work consistently sample_centering.trigger_flow(wait=True) logging.getLogger("HWR").debug("Optical centering finished") From 9e2b7350eaaeae5314223a7a76eec1aa2815af0e Mon Sep 17 00:00:00 2001 From: Francisco Hernandez Vivanco Date: Tue, 2 Dec 2025 10:18:15 +1100 Subject: [PATCH 17/17] version bump --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 246a4b1b73..4bec529277 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "mxcubecore" -version = "1.195.0.47" +version = "1.195.0.48" license = "LGPL-3.0-or-later" description = "Core libraries for the MXCuBE application" authors = ["The MXCuBE collaboration "]