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). More details on this in a later notebook.

• define time-dependent coil voltages, plasma current density profile parameters, and plasma resistivity.

• 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.

Import packages

import numpy as np
import matplotlib.pyplot as plt
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('data/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 2.93e-09 (vs. requested 1.00e-08) reached in 24/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. Use automatic_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 or the later notebook on “Growth Rates” for more details. For example, one may explicitly set your own resistance and inductance matries for the tokamak, rather than the geometrical value calculated internally in FreeGSNKE.

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 (when using the linear solver or calculating linear growth rates).

from freegsnke import nonlinear_solve

stepping = nonlinear_solve.nl_solver(
    eq=eq, 
    profiles=profiles, 
    GSStaticSolver=GSStaticSolver,
    full_timestep=5e-4, 
    plasma_resistivity=1e-6,
    max_mode_frequency=10**2.5,
)

-----
Checking that the provided 'eq' and 'profiles' are a GS solution...
Forward static solve SUCCESS. Tolerance 2.93e-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' options 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'
      24 recovered that couple with the plasma more than 'threshold_dIy_dI'
      4 removed that couple with the plasma less than 'min_dIy_dI'
      total selected = 36 (out of 138)
   Total number of modes = 48 (12 active coils + 36 passive structures)
      (Note: some additional modes may be removed after Jacobian calculation)
-----
Building the 3051 x 49 Jacobian (dIy/dI) of plasma current density (inside the LCFS) with respect to all metal currents and the total plasma current.

Building the 3051 x 3 Jacobian (dIy/dtheta) of plasma current density (inside the LCFS) with respect to all plasma current density profile parameters within Jtor.

-----
Further mode reduction:
   1 previously included modes couple with the plasma less than 'min_dIy_dI' (following Jacobian calculation)
   Final number of modes = 47 (12 active coils + 35 passive structures)
   Re-sizing the Jacobian matrix to account for any removed modes (if required).
-----
Stability paramters:

   Deformable plasma metrics:
      Growth rate = [105.74938499] [1/s]
      Instability timescale = [0.00945632] [s]
      Inductive stability margin = [0.45695402+0.j]
   Rigid plasma metrics:
      Leuer parameter (ratio of stabilsing to de-stabilising force gradients):
          between all metals and all metals = 1.7726884043188813
          between all metals and active metals = 1.7726884043188813
          between passive metals and active metals = 1.7726737311867975
-----
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.coil_resist)[:stepping.evol_metal_curr.n_active_coils] 

To start, the solver is prepared by setting the initial conditions.

stepping.initialize_from_ICs(eq, profiles)

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 = 50

# initialising some variables for iteration and logging
counter = 0
t = 0

history_times = [t]
history_currents = [stepping.currents_vec]
history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]
history_o_points = [stepping.eq1.opt[0]]
history_elongation = [stepping.eq1.geometricElongation()]
history_triangularity = [stepping.eq1.triangularity()]
history_squareness = [stepping.eq1.squareness()[1]]
history_area = [stepping.eq1.separatrix_area()]
history_length = [stepping.eq1.separatrix_length()]

Set time-dependent active coil voltages, profile parameters, and plasma resistivity

To simulate the VDE in the example, we will evolve a plasma in absence of any control policy or current drive. 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 over time). For more realistic plasma simulation, these active voltages would be truly time-dependent and determined by a control policy.

In this example, the profile parameters and the plasma resistivity will also remain constant over time. In the next example, we will show how to vary these.

voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] 

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. This makes use of the Jacobians calculated when the nonlinear solver object was initialised. Note that the linear solution will only hold for equilibrium “close to” that which was used to initialise the nonlinear solver object.

# initialise the solver with the initial equilibrium/profiles
stepping.initialize_from_ICs(eq, profiles)

