diff --git a/src/aare/common/find_xtal.py b/src/aare/common/find_xtal.py index 23468175..6433fa10 100644 --- a/src/aare/common/find_xtal.py +++ b/src/aare/common/find_xtal.py @@ -314,4 +314,300 @@ def has_sufficient_low_res_spots( ) return False - return True \ No newline at end of file + return True + + +def compute_crystal_score_array( + scan_results: List, + w_bkg: float = 0.25, + w_low_res: float = 0.55, + w_indexed: float = 0.20, +) -> np.ndarray: + """Combine bkg (25%), spots_low_res (55%), and spots_indexed (20%) into a 0–100 score. + + Each field is min-max normalised to [0, 100] within the grid before weighting, + so the final score is the probability (0–100) that a pixel belongs to a crystal. + """ + arr_bkg = rebuild_array_from_scan_results(scan_results, "bkg") + arr_low = rebuild_array_from_scan_results(scan_results, "spots_low_res") + arr_idx = rebuild_array_from_scan_results(scan_results, "spots_indexed") + + def _norm(a: np.ndarray) -> np.ndarray: + mn, mx = float(a.min()), float(a.max()) + if mx == mn: + return np.zeros_like(a, dtype=float) + return (a - mn) / (mx - mn) * 100.0 + + return w_bkg * _norm(arr_bkg) + w_low_res * _norm(arr_low) + w_indexed * _norm(arr_idx) + + +def _max_cell(arr: np.ndarray) -> tuple[int, int]: + """Return (nx, ny) of the cell with the highest value.""" + idx = np.unravel_index(np.argmax(arr), arr.shape) + return int(idx[0]), int(idx[1]) + + +def _draw_panel( + ax, + arr: np.ndarray, + mask: np.ndarray, + label: str, + threshold: float, + grid_size_mm: Optional[tuple[float, float]] = None, + cbar_label: str = "spots_low_res", +) -> None: + """Shared helper: heatmap + crystal contour + max-cell square on one Axes. + + If grid_size_mm=(step_x_mm, step_y_mm) is provided, the crystal size in µm + is computed via get_xtal_size and shown in the panel title. + """ + import matplotlib.patches as mpatches + import matplotlib.pyplot as plt + from aare.common.coordinate import Coordinate + + n_nx, n_ny = arr.shape + max_nx, max_ny = _max_cell(arr) + n_cells = int(mask.sum()) + + im = ax.imshow( + arr.T, + origin="lower", + cmap="viridis", + aspect="equal", + extent=[-0.5, n_nx - 0.5, -0.5, n_ny - 0.5], + ) + plt.colorbar(im, ax=ax, label=cbar_label, fraction=0.046, pad=0.04) + ax.contour( + np.arange(n_nx), np.arange(n_ny), mask.T.astype(float), + levels=[0.5], + colors="cyan", + linewidths=1.8, + ) + rect = mpatches.Rectangle( + (max_nx - 0.5, max_ny - 0.5), 1, 1, + linewidth=2, edgecolor="red", facecolor="none", + ) + ax.add_patch(rect) + + size_line = "" + if grid_size_mm is not None and n_cells > 0: + r = RasterGridRequest( + exp_time_s=0.0, + n_x=n_nx, + n_y=n_ny, + grid_size_mm=Coordinate(x=grid_size_mm[0], y=grid_size_mm[1]), + smargon_top_left=None, + ) + xtal_size = get_xtal_size(CrystalSize(x=0, y=0, z=0), mask.astype(float), r) + size_line = f"\n{xtal_size.x:.0f}×{xtal_size.y:.0f} µm" + + ax.set_title( + f"{label}\nthresh≈{threshold:.1f} | {n_cells} cells{size_line}", + fontsize=8, + ) + ax.set_xlabel("nx", fontsize=7) + ax.set_ylabel("ny", fontsize=7) + ax.tick_params(labelsize=6) + + +# ── Clustering / thresholding methods ───────────────────────────────────────── + +def crystal_mask_corner_background(arr: np.ndarray) -> tuple[np.ndarray, float]: + """Threshold = max(far-corner value, floor=10). + + Simple, parameter-free. Returns too many cells when corners are zero + because any nonzero value passes. + """ + n_nx, n_ny = arr.shape + corners = [arr[0, 0], arr[n_nx - 1, 0], arr[0, n_ny - 1], arr[n_nx - 1, n_ny - 1]] + threshold = max(float(np.max(corners)), 10.0) + return arr > threshold, threshold + + +def crystal_mask_otsu_nonzero(arr: np.ndarray) -> tuple[np.ndarray, float]: + """Otsu threshold computed only on the nonzero values. + + Finds the natural gap in the signal distribution. + Ignores the large mass of background zeros so the threshold is + placed within the diffraction-signal population. + """ + from skimage.filters import threshold_otsu + + nonzero = arr[arr > 0] + if nonzero.size == 0: + return np.zeros_like(arr, dtype=bool), 0.0 + threshold = float(threshold_otsu(nonzero)) + return arr > threshold, threshold + + +def crystal_mask_mean_sigma(arr: np.ndarray, n_sigma: float = 0.5) -> tuple[np.ndarray, float]: + """Threshold = mean + n_sigma * std of nonzero values. + + n_sigma=0.5 keeps cells within ~1/2 std above average signal. + Raise n_sigma to tighten the region around the strongest-diffracting core. + """ + nonzero = arr[arr > 0] + if nonzero.size == 0: + return np.zeros_like(arr, dtype=bool), 0.0 + threshold = float(nonzero.mean() + n_sigma * nonzero.std()) + return arr > threshold, threshold + + +def crystal_mask_signal_percentile(arr: np.ndarray, percentile: float = 60.0) -> tuple[np.ndarray, float]: + """Threshold = given percentile of nonzero values. + + percentile=60 keeps the top 40 % of signal cells; raise to sharpen. + Unlike mean+sigma this is robust to heavy-tailed distributions. + """ + nonzero = arr[arr > 0] + if nonzero.size == 0: + return np.zeros_like(arr, dtype=bool), 0.0 + threshold = float(np.percentile(nonzero, percentile)) + return arr > threshold, threshold + + +def crystal_mask_dbscan(arr: np.ndarray, min_signal: float = 10.0, + eps: float = 1.5, min_samples: int = 3) -> tuple[np.ndarray, float]: + """DBSCAN spatial clustering on cells with signal > min_signal. + + Groups adjacent diffracting cells into clusters; the largest cluster + (by total signal weight) is labelled as the crystal. + eps=1.5 connects cells that are one grid step apart (including diagonal). + """ + from sklearn.cluster import DBSCAN + + ys, xs = np.where(arr > min_signal) + if len(ys) == 0: + return np.zeros_like(arr, dtype=bool), min_signal + + coords = np.column_stack([ys, xs]).astype(float) + labels = DBSCAN(eps=eps, min_samples=min_samples).fit_predict(coords) + + best_label, best_weight = -1, -1.0 + for lbl in set(labels): + if lbl == -1: + continue + weight = float(arr[ys[labels == lbl], xs[labels == lbl]].sum()) + if weight > best_weight: + best_label, best_weight = lbl, weight + + mask = np.zeros_like(arr, dtype=bool) + if best_label != -1: + sel = labels == best_label + mask[ys[sel], xs[sel]] = True + return mask, min_signal + + +def crystal_mask_kmeans(arr: np.ndarray, n_clusters: int = 3) -> tuple[np.ndarray, float]: + """K-means on signal values (1-D feature). + + Splits cells into n_clusters groups; the cluster with the highest + centroid is the crystal. n_clusters=3 separates background / fringe / crystal. + """ + from sklearn.cluster import KMeans + + flat = arr.flatten().reshape(-1, 1) + km = KMeans(n_clusters=n_clusters, random_state=0, n_init="auto").fit(flat) + crystal_label = int(np.argmax(km.cluster_centers_.flatten())) + threshold = float(sorted(km.cluster_centers_.flatten())[-2]) + mask = (km.labels_.reshape(arr.shape) == crystal_label) + return mask, threshold + + +CRYSTAL_METHOD_MAP: dict = { + "corner": ("Corner background", crystal_mask_corner_background), + "otsu": ("Otsu (nonzero)", crystal_mask_otsu_nonzero), + "mean_sigma": ("Mean + 0.5σ", lambda arr: crystal_mask_mean_sigma(arr, n_sigma=0.5)), + "percentile": ("Top-40% signal", lambda arr: crystal_mask_signal_percentile(arr, percentile=60.0)), + "dbscan": ("DBSCAN (spatial)", crystal_mask_dbscan), + "kmeans": ("K-means (k=3)", crystal_mask_kmeans), +} + + +def compare_crystal_methods( + results: List, + arr: Optional[np.ndarray] = None, + grid_size_mm: Optional[tuple[float, float]] = None, +) -> None: + """Plot each crystal-detection method side-by-side for visual comparison. + + Cyan contour = detected crystal region. + Red square = cell with maximum spots_low_res. + grid_size_mm = (step_x_mm, step_y_mm) — when provided, crystal size in µm + is calculated via get_xtal_size and shown in each panel title. + """ + import matplotlib.pyplot as plt + + if arr is None: + arr = rebuild_array_from_scan_results(results, "spots_low_res") + + methods = [ + ("Corner background", *crystal_mask_corner_background(arr)), + ("Otsu (nonzero)", *crystal_mask_otsu_nonzero(arr)), + ("Mean + 0.5σ", *crystal_mask_mean_sigma(arr, n_sigma=0.5)), + ("Top-40% signal", *crystal_mask_signal_percentile(arr, percentile=60.0)), + ("DBSCAN (spatial)", *crystal_mask_dbscan(arr)), + ("K-means (k=3)", *crystal_mask_kmeans(arr)), + ] + + fig, axes = plt.subplots(2, 3, figsize=(14, 9)) + for ax, (label, mask, threshold) in zip(axes.flat, methods): + _draw_panel(ax, arr, mask, label, threshold, grid_size_mm=grid_size_mm) + + fig.suptitle("Crystal region detection — method comparison\n" + "cyan = crystal outline | red square = max spots_low_res cell", + fontsize=10) + plt.tight_layout() + plt.show() + + +def compare_crystal_methods_scored( + results: List, + score_arr: Optional[np.ndarray] = None, + method: Optional[str] = None, + grid_size_mm: Optional[tuple[float, float]] = None, + w_bkg: float = 0.20, + w_low_res: float = 0.60, + w_indexed: float = 0.20, +) -> None: + """Plot each crystal-detection method applied to the composite 0–100 crystal score. + + The score combines bkg (w_bkg), spots_low_res (w_low_res), and spots_indexed + (w_indexed), each min-max normalised. All six methods are shown side-by-side. + + Args: + results: scan result list (used to build score if score_arr is None) + score_arr: pre-computed score array; computed from results when None + method: if given (one of 'corner','otsu','mean_sigma','percentile', + 'dbscan','kmeans'), that panel is highlighted with a yellow border + grid_size_mm: (step_x_mm, step_y_mm) for crystal-size annotation + w_bkg, w_low_res, w_indexed: composite score weights (must sum to 1.0) + """ + import matplotlib.pyplot as plt + + if score_arr is None: + score_arr = compute_crystal_score_array(results, w_bkg=w_bkg, w_low_res=w_low_res, w_indexed=w_indexed) + + method_entries = [ + (name, label, fn) + for name, (label, fn) in CRYSTAL_METHOD_MAP.items() + ] + + fig, axes = plt.subplots(2, 3, figsize=(15, 9)) + for ax, (name, label, fn) in zip(axes.flat, method_entries): + mask, threshold = fn(score_arr) + _draw_panel(ax, score_arr, mask, label, threshold, + grid_size_mm=grid_size_mm, cbar_label="crystal score (0–100)") + if method and name == method: + for spine in ax.spines.values(): + spine.set_edgecolor("yellow") + spine.set_linewidth(3) + + weight_note = f"bkg×{w_bkg:.0%} + spots_low_res×{w_low_res:.0%} + spots_indexed×{w_indexed:.0%}" + fig.suptitle( + f"Crystal region detection on composite score [{weight_note}]\n" + "cyan = crystal outline | red square = max-score cell | score range 0–100", + fontsize=10, + ) + plt.tight_layout() + plt.show() \ No newline at end of file diff --git a/src/aare/common/raster_grid.py b/src/aare/common/raster_grid.py index 7c4327ff..676cf2d9 100644 --- a/src/aare/common/raster_grid.py +++ b/src/aare/common/raster_grid.py @@ -32,6 +32,10 @@ class RasterGridRequest(BaseModel): visible: bool = True + # When True, skip all hardware and generate simulated diffraction data. + # The file_prefix will have "_SIMU" appended so the DB record is identifiable. + no_beam: bool = False + def get_image_number(self) -> int: return self.n_x * self.n_y diff --git a/src/aare/common/simulate_raster.py b/src/aare/common/simulate_raster.py new file mode 100644 index 00000000..5cc51e20 --- /dev/null +++ b/src/aare/common/simulate_raster.py @@ -0,0 +1,214 @@ +""" +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 jfjoch_client import ScanResult +from jfjoch_client.models.scan_result_images_inner import ScanResultImagesInner + +from aare.common.raster_grid import RasterGridRequest + +# --------------------------------------------------------------------------- +# 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 row-major order + images: list[ScanResultImagesInner] = [] + for xi in range(n_x): + for yi in range(n_y): + img_number = xi * n_y + yi + 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, + )) + + return ScanResult(file_prefix=request.file_prefix, images=images) diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 14bd5164..8e1c9210 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -33,6 +33,7 @@ from aare.common.models import ( SimpleScanParameters, MLBoxModel, FluorescenceSpectrumParameterModel, FluorescenceSpectrumOutputModel) from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem +from aare.common.simulate_raster import generate_no_beam_scan_result from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan from aare.common.sample_geometry import SampleGeometryModel from aare.daq.spreadsheetupdater import beamline @@ -917,6 +918,30 @@ class AareDAQ: if self.sample is not None and self.sample.db_id is not None: self.__aare.create_gridscan_run(self.sample, request, status) + if request.no_beam: + request = copy.deepcopy(request) + if request.file_prefix: + request.file_prefix = f"{request.file_prefix}_SIMU" + logger.info("No beam mode: skipping hardware, generating simulated raster result.") + scan_result = generate_no_beam_scan_result(request) + sample_id = self.sample.db_id if self.sample and self.sample.db_id is not None else None + if sample_id: + self.__set_state(BeamlineStateEnum.SampleAlignment) + self.save_screenshot_db(sample_id, f"{sample_id}_post_raster_{request.omega_deg}deg") + self.__aare.ingest_gridscan( + sample=self.sample, + raster_result=scan_result, + raster_request=request, + geom=self.sample_geometry, + com=None, + beam_mark_pxl=self.__cfg.get_beam_mark(self.zoom), + ) + return CompletedRasterGridElem( + request=request, + result=scan_result, + centre_of_mass=None, + ) + if not self.__cfg.simulated_detector: logger.info("initialise detector") self.__jfjoch.measure_raster(request, status)