Models: added DAQOperation Enum and LoopCenteringResult dataclass. Corrected some whitespace formatting

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
+    confirmation_code: str
+
+
+@dataclass
+class LoopCenteringResult:
+    success: bool
+    comment: str | None = None
+    error: Exception | None = None