# loop over time steps
while counter<max_count:
    print(f'Step: {counter}/{max_count-1}')
    print(f'--- t = {t:.2e}')

    # carry out the time step (profile parameters and resistivity are constant at each time step)
    stepping.nlstepper(
        active_voltage_vec=voltages,   # same voltages used at each time step
        linear_only=True,
        verbose=False,
        max_solving_iterations=50,
    )

    # 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(stepping.eq1.create_auxiliary_equilibrium())
    history_o_points.append(stepping.eq1.opt[0])
    history_elongation.append(stepping.eq1.geometricElongation())
    history_triangularity.append(stepping.eq1.triangularity())
    history_squareness.append(stepping.eq1.squareness()[1])
    history_area.append(stepping.eq1.separatrix_area())
    history_length.append(stepping.eq1.separatrix_length())

# transform lists to arrays
history_currents = np.array(history_currents)
history_times = np.array(history_times)
history_o_points = np.array(history_o_points)
history_elongation = np.array(history_elongation)
history_triangularity = np.array(history_triangularity)
history_squareness = np.array(history_squareness)
history_area = np.array(history_area)
history_length = np.array(history_length)

Step: 0/49
--- t = 0.00e+00

Step: 1/49
--- t = 5.00e-04

Step: 2/49
--- t = 1.00e-03

Step: 3/49
--- t = 1.50e-03

Step: 4/49
--- t = 2.00e-03

Step: 5/49
--- t = 2.50e-03

Step: 6/49
--- t = 3.00e-03

Step: 7/49
--- t = 3.50e-03

Step: 8/49
--- t = 4.00e-03

Step: 9/49
--- t = 4.50e-03

Step: 10/49
--- t = 5.00e-03

Step: 11/49
--- t = 5.50e-03

Step: 12/49
--- t = 6.00e-03

Step: 13/49
--- t = 6.50e-03

Step: 14/49
--- t = 7.00e-03

Step: 15/49
--- t = 7.50e-03

Step: 16/49
--- t = 8.00e-03

Step: 17/49
--- t = 8.50e-03

Step: 18/49
--- t = 9.00e-03

Step: 19/49
--- t = 9.50e-03

Step: 20/49
--- t = 1.00e-02

Step: 21/49
--- t = 1.05e-02

Step: 22/49
--- t = 1.10e-02

Step: 23/49
--- t = 1.15e-02

Step: 24/49
--- t = 1.20e-02

Step: 25/49
--- t = 1.25e-02

Step: 26/49
--- t = 1.30e-02

Step: 27/49
--- t = 1.35e-02

Step: 28/49
--- t = 1.40e-02

Step: 29/49
--- t = 1.45e-02

Step: 30/49
--- t = 1.50e-02

Step: 31/49
--- t = 1.55e-02

Step: 32/49
--- t = 1.60e-02

Step: 33/49
--- t = 1.65e-02

Step: 34/49
--- t = 1.70e-02

Step: 35/49
--- t = 1.75e-02

Step: 36/49
--- t = 1.80e-02

Step: 37/49
--- t = 1.85e-02

Step: 38/49
--- t = 1.90e-02

Step: 39/49
--- t = 1.95e-02

Step: 40/49
--- t = 2.00e-02

Step: 41/49
--- t = 2.05e-02

Step: 42/49
--- t = 2.10e-02

Step: 43/49
--- t = 2.15e-02

Step: 44/49
--- t = 2.20e-02

Step: 45/49
--- t = 2.25e-02

Step: 46/49
--- t = 2.30e-02

Step: 47/49
--- t = 2.35e-02

Step: 48/49
--- t = 2.40e-02

Step: 49/49
--- t = 2.45e-02

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 = [stepping.eq1.create_auxiliary_equilibrium()]
history_o_points_nl = [stepping.eq1.opt[0]]
history_elongation_nl = [stepping.eq1.geometricElongation()]
history_triangularity_nl = [stepping.eq1.triangularity()]
history_squareness_nl = [stepping.eq1.squareness()[1]]
history_area_nl = [stepping.eq1.separatrix_area()]
history_length_nl = [stepping.eq1.separatrix_length()]

