Source code for plasmapy.plasma.sources.openpmd_hdf5

"""Functionality for reading in HDF5 files following the OpenPMD standard."""

from types import TracebackType

__all__ = ["HDF5Reader"]

from pathlib import Path

import astropy.units as u
import h5py
import numpy as np
from packaging.version import Version

from plasmapy.plasma.exceptions import DataStandardError
from plasmapy.plasma.plasma_base import GenericPlasma

_OUTDATED_VERSION = "1.1.0"
_NEWER_VERSION = "2.0.0"

# This is the order what OpenPMD uses to store unit
# dimensions for a record.
_UNITS = (u.meter, u.kilogram, u.second, u.ampere, u.Kelvin, u.mol, u.candela)


def _fetch_units(openPMD_dims):
    """Converts a collection of OpenPMD dimensions to astropy.units."""

    units = u.dimensionless_unscaled
    for factor, unit in zip(openPMD_dims, _UNITS, strict=False):
        units *= unit**factor
    units, *_ = units.compose()
    return units


def _valid_version(openPMD_version, outdated=_OUTDATED_VERSION, newer=_NEWER_VERSION):
    """Checks if the passed version is supported or not."""

    parsed_version = Version(openPMD_version)
    outdated_version = Version(outdated)
    newer_version = Version(newer)
    return outdated_version <= parsed_version < newer_version


[docs] class HDF5Reader(GenericPlasma): """ Core class for accessing various attributes on HDF5 files that are based on `OpenPMD <https://www.openpmd.org>`__ standards. Parameters ---------- hdf5 : `str` Path to HDF5 file. **kwargs Any keyword accepted by `~plasmapy.plasma.plasma_base.GenericPlasma`. """ def __init__(self, hdf5, **kwargs) -> None: super().__init__(**kwargs) if not Path(hdf5).is_file(): raise FileNotFoundError(f"Could not find file: '{hdf5}'") h5 = h5py.File(hdf5, "r") self.h5 = h5 self._check_valid_openpmd_version() self.subname = next(iter(self.h5["data"])) def __enter__(self): return self.h5
[docs] def close(self) -> None: self.h5.close()
def __exit__( self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, ): self.h5.close() def _check_valid_openpmd_version(self) -> bool: try: openPMD_version = self.h5.attrs["openPMD"].decode("utf-8") if _valid_version(openPMD_version): return True else: raise DataStandardError( f"We currently only support HDF5 versions" f"starting from v{_OUTDATED_VERSION} and " f"lower than v{_NEWER_VERSION}. You can " f"however convert your HDF5 to a supported " f"version. For more information; see " f"https://github.com/openPMD/openPMD-updater" ) except KeyError as ex: raise DataStandardError( "Input HDF5 file does not go on with standards defined by OpenPMD" ) from ex @property def electric_field(self): """ An :math:`(x, y, z)` array containing electric field data. Returns ------- `~astropy.units.Quantity` """ path = f"data/{self.subname}/fields/E" if path in self.h5: units = _fetch_units(self.h5[path].attrs["unitDimension"]) axes = [self.h5[path][axis] for axis in self.h5[path]] return np.array(axes) * units else: raise AttributeError("No electric field data available in HDF5 file") @property def charge_density(self): """ An array containing charge density data. Returns ------- `~astropy.units.Quantity` """ path = f"data/{self.subname}/fields/rho" if path in self.h5: units = _fetch_units(self.h5[path].attrs["unitDimension"]) return np.array(self.h5[path]) * units else: raise AttributeError("No charge density data available in HDF5 file") @property def magnetic_field(self): """ An array containing magnetic field data. Returns ------- `~astropy.units.Quantity` """ path = f"data/{self.subname}/fields/B" if path in self.h5: units = _fetch_units(self.h5[path].attrs["unitDimension"]) axes = [self.h5[path][axis] for axis in self.h5[path]] return np.array(axes) * units else: raise AttributeError("No magnetic field data available in HDF5 file") @property def electric_current(self): """ An array containing electric current data. Returns ------- `~astropy.units.Quantity` """ path = f"data/{self.subname}/fields/J" if path in self.h5: units = _fetch_units(self.h5[path].attrs["unitDimension"]) axes = [self.h5[path][axis] for axis in self.h5[path]] return np.array(axes) * units else: raise AttributeError("No electric current data available in HDF5 file")
[docs] @classmethod def is_datasource_for(cls, **kwargs): if "hdf5" not in kwargs: return False hdf5 = kwargs.get("hdf5") openPMD = kwargs.get("openPMD") if not Path(hdf5).is_file(): raise FileNotFoundError(f"Could not find file: '{hdf5}'") if "openPMD" not in kwargs: h5 = h5py.File(hdf5, "r") try: openPMD = h5.attrs["openPMD"] except KeyError: openPMD = False return openPMD