259 lines
7.5 KiB
Python
259 lines
7.5 KiB
Python
from enum import Enum
|
|
from typing import Annotated, Literal, Tuple, List, Optional
|
|
|
|
from aareDBclient.models.data_collection_parameters import DataCollectionParameters
|
|
from pydantic import BaseModel, Field
|
|
|
|
from aaredaqlib.coordinate import Coordinate
|
|
from aaredaqlib.diffraction_geometry import DiffractionGeometry
|
|
from aaredaqlib.sample_geometry import SampleGeometryModel
|
|
|
|
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)]
|
|
|
|
|
|
class SampleDewarAddress(BaseModel):
|
|
puck: DewarAddress
|
|
pin: Annotated[int, Field(ge=1, le=16)]
|
|
|
|
|
|
# From the database for puck loading
|
|
class PuckInfo(BaseModel):
|
|
db_id: int
|
|
puck_name: str
|
|
dewar_name: str
|
|
user: str = ""
|
|
location: Optional[DewarAddress] = None
|
|
|
|
|
|
# From TELL to database after loading
|
|
class PuckLoadedInfo(BaseModel):
|
|
puck_name: str
|
|
location: DewarAddress
|
|
|
|
|
|
# From database to TELL after loading
|
|
class SampleShortInfo(BaseModel):
|
|
db_id: int
|
|
puck_name: str
|
|
dewar_name: str
|
|
sample_name: str
|
|
aaredb_params: Optional[DataCollectionParameters] = None
|
|
user: str = ""
|
|
pin: Annotated[int, Field(ge=1, le=16)]
|
|
location: DewarAddress | None = None
|
|
priority: Optional[float] = 1.0
|
|
comment: Optional[str] = None
|
|
mount_count: int = 0
|
|
def tell_address(self) -> SampleDewarAddress:
|
|
return SampleDewarAddress(puck=self.location, pin=self.pin)
|
|
|
|
def loc_str(self) -> str:
|
|
if self.location is None:
|
|
return "-"
|
|
else:
|
|
return f"{self.location.segment}{self.location.pos}-{self.pin}"
|
|
|
|
def loc_str_sort(self) -> str:
|
|
if self.location is None:
|
|
return ""
|
|
else:
|
|
return f"{self.location.segment}{self.location.pos}-{self.pin:02d}"
|
|
|
|
|
|
class SampleShortInfoList(BaseModel):
|
|
s: List[SampleShortInfo]
|
|
|
|
|
|
class BeamMarkCoeffModel(BaseModel):
|
|
"""
|
|
Model to calculate beam center for a given zoom level based on a quadratic approximation
|
|
For both x and y it contains tuple of coefficients a, b, c (a*zoom^2 + b*zoom + c)
|
|
Default is beam center in 1000,1000 at any zoom level
|
|
"""
|
|
|
|
coeff_x: Tuple[float, float, float] = (0, 0, 1000.0)
|
|
coeff_y: Tuple[float, float, float] = (0, 0, 1000.0)
|
|
|
|
def apply(self, zoom: float):
|
|
return Coordinate(
|
|
x=self.coeff_x[0] * zoom**2 + self.coeff_x[1] * zoom + self.coeff_x[2],
|
|
y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2],
|
|
)
|
|
|
|
class BoundingBoxModel(BaseModel):
|
|
top_x: float
|
|
top_y: float
|
|
bottom_x: float
|
|
bottom_y: float
|
|
|
|
|
|
class LoopCenteringZoomModelElem(BaseModel):
|
|
zoom_value: float
|
|
sam_cam_gain: float
|
|
sam_cam_exp: float
|
|
|
|
|
|
class LoopCenteringZoomModel(BaseModel):
|
|
z: List[LoopCenteringZoomModelElem]
|
|
|
|
# zoom gain exp (X06DA)
|
|
# 1 50 0.05
|
|
# 200 50 0.05
|
|
# 500 70 0.05
|
|
# 700 150 0.05
|
|
# 800 100 0.10
|
|
# 1000 200 0.10
|
|
|
|
|
|
def def_loop_centering_zoom() -> LoopCenteringZoomModel:
|
|
return LoopCenteringZoomModel(
|
|
z=[
|
|
LoopCenteringZoomModelElem(
|
|
zoom_value=280, sam_cam_gain=100.0, sam_cam_exp=0.05
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
class BeamlineStateEnum(Enum):
|
|
Maintenance = 1
|
|
SampleExchange = 2
|
|
SampleAlignment = 3
|
|
DataCollection = 4
|
|
DewarTransfer = 5
|
|
XrayFluorescence = 6
|
|
BeamLocation = 7
|
|
Moving = 8
|
|
RobotSampleExchange = 9
|
|
# Busy state is used whenever transition between two states happen or mounting/data collection procedure happens
|
|
|
|
|
|
class SessionsStateEnum(Enum):
|
|
OwnedByYou = 1
|
|
OwnedByElse = 2
|
|
Vacant = 3
|
|
|
|
class SampleCameraSettings(BaseModel):
|
|
gain: float
|
|
exposure: float
|
|
|
|
|
|
class ZoomModel(BaseModel):
|
|
z:dict[float, SampleCameraSettings]
|
|
|
|
def get_camera_settings(self, zoom_value: float) -> SampleCameraSettings | tuple[float, float]:
|
|
if not self.z:
|
|
raise ValueError("No zoom data available")
|
|
|
|
if zoom_value in self.z:
|
|
elem = self.z[zoom_value]
|
|
return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)
|
|
|
|
sorted_zooms = sorted(self.z.keys())
|
|
|
|
if zoom_value <= sorted_zooms[0]:
|
|
elem = self.z[sorted_zooms[0]]
|
|
return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)
|
|
|
|
if zoom_value >= sorted_zooms[-1]:
|
|
elem = self.z[sorted_zooms[-1]]
|
|
return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)
|
|
|
|
for i in range(len(sorted_zooms) - 1):
|
|
if sorted_zooms[i] <= zoom_value <= sorted_zooms[i + 1]:
|
|
lower_zoom = sorted_zooms[i]
|
|
upper_zoom = sorted_zooms[i + 1]
|
|
|
|
lower_elem = self.z[lower_zoom]
|
|
upper_elem = self.z[upper_zoom]
|
|
|
|
t = (zoom_value - lower_zoom) / (upper_zoom - lower_zoom)
|
|
|
|
interpolated_gain = lower_elem.gain + t * (upper_elem.gain - lower_elem.gain)
|
|
interpolated_exp = lower_elem.exposure + t * (upper_elem.exposure - lower_elem.exposure)
|
|
|
|
return SampleCameraSettings(gain=interpolated_gain, exposure=interpolated_exp)
|
|
|
|
closest = min(self.z.keys(), key=lambda x: abs(x - zoom_value))
|
|
elem = self.z[closest]
|
|
return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)
|
|
|
|
def def_zoom() -> ZoomModel:
|
|
return ZoomModel(
|
|
z={
|
|
1: SampleCameraSettings(gain=50.0, exposure=0.05),
|
|
280: SampleCameraSettings(gain=100.0, exposure=0.05),
|
|
500: SampleCameraSettings(gain=100.0, exposure=0.05),
|
|
700: SampleCameraSettings(gain=150.0, exposure=0.05),
|
|
800: SampleCameraSettings(gain=100.0, exposure=0.10),
|
|
1000: SampleCameraSettings(gain=200.0, exposure=0.10)
|
|
}
|
|
)
|
|
|
|
|
|
class AutofocusSettings(BaseModel):
|
|
center_x_pxl: float
|
|
center_y_pxl: float
|
|
radius_pxl: float
|
|
z_range_um: float
|
|
z_steps: int
|
|
|
|
class BeamlineStatus(BaseModel):
|
|
name: str
|
|
ring_current_mA: float
|
|
light: Annotated[float, Field(ge=0.0, le=100.0)]
|
|
cryojet_K: float
|
|
shutter_open: bool
|
|
flux_ph_s: float
|
|
sample_camera: SampleCameraSettings
|
|
transmission: Annotated[float, Field(ge=0.0, le=1.0)] | None
|
|
zoom: float
|
|
|
|
class SessionStatus(BaseModel):
|
|
session: SessionsStateEnum = SessionsStateEnum.Vacant
|
|
current_pgroup: str | None = None
|
|
staff: bool = False
|
|
|
|
class DAQStatusModel(BaseModel):
|
|
geom: SampleGeometryModel
|
|
diffraction: DiffractionGeometry
|
|
bl: BeamlineStatus
|
|
state: BeamlineStateEnum
|
|
busy: bool
|
|
sample: SampleShortInfo | None = None
|
|
session: SessionStatus
|
|
box: BoundingBoxModel | None = None
|
|
|
|
class BeamlineSettingsModel(BaseModel):
|
|
dtz_max: float | None = 1600.0
|
|
dtz_min: float | None = 120.0
|
|
dtz_collection: float | None = 130.0
|
|
dtz_park: float | None = 150.0
|
|
dtz_wash_sample_distance: float | None = 150.0
|
|
dtz_bsz_safety_margin: float | None = 50.0
|
|
bsz: float | None = 25.0
|
|
|
|
camera_max_magnification: float | None = 1.0
|
|
camera_min_magnification: float | None = 500.0
|
|
camera_translation_factor_a: float | None = 0.00253
|
|
camera_translation_factor_b: float | None = 255.0
|
|
|
|
class CryojetSettingsModel(BaseModel):
|
|
cryojet_park_position: float | None = 12.0
|
|
cryojet_measurement_position: float | None = 5.0
|
|
cryojet_in_use: bool | None = True
|