Source code for coseda.nexus.groups

"""Helpers to ensure NeXus parent groups exist with NX_class attributes."""

from __future__ import annotations

import h5py
import numpy as np

from coseda.nexus.paths import (
    BEAM_DOSE_PATH,
    ENTRY_DEFINITION_PATH,
    SAMPLE_GROUP_PATH,
    SAMPLE_TRANSFORMATIONS_PATH,
    SAMPLE_TRANSFORM_PATH,
)




def ensure_nexus_parents(h5file: h5py.File) -> None:
    """Create core NeXus parent groups and set NX_class attributes."""
    entry = h5file.require_group("entry")
    entry.attrs["NX_class"] = "NXentry"

    definition_name = ENTRY_DEFINITION_PATH.rsplit("/", 1)[-1]
    if definition_name not in entry:
        dtype = h5py.string_dtype(encoding="utf-8")
        entry.create_dataset(definition_name, data="NXcserialed", dtype=dtype)

    sample = h5file.require_group(SAMPLE_GROUP_PATH.lstrip("/"))
    sample.attrs["NX_class"] = "NXsample"
    trans = h5file.require_group(SAMPLE_TRANSFORMATIONS_PATH.lstrip("/"))
    trans.attrs["NX_class"] = "NXtransformations"
    if SAMPLE_TRANSFORM_PATH.lstrip("/") not in h5file:
        ds = trans.create_dataset("sample_pos", data=0.0)
        ds.attrs["transformation_type"] = "translation"
        ds.attrs["vector"] = np.array([0.0, 0.0, 0.0], dtype=np.float32)
        ds.attrs["units"] = "m"
        ds.attrs["depends_on"] = "."
    if "depends_on" not in sample.attrs:
        sample.attrs["depends_on"] = SAMPLE_TRANSFORM_PATH

    instrument = entry.require_group("instrument")
    instrument.attrs["NX_class"] = "NXinstrument"

    detector = instrument.require_group("detector")
    detector.attrs["NX_class"] = "NXdetector"

    beam = instrument.require_group("beam")
    beam.attrs["NX_class"] = "NXbeam"

    dose_name = BEAM_DOSE_PATH.rsplit("/", 1)[-1]
    if dose_name not in beam:
        ds = beam.create_dataset(dose_name, data=float("nan"))
        ds.attrs["units"] = "e-/A^2"
        ds.attrs["long_name"] = "Electron dose at sample"