597 lines
20 KiB
Python
597 lines
20 KiB
Python
from pathlib import Path
|
|
import re
|
|
from enum import Enum
|
|
from typing import Annotated, Literal, Tuple, List, Optional
|
|
|
|
from pydantic import BaseModel, Field, field_validator
|
|
|
|
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
|
|
|
|
class DataCollectionParameters(BaseModel):
|
|
directory: Optional[str] = None
|
|
oscillation: Optional[float] = None # Only accept positive float
|
|
exposure: Optional[float] = None # Only accept positive floats between 0 and 1
|
|
totalrange: Optional[int] = None # Only accept positive integers between 0 and 360
|
|
transmission: Optional[
|
|
int
|
|
] = None # Only accept positive integers between 0 and 100
|
|
targetresolution: Optional[float] = None # Only accept positive float
|
|
aperture: Optional[str] = None # Optional string field
|
|
datacollectiontype: Optional[
|
|
str
|
|
] = None # Only accept "standard", other types might be added later
|
|
processingpipeline: Optional[
|
|
str
|
|
] = "" # Only accept "gopy", "autoproc", "xia2dials"
|
|
spacegroupnumber: Optional[
|
|
int
|
|
] = None # Only accept positive integers between 1 and 230
|
|
cellparameters: Optional[
|
|
str
|
|
] = None # Must be a set of six positive floats or integers
|
|
rescutkey: Optional[str] = None # Only accept "is" or "cchalf"
|
|
rescutvalue: Optional[
|
|
float
|
|
] = None # Must be a positive float if rescutkey is provided
|
|
userresolution: Optional[float] = None
|
|
pdbid: Optional[
|
|
str
|
|
] = "" # Accepts either the format of the protein data bank code or {provided}
|
|
autoprocfull: Optional[bool] = None
|
|
procfull: Optional[bool] = None
|
|
adpenabled: Optional[bool] = None
|
|
noano: Optional[bool] = None
|
|
ffcscampaign: Optional[bool] = None
|
|
trustedhigh: Optional[float] = None # Should be a float between 0 and 2.0
|
|
autoprocextraparams: Optional[str] = None # Optional string field
|
|
chiphiangles: Optional[float] = None # Optional float field between 0 and 30
|
|
dose: Optional[float] = None # Optional float field
|
|
cloud: bool = True
|
|
pdbmodel: Optional[str] = None
|
|
|
|
def to_dict(self):
|
|
"""Convert the model instance to a dictionary."""
|
|
return self.dict(
|
|
exclude_unset=True
|
|
) # Use this built-in method for serialization
|
|
|
|
class Config:
|
|
from_attributes = True
|
|
|
|
@field_validator("directory", mode="after")
|
|
@classmethod
|
|
def directory_characters(cls, v):
|
|
|
|
# Default directory value if empty
|
|
if not v: # Handles None or empty cases
|
|
default_value = "{sgPuck}/{sgPosition}"
|
|
return default_value
|
|
|
|
# Strip trailing slashes and store original value for comparison
|
|
v = str(v).strip("/") # Ensure it's a string and no trailing slashes
|
|
original_value = v
|
|
|
|
# Replace spaces with underscores
|
|
v = v.replace(" ", "_")
|
|
|
|
# Validate directory pattern with macros and allowed characters
|
|
valid_macros = [
|
|
"{date}",
|
|
"{prefix}",
|
|
"{sgPuck}",
|
|
"{sgPosition}",
|
|
"{beamline}",
|
|
"{sgPrefix}",
|
|
"{sgPriority}",
|
|
"{protein}",
|
|
"{method}",
|
|
]
|
|
valid_macro_pattern = re.compile(
|
|
"|".join(re.escape(macro) for macro in valid_macros)
|
|
)
|
|
|
|
# Check if the value contains valid macros
|
|
allowed_chars_pattern = "[a-z0-9_.+-/]"
|
|
v_without_macros = valid_macro_pattern.sub("macro", v)
|
|
|
|
allowed_path_pattern = re.compile(
|
|
f"^(({allowed_chars_pattern}+|macro)*/*)*$", re.IGNORECASE
|
|
)
|
|
if not allowed_path_pattern.match(v_without_macros):
|
|
raise ValueError(
|
|
f"'{v}' is not valid. Value must be a valid path or macro."
|
|
)
|
|
return v
|
|
|
|
@field_validator("aperture", mode="before")
|
|
@classmethod
|
|
def aperture_selection(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = int(float(v))
|
|
if v not in {1, 2, 3}:
|
|
raise ValueError(f" '{v}' is not valid. Value must be 1, 2, or 3.")
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid. Value must be 1, 2, or 3."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("oscillation", mode="before")
|
|
@classmethod
|
|
def positive_float_validator(cls, v):
|
|
if v is None:
|
|
return None
|
|
try:
|
|
v = float(v)
|
|
if v <= 0:
|
|
raise ValueError(f"'{v}' is not valid. Value must be a positive float.")
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f"'{v}' is not valid. Value must be a positive float."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("exposure", mode="before")
|
|
@classmethod
|
|
def exposure_in_range(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = float(v)
|
|
if not (0 <= v <= 1):
|
|
raise ValueError(
|
|
f" '{v}' is not valid. Value must be a float between 0 and 1."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid. Value must be a float between 0 and 1."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("totalrange", mode="before")
|
|
@classmethod
|
|
def totalrange_in_range(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = int(v)
|
|
if not (0 <= v <= 360):
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be an integer between 0 and 360."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be an integer between 0 and 360."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("transmission", mode="before")
|
|
@classmethod
|
|
def transmission_fraction(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = int(v)
|
|
if not (0 <= v <= 100):
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be an integer between 0 and 100."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be an integer between 0 and 100."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("datacollectiontype", mode="before")
|
|
@classmethod
|
|
def datacollectiontype_allowed(cls, v):
|
|
allowed = {"standard"} # Other types of data collection might be added later
|
|
if v and v.lower() not in allowed:
|
|
raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.")
|
|
return v
|
|
|
|
@field_validator("processingpipeline", mode="before")
|
|
@classmethod
|
|
def processingpipeline_allowed(cls, v):
|
|
allowed = {"aareproc", "autoproc"}
|
|
if v and v.lower() not in allowed:
|
|
raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.")
|
|
return v
|
|
|
|
@field_validator("spacegroupnumber", mode="before")
|
|
@classmethod
|
|
def spacegroupnumber_allowed(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = int(v)
|
|
if not (1 <= v <= 230):
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be an integer between 1 and 230."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be an integer between 1 and 230."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("cellparameters", mode="before")
|
|
@classmethod
|
|
def cellparameters_format(cls, v):
|
|
if v:
|
|
# Replace commas with spaces, then split on whitespace
|
|
tokens = v.replace(",", " ").split()
|
|
try:
|
|
values = [float(i) for i in tokens]
|
|
except ValueError:
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
" Value must be a set of six positive floats"
|
|
" or integers (separated by space or comma)."
|
|
)
|
|
if len(values) != 6 or any(val <= 0 for val in values):
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
" Value must be a set of six positive floats"
|
|
" or integers (separated by space or comma)."
|
|
)
|
|
return v
|
|
|
|
# @field_validator("rescutkey", "rescutvalue", mode="before")
|
|
# @classmethod
|
|
# def rescutkey_value_pair(cls, values):
|
|
# rescutkey = values.get("rescutkey")
|
|
# rescutvalue = values.get("rescutvalue")
|
|
# if rescutkey and rescutvalue:
|
|
# if rescutkey not in {"is", "cchalf"}:
|
|
# raise ValueError("Rescutkey must be either 'is' or 'cchalf'")
|
|
# if not isinstance(rescutvalue, float) or rescutvalue <= 0:
|
|
# raise ValueError(
|
|
# "Rescutvalue must be a positive float if rescutkey is provided"
|
|
# )
|
|
# return values
|
|
|
|
@field_validator("trustedhigh", mode="before")
|
|
@classmethod
|
|
def trustedhigh_allowed(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = float(v)
|
|
if not (0 <= v <= 2.0):
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be a float between 0 and 2.0."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid." f"Value must be a float between 0 and 2.0."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("chiphiangles", mode="before")
|
|
@classmethod
|
|
def chiphiangles_allowed(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = float(v)
|
|
if not (0 <= v <= 30):
|
|
raise ValueError(
|
|
f" '{v}' is not valid."
|
|
f"Value must be a float between 0 and 30."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid. Value must be a float between 0 and 30."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("dose", mode="before")
|
|
@classmethod
|
|
def dose_positive(cls, v):
|
|
if v is not None:
|
|
try:
|
|
v = float(v)
|
|
if v <= 0:
|
|
raise ValueError(
|
|
f" '{v}' is not valid. Value must be a positive float."
|
|
)
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(
|
|
f" '{v}' is not valid. Value must be a positive float."
|
|
) from e
|
|
return v
|
|
|
|
@field_validator("pdbmodel", mode="after")
|
|
@classmethod
|
|
def validate_filepath(cls, v):
|
|
if v is None:
|
|
return v
|
|
v_str = str(v) # Ensure v is a string for further checks
|
|
if any(c in v_str for c in '<>:"|?*'):
|
|
raise ValueError("File path contains invalid characters.")
|
|
path = Path(v_str)
|
|
if not path.parts:
|
|
raise ValueError("Not a valid path.")
|
|
return v_str # Return as string for JSON serialization
|
|
|
|
@field_validator("cloud", mode="before")
|
|
@classmethod
|
|
def coerce_cloud_default(cls, v):
|
|
if v in ("", None):
|
|
return True
|
|
if isinstance(v, bool):
|
|
return v
|
|
v_str = str(v).strip().lower()
|
|
if v_str in {"true", "yes", "1"}:
|
|
return True
|
|
if v_str in {"false", "no", "0"}:
|
|
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
|
|
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 ReferencePuckInfo(BaseModel):
|
|
# test_sample = SampleShortInfo(db_id=-1, puck_name="test_puck", dewar_name="test_dewar", sample_name="test_sample",
|
|
# pin=11, location=DewarAddress(segment="X", pos=1))
|
|
|
|
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=1, sam_cam_gain=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)
|
|
print('interpolating zoom')
|
|
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=0, exposure=0.05),
|
|
280: SampleCameraSettings(gain=0, exposure=0.05),
|
|
500: SampleCameraSettings(gain=0, exposure=0.05),
|
|
700: SampleCameraSettings(gain=0, exposure=0.05),
|
|
800: SampleCameraSettings(gain=0, exposure=0.10),
|
|
1000: SampleCameraSettings(gain=0, exposure=0.20)
|
|
}
|
|
)
|
|
|
|
def def_bl_zoom() -> ZoomModel:
|
|
return ZoomModel(
|
|
z={
|
|
1: SampleCameraSettings(gain=0, exposure=0.002),
|
|
280: SampleCameraSettings(gain=0, exposure=0.002),
|
|
500: SampleCameraSettings(gain=0, exposure=0.002),
|
|
700: SampleCameraSettings(gain=0, exposure=0.002),
|
|
800: SampleCameraSettings(gain=0, exposure=0.002),
|
|
1000: SampleCameraSettings(gain=0, exposure=0.005)
|
|
}
|
|
)
|
|
|
|
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
|