EQDSK support for geometry

[fusionby2030]

TORAX has a lot of potential and I have been excited to use it. However, I believe an outstanding issue is the lack of generalized geometry support. To this end, I would like to help add EQDSK support. Many EQDSK reading and writing tools exist already (most derived from Ben Dudson, c.f. TSVV-15's or freegs), which are able to handle the COCOS transformations given an g/f/-EQDSK file. The main trick is taking equilibrium from EQDSK and calculating the flux surface averaged quantities. To this end, I append the following code which does this and compare it to a CHEASE run.

from eqdsk import EQDSKInterface # https://github.com/Fusion-Power-Plant-Framework/eqdsk
from scipy.integrate import cumulative_trapezoid
from scipy.interpolate import interp1d, RectBivariateSpline
from contourpy import contour_generator 
import numpy as np 

eq_fpath = './eqdsk_cocos02.eqdsk'
eq       = EQDSKInterface.from_file(eq_fpath, no_cocos=True)
eqfile   = eq.__dict__

vaccum_permativity = 1.25663706212e-6

# ---- Scalars 
RMAJ = eqfile['xmag']
RMIN = eqfile['xbdry'].max() - RMAJ
B0   = eqfile['bcentre']
# ---- Grids
# Psi Uniform Grid
psi_eqdsk_1dgrid = np.linspace(eqfile['psimag'], eqfile['psibdry'], eqfile['nx'])

# X and Z 1D grid, also Meshgrid
X_1D = np.linspace(eqfile['xgrid1'], eqfile['xgrid1'] + eqfile['xdim'], eqfile['nx'])
Z_1D = np.linspace(eqfile['zmid'] - eqfile['zdim']/2, eqfile['zmid'] + eqfile['zdim']/2, eqfile['nz'])
X, Z = np.meshgrid(X_1D, Z_1D, indexing='ij')

# Plasma Bndry (Last closed flux surface)
Xlcfs, Zlcfs = eqfile['xbdry'], eqfile['zbdry']

# Psi 2D grid defined on the Meshgrid
psi_eqdsk_2dgrid  = eqfile['psi']
# normalised Psi w.r.t boundary
psin_eqdsk_2dgrid = (psi_eqdsk_2dgrid - eqfile['psimag']) / (eqfile['psibdry'] - eqfile['psimag'])

# Create mask for the confined region, i.e.,Xlcfs.min() < X < Xlcfs.max(), Zlcfs.min() < Z < Zlcfs.max()
offset = 0.01
mask = (X > Xlcfs.min() - offset) & (X < Xlcfs.max() + offset) & (Z > Zlcfs.min() - offset) & (Z < Zlcfs.max() + offset)
masked_psi_eqdsk_2dgrid = np.ma.masked_where(~mask, psi_eqdsk_2dgrid)

# q on uniform grid (pressure, etc., also defined here)
q_eqdsk_1dgrid = eqfile['qpsi']

# ---- Interpolations
q_interp        = interp1d(psi_eqdsk_1dgrid, eqfile['qpsi'], kind='cubic')
psi_spline_fit  = RectBivariateSpline(X_1D, Z_1D, psi_eqdsk_2dgrid, kx=3, ky=3, s=0)
F_interp        = interp1d(psi_eqdsk_1dgrid, eqfile['fpol'], kind='cubic') # toroidal field flux function 


# -----------------------------------------------------------
# --------- Make flux surface contours ---------
# -----------------------------------------------------------

epsilon = 1e-7 # numerics avoidance for 0.0 
# Create contours on the following psi grid 
psi_interpolant  = np.linspace(eqfile['psimag'] + epsilon, eqfile['psibdry'] - epsilon, 100)

surfaces = [] 
cg_psi = contour_generator(X, Z, masked_psi_eqdsk_2dgrid)

for i, _psi in enumerate(psi_interpolant): 
    vertices = cg_psi.create_contour(_psi)
    x_surface, z_surface = vertices[0].T[0], vertices[0].T[1]
    surfaces.append((x_surface, z_surface))
   

# -----------------------------------------------------------
# --------- Compute Flux surface averages and 1D profiles ---------
# --- Area, Volume, R_inboard, R_outboard
# --- FSA: <1/R^2>, <Bp^2>, <|grad(psi)|>, <|grad(psi)|^2>
# --- Toroidal plasma current
# --- Integral dl/Bp
# -----------------------------------------------------------

# helper function to calculate area of a polygon given curve along x, z
def calculate_area(x, z):
    # Gauss-shoelace formula (https://en.wikipedia.org/wiki/Shoelace_formula)
    # could likely use the determinant approach
    #  0.5 * np.abs(np.dot(x, np.roll(z, 1)) - np.dot(z, np.roll(x, 1)))
    n = len(x)
    area = 0.0
    for i in range(n):
        j = (i + 1) % n  # roll over at n
        area += x[i] * z[j]
        area -= z[i] * x[j]
    area = abs(area) / 2.0
    return area 

# Staging profiles
areas, volumes                = [], []
R_inboard, R_outboard         = [], []
flux_surf_avg_1_over_R2_eqdsk = [] # <1/R**2>
flux_surf_avg_Bp2_eqdsk       = [] # <Bp**2>
flux_surf_avg_RBp_eqdsk       = [] # <|grad(psi)|>
flux_surf_avg_R2Bp2_eqdsk     = [] # <|grad(psi)|**2>
int_dl_over_Bp_eqdsk          = [] # int(Rdl / | grad(psi) |)
Ip_eqdsk                      = [] # Toroidal plasma current


# ---- Compute 
for n, (x_surface, z_surface) in enumerate(surfaces):
    surface_poloidal_angle = ((np.arctan2(z_surface - eqfile['zmag'], x_surface - eqfile['xmag'])) + np.pi)
    
    # dl, line elements on which we will integrate 
    surface_dl = np.sqrt(np.gradient(x_surface) ** 2 + np.gradient(z_surface) ** 2)
    
    # calculating gradient of psi in 2D
    surface_dpsi_x = psi_spline_fit.ev(x_surface, z_surface, dx=1)
    surface_dpsi_z = psi_spline_fit.ev(x_surface, z_surface, dy=1)
    surface_abs_grad_psi = np.sqrt(surface_dpsi_x**2 + surface_dpsi_z**2)
    
    # Poloidal field strength Bp = |grad(psi)| / R
    surface__Bpol = surface_abs_grad_psi / x_surface 
    surface_int_dl_over_bpol = np.sum(surface_dl / surface__Bpol) # This is denominator of all FSA 
    
    # plasma current 
    surface_int_bpol_dl = np.sum(surface__Bpol * surface_dl) 
    
    # 4 FSA, < 1/ R^2>, < | grad psi | >, < B_pol^2>, < | grad psi |^2 >
    # where FSA(G) = int (G dl / Bpol) / (int (dl / Bpol))
    surface_FSA_int_one_over_r2 = np.sum(1 / x_surface ** 2 * surface_dl / surface__Bpol) / surface_int_dl_over_bpol
    surface_FSA_abs_grad_psi = np.sum(surface_abs_grad_psi * surface_dl / surface__Bpol) / surface_int_dl_over_bpol
    surface_FSA_Bpol_sqrd = np.sum(surface__Bpol * surface_dl) / surface_int_dl_over_bpol
    surface_FSA_abs_grad_psi2 = np.sum(surface_abs_grad_psi ** 2 * surface_dl / surface__Bpol) / surface_int_dl_over_bpol
    
    # volumes and areas
    area = calculate_area(x_surface, z_surface)
    volume = area * 2 * np.pi * RMAJ
    
    # Append to lists 
    areas.append(area)
    volumes.append(volume)
    R_inboard.append(x_surface.min())
    R_outboard.append(x_surface.max())
    int_dl_over_Bp_eqdsk.append(surface_int_dl_over_bpol)
    flux_surf_avg_1_over_R2_eqdsk.append(surface_FSA_int_one_over_r2)
    flux_surf_avg_RBp_eqdsk.append(surface_FSA_abs_grad_psi)
    flux_surf_avg_R2Bp2_eqdsk.append(surface_FSA_abs_grad_psi2)
    flux_surf_avg_Bp2_eqdsk.append(surface_FSA_Bpol_sqrd)
    Ip_eqdsk.append(surface_int_bpol_dl / vaccum_permativity) 



# -----------------------------------------------------------
# ---------  Compute 1D quantties  ---------
# --- q-profile
# --- Toroidal flux
# -----------------------------------------------------------

# q-profile on interpolation 
q_profile = q_interp(psi_interpolant)

# toroidal flux 
Phi_eqdsk = cumulative_trapezoid(q_profile, psi_interpolant, initial=0.0) * 2*np.pi 

# toroidal field flux function, T=RBphi
F_eqdsk = F_interp(psi_interpolant) 

# -----------------------------------------------------------
# --------- Compute scalars ---------
# --- Upper and lower triangularity 
# -----------------------------------------------------------

idx_upperextent = np.argmax(Zlcfs)
idx_lowerextent = np.argmin(Zlcfs)

X_upperextent = Xlcfs[idx_upperextent]
X_lowerextent = Xlcfs[idx_lowerextent]

delta_upper_face = (RMAJ - X_upperextent) / RMIN 
delta_lower_face = (RMAJ - X_lowerextent) / RMIN

I have compared the above transformation with that from CHEASE when ran with the same equilibrium.

I expect that with modifications (e.g., porting to JAX) this could be implemented relatively easily into the geometry.py, but I expect it would take much less time for the TORAX team to implement than it would for me to do this, which is why I only supply the above code to help inspire and start the discussion. I attach the following files to reproduce this above (both have filenames that end in .txt which is the only supported version filetype to upload):

• eqdsk_cocos02.eqdsk, the G-EQDSK file used above (arbitrary tokamak) eqdsk_cocos02.eqdsk.txt
• chease.cols, the CHEASE file used in comparison chease.cols.txt

[jcitrin]

Hi Adam. Thanks for this, it's really great and very much appreciated!

Quickly looking through it, the approach looks very good and no problem to depend on the eqdsk and contour_generator libraries, due to their permissive licenses.

Thankfully, there is no need to port any of this code to JAX since all geometry processing is done in the initialization phase. You can take a look e.g. at

torax/torax/geometry.py

Line 832 in dbde9b4

def from_fbt(

, which is one of the class methods for generating the geometry structures based on a specific backend (FBT in this case).

What would be needed for EQDSK, is to populate the StandardGeometryIntermediates class using a similar constructor method, and then propagate some of the logic elsewhere in TORAX, e.g. add the extra EQDSK option to the various geometry_provider methods in build_sim.py, reusing the existing patterns.

We'd be happy to accept a PR along these lines, and can actively help/co-develop if you prefer. Can also discuss over video call.

Regarding the comparisons, it's already looking very good. I'm wondering though if there is an explanation for some of the systematic differences, e.g. for <1/R^2>, Rout, Rin.

Ultimately it would also be good to have an integration benchmark against another integrated modelling tool using the same input EQDSK file, with same simple sources, transport, and boundary conditions.