Source code for psi_io.data

"""
Module for fetching HDF5 assets used through examples.

This module uses the ``pooch`` library to manage the downloading and caching of
HDF4 and HDF5 files that adhere to PSI data conventions. It defines functions to
fetch specific example datasets, including 1D radial scale data, 2D coronal hole
maps, 3D radial magnetic field data, magnetic fieldline data, and synchronic maps
used in coronal and heliospheric magnetic field modeling.

Currently, these files are hosted on the PredSci documentation website:
at https://www.predsci.com/doc/assets/ and are primarily intended for use in
building examples in the PSI I/O and mapflpy packages.
"""


from __future__ import annotations

import inspect
from functools import wraps
from typing import Callable, ParamSpec, TypeVar

from psi_io.psi_io import HdfExtType, HDFEXT

try:
    import pooch
except ImportError as e:
    raise ImportError(
        "Missing the optional 'pooch' dependency required for data fetching. "
        "Please install it via pip or conda to access the necessary datasets."
    ) from e


REGISTRY = {
	"h4h5-files/rscale.h5": "sha256:60a0cbcd4dc69f7d250cbbdddd6fc3680f09d87c1e4cee6a79d8ec3731533718",
    "h4h5-files/chmap.h5": "sha256:668b5fe7e86903e6af4effdf65e3d2dd499a1217e93ca60d8b54b68941b6f1f7",
    "h4h5-files/fieldline.h5": "sha256:a5b2a1cc0c458d0d9510d8eacc93d3b4a2cc7e99e0a3f86cd3d6b164e74f370d",
    "h4h5-files/br.h5": "sha256:2038dc8e67303cf0b31414d532352b40e8c75ebd8917bc8b68614cf4e7b24055",
    "h4h5-files/rscale.hdf": "sha256:1c15bd669fc5a92dfdda7dc23703294c23f0a09440599fd5c30cf7a0e1a6f3c4",
    "h4h5-files/chmap.hdf": "sha256:fa2f1134aa4f1c9c0dd729b4e8f23f480bea5cb178e44e8da01bdffad09a2225",
    "h4h5-files/fieldline.hdf": "sha256:a4149783780e1ce44a8fe76a83c674e0a3082cd78c6a635b6c8e860e0fdd3891",
    "h4h5-files/br.hdf": "sha256:3a4b3174e5d6f45244bd25826890486b5659196b8fe093541c542375a88cdf52",
    "h4h5-files/synchronic_map.h5": "sha256:170794a5a19684246339ca9782a2b89066b89661400ec48bb6fc0a082e0a2450"
}
"""Registry of available magnetic field files with their SHA256 hashes. 

This registry is used by the pooch fetcher to verify the integrity of
downloaded files, and is primarily intended for building sphinx-gallery
examples that require MHD data files.
"""


BASE_URL = "https://www.predsci.com/doc/assets/"
"""Base URL hosting magnetic field file assets.
"""


FETCHER = pooch.create(
    path=pooch.os_cache("psi"),
    base_url=BASE_URL,
    registry=REGISTRY,
    env="PSI_IO_CACHE",
)
"""Pooch fetcher for downloading and caching magnetic field files.

.. note::
    The cache directory can be overridden by setting the ``PSI_IO_CACHE``
    environment variable to a desired path. Otherwise, the default cache
    directory is platform-dependent, as determined by :func:`pooch.os_cache`.
    
.. note::
    The default (os-dependent) cache directory stores assets under a
    subdirectory named ``psi``. The reason for this naming choice – as opposed
    to ``psi_io`` – is to maintain consistency with other PredSci packages
    that utilize the same asset hosting and caching mechanism.
"""

_P = ParamSpec("_P")
_R = TypeVar("_R")



def check_hdf_type(func: Callable[_P, _R]) -> Callable[_P, _R]:
    """Validate the ``hdf`` keyword argument of a data-fetch function.

    This decorator inspects the bound ``hdf`` parameter of the wrapped function
    and raises :exc:`ValueError` if the value is not one of the supported HDF
    file extensions (``'.h5'`` or ``'.hdf'``).

    Parameters
    ----------
    func : Callable
        The data-fetch function to wrap.  Must accept an ``hdf`` keyword
        argument that names the desired file extension.

    Returns
    -------
    out : Callable
        The wrapped function with ``hdf`` validation applied.

    Raises
    ------
    ValueError
        If the ``hdf`` argument is not a member of :data:`~psi_io.psi_io.HDFEXT`.

    Examples
    --------
    >>> from psi_io import data
    >>> data.get_3d_data.__wrapped__  # doctest: +ELLIPSIS
    <function get_3d_data at 0x...>
    """
    sig = inspect.signature(func)

    @wraps(func)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()

        hdf = bound.arguments["hdf"]  # assumes the param is literally named "hdf"
        if hdf not in HDFEXT:
            raise ValueError(f"Invalid HDF type {hdf!r}. Must be in {sorted(HDFEXT)}.")

        return func(*bound.args, **bound.kwargs)

    return wrapper





