diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 62f6b596..27c591b8 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -19,7 +19,7 @@ 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, MLBoxPredictionsResult +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 @@ -40,7 +40,7 @@ from aare.common.models import ( SampleShortInfoList, AutofocusSettings, DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, ZoomModeEnum, SimpleScanParameters, MLBoxModel, FluorescenceSpectrumParameterModel, - FluorescenceSpectrumOutputModel, DAQOperation, LoopCenteringResult, MLBoxType) + FluorescenceSpectrumOutputModel, DAQOperation) from aare.common.automation_models import ( AutomationProgress, StepState, @@ -50,7 +50,9 @@ from aare.common.automation_models import ( 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.spreadsheetupdater import beamline +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 @@ -375,6 +377,31 @@ class AareDAQ: ) return False + 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 _execute_loop_centering(self, sample: SampleShortInfo) -> bool: """ Execute loop centering sequence. @@ -383,7 +410,8 @@ class AareDAQ: True if successful, False otherwise """ result = None - additional_comment=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 @@ -393,7 +421,9 @@ class AareDAQ: try: self.__aare.send_sample_event(sample, SampleEventType.CENTERING) - result = self.__loop_center_sequence(sample.db_id) + + service = self._create_loop_centering_service() + result = service.run(sample_id=sample.db_id) if not result.success: self._handle_operation_error( @@ -401,7 +431,7 @@ class AareDAQ: 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 "" + additional_comment=result.comment if result.comment is not None else "", ) return False @@ -418,7 +448,7 @@ class AareDAQ: sample=sample, error=e, event_type=SampleEventType.ALCFAILED, - additional_comment=additional_comment + additional_comment=additional_comment, ) return False @@ -1937,91 +1967,6 @@ class AareDAQ: omega_deg=geom.omega_deg ) - def __ml_loop_centre_box( - self, - sample_id: int | None = None, - filename: str | None = None, - boxes=None, - bundle_image: np.ndarray | None = None, - ) -> tuple[SmargonCoordinate | None, int | None, list[int] | None]: - if boxes is None: - prediction_result: MLBoxPredictionsResult = self.__mlbox.predict_all_best( - overlap_with_pin=0.5, - confidence_min=0.3, - return_image=True, - return_bundle_meta=True - ) - boxes = prediction_result.predictions - bundle_image = prediction_result.image - log_ml_bundle_meta( - logger, - f"ml_loop_centre_box:{filename or 'unnamed'}", - target_point=prediction_result.target_point, - focus=prediction_result.focus, - ) - - if boxes is None: - 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, None, None - - classes: list[int] = [] - pin = None - centre_x, centre_y = None, None - for box in boxes.boxes.values(): - if box and box.cls is not None: - classes.append(int(box.cls.value)) - if int(box.cls.value) == 1: - pin = box - - best_box = self.__mlbox.get_preferred_class_box(boxes, (2, 3, 0, 1)) - if best_box is None or best_box.box is None or best_box.cls is None: - return None, None, classes if classes else None - - cls = int(best_box.cls.value) - x1 = best_box.box.top_x - y1 = best_box.box.top_y - x2 = best_box.box.bottom_x - y2 = best_box.box.bottom_y - - if filename is not None and bundle_image is not None: - self.__aare.upload_image(sample_id, filename, bundle_image) - - geom = self.sample_geometry - - if cls == 0: # loop_all - if y1 + y2 <= x1 + x2: - centre_y = y1 + (y2 - y1) / 2 - centre_x = x1 - elif pin: - position_dict = self.__mlbox.check_box_relation(pin, best_box) - if position_dict["overlap_y"] and position_dict["overlap_x"]: - centre_y = y1 + (y2 - y1) / 2 - centre_x = x1 - cls = 1 - logger.debug("significant overlap between pin and loop_all not picked up by ML") - else: - centre_y = y2 if position_dict["top"] else y1 - centre_x = x1 + (x2 - x1) / 2 - else: - centre_y = y1 - centre_x = x1 + (x2 - x1) / 2 - - elif cls == 1: # pin - centre_y = y1 + (y2 - y1) / 2 - centre_x = x1 - - elif cls == 2 or cls == 3: # crystal or loop_face - centre_y = y1 + (y2 - y1) / 2 - centre_x = x1 + (x2 - x1) / 2 - - else: - logger.debug(f"unknown box class {cls}") - return None, cls, classes if classes else None - - coord = geom.picture_to_smargon(Coordinate(x=centre_x, y=centre_y)) - return SmargonCoordinate(sh_mm=coord), cls, classes - 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. @@ -2225,43 +2170,6 @@ class AareDAQ: self._emit_face_detection_progress(result) return result - @staticmethod - def _select_smargon_target( - calculated_target: SmargonCoordinate | None, - predicted_target: SmargonCoordinate | None, - tolerance_um: float = 500.0, - ) -> SmargonCoordinate | None: - if calculated_target is None: - return predicted_target - if predicted_target is None: - return calculated_target - - tolerance_mm = tolerance_um / 1000.0 - dx = abs(calculated_target.sh_mm.x - predicted_target.sh_mm.x) - dy = abs(calculated_target.sh_mm.y - predicted_target.sh_mm.y) - dz = abs(calculated_target.sh_mm.z - predicted_target.sh_mm.z) - - if dx <= tolerance_mm and dy <= tolerance_mm and dz <= tolerance_mm: - logger.info( - "Using prediction target because it is within %.0f um of calculated target " - "(dx=%.4f mm, dy=%.4f mm, dz=%.4f mm)", - tolerance_um, - dx, - dy, - dz, - ) - return predicted_target - - logger.info( - "Keeping calculated target because prediction target is outside %.0f um " - "(dx=%.4f mm, dy=%.4f mm, dz=%.4f mm)", - tolerance_um, - dx, - dy, - dz, - ) - return calculated_target - def _auto_center_line_scan_top_left( self, *, @@ -2415,169 +2323,6 @@ class AareDAQ: ) return top_left, n_y - #TODO move wait_screenshot_sleep to save_screenshot_db!!! - @log_timing(logger, "Loop center sequence") - def __loop_center_sequence( - self, - sample_id: Optional[int] = None, - trace_all_alc_moves: bool = False, - wait_screenshot_sleep_sec: float = 0.1, - ) -> LoopCenteringResult: - - found_classes_count: dict[int, int] = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0} - - try: - for zoom_iter, zoom_value in enumerate([200]): - max_attempt = 2 - attempt = 0 - - base_angles = (0, 90) if (zoom_iter % 2 == 0) else (90, 0) - if sample_id is not None and zoom_iter == 0: - logger.info(f"submitting to db loop center sequence for sample {sample_id}, zoom={zoom_value}") - time.sleep(wait_screenshot_sleep_sec) - self.save_screenshot_db(sample_id, f"pre_alc") - - while attempt < max_attempt: - self.zoom = zoom_value - - found_flag = False - found_angle: int | None = None - targets_found_this_attempt = 0 - - for angle in base_angles: - logger.debug(f"Moving to new omega: {angle}") - time_to_move_aerotech = time.perf_counter() - self.__devs.aerotech_omega = angle - log_duration( - logger, - "Completed Aerotech move during loop centering", - time.perf_counter() - time_to_move_aerotech, - extra={"angle_deg": angle, "zoom": zoom_value}, - ) - - filename = f"{sample_id}_{angle}_{zoom_value:.0f}" if sample_id is not None else None - - try: - time_to_get_pred = time.perf_counter() - prediction_result: MLBoxPredictionsResult = self.__mlbox.predict_all_best( - overlap_with_pin=0.5, - confidence_min=None, - return_image=True, - return_bundle_meta=True - ) - log_duration( - logger, - "Completed ML prediction during loop centering", - time.perf_counter() - time_to_get_pred, - extra={"angle_deg": angle, "zoom": zoom_value}, - ) - boxes = prediction_result.predictions - pred_target_point = prediction_result.target_point - - pred_target_point_smargon = None - if pred_target_point is not None: - pred_target_point_smargon = SmargonCoordinate(sh_mm=self.sample_geometry.picture_to_smargon( - Coordinate(x=pred_target_point[0], y=pred_target_point[1]) - )) - - bundle_image = prediction_result.image - log_ml_bundle_meta( - logger, - f"loop_center_angle_{angle}_zoom_{zoom_value:.0f}", - target_point=prediction_result.target_point, - focus=prediction_result.focus, - ) - - target, cls, classes = self.__ml_loop_centre_box( - sample_id=sample_id, - filename=filename, - boxes=boxes, - bundle_image=bundle_image, - ) - if cls in [MLBoxType.LOOP_ALL.value, MLBoxType.LOOP_FACE.value, MLBoxType.CRYSTAL.value]: - target = self._select_smargon_target( - calculated_target=target, - predicted_target=pred_target_point_smargon, - tolerance_um=500.0) - logger.debug(f"calculated target: {target} compares to prediction: {pred_target_point_smargon}") - except Exception as e: - logger.error(f"Error getting ML box for angle {angle}") - logger.error(f"Exception: {e}") - target, cls, classes = None, None, None - - if target is None: - logger.debug("no target found") - continue - - if classes: - for c in classes: - found_classes_count[int(c)] = found_classes_count.get(int(c), 0) + 1 - logger.debug(f"classes found: {classes}") - logger.debug(f"class found: {cls}") - if cls is not None and cls != MLBoxType.PIN.value and cls != MLBoxType.NEEDLE.value and cls != MLBoxType.ICE.value: - targets_found_this_attempt += 1 - found_angle = angle - - time_to_move_smargon = time.perf_counter() - self.__devs.smargon_pos = target - self.__devs.smargon_wait(60) - log_duration( - logger, - "Completed Smargon move during loop centering", - time.perf_counter() - time_to_move_smargon, - extra={"angle_deg": angle, "zoom": zoom_value}, - ) - if sample_id is not None: - if trace_all_alc_moves: - self._append_smargon_trace( - sample_id=sample_id, - event=f"alc_move_zoom_{zoom_value:.0f}_angle_{angle}" - ) - self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{zoom_value:.0f}") - - logger.debug(f"found a target at angle {found_angle} in attempt {attempt + 1}") - base_angles = (0, 90) - logger.debug(f"new base angles: {base_angles}") - attempt += 1 - logger.debug(f"attempt {attempt} of {max_attempt}") - if attempt >= max_attempt: - if targets_found_this_attempt >= len(base_angles): - logger.debug( - f"sucessfully found {targets_found_this_attempt} targets in attempt {attempt + 1} ") - break - logger.error(f"{attempt} exceeds max attempts {max_attempt}") - raise LoopCenteringFailed - - #i += 1 - logger.debug("alc success") - logger.debug(f"current sample: {self.__cfg.current_sample}, sample_id of scan: {sample_id}") - if sample_id is not None: - logger.info(f"sample {sample_id} centered") - self._append_smargon_trace(sample_id=sample_id, event="alc_success") - return LoopCenteringResult(success=True) - - except Exception as e: - total_detections = sum(found_classes_count.values()) - if total_detections == 0: - alc_comment="No objects detected" - else: - alc_comment = ( - "Failed to centre but detected objects - " - f"Crystal: {found_classes_count.get(2, 0)}, " - f"Loop_face: {found_classes_count.get(3, 0)}, " - f"Loop_all: {found_classes_count.get(0, 0)}, " - f"Pin: {found_classes_count.get(1, 0)}, " - f"Ice: {found_classes_count.get(4, 0)}, " - f"Needle: {found_classes_count.get(5, 0)}" - ) - logger.error(f"ALC exception: {alc_comment}") - logger.error(traceback.format_exc()) - logger.error(f"Error in loop centering: {e}") - return LoopCenteringResult( - success=False, - comment=alc_comment, - error=e - ) @log_timing(logger, "Auto loop center") def auto_loop_center(self, sample:Optional[SampleShortInfo]=None) -> float: @@ -2638,21 +2383,29 @@ class AareDAQ: raise RuntimeError("No inference image available from aarelc-infer") return image - def save_screenshot(self, filename: str): - #time.sleep(0.2) # Wait 200 ms to ensure camera image is stable + 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): + 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(0.2) # Wait 200 ms to ensure camera image is stable + 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) diff --git a/src/aare/daq/operations/loop_centering/__init__.py b/src/aare/daq/operations/loop_centering/__init__.py new file mode 100644 index 00000000..1d82c710 --- /dev/null +++ b/src/aare/daq/operations/loop_centering/__init__.py @@ -0,0 +1,17 @@ +from aare.daq.operations.loop_centering.models import ( + LoopCenteringContext, + AngleAnalysis, + AttemptSummary, + LoopCenteringSettings, +) +from aare.daq.operations.loop_centering.service import LoopCenteringService +from aare.daq.operations.loop_centering.analyzer import LoopCenteringAnalyzer + +__all__ = [ + "LoopCenteringContext", + "AngleAnalysis", + "AttemptSummary", + "LoopCenteringSettings", + "LoopCenteringService", + "LoopCenteringAnalyzer", +] \ No newline at end of file diff --git a/src/aare/daq/operations/loop_centering/analyzer.py b/src/aare/daq/operations/loop_centering/analyzer.py new file mode 100644 index 00000000..f1848d01 --- /dev/null +++ b/src/aare/daq/operations/loop_centering/analyzer.py @@ -0,0 +1,188 @@ +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.logger_events import log_ml_bundle_meta +from aare.common.models import MLBoxType + +from aare.daq.operations.loop_centering.models import ( + AngleAnalysis, + LoopCenteringContext, +) + + +class LoopCenteringAnalyzer: + def __init__(self, *, context: LoopCenteringContext, logger): + self.ctx = context + self.logger = logger + + def is_ignore_only_classes(self, classes: list[int] | None) -> bool: + if not classes: + return False + + ignore_classes = { + MLBoxType.PIN.value, + MLBoxType.NEEDLE.value, + MLBoxType.ICE.value, + } + found = {int(cls) for cls in classes if cls is not None} + return bool(found) and found.issubset(ignore_classes) + + def has_valid_target_classes(self, classes: list[int] | None) -> bool: + if not classes: + return False + + valid_classes = { + MLBoxType.LOOP_ALL.value, + MLBoxType.LOOP_FACE.value, + MLBoxType.CRYSTAL.value, + } + return any(int(cls) in valid_classes for cls in classes if cls is not None) + + def select_smargon_target( + self, + *, + calculated_target: SmargonCoordinate | None, + predicted_target: SmargonCoordinate | None, + tolerance_um: float | None = None, + ) -> SmargonCoordinate | None: + if calculated_target is None: + return predicted_target + if predicted_target is None: + return calculated_target + + if tolerance_um is None: + tolerance_um = self.ctx.settings.target_tolerance_um + + tolerance_mm = tolerance_um / 1000.0 + dx = abs(calculated_target.sh_mm.x - predicted_target.sh_mm.x) + dy = abs(calculated_target.sh_mm.y - predicted_target.sh_mm.y) + dz = abs(calculated_target.sh_mm.z - predicted_target.sh_mm.z) + + if dx <= tolerance_mm and dy <= tolerance_mm and dz <= tolerance_mm: + self.logger.info( + "Using prediction target because it is within %.0f um of calculated target " + "(dx=%.4f mm, dy=%.4f mm, dz=%.4f mm)", + tolerance_um, + dx, + dy, + dz, + ) + return predicted_target + + self.logger.info( + "Keeping calculated target because prediction target is outside %.0f um " + "(dx=%.4f mm, dy=%.4f mm, dz=%.4f mm)", + tolerance_um, + dx, + dy, + dz, + ) + return calculated_target + + def _extract_prediction_target( + self, + *, + target_point: tuple[float, float] | None, + ) -> SmargonCoordinate | None: + if target_point is None: + return None + + geom = self.ctx.sample_geometry_provider() + return SmargonCoordinate( + sh_mm=geom.picture_to_smargon(Coordinate(x=target_point[0], y=target_point[1])) + ) + + def _interpret_ml_loop_centre_box(self, boxes) -> tuple[SmargonCoordinate | None, int | None, list[int]]: + if boxes is None: + return None, None, [] + + classes: list[int] = [] + pin = None + centre_x, centre_y = None, None + + for box in boxes.boxes.values(): + if box and box.cls is not None: + classes.append(int(box.cls.value)) + if box.cls == MLBoxType.PIN: + pin = box + + best_box = self.ctx.mlbox.get_preferred_class_box( + boxes, + (MLBoxType.CRYSTAL, MLBoxType.LOOP_FACE, MLBoxType.LOOP_ALL, MLBoxType.PIN), + ) + if best_box is None or best_box.box is None or best_box.cls is None: + return None, None, classes + + cls = int(best_box.cls.value) + x1 = best_box.box.top_x + y1 = best_box.box.top_y + x2 = best_box.box.bottom_x + y2 = best_box.box.bottom_y + + geom = self.ctx.sample_geometry_provider() + + if cls == MLBoxType.LOOP_ALL.value: + if y1 + y2 <= x1 + x2: + centre_y = y1 + (y2 - y1) / 2 + centre_x = x1 + elif pin: + position_dict = self.ctx.mlbox.check_box_relation(pin, best_box) + if position_dict["overlap_y"] and position_dict["overlap_x"]: + centre_y = y1 + (y2 - y1) / 2 + centre_x = x1 + cls = MLBoxType.PIN.value + else: + centre_y = y2 if position_dict["top"] else y1 + centre_x = x1 + (x2 - x1) / 2 + else: + centre_y = y1 + centre_x = x1 + (x2 - x1) / 2 + + elif cls == MLBoxType.PIN.value: + centre_y = y1 + (y2 - y1) / 2 + centre_x = x1 + + elif cls in (MLBoxType.CRYSTAL.value, MLBoxType.LOOP_FACE.value): + centre_y = y1 + (y2 - y1) / 2 + centre_x = x1 + (x2 - x1) / 2 + + else: + return None, cls, classes + + coord = geom.picture_to_smargon(Coordinate(x=centre_x, y=centre_y)) + return SmargonCoordinate(sh_mm=coord), cls, classes + + def analyze_angle( + self, + *, + angle_deg: int, + zoom_value: float, + sample_id: int | None, + ) -> AngleAnalysis: + prediction_result = self.ctx.get_predictions() + log_ml_bundle_meta( + self.logger, + f"loop_center_angle_{angle_deg}_zoom_{zoom_value:.0f}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) + + predicted_target = self._extract_prediction_target( + target_point=prediction_result.target_point, + ) + calculated_target, selected_class, classes = self._interpret_ml_loop_centre_box( + prediction_result.predictions + ) + + analysis = AngleAnalysis( + angle_deg=angle_deg, + classes=classes, + selected_class=selected_class, + calculated_target=calculated_target, + predicted_target=predicted_target, + ) + analysis.ignore_only = self.is_ignore_only_classes(analysis.classes) + analysis.has_valid_target = self.has_valid_target_classes(analysis.classes) + analysis.final_target = self.select_smargon_target( + calculated_target=analysis.calculated_target, + predicted_target=analysis.predicted_target, + ) + return analysis \ No newline at end of file diff --git a/src/aare/daq/operations/loop_centering/models.py b/src/aare/daq/operations/loop_centering/models.py new file mode 100644 index 00000000..d64cd2b0 --- /dev/null +++ b/src/aare/daq/operations/loop_centering/models.py @@ -0,0 +1,51 @@ +from dataclasses import dataclass, field +from typing import Callable + +from aare.common.coordinate import SmargonCoordinate +from aare.common.sample_geometry import SampleGeometryModel +from aare.daq.config import BeamlineConfig +from aare.daq.devices import BeamlineDevices +from aare.daq.mlbox import MlBox, MLBoxPredictionsResult + + +@dataclass +class AngleAnalysis: + angle_deg: int + classes: list[int] = field(default_factory=list) + selected_class: int | None = None + calculated_target: SmargonCoordinate | None = None + predicted_target: SmargonCoordinate | None = None + final_target: SmargonCoordinate | None = None + moved: bool = False + ignore_only: bool = False + has_valid_target: bool = False + + +@dataclass +class AttemptSummary: + attempt_number: int + first_pass: list[AngleAnalysis] = field(default_factory=list) + correction_pass: list[AngleAnalysis] = field(default_factory=list) + + +@dataclass +class LoopCenteringSettings: + overlap_with_pin: float = 0.5 + confidence_min: float | None = None + zoom_value: float = 200.0 + max_attempts: int = 2 + target_tolerance_um: float = 500.0 + angles: tuple[int, int] = (0, 90) + + +@dataclass +class LoopCenteringContext: + cfg: BeamlineConfig + devs: BeamlineDevices + mlbox: MlBox + settings: LoopCenteringSettings + + sample_geometry_provider: Callable[[], SampleGeometryModel] + save_screenshot_db: Callable[[int, str, float], None] + append_smargon_trace: Callable[..., None] + get_predictions: Callable[[], MLBoxPredictionsResult] \ No newline at end of file diff --git a/src/aare/daq/operations/loop_centering/service.py b/src/aare/daq/operations/loop_centering/service.py new file mode 100644 index 00000000..fd28a9ff --- /dev/null +++ b/src/aare/daq/operations/loop_centering/service.py @@ -0,0 +1,206 @@ +import time +import traceback + +from aare.common.exception_handler import LoopCenteringFailed +from aare.common.logger_events import log_duration +from aare.common.models import LoopCenteringResult, MLBoxType +from aare.devices.area_detector import AutoEnum + +from aare.daq.operations.loop_centering.analyzer import LoopCenteringAnalyzer +from aare.daq.operations.loop_centering.models import ( + AngleAnalysis, + AttemptSummary, + LoopCenteringContext, +) + + +class LoopCenteringService: + def __init__(self, *, context: LoopCenteringContext, logger): + self.ctx = context + self.logger = logger + self.analyzer = LoopCenteringAnalyzer(context=context, logger=logger) + + def _run_angle_pass( + self, + *, + angle: int, + zoom_value: float, + sample_id: int | None, + trace_all_alc_moves: bool, + wait_screenshot_sleep_sec: float = 0.2, + ) -> AngleAnalysis: + self.logger.debug(f"Moving to new omega: {angle}") + time_to_move_aerotech = time.perf_counter() + self.ctx.devs.aerotech_omega = angle + log_duration( + self.logger, + "Completed Aerotech move during loop centering", + time.perf_counter() - time_to_move_aerotech, + extra={"angle_deg": angle, "zoom": zoom_value}, + ) + + analysis = self.analyzer.analyze_angle( + angle_deg=angle, + zoom_value=zoom_value, + sample_id=sample_id, + ) + + if analysis.has_valid_target and not analysis.ignore_only and analysis.final_target is not None: + time_to_move_smargon = time.perf_counter() + self.ctx.devs.smargon_pos = analysis.final_target + self.ctx.devs.smargon_wait(60) + log_duration( + self.logger, + "Completed Smargon move during loop centering", + time.perf_counter() - time_to_move_smargon, + extra={"angle_deg": angle, "zoom": zoom_value}, + ) + analysis.moved = True + + if sample_id is not None and trace_all_alc_moves: + self.ctx.append_smargon_trace( + sample_id=sample_id, + event=f"alc_move_zoom_{zoom_value:.0f}_angle_{angle}", + ) + else: + self.logger.debug(f"Skipping move at angle {angle}; classes={analysis.classes}") + + if sample_id is not None and analysis.selected_class in ( + MLBoxType.LOOP_ALL.value, + MLBoxType.LOOP_FACE.value, + MLBoxType.CRYSTAL.value, + ): + self.ctx.save_screenshot_db( + sample_id, + f"{sample_id}_{angle}_{zoom_value:.0f}", + wait_screenshot_sleep_sec, + ) + + return analysis + + def run( + self, + *, + sample_id: int | None, + trace_all_alc_moves: bool = False, + wait_screenshot_sleep_sec: float = 0.2, + max_attempts: int | None = None, + ) -> LoopCenteringResult: + settings = self.ctx.settings + found_classes_count: dict[int, int] = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0} + angles = settings.angles + max_attempts = settings.max_attempts if max_attempts is None else max_attempts + failure_reason: str | None = None + attempts: list[AttemptSummary] = [] + + try: + zoom_value = settings.zoom_value + + if abs(self.ctx.devs.zoom - zoom_value) > 1e-6: + self.ctx.devs.samcam_auto(AutoEnum.AUTO) + self.ctx.devs.zoom = zoom_value + time.sleep(0.2) + self.ctx.devs.samcam_auto(AutoEnum.ONCE) + + if sample_id is not None: + self.logger.info( + f"submitting to db loop center sequence for sample {sample_id}, zoom={zoom_value}" + ) + self.ctx.save_screenshot_db(sample_id, "pre_alc", wait_screenshot_sleep_sec) + + for attempt_number in range(1, max_attempts + 1): + self.logger.info(f"Starting ALC attempt {attempt_number}/{max_attempts}") + attempt = AttemptSummary(attempt_number=attempt_number) + + for angle in angles: + result = self._run_angle_pass( + angle=angle, + zoom_value=zoom_value, + sample_id=sample_id, + trace_all_alc_moves=trace_all_alc_moves, + wait_screenshot_sleep_sec=wait_screenshot_sleep_sec, + ) + attempt.first_pass.append(result) + + for c in result.classes: + found_classes_count[int(c)] = found_classes_count.get(int(c), 0) + 1 + + if attempt_number == 1: + valid_seen_first_pass = any(result.has_valid_target for result in attempt.first_pass) + if not valid_seen_first_pass: + failure_reason = ( + "ALC failed on attempt 1: no loop face, loop all or crystal " + "detected at 0 and 90 degrees" + ) + raise LoopCenteringFailed(failure_reason) + + for angle in angles: + result = self._run_angle_pass( + angle=angle, + zoom_value=zoom_value, + sample_id=sample_id, + trace_all_alc_moves=trace_all_alc_moves, + wait_screenshot_sleep_sec=wait_screenshot_sleep_sec, + ) + attempt.correction_pass.append(result) + + for c in result.classes: + found_classes_count[int(c)] = found_classes_count.get(int(c), 0) + 1 + + attempts.append(attempt) + + valid_seen_correction = any(result.has_valid_target for result in attempt.correction_pass) + + self.logger.debug( + f"ALC attempt {attempt_number}: valid_seen_correction={valid_seen_correction}" + ) + + if valid_seen_correction: + if sample_id is not None: + self.logger.info(f"sample {sample_id} centered") + self.ctx.append_smargon_trace(sample_id=sample_id, event="alc_success") + return LoopCenteringResult(success=True) + + failure_reason = f"ALC exceeded max attempts ({max_attempts})" + raise LoopCenteringFailed(failure_reason) + + except Exception as e: + total_detections = sum(found_classes_count.values()) + + if failure_reason is not None and total_detections > 0: + alc_comment = ( + f"{failure_reason}. " + f"Detected objects - " + f"Crystal: {found_classes_count.get(2, 0)}, " + f"Loop_face: {found_classes_count.get(3, 0)}, " + f"Loop_all: {found_classes_count.get(0, 0)}, " + f"Pin: {found_classes_count.get(1, 0)}, " + f"Ice: {found_classes_count.get(4, 0)}, " + f"Needle: {found_classes_count.get(5, 0)}" + ) + + elif failure_reason is not None: + alc_comment = failure_reason + + elif total_detections == 0: + alc_comment = "No objects detected" + + else: + alc_comment = ( + "Failed to centre but detected objects - " + f"Crystal: {found_classes_count.get(2, 0)}, " + f"Loop_face: {found_classes_count.get(3, 0)}, " + f"Loop_all: {found_classes_count.get(0, 0)}, " + f"Pin: {found_classes_count.get(1, 0)}, " + f"Ice: {found_classes_count.get(4, 0)}, " + f"Needle: {found_classes_count.get(5, 0)}" + ) + + self.logger.error(f"ALC exception: {alc_comment}") + self.logger.error(traceback.format_exc()) + self.logger.error(f"Error in loop centering: {e}") + return LoopCenteringResult( + success=False, + comment=alc_comment, + error=e, + ) \ No newline at end of file diff --git a/tests/unit/daq/operations/loop_centering/test_loop_centering_analyzer.py b/tests/unit/daq/operations/loop_centering/test_loop_centering_analyzer.py new file mode 100644 index 00000000..8b019156 --- /dev/null +++ b/tests/unit/daq/operations/loop_centering/test_loop_centering_analyzer.py @@ -0,0 +1,154 @@ +import types + +import pytest + +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.models import MLBoxType, MLOutputModel +from aare.daq.mlbox import MLBoxPredictionsResult +from aare.daq.operations.loop_centering.analyzer import LoopCenteringAnalyzer +from aare.daq.operations.loop_centering.models import ( + LoopCenteringContext, + LoopCenteringSettings, +) + + +class DummyLogger: + def info(self, *args, **kwargs): + pass + + def debug(self, *args, **kwargs): + pass + + def error(self, *args, **kwargs): + pass + + +class DummyGeometry: + def picture_to_smargon(self, coord: Coordinate) -> Coordinate: + return Coordinate(x=coord.x / 100.0, y=coord.y / 100.0, z=0.0) + + +class DummyMlBox: + def get_preferred_class_box(self, boxes, preferred_class=None): + for cls in preferred_class: + box = boxes.get_best_for_class(cls) + if box: + return box + return None + + def check_box_relation(self, box0, box1): + return {"left": True, "top": False, "overlap_x": False, "overlap_y": False} + + +def _coord_target(x: float, y: float, z: float = 0.0) -> SmargonCoordinate: + return SmargonCoordinate(sh_mm=Coordinate(x=x, y=y, z=z)) + + +@pytest.fixture +def analyzer(): + ctx = LoopCenteringContext( + cfg=types.SimpleNamespace(), + devs=types.SimpleNamespace(), + mlbox=DummyMlBox(), + settings=LoopCenteringSettings(), + sample_geometry_provider=lambda: DummyGeometry(), + save_screenshot_db=lambda *args, **kwargs: None, + append_smargon_trace=lambda *args, **kwargs: None, + get_predictions=lambda: MLBoxPredictionsResult( + predictions=None, + image=None, + target_point=None, + focus=None, + ), + ) + return LoopCenteringAnalyzer(context=ctx, logger=DummyLogger()) + + +def test_is_ignore_only_classes_true(analyzer): + assert analyzer.is_ignore_only_classes([ + MLBoxType.PIN.value, + MLBoxType.ICE.value, + MLBoxType.NEEDLE.value, + ]) + + +def test_is_ignore_only_classes_false_when_loop_present(analyzer): + assert not analyzer.is_ignore_only_classes([ + MLBoxType.PIN.value, + MLBoxType.LOOP_FACE.value, + ]) + + +def test_has_valid_target_classes_true(analyzer): + assert analyzer.has_valid_target_classes([ + MLBoxType.PIN.value, + MLBoxType.CRYSTAL.value, + ]) + + +def test_has_valid_target_classes_false(analyzer): + assert not analyzer.has_valid_target_classes([ + MLBoxType.PIN.value, + MLBoxType.ICE.value, + ]) + + +def test_select_smargon_target_prefers_prediction_when_within_tolerance(analyzer): + calculated = _coord_target(1.0, 2.0, 3.0) + predicted = _coord_target(1.0001, 2.0001, 3.0001) + + result = analyzer.select_smargon_target( + calculated_target=calculated, + predicted_target=predicted, + tolerance_um=500.0, + ) + + assert result is predicted + + +def test_select_smargon_target_keeps_calculated_when_outside_tolerance(analyzer): + calculated = _coord_target(1.0, 2.0, 3.0) + predicted = _coord_target(2.0, 3.0, 4.0) + + result = analyzer.select_smargon_target( + calculated_target=calculated, + predicted_target=predicted, + tolerance_um=500.0, + ) + + assert result is calculated + + +def test_analyze_angle_with_no_predictions_returns_empty_analysis(analyzer): + analysis = analyzer.analyze_angle(angle_deg=0, zoom_value=200.0, sample_id=1) + + assert analysis.angle_deg == 0 + assert analysis.classes == [] + assert analysis.selected_class is None + assert analysis.has_valid_target is False + assert analysis.ignore_only is False + assert analysis.calculated_target is None + assert analysis.predicted_target is None + assert analysis.final_target is None + + +def test_analyze_angle_with_crystal_prediction_returns_valid_target(analyzer): + boxes = MLOutputModel() + boxes.add_box(MLBoxType.CRYSTAL, (10, 20, 30, 40), conf=0.9) + + analyzer.ctx.get_predictions = lambda: MLBoxPredictionsResult( + predictions=boxes, + image=None, + target_point=(25.0, 30.0), + focus=1.0, + ) + + analysis = analyzer.analyze_angle(angle_deg=90, zoom_value=200.0, sample_id=1) + + assert analysis.selected_class == MLBoxType.CRYSTAL.value + assert MLBoxType.CRYSTAL.value in analysis.classes + assert analysis.has_valid_target is True + assert analysis.ignore_only is False + assert analysis.calculated_target is not None + assert analysis.predicted_target is not None + assert analysis.final_target is not None \ No newline at end of file diff --git a/tests/unit/daq/operations/loop_centering/test_loop_centering_service.py b/tests/unit/daq/operations/loop_centering/test_loop_centering_service.py new file mode 100644 index 00000000..1720c679 --- /dev/null +++ b/tests/unit/daq/operations/loop_centering/test_loop_centering_service.py @@ -0,0 +1,124 @@ +import types + +import pytest + +from aare.common.models import LoopCenteringResult +from aare.daq.operations.loop_centering.models import ( + AngleAnalysis, + LoopCenteringContext, + LoopCenteringSettings, +) +from aare.daq.operations.loop_centering.service import LoopCenteringService + + +class DummyLogger: + def info(self, *args, **kwargs): + pass + + def debug(self, *args, **kwargs): + pass + + def error(self, *args, **kwargs): + pass + + +@pytest.fixture +def context(): + devs = types.SimpleNamespace( + zoom=100.0, + aerotech_omega=0.0, + smargon_pos=None, + samcam_auto=lambda *args, **kwargs: None, + smargon_wait=lambda *args, **kwargs: None, + ) + screenshots = [] + traces = [] + + ctx = LoopCenteringContext( + cfg=types.SimpleNamespace(), + devs=devs, + mlbox=types.SimpleNamespace(), + settings=LoopCenteringSettings(), + sample_geometry_provider=lambda: types.SimpleNamespace(), + save_screenshot_db=lambda sample_id, filename, wait=0.0: screenshots.append((sample_id, filename, wait)), + append_smargon_trace=lambda **kwargs: traces.append(kwargs), + get_predictions=lambda: None, + ) + ctx._screenshots = screenshots + ctx._traces = traces + return ctx + + +def test_service_fails_when_attempt_one_has_no_valid_targets(monkeypatch, context): + service = LoopCenteringService(context=context, logger=DummyLogger()) + + monkeypatch.setattr( + service, + "_run_angle_pass", + lambda **kwargs: AngleAnalysis( + angle_deg=kwargs["angle"], + classes=[], + has_valid_target=False, + ignore_only=False, + ), + ) + + result = service.run(sample_id=1) + + assert isinstance(result, LoopCenteringResult) + assert result.success is False + assert "attempt 1" in result.comment.lower() + + +def test_service_succeeds_when_correction_pass_has_valid_target(monkeypatch, context): + service = LoopCenteringService(context=context, logger=DummyLogger()) + + call_count = {"count": 0} + + def fake_run_angle_pass(**kwargs): + call_count["count"] += 1 + if call_count["count"] <= 2: + return AngleAnalysis( + angle_deg=kwargs["angle"], + classes=[MLBoxType.CRYSTAL.value] if kwargs["angle"] == 0 else [], + has_valid_target=True if kwargs["angle"] == 0 else False, + ignore_only=False, + ) + return AngleAnalysis( + angle_deg=kwargs["angle"], + classes=[MLBoxType.CRYSTAL.value], + has_valid_target=True, + ignore_only=False, + moved=True, + ) + + from aare.common.models import MLBoxType + monkeypatch.setattr(service, "_run_angle_pass", fake_run_angle_pass) + + result = service.run(sample_id=2) + + assert result.success is True + + +def test_service_uses_settings_values(monkeypatch, context): + context.settings.zoom_value = 250.0 + context.settings.max_attempts = 3 + context.settings.angles = (0, 90) + + service = LoopCenteringService(context=context, logger=DummyLogger()) + + monkeypatch.setattr( + service, + "_run_angle_pass", + lambda **kwargs: AngleAnalysis( + angle_deg=kwargs["angle"], + classes=[], + has_valid_target=False, + ignore_only=False, + ), + ) + + result = service.run(sample_id=3) + + assert context.devs.zoom == 250.0 + assert result.success is False \ No newline at end of file diff --git a/tests/unit/daq/test_aare_daq_loop_centering.py b/tests/unit/daq/test_aare_daq_loop_centering.py new file mode 100644 index 00000000..dc6fbe9a --- /dev/null +++ b/tests/unit/daq/test_aare_daq_loop_centering.py @@ -0,0 +1,86 @@ +import types +from unittest.mock import MagicMock + +from aare.common.models import SampleShortInfo +from aare.common.exception_handler import LoopCenteringFailed +from aare.daq.daq import AareDAQ + + +def test_execute_loop_centering_rejects_invalid_sample(): + daq = object.__new__(AareDAQ) + daq._AareDAQ__set_state = MagicMock() + daq._AareDAQ__devs = MagicMock() + daq._AareDAQ__aare = MagicMock() + daq._handle_operation_error = MagicMock() + daq._create_loop_centering_service = MagicMock() + + result = daq._execute_loop_centering(None) + + assert result is False + daq._create_loop_centering_service.assert_not_called() + + +def test_execute_loop_centering_handles_service_failure(): + daq = object.__new__(AareDAQ) + daq._AareDAQ__set_state = MagicMock() + daq._AareDAQ__devs = MagicMock() + daq._AareDAQ__aare = MagicMock() + daq._handle_operation_error = MagicMock() + daq.save_screenshot_db = MagicMock() + + sample = SampleShortInfo( + db_id=1, + sample_name="S1", + puck_name="P1", + dewar_name="D1", + run_number=1, + user="p1", + pin=1, + location=None, + ) + + service = MagicMock() + service.run.return_value = types.SimpleNamespace( + success=False, + comment="failed", + error=LoopCenteringFailed("failed"), + ) + daq._create_loop_centering_service = MagicMock(return_value=service) + + result = daq._execute_loop_centering(sample) + + assert result is False + daq._handle_operation_error.assert_called_once() + + +def test_execute_loop_centering_handles_service_success(): + daq = object.__new__(AareDAQ) + daq._AareDAQ__set_state = MagicMock() + daq._AareDAQ__devs = MagicMock() + daq._AareDAQ__aare = MagicMock() + daq._handle_operation_error = MagicMock() + daq.save_screenshot_db = MagicMock() + + sample = SampleShortInfo( + db_id=1, + sample_name="S1", + puck_name="P1", + dewar_name="D1", + run_number=1, + user="p1", + pin=1, + location=None, + ) + + service = MagicMock() + service.run.return_value = types.SimpleNamespace( + success=True, + comment=None, + error=None, + ) + daq._create_loop_centering_service = MagicMock(return_value=service) + + result = daq._execute_loop_centering(sample) + + assert result is True + daq.save_screenshot_db.assert_called_once()