Example: Evolutive equilibrium calculations
This notebook will demonstrate the core use-case for FreeGSNKE: simulating the (time-dependent) evolution of Grad-Shafranov (GS) equilibria. In particular we will simulate a vertical displacement event (VDE).
To do this, we need to:
build the tokamak machine.
instatiate a GS equilibrium (to be used as an initial condition for the evolutive solver).
calculate a vertical instability growth rate for this equilibrium and carry out passive structure mode removal via a normal mode decomposition (i.e. removing modes that have little effect on the evolution).
define time-dependent plasma current density profile parameters and coil voltages.
evolve the active coil currents, the total plasma current, and the equilbirium using these profile parameters and voltages by solving the circuit equations alongside the GS equation.
Refer to the paper by Amorisco et al. (2024) for more details.
We should note that here we will use fixed (time-independent) profile parameters and voltages to simulate a VDE, however, truly time-dependent parameters would be required to simulate a plasma shot (see future notebooks).
Import packages
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from IPython.display import display, clear_output
import pickle
Build the machine
# build machine
from freegsnke import build_machine
tokamak = build_machine.tokamak(
active_coils_path=f"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle",
passive_coils_path=f"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle",
limiter_path=f"../machine_configs/MAST-U/MAST-U_like_limiter.pickle",
wall_path=f"../machine_configs/MAST-U/MAST-U_like_wall.pickle",
)
Active coils --> built from pickle file.
Passive structures --> built from pickle file.
Limiter --> built from pickle file.
Wall --> built from pickle file.
Magnetic probes --> none provided.
Resistance (R) and inductance (M) matrices --> built using actives (and passives if present).
Tokamak built.
Diverted plasma example
Instantiate an equilibrium
from freegsnke import equilibrium_update
eq = equilibrium_update.Equilibrium(
tokamak=tokamak,
Rmin=0.1, Rmax=2.0, # Radial range
Zmin=-2.2, Zmax=2.2, # Vertical range
nx=65, # Number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)
ny=129, # Number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)
# psi=plasma_psi
)
Instatiate a profile object
from freegsnke.jtor_update import ConstrainPaxisIp
profiles = ConstrainPaxisIp(
eq=eq, # equilibrium object
paxis=8.1e3, # profile object
Ip=6.2e5, # plasma current
fvac=0.5, # fvac = rB_{tor}
alpha_m=1.8, # profile function parameter
alpha_n=1.2 # profile function parameter
)
Set coil currents
Here we set coil currents that create a diverted plasma (as seen in prior notebooks).
with open('simple_diverted_currents_PaxisIp.pk', 'rb') as f:
current_values = pickle.load(f)
for key in current_values.keys():
eq.tokamak.set_coil_current(coil_label=key, current_value=current_values[key])
Instatiate the solver
from freegsnke import GSstaticsolver
GSStaticSolver = GSstaticsolver.NKGSsolver(eq)
Call forward solver to find equilibrium
GSStaticSolver.solve(eq=eq,
profiles=profiles,
constrain=None,
target_relative_tolerance=1e-8,
verbose=0
)
Forward static solve SUCCESS. Tolerance 8.07e-10 (vs. requested 1.00e-08) reached in 26/100 iterations.
Plot the initial equilibrium
fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)
ax1.grid(True, which='both')
eq.plot(axis=ax1, show=False)
eq.tokamak.plot(axis=ax1, show=False)
ax1.set_xlim(0.1, 2.15)
ax1.set_ylim(-2.25, 2.25)
plt.tight_layout()

