Feature/add separatrix and magnetic axes in plotequilibrium

README.md

@@ -60,11 +60,20 @@ plotequilibrium --uri "imas:mdsplus?user=public;pulse=134174;run=117;database=IT
 
 ## Requirements
 
-- Python ≥ 3.8
-- IMAS Python Access Layer (`imas-python`)
+- Python ≥ 3.10
+
+### Installed automatically via pip
 - NumPy, Matplotlib, Pandas
 - Rich (for enhanced terminal output)
 
+### Requires separate installation
+- **Tkinter** — usually bundled with Python but may require system packages:
+  - Linux (Debian/Ubuntu): `sudo apt install python3-tk`
+  - Linux (RHEL/CentOS/Rocky): `sudo dnf install python3-tkinter`
+  - Windows: included in the [python.org](https://www.python.org/downloads/) installer ("tcl/tk and IDLE" component, enabled by default)
+  - macOS (python.org installer): included by default
+  - macOS (Homebrew): `brew install python-tk` (or `brew install python-tk@3.x` for a specific version)
+
 ## Documentation
 
 Full documentation is available at the project repository. Each tool includes built-in help:

idstools/compute/equilibrium.py

@@ -16,6 +16,8 @@
 
 from idstools.database import DBMaster
 
+_IDS_VALID_THRESHOLD = abs(imas.ids_defs.EMPTY_FLOAT)
+
 logger = logging.getLogger("module")
 
 
@@ -285,6 +287,331 @@ def get_ip(self) -> list:
             for time_index in range(len(self.ids.time_slice))
         ]
 
