import copy import secrets import time from datetime import datetime from math import ceil, floor 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.daq import workflows 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.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, sample_log_context, ) from aare.common.models import ( SampleShortInfo, PuckLoadedInfo, SampleShortInfoList, AutofocusSettings, DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, ZoomModeEnum, SimpleScanParameters, FluorescenceSpectrumParameterModel, FluorescenceSpectrumOutputModel, DAQOperation) from aare.common.automation_models import ( AutomationProgress, StepState, StepStatus, WorkflowStateKind, ) from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem, grid_to_image_id from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan from aare.common.sample_geometry import SampleGeometryModel from aare.daq.operations.face_detection import FaceDetectionContext, FaceDetectionService from aare.daq.operations.loop_centering import LoopCenteringService, LoopCenteringContext from aare.daq.operations.loop_centering.models import LoopCenteringSettings from aare.devices.area_detector import AutoEnum from aare.devices.jfjoch import JFJochWrapper from aare.devices.mx_lib import clean_filename from aare.common.exception_handler import ( TransformationInvalidException, LoopCenteringFailed, MountingFailed, WarningTellException, CriticalTellException, AXCFailed, SmargonCommunicationError, TellCommunicationError, JFJochCommunicationError, AerotechCommunicationError, MagnetPositionSensorErorr, UnmountingFailed, DataCollectionException, RasterScanException, TellMountFailedException ) from aare.devices.tell_client import TellEventValueEnum logger = setup_logger("aareDAQ") #TODO tidy up DAQ - migrate functions itno different scripts, to reduce size? class AareDAQ: """ Main Data Acquisition class for the Aare system. This class orchestrates interactions between various beamline devices, the Aare database, ML services, and data collection workflows. """ MIN_SPOTS_LOW_RES_THRESHOLD = 10.0 def __init__(self, cfg: BeamlineConfig, bl: MXBeamline): self.last_time = 0.0 self.__cfg = cfg self.__devs = BeamlineDevices(bl) self.__mlbox = MlBox(bl) self.__jfjoch = JFJochWrapper(bl) self.__bl = bl.value.upper() self.__aare = AareWrapper(bl) self.__saved_box = None self._smargon_trace_path = Path("/sls/mx/applications/logs") / "smargon_trace.csv" self._face_detection_progress_cb: Callable[[dict], None] | None = None self._automation_progress_cb: Callable[[AutomationProgress], None] | None = None self._last_sample_sync_ts = 0.0 self._sample_sync_min_interval_s = 2.0 def _is_hardware_failure(self, error: Exception) -> bool: return isinstance( error, ( CriticalTellException, TellCommunicationError, SmargonCommunicationError, AerotechCommunicationError, JFJochCommunicationError, MountingFailed, UnmountingFailed, MagnetPositionSensorErorr, TransformationInvalidException, DataCollectionException, RasterScanException, ), ) def _run_noncritical( self, action: Callable[[], object], *, description: str, sample: SampleShortInfo | None = None, ) -> object | None: try: return action() except Exception as e: logger.warning( f"Non-critical failure during {description}: {e}", extra={ "sample_name": getattr(sample, "sample_name", None), "sample_id": getattr(sample, "db_id", None), }, exc_info=True, ) return None def _run_critical( self, action: Callable[[], object], *, description: str, sample: SampleShortInfo | None = None, ) -> object: try: return action() except Exception as e: if isinstance(e, WarningTellException): logger.warning( f"Non-critical warning during {description}: {e}", extra={ "sample_name": getattr(sample, "sample_name", None), "sample_id": getattr(sample, "db_id", None), }, exc_info=True, ) return None raise def set_face_detection_progress_callback(self, cb: Callable[[dict], None] | None) -> None: self._face_detection_progress_cb = cb def set_automation_progress_callback( self, cb: Callable[[AutomationProgress], None] | None, ) -> None: self._automation_progress_cb = cb def _emit_face_detection_progress(self, payload: dict) -> None: if self._face_detection_progress_cb is None: return try: self._face_detection_progress_cb(payload) except Exception as e: logger.warning(f"Failed to emit face detection progress: {e}") def _emit_automation_progress(self, progress: AutomationProgress) -> None: if self._automation_progress_cb is None: return try: self._automation_progress_cb(progress) except Exception as e: logger.warning(f"Failed to emit automation progress: {e}") def _new_automation_progress(self) -> AutomationProgress: return AutomationProgress( current_step=None, steps=[ StepState(step=WorkflowStateKind.MOUNT, status=StepStatus.PENDING), StepState(step=WorkflowStateKind.LOOP_CENTRE, status=StepStatus.PENDING), StepState(step=WorkflowStateKind.RASTER, status=StepStatus.PENDING), StepState(step=WorkflowStateKind.DATA_COLLECTION, status=StepStatus.PENDING), StepState(step=WorkflowStateKind.FINAL, status=StepStatus.PENDING), ], finished=False, success=None, ) @staticmethod def _step_display_name(step: WorkflowStateKind) -> str: labels = { WorkflowStateKind.MOUNT: "Mount", WorkflowStateKind.LOOP_CENTRE: "Center", WorkflowStateKind.RASTER: "Raster", WorkflowStateKind.DATA_COLLECTION: "Collect", WorkflowStateKind.FINAL: "Paused/Finished", } return labels.get(step, str(step.value)) def _set_progress_step( self, progress: AutomationProgress, step: WorkflowStateKind, status: StepStatus, message: str = "", *, make_current: bool = False, ) -> None: for item in progress.steps: if item.step == step: item.status = status item.message = message break if make_current: progress.current_step = self._step_display_name(step) self._emit_automation_progress(progress) def _mark_progress_running( self, progress: AutomationProgress, step: WorkflowStateKind, message: str = "", ) -> None: self._set_progress_step( progress, step, StepStatus.RUNNING, message, make_current=True, ) def _mark_progress_success( self, progress: AutomationProgress, step: WorkflowStateKind, message: str = "", ) -> None: self._set_progress_step(progress, step, StepStatus.SUCCESS, message) def _mark_progress_failed( self, progress: AutomationProgress, step: WorkflowStateKind, message: str = "", ) -> None: self._set_progress_step( progress, step, StepStatus.FAILED, message, make_current=True, ) self._set_progress_step(progress, WorkflowStateKind.FINAL, StepStatus.FAILED, message) progress.current_step = self._step_display_name(WorkflowStateKind.FINAL) progress.finished = True progress.success = False self._emit_automation_progress(progress) def _mark_progress_finished( self, progress: AutomationProgress, success: bool, message: str = "", ) -> None: final_status = StepStatus.SUCCESS if success else StepStatus.FAILED self._set_progress_step(progress, WorkflowStateKind.FINAL, final_status, message) progress.current_step = self._step_display_name(WorkflowStateKind.FINAL) progress.finished = True progress.success = success self._emit_automation_progress(progress) #-------------------------------------------- # Operation Services #-------------------------------------------- def _create_loop_centering_settings(self) -> LoopCenteringSettings: return LoopCenteringSettings() def _create_loop_centering_service(self) -> LoopCenteringService: settings = self._create_loop_centering_settings() return LoopCenteringService( context=LoopCenteringContext( cfg=self.__cfg, 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, ), ), logger=logger, ) def _create_face_detection_service(self) -> FaceDetectionService: return FaceDetectionService( context=FaceDetectionContext( cfg=self.__cfg, devs=self.__devs, mlbox=self.__mlbox, sample_geometry_provider=lambda: self.sample_geometry, emit_progress=self._emit_face_detection_progress, ), logger=logger, ) # # def _create_mounting_service(self): # return MountingService( # context=MountingContext( # cfg=self.__cfg, # devs=self.__devs, # mount_position=ABR_POS_MOUNT, # ), # logger=logger, # ) # # def _create_raster_service(self): # return RasterService( # context=RasterContext( # cfg=self.__cfg, # devs=self.__devs, # 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, # ), # logger=logger, # ) #-------------------------------------------- # Operation Handlers #-------------------------------------------- #TODO make sure this is implemented currectly def _handle_operation_error( self, operation: DAQOperation, sample: Optional[SampleShortInfo], error: Exception, event_type: SampleEventType = SampleEventType.FAILED, additional_comment: Optional[str] = None, ) -> None: """ Centralized databse maessage error handling for all operations. Args: operation: Name of the operation that failed (e.g., "mount", "alc", "raster") sample: The sample being processed error: The exception that occurred event_type: Type of event to send to database - "mount", "alc", "axc", or "generic" """ sample_name = getattr(sample, "sample_name", None) sample_id = getattr(sample, "db_id", None) logger.exception( "Operation failed", extra={ "operation": operation.value, "event_type": str(event_type), "sample_name": sample_name, "sample_id": sample_id, "additional_comment": additional_comment, }, ) if sample is None: logger.error(f"Error in {operation.value}: {error}") return comment = f"Error in {operation.value}: {error}" if additional_comment is not None: comment += f" Additional comment: {additional_comment}" try: self.__aare.send_sample_event( sample, event_type, comment=comment, ) except Exception as db_error: logger.exception( "Failed to report operation error to database", extra={ "operation": operation.value, "sample_name": sample_name, "sample_id": sample_id, }, ) #todo make sure these functions are correctly implemented! #Operation handlers should handle database communication and beamline state changes, #where possible/between operations. For example in raster, we may use XtalSnapshot to take a screenshot, then #return to Datacollection. #Busy stats are handled by the public function call i.e. sample() or by automation i.e. measure() def _execute_mount_and_prepare(self, sample: SampleShortInfo) -> bool: """ Operation handler for executing mounting and take screenshot. Returns: True if successful, False otherwise """ try: previous_sample = self.sample if previous_sample is None or previous_sample.db_id is None: try: previous_sample = self.sync_current_sample_from_tell( force=True, clear_cached_on_empty=False, ) except Exception as sync_error: logger.warning(f"Failed to reconcile previous sample from TELL before mount: {sync_error}") if previous_sample is not None and previous_sample.db_id is not None: self.__aare.send_sample_event(previous_sample, SampleEventType.UNMOUNTING) self.__set_state(BeamlineStateEnum.RobotSampleExchange) if sample is not None and sample.db_id is not None: self.__aare.send_sample_event(sample, SampleEventType.MOUNTING) self.__mount(sample) if previous_sample is not None and previous_sample.db_id is not None: self.__aare.send_sample_event(previous_sample, SampleEventType.UNMOUNTED) if sample is not None and sample.db_id is not None: self.__aare.send_sample_event(sample, SampleEventType.MOUNTED) self.save_screenshot_db(sample.db_id, f"{sample.db_id}_mounted") self.__set_state(BeamlineStateEnum.SampleAlignment) return True except Exception as e: logger.error(f"Mount failed: {e}") try: self.__set_state(BeamlineStateEnum.SampleAlignment) except Exception: logger.exception("Failed to set state to SampleAlignment") finally: self._handle_operation_error( operation=DAQOperation.MOUNT, sample=sample, error=e, event_type=SampleEventType.MOUNTFAILED ) return False def _execute_loop_centering(self, sample: SampleShortInfo) -> bool: """ Execute loop centering sequence. Returns: True if successful, False otherwise """ result = None additional_comment = None if sample is None or sample.db_id is None or sample.db_id < 0: logger.error("Loop centering failed: no sample provided") return False self.__set_state(BeamlineStateEnum.SampleAlignment) self.__devs.lamp_light = 2.5 try: self.__aare.send_sample_event(sample, SampleEventType.CENTERING) service = self._create_loop_centering_service() result = service.run(sample_id=sample.db_id) if not result.success: self._handle_operation_error( operation=DAQOperation.LOOP_CENTERING, sample=sample, error=result.error or Exception("Loop centering failed"), event_type=SampleEventType.ALCFAILED, additional_comment=result.comment if result.comment is not None else "", ) return False self.save_screenshot_db(sample.db_id, "loop_centering") self.__aare.send_sample_event(sample, SampleEventType.CENTERED) return True except Exception as e: logger.error(f"Loop centering failed: {e}") if result is not None and result.error is not None: additional_comment = result.comment if result.comment is not None else "" self._handle_operation_error( operation=DAQOperation.LOOP_CENTERING, sample=sample, error=e, event_type=SampleEventType.ALCFAILED, additional_comment=additional_comment, ) return False def _execute_face_detection( self, *, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3, report_error: bool = True, ): """ Execute face detection sequence through the face detection service. Returns: FaceDetectionResult """ self.__set_state(BeamlineStateEnum.SampleAlignment) result = self._create_face_detection_service().run( steps=steps, step_size=step_size, face_min_ratio=face_min_ratio, ) if not result.success and report_error: self._handle_operation_error( operation=DAQOperation.FACE_CENTERING, sample=self.sample, error=result.error or Exception("Face detection failed"), additional_comment=result.comment, ) return result def _execute_raster_sequence(self, grid_request: RasterGridRequest, auto_center: bool = False) -> CompletedRasterGrid | None: """ Execute raster scan with optional auto-centering. Args: grid_request: Raster grid parameters auto_center: If True, use auto-centering; if False, use direct raster Returns: CompletedRasterGrid result or None if failed """ try: logger.info( "Starting raster sequence", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(grid_request), {"auto_center": auto_center}, ), ) 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) if result is None: logger.error( "Raster sequence returned no result after auto-centering", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(grid_request), ), ) else: if not self.__cfg.simulated_detector: logger.info(f"initialise detector for raster") status = self.status self.__jfjoch.measure_raster(grid_request, status) logger.info("detector initialised") else: logger.info("Simulated detector mode enabled; using fake raster result.") self.__setup_datacollection(request=grid_request) self.__set_state(BeamlineStateEnum.DataCollection) raster_result = self.__raster(grid_request) result = CompletedRasterGrid(r=[raster_result]) self.__set_state(BeamlineStateEnum.SampleAlignment) if result is not None: logger.info( "Raster sequence completed", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(grid_request), { "auto_center": auto_center, "result_count": len(result.r) if hasattr(result, "r") and result.r is not None else None, }, ), ) else: if self.sample is not None and self.sample.db_id is not None: self.__aare.send_sample_event(self.sample, SampleEventType.RASTERINGFAILED) logger.error( "Raster sequence returned no result", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(grid_request), ), ) return result except JFJochCommunicationError as e: logger.exception( "Raster sequence failed due to JFJoch communication error", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(grid_request), ), ) self._handle_operation_error( operation=DAQOperation.RASTER, sample=self.sample, error=e, event_type=SampleEventType.RASTERINGFAILED, additional_comment=f"JFJoch communication error: {e}" ) return None except Exception as e: logger.exception( "Raster sequence failed", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(grid_request), {"auto_center": auto_center}, ), ) self._handle_operation_error( operation=DAQOperation.RASTER, sample=self.sample, error=e, event_type=SampleEventType.RASTERINGFAILED ) return None def _execute_rotation_sequence(self, rotation_request: RotationScanRequest) -> CompletedRotationScan | None: """ Execute rotation scan. Args: rotation_request: Rotation scan parameters Returns: CompletedRotationScan result or None if failed """ try: status= self.status if self.__cfg.simulated_detector: logger.info("Simulated detector mode enabled; skipping JFJoch start.") else: self.__jfjoch.measure_rotation(rotation_request, status, self.__cfg.xrf) self.__setup_datacollection(request=rotation_request) if self.sample is not None and self.sample.db_id is not None: self.__aare.send_sample_event(self.sample, SampleEventType.COLLECTING) self.__set_state(BeamlineStateEnum.DataCollection) result = self.__rotation(rotation_request) self.__set_state(BeamlineStateEnum.SampleAlignment) if self.sample is not None and self.sample.db_id is not None: self.save_screenshot_db(self.sample.db_id, "scan_preview") self.__aare.send_sample_event(self.sample, SampleEventType.COLLECTED) self.__aare.ingest_scan(sample=self.sample, result=result.result, geom=self.sample_geometry, beam_mark_pxl=self.__cfg.get_beam_mark(self.zoom)) return result except JFJochCommunicationError as e: logger.error(f"Rotation sequence failed due to JFJoch Communication error: {e}") self._handle_operation_error( operation=DAQOperation.ROTATION, sample=self.sample, error=e, event_type=SampleEventType.COLLECTIONFAILED, additional_comment=f"JFJoch communication error: {e}" ) return None except Exception as e: logger.error(f"Rotation sequence failed: {e}") self._handle_operation_error( operation=DAQOperation.ROTATION, sample=self.sample, error=e, event_type=SampleEventType.COLLECTIONFAILED ) return None def _append_smargon_trace(self, *, sample_id: int | None, event: str) -> None: try: path = self._smargon_trace_path path.parent.mkdir(parents=True, exist_ok=True) is_new_file = not path.exists() or path.stat().st_size == 0 pos = self.smargon sh = pos.sh_mm with path.open("a", encoding="utf-8", buffering=1) as f: if is_new_file: f.write( "timestamp,event,sample_id,omega_deg,zoom," "shx_mm,shy_mm,shz_mm,phi_deg,chi_deg\n" ) f.write( f"{datetime.now().isoformat(timespec='milliseconds')}," f"{event}," f"{'' if sample_id is None else sample_id}," f"{self.omega:.3f}," f"{self.zoom:.3f}," f"{sh.x:.5f}," f"{sh.y:.5f}," f"{sh.z:.5f}," f"{pos.phi_deg:.5f}," f"{pos.chi_deg:.5f}\n" ) f.flush() except Exception as e: logger.warning(f"Failed to append smargon trace: {e}") def _sample_matches_mounted_address( self, sample: SampleShortInfo | None, mounted_address, ) -> bool: if sample is None or sample.location is None or mounted_address is None: return False return ( sample.location.segment == mounted_address.puck.segment and sample.location.pos == mounted_address.puck.pos and sample.pin == mounted_address.pin ) def _find_sample_by_mounted_address(self, mounted_address) -> SampleShortInfo | None: for sample in self.__cfg.spreadsheet.s: if self._sample_matches_mounted_address(sample, mounted_address): return sample for sample in self.__cfg.reference_tools.s: if self._sample_matches_mounted_address(sample, mounted_address): return sample return None def _placeholder_sample_from_mounted_address(self, mounted_address) -> SampleShortInfo: return SampleShortInfo( db_id=-1, puck_name="", dewar_name="", sample_name=f"Mounted sample {mounted_address.puck.segment}{mounted_address.puck.pos}-{mounted_address.pin}", run_number=0, user="", pin=mounted_address.pin, location=mounted_address.puck, ) def _find_sample_by_mounted_address(self, mounted_address) -> SampleShortInfo | None: for sample in self.__cfg.spreadsheet.s: if self._sample_matches_mounted_address(sample, mounted_address): return sample for sample in self.__cfg.reference_tools.s: if self._sample_matches_mounted_address(sample, mounted_address): return sample return None def _placeholder_sample_from_mounted_address(self, mounted_address) -> SampleShortInfo: return SampleShortInfo( db_id=-1, puck_name="", dewar_name="", sample_name=f"Mounted sample {mounted_address.puck.segment}{mounted_address.puck.pos}-{mounted_address.pin}", run_number=0, user="", pin=mounted_address.pin, location=mounted_address.puck, ) def sync_current_sample_from_tell( self, force: bool = False, clear_cached_on_empty: bool = True, ) -> SampleShortInfo | None: current_sample = self.__cfg.current_sample if current_sample is not None and current_sample.location is None: return current_sample now = time.monotonic() if not force and (now - self._last_sample_sync_ts) < self._sample_sync_min_interval_s: return current_sample self._last_sample_sync_ts = now mounted_address = self.__devs.tell.get_mounted_sample() if mounted_address is None: if current_sample is not None and current_sample.location is not None: if clear_cached_on_empty: logger.warning("TELL reports no mounted sample; clearing cached current_sample") self.__cfg.current_sample = None else: logger.warning( "TELL reports no mounted sample; keeping cached current_sample to avoid losing context" ) return self.__cfg.current_sample if self._sample_matches_mounted_address(current_sample, mounted_address): return current_sample resolved_sample = self._find_sample_by_mounted_address(mounted_address) if resolved_sample is None: resolved_sample = self._placeholder_sample_from_mounted_address(mounted_address) logger.warning( "Mounted sample from TELL was not found in known sample lists; using placeholder", extra={"mounted_address": str(mounted_address)}, ) else: logger.info( f"Reconciled cached sample from TELL to {resolved_sample.sample_name}", extra={"db_id": resolved_sample.db_id}, ) self.__cfg.current_sample = resolved_sample return resolved_sample @property def state(self) -> BeamlineStateEnum: return self.__cfg.state @property def busy(self) -> bool: return self.__cfg.state_busy @state.setter def state(self, target: BeamlineStateEnum): if target == BeamlineStateEnum.Moving: logger.error(f"Cannot explicitly move to busy state", extra={"target":target, "state":self.__cfg.state}) raise Exception("Cannot explicitly move to busy state") start = time.perf_counter() self.__cfg.try_set_busy(timeout=300) self.__set_state(target) self.__cfg.state_busy = False end = time.perf_counter() self.last_time = end - start def spreadsheet_params(self) -> tuple[Optional[SimpleScanParameters], str|None]: file_prefix = None if self.status.sample is None: return None, file_prefix logger.debug(f"generate params for sample: {self.status.sample}") aaredb_params = self.status.sample.aaredb_params if hasattr(self.status.sample, "aaredb_params") else None if aaredb_params is None: return None, file_prefix #if aaredb_params.directory: # file_prefix = aaredb_params.directory if ( getattr(aaredb_params, 'exposure', None) is None and getattr(aaredb_params, 'transmission', None) is None and getattr(aaredb_params, 'oscillation', None) is None and getattr(aaredb_params, 'totalrange', None) is None and getattr(aaredb_params, 'targetresolution', None) is None ): return None, file_prefix params = SimpleScanParameters() if (exp := getattr(aaredb_params, 'exposure', None)) is not None: params.exp_time_s = exp if (trans := getattr(aaredb_params, 'transmission', None)) is not None: logger.debug(f"transmission: {trans}") params.transmission = trans / 100.0 if trans > 1.0 else trans if (res := getattr(aaredb_params, 'targetresolution', None)) is not None: logger.debug(f"resolution: {res}") logger.debug(f"requested dtz: {self.diffraction_geometry.calc_dtz_mm(res)} ") new_res = 1 / ((1 / res) + 0.1) corrected_dtz = self.diffraction_geometry.calc_dtz_mm(new_res) logger.debug(f"corrected dtz: {corrected_dtz}") if corrected_dtz < 108: corrected_dtz = 108 params.dtz = round(corrected_dtz) osc = getattr(aaredb_params, 'oscillation', None) total = getattr(aaredb_params, 'totalrange', None) if osc is not None: osc = abs(osc) if osc > 0: params.incr_omega_deg = osc if total is not None and abs(total) > 0: params.steps = round(abs(total) / osc) else: params.steps = round(360.0 / osc) elif total is not None and abs(total) > 0: default_osc = SimpleScanParameters().incr_omega_deg params.incr_omega_deg = default_osc params.steps = round(abs(total) / default_osc) #if file_prefix is not None: # params.file_prefix = file_prefix return params, file_prefix @property def omega(self) -> float: return self.__devs.aerotech_omega def __omega(self, val: float): self.__saved_box = None logger.info(f"Set omega to {val}") if -2000 < val < 2000: try: self.__devs.aerotech_omega = val self.__cfg.state_busy = False except Exception as e: logger.error(f"Omega error: {e}") self.__cfg.state_busy = False logger.error("Omega move timed out") else: self.__cfg.state_busy = False logger.error("Omega has to be between -2000 and 2000 degrees") raise ValueError("Omega has to be between -2000 and 2000 degrees (for now)") @omega.setter def omega(self, val: float): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) self.__omega(val) def omega_rel(self, val: float): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) curr_omega = self.__devs.aerotech_omega self.__omega(curr_omega + val) @property def zoom(self) -> float: return self.__devs.zoom @zoom.setter def zoom(self, val: float): self.__saved_box = None self.__devs.samcam_auto(AutoEnum.AUTO) self.__devs.zoom = val time.sleep(0.2) self.__devs.samcam_auto(AutoEnum.ONCE) @property def front_light(self) -> float: val = self.__devs.lamp_light if val <= 1.0: return 0 elif val >= 2.5: return 100.0 else: return (val - 1.0) / 1.5 * 100.0 @front_light.setter def front_light(self, f: float): conv = (f / 100.0 * 1.5) + 1.0 logger.info(f"Light {f} -> {conv}") self.__devs.lamp_light = conv @property def back_light(self) -> float: val = self.__devs.back_light if val <= 0.0: return 0.0 elif val <= 3.0: return val / 3.0 * 100.0 else: return 100.0 @back_light.setter def back_light(self, f: float): conv = f / 100.0 * 3.0 logger.info(f"Backlight {f} -> {conv}") self.__devs.back_light = conv @property def sample(self) -> SampleShortInfo | None: return self.__cfg.current_sample @property def samcam_settings(self) -> SampleCameraSettings: return self.__devs.samcam_settings @samcam_settings.setter def samcam_settings(self, s: SampleCameraSettings): self.__devs.samcam_settings = s def tweak_abr_meas_pos(self, c: AerotechCoordinate): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: new_meas_pos = AerotechCoordinate(at_mm=self.__cfg.abr_meas_pos.at_mm + c.at_mm) self.__cfg.abr_meas_pos = new_meas_pos self.__devs.aerotech_pos = new_meas_pos self.__saved_box = None self.__cfg.state_busy = False except Exception: self.__cfg.state_busy = False raise def save_abr_meas_pos(self): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: self.__cfg.abr_meas_pos = AerotechCoordinate(at_mm=self.__devs.aerotech_pos.at_mm, omega_deg=0.0) self.__cfg.state_busy = False except Exception: self.__cfg.state_busy = False raise def goto_abr_meas_pos(self): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: self.__devs.aerotech_pos = self.__cfg.abr_meas_pos self.__cfg.state_busy = False except Exception: self.__cfg.state_busy = False raise def create_sample(self, target: SampleShortInfo): self.__cfg.try_set_busy(timeout=360) try: curr_sample = self.__cfg.current_sample if curr_sample is not None and curr_sample.location is not None: raise Exception("Sample from TELL is loaded") self.__aare.create_manual_sample(target) self.__cfg.current_sample = target self.__cfg.state_busy = False except Exception: self.__cfg.state_busy = False raise def park_and_dry(self): self.__cfg.try_set_busy(timeout=360) try: self.__devs.tell.check_enable_motion() self.__devs.tell.wait_not_busy() self.__devs.tell.set_in_mount_position(True) self.__devs.tell.unmount(wait=True, timeout=60.0) self.__devs.tell.dry(wait_cold=-1, wait=True, timeout=360.0) self.__cfg.current_sample = None self.__cfg.state_busy = False except Exception as e: self.__cfg.state_busy = False logger.error(f"Failed to park and dry: {e}") raise def blower_control(self): try: self.__devs.tell.toggle_blower() except Exception as e: logger.error(f"Failed to turn off blower: {e}") def __magnet_position_sensor_check(self, timeout: float = 1.0, repeat: bool = True): #TODO check this works, add beamstop z controls and test. # logger.info("checking beamstop") # if self.__devs.beamstop_z.value < 24.0: # raise Exception("Beamstop Z below 24.0 mm - potentially unsafe with mounting") if self.__devs.magnet_position_sensor.value != 0: logger.warning("Goniometer is not in position based on magnet position sensor readout") for i in range(round(timeout * 10.0)): if self.__devs.magnet_position_sensor.value == 0: return time.sleep(0.1) logger.error("Goniometer didn't reach position based on magnet position sensor readout") raise Exception("Goniometer is not in position based on magnet position sensor readout") def __mount_fail_handler( self, mount_error: bool = False, dry_after_fail_count: int = 3, end_automation_after_fail_count: int = 5 ): if mount_error: count = self.__cfg.record_mount_failure() if count > end_automation_after_fail_count: logger.error(f"Mounting continued to fail after stopping automation, contact your local contact") logger.debug(f"Clearing mount fail count") self.__cfg.record_mount_success() raise CriticalTellException(f"Mounting continued to fail after stopping automation, contact your local contact.") elif count == end_automation_after_fail_count: logger.error(f"Mounting failed {count} times, stopping automation") raise CriticalTellException(f"Mount failed {count} times in a row, stopping automation.") elif count == dry_after_fail_count: logger.warning(f"Mount failed {count} times in a row; drying") try: self.park_and_dry() except Exception as e: logger.exception(f"Failed to dry after mount failure: {e}") else: return else: logger.debug("Mounting succeeded, resetting mount failure counter") self.__cfg.record_mount_success() def __mount_handler(self, target : SampleShortInfo): value = self.__devs.tell.mount(address=target.tell_address(), force=True, auto_unmount=True, read_dm=False, wait=True, timeout=360.0) logger.debug(f"Mount response: {value.value}") if value == TellEventValueEnum.NO_PIN_IN_GRIPPER: raise MountingFailed("No Pin in Gripper") elif value == TellEventValueEnum.PIN_IS_LOST_GRIPPER: logger.error("Pin is lost") raise MountingFailed("Pin is lost") elif value == TellEventValueEnum.ROBOT_CLEAR_AFTER_MOUNT or value == TellEventValueEnum.DRY or value == TellEventValueEnum.COLD: logger.info(f"Robot is {value.value} - freeing beamline for user") else: logger.info(f"Mount response: {value.value} is not currently handled") return def __mount(self, target: SampleShortInfo | None): self.__devs.smargon_move_home() self.__devs.aerotech_pos = ABR_POS_MOUNT #collimator should be down!!! self.__magnet_position_sensor_check(timeout=360.0) self.__devs.tell.check_enable_motion() self.__devs.tell.wait_not_busy() self.__devs.tell.set_in_mount_position(True) previous_sample = self.__cfg.current_sample if target is None: if previous_sample is not None: logger.debug(f"Unmounting sample: {previous_sample}") try: self.__devs.tell.unmount(wait=True, timeout=60.0) self.__cfg.current_sample = target except Exception as e: raise UnmountingFailed(f"Failed to unmount: {e}") else: try: #TODO: how to differentiate between mounting and unmounting fails # during a typucal mount call as this is handled by Tell? log_msg = f"Mounting sample: {target}" if previous_sample is not None: log_msg = (f"Unmounting sample: {previous_sample} \n " f"and {log_msg}") logger.debug(log_msg) self.__mount_handler(target) self.__mount_fail_handler(mount_error=False) self.__cfg.current_sample = target except CriticalTellException as e: logger.error(f"Critical Tell error: {e}") self.__mount_fail_handler(mount_error=True) raise except Exception as e: logger.error(f"Mount failed: {e}") self.__mount_fail_handler(mount_error=True) raise def recovery_unmount_sample(self) -> None: self.__cfg.try_set_busy(timeout=360) try: self.__set_state(BeamlineStateEnum.RobotSampleExchange) self.__devs.tell.check_enable_motion() self.__devs.tell.wait_not_busy() self.__devs.tell.set_in_mount_position(True) self.__devs.tell.unmount(wait=True, timeout=360.0) self.__cfg.current_sample = None self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False except Exception: self.__cfg.state_busy = False raise @sample.setter def sample(self, target: SampleShortInfo | None): self.__cfg.try_set_busy(timeout=360) try: logger.debug(f"Mount target {target}") if target is None: logger.debug("Unmounting sample") type = "Unmount" else: logger.debug(f"Mounting sample: {target}") type = "Mount" if not self._execute_mount_and_prepare(target): raise MountingFailed(f"Failed to {type} sample {target or self.__cfg.current_sample}") logger.info(f"Sample mounted: {target}") self.__cfg.state_busy = False except Exception as e: self.__cfg.state_busy = False logger.debug(f"Failed to mount sample: {e}") raise @property def camera_image(self) -> np.ndarray | None: image = self.__devs.samcam_get_image(gray=False) return image @property def camera_image_gray(self) -> np.ndarray | None: image = self.__devs.samcam_get_image(gray=True) return image def list_loaded_pucks(self) -> List[PuckLoadedInfo]: return [] def __auto_focus(self, settings: AutofocusSettings, settle_time_s: float = 1.0) -> float: #TODO uses old code change """ Scan smargon Z and find the position with maximum focus measure. Args: settings: AutofocusSettings with center, radius, range, and steps settle_time_s: Time to wait after each move before capturing image Returns: Best Z position (mm) found during the scan """ raise NotImplementedError("Autofocus not implemented yet") def auto_focus(self, settings: AutofocusSettings) -> float: """ Public autofocus method. Only allowed in SampleAlignment state. Args: settings: AutofocusSettings with center, radius, range, and steps Returns: Best Z position (mm) found during the scan """ self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: best_z = self.__auto_focus(settings) self.__cfg.state_busy = False return best_z except Exception as e: logger.error(f"Autofocus failed: {e}") self.__cfg.state_busy = False raise 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: request_omega = getattr(request, "start_omega_deg", None) if request_omega is not None: logger.info(f"requesting omega to move to {request_omega}") self.__devs.aerotech_omega = request_omega if request.dtz is not None: logger.info(f'requesting dtz to move to {request.dtz}') self.__cfg.dtz = request.dtz if self.sample is not None and self.sample.db_id is not None: sample_id = self.sample.db_id if screening: screenshot_name = f"{sample_id}_before_screening" elif type(request) is RasterGridRequest: screenshot_name = f"{sample_id}_before_raster" else: screenshot_name = f"{sample_id}_before_data_collection" # self.auto_exposure() # time.sleep(0.2) self.save_screenshot_db(sample_id, screenshot_name) if request.transmission is not None: logger.info(f'requesting transmission to move to {request.transmission}') self.__devs.transmission = request.transmission start_pos = getattr(request, "start", None) if start_pos is not None: logger.info(f'requesting smargon to move to {start_pos}') self.__devs.smargon_pos = start_pos else: smargon_top_left = getattr(request, "smargon_top_left", None) has_valid_smargon_target = ( smargon_top_left is not None and getattr(smargon_top_left, "sh_mm", None) is not None ) is_zero_placeholder = ( has_valid_smargon_target and abs(smargon_top_left.sh_mm.x) < 1e-6 and abs(smargon_top_left.sh_mm.y) < 1e-6 and abs(smargon_top_left.sh_mm.z) < 1e-6 and abs(smargon_top_left.phi_deg) < 1e-6 and abs(smargon_top_left.chi_deg) < 1e-6 ) if has_valid_smargon_target and not is_zero_placeholder: logger.info(f'requesting smargon to move to {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, ) ) else: logger.debug("No explicit smargon target in request; skipping pre-datacollection smargon move") self.__devs.smargon_wait(timeout=180) #todo ADD TRANSMISSION #if request.transmission is not None: # self.__devs.transmission.wait() return def _build_fake_scan_result(self, *, file_prefix: str | None, image_count: int) -> ScanResult: images = [ ScanResultImagesInner( number=i, efficiency=1.0, bkg=0.0, spots=0, spots_low_res=0, spots_indexed=0, index=0, b=0.0, ) for i in range(max(1, image_count)) ] 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, ) 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 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 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, ) def __raster(self, request: RasterGridRequest, wait_for_screenshot: float | None = 0.3) -> CompletedRasterGridElem: 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)}, ), ) total_time = request.exp_time_s * request.n_x * request.n_y + request.n_y * 0.3 try: if self.sample is not None and self.sample.db_id is not None: self.__aare.create_gridscan_run(self.sample, request, status) self.__jfjoch.wait_till_running(timeout=60.0) 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 * 2, 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(360)) 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 grid_centre_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": grid_centre_offset.x, "centre_offset_y_mm": grid_centre_offset.y, "centre_offset_z_mm": grid_centre_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 {grid_centre_offset}") grid_centre_smargon = SmargonCoordinate( sh_mm=request.smargon_top_left.sh_mm + grid_centre_offset, phi_deg=request.smargon_top_left.phi_deg, chi_deg=request.smargon_top_left.chi_deg ) self.__devs.smargon_pos = grid_centre_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": grid_centre_smargon.sh_mm.x, "centre_sh_y_mm": grid_centre_smargon.sh_mm.y, "centre_sh_z_mm": grid_centre_smargon.sh_mm.z, "centre_phi_deg": grid_centre_smargon.phi_deg, "centre_chi_deg": grid_centre_smargon.chi_deg, }, ), ) if self.__cfg.simulated_detector: scan_result = self._build_fake_scan_result( file_prefix=request.file_prefix, image_count=request.n_x * request.n_y, ) com = None else: scan_result = self.__jfjoch.wait_till_done(60) 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}, ), ) 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=None, ) except Exception as e: logger.exception( "Failed during raster", 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. Args: request: RasterGridRequest parameters for the scan. auto_center: Boolean flag indicating if this is part of an automated sequence. Returns: CompletedRasterGrid result. """ logger.info( "Received raster scan request", extra=merge_log_context( sample_log_context(self.sample), raster_request_log_context(request), ), ) self.__cfg.try_set_busy(timeout=ceil(360)) try: result = self._execute_raster_sequence(request, auto_center=auto_center) if result is None: raise RasterScanException("Raster scan failed") self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False return result except Exception as e: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False raise def __rotation(self, request: RotationScanRequest) -> CompletedRotationScan: omega_start = self.omega status = self.status #self.__aare.create_rotation_run(self.sample, request, status) if request.exp_time_s < 0.004: logger.warning("Exposure time too short for PXII rotation scan") request.exp_time_s = 0.004 total_time = request.exp_time_s * request.steps if self.sample is not None and self.sample.db_id is not None: self.__aare.create_rotation_run(self.sample, request, status) try: self.__jfjoch.wait_till_running(timeout=60.0) if request.screening: self.__devs.aerotech.screening_scan( rotation_deg=request.steps*request.incr_omega_deg, wedge_deg=request.wedge_omega_deg, time_sec=total_time, steps=request.steps, run_async=True, ) else: self.__devs.aerotech.rotation_scan( rotation_deg=request.steps*request.incr_omega_deg, time_sec=total_time, start_pos_deg=request.start_omega_deg, run_async=True, ) #Is this for helical scans...? do we do smargon scans? if request.start is not None and request.end is not None: smargon_time_step = request.exp_time_s / float(request.steps) pos_step = (request.end.sh_mm - request.start.sh_mm) * (1.0 / float(request.steps)) for i in range(request.steps): self.__devs.smargon.target = SmargonCoordinate( sh_mm=request.start.sh_mm + pos_step * i ) time.sleep(smargon_time_step) self.__devs.aerotech.wait_till_done(timeout=int(round(total_time + 60,0))) self.__devs.aerotech_omega = omega_start if self.__cfg.simulated_detector: logger.warning("Detector in simulation mode, returning fake zero rotation result.") return self._build_fake_rotation_result(request) else: scan_result = self.__jfjoch.wait_till_done(60) return CompletedRotationScan( request=copy.deepcopy(request), result=scan_result, ) except JFJochCommunicationError as e: logger.error(f"Exception during rotation scan related to JFJoch: {e}") raise except Exception as e: logger.error(f"Exception during rotation scan: {e}") raise def measure_rotation(self, request: RotationScanRequest) -> CompletedRotationScan: """ Execute a rotation scan. Args: request: RotationScanRequest parameters. Returns: CompletedRotationScan result. """ total_time = request.exp_time_s * request.steps logger.info( "Received rotation scan request", extra=merge_log_context( sample_log_context(self.sample), rotation_request_log_context(request, total_time_s=total_time), ), ) self.__cfg.try_set_busy(timeout=ceil(total_time + 360)) try: result = self._execute_rotation_sequence(request) if result is None: raise DataCollectionException("Rotation scan failed, no result returned") self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False return result except Exception as e: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False raise @property def dtz(self) -> float: tmp = self.__cfg.dtz if tmp is None: return 150.0 else: return tmp @dtz.setter def dtz(self, val: float): self.__cfg.try_set_busy(timeout=360) state = self.__cfg.state if val < self.__devs.dtz_low or val > self.__devs.dtz_high: self.__cfg.state_busy = False raise RuntimeError(f"dtz={val} outside limits {self.__devs.dtz_low} to {self.__devs.dtz_high}") if state == BeamlineStateEnum.DataCollection: self.__cfg.state_busy = False raise RuntimeError("Cannot set dtz during data collection") elif state == BeamlineStateEnum.SampleAlignment: self.__devs.set_dtz(val, wait=False) self.__cfg.dtz = val self.__cfg.state_busy = False @property def smargon(self) -> SmargonCoordinate: return self.__devs.smargon_pos @smargon.setter def smargon(self, sc: SmargonCoordinate): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: self.__saved_box = None self.__devs.smargon_pos = sc self.__devs.smargon_wait() self.__cfg.state_busy = False except Exception as e: self.__cfg.state_busy = False raise pass def mark_beam(self, x_pxl: float, y_pxl: float): self.__cfg.set_busy(BeamlineStateEnum.BeamLocation) try: self.__cfg.mark_beam(x_pxl, y_pxl, self.__devs.zoom) self.__cfg.state_busy = False except Exception as e: self.__cfg.state_busy = False raise def clear_mark_beam(self): self.__cfg.set_busy(BeamlineStateEnum.BeamLocation) self.__cfg.clear_mark_beam() self.__cfg.state_busy = False @property def sample_geometry(self) -> SampleGeometryModel: zoom = self.__devs.zoom aerotech_pos_ref = self.__cfg.abr_meas_pos.at_mm aerotech_pos = self.__devs.aerotech_pos.at_mm sample_geom = SampleGeometryModel( beam_location_pxl=self.__cfg.beam_mark_coeff.apply(zoom), pixel_in_mm=self.__cfg.pixel_to_mm(zoom), omega_deg=self.__devs.aerotech_omega, smargon=self.__devs.smargon_pos, beam_size_mm=self.__cfg.beam_size_mm, aerotech=aerotech_pos - aerotech_pos_ref, aerotech_meas=aerotech_pos ) return sample_geom @property def beam_center(self) -> Tuple[float, float]: return self.__cfg.beam_center @beam_center.setter def beam_center(self, val: Tuple[float, float]): self.__cfg.beam_center = val @property def beam_size_mm(self) -> Coordinate: return self.__cfg.beam_size_mm @beam_size_mm.setter def beam_size_mm(self, val: Coordinate): self.__cfg.beam_size_mm = val def get_beam_mark(self): return self.__cfg.get_beam_mark(self.__devs.zoom) 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))) 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.02, 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. Args: sample_id: Optional sample ID. filename: Optional filename for saving/logging. Returns: RasterGridRequest object representing the found bounding box, or None if failed. """ try: self.__cfg.try_set_busy(timeout=360) r = self.__ml_bounding_box(sample_id, filename) self.__cfg.state_busy = False return r except Exception: self.__cfg.state_busy = False raise def face_detection(self, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3) -> dict: """ Perform a face detection sequence by rotating the sample and using ML to find the flat face. Args: steps: Number of rotation steps. Default is 14. step_size: Size of each rotation step in degrees. Default is 15. face_min_ratio: Minimum ratio of loopface count to loop_all count to use loop_face over loop_all. Returns: Dictionary containing face detection results, including found samples and fits. """ self.__cfg.try_set_busy(timeout=360) try: result = self._execute_face_detection( steps=steps, step_size=step_size, face_min_ratio=face_min_ratio, report_error=True, ) return result.payload 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: """ Automatically center the loop using ML-based detection. This performs a multi-step sequence including rotation and centering. Args: sample: the current SampleShortInfo. If None, the current sample (self.sample) is used. Returns: Time taken for the centering process (seconds) Raises: LoopCenteringFailed: If the centering sequence fails """ start = time.perf_counter() try: self.__cfg.try_set_busy(timeout=360) if sample is None: if self.sample is None: raise LoopCenteringFailed("No sample available for loop centering." "If you have mounted a manual sample," "please add it in the manual sample panel.") sample = self.sample if not self._execute_loop_centering(sample): raise LoopCenteringFailed self.__cfg.state_busy = False except Exception as e: self.__cfg.zoom_mode = ZoomModeEnum.User self.__cfg.state_busy = False raise finally: self.__cfg.zoom_mode = ZoomModeEnum.User end = time.perf_counter() return end - start def _default_screenshot_message(self, sample_id: int) -> str: omega_value = self.omega zoom_value = self.zoom samcam = self.samcam_settings return ( f"sample_id: {sample_id} " f"zoom: {zoom_value} " f"exp:{samcam.exposure} " f"gain:{samcam.gain} " 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. Args: 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) def save_screenshot_db(self, sample_id: int, filename: str, settle_time_s: float = 0.2): """ Capture a screenshot and upload it to the database for a specific sample. Args: sample_id: Database ID of the sample. 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, ) 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: 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, ) @property def sample_spreadsheet(self) -> SampleShortInfoList: return self.__cfg.spreadsheet @property def reference_tools(self) -> SampleShortInfoList: return self.__cfg.reference_tools def sample_spreadsheet_user(self, pgroup: str) -> SampleShortInfoList: # sample = copy.deepcopy(self.sample_spreadsheet) #switch to copy if too heavy! # sample.s = list(filter(lambda x: x.user == pgroup, sample.s)) # return sample return SampleShortInfoList(s=[x for x in self.sample_spreadsheet.s if x.user == pgroup]) def get_auto_raster_params(self) -> SimpleScanParameters: if self.status.sample is None: return SimpleScanParameters(dtz=150, exp_time_s=0.04, transmission=1.0) aaredb_params = self.status.sample.aaredb_params if hasattr(self.status.sample, "aaredb_params") else None if aaredb_params is None: return SimpleScanParameters(dtz=150, exp_time_s=0.04, transmission=1.0) params = SimpleScanParameters() # Exposure exp = getattr(aaredb_params, 'exposure', None) if exp is not None: params.exp_time_s = exp else: params.exp_time_s = 0.04 # Default # Transmission trans = getattr(aaredb_params, 'transmission', None) if trans is not None: logger.debug(f"transmission: {trans}") params.transmission = trans / 100.0 if trans > 1.0 else trans else: params.transmission = 1.0 # Default # Resolution and DTZ res = getattr(aaredb_params, 'targetresolution', None) if res is not None: logger.debug(f"resolution: {res}") try: params.dtz = self.diffraction_geometry.calc_dtz_mm(res) logger.debug(f"dtz: {params.dtz}") except Exception as e: logger.error(f"Failed to calculate dtz for resolution {res}: {e}") params.dtz = 150 # Fallback default # Adjust exposure time based on resolution if res <= 1.5: params.exp_time_s = 0.02 elif 1.5 < res <= 3.0: params.exp_time_s = 0.04 else: params.exp_time_s = 0.08 else: params.dtz = 150 params.exp_time_s = 0.04 return params def get_collection_params(self, prefer_smart: bool = False) -> tuple[SimpleScanParameters, str]: spreadsheet_params, file_prefix = self.spreadsheet_params() logger.debug(f"spreadsheet_params: {spreadsheet_params}") smart_params = self.__cfg.auto_params default_params = SimpleScanParameters(exp_time_s=0.04, dtz=110, incr_omega_deg=0.2) #if file_prefix is not None: # default_params.file_prefix = file_prefix #self.__aare.send_msg_to_db(self.sample,event_type=SampleEventType(''), comment=f'smart_params: {smart_params}') if prefer_smart: if smart_params: #if file_prefix: # smart_params.file_prefix = file_prefix return smart_params, "smart_params" if spreadsheet_params: return spreadsheet_params, "spreadsheet_params" else: if spreadsheet_params: return spreadsheet_params, "spreadsheet_params" if smart_params: #if file_prefix: # smart_params.file_prefix = file_prefix return smart_params, "smart_params" return default_params, "defaults" def _end_operation(self, start, operation:Optional[DAQOperation]=DAQOperation.AUTOMATION, error: bool = False) -> float: """ End an operation and return elapsed time. Returns: (elapsed_seconds, end_time_perf_counter) """ if operation is DAQOperation.AUTOMATION: msg = "Ended automation operation" else: msg = f"Ended automation operation after {operation.value}" if error: msg += f"with an error" self.__set_state(BeamlineStateEnum.RobotSampleExchange) else: msg += f"successfully" logger.error(f"{msg}, time taken {time.perf_counter() - start} seconds.") self.__cfg.state_busy = False end = time.perf_counter() return end - start @log_timing(logger, "Measure Sequence") def measure(self, sample: SampleShortInfo) -> float: """ Main measurement sequence for a given sample. Includes mounting, loop centering, face detection, auto-raster, and data collection. Args: sample: SampleShortInfo object containing sample details Returns: Total time taken for the measurement (seconds) """ start = time.perf_counter() progress = self._new_automation_progress() self._emit_automation_progress(progress) formatted_date = datetime.now().strftime('%Y%m%d') sample_prefix = "{}/{}/{:02d}/{}".format( formatted_date, sample.puck_name, sample.pin, sample.sample_name ) try: logger.info(f"setting busy at {time.perf_counter() - start}") self.__cfg.try_set_busy(timeout=360) logger.info(f"set busy at {time.perf_counter() - start}") self._mark_progress_running(progress, WorkflowStateKind.MOUNT, "Mounting sample") if not self._execute_mount_and_prepare(sample): self._mark_progress_failed(progress, WorkflowStateKind.MOUNT, "Mount failed") return self._end_operation(start, DAQOperation.MOUNT, error=True) self._mark_progress_success(progress, WorkflowStateKind.MOUNT, "Mount complete") self.__set_state(BeamlineStateEnum.SampleAlignment) logger.info(f"mounting done at {time.perf_counter() - start}") self._mark_progress_running(progress, WorkflowStateKind.LOOP_CENTRE, "Centering sample") if not self._execute_loop_centering(sample): self._mark_progress_failed(progress, WorkflowStateKind.LOOP_CENTRE, "Centering failed") return self._end_operation(start, DAQOperation.LOOP_CENTERING, error=True) logger.info(f"Loop Centering done at {time.perf_counter() - start}") face_detection_result = self._execute_face_detection( steps=7, step_size=30, face_min_ratio=0.3, report_error=True, ) if not face_detection_result.success: logger.warning( "Face detection failed during automation; continuing with latest payload", extra=merge_log_context( sample_log_context(sample), {"comment": face_detection_result.comment}, ), ) logger.info(f"Face Detection done at {time.perf_counter() - start}") self._mark_progress_success(progress, WorkflowStateKind.LOOP_CENTRE, "Centering complete") self._mark_progress_running(progress, WorkflowStateKind.RASTER, "Running raster") hex_string = secrets.token_hex(3) raster_params = self.get_auto_raster_params() geom = self.sample_geometry raster_grid = RasterGridRequest( exp_time_s=raster_params.exp_time_s, file_prefix=sample_prefix + f"_{hex_string}", smargon_top_left=SmargonCoordinate(), n_x=1, n_y=1, dtz=raster_params.dtz, grid_size_mm=Coordinate(x=geom.beam_size_mm.x * 0.5, y=geom.beam_size_mm.y * 0.5), omega_deg=self.omega, transmission=raster_params.transmission, ) raster_result = self._execute_raster_sequence(raster_grid, auto_center=True) if raster_result is None: logger.error( "Raster result was None during automation", extra=merge_log_context( sample_log_context(sample), raster_request_log_context(raster_grid), ), ) self._mark_progress_failed(progress, WorkflowStateKind.RASTER, "Raster failed") return self._end_operation(start, DAQOperation.RASTER, error=True) self._mark_progress_success(progress, WorkflowStateKind.RASTER, "Raster complete") self.__set_state(BeamlineStateEnum.DataCollection) logger.info(f"Raster scans completed at {time.perf_counter() - start}") self._mark_progress_running(progress, WorkflowStateKind.DATA_COLLECTION, "Collecting data") params, source = self.get_collection_params(prefer_smart=False) logger.info(f"Using {source} for data collection: {params}") self.__devs.dtz = params.dtz start_omega = self.omega rotation_request = RotationScanRequest( start_omega_deg=start_omega, dtz=params.dtz, file_prefix="data/" + sample_prefix + f"_{hex_string}", exp_time_s=params.exp_time_s, incr_omega_deg=params.incr_omega_deg, steps=params.steps, transmission=params.transmission, ) rotation_result = self._execute_rotation_sequence(rotation_request) if rotation_result is None: logger.error("Rotation result was None") self._mark_progress_failed(progress, WorkflowStateKind.DATA_COLLECTION, "Collection failed") return self._end_operation(start, DAQOperation.ROTATION, error=True) self._mark_progress_success(progress, WorkflowStateKind.DATA_COLLECTION, "Collection complete") logger.info(f"Rotation scan done at {time.perf_counter() - start}") except Exception as e: logger.error(f"Error in measure: {e}") if progress.current_step is not None: current_kind = next( ( item.step for item in progress.steps if self._step_display_name(item.step) == progress.current_step ), WorkflowStateKind.FINAL, ) self._mark_progress_failed(progress, current_kind, str(e)) else: self._mark_progress_finished(progress, False, str(e)) self._handle_operation_error( operation=DAQOperation.AUTOMATION, sample=sample, error=e, event_type=SampleEventType.FAILED) raise Exception(f"Critical Error in automation: {e}") from e self._mark_progress_finished(progress, True, "Automation complete") return self._end_operation(start, DAQOperation.AUTOMATION, error=False) def __set_state(self, target: BeamlineStateEnum): """__set_state assumes that beamline is already in busy state it will apply a proper transformation and change state afterward specifically: 1. If target state is maintenance, just go there 2. If target state is same as current, nothing will happen 3. If target state cannot be reached, exception is raised and current state is kept 4. If exception is raised during transformation, state is set to maintenance 5. If transformation goes OK, target state is set Busy state will be cleared only, if exception is raised. """ curr_state = self.__cfg.state if not self.__cfg.state_busy: raise Exception("Beamline should be busy") if target == BeamlineStateEnum.Maintenance: self.__cfg.state = BeamlineStateEnum.Maintenance elif target != curr_state: self.__cfg.state = BeamlineStateEnum.Moving try: match curr_state: case BeamlineStateEnum.Maintenance: if target == BeamlineStateEnum.SampleExchange: workflows.m2se(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.SampleExchange: if target == BeamlineStateEnum.SampleAlignment: workflows.se2sa(self.__devs, self.__cfg) elif target == BeamlineStateEnum.BeamLocation: workflows.se2sa(self.__devs, self.__cfg) workflows.sa2bl(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.DewarTransfer: if target == BeamlineStateEnum.SampleAlignment: workflows.dh2sa(self.__devs, self.__cfg) elif target == BeamlineStateEnum.RobotSampleExchange: workflows.dh2sa(self.__devs, self.__cfg) workflows.sa2rse(self.__devs, self.__cfg) elif target == BeamlineStateEnum.SampleExchange: workflows.dh2sa(self.__devs, self.__cfg) workflows.sa2se(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.DataCollection: if target == BeamlineStateEnum.SampleAlignment: workflows.dc2sa(self.__devs, self.__cfg) elif target == BeamlineStateEnum.RobotSampleExchange: workflows.dc2rse(self.__devs, self.__cfg) elif target == BeamlineStateEnum.XtalSnapshot: workflows.sa2xtal_snapshot(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.BeamLocation: if target == BeamlineStateEnum.SampleAlignment: workflows.bl2sa(self.__devs, self.__cfg) elif target == BeamlineStateEnum.SampleExchange: workflows.bl2sa(self.__devs, self.__cfg) workflows.sa2se(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.SampleAlignment: if target == BeamlineStateEnum.DewarTransfer: workflows.sa2dh(self.__devs, self.__cfg) elif target == BeamlineStateEnum.SampleExchange: workflows.sa2se(self.__devs, self.__cfg) elif target == BeamlineStateEnum.RobotSampleExchange: workflows.sa2rse(self.__devs, self.__cfg) elif target == BeamlineStateEnum.DataCollection: workflows.sa2dc(self.__devs, self.__cfg) elif target == BeamlineStateEnum.XrayFluorescence: workflows.sa2xrf(self.__devs, self.__cfg) elif target == BeamlineStateEnum.BeamLocation: workflows.sa2bl(self.__devs, self.__cfg) elif target == BeamlineStateEnum.XtalSnapshot: workflows.sa2xtal_snapshot(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.XrayFluorescence: if target == BeamlineStateEnum.SampleAlignment: workflows.xrf2sa(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.RobotSampleExchange: if target == BeamlineStateEnum.SampleAlignment: workflows.rse2sa(self.__devs, self.__cfg) else: raise TransformationInvalidException() case BeamlineStateEnum.XtalSnapshot: if target == BeamlineStateEnum.SampleAlignment: workflows.xtal_snapshot2sa(self.__devs, self.__cfg) elif target == BeamlineStateEnum.SampleExchange: workflows.xtal_snapshot2se(self.__devs, self.__cfg) elif target == BeamlineStateEnum.RobotSampleExchange: workflows.xtal_snapshot2rse(self.__devs, self.__cfg) elif target == BeamlineStateEnum.DataCollection: workflows.xtal_snapshot2dc(self.__devs, self.__cfg) elif target == BeamlineStateEnum.XrayFluorescence: workflows.xtal_snapshot2xrf(self.__devs, self.__cfg) elif target == BeamlineStateEnum.DewarTransfer: workflows.xtal_snapshot2dh(self.__devs, self.__cfg) elif target == BeamlineStateEnum.BeamLocation: workflows.xtal_snapshot2bl(self.__devs, self.__cfg) else: raise TransformationInvalidException() self.__cfg.state = target except TransformationInvalidException as e: self.__cfg.state = curr_state self.__cfg.state_busy = False raise except Exception as e: self.__cfg.state = BeamlineStateEnum.Maintenance self.__cfg.state_busy = False raise @property def shutter(self) -> bool: return self.__devs.shutter @shutter.setter def shutter(self, v: bool): self.__devs.shutter = v @property def diffraction_geometry(self) -> DiffractionGeometry: det_cfg = self.__jfjoch.detector() return DiffractionGeometry( energy_keV=self.__devs.energy_kev, dtz_mm=self.__devs.dtz, detector_size_pxl=(det_cfg.width, det_cfg.height), pixel_size_mm=det_cfg.pixel_size_mm, beam_center_pxl=self.__cfg.beam_center, detector_description=det_cfg.description, detector_serial_number=det_cfg.serial_number, poni_rot1_rad=-0.001396263, poni_rot2_rad=-0.003839724, ) @property def beamline_status(self) -> BeamlineStatus: return BeamlineStatus( ring_current_mA=self.__devs.ring_current, front_light=self.front_light, back_light=self.back_light, cryojet_K=self.__devs.cryojet_temp, shutter_open=self.__devs.shutter, exp_shutter_open=self.__devs.shutter, flux_ph_s=self.__devs.full_flux, sample_camera=self.__devs.samcam_settings, name=self.__bl, transmission=self.__devs.transmission, zoom=self.__devs.zoom, commissioning_mode=self.__cfg.commissioning_mode, dtz_min=20, dtz_max=1000, ) def _safe_sample(self) -> tuple[SampleShortInfo | None, bool, str | None]: """ Return (sample, tell_connected, tell_error) without raising. """ try: return self.sync_current_sample_from_tell(), True, None except TellCommunicationError as e: return self.__cfg.current_sample, False, str(e) except Exception as e: # Keep status flowing even if Tell code throws something unexpected return self.__cfg.current_sample, False, f"TELL unavailable: {e}" def _aerotech_status(self) -> tuple[bool, str | None]: aerotech_ok = True aerotech_err: str | None = None try: _ = self.__devs.aerotech.status() except Exception as e: aerotech_ok = False aerotech_err = f"Cannot connect to Aerotech: {e}" return aerotech_ok, aerotech_err def _safe_geom(self) -> tuple[SampleGeometryModel, bool, str | None, bool, str | None]: """ Return (geom, smargon_connected, smargon_error, aerotech_connected, aerotech_error) without raising. Uses a conservative fallback geometry if Smargon access fails. """ aerotech_connected, smargon_connected = True, True aerotech_error, smargon_error = None, None try: return self.sample_geometry, smargon_connected, smargon_error, aerotech_connected, aerotech_error except SmargonCommunicationError as e: smargon_error = f"Cannot connect to Smargon: {e}" smargon_connected = False except AerotechCommunicationError as e: aerotech_error = f"Cannot connect to Aerotech: {e}" aerotech_connected = False except Exception as e: smargon_error = f"Safe geometry failed: {e}" aerotech_error = f"Safe geometry failed: {e}" aerotech_connected = False smargon_connected = False zoom = self.__devs.zoom fallback = SampleGeometryModel( beam_location_pxl=self.__cfg.beam_mark_coeff.apply(zoom), pixel_in_mm=self.__cfg.pixel_to_mm(zoom), omega_deg=0.0, smargon=SmargonCoordinate( sh_mm=Coordinate(x=0.0, y=0.0, z=0.0), phi_deg=0.0, chi_deg=0.0, ), beam_size_mm=self.__cfg.beam_size_mm, aerotech=Coordinate(x=0.0, y=0.0, z=0.0), aerotech_meas=Coordinate(x=0.0, y=0.0, z=0.0), ) return fallback, smargon_connected, smargon_error, aerotech_connected, aerotech_error def _safe_beamline_status(self) -> BeamlineStatus: try: return self.beamline_status except Exception: return BeamlineStatus( name=self.__bl, ring_current_mA=0.0, front_light=0.0, back_light=0.0, cryojet_K=0.0, shutter_open=False, exp_shutter_open=None, flux_ph_s=0.0, sample_camera=SampleCameraSettings(gain=0.0, exposure=0.0), transmission=None, zoom=self.__devs.zoom, commissioning_mode=self.__cfg.commissioning_mode, dtz_min=20, dtz_max=1000, ) def _safe_diffraction_geometry(self) -> DiffractionGeometry: try: return self.diffraction_geometry except Exception: # Must satisfy pydantic constraints in DiffractionGeometry return DiffractionGeometry( energy_keV=12.4, dtz_mm=150.0, detector_size_pxl=(1, 1), pixel_size_mm=0.15, beam_center_pxl=(0.0, 0.0), detector_description="unavailable", detector_serial_number="unavailable", poni_rot1_rad=0.0, poni_rot2_rad=0.0, ) @property def status(self) -> DAQStatusModel: safe_sample, tell_ok, tell_err = self._safe_sample() safe_geom, smargon_ok, smargon_err, aerotech_ok, aerotech_err = self._safe_geom() return DAQStatusModel( state=self.state, busy=self.busy, geom=safe_geom, bl=self._safe_beamline_status(), sample=safe_sample, session=SessionStatus( current_pgroup=self.__cfg.pgroup, session=self.__cfg.session_state(0), # 0 is dummy session staff=False ), diffraction=self._safe_diffraction_geometry(), box=self.__saved_box, last_best_res=self.__cfg.last_best_res, last_best_b_factor=self.__cfg.last_best_b_factor, crystal_size=self.__cfg.crystal_size, tell_connected=tell_ok, tell_error=tell_err, smargon_connected=smargon_ok, smargon_error=smargon_err, aerotech_connected=aerotech_ok, aerotech_error=aerotech_err, ) def cancel(self): if self.__cfg.state == BeamlineStateEnum.DataCollection: self.__devs.aerotech.cancel() self.__jfjoch.cancel() def anneal(self, time_s: float): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: self.__devs.anneal(time_s) finally: self.__cfg.state_busy = False def fluorimeter_take_spectrum(self, fm: FluorescenceSpectrumParameterModel) -> FluorescenceSpectrumOutputModel: self.__cfg.try_set_busy(timeout=360) try: self.__set_state(BeamlineStateEnum.XrayFluorescence) if fm.transmission is not None: self.__devs.transmission = fm.transmission # TODO: Fill self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False return None except Exception as e: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__cfg.state_busy = False raise