Convert core/edge to plasma IDS

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

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",
+]

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

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)

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