From d43418e00341b4398f682b8392bda0f2e49ba7df Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 22 May 2026 15:26:10 +0200 Subject: [PATCH] DAQ: operations refactor. added providers, updated service creators --- src/aare/daq/daq.py | 1153 ++++++++----------------------------------- 1 file changed, 215 insertions(+), 938 deletions(-) diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 54e280a7..42d40271 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -3,14 +3,13 @@ import json import secrets import time from datetime import datetime, timezone -from math import ceil, floor +from math import ceil from pathlib import Path from typing import List, Tuple, Optional, Callable import cv2 import numpy as np from aareDB import SampleEventType -from jfjoch_client.exceptions import NotFoundException from jfjoch_client import ScanResult, ScanResultImagesInner from aare.common.tell_models import TellStateModel, TellPhaseEnum @@ -20,16 +19,13 @@ from aare.daq.aaredb import AareWrapper from aare.daq.config import BeamlineConfig, ABR_POS_MOUNT from aare.daq.config import BeamlineStateEnum from aare.daq.devices import BeamlineDevices -from aare.daq.mlbox import MlBox, MLBoxPredictionResult +from aare.daq.mlbox import MlBox from aare.common.beamline import MXBeamline from aare.common.coordinate import Coordinate, SmargonCoordinate, AerotechCoordinate from aare.common.diffraction_geometry import DiffractionGeometry from aare.common.logger_config import setup_logger from aare.common.logger_events import ( - geom_log_context, - log_duration, log_timing, - log_ml_bundle_meta, merge_log_context, raster_request_log_context, rotation_request_log_context, @@ -48,17 +44,30 @@ from aare.common.automation_models import ( StepStatus, WorkflowStateKind, ) -from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem, grid_to_image_id -from aare.common.simulate_raster import generate_no_beam_scan_result -from aare.common.find_xtal import rebuild_array_from_scan_results, raster_centre_of_mass, create_quality_filtered_array +from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan from aare.common.sample_geometry import SampleGeometryModel + from aare.daq.operations.face_detection import FaceDetectionContext, FaceDetectionService, FaceDetectionResult from aare.daq.operations.loop_centering import LoopCenteringService, LoopCenteringContext from aare.daq.operations.loop_centering.models import LoopCenteringSettings -from aare.daq.operations.mounting import MountingContext, MountingService -from aare.daq.operations.mounting.models import MountingResult - +from aare.daq.operations.mounting.service import MountingService +from aare.daq.operations.mounting.models import MountingResult, MountingContext +from aare.daq.operations.raster.models import RasterContext +from aare.daq.operations.raster.service import RasterService +from aare.daq.operations.common.ml_bounding_box import get_ml_bounding_box +from aare.daq.operations.screenshot.service import ScreenshotService +from aare.daq.operations.common.simulate_scan_result import ( + build_fake_rotation_result, +) +from aare.daq.operations.common.runtime import ( + DAQRuntimeState, + FaceDetectionProgressEmitter, + OperationServices, + PredictionProvider, + StateController, + TraceWriter, +) from aare.devices.area_detector import AutoEnum from aare.devices.jfjoch import JFJochWrapper from aare.devices.mx_lib import clean_filename @@ -71,7 +80,6 @@ from aare.common.exception_handler import ( MountingFailed, WarningTellException, CriticalTellException, - AXCFailed, SmargonCommunicationError, TellCommunicationError, BECCommunicationError, @@ -81,14 +89,111 @@ from aare.common.exception_handler import ( UnmountingFailed, DataCollectionException, RasterScanException, - TellMountFailedException, BeamlineBusyTimeoutException, BeamlineBusyException, AutoRasterSampleSkipped ) -from aare.devices.tell_client import TellEventValueEnum logger = setup_logger("aareDAQ") + + +class _DAQSampleProvider: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + @property + def sample(self) -> SampleShortInfo | None: + return self._daq.sample + + +class _DAQSampleGeometryProvider: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + @property + def sample_geometry(self) -> SampleGeometryModel: + return self._daq.sample_geometry + + +class _DAQNonCriticalRunner: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + def __call__( + self, + action, + *, + description: str, + sample: SampleShortInfo | None = None, + ) -> object | None: + return self._daq._run_noncritical(action, description=description, sample=sample) + + +class _DAQScreenshotSampleProvider: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + @property + def sample(self) -> SampleShortInfo | None: + return self._daq.sample + + +class _DAQPGroupProvider: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + @property + def pgroup(self) -> str | None: + return self._daq._AareDAQ__cfg.pgroup + + +class _DAQStatusProvider: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + @property + def status(self) -> DAQStatusModel: + return self._daq.status + + +class _DAQStateSetter: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + def set_state(self, target: BeamlineStateEnum) -> None: + self._daq._AareDAQ__set_state(target) + + +class _DAQTraceAppender: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + def append_smargon_trace(self, *, sample_id: int | None, event: str) -> None: + self._daq._append_smargon_trace(sample_id=sample_id, event=event) + + +class _LoopCenteringPredictionGetter: + def __init__(self, daq: "AareDAQ", settings: LoopCenteringSettings): + self._daq = daq + self._settings = settings + + def get_predictions(self): + return self._daq._AareDAQ__mlbox.predict_all_best( + overlap_with_pin=self._settings.overlap_with_pin, + confidence_min=self._settings.confidence_min, + return_image=True, + return_bundle_meta=True, + ) + + +class _FaceDetectionProgressReporter: + def __init__(self, daq: "AareDAQ"): + self._daq = daq + + def emit_progress(self, payload: dict) -> None: + self._daq._emit_face_detection_progress(payload) + + #TODO tidy up DAQ - migrate functions into different scripts, to reduce size? #TODO investigate using a state machine within each operation to reduce callbacks? class AareDAQ: @@ -125,6 +230,16 @@ class AareDAQ: self._automation_last_sample_name = "" self._automation_samples_in_queue = 0 self._last_mount_error_message = "" + + self._screenshot_service = ScreenshotService( + mlbox=self.__mlbox, + aare=self.__aare, + logger=logger, + run_noncritical=_DAQNonCriticalRunner(self), + sample_provider=_DAQScreenshotSampleProvider(self), + pgroup_provider=_DAQPGroupProvider(self), + ) + self._create_mounting_service().reset_mount_failure_counter("DAQ startup") def _is_hardware_failure(self, error: Exception) -> bool: @@ -528,11 +643,22 @@ class AareDAQ: progress.success = success self._emit_automation_progress(progress) - #-------------------------------------------- # Operation Services #-------------------------------------------- + def _build_runtime_state(self) -> DAQRuntimeState: + return DAQRuntimeState( + sample_provider=_DAQSampleProvider(self), + sample_geometry_provider=_DAQSampleGeometryProvider(self), + status_provider=_DAQStatusProvider(self), + ) + + def _build_operation_services(self) -> OperationServices: + return OperationServices( + screenshots=self._screenshot_service, + ) + def _create_loop_centering_settings(self) -> LoopCenteringSettings: return LoopCenteringSettings() @@ -545,14 +671,13 @@ class AareDAQ: devs=self.__devs, mlbox=self.__mlbox, settings=settings, - sample_geometry_provider=lambda: self.sample_geometry, - save_screenshot_db=self.save_screenshot_db, - append_smargon_trace=self._append_smargon_trace, - get_predictions=lambda: self.__mlbox.predict_all_best( - overlap_with_pin=settings.overlap_with_pin, - confidence_min=settings.confidence_min, - return_image=True, - return_bundle_meta=True, + runtime=self._build_runtime_state(), + services=self._build_operation_services(), + trace_writer=TraceWriter( + appender=_DAQTraceAppender(self), + ), + prediction_provider=PredictionProvider( + getter=_LoopCenteringPredictionGetter(self, settings), ), ), logger=logger, @@ -564,8 +689,10 @@ class AareDAQ: cfg=self.__cfg, devs=self.__devs, mlbox=self.__mlbox, - sample_geometry_provider=lambda: self.sample_geometry, - emit_progress=self._emit_face_detection_progress, + runtime=self._build_runtime_state(), + progress_emitter=FaceDetectionProgressEmitter( + reporter=_FaceDetectionProgressReporter(self), + ), ), logger=logger, ) @@ -588,33 +715,33 @@ class AareDAQ: mlbox=self.__mlbox, jfjoch=self.__jfjoch, aare=self.__aare, - sample_provider=lambda: self.sample, - sample_geometry_provider=lambda: self.sample_geometry, - status_provider=lambda: self.status, - set_state=self.__set_state, - save_screenshot_db=self.save_screenshot_db, - upload_raster_diffraction_preview=self._upload_raster_diffraction_preview, - auto_center_line_scan_top_left=self._auto_center_line_scan_top_left, - ml_bounding_box=self.__ml_bounding_box, - ), - logger=logger, - ) - - def _create_rotation_service(self): - return RotationService( - context=RotationContext( - cfg=self.__cfg, - devs=self.__devs, - jfjoch=self.__jfjoch, - aare=self.__aare, - sample_provider=lambda: self.sample, - sample_geometry_provider=lambda: self.sample_geometry, - status_provider=lambda: self.status, - set_state=self.__set_state, - save_screenshot_db=self.save_screenshot_db, + runtime=self._build_runtime_state(), + state_controller=StateController( + setter=_DAQStateSetter(self), + ), + services=self._build_operation_services(), + auto_raster_max_images=self.AUTO_RASTER_MAX_IMAGES, + auto_raster_min_cell_size_mm=self.AUTO_RASTER_MIN_CELL_SIZE_MM, + auto_raster_skip_if_exceed_max_image_threshold=self.AUTO_RASTER_SKIP_IF_EXCEED_MAX_IMAGE_THRESHOLD, ), logger=logger, ) + # + # def _create_rotation_service(self): + # return RotationService( + # context=RotationContext( + # cfg=self.__cfg, + # devs=self.__devs, + # jfjoch=self.__jfjoch, + # aare=self.__aare, + # sample_provider=lambda: self.sample, + # sample_geometry_provider=lambda: self.sample_geometry, + # status_provider=lambda: self.status, + # set_state=self.__set_state, + # save_screenshot_db=self.save_screenshot_db, + # ), + # logger=logger, + # ) #-------------------------------------------- @@ -937,11 +1064,13 @@ class AareDAQ: ), ) + raster_service = self._create_raster_service() + if auto_center: setup_request = copy.deepcopy(grid_request) setup_request.smargon_top_left = None self.__setup_datacollection(request=setup_request) - result = self.__auto_center(grid_request) + result = raster_service.execute_auto_center(grid_request) if result is None: logger.error( "Raster sequence returned no result after auto-centering", @@ -962,7 +1091,7 @@ class AareDAQ: self.__setup_datacollection(request=grid_request) self.__set_state(BeamlineStateEnum.DataCollection) - raster_result = self.__raster(grid_request) + raster_result = raster_service.execute(grid_request) result = CompletedRasterGrid(r=[raster_result]) self.__set_state(BeamlineStateEnum.SampleAlignment) @@ -1449,13 +1578,18 @@ class AareDAQ: self.__cfg.state_busy = False raise - def _execute_dry(self, park=True): - self._create_mounting_service().dry(park=park) + def _execute_dry(self, park:bool=True, unmount:bool=False): + self._create_mounting_service().dry(park=park, unmount=unmount) - def park_and_dry(self, park = True): + def park_and_dry(self, park = True, unmount: bool = False): self.__cfg.try_set_busy(timeout=360) try: - self._execute_dry(park=park) + self.__set_state(BeamlineStateEnum.RobotSampleExchange) + except TransformationInvalidException as e: + logger.error(f"Failed to go to robot sample exchange: {e}") + logger.warning(f"trying to day and park without unmounting first") + try: + self._execute_dry(park=park, unmount=unmount) self.__cfg.state_busy = False except Exception as e: self.__cfg.state_busy = False @@ -1582,177 +1716,6 @@ class AareDAQ: def auto_exposure(self): self.__devs.samcam_auto(AutoEnum.ONCE) - def __auto_center(self, request: RasterGridRequest) -> CompletedRasterGrid | None: - # TODO do we need to handle the two grid scans differently? - - sample = self.sample - - if sample is None: - raise Exception("Sample must be mounted to auto center") - - old_prefix = request.file_prefix - geom = self.sample_geometry - - logger.info( - "Starting auto-center raster workflow", - extra=merge_log_context( - sample_log_context(sample), - geom_log_context(geom, prefix="current_"), - { - "omega_deg": geom.omega_deg, - "beam_x_pxl": geom.beam_location_pxl.x, - "beam_y_pxl": geom.beam_location_pxl.y, - "pixel_in_mm": geom.pixel_in_mm, - "file_prefix": old_prefix, - }, - ), - ) - self.__aare.send_sample_event(self.sample, SampleEventType.RASTERING, - comment=f"Raster at {geom.omega_deg:.1f} deg") - r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg") - - if r is None: - logger.warning( - "No ML bounding box found at primary angle during auto-center raster", - extra=merge_log_context( - sample_log_context(sample), - { - "omega_deg": geom.omega_deg, - "file_prefix": old_prefix, - }, - ), - ) - self.__devs.aerotech_omega = geom.omega_deg + 90.0 - time.sleep(0.2) - r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg + 90.0:.2f}deg") - - if r is not None: - logger.info( - "ML bounding box found for auto-center raster", - extra=merge_log_context( - sample_log_context(sample), - { - "ml_omega_deg": r.omega_deg, - "ml_n_x": r.n_x, - "ml_n_y": r.n_y, - "ml_grid_size_x_mm": getattr(r.grid_size_mm, "x", None), - "ml_grid_size_y_mm": getattr(r.grid_size_mm, "y", None), - "ml_top_left_x_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "x", None), - "ml_top_left_y_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "y", None), - "ml_top_left_z_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "z", None), - }, - ), - ) - - geom = self.sample_geometry - grid = copy.deepcopy(request) - grid.smargon_top_left = r.smargon_top_left - grid.grid_size_mm = r.grid_size_mm - grid.n_x = r.n_x - grid.n_y = r.n_y - grid.file_prefix = f"{old_prefix}_{grid.omega_deg}deg" - grid.omega_deg = geom.omega_deg - - status = self.status - - if not self.__cfg.simulated_detector: - logger.info("initialise detector") - self.__jfjoch.measure_raster(grid, status) - logger.info("detector initialised") - else: - logger.info("Simulated detector mode enabled; using fake raster result.") - - self.__set_state(BeamlineStateEnum.DataCollection) - - logger.info( - "Running first auto-center raster", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(grid), - { - "top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None), - }, - ), - ) - res1 = self.__raster(grid) - - if res1 is None: - self.__aare.send_sample_event(self.sample, SampleEventType.RASTERINGFAILED, - comment=f"No ML Box detected {geom.omega_deg:.1f} deg") - - grid.omega_deg += 90 - self.__aare.send_sample_event(self.sample, SampleEventType.RASTERING, - comment=f"Raster at {geom.omega_deg:.1f} deg") - self.__devs.aerotech_omega = grid.omega_deg - - grid.n_x = 1 - grid.file_prefix = f"{old_prefix}_{grid.omega_deg}deg" - geom = self.sample_geometry - grid.grid_size_mm = Coordinate(x=geom.beam_size_mm.x, y=geom.beam_size_mm.y * 0.25) - grid.smargon_top_left, grid.n_y = self._auto_center_line_scan_top_left( - omega_deg=grid.omega_deg, - file_prefix=grid.file_prefix, - grid_size_mm=grid.grid_size_mm, - default_n_y=50, - y_retarget_threshold_mm=max( - geom.beam_size_mm.y * 2.0, - grid.grid_size_mm.y * 4.0, - ), - y_padding_fraction_each_side=0.10, - ) - - logger.info( - "Prepared second auto-center raster", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(grid), - { - "top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None), - }, - ), - ) - status = self.status - if not self.__cfg.simulated_detector: - logger.info(f"initialise detector for raster at {grid.omega_deg}") - self.__jfjoch.measure_raster(grid, status) - logger.info("detector initialised") - else: - logger.info("Simulated detector mode enabled; using fake raster result.") - - self.__set_state(BeamlineStateEnum.DataCollection) - - logger.info( - "Running second auto-center raster", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(grid), - ), - ) - res2 = self.__raster(grid) - - if res2 is None: - self.__aare.send_sample_event(self.sample, SampleEventType.RASTERINGFAILED, - comment=f"No ML Box detected {geom.omega_deg:.1f} deg") - - return CompletedRasterGrid(r=[res1, res2]) - else: - logger.error( - "Auto-center raster aborted because no ML bounding box was found at either angle", - extra=merge_log_context( - sample_log_context(sample), - { - "primary_omega_deg": geom.omega_deg, - "secondary_omega_deg": geom.omega_deg + 90.0, - "file_prefix": old_prefix, - }, - ), - ) - return None - def __setup_datacollection(self, request: RasterGridRequest | RotationScanRequest, screening: bool = False): request_omega = getattr(request, "omega_deg", None) if request_omega is None: @@ -1820,343 +1783,18 @@ class AareDAQ: # self.__devs.transmission.wait() return - def _build_fake_scan_result(self, *, file_prefix: str | None, image_count: int, - rotation: bool = False) -> ScanResult: - total_images = max(1, image_count) - - start_angle = 0.0 - if rotation: - try: - start_angle = float(self.omega) - except Exception: - start_angle = 0.0 - - angle_step = 360.0 / total_images if rotation else 0.0 - - images = [ - ScanResultImagesInner( - number=i, - efficiency=1.0, - bkg=0.0, - spots=0, - spots_low_res=0, - spots_indexed=0, - index=0, - b=0.0, - angle=start_angle + i * angle_step if rotation else None, - ) - for i in range(total_images) - ] - return ScanResult(file_prefix=file_prefix, images=images) - def _build_fake_rotation_result(self, request: RotationScanRequest) -> CompletedRotationScan: - result = self._build_fake_scan_result( - file_prefix=request.file_prefix, - image_count=request.steps, - rotation=True - ) - return CompletedRotationScan( - request=copy.deepcopy(request), - result=result, - ) - - def _build_fake_raster_result(self, request: RasterGridRequest) -> CompletedRasterGridElem: - result = self._build_fake_scan_result( - file_prefix=request.file_prefix, - image_count=request.n_x * request.n_y, - ) - return CompletedRasterGridElem( - request=copy.deepcopy(request), - result=result, - centre_of_mass=None, - ) - - @staticmethod - def _scan_result_image_count(scan_result: ScanResult | None) -> int | None: - if scan_result is None or scan_result.images is None: - return None - return len(scan_result.images) - - def _upload_raster_diffraction_preview( - self, - *, - sample_id: int, - filename: str, - image_id: int, - scan_result: ScanResult | None, - request: RasterGridRequest, - ) -> None: - image_count = self._scan_result_image_count(scan_result) - - if image_count is not None and not (0 <= image_id < image_count): - logger.warning( - "Skipping diffraction image upload because image id is outside scan result range", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - { - "diffraction_image_id": image_id, - "image_count": image_count, - }, - ), - ) - return - + start_angle = 0.0 try: - diffraction_image = self.__jfjoch.get_diffraction_image(image_id) - except NotFoundException: - logger.warning( - "JFJoch diffraction preview image was not found after raster; continuing without upload", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - { - "diffraction_image_id": image_id, - "image_count": image_count, - }, - ), - exc_info=True, - ) - return + start_angle = float(self.omega) + except Exception: + pass - self.__aare.upload_jpg(sample_id, filename, diffraction_image) - - @staticmethod - def _grid_image_id_from_centre_offset( - *, - x_mm: float, - y_mm: float, - request: RasterGridRequest, - ) -> int: - if request.grid_size_mm.x <= 0 or request.grid_size_mm.y <= 0: - raise ValueError("grid_size_mm must be positive") - if request.n_x < 1 or request.n_y < 1: - raise ValueError("Raster grid dimensions must be >= 1") - - grid_x = floor(x_mm / request.grid_size_mm.x) + 1 - grid_y = floor(y_mm / request.grid_size_mm.y) + 1 - - grid_x = min(max(1, grid_x), request.n_x) - grid_y = min(max(1, grid_y), request.n_y) - - return grid_to_image_id( - grid_x=grid_x, - grid_y=grid_y, - number_of_cols=request.n_x, + return build_fake_rotation_result( + request, + start_angle=start_angle, ) - def __raster(self, request: RasterGridRequest, wait_for_screenshot: float | None = 0.3) -> CompletedRasterGridElem: - - total_time = request.exp_time_s * request.n_x * request.n_y + request.n_y * 0.3 - - if total_time > 1200: - raise RasterScanException(f"Raster scan is too long {total_time}s > 20min") - - status = self.status - - smargon_top_left =request.smargon_top_left - self.__devs.set_smargon_pos( - SmargonCoordinate( - sh_mm=smargon_top_left.sh_mm, - phi_deg=smargon_top_left.phi_deg, - chi_deg=smargon_top_left.chi_deg, - ) - ) - - logger.info( - "Starting raster acquisition", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - {"state": getattr(status, "state", None)}, - ), - ) - - try: - if self.sample is not None and self.sample.db_id is not None: - self.__aare.create_gridscan_run(self.sample, request, status) - - if not self.__cfg.simulated_detector: - try: - self.__jfjoch.wait_till_running(timeout=60.0) - except Exception as e: - self._raise_if_critical_jfjoch_detector_error(e, command="wait_till_running") - raise - else: - logger.info("Simulated detector mode enabled; faking jfjoch intilalisation.") - logger.debug(f"Starting grid scan with {request.n_x}x{request.n_y} points, exp time {request.exp_time_s}s") - self.__devs.aerotech.grid_scan( - grid_elem_size_y_um=request.grid_size_mm.y * 1000, - grid_elem_size_x_um=request.grid_size_mm.x * 1000, - grid_elem_count_x=request.n_x, - grid_elem_count_y=request.n_y, - time_sec=request.exp_time_s, - run_async=True, - ) - - self.__devs.aerotech.wait_till_done(timeout=int(round(total_time + total_time * 0.1 + 60, 0))) - # go back to aerotech x,y,z home not U home (0 degrees). - if isinstance(self.__cfg.abr_meas_pos, Coordinate): - coord = self.__cfg.abr_meas_pos - else: - coord = self.__cfg.abr_meas_pos.at_mm - self.__devs.aerotech_pos = AerotechCoordinate(at_mm=coord, omega_deg=self.__devs.aerotech_omega) - self.__devs.aerotech.wait_till_done(timeout=int(60)) - - if self.__cfg.simulated_detector: - logger.info("Simulated detector mode enabled; using fake raster result.") - scan_result = generate_no_beam_scan_result(request) - result_array = create_quality_filtered_array(scan_result.images, - 'spots_low_res', min_spots=None, - min_efficiency=1.0, min_background=None, - min_low_res_spots=None) - - com = raster_centre_of_mass(result_array, scan_result.images) - target_coor = com.get_com_mm(request) - target_coor_offset = self.sample_geometry.smargon_nudge(target_coor) - target_smargon = SmargonCoordinate( - sh_mm=request.smargon_top_left.sh_mm + target_coor_offset, - phi_deg=request.smargon_top_left.phi_deg, - chi_deg=request.smargon_top_left.chi_deg - ) - - else: - #TODO how to handle no beam grids - do we always want to go to the center? - #TODO How to handle grids with no/minimal diffraction - use bkg instead? - try: - scan_result = self.__jfjoch.wait_till_done(60) - except Exception as e: - self._raise_if_critical_jfjoch_detector_error(e, command="wait_till_done") - raise - com = None - if scan_result is None: - logger.error( - "JFJoch returned no ScanResult for raster", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - {"exp_time_s": request.exp_time_s}, - ), - ) - if request.n_x == 1: - x = request.grid_size_mm.x / 2.0 - else: - x = ((request.n_x - 1) * request.grid_size_mm.x) / 2.0 - y = ((request.n_y - 1) * request.grid_size_mm.y) / 2.0 - target_coor_offset = self.sample_geometry.smargon_nudge(Coordinate(x=x, y=y)) - - logger.info( - "Calculated raster centre offset", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - { - "centre_offset_x_mm": target_coor_offset.x, - "centre_offset_y_mm": target_coor_offset.y, - "centre_offset_z_mm": target_coor_offset.z, - "grid_half_width_x_mm": x, - "grid_half_height_y_mm": y, - "top_left_x_mm": getattr(request.smargon_top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(request.smargon_top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(request.smargon_top_left.sh_mm, "z", None), - }, - ), - ) - - logger.info(f"moving Smargon to grid centre offset {target_coor_offset}") - target_smargon = SmargonCoordinate( - sh_mm=request.smargon_top_left.sh_mm + target_coor_offset, - phi_deg=request.smargon_top_left.phi_deg, - chi_deg=request.smargon_top_left.chi_deg - ) - - self.__devs.smargon_pos = target_smargon - self.__devs.smargon_wait(timeout=180) - - logger.info( - "Moved Smargon to raster centre", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - { - "centre_sh_x_mm": target_smargon.sh_mm.x, - "centre_sh_y_mm": target_smargon.sh_mm.y, - "centre_sh_z_mm": target_smargon.sh_mm.z, - "centre_phi_deg": target_smargon.phi_deg, - "centre_chi_deg": target_smargon.chi_deg, - }, - ), - ) - - - sample_id = self.sample.db_id if self.sample and self.sample.db_id is not None else None - if sample_id: - logger.debug(f"moving to XtalSnapshot to take a screenshot of the sample") - self.__set_state(BeamlineStateEnum.XtalSnapshot) - if wait_for_screenshot and wait_for_screenshot > 0: - time.sleep(wait_for_screenshot) - self.save_screenshot_db(sample_id, f"{sample_id}_post_raster_{int(request.omega_deg)}deg") - self.__aare.ingest_gridscan( - sample=self.sample, - raster_result=scan_result, - raster_request=request, - geom=self.sample_geometry, - com=None, - beam_mark_pxl=self.__cfg.get_beam_mark(self.zoom), - ) - if com is not None and com.max_image is not None: - diffraction_image_filename = f"{self.sample.db_id}_best_diffraction_from_raster_image_{com.max_image}" - diffraction_image_id = com.max_image - else: - diffraction_image_filename = f"{self.sample.db_id}_diffraction_image_near_grid_scan" - diffraction_image_id = self._grid_image_id_from_centre_offset( - x_mm=x, - y_mm=y, - request=request, - ) - - self._upload_raster_diffraction_preview( - sample_id=self.sample.db_id, - filename=diffraction_image_filename, - image_id=diffraction_image_id, - scan_result=scan_result, - request=request, - ) - self.__aare.send_sample_event( - self.sample, - event_type=SampleEventType.RASTERED, - comment=f"Raster completed at {request.omega_deg:.1f} deg" - ) - - logger.info( - "Raster finished", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - { - "scan_result_is_none": scan_result is None, - "sample_id": sample_id, - }, - ), - ) - - return CompletedRasterGridElem( - request=copy.deepcopy(request), - result=scan_result, - centre_of_mass=com, - ) - - except Exception as e: - logger.exception( - f"Failed during raster: {e}", - extra=merge_log_context( - sample_log_context(self.sample), - raster_request_log_context(request), - {"exp_time_s": request.exp_time_s}, - ), - ) - raise - def measure_raster(self, request: RasterGridRequest, auto_center: bool) -> CompletedRasterGrid: """ Execute a raster scan. @@ -2383,188 +2021,6 @@ class AareDAQ: def get_beam_mark(self): return self.__cfg.get_beam_mark(self.__devs.zoom) - @classmethod - def _scale_auto_raster_grid( - cls, - *, - n_x: int, - n_y: int, - grid_size: Coordinate, - skip: bool = False, - ) -> tuple[int, int, Coordinate]: - """ - Limit auto-raster image count while preserving the physical raster footprint. - - If the ML box would generate too many images, reduce the number of cells and - increase the cell size proportionally. The total raster width/height stays - the same, but the scan is sampled more coarsely. - """ - n_x = max(1, int(n_x)) - n_y = max(1, int(n_y)) - - image_count = n_x * n_y - if image_count <= cls.AUTO_RASTER_MAX_IMAGES: - return n_x, n_y, grid_size - - if skip: - raise AutoRasterSampleSkipped( - f"Auto-raster grid has {image_count} images, which exceeds " - f"the automation limit of {cls.AUTO_RASTER_MAX_IMAGES}; skipping sample" - ) - - physical_size_x_mm = n_x * grid_size.x - physical_size_y_mm = n_y * grid_size.y - - scale = (image_count / cls.AUTO_RASTER_MAX_IMAGES) ** 0.5 - scaled_n_x = max(1, int(floor(n_x / scale))) - scaled_n_y = max(1, int(floor(n_y / scale))) - - while scaled_n_x * scaled_n_y > cls.AUTO_RASTER_MAX_IMAGES: - if scaled_n_x >= scaled_n_y and scaled_n_x > 1: - scaled_n_x -= 1 - elif scaled_n_y > 1: - scaled_n_y -= 1 - else: - break - - scaled_grid_size = Coordinate( - x=max(cls.AUTO_RASTER_MIN_CELL_SIZE_MM, physical_size_x_mm / scaled_n_x), - y=max(cls.AUTO_RASTER_MIN_CELL_SIZE_MM, physical_size_y_mm / scaled_n_y), - ) - - return scaled_n_x, scaled_n_y, scaled_grid_size - - def __ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None) -> RasterGridRequest | None: - time.sleep(0.2) # Just to be sure image is stable - prediction_result: MLBoxPredictionResult = self.__mlbox.predict( - preferred_class=(3, 0), - return_image=True, - return_bundle_meta=True, - ) - m = prediction_result.box - bundle_image = prediction_result.image - log_ml_bundle_meta( - logger, - f"ml_bounding_box:{filename or 'unnamed'}", - target_point=prediction_result.target_point, - focus=prediction_result.focus, - ) - - if m is None: - logger.warning( - "ML bounding box returned no detection", - extra={ - "sample_id": sample_id, - "ml_image_name": filename, - "target_point": prediction_result.target_point, - }, - ) - if filename is not None and bundle_image is not None: - self.__aare.upload_image(sample_id, f"{filename}_no_detection", bundle_image) - return None - x1, y1, x2, y2 = m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y - - if filename is not None and bundle_image is not None: - upload_image = bundle_image.copy() - cv2.rectangle(upload_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) - self.__aare.upload_image(sample_id, filename, upload_image) - - geom = self.sample_geometry - logger.info( - "ML bounding box selected", - extra=merge_log_context( - sample_log_context(self.sample), - { - "sample_id": sample_id, - "ml_image_name": filename, - }, - geom_log_context(geom), - { - "box_x1": x1, - "box_y1": y1, - "box_x2": x2, - "box_y2": y2, - "target_point": prediction_result.target_point, - }, - ), - ) - - start_coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1)) - grid_size = Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.8) - n_x = max(1, abs(ceil((x2 - x1) * geom.pixel_in_mm / grid_size.x))) - n_y = max(1, abs(ceil((y2 - y1) * geom.pixel_in_mm / grid_size.y))) - - original_n_x = n_x - original_n_y = n_y - original_grid_size = grid_size - #TODO should we have seperate handling for manual raster? - n_x, n_y, grid_size = self._scale_auto_raster_grid( - n_x=n_x, - n_y=n_y, - grid_size=grid_size, - skip = self.AUTO_RASTER_SKIP_IF_EXCEED_MAX_IMAGE_THRESHOLD - ) - - if (n_x, n_y, grid_size.x, grid_size.y) != ( - original_n_x, - original_n_y, - original_grid_size.x, - original_grid_size.y, - ): - logger.info( - "Scaled ML raster grid to stay within auto-raster image limit", - extra=merge_log_context( - sample_log_context(self.sample), - { - "sample_id": sample_id, - "ml_image_name": filename, - "max_images": self.AUTO_RASTER_MAX_IMAGES, - "original_n_x": original_n_x, - "original_n_y": original_n_y, - "original_image_count": original_n_x * original_n_y, - "original_grid_size_x_mm": original_grid_size.x, - "original_grid_size_y_mm": original_grid_size.y, - "scaled_n_x": n_x, - "scaled_n_y": n_y, - "scaled_image_count": n_x * n_y, - "scaled_grid_size_x_mm": grid_size.x, - "scaled_grid_size_y_mm": grid_size.y, - }, - ), - ) - - logger.info( - "Converted ML bounding box to raster request", - extra=merge_log_context( - sample_log_context(self.sample), - { - "sample_id": sample_id, - "ml_image_name": filename, - "start_sh_x_mm": start_coord.x, - "start_sh_y_mm": start_coord.y, - "start_sh_z_mm": start_coord.z, - "grid_size_x_mm": grid_size.x, - "grid_size_y_mm": grid_size.y, - "n_x": n_x, - "n_y": n_y, - "smargon_phi_deg": geom.smargon.phi_deg, - "smargon_chi_deg": geom.smargon.chi_deg, - }, - ), - ) - - return RasterGridRequest( - exp_time_s=0.01, #TODO different settings for different beamlines - transmission=1.0, - smargon_top_left=SmargonCoordinate(chi_deg=geom.smargon.chi_deg, - phi_deg=geom.smargon.phi_deg, - sh_mm=start_coord), - n_x=n_x, - n_y=n_y, - grid_size_mm=grid_size, - omega_deg=geom.omega_deg - ) - def ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None) -> RasterGridRequest | None: """ Request an ML-based bounding box for the sample. @@ -2578,7 +2034,17 @@ class AareDAQ: """ try: self.__cfg.try_set_busy(timeout=360) - r = self.__ml_bounding_box(sample_id, filename) + r = get_ml_bounding_box( + mlbox=self.__mlbox, + sample=self.sample, + sample_geometry=self.sample_geometry, + filename=filename, + upload_image=self.__aare.upload_image, + logger=logger, + max_images=self.AUTO_RASTER_MAX_IMAGES, + min_cell_size_mm=self.AUTO_RASTER_MAX_IMAGES, + skip_if_exceed_max_image_threshold=self.AUTO_RASTER_SKIP_IF_EXCEED_MAX_IMAGE_THRESHOLD + ) self.__cfg.state_busy = False return r except Exception: @@ -2609,160 +2075,6 @@ class AareDAQ: finally: self.__cfg.state_busy = False - def _auto_center_line_scan_top_left( - self, - *, - omega_deg: float, - file_prefix: str | None, - grid_size_mm: Coordinate, - default_n_y: int = 50, - y_retarget_threshold_mm: float | None = None, - y_padding_fraction_each_side: float = 0.10, - ) -> tuple[SmargonCoordinate, int]: - geom = self.sample_geometry - beam_x_pxl = geom.beam_location_pxl.x - beam_y_pxl = geom.beam_location_pxl.y - - line_scan_centre = geom.picture_to_smargon( - Coordinate(x=beam_x_pxl, y=beam_y_pxl) - ) - n_y = default_n_y - - prediction_result: MLBoxPredictionResult = self.__mlbox.predict( - preferred_class=(3, 0), - return_image=False, - return_bundle_meta=True, - ) - log_ml_bundle_meta( - logger, - f"auto_center_line_scan_{omega_deg:.2f}deg", - target_point=prediction_result.target_point, - focus=prediction_result.focus, - ) - target_point = prediction_result.target_point - prediction_box = prediction_result.box - - if target_point is not None: - target_y_pxl = target_point[1] - y_delta_mm = abs(target_y_pxl - beam_y_pxl) * geom.pixel_in_mm - - if y_retarget_threshold_mm is None: - y_retarget_threshold_mm = geom.beam_size_mm.y * 2.0 - - if y_delta_mm > y_retarget_threshold_mm: - line_scan_centre = geom.picture_to_smargon( - Coordinate(x=beam_x_pxl, y=target_y_pxl) - ) - logger.info( - "Using ML target y for second auto-center raster", - extra=merge_log_context( - sample_log_context(self.sample), - { - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "beam_y_pxl": beam_y_pxl, - "target_y_pxl": target_y_pxl, - "y_delta_mm": y_delta_mm, - "threshold_mm": y_retarget_threshold_mm, - "target_sh_x_mm": line_scan_centre.x, - "target_sh_y_mm": line_scan_centre.y, - "target_sh_z_mm": line_scan_centre.z, - }, - ), - ) - else: - logger.info( - "Keeping beam-centred y line scan because ML target y shift is small", - extra=merge_log_context( - sample_log_context(self.sample), - { - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "beam_y_pxl": beam_y_pxl, - "target_y_pxl": target_y_pxl, - "y_delta_mm": y_delta_mm, - "threshold_mm": y_retarget_threshold_mm, - }, - ), - ) - else: - logger.info( - "No ML target point for second auto-center raster; using beam-centred line scan", - extra=merge_log_context( - sample_log_context(self.sample), - { - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "beam_x_pxl": beam_x_pxl, - "beam_y_pxl": beam_y_pxl, - }, - ), - ) - - if prediction_box is not None and prediction_box.box is not None and grid_size_mm.y > 0: - box_height_pxl = abs(prediction_box.box.bottom_y - prediction_box.box.top_y) - padded_height_mm = box_height_pxl * geom.pixel_in_mm * (1.0 + 2.0 * y_padding_fraction_each_side) - n_y = max(1, int(ceil(padded_height_mm / grid_size_mm.y))) - logger.info( - "Computed second auto-center raster y size from ML box height", - extra=merge_log_context( - sample_log_context(self.sample), - { - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "box_height_pxl": box_height_pxl, - "pixel_in_mm": geom.pixel_in_mm, - "grid_size_y_mm": grid_size_mm.y, - "padding_fraction_each_side": y_padding_fraction_each_side, - "padded_height_mm": padded_height_mm, - "computed_n_y": n_y, - }, - ), - ) - else: - logger.info( - "Using default y size for second auto-center raster", - extra=merge_log_context( - sample_log_context(self.sample), - { - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "default_n_y": default_n_y, - "has_prediction_box": prediction_box is not None and prediction_box.box is not None, - }, - ), - ) - - offset = Coordinate( - x=-grid_size_mm.x / 2.0, - y=-(n_y - 1) * grid_size_mm.y / 2.0, - ) - top_left = SmargonCoordinate( - sh_mm=line_scan_centre + geom.smargon_nudge(offset), - phi_deg=geom.smargon.phi_deg, - chi_deg=geom.smargon.chi_deg, - ) - - logger.info( - "Prepared second auto-center raster top-left from helper", - extra=merge_log_context( - sample_log_context(self.sample), - { - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "n_y": n_y, - "offset_x_mm": offset.x, - "offset_y_mm": offset.y, - "offset_z_mm": offset.z, - "top_left_x_mm": getattr(top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(top_left.sh_mm, "z", None), - }, - ), - ) - - return top_left, n_y - @log_timing(logger, "Auto loop center") def auto_loop_center(self, sample:Optional[SampleShortInfo]=None) -> float: """ @@ -2816,12 +2128,6 @@ class AareDAQ: f"omega:{omega_value:.2f}" ) - def _get_inference_image(self) -> np.ndarray: - image = self.__mlbox.get_latest_image() - if image is None: - raise RuntimeError("No inference image available from aarelc-infer") - return image - def save_screenshot(self, filename: str, settle_time_s: float = 0.2) -> None: """ Capture a screenshot and write it locally. @@ -2830,10 +2136,7 @@ class AareDAQ: filename: Name to give to the uploaded image. settle_time_s: float time to wait before taking the screenshot. """ - time.sleep(settle_time_s) # Wait 200 ms to ensure camera image is stable - bgr_image = self._get_inference_image() - logger.debug(f"saving screenshot {filename} from inference image") - cv2.imwrite(f"/sls/mx/applications/logs/{filename}.jpg", bgr_image) + self._screenshot_service.save_local(filename, settle_time_s) def save_screenshot_db(self, sample_id: int, filename: str, settle_time_s: float = 0.2): """ @@ -2844,43 +2147,17 @@ class AareDAQ: filename: Name to give to the uploaded image. settle_time_s: float time to wait before taking the screenshot. """ - time.sleep(settle_time_s) # Wait 200 ms to ensure camera image is stable - def _upload() -> None: - bgr_image = self._get_inference_image() - self.__aare.upload_image(sample_id, filename, bgr_image) - - self._run_noncritical( - _upload, - description=f"screenshot upload '{filename}'", - sample=self.sample, - ) + self._screenshot_service.save_to_db(sample_id, filename, settle_time_s) def send_screenshot_db(self, filename: str | None = None, message: str | None = None) -> None: sample = self.sample - if sample is None or sample.db_id is None or sample.db_id < 0: + if sample is None: raise ValueError("No sample with a valid sample_id is mounted.") - sample_id = sample.db_id - bgr_image = self._get_inference_image() - - if filename: - filename = clean_filename(filename) - pgroup = self.__cfg.pgroup - if not pgroup: - raise ValueError("No active pgroup set; cannot save screenshot to photos directory.") - - photos_dir = Path("/sls/mx/data") / pgroup / "raw" / "photos" - photos_dir = photos_dir / str(sample_id) - photos_dir.mkdir(parents=True, exist_ok=True) - photo_path = photos_dir / f"{filename}.jpeg" - cv2.imwrite(str(photo_path), bgr_image) - - upload_name = filename or f"{sample_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" - final_message = (message or "").strip() or self._default_screenshot_message(sample_id) - self._run_noncritical( - lambda: self.__aare.upload_image(sample_id, upload_name, bgr_image, message=final_message), - description=f"send screenshot '{upload_name}'", - sample=sample, + self._screenshot_service.send_to_db( + filename=filename, + message=message, + default_message=self._default_screenshot_message(sample.db_id), ) @property