Autoformat

freegs/__init__.py

@@ -48,3 +48,13 @@
 
 __version__ = metadata(__package__)["Version"]
 __author__ = metadata(__package__)["Author"]
+
+__all__ = [
+    "Equilibrium",
+    "jtor",
+    "machine",
+    "control",
+    "solve",
+    "OutputFile",
+    "plotting",
+]

freegs/boundary.py

@@ -159,7 +159,7 @@ def freeBoundaryHagenow(eq, Jtor, psi):
         sum([weight * psi_fixed[:, -(1 + index)] for index, weight in coeffs]) / dZ
     )  # Upper boundary
 
-    dd = sqrt(dR ** 2 + dZ ** 2)  # Diagonal spacing
+    dd = sqrt(dR**2 + dZ**2)  # Diagonal spacing
 
     # Left down corner
     dUdn_L[0] = dUdn_D[0] = (

freegs/coil.py

@@ -194,7 +194,7 @@ def getForces(self, equilibrium):
         # Force per unit length.
         # In cgs units f = I^2/(c^2 * R) * (ln(8*R/a) - 1 + xi/2)
         # In SI units f = mu0 * I^2 / (4*pi*R) * (ln(8*R/a) - 1 + xi/2)
-        self_fr = (mu0 * total_current ** 2 / (4.0 * np.pi * self.R)) * (
+        self_fr = (mu0 * total_current**2 / (4.0 * np.pi * self.R)) * (
             np.log(8.0 * self.R / minor_radius) - 1 + self_inductance / 2.0
         )
 

freegs/control.py

@@ -8,7 +8,6 @@
 from numpy.linalg import inv, norm
 import numpy as np
 from scipy import optimize
-from . import critical
 
 
 class constrain(object):
@@ -130,7 +129,7 @@ def __call__(self, eq):
 
         # Calculate the change in coil current
         self.current_change = dot(
-            inv(dot(transpose(A), A) + self.gamma ** 2 * eye(ncontrols)),
+            inv(dot(transpose(A), A) + self.gamma**2 * eye(ncontrols)),
             dot(transpose(A), b),
         )
 
@@ -179,7 +178,11 @@ def max_total_currents(x):
         if self.current_change.shape[0] > 0:
             x0 = self.current_change
             sol = optimize.minimize(
-                objective, x0, method="SLSQP", bounds=current_change_bnds, constraints=cons
+                objective,
+                x0,
+                method="SLSQP",
+                bounds=current_change_bnds,
+                constraints=cons,
             )
 
             self.current_change = sol.x

freegs/critical.py

@@ -72,13 +72,13 @@ def find_critical(R, Z, psi, discard_xpoints=True):
     f = interpolate.RectBivariateSpline(R[:, 0], Z[0, :], psi)
 
     # Find candidate locations, based on minimising Bp^2
-    Bp2 = (f(R, Z, dx=1, grid=False) ** 2 + f(R, Z, dy=1, grid=False) ** 2) / R ** 2
+    Bp2 = (f(R, Z, dx=1, grid=False) ** 2 + f(R, Z, dy=1, grid=False) ** 2) / R**2
 
     # Get grid resolution, which determines a reasonable tolerance
     # for the Newton iteration search area
     dR = R[1, 0] - R[0, 0]
     dZ = Z[0, 1] - Z[0, 0]
-    radius_sq = 9 * (dR ** 2 + dZ ** 2)
+    radius_sq = 9 * (dR**2 + dZ**2)
 
     # Find local minima
 
@@ -100,7 +100,6 @@ def find_critical(R, Z, psi, discard_xpoints=True):
                 and (Bp2[i, j] < Bp2[i, j + 1])
                 and (Bp2[i, j] < Bp2[i, j - 1])
             ):
-
                 # Found local minimum
 
                 R0 = R[i, j]
@@ -113,11 +112,10 @@ def find_critical(R, Z, psi, discard_xpoints=True):
 
                 count = 0
                 while True:
-
                     Br = -f(R1, Z1, dy=1, grid=False) / R1
                     Bz = f(R1, Z1, dx=1, grid=False) / R1
 
-                    if Br ** 2 + Bz ** 2 < 1e-6:
+                    if Br**2 + Bz**2 < 1e-6:
                         # Found a minimum. Classify as either
                         # O-point or X-point
 
@@ -133,7 +131,7 @@ def find_critical(R, Z, psi, discard_xpoints=True):
                             (psi[i + 2, j + 2] - psi[i + 2, j - 2]) / (4.0 * dZ)
                             - (psi[i - 2, j + 2] - psi[i - 2, j - 2]) / (4.0 * dZ)
                         ) / (4.0 * dR)
