You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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
The text was updated successfully, but these errors were encountered:
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
, 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.
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.
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.txtchease.cols
, the CHEASE file used in comparison chease.cols.txtThe text was updated successfully, but these errors were encountered: