From 8f75b466b7427ffd1bb129de652a6fb3c23b14ef Mon Sep 17 00:00:00 2001 From: appleb_m Date: Mon, 20 Apr 2026 11:47:14 +0200 Subject: [PATCH] Models: added DAQOperation Enum and LoopCenteringResult dataclass. Corrected some whitespace formatting --- src/aare/common/models.py | 48 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/src/aare/common/models.py b/src/aare/common/models.py index 2f6ca491..8f0c6be0 100644 --- a/src/aare/common/models.py +++ b/src/aare/common/models.py @@ -2,7 +2,7 @@ from pathlib import Path import re from enum import Enum from typing import Annotated, Literal, Tuple, List, Optional - +from dataclasses import dataclass from pydantic import BaseModel, Field, field_validator, AfterValidator from aare.common.coordinate import Coordinate, positive_coords @@ -12,18 +12,21 @@ from jfjoch_client.models.scan_result import ScanResult from aare.common.beamline import MXBeamline + class StagePositionEnum(Enum): MEASURE = 0 PARK = 1 DOWN = 2 UNKNOWN = 3 + class TokenData(BaseModel): sub: str # Username pgroups: List[str] session: int staff: bool = False + class DewarAddress(BaseModel): segment: Literal["A", "B", "C", "D", "E", "F", "X", "R"] pos: Annotated[int, Field(ge=1, le=5)] @@ -48,6 +51,7 @@ class PuckLoadedInfo(BaseModel): puck_name: str location: DewarAddress + class DataCollectionParameters(BaseModel): directory: Optional[str] = None oscillation: Optional[float] = None # Only accept positive float @@ -374,6 +378,7 @@ class DataCollectionParameters(BaseModel): return False raise ValueError("cloud must be blank for default, or True/False") + # From database to TELL after loading class SampleShortInfo(BaseModel): db_id: int @@ -407,9 +412,11 @@ class SampleShortInfo(BaseModel): def from_dict(cls, data: dict): return cls(**data) + class SampleShortInfoList(BaseModel): s: List[SampleShortInfo] + class BeamMarkCoeffModel(BaseModel): """ Model to calculate beam center for a given zoom level based on a quadratic approximation @@ -426,17 +433,20 @@ class BeamMarkCoeffModel(BaseModel): y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2], ) + class FluorescenceSpectrumParameterModel(BaseModel): erase: bool = True acq_time_s: float transmission: Annotated[float, Field(ge=0,le=1)] | None = None + class FluorescenceSpectrumOutputModel(BaseModel): bkg: list[float] | None = None spectrum: list[float] energy_eV: list[float] average_dead_time: Annotated[float, Field(ge=0,le=1)] + class MLBoxType(Enum): Loop_all = 0 Pin = 1 @@ -445,12 +455,14 @@ class MLBoxType(Enum): Ice = 4 Needle = 5 + class BoundingBoxModel(BaseModel): top_x: float top_y: float bottom_x: float bottom_y: float + class MLBoxModel(BaseModel): cls: MLBoxType box: BoundingBoxModel @@ -465,6 +477,7 @@ class MLBoxModel(BaseModel): conf=float(conf) ) + class MLOutputModel(BaseModel): boxes: dict[str, MLBoxModel] = {} @@ -557,6 +570,17 @@ class BeamlineStateEnum(Enum): # Busy state is used whenever transition between two states happen or mounting/data collection procedure happens +class DAQOperation(str, Enum): + AUTOMATION = "automation" + MOUNT = "mount" + UNMOUNT = "unmount" + LOOP_CENTERING = "loop_centering" + FACE_CENTERING = "face_centering" + RASTER = "raster" + ROTATION = "rotation" + MEASURE = "measure" + + class SessionsStateEnum(Enum): Vacant = 0 OwnedByYou = 1 @@ -564,15 +588,18 @@ class SessionsStateEnum(Enum): PendingYouToElse = 3 PendingElseToYou = 4 + class SampleCameraSettings(BaseModel): gain: float exposure: float + class ZoomModeEnum(Enum): User = 1 LoopCenter = 2 BeamLocation = 3 + class ZoomModel(BaseModel): z:dict[float, SampleCameraSettings] @@ -683,7 +710,6 @@ def zoom_manager(mode: ZoomModeEnum = ZoomModeEnum.User, beamline: MXBeamline = raise ValueError(f"Invalid zoom mode: {mode}") - class AutofocusSettings(BaseModel): center_x_pxl: float | None # Use beam center center_y_pxl: float | None # Use beam center @@ -691,6 +717,7 @@ class AutofocusSettings(BaseModel): z_range_um: float z_steps: int + class BeamlineStatus(BaseModel): name: str ring_current_mA: float @@ -707,16 +734,19 @@ class BeamlineStatus(BaseModel): dtz_min: float dtz_max: float + class SessionStatus(BaseModel): session: SessionsStateEnum = SessionsStateEnum.Vacant current_pgroup: str | None = None staff: bool = False + class CrystalSize(BaseModel): x: float = 0.0 y: float = 0.0 z: float = 0.0 + class DAQStatusModel(BaseModel): geom: SampleGeometryModel diffraction: DiffractionGeometry @@ -739,6 +769,7 @@ class DAQStatusModel(BaseModel): aerotech_connected: bool = True aerotech_error: str | None = None + class BeamlineSettingsModel(BaseModel): dtz_max: float | None = 1600.0 dtz_min: float | None = 120.0 @@ -753,11 +784,13 @@ class BeamlineSettingsModel(BaseModel): camera_translation_factor_a: float | None = 0.00253 camera_translation_factor_b: float | None = 512.0 + class CryojetSettingsModel(BaseModel): cryojet_park_position: float | None = 12.0 cryojet_measurement_position: float | None = 5.0 cryojet_in_use: bool | None = True + class SimpleStrategyInputModel(BaseModel): last_best_res: float | None = None last_best_b_factor: float | None = None @@ -792,6 +825,7 @@ class SimpleScanParameters(BaseModel): d_vis: float | None = None d_tar: float | None = None + class ScanResultPayloadModel(BaseModel): result: ScanResult sample_id: int @@ -799,5 +833,13 @@ class ScanResultPayloadModel(BaseModel): beam_mark_pxl: tuple[float, float] beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)] + class RecoveryActionRequest(BaseModel): - confirmation_code: str \ No newline at end of file + confirmation_code: str + + +@dataclass +class LoopCenteringResult: + success: bool + comment: str | None = None + error: Exception | None = None \ No newline at end of file