Plotting updates

freegs4e/equilibrium.py

@@ -37,6 +37,7 @@
     GSsparse4thOrder,
     mu0,
 )
+from .plotting import plotEquilibrium
 
 
 class Equilibrium:
@@ -2737,29 +2738,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/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)