from coseda.logging_utils import log_print
from coseda.centerrefinement.cancellation import RefinementCancelled, check_cancel
from numba import njit, prange
@njit(parallel=True)
def quadrant_variance_numba(img, cx, cy, half, deltas_x, deltas_y):
Ny, Nx = deltas_y.shape[0], deltas_x.shape[0]
H, W = img.shape
out = np.empty((Ny, Nx), dtype=np.float64)
for iy in prange(Ny):
dy = deltas_y[iy]
for ix in range(Nx):
dx = deltas_x[ix]
cx_dx = cx + dx
cy_dy = cy + dy
s0 = s1 = s2 = s3 = 0.0
for m in range(-half, half+1):
if m == 0: continue
py = cy_dy + m
y0 = int(np.floor(py)); wy = py - y0
if y0 < 0: y0, wy = 0, 0.0
elif y0 >= H-1: y0, wy = H-2, 1.0
y1 = y0 + 1
for n in range(-half, half+1):
if n == 0: continue
px = cx_dx + n
x0 = int(np.floor(px)); wx = px - x0
if x0 < 0: x0, wx = 0, 0.0
elif x0 >= W-1: x0, wx = W-2, 1.0
x1 = x0 + 1
v00 = img[y0, x0]; v10 = img[y0, x1]
v01 = img[y1, x0]; v11 = img[y1, x1]
val = (1-wy)*(1-wx)*v00 + (1-wy)*wx*v10 + wy*(1-wx)*v01 + wy*wx*v11
if m < 0:
if n < 0: s0 += val
else: s1 += val
else:
if n < 0: s2 += val
else: s3 += val
mean = 0.25 * (s0 + s1 + s2 + s3)
out[iy, ix] = 0.25 * ((s0-mean)**2 + (s1-mean)**2 + (s2-mean)**2 + (s3-mean)**2)
return out
import numpy as np
from numba import njit, prange
import h5py
import tifffile
from tifffile import TiffFileError
from tqdm.notebook import tqdm
from coseda.dask_client_manager import DaskClientManager
import matplotlib.pyplot as plt
import os
from datetime import datetime
from coseda.io import handle_input, read_config, config_to_paths
from coseda.logging_utils import log_result, log_start, shoutout
from coseda.centerfinding.findcenters_beamstop import update_detector_shifts
from coseda.nexus.process import write_nxprocess_centerrefinement
[docs]
def resolve_loess_span(loess_span, loess_frac, total_frames, default_frac=0.1):
"""
Determine the LOESS span (in frame units) from either an explicit span or a fraction.
Returns a tuple (span, fraction_used). If both are missing and default_frac is None or <=0,
span will be None to indicate smoothing should be skipped.
"""
if loess_span is not None:
try:
span_value = float(loess_span)
except (TypeError, ValueError):
raise ValueError(f"Invalid LOESS span value: {loess_span}")
if span_value <= 0:
return None, loess_frac
return span_value, loess_frac
frac = loess_frac if loess_frac is not None else default_frac
if frac is None or frac <= 0:
return None, frac
if total_frames is None:
raise ValueError("total_frames must be provided when deriving LOESS span from a fraction.")
window = max(1, int(np.ceil(frac * total_frames)))
span_value = max(1.0, window / 2.0)
return span_value, frac
# Outputs direct refinement results to numpy files.
[docs]
def output_direct(
indices,
coarse,
fine,
outputfolder_path,
logfile_path,
loess_span: float = None,
loess_frac: float = None,
total_frames: int = None,
iteration: int = 0
):
"""Outputs direct refinement results to numpy files."""
import os
import numpy as np
# iteration suffix for filenames
suffix = f"_it{iteration:03d}"
# ensure output folder exists
os.makedirs(outputfolder_path, exist_ok=True)
# save raw direct outputs for debug
np.save(os.path.join(outputfolder_path, f'direct_refine_indices{suffix}.npy'), indices)
np.save(os.path.join(outputfolder_path, f'direct_refine_coarse{suffix}.npy'), coarse)
np.save(os.path.join(outputfolder_path, f'direct_refine_fine{suffix}.npy'), fine)
# also save beamstop-style .npz files for GUI preview
# compute true deviations
deviations = fine - coarse
indices = np.asarray(indices)
if deviations.shape[0] != indices.shape[0]:
indices_for_fit = np.arange(deviations.shape[0])
else:
indices_for_fit = indices
from coseda.centerrefinement.refinecenters_direct import lowess_xunit, resolve_loess_span
effective_span, fraction_used = resolve_loess_span(
loess_span,
loess_frac,
total_frames if total_frames is not None else deviations.shape[0]
)
if effective_span is None or deviations.size == 0:
# Fallback: no smoothing applied
zero_dev = np.zeros(deviations.shape[0], dtype=float)
loess_x = np.column_stack((indices_for_fit, zero_dev))
loess_y = np.column_stack((indices_for_fit, zero_dev))
smoothed_dev_x = loess_x[:, 1]
smoothed_dev_y = loess_y[:, 1]
else:
try:
loess_x = lowess_xunit(indices_for_fit, deviations[:, 0], indices_for_fit, effective_span)
loess_y = lowess_xunit(indices_for_fit, deviations[:, 1], indices_for_fit, effective_span)
smoothed_dev_x = loess_x[:, 1]
smoothed_dev_y = loess_y[:, 1]
except ValueError as exc:
log_result(logfile_path, "Failed to compute LOESS smoothing; falling back to unsmoothed values.", str(exc))
zero_dev = np.zeros(deviations.shape[0], dtype=float)
loess_x = np.column_stack((indices_for_fit, zero_dev))
loess_y = np.column_stack((indices_for_fit, zero_dev))
smoothed_dev_x = loess_x[:, 1]
smoothed_dev_y = loess_y[:, 1]
# compute smoothed centers
# assemble smoothed center coordinates
smoothed_centers_x = fine[:, 0] + smoothed_dev_x
smoothed_centers_y = fine[:, 1] + smoothed_dev_y
smoothed_centers = np.vstack((smoothed_centers_x, smoothed_centers_y)).T
np.savez(
os.path.join(outputfolder_path, f'lowess_and_centers{suffix}.npz'),
lowess_x=loess_x,
lowess_y=loess_y,
existing_center_x=smoothed_centers_x,
existing_center_y=smoothed_centers_y,
existing_indices=indices_for_fit,
loess_span=effective_span,
loess_frac=fraction_used
)
# save full refinement results
np.savez(
os.path.join(outputfolder_path, f'refinement_results{suffix}.npz'),
indices=indices,
coarse=fine,
fine=smoothed_centers
)
# deviations as fine minus coarse
deviations = fine - coarse
np.savez(
os.path.join(outputfolder_path, f'deviations{suffix}.npz'),
indices=indices,
deviations=deviations
)
# log the output
log_result(logfile_path, f"Saved direct refinement results to {outputfolder_path}", None)
# Custom LOESS function using a fixed window span in x-units
[docs]
def lowess_xunit(x, y, xout, span):
"""
Custom LOESS smoothing using a fixed window span in x-units.
x, y: original data arrays.
xout: points at which to evaluate the fit (subset of x).
span: half-window width in the same units as x.
Returns array of shape (len(xout), 2) with columns (xout, smoothed y).
"""
fitted = np.zeros((len(xout), 2), dtype=float)
for i, xi in enumerate(xout):
mask = np.abs(x - xi) <= span
xi_sel = x[mask]; yi_sel = y[mask]
if len(xi_sel) < 2:
fitted[i] = (xi, yi_sel.mean() if len(yi_sel)>0 else 0.0)
continue
# tricube weights
dist = np.abs(xi_sel - xi) / span
w = (1 - dist**3) ** 3
coeffs = np.polyfit(xi_sel, yi_sel, 1, w=w)
yi = np.polyval(coeffs, xi)
fitted[i] = (xi, yi)
return fitted
[docs]
def refine_direct_single_frame(
frameindex: int,
frame: np.ndarray,
x0: float,
y0: float,
sigma_x: float,
sigma_y: float,
threshold_frac: float,
box1_half: int,
step1: float,
step2: float,
R_min: float = None,
R_max: float = None,
debug_plot_path: str = None,
) -> dict:
"""
Perform adaptive two-stage quadrant-variance refinement on a single frame.
Returns a dict with keys 'coarse' and 'fine' containing (x,y) tuples.
"""
import numpy as np
from scipy.ndimage import map_coordinates
if R_min is None:
R_min = 0.5 # fallback minimum range for good fit
if R_max is None:
R_max = 10 # fallback maximum range for bad fit
box2_half = box1_half # fine pass uses same half-size
# compute fine-pass half-range to cover quantization error
R2 = max(2 * step1, 10 * step2)
# shrink working window: crop frame to ROI around the initial center
import math
margin = int(math.ceil(max(3*sigma_x, 3*sigma_y, R2) + box1_half))
y0i = int(round(y0)); x0i = int(round(x0))
y1_roi = max(0, y0i - margin); y2_roi = min(frame.shape[0], y0i + margin + 1)
x1_roi = max(0, x0i - margin); x2_roi = min(frame.shape[1], x0i + margin + 1)
frame_roi = frame[y1_roi:y2_roi, x1_roi:x2_roi]
# adjust initial coordinates to ROI-local
x0_local = x0 - x1_roi
y0_local = y0 - y1_roi
# Vectorized interpolation and quadrant variance computation
def interpolate_roi(img, xc, yc, half, R, step):
rel = np.arange(-(half + R), half + R + step, step)
RX, RY = np.meshgrid(rel, rel)
X_full = xc + RX
Y_full = yc + RY
from scipy.ndimage import map_coordinates
interpolated_roi = map_coordinates(img, [Y_full.ravel(), X_full.ravel()], order=1)
interpolated_roi = interpolated_roi.reshape(RY.shape)
return interpolated_roi, rel
@njit(parallel=True)
def quadrant_var_cost(interpolated_roi, center, box, step, dys, dxs):
n_dy, n_dx = len(dys), len(dxs)
costs = np.zeros((n_dy, n_dx))
for idx_y in prange(n_dy):
yc = center + int(np.round(dys[idx_y] / step))
for idx_x in range(n_dx):
xc = center + int(np.round(dxs[idx_x] / step))
q1_sum = interpolated_roi[yc-box:yc, xc-box:xc].sum()
q2_sum = interpolated_roi[yc-box:yc, xc+1:xc+box+1].sum()
q3_sum = interpolated_roi[yc+1:yc+box+1, xc-box:xc].sum()
q4_sum = interpolated_roi[yc+1:yc+box+1, xc+1:xc+box+1].sum()
mean_quad = (q1_sum + q2_sum + q3_sum + q4_sum) / 4.0
var_quad = ((q1_sum - mean_quad)**2 + (q2_sum - mean_quad)**2 +
(q3_sum - mean_quad)**2 + (q4_sum - mean_quad)**2) / 4.0
costs[idx_y, idx_x] = var_quad
return costs
# Adaptive coarse pass (±2σ, fallback ±3σ) using vectorized interpolation
R1_x = np.clip(2 * sigma_x, R_min, R_max)
R1_y = np.clip(2 * sigma_y, R_min, R_max)
max_R1 = max(R1_x, R1_y)
dx1 = np.arange(-R1_x, R1_x + step1/2, step1)
dy1 = np.arange(-R1_y, R1_y + step1/2, step1)
# Interpolate ROI for coarse grid
interpolated_roi1, rel1 = interpolate_roi(frame_roi, x0_local, y0_local, box1_half, max_R1, step1)
box = int(box1_half / step1)
center = len(rel1) // 2
cost1 = quadrant_var_cost(interpolated_roi1, center, box, step1, dy1, dx1)
i1, j1 = np.unravel_index(np.argmin(cost1), cost1.shape)
best_dx, best_dy = dx1[j1], dy1[i1]
# fallback to ±3σ if on edge
if abs(best_dx) >= (1 - threshold_frac) * R1_x or abs(best_dy) >= (1 - threshold_frac) * R1_y:
R1_x = np.clip(3 * sigma_x, R_min, R_max)
R1_y = np.clip(3 * sigma_y, R_min, R_max)
max_R1 = max(R1_x, R1_y)
dx1 = np.arange(-R1_x, R1_x + step1/2, step1)
dy1 = np.arange(-R1_y, R1_y + step1/2, step1)
interpolated_roi1, rel1 = interpolate_roi(frame_roi, x0_local, y0_local, box1_half, max_R1, step1)
box = int(box1_half / step1)
center = len(rel1) // 2
cost1 = quadrant_var_cost(interpolated_roi1, center, box, step1, dy1, dx1)
i1, j1 = np.unravel_index(np.argmin(cost1), cost1.shape)
best_dx, best_dy = dx1[j1], dy1[i1]
# if still on edge after 3σ fallback, skip this frame
if abs(best_dx) >= (1 - threshold_frac) * R1_x or abs(best_dy) >= (1 - threshold_frac) * R1_y:
return None
# convert coarse result back to global frame
x1_local = x0_local + best_dx
y1_local = y0_local + best_dy
x1 = x1_local + x1_roi
y1 = y1_local + y1_roi
# Fine pass (small ±quantization error window), vectorized interpolation
dx2 = np.arange(-R2, R2 + step2/2, step2)
dy2 = np.arange(-R2, R2 + step2/2, step2)
x1_local = x1 - x1_roi
y1_local = y1 - y1_roi
interpolated_roi2, rel2 = interpolate_roi(frame_roi, x1_local, y1_local, box2_half, R2, step2)
box = int(box2_half / step2)
center = len(rel2) // 2
cost2 = quadrant_var_cost(interpolated_roi2, center, box, step2, dy2, dx2)
i2, j2 = np.unravel_index(np.argmin(cost2), cost2.shape)
xf_local = x1_local + dx2[j2]
yf_local = y1_local + dy2[i2]
xf = xf_local + x1_roi
yf = yf_local + y1_roi
result = {'index': frameindex, 'coarse': (x1, y1), 'fine': (xf, yf)}
# optional debug plot
if debug_plot_path:
import matplotlib.pyplot as plt
# recreate debug panels
fig, axs = plt.subplots(1, 3, figsize=(15, 4))
# adaptive coarse cost
im0 = axs[0].imshow(cost1, origin='lower',
extent=[dx1.min(), dx1.max(), dy1.min(), dy1.max()])
axs[0].scatter(0, 0, c='white', marker='x')
axs[0].scatter(best_dx, best_dy, c='red', marker='o')
axs[0].set_title('Adaptive Coarse Cost')
# fine cost
im1 = axs[1].imshow(cost2, origin='lower',
extent=[-R2, R2, -R2, R2])
axs[1].scatter(0, 0, c='white', marker='x')
axs[1].scatter(dx2[j2], dy2[i2], c='red', marker='o')
axs[1].set_title('Fine Cost')
# final full-beam patch
rel = np.arange(-box2_half, box2_half+1)
RX, RY = np.meshgrid(rel, rel)
patch = map_coordinates(frame, [(y1 + RY).ravel(), (x1 + RX).ravel()], order=1).reshape(RX.shape)
axs[2].imshow(patch, cmap='gray', origin='lower')
axs[2].axvline(box2_half, ls='--', c='red')
axs[2].axhline(box2_half, ls='--', c='red')
axs[2].set_title('Final Patch')
plt.tight_layout()
fig.savefig(debug_plot_path)
plt.close(fig)
return result
[docs]
def refine_direct_batch(
h5file_path: str,
frame_indices: np.ndarray,
sigma_x: float,
sigma_y: float,
threshold_frac: float,
box1_half: int,
step1: float,
step2: float,
R_min: float = None,
R_max: float = None,
debug_plot_dir: str = None,
stop_event=None,
):
# ——— LOAD HDF5 ONCE ———
with h5py.File(h5file_path, 'r') as f:
frames = f['entry/data/images'][frame_indices, :, :]
batch_x = f['entry/data/center_x'][frame_indices]
batch_y = f['entry/data/center_y'][frame_indices]
# precompute shift grids
min_r = R_min if R_min is not None else 0.5
max_r = R_max if R_max is not None else 10.0
rxc = np.clip(2*sigma_x, min_r, max_r)
ryc = np.clip(2*sigma_y, min_r, max_r)
dx1 = np.arange(-rxc, rxc + step1/2, step1)
dy1 = np.arange(-ryc, ryc + step1/2, step1)
radius_fine = max(2 * step1, 10 * step2)
dx2 = np.arange(-radius_fine, radius_fine + step2/2, step2)
dy2 = np.arange(-radius_fine, radius_fine + step2/2, step2)
# JIT warm-up on first frame
_ = quadrant_variance_numba(frames[0], batch_x[0], batch_y[0], box1_half, dx1, dy1)
cost_dummy = quadrant_variance_numba(frames[0], batch_x[0], batch_y[0], box1_half, dx1, dy1)
i_dummy, j_dummy = np.unravel_index(np.argmin(cost_dummy), cost_dummy.shape)
cx_dummy = batch_x[0] + dx1[j_dummy]
cy_dummy = batch_y[0] + dy1[i_dummy]
_ = quadrant_variance_numba(frames[0], cx_dummy, cy_dummy, box1_half, dx2, dy2)
valid_indices = []
coarse_results = []
fine_results = []
# process each frame
for i, idx in enumerate(frame_indices):
check_cancel(stop_event, "direct frame refinement")
img = frames[i]
x0 = batch_x[i]; y0 = batch_y[i]
# coarse
cost1 = quadrant_variance_numba(img, x0, y0, box1_half, dx1, dy1)
i1, j1 = np.unravel_index(np.argmin(cost1), cost1.shape)
cx = x0 + dx1[j1]; cy = y0 + dy1[i1]
# fine
cost2 = quadrant_variance_numba(img, cx, cy, box1_half, dx2, dy2)
i2, j2 = np.unravel_index(np.argmin(cost2), cost2.shape)
xf = cx + dx2[j2]; yf = cy + dy2[i2]
valid_indices.append(idx)
coarse_results.append((cx, cy))
fine_results.append((xf, yf))
return {
'indices': np.array(valid_indices),
'coarse': np.array(coarse_results),
'fine': np.array(fine_results),
}
[docs]
def quadrant_var_cost(img, xc, yc, half, dxs, dys):
from scipy.ndimage import map_coordinates
costs = np.zeros((dys.size, dxs.size))
rel = np.arange(-half, half + 1)
RX, RY = np.meshgrid(rel, rel)
for i, dy in enumerate(dys):
for j, dx in enumerate(dxs):
X = xc + dx + RX
Y = yc + dy + RY
patch = map_coordinates(img, [Y.ravel(), X.ravel()], order=1).reshape(RX.shape)
m = half
q1, q2 = patch[:m, :m], patch[:m, m+1:]
q3, q4 = patch[m+1:, :m], patch[m+1:, m+1:]
costs[i, j] = np.var([q.sum() for q in (q1, q2, q3, q4)])
return costs
[docs]
def refine_direct_all(
h5file_path: str,
all_indices: np.ndarray,
sigma_x: float,
sigma_y: float,
threshold_frac: float,
box1_half: int,
step1: float,
step2: float,
batch_size: int = 500,
R_min: float = None,
R_max: float = None,
debug_plot_dir: str = None,
usedask: bool = False,
stop_event=None,
):
"""
Driver: divides frames into batches, optionally calls refine_direct_batch for each batch in parallel with Dask.
Collects and returns all valid results.
"""
import dask
check_cancel(stop_event, "direct batch preparation")
# Initialize Dask client if usedask is True
if usedask:
dask_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('')
# warm up Numba on each worker
def _warmup_worker():
import numpy as _np
# dummy small array
dummy = _np.zeros((2*box1_half+1, 2*box1_half+1))
_ = quadrant_variance_numba(dummy, 0.0, 0.0, box1_half, _np.array([0.0]), _np.array([0.0]))
dask_client.run(_warmup_worker)
else:
dask_client = None
all_valid_indices = []
all_coarse = []
all_fine = []
n_batches = int(np.ceil(len(all_indices) / batch_size))
tasks = []
for ibatch in range(n_batches):
check_cancel(stop_event, "direct batch loop")
i0 = ibatch * batch_size
i1 = min((ibatch+1) * batch_size, len(all_indices))
batch_indices = all_indices[i0:i1]
log_print(f"Processing batch {ibatch+1}/{n_batches} (frames {i0} to {i1-1}) ...")
if usedask and dask_client is not None:
# submit as Dask task
task = dask_client.submit(
refine_direct_batch,
h5file_path,
batch_indices,
sigma_x,
sigma_y,
threshold_frac,
box1_half,
step1,
step2,
R_min,
R_max,
debug_plot_dir,
)
tasks.append(task)
else:
# run serially
batch_res = refine_direct_batch(
h5file_path,
batch_indices,
sigma_x,
sigma_y,
threshold_frac,
box1_half,
step1,
step2,
R_min=R_min,
R_max=R_max,
debug_plot_dir=debug_plot_dir,
stop_event=stop_event,
)
all_valid_indices.append(batch_res['indices'])
all_coarse.append(batch_res['coarse'])
all_fine.append(batch_res['fine'])
# Collect Dask results if needed
if usedask and dask_client is not None and tasks:
# allow mid-flight cancellation by the user
check_cancel(stop_event, "direct dask gather")
try:
batch_results = dask_client.gather(tasks)
finally:
if stop_event is not None and stop_event.is_set():
for t in tasks:
try:
t.cancel()
except Exception:
pass
for batch_res in batch_results:
check_cancel(stop_event, "direct dask gather")
all_valid_indices.append(batch_res['indices'])
all_coarse.append(batch_res['coarse'])
all_fine.append(batch_res['fine'])
# Concatenate results from all batches
if all_valid_indices:
all_valid_indices = np.concatenate(all_valid_indices)
all_coarse = np.vstack(all_coarse)
all_fine = np.vstack(all_fine)
else:
all_valid_indices = np.empty((0,), dtype=int)
all_coarse = np.empty((0, 2))
all_fine = np.empty((0, 2))
return {
'indices': all_valid_indices,
'coarse': all_coarse,
'fine': all_fine,
}
[docs]
def refine_and_update_centers(
h5file_path: str,
sigma_x: float,
sigma_y: float,
threshold_frac: float,
box1_half: int,
step1: float,
step2: float,
batch_size: int = 500,
R_min: float = None,
R_max: float = None,
debug_plot_dir: str = None,
logfile_path: str = None,
loess_span: float = None,
loess_frac: float = None,
usedask: bool = False,
convergence_threshold: float = None,
skip_fit: bool = False,
stop_event=None,
):
"""
Runs refine_direct_all over all frames, updates center_x and center_y datasets
with the new fine centers (for all valid frames).
"""
check_cancel(stop_event, "direct refine_and_update setup")
# Get total number of frames and HDF5 index map
with h5py.File(h5file_path, 'r') as f:
n_frames = f['entry/data/images'].shape[0]
index_map = f['entry/data/index'][:] # orig_frame -> new_pos or -1
# Map original frames to new-image positions, keep only valid
raw_indices_map = np.where(index_map != -1)[0]
all_indices = index_map[raw_indices_map] # new-image indices
# Run the direct refinement over all frames
result = refine_direct_all(
h5file_path,
all_indices,
sigma_x,
sigma_y,
threshold_frac,
box1_half,
step1,
step2,
batch_size=batch_size,
R_min=R_min,
R_max=R_max,
debug_plot_dir=debug_plot_dir,
usedask=usedask,
stop_event=stop_event,
)
# Store original-frame indices for GUI export
result['raw_indices'] = raw_indices_map
indices = result['indices']
fine_centers = result['fine'] # shape (N,2), columns (x,y)
# keep copy of original fine centers for convergence check
orig_centers = fine_centers.copy()
# store pre-smoothing results
indices_pre = indices.copy()
coarse_pre = result['coarse']
check_cancel(stop_event, "direct LOESS span resolve")
effective_loess_span, effective_loess_frac = resolve_loess_span(
loess_span,
loess_frac,
n_frames
)
if skip_fit is not True:
# optionally smooth fine centers with LOESS, preserving correct spacing over gaps
if effective_loess_span is not None:
# map current image indices back to original frame numbers
with h5py.File(h5file_path, 'r') as f:
index_map = f['entry/data/index'][:] # orig_frame -> new_pos or -1
# invert mapping: new image position -> original frame number
new_to_orig = np.empty(n_frames, dtype=int)
for orig_idx, new_idx in enumerate(index_map):
if new_idx != -1:
new_to_orig[new_idx] = orig_idx
# original frame numbers for refined frames and for all frames
orig_indices = new_to_orig[indices] # x-values for existing points
orig_all = new_to_orig # xout for every frame
# apply LOESS on original frame timeline
loess_fit_x = lowess_xunit(orig_indices, fine_centers[:, 0], orig_all, effective_loess_span)
smooth_x = loess_fit_x[:, 1]
loess_fit_y = lowess_xunit(orig_indices, fine_centers[:, 1], orig_all, effective_loess_span)
smooth_y = loess_fit_y[:, 1]
# assemble smoothed centers and reset indices to full image range
fine_centers = np.vstack((smooth_x, smooth_y)).T
indices = np.arange(n_frames)
# build full-length coarse array matching smoothed fine_centers
coarse_full = np.zeros((n_frames, 2))
coarse_full[indices_pre] = coarse_pre
# update result dict for output_direct
result['indices'] = indices
result['coarse'] = coarse_full
result['fine'] = fine_centers
# compute convergence if threshold provided
if convergence_threshold is not None:
dx = fine_centers[:, 0] - orig_centers[:, 0]
dy = fine_centers[:, 1] - orig_centers[:, 1]
max_dx = np.max(np.abs(dx))
max_dy = np.max(np.abs(dy))
converged = (max_dx < convergence_threshold) and (max_dy < convergence_threshold)
# report convergence metrics
log_start(logfile_path,
f"max_dx: {max_dx:.4f}; max_dy: {max_dy:.4f}; threshold: {convergence_threshold}; converged: {converged}")
else:
converged = False
else:
log_start(logfile_path, "Skipping LOESS smoothing and convergence check.")
converged = True
# Update the centers in the HDF5 file
with h5py.File(h5file_path, 'r+') as f:
center_x = f['entry/data/center_x']
center_y = f['entry/data/center_y']
center_x[indices] = fine_centers[:, 0]
center_y[indices] = fine_centers[:, 1]
log_result(logfile_path, f"Updated centers for {len(indices)} frames.", None)
# annotate result with smoothing metadata
result['loess_span_used'] = effective_loess_span
result['loess_frac_used'] = effective_loess_frac
result['total_frames'] = n_frames
return result, converged
[docs]
def refinecenters_direct_file(configfile, client=None, batch_size=50, stop_event=None):
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
"""Main center refinement function."""
log_start(logfile_path, f"refinecenters_direct_file called with configfile: {configfile}")
# Generate timestamped graphics output directory
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
graphics_subfolder = os.path.join(outputfolder_path, f'refinecenters_{timestamp}')
os.makedirs(graphics_subfolder, exist_ok=True)
# Get config
config = read_config(configfile)
# Read direct refinement parameters from config
threshold_frac = float(config.get('Parameters', 'centerrefinement_direct_threshold_frac'))
box1_half = int(config.get('Parameters', 'centerrefinement_direct_box1_half'))
step1 = float(config.get('Parameters', 'centerrefinement_direct_step1'))
step2 = float(config.get('Parameters', 'centerrefinement_direct_step2'))
sigma_x = float(config.get('Parameters', 'centerrefinement_direct_sigma_x'))
sigma_y = float(config.get('Parameters', 'centerrefinement_direct_sigma_y'))
# optional LOESS span (frames units)
loess_span = float(config.get('Parameters', 'centerrefinement_direct_loess_span')) \
if config.has_option('Parameters', 'centerrefinement_direct_loess_span') else None
loess_frac = float(config.get('Parameters', 'centerrefinement_direct_loess_frac')) \
if config.has_option('Parameters', 'centerrefinement_direct_loess_frac') else None
# optional batch size override
if config.has_option('Parameters', 'centerrefinement_direct_batch_size'):
batch_size = int(config.get('Parameters', 'centerrefinement_direct_batch_size'))
# optional number of direct refinement cycles
cycles = int(config.get('Parameters', 'centerrefinement_direct_cycles')) \
if config.has_option('Parameters', 'centerrefinement_direct_cycles') else 1
# optional convergence threshold override
convergence_threshold = float(config.get('Parameters', 'centerrefinement_direct_convergence_threshold')) \
if config.has_option('Parameters', 'centerrefinement_direct_convergence_threshold') else None
# use Dask if a client was provided
usedask = client is not None
# Dump direct refinement parameters to paramdump.txt
param_file = os.path.join(graphics_subfolder, 'paramdump.txt')
with open(param_file, 'w') as pf:
pf.write(f"centerrefinement_direct_threshold_frac={threshold_frac}\n")
pf.write(f"centerrefinement_direct_box1_half={box1_half}\n")
pf.write(f"centerrefinement_direct_step1={step1}\n")
pf.write(f"centerrefinement_direct_step2={step2}\n")
pf.write(f"centerrefinement_direct_sigma_x={sigma_x}\n")
pf.write(f"centerrefinement_direct_sigma_y={sigma_y}\n")
pf.write(f"centerrefinement_direct_loess_span={loess_span}\n")
pf.write(f"centerrefinement_direct_loess_frac={loess_frac}\n")
pf.write(f"centerrefinement_direct_cycles={cycles}\n")
pf.write(f"centerrefinement_direct_convergence_threshold={convergence_threshold}\n")
framepath = config.get('Paths', 'framepath')
pixels_per_meter = float(config.get('AcquisitionDetails', 'pixels_per_meter'))
try:
skip_fit = config.getboolean('Parameters', 'centerrefinement_direct_skip_fit')
except:
skip_fit = False
# Log the start of refinement
log_start(logfile_path, "starting center refinement")
# run direct center refinement cycles with convergence check
result = None
converged = False
cycle = 0
while cycle < cycles and not converged:
check_cancel(stop_event, "direct cycle loop")
cycle += 1
log_start(logfile_path, f"Starting direct refinement cycle {cycle}/{cycles}")
result, converged = refine_and_update_centers(
h5file_path,
sigma_x,
sigma_y,
threshold_frac,
box1_half,
step1,
step2,
batch_size=batch_size,
R_min=None,
R_max=None,
debug_plot_dir=None,
logfile_path=logfile_path,
loess_span=loess_span,
loess_frac=loess_frac,
usedask=usedask,
convergence_threshold=convergence_threshold,
skip_fit=skip_fit,
stop_event=stop_event
)
if skip_fit is False:
log_start(logfile_path, f"Completed direct refinement cycle {cycle}/{cycles}, converged={converged}")
else:
log_start(logfile_path, f"Completed direct refinement, skipping fit.")
# save iteration outputs
output_direct(
result['raw_indices'],
result['coarse'],
result['fine'],
graphics_subfolder,
logfile_path,
loess_span=result.get('loess_span_used'),
loess_frac=result.get('loess_frac_used'),
total_frames=result.get('total_frames'),
iteration=cycle-1
)
log_start(logfile_path, f"Completed direct refinement: updated {len(result['indices'])} frames.")
# Update detector shifts one final time
update_detector_shifts(h5file_path, logfile_path, framepath, pixels_per_meter)
# Output final results
output_direct(
result['raw_indices'],
result['coarse'],
result['fine'],
graphics_subfolder,
logfile_path,
loess_span=result.get('loess_span_used'),
loess_frac=result.get('loess_frac_used'),
total_frames=result.get('total_frames'),
iteration=cycle-1
)
try:
with h5py.File(h5file_path, "r+") as workingfile:
write_nxprocess_centerrefinement(
workingfile,
program="coseda.centerrefinement.refinecenters_direct",
method="direct",
input_path=h5file_path,
output_path=h5file_path,
parameters={
"threshold_frac": threshold_frac,
"box1_half": box1_half,
"step1": step1,
"step2": step2,
"sigma_x": sigma_x,
"sigma_y": sigma_y,
"loess_span": loess_span if loess_span is not None else "none",
"loess_frac": loess_frac if loess_frac is not None else "none",
"batch_size": batch_size,
"cycles": cycles,
"convergence_threshold": convergence_threshold if convergence_threshold is not None else "none",
"pixels_per_meter": pixels_per_meter,
"usedask": usedask,
"converged": converged,
},
)
except Exception as exc:
log_start(logfile_path, f"Failed to write NXprocess center refinement: {exc}")
[docs]
def refinecenters_direct_batch(input_path, usedask=True, batch_size=50, stop_event=None):
"""Process all configuration files for center refinement."""
configfiles, input_path = handle_input(input_path)
# Determine logfile_path for batch logging using first config file
_, _, _, logfile_path, _, _ = config_to_paths(configfiles[0])
log_start(logfile_path, f"Following .ini files were found and will be processed: {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_start(logfile_path, f"Click here to monitor progress: {dashboard_address}")
else:
client = None
for configfile in configfiles:
check_cancel(stop_event, "direct batch file loop")
# Retrieve logfile_path for logging for each file
_, _, _, logfile_path, _, _ = config_to_paths(configfile)
log_start(logfile_path, f"Starting direct refinement for {configfile}")
refinecenters_direct_file(configfile, client=client, batch_size=batch_size, stop_event=stop_event)
if fileocount < len(configfiles):
log_start(logfile_path, f"Proceeding to next file ({fileocount+1}/{len(configfiles)})")
fileocount += 1
log_start(logfile_path, f"Finished centerrefinement for all files ({fileocount-1}/{len(configfiles)})")
[docs]
def write_centerrefinement_direct_settings(
input_path,
threshold_frac,
box1_half,
step1,
step2,
sigma_x,
sigma_y,
loess_span=None,
loess_frac=None,
batch_size=None,
cycles=None,
convergence_threshold=None,
):
"""
Update .ini files to set parameters for direct center refinement.
"""
configfiles, input_path = handle_input(input_path)
for configfile in configfiles:
shoutout(configfile)
config = read_config(configfile)
# Set direct refinement parameters
config.set('Parameters', 'centerrefinement_direct_threshold_frac', f'{threshold_frac}')
config.set('Parameters', 'centerrefinement_direct_box1_half', f'{box1_half}')
config.set('Parameters', 'centerrefinement_direct_step1', f'{step1}')
config.set('Parameters', 'centerrefinement_direct_step2', f'{step2}')
config.set('Parameters', 'centerrefinement_direct_sigma_x', f'{sigma_x}')
config.set('Parameters', 'centerrefinement_direct_sigma_y', f'{sigma_y}')
if loess_span is not None:
config.set('Parameters', 'centerrefinement_direct_loess_span', f'{loess_span}')
if loess_frac is not None:
config.set('Parameters', 'centerrefinement_direct_loess_frac', f'{loess_frac}')
if batch_size is not None:
config.set('Parameters', 'centerrefinement_direct_batch_size', f'{batch_size}')
if cycles is not None:
config.set('Parameters', 'centerrefinement_direct_cycles', f'{cycles}')
if convergence_threshold is not None:
config.set('Parameters', 'centerrefinement_direct_convergence_threshold', f'{convergence_threshold}')
# Write the updated config file
with open(configfile, 'w') as cfgfile:
config.write(cfgfile)
# Log the parameter update
_, _, _, logfile_path, _, _ = config_to_paths(configfile)
msg = (
f"parameters set for direct center refinement, "
f"threshold_frac = {threshold_frac}, box1_half = {box1_half}, "
f"step1 = {step1}, step2 = {step2}, "
f"sigma_x = {sigma_x}, sigma_y = {sigma_y}"
)
if cycles is not None:
msg += f", cycles = {cycles}"
if convergence_threshold is not None:
msg += f", convergence_threshold = {convergence_threshold}"
if loess_span is not None:
msg += f", loess_span = {loess_span}"
if loess_frac is not None:
msg += f", loess_frac = {loess_frac}"
if batch_size is not None:
msg += f", batch_size = {batch_size}"
log_start(logfile_path, msg)