# loop over the time steps
while counter<max_count:

    print(f'Step: {counter}/{max_count-1}')
    print(f'--- t = {t:.2e}')
    
    # carry out the time step
    stepping.nlstepper(
        active_voltage_vec=voltages, 
        linear_only=False,
        verbose=False,
        max_solving_iterations=50,
    )

    # 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(stepping.eq1.create_auxiliary_equilibrium())
    history_o_points_nl.append(stepping.eq1.opt[0])
    history_elongation_nl.append(stepping.eq1.geometricElongation())
    history_triangularity_nl.append(stepping.eq1.triangularity())
    history_squareness_nl.append(stepping.eq1.squareness()[1])
    history_area_nl.append(stepping.eq1.separatrix_area())
    history_length_nl.append(stepping.eq1.separatrix_length())

# 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)
history_elongation_nl = np.array(history_elongation_nl)
history_triangularity_nl = np.array(history_triangularity_nl)
history_squareness_nl = np.array(history_squareness_nl)
history_area_nl = np.array(history_area_nl)
history_length_nl = np.array(history_length_nl)

Step: 0/49
--- t = 0.00e+00

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.95e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.44e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 1/49
--- t = 5.00e-04

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.13e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.14e-03 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 2/49
--- t = 1.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.28e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.60e-03 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 3/49
--- t = 1.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.33e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.58e-03 (vs. requested 3.00e-03).
   Iterations taken: 4/50.

Step: 4/49
--- t = 2.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.94e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.06e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 5/49
--- t = 2.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.04e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.35e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 6/49
--- t = 3.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.91e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.76e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 7/49
--- t = 3.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.02e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.98e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 8/49
--- t = 4.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.71e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.26e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 9/49
--- t = 4.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.65e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.48e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 10/49
--- t = 5.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.89e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.40e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 11/49
--- t = 5.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.66e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.20e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 12/49
--- t = 6.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.47e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.10e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 13/49
--- t = 6.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.05e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.53e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 14/49
--- t = 7.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.13e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.27e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 15/49
--- t = 7.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 5.99e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.08e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 16/49
--- t = 8.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.15e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.80e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 17/49
--- t = 8.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.11e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.66e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 18/49
--- t = 9.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.68e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.84e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 19/49
--- t = 9.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.25e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.96e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 20/49
--- t = 1.00e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.10e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.12e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 21/49
--- t = 1.05e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.08e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.66e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 22/49
--- t = 1.10e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 6.42e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.13e-03 (vs. requested 3.00e-03).
   Iterations taken: 12/50.

Step: 23/49
--- t = 1.15e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.95e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.20e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 24/49
--- t = 1.20e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.42e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.37e-03 (vs. requested 3.00e-03).
   Iterations taken: 12/50.

Step: 25/49
--- t = 1.25e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.97e-05 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.20e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 26/49
--- t = 1.30e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 8.57e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.37e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 27/49
--- t = 1.35e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.78e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.54e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 28/49
--- t = 1.40e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.23e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.00e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 29/49
--- t = 1.45e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.92e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.26e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 30/49
--- t = 1.50e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.85e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.56e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 31/49
--- t = 1.55e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.83e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.90e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 32/49
--- t = 1.60e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.78e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.50e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 33/49
--- t = 1.65e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.30e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.96e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 34/49
--- t = 1.70e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.21e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.58e-03 (vs. requested 3.00e-03).
   Iterations taken: 11/50.

Step: 35/49
--- t = 1.75e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 6.90e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.46e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 36/49
--- t = 1.80e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 5.92e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.73e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 37/49
--- t = 1.85e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.07e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.90e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 38/49
--- t = 1.90e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.12e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.80e-03 (vs. requested 3.00e-03).
   Iterations taken: 12/50.

Step: 39/49
--- t = 1.95e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.87e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.83e-03 (vs. requested 3.00e-03).
   Iterations taken: 11/50.

Step: 40/49
--- t = 2.00e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.63e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.62e-03 (vs. requested 3.00e-03).
   Iterations taken: 12/50.

Step: 41/49
--- t = 2.05e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.97e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.67e-03 (vs. requested 3.00e-03).
   Iterations taken: 11/50.

