"""Aurora-plot helpers for peak saturation inspection."""
from __future__ import annotations
import configparser
import os
from typing import Callable, Iterable, Optional
import h5py
import numpy as np
AURORA_WINDOW_SIZE = 5
AURORA_BINS = 420
AURORA_OUTPUT_SUBFOLDER = "aurora_plots"
BIT_DEPTH_ATTR_CANDIDATES = ("bit_depth", "bitDepth", "BitDepth")
def _read_ini_section(ini_path: Optional[str], section: str) -> dict:
if not ini_path:
return {}
parser = configparser.ConfigParser(
interpolation=None,
inline_comment_prefixes=(";", "#"),
)
try:
parser.read(ini_path)
except Exception:
return {}
if parser.has_section(section):
return dict(parser.items(section))
return {}
def _float_or_none(value) -> Optional[float]:
try:
out = float(value)
except (TypeError, ValueError):
return None
return out if np.isfinite(out) else None
[docs]
def aurora_points_cache_path(h5file_path: str, window_size: int = AURORA_WINDOW_SIZE) -> str:
"""Return the per-HDF5 cache path for aurora point data."""
folder = os.path.join(os.path.dirname(os.path.abspath(h5file_path)), AURORA_OUTPUT_SUBFOLDER)
stem = os.path.splitext(os.path.basename(h5file_path))[0]
size = int(max(1, window_size))
return os.path.join(folder, f"{stem}_aurora_peakmax_{size}x{size}.npy")
[docs]
def image_bit_depth_label(h5file_path: str) -> str:
"""Return a display label for the HDF5 image dataset bit depth."""
try:
with h5py.File(h5file_path, "r") as h5file:
ds = h5file.get("entry/data/images")
if ds is None:
return "unknown"
attr_depth = None
for key in BIT_DEPTH_ATTR_CANDIDATES:
if key in ds.attrs:
attr_depth = ds.attrs[key]
break
if isinstance(attr_depth, bytes):
attr_depth = attr_depth.decode(errors="replace")
if isinstance(attr_depth, np.ndarray):
attr_depth = attr_depth.item() if attr_depth.shape == () else attr_depth.tolist()
dtype = np.dtype(ds.dtype)
dtype_label = str(dtype)
if attr_depth not in (None, ""):
return f"{attr_depth}-bit ({dtype_label})"
if dtype.kind in "iu":
return f"{dtype.itemsize * 8}-bit ({dtype_label})"
if dtype.kind == "f":
return f"{dtype.itemsize * 8}-bit float ({dtype_label})"
return dtype_label
except Exception:
return "unknown"
def _fixed_center_from_ini(ini_path: Optional[str]):
params = _read_ini_section(ini_path, "Parameters")
if not params:
return None
x_raw = params.get("peakfinding_x0", params.get("pf9_x0"))
y_raw = params.get("peakfinding_y0", params.get("pf9_y0"))
x_val = _float_or_none(x_raw)
y_val = _float_or_none(y_raw)
if x_val is None or y_val is None:
return None
return x_val, y_val
def _sample_peak_max_intensities(image: np.ndarray, xs: np.ndarray, ys: np.ndarray, radius: int):
xi = np.rint(xs).astype(np.int64)
yi = np.rint(ys).astype(np.int64)
height, width = image.shape[:2]
in_bounds = (xi >= 0) & (xi < width) & (yi >= 0) & (yi < height)
if not np.any(in_bounds):
return np.array([], dtype=np.float64), in_bounds
xi_valid = xi[in_bounds]
yi_valid = yi[in_bounds]
values = np.full(xi_valid.shape, -np.inf, dtype=np.float64)
radius = max(0, int(radius))
for dy in range(-radius, radius + 1):
yj = yi_valid + dy
y_ok = (yj >= 0) & (yj < height)
if not np.any(y_ok):
continue
for dx in range(-radius, radius + 1):
xj = xi_valid + dx
ok = y_ok & (xj >= 0) & (xj < width)
if np.any(ok):
values[ok] = np.maximum(values[ok], image[yj[ok], xj[ok]])
return values, in_bounds
def _iter_aurora_points(
h5file: h5py.File,
total_frames: int,
chunk_size: int,
window_radius: int,
fixed_center,
):
data_group = h5file["entry/data"]
n_peaks_ds = data_group["nPeaks"]
peak_x_ds = data_group["peakXPosRaw"]
peak_y_ds = data_group["peakYPosRaw"]
image_ds = data_group["images"]
center_x_ds = data_group.get("center_x")
center_y_ds = data_group.get("center_y")
for start in range(0, total_frames, chunk_size):
end = min(start + chunk_size, total_frames)
n_peaks_chunk = np.asarray(n_peaks_ds[start:end], dtype=np.int64)
peak_x_chunk = np.asarray(peak_x_ds[start:end, :], dtype=np.float64)
peak_y_chunk = np.asarray(peak_y_ds[start:end, :], dtype=np.float64)
image_chunk = np.asarray(image_ds[start:end, :, :])
if fixed_center is None:
center_x_chunk = np.asarray(center_x_ds[start:end], dtype=np.float64)
center_y_chunk = np.asarray(center_y_ds[start:end], dtype=np.float64)
else:
chunk_len = end - start
center_x_chunk = np.full(chunk_len, fixed_center[0], dtype=np.float64)
center_y_chunk = np.full(chunk_len, fixed_center[1], dtype=np.float64)
radius_parts = []
intensity_parts = []
for local_idx, n_peaks in enumerate(n_peaks_chunk):
n_peaks = int(min(max(n_peaks, 0), peak_x_chunk.shape[1], peak_y_chunk.shape[1]))
if n_peaks <= 0:
continue
cx = center_x_chunk[local_idx]
cy = center_y_chunk[local_idx]
if not (np.isfinite(cx) and np.isfinite(cy)):
continue
xs = peak_x_chunk[local_idx, :n_peaks]
ys = peak_y_chunk[local_idx, :n_peaks]
valid = np.isfinite(xs) & np.isfinite(ys) & (xs >= 0) & (ys >= 0)
if not np.any(valid):
continue
xs = xs[valid]
ys = ys[valid]
intensities, in_bounds = _sample_peak_max_intensities(
image_chunk[local_idx], xs, ys, window_radius
)
xs = xs[in_bounds]
ys = ys[in_bounds]
radii = np.hypot(xs - cx, ys - cy)
valid_points = (
np.isfinite(radii)
& np.isfinite(intensities)
& (radii > 0)
& (intensities > 0)
)
if np.any(valid_points):
radius_parts.append(radii[valid_points])
intensity_parts.append(intensities[valid_points])
if radius_parts:
yield start, end, np.column_stack((
np.concatenate(radius_parts),
np.concatenate(intensity_parts),
))
else:
yield start, end, np.empty((0, 2), dtype=np.float64)
[docs]
def compute_aurora_points(
h5file_path: str,
ini_path: Optional[str] = None,
window_size: int = AURORA_WINDOW_SIZE,
chunk_size: int = 8,
progress_callback: Optional[Callable[[str, int, int], None]] = None,
) -> np.ndarray:
"""Compute raw aurora points as Nx2 columns: radius_px, peak_max_pixel."""
window_size = int(max(1, window_size))
window_radius = window_size // 2
chunk_size = int(max(1, chunk_size))
with h5py.File(h5file_path, "r") as h5file:
data_group = h5file.get("entry/data")
if data_group is None:
raise KeyError("Missing group '/entry/data'.")
required = ("images", "nPeaks", "peakXPosRaw", "peakYPosRaw")
missing = [name for name in required if name not in data_group]
if missing:
raise KeyError(f"Missing dataset(s) in /entry/data: {', '.join(missing)}")
fixed_center = None
if "center_x" not in data_group or "center_y" not in data_group:
fixed_center = _fixed_center_from_ini(ini_path)
if fixed_center is None:
raise KeyError(
"Missing '/entry/data/center_x' or '/entry/data/center_y', "
"and no numeric peakfinding_x0/peakfinding_y0 fallback was found in the INI."
)
total_frames = min(
int(data_group["images"].shape[0]),
int(data_group["nPeaks"].shape[0]),
int(data_group["peakXPosRaw"].shape[0]),
int(data_group["peakYPosRaw"].shape[0]),
)
if fixed_center is None:
total_frames = min(
total_frames,
int(data_group["center_x"].shape[0]),
int(data_group["center_y"].shape[0]),
)
if total_frames <= 0:
raise ValueError("No frames are available for aurora plotting.")
chunks = []
for start, end, points in _iter_aurora_points(
h5file,
total_frames,
chunk_size,
window_radius,
fixed_center,
):
if progress_callback:
progress_callback("Calculating peak maxima", end, total_frames)
if points.size:
chunks.append(points.astype(np.float32, copy=False))
if not chunks:
raise ValueError("No valid peaks were found for aurora plotting.")
return np.concatenate(chunks, axis=0)
[docs]
def load_or_compute_aurora_points(
h5file_path: str,
ini_path: Optional[str] = None,
window_size: int = AURORA_WINDOW_SIZE,
force_recompute: bool = False,
progress_callback: Optional[Callable[[str, int, int], None]] = None,
) -> tuple[np.ndarray, str, bool]:
"""Load cached aurora points or compute and cache them beside the HDF5 file."""
cache_path = aurora_points_cache_path(h5file_path, window_size)
if not force_recompute and os.path.isfile(cache_path):
points = np.load(cache_path)
if points.ndim == 2 and points.shape[1] == 2:
return points.astype(np.float32, copy=False), cache_path, True
points = compute_aurora_points(
h5file_path,
ini_path=ini_path,
window_size=window_size,
progress_callback=progress_callback,
)
os.makedirs(os.path.dirname(cache_path), exist_ok=True)
np.save(cache_path, points.astype(np.float32, copy=False))
return points, cache_path, False
[docs]
def compute_aurora_histogram_for_files(
file_specs: Iterable[dict],
window_size: int = AURORA_WINDOW_SIZE,
bins: int = AURORA_BINS,
force_recompute: bool = False,
progress_callback: Optional[Callable[[str, int, int], None]] = None,
) -> dict:
"""Build an aurora histogram from one or more HDF5/INI file specs."""
specs = [
spec for spec in file_specs
if spec.get("hdf5_path") and os.path.exists(spec.get("hdf5_path"))
]
if not specs:
raise ValueError("No HDF5 files are available for aurora plotting.")
bins = int(max(50, min(bins, 1200)))
x_parts = []
y_parts = []
bit_depth_labels = []
cache_paths = []
cache_hits = 0
total_files = len(specs)
for index, spec in enumerate(specs, 1):
h5_path = spec["hdf5_path"]
ini_path = spec.get("ini_path")
bit_depth_labels.append(image_bit_depth_label(h5_path))
def per_file_progress(phase, current, total, file_index=index, file_total=total_files):
if progress_callback:
name = os.path.basename(h5_path)
progress_callback(f"{phase}: {name}", file_index - 1 + current / max(total, 1), file_total)
points, cache_path, from_cache = load_or_compute_aurora_points(
h5_path,
ini_path=ini_path,
window_size=window_size,
force_recompute=force_recompute,
progress_callback=per_file_progress,
)
cache_paths.append(cache_path)
cache_hits += int(from_cache)
if progress_callback and from_cache:
progress_callback(f"Loaded cache: {os.path.basename(h5_path)}", index, total_files)
x_values = points[:, 0]
y_values = points[:, 1]
valid = (
np.isfinite(x_values)
& np.isfinite(y_values)
& (x_values > 0)
& (y_values > 0)
)
if np.any(valid):
x_parts.append(np.asarray(x_values[valid], dtype=np.float64))
y_parts.append(np.asarray(y_values[valid], dtype=np.float64))
if not x_parts:
raise ValueError("No plottable aurora points were found.")
x = np.concatenate(x_parts)
y = np.concatenate(y_parts)
x_edges = np.linspace(0.0, np.nextafter(float(np.max(x)), np.inf), bins + 1)
y_edges = np.linspace(0.0, np.nextafter(float(np.max(y)), np.inf), bins + 1)
hist, _, _ = np.histogram2d(x, y, bins=(x_edges, y_edges))
unique_bit_depths = sorted(set(bit_depth_labels))
bit_depth_label = unique_bit_depths[0] if len(unique_bit_depths) == 1 else "mixed: " + ", ".join(unique_bit_depths)
return {
"hist": hist,
"x_edges": x_edges,
"y_edges": y_edges,
"total_peaks": int(x.size),
"file_count": total_files,
"cache_paths": cache_paths,
"cache_hits": cache_hits,
"bit_depth_label": bit_depth_label,
"x_label": "Radius from beam center (px)",
"y_label": f"Peak max pixel intensity ({window_size}x{window_size})",
}
[docs]
def compute_aurora_histogram(
h5file_path: str,
ini_path: Optional[str] = None,
window_size: int = AURORA_WINDOW_SIZE,
bins: int = AURORA_BINS,
force_recompute: bool = False,
progress_callback: Optional[Callable[[str, int, int], None]] = None,
) -> dict:
"""Compatibility wrapper for a single-file aurora histogram."""
return compute_aurora_histogram_for_files(
[{"hdf5_path": h5file_path, "ini_path": ini_path}],
window_size=window_size,
bins=bins,
force_recompute=force_recompute,
progress_callback=progress_callback,
)