Source code for coseda.nexus.peaks

"""NeXus-friendly NXdata views for peakfinding outputs."""

from __future__ import annotations

from typing import Optional

import numpy as np
import h5py

from coseda.nexus.paths import (
    NPEAKS_PATH,
    PEAK_INTENSITY_PATH,
    PEAK_X_PATH,
    PEAK_Y_PATH,
    PEAKS_NXDATA_PATH,
    PEAK_COUNTS_NXDATA_PATH,
)


def _set_attr_if_missing(attrs, key: str, value) -> None:
    if key not in attrs:
        attrs[key] = value


def _ensure_axis_dataset(group: h5py.Group, name: str, length: int) -> None:
    if name in group:
        if group[name].shape != (length,):
            del group[name]
        else:
            group[name][...] = np.arange(length, dtype=group[name].dtype)
            return
    group.create_dataset(name, data=np.arange(length, dtype=np.int32))


def _link_dataset(group: h5py.Group, name: str, target: h5py.Dataset) -> None:
    if name in group:
        del group[name]
    group[name] = target


def _ensure_peaks_group(h5file: h5py.File) -> Optional[int]:
    if not all(path in h5file for path in (PEAK_INTENSITY_PATH, PEAK_X_PATH, PEAK_Y_PATH)):
        return None

    intensity_ds = h5file[PEAK_INTENSITY_PATH]
    if intensity_ds.ndim != 2:
        return None

    num_frames, peaks_per_frame = intensity_ds.shape

    peaks_group = h5file.require_group(PEAKS_NXDATA_PATH.lstrip("/"))
    peaks_group.attrs["NX_class"] = "NXdata"
    peaks_group.attrs["signal"] = "peakTotalIntensity"
    peaks_group.attrs["axes"] = ["frame", "peak_index"]

    _link_dataset(peaks_group, "peakTotalIntensity", intensity_ds)
    _link_dataset(peaks_group, "peakXPosRaw", h5file[PEAK_X_PATH])
    _link_dataset(peaks_group, "peakYPosRaw", h5file[PEAK_Y_PATH])
    _ensure_axis_dataset(peaks_group, "frame", num_frames)
    _ensure_axis_dataset(peaks_group, "peak_index", peaks_per_frame)

    return num_frames


def _ensure_peak_counts_group(h5file: h5py.File, num_frames: Optional[int]) -> None:
    if NPEAKS_PATH not in h5file:
        return
    n_peaks_ds = h5file[NPEAKS_PATH]
    if n_peaks_ds.ndim != 1:
        return
    counts_group = h5file.require_group(PEAK_COUNTS_NXDATA_PATH.lstrip("/"))
    counts_group.attrs["NX_class"] = "NXdata"
    counts_group.attrs["signal"] = "nPeaks"
    counts_group.attrs["axes"] = ["frame"]
    _link_dataset(counts_group, "nPeaks", n_peaks_ds)
    axis_len = n_peaks_ds.shape[0] if num_frames is None else num_frames
    _ensure_axis_dataset(counts_group, "frame", axis_len)


def _annotate_peak_datasets(h5file: h5py.File) -> None:
    if NPEAKS_PATH in h5file:
        ds = h5file[NPEAKS_PATH]
        _set_attr_if_missing(ds.attrs, "long_name", "Number of peaks per frame")
        _set_attr_if_missing(ds.attrs, "units", "count")
    if PEAK_X_PATH in h5file:
        ds = h5file[PEAK_X_PATH]
        _set_attr_if_missing(ds.attrs, "long_name", "Peak X position (pixels)")
        _set_attr_if_missing(ds.attrs, "units", "pixel")
    if PEAK_Y_PATH in h5file:
        ds = h5file[PEAK_Y_PATH]
        _set_attr_if_missing(ds.attrs, "long_name", "Peak Y position (pixels)")
        _set_attr_if_missing(ds.attrs, "units", "pixel")
    if PEAK_INTENSITY_PATH in h5file:
        ds = h5file[PEAK_INTENSITY_PATH]
        _set_attr_if_missing(ds.attrs, "long_name", "Peak total intensity")
        _set_attr_if_missing(ds.attrs, "units", "adu")


[docs] def ensure_peak_nxdata(h5file: h5py.File) -> None: """ Create NXdata views for peakfinding outputs without moving legacy datasets. """ _annotate_peak_datasets(h5file) num_frames = _ensure_peaks_group(h5file) _ensure_peak_counts_group(h5file, num_frames)