Add defaults for ITER and WEST machine presets

waveform_editor/gui/shape_editor/coil_currents.py

@@ -8,7 +8,15 @@
 from panel.viewable import Viewer
 
 from waveform_editor.gui.util import set_xml_parameter
-from waveform_editor.settings import settings
+from waveform_editor.settings import NiceSettings, settings
+
+# The NICE current group of each coil of a machine, in the order the coils appear in
+# its pf_active. Coils that share a group carry the same current, such as the four
+# divertor coils in each of WEST's divertor circuits.
+MACHINE_COIL_GROUPS = {
+    NiceSettings.PRESET_ITER: [0, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+    NiceSettings.PRESET_WEST: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9, 10, 10, 10, 10],
+}
 
 
 class CoilCurrentEntry(param.Parameterized):
@@ -265,6 +273,27 @@ def update_xml(self, xml_params: ET.Element):
         self._update_fixed_coils_in_xml(xml_params)
         self._update_penalization_in_xml(xml_params)
 
+    def write_coil_groups(self, xml_params: ET.Element):
+        """Write how the coils of the machine of the selected preset are grouped, as
+        the parameter file holds the groups of ITER.
+
+        Args:
+            xml_params: XML representing configuration parameters, which are updated
+                in-place.
+        """
+        groups = MACHINE_COIL_GROUPS.get(self.nice_settings.machine_preset)
+        if groups is None:
+            return
+        if len(groups) != len(self.coils):
+            raise ValueError(
+                f"The {self.nice_settings.machine_preset} machine description has "
+                f"{len(self.coils)} coils, but {len(groups)} coils are grouped for it."
+            )
+        set_xml_parameter(xml_params, "n_coil_group_index", len(groups))
+        set_xml_parameter(
+            xml_params, "coil_group_index", " ".join(str(group) for group in groups)
+        )
+
     def _update_fixed_coils_in_xml(self, xml_params: ET.Element):
         """Update XML parameters indicating which coils are fixed based on
         UI checkboxes.

waveform_editor/gui/shape_editor/nice_plotter.py

@@ -59,27 +59,37 @@ class NicePlotter(Viewer):
     _points_toolbar = None
     _syncing_points = False
 
-    # Bokeh cannot hold a data aspect while it resizes, so fix the size instead
-    R_RANGE = (0, 13)
-    Z_RANGE = (-10, 10)
     FRAME_HEIGHT = 700
+    # The r and z range to plot per machine, wide enough to hold its coils
+    MACHINE_RANGES = {
+        NiceSettings.PRESET_ITER: ((0, 13), (-10, 10)),
+        NiceSettings.PRESET_WEST: ((0, 4.8), (-2.6, 2.6)),
+    }
+    R_RANGE, Z_RANGE = MACHINE_RANGES[NiceSettings.PRESET_ITER]
     FRAME_WIDTH = round(
         FRAME_HEIGHT * (R_RANGE[1] - R_RANGE[0]) / (Z_RANGE[1] - Z_RANGE[0])
     )
 
     def __init__(self, **params):
         super().__init__(**params)
+        self.nice_settings = settings.nice
+        self._figure = None
+        self.nice_settings.param.watch(self._apply_machine_ranges, "machine_preset")
+        r_range, z_range = self.MACHINE_RANGES.get(
+            self.nice_settings.machine_preset, (self.R_RANGE, self.Z_RANGE)
+        )
+        span_r, span_z = r_range[1] - r_range[0], z_range[1] - z_range[0]
         self.DEFAULT_OPTS = hv.opts.Overlay(
-            xlim=self.R_RANGE,
-            ylim=self.Z_RANGE,
-            frame_width=self.FRAME_WIDTH,
+            xlim=r_range,
+            ylim=z_range,
+            frame_width=round(self.FRAME_HEIGHT * span_r / span_z),
             frame_height=self.FRAME_HEIGHT,
             title="",
             xlabel="r [m]",
             ylabel="z [m]",
             fontsize={"labels": 15, "ticks": 11},
+            hooks=[self._capture_figure],
         )
-        self.nice_settings = settings.nice
         self.CONTOUR_OPTS = hv.opts.Contours(
             cmap="viridis",
             colorbar=True,
@@ -201,6 +211,22 @@ def _on_point_draw(self, data):
         finally:
             self._syncing_points = False
 
+    def _capture_figure(self, plot, element):
+        """Keep hold of the figure, so its ranges can follow the machine preset."""
+        self._figure = plot.state
+
+    def _apply_machine_ranges(self, event=None):
+        """Show the machine of the selected preset, coils and all."""
+        if self._figure is None:
+            return
+        r_range, z_range = self.MACHINE_RANGES.get(
+            self.nice_settings.machine_preset, (self.R_RANGE, self.Z_RANGE)
+        )
+        self._figure.x_range.start, self._figure.x_range.end = r_range
+        self._figure.y_range.start, self._figure.y_range.end = z_range
+        span_r, span_z = r_range[1] - r_range[0], z_range[1] - z_range[0]
+        self._figure.frame_width = round(self.FRAME_HEIGHT * span_r / span_z)
+
     @pn.depends(
         "plasma_shape.shape_updated", "show_desired_shape", "nice_settings.mode"
     )

waveform_editor/gui/shape_editor/plasma_properties.py

@@ -10,6 +10,7 @@
     EquilibriumInput,
     FormattedEditableFloatSlider,
 )
+from waveform_editor.settings import NiceSettings, settings
 from waveform_editor.shape_editor.plasma_properties_calc import (
     compute_profiles_from_params,
 )
@@ -118,6 +119,28 @@ def __panel__(self):
         )
 
 
+MACHINE_PROPERTIES = {
+    NiceSettings.PRESET_ITER: {
+        "ip": -1.5e7,
+        "r0": 6.2,
+        "b0": -5.3,
+        "alpha": 1.0,
+        "beta": 0.65,
+        "gamma": 1.05,
+    },
+    NiceSettings.PRESET_WEST: {
+        "ip": -4.0e5,
+        "r0": 2.42,
+        "b0": -3.76,
+        "alpha": 1.0,
+        "beta": 0.85,
+        "gamma": 1.2,
+    },
+}
+# The defaults, before a machine preset is chosen
+DEFAULT_PROPERTIES = MACHINE_PROPERTIES[NiceSettings.PRESET_ITER]
+
+
 class PlasmaProfiles(Viewer):
     """Widget for plasma profile source: starts parametric, populates from IDS on load.
 
