From 98d772e5101e669470e18714f90ce36e3f7ab592 Mon Sep 17 00:00:00 2001 From: Sebbe Blokhuizen Date: Fri, 2 Oct 2026 15:18:41 +0200 Subject: [PATCH 1/3] add calculation and display of WEST clearance gaps --- waveform_editor/gui/shape_editor/metrics.py | 11 +- .../gui/shape_editor/shape_editor.py | 45 ++++++ waveform_editor/shape_editor/west_gaps.py | 131 ++++++++++++++++++ 3 files changed, 184 insertions(+), 3 deletions(-) create mode 100644 waveform_editor/shape_editor/west_gaps.py diff --git a/waveform_editor/gui/shape_editor/metrics.py b/waveform_editor/gui/shape_editor/metrics.py index 7e63cbc7..024ac8ad 100644 --- a/waveform_editor/gui/shape_editor/metrics.py +++ b/waveform_editor/gui/shape_editor/metrics.py @@ -10,6 +10,8 @@ class Metrics(Viewer): """Chips row showing equilibrium metrics below the flux map.""" metrics = param.Dict(default={}) + # Chips of the machine in use, as (symbol, unit, full name) by name + machine_metrics = param.Dict(default={}) ELONGATION = "elongation" TRIANGULARITY = "triangularity" @@ -35,15 +37,18 @@ class Metrics(Viewer): def __init__(self, **params): super().__init__(**params) self._pane = pn.pane.HTML( - pn.bind(self._render, self.param.metrics), + pn.bind(self._render, self.param.metrics, self.param.machine_metrics), sizing_mode="stretch_width", max_width=NicePlotter.FRAME_WIDTH, stylesheets=STYLES, ) - def _render(self, metrics=None): + def _render(self, metrics=None, machine_metrics=None): chips = [] - for key, (symbol, unit, tooltip) in self.METRICS.items(): + for key, (symbol, unit, tooltip) in { + **self.METRICS, + **self.machine_metrics, + }.items(): val = metrics.get(key, "—") if metrics else "—" if isinstance(val, float): val = f"{val:.4g}" diff --git a/waveform_editor/gui/shape_editor/shape_editor.py b/waveform_editor/gui/shape_editor/shape_editor.py index c43036a9..89dc95ef 100644 --- a/waveform_editor/gui/shape_editor/shape_editor.py +++ b/waveform_editor/gui/shape_editor/shape_editor.py @@ -21,9 +21,21 @@ from waveform_editor.gui.util import set_xml_parameter from waveform_editor.settings import NiceSettings, settings from waveform_editor.shape_editor.nice_integration import NiceIntegration +from waveform_editor.shape_editor.west_gaps import compute_gaps logger = logging.getLogger(__name__) +# TODO: add the gaps for ITER used in DINA: gap1, gap2, gap4, gap5, Rmin, Rmax +# The gaps of a WEST plasma, as (symbol, unit, full name) by name +WEST_GAP_METRICS = { + "UROG": ("UROG", "cm", "Upper radial outer gap"), + "EROG": ("EROG", "cm", "Equatorial radial outer gap"), + "LROG": ("LROG", "cm", "Lower radial outer gap"), + "dXlow": ("dXlow", "cm", "Distance of the lower x-point to the divertor"), + "dXup": ("dXup", "cm", "Distance of the upper x-point to the divertor"), + "dbaffle": ("dbaffle", "cm", "Distance of the plasma to the baffle"), +} + # NICE reads the desired boundary into an array of this fixed size # (MAX_PLASMA_BOUNDARY_POINTS in its solver_structs.h) @@ -171,6 +183,8 @@ def __init__(self, main_gui): ) self.metrics = Metrics() + self.nice_settings.param.watch(self._update_machine_metrics, "machine_preset") + self._update_machine_metrics() self._active_tab = 0 options = pn.bind( self._create_options_tabs, @@ -471,6 +485,36 @@ async def submit(self, event=None): ): self._add_to_history() + def _update_machine_metrics(self, event=None): + """Show the chips of the machine of the selected preset, also before a run + has filled them in.""" + self.metrics.machine_metrics = ( + WEST_GAP_METRICS + if self.nice_settings.machine_preset == NiceSettings.PRESET_WEST + else {} + ) + + def _west_gaps(self, time_slice): + """The gaps of the plasma to the parts of WEST it is kept away from. They are + not in the equilibrium, so they are computed from its boundary. + + Args: + time_slice: The time slice NICE returned. + + Returns: + Dict of gap name to distance in centimetres, empty for another machine. + """ + if self.nice_settings.machine_preset != NiceSettings.PRESET_WEST: + return {} + x_points = [ + (float(node.r), float(node.z)) + for node in time_slice.contour_tree.node + if int(node.critical_type) == 1 + ] + outline = time_slice.boundary.outline + gaps = compute_gaps(outline.r, outline.z, x_points) + return {name: gap * 100 for name, gap in gaps.items()} + def _update_metrics(self): eq = self.communicator.equilibrium global_quantities = eq.time_slice[0].global_quantities @@ -485,6 +529,7 @@ def _update_metrics(self): self.metrics.VERTICAL: float(boundary.geometric_axis.z), self.metrics.MINOR_RADIUS: float(boundary.minor_radius), self.metrics.Q95: float(global_quantities.q_95), + **self._west_gaps(eq.time_slice[0]), } @param.depends( diff --git a/waveform_editor/shape_editor/west_gaps.py b/waveform_editor/shape_editor/west_gaps.py new file mode 100644 index 00000000..23b1c7ac --- /dev/null +++ b/waveform_editor/shape_editor/west_gaps.py @@ -0,0 +1,131 @@ +"""The gaps between a WEST plasma and the parts of the machine it is kept away from. + +They are not stored in the machine description, so they are computed the way FEEQS +does, in Projects/WEST/Lib/plot_plasma_gaps_etc.m, from the geometry of WEST. +""" + +import numpy as np + +# The arc the outer radial gaps are measured to, which passes through r=3 m on the +# midplane, as (centre r, radius) +OUTER_ARC = (2.2, 0.8) +# The heights the upper and lower outer radial gaps are measured at +OUTER_GAP_HEIGHT = 0.25 +# The divertor targets, as the two points of the line through each of them +LOWER_DIVERTOR = ((1.909, -0.5796), (2.362, -0.7624)) +UPPER_DIVERTOR = ((1.9009, 0.5824), (2.446, 0.7995)) +# The corner of the baffle the plasma is kept away from +BAFFLE = (2.381, -0.6757) + + +def compute_gaps(outline_r, outline_z, x_points): + """The gaps of a WEST plasma, in metres. + + Args: + outline_r: Radial coordinates of the plasma boundary. + outline_z: Height coordinates of the plasma boundary. + x_points: The (r, z) of each x-point of the equilibrium. + + Returns: + Dict of gap name to distance, holding only the gaps that the boundary and the + x-points given allow to be computed. + """ + r, z = np.asarray(outline_r, dtype=float), np.asarray(outline_z, dtype=float) + gaps = {} + centre_r, radius = OUTER_ARC + for name, height in ( + ("UROG", OUTER_GAP_HEIGHT), + ("EROG", 0.0), + ("LROG", -OUTER_GAP_HEIGHT), + ): + arc_r = centre_r + np.sqrt(radius**2 - height**2) + boundary_r = _outboard_radius(r, z, height) + if boundary_r is not None: + gaps[name] = arc_r - boundary_r + + for name, divertor, below in ( + ("dXlow", LOWER_DIVERTOR, True), + ("dXup", UPPER_DIVERTOR, False), + ): + x_point = _x_point(x_points, below) + if x_point is not None: + gaps[name] = _distance_to_line(x_point, *divertor) + + gaps["dbaffle"] = _distance_to_outline(BAFFLE, r, z) + return gaps + + +def _outboard_radius(r, z, height): + """The radius of the outboard side of the boundary at a height. + + Args: + r: Radial coordinates of the boundary. + z: Height coordinates of the boundary. + height: The height to take the boundary at. + + Returns: + The radius, or None when the boundary does not reach that height. + """ + outboard = r > r.mean() + r, z = r[outboard], z[outboard] + if not len(z) or not z.min() <= height <= z.max(): + return None + order = np.argsort(z) + return float(np.interp(height, z[order], r[order])) + + +def _x_point(x_points, below): + """The x-point below or above the midplane. + + Args: + x_points: The (r, z) of each x-point. + below: Whether to take the x-point below the midplane. + + Returns: + The (r, z) of the x-point, or None when there is none on that side. + """ + on_side = [point for point in x_points if (point[1] < 0) == below] + return on_side[0] if on_side else None + + +def _distance_to_line(point, start, end): + """The distance of a point to the line through two points. + + Args: + point: The (r, z) to measure from. + start: The (r, z) of a point of the line. + end: The (r, z) of another point of the line. + + Returns: + The distance. + """ + point, start, end = (np.asarray(p, dtype=float) for p in (point, start, end)) + along = (end - start) / np.linalg.norm(end - start) + to_point = point - start + return float(np.linalg.norm(to_point - np.dot(to_point, along) * along)) + + +def _distance_to_outline(point, r, z): + """The distance of a point to the closest of the segments of an outline. + + Args: + point: The (r, z) to measure from. + r: Radial coordinates of the outline. + z: Height coordinates of the outline. + + Returns: + The distance. + """ + starts = np.column_stack([r, z]) + ends = np.roll(starts, -1, axis=0) + segments = ends - starts + lengths = np.sum(segments**2, axis=1) + # How far along each segment the point is, kept within the segment + along = np.clip( + np.sum((np.asarray(point, dtype=float) - starts) * segments, axis=1) + / np.where(lengths > 0, lengths, 1), + 0, + 1, + ) + closest = starts + along[:, None] * segments + return float(np.min(np.linalg.norm(closest - point, axis=1))) From f2e2e76130e778dc06937bc1f41a2a9507889cd8 Mon Sep 17 00:00:00 2001 From: Sebbe Blokhuizen Date: Mon, 5 Oct 2026 14:04:28 +0200 Subject: [PATCH 2/3] add additional gaps functionality --- tests/test_iter_gaps.py | 310 ++++++++++++++++++ .../gui/shape_editor/nice_plotter.py | 143 ++++++++ .../gui/shape_editor/plasma_shape.py | 209 ++++++++++-- .../gui/shape_editor/settings_modal.py | 1 + .../gui/shape_editor/shape_editor.py | 94 ++++-- waveform_editor/shape_editor/iter_gaps.py | 260 +++++++++++++++ .../shape_editor/plasma_shape_calc.py | 36 ++ waveform_editor/shape_editor/west_gaps.py | 177 ++++++---- 8 files changed, 1121 insertions(+), 109 deletions(-) create mode 100644 tests/test_iter_gaps.py create mode 100644 waveform_editor/shape_editor/iter_gaps.py diff --git a/tests/test_iter_gaps.py b/tests/test_iter_gaps.py new file mode 100644 index 00000000..4b4a6af0 --- /dev/null +++ b/tests/test_iter_gaps.py @@ -0,0 +1,310 @@ +import math + +import numpy as np + +from waveform_editor.gui.shape_editor.shape_editor import ITER_GAP_METRICS +from waveform_editor.shape_editor.iter_gaps import ( + GAP_POINTS, + _distance_to_outline, + compute_gaps, +) + + +def test_gap_points_defined(): + assert set(GAP_POINTS.keys()) == {"gap1", "gap2", "gap4", "gap5"} + assert GAP_POINTS["gap1"] == (4.2230, -3.7920) + assert GAP_POINTS["gap2"] == (5.5650, -4.4040) + assert GAP_POINTS["gap4"] == (7.5095, 2.9971) + assert GAP_POINTS["gap5"] == (5.3315, 4.5804) + + +def test_iter_gap_metrics_keys(): + expected_keys = {"gap1", "gap2", "gap4", "gap5", "Rmin", "Rmax"} + assert set(ITER_GAP_METRICS.keys()) == expected_keys + for key in ("gap1", "gap2", "gap4", "gap5"): + assert ITER_GAP_METRICS[key][1] == "cm" + for key in ("Rmin", "Rmax"): + assert ITER_GAP_METRICS[key][1] == "m" + + +def test_compute_gaps_empty(): + assert compute_gaps(None, None) == {} + assert compute_gaps([], []) == {} + assert compute_gaps(np.array([]), np.array([])) == {} + + +def test_distance_to_outline_projection(): + # Outline is a square: (0, 0) -> (2, 0) -> (2, 2) -> (0, 2) + r = [0.0, 2.0, 2.0, 0.0] + z = [0.0, 0.0, 2.0, 2.0] + + # Point directly above bottom segment: closest point is (1, 0), dist is 0.5 + assert math.isclose(_distance_to_outline((1.0, 0.5), r, z), 0.5) + + # Point outside corner (3, 3): closest point is (2, 2), dist is sqrt(2) + assert math.isclose(_distance_to_outline((3.0, 3.0), r, z), math.sqrt(2)) + + +def test_compute_gaps_circular_plasma(): + theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) + r0, z0, a = 6.2, 0.0, 2.0 + outline_r = r0 + a * np.cos(theta) + outline_z = z0 + a * np.sin(theta) + + gaps = compute_gaps(outline_r, outline_z) + assert set(gaps.keys()) == {"gap1", "gap2", "gap4", "gap5", "Rmin", "Rmax"} + + assert math.isclose(gaps["Rmin"], 4.2, abs_tol=1e-3) + assert math.isclose(gaps["Rmax"], 8.2, abs_tol=1e-3) + + # Circular plasma has no divertor legs (min Z = -2.0 > -3.6) + assert gaps["gap1"] is None + assert gaps["gap2"] is None + for gap_name in ("gap4", "gap5"): + assert isinstance(gaps[gap_name], float) + assert gaps[gap_name] > 0 + + +def test_compute_gaps_with_separatrix_contour(): + from waveform_editor.shape_editor.iter_gaps import compute_gap_geometry + + theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) + r0, z0, a = 6.2, 0.0, 2.0 + outline_r = r0 + a * np.cos(theta) + outline_z = z0 + a * np.sin(theta) + + # Synthetic divertor legs passing near gap1 (4.223, -3.792) and gap2 (5.565, -4.404) + inner_leg = np.array([[4.25, -3.75], [4.20, -3.85]]) + outer_leg = np.array([[5.55, -4.35], [5.58, -4.45]]) + contour = [inner_leg, outer_leg] + magnetic_axis = (6.2, 0.5) + + gaps = compute_gaps( + outline_r, outline_z, separatrix_contour=contour, magnetic_axis=magnetic_axis + ) + for name in ("gap1", "gap2", "gap4", "gap5", "Rmin", "Rmax"): + assert isinstance(gaps[name], float) + assert abs(gaps["gap1"]) < 0.1 # within 10 cm of inner leg + assert abs(gaps["gap2"]) < 0.1 # within 10 cm of outer leg + + items = compute_gap_geometry( + outline_r, outline_z, separatrix_contour=contour, magnetic_axis=magnetic_axis + ) + assert len(items) == 6 + symbols = {item["symbol"] for item in items} + assert symbols == {"g₁", "g₂", "g₄", "g₅", "Rₘᵢₙ", "Rₘₐₓ"} + + +def test_compute_gap_geometry_empty(): + from waveform_editor.shape_editor.iter_gaps import compute_gap_geometry + + assert compute_gap_geometry(None, None) == [] + assert compute_gap_geometry([], []) == [] + + +def test_compute_gap_geometry_iter(): + from waveform_editor.shape_editor.iter_gaps import compute_gap_geometry + + theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) + r0, z0, a = 6.2, 0.0, 2.0 + outline_r = r0 + a * np.cos(theta) + outline_z = z0 + a * np.sin(theta) + + # Without divertor legs: all 6 gaps are still returned + # so strike point locations are visible + items = compute_gap_geometry(outline_r, outline_z) + assert len(items) == 6 + symbols = {item["symbol"] for item in items} + assert symbols == {"g₁", "g₂", "g₄", "g₅", "Rₘᵢₙ", "Rₘₐₓ"} + + for item in items: + expected_dist = math.hypot( + item["r_target"] - item["r_orig"], + item["z_target"] - item["z_orig"], + ) + assert math.isclose(item["distance"], expected_dist, rel_tol=1e-5) + assert item["distance"] > 0 + assert "cm" in item["distance_str"] or "m" in item["distance_str"] + + +def test_compute_gap_geometry_west(): + from waveform_editor.shape_editor.west_gaps import compute_gap_geometry + + theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) + r0, z0, a = 2.5, 0.0, 0.4 + outline_r = r0 + a * np.cos(theta) + outline_z = z0 + a * np.sin(theta) + x_points = [(2.2, -0.6), (2.2, 0.6)] + + items = compute_gap_geometry(outline_r, outline_z, x_points) + assert len(items) == 6 + symbols = {item["symbol"] for item in items} + assert symbols == {"UROG", "EROG", "LROG", "dXlow", "dXup", "dbaffle"} + + for item in items: + assert item["distance"] > 0 + assert "cm" in item["distance_str"] + + +def test_nice_plotter_gap_visualization(): + import holoviews as hv + + hv.extension("bokeh") + import imas + + from waveform_editor.gui.shape_editor.nice_plotter import NicePlotter + from waveform_editor.gui.shape_editor.plasma_properties import PlasmaProperties + from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape + from waveform_editor.settings import NiceSettings, settings + from waveform_editor.shape_editor.nice_integration import NiceIntegration + + factory = imas.IDSFactory() + comm = NiceIntegration( + factory, on_output=lambda x: None, on_run_finished=lambda x: None + ) + shape = PlasmaShape() + props = PlasmaProperties() + plotter = NicePlotter( + communicator=comm, plasma_shape=shape, plasma_properties=props + ) + + assert hasattr(plotter, "show_gaps") + assert plotter.show_gaps is True + + # When no shape or equilibrium is loaded, overlay is empty + empty_overlay = plotter._plot_clearance_gaps() + assert isinstance(empty_overlay, hv.Overlay) + + # Set up a parameterized shape for ITER + settings.nice.machine_preset = NiceSettings.PRESET_ITER + shape.outline_r = [5.0, 7.0, 7.0, 5.0] + shape.outline_z = [-1.0, -1.0, 1.0, 1.0] + + overlay = plotter._plot_clearance_gaps() + assert isinstance(overlay, hv.Overlay) + # Check that overlay contains Points and Segments + types = [type(el) for el in overlay.values()] + assert hv.Points in types + assert hv.Segments in types + + # Toggle off show_gaps + plotter.show_gaps = False + off_overlay = plotter._plot_clearance_gaps() + assert isinstance(off_overlay, hv.Overlay) + + +def test_auto_enable_warm_start_on_direct_mode(): + import holoviews as hv + + hv.extension("bokeh") + from waveform_editor.configuration import WaveformConfiguration + from waveform_editor.gui.shape_editor.shape_editor import ShapeEditor + from waveform_editor.settings import NiceSettings + + class MockMainGui: + def __init__(self): + self.config = WaveformConfiguration() + + editor = ShapeEditor(MockMainGui()) + editor.nice_settings.mode = NiceSettings.INVERSE_MODE + editor.use_previous_run = False + editor.communicator.converged = False + + # 1. Switch to direct mode without any previous run -> remains False + editor.nice_settings.mode = NiceSettings.DIRECT_MODE + assert editor.use_previous_run is False + + # 2. Switch back to inverse mode + editor.nice_settings.mode = NiceSettings.INVERSE_MODE + assert editor.use_previous_run is False + + # 3. Previous run converges + editor.communicator.converged = True + + # 4. Switch to direct mode with converged run -> automatically True + editor.nice_settings.mode = NiceSettings.DIRECT_MODE + assert editor.use_previous_run is True + + +def test_target_gap_pills_in_plasma_shape(): + from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape + + shape = PlasmaShape() + shape.machine_preset = "ITER" + shape.show_gaps = True + + # No shape loaded yet -> empty string + assert shape._render_target_gap_pills() == "" + + # Load parameterized shape + shape.input_mode = shape.PARAMETERIZED_INPUT + shape._load_shape_from_params() + assert shape.has_shape is True + + # With show_gaps=True and ITER -> pills rendered + html = shape._render_target_gap_pills() + assert '
' in html + assert "g₄" in html + assert "g₅" in html + assert "Rₘᵢₙ" in html + assert "Rₘₐₓ" in html + + # With show_gaps=False -> no pills rendered + shape.show_gaps = False + assert shape._render_target_gap_pills() == "" + + # Switch to WEST -> WEST pills rendered + shape.show_gaps = True + shape.machine_preset = "WEST" + west_html = shape._render_target_gap_pills() + assert '
' in west_html + assert "UROG" in west_html + assert "EROG" in west_html + + +def test_gaps_input_mode_default_machine_gaps(): + from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape + from waveform_editor.shape_editor.iter_gaps import get_default_iter_gaps + from waveform_editor.shape_editor.west_gaps import get_default_west_gaps + + # Test default gaps generators directly + iter_gaps = get_default_iter_gaps() + assert len(iter_gaps) == 6 + assert iter_gaps[0].name == "Inner divertor leg (g₁)" + assert iter_gaps[1].name == "Outer divertor leg (g₂)" + assert iter_gaps[0].value == 0.0 + assert iter_gaps[1].value == 0.0 + + west_gaps = get_default_west_gaps() + assert len(west_gaps) == 6 + assert any("UROG" in g.name for g in west_gaps) + assert any("dXlow" in g.name for g in west_gaps) + + # Test PlasmaShape in Gaps mode + shape = PlasmaShape() + shape.machine_preset = "ITER" + shape.input_mode = shape.GAP_INPUT + + # Default gap source is Default Gaps + assert shape.gap_source == shape.GAP_SOURCE_DEFAULT + assert shape.has_shape is True + assert len(shape.gaps) == 6 + assert len(shape.gap_ui) == 6 + + # Test changing a gap slider value + orig_val = shape.gap_ui[0].value + shape.gap_ui[0].value = 0.05 + assert math.isclose(shape.gaps[0].value, 0.05) + + # Test reset defaults + shape._on_reset_default_gaps() + assert math.isclose(shape.gaps[0].value, orig_val) + + # Test switching to WEST preset + shape.machine_preset = "WEST" + assert len(shape.gaps) == 6 + assert any("UROG" in g.name for g in shape.gaps) + + # Test switching to IDS mode + shape.gap_source = shape.GAP_SOURCE_IDS + assert shape.gap_source == shape.GAP_SOURCE_IDS diff --git a/waveform_editor/gui/shape_editor/nice_plotter.py b/waveform_editor/gui/shape_editor/nice_plotter.py index 89e42389..bf0ab71c 100644 --- a/waveform_editor/gui/shape_editor/nice_plotter.py +++ b/waveform_editor/gui/shape_editor/nice_plotter.py @@ -14,7 +14,16 @@ from waveform_editor.gui.shape_editor.plasma_properties import PlasmaProperties from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape from waveform_editor.settings import NiceSettings, settings +from waveform_editor.shape_editor.iter_gaps import ( + compute_gap_geometry as compute_iter_gap_geometry, +) +from waveform_editor.shape_editor.iter_gaps import ( + extract_separatrix_from_time_slice, +) from waveform_editor.shape_editor.nice_integration import NiceIntegration +from waveform_editor.shape_editor.west_gaps import ( + compute_gap_geometry as compute_west_gap_geometry, +) matplotlib.use("Agg") logger = logging.getLogger(__name__) @@ -123,6 +132,7 @@ class NicePlotter(Viewer): show_xo = param.Boolean(default=True, label="Show x-point and o-point") show_separatrix = param.Boolean(default=True, label="Show separatrix") show_desired_shape = param.Boolean(default=True, label="Show desired shape") + show_gaps = param.Boolean(default=True, label="Show gaps") # Renderer of the editable points, set once the plot is first rendered _points_renderer = None @@ -211,6 +221,7 @@ def __init__(self, **params): hv.DynamicMap(self._plot_passive_structures), hv.DynamicMap(self._plot_iron_core), hv.DynamicMap(self._plot_plasma_shape), + hv.DynamicMap(self._plot_clearance_gaps), self.editable_points, ] # Lets the weighted points be added, dragged and deleted on the plot @@ -736,5 +747,137 @@ def _plot_xo_points(self): ) return o_scatter * x_scatter + @pn.depends( + "communicator.equilibrium", + "show_gaps", + "nice_settings.machine_preset", + "plasma_shape.shape_updated", + "plasma_shape.input_mode", + ) + def _plot_clearance_gaps(self): + """Plots the clearance gaps for the active machine preset, showing both the + measurement reference points on the wall and the distance lines to the + plasma boundary. + + Returns: + Holoviews overlay with points, target intercepts, dashed distance + segments, and text labels. + """ + if not self.show_gaps: + return self._empty_gaps_overlay() + + eq = self.communicator.equilibrium + x_points = [] + if eq is not None and len(eq.time_slice) > 0: + outline = eq.time_slice[0].boundary.outline + r, z = outline.r, outline.z + x_points = [ + (float(node.r), float(node.z)) + for node in eq.time_slice[0].contour_tree.node + if int(node.critical_type) == 1 + ] + elif not self.nice_settings.is_direct_mode and self.plasma_shape.has_shape: + r, z = self.plasma_shape.outline_r, self.plasma_shape.outline_z + else: + return self._empty_gaps_overlay() + + if r is None or len(r) == 0: + return self._empty_gaps_overlay() + + gap_items = [] + preset = self.nice_settings.machine_preset + if preset == NiceSettings.PRESET_ITER: + ts = eq.time_slice[0] if eq is not None and len(eq.time_slice) > 0 else None + ma = ( + getattr(getattr(ts, "global_quantities", None), "magnetic_axis", None) + if ts + else None + ) + gap_items = compute_iter_gap_geometry( + r, + z, + separatrix_contour=extract_separatrix_from_time_slice(ts) + if ts + else None, + magnetic_axis=(float(ma.r), float(ma.z)) if ma else None, + ) + elif preset == NiceSettings.PRESET_WEST: + gap_items = compute_west_gap_geometry(r, z, x_points) + + if not gap_items: + return self._empty_gaps_overlay() + + return self._build_gaps_overlay(gap_items) + + def _build_gaps_overlay(self, items): + vdims = ["name", "symbol", "distance_str"] + pts = hv.Points( + [{"r": it["r_orig"], "z": it["z_orig"], **it} for it in items], + kdims=["r", "z"], + vdims=vdims, + ).opts( + color="#e66101", + size=8, + marker="diamond", + hover_tooltips=[ + ("Gap", "@name (@symbol)"), + ("Location (R, Z)", "@r{0.000} m, @z{0.000} m"), + ("Distance", "@distance_str"), + ], + show_legend=False, + ) + tgts = hv.Points( + [{"r": it["r_target"], "z": it["z_target"], **it} for it in items], + kdims=["r", "z"], + vdims=vdims, + ).opts( + color="#e66101", + size=5, + marker="circle", + hover_tooltips=[("Gap", "@name (@symbol)"), ("Distance", "@distance_str")], + show_legend=False, + ) + segs = hv.Segments( + [ + { + "r0": it["r_orig"], + "z0": it["z_orig"], + "r1": it["r_target"], + "z1": it["z_target"], + **it, + } + for it in items + ], + kdims=["r0", "z0", "r1", "z1"], + vdims=vdims, + ).opts( + color="#e66101", + line_dash="dashed", + line_width=2, + hover_tooltips=[("Gap", "@name (@symbol)"), ("Distance", "@distance_str")], + show_legend=False, + ) + lbls = _no_hover( + hv.Labels( + [ + {"r": it["r_orig"], "z": it["z_orig"], "text": it["symbol"]} + for it in items + ], + kdims=["r", "z"], + vdims=["text"], + ).opts( + text_font_size="10pt", + text_color="#e66101", + text_baseline="bottom", + text_align="center", + text_font_style="bold", + ) + ) + return pts * tgts * segs * lbls + + def _empty_gaps_overlay(self): + """Empty overlay for when gaps are disabled or cannot be computed.""" + return self._build_gaps_overlay([]) + def __panel__(self): return self.panel_layout diff --git a/waveform_editor/gui/shape_editor/plasma_shape.py b/waveform_editor/gui/shape_editor/plasma_shape.py index 3b5c9fbc..199dce50 100644 --- a/waveform_editor/gui/shape_editor/plasma_shape.py +++ b/waveform_editor/gui/shape_editor/plasma_shape.py @@ -6,12 +6,22 @@ from waveform_editor.gui.util import ( CARD_CSS, + STYLES, EquilibriumInput, FixedWidthEditableIntSlider, FormattedEditableFloatSlider, WarningIndicator, ) from waveform_editor.settings import NiceSettings, settings +from waveform_editor.shape_editor.iter_gaps import ( + GAP_METADATA as ITER_GAP_METADATA, +) +from waveform_editor.shape_editor.iter_gaps import ( + compute_gaps as compute_iter_gaps, +) +from waveform_editor.shape_editor.iter_gaps import ( + get_default_iter_gaps, +) from waveform_editor.shape_editor.plasma_shape_calc import ( Gap, apply_point_weights, @@ -19,6 +29,13 @@ compute_outline_from_params, update_outline_from_gaps, ) +from waveform_editor.shape_editor.west_gaps import ( + WEST_METADATA, + get_default_west_gaps, +) +from waveform_editor.shape_editor.west_gaps import ( + compute_gaps as compute_west_gaps, +) # A shape to start from for each machine, taken from one of its discharges, as # (value, slider range) per parameter @@ -411,7 +428,7 @@ def __panel__(self): class PlasmaShape(Viewer): PARAMETERIZED_INPUT = "Parameterized" EQUILIBRIUM_INPUT = "Equilibrium IDS outline" - GAP_INPUT = "Equilibrium IDS Gaps" + GAP_INPUT = "Gaps" WEIGHTED_POINTS_INPUT = "Weighted Points" input_mode = param.ObjectSelector( default=EQUILIBRIUM_INPUT, @@ -423,6 +440,13 @@ class PlasmaShape(Viewer): ], label="Shape input mode", ) + GAP_SOURCE_DEFAULT = "Default Gaps" + GAP_SOURCE_IDS = "Equilibrium IDS" + gap_source = param.ObjectSelector( + default=GAP_SOURCE_DEFAULT, + objects=[GAP_SOURCE_DEFAULT, GAP_SOURCE_IDS], + label="Gap Source", + ) input_outline = param.ClassSelector( class_=EquilibriumInput, default=EquilibriumInput() ) @@ -438,6 +462,10 @@ class PlasmaShape(Viewer): has_shape = param.Boolean(doc="Whether a plasma shape is loaded.") shape_updated = param.Event(doc="Triggered whenever the plasma shape updates.") + show_gaps = param.Boolean(default=True, doc="Whether clearance gaps are displayed.") + machine_preset = param.String( + default="", doc="Active machine preset (e.g. ITER, WEST)." + ) def __init__(self): super().__init__() @@ -454,6 +482,24 @@ def _indicator(tooltip): self.weighted_points_indicator = _indicator( "At least 1 point is required to define a plasma shape" ) + self.gap_source_radio = pn.widgets.RadioButtonGroup( + options={ + "Default Gaps": self.GAP_SOURCE_DEFAULT, + "Equilibrium IDS": self.GAP_SOURCE_IDS, + }, + value=self.gap_source, + button_type="primary", + sizing_mode="stretch_width", + margin=(5, 20, 5, 20), + ) + self.gap_source_radio.link(self, value="gap_source", bidirectional=True) + self.reset_default_gaps_button = pn.widgets.Button( + name="Reset Defaults", + button_type="default", + sizing_mode="stretch_width", + margin=(5, 20, 5, 20), + ) + self.reset_default_gaps_button.on_click(self._on_reset_default_gaps) self._mode_config = { self.EQUILIBRIUM_INPUT: ( self._load_shape_from_ids, @@ -469,9 +515,7 @@ def _indicator(tooltip): ), self.GAP_INPUT: ( self._load_shape_from_gaps, - lambda: pn.Row( - self.input_gaps, self.gap_indicator, margin=(10, 20, 0, 20) - ), + self._panel_gap_options, ), self.WEIGHTED_POINTS_INPUT: ( self._load_shape_from_weighted_points, @@ -485,7 +529,7 @@ def _indicator(tooltip): options={ "Equilibrium\nIDS Outline": self.EQUILIBRIUM_INPUT, "Parameterized": self.PARAMETERIZED_INPUT, - "Equilibrium\nIDS Gaps": self.GAP_INPUT, + "Gaps": self.GAP_INPUT, "Weighted\nPoints": self.WEIGHTED_POINTS_INPUT, }, value=self.input_mode, @@ -495,7 +539,18 @@ def _indicator(tooltip): stylesheets=[CARD_CSS], ) self.radio_box.link(self, value="input_mode", bidirectional=True) - self.panel = pn.Column(self.radio_box, self._panel_shape_options, self.gap_ui) + self.gap_pills_pane = pn.pane.HTML( + self._render_target_gap_pills, + sizing_mode="stretch_width", + margin=(5, 20, 0, 20), + stylesheets=STYLES, + ) + self.panel = pn.Column( + self.radio_box, + self.gap_pills_pane, + self._panel_shape_options, + self.gap_ui, + ) self.outline_r = None self.outline_z = None self.gaps = [] @@ -518,6 +573,8 @@ def _indicator(tooltip): "input_gaps.load", "weighted_points_table.param", "input_mode", + "gap_source", + "machine_preset", watch=True, ) def _set_plasma_shape(self): @@ -535,6 +592,32 @@ def _set_plasma_shape(self): self.has_shape = False self.param.trigger("shape_updated") + @pn.depends("gap_source", "machine_preset") + def _panel_gap_options(self): + """Render controls for Gaps mode: gap source toggle and appropriate loaders.""" + preset = self.machine_preset or "ITER" + if self.gap_source == self.GAP_SOURCE_DEFAULT: + return pn.Column( + self.gap_source_radio, + pn.Row( + pn.pane.Markdown( + f"**{preset} Default Gap Definitions**", + margin=(8, 10, 0, 20), + ), + self.reset_default_gaps_button, + sizing_mode="stretch_width", + ), + ) + return pn.Column( + self.gap_source_radio, + pn.Row(self.input_gaps, self.gap_indicator, margin=(5, 20, 0, 20)), + ) + + def _on_reset_default_gaps(self, event=None): + """Reset default gaps back to their machine defaults.""" + self._load_shape_from_gaps() + self.param.trigger("shape_updated") + def _load_shape_from_ids(self): """Load plasma boundary outline from IDS equilibrium input.""" if not self.input_outline.uri: @@ -555,10 +638,16 @@ def _load_shape_from_ids(self): self.outline_r = self.outline_z = None def _load_shape_from_gaps(self): - """Load plasma boundary outline from IDS equilibrium gap definitions.""" + """Load plasma boundary outline from default machine gaps or IDS equilibrium.""" self.gaps = [] - if self.input_gaps.uri: + if self.gap_source == self.GAP_SOURCE_DEFAULT: + self.gaps = ( + get_default_west_gaps() + if self.machine_preset == "WEST" + else get_default_iter_gaps() + ) + elif self.gap_source == self.GAP_SOURCE_IDS and self.input_gaps.uri: try: with imas.DBEntry(self.input_gaps.uri, "r") as entry: equilibrium = entry.get_slice( @@ -572,16 +661,16 @@ def _load_shape_from_gaps(self): "The equilibrium IDS does not have any gaps" ) else: - for gap in input_gaps: - self.gaps.append( - Gap( - r=gap.r, - z=gap.z, - name=gap.name, - angle=gap.angle, - value=gap.value, - ) + self.gaps = [ + Gap( + r=float(g.r), + z=float(g.z), + name=str(g.name), + angle=float(g.angle), + value=float(g.value), ) + for g in input_gaps + ] except Exception as e: pn.state.notifications.error( f"Could not load gaps from {self.input_gaps.uri}: {str(e)}" @@ -609,11 +698,23 @@ def _create_gap_ui(self): new_gap_ui = [] for i, gap in enumerate(self.gaps): + name_lower = gap.name.lower() + is_divertor = any( + k in name_lower for k in ("divertor", "g1", "g2", "dxlow", "dxup") + ) + start_val = -0.5 if is_divertor else 0.0 + end_val = 0.5 if is_divertor else 1.0 + val = float(gap.value) + if val < start_val: + start_val = round(val - 0.1, 2) + if val > end_val: + end_val = round(val + 0.1, 2) + value_input = FormattedEditableFloatSlider( name=f"Gap {i}: {gap.name} Value [m]", - value=float(gap.value), - start=0, - end=1, + value=val, + start=start_val, + end=end_val, step=0.01, stretch_width=True, ) @@ -682,5 +783,73 @@ def _panel_shape_options(self): _, panel_factory = self._mode_config[self.input_mode] return panel_factory() + @param.depends( + "show_gaps", "machine_preset", "shape_updated", "has_shape", "input_mode" + ) + def _render_target_gap_pills(self): + """Render target shape metric pills when show_gaps is enabled.""" + if ( + not self.show_gaps + or not self.has_shape + or self.outline_r is None + or self.outline_z is None + or len(self.outline_r) == 0 + ): + return "" + + chips = [] + + def _chip(symbol, display, tooltip): + return ( + f'' + f'{symbol}' + f'{display}' + ) + + if self.machine_preset == "ITER": + gaps = compute_iter_gaps(self.outline_r, self.outline_z) + if self.input_mode == self.GAP_INPUT and self.gaps: + for g in self.gaps: + nl = str(g.name).lower() + if "gap1" in nl or "g1" in nl or "inner" in nl: + gaps["gap1"] = float(g.value) + elif "gap2" in nl or "g2" in nl or "outer" in nl: + gaps["gap2"] = float(g.value) + + for key in ("gap1", "gap2", "gap4", "gap5", "Rmin", "Rmax"): + symbol, unit, tooltip = ITER_GAP_METADATA[key] + val = gaps.get(key) + display = ( + "—" + if val is None + else f"{val:.4g} m" + if unit == "m" + else f"{val * 100:.3g} cm" + ) + chips.append(_chip(symbol, display, tooltip)) + + elif self.machine_preset == "WEST": + x_points = [] + if self.input_mode == self.PARAMETERIZED_INPUT: + x_points = [(self.shape_params.rx, self.shape_params.zx)] + elif self.input_mode == self.GAP_INPUT and self.gaps: + x_points = [ + (g.r_sep, g.z_sep) + for g in self.gaps + if any(k in g.name.lower() for k in ("dxlow", "dxup")) + ] + west_gaps = compute_west_gaps(self.outline_r, self.outline_z, x_points) + if self.input_mode == self.GAP_INPUT and self.gaps: + for g in self.gaps: + for key in WEST_METADATA: + if key.lower() in g.name.lower(): + west_gaps[key] = float(g.value) + for key, (symbol, _, tooltip) in WEST_METADATA.items(): + val = west_gaps.get(key) + if val is not None: + chips.append(_chip(symbol, f"{val * 100:.3g} cm", tooltip)) + + return f'
{"".join(chips)}
' if chips else "" + def __panel__(self): return self.panel diff --git a/waveform_editor/gui/shape_editor/settings_modal.py b/waveform_editor/gui/shape_editor/settings_modal.py index 13221747..b785f25b 100644 --- a/waveform_editor/gui/shape_editor/settings_modal.py +++ b/waveform_editor/gui/shape_editor/settings_modal.py @@ -229,6 +229,7 @@ def _build_modal(self): "show_xo", "show_separatrix", "show_desired_shape", + "show_gaps", ], show_name=False, widgets={ diff --git a/waveform_editor/gui/shape_editor/shape_editor.py b/waveform_editor/gui/shape_editor/shape_editor.py index 89dc95ef..f82955e9 100644 --- a/waveform_editor/gui/shape_editor/shape_editor.py +++ b/waveform_editor/gui/shape_editor/shape_editor.py @@ -20,22 +20,25 @@ from waveform_editor.gui.shape_editor.waveform_sync import WaveformSync from waveform_editor.gui.util import set_xml_parameter from waveform_editor.settings import NiceSettings, settings +from waveform_editor.shape_editor.iter_gaps import ( + GAP_METADATA as ITER_GAP_METRICS, +) +from waveform_editor.shape_editor.iter_gaps import ( + compute_gaps as compute_iter_gaps, +) +from waveform_editor.shape_editor.iter_gaps import ( + extract_separatrix_from_time_slice, +) from waveform_editor.shape_editor.nice_integration import NiceIntegration -from waveform_editor.shape_editor.west_gaps import compute_gaps +from waveform_editor.shape_editor.west_gaps import ( + WEST_METADATA as WEST_GAP_METRICS, +) +from waveform_editor.shape_editor.west_gaps import ( + compute_gaps as compute_west_gaps, +) logger = logging.getLogger(__name__) -# TODO: add the gaps for ITER used in DINA: gap1, gap2, gap4, gap5, Rmin, Rmax -# The gaps of a WEST plasma, as (symbol, unit, full name) by name -WEST_GAP_METRICS = { - "UROG": ("UROG", "cm", "Upper radial outer gap"), - "EROG": ("EROG", "cm", "Equatorial radial outer gap"), - "LROG": ("LROG", "cm", "Lower radial outer gap"), - "dXlow": ("dXlow", "cm", "Distance of the lower x-point to the divertor"), - "dXup": ("dXup", "cm", "Distance of the upper x-point to the divertor"), - "dbaffle": ("dbaffle", "cm", "Distance of the plasma to the baffle"), -} - # NICE reads the desired boundary into an array of this fixed size # (MAX_PLASMA_BOUNDARY_POINTS in its solver_structs.h) @@ -90,6 +93,16 @@ def __init__(self, main_gui): ) self.run_select.param.watch(self._restore_run, "value") self.nice_settings = settings.nice + self.plasma_shape.machine_preset = self.nice_settings.machine_preset + self.plasma_shape.show_gaps = self.nice_plotter.show_gaps + self.nice_settings.param.watch( + lambda e: setattr(self.plasma_shape, "machine_preset", e.new), + "machine_preset", + ) + self.nice_plotter.param.watch( + lambda e: setattr(self.plasma_shape, "show_gaps", e.new), + "show_gaps", + ) self.xml_text = ( importlib.resources.files("waveform_editor.shape_editor.xml_param") @@ -301,8 +314,11 @@ def _restore_run(self, event): _, equilibrium, pf_active = self.run_history[event.new] self.communicator.equilibrium = equilibrium self.communicator.pf_active = pf_active + self.communicator.converged = True self.coil_currents.sync_ui_with_pf_active(pf_active) self._update_metrics() + if self.nice_settings.mode == NiceSettings.DIRECT_MODE: + self.use_previous_run = True @param.depends("nice_settings.md_pf_active.uri", watch=True) def _load_pf_active(self): @@ -488,31 +504,49 @@ async def submit(self, event=None): def _update_machine_metrics(self, event=None): """Show the chips of the machine of the selected preset, also before a run has filled them in.""" + preset = self.nice_settings.machine_preset self.metrics.machine_metrics = ( WEST_GAP_METRICS - if self.nice_settings.machine_preset == NiceSettings.PRESET_WEST + if preset == NiceSettings.PRESET_WEST + else ITER_GAP_METRICS + if preset == NiceSettings.PRESET_ITER else {} ) - def _west_gaps(self, time_slice): - """The gaps of the plasma to the parts of WEST it is kept away from. They are - not in the equilibrium, so they are computed from its boundary. - - Args: - time_slice: The time slice NICE returned. + def _iter_gaps(self, time_slice): + """The gaps of an ITER plasma in centimetres (Rmin/Rmax in metres).""" + if self.nice_settings.machine_preset != NiceSettings.PRESET_ITER: + return {} + outline = time_slice.boundary.outline + if outline.r is None or len(outline.r) == 0: + return {} + ma = getattr( + getattr(time_slice, "global_quantities", None), "magnetic_axis", None + ) + axis = (float(ma.r), float(ma.z)) if ma else None + gaps = compute_iter_gaps( + outline.r, + outline.z, + separatrix_contour=extract_separatrix_from_time_slice(time_slice), + magnetic_axis=axis, + ) + return { + name: val * 100 if name in ("gap1", "gap2", "gap4", "gap5") else val + for name, val in gaps.items() + if val is not None + } - Returns: - Dict of gap name to distance in centimetres, empty for another machine. - """ + def _west_gaps(self, time_slice): + """The gaps of a WEST plasma in centimetres.""" if self.nice_settings.machine_preset != NiceSettings.PRESET_WEST: return {} x_points = [ (float(node.r), float(node.z)) - for node in time_slice.contour_tree.node + for node in getattr(getattr(time_slice, "contour_tree", None), "node", []) if int(node.critical_type) == 1 ] outline = time_slice.boundary.outline - gaps = compute_gaps(outline.r, outline.z, x_points) + gaps = compute_west_gaps(outline.r, outline.z, x_points) return {name: gap * 100 for name, gap in gaps.items()} def _update_metrics(self): @@ -530,8 +564,20 @@ def _update_metrics(self): self.metrics.MINOR_RADIUS: float(boundary.minor_radius), self.metrics.Q95: float(global_quantities.q_95), **self._west_gaps(eq.time_slice[0]), + **self._iter_gaps(eq.time_slice[0]), } + @param.depends("nice_settings.mode", watch=True) + def _enable_warm_start_on_direct_mode(self): + """Automatically enable warm start when switching to NICE direct mode if + a previous run is available to warm start from.""" + if self.nice_settings.mode == NiceSettings.DIRECT_MODE and ( + self.communicator.converged or bool(self.run_history) + ): + if not self.communicator.converged and self.run_history: + self.communicator.converged = True + self.use_previous_run = True + @param.depends( "nice_settings.mode", "nice_settings.md_pf_active.uri", diff --git a/waveform_editor/shape_editor/iter_gaps.py b/waveform_editor/shape_editor/iter_gaps.py new file mode 100644 index 00000000..fd279ca2 --- /dev/null +++ b/waveform_editor/shape_editor/iter_gaps.py @@ -0,0 +1,260 @@ +"""The gaps between an ITER plasma and the first wall, as defined in DINA. + +The gap definitions and measurement points are taken from DINA-IMAS: +- gaps 1, 2, 4, 5 are the minimum distances from fixed wall coordinates to the + plasma boundary outline. +- Rmin and Rmax are the inboard and outboard radial extents (min and max R) of + the plasma boundary. +""" + +import math + +import numpy as np + +from waveform_editor.shape_editor.plasma_shape_calc import Gap, closest_outline_point + +# Fixed gap measurement points on the wall, in metres (r, z): +GAP_POINTS = { + "gap1": (4.2230, -3.7920), + "gap2": (5.5650, -4.4040), + "gap4": (7.5095, 2.9971), + "gap5": (5.3315, 4.5804), +} + +WALL_REFERENCE_POINTS = { + "Rmin": (4.0599, 0.7777), + "Rmax": (8.2806, 0.4665), +} + +GAP_METADATA = { + "gap1": ("g₁", "cm", "Inner divertor leg gap"), + "gap2": ("g₂", "cm", "Outer divertor leg gap"), + "gap4": ("g₄", "cm", "Point at 2 o'clock gap"), + "gap5": ("g₅", "cm", "Uppest boundary point gap"), + "Rmin": ("Rₘᵢₙ", "m", "Inboard boundary radius"), + "Rmax": ("Rₘₐₓ", "m", "Outboard boundary radius"), +} + + +def _distance_to_outline(point, r, z, is_closed=True): + return closest_outline_point(point, r, z, is_closed=is_closed)[1] + + +def extract_separatrix_from_time_slice(time_slice): + """Extract separatrix isocontour segments from an equilibrium time slice.""" + if time_slice is None: + return [] + + b_psi = getattr( + getattr(time_slice, "global_quantities", None), "psi_boundary", None + ) + if b_psi is None or float(b_psi) == -9e40: + b_psi = getattr(getattr(time_slice, "boundary", None), "psi", None) + if b_psi is None or float(b_psi) == -9e40: + return [] + b_psi = float(b_psi) + + import matplotlib.pyplot as plt + + def _contour_segs(cs): + segs = [np.asarray(seg) for seg in cs.allsegs[0] if len(seg) > 1] + plt.close(fig) + return segs + + # 1. Try GGD (NICE output) + if hasattr(time_slice, "ggd") and len(time_slice.ggd) > 0: + try: + g = time_slice.ggd[0] + r, z, psi = g.r[0].values, g.z[0].values, g.psi[0].values + if r and z and psi: + fig, ax = plt.subplots() + segs = _contour_segs(ax.tricontour(r, z, psi, levels=[b_psi])) + if segs: + return segs + except Exception: + pass + + # 2. Try profiles_2d (IMAS grid) + if hasattr(time_slice, "profiles_2d") and len(time_slice.profiles_2d) > 0: + try: + p2d = time_slice.profiles_2d[0] + r, z, psi = ( + np.asarray(p2d.grid.dim1), + np.asarray(p2d.grid.dim2), + np.asarray(p2d.psi), + ) + if len(r) > 0 and len(z) > 0 and psi.size > 0: + psi_grid = psi.T if psi.shape == (len(r), len(z)) else psi + fig, ax = plt.subplots() + segs = _contour_segs(ax.contour(r, z, psi_grid, levels=[b_psi])) + if segs: + return segs + except Exception: + pass + + return [] + + +def _normalize_contour_segments(separatrix_contour): + """Normalize contour segment representations into a list of (N, 2) arrays.""" + if separatrix_contour is None: + return [] + if ( + isinstance(separatrix_contour, (tuple, list)) + and len(separatrix_contour) == 2 + and isinstance(separatrix_contour[0], (list, np.ndarray)) + and len(separatrix_contour[0]) > 0 + and not isinstance(separatrix_contour[0][0], (list, np.ndarray, dict)) + ): + r, z = ( + np.asarray(separatrix_contour[0], float), + np.asarray(separatrix_contour[1], float), + ) + return [np.column_stack([r, z])] if len(r) > 1 else [] + + segments = [] + for seg in separatrix_contour: + arr = ( + np.column_stack([seg["x"], seg["y"]]) + if isinstance(seg, dict) and "x" in seg + else np.asarray(seg, float) + ) + if arr.ndim == 2 and arr.shape[1] >= 2 and len(arr) > 1: + segments.append(arr[:, :2]) + return segments + + +def _dina_divertor_sign(k, r_target, z_target, r_orig, z_orig, magnetic_axis): + """Calculate the DINA sign for divertor gaps 1 and 2.""" + if magnetic_axis is None: + return 1.0 + rmag, zmag = magnetic_axis + vecpro = (r_target - rmag) * (z_orig - zmag) - (z_target - zmag) * (r_orig - rmag) + s_vecpro = 1.0 if vecpro >= 0 else -1.0 + p = k - 1.5 + return (p / abs(p)) * s_vecpro + + +def _closest_point_on_segments(point, segments): + """Find the closest point to `point` across multiple polyline segments.""" + min_dist, best_target = float("inf"), None + for seg in segments: + target, dist = closest_outline_point( + point, seg[:, 0], seg[:, 1], is_closed=False + ) + if dist < min_dist: + min_dist, best_target = dist, target + return best_target, min_dist + + +def compute_gaps(outline_r, outline_z, separatrix_contour=None, magnetic_axis=None): + """The gaps and radial extent of an ITER plasma, in metres.""" + if outline_r is None or outline_z is None or len(outline_r) == 0: + return {} + + r = np.asarray(outline_r, dtype=float) + z = np.asarray(outline_z, dtype=float) + + gaps = { + "gap4": closest_outline_point(GAP_POINTS["gap4"], r, z)[1], + "gap5": closest_outline_point(GAP_POINTS["gap5"], r, z)[1], + "Rmin": float(np.min(r)), + "Rmax": float(np.max(r)), + } + + segments = _normalize_contour_segments(separatrix_contour) + has_legs = len(segments) > 0 or float(np.min(z)) <= -3.6 + if not segments and has_legs: + segments = [np.column_stack([r, z])] + + if has_legs: + for k, name in enumerate(("gap1", "gap2"), start=1): + pt = GAP_POINTS[name] + target, dist = _closest_point_on_segments(pt, segments) + if target is not None: + sign = _dina_divertor_sign( + k, target[0], target[1], pt[0], pt[1], magnetic_axis + ) + gaps[name] = dist * sign + else: + gaps[name] = None + else: + gaps["gap1"] = None + gaps["gap2"] = None + + return gaps + + +def compute_gap_geometry( + outline_r, outline_z, separatrix_contour=None, magnetic_axis=None +): + """Compute measuring points and distance segments for ITER gaps.""" + if outline_r is None or outline_z is None or len(outline_r) == 0: + return [] + + r, z = np.asarray(outline_r, float), np.asarray(outline_z, float) + items = [] + + segments = _normalize_contour_segments(separatrix_contour) or [ + np.column_stack([r, z]) + ] + for k, name in enumerate(("gap1", "gap2"), start=1): + r_orig, z_orig = GAP_POINTS[name] + target, dist = _closest_point_on_segments((r_orig, z_orig), segments) + if target is not None: + sign = _dina_divertor_sign( + k, target[0], target[1], r_orig, z_orig, magnetic_axis + ) + symbol, _, full_name = GAP_METADATA[name] + items.append( + { + "name": full_name, + "symbol": symbol, + "r_orig": r_orig, + "z_orig": z_orig, + "r_target": target[0], + "z_target": target[1], + "distance": dist, + "distance_str": f"{dist * sign * 100:+.2f} cm", + } + ) + + ref_points = {**GAP_POINTS, **WALL_REFERENCE_POINTS} + min_r, max_r = float(np.min(r)), float(np.max(r)) + for name in ("gap4", "gap5", "Rmin", "Rmax"): + r_orig, z_orig = ref_points[name] + (r_target, z_target), dist = closest_outline_point((r_orig, z_orig), r, z) + symbol, _, full_name = GAP_METADATA[name] + dist_str = ( + f"Rₘᵢₙ = {min_r:.4g} m (gap: {dist * 100:.3g} cm)" + if name == "Rmin" + else f"Rₘₐₓ = {max_r:.4g} m (gap: {dist * 100:.3g} cm)" + if name == "Rmax" + else f"{dist * 100:.3g} cm" + ) + items.append( + { + "name": full_name, + "symbol": symbol, + "r_orig": r_orig, + "z_orig": z_orig, + "r_target": r_target, + "z_target": z_target, + "distance": dist, + "distance_str": dist_str, + } + ) + + return items + + +def get_default_iter_gaps(): + """Return default Gap definitions for ITER based on DINA coordinates.""" + return [ + Gap("Inner divertor leg (g₁)", 4.2230, -3.7920, math.radians(-115.0), 0.0), + Gap("Outer divertor leg (g₂)", 5.5650, -4.4040, math.radians(-150.0), 0.0), + Gap("Point at 2 o'clock (g₄)", 7.5095, 2.9971, math.radians(135.0), 0.228), + Gap("Uppest boundary point (g₅)", 5.3315, 4.5804, math.radians(90.0), 0.596), + Gap("Inboard mid-plane (Rmin)", 4.0599, 0.7777, math.radians(0.0), 0.150), + Gap("Outboard mid-plane (Rmax)", 8.2806, 0.4665, math.radians(180.0), 0.080), + ] diff --git a/waveform_editor/shape_editor/plasma_shape_calc.py b/waveform_editor/shape_editor/plasma_shape_calc.py index 1b69b72a..6da07e20 100644 --- a/waveform_editor/shape_editor/plasma_shape_calc.py +++ b/waveform_editor/shape_editor/plasma_shape_calc.py @@ -191,3 +191,39 @@ def update_outline_from_gaps(gaps): if not gaps: return None, None return [gap.r_sep for gap in gaps], [gap.z_sep for gap in gaps] + + +def closest_outline_point(point, r, z, is_closed=True): + """Find closest point on outline to a given (r, z) coordinate.""" + starts = np.column_stack([r, z]) + if len(starts) == 0: + return (0.0, 0.0), float("inf") + if len(starts) == 1: + return (float(starts[0, 0]), float(starts[0, 1])), float( + np.linalg.norm(starts[0] - point) + ) + ends = np.roll(starts, -1, axis=0) if is_closed else starts[1:] + if not is_closed: + starts = starts[:-1] + segs = ends - starts + lens = np.sum(segs**2, axis=1) + along = np.clip( + np.sum((np.asarray(point, float) - starts) * segs, axis=1) + / np.where(lens > 0, lens, 1), + 0, + 1, + ) + closest = starts + along[:, None] * segs + dists = np.linalg.norm(closest - point, axis=1) + idx = int(np.argmin(dists)) + return (float(closest[idx, 0]), float(closest[idx, 1])), float(dists[idx]) + + +def project_point_to_line(point, start, end): + """Project a point onto the line passing through start and end.""" + start, end = np.asarray(start, float), np.asarray(end, float) + along = (end - start) / np.linalg.norm(end - start) + to_pt = np.asarray(point, float) - start + proj = start + np.dot(to_pt, along) * along + dist = float(np.linalg.norm(to_pt - np.dot(to_pt, along) * along)) + return (float(proj[0]), float(proj[1])), dist diff --git a/waveform_editor/shape_editor/west_gaps.py b/waveform_editor/shape_editor/west_gaps.py index 23b1c7ac..06e0f7cf 100644 --- a/waveform_editor/shape_editor/west_gaps.py +++ b/waveform_editor/shape_editor/west_gaps.py @@ -4,8 +4,16 @@ does, in Projects/WEST/Lib/plot_plasma_gaps_etc.m, from the geometry of WEST. """ +import math + import numpy as np +from waveform_editor.shape_editor.plasma_shape_calc import ( + Gap, + closest_outline_point, + project_point_to_line, +) + # The arc the outer radial gaps are measured to, which passes through r=3 m on the # midplane, as (centre r, radius) OUTER_ARC = (2.2, 0.8) @@ -17,6 +25,15 @@ # The corner of the baffle the plasma is kept away from BAFFLE = (2.381, -0.6757) +WEST_METADATA = { + "UROG": ("UROG", "cm", "Upper radial outer gap"), + "EROG": ("EROG", "cm", "Equatorial radial outer gap"), + "LROG": ("LROG", "cm", "Lower radial outer gap"), + "dXlow": ("dXlow", "cm", "Distance of the lower x-point to the divertor"), + "dXup": ("dXup", "cm", "Distance of the upper x-point to the divertor"), + "dbaffle": ("dbaffle", "cm", "Distance of the plasma to the baffle"), +} + def compute_gaps(outline_r, outline_z, x_points): """The gaps of a WEST plasma, in metres. @@ -27,8 +44,7 @@ def compute_gaps(outline_r, outline_z, x_points): x_points: The (r, z) of each x-point of the equilibrium. Returns: - Dict of gap name to distance, holding only the gaps that the boundary and the - x-points given allow to be computed. + Dict of gap name to distance in metres. """ r, z = np.asarray(outline_r, dtype=float), np.asarray(outline_z, dtype=float) gaps = {} @@ -38,10 +54,9 @@ def compute_gaps(outline_r, outline_z, x_points): ("EROG", 0.0), ("LROG", -OUTER_GAP_HEIGHT), ): - arc_r = centre_r + np.sqrt(radius**2 - height**2) boundary_r = _outboard_radius(r, z, height) if boundary_r is not None: - gaps[name] = arc_r - boundary_r + gaps[name] = centre_r + np.sqrt(radius**2 - height**2) - boundary_r for name, divertor, below in ( ("dXlow", LOWER_DIVERTOR, True), @@ -49,23 +64,84 @@ def compute_gaps(outline_r, outline_z, x_points): ): x_point = _x_point(x_points, below) if x_point is not None: - gaps[name] = _distance_to_line(x_point, *divertor) + gaps[name] = project_point_to_line(x_point, *divertor)[1] - gaps["dbaffle"] = _distance_to_outline(BAFFLE, r, z) + gaps["dbaffle"] = closest_outline_point(BAFFLE, r, z)[1] return gaps -def _outboard_radius(r, z, height): - """The radius of the outboard side of the boundary at a height. +def compute_gap_geometry(outline_r, outline_z, x_points): + """Compute measuring points and distance segments for WEST gaps.""" + if outline_r is None or outline_z is None or len(outline_r) == 0: + return [] - Args: - r: Radial coordinates of the boundary. - z: Height coordinates of the boundary. - height: The height to take the boundary at. + r, z = np.asarray(outline_r, dtype=float), np.asarray(outline_z, dtype=float) + items = [] + centre_r, radius = OUTER_ARC + for name, height in ( + ("UROG", OUTER_GAP_HEIGHT), + ("EROG", 0.0), + ("LROG", -OUTER_GAP_HEIGHT), + ): + arc_r = float(centre_r + np.sqrt(radius**2 - height**2)) + boundary_r = _outboard_radius(r, z, height) + if boundary_r is not None: + dist = arc_r - boundary_r + symbol, _, full_name = WEST_METADATA[name] + items.append( + { + "name": full_name, + "symbol": symbol, + "r_orig": arc_r, + "z_orig": height, + "r_target": boundary_r, + "z_target": height, + "distance": dist, + "distance_str": f"{dist * 100:.3g} cm", + } + ) - Returns: - The radius, or None when the boundary does not reach that height. - """ + for name, divertor, below in ( + ("dXlow", LOWER_DIVERTOR, True), + ("dXup", UPPER_DIVERTOR, False), + ): + x_point = _x_point(x_points, below) + if x_point is not None: + proj, dist = project_point_to_line(x_point, *divertor) + symbol, _, full_name = WEST_METADATA[name] + items.append( + { + "name": full_name, + "symbol": symbol, + "r_orig": proj[0], + "z_orig": proj[1], + "r_target": float(x_point[0]), + "z_target": float(x_point[1]), + "distance": dist, + "distance_str": f"{dist * 100:.3g} cm", + } + ) + + (target_r, target_z), dist = closest_outline_point(BAFFLE, r, z) + symbol, _, full_name = WEST_METADATA["dbaffle"] + items.append( + { + "name": full_name, + "symbol": symbol, + "r_orig": BAFFLE[0], + "z_orig": BAFFLE[1], + "r_target": target_r, + "z_target": target_z, + "distance": dist, + "distance_str": f"{dist * 100:.3g} cm", + } + ) + + return items + + +def _outboard_radius(r, z, height): + """The radius of the outboard side of the boundary at a height.""" outboard = r > r.mean() r, z = r[outboard], z[outboard] if not len(z) or not z.min() <= height <= z.max(): @@ -75,57 +151,28 @@ def _outboard_radius(r, z, height): def _x_point(x_points, below): - """The x-point below or above the midplane. - - Args: - x_points: The (r, z) of each x-point. - below: Whether to take the x-point below the midplane. - - Returns: - The (r, z) of the x-point, or None when there is none on that side. - """ + """The x-point below or above the midplane.""" on_side = [point for point in x_points if (point[1] < 0) == below] return on_side[0] if on_side else None -def _distance_to_line(point, start, end): - """The distance of a point to the line through two points. - - Args: - point: The (r, z) to measure from. - start: The (r, z) of a point of the line. - end: The (r, z) of another point of the line. - - Returns: - The distance. - """ - point, start, end = (np.asarray(p, dtype=float) for p in (point, start, end)) - along = (end - start) / np.linalg.norm(end - start) - to_point = point - start - return float(np.linalg.norm(to_point - np.dot(to_point, along) * along)) - - -def _distance_to_outline(point, r, z): - """The distance of a point to the closest of the segments of an outline. - - Args: - point: The (r, z) to measure from. - r: Radial coordinates of the outline. - z: Height coordinates of the outline. - - Returns: - The distance. - """ - starts = np.column_stack([r, z]) - ends = np.roll(starts, -1, axis=0) - segments = ends - starts - lengths = np.sum(segments**2, axis=1) - # How far along each segment the point is, kept within the segment - along = np.clip( - np.sum((np.asarray(point, dtype=float) - starts) * segments, axis=1) - / np.where(lengths > 0, lengths, 1), - 0, - 1, - ) - closest = starts + along[:, None] * segments - return float(np.min(np.linalg.norm(closest - point, axis=1))) +def get_default_west_gaps(): + """Return default Gap definitions for WEST based on machine geometry.""" + return [ + Gap("Upper radial outer gap (UROG)", 2.9599, 0.25, math.pi, 0.06), + Gap("Equatorial radial outer gap (EROG)", 3.0000, 0.0, math.pi, 0.08), + Gap("Lower radial outer gap (LROG)", 2.9599, -0.25, math.pi, 0.06), + Gap( + "Lower x-point to divertor (dXlow)", + 2.1350, + -0.6710, + -math.radians(68.0), + 0.05, + ), + Gap( + "Upper x-point to divertor (dXup)", 2.1730, 0.6910, math.radians(68.0), 0.10 + ), + Gap( + "Distance to baffle (dbaffle)", 2.3810, -0.6757, -math.radians(135.0), 0.04 + ), + ] From a30fd22cebf28f7d132871a43bd67c7f602aec11 Mon Sep 17 00:00:00 2001 From: Sebbe Blokhuizen Date: Tue, 6 Oct 2026 09:44:47 +0200 Subject: [PATCH 3/3] add DINA gaps --- tests/test_iter_gaps.py | 468 ++++++++++++++---- .../gui/shape_editor/nice_plotter.py | 102 ++-- .../gui/shape_editor/plasma_shape.py | 110 ++-- .../gui/shape_editor/settings_modal.py | 16 + .../gui/shape_editor/shape_editor.py | 65 +-- waveform_editor/gui/styles/property_card.css | 3 + waveform_editor/shape_editor/iter_gaps.py | 285 +++++------ .../shape_editor/plasma_shape_calc.py | 107 ++++ waveform_editor/shape_editor/west_gaps.py | 266 ++++++---- 9 files changed, 940 insertions(+), 482 deletions(-) diff --git a/tests/test_iter_gaps.py b/tests/test_iter_gaps.py index 4b4a6af0..3e1d184f 100644 --- a/tests/test_iter_gaps.py +++ b/tests/test_iter_gaps.py @@ -3,6 +3,7 @@ import numpy as np from waveform_editor.gui.shape_editor.shape_editor import ITER_GAP_METRICS +from waveform_editor.settings import NiceSettings, settings from waveform_editor.shape_editor.iter_gaps import ( GAP_POINTS, _distance_to_outline, @@ -110,21 +111,55 @@ def test_compute_gap_geometry_iter(): outline_r = r0 + a * np.cos(theta) outline_z = z0 + a * np.sin(theta) - # Without divertor legs: all 6 gaps are still returned - # so strike point locations are visible + # Without divertor legs the strike points are still shown, but not measured: a + # closed boundary has no legs to measure them to items = compute_gap_geometry(outline_r, outline_z) assert len(items) == 6 - symbols = {item["symbol"] for item in items} - assert symbols == {"g₁", "g₂", "g₄", "g₅", "Rₘᵢₙ", "Rₘₐₓ"} - - for item in items: + by_key = {item["key"]: item for item in items} + for key in ("gap1", "gap2"): + assert by_key[key]["distance"] is None + assert by_key[key]["r_target"] == by_key[key]["r_orig"] + for key in ("gap4", "gap5", "Rmin", "Rmax"): + item = by_key[key] expected_dist = math.hypot( item["r_target"] - item["r_orig"], item["z_target"] - item["z_orig"], ) assert math.isclose(item["distance"], expected_dist, rel_tol=1e-5) assert item["distance"] > 0 - assert "cm" in item["distance_str"] or "m" in item["distance_str"] + + # Rmin and Rmax are measured from r=0, at the height of the extreme point + for key in ("Rmin", "Rmax"): + assert by_key[key]["r_orig"] == 0.0 + assert by_key[key]["z_orig"] == by_key[key]["z_target"] + assert math.isclose(by_key["Rmin"]["distance"], 4.2, abs_tol=1e-3) + assert math.isclose(by_key["Rmax"]["distance"], 8.2, abs_tol=1e-3) + + +def test_iter_closed_separatrix_has_no_divertor_gaps(): + from waveform_editor.shape_editor.iter_gaps import compute_gaps + + theta = np.linspace(0, 2 * np.pi, 200) + r, z = 6.2 + 2.0 * np.cos(theta), 2.0 * np.sin(theta) + # A limited plasma: its boundary contour closes on itself + gaps = compute_gaps(r, z, separatrix_contour=[np.column_stack([r, z])]) + assert gaps["gap1"] is None + assert gaps["gap2"] is None + + +def test_iter_gaps_4_5_only_measure_towards_the_axis(): + from waveform_editor.shape_editor.plasma_shape_calc import ( + closest_outline_point, + closest_point_facing, + ) + + # A branch 0.1 behind the point, away from the axis below it, and the plasma 1 + # towards the axis: DINA only takes the latter for gaps 4 and 5 + r = [-1.0, 1.0, 1.0, -1.0] + z = [0.1, 0.1, -1.0, -1.0] + assert math.isclose(closest_outline_point((0, 0), r, z, is_closed=False)[1], 0.1) + target, distance = closest_point_facing((0, 0), r, z, (0, -10), is_closed=False) + assert math.isclose(distance, 1.0) def test_compute_gap_geometry_west(): @@ -137,14 +172,69 @@ def test_compute_gap_geometry_west(): x_points = [(2.2, -0.6), (2.2, 0.6)] items = compute_gap_geometry(outline_r, outline_z, x_points) - assert len(items) == 6 - symbols = {item["symbol"] for item in items} - assert symbols == {"UROG", "EROG", "LROG", "dXlow", "dXup", "dbaffle"} - + by_key = {item["key"]: item for item in items} + assert set(by_key) == { + "UROG", + "EROG", + "LROG", + "RIG", + "TIG", + "TOG", + "dXlow", + "dXup", + "dbaffle", + } for item in items: assert item["distance"] > 0 assert "cm" in item["distance_str"] + # The FEEQS definitions, worked out for a circle + def circle_r(height): + return r0 + math.sqrt(a**2 - height**2) + + def circle_z(radius): + return math.sqrt(a**2 - (radius - r0) ** 2) + + wall_up = 3 - 0.8 + math.sqrt(0.8**2 - 0.25**2) + expected = { + "UROG": wall_up - circle_r(0.25), + "EROG": 3.0 - circle_r(0.0), + "LROG": wall_up - circle_r(0.25), + "RIG": (r0 - a) - 1.834, + "TIG": 0.672 - circle_z(2.132), + "TOG": 0.749 - circle_z(2.456), + } + for key, value in expected.items(): + assert math.isclose(by_key[key]["distance"], value, abs_tol=1e-4), key + + +def test_west_x_points_outside_the_feeqs_region_are_ignored(): + from waveform_editor.shape_editor.west_gaps import compute_gaps + + theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) + r, z = 2.5 + 0.4 * np.cos(theta), 0.4 * np.sin(theta) + # Below the lower divertor, where FEEQS does not look for x-points + gaps = compute_gaps(r, z, [(2.2, -0.95), (2.2, -0.6)]) + assert math.isclose(gaps["dXlow"], compute_gaps(r, z, [(2.2, -0.6)])["dXlow"]) + + +def test_west_first_separatrix_is_closest_in_flux_to_the_axis(): + from waveform_editor.shape_editor.west_gaps import compute_gaps + + theta = np.linspace(0, 2 * np.pi, 200, endpoint=False) + r, z = 2.5 + 0.4 * np.cos(theta), 0.4 * np.sin(theta) + + def contour(level): + # Separatrices as circles, the one of the upper x-point slightly larger + radius = 0.4 if level == -0.10 else 0.43 + t = np.linspace(0, 2 * np.pi, 400) + return [np.column_stack([2.5 + radius * np.cos(t), radius * np.sin(t)])] + + # psi decreasing outwards from the axis; the lower x-point is the first + x_points = [(2.2, 0.6, -0.20), (2.2, -0.6, -0.10)] + gaps = compute_gaps(r, z, x_points, magnetic_axis=(2.5, 0.0, 0.0), contour=contour) + assert math.isclose(gaps["dRsep"], 0.03, abs_tol=1e-3) + def test_nice_plotter_gap_visualization(): import holoviews as hv @@ -169,28 +259,73 @@ def test_nice_plotter_gap_visualization(): ) assert hasattr(plotter, "show_gaps") + assert hasattr(plotter, "show_desired_gaps") + assert hasattr(plotter, "show_result_gaps") assert plotter.show_gaps is True + assert plotter.show_desired_gaps is True + assert plotter.show_result_gaps is False # When no shape or equilibrium is loaded, overlay is empty empty_overlay = plotter._plot_clearance_gaps() assert isinstance(empty_overlay, hv.Overlay) # Set up a parameterized shape for ITER - settings.nice.machine_preset = NiceSettings.PRESET_ITER - shape.outline_r = [5.0, 7.0, 7.0, 5.0] - shape.outline_z = [-1.0, -1.0, 1.0, 1.0] - - overlay = plotter._plot_clearance_gaps() - assert isinstance(overlay, hv.Overlay) - # Check that overlay contains Points and Segments - types = [type(el) for el in overlay.values()] - assert hv.Points in types - assert hv.Segments in types - - # Toggle off show_gaps - plotter.show_gaps = False - off_overlay = plotter._plot_clearance_gaps() - assert isinstance(off_overlay, hv.Overlay) + orig_preset = settings.nice.machine_preset + orig_mode = settings.nice.mode + try: + settings.nice.machine_preset = NiceSettings.PRESET_ITER + settings.nice.mode = NiceSettings.INVERSE_MODE + shape.has_shape = True + shape.outline_r = [5.0, 7.0, 7.0, 5.0] + shape.outline_z = [-1.0, -1.0, 1.0, 1.0] + + overlay = plotter._plot_clearance_gaps() + assert isinstance(overlay, hv.Overlay) + types = [type(el) for el in overlay.values()] + assert hv.Points in types + assert hv.Segments in types + for el in overlay.values(): + if isinstance(el, (hv.Points, hv.Segments)): + assert el.opts.get("style").kwargs["color"] == "blue" + elif isinstance(el, hv.Labels): + assert el.opts.get("style").kwargs["text_color"] == "blue" + + # Set up an equilibrium result, switch to result gaps + eq = factory.equilibrium() + eq.time_slice.resize(1) + eq.time_slice[0].boundary.outline.r = np.array([5.0, 7.0, 7.0, 5.0]) + eq.time_slice[0].boundary.outline.z = np.array([-1.0, -1.0, 1.0, 1.0]) + eq.code.output_flag = [0] + comm.equilibrium = eq + plotter.show_result_gaps = True + plotter.show_desired_gaps = False + + eq_overlay = plotter._plot_clearance_gaps() + assert isinstance(eq_overlay, hv.Overlay) + for el in eq_overlay.values(): + if isinstance(el, (hv.Points, hv.Segments)): + assert el.opts.get("style").kwargs["color"] == "red" + elif isinstance(el, hv.Labels): + assert el.opts.get("style").kwargs["text_color"] == "red" + + # Turn on both desired gaps and result gaps simultaneously + plotter.show_desired_gaps = True + both_overlay = plotter._plot_clearance_gaps() + assert isinstance(both_overlay, hv.Overlay) + colors = { + el.opts.get("style").kwargs.get("color") + for el in both_overlay.values() + if isinstance(el, (hv.Points, hv.Segments)) + } + assert "blue" in colors and "red" in colors + + # Toggle off main show_gaps -> overlay is empty + plotter.show_gaps = False + off_overlay = plotter._plot_clearance_gaps() + assert isinstance(off_overlay, hv.Overlay) + finally: + settings.nice.machine_preset = orig_preset + settings.nice.mode = orig_mode def test_auto_enable_warm_start_on_direct_mode(): @@ -206,60 +341,73 @@ def __init__(self): self.config = WaveformConfiguration() editor = ShapeEditor(MockMainGui()) - editor.nice_settings.mode = NiceSettings.INVERSE_MODE - editor.use_previous_run = False - editor.communicator.converged = False + orig_mode = editor.nice_settings.mode + try: + editor.nice_settings.mode = NiceSettings.INVERSE_MODE + editor.use_previous_run = False + editor.communicator.converged = False - # 1. Switch to direct mode without any previous run -> remains False - editor.nice_settings.mode = NiceSettings.DIRECT_MODE - assert editor.use_previous_run is False + # 1. Switch to direct mode without any previous run -> remains False + editor.nice_settings.mode = NiceSettings.DIRECT_MODE + assert editor.use_previous_run is False - # 2. Switch back to inverse mode - editor.nice_settings.mode = NiceSettings.INVERSE_MODE - assert editor.use_previous_run is False + # 2. Switch back to inverse mode + editor.nice_settings.mode = NiceSettings.INVERSE_MODE + assert editor.use_previous_run is False - # 3. Previous run converges - editor.communicator.converged = True + # 3. Previous run converges + editor.communicator.converged = True - # 4. Switch to direct mode with converged run -> automatically True - editor.nice_settings.mode = NiceSettings.DIRECT_MODE - assert editor.use_previous_run is True + # 4. Switch to direct mode with converged run -> automatically True + editor.nice_settings.mode = NiceSettings.DIRECT_MODE + assert editor.use_previous_run is True + finally: + editor.nice_settings.mode = orig_mode def test_target_gap_pills_in_plasma_shape(): from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape shape = PlasmaShape() - shape.machine_preset = "ITER" - shape.show_gaps = True - - # No shape loaded yet -> empty string - assert shape._render_target_gap_pills() == "" - - # Load parameterized shape - shape.input_mode = shape.PARAMETERIZED_INPUT - shape._load_shape_from_params() - assert shape.has_shape is True - - # With show_gaps=True and ITER -> pills rendered - html = shape._render_target_gap_pills() - assert '
' in html - assert "g₄" in html - assert "g₅" in html - assert "Rₘᵢₙ" in html - assert "Rₘₐₓ" in html - - # With show_gaps=False -> no pills rendered - shape.show_gaps = False - assert shape._render_target_gap_pills() == "" - - # Switch to WEST -> WEST pills rendered - shape.show_gaps = True - shape.machine_preset = "WEST" - west_html = shape._render_target_gap_pills() - assert '
' in west_html - assert "UROG" in west_html - assert "EROG" in west_html + # The preset is shared, so it is set back for the other tests + orig_preset = settings.nice.machine_preset + try: + settings.nice.machine_preset = NiceSettings.PRESET_ITER + + # No shape loaded yet -> empty string + assert shape._render_target_gap_pills() == "" + + # Load parameterized shape + shape.input_mode = shape.PARAMETERIZED_INPUT + shape._load_shape_from_params() + assert shape.has_shape is True + + # Target gap pills rendered whenever shape is loaded + html = shape._render_target_gap_pills() + assert '
' in html + assert "g₄" in html + assert "g₅" in html + assert "Rₘᵢₙ" in html + assert "Rₘₐₓ" in html + + # Switch to WEST -> WEST pills rendered, for a WEST sized shape + settings.nice.machine_preset = NiceSettings.PRESET_WEST + shape.shape_params.param.update( + a=0.46, + center_r=2.54, + center_z=-0.02, + kappa=1.31, + delta=0.38, + rx=2.23, + zx=-0.62, + ) + shape._load_shape_from_params() + west_html = shape._render_target_gap_pills() + assert '
' in west_html + for symbol in ("UROG", "EROG", "LROG", "RIG", "TIG", "TOG", "dXlow", "dbaffle"): + assert symbol in west_html + finally: + settings.nice.machine_preset = orig_preset def test_gaps_input_mode_default_machine_gaps(): @@ -282,29 +430,157 @@ def test_gaps_input_mode_default_machine_gaps(): # Test PlasmaShape in Gaps mode shape = PlasmaShape() - shape.machine_preset = "ITER" - shape.input_mode = shape.GAP_INPUT - - # Default gap source is Default Gaps - assert shape.gap_source == shape.GAP_SOURCE_DEFAULT - assert shape.has_shape is True - assert len(shape.gaps) == 6 - assert len(shape.gap_ui) == 6 - - # Test changing a gap slider value - orig_val = shape.gap_ui[0].value - shape.gap_ui[0].value = 0.05 - assert math.isclose(shape.gaps[0].value, 0.05) - - # Test reset defaults - shape._on_reset_default_gaps() - assert math.isclose(shape.gaps[0].value, orig_val) - - # Test switching to WEST preset - shape.machine_preset = "WEST" - assert len(shape.gaps) == 6 - assert any("UROG" in g.name for g in shape.gaps) - - # Test switching to IDS mode - shape.gap_source = shape.GAP_SOURCE_IDS - assert shape.gap_source == shape.GAP_SOURCE_IDS + # The preset is shared, so it is set back for the other tests + orig_preset = settings.nice.machine_preset + try: + settings.nice.machine_preset = NiceSettings.PRESET_ITER + shape.input_mode = shape.GAP_INPUT + + # Default gap source is Default Gaps + assert shape.gap_source == shape.GAP_SOURCE_DEFAULT + assert shape.has_shape is True + assert len(shape.gaps) == 6 + assert len(shape.gap_ui) == 6 + + # Test changing a gap slider value + orig_val = shape.gap_ui[0].value + shape.gap_ui[0].value = 0.05 + assert math.isclose(shape.gaps[0].value, 0.05) + + # Test reset defaults + shape._on_reset_default_gaps() + assert math.isclose(shape.gaps[0].value, orig_val) + + # Test switching to WEST preset + settings.nice.machine_preset = NiceSettings.PRESET_WEST + assert len(shape.gaps) == 6 + assert any("UROG" in g.name for g in shape.gaps) + + # Test switching to IDS mode + shape.gap_source = shape.GAP_SOURCE_IDS + assert shape.gap_source == shape.GAP_SOURCE_IDS + pane_ids = shape.panel[2]._pane + assert "Reset Defaults" not in str(pane_ids.objects[1]) + assert shape.has_shape is False + assert len(shape.gaps) == 0 + + # Test loading equilibrium IDS with gaps + from unittest.mock import MagicMock, patch + + mock_gap1 = MagicMock(r=4.0, z=-3.0, name="Gap1", angle=1.5, value=0.1) + mock_gap2 = MagicMock(r=5.0, z=-4.0, name="Gap2", angle=2.0, value=0.2) + mock_slice = MagicMock() + mock_slice.boundary.gap = [mock_gap1, mock_gap2] + mock_eq = MagicMock() + mock_eq.time_slice = [mock_slice] + mock_eq.time.tolist.return_value = [0.0] + + mock_entry = MagicMock() + mock_entry.__enter__.return_value = mock_entry + mock_entry.get.return_value = mock_eq + mock_entry.get_slice.return_value = mock_eq + + with patch("imas.DBEntry", return_value=mock_entry): + shape.input_gaps.uri = "imas:test" + shape.input_gaps.param.trigger("load") + + assert len(shape.gaps) == 2 + assert shape.has_shape is True + assert len(shape.gap_ui.objects) == 2 + + # Test switching back to Default Gaps + shape.gap_source = shape.GAP_SOURCE_DEFAULT + pane_def = shape.panel[2]._pane + assert "Reset Defaults" in str(pane_def.objects[1]) + assert len(shape.gaps) == 6 + assert shape.has_shape is True + finally: + # The shape keeps following the shared preset after this test, and would + # read this URI without the mock as soon as a preset without defaults is set + shape.input_gaps.uri = "" + settings.nice.machine_preset = orig_preset + + +def test_gaps_auto_switch_on_run_finished(): + import holoviews as hv + + hv.extension("bokeh") + from waveform_editor.configuration import WaveformConfiguration + from waveform_editor.gui.shape_editor.shape_editor import ShapeEditor + + class MockMainGui: + def __init__(self): + self.config = WaveformConfiguration() + + editor = ShapeEditor(MockMainGui()) + plotter = editor.nice_plotter + + # Initially, show_gaps=True, show_desired_gaps=True, show_result_gaps=False + assert plotter.show_gaps is True + assert plotter.show_desired_gaps is True + assert plotter.show_result_gaps is False + + # Simulate NICE run finished successfully + editor._on_nice_run_finished(True) + assert plotter.show_result_gaps is True + assert plotter.show_desired_gaps is False + + # If show_gaps is turned off, subsettings are preserved and not auto-switched + plotter.show_gaps = False + plotter.show_desired_gaps = True + plotter.show_result_gaps = False + editor._on_nice_run_finished(True) + assert plotter.show_desired_gaps is True + assert plotter.show_result_gaps is False + + +def test_settings_modal_gap_subsettings(): + import holoviews as hv + + hv.extension("bokeh") + from waveform_editor.configuration import WaveformConfiguration + from waveform_editor.gui.shape_editor.settings_modal import SettingsModal + from waveform_editor.gui.shape_editor.shape_editor import ShapeEditor + + class MockMainGui: + def __init__(self): + self.config = WaveformConfiguration() + + editor = ShapeEditor(MockMainGui()) + modal = SettingsModal(nice_plotter=editor.nice_plotter) + + display_tab = modal.tabs[0] + # The gap options are nested under "Show gaps", in the visibility section + gap_options = display_tab[1][1] + + # Shown while show_gaps is on, hidden while it is off + assert gap_options.visible is True + editor.nice_plotter.show_gaps = False + assert gap_options.visible is False + editor.nice_plotter.show_gaps = True + assert gap_options.visible is True + + +def test_custom_machine_only_takes_gaps_from_an_ids(): + from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape + + # The preset is shared, so it is set back for the other tests + orig_preset = settings.nice.machine_preset + try: + settings.nice.machine_preset = NiceSettings.PRESET_ITER + shape = PlasmaShape() + shape.input_mode = shape.GAP_INPUT + assert len(shape.gaps) == 6 + + # A custom machine has no default gaps, so there is no choice of source + settings.nice.machine_preset = NiceSettings.PRESET_CUSTOM + assert shape._active_gap_source == shape.GAP_SOURCE_IDS + assert shape.gaps == [] + assert shape.gap_source_radio not in shape._panel_gap_options().objects + + # The choice for a machine with defaults is kept + settings.nice.machine_preset = NiceSettings.PRESET_WEST + assert shape._active_gap_source == shape.GAP_SOURCE_DEFAULT + assert any("UROG" in gap.name for gap in shape.gaps) + finally: + settings.nice.machine_preset = orig_preset diff --git a/waveform_editor/gui/shape_editor/nice_plotter.py b/waveform_editor/gui/shape_editor/nice_plotter.py index bf0ab71c..a9e8baff 100644 --- a/waveform_editor/gui/shape_editor/nice_plotter.py +++ b/waveform_editor/gui/shape_editor/nice_plotter.py @@ -14,16 +14,8 @@ from waveform_editor.gui.shape_editor.plasma_properties import PlasmaProperties from waveform_editor.gui.shape_editor.plasma_shape import PlasmaShape from waveform_editor.settings import NiceSettings, settings -from waveform_editor.shape_editor.iter_gaps import ( - compute_gap_geometry as compute_iter_gap_geometry, -) -from waveform_editor.shape_editor.iter_gaps import ( - extract_separatrix_from_time_slice, -) +from waveform_editor.shape_editor import iter_gaps, west_gaps from waveform_editor.shape_editor.nice_integration import NiceIntegration -from waveform_editor.shape_editor.west_gaps import ( - compute_gap_geometry as compute_west_gap_geometry, -) matplotlib.use("Agg") logger = logging.getLogger(__name__) @@ -133,6 +125,8 @@ class NicePlotter(Viewer): show_separatrix = param.Boolean(default=True, label="Show separatrix") show_desired_shape = param.Boolean(default=True, label="Show desired shape") show_gaps = param.Boolean(default=True, label="Show gaps") + show_desired_gaps = param.Boolean(default=True, label="Show desired shape gaps") + show_result_gaps = param.Boolean(default=False, label="Show resulting shape gaps") # Renderer of the editable points, set once the plot is first rendered _points_renderer = None @@ -750,6 +744,8 @@ def _plot_xo_points(self): @pn.depends( "communicator.equilibrium", "show_gaps", + "show_desired_gaps", + "show_result_gaps", "nice_settings.machine_preset", "plasma_shape.shape_updated", "plasma_shape.input_mode", @@ -766,57 +762,61 @@ def _plot_clearance_gaps(self): if not self.show_gaps: return self._empty_gaps_overlay() - eq = self.communicator.equilibrium - x_points = [] - if eq is not None and len(eq.time_slice) > 0: - outline = eq.time_slice[0].boundary.outline - r, z = outline.r, outline.z - x_points = [ - (float(node.r), float(node.z)) - for node in eq.time_slice[0].contour_tree.node - if int(node.critical_type) == 1 - ] - elif not self.nice_settings.is_direct_mode and self.plasma_shape.has_shape: - r, z = self.plasma_shape.outline_r, self.plasma_shape.outline_z - else: - return self._empty_gaps_overlay() + preset = self.nice_settings.machine_preset + overlays = [] - if r is None or len(r) == 0: - return self._empty_gaps_overlay() + if ( + self.show_desired_gaps + and not self.nice_settings.is_direct_mode + and self.plasma_shape.has_shape + ): + r_des = self.plasma_shape.outline_r + z_des = self.plasma_shape.outline_z + if r_des is not None and len(r_des) > 0: + if preset == NiceSettings.PRESET_ITER: + des_items = iter_gaps.compute_gap_geometry(r_des, z_des) + elif preset == NiceSettings.PRESET_WEST: + des_items = west_gaps.compute_gap_geometry( + r_des, z_des, self.plasma_shape.desired_x_points + ) + else: + des_items = [] + if des_items: + overlays.append(self._build_gaps_overlay(des_items, color="blue")) - gap_items = [] - preset = self.nice_settings.machine_preset - if preset == NiceSettings.PRESET_ITER: - ts = eq.time_slice[0] if eq is not None and len(eq.time_slice) > 0 else None - ma = ( - getattr(getattr(ts, "global_quantities", None), "magnetic_axis", None) - if ts - else None - ) - gap_items = compute_iter_gap_geometry( - r, - z, - separatrix_contour=extract_separatrix_from_time_slice(ts) - if ts - else None, - magnetic_axis=(float(ma.r), float(ma.z)) if ma else None, - ) - elif preset == NiceSettings.PRESET_WEST: - gap_items = compute_west_gap_geometry(r, z, x_points) + eq = self.communicator.equilibrium + if self.show_result_gaps and eq is not None and len(eq.time_slice) > 0: + ts = eq.time_slice[0] + if len(ts.boundary.outline.r) > 0: + if preset == NiceSettings.PRESET_ITER: + res_items = iter_gaps.compute_gap_geometry( + **iter_gaps.gap_inputs(ts) + ) + elif preset == NiceSettings.PRESET_WEST: + res_items = west_gaps.compute_gap_geometry( + **west_gaps.gap_inputs(ts) + ) + else: + res_items = [] + if res_items: + overlays.append(self._build_gaps_overlay(res_items, color="red")) - if not gap_items: + if not overlays: return self._empty_gaps_overlay() - return self._build_gaps_overlay(gap_items) + combined = overlays[0] + for ov in overlays[1:]: + combined = combined * ov + return combined - def _build_gaps_overlay(self, items): + def _build_gaps_overlay(self, items, color="blue"): vdims = ["name", "symbol", "distance_str"] pts = hv.Points( [{"r": it["r_orig"], "z": it["z_orig"], **it} for it in items], kdims=["r", "z"], vdims=vdims, ).opts( - color="#e66101", + color=color, size=8, marker="diamond", hover_tooltips=[ @@ -831,7 +831,7 @@ def _build_gaps_overlay(self, items): kdims=["r", "z"], vdims=vdims, ).opts( - color="#e66101", + color=color, size=5, marker="circle", hover_tooltips=[("Gap", "@name (@symbol)"), ("Distance", "@distance_str")], @@ -851,7 +851,7 @@ def _build_gaps_overlay(self, items): kdims=["r0", "z0", "r1", "z1"], vdims=vdims, ).opts( - color="#e66101", + color=color, line_dash="dashed", line_width=2, hover_tooltips=[("Gap", "@name (@symbol)"), ("Distance", "@distance_str")], @@ -867,7 +867,7 @@ def _build_gaps_overlay(self, items): vdims=["text"], ).opts( text_font_size="10pt", - text_color="#e66101", + text_color=color, text_baseline="bottom", text_align="center", text_font_style="bold", diff --git a/waveform_editor/gui/shape_editor/plasma_shape.py b/waveform_editor/gui/shape_editor/plasma_shape.py index 199dce50..37543196 100644 --- a/waveform_editor/gui/shape_editor/plasma_shape.py +++ b/waveform_editor/gui/shape_editor/plasma_shape.py @@ -65,6 +65,13 @@ } +# The gaps a machine is designed with; a custom machine only takes gaps from an IDS +MACHINE_DEFAULT_GAPS = { + NiceSettings.PRESET_ITER: get_default_iter_gaps, + NiceSettings.PRESET_WEST: get_default_west_gaps, +} + + class PlasmaShapeParams(Viewer): """Helper class containing parameters to parameterize the plasma shape.""" @@ -462,13 +469,10 @@ class PlasmaShape(Viewer): has_shape = param.Boolean(doc="Whether a plasma shape is loaded.") shape_updated = param.Event(doc="Triggered whenever the plasma shape updates.") - show_gaps = param.Boolean(default=True, doc="Whether clearance gaps are displayed.") - machine_preset = param.String( - default="", doc="Active machine preset (e.g. ITER, WEST)." - ) + nice_settings = param.ClassSelector(class_=NiceSettings) def __init__(self): - super().__init__() + super().__init__(nice_settings=settings.nice) def _indicator(tooltip): return WarningIndicator( @@ -507,6 +511,7 @@ def _indicator(tooltip): self.input_outline, self.outline_indicator, margin=(10, 20, 0, 20), + sizing_mode="stretch_width", ), ), self.PARAMETERIZED_INPUT: ( @@ -520,11 +525,16 @@ def _indicator(tooltip): self.WEIGHTED_POINTS_INPUT: ( self._load_shape_from_weighted_points, lambda: pn.Row( - self.weighted_points_table, self.weighted_points_indicator + self.weighted_points_table, + self.weighted_points_indicator, + sizing_mode="stretch_width", ), ), } - self.gap_ui = pn.Column(visible=self.param.input_mode.rx() == self.GAP_INPUT) + self.gap_ui = pn.Column( + visible=self.param.input_mode.rx() == self.GAP_INPUT, + sizing_mode="stretch_width", + ) self.radio_box = pn.widgets.RadioButtonGroup( options={ "Equilibrium\nIDS Outline": self.EQUILIBRIUM_INPUT, @@ -548,8 +558,9 @@ def _indicator(tooltip): self.panel = pn.Column( self.radio_box, self.gap_pills_pane, - self._panel_shape_options, + pn.panel(self._panel_shape_options, sizing_mode="stretch_width"), self.gap_ui, + sizing_mode="stretch_width", ) self.outline_r = None self.outline_z = None @@ -562,7 +573,7 @@ def _indicator(tooltip): self.param_weights = None self.shape_params.extra_points_table = self.weighted_points_table # Last, because applying a shape fires the watchers above - settings.nice.param.watch( + self.nice_settings.param.watch( self.shape_params.apply_machine_shape, "machine_preset" ) self.shape_params.apply_machine_shape() @@ -574,7 +585,7 @@ def _indicator(tooltip): "weighted_points_table.param", "input_mode", "gap_source", - "machine_preset", + "nice_settings.machine_preset", watch=True, ) def _set_plasma_shape(self): @@ -592,25 +603,40 @@ def _set_plasma_shape(self): self.has_shape = False self.param.trigger("shape_updated") - @pn.depends("gap_source", "machine_preset") + @property + def _active_gap_source(self): + """Where the gaps come from: the chosen source, except for a machine without + default gaps, which can only take them from an IDS. The choice is kept, for + when a machine with defaults is selected again.""" + if self.nice_settings.machine_preset in MACHINE_DEFAULT_GAPS: + return self.gap_source + return self.GAP_SOURCE_IDS + def _panel_gap_options(self): """Render controls for Gaps mode: gap source toggle and appropriate loaders.""" - preset = self.machine_preset or "ITER" - if self.gap_source == self.GAP_SOURCE_DEFAULT: - return pn.Column( - self.gap_source_radio, - pn.Row( - pn.pane.Markdown( - f"**{preset} Default Gap Definitions**", - margin=(8, 10, 0, 20), - ), - self.reset_default_gaps_button, - sizing_mode="stretch_width", + preset = self.nice_settings.machine_preset + if self._active_gap_source == self.GAP_SOURCE_DEFAULT: + content = pn.Row( + pn.pane.Markdown( + f"**{preset} Default Gap Definitions**", + margin=(8, 10, 0, 20), ), + self.reset_default_gaps_button, + sizing_mode="stretch_width", + ) + else: + content = pn.Row( + self.input_gaps, + self.gap_indicator, + margin=(10, 20, 0, 20), + sizing_mode="stretch_width", ) + has_defaults = preset in MACHINE_DEFAULT_GAPS return pn.Column( - self.gap_source_radio, - pn.Row(self.input_gaps, self.gap_indicator, margin=(5, 20, 0, 20)), + # Without defaults there is nothing to choose between + *([self.gap_source_radio] if has_defaults else []), + content, + sizing_mode="stretch_width", ) def _on_reset_default_gaps(self, event=None): @@ -641,13 +667,9 @@ def _load_shape_from_gaps(self): """Load plasma boundary outline from default machine gaps or IDS equilibrium.""" self.gaps = [] - if self.gap_source == self.GAP_SOURCE_DEFAULT: - self.gaps = ( - get_default_west_gaps() - if self.machine_preset == "WEST" - else get_default_iter_gaps() - ) - elif self.gap_source == self.GAP_SOURCE_IDS and self.input_gaps.uri: + if self._active_gap_source == self.GAP_SOURCE_DEFAULT: + self.gaps = MACHINE_DEFAULT_GAPS[self.nice_settings.machine_preset]() + elif self.input_gaps.uri: try: with imas.DBEntry(self.input_gaps.uri, "r") as entry: equilibrium = entry.get_slice( @@ -723,6 +745,13 @@ def _create_gap_ui(self): self.gap_ui.extend(new_gap_ui) + @property + def desired_x_points(self): + """The x-points of the desired shape, which only a parameterized shape has.""" + if self.input_mode != self.PARAMETERIZED_INPUT: + return [] + return [(self.shape_params.rx, self.shape_params.zx)] + @property def uses_weighted_points(self): """Whether the weighted points table feeds the current input mode.""" @@ -778,19 +807,18 @@ def _append_extra_points(self): self.outline_r = self.outline_r + extra_r self.outline_z = self.outline_z + extra_z - @param.depends("input_mode") + @param.depends("input_mode", "gap_source", "nice_settings.machine_preset") def _panel_shape_options(self): _, panel_factory = self._mode_config[self.input_mode] return panel_factory() @param.depends( - "show_gaps", "machine_preset", "shape_updated", "has_shape", "input_mode" + "nice_settings.machine_preset", "shape_updated", "has_shape", "input_mode" ) def _render_target_gap_pills(self): - """Render target shape metric pills when show_gaps is enabled.""" + """Render target shape gap metric pills below the input mode selector.""" if ( - not self.show_gaps - or not self.has_shape + not self.has_shape or self.outline_r is None or self.outline_z is None or len(self.outline_r) == 0 @@ -806,7 +834,7 @@ def _chip(symbol, display, tooltip): f'{display}' ) - if self.machine_preset == "ITER": + if self.nice_settings.machine_preset == NiceSettings.PRESET_ITER: gaps = compute_iter_gaps(self.outline_r, self.outline_z) if self.input_mode == self.GAP_INPUT and self.gaps: for g in self.gaps: @@ -828,11 +856,9 @@ def _chip(symbol, display, tooltip): ) chips.append(_chip(symbol, display, tooltip)) - elif self.machine_preset == "WEST": - x_points = [] - if self.input_mode == self.PARAMETERIZED_INPUT: - x_points = [(self.shape_params.rx, self.shape_params.zx)] - elif self.input_mode == self.GAP_INPUT and self.gaps: + elif self.nice_settings.machine_preset == NiceSettings.PRESET_WEST: + x_points = self.desired_x_points + if self.input_mode == self.GAP_INPUT and self.gaps: x_points = [ (g.r_sep, g.z_sep) for g in self.gaps diff --git a/waveform_editor/gui/shape_editor/settings_modal.py b/waveform_editor/gui/shape_editor/settings_modal.py index b785f25b..0b257b73 100644 --- a/waveform_editor/gui/shape_editor/settings_modal.py +++ b/waveform_editor/gui/shape_editor/settings_modal.py @@ -212,6 +212,21 @@ def _build_modal(self): "Heatmap Detail", ["heatmap_alpha"], self.nice_plotter.param.show_heatmap ) + # Which gaps "Show gaps" shows, nested under it and only there while it is on + gap_options = pn.Column( + pn.widgets.Checkbox.from_param( + self.nice_plotter.param.show_desired_gaps, + name="Desired shape", + visible=self.nice_settings.param.is_inverse_mode.rx(), + ), + pn.widgets.Checkbox.from_param( + self.nice_plotter.param.show_result_gaps, name="Resulting shape" + ), + visible=self.nice_plotter.param.show_gaps, + styles={"border-left": "2px solid #dee2e6", "padding-left": "8px"}, + margin=(0, 0, 4, 22), + ) + display_content = pn.Column( self._section_label("Visibility"), self._settings_section( @@ -238,6 +253,7 @@ def _build_modal(self): }, }, ), + gap_options, ), self._contour_detail, self._heatmap_detail, diff --git a/waveform_editor/gui/shape_editor/shape_editor.py b/waveform_editor/gui/shape_editor/shape_editor.py index f82955e9..29b0467b 100644 --- a/waveform_editor/gui/shape_editor/shape_editor.py +++ b/waveform_editor/gui/shape_editor/shape_editor.py @@ -20,22 +20,10 @@ from waveform_editor.gui.shape_editor.waveform_sync import WaveformSync from waveform_editor.gui.util import set_xml_parameter from waveform_editor.settings import NiceSettings, settings -from waveform_editor.shape_editor.iter_gaps import ( - GAP_METADATA as ITER_GAP_METRICS, -) -from waveform_editor.shape_editor.iter_gaps import ( - compute_gaps as compute_iter_gaps, -) -from waveform_editor.shape_editor.iter_gaps import ( - extract_separatrix_from_time_slice, -) +from waveform_editor.shape_editor import iter_gaps, west_gaps +from waveform_editor.shape_editor.iter_gaps import GAP_METADATA as ITER_GAP_METRICS from waveform_editor.shape_editor.nice_integration import NiceIntegration -from waveform_editor.shape_editor.west_gaps import ( - WEST_METADATA as WEST_GAP_METRICS, -) -from waveform_editor.shape_editor.west_gaps import ( - compute_gaps as compute_west_gaps, -) +from waveform_editor.shape_editor.west_gaps import WEST_METADATA as WEST_GAP_METRICS logger = logging.getLogger(__name__) @@ -93,16 +81,6 @@ def __init__(self, main_gui): ) self.run_select.param.watch(self._restore_run, "value") self.nice_settings = settings.nice - self.plasma_shape.machine_preset = self.nice_settings.machine_preset - self.plasma_shape.show_gaps = self.nice_plotter.show_gaps - self.nice_settings.param.watch( - lambda e: setattr(self.plasma_shape, "machine_preset", e.new), - "machine_preset", - ) - self.nice_plotter.param.watch( - lambda e: setattr(self.plasma_shape, "show_gaps", e.new), - "show_gaps", - ) self.xml_text = ( importlib.resources.files("waveform_editor.shape_editor.xml_param") @@ -395,11 +373,16 @@ def _fill_equilibrium(self, equilibrium): def _on_nice_run_finished(self, success): if success: - pn.state.notifications.success("NICE run complete.") + if pn.state.notifications: + pn.state.notifications.success("NICE run complete.") + if self.nice_plotter.show_gaps: + self.nice_plotter.show_result_gaps = True + self.nice_plotter.show_desired_gaps = False else: - pn.state.notifications.error( - "NICE did not converge. Check the terminal for details." - ) + if pn.state.notifications: + pn.state.notifications.error( + "NICE did not converge. Check the terminal for details." + ) def _has_valid_boundary(self): """Check that the desired boundary fits in the fixed size array NICE reads it @@ -517,21 +500,9 @@ def _iter_gaps(self, time_slice): """The gaps of an ITER plasma in centimetres (Rmin/Rmax in metres).""" if self.nice_settings.machine_preset != NiceSettings.PRESET_ITER: return {} - outline = time_slice.boundary.outline - if outline.r is None or len(outline.r) == 0: - return {} - ma = getattr( - getattr(time_slice, "global_quantities", None), "magnetic_axis", None - ) - axis = (float(ma.r), float(ma.z)) if ma else None - gaps = compute_iter_gaps( - outline.r, - outline.z, - separatrix_contour=extract_separatrix_from_time_slice(time_slice), - magnetic_axis=axis, - ) + gaps = iter_gaps.compute_gaps(**iter_gaps.gap_inputs(time_slice)) return { - name: val * 100 if name in ("gap1", "gap2", "gap4", "gap5") else val + name: val * 100 if ITER_GAP_METRICS[name][1] == "cm" else val for name, val in gaps.items() if val is not None } @@ -540,13 +511,7 @@ def _west_gaps(self, time_slice): """The gaps of a WEST plasma in centimetres.""" if self.nice_settings.machine_preset != NiceSettings.PRESET_WEST: return {} - x_points = [ - (float(node.r), float(node.z)) - for node in getattr(getattr(time_slice, "contour_tree", None), "node", []) - if int(node.critical_type) == 1 - ] - outline = time_slice.boundary.outline - gaps = compute_west_gaps(outline.r, outline.z, x_points) + gaps = west_gaps.compute_gaps(**west_gaps.gap_inputs(time_slice)) return {name: gap * 100 for name, gap in gaps.items()} def _update_metrics(self): diff --git a/waveform_editor/gui/styles/property_card.css b/waveform_editor/gui/styles/property_card.css index 7309eeaf..cf57e62e 100644 --- a/waveform_editor/gui/styles/property_card.css +++ b/waveform_editor/gui/styles/property_card.css @@ -3,11 +3,14 @@ border-radius: 8px; padding: 8px 12px 12px 12px; display: block; + width: 100%; + box-sizing: border-box; } :host(.ids-source-card) { border-color: #2a5a8c; background: rgba(13, 110, 253, 0.08); + flex: 1; } .bk-btn.bk-btn-primary.bk-active { diff --git a/waveform_editor/shape_editor/iter_gaps.py b/waveform_editor/shape_editor/iter_gaps.py index fd279ca2..aff0b5f1 100644 --- a/waveform_editor/shape_editor/iter_gaps.py +++ b/waveform_editor/shape_editor/iter_gaps.py @@ -1,17 +1,25 @@ """The gaps between an ITER plasma and the first wall, as defined in DINA. -The gap definitions and measurement points are taken from DINA-IMAS: -- gaps 1, 2, 4, 5 are the minimum distances from fixed wall coordinates to the - plasma boundary outline. -- Rmin and Rmax are the inboard and outboard radial extents (min and max R) of - the plasma boundary. +The definitions follow DINA-IMAS (src/scenario/g_gaps_rus.f, tools/GUI/captions.py): +- gaps 1 and 2 are the minimum distances from the reference strike points to the + separatrix, which DINA only measures in the divertor phase, to the divertor legs; +- gaps 4 and 5 are the minimum distances from points on the wall to the boundary, + towards the magnetic axis; +- Rmin and Rmax are the radii of the innermost and outermost boundary points, measured + from r=0 at the height of those points. """ import math import numpy as np -from waveform_editor.shape_editor.plasma_shape_calc import Gap, closest_outline_point +from waveform_editor.shape_editor.plasma_shape_calc import ( + Gap, + closest_outline_point, + closest_point_facing, + geometric_centre, + psi_contour, +) # Fixed gap measurement points on the wall, in metres (r, z): GAP_POINTS = { @@ -21,11 +29,6 @@ "gap5": (5.3315, 4.5804), } -WALL_REFERENCE_POINTS = { - "Rmin": (4.0599, 0.7777), - "Rmax": (8.2806, 0.4665), -} - GAP_METADATA = { "gap1": ("g₁", "cm", "Inner divertor leg gap"), "gap2": ("g₂", "cm", "Outer divertor leg gap"), @@ -40,59 +43,39 @@ def _distance_to_outline(point, r, z, is_closed=True): return closest_outline_point(point, r, z, is_closed=is_closed)[1] +def gap_inputs(time_slice): + """The arguments to compute the gaps of a solved equilibrium with. + + Args: + time_slice: The equilibrium time slice. + + Returns: + Dict of the keyword arguments of compute_gaps and compute_gap_geometry. + """ + axis = time_slice.global_quantities.magnetic_axis + return { + "outline_r": time_slice.boundary.outline.r, + "outline_z": time_slice.boundary.outline.z, + "separatrix_contour": extract_separatrix_from_time_slice(time_slice), + "magnetic_axis": (float(axis.r), float(axis.z)), + } + + def extract_separatrix_from_time_slice(time_slice): - """Extract separatrix isocontour segments from an equilibrium time slice.""" + """The contour of the boundary flux of an equilibrium, legs included.""" if time_slice is None: return [] - - b_psi = getattr( - getattr(time_slice, "global_quantities", None), "psi_boundary", None - ) - if b_psi is None or float(b_psi) == -9e40: - b_psi = getattr(getattr(time_slice, "boundary", None), "psi", None) - if b_psi is None or float(b_psi) == -9e40: + psi_boundary = time_slice.global_quantities.psi_boundary + if not _is_filled(psi_boundary): + psi_boundary = time_slice.boundary.psi + if not _is_filled(psi_boundary): return [] - b_psi = float(b_psi) - - import matplotlib.pyplot as plt - - def _contour_segs(cs): - segs = [np.asarray(seg) for seg in cs.allsegs[0] if len(seg) > 1] - plt.close(fig) - return segs - - # 1. Try GGD (NICE output) - if hasattr(time_slice, "ggd") and len(time_slice.ggd) > 0: - try: - g = time_slice.ggd[0] - r, z, psi = g.r[0].values, g.z[0].values, g.psi[0].values - if r and z and psi: - fig, ax = plt.subplots() - segs = _contour_segs(ax.tricontour(r, z, psi, levels=[b_psi])) - if segs: - return segs - except Exception: - pass - - # 2. Try profiles_2d (IMAS grid) - if hasattr(time_slice, "profiles_2d") and len(time_slice.profiles_2d) > 0: - try: - p2d = time_slice.profiles_2d[0] - r, z, psi = ( - np.asarray(p2d.grid.dim1), - np.asarray(p2d.grid.dim2), - np.asarray(p2d.psi), - ) - if len(r) > 0 and len(z) > 0 and psi.size > 0: - psi_grid = psi.T if psi.shape == (len(r), len(z)) else psi - fig, ax = plt.subplots() - segs = _contour_segs(ax.contour(r, z, psi_grid, levels=[b_psi])) - if segs: - return segs - except Exception: - pass + return psi_contour(time_slice, float(psi_boundary)) - return [] + +def _is_filled(value): + """Whether an IDS float is filled.""" + return value is not None and float(value) != -9e40 def _normalize_contour_segments(separatrix_contour): @@ -135,126 +118,122 @@ def _dina_divertor_sign(k, r_target, z_target, r_orig, z_orig, magnetic_axis): return (p / abs(p)) * s_vecpro -def _closest_point_on_segments(point, segments): - """Find the closest point to `point` across multiple polyline segments.""" - min_dist, best_target = float("inf"), None - for seg in segments: - target, dist = closest_outline_point( - point, seg[:, 0], seg[:, 1], is_closed=False - ) - if dist < min_dist: - min_dist, best_target = dist, target - return best_target, min_dist +def is_diverted(segments): + """Whether a separatrix has divertor legs, which leave the domain, so that a piece + of its contour does not close on itself.""" + return any(np.linalg.norm(seg[0] - seg[-1]) > 1e-3 for seg in segments) def compute_gaps(outline_r, outline_z, separatrix_contour=None, magnetic_axis=None): - """The gaps and radial extent of an ITER plasma, in metres.""" - if outline_r is None or outline_z is None or len(outline_r) == 0: - return {} - - r = np.asarray(outline_r, dtype=float) - z = np.asarray(outline_z, dtype=float) - - gaps = { - "gap4": closest_outline_point(GAP_POINTS["gap4"], r, z)[1], - "gap5": closest_outline_point(GAP_POINTS["gap5"], r, z)[1], - "Rmin": float(np.min(r)), - "Rmax": float(np.max(r)), - } - - segments = _normalize_contour_segments(separatrix_contour) - has_legs = len(segments) > 0 or float(np.min(z)) <= -3.6 - if not segments and has_legs: - segments = [np.column_stack([r, z])] - - if has_legs: - for k, name in enumerate(("gap1", "gap2"), start=1): - pt = GAP_POINTS[name] - target, dist = _closest_point_on_segments(pt, segments) - if target is not None: - sign = _dina_divertor_sign( - k, target[0], target[1], pt[0], pt[1], magnetic_axis - ) - gaps[name] = dist * sign - else: - gaps[name] = None - else: - gaps["gap1"] = None - gaps["gap2"] = None - - return gaps + """The gaps and radial extent of an ITER plasma, in metres. + + Args: + outline_r: Radial coordinates of the plasma boundary. + outline_z: Height coordinates of the plasma boundary. + separatrix_contour: The separatrix with its legs, as pieces of (N, 2) arrays, + which only a solved equilibrium has. Without legs there are no gaps 1, 2. + magnetic_axis: The (r, z) of the magnetic axis, the centre of the outline is + used without one. + + Returns: + Dict of gap name to distance in metres, None for gaps 1, 2 without legs. + """ + items = compute_gap_geometry( + outline_r, outline_z, separatrix_contour, magnetic_axis + ) + return {item["key"]: item["distance"] for item in items} def compute_gap_geometry( outline_r, outline_z, separatrix_contour=None, magnetic_axis=None ): - """Compute measuring points and distance segments for ITER gaps.""" + """The gaps of an ITER plasma, as computed in DINA, together with the points they + are measured between. + + Args: + outline_r: Radial coordinates of the plasma boundary. + outline_z: Height coordinates of the plasma boundary. + separatrix_contour: The separatrix with its legs, as pieces of (N, 2) arrays, + which only a solved equilibrium has. Without legs there are no gaps 1, 2. + magnetic_axis: The (r, z) of the magnetic axis, the centre of the outline is + used without one. + + Returns: + List of dicts, one per gap. + """ if outline_r is None or outline_z is None or len(outline_r) == 0: return [] - r, z = np.asarray(outline_r, float), np.asarray(outline_z, float) + axis = magnetic_axis[:2] if magnetic_axis is not None else geometric_centre(r, z) items = [] - segments = _normalize_contour_segments(separatrix_contour) or [ - np.column_stack([r, z]) - ] - for k, name in enumerate(("gap1", "gap2"), start=1): - r_orig, z_orig = GAP_POINTS[name] - target, dist = _closest_point_on_segments((r_orig, z_orig), segments) - if target is not None: - sign = _dina_divertor_sign( - k, target[0], target[1], r_orig, z_orig, magnetic_axis - ) - symbol, _, full_name = GAP_METADATA[name] - items.append( - { - "name": full_name, - "symbol": symbol, - "r_orig": r_orig, - "z_orig": z_orig, - "r_target": target[0], - "z_target": target[1], - "distance": dist, - "distance_str": f"{dist * sign * 100:+.2f} cm", - } - ) - - ref_points = {**GAP_POINTS, **WALL_REFERENCE_POINTS} - min_r, max_r = float(np.min(r)), float(np.max(r)) - for name in ("gap4", "gap5", "Rmin", "Rmax"): - r_orig, z_orig = ref_points[name] - (r_target, z_target), dist = closest_outline_point((r_orig, z_orig), r, z) - symbol, _, full_name = GAP_METADATA[name] - dist_str = ( - f"Rₘᵢₙ = {min_r:.4g} m (gap: {dist * 100:.3g} cm)" - if name == "Rmin" - else f"Rₘₐₓ = {max_r:.4g} m (gap: {dist * 100:.3g} cm)" - if name == "Rmax" - else f"{dist * 100:.3g} cm" - ) + def add(key, orig, target, distance, distance_str): + symbol, _, name = GAP_METADATA[key] items.append( { - "name": full_name, + "key": key, + "name": name, "symbol": symbol, - "r_orig": r_orig, - "z_orig": z_orig, - "r_target": r_target, - "z_target": z_target, - "distance": dist, - "distance_str": dist_str, + "r_orig": float(orig[0]), + "z_orig": float(orig[1]), + "r_target": float(target[0]), + "z_target": float(target[1]), + "distance": distance, + "distance_str": distance_str, } ) + # A closed boundary has no legs, and the distance from a strike point to it would + # be to the wrong part of the plasma, so the reference point is shown on its own + segments = _normalize_contour_segments(separatrix_contour) + diverted = is_diverted(segments) + for k, key in enumerate(("gap1", "gap2"), start=1): + point = GAP_POINTS[key] + best = None + for seg in segments if diverted else []: + target, distance = closest_outline_point( + point, seg[:, 0], seg[:, 1], is_closed=False + ) + if best is None or distance < best[1]: + best = target, distance + if best is None: + add(key, point, point, None, "— (no divertor legs)") + continue + target, distance = best + distance *= _dina_divertor_sign(k, *target, *point, axis) + add(key, point, target, distance, f"{distance * 100:+.3g} cm") + + for key in ("gap4", "gap5"): + point = GAP_POINTS[key] + target, distance = closest_point_facing(point, r, z, axis) + if target is not None: + add(key, point, target, distance, f"{distance * 100:.3g} cm") + + for key, index in (("Rmin", int(np.argmin(r))), ("Rmax", int(np.argmax(r)))): + extreme = (r[index], z[index]) + add(key, (0.0, z[index]), extreme, float(r[index]), f"{r[index]:.4g} m") + return items +def _dina_divertor_sign(k, r_target, z_target, r_orig, z_orig, magnetic_axis): + """The sign DINA gives gaps 1 and 2 (g_gaps_rus.f), from the side of the axis the + closest point is on.""" + rmag, zmag = magnetic_axis + vecpro = (r_target - rmag) * (z_orig - zmag) - (z_target - zmag) * (r_orig - rmag) + s_vecpro = 1.0 if vecpro >= 0 else -1.0 + p = k - 1.5 + return (p / abs(p)) * s_vecpro + + def get_default_iter_gaps(): """Return default Gap definitions for ITER based on DINA coordinates.""" return [ - Gap("Inner divertor leg (g₁)", 4.2230, -3.7920, math.radians(-115.0), 0.0), + Gap("Inner divertor leg (g₁)", 4.2230, -3.7920, math.radians(-65.0), 0.0), Gap("Outer divertor leg (g₂)", 5.5650, -4.4040, math.radians(-150.0), 0.0), Gap("Point at 2 o'clock (g₄)", 7.5095, 2.9971, math.radians(135.0), 0.228), Gap("Uppest boundary point (g₅)", 5.3315, 4.5804, math.radians(90.0), 0.596), - Gap("Inboard mid-plane (Rmin)", 4.0599, 0.7777, math.radians(0.0), 0.150), - Gap("Outboard mid-plane (Rmax)", 8.2806, 0.4665, math.radians(180.0), 0.080), + # Measured from r=0, at the height of the innermost and outermost point + Gap("Inboard mid-plane (Rmin)", 0.0, 0.7777, 0.0, 4.2099), + Gap("Outboard mid-plane (Rmax)", 0.0, 0.4665, 0.0, 8.2006), ] diff --git a/waveform_editor/shape_editor/plasma_shape_calc.py b/waveform_editor/shape_editor/plasma_shape_calc.py index 6da07e20..3cb22d6d 100644 --- a/waveform_editor/shape_editor/plasma_shape_calc.py +++ b/waveform_editor/shape_editor/plasma_shape_calc.py @@ -227,3 +227,110 @@ def project_point_to_line(point, start, end): proj = start + np.dot(to_pt, along) * along dist = float(np.linalg.norm(to_pt - np.dot(to_pt, along) * along)) return (float(proj[0]), float(proj[1])), dist + + +def closest_point_facing(point, r, z, axis, is_closed=True): + """The closest point on an outline that lies on the side of the axis, as DINA + measures its gaps 4 and 5: the foot must satisfy (foot - point).(axis - point) >= 0, + so a branch behind the reference point is never taken. + + Args: + point: The (r, z) reference point. + r: Radial coordinates of the outline. + z: Height coordinates of the outline. + axis: The (r, z) of the magnetic axis. + is_closed: Whether the outline closes on itself. + + Returns: + Tuple of ((r, z) of the closest point, distance), or (None, inf) if no point + faces the axis. + """ + starts = np.column_stack([r, z]).astype(float) + ends = np.roll(starts, -1, axis=0) if is_closed else starts[1:] + if not is_closed: + starts = starts[:-1] + point = np.asarray(point, float) + segs = ends - starts + lens = np.sum(segs**2, axis=1) + along = np.clip( + np.sum((point - starts) * segs, axis=1) / np.where(lens > 0, lens, 1), 0, 1 + ) + feet = starts + along[:, None] * segs + facing = (feet - point) @ (np.asarray(axis, float) - point) >= 0 + if not facing.any(): + return None, float("inf") + dists = np.where(facing, np.linalg.norm(feet - point, axis=1), np.inf) + idx = int(np.argmin(dists)) + return (float(feet[idx, 0]), float(feet[idx, 1])), float(dists[idx]) + + +def interpolate_branch(r, z, axis, at, theta_range, along="z"): + """Interpolate one branch of an outline, selected by its angle about the axis, as + FEEQS measures the WEST gaps (pchip, in coordinates relative to the axis). + + Args: + r: Radial coordinates of the outline. + z: Height coordinates of the outline. + axis: The (r, z) the angles are taken about, the magnetic axis in FEEQS. + at: The absolute height (along="z") or radius (along="r") to evaluate at. + theta_range: (low, high) angles in radians of the points that form the branch. + A low above high wraps around pi, for the inboard side. + along: "z" to find the radius at a height, "r" to find the height at a radius. + + Returns: + The radius (or height) of the branch, or None if it does not reach `at`. + """ + from scipy.interpolate import PchipInterpolator + + r0, z0 = axis + d_r, d_z = np.asarray(r, float) - r0, np.asarray(z, float) - z0 + theta = np.arctan2(d_z, d_r) + low, high = theta_range + on_branch = ( + (theta > low) & (theta < high) if low < high else (theta > low) | (theta < high) + ) + x, y = (d_z, d_r) if along == "z" else (d_r, d_z) + x, y = x[on_branch], y[on_branch] + x, unique = np.unique(x, return_index=True) + y = y[unique] + target = at - (z0 if along == "z" else r0) + if len(x) < 2 or not x[0] <= target <= x[-1]: + return None + return float(PchipInterpolator(x, y)(target)) + (r0 if along == "z" else z0) + + +def geometric_centre(r, z): + """The centre of the bounding box of an outline.""" + return (float(np.min(r) + np.max(r)) / 2, float(np.min(z) + np.max(z)) / 2) + + +def psi_contour(time_slice, level): + """The contour of the poloidal flux of an equilibrium at a level, from the GGD + that NICE fills, or else from profiles_2d. + + Args: + time_slice: The equilibrium time slice. + level: The flux to contour. + + Returns: + List of (N, 2) arrays of (r, z), one per piece of the contour. + """ + import matplotlib.pyplot as plt + + fig, ax = plt.subplots() + try: + if len(time_slice.ggd) > 0 and len(time_slice.ggd[0].psi) > 0: + ggd = time_slice.ggd[0] + r, z = ggd.r[0].values, ggd.z[0].values + contour = ax.tricontour(r, z, ggd.psi[0].values, levels=[level]) + elif len(time_slice.profiles_2d) > 0: + p2d = time_slice.profiles_2d[0] + r, z = np.asarray(p2d.grid.dim1), np.asarray(p2d.grid.dim2) + psi = np.asarray(p2d.psi) + psi = psi.T if psi.shape == (len(r), len(z)) else psi + contour = ax.contour(r, z, psi, levels=[level]) + else: + return [] + return [np.asarray(seg) for seg in contour.allsegs[0] if len(seg) > 1] + finally: + plt.close(fig) diff --git a/waveform_editor/shape_editor/west_gaps.py b/waveform_editor/shape_editor/west_gaps.py index 06e0f7cf..617684fe 100644 --- a/waveform_editor/shape_editor/west_gaps.py +++ b/waveform_editor/shape_editor/west_gaps.py @@ -7,11 +7,15 @@ import math import numpy as np +from matplotlib.path import Path from waveform_editor.shape_editor.plasma_shape_calc import ( Gap, closest_outline_point, + geometric_centre, + interpolate_branch, project_point_to_line, + psi_contour, ) # The arc the outer radial gaps are measured to, which passes through r=3 m on the @@ -19,141 +23,223 @@ OUTER_ARC = (2.2, 0.8) # The heights the upper and lower outer radial gaps are measured at OUTER_GAP_HEIGHT = 0.25 +# The inner wall the radial inner gap is measured to, on the midplane +RIG_POINT = (1.834, 0.0) +# The points above the plasma the top inner and outer gaps are measured to +TIG_POINT = (2.132, 0.672) +TOG_POINT = (2.456, 0.749) # The divertor targets, as the two points of the line through each of them LOWER_DIVERTOR = ((1.909, -0.5796), (2.362, -0.7624)) UPPER_DIVERTOR = ((1.9009, 0.5824), (2.446, 0.7995)) # The corner of the baffle the plasma is kept away from BAFFLE = (2.381, -0.6757) +# FEEQS only takes the x-points inside this region +X_SEARCH_REGION = Path( + [ + (1.956, -0.6582), + (2.315, -0.8031), + (2.6, -0.8031), + (2.6, 0.8018), + (2.304, 0.8018), + (1.956, 0.6611), + ] +) WEST_METADATA = { "UROG": ("UROG", "cm", "Upper radial outer gap"), "EROG": ("EROG", "cm", "Equatorial radial outer gap"), "LROG": ("LROG", "cm", "Lower radial outer gap"), + "RIG": ("RIG", "cm", "Radial inner gap"), + "TIG": ("TIG", "cm", "Top inner gap"), + "TOG": ("TOG", "cm", "Top outer gap"), "dXlow": ("dXlow", "cm", "Distance of the lower x-point to the divertor"), "dXup": ("dXup", "cm", "Distance of the upper x-point to the divertor"), "dbaffle": ("dbaffle", "cm", "Distance of the plasma to the baffle"), + "dRsep": ( + "dRsep", + "cm", + "Distance between the second and first separatrix on the outboard midplane", + ), } -def compute_gaps(outline_r, outline_z, x_points): +def gap_inputs(time_slice): + """The arguments to compute the gaps of a solved equilibrium with. + + Args: + time_slice: The equilibrium time slice. + + Returns: + Dict of the keyword arguments of compute_gaps and compute_gap_geometry. + """ + axis = time_slice.global_quantities.magnetic_axis + return { + "outline_r": time_slice.boundary.outline.r, + "outline_z": time_slice.boundary.outline.z, + "x_points": [ + (float(node.r), float(node.z), float(node.psi)) + for node in time_slice.contour_tree.node + if int(node.critical_type) == 1 + ], + "magnetic_axis": ( + float(axis.r), + float(axis.z), + float(time_slice.global_quantities.psi_axis), + ), + "contour": lambda level: psi_contour(time_slice, level), + } + + +def compute_gaps(outline_r, outline_z, x_points, magnetic_axis=None, contour=None): """The gaps of a WEST plasma, in metres. Args: outline_r: Radial coordinates of the plasma boundary. outline_z: Height coordinates of the plasma boundary. - x_points: The (r, z) of each x-point of the equilibrium. + x_points: The (r, z) or (r, z, psi) of each x-point of the equilibrium. + magnetic_axis: The (r, z, psi) of the magnetic axis, the centre of the + outline is used without one. + contour: Function returning the contour of the flux at a level, as a list + of (N, 2) arrays, for the gaps measured to the separatrix with its legs. Returns: Dict of gap name to distance in metres. """ - r, z = np.asarray(outline_r, dtype=float), np.asarray(outline_z, dtype=float) - gaps = {} - centre_r, radius = OUTER_ARC - for name, height in ( - ("UROG", OUTER_GAP_HEIGHT), - ("EROG", 0.0), - ("LROG", -OUTER_GAP_HEIGHT), - ): - boundary_r = _outboard_radius(r, z, height) - if boundary_r is not None: - gaps[name] = centre_r + np.sqrt(radius**2 - height**2) - boundary_r + items = compute_gap_geometry(outline_r, outline_z, x_points, magnetic_axis, contour) + return {item["key"]: item["distance"] for item in items} - for name, divertor, below in ( - ("dXlow", LOWER_DIVERTOR, True), - ("dXup", UPPER_DIVERTOR, False), - ): - x_point = _x_point(x_points, below) - if x_point is not None: - gaps[name] = project_point_to_line(x_point, *divertor)[1] - gaps["dbaffle"] = closest_outline_point(BAFFLE, r, z)[1] - return gaps +def compute_gap_geometry( + outline_r, outline_z, x_points, magnetic_axis=None, contour=None +): + """The gaps of a WEST plasma, as computed in FEEQS, together with the points they + are measured between. + Args: + outline_r: Radial coordinates of the plasma boundary. + outline_z: Height coordinates of the plasma boundary. + x_points: The (r, z) or (r, z, psi) of each x-point of the equilibrium. + magnetic_axis: The (r, z, psi) of the magnetic axis, the centre of the + outline is used without one. + contour: Function returning the contour of the flux at a level, as a list + of (N, 2) arrays, for the gaps measured to the separatrix with its legs. -def compute_gap_geometry(outline_r, outline_z, x_points): - """Compute measuring points and distance segments for WEST gaps.""" + Returns: + List of dicts, one per gap that could be measured. + """ if outline_r is None or outline_z is None or len(outline_r) == 0: return [] - r, z = np.asarray(outline_r, dtype=float), np.asarray(outline_z, dtype=float) + axis = magnetic_axis[:2] if magnetic_axis is not None else geometric_centre(r, z) + psi_axis = magnetic_axis[2] if magnetic_axis is not None else None + x_points = _select_x_points(x_points, psi_axis) items = [] + + def add(key, orig, target, distance): + symbol, _, name = WEST_METADATA[key] + items.append( + { + "key": key, + "name": name, + "symbol": symbol, + "r_orig": float(orig[0]), + "z_orig": float(orig[1]), + "r_target": float(target[0]), + "z_target": float(target[1]), + "distance": float(distance), + "distance_str": f"{distance * 100:.3g} cm", + } + ) + + # The outer radial gaps, each on its own quarter of the boundary centre_r, radius = OUTER_ARC - for name, height in ( - ("UROG", OUTER_GAP_HEIGHT), - ("EROG", 0.0), - ("LROG", -OUTER_GAP_HEIGHT), + for key, height, theta_range in ( + ("UROG", OUTER_GAP_HEIGHT, (0, math.pi / 2)), + ("EROG", 0.0, (-math.pi / 2, math.pi / 2)), + ("LROG", -OUTER_GAP_HEIGHT, (-math.pi / 2, 0)), ): - arc_r = float(centre_r + np.sqrt(radius**2 - height**2)) - boundary_r = _outboard_radius(r, z, height) + wall_r = centre_r + math.sqrt(radius**2 - height**2) + boundary_r = interpolate_branch(r, z, axis, height, theta_range) if boundary_r is not None: - dist = arc_r - boundary_r - symbol, _, full_name = WEST_METADATA[name] - items.append( - { - "name": full_name, - "symbol": symbol, - "r_orig": arc_r, - "z_orig": height, - "r_target": boundary_r, - "z_target": height, - "distance": dist, - "distance_str": f"{dist * 100:.3g} cm", - } - ) - - for name, divertor, below in ( + add(key, (wall_r, height), (boundary_r, height), wall_r - boundary_r) + + # The radial inner gap, on the inboard side within pi/16 of the midplane. FEEQS + # reports it as wall minus boundary, which is negative; it is a clearance here. + rig_r, rig_z = RIG_POINT + wedge = (math.pi - math.pi / 16, -math.pi + math.pi / 16) + boundary_r = interpolate_branch(r, z, axis, rig_z, wedge) + if boundary_r is not None: + add("RIG", RIG_POINT, (boundary_r, rig_z), boundary_r - rig_r) + + # The top gaps, measured down to the upper half of the boundary + for key, (wall_r, wall_z) in (("TIG", TIG_POINT), ("TOG", TOG_POINT)): + boundary_z = interpolate_branch(r, z, axis, wall_r, (0, math.pi), along="r") + if boundary_z is not None: + add(key, (wall_r, wall_z), (wall_r, boundary_z), wall_z - boundary_z) + + for key, divertor, below in ( ("dXlow", LOWER_DIVERTOR, True), ("dXup", UPPER_DIVERTOR, False), ): - x_point = _x_point(x_points, below) - if x_point is not None: - proj, dist = project_point_to_line(x_point, *divertor) - symbol, _, full_name = WEST_METADATA[name] - items.append( - { - "name": full_name, - "symbol": symbol, - "r_orig": proj[0], - "z_orig": proj[1], - "r_target": float(x_point[0]), - "z_target": float(x_point[1]), - "distance": dist, - "distance_str": f"{dist * 100:.3g} cm", - } - ) - - (target_r, target_z), dist = closest_outline_point(BAFFLE, r, z) - symbol, _, full_name = WEST_METADATA["dbaffle"] - items.append( - { - "name": full_name, - "symbol": symbol, - "r_orig": BAFFLE[0], - "z_orig": BAFFLE[1], - "r_target": target_r, - "z_target": target_z, - "distance": dist, - "distance_str": f"{dist * 100:.3g} cm", - } - ) + on_side = [point for point in x_points if (point[1] < 0) == below] + if on_side: + foot, distance = project_point_to_line(on_side[0][:2], *divertor) + add(key, foot, on_side[0][:2], distance) + + # FEEQS measures the baffle to the separatrix with its legs, which only a solved + # equilibrium has; a desired boundary is its own separatrix + separatrix, is_closed = [np.column_stack([r, z])], True + if contour is not None and x_points and len(x_points[0]) > 2: + pieces = contour(x_points[0][2]) + if pieces: + # Contour pieces are open polylines; a closed one repeats its first point + separatrix, is_closed = pieces, False + best = None + for piece in separatrix: + target, distance = closest_outline_point( + BAFFLE, piece[:, 0], piece[:, 1], is_closed=is_closed + ) + if best is None or distance < best[1]: + best = target, distance + add("dbaffle", BAFFLE, best[0], best[1]) + + # The distance between the separatrices of a double null on the outboard midplane + if contour is not None and len(x_points) == 2 and len(x_points[1]) > 2: + radii = [ + _outboard_midplane_radius(contour(point[2]), axis) for point in x_points + ] + if None not in radii: + first, second = radii + add("dRsep", (first, 0.0), (second, 0.0), second - first) return items -def _outboard_radius(r, z, height): - """The radius of the outboard side of the boundary at a height.""" - outboard = r > r.mean() - r, z = r[outboard], z[outboard] - if not len(z) or not z.min() <= height <= z.max(): - return None - order = np.argsort(z) - return float(np.interp(height, z[order], r[order])) +def _select_x_points(x_points, psi_axis): + """The x-points FEEQS takes: those inside its search region, at most two, the + one of the first separatrix first. + + The first separatrix is the one closest in flux to the magnetic axis, which FEEQS + finds by sorting on psi in its own sign convention. + """ + inside = [ + tuple(float(v) for v in point) + for point in x_points + if X_SEARCH_REGION.contains_point(point[:2]) + ] + if psi_axis is not None and all(len(point) > 2 for point in inside): + inside.sort(key=lambda point: abs(point[2] - psi_axis)) + return inside[:2] -def _x_point(x_points, below): - """The x-point below or above the midplane.""" - on_side = [point for point in x_points if (point[1] < 0) == below] - return on_side[0] if on_side else None +def _outboard_midplane_radius(pieces, axis): + """The radius at z=0 of the outboard side of a contour.""" + if not pieces: + return None + points = np.concatenate(pieces) + return interpolate_branch( + points[:, 0], points[:, 1], axis, 0.0, (-math.pi / 2, math.pi / 2) + ) def get_default_west_gaps():