Merge pull request #47 from bendudson/chris-updates

Some bug fixes, tidying

freegs/equilibrium.py

@@ -128,7 +128,7 @@ def __init__(
             nx, ny, generator, nlevels=1, ncycle=1, niter=2, direct=True
         )
 
-    def setSolverVcycle(nlevels=1, ncycle=1, niter=1, direct=True):
+    def setSolverVcycle(self, nlevels=1, ncycle=1, niter=1, direct=True):
         """
         Creates a new linear solver, based on the multigrid code
 
@@ -138,7 +138,7 @@ def setSolverVcycle(nlevels=1, ncycle=1, niter=1, direct=True):
         direct   - Use a direct solver at the coarsest level?
 
         """
-        generator = GSsparse(Rmin, Rmax, Zmin, Zmax)
+        generator = GSsparse(self.Rmin, self.Rmax, self.Zmin, self.Zmax)
         nx, ny = self.R.shape
 
         self._solver = multigrid.createVcycle(
@@ -151,7 +151,7 @@ def setSolverVcycle(nlevels=1, ncycle=1, niter=1, direct=True):
             direct=direct,
         )
 
-    def setSolver(solver):
+    def setSolver(self, solver):
         """
         Sets the linear solver to use. The given object/function must have a __call__ method
         which takes two inputs
@@ -626,11 +626,12 @@ def effectiveElongation(self, R_wall_inner, R_wall_outer, npoints=300):
 
     def internalInductance1(self, npoints=300):
         """Calculates li1 plasma internal inductance"""
-        # Produce array of Bpol^2 in (R,Z)
-        B_polvals_2 = self.Bz(self.R, self.Z) ** 2 + self.Br(R, Z) ** 2
 
         R = self.R
         Z = self.Z
+        # Produce array of Bpol^2 in (R,Z)
+        B_polvals_2 = self.Bz(R, Z) ** 2 + self.Br(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
@@ -653,8 +654,6 @@ def internalInductance2(self):
 
         R = self.R
         Z = self.Z
-        psi = self.psi()
-
         # Produce array of Bpol in (R,Z)
         B_polvals_2 = self.Br(R, Z) ** 2 + self.Bz(R, Z) ** 2
 
@@ -675,8 +674,6 @@ def internalInductance3(self, R_wall_inner, R_wall_outer, npoints=300):
 
         R = self.R
         Z = self.Z
-        psi = self.psi()
-
         # Produce array of Bpol in (R,Z)
         B_polvals_2 = self.Br(R, Z) ** 2 + self.Bz(R, Z) ** 2
 
@@ -722,8 +719,6 @@ def poloidalBeta2(self):
         field_integral_pol = romb(romb(B_polvals_2 * dV))
         return 2 * mu0 * pressure_integral / field_integral_pol
 
-        return poloidal_beta
-
     def toroidalBeta(self):
         """Calculate plasma toroidal beta by integrating the thermal pressure
         and toroidal magnetic field pressure over the plasma volume.

freegs/jtor.py

@@ -273,8 +273,6 @@ def ffprime(self, pn):
         shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n
         return mu0 * self.L * (1 - self.Beta0) * self.Raxis * shape
 
-        return Jtor, pprime, ffprime
-
     def fvac(self):
         return self._fvac
 
@@ -406,8 +404,6 @@ def ffprime(self, pn):
         shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n
         return mu0 * self.L * (1 - self.Beta0) * self.Raxis * shape
 
-        return Jtor, pprime, ffprime
-
     def fvac(self):
         return self._fvac
 

freegs/machine.py

@@ -230,7 +230,7 @@ def plot(self, axis=None, show=False):
         for label, coil, multiplier in self.coils:
             axis = coil.plot(axis=axis, show=False)
         if show:
-            import matplotlib.pyplot
+            import matplotlib.pyplot as plt
 
             plt.show()
         return axis

freegs/test_equilibrium.py

@@ -48,3 +48,11 @@ def test_fixed_boundary_psi():
 
     assert eq.psi_bndry == 0.0
     assert eq.poloidalBeta() > 0.0
+
+
+def test_setSolverVcycle():
+    eq = equilibrium.Equilibrium(Rmin=0.1, Rmax=2.0, Zmin=-1.0, Zmax=1.0, nx=65, ny=65)
+
+    oldsolver = eq._solver
+    eq.setSolverVcycle(nlevels=2, ncycle=1, niter=5)
+    assert eq._solver != oldsolver