diff --git a/freegs4e/coil.py b/freegs4e/coil.py index e7a70f5..f89a5c7 100644 --- a/freegs4e/coil.py +++ b/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 diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 59fadac..202b007 100644 --- a/freegs4e/equilibrium.py +++ b/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): """ diff --git a/freegs4e/gradshafranov.py b/freegs4e/gradshafranov.py index 37a1c80..193a332 100644 --- a/freegs4e/gradshafranov.py +++ b/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 + ) diff --git a/freegs4e/machine.py b/freegs4e/machine.py index 6b9cbab..ae9c6ec 100644 --- a/freegs4e/machine.py +++ b/freegs4e/machine.py @@ -168,6 +168,42 @@ def createBzGreensVec(self, R, Z): pgreen += multiplier * coil.createBzGreensVec(R, Z) return pgreen + def createdBrdrGreensVec(self, R, Z): + """ + Calculate dBrdr Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createdBrdrGreensVec(R, Z) + return pgreen + + def createdBzdzGreensVec(self, R, Z): + """ + Calculate dBzdz Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createdBzdzGreensVec(R, Z) + return pgreen + + def createdBrdzGreensVec(self, R, Z): + """ + Calculate dBrdz Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createdBrdzGreensVec(R, Z) + return pgreen + + def createdBzdrGreensVec(self, R, Z): + """ + Calculate dBzdr Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createdBzdrGreensVec(R, Z) + return pgreen + def calcPsiFromGreens(self, pgreen): """ Calculate poloidal flux from Greens functions. @@ -283,7 +319,7 @@ def controlBr(self, R, Z): def controlBz(self, R, Z): """ - Calculate vertical magnetic field Br at (R,Z) due to a single unit [Amp] of current genreated by + Calculate vertical magnetic field Bz at (R,Z) due to a single unit [Amp] of current genreated by the coils in the circuit. Parameters @@ -296,7 +332,7 @@ def controlBz(self, R, Z): Returns ------- np.array - Vertical magnetic field Br at (R,Z) due to a single unit of current generated by the coils. + Vertical magnetic field Bz at (R,Z) due to a single unit of current generated by the coils. """ result = 0.0 @@ -304,6 +340,98 @@ def controlBz(self, R, Z): result += multiplier * coil.controlBz(R, Z) return result + def controldBrdr(self, R, Z): + """ + Calculate magnetic field dBrdr at (R,Z) due to a single unit [Amp] of current genreated by + the coils in the circuit. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBrdr at (R,Z) due to a single unit of current generated by the coils. + """ + + result = 0.0 + for label, coil, multiplier in self.coils: + result += multiplier * coil.controldBrdr(R, Z) + return result + + def controldBzdz(self, R, Z): + """ + Calculate magnetic field dBzdz at (R,Z) due to a single unit [Amp] of current genreated by + the coils in the circuit. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBzdz at (R,Z) due to a single unit of current generated by the coils. + """ + + result = 0.0 + for label, coil, multiplier in self.coils: + result += multiplier * coil.controldBzdz(R, Z) + return result + + def controldBrdz(self, R, Z): + """ + Calculate magnetic field dBrdz at (R,Z) due to a single unit [Amp] of current genreated by + the coils in the circuit. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBrdz at (R,Z) due to a single unit of current generated by the coils. + """ + + result = 0.0 + for label, coil, multiplier in self.coils: + result += multiplier * coil.controldBrdz(R, Z) + return result + + def controldBzdr(self, R, Z): + """ + Calculate magnetic field dBzdr at (R,Z) due to a single unit [Amp] of current genreated by + the coils in the circuit. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBzdr at (R,Z) due to a single unit of current generated by the coils. + """ + + result = 0.0 + for label, coil, multiplier in self.coils: + result += multiplier * coil.controldBzdr(R, Z) + return result + def getForces(self, equilibrium): """ Calculate the forces on the coils in the circuit. @@ -624,6 +752,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 poloidal flux from Greens functions. @@ -679,6 +831,82 @@ def Bz(self, R, Z): return self.controlBz(R, Z) * self.current + def dBrdr(self, R, Z): + """ + Calculate magnetic field dBrdr at (R,Z). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field at (R,Z) locations [T]. + """ + + return self.controldBrdr(R, Z) * self.current + + def dBzdz(self, R, Z): + """ + Calculate magnetic field dBzdz at (R,Z). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field at (R,Z) locations [T]. + """ + + return self.controldBzdz(R, Z) * self.current + + def dBrdz(self, R, Z): + """ + Calculate magnetic field dBrdz at (R,Z). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field at (R,Z) locations [T]. + """ + + return self.controldBrdz(R, Z) * self.current + + def dBzdr(self, R, Z): + """ + Calculate magnetic field dBzdr at (R,Z). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field at (R,Z) locations [T]. + """ + + return self.controldBzdr(R, Z) * self.current + def controlPsi(self, R, Z): """ Calculate poloidal flux at (R,Z) due to a single unit [Amp] of current genreated by @@ -748,6 +976,98 @@ def controlBz(self, R, Z): result += GreensBz(self.Rs, Zs, R, Z) return result + def controldBrdr(self, R, Z): + """ + Calculate radial magnetic field dBrdr at (R,Z) due to a single unit [Amp] of current genreated by + the solenoid. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Radial magnetic field dBrdr at (R,Z) due to a single unit of current generated by the solenoid. + """ + + result = 0.0 + for Zs in linspace(self.Zsmin, self.Zsmax, self.Ns): + result += GreensdBrdr(self.Rs, Zs, R, Z) + return result + + def controldBzdz(self, R, Z): + """ + Calculate vertical magnetic field dBzdz at (R,Z) due to a single unit [Amp] of current genreated by + the solenoid. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Vertical magnetic field dBzdz at (R,Z) due to a single unit of current generated by the solenoid. + """ + + result = 0.0 + for Zs in linspace(self.Zsmin, self.Zsmax, self.Ns): + result += GreensdBzdz(self.Rs, Zs, R, Z) + return result + + def controldBrdz(self, R, Z): + """ + Calculate radial magnetic field dBrdz at (R,Z) due to a single unit [Amp] of current genreated by + the solenoid. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Radial magnetic field dBrdz at (R,Z) due to a single unit of current generated by the solenoid. + """ + + result = 0.0 + for Zs in linspace(self.Zsmin, self.Zsmax, self.Ns): + result += GreensdBrdz(self.Rs, Zs, R, Z) + return result + + def controldBzdr(self, R, Z): + """ + Calculate vertical magnetic field dBzdr at (R,Z) due to a single unit [Amp] of current genreated by + the solenoid. + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Vertical magnetic field dBzdr at (R,Z) due to a single unit of current generated by the solenoid. + """ + + result = 0.0 + for Zs in linspace(self.Zsmin, self.Zsmax, self.Ns): + result += GreensdBzdr(self.Rs, Zs, R, Z) + return result + def getForces(self, equilibrium): """ Calculate the forces on the solenoid. @@ -1167,6 +1487,66 @@ def createBzGreensVec(self, R, Z, coils=None): i += 1 return pgreen + def createdBrdrGreensVec(self, R, Z, coils=None): + """ + Pre-computes the dBrdr Greens functions + and puts into arrays for each coil. + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createdBrdrGreensVec(R, Z) + i += 1 + return pgreen + + def createdBzdzGreensVec(self, R, Z, coils=None): + """ + Pre-computes the dBzdz Greens functions + and puts into arrays for each coil. + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createdBzdzGreensVec(R, Z) + i += 1 + return pgreen + + def createdBrdzGreensVec(self, R, Z, coils=None): + """ + Pre-computes the dBrdz Greens functions + and puts into arrays for each coil. + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createdBrdzGreensVec(R, Z) + i += 1 + return pgreen + + def createdBzdrGreensVec(self, R, Z, coils=None): + """ + Pre-computes the dBzdr Greens functions + and puts into arrays for each coil. + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createdBzdrGreensVec(R, Z) + i += 1 + return pgreen + def calcPsiFromGreens(self, pgreen): """ Calculate poloidal flux from Greens functions. @@ -1282,6 +1662,98 @@ def controlBz(self, R, Z): coil.controlBz(R, Z) for label, coil in self.coils if coil.control ] + def controldBrdr(self, R, Z): + """ + Calculate magnetic field dBrdr at (R,Z) due to a single unit [Amp] of current genreated by + coils in the tokamak (that have control=True). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBrdr at (R,Z) due to a single unit of current generated by the solenoid. + """ + return [ + coil.controldBrdr(R, Z) + for label, coil in self.coils + if coil.control + ] + + def controldBzdz(self, R, Z): + """ + Calculate magnetic field dBzdz at (R,Z) due to a single unit [Amp] of current genreated by + coils in the tokamak (that have control=True). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBzdz at (R,Z) due to a single unit of current generated by the solenoid. + """ + return [ + coil.controldBzdz(R, Z) + for label, coil in self.coils + if coil.control + ] + + def controldBrdz(self, R, Z): + """ + Calculate magnetic field dBrdz at (R,Z) due to a single unit [Amp] of current genreated by + coils in the tokamak (that have control=True). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBrdz at (R,Z) due to a single unit of current generated by the solenoid. + """ + return [ + coil.controldBrdz(R, Z) + for label, coil in self.coils + if coil.control + ] + + def controldBzdr(self, R, Z): + """ + Calculate magnetic field dBzdr at (R,Z) due to a single unit [Amp] of current genreated by + coils in the tokamak (that have control=True). + + Parameters + ---------- + R : float or np.array + Radial position(s) to evaluate at. + Z : float or np.array + Vertical position(s) to evaluate at. + + Returns + ------- + np.array + Magnetic field dBzdr at (R,Z) due to a single unit of current generated by the solenoid. + """ + return [ + coil.controldBzdr(R, Z) + for label, coil in self.coils + if coil.control + ] + def controlPsi(self, R, Z): """ Calculate poloidal flux at (R,Z) due to a single unit [Amp] of current genreated by diff --git a/freegs4e/multi_coil.py b/freegs4e/multi_coil.py index 12172d1..7365626 100644 --- a/freegs4e/multi_coil.py +++ b/freegs4e/multi_coil.py @@ -131,7 +131,16 @@ def __init__( self._Z_centre = np.mean(self.Zfil) def copy(self): - return type(self)( + """ + Create a copy of this object. + + Returns + ------- + new_obj : type(self) + A new instance of the same class with copied and optionally modified attributes. + """ + + new_obj = type(self)( np.copy(self.Rfil), np.copy(self.Zfil), self.current, @@ -142,6 +151,18 @@ def copy(self): self.area, ) + # copy any additional attributes if they exist on the original + extra_attrs = ["rectangle", "resistivity", "dR", "dZ"] + for key in extra_attrs: + try: + value = getattr(self, key) + setattr(new_obj, key, value) + except AttributeError: + # Skip if the attribute doesn't exist on the original + pass + + return new_obj + def controlPsi(self, R, Z): """ Calculate poloidal flux at (R,Z) due to a unit current @@ -260,11 +281,29 @@ def from_numpy_array(cls, value): def plot(self, axis=None, show=False): """ - Plot the coil including turn locations + Plot the coil geometry including the turn (filament) locations. + + This method attempts to plot the coil using the class's `plot_nke` method. + If that fails (e.g., due to a missing method or plotting error), it falls + back to plotting the coil turn positions directly using `matplotlib`. + + Parameters + ---------- + axis : matplotlib.axes.Axes, optional + The Matplotlib axis to plot on. If None, a new figure and axis are created. + show : bool, optional + Whether to call `plt.show()` after plotting. Default is False. + + Returns + ------- + axis : matplotlib.axes.Axes + The axis object containing the plot. + """ + try: axis = self.plot_nke(axis, show) - except: + except Exception as e: if axis is None: fig = plt.figure() axis = fig.add_subplot(111) diff --git a/freegs4e/plotting.py b/freegs4e/plotting.py index 4c7cfb7..0f4cb08 100644 --- a/freegs4e/plotting.py +++ b/freegs4e/plotting.py @@ -1,5 +1,5 @@ """ -Plotting using matplotlib +Plot aspects of the tokamak and equilibrium. Copyright 2024 Nicola C. Amorisco, Adriano Agnello, George K. Holt, Ben Dudson. @@ -18,13 +18,42 @@ """ +import matplotlib.pyplot as plt +import numpy as np from numpy import amax, amin, linspace from . import critical def plotCoils(coils, axis=None): - import matplotlib.pyplot as plt + """ + Plot the geometry of magnetic field coils. + + This function prepares a plot of the coil set used in a tokamak or other + magnetic confinement configuration. It creates a Matplotlib axis if one + is not provided, allowing additional geometric or field features to be + overlaid later. + + Parameters + ---------- + coils : object + Coil geometry object containing spatial information about the magnetic + coils (e.g., current centerlines or cross-sections). The exact structure + of `coils` depends on the specific equilibrium or control framework. + axis : matplotlib.axes.Axes, optional + Axis object on which to plot. If `None`, a new figure and axis are created. + + Returns + ------- + axis : matplotlib.axes.Axes + The matplotlib axis prepared for coil plotting. + + Notes + ----- + - This function currently initializes the plot but does not draw coil shapes. + It can be extended to plot coil positions or outlines using data from + `coils`. + """ if axis is None: fig = plt.figure() @@ -35,14 +64,44 @@ def plotCoils(coils, axis=None): def plotConstraints(control, axis=None, show=True): """ - Plots constraints used for coil current control + Plot magnetic control constraints such as X-points and isoflux surfaces. - axis - Specify the axis on which to plot - show - Call matplotlib.pyplot.show() before returning + This function visualizes the spatial constraints used for coil current + optimization or magnetic equilibrium control, including X-point locations + and isoflux surface definitions. These constraints are typically used in + inverse equilibrium solvers or feedback control algorithms. - """ + Parameters + ---------- + control : object + Control data object containing constraint information with attributes: + - `xpoints` : list of (R, Z) tuples + Target locations for X-point constraints. + - `isoflux` : list of (R1, Z1, R2, Z2) tuples + Pairs of points defining line segments that should lie on the same + poloidal flux surface. + axis : matplotlib.axes.Axes, optional + Axis object on which to plot. If `None`, a new figure and axis are created. + show : bool, default=True + If True, call `matplotlib.pyplot.show()` before returning. - import matplotlib.pyplot as plt + Returns + ------- + axis : matplotlib.axes.Axes + The matplotlib axis containing the plotted constraints. + + Notes + ----- + - X-point constraints are plotted as blue crosses (`"bx"`). + - Isoflux constraints are plotted as dashed blue line segments with triangle + markers (`":b^"`). + - A legend is displayed automatically if `show=True`. + + Examples + -------- + axis = plotConstraints(control) + axis.set_title("Coil Control Constraints") + """ if axis is None: fig = plt.figure() @@ -71,42 +130,117 @@ def plotConstraints(control, axis=None, show=True): def plotIOConstraints(control, axis=None, show=True): """ - Plots constraints used for coil current control + Plot magnetic control constraints for equilibrium reconstruction. - axis - Specify the axis on which to plot - show - Call matplotlib.pyplot.show() before returning + This function visualizes the geometric constraints used in coil current + control or equilibrium reconstruction, including magnetic null points and + isoflux target locations. These constraints are typically used to ensure + that a magnetic configuration satisfies specified plasma shape or field + conditions. - """ - - import matplotlib - import matplotlib.pyplot as plt + Parameters + ---------- + control : object + An object containing magnetic control data with attributes: + - `null_points` : tuple of array-like + Coordinates (R, Z) of magnetic null points, if defined. + - `isoflux_set` : list of array-like + A list of isoflux constraint points, where each element is a pair + of arrays (R, Z) defining a set of target locations for equal + poloidal flux. + axis : matplotlib.axes.Axes, optional + Axis object on which to plot. If `None`, a new figure and axis are created. + show : bool, default=True + If True, call `matplotlib.pyplot.show()` before returning. - cmap = matplotlib.cm.get_cmap("gnuplot") + Returns + ------- + axis : matplotlib.axes.Axes + The matplotlib axis containing the plotted constraints. + + Notes + ----- + - Null points are plotted as purple downward tick markers (`"1"`). + - Isoflux constraints are plotted as plus markers (`"+"`) in distinct colors. + - A legend is automatically displayed in the upper-right corner if `show=True`. + + Examples + -------- + axis = plotIOConstraints(control) + axis.set_title("Magnetic Control Constraints: Null Points and Isoflux Sets") + """ + # axes if axis is None: fig = plt.figure() axis = fig.add_subplot(111) - # Locations of the X-points + # locations of the null points if control.null_points is not None: axis.plot( control.null_points[0], control.null_points[1], "1", - color="purple", + color="m", markersize=10, + markeredgewidth=2.0, ) axis.plot( - [], [], "1", color="purple", markersize=10, label="Null-points" + [], + [], + "1", + color="m", + markersize=10, + markeredgewidth=2.0, + label="Null points", ) - # Isoflux surfaces + # locations of the 2nd order null points + if control.null_points_2nd_order is not None: + axis.plot( + control.null_points_2nd_order[0], + control.null_points_2nd_order[1], + "1", + color="orange", + markersize=10, + markeredgewidth=2.0, + ) + axis.plot( + [], + [], + "1", + color="orange", + markersize=10, + markeredgewidth=2.0, + label="2nd order null points", + ) + + # plot isoflux constraints if control.isoflux_set is not None: - color = cmap(np.random.random()) + color = [ + "tab:brown", + "tab:orange", + "tab:pink", + "tab:gray", + "tab:olive", + ] for i, isoflux in enumerate(control.isoflux_set): - axis.plot(isoflux[0], isoflux[1], "+", color=color, markersize=10) axis.plot( - [], [], "+", color=color, label=f"Isoflux_{i}", markersize=10 + isoflux[0], + isoflux[1], + "+", + color=color[i % len(color)], + markersize=10, + markeredgewidth=2.0, + ) + axis.plot( + [], + [], + "+", + color=color[i % len(color)], + label=f"Isoflux set ({i})", + markersize=10, + markeredgewidth=2.0, ) if show: @@ -117,126 +251,204 @@ def plotIOConstraints(control, axis=None, show=True): def plotEquilibrium( - eq, axis=None, show=True, oxpoints=True, wall=True, limiter=True + eq, + axis=None, + xpoints=True, + opoints=True, + wall=True, + limiter=True, + legend=False, + show=True, ): """ - Plot the equilibrium flux surfaces + Plot poloidal magnetic ψ(R, Z) flux surfaces and key geometric features of an equilibrium in + FreeGSNKE. It can also overlay the separatrix, magnetic X- and + O-points, the plasma-facing wall, and the limiter geometry. - axis - Specify the axis on which to plot - show - Call matplotlib.pyplot.show() before returning - oxpoints - Plot X points as red circles, O points as green circles - wall - Plot the wall (limiter) + Parameters + ---------- + eq : object + An equilibrium object containing fields `R`, `Z`, and `psi()`, as well + as profile data (`_profiles.opt`, `_profiles.xpt`) and tokamak geometry + (`tokamak.wall`, `tokamak.limiter`). + axis : matplotlib.axes.Axes, optional + Axis object to plot on. If `None`, a new figure and axis are created. + xpoints : bool, default=True + If True, plot magnetic X-points (red 'x' markers). + opoints : bool, default=True + If True, plot magnetic O-points (green '*' markers). + wall : bool, default=True + If True, plot the tokamak wall outline as a solid black curve. + limiter : bool, default=True + If True, plot the limiter outline as a dashed black curve. + 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 + ----- + - The function assumes that the equilibrium has already been solved; + if not, it will print a warning message. + - The separatrix is plotted using a solid red line for diverted equilibria + and a dashed red line for limited equilibria. + + Examples + -------- + axis = plotEquilibrium(eq, legend=True) + axis.set_title("Tokamak Equilibrium Flux Surfaces") """ - import matplotlib.pyplot as plt - - psi = eq.psi() - + # extract data + try: + psi = eq.psi() + opt = eq._profiles.opt + xpt = eq._profiles.xpt + except AttributeError as e: + raise RuntimeError( + "This equilibrium has not been solved: please solve for an equilibrium first!" + ) from e + + # axes if axis is None: fig = plt.figure() axis = fig.add_subplot(111) - - levels = linspace(amin(psi), amax(psi), 50) - - axis.contour(eq.R, eq.Z, psi, levels=levels) axis.set_aspect("equal") axis.set_xlabel("Major radius [m]") axis.set_ylabel("Height [m]") - try: - eq._profiles - - if oxpoints: - # Add O- and X-points - # opt, xpt = critical.find_critical(eq.R, eq.Z, psi) - opt = eq._profiles.opt - xpt = eq._profiles.xpt - - for r, z, _ in xpt: - axis.plot(r, z, "rx") - for r, z, _ in opt: - axis.plot(r, z, "gx") - - if xpt is not []: - psi_bndry = xpt[0][2] - if eq._profiles.flag_limiter: - axis.contour( - eq.R, - eq.Z, - psi, - levels=[eq._profiles.psi_bndry], - colors="k", - ) - axis.contour( - eq.R, - eq.Z, - psi, - levels=[psi_bndry], - colors="r", - linestyles="dashed", - ) - # cs = plt.contour(eq.R, eq.Z, psi, levels=[eq._profiles.psi_bndry], alpha=0) - # paths = cs.collections[0].get_paths() - # for path in paths: - # vertices = path.vertices - # if np.sum(vertices[0] == vertices[-1])>1: - # axis.plot(vertices[:,0], vertices[:,1], 'k') - - else: - axis.contour( - eq.R, eq.Z, psi, levels=[psi_bndry], colors="r" - ) - - # Add legend - axis.plot([], [], "rx", label="X-points") - axis.plot([], [], "r", label="Separatrix") - if opt is not []: - axis.plot([], [], "gx", label="O-points") - - except: - print( - "This equilibrium has not been solved: the separatrix can not be drawn." + # plot flux contours + levels = linspace(amin(psi), amax(psi), 50) + axis.contour(eq.R, eq.Z, psi, levels=levels) + + # plot separatrix (from primary X-point) + if eq._profiles.flag_limiter: + label = "LCFS (limited plasma)" + else: + label = "Separatrix (primary X-point)" + colour = "r" + style = "solid" + axis.contour( + eq.R, eq.Z, psi, levels=[eq.psi_bndry], colors=colour, linestyles=style + ) + axis.plot([], [], colour, label=label, linestyle=style) + + # plot extra separatrix (LCFS) if plasma limited + if eq._profiles.flag_limiter: + colour = "k" + style = "dashed" + axis.contour( + eq.R, + eq.Z, + psi, + levels=[xpt[0][2]], + colors=colour, + linestyles=style, + ) + axis.plot( + [], + [], + colour, + label="Separatrix (primary X-point)", + linestyle=style, + ) + + # plot x point + if xpoints: + for r, z, _ in xpt: + axis.plot(r, z, "rx", markersize=9) + axis.plot( + xpt[0][0], xpt[0][1], "rx", markersize=9, markeredgewidth=2.5 + ) + axis.plot([], [], "rx", markersize=9, label="X-points") + axis.plot( + [], + [], + "rx", + markersize=9, + markeredgewidth=2.5, + label="X-point (primary)", ) - print("Please solve first for a plot of the critical points.") + # plot op points + if opoints: + for r, z, _ in opt: + axis.plot(r, z, "g2", markersize=9) + axis.plot([], [], "g2", markersize=9, label="O-points") + + # plot wall if wall and eq.tokamak.wall and len(eq.tokamak.wall.R): axis.plot( list(eq.tokamak.wall.R) + [eq.tokamak.wall.R[0]], list(eq.tokamak.wall.Z) + [eq.tokamak.wall.Z[0]], "k", ) + + # plot limiter if limiter and eq.tokamak.limiter and len(eq.tokamak.limiter.R): axis.plot( list(eq.tokamak.limiter.R) + [eq.tokamak.limiter.R[0]], list(eq.tokamak.limiter.Z) + [eq.tokamak.limiter.Z[0]], - "k--", - lw=0.5, + "k:", ) + if legend: + axis.legend(loc="upper right") + if show: - plt.legend() plt.show() return axis -import numpy as np - - def make_broad_mask(mask, layer_size=1): - """Enlarges a mask with the points just outside the input, with a width=`layer_size` + """ + Expand a binary mask by adding a border layer of specified thickness. + + This function broadens a 2D boolean or binary mask by marking all points + that lie within a specified Manhattan distance (`layer_size`) from the + original masked region. The result is a new mask that includes both the + original region and its expanded boundary. Parameters ---------- + mask : np.ndarray + 2D boolean or binary array representing the region to be expanded. layer_size : int, optional - Width of the layer outside the limiter, by default 3 + Number of pixels by which to expand the mask in all directions. + Default is 1. Returns ------- layer_mask : np.ndarray - Mask of the points outside the limiter within a distance of `layer_size` from the limiter + Boolean array of the same shape as `mask`, where `True` values + correspond to the original region plus an outer layer of width + `layer_size`. + + Notes + ----- + - The expansion is performed by convolving the mask with a square kernel + of size `(2 * layer_size + 1)`, effectively performing a morphological + dilation. + - This operation is useful for defining buffer regions (e.g., just outside + a limiter or wall in equilibrium simulations). + + Examples + -------- + mask = np.array([[0, 1, 0], + ... [0, 1, 0], + ... [0, 0, 0]], dtype=bool) + make_broad_mask(mask, layer_size=1) + array([[ True, True, True], + [ True, True, True], + [ False, True, False]]) """ + nx, ny = np.shape(mask) layer_mask = np.zeros( np.array([nx, ny]) + 2 * np.array([layer_size, layer_size]) @@ -256,17 +468,51 @@ def plotProbes( probes, axis=None, show=True, floops=True, pickups=True, pickups_scale=0.05 ): """ - Plot the fluxloops and pickup coils. + Plot magnetic diagnostics including flux loops and pickup coils. - axis - Specify the axis on which to plot - show - Call matplotlib.pyplot.show() before returning - floops - Plot the floops - pickups - Plot the pickups + This function visualises diagnostic probe locations used for equilibrium + reconstruction, showing the spatial distribution of flux loops and pickup + coils. Pickup coil orientations are indicated by small line segments + scaled by a user-defined factor. - """ + Parameters + ---------- + probes : object + An object containing probe geometry data with the following attributes: + - `floop_pos` : ndarray of shape (N, 2) + Positions of flux loops in (R, Z) coordinates. + - `pickup_pos` : ndarray of shape (M, 3) + Positions of pickup coils in (R, φ, Z) coordinates (only R and Z used). + - `pickup_or` : ndarray of shape (M, 3) + Orientation vectors of the pickup coils. + axis : matplotlib.axes.Axes, optional + Axis object on which to plot. If `None`, a new figure and axis are created. + show : bool, default=True + If True, call `matplotlib.pyplot.show()` before returning. + floops : bool, default=True + If True, plot flux loop probe locations as orange diamonds. + pickups : bool, default=True + If True, plot pickup coil locations as brown circles and their orientations + as short brown line segments. + pickups_scale : float, default=0.05 + Scaling factor applied to pickup coil orientation vectors for visualization. - import matplotlib as mpl - import matplotlib.pyplot as plt + Returns + ------- + axis : matplotlib.axes.Axes + The matplotlib axis containing the plotted probes. + + Notes + ----- + - The function assumes `probes.floop_pos`, `probes.pickup_pos`, and + `probes.pickup_or` are NumPy arrays with appropriate dimensions. + - The orientation lines are plotted using (R, Z) components only. + + Examples + -------- + axis = plotProbes(probes, pickups_scale=0.1) + axis.set_title("Magnetic Diagnostics: Flux Loops and Pickup Coils") + """ # create axis if none exists if axis is None: @@ -309,20 +555,8 @@ def plotProbes( s=3, ) - # # pickup orientation - - # # Calculate the angle in radians and convert to degrees - # angle_radians = np.arctan2(probes.pickup_or[:, 2], probes.pickup_or[:, 0]) - # angle_degrees = np.degrees(angle_radians) - - # # make a markerstyle class instance and modify its transform prop - # for i in range(0, len(probes.pickup_pos[:, 0])): - # t = mpl.markers.MarkerStyle(marker=u'$\u21A6$') # "$>$" u'$\u2192$ - # t._transform = t.get_transform().rotate_deg(angle_degrees[i]) - # axis.scatter(probes.pickup_pos[i, 0], probes.pickup_pos[i, 2], marker=t, s=40, color="brown") - if show: - plt.legend() + plt.legend(loc="upper right") plt.show() return axis diff --git a/pyproject.toml b/pyproject.toml index 524ee9b..479e338 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ Homepage = "https://github.com/FusionComputingLab/freegs4e" [tool.isort] profile = "black" +line_length = 79 [tool.black] line-length = 79 \ No newline at end of file