Merge pull request #52 from bendudson/geqdsk-read-fixes

Geqdsk read fixes

freegs/__init__.py

@@ -42,3 +42,5 @@
 from .picard import solve
 
 from .dump import OutputFile
+
+from . import plotting

freegs/control.py

@@ -190,7 +190,7 @@ class ConstrainPsi2D(object):
     Ignores constant offset differences between psi array
     """
 
-    def __init__(self, target_psi, weights=1.0):
+    def __init__(self, target_psi, weights=None):
         """
         target_psi : 2D (R,Z) array
             Must be the same size as the equilibrium psi
@@ -201,6 +201,9 @@ def __init__(self, target_psi, weights=1.0):
             Set points to zero to ignore them in fitting.
 
         """
+        if weights is None:
+            weights = np.full(target_psi.shape, 1.0)
+
         # Remove the average so constant offsets are ignored
         self.target_psi = target_psi - np.average(target_psi, weights=weights)
 
@@ -283,10 +286,15 @@ def psinorm_difference(self, currents, eq):
 
         opt, xpt = critical.find_critical(eq.R, eq.Z, psi)
         if not opt:
+            print("No O-points found!")
+            print(opt, xpt)
+            eq.plot()
             raise ValueError("No O-points found!")
         psi_axis = opt[0][2]
 
         if not xpt:
+            print("No X-points found!")
+            eq.plot()
             raise ValueError("No X-points found")
         psi_bndry = xpt[0][2]
 

freegs/geqdsk.py

@@ -165,10 +165,12 @@ def read(
     ntheta=8,
     show=False,
     axis=None,
+    pause=0.0001,
     cocos=1,
     domain=None,
     blend=0.0,
     fit_sol=False,
+    maxits=50,
 ):
     """
     Reads a G-EQDSK format file
@@ -184,6 +186,8 @@ def read(
         Set to true to show solution in a new figure
     axis : Matplotlib Axis object
        Set to an axis for plotting. Implies show=True
+    pause : float
+       Delay between output plots. If negative, waits for window to be closed
     cocos : integer
         COordinate COnventions. Not fully handled yet,
         only whether psi is divided by 2pi or not.
@@ -203,6 +207,9 @@ def read(
         outside the separatrix.
         If True, the whole domain is used in the fitting.
         This is useful if the locations of strike points need to be constrained.
+    maxits : integer
+        Maximum number of iterations. Set to None for no limit.
+        If this limit is exceeded then a RuntimeError is raised.
 
     A nonlinear solve will be performed, using Picard iteration
 
@@ -319,10 +326,6 @@ def q_func(psinorm):
     # 1 = plasma boundary
     psi_norm = clip((psi - psi_axis) / (psi_bndry - psi_axis), 0.0, 1.1)
 
-    # Create masking function: 1 inside plasma, 0 outside
-    mask = np.ones(psi.shape)
-    mask[psi_norm > 1.0 - 1e-6] = 0.0  # Ignore areas outside the plasma
-
     # Create an Equilibrium object
     eq = Equilibrium(
         tokamak=machine,
@@ -337,6 +340,10 @@ def q_func(psinorm):
     dR = eq.R[1, 0] - eq.R[0, 0]
     dZ = eq.Z[0, 1] - eq.Z[0, 0]
 
+    # Create masking function: 1 inside plasma, 0 outside
+    opoint, xpoint = critical.find_critical(eq.R, eq.Z, psi)
+    mask = critical.core_mask(eq.R, eq.Z, psi, opoint, xpoint, psi_bndry)
+
     # Toroidal current
     Jtor = eq.R * pprime_func(psi_norm) + ffprime_func(psi_norm) / (eq.R * mu0)
     Jtor *= mask
@@ -378,8 +385,8 @@ def q_func(psinorm):
         psi_norm = clip((psi - psi_axis) / (psi_bndry - psi_axis), 0.0, 1.0)
 
         # Create masking function: 1 inside plasma, 0 outside
-        mask = np.ones(psi.shape)
-        mask[psi_norm > 1.0 - 1e-6] = 0.0  # Ignore areas outside the plasma
+        opoint, xpoint = critical.find_critical(eq.R, eq.Z, psi)
+        mask = critical.core_mask(eq.R, eq.Z, psi, opoint, xpoint, psi_bndry)
 
     # Note: Here we have
     #   eq : Equilibrium object
@@ -454,8 +461,10 @@ def q_func(psinorm):
         controlsystem,
         show=show,
         axis=axis,
+        pause=pause,
         rtol=rtol,
-        blend=0.5,
+        blend=blend,
+        maxits=maxits,
     )
 
     print(