-                        D = d2dr2 * d2dz2 - d2drdz ** 2
+                        D = d2dr2 * d2dz2 - d2drdz**2
 
                         if D < 0.0:
                             # Found X-point
@@ -368,12 +366,15 @@ def find_psisurface(eq, psifunc, r0, z0, r1, z1, psival=1.0, n=100, axis=None):
             # Changed 1.0 to psival in f
             # make f gradient to psival surface
             f = (pnorm[ind] - psival) / (pnorm[ind] - pnorm[ind - 1])
-            
+
             # Interpolate between points
             r = (1.0 - f) * r[ind] + f * r[ind - 1]
             z = (1.0 - f) * z[ind] + f * z[ind - 1]
 
-            if f > 1.0: warn(f"find_psisurface has encountered an extrapolation. This will probably result in a point where you don't expect it.")
+            if f > 1.0:
+                warn(
+                    "find_psisurface has encountered an extrapolation. This will probably result in a point where you don't expect it."
+                )
 
     if axis is not None:
         axis.plot(r, z, "bo")
@@ -527,17 +528,17 @@ def find_safety(
     fpol = eq.fpol(psirange[:]).reshape(npsi, 1)
     Br = eq.Br(r, z)
     Bz = eq.Bz(r, z)
-    Bthe = sqrt(Br ** 2 + Bz ** 2)
+    Bthe = sqrt(Br**2 + Bz**2)
 
     # Differentiate location w.r.t. index
     dr_di = (np.roll(r, 1, axis=1) - np.roll(r, -1, axis=1)) / 2.0
     dz_di = (np.roll(z, 1, axis=1) - np.roll(z, -1, axis=1)) / 2.0
 
     # Distance between points
-    dl = sqrt(dr_di ** 2 + dz_di ** 2)
+    dl = sqrt(dr_di**2 + dz_di**2)
 
     # Integrand - Btor/(R*Bthe) = Fpol/(R**2*Bthe)
-    qint = fpol / (r ** 2 * Bthe)
+    qint = fpol / (r**2 * Bthe)
 
     # Integral
     q = sum(qint * dl, axis=1) / (2 * pi)

freegs/equilibrium.py

@@ -113,7 +113,7 @@ def __init__(
         if psi is None:
             # Starting guess for psi
             xx, yy = meshgrid(linspace(0, 1, nx), linspace(0, 1, ny), indexing="ij")
-            psi = exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4 ** 2)
+            psi = exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4**2)
 
             psi[0, :] = 0.0
             psi[:, 0] = 0.0
@@ -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):
@@ -1174,9 +1186,6 @@ def pressure_ave(self):
         R = self.R
         Z = self.Z
 
-        # Produce array of Btor in (R,Z)
-        B_torvals_2 = self.Btor(R, Z) ** 2
-
         dR = R[1, 0] - R[0, 0]
         dZ = Z[0, 1] - Z[0, 0]
         dV = 2.0 * np.pi * R * dR * dZ
@@ -1362,11 +1371,9 @@ def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None):
 
 
 if __name__ == "__main__":
-
     # Test the different spline interpolation routines
 
     from numpy import ravel
-    import matplotlib.pyplot as plt
 
     import machine
 

freegs/fieldtracer.py