@@ -126,9 +149,15 @@ class PlasmaProfiles(Viewer):
     to an "edited" badge. Clicking Reset restores the IDS profiles.
     """
 
-    alpha = param.Number(default=0.5, softbounds=[0.5, 2], step=0.01)
-    beta = param.Number(default=0.5, softbounds=[0.5, 2], step=0.01)
-    gamma = param.Number(default=1.0, softbounds=[0.5, 2], step=0.01)
+    alpha = param.Number(
+        default=DEFAULT_PROPERTIES["alpha"], softbounds=[0.5, 2], step=0.01
+    )
+    beta = param.Number(
+        default=DEFAULT_PROPERTIES["beta"], softbounds=[0, 2], step=0.01
+    )
+    gamma = param.Number(
+        default=DEFAULT_PROPERTIES["gamma"], softbounds=[0.5, 2], step=0.01
+    )
     changed = param.Event()
     input_state = param.String(default="manual", precedence=-1)
 
@@ -337,16 +366,16 @@ def __init__(self):
         super().__init__()
         self._ip = PropertyInput(
             "Plasma current [A]",
-            default_value=-1.5e7,
+            default_value=DEFAULT_PROPERTIES["ip"],
             step=1e6,
         )
         self._r0 = PropertyInput(
             "Reference major radius [m]",
-            default_value=6.2,
+            default_value=DEFAULT_PROPERTIES["r0"],
         )
         self._b0 = PropertyInput(
             "Toroidal magnetic field [T]",
-            default_value=-5.3,
+            default_value=DEFAULT_PROPERTIES["b0"],
         )
         self._profiles = PlasmaProfiles()
 
@@ -363,8 +392,24 @@ def __init__(self):
         for widget in [self._ip, self._r0, self._b0, self._profiles]:
             widget.param.watch(self._load_plasma_properties, "changed")
 
+        settings.nice.param.watch(self.apply_machine_properties, "machine_preset")
+        self.apply_machine_properties()
         self._load_plasma_properties()
 
+    def apply_machine_properties(self, event=None):
+        """Start from the properties the machine of the selected preset is run at."""
+        properties = MACHINE_PROPERTIES.get(settings.nice.machine_preset)
+        if properties is None:
+            return
+        self._ip.value = properties["ip"]
+        self._r0.value = properties["r0"]
+        self._b0.value = properties["b0"]
+        self._profiles.param.update(
+            alpha=properties["alpha"],
+            beta=properties["beta"],
+            gamma=properties["gamma"],
+        )
+
     _IDS_SCALAR_PATHS = [
         ("_ip", "time_slice(0)/global_quantities/ip"),
         ("_r0", "vacuum_toroidal_field/r0"),

waveform_editor/gui/shape_editor/plasma_shape.py

@@ -11,6 +11,7 @@
     FormattedEditableFloatSlider,
     WarningIndicator,
 )
+from waveform_editor.settings import NiceSettings, settings
 from waveform_editor.shape_editor.plasma_shape_calc import (
     Gap,
     apply_point_weights,
@@ -19,6 +20,33 @@
     update_outline_from_gaps,
 )
 
+# A shape to start from for each machine, taken from one of its discharges, as
+# (value, slider range) per parameter
+MACHINE_SHAPES = {
+    NiceSettings.PRESET_ITER: {
+        "a": (1.9, (1, 2)),
+        "center_r": (6.2, (5, 7)),
+        "center_z": (0.545, (0, 1.5)),
+        "kappa": (1.8, (0, 3)),
+        "delta": (0.43, (-1, 1)),
+        "rx": (5.089, (4.5, 6)),
+        "zx": (-3.346, (-4, -2)),
+        "rx_upper": (5.089, (4.5, 6)),
+        "zx_upper": (3.346, (2, 4)),
+    },
+    NiceSettings.PRESET_WEST: {
+        "a": (0.46, (0.2, 0.8)),
+        "center_r": (2.54, (2, 3)),
+        "center_z": (-0.02, (-0.5, 0.5)),
+        "kappa": (1.31, (0, 3)),
+        "delta": (0.38, (-1, 1)),
+        "rx": (2.23, (1.8, 3.2)),
+        "zx": (-0.62, (-1.2, 0)),
+        "rx_upper": (2.23, (1.8, 3.2)),
+        "zx_upper": (0.62, (0, 1.2)),
+    },
+}
+
 
 class PlasmaShapeParams(Viewer):
     """Helper class containing parameters to parameterize the plasma shape."""
@@ -99,14 +127,34 @@ class PlasmaShapeParams(Viewer):
     )
     extra_points_table = param.Parameter()
 
+    def apply_machine_shape(self, event=None):
+        """Start from the shape of the machine of the selected preset."""
+        shape = MACHINE_SHAPES.get(settings.nice.machine_preset)
+        if shape is None:  # a custom machine, whose shape we cannot know
+            return
+        for name, (_, slider_range) in shape.items():
+            if name in self.param:
+                self.param[name].softbounds = slider_range
+            if name in getattr(self, "_sliders", {}):
+                self._sliders[name].start, self._sliders[name].end = slider_range
+        self.param.update(
+            **{name: value for name, (value, _) in shape.items() if name in self.param}
+        )
+
     def __panel__(self):
+        # The sliders of this panel, so that their range can follow the machine
+        self._sliders = {}
+
         def _slider(n):
             p = getattr(self.param, n)
             if isinstance(self.param[n], param.Boolean):
                 return pn.widgets.Checkbox.from_param(p)
             if isinstance(self.param[n], param.Integer):
                 return FixedWidthEditableIntSlider.from_param(p, stretch_width=True)
-            return FormattedEditableFloatSlider.from_param(p, stretch_width=True)
+            self._sliders[n] = FormattedEditableFloatSlider.from_param(
+                p, stretch_width=True
+            )
+            return self._sliders[n]
 
         def _group(title, *children, visible=True):
             return pn.Column(
@@ -458,6 +506,11 @@ def _indicator(tooltip):
         self.param_z = None
         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.shape_params.apply_machine_shape, "machine_preset"
+        )
+        self.shape_params.apply_machine_shape()
 
     @pn.depends(
         "shape_params.param",

waveform_editor/gui/shape_editor/settings_modal.py

@@ -85,6 +85,7 @@ def __init__(self, nice_plotter: NicePlotter, **params):
         self.nice_settings.param.watch(
             self._update_md_inputs_visibility, ["machine_preset"]
         )
+        self._update_md_inputs_visibility(None)
 
     def _parameters_section(self):
         """A table of every NICE parameter, which writes edits to the settings."""

waveform_editor/gui/shape_editor/shape_editor.py

@@ -418,6 +418,7 @@ async def submit(self, event=None):
 
         self.coil_currents.fill_pf_active(self.pf_active)
         xml_params = ET.fromstring(self.xml_text)
+        self.coil_currents.write_coil_groups(xml_params)
         if not self._apply_xml_parameters(xml_params):
             return