Plasma shape input from weighted points

waveform_editor/gui/shape_editor/nice_plotter.py

@@ -148,6 +148,8 @@ def _plot_plasma_shape(self):
 
         if self.plasma_shape.input_mode == self.plasma_shape.GAP_INPUT:
             return self._plot_gaps(r, z)
+        elif self.plasma_shape.input_mode == self.plasma_shape.WEIGHTED_POINTS_INPUT:
+            return self._plot_weighted_points(r, z)
         else:
             return self._plot_outline_shape(r, z)
 
@@ -187,6 +189,18 @@ def _plot_gaps(self, r, z):
             )
         return hv.Overlay(plot_elements)
 
+    def _plot_weighted_points(self, r, z):
+        """Plots weighted points as scatterplot.
+
+        Args:
+            r: Radial coordinates of the points.
+            z: Height coordinates of the points.
+
+        Returns:
+            Holoviews overlay with scatter plot of the points.
+        """
+        return hv.Overlay([hv.Scatter((r, z)).opts(color="blue", size=8, marker="o")])
+
     @pn.depends("pf_active", "show_coils", "communicator.pf_active")
     def _plot_coil_rectangles(self):
         """Creates rectangular and path overlays for PF coils.

waveform_editor/gui/shape_editor/plasma_shape.py

@@ -1,4 +1,5 @@
 import imas
+import pandas as pd
 import panel as pn
 import param
 from panel.viewable import Viewer
@@ -50,13 +51,140 @@ def __panel__(self):
         return pn.Param(self.param, widgets=widgets, show_name=False)
 
 
+class WeightedPointsTable(param.Parameterized):
+    """Widget for managing weighted points table for defining plasma shape."""
+
+    COL_R = "R [m]"
+    COL_Z = "Z [m]"
+    COL_WEIGHT = "weight"
+    COL_DELETE = "🗑️"
+
+    points = param.DataFrame(default=pd.DataFrame(columns=[COL_R, COL_Z, COL_WEIGHT]))
+
+    def __init__(self):
+        super().__init__()
+        initial_df = pd.DataFrame(
+            columns=(self.COL_DELETE, self.COL_R, self.COL_Z, self.COL_WEIGHT)
+        )
+        self._tabulator = pn.widgets.Tabulator(
+            value=initial_df,
+            editors={
+                self.COL_DELETE: None,
+                self.COL_R: {"type": "number"},
+                self.COL_Z: {"type": "number"},
+                self.COL_WEIGHT: {"type": "number"},
+            },
+            layout="fit_data_fill",
+            sizing_mode="stretch_width",
+            show_index=False,
+        )
+        self._tabulator.on_click(self._on_delete_click)
+        self._tabulator.on_edit(self._on_edit)
+        self.param.watch(self._update_tabulator, "points", onlychanged=True)
+        self._update_tabulator()
+
+    def _update_tabulator(self, event=None):
+        """Update the Tabulator to reflect the current DataFrame."""
+        df = self.points
+        data = []
+        for _, row in df.iterrows():
+            data.append(
+                (
+                    self.COL_DELETE,
+                    row[self.COL_R],
+                    row[self.COL_Z],
+                    row[self.COL_WEIGHT],
+                )
+            )
+
+        # Add empty row if last data row has R and Z values filled
+        if len(df) == 0 or (
+            pd.notna(df.iloc[-1][self.COL_R])
+            and pd.notna(df.iloc[-1][self.COL_Z])
+            and df.iloc[-1][self.COL_R] != ""
+            and df.iloc[-1][self.COL_Z] != ""
+        ):
+            data.append((self.COL_DELETE, "", "", 1))
+
+        new_df = pd.DataFrame(
+            data, columns=(self.COL_DELETE, self.COL_R, self.COL_Z, self.COL_WEIGHT)
+        )
+        # Convert columns to allow mixed types
+        new_df[self.COL_R] = new_df[self.COL_R].astype(object)
+        new_df[self.COL_Z] = new_df[self.COL_Z].astype(object)
+        self._tabulator.value = new_df
+
+    def _on_delete_click(self, event):
+        if event.column == self.COL_DELETE:
+            n_data_rows = len(self.points)
+            if event.row < n_data_rows:
+                df = self.points.copy()
+                df = df.drop(index=event.row).reset_index(drop=True)
+                self.param.update(points=df)
+
+    def _on_edit(self, event):
+        """Handle edits in the weighted points tabulator."""
+        df = self.points.copy()
+        is_empty_row = event.row >= len(df)
+
+        # Convert columns to object dtype to allow mixed types
+        for col in [self.COL_R, self.COL_Z, self.COL_WEIGHT]:
+            if col in df.columns:
+                df[col] = df[col].astype(object)
+
+        if is_empty_row:
+            new_row = {self.COL_R: "", self.COL_Z: "", self.COL_WEIGHT: 1}
+            new_row[event.column] = event.value
+            new_df = pd.DataFrame([new_row])
+            df = pd.concat([df, new_df], ignore_index=True)
+        else:
+            df.at[event.row, event.column] = event.value
+
+        self.param.update(points=df)
+
+    def get_outline_coordinates(self):
+        """Generate outline coordinates from weighted points.
+
+        Returns:
+            tuple: (outline_r, outline_z) lists of coordinates, or (None, None) if empty
+        """
+        if self.points.empty:
+            return None, None
+
+        # Filter out rows with empty R or Z
+        valid_df = self.points.dropna(subset=[self.COL_R, self.COL_Z])
+        valid_df = valid_df[(valid_df[self.COL_R] != "") & (valid_df[self.COL_Z] != "")]
+
+        if valid_df.empty:
+            return None, None
+
+        # Duplicate points according to their weight
+        outline_r = []
+        outline_z = []
+        for _, row in valid_df.iterrows():
+            weight = int(row[self.COL_WEIGHT]) if pd.notna(row[self.COL_WEIGHT]) else 1
+            outline_r.extend([float(row[self.COL_R])] * weight)
+            outline_z.extend([float(row[self.COL_Z])] * weight)
+
+        return outline_r, outline_z
+
+    def __panel__(self):
+        return self._tabulator
+
+
 class PlasmaShape(Viewer):
     PARAMETERIZED_INPUT = "Parameterized"
     EQUILIBRIUM_INPUT = "Equilibrium IDS outline"
     GAP_INPUT = "Equilibrium IDS Gaps"
+    WEIGHTED_POINTS_INPUT = "Weighted Points"
     input_mode = param.ObjectSelector(
         default=EQUILIBRIUM_INPUT,
-        objects=[EQUILIBRIUM_INPUT, PARAMETERIZED_INPUT, GAP_INPUT],
+        objects=[
+            EQUILIBRIUM_INPUT,
+            PARAMETERIZED_INPUT,
+            GAP_INPUT,
+            WEIGHTED_POINTS_INPUT,
+        ],
         label="Shape input mode",
     )
     input_outline = param.ClassSelector(
@@ -65,6 +193,9 @@ class PlasmaShape(Viewer):
     input_gaps = param.ClassSelector(
         class_=EquilibriumInput, default=EquilibriumInput()
     )
+    weighted_points_table = param.ClassSelector(
+        class_=WeightedPointsTable, default=WeightedPointsTable()
+    )
     shape_params = param.ClassSelector(
         class_=PlasmaShapeParams, default=PlasmaShapeParams()
     )
@@ -88,6 +219,7 @@ def __init__(self):
         "shape_params.param",
         "input_outline.param",
         "input_gaps.param",
+        "weighted_points_table.param",
         "input_mode",
         watch=True,
     )
@@ -102,6 +234,8 @@ def _set_plasma_shape(self):
             self._load_shape_from_params()
         elif self.input_mode == self.GAP_INPUT:
             self._load_shape_from_gaps()
+        elif self.input_mode == self.WEIGHTED_POINTS_INPUT:
+            self._load_shape_from_weighted_points()
 
         if self.outline_r and self.outline_z:
             self.has_shape = True
@@ -196,6 +330,12 @@ def _create_gap_ui(self):
 
         self.gap_ui.extend(new_gap_ui)
 
+    def _load_shape_from_weighted_points(self):
+        """Load plasma boundary outline from weighted points."""
+        self.outline_r, self.outline_z = (
+            self.weighted_points_table.get_outline_coordinates()
+        )
+
     def _load_shape_from_params(self):
         """Compute plasma boundary outline from parameterized shape inputs."""
         self.outline_r, self.outline_z = compute_outline_from_params(
@@ -217,6 +357,8 @@ def _panel_shape_options(self):
             return pn.Row(pn.Param(self.input_outline, show_name=False), self.indicator)
         elif self.input_mode == self.GAP_INPUT:
             return pn.Row(pn.Param(self.input_gaps, show_name=False), self.indicator)
+        elif self.input_mode == self.WEIGHTED_POINTS_INPUT:
+            return self.weighted_points_table
 
     def __panel__(self):
         return self.panel