Merge pull request #15 from bendudson/connection-length

Calculate connection length

10-mastu-connection.py

@@ -0,0 +1,32 @@
+# Read MAST-U GEQDSK file, calculate connection length
+
+from freegs import geqdsk
+from freegs import machine
+
+tokamak = machine.MASTU()
+
+with open("mast-upgrade.geqdsk") as f:
+    eq = geqdsk.read(f, tokamak, show=False)
+
+from freegs import fieldtracer
+import matplotlib.pyplot as plt
+
+# Plot equilibrium
+axis = eq.plot(show=False)
+
+# Trace field lines both directions along the magnetic field
+# By passing axis 
+forward, backward = fieldtracer.traceFieldLines(eq, axis=axis, nturns=50)
+
+plt.savefig("mast-upgrade-fieldtrace.pdf")
+plt.show()
+
+# Plot field line length from midplane to target
+plt.plot(forward.R[0,:], forward.length[-1,:], label="Forward")
+plt.plot(backward.R[0,:], backward.length[-1,:], label="Backward")
+plt.legend()
+plt.xlabel("Starting major radius [m]")
+plt.ylabel("Parallel connection length [m]")
+
+plt.savefig("mast-upgrade-lpar.pdf")
+plt.show()

docs/diagnostics.rst

@@ -146,3 +146,78 @@ This can be changed e.g::
   tokamak["P1L"].area = machine.AreaCurrentLimit(1e9)
 
 would set the current limit for coil "P1L" to 1e9 Amps per square meter.
+
+Field line connection length
+----------------------------
+
+Example: ``10-mastu-connection.py``. Requires the file ``mast-upgrade.geqdsk``
+which is created by running ``08-mast-upgrade.py``.
+
+To calculate the distance along magnetic field lines from the outboard midplane
+to the walls in an equilibrium ``eq``, the simplest way is::
+
+  from freegs import fieldtracer
+  forward, backward = fieldtracer.traceFieldLines(eq)
+
+
+To also plot the field lines on top of the equilibrium::
+  
+  axis = eq.plot(show=False)
+  forward, backward = fieldtracer.traceFieldLines(eq, axis=axis)
+  plt.show()
+  
+This will display the poloidal cross-section of the plasma, and plot field lines
+traced in both directions along the magnetic field from the outboard midplane.
+
+To plot the distances along the magnetic field from midplane to target as a
+function of the starting radius::
+
+  plt.plot(forward.R[0,:], forward.length[-1,:], label="Forward")
+  plt.plot(backward.R[0,:], backward.length[-1,:], label="Backward")
+  plt.legend()
+  plt.xlabel("Starting major radius [m]")
+  plt.ylabel("Parallel connection length [m]")
+  
+  plt.show()
+
+Here ``forward.R`` and ``forward.length`` are 2D arrays, where the first index
+is the point along the magnetic field (0 = start, -1 = end), and the second
+index is the field line number. There is also ``forward.Z`` with the height in meters.
+
+The output can be customised by passing keywords to ``traceFieldLines``:
+``solwidth`` sets the width of the starting region at the outboard midplane;
+``nlines`` is the number of field lines to follow in each direction;
+``nturns`` the number of times around the torus to follow the field;
+``npoints`` is the number of points along each field line.
+
+For more control over which field lines are followed, the ``FieldTracer`` class
+does the actual field line following::
+
+  from freegs import fieldtracer
+  import numpy as np
+  
+  tracer = fieldtracer.FieldTracer(eq)
+
+  result = tracer.follow([1.35], [0.0], np.linspace(0.0, 2*np.pi, 20))
+
+This follows a magnetic field in the direction of B, starting at ``R=1.35``m,
+``Z=0.0``m, outputting positions at 20 toroidal angles between 0 and 2pi
+i.e. one toroidal turn. The R and Z starting locations should be an array or
+list with the same shape.
+
+The ``result`` is an array: The first index is the angle (size 20 here), and the
+last index has size 3 (R, Z, length). Between the first and last indices the
+result has the same shape as the R and Z starting positions. In the above code
+``result`` has size ``(20, 1, 3)``. To plot the field line on top of the
+equilibrium::
+
+  import matplotlib.pyplot as plt
+  
+  eq.plot(show=False)
+  plt.plot(result[:,0,0], result[:,0,1])
+  plt.show()
+
+The direction to follow along the field can be reversed by passing
+``backward=True`` keyword to ``tracer.follow``.
+
+