Step: 42/49
--- t = 2.10e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.56e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.91e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 43/49
--- t = 2.15e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.59e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.84e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 44/49
--- t = 2.20e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.38e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.65e-03 (vs. requested 3.00e-03).
   Iterations taken: 11/50.

Step: 45/49
--- t = 2.25e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 8.18e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.50e-03 (vs. requested 3.00e-03).
   Iterations taken: 11/50.

Step: 46/49
--- t = 2.30e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.63e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.74e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 47/49
--- t = 2.35e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.08e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.97e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 48/49
--- t = 2.40e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.75e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.46e-03 (vs. requested 3.00e-03).
   Iterations taken: 12/50.

/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/freegs4e/critical.py:1045: UserWarning: Theta grid too close to X-point, shifting by half-step
  warnings.warn("Theta grid too close to X-point, shifting by half-step")

Step: 49/49
--- t = 2.45e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.85e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.49e-03 (vs. requested 3.00e-03).
   Iterations taken: 11/50.

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(3, 3, figsize=(15, 10), 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('Magnetic axis, $R$ [$m$]')
axs_flat[0].legend()
axs_flat[0].grid()

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(Magnetic axis, $Z$) [$m$]')
axs_flat[1].grid()

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('Magnetic axis, $\psi$')
axs_flat[2].grid()

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[3].grid()

axs_flat[4].plot(history_times, history_elongation,'k+')
axs_flat[4].plot(history_times_nl, history_elongation_nl,'rx')
axs_flat[4].set_xlabel('Time')
axs_flat[4].set_ylabel('Geometric elongation')
axs_flat[4].grid()

axs_flat[5].plot(history_times, history_triangularity,'k+')
axs_flat[5].plot(history_times_nl, history_triangularity_nl,'rx')
axs_flat[5].set_xlabel('Time')
axs_flat[5].set_ylabel('Triangularity')
axs_flat[5].grid()

axs_flat[6].plot(history_times, history_squareness,'k+')
axs_flat[6].plot(history_times_nl, history_squareness_nl,'rx')
axs_flat[6].set_xlabel('Time')
axs_flat[6].set_ylabel('Squarenss')
axs_flat[6].grid()

axs_flat[7].plot(history_times, history_area,'k+')
axs_flat[7].plot(history_times_nl, history_area_nl,'rx')
axs_flat[7].set_xlabel('Time')
axs_flat[7].set_ylabel('LCFS area')
axs_flat[7].grid()

axs_flat[8].plot(history_times, history_length,'k+')
axs_flat[8].plot(history_times_nl, history_length_nl,'rx')
axs_flat[8].set_xlabel('Time')
axs_flat[8].set_ylabel('LCFS circumference')
axs_flat[8].grid()

# 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

In this example, we examine a limiter configuration plasma and show how to use time-varying plasma profile parameters and plasma resisitivity. Note there is also the option to specify time-dependent active coil resistances. For example, if you want to switch a coil “off” during a simulation you can set its resistance to an arbitrarily high number. See the optional parameter custom_active_coil_resistances.

from freegsnke.jtor_update import ConstrainPaxisIp
from freegsnke import equilibrium_update, GSstaticsolver

# 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
)  

# load coil currents
with open('data/simple_limited_currents_PaxisIp.pk', 'rb') as f:
    current_values = pickle.load(f)

for key in current_values.keys():
    eq.tokamak.set_coil_current(key, current_values[key])


# define profiles
profiles = ConstrainPaxisIp(
    eq=eq,
    paxis=6e3,
    Ip=4e5,
    fvac=0.5,
    alpha_m=1.8,
    alpha_n=1.2
)

# solve
GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    
GSStaticSolver.solve(
    eq, 
    profiles, 
    target_relative_tolerance=1e-8, 
    verbose=False
    )

# visualise:
# 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. 
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 3.83e-09 (vs. requested 1.00e-08) reached in 27/100 iterations.

from freegsnke import nonlinear_solve

