Source code for coseda.nexus.images

"""NeXus-friendly NXdata view for image stacks."""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np
import h5py

from coseda.nexus.paths import IMAGE_DATA_PATH, IMAGE_NXDATA_PATH

_COORDINATE_NOTE = "detector_pixel (origin=center of lower-left pixel, +y up)"


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, units: str, long_name: str) -> None:
    if name in group:
        if group[name].shape != (length,):
            del group[name]
        else:
            group[name][...] = np.arange(length, dtype=group[name].dtype)
            _set_attr_if_missing(group[name].attrs, "units", units)
            _set_attr_if_missing(group[name].attrs, "long_name", long_name)
            return
    ds = group.create_dataset(name, data=np.arange(length, dtype=np.int32))
    ds.attrs["units"] = units
    ds.attrs["long_name"] = long_name


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


def _get_image_shape(h5file: h5py.File) -> Optional[Tuple[int, ...]]:
    if IMAGE_DATA_PATH not in h5file:
        return None
    return tuple(h5file[IMAGE_DATA_PATH].shape)




def ensure_image_nxdata(h5file: h5py.File) -> None:
    """
    Create an NXdata view for the image dataset without moving the legacy data.
    """
    shape = _get_image_shape(h5file)
    if not shape:
        return

    images_ds = h5file[IMAGE_DATA_PATH]
    if images_ds.ndim not in (2, 3):
        return

    nx_group = h5file.require_group(IMAGE_NXDATA_PATH.lstrip("/"))
    nx_group.attrs["NX_class"] = "NXdata"
    nx_group.attrs["signal"] = "images"
    nx_group.attrs["coordinate_system"] = _COORDINATE_NOTE

    _link_dataset(nx_group, "images", images_ds)

    if images_ds.ndim == 3:
        num_frames, height, width = images_ds.shape
        nx_group.attrs["axes"] = ["frame", "pixel_y", "pixel_x"]
        _ensure_axis_dataset(nx_group, "frame", num_frames, "index", "Frame index")
    else:
        height, width = images_ds.shape
        nx_group.attrs["axes"] = ["pixel_y", "pixel_x"]

    _ensure_axis_dataset(nx_group, "pixel_y", height, "pixel", "Detector pixel Y")
    _ensure_axis_dataset(nx_group, "pixel_x", width, "pixel", "Detector pixel X")