freegs/equilibrium.py

@@ -91,14 +91,14 @@ def __init__(self, tokamak=machine.EmptyTokamak(),
             psi[-1, :] = 0.0
             psi[:, -1] = 0.0
 
-        self._updatePlasmaPsi(psi)
-
         # Calculate coil Greens functions. This is an optimisation,
         # used in self.psi() to speed up calculations
         self._pgreen = tokamak.createPsiGreens(self.R, self.Z)
 
         self._current = current  # Plasma current
-
+
+        self._updatePlasmaPsi(psi)  # Needs to be after _pgreen
+
         # Create the solver
         generator = GSsparse(Rmin, Rmax, Zmin, Zmax)
         self._solver = multigrid.createVcycle(nx, ny,
@@ -172,73 +172,33 @@ def poloidalBeta(self):
         Return the poloidal beta 
         betap = (8pi/mu0) * int(p)dRdZ / Ip^2
         """
-
-        # Total poloidal flux (plasma + coils)
-        psi = self.psi()
-
-        # Analyse the equilibrium, finding O- and X-points
-        opt, xpt = critical.find_critical(self.R, self.Z, psi)
-        if not opt:
-            raise ValueError("No O-points found!")
-        psi_axis = opt[0][2]
-
-        if xpt:
-            psi_bndry = xpt[0][2]
-            mask = critical.core_mask(self.R, self.Z, psi, opt, xpt)
-        else:
-            # No X-points
-            if psi[0,0] > psi_axis:
-                psi_bndry = np.amax(psi)
-            else:
-                psi_bndry = np.amin(psi)
-            mask = None
 
         dR = self.R[1,0] - self.R[0,0]
         dZ = self.Z[0,1] - self.Z[0,0]
 
         # Normalised psi
-        psi_norm = (psi - psi_axis)  / (psi_bndry - psi_axis)
+        psi_norm = (self.psi() - self.psi_axis)  / (self.psi_bndry - self.psi_axis)
 
         # Plasma pressure
         pressure = self.pressure(psi_norm)
-        if mask is not None:
+        if self.mask is not None:
             # If there is a masking function (X-points, limiters)
-            pressure *= mask
+            pressure *= self.mask
 
         # Integrate pressure in 2D
         return ((8.*pi)/mu0)*romb(romb(pressure))*dR*dZ / (self.plasmaCurrent()**2)
 
     def plasmaVolume(self):
         """Calculate the volume of the plasma in m^3"""
 
-        # Total poloidal flux (plasma + coils)
-        psi = self.psi()
-
-        # Analyse the equilibrium, finding O- and X-points
-        opt, xpt = critical.find_critical(self.R, self.Z, psi)
-        if not opt:
-            raise ValueError("No O-points found!")
-        psi_axis = opt[0][2]
-
-        if xpt:
-            psi_bndry = xpt[0][2]
-            mask = critical.core_mask(self.R, self.Z, psi, opt, xpt)
-        else:
-            # No X-points
-            if psi[0,0] > psi_axis:
-                psi_bndry = np.amax(psi)
-            else:
-                psi_bndry = np.amin(psi)
-            mask = None
-
         dR = self.R[1,0] - self.R[0,0]
         dZ = self.Z[0,1] - self.Z[0,0]
 
         # Volume element
         dV = 2.*pi*self.R * dR * dZ
 
-        if mask is not None:   # Only include points in the core
-            dV *= mask
+        if self.mask is not None:   # Only include points in the core
+            dV *= self.mask
 
         # Integrate volume in 2D
         return romb(romb(dV))
@@ -268,18 +228,40 @@ def Bz(self, R, Z):
         Total vertical magnetic field
         """
         return self.plasmaBz(R,Z) + self.tokamak.Bz(R,Z)
-
+
+    def Btor(self, R, Z):
+        """
+        Toroidal magnetic field
+        """
+        # Normalised psi
+        psi_norm = (self.psiRZ(R, Z) - self.psi_axis)  / (self.psi_bndry - self.psi_axis)
+
+        # Get f = R * Btor in the core. May be invalid outside the core
+        fpol = self.fpol(psi_norm)
+
+        if self.mask is not None:
+            # Get the values of the core mask at the requested R,Z locations
+            # This is 1 in the core, 0 outside
+            mask = self.mask_func(R,Z, grid=False)
+            fpol = fpol * mask + (1.0 - mask)*self.fvac()
+
+        return fpol / R
+
     def psi(self):
+        """
+        Total poloidal flux ψ (psi), including contribution from
+        plasma and external coils.
+        """
         #return self.plasma_psi + self.tokamak.psi(self.R, self.Z)
         return self.plasma_psi + self.tokamak.calcPsiFromGreens(self._pgreen)
 
     def psiRZ(self, R, Z):
         """
         Return poloidal flux psi at given (R,Z) location
         """
-        return self.psi_func(R,Z) + self.tokamak.psi(R,Z)
+        return self.psi_func(R, Z, grid=False) + self.tokamak.psi(R,Z)
 
-    def fpol(self,psinorm):
+    def fpol(self, psinorm):
         """
         Return f = R*Bt at specified values of normalised psi
         """
@@ -402,13 +384,36 @@ def solve(self, profiles, Jtor=None, psi=None, psi_bndry=None):
 
     def _updatePlasmaPsi(self, plasma_psi):
         """
-        Sets the plasma psi data, updates spline interpolation coefficients
+        Sets the plasma psi data, updates spline interpolation coefficients.
+        Also updates:
+
+        self.mask        2D (R,Z) array which is 1 in the core, 0 outside
+        self.psi_axis    Value of psi on the magnetic axis
+        self.psi_bndry   Value of psi on plasma boundary
         """
         self.plasma_psi = plasma_psi
 
         # Update spline interpolation
         self.psi_func = interpolate.RectBivariateSpline(self.R[:,0], self.Z[0,:], plasma_psi)
-
+
+        # Update the locations of the X-points, core mask, psi ranges.
+        # Note that this may fail if there are no X-points, so it should not raise an error
+        # Analyse the equilibrium, finding O- and X-points
+        psi = self.psi()
+        opt, xpt = critical.find_critical(self.R, self.Z, psi)
+        if opt:
+            self.psi_axis = opt[0][2]
+
+            if xpt:
+                self.psi_bndry = xpt[0][2]
+                self.mask = critical.core_mask(self.R, self.Z, psi, opt, xpt)
+
+                # Use interpolation to find if a point is in the core.
+                self.mask_func = interpolate.RectBivariateSpline(self.R[:,0], self.Z[0,:], self.mask)
+            else:
+                self.psi_bndry = None
+                self.mask = None        
+
     def plot(self, axis=None, show=True, oxpoints=True):
         """
         Plot the equilibrium flux surfaces
@@ -417,6 +422,11 @@ def plot(self, axis=None, show=True, oxpoints=True):
         show     - Call matplotlib.pyplot.show() before returning
         oxpoints - Plot X points as red circles, O points as green circles
     
+        Returns
+        -------
+
+        axis  object from Matplotlib
+
         """
         from .plotting import plotEquilibrium
         return plotEquilibrium(self, axis=axis, show=show, oxpoints=oxpoints)