+    def get_boundary_data(self, time_slice: int) -> dict:
+        """Return boundary and boundary_separatrix data for a given time slice.
+
+        Reads:
+
+        * ``boundary/outline/r|z``
+        * ``boundary/type``  (0=limiter, 1=diverted)
+        * ``boundary/psi_norm``
+        * ``boundary/geometric_axis/r|z``
+
+        if boundary_separatrix is available:
+
+        * ``boundary_separatrix/outline/r|z``
+        * ``boundary_separatrix/x_point[i]/r|z``
+        * ``boundary_separatrix/strike_point[i]/r|z``
+
+        if contour_tree is available:
+
+        * ``contour_tree/node[i]/critical_type``
+        * ``contour_tree/node[i]/r|z`` for X-points (``critical_type == 1``)
+        * ``contour_tree/node[i]/levelset/r|z`` for separatrix outline
+        * ``constraints/strike_point[i]/position_reconstructed/r|z`` for strike points
+          (fallback to ``position_measured/r|z`` when needed)
+
+        Args:
+            time_slice (int): Index into ``time_slice``.
+
+        Returns:
+            dict with keys:
+
+            * ``"bnd_r"``, ``"bnd_z"``            boundary outline (closed), or ``None``
+            * ``"bnd_type"``                      int or ``None``
+            * ``"bnd_psi_norm"``                  float or ``None``
+            * ``"bnd_geom_r"``, ``"bnd_geom_z"``  geometric axis scalars or ``None``
+            * ``"sep_r"``, ``"sep_z"``            separatrix outline (closed), or ``None``
+            * ``"sep_xpoints"``                   list of (r, z) tuples
+            * ``"sep_strikepoints"``              list of (r, z) tuples
+        """
+
+        def _valid_arr(arr):
+            a = np.asarray(arr, dtype=float)
+            return a.size > 0 and np.any(np.isfinite(a) & (np.abs(a) < _IDS_VALID_THRESHOLD))
+
+        def _valid_scalar(val):
+            try:
+                v = float(val)
+                return np.isfinite(v) and abs(v) < _IDS_VALID_THRESHOLD
+            except Exception:
+                return False
+
+        def _clean(arr):
+            a = np.asarray(arr, dtype=float)
+            a[(~np.isfinite(a)) | (np.abs(a) >= _IDS_VALID_THRESHOLD)] = np.nan
+            return a
+
+        def _read_outline(node):
+            try:
+                r = np.asarray(node.outline.r, dtype=float)
+                z = np.asarray(node.outline.z, dtype=float)
+            except Exception:
+                return None, None
+            if not (_valid_arr(r) and _valid_arr(z)):
+                return None, None
+            r, z = _clean(r), _clean(z)
+            # Insert NaN at large jumps so disconnected arcs are not joined
+            dist = np.sqrt(np.diff(r) ** 2 + np.diff(z) ** 2)
+            median_dist = np.nanmedian(dist)
+            if median_dist > 0:
+                breaks = np.where(dist > 20.0 * median_dist)[0] + 1
+                if len(breaks):
+                    r = np.insert(r, breaks, np.nan)
+                    z = np.insert(z, breaks, np.nan)
+            return r, z
+
+        def _read_points(node, attr):
+            pts = []
+            try:
+                arr = getattr(node, attr)
+            except AttributeError:
+                return pts
+            for pt in arr:
+                try:
+                    r, z = float(pt.r), float(pt.z)
+                except Exception:
+                    continue
+                if _valid_scalar(r) and _valid_scalar(z):
+                    pts.append((r, z))
+            return pts
+
+        def _read_contour_tree(ts_node):
+            """Read separatrix/X-point data from ``time_slice.contour_tree.node``.
+
+            * ``node.critical_type == 1`` for X-points (saddle points)
+            * first valid X-point ``node.levelset.r/z`` as separatrix contour
+            """
+            sep_r = sep_z = None
+            xpoints = []
+
+            try:
+                nodes = ts_node.contour_tree.node
+            except Exception:
+                return sep_r, sep_z, xpoints
+
+            for node in nodes:
+                try:
+                    critical_type = int(node.critical_type)
+                except Exception:
+                    continue
+
+                if critical_type != 1:  # 1 = saddle / X-point
+                    continue
+
+                try:
+                    xr = float(node.r)
+                    xz = float(node.z)
+                except Exception:
+                    xr = xz = None
+
+                if xr is not None and _valid_scalar(xr) and xz is not None and _valid_scalar(xz):
+                    xpoints.append((xr, xz))
+
+                if sep_r is not None and sep_z is not None:
+                    continue
+
+                try:
+                    r = np.asarray(node.levelset.r, dtype=float)
+                    z = np.asarray(node.levelset.z, dtype=float)
+                except Exception:
+                    continue
+
+                if not (_valid_arr(r) and _valid_arr(z)):
+                    continue
+
+                sep_r = _clean(r)
+                sep_z = _clean(z)
+
+            return sep_r, sep_z, xpoints
+
+        result = {
+            "bnd_r": None,
+            "bnd_z": None,
+            "bnd_type": None,
+            "bnd_psi_norm": None,
+            "bnd_geom_r": None,
+            "bnd_geom_z": None,
+            "sep_r": None,
+            "sep_z": None,
+            "sep_xpoints": [],
+            "sep_strikepoints": [],
+        }
+
+        try:
+            ts = self.ids.time_slice[time_slice]
+        except Exception:
+            return result
+
+        # boundary
+        try:
+            bnd = ts.boundary
+            result["bnd_r"], result["bnd_z"] = _read_outline(bnd)
+
+            bnd_type = int(bnd.type)
+            if _valid_scalar(bnd_type):
+                result["bnd_type"] = bnd_type
+        except Exception:
+            pass
+
+        try:
+            psi_norm = float(ts.boundary.psi_norm)
+            if _valid_scalar(psi_norm):
+                result["bnd_psi_norm"] = psi_norm
+        except Exception:
+            pass
+
+        try:
+            gax_r = float(ts.boundary.geometric_axis.r)
+            gax_z = float(ts.boundary.geometric_axis.z)
+            if _valid_scalar(gax_r) and _valid_scalar(gax_z):
+                result["bnd_geom_r"] = gax_r
+                result["bnd_geom_z"] = gax_z
+        except Exception:
+            pass
+
+        # boundary_separatrix (DD3 )
+        if hasattr(ts, "boundary_separatrix"):
+            sep = ts.boundary_separatrix
+            try:
+                result["sep_r"], result["sep_z"] = _read_outline(sep)
+                result["sep_xpoints"] = _read_points(sep, "x_point")
+                result["sep_strikepoints"] = _read_points(sep, "strike_point")
+            except Exception:
+                pass
+
+        #  contour_tree.node (DD4)
+        if hasattr(ts, "contour_tree") and hasattr(ts.contour_tree, "node"):
+            contour_sep_r, contour_sep_z, contour_xpoints = _read_contour_tree(ts)
+
+            if (
+                (result["sep_r"] is None or result["sep_z"] is None)
+                and contour_sep_r is not None
+                and contour_sep_z is not None
+            ):
+                result["sep_r"] = contour_sep_r
+                result["sep_z"] = contour_sep_z
+
+            if not result["sep_xpoints"] and contour_xpoints:
+                result["sep_xpoints"] = contour_xpoints
+
+        return result
+
+    def get_magnetic_axis(self, time_slice: int) -> Union[dict, None]:
+        """Return the magnetic axis position for a given time slice.
+
+        Reads ``global_quantities.magnetic_axis.r/z`` and validates the
+        scalar values.
+
+        Args:
+            time_slice (int): Index into ``time_slice``.
+
+        Returns:
+            dict with scalar keys ``"r"`` and ``"z"`` (floats), or
+            ``None`` if the data are absent or invalid.
+        """
+        try:
+            mag_ax = self.ids.time_slice[time_slice].global_quantities.magnetic_axis
+            r = float(mag_ax.r)
+            z = float(mag_ax.z)
+        except Exception as exc:
+            logger.debug(f"get_magnetic_axis: could not read magnetic_axis – {exc}")
+            return None
+
+        def _valid(val):
+            return np.isfinite(val) and abs(val) < _IDS_VALID_THRESHOLD
+
+        if not (_valid(r) and _valid(z)):
+            logger.debug("get_magnetic_axis: magnetic_axis contains no valid data")
+            return None
+
+        return {"r": r, "z": z}
+
+    def get_current_centre(self, time_slice: int) -> Union[dict, None]:
+        """Return the current centroid position for a given time slice.
+
+        Reads ``global_quantities.current_centre.r/z`` and validates the
+        scalar values.
+
+        Args:
+            time_slice (int): Index into ``time_slice``.
+
+        Returns:
+            dict with scalar keys ``"r"`` and ``"z"`` (floats), or
+            ``None`` if the data are absent or invalid.
+        """
+        try:
+            cc = self.ids.time_slice[time_slice].global_quantities.current_centre
+            r = float(cc.r)
+            z = float(cc.z)
+        except Exception as exc:
+            logger.debug(f"get_current_centre: could not read current_centre – {exc}")
+            return None
+
+        def _valid(val):
+            return np.isfinite(val) and abs(val) < _IDS_VALID_THRESHOLD
+
+        if not (_valid(r) and _valid(z)):
+            logger.debug("get_current_centre: current_centre contains no valid data")
+            return None
+
+        return {"r": r, "z": z}
+
+    def get_scalar_annotation_quantities(self, time_slice: int) -> list:
+        """Return validated scalar global/boundary quantities for annotation display.
+
+        Reads a fixed set of scalar fields from ``global_quantities`` and
+        ``boundary``, validates each value (finite and below the IDS fill
+        value threshold), and returns
+        only those with valid data.
+
+        Args:
+            time_slice (int): Index into ``time_slice``.
+
+        Returns:
+            list of dicts, each with ``"label"`` (LaTeX str) and ``"text"``
+            (formatted value + unit str).  Empty list if nothing is valid.
+        """
+
+        def _valid(val):
+            try:
+                v = float(val)
+                return np.isfinite(v) and abs(v) < _IDS_VALID_THRESHOLD
+            except Exception:
+                return False
+
+        items = []
+        ts = self.ids.time_slice[time_slice]
+        gq = ts.global_quantities
+        bnd = ts.boundary
+
+        _specs = [
+            (lambda: float(gq.ip), lambda v: {"label": "$I_p$", "text": f"{v / 1e6:.3f} MA"}),
+            (
+                lambda: float(
+                    getattr(
+                        gq.magnetic_axis, "b_field_phi" if hasattr(gq.magnetic_axis, "b_field_phi") else "b_field_tor"
+                    )
+                ),
+                lambda v: {"label": r"$B_\phi$(axis)", "text": f"{v:.3f} T"},
+            ),
+            (lambda: float(gq.psi_axis), lambda v: {"label": r"$\psi_{\rm axis}$", "text": f"{v:.4g} Wb"}),
+            (lambda: float(gq.psi_boundary), lambda v: {"label": r"$\psi_{\rm bnd}$", "text": f"{v:.4g} Wb"}),
+            (lambda: float(gq.q_axis), lambda v: {"label": "$q_0$", "text": f"{v:.3f}"}),
+            (lambda: float(gq.q_95), lambda v: {"label": "$q_{95}$", "text": f"{v:.3f}"}),
+            (lambda: float(bnd.minor_radius), lambda v: {"label": "$a$", "text": f"{v:.3f} m"}),
+            (lambda: float(bnd.elongation), lambda v: {"label": r"$\kappa$", "text": f"{v:.3f}"}),
+            (lambda: float(bnd.triangularity), lambda v: {"label": r"$\delta$", "text": f"{v:.3f}"}),
+        ]
+        for getter, formatter in _specs:
+            try:
+                val = getter()
+                if _valid(val):
+                    items.append(formatter(val))
+            except Exception:
+                pass
+        return items
+
     def get_top_view(self, time_slice: int) -> dict:
         """
         The function returns data for plotting the top view of a 2D shape.