# initialise the nonlinear solver object
stepping = nonlinear_solve.nl_solver(
    eq=eq, 
    profiles=profiles, 
    full_timestep=.5e-3, 
    plasma_resistivity=1e-6, 
    GSStaticSolver=GSStaticSolver   
)

-----
Checking that the provided 'eq' and 'profiles' are a GS solution...
Forward static solve SUCCESS. Tolerance 3.83e-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' options 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'
      30 recovered that couple with the plasma more than 'threshold_dIy_dI'
      0 removed that couple with the plasma less than 'min_dIy_dI'
      total selected = 34 (out of 138)
   Total number of modes = 46 (12 active coils + 34 passive structures)
      (Note: some additional modes may be removed after Jacobian calculation)
-----
Building the 3051 x 47 Jacobian (dIy/dI) of plasma current density (inside the LCFS) with respect to all metal currents and the total plasma current.

Building the 3051 x 3 Jacobian (dIy/dtheta) of plasma current density (inside the LCFS) with respect to all plasma current density profile parameters within Jtor.

-----
Further mode reduction:
   1 previously included modes couple with the plasma less than 'min_dIy_dI' (following Jacobian calculation)
   Final number of modes = 45 (12 active coils + 33 passive structures)
   Re-sizing the Jacobian matrix to account for any removed modes (if required).
-----
Stability paramters:

   Deformable plasma metrics:
      Growth rate = [19.98380817+0.j] [1/s]
      Instability timescale = [0.05004051-0.j] [s]
      Inductive stability margin = [0.61319796+0.j]
   Rigid plasma metrics:
      Leuer parameter (ratio of stabilsing to de-stabilising force gradients):
          between all metals and all metals = 2.106272155801492
          between all metals and active metals = 2.106272155801492
          between passive metals and active metals = 2.106254133747708
-----
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!
-----

# number of time steps to simulate
max_count = 50

# determine the (constant) active voltages to be applied at each time step
voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] 

# we define some time-dependent plasma current density profile parameters here
alpha_m = np.tile(profiles.alpha_m, max_count+1)
alpha_m -= (0.1 * np.sin(0.05 * np.pi * np.arange(max_count+1))) # we add some perturbation

alpha_n = np.tile(profiles.alpha_n, max_count+1)
alpha_n += (0.1 * np.sin(0.1 * np.pi * np.arange(max_count+1))) # we add some perturbation

paxis = np.tile(profiles.paxis, max_count+1)
paxis += (0.1 * np.sin(0.01 * np.pi * np.arange(max_count+1))) # we add some perturbation


# plot of profiles
fig, axs = plt.subplots(3, 1, figsize=(8, 6), dpi=80, constrained_layout=True)
axs_flat = axs.flat

axs_flat[0].plot(alpha_m, label='alpha_m')
axs_flat[0].set_xticklabels([])
axs_flat[0].set_ylabel('Parameter value')
axs_flat[0].legend()
axs_flat[0].grid()

axs_flat[1].plot(alpha_n, label='alpha_n')
axs_flat[1].set_xticklabels([])
axs_flat[1].set_ylabel('Parameter value')
axs_flat[1].legend()
axs_flat[1].grid()

axs_flat[2].plot(paxis, label='paxis')
axs_flat[2].set_xlabel('Time Step')
axs_flat[2].set_ylabel('Parameter value')
axs_flat[2].legend()
axs_flat[2].grid()

# we also define a time-dependent plasma resisitivity. 
resisitivity = np.tile(stepping.plasma_resistivity, max_count+1)
resisitivity += (1e-7 * np.sin(0.05 * np.pi * np.arange(max_count+1))) # we add some perturbation

# plot of resisitivity evolution
fig, axs = plt.subplots(1, 1, figsize=(8, 3), dpi=80)
axs.plot(resisitivity)
axs.set_xlabel('Time Step')
axs.set_ylabel('Resisitivity')
axs.grid()

# initialising some variables for iteration and logging
counter = 0
t = 0

# reset the solver object by resetting the initial conditions
stepping.initialize_from_ICs(eq, profiles)

