Plotting updates

freegs4e/coil.py

@@ -30,7 +30,16 @@
 import numpy as np
 from matplotlib.patches import Rectangle
 
-from .gradshafranov import Greens, GreensBr, GreensBz, mu0
+from .gradshafranov import (
+    Greens,
+    GreensBr,
+    GreensBz,
+    GreensdBrdr,
+    GreensdBrdz,
+    GreensdBzdr,
+    GreensdBzdz,
+    mu0,
+)
 
 
 class AreaCurrentLimit:
@@ -166,6 +175,30 @@ def createBzGreensVec(self, R, Z):
         """
         return self.controlBz(R, Z)
 
+    def createdBrdrGreensVec(self, R, Z):
+        """
+        Calculate dBrdr Greens functions
+        """
+        return self.controldBrdr(R, Z)
+
+    def createdBzdzGreensVec(self, R, Z):
+        """
+        Calculate dBzdz Greens functions
+        """
+        return self.controldBzdz(R, Z)
+
+    def createdBrdzGreensVec(self, R, Z):
+        """
+        Calculate dBrdz Greens functions
+        """
+        return self.controldBrdz(R, Z)
+
+    def createdBzdrGreensVec(self, R, Z):
+        """
+        Calculate dBzdr Greens functions
+        """
+        return self.controldBzdr(R, Z)
+
     def calcPsiFromGreens(self, pgreen):
         """
         Calculate plasma psi from Greens functions and current
@@ -192,6 +225,46 @@ def Bz(self, R, Z):
             bz = 0
         return bz
 
+    def dBrdr(self, R, Z):
+        """
+        Calculate magnetic field dBrdr at (R,Z)
+        """
+        if np.abs(self.current) > 1e-5:
+            br = self.controldBrdr(R, Z) * self.current
+        else:
+            br = 0
+        return br
+
+    def dBzdz(self, R, Z):
+        """
+        Calculate magnetic field dBzdz at (R,Z)
+        """
+        if np.abs(self.current) > 1e-5:
+            bz = self.controldBzdz(R, Z) * self.current
+        else:
+            bz = 0
+        return bz
+
+    def dBrdz(self, R, Z):
+        """
+        Calculate magnetic field dBrdz at (R,Z)
+        """
+        if np.abs(self.current) > 1e-5:
+            br = self.controldBrdz(R, Z) * self.current
+        else:
+            br = 0
+        return br
+
+    def dBzdr(self, R, Z):
+        """
+        Calculate magnetic field dBzdr at (R,Z)
+        """
+        if np.abs(self.current) > 1e-5:
+            bz = self.controldBzdr(R, Z) * self.current
+        else:
+            bz = 0
+        return bz
+
     def controlPsi(self, R, Z):
         """
         Calculate poloidal flux at (R,Z) due to a unit current
@@ -210,6 +283,30 @@ def controlBz(self, R, Z):
         """
         return GreensBz(self.R, self.Z, R, Z) * self.turns
 
