96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
from typing import Annotated, List
|
|
|
|
import numpy as np
|
|
from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords
|
|
from aarecommon.math.sample_geometry import SampleGeometryModel
|
|
from jfjoch_client.models.scan_result import ScanResult
|
|
from pydantic import AfterValidator, BaseModel, Field
|
|
|
|
|
|
class RasterGridRequest(BaseModel):
|
|
dtz: float | None = None
|
|
transmission: float | None = None
|
|
file_prefix: str | None = None
|
|
exp_time_s: float
|
|
|
|
# Number of grid elements - 0 is allowed for empty grid
|
|
n_x: Annotated[int, Field(ge=0)]
|
|
n_y: Annotated[int, Field(ge=0)]
|
|
|
|
# Size of grid elements
|
|
grid_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)]
|
|
|
|
# Coordinates of top left corner in Smargon (SH...) coordinates.
|
|
# If None, measure from the current position.
|
|
smargon_top_left: SmargonCoordinate | None
|
|
|
|
# omega angle
|
|
omega_deg: float = 0.0
|
|
|
|
visible: bool = True
|
|
|
|
def get_image_number(self) -> int:
|
|
return self.n_x * self.n_y
|
|
|
|
def grid_size_pxl(self, geom: SampleGeometryModel) -> Coordinate:
|
|
return self.grid_size_mm / geom.pixel_in_mm
|
|
|
|
def start_pxl(self, geom: SampleGeometryModel) -> Coordinate:
|
|
if self.smargon_top_left is None:
|
|
return geom.smargon_to_beamline(geom.smargon.sh_mm)
|
|
return geom.smargon_to_beamline(self.smargon_top_left.sh_mm)
|
|
|
|
|
|
class CenterOfMassModel(BaseModel):
|
|
n_x: float
|
|
n_y: float
|
|
max_image: int | None = None
|
|
|
|
@classmethod
|
|
def model_validate_maybe(cls, n_x: float, n_y: float) -> "CenterOfMassModel | None":
|
|
# Return None if either coordinate is NaN
|
|
if np.isnan(n_x) or np.isnan(n_y):
|
|
return None
|
|
return cls(n_x=float(n_x), n_y=float(n_y))
|
|
|
|
def get_com_mm(self, r: RasterGridRequest) -> Coordinate:
|
|
# defines COM in mm at beam position, relative to the top left corner of grid
|
|
return Coordinate(
|
|
x=(self.n_x + 0.5) * r.grid_size_mm.x, y=(self.n_y + 0.5) * r.grid_size_mm.y
|
|
)
|
|
|
|
def get_com_pxl(
|
|
self, r: RasterGridRequest, geom: SampleGeometryModel
|
|
) -> Coordinate:
|
|
# defines COM in pixels at beam position, relative to the top left corner of grid
|
|
com_mm = self.get_com_mm(r)
|
|
return Coordinate(x=com_mm.x / geom.pixel_in_mm, y=com_mm.y / geom.pixel_in_mm)
|
|
|
|
def com_nan_check(self):
|
|
# True if both values are finite (not NaN)
|
|
return not (np.isnan(self.n_x) or np.isnan(self.n_y))
|
|
|
|
|
|
class CompletedRasterGridElem(BaseModel):
|
|
request: RasterGridRequest
|
|
result: ScanResult
|
|
centre_of_mass: CenterOfMassModel | None
|
|
|
|
|
|
class CompletedRasterGrid(BaseModel):
|
|
r: List[CompletedRasterGridElem]
|
|
|
|
|
|
class RasterPayloadModel(BaseModel):
|
|
request: RasterGridRequest
|
|
result: ScanResult
|
|
sample_id: int
|
|
attach_image: bool = True
|
|
centre_of_mass: CenterOfMassModel | None = None
|
|
raster_score: List[float | None] | None = None
|
|
center_pxl: Coordinate | None
|
|
start_pxl: Coordinate
|
|
cell_size_pxl: Annotated[Coordinate, AfterValidator(positive_coords)]
|
|
beam_mark_pxl: tuple[float, float]
|
|
beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)]
|