Minor fixes

freegs/equilibrium.py

@@ -338,18 +338,30 @@ def q(self, psinorm=None, npsi=100):
         >>> psinorm, q = eq.q()
 
         Note: psinorm = 0 is the magnetic axis, and psinorm = 1 is the separatrix.
-              Calculating q on either of these flux surfaces is problematic,
-              and the results will probably not be accurate.
+              If you request a value close to either of these limits, an extrapolation
+              based on a 1D grid of values from 0.01 to 0.99 will be used. This gives
+              smooth and continuous q-profiles, but comes at an increased computational
+              cost.
         """
         if psinorm is None:
             # An array which doesn't include psinorm = 0 or 1
             psinorm = linspace(1.0 / (npsi + 1), 1.0, npsi, endpoint=False)
             return psinorm, critical.find_safety(self, psinorm=psinorm)
 
-        result = critical.find_safety(self, psinorm=psinorm)
+        elif np.any((psinorm < 0.01) | (psinorm > 0.99)):
+            psinorm_inner = np.linspace(0.01, 0.99, num=npsi)
+            q_inner = critical.find_safety(self, psinorm=psinorm_inner)
+
+            interp = interpolate.interp1d(psinorm_inner, q_inner,
+                                          kind = "quadratic",
+                                          fill_value = "extrapolate")
+            result = interp(psinorm)
+        else:
+            result = critical.find_safety(self, psinorm=psinorm)
+
         # Convert to a scalar if only one result
-        if len(result) == 1:
-            return result[0]
+        if np.size(result) == 1:
+            return float(result)
         return result
 
     def tor_flux(self, psi=None):

freegs/gradshafranov.py

@@ -38,7 +38,7 @@ class GSElliptic:
     Represents the Grad-Shafranov elliptic operator
 
     .. math::
-        \Delta^* = R^2 \nabla\cdot\frac{1}{R^2}\nabla
+        \\Delta^* = R^2 \\nabla\\cdot\\frac{1}{R^2}\\nabla
 
     which is
 

freegs/optimise.py

@@ -27,7 +27,7 @@
 import matplotlib.pyplot as plt
 from freegs.plotting import plotEquilibrium
 
-from math import sqrt
+from numpy import sqrt, inf
 
 # Measures which operate on Equilibrium objects
 
@@ -199,7 +199,7 @@ def solve_and_measure(eq):
             return measure(eq)
         except:
             # Solve failed.
-            return float("inf")
+            return float(inf)
 
     # Call the generic optimiser,
     return optimiser.optimise(

freegs/optimiser.py

@@ -24,7 +24,7 @@
 along with FreeGS.  If not, see <http://www.gnu.org/licenses/>.
 """
 
-import random
+from numpy import random
 import copy
 import bisect
 import sys
@@ -65,7 +65,9 @@ def pickUnique(N, m, e):
     inds = sorted(e)  # Sorted list of indices. Used to avoid clashes
     others = []  # The list of three agents
     for i in range(m):
-        newind = random.randint(0, N - 1 - i - len(e))
+        high = N - 1 - i - len(e)
+        newind = random.randint(0, high) if high > 0 else 0
+
         for ind in inds:
             if newind == ind:
                 newind += 1

freegs/picard.py

@@ -93,7 +93,7 @@ def solve(
     bndry = 0.0
 
     # Set an initial value for bndry_change (set high to prevent immediate convergence)
-    bndry_change = 10.0
+    bndry_change = np.inf
 
     # Plasma assumed to not be limited at first
     has_been_limited = False