history_times = [t]
history_currents = [stepping.currents_vec]
history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]
history_o_points = [stepping.eq1.opt[0]]
history_elongation = [stepping.eq1.geometricElongation()]
history_triangularity = [stepping.eq1.triangularity()]
history_squareness = [stepping.eq1.squareness()[1]]
history_area = [stepping.eq1.separatrix_area()]
history_length = [stepping.eq1.separatrix_length()]

# simulate using the linear solver
while counter<max_count:

    print(f'Step: {counter}/{max_count-1}')
    print(f'--- t = {t:.2e}')
    
    # carry out the time step
    stepping.nlstepper(
        active_voltage_vec = voltages, # same voltages applied at each step
        linear_only = True,
        profiles_parameters = { # different at each step
            "alpha_m": alpha_m[counter],
            "alpha_n": alpha_n[counter],
            "paxis": paxis[counter],
            },
        plasma_resistivity = resisitivity[counter] # different at each step
        )

    # 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(stepping.eq1.create_auxiliary_equilibrium())
    history_o_points.append(stepping.eq1.opt[0])
    history_elongation.append(stepping.eq1.geometricElongation())
    history_triangularity.append(stepping.eq1.triangularity())
    history_squareness.append(stepping.eq1.squareness()[1])
    history_area.append(stepping.eq1.separatrix_area())
    history_length.append(stepping.eq1.separatrix_length())

# transform lists to arrays
history_currents = np.array(history_currents)
history_times = np.array(history_times)
history_o_points = np.array(history_o_points)
history_elongation = np.array(history_elongation)
history_triangularity = np.array(history_triangularity)
history_squareness = np.array(history_squareness)
history_area = np.array(history_area)
history_length = np.array(history_length)

Step: 0/49
--- t = 0.00e+00

Step: 1/49
--- t = 5.00e-04

Step: 2/49
--- t = 1.00e-03

/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/freegs4e/critical.py:1045: UserWarning: Theta grid too close to X-point, shifting by half-step
  warnings.warn("Theta grid too close to X-point, shifting by half-step")

Step: 3/49
--- t = 1.50e-03

Step: 4/49
--- t = 2.00e-03

Step: 5/49
--- t = 2.50e-03

Step: 6/49
--- t = 3.00e-03

Step: 7/49
--- t = 3.50e-03

Step: 8/49
--- t = 4.00e-03

Step: 9/49
--- t = 4.50e-03

Step: 10/49
--- t = 5.00e-03

Step: 11/49
--- t = 5.50e-03

Step: 12/49
--- t = 6.00e-03

Step: 13/49
--- t = 6.50e-03

Step: 14/49
--- t = 7.00e-03

Step: 15/49
--- t = 7.50e-03

Step: 16/49
--- t = 8.00e-03

Step: 17/49
--- t = 8.50e-03

Step: 18/49
--- t = 9.00e-03

Step: 19/49
--- t = 9.50e-03

Step: 20/49
--- t = 1.00e-02

Step: 21/49
--- t = 1.05e-02

Step: 22/49
--- t = 1.10e-02

Step: 23/49
--- t = 1.15e-02

Step: 24/49
--- t = 1.20e-02

Step: 25/49
--- t = 1.25e-02

Step: 26/49
--- t = 1.30e-02

Step: 27/49
--- t = 1.35e-02

Step: 28/49
--- t = 1.40e-02

Step: 29/49
--- t = 1.45e-02

Step: 30/49
--- t = 1.50e-02

Step: 31/49
--- t = 1.55e-02

Step: 32/49
--- t = 1.60e-02

Step: 33/49
--- t = 1.65e-02

Step: 34/49
--- t = 1.70e-02

Step: 35/49
--- t = 1.75e-02

Step: 36/49
--- t = 1.80e-02

Step: 37/49
--- t = 1.85e-02

Step: 38/49
--- t = 1.90e-02

Step: 39/49
--- t = 1.95e-02

Step: 40/49
--- t = 2.00e-02

Step: 41/49
--- t = 2.05e-02

