Merge pull request #16 from bendudson/higher-order

Higher order solver

docs/index.rst

@@ -13,6 +13,7 @@ Welcome to FreeGS's documentation!
    creating_equilibria
    input_and_output
    diagnostics
+   tests
 
 .. automodule:: freegs
 

docs/tests.rst

@@ -0,0 +1,21 @@
+Tests
+=====
+
+Convergence test
+----------------
+
+The ``test-convergence.py`` script solves for the same plasma as the
+``01-freeboundary.py`` example, with four poloidal field coils and an X-point.
+The change in results as the resolution is doubled is plotted as a function of
+grid resolution. This is therefore not testing convergence to a known solution,
+but is a check that the code converges to value, and indication of its
+convergence rate.
+Results are shown below, using the von Hagenow free boundary, and 4th-order solver for
+the poloidal flux :math:`\psi`.
+
+.. image:: freegs_convergence.png
+   :alt: Convergence of flux, magnetic field, volume and plasma current
+
+This indicates that in general convergence is between 2nd and 4th-order. The
+plasma volume is currently calculated by integrating over the poloidal
+cross-section, and could be improved by converting this to a surface integral.

freegs/boundary.py

@@ -132,30 +132,37 @@ def freeBoundaryHagenow(eq, Jtor, psi):
     psi_fixed = eq.callSolver(psi, rhs)
 
     # Differentiate at the boundary
-    # Second-order one-sided differences
+    # Note: normal is out of the domain, so this is -dU/dx
 
-    dUdn_L = (1.5*psi_fixed[0,:] - 2.*psi_fixed[1,:] + 0.5*psi_fixed[2,:])/dR   # left boundary
+    # Fourth-order one-sided differences
+    coeffs = [(0, 25./12),
+              (1, -4.0),
+              (2,  3.0),
+              (3, -16./12),
+              (4, 1./4)]
 
-    dUdn_R = (1.5*psi_fixed[-1,:] - 2.*psi_fixed[-2,:] + 0.5*psi_fixed[-3,:])/dR # Right boundary
+    dUdn_L = sum([weight * psi_fixed[index, :] for index, weight in coeffs]) / dR    # left boundary
 
-    dUdn_D = (1.5*psi_fixed[:,0] - 2.*psi_fixed[:,1] + 0.5*psi_fixed[:,2])/dZ  # Down boundary
+    dUdn_R = sum([weight * psi_fixed[-(1 + index), :] for index, weight in coeffs]) / dR    # Right boundary
 
-    dUdn_U = (1.5*psi_fixed[:,-1] - 2.*psi_fixed[:,-2] + 0.5*psi_fixed[:,-3])/dZ  # Upper boundary
+    dUdn_D = sum([weight * psi_fixed[:, index] for index, weight in coeffs]) / dZ  # Down boundary
+
+    dUdn_U = sum([weight * psi_fixed[:, -(1 + index)] for index, weight in coeffs]) / dZ  # Upper boundary
 
     dd = sqrt(dR**2 + dZ**2) # Diagonal spacing
 
     # Left down corner 
-    dUdn_L[0] = dUdn_D[0] = (1.5*psi_fixed[0,0] - 2.*psi_fixed[1,1] + 0.5*psi_fixed[2,2]) / dd
+    dUdn_L[0] = dUdn_D[0] = sum([weight * psi_fixed[index, index] for index, weight in coeffs]) / dd
 
     # Left upper corner
-    dUdn_L[-1] = dUdn_U[0] = (1.5*psi_fixed[0,-1] - 2.*psi_fixed[1,-2] + 0.5*psi_fixed[2,-3]) / dd
+    dUdn_L[-1] = dUdn_U[0] = sum([weight * psi_fixed[index, -(1 + index)] for index, weight in coeffs]) / dd
 
     # Right down corner
-    dUdn_R[0] = dUdn_D[-1] = (1.5*psi_fixed[-1,0] - 2.*psi_fixed[-2,1] + 0.5*psi_fixed[-3,2]) / dd
+    dUdn_R[0] = dUdn_D[-1] = sum([weight * psi_fixed[-(1 + index), index] for index, weight in coeffs]) / dd
 
     # Right upper corner
-    dUdn_R[-1] = dUdn_U[-1] = (1.5*psi_fixed[-1,-1] - 2.*psi_fixed[-2,-2] + 0.5*psi_fixed[-3,-3]) / dd
-    
+    dUdn_R[-1] = dUdn_U[-1] = sum([weight * psi_fixed[-(1 + index), -(1 + index)] for index, weight in coeffs]) / dd
+
     # Now for each point on the boundary perform a loop integral
 
     eps = 1e-2

