Source code for coseda.nexus.indices

"""Helpers for index mappings between dense and sparse frame streams."""

from __future__ import annotations

import numpy as np
import h5py

from coseda.nexus.paths import IMAGE_DATA_PATH, IMAGE_KEY_PATH, INDEX_PATH




def ensure_image_key(h5file: h5py.File) -> None:
    """
    Ensure /entry/data/image_key exists, mapping current frames to original indices.

    If an index dataset exists and matches the image length, reuse it (unless it
    already contains stripped markers). Otherwise, fall back to arange.
    """
    if IMAGE_DATA_PATH not in h5file:
        return
    images = h5file[IMAGE_DATA_PATH]
    if images.ndim == 3:
        num_frames = images.shape[0]
    elif images.ndim == 2:
        num_frames = 1
    else:
        return

    if IMAGE_KEY_PATH in h5file:
        existing = h5file[IMAGE_KEY_PATH]
        if existing.ndim == 1 and existing.shape[0] == num_frames:
            return
        del h5file[IMAGE_KEY_PATH]

    mapping = None
    if INDEX_PATH in h5file:
        idx_ds = h5file[INDEX_PATH]
        if idx_ds.ndim == 1 and idx_ds.shape[0] == num_frames:
            idx_values = idx_ds[...]
            if np.all(idx_values != -1):
                mapping = idx_values

    if mapping is None:
        mapping = np.arange(num_frames, dtype=np.int32)

    h5file.create_dataset(IMAGE_KEY_PATH, data=mapping, dtype="i4")