From 9cad8ec4fb72e8a7cc701d31bb90929e72f7c0c8 Mon Sep 17 00:00:00 2001 From: perl_d Date: Thu, 3 Sep 2026 09:20:55 +0200 Subject: [PATCH] feat: add Meitian's gridscan analysis util --- src/aarecommon/config/config_models.py | 6 + src/aarecommon/math/jfjoch_gridscan_union.py | 892 +++++++++++++++++++ 2 files changed, 898 insertions(+) create mode 100644 src/aarecommon/config/config_models.py create mode 100644 src/aarecommon/math/jfjoch_gridscan_union.py diff --git a/src/aarecommon/config/config_models.py b/src/aarecommon/config/config_models.py new file mode 100644 index 0000000..2f0a5b3 --- /dev/null +++ b/src/aarecommon/config/config_models.py @@ -0,0 +1,6 @@ +from enum import StrEnum, auto + + +class GridscanAnalysisMode(StrEnum): + FindXtal = auto() + JfjochGridscanUnion = auto() diff --git a/src/aarecommon/math/jfjoch_gridscan_union.py b/src/aarecommon/math/jfjoch_gridscan_union.py new file mode 100644 index 0000000..1dfadc4 --- /dev/null +++ b/src/aarecommon/math/jfjoch_gridscan_union.py @@ -0,0 +1,892 @@ +"""Grid-scan decision from a Jungfraujoch ``scan_result``. + +Turns the per-image spot statistics of a grid scan into the position to collect +from (in grid coordinates), a crystal size, a handful of numbers describing the +scan, and a JPEG for the database. + + from jfjoch_gridscan import analyse + r = analyse(scan_result, dataset_settings.grid_scan, daq=daq_status) + if r.found: + move_to(r.centre.nx, r.centre.ny) # grid coordinates + daq_status.crystal_size = CrystalSize(**r.size.crystal_size()) + db.store(r.jpeg, r.model_dump_json()) + +The beam size comes from the DAQ status (or from a ScanResultPayloadModel, which +carries its own ``beam_size_mm``); it is the one quantity the Jungfraujoch files +never record, and without it every size is an upper bound. + +The protein channel is ``spots - spots_ice - spots_low_res``: spots that are +neither inside an ice band nor above the low-resolution cut, so what remains is +diffraction below the cut -- the requirement that a crystal show spots below +5 A. Nothing derived from indexing is used: indexing fires on ice often enough +that a map built from it carries that error in. + +The centre is the *area centroid of the region enclosed by the 50 % contour*, +not an intensity-weighted centre of the whole map. A weighted mean is dragged +by the weak signal trailing across the rest of the loop and lands on the +shoulder of the contour rather than inside it. + +Input may be jfjoch-client models or plain dicts. Needs numpy, pillow, pydantic; +no scipy, no matplotlib. + +Measured over 194 rasters from one session, 4x3 to 53x48 points: + decision only (jpeg=False) median 3 ms, max 26 ms + plus the grid-cell picture median 25 ms + plus the smoothed picture median 81 ms, max 143 ms +The decision is what has to be fast; pass jpeg=False to move, then call again for +the picture if the database wants one. +""" + +from __future__ import annotations + +import io +from typing import Any, Optional, Sequence + +import numpy as np +from pydantic import BaseModel, Field + +# The structures this exchanges with AareDAQ and Jungfraujoch. Verified against +# jfjoch-client 1.0.0-rc.165. The fallbacks let the module import and its self-test +# run outside AareDAQ; inside it the real classes are always the ones used. +try: + from jfjoch_client.models.grid_scan import GridScan + from jfjoch_client.models.scan_result import ScanResult + from jfjoch_client.models.scan_result_images_inner import ScanResultImagesInner + from jfjoch_client.models.unit_cell import UnitCell +except ImportError: # pragma: no cover + GridScan = ScanResult = ScanResultImagesInner = UnitCell = Any + +try: + from aarecommon.math.coordinate import Coordinate + from aarecommon.math.sample_geometry import SampleGeometryModel +except ImportError: # pragma: no cover + SampleGeometryModel = Any + + class Coordinate(BaseModel): # same shape as aarecommon's + x: float = 0.0 + y: float = 0.0 + z: float = 0.0 + + +__all__ = ["analyse", "GridScanResult", "Centre", "Size", "Counts", "Thresholds"] + +#: DAQStatusModel, or anything else carrying the beam: its ``.geom`` is a +#: SampleGeometryModel, and a ScanResultPayloadModel has ``beam_size_mm`` itself. +BeamSource = Any + +# Flat grey object, orange protein, blue ice, mixed in proportion where both. +_PLATE = np.array([0.953, 0.953, 0.941]) +_GREY = np.array([0.706, 0.712, 0.700]) +_ORANGE = np.array([0.918, 0.478, 0.153]) +_BLUE = np.array([0.176, 0.470, 0.757]) + +_COLOUR_FLOOR = 10.0 # spots the top of the colour scale may not fall below, so +# that an empty scan renders empty instead of saturated +_LEVEL = 0.50 # the contour that defines the crystal +_LEVELS = (0.25, 0.50, 0.75) # drawn on the smoothed picture, as in the reports +_OBJ_LINE = (90, 95, 92) # #5A5F5C +_ORANGE_LINE = (138, 61, 0) # #8A3D00 +_BLUE_LINE = (18, 62, 99) # #123E63 + + +# --------------------------------------------------------------------------- # +# output +# --------------------------------------------------------------------------- # +class Centre(BaseModel): + """Where to collect. Grid coordinates are fractional and 0-based.""" + + nx: float + ny: float + image_number: int = Field(description="Nearest collected image, for addressing") + x_um: Optional[float] = Field( + None, description="Signed offset from the centre of cell (0,0), along the grid's own axes" + ) + y_um: Optional[float] = None + + def offset_mm(self) -> Coordinate: + """Offset in mm from the centre of cell (0,0), as an aarecommon Coordinate. + + In the beam plane and signed, so it adds to wherever the scan started: + ``geom.translate_smargon(r.centre.offset_mm())`` gives the goniometer target. + Zero on an axis the scan did not sample. + """ + return Coordinate(x=(self.x_um or 0.0) / 1000.0, y=(self.y_um or 0.0) / 1000.0, z=0.0) + + +class Size(BaseModel): + """Extent of the 50 % contour. + + ``x_um``/``y_um``/``z_um`` are what was measured, beam included, and are always + present for an axis that was scanned. The ``*_deconv`` fields remove the beam in + quadrature and go None per axis when the feature is not resolved beyond it -- a + one-cell feature on a 20 um grid with a 20 um beam. Keeping both means a crystal + 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") + volume_deconvolved: bool = Field( + False, description="True if volume_pl used the beam-removed axes" + ) + + def crystal_size(self) -> dict: + """``{"x": .., "y": .., "z": ..}`` in um: ``CrystalSize(**r.size.crystal_size())``. + + A dict rather than a CrystalSize because that class lives in the AareDAQ + module this file does not import -- one keyword splat at the call site costs + less than a wrong import path here. + + Beam-removed where the beam resolved the axis, measured where it did not, + and 0.0 for an axis no scan sampled, since CrystalSize has no None. + """ + return { + k: float(d if d is not None else (m if m is not None else 0.0)) + for k, d, m in ( + ("x", self.x_um_deconv, self.x_um), + ("y", self.y_um_deconv, self.y_um), + ("z", self.z_um_deconv, self.z_um), + ) + } + + +class Counts(BaseModel): + """Frame classification and scan-level numbers.""" + + n_images: int + n_protein: int = Field(description="Protein spots over threshold, no ice") + n_protein_ice: int + n_ice: int + 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( + 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( + 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 + 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( + 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( + None, + exclude=True, + repr=False, + description="The picture, in whichever single style was asked for", + ) + jpeg_cells: Optional[bytes] = Field( + None, + exclude=True, + repr=False, + description="The grid-cell picture, when jpeg_style is 'both'", + ) + + +class Thresholds(BaseModel): + """Everything tunable. Defaults come from ~250 loops of PSI beamline data.""" + + min_spots: int = Field(10, description="Protein spots needed before a contour means anything") + object_union_spots: bool = Field( + True, + description="Treat cells with >= min_spots protein as part of the object " + "even where the bkg proxy disagrees. A crystal on a thin tip has " + "low bkg while a metal pin foot dominates the percentile band; " + "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( + None, description="Also call ice if scan_result 'ice' exceeds this" + ) + beam_um: Optional[float] = Field( + None, + description="Beam FWHM in um, used only when no DAQ status is passed. " + "A single number is taken as a square beam", + ) + z_beam_axis: str = Field( + "y", + description="Which beam dimension deconvolves the third axis, since the " + "partner scan that measured it is vertical here", + ) + check_unit_cell: bool = Field( + False, + description="Cross-check the contour against per-image unit cells. " + "Off by default because it needs indexing; it never " + "touches the map or the centre, only reports agreement", + ) + cell_tol: float = Field(0.06, description="Relative tolerance on a, b, c for agreement") + cell_agree_min: float = Field( + 0.20, + description="Warn below this agreement. Measured: ice-only loop 0.01, " + "real crystals 0.33-0.78", + ) + jpeg_quality: int = 82 + jpeg_width: int = 560 + jpeg_style: str = Field( + "both", + description='"smooth" reproduces the report maps -- bicubic ' + 'composite with contours at 0.25/0.50/0.75. "cells" ' + "shows the raw grid, one block per position, where every " + 'pixel is a frame that was actually measured. "both" ' + "fills jpeg with the smooth one and jpeg_cells with the " + 'other. "cells" is about 5x cheaper to draw', + ) + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # +def _get(o: Any, k: str, d=None): + """Field access that works for pydantic models and plain dicts alike.""" + v = o.get(k, d) if isinstance(o, dict) else getattr(o, k, d) + return d if v is None else v + + +def _upsample(a: np.ndarray, f: int) -> np.ndarray: + """Bicubic upsample through pillow -- keeps scipy out of the dependency list.""" + from PIL import Image + + ny, nx = a.shape + if f <= 1 or (ny == 1 and nx == 1): + return a + fy, fx = (f if ny > 1 else 1), (f if nx > 1 else 1) + im = Image.fromarray(np.ascontiguousarray(a, dtype=np.float32), mode="F") + 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: + """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, + not from interpolated pixels, and taking them off the small array keeps three + percentile passes off the full-size image. + """ + r = a if ref is None else ref + lo, hi = np.nanpercentile(r, 2), np.nanpercentile(r, 98) + hi = max(hi, lo + floor) + if not np.isfinite(hi) or hi <= lo: + return np.zeros_like(a) + return np.clip((a - lo) / (hi - lo), 0, 1) ** 0.65 + + +def _crossings(pos: np.ndarray, v: np.ndarray, level: float) -> list: + """Where a profile crosses a level -- the 1D contour.""" + s = np.sign(v - level) + j = np.where(s[:-1] * s[1:] < 0)[0] + a, b = v[j], v[j + 1] + return list(pos[j] + (level - a) * (pos[j + 1] - pos[j]) / (b - a)) + + +def _component(mask: np.ndarray, seed) -> np.ndarray: + """The 4-connected blob of ``mask`` containing ``seed``, by iterated dilation. + + Run at grid resolution, where the array is a few hundred cells and this costs + nothing. A crystal is one connected region: without this the "contour" of an + ice-only loop is a scatter of specks and its size is their bounding box. + """ + cur = np.zeros_like(mask) + cur[seed] = True + while True: + g = cur.copy() + g[1:] |= cur[:-1] + g[:-1] |= cur[1:] + g[:, 1:] |= cur[:, :-1] + g[:, :-1] |= cur[:, 1:] + g &= mask + if g.sum() == cur.sum(): + return cur + cur = g + + +def _blobs(mask: np.ndarray, limit: int = 64): + """Every separate blob the level encloses, largest-signal choice left to caller.""" + rest, out = mask.copy(), [] + while rest.any() and len(out) < limit: + yy, xx = np.nonzero(rest) + c = _component(rest, (yy[0], xx[0])) + rest &= ~c + out.append(c) + return out + + +def _cell_agreement(cells: np.ndarray, tol: float) -> tuple: + """Median unit cell of a set of frames, and the fraction agreeing with it. + + A crystal gives one cell repeated across its frames; ice gives cells that + scatter. This is the only test found to separate the two -- per-frame spot + counts do not -- but it needs indexing, so it is opt-in and advisory only. + Measured over whole scans: ice-only loop 0.01, real crystals 0.33-0.78. + """ + v = np.array( + [ + [ + c.a if not isinstance(c, dict) else c["a"], + c.b if not isinstance(c, dict) else c["b"], + c.c if not isinstance(c, dict) else c["c"], + ] + for c in np.ravel(cells) + if c is not None and (c["a"] if isinstance(c, dict) else c.a) + ], + dtype=float, + ) + if len(v) < 3: + return None, None + med = np.median(v, axis=0) + ok = np.all(np.abs(v - med) <= tol * np.maximum(med, 1e-9), axis=1) + return [round(float(x), 2) for x in med], round(float(ok.mean()), 3) + + +_BEAM_PATHS = ( + ("geom.beam_size_mm", 1000.0), # DAQStatusModel.geom, a SampleGeometryModel + ("beam_size_mm", 1000.0), # ScanResultPayloadModel, or a bare geometry +) + + +def _beam_um(obj: Optional[BeamSource]) -> tuple[Optional[float], Optional[float]]: + """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 + beam a single figure would over-correct one axis and under-correct the other. + A scalar is accepted and taken as square. + + Read from SampleGeometryModel.beam_size_mm, reached either as + DAQStatusModel.geom or directly on a ScanResultPayloadModel. + """ + if obj is None: + return None, None + for path, scale in _BEAM_PATHS: + v = obj + for part in path.split("."): + v = _get(v, part) + if v is None: + break + if v is None: + continue + x = _get(v, "x") + if x is not None: # a Coordinate + bx, by = float(x) * scale, float(_get(v, "y", x)) * scale + elif isinstance(v, (int, float)): # one number: square beam + bx = by = float(v) * scale + elif isinstance(v, (tuple, list)) and len(v) >= 2: + bx, by = float(v[0]) * scale, float(v[1]) * scale + else: + continue + if bx > 0 and by > 0: # a zero beam is no beam, not a point one + return bx, by + return None, None + + +def _deconvolve(fwhm: Optional[float], beam: Optional[float]) -> Optional[float]: + """Quadrature removal of the beam. None when the beam swallows the feature.""" + if fwhm is None or beam is None: + return fwhm + d = fwhm * fwhm - beam * beam + return float(np.sqrt(d)) if d > 0 else None + + +# --------------------------------------------------------------------------- # +# analysis +# --------------------------------------------------------------------------- # +def _grids(images: Sequence[ScanResultImagesInner]): + """Scatter the per-image statistics onto [ny, nx] arrays. + + nx/ny come from the scan_result itself, so the snake ordering, the vertical + flag and the step signs are already applied upstream and are not redone here. + """ + nx = np.array([_get(i, "nx", -1) for i in images], dtype=np.int64) + ny = np.array([_get(i, "ny", -1) for i in images], dtype=np.int64) + if (nx < 0).any() or (ny < 0).any(): + return None + shape = (int(ny.max()) + 1, int(nx.max()) + 1) + + def col(key, dtype=float): + return np.array([_get(i, key, np.nan) for i in images], dtype=dtype) + + spots, low, ice = col("spots"), col("spots_low_res"), col("spots_ice") + # No total spot count: fall back to spots above the low-resolution cut. That + # channel is ice-free on the same argument (the lowest hexagonal ice ring is + # 3.895 A) and tracks the real one at r = 0.79-0.98 on our data. + if np.isnan(spots).all(): + protein, channel = np.nan_to_num(low), "spots_low_res (no total count)" + else: + protein = np.nan_to_num(spots) - np.nan_to_num(ice) - np.nan_to_num(low) + channel = "spots - spots_ice - spots_low_res" + protein = np.clip(protein, 0, None) + + F = {} + for name, v in ( + ("protein", protein), + ("ice", np.nan_to_num(ice)), + ("bkg", col("bkg")), + ("res", col("res")), + ("ring", col("ice")), + ("spots", np.nan_to_num(spots)), + ): + g = np.full(shape, np.nan) + g[ny, nx] = v + F[name] = g + uc = np.empty(shape, dtype=object) + for k, im in enumerate(images): + uc[ny[k], nx[k]] = _get(im, "uc") + F["uc"] = uc + F["number"] = np.full(shape, -1, dtype=np.int64) + F["number"][ny, nx] = col("number", np.int64) + return F, channel + + +def _classify(F: dict, th: Thresholds) -> Counts: + """Protein / protein+ice / ice / noise / blank, per frame.""" + p, i_, s = F["protein"], F["ice"], F["spots"] + ok = np.isfinite(p) + prot = ok & (p >= th.min_spots) + ice = ok & (i_ >= th.ice_spots) + if th.ice_ring is not None: + ice |= ok & np.isfinite(F["ring"]) & (F["ring"] >= th.ice_ring) + any_spot = ok & (np.nan_to_num(s) + p + i_ > 0) + n = int(ok.sum()) + return Counts( + n_images=n, + n_protein=int((prot & ~ice).sum()), + n_protein_ice=int((prot & ice).sum()), + n_ice=int((ice & ~prot).sum()), + n_noise=int((any_spot & ~prot & ~ice).sum()), + n_blank=int((ok & ~any_spot).sum()), + peak_protein_spots=int(np.nanmax(p)) if n else 0, + median_bkg=float(np.nanmedian(F["bkg"])) if np.isfinite(F["bkg"]).any() else None, + ice_fraction=round(float(ice.sum() / n), 3) if n else 0.0, + ) + + +def analyse( + scan_result: ScanResult, + grid_scan: Optional[GridScan] = None, + *, + daq: Optional[BeamSource] = None, + thresholds: Optional[Thresholds] = None, + z_um: Optional[float] = None, + jpeg: bool = True, +) -> GridScanResult: + """Analyse one grid scan. + + ``scan_result``: a jfjoch_client ScanResult, or a ScanResultPayloadModel + wrapping one in ``.result`` -- in which case its ``beam_size_mm`` is used and no + ``daq`` is needed. + + ``grid_scan``: the jfjoch_client GridScan from ``DatasetSettings.grid_scan``. + Only ``step_x_um`` and ``step_y_um`` are read; ``n_fast``, ``vertical`` and + ``snake`` describe the acquisition order, which ``nx``/``ny`` already reflect. + Without it the result is in grid steps and every micrometre field is None. + + ``daq``: a DAQStatusModel, whose ``geom`` is a SampleGeometryModel carrying + ``beam_size_mm``. A SampleGeometryModel or ScanResultPayloadModel is accepted + directly too. The beam is the one quantity the Jungfraujoch files never record, + and without it a size is only an upper bound. + + ``z_um`` is the in-plane size measured by an orthogonal scan of the same loop + -- pass ``line.size.y_um`` from a partner line scan to get a volume, since one + scan cannot know the third axis. + """ + th = thresholds or Thresholds() + # A ScanResultPayloadModel wraps the result and carries the beam with it. + payload = scan_result + inner = _get(scan_result, "result") + if inner is not None and _get(inner, "images") is not None: + scan_result = inner + bx, by = _beam_um(daq) + if bx is None: + bx, by = _beam_um(payload) + if bx is None and th.beam_um is not None: + bx = by = th.beam_um + images = list(_get(scan_result, "images", []) or []) + prefix = _get(scan_result, "file_prefix") + if not images: + return GridScanResult(found=False, reason="no images", file_prefix=prefix) + + built = _grids(images) + if built is None: + return GridScanResult( + found=False, file_prefix=prefix, reason="images carry no nx/ny; not a grid scan" + ) + F, channel = built + ny_n, nx_n = F["protein"].shape + counts = _classify(F, th) + # The step is signed -- negative means the grid runs right to left, or bottom to + # top. A size only wants the magnitude, but an offset that drops the sign sends + # the goniometer the wrong way, so both are kept. + sx_s = float(_get(grid_scan, "step_x_um", 0.0)) or None + sy_s = float(_get(grid_scan, "step_y_um", 0.0)) or None + sx = abs(sx_s) if sx_s else None + sy = abs(sy_s) if sy_s else None + out = GridScanResult( + found=False, file_prefix=prefix, n_fast=nx_n, n_slow=ny_n, channel=channel, counts=counts + ) + + peak = float(np.nanmax(F["protein"])) + if not np.isfinite(peak) or peak < th.min_spots: + # The 50 % contour of a channel peaking at two spots is a contour of + # noise: say there is no crystal rather than centre on nothing. + out.reason = f"peak protein channel {peak:.0f} < {th.min_spots} spots" + if jpeg: + _pictures(F, None, None, out, th, 0.0) + return out + + # Object mask from the background estimate: this finds the loop, not the crystal. + bkg = F["bkg"] + if np.isfinite(bkg).any(): + lo, hi = np.nanpercentile(bkg, 5), np.nanpercentile(bkg, 95) + level = lo + 0.25 * (hi - lo) + else: + bkg, level = np.ones_like(F["protein"]), 0.0 + + line = nx_n == 1 or ny_n == 1 + factor = 24 if line else int(np.clip(round(512 / max(nx_n, ny_n)), 4, 24)) + P = _upsample(np.nan_to_num(F["protein"]), factor) + B = _upsample(np.nan_to_num(bkg, nan=float(np.nanmin(bkg))), factor) + obj_gate = B > level + if th.object_union_spots: + # Bragg spots are direct evidence of sample; bkg is only a proxy + obj_gate |= P >= th.min_spots + inside = (P >= _LEVEL * peak) & obj_gate + if not inside.any(): + out.reason = "no region above the 50 % level inside the object" + if jpeg: + _pictures(F, None, None, out, th, level) + return out + + my, mx = inside.shape + # Sample the contour back onto the collected cells, then keep only the blob + # holding the strongest frame: a crystal is one region, and the extent of a + # scatter of specks is not a size. + ci = np.clip(((np.arange(ny_n) + 0.5) / ny_n * my).astype(int), 0, my - 1) + cj = np.clip(((np.arange(nx_n) + 0.5) / nx_n * mx).astype(int), 0, mx - 1) + cell = inside[np.ix_(ci, cj)] + blobs = _blobs(cell) + n_frag = len(blobs) + if blobs: + # The blob with the most signal integrated over it, not the one holding the + # single hottest frame: a crystal is a region, and one hot frame at the edge + # of the scan should not outvote a larger, fully sampled region. + prot = np.nan_to_num(F["protein"]) + cell = max(blobs, key=lambda c: float(prot[c].sum())) + # Grown by one cell before it masks the fine grid: the cell mask samples + # cell centres, so using it undilated would shave up to half a step off + # every edge of the blob and shrink the size by a couple of cells. + grown = cell.copy() + grown[1:] |= cell[:-1] + grown[:-1] |= cell[1:] + grown[:, 1:] |= cell[:, :-1] + grown[:, :-1] |= cell[:, 1:] + keep = np.repeat(np.repeat(grown, inside.shape[0] // ny_n, 0), inside.shape[1] // nx_n, 1) + if keep.shape == inside.shape: + inside = inside & keep + counts.n_fragments = n_frag + + if line: + # One position wide: the 50 % crossings either side of the peak. The outer + # pair would span every blob on the line, not the one being centred on. + vert = nx_n == 1 + prof = P[:, 0] if vert else P[0, :] + base = B[:, 0] if vert else B[0, :] + pos = (np.arange(len(prof)) + 0.5) / len(prof) * (ny_n if vert else nx_n) + line_gate = base > level + if th.object_union_spots: + line_gate |= prof >= th.min_spots + v = np.where(line_gate, prof, 0.0) + c = _crossings(pos, v, _LEVEL * peak) + top = pos[int(np.argmax(v))] + lo = [x for x in c if x <= top] + hi = [x for x in c if x >= top] + if lo and hi: + a, b = max(lo), min(hi) + mid, span = (a + b) / 2, b - a + else: + mid, span = top, np.nan + gx, gy = (0.0, mid - 0.5) if vert else (mid - 0.5, 0.0) + w_cells = (np.nan, span) if vert else (span, np.nan) + else: + yy, xx = np.nonzero(inside) + gx = (xx.mean() + 0.5) / mx * nx_n - 0.5 + gy = (yy.mean() + 0.5) / my * ny_n - 0.5 + 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))) + num = int(F["number"][iy, ix]) + if num < 0: # that cell was never collected + yy2, xx2 = np.nonzero(F["number"] >= 0) + k = int(np.argmin((yy2 - gy) ** 2 + (xx2 - gx) ** 2)) + num = int(F["number"][yy2[k], xx2[k]]) + + # A 50 % crossing can land half a step outside the first or last cell centre. + # Clamp so the DAQ is never sent off the scanned area, and say when the contour + # reaches a border -- that means the crystal was not fully covered. + gx = float(np.clip(gx, 0, nx_n - 1)) + gy = float(np.clip(gy, 0, ny_n - 1)) + if cell[0, :].any() or cell[-1, :].any() or cell[:, 0].any() or cell[:, -1].any(): + out.warnings.append("contour touches the scan edge; crystal may extend beyond it") + + out.centre = Centre( + nx=round(float(gx), 3), + ny=round(float(gy), 3), + image_number=num, + x_um=round(gx * sx_s, 2) if sx_s else None, + y_um=round(gy * sy_s, 2) if sy_s else None, + ) + + def um(cells, step): + return float(cells * step) if (step and np.isfinite(cells)) else None + + mx_, my_, mz_ = um(w_cells[0], sx), um(w_cells[1], sy), z_um + bz = by if th.z_beam_axis == "y" else bx + dx, dy, dz = (_deconvolve(mx_, bx), _deconvolve(my_, by), _deconvolve(mz_, bz)) + area = ( + inside.sum() / (mx * my) * (nx_n * sx) * (ny_n * sy) if (sx and sy and not line) else None + ) + # Prefer the beam-removed axes for the volume, but fall back to the measured + # ones rather than reporting no volume at all. + trio = (dx, dy, dz) if None not in (dx, dy, dz) else (mx_, my_, mz_) + vol = np.pi / 6 * trio[0] * trio[1] * trio[2] / 1000 if None not in trio else None + rnd = lambda v: None if v is None else round(v, 1) + out.size = Size( + x_um=rnd(mx_), + y_um=rnd(my_), + z_um=rnd(mz_), + x_um_deconv=rnd(dx), + y_um_deconv=rnd(dy), + z_um_deconv=rnd(dz), + beam_x_um=bx, + beam_y_um=by, + area_um2=rnd(area), + equiv_diameter_um=(round(2 * float(np.sqrt(area / np.pi)), 1) if area else None), + volume_pl=rnd(vol), + volume_deconvolved=None not in (dx, dy, dz), + ) + + # what sits under the contour + if cell.any(): + counts.mean_protein_in_contour = round(float(np.nanmean(F["protein"][cell])), 1) + r = F["res"][cell & np.isfinite(F["res"])] + if r.size: + counts.best_resolution_A = round(float(r.min()), 2) + counts.ice_in_contour = round(float(np.mean(F["ice"][cell] >= th.ice_spots)), 3) + if th.check_unit_cell: + out.unit_cell, counts.unit_cell_agreement = _cell_agreement(F["uc"][cell], th.cell_tol) + if ( + counts.unit_cell_agreement is not None + and counts.unit_cell_agreement < th.cell_agree_min + ): + out.warnings.append( + f"only {counts.unit_cell_agreement:.0%} of indexed frames in the " + "contour share a cell; may be ice rather than one crystal" + ) + + # No per-frame spot statistic separates protein Bragg from ice Bragg: on our + # data an ice-only loop produced a contour stronger than several real crystals, + # and every contour tested had 100 % of its frames over any ice-count bar. So + # ice is reported as a number, never used to reject, and the only test that + # discriminates is the optional unit-cell one above. + if n_frag > 3: + out.warnings.append( + f"{n_frag} separate regions at the 50 % level; centred on the strongest" + ) + frac = cell.sum() / max(int(np.isfinite(F["protein"]).sum()), 1) + if frac > 0.4: + out.warnings.append(f"contour covers {frac:.0%} of the scan; likely not a single crystal") + out.found = True + if jpeg: + _pictures(F, inside, (gx, gy), out, th, level) + return out + + +# --------------------------------------------------------------------------- # +# picture +# --------------------------------------------------------------------------- # +def _edge(mask: np.ndarray, width: int = 1) -> np.ndarray: + """The inner boundary of a mask -- a contour line, without matplotlib.""" + core = mask + for _ in range(width): + c = core.copy() + for ax, sh in ((0, 1), (0, -1), (1, 1), (1, -1)): + c &= np.roll(core, sh, ax) + core = c + return mask & ~core + + +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]) + if len(want) > 1: + res.jpeg_cells = _render(F, inside, centre, res, th, obj_level, want[1]) + + +def _render( + F, inside, centre, res: "GridScanResult", th: Thresholds, obj_level: float, style: str +) -> bytes: + """Map, contours, crosshair and a one-line caption, as JPEG. + + In "smooth" style this is the picture from the reports: the channels are + upsampled bicubically and the composite built there, so the map reads as a + continuous object rather than a mosaic of cells. + """ + from PIL import Image, ImageDraw, ImageFont + + ny_n, nx_n = F["protein"].shape + # Scale from the longer axis, or a 115-point line scan renders one column + # 560 px wide and 64000 px tall. A line is then widened into a strip. + scale = int(np.clip(th.jpeg_width / max(nx_n, ny_n, 1), 3, 40)) + sxp = scale if nx_n > 1 else max(scale, 80) + syp = scale if ny_n > 1 else max(scale, 80) + w, h = nx_n * sxp, ny_n * syp + smooth = style == "smooth" + + def grow(a): + a = np.nan_to_num(a, nan=float(np.nanmin(a)) if np.isfinite(a).any() else 0.0) + if not smooth: + return np.repeat(np.repeat(a, syp, 0), sxp, 1) + from PIL import Image as _I + + im = _I.fromarray(np.ascontiguousarray(a, dtype=np.float32), mode="F") + return np.asarray(im.resize((w, h), _I.BICUBIC), dtype=float) + + # Smooth needs the composite at pixel resolution because the contours are drawn + # there. Cells does not: every operation is elementwise, so it composites on the + # grid and expands afterwards, which is several times cheaper. + up_b, up_p, up_i = ( + (grow(F["bkg"]), grow(F["protein"]), grow(F["ice"])) + if smooth + else (np.nan_to_num(F["bkg"]), np.nan_to_num(F["protein"]), np.nan_to_num(F["ice"])) + ) + obj = _norm(up_b, ref=F["bkg"]) if np.isfinite(F["bkg"]).any() else np.ones_like(up_b) + p = _norm(up_p, _COLOUR_FLOOR, ref=F["protein"]) * obj + i_ = _norm(up_i, _COLOUR_FLOOR, ref=F["ice"]) * obj + + base = _PLATE * (1 - obj[..., None]) + _GREY * obj[..., None] + tot = np.clip(p + i_, 0, 1)[..., None] + mix = (_ORANGE * p[..., None] + _BLUE * i_[..., None]) / np.maximum(p + i_, 1e-9)[..., None] + rgb = np.clip(base * (1 - tot) + mix * tot, 0, 1) + if not smooth: + rgb = np.repeat(np.repeat(rgb, syp, 0), sxp, 1) + + if smooth: + # Contours on the levels the reports use, against each channel's own peak + # over the collected grid, and only inside the object. + yy, xx = np.mgrid[0:h, 0:w] + dash = ((xx + yy) % 8) < 4 + if np.isfinite(F["bkg"]).any(): + rgb[_edge(up_b >= obj_level, 2)] = np.array(_OBJ_LINE) / 255 + for chan, peak, colour, dashed in ( + (up_i, np.nanmax(F["ice"]), _BLUE_LINE, True), + (up_p, np.nanmax(F["protein"]), _ORANGE_LINE, False), + ): + if not np.isfinite(peak) or peak <= 0: + continue + for lv in _LEVELS: + e = _edge((chan >= lv * peak) & (obj > 0)) + rgb[e & dash if dashed else e] = np.array(colour) / 255 + elif inside is not None: + m = ( + np.asarray(Image.fromarray(inside.astype(np.uint8) * 255).resize((w, h), Image.NEAREST)) + > 127 + ) + rgb[_edge(m)] = 1.0 + + bar = 34 + canvas = Image.new("RGB", (w, h + bar), (255, 255, 255)) + canvas.paste(Image.fromarray((rgb * 255).astype(np.uint8), "RGB"), (0, 0)) + d = ImageDraw.Draw(canvas) + try: + font = ImageFont.truetype("DejaVuSans.ttf", 13) + except OSError: + font = ImageFont.load_default() + + if centre is not None: + cx, cy = (centre[0] + 0.5) * sxp, (centre[1] + 0.5) * syp + r = max(6, min(sxp, syp) // 2) + d.line([(cx - r, cy), (cx + r, cy)], fill=(15, 15, 20), width=2) + d.line([(cx, cy - r), (cx, cy + r)], fill=(15, 15, 20), width=2) + sz = res.size + um = f" {sz.x_um:.0f} x {sz.y_um:.0f} um" if sz and sz.x_um and sz.y_um else "" + cap = f"centre {centre[0]:.2f}, {centre[1]:.2f} #{res.centre.image_number}{um}" + else: + cap = f"no crystal contour - {res.reason}" + d.text((6, h + 9), cap, fill=(40, 40, 46), font=font) + + buf = io.BytesIO() + canvas.save(buf, "JPEG", quality=th.jpeg_quality) + return buf.getvalue() + + +# --------------------------------------------------------------------------- # +if __name__ == "__main__": + import time + + # Built from the real jfjoch-client classes, so this doubles as the + # integration check: what AareDAQ passes is what is exercised here. + rng = np.random.default_rng(0) + ny_n, nx_n = 30, 30 + yy, xx = np.mgrid[0:ny_n, 0:nx_n] + blob = 90 * np.exp(-(((xx - 11) ** 2 + (yy - 17) ** 2) / 12.0)) + images = [ + ScanResultImagesInner( + number=int(y * nx_n + x), + nx=int(x), + ny=int(y), + efficiency=1.0, + spots=int(blob[y, x] + 12 + rng.integers(0, 4)), + spots_low_res=int(6 + rng.integers(0, 3)), + spots_ice=int(6 * np.exp(-(((x - 24) ** 2 + (y - 8) ** 2) / 20.0))), + bkg=float(40 + 60 * np.exp(-(((x - 15) ** 2 + (y - 15) ** 2) / 200.0))), + res=float(1.8 + 3 / (1 + blob[y, x])), + ) + for y in range(ny_n) + for x in range(nx_n) + ] + scan = ScanResult(file_prefix="selftest", images=images) + grid = GridScan(n_fast=nx_n, step_x_um=10.0, step_y_um=-10.0, snake=True) + + class _Geom(BaseModel): # stands in for SampleGeometryModel + beam_size_mm: Coordinate + + t = time.perf_counter() + r = analyse(scan, grid, daq=_Geom(beam_size_mm=Coordinate(x=0.080, y=0.020)), z_um=70.0) + dt = (time.perf_counter() - t) * 1e3 + print(r.model_dump_json(indent=2)) + print("crystal_size:", r.size.crystal_size()) + print("offset_mm :", r.centre.offset_mm()) + print( + f"\ntrue centre 11.0, 17.0; step_y is negative, so offset y is too" + f" {dt:.1f} ms jpeg {len(r.jpeg) / 1024:.1f} kB" + ) + with open("selftest.jpg", "wb") as fh: + fh.write(r.jpeg)