Files
AareDAQ/src/aare/daq/daq.py
T
appleb_m a00f0b3dff
Build and Publish / test (push) Failing after 1m4s
Build and Publish / Build and Deploy Docs (push) Has been skipped
Build and Publish / build (push) Has been skipped
DAQ: updated face_detection added operations/face_detection fodler including service and models, new tests of face_detection
2026-04-30 12:44:55 +02:00

3027 lines
119 KiB
Python

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