Source code for coseda.peakfinding.peakfinder9

"""Numba-accelerated peakfinder9 implementation with Dask parallelism.

Algorithm
---------
For every candidate pixel (i, j) with value *val*, the following conditions
must all pass (identical to the calNG / CrystFEL pf9 logic):

1. *val* must be the local maximum on the ring:
       val > max(ring_pixels) + minPeakValueOverNeighbors
2. At least ``max(1, 2*window_radius - 1)`` ring pixels must be finite.
3. SNR of the peak pixel:
       val > ring_mean + min_snr_max_pixel * ring_std
4. Integrated peak mass of neighbouring pixels that exceed
   ``ring_mean + min_snr_peak_pixels * ring_std``:
       peak_mass >= min_snr_whole_peak * ring_std

The "ring" is the perimeter of the ``(2r+1)×(2r+1)`` square centred on the
candidate (pixels at Chebyshev distance exactly *r*).  Masked / bad pixels
are excluded via a ``valid`` boolean array (no NaN injection needed).

Implementation
--------------
Two Numba-compiled kernels replace the NumPy/SciPy approach:

``_pf9_compute_candidates``  (``@njit, parallel=True``)
    Iterates every pixel in parallel across rows.  For each valid pixel,
    computes ring mean, std, and maximum entirely in-register (no
    temporary arrays), then flags it as a candidate if all pre-checks pass.

``_pf9_collect_peaks``  (``@njit``)
    Serial pass over the candidate map: computes peak mass and
    centre-of-mass for each candidate and assembles the output arrays.

Output format is identical to the pf8 pipeline
(``nPeaks``, ``peakXPosRaw``, ``peakYPosRaw``, ``peakTotalIntensity``)
so the two algorithms are drop-in replacements for each other.
"""

from __future__ import annotations

import datetime
import os

import h5py
import numpy as np
from dask import delayed
from dask.diagnostics import ProgressBar
from dask.distributed import as_completed
from numba import njit, prange

from coseda.dask_client_manager import DaskClientManager
from coseda.io import config_to_paths, handle_input, read_config
from coseda.logging_utils import log_print, log_result, log_start
from coseda.nexus.paths import get_mask_dataset
from coseda.nexus.peaks import ensure_peak_nxdata
from coseda.nexus.process import write_nxprocess_peakfinding
from coseda.peakfinding.maxres import write_maxres_dataset

MAX_PEAKS = 500


# ---------------------------------------------------------------------------
# Numba-compiled core kernels
# ---------------------------------------------------------------------------

@njit(cache=True, parallel=True)
def _pf9_compute_candidates(
    data,                    # float32 2-D frame
    valid,                   # bool   2-D (True = pixel is usable)
    r,                       # int    window_radius
    min_sigma,               # float64
    min_peak_over_neighbors, # float64
    min_snr_max_pixel,       # float64
    min_valid_ring,          # int    minimum ring pixels required
    ring_mean_out,           # float64 2-D  (output)
    ring_std_out,            # float64 2-D  (output)
    cand_out,                # bool   2-D  (output, pre-zeroed)
):
    """Parallel kernel: ring stats + candidate flagging for every pixel.

    The ring around (ci, cj) is all (ci+di, cj+dj) where
    max(|di|, |dj|) == r  (Chebyshev distance exactly r).
    The skip condition  ``abs(di) < r and abs(dj) < r``  tests the
    interior (distance < r) and is equivalent to max(|di|,|dj|) < r.
    """
    rows, cols = data.shape

    for ci in prange(rows):
        for cj in range(cols):
            if not valid[ci, cj]:
                continue

            val   = np.float64(data[ci, cj])
            rsum  = 0.0
            rsq   = 0.0
            rcnt  = np.int32(0)
            rmax  = -1e38

            for di in range(-r, r + 1):
                for dj in range(-r, r + 1):
                    # skip interior (Chebyshev distance < r)
                    if abs(di) < r and abs(dj) < r:
                        continue
                    ni = ci + di
                    nj = cj + dj
                    if 0 <= ni < rows and 0 <= nj < cols and valid[ni, nj]:
                        v     = np.float64(data[ni, nj])
                        rsum += v
                        rsq  += v * v
                        rcnt += np.int32(1)
                        if v > rmax:
                            rmax = v

            if rcnt < min_valid_ring:
                continue

            rm = rsum / rcnt
            rv = rsq / rcnt - rm * rm
            if rv < 0.0:
                rv = 0.0
            rs = np.sqrt(rv)
            if rs < min_sigma:
                rs = min_sigma

            ring_mean_out[ci, cj] = rm
            ring_std_out[ci, cj]  = rs

            if (val > rmax + min_peak_over_neighbors and
                    val > rm + min_snr_max_pixel * rs):
                cand_out[ci, cj] = True


