"""Peak finding with Diffractem's peakfinder_8 using Dask for batch parallelism."""
from coseda.logging_utils import log_print
import os
import sys
import numpy as np
from dask import delayed
from dask.distributed import Client, LocalCluster, as_completed
from dask.diagnostics import ProgressBar
import h5py
from diffractem.peakfinder8_extension import peakfinder_8
import datetime
from coseda.initialize import find_configfiles
from coseda.io import handle_input, parse_config, config_to_paths, read_config
from coseda.logging_utils import log_start, log_result, shoutout
from coseda.dask_client_manager import DaskClientManager
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
@delayed
def load_and_process_batch(h5file_path, indices, x0_list, y0_list, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, pixmask=None):
"""Load and process a batch of frames with per-frame or fixed centers and optional mask."""
batch_results = []
with h5py.File(h5file_path, 'r') as f:
for i, frame_idx in enumerate(indices):
image_data = f['entry/data/images'][frame_idx]
# Use x0 and y0 from the lists for this specific frame
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
pks = peakfinder_8(500, image_data.astype(np.float32), mask, R, image_data.shape[1], image_data.shape[0], 1, 1, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius)
nPeaks = len(pks[0]) if pks is not None else 0
if nPeaks == 0:
fill = [0] * 500
batch_results.append({
'index': frame_idx,
'nPeaks': 0,
'peakTotalIntensity': np.array(fill),
'peakXPosRaw': np.array(fill),
'peakYPosRaw': np.array(fill),
})
else:
fill = [0] * (500 - nPeaks)
batch_results.append({
'index': frame_idx,
'nPeaks': nPeaks,
'peakTotalIntensity': np.array(pks[2] + fill),
'peakXPosRaw': np.array(pks[0] + fill),
'peakYPosRaw': np.array(pks[1] + fill),
})
return batch_results
[docs]
def findpeaks_dask_file(client, h5file_path, x0, y0, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, batch_size, progress_callback=None, outputfolder_path=None):
"""
Run peakfinder_8 over an HDF5 stack using Dask batches, then write results once.
Supports per-frame centers from `center_x/y` (x0=y0="h5") or fixed centers,
applies an optional `/mask` (2D only), and emits paramdump/peak_counts when
`outputfolder_path` is provided.
"""
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 # Calculate the number of batches
# Check if x0 and y0 are set to "h5"
if x0 == "h5" and y0 == "h5":
# Read center positions from the HDF5 file for each frame
center_x_values = f['entry/data/center_x'][:num_frames]
center_y_values = f['entry/data/center_y'][:num_frames]
log_print('Using centers from file')
log_print(center_x_values)
else:
# Use fixed values for x0 and y0 for all frames
center_x_values = [x0] * num_frames
center_y_values = [y0] * num_frames
# Load pixel mask if present and only support 2D masks
pixelmask = None
pm_ds = get_mask_dataset(f)
if pm_ds is not None:
# Only accept 2D masks matching image dimensions
if pm_ds.ndim == 2 and pm_ds.shape == (shapex, shapey):
# Read into memory and cast
pixelmask = pm_ds[()]
pixelmask = pixelmask.astype(np.int8)
# Log mask usage once per file
log_start(h5file_path, f'Using pixel mask ({pm_ds.name}) from file')
else:
# Ignore unsupported mask formats
log_start(h5file_path, f'Ignoring pixel mask ({pm_ds.name}) with unsupported dimensions: {pm_ds.ndim}D shape {pm_ds.shape}')
# Record peakfinding parameters before processing.
with h5py.File(h5file_path, 'r+') as f:
write_nxprocess_peakfinding(
f,
{
"threshold": threshold,
"min_snr": min_snr,
"min_pix_count": min_pix_count,
"max_pix_count": max_pix_count,
"local_bg_radius": local_bg_radius,
"min_res": min_res,
"max_res": max_res,
"x0": x0,
"y0": y0,
"batch_size": batch_size,
},
program="coseda.peakfinding.findpeaks_dask",
inputs=["/entry/data/images"],
outputs=[
"/entry/data/nPeaks",
"/entry/data/peakTotalIntensity",
"/entry/data/peakXPosRaw",
"/entry/data/peakYPosRaw",
],
)
# Create a list to collect all delayed tasks for processing batches
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)
# Get the appropriate x0 and y0 values for the batch
batch_x0 = center_x_values[start_idx:end_idx]
batch_y0 = center_y_values[start_idx:end_idx]
# Create a delayed task for each batch, passing batch-specific x0 and y0
delayed_task = load_and_process_batch(
h5file_path, indices, batch_x0, batch_y0, threshold, min_snr,
min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, pixmask=pixelmask
)
delayed_tasks.append(delayed_task)
# Submit all tasks to Dask and gather results
futures = client.compute(delayed_tasks)
if progress_callback is None:
# old behavior: console progress bar
with ProgressBar():
results = client.gather(futures)
all_results = [result for batch_results in results for result in batch_results]
else:
# new behavior: GUI progress callback
done = 0
total = len(futures)
all_results = []
for future in as_completed(futures):
batch_results = future.result()
all_results.extend(batch_results)
done += 1
progress_callback(done, total)
# Write results back to the HDF5 file in one go
with h5py.File(h5file_path, 'r+') as f:
# Delete datasets if they already exist
for dataset_name in ['nPeaks', 'peakTotalIntensity', 'peakXPosRaw', 'peakYPosRaw']:
full_name = f'entry/data/{dataset_name}'
if full_name in f:
del f[full_name]
# Create datasets
f.create_dataset('entry/data/nPeaks', shape=(num_frames,), dtype=int)
f.create_dataset('entry/data/peakTotalIntensity', shape=(num_frames, 500), dtype=float)
f.create_dataset('entry/data/peakXPosRaw', shape=(num_frames, 500), dtype=float)
f.create_dataset('entry/data/peakYPosRaw', shape=(num_frames, 500), dtype=float)
# Write all collected results to the datasets
for result in all_results:
idx = result['index']
f['entry/data/nPeaks'][idx] = result['nPeaks']
f['entry/data/peakTotalIntensity'][idx] = result['peakTotalIntensity'][:500]
f['entry/data/peakXPosRaw'][idx] = result['peakXPosRaw'][:500]
f['entry/data/peakYPosRaw'][idx] = result['peakYPosRaw'][:500]
ensure_peak_nxdata(f)
# Save numpy array of [index, nPeaks] for each frame
if outputfolder_path is not None:
indices = [res['index'] for res in all_results]
n_peaks_list = [res['nPeaks'] for res in all_results]
peaks_array = np.vstack((indices, n_peaks_list)).T # shape (num_frames, 2)
np.save(os.path.join(outputfolder_path, 'peak_counts.npy'), peaks_array)
# Dump run parameters to paramdump.txt
param_file = os.path.join(outputfolder_path, 'paramdump.txt')
with open(param_file, 'w') as pf:
pf.write(f'x0={x0}\n')
pf.write(f'y0={y0}\n')
pf.write(f'threshold={threshold}\n')
pf.write(f'min_snr={min_snr}\n')
pf.write(f'min_pix_count={min_pix_count}\n')
pf.write(f'max_pix_count={max_pix_count}\n')
pf.write(f'local_bg_radius={local_bg_radius}\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')
[docs]
def peakfinding_stats(h5file_path, logfile_path):
"""Calculate and log peak count statistics for an HDF5 file."""
with h5py.File(h5file_path, 'r') as f:
# Load the 'nPeaks' dataset
n_peaks = f['entry/data/nPeaks'][:]
num_frames = len(n_peaks)
# Calculate the maximum number of peaks
max_number = n_peaks.max()
# Define the bins with actual ranges
empty_bin = n_peaks == 0
bin_1 = (n_peaks > 0) & (n_peaks <= 0.25 * max_number)
bin_2 = (n_peaks > 0.25 * max_number) & (n_peaks <= 0.5 * max_number)
bin_3 = (n_peaks > 0.5 * max_number) & (n_peaks <= 0.75 * max_number)
bin_4 = (n_peaks > 0.75 * max_number) & (n_peaks <= max_number)
# Count frames in each bin and calculate percentages
empty_count = np.sum(empty_bin)
bin_1_count = np.sum(bin_1)
bin_2_count = np.sum(bin_2)
bin_3_count = np.sum(bin_3)
bin_4_count = np.sum(bin_4)
empty_percentage = (empty_count / num_frames) * 100
bin_1_percentage = (bin_1_count / num_frames) * 100
bin_2_percentage = (bin_2_count / num_frames) * 100
bin_3_percentage = (bin_3_count / num_frames) * 100
bin_4_percentage = (bin_4_count / num_frames) * 100
# Calculate peak ranges for each bin
bin_1_range = (1, int(0.25 * max_number))
bin_2_range = (int(0.25 * max_number) + 1, int(0.5 * max_number))
bin_3_range = (int(0.5 * max_number) + 1, int(0.75 * max_number))
bin_4_range = (int(0.75 * max_number) + 1, max_number)
# Log results using log_start
log_start(logfile_path, f'Peakfinding statistics:\n')
log_start(logfile_path, f'{empty_count} frames ({empty_percentage:.2f}%) are empty')
log_start(logfile_path, f'{bin_1_count} frames ({bin_1_percentage:.2f}%) have between {bin_1_range[0]} and {bin_1_range[1]} peaks')
log_start(logfile_path, f'{bin_2_count} frames ({bin_2_percentage:.2f}%) have between {bin_2_range[0]} and {bin_2_range[1]} peaks')
log_start(logfile_path, f'{bin_3_count} frames ({bin_3_percentage:.2f}%) have between {bin_3_range[0]} and {bin_3_range[1]} peaks')
log_start(logfile_path, f'{bin_4_count} frames ({bin_4_percentage:.2f}%) have between {bin_4_range[0]} and {bin_4_range[1]} peaks')
if empty_percentage > 10:
log_print('')
log_start(logfile_path, "Using the file stripping function can significantly reduce disk space consumption!")
[docs]
def findpeaks_dask_batch(input_path, batch_size=1000, progress_callback=None):
"""Batch entry point: run Dask peak finding for each INI in a folder/file/list."""
configfiles, input_path = handle_input(input_path)
log_print(f'following .ini files were found and will be processed:')
log_print(f'{configfiles}')
filecount = 1
for configfile in configfiles:
# Get paths
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
# Create timestamped graphics folder for peak outputs
timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
graphicsfolder_path = os.path.join(outputfolder_path, f'findpeaks_{timestamp}')
os.makedirs(graphicsfolder_path, exist_ok=True)
# Get config
config = read_config(configfile)
# Check if necessary parameters are defined
for param in ['peakfinding_min_snr','peakfinding_min_pix_count','peakfinding_max_pix_count','peakfinding_local_bg_radius','peakfinding_min_res']:
if not config.has_option('Parameters', param):
with open(f'{logfile_path}', 'a') as file:
file.write(f'{datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]}; Error: {param.split("_", 1)[1]} not defined, peak finding interrupted\n')
raise Exception(f'{param.split("_", 1)[1]} not defined')
# Load parameters
threshold = float(config.get('Parameters', 'peakfinding_threshold'))
min_snr = float(config.get('Parameters', 'peakfinding_min_snr'))
min_pix_count = float(config.get('Parameters', 'peakfinding_min_pix_count'))
max_pix_count = float(config.get('Parameters', 'peakfinding_max_pix_count'))
local_bg_radius = float(config.get('Parameters', 'peakfinding_local_bg_radius'))
min_res = float(config.get('Parameters', 'peakfinding_min_res'))
max_res = float(config.get('Parameters', 'peakfinding_max_res'))
try:
x0_input = config.get('Parameters', 'peakfinding_x0')
y0_input = config.get('Parameters', 'peakfinding_y0')
if x0_input.lower() == "h5" and y0_input.lower() == "h5":
# Set x0 and y0 to the string "h5"
x0 = "h5"
y0 = "h5"
else:
# Convert inputs to integers if they are not "h5"
x0 = int(x0_input)
y0 = int(y0_input)
except Exception as e:
log_start(logfile_path, f"Error in reading beam position: {e}")
# Default to center if an error occurs
x0 = "h5"
y0 = "h5"
client = DaskClientManager.get_client()
dashboard_address, num_workers, threads_per_worker, memory_per_worker = 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 peak finding using dask, threshold = {threshold}, min_snr = {min_snr}, min_pix_count = {min_pix_count}, max_pix_count = {max_pix_count}, local_bg_radius = {local_bg_radius}, min_res = {min_res}, max_res = {max_res}, batch size = {batch_size}')
try:
findpeaks_dask_file(client, h5file_path, x0, y0, threshold, min_snr, min_pix_count, max_pix_count,
local_bg_radius, min_res, max_res, batch_size, progress_callback=progress_callback, outputfolder_path=graphicsfolder_path)
except Exception as e:
log_result(logfile_path, "peakfinding finished without errors", str(e))
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)})")