diff --git a/src/aare/common/find_xtal.py b/src/aare/common/find_xtal.py index 9fefdb2f..6433fa10 100644 --- a/src/aare/common/find_xtal.py +++ b/src/aare/common/find_xtal.py @@ -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 \ 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/daq/daq.py b/src/aare/daq/daq.py index 5e082412..24846c6d 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -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)