@njit(cache=True)
def _pf9_collect_peaks(
    data,                # float32 2-D
    valid,               # bool   2-D
    r,                   # int
    ring_mean,           # float64 2-D
    ring_std,            # float64 2-D
    cand,                # bool   2-D candidate map
    min_snr_peak_pixels, # float64
    min_snr_whole_peak,  # float64
    max_peaks,           # int
):
    """Serial kernel: peak-mass + centre-of-mass for each candidate pixel."""
    rows, cols = data.shape
    xs          = np.empty(max_peaks, dtype=np.float64)
    ys          = np.empty(max_peaks, dtype=np.float64)
    intensities = np.empty(max_peaks, dtype=np.float64)
    n = np.int32(0)

    for ci in range(rows):
        if n >= max_peaks:
            break
        for cj in range(cols):
            if not cand[ci, cj]:
                continue

            rm  = ring_mean[ci, cj]
            rs  = ring_std[ci, cj]
            thr = rm + min_snr_peak_pixels * rs

            peak_mass = 0.0
            com_row   = 0.0
            com_col   = 0.0

            for di in range(-r, r + 1):
                for dj in range(-r, r + 1):
                    ni = ci + di
                    nj = cj + dj
                    if 0 <= ni < rows and 0 <= nj < cols and valid[ni, nj]:
                        pv = np.float64(data[ni, nj])
                        if pv > thr:
                            w          = pv - rm
                            peak_mass += w
                            com_row   += w * ni
                            com_col   += w * nj

            if peak_mass <= 0.0 or peak_mass < min_snr_whole_peak * rs:
                continue

            xs[n]          = com_col / peak_mass
            ys[n]          = com_row / peak_mass
            intensities[n] = peak_mass
            n += 1
            if n >= max_peaks:
                break

    return xs[:n], ys[:n], intensities[:n]


# ---------------------------------------------------------------------------
# Single-frame pf9
# ---------------------------------------------------------------------------

