Source code for ewoksid14.timepix4.t4_write

import h5py
import numpy

from . import t4_camera




def save_nxdata(
    filename: str,
    name: str,
    h3d: numpy.ndarray,
    toa_edges: numpy.ndarray,
    relative_toa: bool = True,
) -> numpy.ndarray:
    """
    Save a 3D histogram in NeXus NXdata format (HDF5) with TOA bin centers.
    The 3D histogram has shape (n_rows, n_cols, n_toa).
    Also creates an additional NXdata group with the histogram summed along time.

    :param filename: Output HDF5 filename.
    :param name: NXentry group name.
    :param h3d: 3D histogram array (n_rows, n_cols, n_toa).
    :param toa_edges: TOA bin edges corresponding to the last axis.
    :returns: 2D histogram
    """
    if relative_toa:
        time_units = "ns"
        factor = 0.5
        time_name = "rel_toa"
        long_name = "Relative TOA (ns)"
    else:
        time_units = "s"
        factor = 0.5e-9
        time_name = "toa"
        long_name = "Absolute TOA (s)"

    y = numpy.arange(t4_camera.n_rows)
    x = numpy.arange(t4_camera.n_cols)

    with h5py.File(filename, "w") as nxroot:
        nxroot.attrs["NX_class"] = "NXroot"
        nxroot.attrs["default"] = name

        # NXentry group
        nxentry = nxroot.create_group(name)
        nxentry.attrs["NX_class"] = "NXentry"
        nxentry.attrs["default"] = "timepix4_hist_2d"

        # 3D histogram
        nxdata_3d = nxentry.create_group("timepix4_hist_3d")
        nxdata_3d.attrs["NX_class"] = "NXdata"

        signal_3d = nxdata_3d.create_dataset("counts", data=h3d)
        signal_3d.attrs["units"] = "counts"
        signal_3d.attrs["interpretation"] = "image"
        nxdata_3d.attrs["signal"] = "counts"

        toa_centers = factor * (toa_edges[:-1] + toa_edges[1:])
        time = nxdata_3d.create_dataset(time_name, data=toa_centers)
        time.attrs.update({"units": time_units, "long_name": long_name})
        dy = nxdata_3d.create_dataset("y", data=y)
        dy.attrs.update({"units": "pixels", "long_name": "Rows (pixels)"})
        dx = nxdata_3d.create_dataset("x", data=x)
        dx.attrs.update({"units": "pixels", "long_name": "Columns (pixels)"})

        nxdata_3d.attrs["axes"] = ["y", "x", time_name]

        # 2D histogram summed over time
        h2d = h3d.sum(axis=2)  # sum over TOA (last axis)
        nxdata_2d = nxentry.create_group("timepix4_hist_2d")
        nxdata_2d.attrs["NX_class"] = "NXdata"

        signal_2d = nxdata_2d.create_dataset("counts", data=h2d)
        signal_2d.attrs["units"] = "counts"
        signal_2d.attrs["interpretation"] = "image"
        nxdata_2d.attrs["signal"] = "counts"

        dy = nxdata_2d.create_dataset("y", data=y)
        dy.attrs.update({"units": "pixels"})
        dx = nxdata_2d.create_dataset("x", data=x)
        dx.attrs.update({"units": "pixels"})
        nxdata_2d.attrs["axes"] = ["y", "x"]

    print(f"Saved {filename}")
    return h2d