Equilibrium DD3to4 conversion for boundary_separatrix

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."""
 
@@ -343,6 +352,13 @@ def _apply_3to4_conversion(self, old: Element, new: Element) -> None:
                 new_path = self.old_to_new.path.get(old_path, old_path)
                 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"
@@ -544,6 +560,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)
@@ -1063,3 +1083,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

@@ -529,3 +529,107 @@ def test_3to4_migrate_deprecated_fields():  # GH#55
     del cp342.profiles_1d[0].ion[0].label
     cp4 = convert_ids(cp342, "4.0.0")
     assert cp4.profiles_1d[0].ion[0].name == "y"
+
+
+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
+    )