@@ -102,7 +102,7 @@ def fieldDirection(self, pos, toroidal_angle, evolving, backward):
         Bz = self._eq.Bz(R, Z)
         Btor = self._eq.Btor(R, Z)
 
-        B = np.sqrt(Br ** 2 + Bz ** 2 + Btor ** 2)
+        B = np.sqrt(Br**2 + Bz**2 + Btor**2)
 
         # Detect when the boundary has been reached
         self.updateEvolving(R, Z, evolving)

freegs/filament_coil.py

@@ -25,7 +25,7 @@
 import numpy as np
 from . import polygons
 from shapely import geometry
-from .coil import Coil, AreaCurrentLimit
+from .coil import Coil
 from .gradshafranov import Greens, GreensBr, GreensBz
 
 
@@ -72,7 +72,6 @@ def populate_with_fils(shape, Nfils):
     nInside = 0
 
     while nInside < Nfils:
-
         nInside = 0
         i = 0  # Row counter used for staggering every other row
         # Create the (non-staggered) grid points for the given d(=dR,dZ)
@@ -83,17 +82,14 @@ def populate_with_fils(shape, Nfils):
         nZ = len(Zgrid)
 
         for i in range(nZ):
-
             zpoint = Zgrid[i]
 
             if i % 2:  # Every other row, stagger R coords by 0.5*dR
                 offset = 0.5 * d
             else:
-
                 offset = 0.0
 
             for j in range(nR):
-
                 rpoint = Rgrid[j] + offset
 
                 point = geometry.Point(rpoint, zpoint)
@@ -131,17 +127,14 @@ def populate_with_fils(shape, Nfils):
     nZ = len(Zgrid)
 
     for i in range(nZ):
-
         zpoint = Zgrid[i]
 
         if i % 2:  # Every other row, stagger R coords by 0.5*dR
             offset = 0.5 * d_opt
         else:
-
             offset = 0.0
 
         for j in range(nR):
-
             rpoint = Rgrid[j] + offset
 
             point = geometry.Point(rpoint, zpoint)
@@ -296,7 +289,7 @@ def controlBz(self, R, Z):
         result = result * self.turns / float(self.npoints)
         return result
 
-    def inShape(self,polygon):
+    def inShape(self, polygon):
         counter = 0
         for r, z in self.points:
             if polygon.contains(geometry.Point(r, z)):

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
 
@@ -71,8 +71,8 @@ def __call__(self, psi, dR, dZ):
 
         b = zeros([nx, ny])
 
-        invdR2 = 1.0 / dR ** 2
-        invdZ2 = 1.0 / dZ ** 2
+        invdR2 = 1.0 / dR**2
+        invdZ2 = 1.0 / dZ**2
 
         for x in range(1, nx - 1):
             R = self.Rmin + dR * x  # Major radius of this point
@@ -92,7 +92,7 @@ def diag(self, dR, dZ):
         Return the diagonal elements
 
         """
-        return -2.0 / dR ** 2 - 2.0 / dZ ** 2
+        return -2.0 / dR**2 - 2.0 / dZ**2
 
 
 class GSsparse:
@@ -122,8 +122,8 @@ def __call__(self, nx, ny):
         # Create a linked list sparse matrix
         A = eye(N, format="lil")
 
-        invdR2 = 1.0 / dR ** 2
-        invdZ2 = 1.0 / dZ ** 2
+        invdR2 = 1.0 / dR**2
+        invdZ2 = 1.0 / dZ**2
 
         for x in range(1, nx - 1):
             R = self.Rmin + dR * x  # Major radius of this point
@@ -210,8 +210,8 @@ def __call__(self, nx, ny):
         # Create a linked list sparse matrix
         A = lil_matrix((N, N))
 
-        invdR2 = 1.0 / dR ** 2
-        invdZ2 = 1.0 / dZ ** 2
+        invdR2 = 1.0 / dR**2
+        invdZ2 = 1.0 / dZ**2
 
         for x in range(1, nx - 1):
             R = self.Rmin + dR * x  # Major radius of this point