#!/usr/bin/env python3
"""Direct beam center finding via quadrant evenness and optional auto-thresholding."""
import h5py
import numpy as np
import configparser
import random
import sys
from coseda.initialize import config_to_paths, read_config
from coseda.logging_utils import log_start
from coseda.centerfinding.findcenters_beamstop import (
get_frame_stack_shape,
resolve_frame_dataset,
update_detector_shifts,
)
from coseda.dask_client_manager import DaskClientManager
from coseda.nexus.process import write_nxprocess_centerfinding
import os
from datetime import datetime
from distributed import as_completed
from coseda.centerrefinement.cancellation import check_cancel, RefinementCancelled
[docs]
def is_within_bounds(r, c, shape, box):
"""
Return True if a box of size `box` at (r, c) fits entirely inside an array of shape `shape`.
"""
rows, cols = shape
offsets = [(-box, -box), (-box, 0), (0, -box), (0, 0)]
for dr, dc in offsets:
rs, cs = r + dr, c + dc
re, ce = rs + box, cs + box
if rs < 0 or cs < 0 or re > rows or ce > cols:
return False
return True
[docs]
def compute_auto_significance_threshold(h5file_path, framepath, sample_count, patch_size, bin_size):
"""
Sample frames, bin to `patch_size` patches, and compute mean brightest/corner ratio.
Picks `sample_count` random frames, bins them coarsely, and for each compares the
brightest patch with the three darkest corners. Returns the mean ratio to serve
as an automatic significance threshold for beam detection.
"""
ratios = []
with h5py.File(h5file_path, 'r') as f:
images, _ = resolve_frame_dataset(f, framepath)
if images.ndim == 2:
return 0.0
total_frames = images.shape[0]
indices = random.sample(range(total_frames), min(sample_count, total_frames))
for idx in indices:
frame = images[idx]
rows, cols = frame.shape
# Bin into patch_size × patch_size
pr = patch_size
br = rows // pr
bc = cols // pr
cropped = frame[:br*pr, :bc*pr]
binned = cropped.reshape(br, pr, bc, pr).sum(axis=(1,3))
# brightest patch
i,j = np.unravel_index(np.argmax(binned), binned.shape)
brightest = binned[i, j]
# corner patches
corners = [binned[0,0], binned[0,-1], binned[-1,0], binned[-1,-1]]
darkest_three = sorted(corners)[:3]
bg = float(np.mean(darkest_three)) if darkest_three else 0.0
if bg > 0:
ratios.append(brightest / bg)
return float(np.mean(ratios)) if ratios else 0.0
[docs]
def find_primary_beam_direct(frame, box_size, search_radius, bin_size, threshold):
"""
Locate the beam by coarse binning, quadrant-evenness refinement, and significance test.
Uses an integral image to find the brightest bin, sweeps a search window to
minimize quadrant variance, then rejects the candidate if its local mean is
not sufficiently brighter than the image corners. Returns (row, col) or None.
"""
rows, cols = frame.shape
# Build integral image with 64-bit ints to avoid overflow
intimg = frame.astype(np.int64).cumsum(axis=0).cumsum(axis=1)
def region_sum(r1, c1, r2, c2):
"""Sum of frame[r1:r2, c1:c2] using integral image"""
total = intimg[r2-1, c2-1]
if r1 > 0:
total -= intimg[r1-1, c2-1]
if c1 > 0:
total -= intimg[r2-1, c1-1]
if r1 > 0 and c1 > 0:
total += intimg[r1-1, c1-1]
return total
# Coarse binning
br = rows // bin_size
bc = cols // bin_size
cropped = frame[:br*bin_size, :bc*bin_size]
binned = cropped.reshape(br, bin_size, bc, bin_size).sum(axis=(1,3))
i,j = np.unravel_index(np.argmax(binned), binned.shape)
brightest_row = i * bin_size + bin_size // 2
brightest_col = j * bin_size + bin_size // 2
# Refinement by minimizing quadrant stddev
best_center = (brightest_row, brightest_col)
min_evenness = float('inf')
for r in range(max(brightest_row - search_radius, 0), min(brightest_row + search_radius + 1, rows - box_size)):
for c in range(max(brightest_col - search_radius, 0), min(brightest_col + search_radius + 1, cols - box_size)):
# skip boxes that don't fully fit
if not is_within_bounds(r, c, frame.shape, box_size):
continue
# Quadrant sums via integral image
q1 = region_sum(r-box_size, c-box_size, r, c)
q2 = region_sum(r-box_size, c, r, c+box_size)
q3 = region_sum(r, c-box_size, r+box_size, c)
q4 = region_sum(r, c, r+box_size, c+box_size)
evenness = np.std([q1, q2, q3, q4])
if evenness < min_evenness:
min_evenness = evenness
best_center = (r, c)
# Significance test vs image corners
r0, c0 = best_center
if not is_within_bounds(r0, c0, frame.shape, box_size):
return None
# box corners for significance
beam_tiles = [
frame[r0-box_size:r0, c0-box_size:c0],
frame[r0-box_size:r0, c0:c0+box_size],
frame[r0:r0+box_size, c0-box_size:c0],
frame[r0:r0+box_size, c0:c0+box_size]
]
beam_mean = float(np.mean([np.mean(tile) for tile in beam_tiles]))
corners = [
frame[:box_size, :box_size],
frame[:box_size, -box_size:],
frame[-box_size:, :box_size],
frame[-box_size:, -box_size:]
]
corner_means = [float(np.mean(c)) for c in corners]
if len(corner_means) < 3:
return None
ref_bg = float(np.mean(sorted(corner_means)[:3]))
if ref_bg <= 0 or beam_mean / ref_bg < threshold:
return None
return best_center
[docs]
def interpolate_centers(centers):
"""
Linearly interpolate missing centers; returns arrays of x and y positions.
If fewer than two valid points are present, fills NaNs with zeros.
"""
n = len(centers)
xs = np.array([c[1] if c is not None else np.nan for c in centers], dtype=float)
ys = np.array([c[0] if c is not None else np.nan for c in centers], dtype=float)
valid = ~np.isnan(xs)
idx = np.arange(n)
if valid.sum() >= 2:
xs = np.interp(idx, idx[valid], xs[valid])
ys = np.interp(idx, idx[valid], ys[valid])
else:
# If too few valid, fill with zeros
xs = np.nan_to_num(xs)
ys = np.nan_to_num(ys)
return xs, ys
[docs]
def process_centers(h5file_path, framepath, interval, box_size, search_radius, bin_size, threshold, outputfolder_path, pixels_per_meter, logfile_path=None, progress_callback=None, stop_event=None):
"""
Detect beam centers on a subsampled set of frames, fit linear drift, and write results.
- Samples original frames by `interval` honoring stripped index mapping
- Runs Dask-parallel `find_primary_beam_direct`
- Saves raw detections, fit parameters, and writes `center_x/y` + detector shifts
"""
with h5py.File(h5file_path, 'r+') as f:
images, resolved_framepath = resolve_frame_dataset(f, framepath)
nframes, _, _, _ = get_frame_stack_shape(f, framepath)
if images.ndim == 2:
raise ValueError(
f"Frame dataset '{resolved_framepath}' is 2D (single image). Direct centerfinding expects a frame stack."
)
if resolved_framepath != (framepath if framepath.startswith("/") else f"/{framepath}"):
log_start(
logfile_path,
f"Resolved framepath '{framepath}' to dataset '{resolved_framepath}' for direct centerfinding.",
)
check_cancel(stop_event, "direct centerfinding setup")
# Load or create index mapping for original frames
if 'entry/data/index' in f:
index_map = f['entry/data/index'][:]
else:
# no index dataset: assume no frames have been stripped
index_map = np.arange(nframes, dtype=int)
orig_nframes = len(index_map)
stripped_nframes = images.shape[0]
# List of original frame indices that exist in stripped dataset
original_valid = np.where(index_map != -1)[0]
# Pick original frames based on interval
sample_orig = original_valid[::interval]
# Convert to stripped indices for processing
indices = [int(index_map[orig]) for orig in sample_orig]
# Create a timestamped output folder for direct results
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
outdir = os.path.join(outputfolder_path, f"findcenters_{ts}")
os.makedirs(outdir, exist_ok=True)
# Dump run parameters to paramdump.txt
param_file = os.path.join(outdir, 'paramdump.txt')
with open(param_file, 'w') as pf:
pf.write(f"framepath={framepath}\n")
pf.write(f"interval={interval}\n")
pf.write(f"box_size={box_size}\n")
pf.write(f"search_radius={search_radius}\n")
pf.write(f"bin_size={bin_size}\n")
pf.write(f"threshold={threshold}\n")
centers = [None] * nframes
# Parallel execution via Dask
client = DaskClientManager.get_client()
futures = {
idx: client.submit(
find_primary_beam_direct,
images[idx],
box_size,
search_radius,
bin_size,
threshold
)
for idx in indices
}
# Gather results with progress reporting
total = len(indices)
done = 0
centers_dict = {}
future_to_idx = {f: idx for idx, f in futures.items()}
try:
for future in as_completed(futures.values()):
check_cancel(stop_event, "direct centerfinding gather")
idx = future_to_idx[future]
result = future.result()
centers_dict[idx] = result
done += 1
if progress_callback:
progress_callback(done, total)
if stop_event is not None and stop_event.is_set():
break
finally:
if stop_event is not None and stop_event.is_set():
for fut in futures.values():
try:
fut.cancel()
except Exception:
pass
if stop_event is not None and stop_event.is_set():
raise RefinementCancelled("Center finding cancelled by user")
# Reconstruct ordered list of centers
centers_list = [centers_dict[idx] for idx in indices]
# Assign results back into full centers list and collect raw detections
centers = [None] * nframes
raw = []
for i, c in enumerate(centers_list):
stripped_idx = indices[i]
orig_idx = sample_orig[i]
centers[stripped_idx] = c
if c is not None:
raw.append([orig_idx, c[1], c[0]])
# Save raw detections for fitting: columns = frame, x, y
raw_arr = np.array(raw, dtype=float)
np.save(os.path.join(outdir, "direct_centers_raw.npy"), raw_arr)
# Fit linear trends (x and y vs frame index)
valid = [(i, c) for i, c in enumerate(centers) if c is not None]
n = nframes
if len(valid) >= 2:
# Use original frame indices corresponding to sample_orig for fitting
orig_idxs = [orig for orig, c in zip(sample_orig, centers_list) if c is not None]
coords = [c for c in centers_list if c is not None]
xs_valid = np.array([c[1] for c in coords], dtype=float)
ys_valid = np.array([c[0] for c in coords], dtype=float)
# fit line using original frame numbers
slope_x, intercept_x = np.polyfit(orig_idxs, xs_valid, 1)
slope_y, intercept_y = np.polyfit(orig_idxs, ys_valid, 1)
# Log fitted line equations
log_start(logfile_path, f"[Direct Fit] x = {slope_x:.6f}*frame + {intercept_x:.2f}; y = {slope_y:.6f}*frame + {intercept_y:.2f}")
# Compute R^2 for x and y fits
xs_pred = slope_x * np.array(orig_idxs) + intercept_x
ys_pred = slope_y * np.array(orig_idxs) + intercept_y
ss_res_x = np.sum((xs_valid - xs_pred)**2)
ss_tot_x = np.sum((xs_valid - xs_valid.mean())**2)
r2_x = 1 - ss_res_x/ss_tot_x if ss_tot_x>0 else 0.0
ss_res_y = np.sum((ys_valid - ys_pred)**2)
ss_tot_y = np.sum((ys_valid - ys_valid.mean())**2)
r2_y = 1 - ss_res_y/ss_tot_y if ss_tot_y>0 else 0.0
# Compute standard deviations
stdev_x = float(np.std(xs_valid))
stdev_y = float(np.std(ys_valid))
# Save fit parameters: [slope, intercept, r2, standard deviation] for x and y
fit_arr = np.array([
[slope_x, intercept_x, r2_x, stdev_x],
[slope_y, intercept_y, r2_y, stdev_y],
], dtype=float)
np.save(os.path.join(outdir, "direct_centers_fit.npy"), fit_arr)
# Build inverse map for all stripped frames
inv_map = np.full(stripped_nframes, -1, dtype=int)
for orig in original_valid:
inv_map[int(index_map[orig])] = orig
all_orig = inv_map
xs = slope_x * all_orig + intercept_x
ys = slope_y * all_orig + intercept_y
else:
# fallback: use first valid or zeros
if valid:
_, (r0, c0) = valid[0]
xs = np.full(n, c0, dtype=float)
ys = np.full(n, r0, dtype=float)
else:
xs = np.zeros(n, dtype=float)
ys = np.zeros(n, dtype=float)
# Write back
if 'entry/data/center_x' in f:
del f['entry/data/center_x']
if 'entry/data/center_y' in f:
del f['entry/data/center_y']
f.create_dataset('entry/data/center_x', data=xs, dtype='float64')
f.create_dataset('entry/data/center_y', data=ys, dtype='float64')
# Update detector shift datasets
update_detector_shifts(h5file_path, logfile_path, framepath, pixels_per_meter)
[docs]
def run_findcenters_direct(configfile, progress_callback=None, stop_event=None):
"""Entry point for direct beam finding given a single dataset INI."""
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
config = read_config(configfile)
sec = 'Parameters'
framepath = config.get('Paths', 'framepath')
interval = config.getint(sec, 'findcentersdirect_interval')
box_size = config.getint(sec, 'findcentersdirect_box_size')
search_radius = config.getint(sec, 'findcentersdirect_search_radius')
bin_size = config.getint(sec, 'findcentersdirect_bin_size')
# significance threshold optional
auto_threshold = True
if config.has_option(sec, 'findcentersdirect_significance_threshold'):
threshold = config.getfloat(sec, 'findcentersdirect_significance_threshold')
auto_threshold = False
else:
sample_ct = config.getint(sec, 'findcentersdirect_auto_sample_count')
patch_sz = config.getint(sec, 'findcentersdirect_auto_patch_size')
log_start(logfile_path, "Computing auto significance threshold")
threshold = compute_auto_significance_threshold(h5file_path, framepath, sample_ct, patch_sz, bin_size)
log_start(logfile_path, f"Auto significance threshold: {threshold}")
pixels_per_meter = config.getfloat('AcquisitionDetails', 'pixels_per_meter')
log_start(logfile_path, f"Running direct center-finder with interval={interval}, ")
process_centers(
h5file_path,
framepath,
interval,
box_size,
search_radius,
bin_size,
threshold,
outputfolder_path,
pixels_per_meter,
logfile_path,
progress_callback=progress_callback,
stop_event=stop_event
)
try:
with h5py.File(h5file_path, "r+") as workingfile:
write_nxprocess_centerfinding(
workingfile,
program="coseda.centerfinding.findcenters_direct",
method="direct",
input_path=h5file_path,
output_path=h5file_path,
parameters={
"interval": interval,
"box_size": box_size,
"search_radius": search_radius,
"bin_size": bin_size,
"threshold": threshold,
"pixels_per_meter": pixels_per_meter,
"auto_threshold": auto_threshold,
},
)
except Exception as exc:
log_start(logfile_path, f"Failed to write NXprocess centerfinding: {exc}")
[docs]
def write_centerfinderdirectsettings(input_path,
interval,
box_size,
search_radius,
bin_size,
auto_sample_count=None,
auto_patch_size=None,
auto_factor=None):
"""
Update INI files under [Parameters] with direct center-finding settings.
"""
from coseda.initialize import handle_input, shoutout, read_config, log_start, config_to_paths
configfiles, _ = handle_input(input_path)
for configfile in configfiles:
shoutout(configfile)
config = read_config(configfile)
section = 'Parameters'
if not config.has_section(section):
config.add_section(section)
# Required parameters
config.set(section, 'findcentersdirect_interval', str(interval))
config.set(section, 'findcentersdirect_box_size', str(box_size))
config.set(section, 'findcentersdirect_search_radius', str(search_radius))
config.set(section, 'findcentersdirect_bin_size', str(bin_size))
# Auto-threshold parameters
if auto_sample_count is not None:
config.set(section, 'findcentersdirect_auto_sample_count', str(auto_sample_count))
if auto_patch_size is not None:
config.set(section, 'findcentersdirect_auto_patch_size', str(auto_patch_size))
if auto_factor is not None:
config.set(section, 'findcentersdirect_auto_factor', str(auto_factor))
# Write back
with open(configfile, 'w', encoding='utf-8') as f:
config.write(f)
# Log the change
_, _, _, logfile_path, _, _ = config_to_paths(configfile)
log_start(
logfile_path,
f"Direct centerfinder settings set: interval={interval}, "
f"box_size={box_size}, search_radius={search_radius}, bin_size={bin_size}"
+ (f", auto_sample_count={auto_sample_count}" if auto_sample_count is not None else "")
+ (f", auto_patch_size={auto_patch_size}" if auto_patch_size is not None else "")
+ (f", auto_factor={auto_factor}" if auto_factor is not None else "")
)