216 lines
7.1 KiB
Python
216 lines
7.1 KiB
Python
"""
|
||
Synthetic raster scan data for no-beam / offline testing.
|
||
|
||
Entry point
|
||
-----------
|
||
generate_no_beam_scan_result(request, seed=None) -> ScanResult
|
||
|
||
The returned ScanResult is structurally identical to a real one, so all
|
||
find_xtal.py functions consume it without modification.
|
||
|
||
Pre-defined seeds
|
||
-----------------
|
||
Each seed fixes cluster geometry so results are reproducible. The six
|
||
scenarios below cover the main cases needed for unit-testing crystal-finding
|
||
algorithms.
|
||
|
||
Seed Scenario
|
||
---- --------
|
||
1 Single crystal, centred – baseline positive detection
|
||
2 Single crystal, off-centre – tests COM accuracy near an edge
|
||
3 Two well-separated crystals – multi-crystal detection
|
||
4 Two overlapping crystals – segmentation challenge
|
||
5 Two clusters at ~90 ° – twinned / differently oriented crystal
|
||
6 Pure background only – true negative (no crystal)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass, field
|
||
from typing import List
|
||
|
||
import numpy as np
|
||
from aarecommon.models.raster_grid import RasterGridRequest
|
||
from jfjoch_client import ScanResult
|
||
from jfjoch_client.models.scan_result_images_inner import ScanResultImagesInner
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cluster geometry descriptor
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@dataclass
|
||
class _ClusterParams:
|
||
"""Fractional coordinates and shape of one elliptical crystal cluster.
|
||
|
||
cx, cy – cluster centre as a fraction of (n_x, n_y) [0..1]
|
||
ax, ay – semi-axes as a fraction of (n_x, n_y)
|
||
theta – rotation of the ellipse in radians
|
||
peak – peak spots_low_res at cluster centre (integer)
|
||
"""
|
||
|
||
cx: float
|
||
cy: float
|
||
ax: float
|
||
ay: float
|
||
theta: float
|
||
peak: int
|
||
|
||
|
||
@dataclass
|
||
class _SeedConfig:
|
||
clusters: List[_ClusterParams] = field(default_factory=list)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Pre-defined seed catalogue
|
||
# ---------------------------------------------------------------------------
|
||
|
||
SEED_CATALOGUE: dict[int, _SeedConfig] = {
|
||
# --- 1: single crystal, centred, compact, strong signal ----------------
|
||
1: _SeedConfig(
|
||
clusters=[_ClusterParams(cx=0.50, cy=0.50, ax=0.15, ay=0.15, theta=0.0, peak=120)]
|
||
),
|
||
# --- 2: single crystal, off-centre, slightly elongated -----------------
|
||
2: _SeedConfig(
|
||
clusters=[_ClusterParams(cx=0.25, cy=0.70, ax=0.12, ay=0.18, theta=0.35, peak=90)]
|
||
),
|
||
# --- 3: two well-separated crystals ------------------------------------
|
||
3: _SeedConfig(
|
||
clusters=[
|
||
_ClusterParams(cx=0.20, cy=0.20, ax=0.12, ay=0.12, theta=0.0, peak=100),
|
||
_ClusterParams(cx=0.75, cy=0.72, ax=0.15, ay=0.10, theta=0.2, peak=80),
|
||
]
|
||
),
|
||
# --- 4: two overlapping crystals (centres ~1 sigma apart) --------------
|
||
4: _SeedConfig(
|
||
clusters=[
|
||
_ClusterParams(cx=0.40, cy=0.45, ax=0.18, ay=0.15, theta=0.0, peak=110),
|
||
_ClusterParams(cx=0.58, cy=0.55, ax=0.16, ay=0.18, theta=0.5, peak=95),
|
||
]
|
||
),
|
||
# --- 5: two clusters with ~90° different orientations (twinned) --------
|
||
5: _SeedConfig(
|
||
clusters=[
|
||
_ClusterParams(cx=0.30, cy=0.40, ax=0.25, ay=0.08, theta=0.0, peak=105),
|
||
_ClusterParams(cx=0.68, cy=0.62, ax=0.08, ay=0.25, theta=0.0, peak=100),
|
||
]
|
||
),
|
||
# --- 6: pure background, no crystal – true negative --------------------
|
||
6: _SeedConfig(clusters=[]),
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Core generation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _gaussian_cluster(
|
||
ix: np.ndarray, iy: np.ndarray, p: _ClusterParams, n_x: int, n_y: int
|
||
) -> np.ndarray:
|
||
"""Return a 2-D Gaussian signal array for one cluster.
|
||
|
||
ix, iy are integer coordinate grids with shape (n_x, n_y).
|
||
"""
|
||
cx = p.cx * (n_x - 1)
|
||
cy = p.cy * (n_y - 1)
|
||
|
||
dx = (ix - cx) / max(p.ax * n_x, 0.5)
|
||
dy = (iy - cy) / max(p.ay * n_y, 0.5)
|
||
|
||
cos_t = np.cos(p.theta)
|
||
sin_t = np.sin(p.theta)
|
||
dx_r = dx * cos_t + dy * sin_t
|
||
dy_r = -dx * sin_t + dy * cos_t
|
||
|
||
return p.peak * np.exp(-2.0 * (dx_r**2 + dy_r**2))
|
||
|
||
|
||
def generate_no_beam_scan_result(request: RasterGridRequest, seed: int | None = None) -> ScanResult:
|
||
"""Build a synthetic ScanResult for no-beam / offline operation.
|
||
|
||
Parameters
|
||
----------
|
||
request:
|
||
The grid request whose n_x, n_y, and file_prefix are used.
|
||
seed:
|
||
Integer 1-6 selects a pre-defined scenario from SEED_CATALOGUE.
|
||
Any other value (or None) draws cluster parameters randomly using
|
||
the seed as an RNG seed (None → fully random).
|
||
|
||
Returns
|
||
-------
|
||
ScanResult
|
||
Populated with one ScanResultImagesInner per grid cell, with
|
||
realistic background noise and optional crystal cluster(s).
|
||
"""
|
||
n_x = max(request.n_x, 1)
|
||
n_y = max(request.n_y, 1)
|
||
|
||
rng = np.random.default_rng(seed)
|
||
|
||
# Coordinate grids
|
||
ix, iy = np.meshgrid(np.arange(n_x), np.arange(n_y), indexing="ij")
|
||
|
||
# Background: Poisson noise, mean ≈ 3 counts
|
||
bkg = rng.poisson(lam=3.0, size=(n_x, n_y)).astype(float)
|
||
|
||
# Crystal signal layer (spots_low_res)
|
||
signal = np.zeros((n_x, n_y), dtype=float)
|
||
|
||
if seed in SEED_CATALOGUE:
|
||
cfg = SEED_CATALOGUE[seed]
|
||
clusters = cfg.clusters
|
||
else:
|
||
# Random fallback: 1 or 2 clusters
|
||
n_clusters = rng.integers(1, 3)
|
||
clusters = [
|
||
_ClusterParams(
|
||
cx=rng.uniform(0.15, 0.85),
|
||
cy=rng.uniform(0.15, 0.85),
|
||
ax=rng.uniform(0.12, 0.30),
|
||
ay=rng.uniform(0.12, 0.30),
|
||
theta=rng.uniform(0, np.pi),
|
||
peak=int(rng.integers(50, 150)),
|
||
)
|
||
for _ in range(n_clusters)
|
||
]
|
||
|
||
for p in clusters:
|
||
signal += _gaussian_cluster(ix, iy, p, n_x, n_y)
|
||
|
||
# Add Poisson noise on top of the crystal signal
|
||
spots_low_res = rng.poisson(lam=np.maximum(signal, 0)).astype(int)
|
||
|
||
# Assemble image list – one entry per grid cell in serpentine order:
|
||
# row 0 left→right, row 1 right→left, row 2 left→right, …
|
||
images: list[ScanResultImagesInner] = []
|
||
img_number = 0
|
||
for xi in range(n_x):
|
||
yi_range = range(n_y)
|
||
for yi in yi_range:
|
||
images.append(
|
||
ScanResultImagesInner(
|
||
number=img_number,
|
||
nx=xi,
|
||
ny=yi,
|
||
efficiency=1.0,
|
||
bkg=float(bkg[xi, yi]),
|
||
spots=int(spots_low_res[xi, yi]),
|
||
spots_low_res=int(spots_low_res[xi, yi]),
|
||
spots_indexed=0,
|
||
spots_ice=0,
|
||
index=0,
|
||
b=0.0,
|
||
res=None,
|
||
pixel_sum=None,
|
||
max=None,
|
||
sat=None,
|
||
err=None,
|
||
)
|
||
)
|
||
img_number += 1
|
||
|
||
return ScanResult(file_prefix=request.file_prefix, images=images)
|