"""Beamstop-based center finding with Friedel-pair clustering and drift fitting."""
from coseda.logging_utils import log_print
import matplotlib
matplotlib.use('Agg')
import os
import sys
import h5py
import numpy as np
import matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture
from sklearn.cluster import DBSCAN
from dask import delayed, compute
from dask.distributed import Client
import configparser
import gc
from scipy.stats import linregress
from datetime import datetime
from dask.distributed import as_completed
from coseda.centerrefinement.cancellation import check_cancel, RefinementCancelled
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.detector_geometry import (
corrected_to_raw_points,
load_detector_geometry,
raw_to_corrected,
raw_to_corrected_points,
validate_geometry_for_frame,
)
from coseda.nexus.logs import ensure_dense_logs
from coseda.nexus.goniometer import ensure_goniometer_transforms
from coseda.nexus.process import write_nxprocess_centerfinding
DEFAULT_CENTERFINDING_DBSCAN_EPS = 5.0
# -----------------------------------------------------------------------------
# HDF5 frame path helpers
def _normalize_h5_path(path: str) -> str:
raw = (path or "").strip()
if not raw:
raise ValueError("Empty HDF5 frame path.")
return raw if raw.startswith("/") else f"/{raw}"
def _decode_h5_attr(value):
if isinstance(value, bytes):
return value.decode("utf-8", errors="ignore")
return value
def _resolve_dataset_from_group(group: h5py.Group, token: str):
token = (token or "").strip()
if not token:
return None
if token in group:
obj = group[token]
if isinstance(obj, h5py.Dataset):
return obj
return None
[docs]
def resolve_frame_dataset(h5file: h5py.File, framepath: str):
"""
Resolve a frame path to a concrete HDF5 dataset.
Accepts either a dataset path (e.g. `/entry/data/images`) or a group path
(e.g. `/entry/data/image_data` NXdata), where the image dataset is resolved
via `signal`, then common child names, then a single 2D/3D child dataset.
"""
normalized = _normalize_h5_path(framepath)
obj = h5file.get(normalized)
if obj is None:
obj = h5file.get(normalized.lstrip("/"))
if obj is None:
raise KeyError(f"Frame path '{framepath}' not found in HDF5 file.")
if isinstance(obj, h5py.Dataset):
if obj.ndim < 2:
raise ValueError(
f"Frame dataset '{obj.name}' must have at least 2 dimensions, got shape {obj.shape}."
)
return obj, obj.name
if not isinstance(obj, h5py.Group):
raise ValueError(
f"Frame path '{framepath}' resolved to unsupported HDF5 object type {type(obj).__name__}."
)
# Try NXdata signal first.
signal = _decode_h5_attr(obj.attrs.get("signal"))
if isinstance(signal, (list, tuple, np.ndarray)) and signal:
signal = _decode_h5_attr(signal[0])
if isinstance(signal, str):
candidate = _resolve_dataset_from_group(obj, signal)
if candidate is None and signal.startswith("/") and signal in h5file:
maybe = h5file[signal]
if isinstance(maybe, h5py.Dataset):
candidate = maybe
if candidate is not None and candidate.ndim >= 2:
return candidate, candidate.name
# Common child names.
for name in ("images", "data"):
candidate = _resolve_dataset_from_group(obj, name)
if candidate is not None and candidate.ndim >= 2:
return candidate, candidate.name
# Fallback: choose a single plausible dataset from this group.
candidates = [
child for child in obj.values()
if isinstance(child, h5py.Dataset) and child.ndim >= 2
]
if len(candidates) == 1:
return candidates[0], candidates[0].name
if len(candidates) > 1:
image_like = [ds for ds in candidates if "image" in ds.name.lower()]
if len(image_like) == 1:
return image_like[0], image_like[0].name
names = ", ".join(ds.name for ds in candidates[:6])
raise ValueError(
f"Frame path '{framepath}' points to group '{obj.name}' with multiple dataset candidates: {names}. "
"Set [Paths] framepath to the image dataset path."
)
raise ValueError(
f"Frame path '{framepath}' points to group '{obj.name}' without a 2D/3D image dataset."
)
[docs]
def get_frame_stack_shape(h5file: h5py.File, framepath: str):
"""Return (num_frames, frame_height, frame_width, resolved_dataset_path)."""
dataset, resolved_path = resolve_frame_dataset(h5file, framepath)
shape = tuple(int(v) for v in dataset.shape)
if dataset.ndim == 2:
return 1, shape[0], shape[1], resolved_path
# For standard image stacks this is (N, H, W).
return shape[0], shape[-2], shape[-1], resolved_path
# -----------------------------------------------------------------------------
# Plotting helpers
[docs]
def plot_deviation_cloud(deviations, mean, current_center, batch_index, itcount, graphicsfolder_path):
"""Scatter the Friedel-pair deviation cloud with the current mean overlaid."""
import os
import matplotlib.pyplot as plt
if deviations.size == 0:
return
plt.figure(figsize=(6, 6))
plt.scatter(deviations[:, 0], deviations[:, 1], label="Deviation Cloud", marker='o', s=0.5)
if mean is not None:
plt.scatter(mean[0], mean[1], color='red',
label=f"Center ({current_center[0]:.3f}, {current_center[1]:.3f})",
marker='x')
plt.xlabel("X-deviation")
plt.ylabel("Y-deviation")
plt.legend(loc='lower right')
plt.title(f"Deviation Cloud (batch {batch_index}/it {itcount})")
filename = os.path.join(graphicsfolder_path,
f'centerfinding_deviations_batch{batch_index}_it{itcount:03d}.png')
plt.savefig(filename)
plt.close()
[docs]
def plot_linear_fit(batch_positions, batch_centers, frame_positions, updated_x, updated_y, graphicsfolder_path):
import os
import matplotlib.pyplot as plt
plt.figure()
xs = batch_positions
ysx = [c[0] for c in batch_centers]
ysy = [c[1] for c in batch_centers]
plt.scatter(xs, ysx, s=5, label='Batch Centers X')
plt.scatter(xs, ysy, s=5, label='Batch Centers Y')
plt.plot(frame_positions, updated_x, label='Fitted Line X')
plt.plot(frame_positions, updated_y, label='Fitted Line Y')
plt.xlabel('Frame Position')
plt.ylabel('Center Value')
plt.legend()
plt.title('Center Finding Linear Fit')
plt.savefig(os.path.join(graphicsfolder_path, 'centerfinding_linearfit.png'))
plt.close()
[docs]
def plot_batch_centers(frame_positions, updated_x, updated_y, graphicsfolder_path):
import os
import matplotlib.pyplot as plt
plt.figure()
plt.scatter(frame_positions, updated_x, s=1, label='Frame Centers X')
plt.scatter(frame_positions, updated_y, s=1, label='Frame Centers Y')
plt.xlabel('Frame Position')
plt.ylabel('Center Value')
plt.legend()
plt.title('Center Finding Batch-Based')
plt.savefig(os.path.join(graphicsfolder_path, 'centerfinding_batch.png'))
plt.close()
# -----------------------------------------------------------------------------
# Function to divide indices into batches, ensuring at least 4 batches, with logging
[docs]
def divide_into_batches(total_indices, batch_size, logfile_path):
"""Split surviving frame indices into batches, enforcing at least four batches."""
total_frames = len(total_indices)
num_batches = (total_frames + batch_size - 1) // batch_size # Calculate initial number of batches
log_start(logfile_path, f'Initial number of batches: {num_batches}')
# Ensure at least 4 batches
if num_batches < 4:
batch_size = (total_frames + 3) // 4 # Adjust batch size to ensure at least 4 batches
num_batches = (total_frames + batch_size - 1) // batch_size
log_start(logfile_path, f'Adjusted batch size to {batch_size} to ensure at least 4 batches')
else:
log_start(logfile_path, f'Using provided batch size: {batch_size}')
log_start(logfile_path, f'Final number of batches: {num_batches}')
batches = [total_indices[i:i + batch_size] for i in range(0, total_frames, batch_size)]
return batches, total_frames, batch_size
# Function to fit Gaussian to the largest cluster of deviations
[docs]
def fit_gaussian_to_largest_cluster(
deviations,
min_samples_fraction,
dbscan_eps=DEFAULT_CENTERFINDING_DBSCAN_EPS,
):
"""Cluster deviations with DBSCAN and fit a 1-component GMM to the largest cluster."""
if deviations.size == 0:
return None, None
# Use DBSCAN for clustering
clustering = DBSCAN(
eps=float(dbscan_eps),
min_samples=max(int(min_samples_fraction * len(deviations)), 2),
).fit(deviations)
cluster_labels = clustering.labels_
# Find the cluster labels (ignoring -1, which are 'noise' points)
unique_labels = set(cluster_labels)
if -1 in unique_labels:
unique_labels.remove(-1)
# Find the largest cluster
cluster_sizes = [np.sum(cluster_labels == i) for i in unique_labels]
if len(cluster_sizes) == 0:
return None, None # Handle the case where no cluster is found
largest_cluster_index = list(unique_labels)[np.argmax(cluster_sizes)]
# Filter deviations belonging to the largest cluster
largest_cluster_deviations = deviations[cluster_labels == largest_cluster_index]
# Fit Gaussian to the largest cluster
gmm = GaussianMixture(n_components=1).fit(largest_cluster_deviations)
mean = gmm.means_[0].astype(np.float32, copy=False)
covariance = gmm.covariances_[0]
gc.collect()
return mean, covariance
# Try eliminating nested loops for Friedel pairs
[docs]
def find_friedel_pairs(peak_pos: np.ndarray,
center: np.ndarray,
tol: float,
block: int = 2048,
max_pairs: int | None = None) -> np.ndarray:
"""
Vectorised Friedel-pair search that bounds memory to ~32MB at block=2048.
Shifts peaks relative to the current center and searches for pairs whose
sum-norm is below `tol` using blockwise matrix operations.
"""
shifted = (peak_pos - center).astype(np.float32, copy=False)
keep = []
kept_count = 0
pair_cap = None
if max_pairs is not None:
try:
pair_cap = int(max_pairs)
except (TypeError, ValueError):
pair_cap = None
if pair_cap is not None and pair_cap <= 0:
pair_cap = None
def _append_pairs(arr: np.ndarray):
nonlocal kept_count
if arr.size == 0:
return False
if pair_cap is None:
keep.append(arr)
kept_count += int(arr.shape[0])
return False
remaining = pair_cap - kept_count
if remaining <= 0:
return True
if arr.shape[0] > remaining:
arr = arr[:remaining]
keep.append(arr)
kept_count += int(arr.shape[0])
return kept_count >= pair_cap
for i0 in range(0, len(shifted), block):
a = shifted[i0:i0 + block] # (b,2)
# in‑block pairs: upper triangle only
inblock = a[:, None, :] + a[None, :, :] # (b,b,2)
m = np.triu(np.linalg.norm(inblock, axis=-1) < tol, 1)
if m.any():
if _append_pairs(inblock[m]):
break
# cross pairs with following blocks
stop_outer = False
for j0 in range(i0 + block, len(shifted), block):
b = shifted[j0:j0 + block] # (b,2)
cross = a[:, None, :] + b[None, :, :] # (b,b,2)
m = np.linalg.norm(cross, axis=-1) < tol
if m.any():
if _append_pairs(cross[m]):
stop_outer = True
break
if stop_outer:
break
return (np.concatenate(keep, axis=0)
if keep else np.empty((0, 2), np.float32))
# Function to refine center with Friedel pairs
[docs]
def refine_center_with_friedel(h5file_path, current_center, tolerance, min_peaks, resolution_limit,
subset_indices, batch_index, logfile_path=None,
detector_geometry=None, max_pairs_per_frame=None):
"""
Extract Friedel-pair deviations for a batch of frames and return concatenated deltas.
- Skips frames with fewer than `min_peaks`
- Applies a radial cutoff (`resolution_limit`) around the current center
- Validates shapes before appending to the deviation list
"""
all_deviations = []
with h5py.File(h5file_path, 'r') as workingfile:
indices_to_process = subset_indices if subset_indices is not None else range(len(workingfile['entry/data/images']))
log_start(logfile_path, f'Batch {batch_index}: Looking for Friedel pairs')
for frame_idx in indices_to_process:
num_peaks = int(workingfile['entry/data/nPeaks'][frame_idx])
if num_peaks < min_peaks:
continue
peak_xpos_raw = workingfile['entry/data/peakXPosRaw'].astype('float32')[frame_idx, :num_peaks]
peak_ypos_raw = workingfile['entry/data/peakYPosRaw'].astype('float32')[frame_idx, :num_peaks]
peak_positions_raw = np.column_stack((peak_xpos_raw, peak_ypos_raw))
if detector_geometry is not None:
peak_positions, _ = raw_to_corrected_points(peak_positions_raw, detector_geometry)
else:
peak_positions = peak_positions_raw
# Filter out peaks based on resolution_limit
distances = np.linalg.norm(peak_positions - current_center, axis=1)
peak_positions = peak_positions[distances < resolution_limit]
deviations = find_friedel_pairs(
peak_positions,
current_center,
tolerance,
max_pairs=max_pairs_per_frame,
)
# Normalize deviations into a 2‑column array
if deviations.size > 0:
# If flat 1D array of even length, reshape into rows of two
if deviations.ndim == 1:
if deviations.shape[0] == 2:
deviations = deviations.reshape(1, 2)
elif deviations.shape[0] % 2 == 0:
deviations = deviations.reshape(deviations.shape[0] // 2, 2)
# Only accept 2D arrays with exactly 2 columns
if deviations.ndim == 2 and deviations.shape[1] == 2:
all_deviations.append(deviations)
else:
log_start(logfile_path,
f'Batch {batch_index}: Skipping deviations with invalid shape {deviations.shape}')
log_start(logfile_path, f'Batch {batch_index}: Finished finding Friedel pairs')
#return np.asarray(all_deviations, dtype=np.float32)
# Concatenate all blocks (or return an empty (0,2) array)
if all_deviations:
return np.concatenate(all_deviations, axis=0).astype(np.float32, copy=False)
else:
return np.empty((0, 2), dtype=np.float32)
# Function to iteratively refine the center using Gaussian fitting on subsets
[docs]
def iterative_gaussian_fitting_on_subset(h5file_path, initial_center, subset_indices, tolerance,
min_peaks, resolution_limit, batch_index,
graphicsfolder_path, logfile_path, min_samples_fraction,
stop_event=None, detector_geometry=None,
max_pairs_per_frame=None,
dbscan_eps=DEFAULT_CENTERFINDING_DBSCAN_EPS):
"""
Iteratively shift the center toward the mean Friedel-pair deviation for one batch.
Stops when deviations fall below 0.4 px in both axes or after `breakcriterion`
iterations; saves per-iteration diagnostics to disk.
"""
current_center = initial_center
mean_x, mean_y = float('inf'), float('inf')
itcount = 0
breakcriterion = 5
while mean_x > 0.4 or mean_y > 0.4:
check_cancel(stop_event, "beamstop batch iteration")
deviations = refine_center_with_friedel(
h5file_path=h5file_path,
current_center=current_center,
tolerance=tolerance,
min_peaks=min_peaks,
resolution_limit=resolution_limit,
subset_indices=subset_indices,
batch_index=batch_index,
logfile_path=logfile_path,
detector_geometry=detector_geometry,
max_pairs_per_frame=max_pairs_per_frame,
)
itcount += 1
# Check if deviations are empty
if deviations.size == 0:
log_start(logfile_path, f"Batch {batch_index}: No deviations found, iteration {itcount}")
break
if itcount == breakcriterion + 1:
log_start(logfile_path, f'Batch {batch_index}: No valid center found after {itcount} iterations, iteration aborted')
break
mean, _ = fit_gaussian_to_largest_cluster(
deviations,
min_samples_fraction,
dbscan_eps=dbscan_eps,
)
if mean is None:
log_start(
logfile_path,
f"Batch {batch_index}: No cluster in deviation cloud (iteration {itcount})",
)
break
current_center = [
current_center[0] + mean[0] / 2,
current_center[1] + mean[1] / 2,
]
mean_x, mean_y = abs(mean[0]), abs(mean[1])
# SAVE *THIS* ITERATION and immediately free memory
base = f"batch{batch_index}_it{itcount:03d}"
np.savez(
os.path.join(graphicsfolder_path, base + ".npz"),
deviations=deviations.astype("float32", copy=False),
mean=mean.astype("float32", copy=False),
current_center=np.asarray(current_center, dtype="float32"),
)
del deviations
if itcount % 3 == 0: # run GC only every 3rd iteration
gc.collect()
try: # on Linux give pages back to the OS
import ctypes, sys
if sys.platform.startswith("linux"):
ctypes.CDLL("libc.so.6").malloc_trim(0)
except Exception:
pass
log_start(logfile_path, f'Batch {batch_index}: Processed in {itcount} iterations, final center = '
f'[{round(current_center[0], 6)}, {round(current_center[1], 6)}], '
f'mean deviation = [{mean_x:.3e}, {mean_y:.3e}]')
return current_center, mean_x, mean_y, itcount
# Function to set the center based on linear fit
[docs]
def set_center_based_on_line_fit(h5file_path, batch_centers, graphicsfolder_path, logfile_path,
framesize, framepath, batch_size, total_frames, pixels_per_meter,
force_linear_fit=False, detector_geometry=None):
"""
Fit a linear drift to batch centers, write per-frame centers to HDF5, and log QA.
Saves fit parameters for GUI display, warns on poor R^2, and optionally falls
back if the fit is weak (unless forced).
"""
log_start(logfile_path, 'Interpolating centers based on linear fit')
# Decide whether to force linear fit despite low R^2
if force_linear_fit:
log_start(logfile_path, 'Force linear fit enabled: ignoring R^2 threshold')
# Load or create index mapping for original frames
import h5py
import numpy as np
with h5py.File(h5file_path, 'r') as f:
grp = f['entry/data']
if 'index' in grp:
index_map = grp['index'][:]
else:
# assume no frames stripped
n_frames, _, _, _ = get_frame_stack_shape(f, framepath)
index_map = np.arange(n_frames, dtype=int)
# Determine original frame positions corresponding to current images
orig_positions = np.where(index_map != -1)[0]
# Filter out invalid centers, using recorded original frame indices
valid_centers = []
valid_positions = []
for avg_orig, center in batch_centers:
if center is not None:
valid_centers.append(center)
valid_positions.append(avg_orig)
if not valid_centers:
log_start(logfile_path, 'No valid centers found for linear regression')
return False # Indicate that linear fit was not successful
# Separate the positions and center coordinates
x = np.array(valid_positions)
y_x = np.array([center[0] for center in valid_centers])
y_y = np.array([center[1] for center in valid_centers])
# Perform linear regression
slope_x, intercept_x, rvalue_x, _, _ = linregress(x, y_x)
slope_y, intercept_y, rvalue_y, _, _ = linregress(x, y_y)
# Warn if R^2 is bad
poor_fit = False
if rvalue_x ** 2 <= 0.25:
log_start(logfile_path, f'Warning: poor quality of x-fit (R^2 = {rvalue_x ** 2:.3f}), review your dataset or consider excluding data')
poor_fit = True
if rvalue_y ** 2 <= 0.25:
log_start(logfile_path, f'Warning: poor quality of y-fit (R^2 = {rvalue_y ** 2:.3f}), review your dataset or consider excluding data')
poor_fit = True
# If the fit is poor, return False to indicate failure
if poor_fit and not force_linear_fit:
return False
# Generate new centers for each frame based on the line fit, mapped to original frames
frame_positions = orig_positions
updated_x = slope_x * frame_positions + intercept_x
updated_y = slope_y * frame_positions + intercept_y
log_start(logfile_path, f'Scope of drift in x,y = [{(max(updated_x)-min(updated_x)):.3f}, {(max(updated_y)-min(updated_y)):.3f}]')
# Use original unstripped frame indices for batch centers
batch_positions = [avg for avg, _ in batch_centers]
frame_positions = np.arange(total_frames)
# Convert batch centers back to raw coordinates for compatibility outputs.
if detector_geometry is not None:
batch_xy = np.array(
[[c[0], c[1]] if c is not None else [np.nan, np.nan] for _, c in batch_centers],
dtype=np.float64,
)
finite_mask = np.isfinite(batch_xy).all(axis=1)
raw_batch_xy = batch_xy.copy()
if np.any(finite_mask):
converted_raw, _ = corrected_to_raw_points(
batch_xy[finite_mask],
detector_geometry,
)
raw_batch_xy[finite_mask] = converted_raw
else:
raw_batch_xy = np.array(
[[c[0], c[1]] if c is not None else [np.nan, np.nan] for _, c in batch_centers],
dtype=np.float64,
)
# Save beamstop raw centers: only mean per-batch centers for line fit.
bp = np.array(batch_positions, dtype=float)
bx = raw_batch_xy[:, 0]
by = raw_batch_xy[:, 1]
raw_array = np.column_stack((bp, bx, by))
np.save(os.path.join(graphicsfolder_path, "beamstop_centers_raw.npy"), raw_array)
if detector_geometry is not None:
corrected_fit = np.column_stack((updated_x, updated_y))
raw_fit, _ = corrected_to_raw_points(
corrected_fit,
detector_geometry,
)
updated_x_raw = raw_fit[:, 0]
updated_y_raw = raw_fit[:, 1]
else:
updated_x_raw = updated_x
updated_y_raw = updated_y
# Save beamstop fit parameters in RAW coordinates so GUI overlays match raw batch-center plots.
raw_valid_mask = np.isfinite(raw_batch_xy).all(axis=1)
raw_valid_positions = np.asarray(batch_positions, dtype=float)[raw_valid_mask]
raw_valid_x = raw_batch_xy[raw_valid_mask, 0]
raw_valid_y = raw_batch_xy[raw_valid_mask, 1]
if raw_valid_positions.size >= 2:
slope_x_raw, intercept_x_raw = np.polyfit(raw_valid_positions, raw_valid_x, 1)
slope_y_raw, intercept_y_raw = np.polyfit(raw_valid_positions, raw_valid_y, 1)
x_pred_raw = slope_x_raw * raw_valid_positions + intercept_x_raw
y_pred_raw = slope_y_raw * raw_valid_positions + intercept_y_raw
ss_res_x_raw = np.sum((raw_valid_x - x_pred_raw) ** 2)
ss_tot_x_raw = np.sum((raw_valid_x - raw_valid_x.mean()) ** 2)
r2_x_raw = 1 - ss_res_x_raw / ss_tot_x_raw if ss_tot_x_raw > 0 else 0.0
ss_res_y_raw = np.sum((raw_valid_y - y_pred_raw) ** 2)
ss_tot_y_raw = np.sum((raw_valid_y - raw_valid_y.mean()) ** 2)
r2_y_raw = 1 - ss_res_y_raw / ss_tot_y_raw if ss_tot_y_raw > 0 else 0.0
else:
slope_x_raw, intercept_x_raw, r2_x_raw = slope_x, intercept_x, rvalue_x ** 2
slope_y_raw, intercept_y_raw, r2_y_raw = slope_y, intercept_y, rvalue_y ** 2
fit_array_raw = np.array([
[slope_x_raw, intercept_x_raw, r2_x_raw],
[slope_y_raw, intercept_y_raw, r2_y_raw]
], dtype=float)
np.save(os.path.join(graphicsfolder_path, "beamstop_centers_fit.npy"), fit_array_raw)
# Save linear fit data for GUI in RAW coordinates.
np.savez(
os.path.join(graphicsfolder_path, "linearfit.npz"),
batch_positions=batch_positions,
updated_x=updated_x_raw,
updated_y=updated_y_raw,
batch_center_x=bx,
batch_center_y=by,
)
# Always save raw batch-center data for GUI even when linear fit succeeds
np.savez(
os.path.join(graphicsfolder_path, "batchcenters.npz"),
frame_positions=np.array(batch_positions),
updated_x=bx,
updated_y=by,
)
# Write the updated centers to the HDF5 file
with h5py.File(h5file_path, 'a') as workingfile:
for ds_name in ['entry/data/center_x', 'entry/data/center_y']:
if ds_name in workingfile:
del workingfile[ds_name]
workingfile.create_dataset('entry/data/center_x', data=updated_x_raw, maxshape=(None,), dtype='float64')
workingfile.create_dataset('entry/data/center_y', data=updated_y_raw, maxshape=(None,), dtype='float64')
update_detector_shifts(
h5file_path,
logfile_path,
framepath,
pixels_per_meter,
detector_geometry=detector_geometry,
)
log_start(logfile_path, 'Interpolated detector shifts written to HDF5 file')
return True # Indicate that linear fit was successful
# Function to set the center based on batch data (used as fallback)
[docs]
def set_center_based_on_batch(h5file_path, batch_centers, graphicsfolder_path, logfile_path,
framepath, batch_size, total_frames, pixels_per_meter,
orig_positions=None, detector_geometry=None):
"""Fallback: write per-batch mean centers to all frames in that batch and persist."""
log_start(logfile_path, 'Setting centers based on batch data')
# Build original frame positions (raw indices) for plotting/alignment
if orig_positions is None:
with h5py.File(h5file_path, 'r') as f:
grp = f['entry/data']
if 'index' in grp:
index_map = grp['index'][:]
orig_positions = np.where(index_map != -1)[0]
else:
n_frames, _, _, _ = get_frame_stack_shape(f, framepath)
orig_positions = np.arange(n_frames, dtype=int)
# Filter out invalid centers (e.g., None values)
valid_centers = [center for _, center in batch_centers if center is not None]
if not valid_centers:
log_start(logfile_path, 'No valid centers found')
return
# Initialize arrays for all frames
updated_x = np.zeros(total_frames)
updated_y = np.zeros(total_frames)
# Set the center for each frame in a batch
for i, (_, center) in enumerate(batch_centers):
if center is not None:
start_idx = i * batch_size
end_idx = min((i + 1) * batch_size, total_frames)
updated_x[start_idx:end_idx] = center[0]
updated_y[start_idx:end_idx] = center[1]
if detector_geometry is not None:
corrected_fit = np.column_stack((updated_x, updated_y))
raw_fit, _ = corrected_to_raw_points(
corrected_fit,
detector_geometry,
)
updated_x_raw = raw_fit[:, 0]
updated_y_raw = raw_fit[:, 1]
else:
updated_x_raw = updated_x
updated_y_raw = updated_y
# Plot the batch centers (removed runtime plotting)
frame_positions = np.array(orig_positions)
# Save batch centers data for GUI immediately
np.savez(
os.path.join(graphicsfolder_path, "batchcenters.npz"),
frame_positions=frame_positions,
updated_x=updated_x_raw,
updated_y=updated_y_raw,
)
# Write the updated centers to the HDF5 file
with h5py.File(h5file_path, 'a') as workingfile:
for ds_name in ['entry/data/center_x', 'entry/data/center_y']:
if ds_name in workingfile:
del workingfile[ds_name]
workingfile.create_dataset('entry/data/center_x', data=updated_x_raw, maxshape=(None,), dtype='float64')
workingfile.create_dataset('entry/data/center_y', data=updated_y_raw, maxshape=(None,), dtype='float64')
update_detector_shifts(
h5file_path,
logfile_path,
framepath,
pixels_per_meter,
detector_geometry=detector_geometry,
)
log_start(logfile_path, 'Detector shifts written to HDF5 file')
# Function to update detector shifts based on new centers
[docs]
def update_detector_shifts(h5file_path, logfile_path, framepath, pixels_per_meter, detector_geometry=None):
"""Compute detector shifts (mm) from center arrays and save to HDF5."""
with h5py.File(h5file_path, 'r+') as workingfile:
updated_center_x = workingfile['entry/data/center_x'][:]
updated_center_y = workingfile['entry/data/center_y'][:]
_, framesize_x, framesize_y, resolved_framepath = get_frame_stack_shape(workingfile, framepath)
if _normalize_h5_path(framepath) != resolved_framepath:
log_start(
logfile_path,
f"Resolved framepath '{framepath}' to dataset '{resolved_framepath}' for detector-shift update.",
)
centers_raw = np.column_stack((updated_center_x, updated_center_y)).astype(np.float64, copy=False)
reference_raw = np.array([[framesize_x / 2.0, framesize_y / 2.0]], dtype=np.float64)
if detector_geometry is not None:
centers_geo, _ = raw_to_corrected_points(centers_raw, detector_geometry)
reference_geo, _ = raw_to_corrected_points(reference_raw, detector_geometry)
center_delta = centers_geo - reference_geo[0]
else:
center_delta = centers_raw - reference_raw[0]
det_shift_x_mm = -(center_delta[:, 0] / pixels_per_meter) * 1000
det_shift_y_mm = -(center_delta[:, 1] / pixels_per_meter) * 1000
# Write the updated detector shifts to the HDF5 file
with h5py.File(h5file_path, 'r+') as workingfile:
for ds_name in ['entry/data/det_shift_x_mm', 'entry/data/det_shift_y_mm']:
if ds_name in workingfile:
del workingfile[ds_name]
workingfile.create_dataset('entry/data/det_shift_x_mm', data=det_shift_x_mm, maxshape=(None,), dtype='float64')
workingfile.create_dataset('entry/data/det_shift_y_mm', data=det_shift_y_mm, maxshape=(None,), dtype='float64')
ensure_dense_logs(workingfile)
ensure_goniometer_transforms(workingfile)
log_start(logfile_path, 'Updated detector shifts written to HDF5 file')
# Delayed function to process a batch
@delayed
def process_batch(h5file_path, initial_center, subset_indices, tolerance, min_peaks,
resolution_limit, batch_index, graphicsfolder_path, logfile_path,
min_samples_fraction, detector_geometry=None, max_pairs_per_frame=None,
dbscan_eps=DEFAULT_CENTERFINDING_DBSCAN_EPS):
"""Delayed wrapper to run iterative fitting for one batch (used by Dask)."""
center, mean_x, mean_y, itcount = iterative_gaussian_fitting_on_subset(
h5file_path=h5file_path,
initial_center=initial_center,
subset_indices=subset_indices,
tolerance=tolerance,
min_peaks=min_peaks,
resolution_limit=resolution_limit,
batch_index=batch_index,
graphicsfolder_path=graphicsfolder_path,
logfile_path=logfile_path,
min_samples_fraction=min_samples_fraction,
stop_event=None, # avoid serializing threading events to workers
detector_geometry=detector_geometry,
max_pairs_per_frame=max_pairs_per_frame,
dbscan_eps=dbscan_eps,
)
return center, mean_x, mean_y, itcount
# Function to determine pixels_per_meter
[docs]
def get_pixels_per_meter(config, logfile_path):
"""
Retrieve the 'pixels_per_meter' value from the configuration.
If 'pixels_per_meter' is not available, compute it using 'pixel_width' and 'pixel_unit'.
"""
# Attempt to get 'pixels_per_meter' directly
if config.has_option('AcquisitionDetails', 'pixels_per_meter'):
try:
pixels_per_meter = float(config.get('AcquisitionDetails', 'pixels_per_meter'))
log_start(logfile_path, f"Retrieved 'pixels_per_meter' directly: {pixels_per_meter}")
return pixels_per_meter
except ValueError:
error_msg = "Invalid value for 'pixels_per_meter'. It must be a number."
log_result(logfile_path, '', error_msg)
raise ValueError(error_msg)
# Fallback to 'pixel_width' and 'pixel_unit'
pixel_width = config.get('AcquisitionDetails', 'pixel_width', fallback=None)
pixel_unit = config.get('AcquisitionDetails', 'pixel_unit', fallback=None)
if pixel_width is not None and pixel_unit is not None:
try:
pixel_width = float(pixel_width)
log_start(logfile_path, f"Retrieved 'pixel_width': {pixel_width}")
except ValueError:
error_msg = "Invalid value for 'pixel_width'. It must be a number."
log_result(logfile_path, '', error_msg)
raise ValueError(error_msg)
if pixel_unit == '1/m':
pixels_per_meter = pixel_width
log_start(logfile_path, f"Using 'pixel_width' as 'pixels_per_meter': {pixels_per_meter}")
return pixels_per_meter
elif pixel_unit == 'm/pixel':
if pixel_width == 0:
error_msg = "'pixel_width' cannot be zero when 'pixel_unit' is 'm/pixel'."
log_result(logfile_path, '', error_msg)
raise ValueError(error_msg)
pixels_per_meter = 1.0 / pixel_width
log_start(logfile_path, f"Computed 'pixels_per_meter' as inverse of 'pixel_width': {pixels_per_meter}")
return pixels_per_meter
else:
error_msg = f"Unknown 'pixel_unit': {pixel_unit}. Expected '1/m' or 'm/pixel'."
log_result(logfile_path, '', error_msg)
raise ValueError(error_msg)
else:
# Identify which parameters are missing
missing = []
if pixel_width is None:
missing.append('pixel_width')
if pixel_unit is None:
missing.append('pixel_unit')
missing_str = ', '.join(missing)
error_msg = f"Missing parameter(s) in 'AcquisitionDetails': {missing_str}"
log_result(logfile_path, '', error_msg)
raise ValueError(error_msg)
# Main function to process batches
[docs]
def quartering(configfile, usedask, batch_size, progress_callback=None, stop_event=None):
"""
Main beamstop center-finding workflow for a single INI/HDF5 pair.
- Validates required parameters
- Maps surviving frames (post-strip) to batches
- Runs Friedel-pair iterative fitting per batch (optionally via Dask)
- Tries linear drift fit, falls back to batch means if poor
"""
check_cancel(stop_event, "centerfinding start")
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
config = read_config(configfile)
detector_geometry = load_detector_geometry(config)
# Create unique findcenters output folder
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
graphicsfolder_path = os.path.join(outputfolder_path, f'findcenters_{timestamp}')
framepath = config.get('Paths', 'framepath')
os.makedirs(graphicsfolder_path, exist_ok=True)
# Check if necessary parameters are defined
for param in ['centerfinding_tolerance', 'centerfinding_min_peaks',
'centerfinding_min_samples_fraction', 'centerfinding_resolution_limit']:
if not config.has_option('Parameters', param):
error = f'{param.split("_", 1)[1]} not defined, peak finding interrupted'
log_result(logfile_path, '', error)
raise Exception(error)
# Load parameters
tolerance = float(config.get('Parameters', 'centerfinding_tolerance'))
min_peaks = float(config.get('Parameters', 'centerfinding_min_peaks'))
resolution_limit = float(config.get('Parameters', 'centerfinding_resolution_limit'))
min_samples_fraction = float(config.get('Parameters', 'centerfinding_min_samples_fraction'))
dbscan_eps = float(
config.get(
'Parameters',
'centerfinding_dbscan_eps',
fallback=str(DEFAULT_CENTERFINDING_DBSCAN_EPS),
)
)
if dbscan_eps <= 0:
raise ValueError("centerfinding_dbscan_eps must be greater than 0.")
max_pairs_per_frame = int(float(config.get('Parameters', 'centerfinding_max_pairs_per_frame', fallback="10000")))
x0_str = config.get('Parameters', 'centerfinding_x0', fallback=None)
y0_str = config.get('Parameters', 'centerfinding_y0', fallback=None)
# Retrieve pixels_per_meter
pixels_per_meter = get_pixels_per_meter(config, logfile_path)
if detector_geometry is not None:
log_start(logfile_path, "Detector geometry enabled for Friedel centerfinding (internal corrected coordinates).")
if usedask:
if sys.platform == "darwin":
log_start(
logfile_path,
"Using macOS Dask safe mode for geometry-enabled Friedel centerfinding (threaded workers).",
)
os.environ["COSEDA_DASK_PROCESSES"] = "0"
try:
DaskClientManager.kill_current_cluster()
except Exception:
pass
else:
log_start(
logfile_path,
"Geometry-enabled centerfinding on non-macOS: keeping default Dask process workers.",
)
# Read linear-fit knobs
force_linear_fit = config.getboolean('Parameters', 'centerfinding_force_linear_fit', fallback=False)
skip_linear_fit = config.getboolean('Parameters', 'centerfinding_skip_linear_fit', fallback=False)
if skip_linear_fit and force_linear_fit:
log_start(logfile_path, 'Both skip_linear_fit and force_linear_fit set; skipping takes precedence')
force_linear_fit = False
# Determine initial center
with h5py.File(h5file_path, 'r') as workingfile:
_, framesize_x, framesize_y, resolved_framepath = get_frame_stack_shape(workingfile, framepath)
if _normalize_h5_path(framepath) != resolved_framepath:
log_start(
logfile_path,
f"Resolved framepath '{framepath}' to dataset '{resolved_framepath}' for centerfinding.",
)
if detector_geometry is not None:
geom_errors, geom_warnings = validate_geometry_for_frame(detector_geometry, int(framesize_x), int(framesize_y))
for msg in geom_warnings:
log_start(logfile_path, f"DetectorGeometry warning: {msg}")
if geom_errors:
for msg in geom_errors:
log_start(logfile_path, f"DetectorGeometry error: {msg}")
raise ValueError(
"DetectorGeometry is incompatible with frame shape; refusing to run with ambiguous geometry mapping."
)
if x0_str and x0_str.lower() != 'none' and y0_str and y0_str.lower() != 'none':
try:
x0 = float(x0_str)
y0 = float(y0_str)
log_start(logfile_path, f"Using provided initial center: [{x0}, {y0}]")
except ValueError:
log_start(logfile_path, 'Invalid initial center provided, using frame center')
x0, y0 = framesize_x / 2, framesize_y / 2
else:
log_start(logfile_path, 'Beam position guess not defined, using frame center')
x0, y0 = framesize_x / 2, framesize_y / 2
initial_center = [x0, y0]
if detector_geometry is not None:
initial_center_corr, _ = raw_to_corrected(initial_center, detector_geometry)
initial_center_fit = initial_center_corr.tolist()
else:
initial_center_fit = initial_center
# Dump run parameters to paramdump.txt
param_file = os.path.join(graphicsfolder_path, 'paramdump.txt')
with open(param_file, 'w') as pf:
pf.write(f"centerfinding_tolerance={tolerance}\n")
pf.write(f"centerfinding_min_peaks={min_peaks}\n")
pf.write(f"centerfinding_min_samples_fraction={min_samples_fraction}\n")
pf.write(f"centerfinding_dbscan_eps={dbscan_eps}\n")
pf.write(f"centerfinding_resolution_limit={resolution_limit}\n")
pf.write(f"centerfinding_x0={x0}\n")
pf.write(f"centerfinding_y0={y0}\n")
pf.write(f"pixels_per_meter={pixels_per_meter}\n")
pf.write(f"force_linear_fit={force_linear_fit}\n")
pf.write(f"skip_linear_fit={skip_linear_fit}\n")
pf.write(f"batch_size={batch_size}\n")
pf.write(f"max_pairs_per_frame={max_pairs_per_frame}\n")
pf.write(f"detector_geometry_enabled={detector_geometry is not None}\n")
# Step 1: load index_map (original→new), build surviving-original-frames list
import numpy as np
with h5py.File(h5file_path, 'r') as workingfile:
grp = workingfile['entry/data']
if 'index' in grp:
index_map = grp['index'][:]
else:
n_frames, _, _, _ = get_frame_stack_shape(workingfile, framepath)
index_map = np.arange(n_frames, dtype=int)
full_indices = np.arange(len(index_map), dtype=int)
kept_mask = index_map != -1
unstripped_indices = full_indices[kept_mask]
total_indices = unstripped_indices.tolist()
log_start(logfile_path, f'Start centerfinding on {len(total_indices)} frames, initial center {initial_center}')
check_cancel(stop_event, "before batch creation")
# Divide into batches
batches, total_frames, batch_size = divide_into_batches(total_indices, batch_size, logfile_path)
batch_centers = []
if usedask:
try:
log_start(logfile_path, 'Parallel processing with Dask enabled')
client = DaskClientManager.get_client()
dashboard_address, _, _, _ = DaskClientManager.get_client_info()
DaskClientManager.log_client_info(logfile_path)
log_print(f'\nClick here to monitor progress: {dashboard_address}\n')
# Steps 2–4: submit tasks via stripped indices, record mean original index
tasks = [
process_batch(
h5file_path=h5file_path,
initial_center=initial_center_fit,
subset_indices=index_map[np.array(batch_indices)].tolist(),
tolerance=tolerance,
min_peaks=min_peaks,
resolution_limit=resolution_limit,
batch_index=i+1,
graphicsfolder_path=graphicsfolder_path,
logfile_path=logfile_path,
min_samples_fraction=min_samples_fraction,
detector_geometry=detector_geometry,
max_pairs_per_frame=max_pairs_per_frame,
dbscan_eps=dbscan_eps,
)
for i, batch_indices in enumerate(batches)
]
futures = client.compute(tasks)
future_to_idx = {f: i for i, f in enumerate(futures)}
results = [None] * len(futures)
try:
completed = 0
for fut in as_completed(futures):
check_cancel(stop_event, "centerfinding dask gather")
res = fut.result()
idx = future_to_idx[fut]
results[idx] = res
completed += 1
if progress_callback:
progress_callback(completed, len(futures))
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:
try:
fut.cancel()
except Exception:
pass
if stop_event is not None and stop_event.is_set():
raise RefinementCancelled("Center finding cancelled by user")
for i, res in enumerate(results):
if res is None:
continue
center, mean_x, mean_y, itcount = res
avg_orig = np.mean(batches[i])
batch_centers.append((avg_orig, center))
if progress_callback:
progress_callback(i+1, len(results))
except Exception as exc:
log_start(logfile_path, f"Dask centerfinding failed ({exc}); falling back to sequential mode.")
usedask = False
if not usedask:
# Sequential processing with stripped-index mapping
for i, batch_indices in enumerate(batches):
check_cancel(stop_event, "centerfinding batch loop")
batch_index = i + 1
stripped = index_map[np.array(batch_indices)].tolist()
center, mean_x, mean_y, itcount = iterative_gaussian_fitting_on_subset(
h5file_path=h5file_path,
initial_center=initial_center_fit,
subset_indices=stripped,
tolerance=tolerance,
min_peaks=min_peaks,
resolution_limit=resolution_limit,
batch_index=batch_index,
graphicsfolder_path=graphicsfolder_path,
logfile_path=logfile_path,
min_samples_fraction=min_samples_fraction,
stop_event=stop_event,
detector_geometry=detector_geometry,
max_pairs_per_frame=max_pairs_per_frame,
dbscan_eps=dbscan_eps,
)
avg_orig = np.mean(batch_indices)
batch_centers.append((avg_orig, center))
if progress_callback:
progress_callback(i+1, len(batches))
# Try linear fit
orig_positions = np.array(total_indices)
if skip_linear_fit:
log_start(logfile_path, 'Skipping linear fit (centerfinding_skip_linear_fit=True); using batch means')
# Persist raw batch centers for GUI (same layout as linear-fit path)
batch_positions = [avg for avg, _ in batch_centers]
bx = [center[0] if center is not None else np.nan for _, center in batch_centers]
by = [center[1] if center is not None else np.nan for _, center in batch_centers]
if detector_geometry is not None:
corr = np.column_stack((np.array(bx, dtype=float), np.array(by, dtype=float)))
finite = np.isfinite(corr).all(axis=1)
if np.any(finite):
raw_vals, _ = corrected_to_raw_points(
corr[finite],
detector_geometry,
)
corr[finite] = raw_vals
bx = corr[:, 0]
by = corr[:, 1]
raw_array = np.column_stack((np.array(batch_positions, dtype=float),
np.array(bx, dtype=float),
np.array(by, dtype=float)))
np.save(os.path.join(graphicsfolder_path, "beamstop_centers_raw.npy"), raw_array)
set_center_based_on_batch(
h5file_path=h5file_path,
batch_centers=batch_centers,
graphicsfolder_path=graphicsfolder_path,
logfile_path=logfile_path,
framepath=framepath,
batch_size=batch_size,
total_frames=total_frames,
pixels_per_meter=pixels_per_meter,
orig_positions=orig_positions,
detector_geometry=detector_geometry,
)
else:
success = set_center_based_on_line_fit(
h5file_path=h5file_path,
batch_centers=batch_centers,
graphicsfolder_path=graphicsfolder_path,
logfile_path=logfile_path,
framesize=framesize_x,
framepath=framepath,
batch_size=batch_size,
total_frames=total_frames,
pixels_per_meter=pixels_per_meter,
force_linear_fit=force_linear_fit,
detector_geometry=detector_geometry,
)
if not success:
log_start(logfile_path, 'Falling back to batch-based centers')
set_center_based_on_batch(
h5file_path=h5file_path,
batch_centers=batch_centers,
graphicsfolder_path=graphicsfolder_path,
logfile_path=logfile_path,
framepath=framepath,
batch_size=batch_size,
total_frames=total_frames,
pixels_per_meter=pixels_per_meter,
orig_positions=orig_positions,
detector_geometry=detector_geometry,
)
try:
with h5py.File(h5file_path, "r+") as workingfile:
write_nxprocess_centerfinding(
workingfile,
program="coseda.centerfinding.findcenters_beamstop",
method="beamstop",
input_path=h5file_path,
output_path=h5file_path,
parameters={
"tolerance": tolerance,
"min_peaks": min_peaks,
"resolution_limit": resolution_limit,
"min_samples_fraction": min_samples_fraction,
"dbscan_eps": dbscan_eps,
"x0": x0,
"y0": y0,
"batch_size": batch_size,
"max_pairs_per_frame": max_pairs_per_frame,
"pixels_per_meter": pixels_per_meter,
"force_linear_fit": force_linear_fit,
"skip_linear_fit": skip_linear_fit,
"usedask": usedask,
},
)
except Exception as exc:
log_start(logfile_path, f"Failed to write NXprocess centerfinding: {exc}")
# Wrapper function to process all configurations
[docs]
def quartering_batch(input_path, usedask=False, batch_size=10000, progress_callback=None, stop_event=None):
"""Run beamstop center finding for each INI in a folder/file/list input."""
configfiles, input_path = handle_input(input_path)
if not configfiles:
log_print("No .ini files found to process.")
return
log_print(f'Following .ini files were found and will be processed:')
for configfile in configfiles:
log_print(f'- {configfile}')
filecount = 1
for configfile in configfiles:
check_cancel(stop_event, "centerfinding batch file loop")
shoutout(configfile)
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
# Get config
config = read_config(configfile)
log_start(logfile_path, f'Processing {os.path.basename(configfile)}')
try:
quartering(configfile, usedask, batch_size, progress_callback=progress_callback, stop_event=stop_event)
log_start(logfile_path, "Center finding completed successfully.")
except MemoryError:
# Handle out-of-memory errors and propagate
log_start(logfile_path, "Processing failed: ran out of memory. Try reducing batch size.")
log_print("Processing failed: ran out of memory. Try reducing the batch size.")
raise
except Exception as e:
log_start(logfile_path, f"Center finding encountered an error: {e}")
raise
if filecount < len(configfiles):
log_print("")
log_print(f"Proceeding to next file ({filecount}/{len(configfiles)})")
log_print("")
filecount += 1
log_print(f"Finished centerfinding for all files ({filecount-1}/{len(configfiles)})")