diff --git a/src/aare/daq/operations/loop_centering/models.py b/src/aare/daq/operations/loop_centering/models.py index d64cd2b0..4313b81a 100644 --- a/src/aare/daq/operations/loop_centering/models.py +++ b/src/aare/daq/operations/loop_centering/models.py @@ -1,11 +1,15 @@ 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 +from aare.daq.mlbox import MlBox +from aare.daq.operations.common.runtime import ( + DAQRuntimeState, + OperationServices, + PredictionProvider, + TraceWriter, +) @dataclass @@ -44,8 +48,7 @@ class LoopCenteringContext: 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 + runtime: DAQRuntimeState + services: OperationServices + trace_writer: TraceWriter + prediction_provider: PredictionProvider \ 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 index c5643ccb..0fe50966 100644 --- a/src/aare/daq/operations/loop_centering/service.py +++ b/src/aare/daq/operations/loop_centering/service.py @@ -58,7 +58,7 @@ class LoopCenteringService: analysis.moved = True if sample_id is not None and trace_all_alc_moves: - self.ctx.append_smargon_trace( + self.ctx.trace_writer.append_smargon_trace( sample_id=sample_id, event=f"alc_move_zoom_{zoom_value:.0f}_angle_{angle}", ) @@ -70,7 +70,7 @@ class LoopCenteringService: MLBoxType.LOOP_FACE.value, MLBoxType.CRYSTAL.value, ): - self.ctx.save_screenshot_db( + self.ctx.services.screenshots.save_to_db( sample_id, f"{sample_id}_{angle}_{zoom_value:.0f}", wait_screenshot_sleep_sec, @@ -106,7 +106,7 @@ class LoopCenteringService: 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) + self.ctx.services.screenshots.save_to_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}") @@ -158,7 +158,7 @@ class LoopCenteringService: 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") + self.ctx.trace_writer.append_smargon_trace(sample_id=sample_id, event="alc_success") return LoopCenteringResult(success=True) failure_reason = f"ALC exceeded max attempts ({max_attempts})"