DAQ: updates to find_xtal.py

This commit is contained in:
2026-05-01 15:42:18 +02:00
parent fcb615255f
commit e27d938236
2 changed files with 371 additions and 63 deletions
+298 -5
View File
@@ -4,13 +4,11 @@ import numpy as np
from scipy import ndimage
from aare.common.models import CrystalSize
from aare.common.logger_events import log_timing
from aare.common.raster_grid import RasterGridRequest, CenterOfMassModel
from aare.common.logger_config import setup_logger
logger = setup_logger('aareDAQ')
@log_timing(logger, "Identify crystal raster")
def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel | None:
images = result.images
if images and any(getattr(img, "spots_low_res", 0) for img in images):
@@ -53,7 +51,6 @@ def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel |
else:
return None
@log_timing(logger, "Rebuild array from scan results")
def rebuild_array_from_scan_results(scan_results: List,
value_field: str,
array_shape: Optional[tuple] = None,
@@ -77,7 +74,7 @@ def rebuild_array_from_scan_results(scan_results: List,
if nx is None or ny is None:
continue
positions.append((int(ny), int(nx))) # Corrected: (row, col) = (ny, nx)
positions.append((int(nx), int(ny))) # Note: (row, col) = (ny, nx)
if not value:
value = 0.0
values.append(float(value))
@@ -317,4 +314,300 @@ def has_sufficient_low_res_spots(
)
return False
return True
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 0100 score.
Each field is min-max normalised to [0, 100] within the grid before weighting,
so the final score is the probability (0100) 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 0100 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 (0100)")
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 0100",
fontsize=10,
)
plt.tight_layout()
plt.show()
+73 -58
View File
@@ -47,6 +47,8 @@ from aare.common.automation_models import (
WorkflowStateKind,
)
from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem, grid_to_image_id
from aare.common.simulate_raster import generate_no_beam_scan_result
from aare.common.find_xtal import rebuild_array_from_scan_results, raster_centre_of_mass, create_quality_filtered_array
from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan
from aare.common.sample_geometry import SampleGeometryModel
from aare.daq.operations.face_detection import FaceDetectionContext, FaceDetectionService, FaceDetectionResult
@@ -659,7 +661,7 @@ class AareDAQ:
),
)
else:
logger.debug(f"Is detector simulated? {self.__cfg.simulated_detector}")
if not self.__cfg.simulated_detector:
logger.info(f"initialise detector for raster")
status = self.status
@@ -1732,7 +1734,10 @@ 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)
self.__jfjoch.wait_till_running(timeout=60.0)
if not self.__cfg.simulated_detector:
self.__jfjoch.wait_till_running(timeout=60.0)
else:
logger.info("Simulated detector mode enabled; faking jfjoch intilalisation.")
self.__devs.aerotech.grid_scan(
grid_elem_size_y_um=request.grid_size_mm.y * 1000,
@@ -1751,62 +1756,23 @@ class AareDAQ:
self.__devs.aerotech_pos = AerotechCoordinate(at_mm=coord, omega_deg=self.__devs.aerotech_omega)
self.__devs.aerotech.wait_till_done(timeout=int(360))
if request.n_x == 1:
x = request.grid_size_mm.x / 2.0
else:
x = ((request.n_x - 1) * request.grid_size_mm.x) / 2.0
y = ((request.n_y - 1) * request.grid_size_mm.y) / 2.0
grid_centre_offset = self.sample_geometry.smargon_nudge(Coordinate(x=x, y=y))
logger.info(
"Calculated raster centre offset",
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"centre_offset_x_mm": grid_centre_offset.x,
"centre_offset_y_mm": grid_centre_offset.y,
"centre_offset_z_mm": grid_centre_offset.z,
"grid_half_width_x_mm": x,
"grid_half_height_y_mm": y,
"top_left_x_mm": getattr(request.smargon_top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(request.smargon_top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(request.smargon_top_left.sh_mm, "z", None),
},
),
)
logger.info(f"moving Smargon to grid centre offset {grid_centre_offset}")
grid_centre_smargon = SmargonCoordinate(
sh_mm=request.smargon_top_left.sh_mm + grid_centre_offset,
phi_deg=request.smargon_top_left.phi_deg,
chi_deg=request.smargon_top_left.chi_deg
)
self.__devs.smargon_pos = grid_centre_smargon
self.__devs.smargon_wait(timeout=180)
logger.info(
"Moved Smargon to raster centre",
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"centre_sh_x_mm": grid_centre_smargon.sh_mm.x,
"centre_sh_y_mm": grid_centre_smargon.sh_mm.y,
"centre_sh_z_mm": grid_centre_smargon.sh_mm.z,
"centre_phi_deg": grid_centre_smargon.phi_deg,
"centre_chi_deg": grid_centre_smargon.chi_deg,
},
),
)
if self.__cfg.simulated_detector:
scan_result = self._build_fake_scan_result(
file_prefix=request.file_prefix,
image_count=request.n_x * request.n_y,
logger.info("Simulated detector mode enabled; using fake raster result.")
scan_result = generate_no_beam_scan_result(request)
result_array = create_quality_filtered_array(scan_result.images,
'spots_low_res', min_spots=None,
min_efficiency=1.0, min_background=None,
min_low_res_spots=None)
com = raster_centre_of_mass(result_array, scan_result.images)
target_coor = com.get_com_mm(request)
target_coor_offset = self.sample_geometry.smargon_nudge(target_coor)
target_smargon = SmargonCoordinate(
sh_mm=request.smargon_top_left.sh_mm + target_coor_offset,
phi_deg=request.smargon_top_left.phi_deg,
chi_deg=request.smargon_top_left.chi_deg
)
com = None
else:
scan_result = self.__jfjoch.wait_till_done(60)
com = None
@@ -1819,6 +1785,56 @@ class AareDAQ:
{"exp_time_s": request.exp_time_s},
),
)
if request.n_x == 1:
x = request.grid_size_mm.x / 2.0
else:
x = ((request.n_x - 1) * request.grid_size_mm.x) / 2.0
y = ((request.n_y - 1) * request.grid_size_mm.y) / 2.0
target_coor_offset = self.sample_geometry.smargon_nudge(Coordinate(x=x, y=y))
logger.info(
"Calculated raster centre offset",
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"centre_offset_x_mm": target_coor_offset.x,
"centre_offset_y_mm": target_coor_offset.y,
"centre_offset_z_mm": target_coor_offset.z,
"grid_half_width_x_mm": x,
"grid_half_height_y_mm": y,
"top_left_x_mm": getattr(request.smargon_top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(request.smargon_top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(request.smargon_top_left.sh_mm, "z", None),
},
),
)
logger.info(f"moving Smargon to grid centre offset {target_coor_offset}")
target_smargon = SmargonCoordinate(
sh_mm=request.smargon_top_left.sh_mm + target_coor_offset,
phi_deg=request.smargon_top_left.phi_deg,
chi_deg=request.smargon_top_left.chi_deg
)
self.__devs.smargon_pos = target_smargon
self.__devs.smargon_wait(timeout=180)
logger.info(
"Moved Smargon to raster centre",
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"centre_sh_x_mm": target_smargon.sh_mm.x,
"centre_sh_y_mm": target_smargon.sh_mm.y,
"centre_sh_z_mm": target_smargon.sh_mm.z,
"centre_phi_deg": target_smargon.phi_deg,
"centre_chi_deg": target_smargon.chi_deg,
},
),
)
sample_id = self.sample.db_id if self.sample and self.sample.db_id is not None else None
if sample_id:
@@ -1874,7 +1890,7 @@ class AareDAQ:
return CompletedRasterGridElem(
request=copy.deepcopy(request),
result=scan_result,
centre_of_mass=None,
centre_of_mass=com,
)
except Exception as e:
@@ -2618,7 +2634,6 @@ class AareDAQ:
if error:
msg += f"with an error"
self.__set_state(BeamlineStateEnum.RobotSampleExchange)
try:
if self.__cfg.state_busy:
self.__set_state(BeamlineStateEnum.RobotSampleExchange)