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 from aaredaqlib.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)] 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 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 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 Crystal = 2 Loop_face = 3 class BoundingBoxModel(BaseModel): top_x: float top_y: float bottom_x: float bottom_y: float class MLBoxModel(BaseModel): cls: MLBoxType box: BoundingBoxModel conf: float @staticmethod def from_tuple(cls:MLBoxType, box_tuple: tuple[float, float, float, float], conf: float) -> "MLBoxModel": x1, y1, x2, y2 = box_tuple return MLBoxModel( cls=cls, box=BoundingBoxModel(top_x=float(x1), top_y=float(y1), bottom_x=float(x2), bottom_y=float(y2)), conf=float(conf) ) class MLOutputModel(BaseModel): boxes: dict[str, MLBoxModel] = {} @staticmethod def get_class_str(cls: MLBoxType) -> str: if cls == MLBoxType.Loop_all: return "Loop_all" if cls == MLBoxType.Pin: return "Pin" if cls == MLBoxType.Crystal: return "Crystal" if cls == MLBoxType.Loop_face: return "Loop_face" return "Unknown" def _next_unique_key(self, base: str) -> str: if base not in self.boxes: return base i = 2 while f"{base}_{i}" in self.boxes: i += 1 return f"{base}_{i}" def add_box(self, cls: MLBoxType, box_tuple: tuple[float, float, float, float], conf: float | None = None) -> str: key_base = self.get_class_str(cls) key = self._next_unique_key(key_base) self.boxes[key] = MLBoxModel.from_tuple(cls, box_tuple, conf) return key def get_box_model(self, key: str) -> MLBoxModel | None: return self.boxes.get(key) def get_box_tuple(self, key: str) -> tuple[float, float, float, float] | None: m = self.get_box_model(key) if not m or not m.box: return None return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y) def get_box_tuple_with_conf(self, key: str) -> tuple[float, float, float, float, float] | None: m = self.get_box_model(key) if not m or not m.box: return None conf = float(m.conf) if m.conf is not None else 0.0 return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, conf) def get_keys_for_class(self, cls: MLBoxType) -> list[str]: base = self.get_class_str(cls) return [k for k, v in self.boxes.items() if v.cls == cls and (k == base or k.startswith(base + "_"))] def get_models_for_class(self, cls: MLBoxType) -> list[MLBoxModel]: keys = self.get_keys_for_class(cls) return [self.boxes[k] for k in keys] def get_tuples_for_class(self, cls: MLBoxType) -> list[tuple[float, float, float, float]]: out: list[tuple[float, float, float, float]] = [] for m in self.get_models_for_class(cls): if m.box: out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y)) return out def get_tuples_with_conf_for_class(self, cls: MLBoxType) -> list[tuple[float, float, float, float, float]]: out: list[tuple[float, float, float, float, float]] = [] for m in self.get_models_for_class(cls): if m.box: out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, float(m.conf))) return out def get_best_for_class(self, cls: MLBoxType) -> MLBoxModel | None: models = self.get_models_for_class(cls) if not models: return None return max(models, key=lambda m: (m.conf if m.conf is not None else 0.0)) 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 ZoomModeEnum(Enum): User = 1 LoopCenter = 2 BeamLocation = 3 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(beamline) -> ZoomModel: print(f'user zoom for {beamline}') if beamline == MXBeamline.X06DA: return ZoomModel( z={ 1: SampleCameraSettings(gain=0, exposure=0.05), 280: SampleCameraSettings(gain=0, exposure=0.05), 500: SampleCameraSettings(gain=0, exposure=0.1), 700: SampleCameraSettings(gain=0, exposure=0.15), 800: SampleCameraSettings(gain=0, exposure=0.2), 1000: SampleCameraSettings(gain=0, exposure=0.25) } ) elif beamline == MXBeamline.X10SA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) elif beamline == MXBeamline.X06SA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) elif beamline == MXBeamline.SIMULATED: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) else: raise ValueError(f"Invalid beamline: {beamline}") def def_bl_zoom(beamline) -> ZoomModel: if beamline == MXBeamline.X06DA: 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) } ) elif beamline == MXBeamline.X10SA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) elif beamline == MXBeamline.X06SA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) elif beamline == MXBeamline.SIMULATED: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) else: raise ValueError(f"Invalid beamline: {beamline}") def def_loop_centering_zoom(beamline) -> ZoomModel: if beamline == MXBeamline.X06DA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) elif beamline == MXBeamline.X10SA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) elif beamline == MXBeamline.X06SA: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) elif beamline == MXBeamline.SIMULATED: return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) else: raise ValueError(f"Invalid beamline: {beamline}") def zoom_manager(mode: ZoomModeEnum = ZoomModeEnum.User, beamline: MXBeamline = None) -> ZoomModel | None: if beamline is None or beamline == MXBeamline.SIMULATED: print('SIMULATED') return def_zoom(beamline) if mode == ZoomModeEnum.BeamLocation: return def_bl_zoom(beamline) elif mode == ZoomModeEnum.LoopCenter: return def_loop_centering_zoom(beamline) elif mode == ZoomModeEnum.User: return def_zoom(beamline) else: raise ValueError(f"Invalid zoom mode: {mode}") 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 exp_shutter_open: bool | None flux_ph_s: float sample_camera: SampleCameraSettings transmission: Annotated[float, Field(ge=0.0, le=1.0)] | None zoom: float commissioning_mode: bool 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 bl: BeamlineStatus state: BeamlineStateEnum busy: bool sample: SampleShortInfo | None = None session: SessionStatus box: BoundingBoxModel | None = None last_best_res: float | None = None last_best_b_factor: float | None = None crystal_size: CrystalSize = CrystalSize(x=0,y=0,z=0) 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 class SimpleStrategyInputModel(BaseModel): last_best_res: float | None = None last_best_b_factor: float | None = None crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0) angular_range: int = 360 start_angle: float = 0 incr_omega_deg:float = 0.2 d_vis: float = 1.5 current_temp_k: float = 100 filename: Optional[str] = None class SimpleScanParameters(BaseModel): filename: Optional[str] = None dtz: float = 150 exp_time_s: float = 0.02 start_omega_deg:float = 0 incr_omega_deg: float = 0.2 steps: int = 1800 transmission: Annotated[float, Field(ge=0.0, le=1.0)] = 1.0 last_best_res: float | None = None last_best_b_factor: float | None = None crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0) flux_ph_s: float | None = None calculated_dose_Mgy: float | None = None calculated_dose_rate_MGy_s: float | None = None xtal_size_dose_rate_MGy_s: float | None = None target_dose_MGy: float | None = None beam_size_x_um: float | None = None beam_size_y_um: float | None = None dose_rate_MGy_s: float | None = None d_vis: float | None = None d_tar: float | None = None