Time evolution
We are now ready to solve the forward time-evolutive problem.
Initialise nonlinear (evolutive) solver
To start, we need to initialise the evolutive solver object, freegsnke.nonlinear_solve.nl_solver
.
The time-evolutive solver requires:
eq
: an equilibrium to inform the solver of the machine and domain properties.profiles
: defined above.full_timestep
: the time step by which time is advanced at every call of the stepper (this can be modified later depending on the growth rate of the equilibrium). An appropriate time step can also be set based on the growth rate calculation. Useautomatic_timestep
to set the time step in units of the (inverse of the) growth rate.plasma_resistivity
: resistivity of the plasma (which here is assumed to be constant during the time evolution but can be made time-dependent if known).min_dIy_dI
: threshold value below which passive structure normal modes are dropped. Modes with norm(d(Iy)/dI)<min_dIy_dI
are dropped, which filters out modes that do not actually couple with the plasma.max_mode_frequency
: threshold value for characteristic frequencies above which passive structure normal modes are dropped (i.e. the fast modes).
Other customisable inputs are available, do see the documentation for more details. For example, one may explicitly set your own resistance and inductance matries for the tokamak machine.
The solver can be used on different equilibria and/or profiles, but these need to have the same machine, domain, and limiter as the one used at the time of the solver instantiation. For different machines, a new time-evolutive solver should be created.
The input equilibrium and profile functions are also used as the expansion point around which the dynamics are linearised.
from freegsnke import nonlinear_solve
stepping = nonlinear_solve.nl_solver(
eq=eq,
profiles=profiles,
full_timestep=5e-4,
plasma_resistivity=1e-6,
min_dIy_dI=0.1,
max_mode_frequency=10**2.5,
)
-----
Checking that the provided 'eq' and 'profiles' are a GS solution...
Forward static solve SUCCESS. Tolerance 8.07e-10 (vs. requested 1.00e-08) reached in 0/100 iterations.
-----
Instantiating nonlinear solver objects...
done.
-----
Identifying mode selection criteria...
'threshold_dIy_dI', 'min_dIy_dI', and 'max_mode_frequency' option selected --> passive structure modes are selected according to these thresholds.
-----
Initial mode selection:
Active coils
total selected = 12 (out of 12)
Passive structures
16 selected with characteristic timescales larger than 'max_mode_frequency'
11 selected that couple with the plasma more than 'threshold_dIy_dI', despite having fast timescales
4 selected that couple with the plasma less than 'min_dIy_dI', despite having slow timescales
total selected = 23 (out of 138)
Total number of modes = 35 (12 active coils + 23 passive structures)
(Note: some additional modes may be removed after Jacobian calculation)
-----
Building the 3051 x 36 Jacobian (dIy/dI) of plasma current density (inside the LCFS) with respect to all metal currents and the total plasma current.
done.
-----
Further mode reduction:
6 previously included modes couple with the plasma less than 'min_dIy_dI' (following Jacobian calculation)
Final number of modes = 29 (12 active coils + 17 passive structures)
Re-sizing the Jacobian matrix to account for any removed modes (if required).
-----
Linear growth calculations:
Growth rate = [102.99389692] [1/s]
Instability timescale = [0.00970931] [s]
Stability margin = [0.39377101+0.j]
-----
Evolutive solver timestep:
Solver timestep 'dt_step' has been set to 0.0005 as requested.
Ensure it is smaller than the growth rate else you may find numerical instability in any subsequent evoltuive simulations!
-----
Set active coil voltages
In this example, we will evolve a plasma in absence of any control policy or current drive.
Just as an example, the following calculates active voltages to be applied to the poloidal field coils (and Solenoid) using \(V = RI\), with current values as defined by the initial equilibrium (i.e. we have constant voltages).
In most FreeGSNKE use cases, these active voltages will be determined by a control policy.
voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.R)[:stepping.evol_metal_curr.n_active_coils]
To start, the solver is prepared by setting the initial conditions.
stepping.initialize_from_ICs(eq, profiles)
done.
Set time steps and storage lists
Now we set the total number of time steps we want to simulate and initialise some variables for logging the evolution.
# number of time steps to simulate
max_count = 45
# initialising some variables for iteration and logging
counter = 0
t = 0
history_times = [t]
history_currents = [stepping.currents_vec]
history_equilibria = [deepcopy(stepping.eq1)]
separatrix = stepping.eq1.separatrix(ntheta=100)
history_width = [np.amax(separatrix[:,0]) - np.amin(separatrix[:,0])]
history_o_points = [stepping.eq1.opt[0]]
history_elongation = [(np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[0]]
Call the solver (linear)
Finally, we call the time-evolutive solver with stepping.nlstepper()
sequentially until we reach the preset end time.
The following demonstrates a solely linear evolution of the plasma by setting linear_only=True
.
# loop over time steps
while counter<max_count:
clear_output(wait=True)
display(f'Step: {counter}/{max_count-1}')
display(f'Time: t = {t:.2e}')
# carry out the time step
stepping.nlstepper(
active_voltage_vec=voltages,
linear_only=True,
verbose=0
)
# store information on the time step
t += stepping.dt_step
history_times.append(t)
counter += 1
# store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)
history_currents.append(stepping.currents_vec)
history_equilibria.append(deepcopy(stepping.eq1))
separatrix = stepping.eq1.separatrix(ntheta=100)
history_width.append(np.amax(separatrix[:,0]) - np.amin(separatrix[:,0]))
history_o_points = np.append(history_o_points, [stepping.eq1.opt[0]], axis=0)
history_elongation.append((np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[-1])
# transform lists to arrays
history_currents = np.array(history_currents)
history_times = np.array(history_times)
history_o_points = np.array(history_o_points)
'Step: 44/44'
'Time: t = 2.20e-02'
fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)
ax1.grid(True, which='both')
history_equilibria[-1].plot(axis=ax1, show=False)
eq.tokamak.plot(axis=ax1, show=False)
ax1.set_xlim(0.1, 2.15)
ax1.set_ylim(-2.25, 2.25)
plt.tight_layout()

Call the solver (nonlinear)
Now we evolve the same plasma according to the full nonlinear dynamics for the same time interval. This is done by using linear_only=False
in the call to the stepper.
We need to re-initialise from the initial conditions and reset the counter, but otherwise the method is identical to the one above.
Note that the full nonlinear evolutive solve is a lot more computationally expensive than solving the linear evolutive problem. As such, the following cell may take several minutes to execute, depending on your hardware.
# reset the solver object by resetting the initial conditions
stepping.initialize_from_ICs(eq, profiles)
# initialising some variables for iteration and logging
counter = 0
t = 0
history_times_nl = [t]
history_currents_nl = [stepping.currents_vec]
history_equilibria_nl = [deepcopy(stepping.eq1)]
separatrix = stepping.eq1.separatrix(ntheta=100)
history_width_nl = [np.amax(separatrix[:,0]) - np.amin(separatrix[:,0])]
history_o_points_nl = [stepping.eq1.opt[0]]
history_elongation_nl = [(np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[0]]
# loop over the time steps
while counter<max_count:
clear_output(wait=True)
display(f'Step: {counter}/{max_count-1}')
display(f'Time: t = {t:.2e}')
# carry out the time step
stepping.nlstepper(
active_voltage_vec=voltages,
linear_only=False,
verbose=0
)
# store information on the time step
t += stepping.dt_step
history_times_nl.append(t)
counter += 1
# store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)
history_currents_nl.append(stepping.currents_vec)
history_equilibria_nl.append(deepcopy(stepping.eq1))
separatrix = stepping.eq1.separatrix(ntheta=100)
history_width_nl.append(np.amax(separatrix[:,0]) - np.amin(separatrix[:,0]))
history_o_points_nl = np.append(history_o_points_nl, [stepping.eq1.opt[0]], axis=0)
history_elongation_nl.append((np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width_nl[-1])
# transform lists to arrays
history_currents_nl = np.array(history_currents_nl)
history_times_nl = np.array(history_times_nl)
history_o_points_nl = np.array(history_o_points_nl)
'Step: 44/44'
'Time: t = 2.20e-02'
Forward static solve SUCCESS. Tolerance 3.93e-07 (vs. requested 2.70e-06) reached in 5/100 iterations.
Forward static solve SUCCESS. Tolerance 1.44e-06 (vs. requested 2.70e-06) reached in 3/100 iterations.
Forward static solve SUCCESS. Tolerance 3.56e-07 (vs. requested 2.70e-06) reached in 4/100 iterations.
Forward static solve SUCCESS. Tolerance 2.80e-07 (vs. requested 2.70e-06) reached in 4/100 iterations.
Forward static solve SUCCESS. Tolerance 3.01e-07 (vs. requested 2.70e-06) reached in 4/100 iterations.
Forward static solve SUCCESS. Tolerance 5.19e-07 (vs. requested 2.70e-06) reached in 4/100 iterations.
Forward static solve SUCCESS. Tolerance 1.02e-06 (vs. requested 2.70e-06) reached in 1/100 iterations.
Plot some time-evolving quantities
The following plots the evolution of a number of tracked values and compares the linear/nonlinear evoltuions against one another.
# Plot evolution of tracked values and compare between linear and non-linear evolution
fig, axs = plt.subplots(2, 3, figsize=(10, 5), dpi=80, constrained_layout=True)
axs_flat = axs.flat
axs_flat[0].plot(history_times, history_o_points[:, 0],'k+', label='linear')
axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')
axs_flat[0].set_xlabel('Time')
axs_flat[0].set_ylabel('O-point $R$')
axs_flat[0].legend()
axs_flat[1].plot(history_times, abs(history_o_points[:, 1]),'k+')
axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')
axs_flat[1].set_yscale('log')
axs_flat[1].set_xlabel('Time')
axs_flat[1].set_ylabel('abs( O-point $Z$ )')
axs_flat[2].plot(history_times, history_o_points[:, 2],'k+')
axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')
axs_flat[2].set_xlabel('Time')
axs_flat[2].set_ylabel('O-point $\Psi$')
axs_flat[3].plot(history_times, history_currents[:,-1]*stepping.plasma_norm_factor,'k+')
axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')
axs_flat[3].set_xlabel('Time')
axs_flat[3].set_ylabel('Plasma current')
axs_flat[4].plot(history_times, history_width,'k+')
axs_flat[4].plot(history_times_nl, history_width_nl,'rx')
axs_flat[4].set_xlabel('Time')
axs_flat[4].set_ylabel('Plasma width')
axs_flat[5].plot(history_times, history_elongation,'k+')
axs_flat[5].plot(history_times_nl, history_elongation_nl,'rx')
axs_flat[5].set_xlabel('Time')
axs_flat[5].set_ylabel('Plasma elongation')
Text(0, 0.5, 'Plasma elongation')

# plot the equilibria at the final time step
fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)
ax1.grid(True, which='both')
history_equilibria[-1].plot(axis=ax1, show=False)
eq.tokamak.plot(axis=ax1, show=False)
ax1.set_xlim(0.1, 2.15)
ax1.set_ylim(-2.25, 2.25)
ax1.set_title("Linear solve")
ax2.grid(True, which='both')
history_equilibria_nl[-1].plot(axis=ax2, show=False)
eq.tokamak.plot(axis=ax2, show=False)
ax2.set_xlim(0.1, 2.15)
ax2.set_ylim(-2.25, 2.25)
ax2.set_title("Nonlinear solve")
plt.tight_layout()

Limited plasma example
Finally, we examine a limiter configuration (there is no fundamental difference from the process described above).
First we re-instantiate the equilibrium object.
eq = equilibrium_update.Equilibrium(
tokamak=tokamak,
Rmin=0.1, Rmax=2.0,
Zmin=-2.2, Zmax=2.2,
nx=65,
ny=129
)
We load a set of coil currents that will give us a limiter plasma in our current tokamak.
import json
with open('limiter_currents.json', 'r') as f:
current_values = json.load(f)
for key in current_values.keys():
eq.tokamak.set_coil_current(key, current_values[key])
Now we set the profiles.
profiles = ConstrainPaxisIp(
eq=eq,
paxis=8.1e3,
Ip=6.2e5,
fvac=0.5,
alpha_m=1.8,
alpha_n=1.2
)
The final part of the setup is to do a forward solve and visualise the limiter equilibrium. The red dashed line is the flux surface through the first X-point, however, the actual last closed flux surface in limiter equilibria is displayed as a full black line.
GSStaticSolver.solve(eq, profiles, target_relative_tolerance=1e-8)
fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)
ax1.grid(True, which='both')
eq.plot(axis=ax1, show=False)
eq.tokamak.plot(axis=ax1, show=False)
ax1.set_xlim(0.1, 2.15)
ax1.set_ylim(-2.25, 2.25)
plt.tight_layout()
Forward static solve SUCCESS. Tolerance 2.51e-09 (vs. requested 1.00e-08) reached in 24/100 iterations.

Finally, we can do the nonlinear evolutive solve as before. First, we re-initialize the solver, which calculates the linearization of the dynamics relevant to this new equilibrium.
from freegsnke import nonlinear_solve
stepping = nonlinear_solve.nl_solver(
eq=eq, profiles=profiles,
full_timestep=.5e-3,
plasma_resistivity=1e-6,
)
-----
Checking that the provided 'eq' and 'profiles' are a GS solution...
Forward static solve SUCCESS. Tolerance 2.51e-09 (vs. requested 1.00e-08) reached in 0/100 iterations.
-----
Instantiating nonlinear solver objects...
done.
-----
Identifying mode selection criteria...
'threshold_dIy_dI', 'min_dIy_dI', and 'max_mode_frequency' option selected --> passive structure modes are selected according to these thresholds.
-----
Initial mode selection:
Active coils
total selected = 12 (out of 12)
Passive structures
4 selected with characteristic timescales larger than 'max_mode_frequency'
23 selected that couple with the plasma more than 'threshold_dIy_dI', despite having fast timescales
0 selected that couple with the plasma less than 'min_dIy_dI', despite having slow timescales
total selected = 27 (out of 138)
Total number of modes = 39 (12 active coils + 27 passive structures)
(Note: some additional modes may be removed after Jacobian calculation)
-----
Building the 3051 x 40 Jacobian (dIy/dI) of plasma current density (inside the LCFS) with respect to all metal currents and the total plasma current.
done.
-----
Further mode reduction:
4 previously included modes couple with the plasma less than 'min_dIy_dI' (following Jacobian calculation)
Final number of modes = 35 (12 active coils + 23 passive structures)
Re-sizing the Jacobian matrix to account for any removed modes (if required).
-----
Linear growth calculations:
Growth rate = [3.71942774e+00+0.j 9.23744806e+04+0.j] [1/s]
Instability timescale = [2.68858564e-01-0.j 1.08255007e-05-0.j] [s]
Stability margin = [1.52179283+0.j]
-----
Evolutive solver timestep:
Solver timestep 'dt_step' has been set to 0.0005 as requested.
Ensure it is smaller than the growth rate else you may find numerical instability in any subsequent evoltuive simulations!
-----
# recalculate the active voltages using the new currents
U_active = (stepping.vessel_currents_vec*stepping.evol_metal_curr.R)[:stepping.evol_metal_curr.n_active_coils]
# number of time steps to simulate
max_count = 45
# reset the solver object by resetting the initial conditions
stepping.initialize_from_ICs(eq, profiles)
# initialising some variables for iteration and logging
counter = 0
t = 0
history_times_nl = [t]
history_currents_nl = [stepping.currents_vec]
history_equilibria_nl = [deepcopy(stepping.eq1)]
separatrix = stepping.eq1.separatrix(ntheta=100)
history_width_nl = [np.amax(separatrix[:,0]) - np.amin(separatrix[:,0])]
history_o_points_nl = [stepping.eq1.opt[0]]
history_elongation_nl = [(np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[0]]
# Simulate the forward evolution
while counter<max_count:
clear_output(wait=True)
display(f'Step: {counter}/{max_count-1}')
display(f'current time t = {t:.2e}')
# carry out the time step
stepping.nlstepper(
active_voltage_vec=U_active,
linear_only=False,
target_relative_tol_currents=0.01,
target_relative_tol_GS=0.01,
)
# store information on the time step
t += stepping.dt_step
history_times_nl.append(t)
counter += 1
# store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)
history_currents_nl.append(stepping.currents_vec)
history_equilibria_nl.append(deepcopy(stepping.eq1))
separatrix = stepping.eq1.separatrix(ntheta=100)
history_width_nl.append(np.amax(separatrix[:,0]) - np.amin(separatrix[:,0]))
history_o_points_nl = np.append(history_o_points_nl, [stepping.eq1.opt[0]], axis=0)
history_elongation_nl.append((np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width_nl[-1])
# transform lists to arrays
history_currents_nl = np.array(history_currents_nl)
history_times_nl = np.array(history_times_nl)
history_o_points_nl = np.array(history_o_points_nl)
'Step: 44/44'
'current time t = 2.20e-02'
Forward static solve SUCCESS. Tolerance 1.16e-07 (vs. requested 4.40e-07) reached in 4/100 iterations.
Forward static solve SUCCESS. Tolerance 9.03e-08 (vs. requested 4.40e-07) reached in 2/100 iterations.
Forward static solve SUCCESS. Tolerance 3.10e-08 (vs. requested 4.40e-07) reached in 2/100 iterations.
# Plot evolution of tracked values
fig, axs = plt.subplots(2, 3, figsize=(10, 5), dpi=80, constrained_layout=True)
axs_flat = axs.flat
axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')
axs_flat[0].set_xlabel('Time')
axs_flat[0].set_ylabel('O-point $R$')
axs_flat[0].legend()
axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')
axs_flat[1].set_yscale('log')
axs_flat[1].set_xlabel('Time')
axs_flat[1].set_ylabel('abs( O-point $Z$ )')
axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')
axs_flat[2].set_xlabel('Time')
axs_flat[2].set_ylabel('O-point $\Psi$')
axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')
axs_flat[3].set_xlabel('Time')
axs_flat[3].set_ylabel('Plasma current')
axs_flat[4].plot(history_times_nl, history_width_nl,'rx')
axs_flat[4].set_xlabel('Time')
axs_flat[4].set_ylabel('Plasma width')
axs_flat[5].plot(history_times_nl, history_elongation_nl,'rx')
axs_flat[5].set_xlabel('Time')
axs_flat[5].set_ylabel('Plasma elongation')
Text(0, 0.5, 'Plasma elongation')

# plot the equilibria at the final time step
fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)
ax1.grid(True, which='both')
history_equilibria_nl[-1].plot(axis=ax1, show=False)
eq.tokamak.plot(axis=ax1, show=False)
ax1.set_xlim(0.1, 2.15)
ax1.set_ylim(-2.25, 2.25)
ax1.set_title('t=0')
ax2.grid(True, which='both')
history_equilibria_nl[-1].plot(axis=ax2, show=False)
eq.tokamak.plot(axis=ax2, show=False)
ax2.set_xlim(0.1, 2.15)
ax2.set_ylim(-2.25, 2.25)
ax2.set_title('t='+str(t))
plt.tight_layout()