Step: 42/49
--- t = 2.10e-02

Step: 43/49
--- t = 2.15e-02

Step: 44/49
--- t = 2.20e-02

Step: 45/49
--- t = 2.25e-02

Step: 46/49
--- t = 2.30e-02

Step: 47/49
--- t = 2.35e-02

Step: 48/49
--- t = 2.40e-02

Step: 49/49
--- t = 2.45e-02

# recalculate the active voltages using the new currents
U_active = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] 

# number of time steps to simulate
max_count = 50

# 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

# reset the solver object by resetting the initial conditions
stepping.initialize_from_ICs(eq, profiles)

history_times_nl = [t]
history_currents_nl = [stepping.currents_vec]
history_equilibria_nl = [stepping.eq1.create_auxiliary_equilibrium()]
history_o_points_nl = [stepping.eq1.opt[0]]
history_elongation_nl = [stepping.eq1.geometricElongation()]
history_triangularity_nl = [stepping.eq1.triangularity()]
history_squareness_nl = [stepping.eq1.squareness()[1]]
history_area_nl = [stepping.eq1.separatrix_area()]
history_length_nl = [stepping.eq1.separatrix_length()]

# simulate using the nonlinear solver
while counter<max_count:

    print(f'Step: {counter}/{max_count-1}')
    print(f'--- t = {t:.2e}')
    
    # carry out the time step
    stepping.nlstepper(
        active_voltage_vec = voltages, # same voltages applied at each step
        linear_only = False,
        profiles_parameters = { # different profile parameters at each step
            "alpha_m": alpha_m[counter],
            "alpha_n": alpha_n[counter],
            "paxis": paxis[counter],
            },
        plasma_resistivity = resisitivity[counter] # different resistivity at each step
        )

    # 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(stepping.eq1.create_auxiliary_equilibrium())
    history_o_points_nl.append(stepping.eq1.opt[0])
    history_elongation_nl.append(stepping.eq1.geometricElongation())
    history_triangularity_nl.append(stepping.eq1.triangularity())
    history_squareness_nl.append(stepping.eq1.squareness()[1])
    history_area_nl.append(stepping.eq1.separatrix_area())
    history_length_nl.append(stepping.eq1.separatrix_length())

# 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)
history_elongation_nl = np.array(history_elongation_nl)
history_triangularity_nl = np.array(history_triangularity_nl)
history_squareness_nl = np.array(history_squareness_nl)
history_area_nl = np.array(history_area_nl)
history_length_nl = np.array(history_length_nl)

Step: 0/49
--- t = 0.00e+00

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.37e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.91e-03 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 1/49
--- t = 5.00e-04

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.77e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.03e-03 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 2/49
--- t = 1.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.89e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.00e-03 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 3/49
--- t = 1.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 5.90e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.33e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 4/49
--- t = 2.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.46e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.21e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 5/49
--- t = 2.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.73e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.05e-03 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 6/49
--- t = 3.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.32e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.81e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 7/49
--- t = 3.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.13e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.25e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 8/49
--- t = 4.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.62e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.11e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 9/49
--- t = 4.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.97e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.29e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 10/49
--- t = 5.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.98e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.55e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 11/49
--- t = 5.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 6.43e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.35e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 12/49
--- t = 6.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.16e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.88e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 13/49
--- t = 6.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 8.11e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.68e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 14/49
--- t = 7.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.35e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.86e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 15/49
--- t = 7.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.50e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.54e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 16/49
--- t = 8.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.53e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.42e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 17/49
--- t = 8.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.74e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.46e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 18/49
--- t = 9.00e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 6.06e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 7.47e-04 (vs. requested 3.00e-03).
   Iterations taken: 3/50.

