diff --git a/docs/source/api.rst b/docs/source/api.rst index 5df6e57..63e8af4 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -7,16 +7,24 @@ This page provides an auto-generated summary of IMAS-Python's API. For more deta and examples, refer to the relevant chapters in the main part of the documentation. -IMAS-Python IDS manipulation ----------------------------- +IMAS-Python public API +---------------------- .. currentmodule:: imas .. autosummary:: + convert_core_edge_plasma.convert_to_plasma_profiles + convert_core_edge_plasma.convert_to_plasma_sources + convert_core_edge_plasma.convert_to_plasma_transport db_entry.DBEntry + ids_convert.convert_ids + ids_data_type.IDSDataType ids_factory.IDSFactory - ids_toplevel.IDSToplevel + ids_identifiers.identifiers + ids_metadata.IDSMetadata + ids_metadata.IDSType ids_primitive.IDSPrimitive - ids_structure.IDSStructure ids_struct_array.IDSStructArray + ids_structure.IDSStructure + ids_toplevel.IDSToplevel \ No newline at end of file diff --git a/imas/__init__.py b/imas/__init__.py index 58a6699..4154b9f 100644 --- a/imas/__init__.py +++ b/imas/__init__.py @@ -1,30 +1,58 @@ # This file is part of IMAS-Python. # You should have received the IMAS-Python LICENSE file with this project. -# isort: skip_file - from packaging.version import Version as _V -from ._version import version as __version__ # noqa: F401 -from ._version import version_tuple # noqa: F401 - # Import logging _first_ -from . import setup_logging +# isort: off +from . import setup_logging # noqa: F401 + +# isort: on -# Import main user API objects in the imas module +# Ensure that `imas.util` is loaded when importing imas +from . import util # noqa: F401 + +# Public API: +from ._version import version as __version__ +from ._version import version_tuple +from .convert_core_edge_plasma import ( + convert_to_plasma_profiles, + convert_to_plasma_sources, + convert_to_plasma_transport, +) from .db_entry import DBEntry -from .ids_factory import IDSFactory from .ids_convert import convert_ids +from .ids_data_type import IDSDataType +from .ids_factory import IDSFactory from .ids_identifiers import identifiers - -# Load the IMAS-Python IMAS AL/DD core -from . import ( - db_entry, - dd_zip, - util, -) +from .ids_metadata import IDSMetadata, IDSType +from .ids_primitive import IDSPrimitive +from .ids_struct_array import IDSStructArray +from .ids_structure import IDSStructure +from .ids_toplevel import IDSToplevel PUBLISHED_DOCUMENTATION_ROOT = "https://imas-python.readthedocs.io/en/latest/" """URL to the published documentation.""" OLDEST_SUPPORTED_VERSION = _V("3.22.0") """Oldest Data Dictionary version that is supported by IMAS-Python.""" + +__all__ = [ + "__version__", + "version_tuple", + "DBEntry", + "IDSDataType", + "IDSFactory", + "IDSMetadata", + "IDSPrimitive", + "IDSStructure", + "IDSStructArray", + "IDSToplevel", + "IDSType", + "convert_ids", + "convert_to_plasma_profiles", + "convert_to_plasma_sources", + "convert_to_plasma_transport", + "identifiers", + "PUBLISHED_DOCUMENTATION_ROOT", + "OLDEST_SUPPORTED_VERSION", +] diff --git a/imas/convert_core_edge_plasma.py b/imas/convert_core_edge_plasma.py new file mode 100644 index 0000000..5a13c5a --- /dev/null +++ b/imas/convert_core_edge_plasma.py @@ -0,0 +1,124 @@ +# This file is part of IMAS-Python. +# You should have received the IMAS-Python LICENSE file with this project. +"""Logic to convert core/edge IDSs to their corresponding plasma ID.""" + +from packaging.version import Version + +from imas.ids_toplevel import IDSToplevel +from imas.ids_factory import IDSFactory +from imas.exception import IDSNameError +from imas.ids_convert import DDVersionMap, NBCPathMap, _copy_structure + + +def convert_to_plasma_profiles( + core_or_edge_profiles: IDSToplevel, *, deepcopy: bool = False +) -> IDSToplevel: + """Convert a core_profiles or edge_profiles IDS to a plasma_profiles IDS. + + The input IDS must use a Data Dictionary version for which the plasma_profiles IDS + exists (3.42.0 or newer). + + Args: + core_or_edge_profiles: The core_profiles or edge_profiles IDS to be converted. + + Keyword Args: + deepcopy: When True, performs a deep copy of all data. When False (default), + numpy arrays are not copied and the converted IDS shares the same underlying + data buffers. + """ + return _convert_to_plasma(core_or_edge_profiles, "profiles", deepcopy) + + +def convert_to_plasma_sources( + core_or_edge_sources: IDSToplevel, *, deepcopy: bool = False +) -> IDSToplevel: + """Convert a core_sources or edge_sources IDS to a plasma_sources IDS. + + The input IDS must use a Data Dictionary version for which the plasma_sources IDS + exists (3.42.0 or newer). + + Args: + core_or_edge_sources: The core_sources or edge_sources IDS to be converted. + + Keyword Args: + deepcopy: When True, performs a deep copy of all data. When False (default), + numpy arrays are not copied and the converted IDS shares the same underlying + data buffers. + """ + return _convert_to_plasma(core_or_edge_sources, "sources", deepcopy) + + +def convert_to_plasma_transport( + core_or_edge_transport: IDSToplevel, *, deepcopy: bool = False +) -> IDSToplevel: + """Convert a core_transport or edge_transport IDS to a plasma_transport IDS. + + The input IDS must use a Data Dictionary version for which the plasma_transport IDS + exists (3.42.0 or newer). + + Args: + core_or_edge_transport: The core_transport or edge_transport IDS to be + converted. + + Keyword Args: + deepcopy: When True, performs a deep copy of all data. When False (default), + numpy arrays are not copied and the converted IDS shares the same underlying + data buffers. + """ + return _convert_to_plasma(core_or_edge_transport, "transport", deepcopy) + + +class _CoreEdgePlasmaMap(DDVersionMap): + """Subclass of DDVersionMap to generate an NBCPathMap that is suitable to copy + between a core/edge IDS and the corresponding plasma IDS.""" + + def __init__(self, source, target, factory): + self.ids_name = source + self.old_version = factory._etree + self.new_version = factory._etree + self.version_old = Version(factory.version) + + self.old_to_new = NBCPathMap() + self.new_to_old = NBCPathMap() + + old_ids_object = factory._etree.find(f"IDS[@name='{source}']") + new_ids_object = factory._etree.find(f"IDS[@name='{target}']") + self._build_map(old_ids_object, new_ids_object) + + +def _convert_to_plasma(source: IDSToplevel, suffix: str, deepcopy: bool) -> IDSToplevel: + # Sanity checks for input data + if not isinstance(source, IDSToplevel): + raise TypeError( + f"First argument to convert_to_plasma_{suffix} must be a core_{suffix} or " + f"edge_{suffix} of type IDSToplevel. Got a type {type(source)} instead." + ) + if source.metadata.name not in [f"core_{suffix}", f"edge_{suffix}"]: + raise ValueError( + f"First argument to convert_to_plasma_{suffix} must be a core_{suffix} or " + f"edge_{suffix} IDS. Got a {source.metadata.name} IDS instead." + ) + if source._lazy: + raise NotImplementedError( + "IDS conversion is not implemented for lazy-loaded IDSs" + ) + + # Construct target plasma_{suffix} IDS + factory: IDSFactory = source._parent + try: + target = factory.new(f"plasma_{suffix}") + except IDSNameError: + raise ValueError( + f"Cannot convert {source.metadata.name} IDS to plasma_{suffix}: the source " + f"IDS uses Data Dictionary version {factory.dd_version} which doesn't have " + f"a plasma_{suffix} IDS. Please convert the source IDS to a supported Data " + "Dictionary version using `imas.convert_ids` and try again." + ) from None + + # Leverage existing logic from ids_convert to do the copying + # First construct a map (to handle missing items in the target IDS) + data_map = _CoreEdgePlasmaMap(source.metadata.name, target.metadata.name, factory) + path_map = data_map.old_to_new # old = core/edge, new = plasma IDS + _copy_structure(source, target, deepcopy, path_map) + + return target diff --git a/imas/test/test_convert_core_edge_plasma.py b/imas/test/test_convert_core_edge_plasma.py new file mode 100644 index 0000000..08d8ca9 --- /dev/null +++ b/imas/test/test_convert_core_edge_plasma.py @@ -0,0 +1,68 @@ +import pytest + +import imas.training +from imas.util import idsdiffgen +from imas.test.test_helpers import fill_with_random_data + + +def assert_equal(core_edge, plasma): + # We only expect the IDS name to be different: + difflist = list(idsdiffgen(core_edge, plasma)) + assert difflist == [("IDS name", core_edge.metadata.name, plasma.metadata.name)] + + +def test_convert_training_core_profiles(): + with imas.training.get_training_db_entry() as entry: + cp = entry.get("core_profiles") + + pp = imas.convert_to_plasma_profiles(cp) + assert_equal(cp, pp) + + +def test_convert_missing_qty(): + cp = imas.IDSFactory("4.1.0").core_profiles() + cp.profiles_1d.resize(1) + cp.profiles_1d[0].ion.resize(1) + cp.profiles_1d[0].ion[0].state.resize(1) + cp.profiles_1d[0].ion[0].state[0].ionization_potential = 0.5 + + pp = imas.convert_to_plasma_profiles(cp) + # check that state[0] is copied, but that it's empty + assert not pp.profiles_1d[0].ion[0].state[0].has_value + + +@pytest.mark.parametrize("idsname", ["core_profiles", "edge_profiles"]) +def test_convert_randomly_filled_profiles(idsname): + ids = imas.IDSFactory("4.1.0").new(idsname) + fill_with_random_data(ids) + + if idsname == "core_profiles": + # ionization_potential doesn't exist in plasma_profiles in DD 4.1.0. This case + # is tested in test_convert_missing_qty. Unset these variables to avoid a diff: + for profiles in list(ids.profiles_1d) + list(ids.profiles_2d): + for ion in profiles.ion: + for state in ion.state: + del state.ionization_potential + del state.ionization_potential_error_upper + del state.ionization_potential_error_lower + + plasma = imas.convert_to_plasma_profiles(ids) + assert_equal(ids, plasma) + + +@pytest.mark.parametrize("idsname", ["core_sources", "edge_sources"]) +def test_convert_randomly_filled_sources(idsname): + ids = imas.IDSFactory("4.1.0").new(idsname) + fill_with_random_data(ids) + + plasma = imas.convert_to_plasma_sources(ids) + assert_equal(ids, plasma) + + +@pytest.mark.parametrize("idsname", ["core_transport", "edge_transport"]) +def test_convert_randomly_filled_transport(idsname): + ids = imas.IDSFactory("4.1.0").new(idsname) + fill_with_random_data(ids) + + plasma = imas.convert_to_plasma_transport(ids) + assert_equal(ids, plasma) diff --git a/setup.cfg b/setup.cfg index 8e5dd29..fe8ea37 100644 --- a/setup.cfg +++ b/setup.cfg @@ -11,9 +11,6 @@ exclude= docs max-line-length = 88 per-file-ignores= - # Ignore import errors in __init__.py (import not at top of file; imported but - # unused) - imas/__init__.py:E402,F401 # Lots of CLASSPATHS in this test file: adhering to line length would be less # readable imas/test/test_dd_helpers.py:E501