Improving plotting routines

- Control system constraints can now be plotted, including
  isoflux pairs

- Plotting equilibrium and constraints includes a legend

- Equilibrium and Control objects have a "plot" member
  which imports then calls the relevant routine in plotting.py

  This should ensure that the code only imports matplotlib
  if a plot function is called.

01-freeboundary.py

@@ -62,5 +62,6 @@
 ##############################################
 # Final plot
 
-from freegs.plotting import plotEquilibrium
-plotEquilibrium(eq)
+axis = eq.plot(show=False)
+constrain.plot(axis=axis, show=True)
+

06-xpoints.py

@@ -29,7 +29,7 @@
 
 ax = plotEquilibrium(eq, show=False)
 plotCoils(tokamak.coils, axis=ax)
-plotConstraints(xpoints, axis=ax)
+plotConstraints(control, axis=ax)
 plt.show()
 
 ##########################################################

freegs/control.py

@@ -116,6 +116,17 @@ def __call__(self, eq):
         print("Current changes: " + str(current_change))
         tokamak.controlAdjust(current_change)
 
+    def plot(self, axis=None, show=True):
+        """
+        Plots constraints used for coil current control
+        
+        axis     - Specify the axis on which to plot
+        show     - Call matplotlib.pyplot.show() before returning
+        
+        """
+        from .plotting import plotConstraints
+        return plotConstraints(self, axis=axis, show=show)
+
 
 def flux_surface(R0, Z0, a, elongation=0.0, triangularity=0.0, indentation=0.0, n=20):
     """

freegs/equilibrium.py

@@ -294,7 +294,18 @@ def _updatePlasmaPsi(self, plasma_psi):
         # Update spline interpolation
         self.psi_func = interpolate.RectBivariateSpline(self.R[:,0], self.Z[0,:], plasma_psi)
 
-
+    def plot(self, axis=None, show=True, oxpoints=True):
+        """
+        Plot the equilibrium flux surfaces
+        
+        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
+    
+        """
+        from .plotting import plotEquilibrium
+        return plotEquilibrium(self, axis=axis, show=show, oxpoints=oxpoints)
+
 
 def refine(eq):
     """

freegs/plotting.py

@@ -25,11 +25,56 @@
 from numpy import linspace, amin, amax
 from . import critical
 
+
+def plotCoils(coils, axis=None):
+    if axis is None:
+        fig = plt.figure()
+        axis = fig.add_subplot(111)
+
+
+    return axis
+
+
+def plotConstraints(control, axis=None, show=True):
+    """
+    Plots constraints used for coil current control
+    
+    axis     - Specify the axis on which to plot
+    show     - Call matplotlib.pyplot.show() before returning
+    
+    """
+
+    if axis is None:
+        fig = plt.figure()
+        axis = fig.add_subplot(111)
+
+    # Locations of the X-points
+    for r,z in control.xpoints:
+        axis.plot(r, z, 'bx')
+
+    if control.xpoints:
+        axis.plot([], [], 'bx', label="X-point constraints")
+
+    # Isoflux surfaces
+    for r1,z1, r2,z2 in control.isoflux:
+        axis.plot([r1,r2], [z1,z2], ':b^')
+
+    if control.isoflux:
+        axis.plot([], [], ':b^', label="Isoflux constraints")
+
+    if show:
+        plt.legend()
+        plt.show()
+
+    return axis
+
 def plotEquilibrium(eq, axis=None, show=True, oxpoints=True):
     """
     Plot the equilibrium flux surfaces
 
-    axis - Specify the axis on which to plot
+    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
     
     """
 
@@ -60,32 +105,16 @@ def plotEquilibrium(eq, axis=None, show=True, oxpoints=True):
         if xpt:
             psi_bndry = xpt[0][2]
             sep_contour=axis.contour(eq.R, eq.Z,psi, levels=[psi_bndry], colors='r')
-
+
+            # Add legend
+            axis.plot([], [], 'ro', label="X-points")
+            axis.plot([], [], 'r', label="Separatrix")
+        if opt:
+            axis.plot([], [], 'go', label="O-points")
+
     if show:
+        plt.legend()
         plt.show()
 
     return axis
 
-def plotCoils(coils, axis=None):
-    if axis is None:
-        fig = plt.figure()
-        axis = fig.add_subplot(111)
-
-
-    return axis
-
-
-def plotConstraints(xpoints, axis=None):
-    """
-    Plots constraints used for coil currents
-    
-    """
-
-    if axis is None:
-        fig = plt.figure()
-        axis = fig.add_subplot(111)
-
-    for xpt in xpoints:
-        axis.plot(xpt[0], xpt[1], 'ro')    
-
-    return axis