from coseda.logging_utils import log_print
import os
import numpy as np
from dask.distributed import Client, as_completed, CancelledError, TimeoutError
import h5py
import matplotlib.pyplot as plt
import datetime
import tempfile
import sys
from coseda.centerrefinement.cancellation import RefinementCancelled, check_cancel
try:
from numba import njit
_NUMBA_AVAILABLE = True
except ImportError: # pragma: no cover
njit = None
_NUMBA_AVAILABLE = False
if _NUMBA_AVAILABLE:
@njit(cache=True)
def _median_abs(arr):
tmp = np.abs(arr.copy())
tmp.sort()
n = tmp.shape[0]
if n % 2 == 1:
return tmp[n // 2]
else:
return 0.5 * (tmp[n // 2 - 1] + tmp[n // 2])
@njit(cache=True)
def _numba_lowess_impl(x, y, frac, iters):
n = x.shape[0]
order = np.argsort(x)
x_sorted = x[order]
y_sorted = y[order]
y_est_sorted = np.zeros(n, dtype=np.float64)
robustness_weights = np.ones(n, dtype=np.float64)
span = max(2, int(np.ceil(frac * n)))
if span >= n:
span = n - 1
if span < 2:
span = 2
for iteration in range(iters):
for i in range(n):
left = i - span // 2
if left < 0:
left = 0
right = left + span - 1
if right >= n:
right = n - 1
left = max(0, right - span + 1)
h_left = x_sorted[i] - x_sorted[left]
h_right = x_sorted[right] - x_sorted[i]
h = h_left if h_left > h_right else h_right
if h <= 0:
h = 1.0
window_size = right - left + 1
weights = np.zeros(window_size, dtype=np.float64)
for j in range(window_size):
dist = np.abs(x_sorted[left + j] - x_sorted[i]) / h
if dist >= 1:
weights[j] = 0.0
else:
w = (1.0 - dist ** 3) ** 3
weights[j] = w * robustness_weights[left + j]
sum_w = np.sum(weights)
if sum_w <= 0:
y_est_sorted[i] = y_sorted[i]
else:
xw = x_sorted[left:right + 1]
yw = y_sorted[left:right + 1]
x_bar = np.sum(weights * xw) / sum_w
y_bar = np.sum(weights * yw) / sum_w
beta_num = 0.0
beta_den = 0.0
for j in range(window_size):
xi = xw[j] - x_bar
beta_num += weights[j] * xi * (yw[j] - y_bar)
beta_den += weights[j] * xi * xi
if beta_den > 0:
slope = beta_num / beta_den
intercept = y_bar - slope * x_bar
y_est_sorted[i] = slope * x_sorted[i] + intercept
else:
y_est_sorted[i] = y_bar
if iteration < iters - 1:
residuals = y_sorted - y_est_sorted
s = _median_abs(residuals)
if s <= 0:
break
for i in range(n):
# Tukey bisquare uses absolute residual magnitude.
r = np.abs(residuals[i]) / (6.0 * s)
if r >= 1:
robustness_weights[i] = 0.0
else:
tmp = 1.0 - r * r
robustness_weights[i] = tmp * tmp
out = np.empty((n, 2), dtype=np.float64)
out[:, 0] = x_sorted
out[:, 1] = y_est_sorted
return out
else:
_numba_lowess_impl = None
def _lowess_axis(values, indices, frac):
"""Helper to run LOWESS on a single axis."""
x = np.ascontiguousarray(indices, dtype=np.float64)
y = np.ascontiguousarray(values, dtype=np.float64)
if _numba_lowess_impl is not None:
return _numba_lowess_impl(x, y, frac, 2)
from statsmodels.nonparametric.smoothers_lowess import lowess as _fallback_lowess
return _fallback_lowess(y, x, frac=frac)
from scipy.stats import linregress
from sklearn.mixture import GaussianMixture
from scipy.interpolate import interp1d
from sklearn.cluster import KMeans, DBSCAN
from coseda.initialize import find_configfiles
from coseda.io import handle_input, parse_config, read_config, config_to_paths
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_points,
validate_geometry_for_frame,
)
from coseda.centerfinding.findcenters_beamstop import update_detector_shifts, find_friedel_pairs
from coseda.nexus.process import write_nxprocess_centerrefinement
# Define Delayed Functions
[docs]
def refine_center_with_friedel_batch(
h5file_path,
min_peaks,
tolerance,
resolution_limit,
indices_subset,
centers_x,
centers_y,
index_map,
detector_geometry=None,
max_pairs_per_frame=None,
):
"""
Process a subset of frames (indices_subset, which are original frame numbers) to find deviations.
"""
all_deviations = []
all_indices = []
with h5py.File(h5file_path, 'r') as workingfile:
for raw_idx in indices_subset:
proc_idx = index_map[raw_idx]
if proc_idx == -1:
continue
num_peaks = workingfile['entry/data/nPeaks'][proc_idx]
if num_peaks < min_peaks:
continue
peak_xpos_raw = workingfile['entry/data/peakXPosRaw'][proc_idx][:num_peaks]
peak_ypos_raw = workingfile['entry/data/peakYPosRaw'][proc_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
current_center = [centers_x[proc_idx], centers_y[proc_idx]]
# Filter out peaks based on resolution_limit
distances = np.linalg.norm(peak_positions - current_center, axis=1)
peak_positions = peak_positions[distances < resolution_limit]
# Use the raw frame number for image_slot
image_slot = raw_idx
deviations = find_friedel_pairs(
peak_positions,
current_center,
tolerance,
max_pairs=max_pairs_per_frame,
)
# Only proceed if we got any pairs back
if deviations.size > 0:
# deviations is a NumPy array of shape (N,2)
all_deviations.append(deviations)
all_indices.append(np.full(deviations.shape[0], raw_idx, dtype=np.int64))
if all_deviations:
return np.concatenate(all_indices), np.concatenate(all_deviations)
return np.array([], dtype=np.int64), np.empty((0, 2), dtype=np.float32)
# Define Processing Functions
[docs]
def refine_center_with_friedel_and_update(
h5file_path, min_peaks, tolerance, resolution_limit, itcount, convergence_threshold,
converged, graphicsfolder_path, logfile_path, client, lowess_frac=0.1, lowess_window=None,
batch_size=500, stop_event=None, detector_geometry=None, max_pairs_per_frame=None,
deviation_aggregation="median",
):
"""Refine centers using Friedel pairs and update the HDF5 file."""
check_cancel(stop_event, "beamstop iteration setup")
log_print(f"[DEBUG] refine_center_with_friedel_and_update called with batch_size={batch_size}, using Dask={client is not None}")
# Open the HDF5 file to read data
with h5py.File(h5file_path, 'r') as workingfile:
existing_center_x = workingfile['entry/data/center_x'][:]
existing_center_y = workingfile['entry/data/center_y'][:]
# Load original-to-current index mapping, or fall back to identity if missing
if 'entry/data/index' in workingfile:
index_map = workingfile['entry/data/index'][:]
else:
# no stripping index: assume all frames present
index_map = np.arange(existing_center_x.shape[0], dtype=int)
existing_indices = np.where(index_map != -1)[0]
index_length = index_map.shape[0]
center_panel_hints = None
if detector_geometry is not None:
raw_centers = np.column_stack((existing_center_x, existing_center_y))
corrected_centers, center_panel_hints = raw_to_corrected_points(raw_centers, detector_geometry)
existing_center_x_fit = corrected_centers[:, 0]
existing_center_y_fit = corrected_centers[:, 1]
else:
existing_center_x_fit = existing_center_x
existing_center_y_fit = existing_center_y
# Divide indices into batches
batches = [existing_indices[i:i + batch_size] for i in range(0, len(existing_indices), batch_size)]
check_cancel(stop_event, "batch preparation")
log_print(f"[DEBUG] Number of batches to process: {len(batches)}")
def _ensure_dask_has_workers(current_client):
"""Return a client that has at least one worker, else None; will attempt restart once."""
try:
if len(current_client.scheduler_info().get("workers", {})) > 0:
return current_client
except Exception as exc: # pragma: no cover - best-effort guard
log_print(f"[WARN] Unable to query Dask scheduler info ({exc}); disabling Dask for this iteration")
return None
log_print("[WARN] Dask client has 0 workers; attempting to restart cluster...")
try:
DaskClientManager.close_client()
restarted = DaskClientManager.get_client()
if restarted and len(restarted.scheduler_info().get("workers", {})) > 0:
log_print("[INFO] Dask cluster restarted with active workers")
return restarted
except Exception as exc: # pragma: no cover
log_print(f"[WARN] Failed to restart Dask cluster ({exc}); falling back to local execution")
return None
results = []
if client is not None:
client = _ensure_dask_has_workers(client)
if client is not None:
# Scatter read-only arrays once to avoid re-serializing them in every task
futures = []
try:
centers_x_future = client.scatter(existing_center_x_fit, broadcast=True)
centers_y_future = client.scatter(existing_center_y_fit, broadcast=True)
index_map_future = client.scatter(index_map, broadcast=True)
log_print("[DEBUG] Submitting tasks to Dask client (streaming futures)...")
for indices_subset in batches:
check_cancel(stop_event, "beamstop dask submission")
fut = client.submit(
refine_center_with_friedel_batch,
h5file_path, min_peaks, tolerance, resolution_limit,
indices_subset, centers_x_future, centers_y_future, index_map_future,
detector_geometry, max_pairs_per_frame,
)
futures.append(fut)
log_print(f"[DEBUG] {len(futures)} futures submitted")
for fut in as_completed(futures):
check_cancel(stop_event, "beamstop dask collection")
results.append(fut.result())
check_cancel(stop_event, "beamstop dask collection")
log_print(f"[DEBUG {datetime.datetime.now().isoformat()}] Gathered {len(results)} results")
except (CancelledError, TimeoutError, Exception) as exc:
log_print(f"[WARN] Dask batch processing failed ({exc}); retrying iteration without Dask")
results = []
client = None
finally:
if stop_event is not None and stop_event.is_set():
for fut in futures:
try:
fut.cancel()
except Exception:
pass
# If Dask was unavailable or returned fewer results than expected, re-run sequentially
if client is None or len(results) < len(batches):
if client is not None:
log_print(f"[WARN] Dask returned {len(results)}/{len(batches)} batches; rerunning sequentially")
results = []
# Execute tasks sequentially without Dask
for indices_subset in batches:
check_cancel(stop_event, "beamstop sequential batch")
result = refine_center_with_friedel_batch(
h5file_path, min_peaks, tolerance, resolution_limit,
indices_subset, existing_center_x_fit, existing_center_y_fit, index_map,
detector_geometry, max_pairs_per_frame,
)
results.append(result)
# Collect deviations and indices
non_empty_results = [res for res in results if res[0].size > 0 and res[1].size > 0]
log_print(f"[DEBUG] {len(non_empty_results)} non-empty result batches")
log_print(f"[DEBUG {datetime.datetime.now().isoformat()}] Starting LOWESS fit and interpolation steps...")
if non_empty_results:
all_indices = np.concatenate([res[0] for res in non_empty_results])
all_deviations = np.concatenate([res[1] for res in non_empty_results])
# Sort the deviations and indices based on indices
sorted_order = np.argsort(all_indices)
all_indices = all_indices[sorted_order]
all_deviations = all_deviations[sorted_order]
else:
all_indices = np.array([])
all_deviations = np.array([])
#print(f"Indices after sorting: {all_indices}")
if all_deviations.size == 0:
# No deviations found; handle accordingly
log_start(logfile_path, f'Iteration {itcount}: No deviations found.')
return all_indices, all_deviations, converged, itcount, existing_center_x, existing_center_y, index_length
if deviation_aggregation in {"pair", "pairs", "none", "notebook"}:
fit_indices = all_indices
fit_deviations = all_deviations
else:
fit_indices, fit_deviations = _aggregate_deviations_per_frame(all_indices, all_deviations)
if fit_deviations.size == 0:
log_start(logfile_path, f'Iteration {itcount}: No frame-wise deviations found.')
return all_indices, all_deviations, converged, itcount, existing_center_x, existing_center_y, index_length
num_points = len(fit_indices)
if num_points == 0:
log_start(logfile_path, f'Iteration {itcount}: No valid indices found.')
return all_indices, all_deviations, converged, itcount, existing_center_x, existing_center_y, index_length
def _compute_lowess_frac():
frac = lowess_frac
if lowess_window is not None and lowess_window > 0:
frac = lowess_window / num_points
if num_points > 0:
min_frac = min(1.0, max(1.0 / num_points, 2.0 / num_points))
frac = max(frac, min_frac)
return min(1.0, frac)
effective_lowess_frac = _compute_lowess_frac()
log_print(f"[DEBUG] Using LOWESS frac={effective_lowess_frac:.6f} (window={lowess_window}) with {num_points} points")
# Use LOWESS smoothing algorithm, parallelizing X & Y fits if Dask client available
if client is not None:
indices_future = client.scatter(fit_indices, broadcast=True)
dev_x_future = client.scatter(fit_deviations[:, 0], broadcast=True)
dev_y_future = client.scatter(fit_deviations[:, 1], broadcast=True)
future_x = client.submit(_lowess_axis, dev_x_future, indices_future, effective_lowess_frac)
future_y = client.submit(_lowess_axis, dev_y_future, indices_future, effective_lowess_frac)
try:
lowess_x, lowess_y = client.gather([future_x, future_y])
except CancelledError as exc:
# If the Dask scheduler forgets the futures (e.g., worker restart or client disconnect),
# fall back to the in-process LOWESS implementation so the refinement can continue.
log_print(f"[WARN] Dask LOWESS futures were cancelled ({exc}); falling back to local LOWESS")
lowess_x, lowess_y = perform_lowess(fit_indices, fit_deviations, frac=effective_lowess_frac)
else:
lowess_x, lowess_y = perform_lowess(fit_indices, fit_deviations, frac=effective_lowess_frac)
log_print(f"[DEBUG {datetime.datetime.now().isoformat()}] Completed LOWESS fit")
# Create interpolation functions for the LOWESS fit
if lowess_x.shape[0] >= 2:
interp_lowess_x = interp1d(lowess_x[:, 0], lowess_x[:, 1], bounds_error=False, fill_value="extrapolate")
raw_interp_x = interp_lowess_x(existing_indices)
elif lowess_x.shape[0] == 1:
raw_interp_x = np.full(existing_indices.shape[0], float(lowess_x[0, 1]), dtype=np.float64)
else:
raw_interp_x = np.zeros(existing_indices.shape[0], dtype=np.float64)
if lowess_y.shape[0] >= 2:
interp_lowess_y = interp1d(lowess_y[:, 0], lowess_y[:, 1], bounds_error=False, fill_value="extrapolate")
raw_interp_y = interp_lowess_y(existing_indices)
elif lowess_y.shape[0] == 1:
raw_interp_y = np.full(existing_indices.shape[0], float(lowess_y[0, 1]), dtype=np.float64)
else:
raw_interp_y = np.zeros(existing_indices.shape[0], dtype=np.float64)
# Update centers based on LOWESS-smoothed deviations, respecting gaps in the raw index map
correction_x = np.zeros_like(existing_center_x_fit)
correction_y = np.zeros_like(existing_center_y_fit)
proc_indices = index_map[existing_indices]
correction_x[proc_indices] = raw_interp_x
correction_y[proc_indices] = raw_interp_y
updated_center_x_fit = existing_center_x_fit + 0.5 * correction_x
updated_center_y_fit = existing_center_y_fit + 0.5 * correction_y
if updated_center_x_fit.size > 1:
updated_center_x_fit[0] = updated_center_x_fit[1] # Ensure the first element is consistent
if updated_center_y_fit.size > 1:
updated_center_y_fit[0] = updated_center_y_fit[1]
if detector_geometry is not None:
corrected_centers = np.column_stack((updated_center_x_fit, updated_center_y_fit))
raw_centers, _ = corrected_to_raw_points(
corrected_centers,
detector_geometry,
panel_hints=center_panel_hints,
)
updated_center_x = raw_centers[:, 0]
updated_center_y = raw_centers[:, 1]
else:
updated_center_x = updated_center_x_fit
updated_center_y = updated_center_y_fit
# Name output file for LOWESS
npz_path = os.path.join(graphicsfolder_path, f'lowess_and_centers_it{itcount:03d}.npz')
np.savez(
npz_path,
lowess_x=lowess_x,
lowess_y=lowess_y,
existing_center_x=existing_center_x,
existing_center_y=existing_center_y,
existing_indices=existing_indices,
index_length=index_length,
)
log_print(f"[DEBUG] Saved LOWESS and centers NPZ to {npz_path}")
# Write the updated centers back to the HDF5 file
with h5py.File(h5file_path, 'r+') as workingfile:
workingfile['entry/data/center_x'][:] = updated_center_x
workingfile['entry/data/center_y'][:] = updated_center_y
log_print(f"[DEBUG] Updated HDF5 centers written to {h5file_path}")
# Log the update
log_start(logfile_path, f'Centers updated (iteration {itcount})')
# Check for convergence
convergence_x = np.all(np.abs(lowess_x[:, 1]) < convergence_threshold)
convergence_y = np.all(np.abs(lowess_y[:, 1]) < convergence_threshold)
if convergence_x and convergence_y:
converged = True
# Output results
log_print("[DEBUG] Calling output() to save deviations")
output(all_indices, all_deviations, graphicsfolder_path, logfile_path, itcount, index_length=index_length)
log_print("[DEBUG] output() complete; iteration done")
return all_indices, all_deviations, converged, itcount, updated_center_x, updated_center_y, index_length
def _aggregate_deviations_per_frame(indices: np.ndarray, deviations: np.ndarray):
"""
Collapse pair-level deviations to one robust deviation per frame.
Using per-frame medians avoids overweighting frames that simply have
more detected Friedel pairs.
"""
if indices.size == 0 or deviations.size == 0:
return np.array([], dtype=np.int64), np.empty((0, 2), dtype=np.float64)
if deviations.ndim != 2 or deviations.shape[1] != 2:
raise ValueError("Deviations must be a 2D array with two columns.")
order = np.argsort(indices)
idx_sorted = np.asarray(indices, dtype=np.int64)[order]
dev_sorted = np.asarray(deviations, dtype=np.float64)[order]
unique_idx, start, counts = np.unique(idx_sorted, return_index=True, return_counts=True)
agg = np.empty((unique_idx.shape[0], 2), dtype=np.float64)
for i, (s, c) in enumerate(zip(start, counts)):
chunk = dev_sorted[s:s + c]
agg[i, 0] = float(np.median(chunk[:, 0]))
agg[i, 1] = float(np.median(chunk[:, 1]))
return unique_idx, agg
[docs]
def plot_lowess_and_centers(lowess_fit_x, lowess_fit_y, original_center_x, original_center_y, existing_indices, itcount, graphicsfolder_path):
"""Plot LOWESS fits and original centers."""
log_print('Plotting LOWESS fit and centers...')
plotname = f'lowess_and_centers_it{itcount:03d}.png'
fig, axes = plt.subplots(2, 1, figsize=(16, 12))
fig.suptitle(f"Center coordinates (it{itcount})")
# Plot for X-coordinate
ax1 = axes[0]
ax1.plot(existing_indices, original_center_x, label="Previous Center X", marker='o', markersize=2, linestyle='-', alpha=0.7)
ax1.set_xlabel("Index")
ax1.set_ylabel("X-coordinate")
ax1.legend(loc='upper left')
ax2 = ax1.twinx()
ax2.scatter(lowess_fit_x[:, 0], lowess_fit_x[:, 1], label="LOWESS fit for X-deviation", color='r', s=10)
ax2.set_ylabel("LOWESS X-deviation")
ax2.legend(loc='upper right')
# Plot for Y-coordinate
ax3 = axes[1]
ax3.plot(existing_indices, original_center_y, label="Previous Center Y", marker='o', markersize=2, linestyle='-', alpha=0.7)
ax3.set_xlabel("Index")
ax3.set_ylabel("Y-coordinate")
ax3.legend(loc='upper left')
ax4 = ax3.twinx()
ax4.scatter(lowess_fit_y[:, 0], lowess_fit_y[:, 1], label="LOWESS fit for Y-deviation", color='g', s=10)
ax4.set_ylabel("LOWESS Y-deviation")
ax4.legend(loc='upper right')
plt.tight_layout()
plt.savefig(os.path.join(graphicsfolder_path, plotname))
plt.close()
log_print(f"Saved LOWESS and centers plot as {plotname}")
[docs]
def output(indices, deviations, graphicsfolder_path, logfile_path, itcount, index_length=None):
"""Save deviations for GUI."""
if index_length is None and indices.size > 0:
index_length = int(indices.max()) + 1
final_path = os.path.join(graphicsfolder_path, f'deviations_it{itcount:03d}.npz')
# Write atomically so GUI readers don't see partially-written NPZ files
tmp_file = None
try:
os.makedirs(graphicsfolder_path, exist_ok=True)
with tempfile.NamedTemporaryFile(
dir=graphicsfolder_path, prefix=f'deviations_it{itcount:03d}_', suffix='.npz', delete=False
) as tmp_file:
np.savez(tmp_file, indices=indices, deviations=deviations, index_length=index_length)
tmp_path = tmp_file.name
os.replace(tmp_path, final_path)
except Exception as exc:
# If atomic write fails (e.g., network FS rename issues), fall back to direct write
log_print(f"[WARN] Atomic write for deviations file failed ({exc}); falling back to direct save")
np.savez(final_path, indices=indices, deviations=deviations, index_length=index_length)
finally:
# Clean up temp file if it still exists
if tmp_file is not None:
try:
os.remove(tmp_file.name)
except OSError:
pass
# Plotting can be done later with plot_deviations(...)
[docs]
def plot_deviations(indices, deviations, graphicsfolder_path, itcount, lowess_frac=0.1):
"""Plot deviations with LOWESS fit."""
log_print('Plotting deviations with LOWESS fit...')
plotname = f'lowess_it{itcount:03d}.png'
# Perform LOWESS fit
z_x, z_y = perform_lowess(indices, deviations, frac=lowess_frac)
plt.figure(figsize=(16, 6))
plt.scatter(indices, deviations[:, 0], label="X-deviation", marker='o', s=3, alpha=0.5)
plt.scatter(indices, deviations[:, 1], label="Y-deviation", marker='x', s=3, alpha=0.5)
# Plot LOWESS fits
plt.plot(z_x[:, 0], z_x[:, 1], label="LOWESS fit for X-deviation", color='r')
plt.plot(z_y[:, 0], z_y[:, 1], label="LOWESS fit for Y-deviation", color='g')
plt.xlabel("Frame")
plt.ylabel("Deviation")
plt.legend(loc='lower right')
plt.title(f"Deviation for individual frames with LOWESS fit (it{itcount})")
plt.savefig(os.path.join(graphicsfolder_path, plotname))
plt.close()
log_print(f"Saved deviations plot as {plotname}")
[docs]
def refine(configfile, client, batch_size=500, stop_event=None):
"""Main center refinement function."""
log_print(f"[DEBUG] refine() called with configfile: {configfile}")
check_cancel(stop_event, "beamstop refine start")
### New version
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
# Get config
config = read_config(configfile)
detector_geometry = load_detector_geometry(config)
timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
graphicsfolder_path = os.path.join(outputfolder_path, f'refinecenters_{timestamp}')
framepath = config.get('Paths', 'framepath')
# Ensure graphics folder exists
os.makedirs(graphicsfolder_path, exist_ok=True)
# Open HDF5 file to get total number of indices
with h5py.File(h5file_path, 'r') as workingfile:
_, framesize_x, framesize_y = workingfile[framepath].shape
if 'entry/data/index' in workingfile:
total_raw_frames = workingfile['entry/data/index'].shape[0]
else:
total_raw_frames = workingfile['entry/data/center_x'].shape[0]
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."
)
framesize = framesize_x
# Check if necessary parameters are defined
required_params = [
'centerrefinement_tolerance',
'centerrefinement_min_peaks',
'centerrefinement_resolution_limit',
'centerrefinement_max_iterations',
'centerrefinement_convergence_threshold'
]
for param in required_params:
if not config.has_option('Parameters', param):
error_message = f"{param.split('_', 1)[1]} not defined, peak finding interrupted"
log_result(logfile_path, '', error_message)
raise Exception(error_message)
# Load parameters
tolerance = float(config.get('Parameters', 'centerrefinement_tolerance'))
min_peaks = float(config.get('Parameters', 'centerrefinement_min_peaks'))
resolution_limit = float(config.get('Parameters', 'centerrefinement_resolution_limit'))
max_pairs_per_frame = int(float(config.get('Parameters', 'centerrefinement_max_pairs_per_frame', fallback="10000")))
max_iterations = int(float(config.get('Parameters', 'centerrefinement_max_iterations')))
convergence_threshold = float(config.get('Parameters', 'centerrefinement_convergence_threshold'))
deviation_aggregation = config.get(
'Parameters',
'centerrefinement_deviation_aggregation',
fallback='median',
).strip().lower()
valid_aggregations = {'median', 'frame_median', 'pair', 'pairs', 'none', 'notebook'}
if deviation_aggregation not in valid_aggregations:
raise ValueError(
"centerrefinement_deviation_aggregation must be one of: "
"median, frame_median, pair, pairs, none, notebook"
)
if deviation_aggregation == 'frame_median':
deviation_aggregation = 'median'
if config.has_option('Parameters', 'centerrefinement_lowess_frac'):
lowess_frac = float(config.get('Parameters', 'centerrefinement_lowess_frac'))
else:
lowess_frac = 0.1
if config.has_option('Parameters', 'centerrefinement_lowess_window'):
lowess_window = int(float(config.get('Parameters', 'centerrefinement_lowess_window')))
if lowess_window <= 0:
lowess_window = None
else:
lowess_window = None
pixels_per_meter = float(config.get('AcquisitionDetails', 'pixels_per_meter'))
if detector_geometry is not None:
log_start(logfile_path, "Detector geometry enabled for Friedel center refinement (internal corrected coordinates).")
if client is not None:
if sys.platform == "darwin":
log_start(
logfile_path,
"Using macOS Dask safe mode for geometry-enabled Friedel center refinement (threaded workers).",
)
os.environ["COSEDA_DASK_PROCESSES"] = "0"
try:
DaskClientManager.kill_current_cluster()
client = DaskClientManager.get_client()
except Exception:
client = None
else:
log_start(
logfile_path,
"Geometry-enabled center refinement on non-macOS: keeping default Dask process workers.",
)
# 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"centerrefinement_tolerance={tolerance}\n")
pf.write(f"centerrefinement_min_peaks={min_peaks}\n")
pf.write(f"centerrefinement_resolution_limit={resolution_limit}\n")
pf.write(f"centerrefinement_max_pairs_per_frame={max_pairs_per_frame}\n")
pf.write(f"centerrefinement_max_iterations={max_iterations}\n")
pf.write(f"centerrefinement_convergence_threshold={convergence_threshold}\n")
pf.write(f"centerrefinement_deviation_aggregation={deviation_aggregation}\n")
pf.write(f"centerrefinement_lowess_frac={lowess_frac}\n")
pf.write(f"centerrefinement_lowess_window={lowess_window}\n")
pf.write(f"pixels_per_meter={pixels_per_meter}\n")
pf.write(f"detector_geometry_enabled={detector_geometry is not None}\n")
# Log the start of refinement
with open(f'{logfile_path}', 'a') as file:
file.write(f"{datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]}; "
f"starting center refinement (tolerance = {tolerance}, min_peaks = {min_peaks}, "
f"resolution_limit = {resolution_limit}, max_iterations = {max_iterations}, "
f"convergence_threshold = {convergence_threshold}, "
f"deviation_aggregation = {deviation_aggregation}, "
f"lowess_frac = {lowess_frac}, lowess_window = {lowess_window})\n")
itcount = 0
converged = False
indices = np.array([])
deviations = np.empty((0, 2))
index_length = total_raw_frames
# Track iteration timing
import datetime as _datetime_debug
while itcount < max_iterations and not converged:
check_cancel(stop_event, "beamstop iteration loop")
# Debug: starting new iteration
log_print(f"[DEBUG] Starting iteration {itcount+1}/{max_iterations} at {_datetime_debug.datetime.now().isoformat()}")
itcount += 1
indices, deviations, converged, itcount, updated_center_x, updated_center_y, index_length = refine_center_with_friedel_and_update(
h5file_path, min_peaks, tolerance, resolution_limit, itcount, convergence_threshold,
converged, graphicsfolder_path, logfile_path, client, lowess_frac=lowess_frac,
lowess_window=lowess_window, batch_size=batch_size, stop_event=stop_event,
detector_geometry=detector_geometry, max_pairs_per_frame=max_pairs_per_frame,
deviation_aggregation=deviation_aggregation,
)
# Debug: completed iteration
log_print(f"[DEBUG] Completed iteration {itcount}/{max_iterations} at {_datetime_debug.datetime.now().isoformat()}")
# Final logging based on convergence
if converged:
log_print(f'Convergence criterion (deviation from LOWESS < {convergence_threshold}) met after {itcount} iteration(s)')
with open(f'{logfile_path}', 'a') as file:
file.write(f"{datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]}; "
f"convergence criterion (deviation from LOWESS < {convergence_threshold}) "
f"for center refinement met after {itcount} iteration(s)\n")
else:
log_print(f'Could not meet convergence criterion ({convergence_threshold}) after {itcount} iterations, refinement terminated')
with open(f'{logfile_path}', 'a') as file:
file.write(f"{datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]}; "
f"could not meet convergence criterion ({convergence_threshold}) "
f"for center refinement after {itcount} iterations, refinement terminated\n")
# Update detector shifts one final time
update_detector_shifts(
h5file_path,
logfile_path,
framepath,
pixels_per_meter,
detector_geometry=detector_geometry,
)
# Output final results
output(indices, deviations, graphicsfolder_path, logfile_path, itcount, index_length=index_length)
try:
with h5py.File(h5file_path, "r+") as workingfile:
write_nxprocess_centerrefinement(
workingfile,
program="coseda.centerrefinement.refinecenters_beamstop",
method="beamstop",
input_path=h5file_path,
output_path=h5file_path,
parameters={
"tolerance": tolerance,
"min_peaks": min_peaks,
"resolution_limit": resolution_limit,
"max_pairs_per_frame": max_pairs_per_frame,
"max_iterations": max_iterations,
"convergence_threshold": convergence_threshold,
"deviation_aggregation": deviation_aggregation,
"lowess_frac": lowess_frac,
"lowess_window": lowess_window if lowess_window is not None else "none",
"pixels_per_meter": pixels_per_meter,
"batch_size": batch_size,
"usedask": client is not None,
"iterations_run": itcount,
"converged": converged,
},
)
except Exception as exc:
log_start(logfile_path, f"Failed to write NXprocess center refinement: {exc}")
[docs]
def refinecenters_batch(input_path, usedask=False, batch_size=500, stop_event=None):
"""Process all configuration files for center refinement."""
configfiles, input_path = handle_input(input_path)
log_print(f'Following .ini files were found and will be processed:')
log_print(f'{configfiles}')
fileocount = 1
# Initialize Dask client if usedask is True
if usedask:
client = DaskClientManager.get_client()
dashboard_address, num_workers, threads_per_worker, memory_per_worker = DaskClientManager.get_client_info()
DaskClientManager.log_client_info(logfile_path=None) # Adjust if you want to log client info
log_print('')
log_print(f'Click here to monitor progress: {dashboard_address}')
log_print('')
else:
client = None
for configfile in configfiles:
check_cancel(stop_event, "beamstop batch file loop")
shoutout(configfile)
### New version
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
# Get config
config = read_config(configfile)
refine(configfile, client=client, batch_size=batch_size, stop_event=stop_event)
if fileocount < len(configfiles):
log_print("")
log_print(f"Proceeding to next file ({fileocount+1}/{len(configfiles)})")
log_print("")
fileocount += 1
log_print(f"Finished centerrefinement for all files ({fileocount-1}/{len(configfiles)})")