Add conversion logic for core/edge -> plasma IDS

Author: maarten-ic

imas/__init__.py

@@ -15,6 +15,11 @@
 from .db_entry import DBEntry
 from .ids_factory import IDSFactory
 from .ids_convert import convert_ids
+from .convert_core_edge_plasma import (
+    convert_to_plasma_profiles,
+    convert_to_plasma_transport,
+    convert_to_plasma_sources,
+)
 from .ids_identifiers import identifiers
 
 # Load the IMAS-Python IMAS AL/DD core

imas/convert_core_edge_plasma.py

@@ -0,0 +1,80 @@
+# 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:
+    return _convert_to_plasma(core_or_edge_profiles, "profiles", deepcopy)
+
+
+def convert_to_plasma_sources(
+    core_or_edge_sources: IDSToplevel, *, deepcopy: bool = False
+) -> IDSToplevel:
+    return _convert_to_plasma(core_or_edge_sources, "sources", deepcopy)
+
+
+def convert_to_plasma_transport(
+    core_or_edge_transport: IDSToplevel, *, deepcopy: bool = False
+) -> IDSToplevel:
+    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."
+        )
+
+    # 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)