Source code for ran.utils._hdf5_load

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Loader for HDF5 files from Matlab."""

from pathlib import Path
from typing import Any

import h5py  # type: ignore[import-untyped]
import numpy as np




def hdf5_load(filename: Path | str) -> dict[str, Any]:
    """
    Load data from an HDF5 file.

    Args:
        filename (Path): Path object pointing to the HDF5 file

    Returns
    -------
        dict: Dictionary containing datasets
    """
    result = {}
    filename = Path(filename)
    with h5py.File(filename, "r") as f:
        for name, item in f["/"].items():
            if not isinstance(item, h5py.Dataset):
                # Current implementation only supports datasets (not groups)
                raise TypeError(f"Expected dataset, got {type(item)}")

            data = item[()]

            # Check if this is a complex number dataset (has 're' and 'im' fields)
            if isinstance(data, np.ndarray) and data.dtype.names is not None:
                if set(data.dtype.names) == {"re", "im"}:
                    # Reconstruct complex array
                    result[name] = np.transpose(data["re"] + 1j * data["im"])
                else:
                    # Regular structured array
                    result[name] = data
            else:
                # Regular array
                result[name] = np.transpose(data)

    return result