+    def controldBrdr(self, R, Z):
+        """
+        Calculate magnetic field dBrdr at (R,Z) due to a unit current
+        """
+        return GreensdBrdr(self.R, self.Z, R, Z) * self.turns
+
+    def controldBzdz(self, R, Z):
+        """
+        Calculate magnetic field dBzdz at (R,Z) due to a unit current
+        """
+        return GreensdBzdz(self.R, self.Z, R, Z) * self.turns
+
+    def controldBrdz(self, R, Z):
+        """
+        Calculate magnetic field dBrdz at (R,Z) due to a unit current
+        """
+        return GreensdBrdz(self.R, self.Z, R, Z) * self.turns
+
+    def controldBzdr(self, R, Z):
+        """
+        Calculate magnetic field dBzdr at (R,Z) due to a unit current
+        """
+        return GreensdBzdr(self.R, self.Z, R, Z) * self.turns
+
     def getForces(self, equilibrium):
         """
         Calculate forces on the coils in Newtons

freegs4e/equilibrium.py

@@ -37,6 +37,7 @@
     GSsparse4thOrder,
     mu0,
 )
+from .plotting import plotEquilibrium
 
 
 class Equilibrium:
@@ -2758,29 +2759,69 @@ def _updatePlasmaPsi(self, plasma_psi):
             self.psi_bndry = None
             self.mask = None
 
-    def plot(self, axis=None, show=True, oxpoints=True):
-        """
-        Plot the equilibrium.
-
-        Parameters
-        ----------
-        axis: object
-            Matplotlib axis object.
-        show: bool
-            Calls matplotlib.pyplot.show() before returning.
-        oxpoints: bool
-            Plot X points and O points.
-
-        Returns
-        -------
-        axis
-            Matplotlib axis object.
-
+    def plot(
+        self,
+        axis=None,
+        xpoints=True,
+        opoints=True,
+        wall=True,
+        limiter=True,
+        legend=False,
+        show=True,
+    ):
         """
-
-        from .plotting import plotEquilibrium
-
-        return plotEquilibrium(self, axis=axis, show=show, oxpoints=oxpoints)
+        Plot the magnetic equilibrium flux surfaces and key geometric features.
+
+        This method provides a convenient interface to visualize the equilibrium
+        stored in the current object. It calls the `plotEquilibrium` function to
+        display magnetic flux contours, separatrix, magnetic nulls (X- and O-points),
+        and the tokamak boundary geometry (wall and limiter).
+
+        Parameters
+        ----------
+        axis : matplotlib.axes.Axes, optional
+            Axis object on which to plot. If `None`, a new figure and axis are created.
+        xpoints : bool, default=True
+            If True, plot magnetic X-points as red 'x' markers.
+        opoints : bool, default=True
+            If True, plot magnetic O-points as green '*' markers.
+        wall : bool, default=True
+            If True, plot the tokamak wall outline as a solid black line.
+        limiter : bool, default=True
+            If True, plot the limiter outline as a dashed black line.
+        legend : bool, default=False
+            If True, display a legend describing plotted features.
+        show : bool, default=True
+            If True, call `matplotlib.pyplot.show()` before returning.
+
+        Returns
+        -------
+        axis : matplotlib.axes.Axes
+            The matplotlib axis containing the plotted equilibrium.
+
+        Notes
+        -----
+        - This method is a wrapper around `plotEquilibrium(self, ...)`.
+        - The equilibrium must be solved before plotting; otherwise, a warning
+        will be printed.
+
+        Examples
+        --------
+            eq.solve()
+            axis = eq.plot(legend=True)
+            axis.set_title("Plasma Equilibrium Flux Surfaces")
+        """
+
+        return plotEquilibrium(
+            self,
+            axis=axis,
+            xpoints=xpoints,
+            opoints=opoints,
+            wall=wall,
+            limiter=limiter,
+            legend=legend,
+            show=show,
+        )
 
     def _separatrix_metrics(self):
         """

freegs4e/gradshafranov.py

@@ -627,3 +627,67 @@ def GreensdBrdr(Rc, Zc, R, Z, eps=2e-3):
     return (GreensBr(Rc, Zc, R + eps, Z) - GreensBr(Rc, Zc, R - eps, Z)) / (
         2.0 * eps
     )
+
+
+def GreensdBrdz(Rc, Zc, R, Z, eps=2e-3):
+    """
+    Calculate vertical derivative of radial magnetic field at (R,Z) due to a
+    single unit of current at (Rc,Zc) using Greens function for the
+    elliptic operator above:
+
+        dBr/dZ (R,Z) = (Br(R, Z + eps) - Br(R, Z - eps))/ 2 * eps.
+
+    Parameters
+    ----------
+    Rc : float
+        Radial position where current is located [m].
+    Zc : float
+        Vertical position where current is located [m].
+    R : float
+        Radial position where poloidal flux is to be calcualted [m].
+    Z : float
+        Vertical position where poloidal flux is to be calcualted [m].
+    eps : float
+        Small step size for numerical differentiation in the radial direction [m].
+
+    Returns
+    -------
+    float
+        Value of the derivative at (R,Z) [T/m].
+    """
+
+    return (GreensBr(Rc, Zc, R, Z + eps) - GreensBr(Rc, Zc, R, Z - eps)) / (
+        2.0 * eps
+    )
+
+
+def GreensdBzdr(Rc, Zc, R, Z, eps=2e-3):
+    """
+    Calculate radial derivative of vertical magnetic field at (R,Z) due to a
+    single unit of current at (Rc,Zc) using Greens function for the
+    elliptic operator above:
+
+        dBz/dR (R,Z) = (Bz(R + eps, Z) - Bz(R - eps, Z))/ 2 * eps.
+
+    Parameters
+    ----------
+    Rc : float
+        Radial position where current is located [m].
+    Zc : float
+        Vertical position where current is located [m].
+    R : float
+        Radial position where poloidal flux is to be calcualted [m].
+    Z : float
+        Vertical position where poloidal flux is to be calcualted [m].
+    eps : float
+        Small step size for numerical differentiation in the vertical direction [m].
+
+    Returns
+    -------
+    float
+        Value of the derivative at (R,Z) [T/m].
+    """
+
+    return (GreensBz(Rc, Zc, R + eps, Z) - GreensBz(Rc, Zc, R - eps, Z)) / (
+        2.0 * eps
+    )