From 0eba3d8a7b227f638ddec647371c052db5108eed Mon Sep 17 00:00:00 2001
From: Chris Marsden <48368617+chrismarsden7@users.noreply.github.com>
Date: Mon, 8 Aug 2022 18:57:09 +0100
Subject: [PATCH] FIxes to geqdsk reader for limited plasmas

Limited plasmas now readable and reproducable from geqdsk files.
---
 02-read-geqdsk.py     |  2 --
 12-limited.py         | 29 +++++++++++++++++++++++++++--
 freegs/equilibrium.py |  6 +++---
 freegs/geqdsk.py      | 20 ++++++++++++--------
 4 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/02-read-geqdsk.py b/02-read-geqdsk.py
index 0dff61a..bd4b043 100644
--- a/02-read-geqdsk.py
+++ b/02-read-geqdsk.py
@@ -2,10 +2,8 @@
 from freegs import machine
 from freegs.plotting import plotEquilibrium
 
-#tokamak = machine.MAST_sym()
 tokamak = machine.TestTokamak()
 
-#with open("g014220.00200") as f:
 with open("lsn.geqdsk") as f:
     eq = geqdsk.read(f, tokamak, show=True)
 
diff --git a/12-limited.py b/12-limited.py
index 18fa8c7..b239439 100644
--- a/12-limited.py
+++ b/12-limited.py
@@ -2,6 +2,8 @@
 
 import freegs
 import freegs.critical as critical
+from freegs import geqdsk
+from freegs.plotting import plotEquilibrium
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -87,12 +89,26 @@
 ax.contour(eq.R,eq.Z,eq.psiN(),levels=[1.0],colors='orange')
 ax.plot(eq.tokamak.wall.R,eq.tokamak.wall.Z,color='k')
 opt, xpt = critical.find_critical(eq.R,eq.Z,eq.psi())
-mask = critical.core_mask(eq.R,eq.Z,eq.psi(),opt,xpt,eq.psi_bndry)
-ax.contourf(eq.R,eq.Z,mask)
+isoflux = np.array(
+    freegs.critical.find_separatrix(eq, ntheta=101, opoint=opt, xpoint=xpt, psi=eq.psi())
+)
+ind = np.argmin(isoflux[:, 1])
+rbdry = np.roll(isoflux[:, 0][::-1], -ind)
+rbdry = np.append(rbdry,rbdry[0])
+zbdry = np.roll(isoflux[:, 1][::-1], -ind)
+zbdry = np.append(zbdry, zbdry[0])
+ax.plot(rbdry,zbdry,'rx')
+ax.contourf(eq.R,eq.Z,eq.mask)
 ax.set_aspect('equal')
 ax.set_xlabel('R(m)')
 ax.set_ylabel('Z(m)')
 plt.show()
+
+from freegs import geqdsk
+
+with open("limited.geqdsk", "w") as f:
+    geqdsk.write(eq, f)
+
 #########################################
 # Create the machine, which specifies coil locations
 # and equilibrium, specifying the domain to solve over
@@ -211,3 +227,12 @@
 ax.set_ylabel('fpol')
 ax.legend()
 plt.show()
+
+#########################
+# Load in the geqdsk of the limited plasma
+tokamak = freegs.machine.TestTokamakLimited()
+
+with open("limited.geqdsk") as f:
+    eq3 = geqdsk.read(f, tokamak, show=True)
+
+plotEquilibrium(eq3)
diff --git a/freegs/equilibrium.py b/freegs/equilibrium.py
index 9c049bf..e788b8f 100644
--- a/freegs/equilibrium.py
+++ b/freegs/equilibrium.py
@@ -63,7 +63,8 @@ def __init__(
         boundary=freeBoundary,
         psi=None,
         current=0.0,
-        order=4
+        order=4,
+        check_limited = False
     ):
         """Initialises a plasma equilibrium
 
@@ -99,7 +100,7 @@ def __init__(
         self.Z_1D = linspace(Zmin, Zmax, ny)
         self.R, self.Z = meshgrid(self.R_1D, self.Z_1D, indexing="ij")
 
-        self.check_limited = False
+        self.check_limited = check_limited
         self.is_limited = False
         self.Rlim = None
         self.Zlim = None
@@ -511,7 +512,6 @@ def _updateBoundaryPsi(self,psi=None):
                 '''
 
                 if xpt:
-                    #print('Opoint present and checking limited and xpoint present')
                     limit_args = np.ravel(np.argwhere(abs(Zlimit)<abs(0.75*xpt[0][1])))
                     Rlimit = Rlimit[limit_args]
                     Zlimit = Zlimit[limit_args]
diff --git a/freegs/geqdsk.py b/freegs/geqdsk.py
index cc834e9..60a40a4 100755
--- a/freegs/geqdsk.py
+++ b/freegs/geqdsk.py
@@ -328,14 +328,16 @@ def q_func(psinorm):
         Zmax=data["zmid"] + 0.5 * data["zdim"],
         nx=nx,
         ny=ny,  # Number of grid points
+        check_limited = True,
+        psi = psi
     )
     # Grid spacing
     dR = eq.R[1, 0] - eq.R[0, 0]
     dZ = eq.Z[0, 1] - eq.Z[0, 0]
 
-    # Create masking function: 1 inside plasma, 0 outside
-    opoint, xpoint = critical.find_critical(eq.R, eq.Z, psi)
-    mask = critical.core_mask(eq.R, eq.Z, psi, opoint, xpoint, psi_bndry)
+    psi_bndry = eq.psi_bndry
+    psi_axis = eq.psi_axis
+    mask = eq.mask
 
     # Toroidal current
     Jtor = eq.R * pprime_func(psi_norm) + ffprime_func(psi_norm) / (eq.R * mu0)
@@ -366,21 +368,21 @@ def q_func(psinorm):
         # Create a new Equilibrium object
         # (replacing previous 'eq')
         eq = Equilibrium(
-            tokamak=machine, Rmin=Rmin, Rmax=Rmax, Zmin=Zmin, Zmax=Zmax, nx=nx, ny=ny
+            tokamak=machine, Rmin=Rmin, Rmax=Rmax, Zmin=Zmin, Zmax=Zmax, nx=nx, ny=ny, check_limited=True
         )
 
         # Interpolate Jtor and psi onto new grid
         Jtor = Jtor_func(eq.R, eq.Z, grid=False)
         psi = psi_func(eq.R, eq.Z, grid=False)
 
+        psi_bndry = eq.psi_bndry
+        psi_axis = eq.psi_axis
+        mask = eq.mask
+
         # Update the mask function by calculating normalised psi
         # on the new grid
         psi_norm = clip((psi - psi_axis) / (psi_bndry - psi_axis), 0.0, 1.0)
 
-        # Create masking function: 1 inside plasma, 0 outside
-        opoint, xpoint = critical.find_critical(eq.R, eq.Z, psi)
-        mask = critical.core_mask(eq.R, eq.Z, psi, opoint, xpoint, psi_bndry)
-
     # Note: Here we have
     #   eq : Equilibrium object
     #   Jtor : 2D array (nx,ny) Toroidal current density
@@ -390,6 +392,7 @@ def q_func(psinorm):
 
     # Perform a linear solve to calculate psi
     # using known Jtor
+    eq.check_limited = True
     eq.solve(profiles, Jtor=Jtor)
 
     print(
@@ -456,6 +459,7 @@ def q_func(psinorm):
         rtol=rtol,
         blend=blend,
         maxits=maxits,
+        check_limited=True
     )
 
     print(