freegs/equilibrium.py

@@ -12,7 +12,7 @@
 from . import critical
 
 # Operators which define the G-S equation
-from .gradshafranov import mu0, GSsparse
+from .gradshafranov import mu0, GSsparse, GSsparse4thOrder
 
 # Multigrid solver
 from . import multigrid
@@ -49,7 +49,7 @@ def __init__(self, tokamak=machine.EmptyTokamak(),
                  Zmin=-1.0, Zmax=1.0,
                  nx=65, ny=65,
                  boundary=freeBoundary,
-                 psi=None, current=0.0):
+                 psi=None, current=0.0, order=4):
         """Initialises a plasma equilibrium
 
         Rmin, Rmax  - Range of major radius R [m]
@@ -64,6 +64,9 @@ def __init__(self, tokamak=machine.EmptyTokamak(),
               surfaces as starting guess
 
         current - Plasma current (default = 0.0)
+
+        order - The order of the differential operators to use.
+                Valid values are 2 or 4.
         """
 
         self.tokamak = tokamak
@@ -100,7 +103,14 @@ def __init__(self, tokamak=machine.EmptyTokamak(),
         self._updatePlasmaPsi(psi)  # Needs to be after _pgreen
 
         # Create the solver
-        generator = GSsparse(Rmin, Rmax, Zmin, Zmax)
+        if order == 2:
+            generator = GSsparse(Rmin, Rmax, Zmin, Zmax)
+        elif order == 4:
+            generator = GSsparse4thOrder(Rmin, Rmax, Zmin, Zmax)
+        else:
+            raise ValueError("Invalid choice of order ({}). Valid values are 2 or 4.".format(order))
+        self.order = order
+
         self._solver = multigrid.createVcycle(nx, ny,
                                               generator,
                                               nlevels=1,
@@ -454,22 +464,33 @@ def print_forces(forces, prefix=""):
 
         print_forces(self.getForces())
 
-def refine(eq):
+def refine(eq, nx=None, ny=None):
     """
     Double grid resolution, returning a new equilibrium
+
     
     """
     # Interpolate the plasma psi
-    plasma_psi = multigrid.interpolate(eq.plasma_psi)
-    nx, ny = plasma_psi.shape
+    #plasma_psi = multigrid.interpolate(eq.plasma_psi)
+    #nx, ny = plasma_psi.shape
 
+    # By default double the number of intervals
+    if not nx:
+        nx = 2*(eq.R.shape[0] - 1) + 1
+    if not ny:
+        ny = 2*(eq.R.shape[1] - 1) + 1
+
     result = Equilibrium(tokamak=eq.tokamak,
                          Rmin = eq.Rmin,
                          Rmax = eq.Rmax,
                          Zmin = eq.Zmin,
                          Zmax = eq.Zmax,
+                         boundary=eq._applyBoundary,
+                         order = eq.order,
                          nx=nx, ny=ny)
-
+
+    plasma_psi = eq.psi_func(result.R, result.Z, grid=False)
+
     result._updatePlasmaPsi(plasma_psi)
 
     if hasattr(eq, "_profiles"):

freegs/gradshafranov.py

@@ -38,6 +38,10 @@ class GSElliptic:
     
     \Delta^* = R^2 \nabla\cdot\frac{1}{R^2}\nabla
     
+    which is
+
+    d^2/dR^2 + d^2/dZ^2 - (1/R)*d/dR
+    
     """
 
     def __init__(self, Rmin):
@@ -137,6 +141,121 @@ def __call__(self, nx, ny):
 
         # Convert to Compressed Sparse Row (CSR) format
         return A.tocsr()
+
+class GSsparse4thOrder:
+    """
+    Calculates sparse matrices for the Grad-Shafranov operator using 4th-order
+    finite differences. 
+    """
+
+    # Coefficients for first derivatives
+    # (index offset, weight)
+
+    centred_1st = [(-2,  1./12),
+                   (-1, -8./12),
+                   ( 1,  8./12),
+                   ( 2, -1./12)]
+
+    offset_1st = [(-1,  -3./12),
+                  ( 0, -10./12),
+                  ( 1,  18./12),
+                  ( 2,  -6./12),
+                  ( 3,   1./12)]
+
+    # Coefficients for second derivatives
+    # (index offset, weight)
+    centred_2nd = [(-2, -1./12),
+                   (-1, 16./12),
+                   ( 0, -30./12),
+                   ( 1, 16./12),
+                   ( 2, -1./12)]
+
+    offset_2nd = [(-1,  10./12),
+                  ( 0, -15./12),
+                  ( 1,  -4./12),
+                  ( 2,  14./12),
+                  ( 3,  -6./12),
+                  ( 4,   1./12)]
+
+    def __init__(self, Rmin, Rmax, Zmin, Zmax):
+        self.Rmin = Rmin
+        self.Rmax = Rmax
+        self.Zmin = Zmin
+        self.Zmax = Zmax
+
+    def __call__(self, nx, ny):
+        """
+        Create a sparse matrix with given resolution
+        
+        """
+
+        # Calculate grid spacing
+        dR = (self.Rmax - self.Rmin)/(nx - 1)
+        dZ = (self.Zmax - self.Zmin)/(ny - 1)
+
+        # Total number of points, including boundaries
+        N = nx * ny
+
+        # Create a linked list sparse matrix
+        A = lil_matrix((N,N))
+
+        invdR2 = 1./dR**2
+        invdZ2 = 1./dZ**2
+
+        for x in range(1,nx-1):
+            R = self.Rmin + dR*x  # Major radius of this point
+            for y in range(1,ny-1):
+                row = x*ny + y
+
+                # d^2 / dZ^2
+                if y == 1:
+                    # One-sided derivatives in Z
+                    for offset, weight in self.offset_2nd:
+                        A[row, row + offset] += weight * invdZ2
+                elif y == ny-2:
+                    # One-sided, reversed direction.
+                    # Note that for second derivatives the sign of the weights doesn't change
+                    for offset, weight in self.offset_2nd:
+                        A[row, row - offset] += weight * invdZ2
+                else:
+                    # Central differencing
+                    for offset, weight in self.centred_2nd:
+                        A[row, row + offset] += weight * invdZ2
+
+                # d^2 / dR^2 - (1/R) d/dR
+
+                if x == 1:
+                    for offset, weight in self.offset_2nd:
+                        A[row, row + offset*ny] += weight * invdR2
+
+                    for offset, weight in self.offset_1st:
+                        A[row, row + offset*ny] -= weight / (R * dR)
+
+                elif x == nx-2:
+                    for offset, weight in self.offset_2nd:
+                        A[row, row - offset*ny] += weight * invdR2
+
+                    for offset, weight in self.offset_1st:
+                        A[row, row - offset*ny] += weight / (R * dR)
+                else:
+                    for offset, weight in self.centred_2nd:
+                        A[row, row + offset*ny] += weight * invdR2
+
+                    for offset, weight in self.centred_1st:
+                        A[row, row + offset*ny] -= weight / (R * dR)
+
+        # Set boundary rows
+        for x in range(nx):
+            for y in [0, ny-1]:
+                row = x*ny + y
+                A[row, row] = 1.0
+        for x in [0, nx-1]:
+            for y in range(ny):
+                row = x*ny + y
+                A[row, row] = 1.0
+
+        # Convert to Compressed Sparse Row (CSR) format
+        return A.tocsr()
 
 def Greens(Rc, Zc, R, Z):
     """

freegs/picard.py

@@ -21,7 +21,7 @@
 
 from numpy import amin, amax
 
-def solve(eq, profiles, constrain=None, rtol=1e-3, blend=0.0,
+def solve(eq, profiles, constrain=None, rtol=1e-3, atol=1e-10, blend=0.0,
           show=False, axis=None, pause=0.0001, psi_bndry=None):
     """
     Perform Picard iteration to find solution to the Grad-Shafranov equation
@@ -30,6 +30,7 @@ def solve(eq, profiles, constrain=None, rtol=1e-3, blend=0.0,
     profiles - A Profile object for toroidal current (jtor.py)
 
     rtol     - Relative tolerance (change in psi)/( max(psi) - min(psi) )
+    atol     - Absolute tolerance, change in psi
     blend    - Blending of previous and next psi solution
                psi{n+1} <- psi{n+1} * (1-blend) + blend * psi{n}
 
@@ -94,7 +95,7 @@ def solve(eq, profiles, constrain=None, rtol=1e-3, blend=0.0,
         print("Maximum change in psi: %e. Relative: %e" % (psi_maxchange, psi_relchange))
 
         # Check if the relative change in psi is small enough
-        if psi_relchange < rtol:
+        if (psi_maxchange < atol) or (psi_relchange < rtol):
             break
 
         # Adjust the coil currents