def file_ids() -> list[str]:
    """List all available magnetic field files in the registry.

    Returns
    -------
    out : list[str]
        File names (registry keys) for all available magnetic field assets.

    Examples
    --------
    >>> from psi_io import data
    >>> ids = data.file_ids()
    >>> isinstance(ids, list) and len(ids) > 0
    True
    """
    return list(FETCHER.registry.keys())





@check_hdf_type
def get_1d_data(hdf: HdfExtType = ".h5") -> str:
    """Fetch the radial scale (1D) data file.

    Parameters
    ----------
    hdf : HdfExtType, optional
        The HDF file format to fetch.  Accepted values are ``'.h5'`` for HDF5
        and ``'.hdf'`` for HDF4.  Default is ``'.h5'``.

    Returns
    -------
    out : str
        Local path to the downloaded (and cached) radial scale data file.

    Examples
    --------
    >>> from psi_io import data
    >>> path = data.get_1d_data()
    >>> path.endswith(".h5")
    True
    """
    filename = f"h4h5-files/rscale{hdf}"
    return FETCHER.fetch(filename)





@check_hdf_type
def get_2d_data(hdf: HdfExtType = ".h5") -> str:
    """Fetch the coronal hole map (2D) data file.

    Parameters
    ----------
    hdf : HdfExtType, optional
        The HDF file format to fetch.  Accepted values are ``'.h5'`` for HDF5
        and ``'.hdf'`` for HDF4.  Default is ``'.h5'``.

    Returns
    -------
    out : str
        Local path to the downloaded (and cached) coronal hole map data file.

    Examples
    --------
    >>> from psi_io import data
    >>> path = data.get_2d_data()
    >>> path.endswith(".h5")
    True
    """
    filename = f"h4h5-files/chmap{hdf}"
    return FETCHER.fetch(filename)





@check_hdf_type
def get_3d_data(hdf: HdfExtType = ".h5") -> str:
    """Fetch the radial magnetic field (3D) data file.

    Parameters
    ----------
    hdf : HdfExtType, optional
        The HDF file format to fetch.  Accepted values are ``'.h5'`` for HDF5
        and ``'.hdf'`` for HDF4.  Default is ``'.h5'``.

    Returns
    -------
    out : str
        Local path to the downloaded (and cached) radial magnetic field data
        file.

    Examples
    --------
    >>> from psi_io import data
    >>> path = data.get_3d_data()
    >>> path.endswith(".h5")
    True
    """
    filename = f"h4h5-files/br{hdf}"
    return FETCHER.fetch(filename)





@check_hdf_type
def get_fieldline_data(hdf: HdfExtType = ".h5") -> str:
    """Fetch the magnetic fieldline (2D) data file.

    .. warning::
        Unlike the other example data files, fieldline data files do not
        contain scale datasets.

    Parameters
    ----------
    hdf : HdfExtType, optional
        The HDF file format to fetch.  Accepted values are ``'.h5'`` for HDF5
        and ``'.hdf'`` for HDF4.  Default is ``'.h5'``.

    Returns
    -------
    out : str
        Local path to the downloaded (and cached) magnetic fieldline data file.

    Examples
    --------
    >>> from psi_io import data
    >>> path = data.get_fieldline_data()
    >>> path.endswith(".h5")
    True
    """
    filename = f"h4h5-files/fieldline{hdf}"
    return FETCHER.fetch(filename)





@check_hdf_type
def get_synchronic_map_data(hdf: HdfExtType = ".h5") -> str:
    """Fetch the synchronic map data file.

    .. warning::
        Synchronic map data is only available in HDF5 format. Furthermore,
        unlike the other example data files, synchronic map data files contain
        additional datasets beyond the primary data and scales.

    Parameters
    ----------
    hdf : HdfExtType, optional
        The HDF file format to fetch.  Accepted values are ``'.h5'`` for HDF5
        and ``'.hdf'`` for HDF4.  Default is ``'.h5'``.

    Returns
    -------
    out : str
        Local path to the downloaded (and cached) synchronic map data file.

    Examples
    --------
    >>> from psi_io import data
    >>> path = data.get_synchronic_map_data()
    >>> path.endswith(".h5")
    True
    """
    if hdf == ".hdf":
        raise NotImplemented("Synchronic map data is only available in HDF5 format.")
    filename = f"h4h5-files/synchronic_map{hdf}"
    return FETCHER.fetch(filename)