From a09186a772899fc013401e3fe3b53f4b8bca7d1d Mon Sep 17 00:00:00 2001 From: Maarten Sebregts Date: Mon, 25 Aug 2025 14:33:11 +0200 Subject: [PATCH 1/4] Initial implementation of equilibrium DD3to4 conversion for boundary_separatrix --- imas/ids_convert.py | 70 +++++++++++++++++++++++++++- imas/test/test_ids_convert.py | 88 +++++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 1 deletion(-) diff --git a/imas/ids_convert.py b/imas/ids_convert.py index 75359f8..9fa52ac 100644 --- a/imas/ids_convert.py +++ b/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,51 @@ 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): + 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 + ts4.contour_tree.node.resize(n_nodes) + # Magnetic axis (primary O-point) + node = ts4.contour_tree.node + node[0].critical_type = 0 # minimum (?) + node[0].r = ts3.global_quantities.magnetic_axis.r + node[0].z = ts3.global_quantities.magnetic_axis.z + node[0].psi = -ts3.global_quantities.psi_axis # COCOS change + + # X-points + if n_nodes >= 2: + if ts3.boundary_separatrix.type == 0: # limiter plasma + node[1].critical_type = 2 # maximum (?) + 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 + # TODO: what if there are multiple x-points? + 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 + # TODO: what if there are multiple x-points? + 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) diff --git a/imas/test/test_ids_convert.py b/imas/test/test_ids_convert.py index 826a797..f5e7806 100644 --- a/imas/test/test_ids_convert.py +++ b/imas/test/test_ids_convert.py @@ -529,3 +529,91 @@ 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.0 + # 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, 2][i] + # Test magnetic axis + assert node[0].critical_type == 0 + assert node[0].r == node[0].z == 2.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 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 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]) + + # not deepcopied, should share numpy arrays + assert ( + eq342.time_slice[1].boundary_separatrix.outline.r.value + is eq4.time_slice[1].contour_tree.node[1].levelset.r.value + ) From d3f37939c655f9a96c9154c5a501138d4e2d16f1 Mon Sep 17 00:00:00 2001 From: Maarten Sebregts Date: Thu, 28 Aug 2025 14:17:18 +0200 Subject: [PATCH 2/4] Update conversion logic to fill equilibrium contour_tree - Handle multiple x-points in boundary_[secondary_]separatrix - Determine if O-point at magnetic axis is a local minimum or maximum of the psi map --- imas/ids_convert.py | 28 ++++++++++++++++++++++------ imas/test/test_ids_convert.py | 26 +++++++++++++++++++++----- 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/imas/ids_convert.py b/imas/ids_convert.py index 9fa52ac..4d62246 100644 --- a/imas/ids_convert.py +++ b/imas/ids_convert.py @@ -1098,10 +1098,14 @@ def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: boo and ts3.boundary_secondary_separatrix.psi.has_value ): n_nodes = 3 - ts4.contour_tree.node.resize(n_nodes) - # Magnetic axis (primary O-point) node = ts4.contour_tree.node - node[0].critical_type = 0 # minimum (?) + node.resize(n_nodes) + # Magnetic axis (primary O-point) + axis_is_psi_minimum = ( + # Note the sign flip for psi due to the COCOS change between DD3 and DD4! + -ts3.global_quantities.psi_axis < -ts3.global_quantities.psi_boundary + ) + node[0].critical_type = 0 if axis_is_psi_minimum else 2 node[0].r = ts3.global_quantities.magnetic_axis.r node[0].z = ts3.global_quantities.magnetic_axis.z node[0].psi = -ts3.global_quantities.psi_axis # COCOS change @@ -1109,7 +1113,7 @@ def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: boo # X-points if n_nodes >= 2: if ts3.boundary_separatrix.type == 0: # limiter plasma - node[1].critical_type = 2 # maximum (?) + 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: @@ -1117,7 +1121,13 @@ def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: boo 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 - # TODO: what if there are multiple x-points? + # 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) @@ -1127,7 +1137,13 @@ def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: boo 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 - # TODO: what if there are multiple x-points? + # 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) diff --git a/imas/test/test_ids_convert.py b/imas/test/test_ids_convert.py index f5e7806..f51d4ba 100644 --- a/imas/test/test_ids_convert.py +++ b/imas/test/test_ids_convert.py @@ -569,7 +569,7 @@ def test_3to4_equilibrium_boundary(): ts.boundary_separatrix.gap.resize(1) if i == 3: # Fill second_separatrix - ts.boundary_secondary_separatrix.psi = -1.0 + 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] @@ -588,10 +588,11 @@ def test_3to4_equilibrium_boundary(): 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, 2][i] + 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 @@ -602,6 +603,7 @@ def test_3to4_equilibrium_boundary(): 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 @@ -609,11 +611,25 @@ def test_3to4_equilibrium_boundary(): 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 - assert ( - eq342.time_slice[1].boundary_separatrix.outline.r.value - is eq4.time_slice[1].contour_tree.node[1].levelset.r.value + 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 ) From 83ff9451290d25410a4d43f50600368c2753593d Mon Sep 17 00:00:00 2001 From: Maarten Sebregts Date: Thu, 28 Aug 2025 14:34:41 +0200 Subject: [PATCH 3/4] Fix formatting --- imas/ids_convert.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/imas/ids_convert.py b/imas/ids_convert.py index 4d62246..77163a1 100644 --- a/imas/ids_convert.py +++ b/imas/ids_convert.py @@ -1101,14 +1101,14 @@ def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: boo node = ts4.contour_tree.node node.resize(n_nodes) # Magnetic axis (primary O-point) - axis_is_psi_minimum = ( - # Note the sign flip for psi due to the COCOS change between DD3 and DD4! - -ts3.global_quantities.psi_axis < -ts3.global_quantities.psi_boundary - ) + 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 = ts3.global_quantities.magnetic_axis.r - node[0].z = ts3.global_quantities.magnetic_axis.z - node[0].psi = -ts3.global_quantities.psi_axis # COCOS change + 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: From fb8476d9dfd946d6dd75e415358326b3a8395c7e Mon Sep 17 00:00:00 2001 From: Maarten Sebregts Date: Thu, 28 Aug 2025 14:57:22 +0200 Subject: [PATCH 4/4] Do not fill contour tree in conversion logic when there is no magnetic axis --- imas/ids_convert.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/imas/ids_convert.py b/imas/ids_convert.py index 77163a1..a400a31 100644 --- a/imas/ids_convert.py +++ b/imas/ids_convert.py @@ -1090,6 +1090,9 @@ def _equilibrium_boundary_3to4(eq3: IDSToplevel, eq4: IDSToplevel, deepcopy: boo # 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