[docs] def find_peaks_pf9( image: np.ndarray, mask: np.ndarray | None, window_radius: int, min_sigma: float, min_peak_over_neighbors: float, min_snr_max_pixel: float, min_snr_peak_pixels: float, min_snr_whole_peak: float, max_peaks: int = MAX_PEAKS, ) -> tuple[list[float], list[float], list[float]]: """Run peakfinder9 on one 2-D frame. Parameters ---------- image : 2-D array-like Raw detector frame. mask : 2-D int8 array or None ``1`` = valid pixel, ``0`` = masked. Same convention as pf8. window_radius : int Half-side of the square background ring (``windowRadius`` in the pf9 parameter table). min_sigma : float Floor applied to the background std before SNR comparisons. min_peak_over_neighbors : float Candidate pixel value must exceed *every* ring pixel by this many ADU. min_snr_max_pixel : float Candidate pixel must be at least this many sigma above the ring mean. min_snr_peak_pixels : float Pixels in the neighbourhood are counted toward peak mass when they are more than this many sigma above the ring mean. min_snr_whole_peak : float Integrated peak mass divided by ring std must exceed this value. max_peaks : int Upper bound on the number of peaks returned per frame. Returns ------- x_list : list of float fast-scan (column) centre-of-mass y_list : list of float slow-scan (row) centre-of-mass intensity_list : list of float peak mass = Σ (val − mean) over peak pixels """ r = int(window_radius) data = np.asarray(image, dtype=np.float32) # Build boolean valid map (avoids NaN injection into data) valid = np.isfinite(data) if mask is not None: valid = valid & (np.asarray(mask) != 0) valid = np.ascontiguousarray(valid) data = np.ascontiguousarray(data) rows, cols = data.shape min_valid_ring = max(1, 2 * r - 1) ring_mean = np.zeros((rows, cols), dtype=np.float64) ring_std = np.zeros((rows, cols), dtype=np.float64) cand = np.zeros((rows, cols), dtype=np.bool_) _pf9_compute_candidates( data, valid, r, float(min_sigma), float(min_peak_over_neighbors), float(min_snr_max_pixel), min_valid_ring, ring_mean, ring_std, cand, ) xs, ys, intensities = _pf9_collect_peaks( data, valid, r, ring_mean, ring_std, cand, float(min_snr_peak_pixels), float(min_snr_whole_peak), max_peaks, ) return list(xs), list(ys), list(intensities)
# --------------------------------------------------------------------------- # Multi-scale entry point # --------------------------------------------------------------------------- def _merge_peaks( peaks_per_scale: list, merge_distance: float, max_peaks: int = MAX_PEAKS, ) -> tuple[list[float], list[float], list[float]]: """Deduplicate peaks from multiple window-radius runs. All peaks from every scale are pooled and sorted by intensity (strongest first). A peak is accepted only if no already-accepted peak lies within *merge_distance* pixels of it. This keeps the strongest detection of each physical spot regardless of which radius found it. """ all_peaks: list[tuple[float, float, float]] = [] for xs, ys, intensities in peaks_per_scale: for x, y, i in zip(xs, ys, intensities): all_peaks.append((i, x, y)) all_peaks.sort(reverse=True) # highest intensity first acc_x: list[float] = [] acc_y: list[float] = [] acc_i: list[float] = [] d2 = merge_distance ** 2 for intens, x, y in all_peaks: if len(acc_x) >= max_peaks: break too_close = any((x - ax) ** 2 + (y - ay) ** 2 < d2 for ax, ay in zip(acc_x, acc_y)) if not too_close: acc_x.append(x) acc_y.append(y) acc_i.append(intens) return acc_x, acc_y, acc_i
[docs] def find_peaks_pf9_multiscale( image: np.ndarray, mask: np.ndarray | None, window_radii, # int OR list[int] min_sigma: float, min_peak_over_neighbors: float, min_snr_max_pixel: float, min_snr_peak_pixels: float, min_snr_whole_peak: float, max_peaks: int = MAX_PEAKS, ) -> tuple[list[float], list[float], list[float]]: """Run pf9 with one or more window radii and merge the results. When *window_radii* is a single ``int`` this is identical to ``find_peaks_pf9``. When it is a list, each radius is run independently and the resulting peak lists are merged with :func:`_merge_peaks` (strongest peaks kept, duplicates within ``max(radii) + 1`` pixels suppressed). This lets you cover both dense patterns (small *r* keeps background ring outside neighbouring peaks) and sparse patterns (large *r* gives a more stable background estimate) in the same run. """ if isinstance(window_radii, (int, np.integer)): return find_peaks_pf9( image, mask, int(window_radii), min_sigma, min_peak_over_neighbors, min_snr_max_pixel, min_snr_peak_pixels, min_snr_whole_peak, max_peaks, ) radii = [int(r) for r in window_radii] peaks_per_scale = [ find_peaks_pf9( image, mask, r, min_sigma, min_peak_over_neighbors, min_snr_max_pixel, min_snr_peak_pixels, min_snr_whole_peak, max_peaks, ) for r in radii ] merge_dist = float(max(radii)) + 1.0 return _merge_peaks(peaks_per_scale, merge_distance=merge_dist, max_peaks=max_peaks)
# --------------------------------------------------------------------------- # Dask-delayed batch helper # --------------------------------------------------------------------------- @delayed def _load_and_process_batch_pf9( h5file_path, indices, x0_list, y0_list, window_radius, min_sigma, min_peak_over_neighbors, min_snr_max_pixel, min_snr_peak_pixels, min_snr_whole_peak, min_res, max_res, pixmask=None, ): """Dask-delayed: load and run pf9 on one batch of frames.""" batch_results = [] with h5py.File(h5file_path, 'r', locking=False) as f: for i, frame_idx in enumerate(indices): image_data = f['entry/data/images'][frame_idx] x0 = x0_list[i] y0 = y0_list[i] X, Y = np.meshgrid( range(image_data.shape[1]), range(image_data.shape[0]) ) R = np.sqrt((X - x0) ** 2 + (Y - y0) ** 2).astype(np.float32) mask = np.ones_like(image_data, dtype=np.int8) mask[R > max_res] = 0 mask[R < min_res] = 0 if pixmask is not None: mask = mask * pixmask xs, ys, intensities = find_peaks_pf9_multiscale( image_data, mask, window_radius, min_sigma, min_peak_over_neighbors, min_snr_max_pixel, min_snr_peak_pixels, min_snr_whole_peak, ) nPeaks = len(xs) fill = [0] * (MAX_PEAKS - nPeaks) batch_results.append({ 'index': frame_idx, 'nPeaks': nPeaks, 'peakXPosRaw': np.array(xs + fill, dtype=float), 'peakYPosRaw': np.array(ys + fill, dtype=float), 'peakTotalIntensity': np.array(intensities + fill, dtype=float), }) return batch_results # --------------------------------------------------------------------------- # File-level Dask entry point # ---------------------------------------------------------------------------
[docs] def findpeaks9_dask_file( client, h5file_path: str, x0, y0, window_radius: int, min_sigma: float, min_peak_over_neighbors: float, min_snr_max_pixel: float, min_snr_peak_pixels: float, min_snr_whole_peak: float, min_res: float, max_res: float, batch_size: int, progress_callback=None, outputfolder_path: str | None = None, ) -> None: """Run pf9 over an HDF5 frame stack using Dask, then write results. Output datasets are identical to the pf8 pipeline so both algorithms are interchangeable downstream: * ``/entry/data/nPeaks`` * ``/entry/data/peakXPosRaw`` * ``/entry/data/peakYPosRaw`` * ``/entry/data/peakTotalIntensity`` Parameters ---------- client : dask.distributed.Client h5file_path : str x0, y0 : int or ``"h5"`` Beam centre. Pass ``"h5"`` for both to read per-frame centres from ``entry/data/center_x`` / ``center_y``. window_radius : int min_sigma : float min_peak_over_neighbors : float min_snr_max_pixel : float min_snr_peak_pixels : float min_snr_whole_peak : float min_res, max_res : float Radial resolution ring limits in pixels. batch_size : int progress_callback : callable(done, total) or None outputfolder_path : str or None If given, saves ``peak_counts.npy`` and ``paramdump.txt`` there. """ with h5py.File(h5file_path, 'r') as f: num_frames = f['entry/data/images'].shape[0] shapex = f['entry/data/images'].shape[1] shapey = f['entry/data/images'].shape[2] num_batches = (num_frames + batch_size - 1) // batch_size if x0 == 'h5' and y0 == 'h5': center_x_values = list(f['entry/data/center_x'][:num_frames]) center_y_values = list(f['entry/data/center_y'][:num_frames]) log_print('Using centers from file') else: center_x_values = [x0] * num_frames center_y_values = [y0] * num_frames pixelmask = None pm_ds = get_mask_dataset(f) if pm_ds is not None: if pm_ds.ndim == 2 and pm_ds.shape == (shapex, shapey): pixelmask = pm_ds[()].astype(np.int8) log_start(h5file_path, f'Using pixel mask ({pm_ds.name}) from file') else: log_start(h5file_path, f'Ignoring pixel mask ({pm_ds.name}) with ' f'unsupported dimensions: {pm_ds.ndim}D ' f'shape {pm_ds.shape}') with h5py.File(h5file_path, 'r+') as f: write_nxprocess_peakfinding( f, { 'window_radius': window_radius, 'min_sigma': min_sigma, 'min_peak_over_neighbors': min_peak_over_neighbors, 'min_snr_max_pixel': min_snr_max_pixel, 'min_snr_peak_pixels': min_snr_peak_pixels, 'min_snr_whole_peak': min_snr_whole_peak, 'min_res': min_res, 'max_res': max_res, 'x0': x0, 'y0': y0, 'batch_size': batch_size, }, program='coseda.peakfinding.peakfinder9', inputs=['/entry/data/images'], outputs=[ '/entry/data/nPeaks', '/entry/data/peakTotalIntensity', '/entry/data/peakXPosRaw', '/entry/data/peakYPosRaw', ], ) # Build delayed task graph delayed_tasks = [] for batch_idx in range(num_batches): start_idx = batch_idx * batch_size end_idx = min(start_idx + batch_size, num_frames) indices = range(start_idx, end_idx) batch_x0 = center_x_values[start_idx:end_idx] batch_y0 = center_y_values[start_idx:end_idx] delayed_tasks.append(_load_and_process_batch_pf9( h5file_path, indices, batch_x0, batch_y0, window_radius, min_sigma, min_peak_over_neighbors, min_snr_max_pixel, min_snr_peak_pixels, min_snr_whole_peak, min_res, max_res, pixmask=pixelmask, )) futures = client.compute(delayed_tasks) if progress_callback is None: with ProgressBar(): results = client.gather(futures) all_results = [r for batch in results for r in batch] else: done = 0 total = len(futures) all_results = [] for future in as_completed(futures): all_results.extend(future.result()) done += 1 progress_callback(done, total) # Write results with h5py.File(h5file_path, 'r+') as f: for ds in ['nPeaks', 'peakTotalIntensity', 'peakXPosRaw', 'peakYPosRaw']: full = f'entry/data/{ds}' if full in f: del f[full] f.create_dataset('entry/data/nPeaks', shape=(num_frames,), dtype=int) f.create_dataset('entry/data/peakTotalIntensity', shape=(num_frames, MAX_PEAKS), dtype=float) f.create_dataset('entry/data/peakXPosRaw', shape=(num_frames, MAX_PEAKS), dtype=float) f.create_dataset('entry/data/peakYPosRaw', shape=(num_frames, MAX_PEAKS), dtype=float) for result in all_results: idx = result['index'] f['entry/data/nPeaks'][idx] = result['nPeaks'] f['entry/data/peakTotalIntensity'][idx] = result['peakTotalIntensity'][:MAX_PEAKS] f['entry/data/peakXPosRaw'][idx] = result['peakXPosRaw'][:MAX_PEAKS] f['entry/data/peakYPosRaw'][idx] = result['peakYPosRaw'][:MAX_PEAKS] ensure_peak_nxdata(f) if outputfolder_path is not None: indices_out = [r['index'] for r in all_results] npeaks_out = [r['nPeaks'] for r in all_results] np.save( os.path.join(outputfolder_path, 'peak_counts.npy'), np.vstack((indices_out, npeaks_out)).T, ) with open(os.path.join(outputfolder_path, 'paramdump.txt'), 'w') as pf: pf.write(f'x0={x0}\n') pf.write(f'y0={y0}\n') pf.write(f'window_radius={window_radius}\n') pf.write(f'min_sigma={min_sigma}\n') pf.write(f'min_peak_over_neighbors={min_peak_over_neighbors}\n') pf.write(f'min_snr_max_pixel={min_snr_max_pixel}\n') pf.write(f'min_snr_peak_pixels={min_snr_peak_pixels}\n') pf.write(f'min_snr_whole_peak={min_snr_whole_peak}\n') pf.write(f'min_res={min_res}\n') pf.write(f'max_res={max_res}\n') pf.write(f'batch_size={batch_size}\n')
# --------------------------------------------------------------------------- # Batch CLI / settings entry point # ---------------------------------------------------------------------------
[docs] def findpeaks9_dask_batch( input_path: str, batch_size: int = 1000, progress_callback=None, ) -> None: """Batch entry point: run pf9 for each INI file in *input_path*. Reads the following keys from the ``[Parameters]`` section of each INI: ============================== ======= ===================================== Key Type Description ============================== ======= ===================================== ``pf9_window_radius`` int Half-side of the background ring ``pf9_min_sigma`` float Minimum background std ``pf9_min_peak_over_neighbors`` float ADU excess over ring neighbours ``pf9_min_snr_max_pixel`` float SNR threshold on the max pixel ``pf9_min_snr_peak_pixels`` float SNR threshold for peak-mass pixels ``pf9_min_snr_whole_peak`` float Integrated peak-mass SNR threshold ``pf9_min_res`` float Inner radial mask limit (pixels) ``pf9_max_res`` float Outer radial mask limit (pixels) ``pf9_x0`` int|h5 Beam centre X (or ``"h5"``) ``pf9_y0`` int|h5 Beam centre Y (or ``"h5"``) ============================== ======= ===================================== """ configfiles, input_path = handle_input(input_path) log_print('following .ini files were found and will be processed:') log_print(str(configfiles)) filecount = 1 for configfile in configfiles: outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = \ config_to_paths(configfile) timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') graphicsfolder_path = os.path.join( outputfolder_path, f'findpeaks9_{timestamp}' ) os.makedirs(graphicsfolder_path, exist_ok=True) config = read_config(configfile) required = [ 'pf9_window_radius', 'pf9_min_sigma', 'pf9_min_peak_over_neighbors', 'pf9_min_snr_max_pixel', 'pf9_min_snr_peak_pixels', 'pf9_min_snr_whole_peak', 'pf9_min_res', 'pf9_max_res', ] for param in required: if not config.has_option('Parameters', param): with open(logfile_path, 'a') as lf: lf.write( f'{datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]}; ' f'Error: {param} not defined, pf9 peak finding interrupted\n' ) raise Exception(f'{param} not defined') _wr_raw = config.get('Parameters', 'pf9_window_radius').strip() window_radius = ( [int(x.strip()) for x in _wr_raw.split(',')] if ',' in _wr_raw else int(_wr_raw) ) min_sigma = float(config.get('Parameters', 'pf9_min_sigma')) min_peak_over_neighbors = float(config.get('Parameters', 'pf9_min_peak_over_neighbors')) min_snr_max_pixel = float(config.get('Parameters', 'pf9_min_snr_max_pixel')) min_snr_peak_pixels = float(config.get('Parameters', 'pf9_min_snr_peak_pixels')) min_snr_whole_peak = float(config.get('Parameters', 'pf9_min_snr_whole_peak')) min_res = float(config.get('Parameters', 'pf9_min_res')) max_res = float(config.get('Parameters', 'pf9_max_res')) try: x0_input = config.get('Parameters', 'pf9_x0') y0_input = config.get('Parameters', 'pf9_y0') if x0_input.lower() == 'h5' and y0_input.lower() == 'h5': x0 = y0 = 'h5' else: x0 = int(x0_input) y0 = int(y0_input) except Exception as e: log_start(logfile_path, f'Error reading beam position: {e}') x0 = y0 = 'h5' client = DaskClientManager.get_client() dashboard_address, _, _, _ = DaskClientManager.get_client_info() DaskClientManager.log_client_info(logfile_path) log_print('') log_print(f'Click here to monitor progress: {dashboard_address}') log_print('') log_start( logfile_path, f'start pf9 peak finding, window_radius={window_radius}, ' f'min_sigma={min_sigma}, ' f'min_peak_over_neighbors={min_peak_over_neighbors}, ' f'min_snr_max_pixel={min_snr_max_pixel}, ' f'min_snr_peak_pixels={min_snr_peak_pixels}, ' f'min_snr_whole_peak={min_snr_whole_peak}, ' f'min_res={min_res}, max_res={max_res}, ' f'batch_size={batch_size}', ) try: findpeaks9_dask_file( client, h5file_path, x0, y0, window_radius, min_sigma, min_peak_over_neighbors, min_snr_max_pixel, min_snr_peak_pixels, min_snr_whole_peak, min_res, max_res, batch_size, progress_callback=progress_callback, outputfolder_path=graphicsfolder_path, ) except Exception as e: log_result(logfile_path, 'pf9 peak finding finished with error', str(e)) # Reuse the shared stats function from the pf8 module from coseda.peakfinding.findpeaks_dask import peakfinding_stats peakfinding_stats(h5file_path, logfile_path) try: write_maxres_dataset( h5file_path=h5file_path, center_x=x0, center_y=y0, logfile_path=logfile_path, ) except Exception as exc: log_start(logfile_path, f"Warning: Failed to update /entry/data/maxres: {exc}") if filecount < len(configfiles): log_print('') log_print( f'proceeding to next task (file {filecount}/{len(configfiles)})' ) log_print('') filecount += 1 log_print(f'batch finished ({filecount - 1}/{len(configfiles)})')