feat: ship the grid-scan decision to AareDB (GridScanDecision on RasterPayloadModel) #29

Merged
perl_d merged 3 commits from feat/gridscan-decision into main 2026-09-03 12:27:25 +02:00
5 changed files with 159 additions and 41 deletions
+1 -1
View File
@@ -10,7 +10,7 @@ readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"pyyaml",
"jfjoch-client",
"jfjoch-client>=1.0.0rc166",
"pydantic",
"numpy",
"scipy",
+50 -40
View File
@@ -40,7 +40,8 @@ the picture if the database wants one.
from __future__ import annotations
import io
from typing import Any, Optional, Sequence
from collections.abc import Sequence
from typing import Any
import numpy as np
from pydantic import BaseModel, Field
@@ -68,7 +69,7 @@ except ImportError: # pragma: no cover
z: float = 0.0
__all__ = ["analyse", "GridScanResult", "Centre", "Size", "Counts", "Thresholds"]
__all__ = ["Centre", "Counts", "GridScanResult", "Size", "Thresholds", "analyse"]
#: DAQStatusModel, or anything else carrying the beam: its ``.geom`` is a
#: SampleGeometryModel, and a ScanResultPayloadModel has ``beam_size_mm`` itself.
@@ -98,10 +99,10 @@ class Centre(BaseModel):
nx: float
ny: float
image_number: int = Field(description="Nearest collected image, for addressing")
x_um: Optional[float] = Field(
x_um: float | None = Field(
None, description="Signed offset from the centre of cell (0,0), along the grid's own axes"
)
y_um: Optional[float] = None
y_um: float | None = None
def offset_mm(self) -> Coordinate:
"""Offset in mm from the centre of cell (0,0), as an aarecommon Coordinate.
@@ -123,17 +124,17 @@ class Size(BaseModel):
narrower than the beam still reports a size instead of nothing.
"""
x_um: Optional[float] = None
y_um: Optional[float] = None
z_um: Optional[float] = Field(None, description="From an orthogonal scan, if passed")
x_um_deconv: Optional[float] = None
y_um_deconv: Optional[float] = None
z_um_deconv: Optional[float] = None
beam_x_um: Optional[float] = Field(None, description="Beam FWHM used to deconvolve x")
beam_y_um: Optional[float] = Field(None, description="Beam FWHM used to deconvolve y")
area_um2: Optional[float] = None
equiv_diameter_um: Optional[float] = None
volume_pl: Optional[float] = Field(None, description="Ellipsoid; needs all three axes")
x_um: float | None = None
y_um: float | None = None
z_um: float | None = Field(None, description="From an orthogonal scan, if passed")
x_um_deconv: float | None = None
y_um_deconv: float | None = None
z_um_deconv: float | None = None
beam_x_um: float | None = Field(None, description="Beam FWHM used to deconvolve x")
beam_y_um: float | None = Field(None, description="Beam FWHM used to deconvolve y")
area_um2: float | None = None
equiv_diameter_um: float | None = None
volume_pl: float | None = Field(None, description="Ellipsoid; needs all three axes")
volume_deconvolved: bool = Field(
False, description="True if volume_pl used the beam-removed axes"
)
@@ -168,44 +169,50 @@ class Counts(BaseModel):
n_noise: int = Field(description="Spots present but neither protein nor ice")
n_blank: int
peak_protein_spots: int
mean_protein_in_contour: Optional[float] = None
best_resolution_A: Optional[float] = Field(None, description="Best res inside the contour")
median_bkg: Optional[float] = None
ice_in_contour: Optional[float] = Field(
mean_protein_in_contour: float | None = None
best_resolution_A: float | None = Field(None, description="Best res inside the contour")
median_bkg: float | None = None
ice_in_contour: float | None = Field(
None, description="Fraction of frames inside the contour showing ice"
)
n_fragments: int = Field(
0, description="Separate blobs at the 50 % level; >1 means more than one region"
)
ice_fraction: float = Field(description="Fraction of frames showing ice")
unit_cell_agreement: Optional[float] = Field(
unit_cell_agreement: float | None = Field(
None, description="Fraction of indexed frames in the contour sharing one cell"
)
class GridScanResult(BaseModel):
found: bool
reason: Optional[str] = Field(None, description="Why not, when found is False")
file_prefix: Optional[str] = None
reason: str | None = Field(None, description="Why not, when found is False")
file_prefix: str | None = None
n_fast: int = 0
n_slow: int = 0
channel: str = Field("", description="Which spot channel built the map")
centre: Optional[Centre] = None
size: Optional[Size] = None
counts: Optional[Counts] = None
unit_cell: Optional[list[float]] = Field(
centre: Centre | None = None
size: Size | None = None
counts: Counts | None = None
unit_cell: list[float] | None = Field(
None, description="Median cell of the contour, when check_unit_cell is on"
)
warnings: list[str] = Field(
default_factory=list, description="Reasons to distrust a found=True result"
)
jpeg: Optional[bytes] = Field(
contour_cells: list[tuple[int, int]] | None = Field(
None,
description="(nx, ny) of every collected cell inside the chosen 50 % "
"blob -- the decision's own footprint, for drawing it over the sample "
"image without recomputing the map",
)
jpeg: bytes | None = Field(
None,
exclude=True,
repr=False,
description="The picture, in whichever single style was asked for",
)
jpeg_cells: Optional[bytes] = Field(
jpeg_cells: bytes | None = Field(
None,
exclude=True,
repr=False,
@@ -225,10 +232,10 @@ class Thresholds(BaseModel):
"the bkg-only gate then rejects the true crystal",
)
ice_spots: int = Field(5, description="spots_ice at or above this counts as ice")
ice_ring: Optional[float] = Field(
ice_ring: float | None = Field(
None, description="Also call ice if scan_result 'ice' exceeds this"
)
beam_um: Optional[float] = Field(
beam_um: float | None = Field(
None,
description="Beam FWHM in um, used only when no DAQ status is passed. "
"A single number is taken as a square beam",
@@ -284,7 +291,7 @@ def _upsample(a: np.ndarray, f: int) -> np.ndarray:
return np.asarray(im.resize((nx * fx, ny * fy), Image.BICUBIC), dtype=float)
def _norm(a: np.ndarray, floor: float = 0.0, ref: Optional[np.ndarray] = None) -> np.ndarray:
def _norm(a: np.ndarray, floor: float = 0.0, ref: np.ndarray | None = None) -> np.ndarray:
"""Normalise for colour against the channel's own range, but never below floor.
``ref`` supplies the range: the bounds come from the frames actually measured,
@@ -372,7 +379,7 @@ _BEAM_PATHS = (
)
def _beam_um(obj: Optional[BeamSource]) -> tuple[Optional[float], Optional[float]]:
def _beam_um(obj: BeamSource | None) -> tuple[float | None, float | None]:
"""Beam FWHM (x, y) in micrometres, dug out of a DAQ status or scan payload.
The beam is a Coordinate, not one number: on a beamline with an 80 x 20 um
@@ -406,7 +413,7 @@ def _beam_um(obj: Optional[BeamSource]) -> tuple[Optional[float], Optional[float
return None, None
def _deconvolve(fwhm: Optional[float], beam: Optional[float]) -> Optional[float]:
def _deconvolve(fwhm: float | None, beam: float | None) -> float | None:
"""Quadrature removal of the beam. None when the beam swallows the feature."""
if fwhm is None or beam is None:
return fwhm
@@ -489,11 +496,11 @@ def _classify(F: dict, th: Thresholds) -> Counts:
def analyse(
scan_result: ScanResult,
grid_scan: Optional[GridScan] = None,
grid_scan: GridScan | None = None,
*,
daq: Optional[BeamSource] = None,
thresholds: Optional[Thresholds] = None,
z_um: Optional[float] = None,
daq: BeamSource | None = None,
thresholds: Thresholds | None = None,
z_um: float | None = None,
jpeg: bool = True,
) -> GridScanResult:
"""Analyse one grid scan.
@@ -610,6 +617,9 @@ def analyse(
if keep.shape == inside.shape:
inside = inside & keep
counts.n_fragments = n_frag
if cell.any():
yy_c, xx_c = np.nonzero(cell)
out.contour_cells = [(int(x), int(y)) for y, x in zip(yy_c, xx_c)]
if line:
# One position wide: the 50 % crossings either side of the peak. The outer
@@ -640,7 +650,7 @@ def analyse(
w_cells = ((xx.max() - xx.min() + 1) / mx * nx_n, (yy.max() - yy.min() + 1) / my * ny_n)
# the nearest image actually collected, for the DAQ to address
iy, ix = int(round(np.clip(gy, 0, ny_n - 1))), int(round(np.clip(gx, 0, nx_n - 1)))
iy, ix = round(np.clip(gy, 0, ny_n - 1)), round(np.clip(gx, 0, nx_n - 1))
num = int(F["number"][iy, ix])
if num < 0: # that cell was never collected
yy2, xx2 = np.nonzero(F["number"] >= 0)
@@ -742,7 +752,7 @@ def _edge(mask: np.ndarray, width: int = 1) -> np.ndarray:
return mask & ~core
def _pictures(F, inside, centre, res: "GridScanResult", th: Thresholds, obj_level: float) -> None:
def _pictures(F, inside, centre, res: GridScanResult, th: Thresholds, obj_level: float) -> None:
"""Fill res.jpeg, and res.jpeg_cells too when both styles were asked for."""
want = ("smooth", "cells") if th.jpeg_style == "both" else (th.jpeg_style,)
res.jpeg = _render(F, inside, centre, res, th, obj_level, want[0])
@@ -751,7 +761,7 @@ def _pictures(F, inside, centre, res: "GridScanResult", th: Thresholds, obj_leve
def _render(
F, inside, centre, res: "GridScanResult", th: Thresholds, obj_level: float, style: str
F, inside, centre, res: GridScanResult, th: Thresholds, obj_level: float, style: str
) -> bytes:
"""Map, contours, crosshair and a one-line caption, as JPEG.
@@ -0,0 +1,36 @@
"""The grid-scan centring decision, as taken at the beamline.
The DAQ runs the analysis (``aarecommon.math.jfjoch_gridscan_union``) to pick
the collection position; this wraps its result verbatim with provenance and
rides on ``RasterPayloadModel.decision`` into AareDB, whose Results viewer
displays THE decision instead of recomputing one -- two implementations of the
same analysis drifting apart is how a UI ends up contradicting the beamline.
The pictures are excluded from the result's serialization by the model itself;
the DAQ ships them through the ordinary image pipeline.
"""
from typing import Any
from pydantic import BaseModel, Field
from aarecommon.math.jfjoch_gridscan_union import GridScanResult
GRIDSCAN_DECISION_CONTRACT = 1
class GridScanDecision(BaseModel):
contract: int = GRIDSCAN_DECISION_CONTRACT
algorithm: str = Field(
description="Which analysis took the decision, e.g. a "
"GridscanAnalysisMode value such as 'jfjoch_gridscan_union'"
)
algorithm_version: str | None = Field(
None, description="aarecommon version (or commit) the DAQ ran"
)
thresholds: dict[str, Any] | None = Field(
None,
description="Thresholds.model_dump() used at decision time, so a "
"stored decision states its own tuning",
)
result: GridScanResult
+4
View File
@@ -6,6 +6,7 @@ from pydantic import AfterValidator, BaseModel, Field
from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords
from aarecommon.math.sample_geometry import SampleGeometryModel
from aarecommon.models.gridscan_decision import GridScanDecision
class RasterGridRequest(BaseModel):
@@ -92,3 +93,6 @@ class RasterPayloadModel(BaseModel):
cell_size_pxl: Annotated[Coordinate, AfterValidator(positive_coords)]
beam_mark_pxl: tuple[float, float]
beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)]
# The centring decision as taken at the beamline (None from DAQs that
# have not adopted it yet); AareDB stores and displays it verbatim.
decision: GridScanDecision | None = None
+68
View File
@@ -0,0 +1,68 @@
"""GridScanDecision wire contract: wraps the analysis result verbatim."""
import json
import numpy as np
from aarecommon.math.jfjoch_gridscan_union import Thresholds, analyse
from aarecommon.models.gridscan_decision import GridScanDecision
from aarecommon.models.raster_grid import RasterPayloadModel
def _synthetic_scan(ny_n=20, nx_n=20):
"""A gaussian protein blob on a flat loop, as plain dicts (jfjoch-free)."""
rng = np.random.default_rng(0)
yy, xx = np.mgrid[0:ny_n, 0:nx_n]
blob = 80 * np.exp(-(((xx - 7) ** 2 + (yy - 12) ** 2) / 8.0))
images = [
{
"number": int(y * nx_n + x),
"nx": int(x),
"ny": int(y),
"spots": int(blob[y, x] + 10 + rng.integers(0, 3)),
"spots_low_res": int(4 + rng.integers(0, 2)),
"spots_ice": 0,
"bkg": float(30 + 40 * np.exp(-(((x - 9) ** 2 + (y - 10) ** 2) / 90.0))),
"res": 2.0,
}
for y in range(ny_n)
for x in range(nx_n)
]
return {"file_prefix": "decision-test", "images": images}
def test_decision_round_trips_without_pictures(request):
th = Thresholds.model_validate({})
r = analyse(_synthetic_scan(), {"step_x_um": 10.0, "step_y_um": 10.0}, thresholds=th)
assert r.found
assert r.contour_cells, "the chosen 50 % blob must ship its cells"
assert all(0 <= x < r.n_fast and 0 <= y < r.n_slow for x, y in r.contour_cells)
# the centre lies within the contour's bounding box
assert r.centre is not None
xs = [x for x, _ in r.contour_cells]
ys = [y for _, y in r.contour_cells]
assert min(xs) - 1 <= r.centre.nx <= max(xs) + 1
assert min(ys) - 1 <= r.centre.ny <= max(ys) + 1
d = GridScanDecision(
algorithm="jfjoch_gridscan_union",
algorithm_version="test",
thresholds=th.model_dump(),
result=r,
)
wire = json.loads(d.model_dump_json())
assert wire["contract"] == 1
assert "jpeg" not in wire["result"], "pictures never ride the JSON"
assert wire["result"]["centre"]["image_number"] == r.centre.image_number
assert wire["thresholds"]["object_union_spots"] is True
back = GridScanDecision.model_validate(wire)
assert back.result.centre is not None
assert back.result.contour_cells is not None
assert back.result.centre.nx == r.centre.nx
assert [tuple(c) for c in back.result.contour_cells] == r.contour_cells
def test_payload_decision_is_optional():
fields = RasterPayloadModel.model_fields
assert "decision" in fields
assert fields["decision"].default is None