Step: 19/49
--- t = 9.50e-03

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 6.98e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.83e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 20/49
--- t = 1.00e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.64e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.32e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 21/49
--- t = 1.05e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.87e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.93e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 22/49
--- t = 1.10e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.60e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.64e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 23/49
--- t = 1.15e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.78e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 6.91e-04 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 24/49
--- t = 1.20e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 5.43e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.54e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 25/49
--- t = 1.25e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.73e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.03e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 26/49
--- t = 1.30e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.62e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.84e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 27/49
--- t = 1.35e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.09e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.10e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 28/49
--- t = 1.40e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.44e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.38e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 29/49
--- t = 1.45e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.03e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.12e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 30/49
--- t = 1.50e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.81e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.60e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 31/49
--- t = 1.55e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.28e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.82e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 32/49
--- t = 1.60e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 3.71e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.24e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 33/49
--- t = 1.65e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.02e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.98e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 34/49
--- t = 1.70e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.61e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.77e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 35/49
--- t = 1.75e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 5.72e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.39e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 36/49
--- t = 1.80e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 9.61e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.75e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 37/49
--- t = 1.85e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.48e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.32e-03 (vs. requested 3.00e-03).
   Iterations taken: 6/50.

Step: 38/49
--- t = 1.90e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.61e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.14e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 39/49
--- t = 1.95e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.40e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.22e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 40/49
--- t = 2.00e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.72e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.40e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 41/49
--- t = 2.05e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.79e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.91e-03 (vs. requested 3.00e-03).
   Iterations taken: 5/50.

Step: 42/49
--- t = 2.10e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.55e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.27e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 43/49
--- t = 2.15e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 7.34e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.22e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 44/49
--- t = 2.20e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.17e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.72e-03 (vs. requested 3.00e-03).
   Iterations taken: 8/50.

Step: 45/49
--- t = 2.25e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.94e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 1.19e-03 (vs. requested 3.00e-03).
   Iterations taken: 7/50.

Step: 46/49
--- t = 2.30e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 4.22e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.10e-03 (vs. requested 3.00e-03).
   Iterations taken: 9/50.

Step: 47/49
--- t = 2.35e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 2.09e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.20e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 48/49
--- t = 2.40e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 1.62e-03 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.38e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

Step: 49/49
--- t = 2.45e-02

Forward evolutive solve SUCCESS.
   Last max. relative currents change: 5.73e-04 (vs. requested 5.00e-03).
   Last max. relative flux change: 2.13e-03 (vs. requested 3.00e-03).
   Iterations taken: 10/50.

# Plot evolution of tracked values and compare between linear and non-linear evolution
fig, axs = plt.subplots(3, 3, figsize=(15, 10), 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('Magnetic axis, $R$ [$m$]')
axs_flat[0].legend()
axs_flat[0].grid()

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(Magnetic axis, $Z$) [$m$]')
axs_flat[1].grid()

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('Magnetic axis, $\psi$')
axs_flat[2].grid()

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[3].grid()

axs_flat[4].plot(history_times, history_elongation,'k+')
axs_flat[4].plot(history_times_nl, history_elongation_nl,'rx')
axs_flat[4].set_xlabel('Time')
axs_flat[4].set_ylabel('Geometric elongation')
axs_flat[4].grid()

axs_flat[5].plot(history_times, history_triangularity,'k+')
axs_flat[5].plot(history_times_nl, history_triangularity_nl,'rx')
axs_flat[5].set_xlabel('Time')
axs_flat[5].set_ylabel('Triangularity')
axs_flat[5].grid()

axs_flat[6].plot(history_times, history_squareness,'k+')
axs_flat[6].plot(history_times_nl, history_squareness_nl,'rx')
axs_flat[6].set_xlabel('Time')
axs_flat[6].set_ylabel('Squarenss')
axs_flat[6].grid()

axs_flat[7].plot(history_times, history_area,'k+')
axs_flat[7].plot(history_times_nl, history_area_nl,'rx')
axs_flat[7].set_xlabel('Time')
axs_flat[7].set_ylabel('LCFS area')
axs_flat[7].grid()

axs_flat[8].plot(history_times, history_length,'k+')
axs_flat[8].plot(history_times_nl, history_length_nl,'rx')
axs_flat[8].set_xlabel('Time')
axs_flat[8].set_ylabel('LCFS circumference')
axs_flat[8].grid()

# 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()