Merge pull request #62 from maarten-ic/feature/dd3to4-boundary-conversion

Equilibrium DD3to4 conversion for boundary_separatrix

Author: olivhoenen

docs/source/multi-dd.rst

@@ -166,6 +166,8 @@ explicit conversion mechanisms.
   Changed definition of ``space/coordinates_type`` in GGD grids, Yes, No
   Migrate obsolescent ``ids_properties/source`` to ``ids_properties/provenance``, Yes, No
   Convert the multiple time-bases in the ``pulse_schedule`` IDS [#ps3to4]_, Yes, No
+  Convert name + identifier -> description + name, Yes, Yes
+  Convert equilibrium ``boundary\_[secondary\_]separatrix`` to ``contour_tree`` [#contourtree]_, Yes, No
 
 .. [#rename] Quantities which have been renamed between the two DD versions. For
   example, the ``ec/beam`` Array of Structures in the ``pulse_schedule`` IDS,
@@ -205,6 +207,12 @@ explicit conversion mechanisms.
     interpolation otherwise. See also:
     https://github.com/iterorganization/IMAS-Python/issues/21.
 
+.. [#contourtree] Fills the `contour_tree
+    <https://imas-data-dictionary.readthedocs.io/en/latest/generated/ids/equilibrium.html#equilibrium-time_slice-contour_tree>`__
+    in the ``equilibrium`` IDS based on data in the ``boundary_separatrix`` and
+    ``boundary_secondary_separatrix`` structures from DD3. See also:
+    https://github.com/iterorganization/IMAS-Python/issues/60. 
+
 .. _`DD background`:
 
 Background information

imas/db_entry.py

@@ -608,7 +608,7 @@ def _get(
                 logger.warning(
                     "On-disk data is stored in DD %s which has a different major "
                     "version than the requested DD version (%s). IMAS-Python will "
-                    "convert the data automatically, but this does not cover all"
+                    "convert the data automatically, but this does not cover all "
                     "changes. "
                     "See %s/multi-dd.html#conversion-of-idss-between-dd-versions",
                     dd_version,

imas/ids_convert.py

@@ -7,7 +7,7 @@
 import logging
 from functools import lru_cache, partial
 from pathlib import Path
-from typing import Callable, Dict, Iterator, Optional, Set, Tuple
+from typing import Callable, Dict, Iterator, List, Optional, Set, Tuple
 from xml.etree.ElementTree import Element, ElementTree
 
 import numpy
@@ -87,6 +87,15 @@ def __init__(self) -> None:
         converted.
         """
 
+        self.post_process_ids: List[
+            Callable[[IDSToplevel, IDSToplevel, bool], None]
+        ] = []
+        """Postprocess functions to be applied to the whole IDS.
+
+        These postprocess functions should be applied to the whole IDS after all data is
+        converted. The arguments supplied are: source IDS, target IDS, deepcopy boolean.
+        """
+
         self.ignore_missing_paths: Set[str] = set()
         """Set of paths that should not be logged when data is present."""
 
@@ -362,6 +371,13 @@ def _apply_3to4_conversion(self, old: Element, new: Element) -> None:
             self.new_to_old.post_process[new_path] = _cocos_change
             self.old_to_new.post_process[old_path] = _cocos_change
 
+        # Convert equilibrium boundary_separatrix and populate contour_tree
+        if self.ids_name == "equilibrium":
+            self.old_to_new.post_process_ids.append(_equilibrium_boundary_3to4)
+            self.old_to_new.ignore_missing_paths |= {
+                "time_slice/boundary_separatrix",
+                "time_slice/boundary_secondary_separatrix",
+            }
         # Definition change for pf_active circuit/connections
         if self.ids_name == "pf_active":
             path = "circuit/connections"
@@ -603,6 +619,10 @@ def convert_ids(
     else:
         _copy_structure(toplevel, target, deepcopy, rename_map)
 
+    # Global post-processing functions
+    for callback in rename_map.post_process_ids:
+        callback(toplevel, target, deepcopy)
+
     logger.info("Conversion of IDS %s finished.", ids_name)
     if provenance_origin_uri:
         _add_provenance_entry(target, toplevel._version, provenance_origin_uri)
@@ -1229,3 +1249,70 @@ def _pulse_schedule_resample_callback(timebase, item: IDSBase, target_item: IDSB
             assume_sorted=True,
         )(timebase)
         target_item.value = value.astype(numpy.int32) if is_integer else value
+
+
+def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: bool):
+    """Convert DD3 boundary[[_secondary]_separatrix] to DD4 contour_tree"""
+    # Implement https://github.com/iterorganization/IMAS-Python/issues/60
+    copy = numpy.copy if deepcopy else lambda x: x
+    for ts3, ts4 in zip(eq3.time_slice, eq4.time_slice):
+        if not ts3.global_quantities.psi_axis.has_value:
+            # No magnetic axis, assume no boundary either:
+            continue
+        n_nodes = 1  # magnetic axis
+        if ts3.boundary_separatrix.psi.has_value:
+            n_nodes = 2
+            if (  # boundary_secondary_separatrix is introduced in DD 3.32.0
+                hasattr(ts3, "boundary_secondary_separatrix")
+                and ts3.boundary_secondary_separatrix.psi.has_value
+            ):
+                n_nodes = 3
+        node = ts4.contour_tree.node
+        node.resize(n_nodes)
+        # Magnetic axis (primary O-point)
+        gq = ts3.global_quantities
+        # Note the sign flip for psi due to the COCOS change between DD3 and DD4!
+        axis_is_psi_minimum = -gq.psi_axis < -gq.psi_boundary
+
+        node[0].critical_type = 0 if axis_is_psi_minimum else 2
+        node[0].r = gq.magnetic_axis.r
+        node[0].z = gq.magnetic_axis.z
+        node[0].psi = -gq.psi_axis  # COCOS change
+
+        # X-points
+        if n_nodes >= 2:
+            if ts3.boundary_separatrix.type == 0:  # limiter plasma
+                node[1].critical_type = 2 if axis_is_psi_minimum else 0
+                node[1].r = ts3.boundary_separatrix.active_limiter_point.r
+                node[1].z = ts3.boundary_separatrix.active_limiter_point.z
+            else:
+                node[1].critical_type = 1  # saddle-point (x-point)
+                if len(ts3.boundary_separatrix.x_point):
+                    node[1].r = ts3.boundary_separatrix.x_point[0].r
+                    node[1].z = ts3.boundary_separatrix.x_point[0].z
+                # Additional x-points. N.B. levelset is only stored on the first node
+                for i in range(1, len(ts3.boundary_separatrix.x_point)):
+                    node.resize(len(node) + 1, keep=True)
+                    node[-1].critical_type = 1
+                    node[-1].r = ts3.boundary_separatrix.x_point[i].r
+                    node[-1].z = ts3.boundary_separatrix.x_point[i].z
+                    node[-1].psi = -ts3.boundary_separatrix.psi
+            node[1].psi = -ts3.boundary_separatrix.psi  # COCOS change
+            node[1].levelset.r = copy(ts3.boundary_separatrix.outline.r)
+            node[1].levelset.z = copy(ts3.boundary_separatrix.outline.z)
+
+        if n_nodes >= 3:
+            node[2].critical_type = 1  # saddle-point (x-point)
+            if len(ts3.boundary_secondary_separatrix.x_point):
+                node[2].r = ts3.boundary_secondary_separatrix.x_point[0].r
+                node[2].z = ts3.boundary_secondary_separatrix.x_point[0].z
+                # Additional x-points. N.B. levelset is only stored on the first node
+                for i in range(1, len(ts3.boundary_secondary_separatrix.x_point)):
+                    node.resize(len(node) + 1, keep=True)
+                    node[-1].critical_type = 1
+                    node[-1].r = ts3.boundary_secondary_separatrix.x_point[i].r
+                    node[-1].z = ts3.boundary_secondary_separatrix.x_point[i].z
+                    node[-1].psi = -ts3.boundary_secondary_separatrix.psi
+            node[2].psi = -ts3.boundary_secondary_separatrix.psi  # COCOS change
+            node[2].levelset.r = copy(ts3.boundary_secondary_separatrix.outline.r)
+            node[2].levelset.z = copy(ts3.boundary_secondary_separatrix.outline.z)

imas/test/test_ids_convert.py

@@ -592,3 +592,107 @@ def test_3to4_cocos_hardcoded_paths():
 
     eq4 = convert_ids(eq, "4.0.0")
     assert eq4.time_slice[0].boundary.psi == -3.141
+
+
+def test_3to4_equilibrium_boundary():
+    eq342 = IDSFactory("3.42.0").equilibrium()
+    eq342.time_slice.resize(5)
+
+    for i, ts in enumerate(eq342.time_slice):
+        # Always fill boundary and magnetic axis
+        ts.boundary.psi = 1
+        ts.boundary.outline.r = [1.0, 3.0, 2.0, 1.0]
+        ts.boundary.outline.z = [1.0, 2.0, 3.0, 1.0]
+        ts.global_quantities.psi_axis = 1.0
+        ts.global_quantities.magnetic_axis.r = 2.0
+        ts.global_quantities.magnetic_axis.z = 2.0
+
+        if i > 0:
+            # Fill separatrix
+            ts.boundary_separatrix.psi = -1.0
+            # Use limiter for time_slice[1], otherwise divertor:
+            if i == 1:
+                ts.boundary_separatrix.type = 0
+                ts.boundary_separatrix.active_limiter_point.r = 3.0
+                ts.boundary_separatrix.active_limiter_point.z = 2.0
+            else:
+                ts.boundary_separatrix.type = 1
+            ts.boundary_separatrix.outline.r = [1.0, 3.0, 2.0, 1.0]
+            ts.boundary_separatrix.outline.z = [1.0, 2.0, 3.0, 1.0]
+            ts.boundary_separatrix.x_point.resize(1)
+            ts.boundary_separatrix.x_point[0].r = 1.0
+            ts.boundary_separatrix.x_point[0].z = 1.0
+            # These are not part of the conversion:
+            ts.boundary_separatrix.strike_point.resize(2)
+            ts.boundary_separatrix.closest_wall_point.r = 1.0
+            ts.boundary_separatrix.closest_wall_point.z = 1.0
+            ts.boundary_separatrix.closest_wall_point.distance = 0.2
+            ts.boundary_separatrix.dr_dz_zero_point.r = 3.0
+            ts.boundary_separatrix.dr_dz_zero_point.z = 2.0
+            ts.boundary_separatrix.gap.resize(1)
+        if i == 3:
+            # Fill second_separatrix
+            ts.boundary_secondary_separatrix.psi = -1.1
+            # Use limiter for time_slice[1], otherwise divertor:
+            ts.boundary_secondary_separatrix.outline.r = [0.9, 3.1, 2.1, 0.9]
+            ts.boundary_secondary_separatrix.outline.z = [0.9, 2.1, 3.1, 0.9]
+            ts.boundary_secondary_separatrix.x_point.resize(1)
+            ts.boundary_secondary_separatrix.x_point[0].r = 2.1
+            ts.boundary_secondary_separatrix.x_point[0].z = 3.1
+            # These are not part of the conversion:
+            ts.boundary_secondary_separatrix.distance_inner_outer = 0.1
+            ts.boundary_secondary_separatrix.strike_point.resize(2)
+        if i == 4:
+            ts.boundary_separatrix.x_point.resize(2, keep=True)
+            ts.boundary_separatrix.x_point[1].r = 2.0
+            ts.boundary_separatrix.x_point[1].z = 3.0
+
+    eq4 = convert_ids(eq342, "4.0.0")
+    assert len(eq4.time_slice) == 5
+    for i, ts in enumerate(eq4.time_slice):
+        node = ts.contour_tree.node
+        assert len(node) == [1, 2, 2, 3, 3][i]
+        # Test magnetic axis
+        assert node[0].critical_type == 0
+        assert node[0].r == node[0].z == 2.0
+        assert node[0].psi == -1.0
+        assert len(node[0].levelset.r) == len(node[0].levelset.z) == 0
+        # boundary_separatrix
+        if i == 1:  # node[1] is boundary for limiter plasma
+            assert node[1].critical_type == 2
+            assert node[1].r == 3.0
+            assert node[1].z == 2.0
+        elif i > 1:  # node[1] is boundary for divertor plasma
+            assert node[1].critical_type == 1
+            assert node[1].r == node[1].z == 1.0
+        if i > 0:
+            assert node[1].psi == 1.0
+            assert numpy.array_equal(node[1].levelset.r, [1.0, 3.0, 2.0, 1.0])
+            assert numpy.array_equal(node[1].levelset.z, [1.0, 2.0, 3.0, 1.0])
+        # boundary_secondary_separatrix
+        if i == 3:
+            assert node[2].critical_type == 1
+            assert node[2].r == 2.1
+            assert node[2].z == 3.1
+            assert node[2].psi == 1.1
+            assert numpy.array_equal(node[2].levelset.r, [0.9, 3.1, 2.1, 0.9])
+            assert numpy.array_equal(node[2].levelset.z, [0.9, 2.1, 3.1, 0.9])
+        # Second x-point from boundary_separatrix
+        if i == 4:
+            assert node[2].critical_type == 1
+            assert node[2].r == 2.0
+            assert node[2].z == 3.0
+            assert node[2].psi == node[1].psi == 1.0
+            # Levelset is only filled for the main x-point (node[1])
+            assert not node[2].levelset.r.has_value
+            assert not node[2].levelset.z.has_value
+
+    # not deepcopied, should share numpy arrays
+    slice1_outline_r = eq342.time_slice[1].boundary_separatrix.outline.r.value
+    assert slice1_outline_r is eq4.time_slice[1].contour_tree.node[1].levelset.r.value
+
+    # deepcopy should create a copy of the numpy arrays
+    eq4_cp = convert_ids(eq342, "4.0.0", deepcopy=True)
+    assert not numpy.may_share_memory(
+        slice1_outline_r, eq4_cp.time_slice[1].contour_tree.node[1].levelset.r.value
+    )