In this example we show how to compute equilibria in ITER with L-mode like profiles for:

The "inverse" case where we have a desired shape, plasma current and pressure, but the required coil currents are unkown
The "forward" case where we have already have coil currents, plasma current, and known position for the desired equilibrium

This example utilizes the mesh built in TokaMaker Meshing Example: Building a mesh for ITER.

Note: Running this example requires the h5py python package, which is installable using pip or other standard methods.

import os
import sys
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize']=(6,6)
plt.rcParams['font.weight']='bold'
plt.rcParams['axes.labelweight']='bold'
plt.rcParams['lines.linewidth']=2
plt.rcParams['lines.markeredgewidth']=2
%matplotlib inline
%config InlineBackend.figure_format = "retina"

Load TokaMaker library

To load the TokaMaker python module we need to tell python where to the module is located. This can be done either through the PYTHONPATH environment variable or using within a script using sys.path.append() as below, where we look for the environement variable OFT_ROOTPATH to provide the path to where the OpenFUSIONToolkit is installed (/Applications/OFT on macOS).

For meshing we will use the gs_Domain class to build a 2D triangular grid suitable for Grad-Shafranov equilibria. This class uses the triangle code through a python wrapper.

tokamaker_python_path = os.getenv('OFT_ROOTPATH')
if tokamaker_python_path is not None:
    sys.path.append(os.path.join(tokamaker_python_path,'python'))
from OpenFUSIONToolkit import OFT_env
from OpenFUSIONToolkit.TokaMaker import TokaMaker
from OpenFUSIONToolkit.TokaMaker.meshing import load_gs_mesh
from OpenFUSIONToolkit.TokaMaker.util import create_power_flux_fun

Compute equilibria

Initialize TokaMaker object

We now create a OFT_env instance for execution using two threads and a TokaMaker instance to use for equilibrium calculations. Note at present only a single TokaMaker instance can be used per python kernel, so this command should only be called once in a given Jupyter notebook or python script. In the future this restriction may be relaxed.

myOFT = OFT_env(nthreads=2)

mygs = TokaMaker(myOFT)

#----------------------------------------------
Open FUSION Toolkit Initialized
Development branch:   v1_beta6
Revision id:          681e857
Parallelization Info:
  # of MPI tasks      =    1
  # of NUMA nodes     =    1
  # of OpenMP threads =    2
Fortran input file    = /var/folders/52/n5qxh27n4w19qxzqygz2btbw0000gn/T/oft_64673/oftpyin
XML input file        = none
Integer Precisions    =    4   8
Float Precisions      =    4   8  16
Complex Precisions    =    4   8
LA backend            = native
#----------------------------------------------

Load mesh into TokaMaker

Now we load the mesh generated in TokaMaker Meshing Example: Building a mesh for ITER using load_gs_mesh() and setup_mesh. Then we use setup_regions(), passing the conductor and coil dictionaries for the mesh, to define the different region types. Finally, we call setup() to setup the required solver objects. During this call we can specify the desired element order (min=2, max=4) and the toroidal field through F0 = B0*R0, where B0 is the toroidal field at a reference location R0.

mesh_pts,mesh_lc,mesh_reg,coil_dict,cond_dict = load_gs_mesh('ITER_mesh.h5')
mygs.setup_mesh(mesh_pts, mesh_lc, mesh_reg)
mygs.setup_regions(cond_dict=cond_dict,coil_dict=coil_dict)
mygs.setup(order = 2, F0 = 5.3*6.2)

**** Loading OFT surface mesh

**** Generating surface grid level  1
  Generating boundary domain linkage
  Mesh statistics:
    Area         =  2.859E+02
    # of points  =    4757
    # of edges   =   14156
    # of cells   =    9400
    # of boundary points =     112
    # of boundary edges  =     112
    # of boundary cells  =     112
  Resolution statistics:
    hmin =  9.924E-03
    hrms =  2.826E-01
    hmax =  8.466E-01
  Surface grounded at vertex     870


**** Creating Lagrange FE space
  Order  =    2
  Minlev =   -1

 Computing flux BC matrix 
 Inverting real matrix
   Time =    1.0910000000000000E-003

Define a vertical stability coil

Like many elongated equilibria, the equilibrium we seek to compute below is vertically unstable. In this case we use the actual ITER Vertical Stability Coil (VSC) in order to help with convergence using the set_coil_vsc() method.

Note: While ITER has a "real" VSC, this is not required and this functionality can instead be used to define a "virtual" VSC by pairing coils in a way that are not necessarily paired experimentally.

mygs.set_coil_vsc({'VS': 1.0})

Define hard limits on coil currents

Hard limits on coil currents can be set using set_coil_bounds(). In this case we just the simple and approximate bi-directional limit of 50 MA in each coil.

Bounds are specified using a dictionary of 2 element lists, containing the minimum and maximum bound, where the dictionary key corresponds to the coil names, which are available in mygs.coil_sets

coil_bounds = {key: [-50.E6, 50.E6] for key in mygs.coil_sets}

mygs.set_coil_bounds(coil_bounds)

Compute Inverse Equilibrium

Define global quantities and targets

For the inverse case we define a target for the plasma current and the peak plasma pressure, which occurs on the magnetic axis.

Note: These constraints can be considered "hard" constraints, where they will be matched to good tolerance as long as the calculation converges.

Ip_target=15.6E6
P0_target=6.2E5
mygs.set_targets(Ip=Ip_target, pax=P0_target)

Define shape targets

In order to constrain the shape of the plasma we can utilize two types of constraints:

isoflux points, which are points we want to lie on the same flux surface (eg. the LCFS)
saddle points, where we want the poloidal magnetic field to vanish (eg. X-points). While one can also use this constraint to enforce a magnetic axis location, instead set_targets() should be used with arguments R0 and V0.

Note: These constraints can be considered "soft" constraints, where the calculation attempts to minimize error in satisfying these constraints subject to other constraints and regularization.

Here we define a handful of isoflux points that we want to lie on the LCFS of the target equilibrium. Additionally, we define a single X-point and set it as a saddle constraint as well as adding it to the list of isoflux points.

isoflux_pts = np.array([
    [ 8.20,  0.41],
    [ 8.06,  1.46],
    [ 7.51,  2.62],
    [ 6.14,  3.78],
    [ 4.51,  3.02],
    [ 4.26,  1.33],
    [ 4.28,  0.08],
    [ 4.49, -1.34],
    [ 7.28, -1.89],
    [ 8.00, -0.68]
])
x_point = np.array([[5.125, -3.4],])
mygs.set_isoflux(np.vstack((isoflux_pts,x_point)))
mygs.set_saddles(x_point)

Define coil regularization matrix

In general, for a given coil set a given plasma shape cannot be exactly reproduced, which generally yields large amplitude coil currents if no constraint on the coil currents is applied. As a result, it is useful to include regularization terms for the coils to balance minimization of the shape error with the amplitude of current in the coils. In TokaMaker these regularization terms have the general form, where each term corresponds to a set of coil coefficients, target value, and weight. The coil_reg_term() method is provided to aid in defining these terms.

In this case, one regularization term is added for each coil with a single unit coefficient for that coil and target of zero with modest weights. This regularization acts to penalize the amplitude of current in each coil, acting to balance coil current with error in the shape targets. Note that the weight on the VSC virtual coil (#VSC) defined above is set high to prevent interaction with the real VS coil set (see below for further information).

# Set regularization weights
regularization_terms = []
for name, coil in mygs.coil_sets.items():
    # Set zero target current and different small weights to help conditioning of fit
    if name.startswith('CS'):
        if name.startswith('CS1'):
            regularization_terms.append(mygs.coil_reg_term({name: 1.0},target=0.0,weight=2.E-2))
        else:
            regularization_terms.append(mygs.coil_reg_term({name: 1.0},target=0.0,weight=1.E-2))
    elif name.startswith('PF'):
        regularization_terms.append(mygs.coil_reg_term({name: 1.0},target=0.0,weight=1.E-2))
    elif name.startswith('VS'):
        regularization_terms.append(mygs.coil_reg_term({name: 1.0},target=0.0,weight=1.E-2))
# Disable VSC virtual coil
regularization_terms.append(mygs.coil_reg_term({'#VSC': 1.0},target=0.0,weight=1.E2))
 
# Pass regularization terms to TokaMaker
mygs.set_coil_reg(reg_terms=regularization_terms)

Define flux functions

Although TokaMaker has a "default" profile for the F*F' and P' terms this should almost never be used and one should instead choose an appropriate flux function for their application. In this case we use an L-mode-like profile of the form \(((1-\hat{\psi})^{\alpha})^{\gamma}\), using create_power_flux_fun(), where \(\alpha\) and \(\gamma\) are set differently for F*F' and P' to provide peaked and broad profiles respectively. Within TokaMaker this profile is represented as a piecewise linear function, which can be set up using the dictionary approach shown below.

# Set profiles
ffp_prof = create_power_flux_fun(40,1.5,2.0)
pp_prof = create_power_flux_fun(40,4.0,1.0)
 
fig, ax = plt.subplots(2,1,sharex=True)
# Plot F*F'
ax[0].plot(ffp_prof['x'],ffp_prof['y'])
ax[0].set_ylabel("FF'")
# Plot P'
ax[1].plot(pp_prof['x'],pp_prof['y'])
ax[1].set_ylabel("P'")
_ = ax[-1].set_xlabel(r"$\hat{\psi}$")
 
mygs.set_profiles(ffp_prof=ffp_prof,pp_prof=pp_prof)

Compute equilibrium

We can now compute a free-boundary equilibrium using these constraints. Note that before running a calculation for the first time we must initialize the flux function \(\psi\), which can be done using init_psi(). This subroutine initializes the flux using the specified Ip_target from above, which is evenly distributed over the entire plasma region or only with a boundary defined using a center point (R,Z), minor radius (a), and elongation and triangularity. Coil currents are also initialized at this point using the constraints above and this uniform plasma current initialization.

solve() is then called to compute a self-consitent Grad-Shafranov equilibrium. If the result variable (err_flag) is zero then the solution has converged to the desired tolerance ( \(10^{-6}\) by default).

R0 = 6.3
Z0 = 0.5
a = 2.0
kappa = 1.4
delta = 0.0
mygs.init_psi(R0, Z0, a, kappa, delta)
mygs.solve()

Starting non-linear GS solver
     1  5.7965E+00  1.5693E-01  6.4686E-01  6.4422E+00  5.3430E-01  9.8316E-04
     2  1.5480E+01  9.5055E-02  2.7469E-01  6.4072E+00  5.3290E-01  1.5938E-03
     3  1.9148E+01  7.9379E-02  1.2986E-01  6.3868E+00  5.3264E-01  1.6630E-03
     4  2.0951E+01  7.3407E-02  6.7604E-02  6.3757E+00  5.3235E-01  1.6305E-03
     5  2.1908E+01  7.0680E-02  3.6983E-02  6.3697E+00  5.3227E-01  1.5882E-03
     6  2.2433E+01  6.9312E-02  2.0669E-02  6.3664E+00  5.3231E-01  1.5557E-03
     7  2.2724E+01  6.8590E-02  1.1653E-02  6.3646E+00  5.3240E-01  1.5342E-03
     8  2.2888E+01  6.8198E-02  6.5944E-03  6.3636E+00  5.3250E-01  1.5207E-03
     9  2.2980E+01  6.7980E-02  3.7379E-03  6.3631E+00  5.3258E-01  1.5126E-03
    10  2.3032E+01  6.7859E-02  2.1210E-03  6.3628E+00  5.3265E-01  1.5079E-03
    11  2.3061E+01  6.7791E-02  1.2045E-03  6.3626E+00  5.3270E-01  1.5051E-03
    12  2.3078E+01  6.7753E-02  6.8471E-04  6.3625E+00  5.3274E-01  1.5035E-03
    13  2.3087E+01  6.7731E-02  3.8962E-04  6.3625E+00  5.3277E-01  1.5025E-03
    14  2.3093E+01  6.7719E-02  2.2195E-04  6.3624E+00  5.3279E-01  1.5020E-03
    15  2.3096E+01  6.7712E-02  1.2660E-04  6.3624E+00  5.3280E-01  1.5017E-03
    16  2.3097E+01  6.7708E-02  7.2329E-05  6.3624E+00  5.3281E-01  1.5015E-03
    17  2.3098E+01  6.7705E-02  4.1395E-05  6.3624E+00  5.3281E-01  1.5014E-03
    18  2.3099E+01  6.7704E-02  2.3741E-05  6.3624E+00  5.3282E-01  1.5013E-03
    19  2.3099E+01  6.7703E-02  1.3650E-05  6.3624E+00  5.3282E-01  1.5013E-03
    20  2.3099E+01  6.7703E-02  7.8710E-06  6.3624E+00  5.3282E-01  1.5013E-03
    21  2.3100E+01  6.7703E-02  4.5542E-06  6.3624E+00  5.3282E-01  1.5013E-03
    22  2.3100E+01  6.7702E-02  2.6456E-06  6.3624E+00  5.3283E-01  1.5013E-03
    23  2.3100E+01  6.7702E-02  1.5440E-06  6.3624E+00  5.3283E-01  1.5013E-03
    24  2.3100E+01  6.7702E-02  9.0578E-07  6.3624E+00  5.3283E-01  1.5013E-03
 Timing:  0.20443300000624731     
   Source:     9.1538999869953841E-002
   Solve:      6.0522999963723123E-002
   Boundary:   5.4930002079345286E-003
   Other:      4.6877999964635819E-002

Plot equilibrium

Flux surfaces of the computed equilibrium can be plotted using the plot_psi() method. The additional plotting methods plot_machine() and plot_constraints() are also used to show context and other information. Each method has a large number of optional arguments for formatting and other options.

fig, ax = plt.subplots(1,1)
mygs.plot_machine(fig,ax,coil_colormap='seismic',coil_symmap=True,coil_scale=1.E-6,coil_clabel=r'$I_C$ [MA]')
mygs.plot_psi(fig,ax,xpoint_color=None,vacuum_nlevels=4)
mygs.plot_constraints(fig,ax,isoflux_color='tab:red',isoflux_marker='o')

Print equilibrium information and coil currents

Basic parameters can be displayed using the print_info() method. For access to these quantities as variables instead the get_stats() can be used.

The final coil currents can also be retrieved using the get_coil_currents() method, which are all within the approximate coil limits imposed above.

mygs.print_info()
 
print()
print("Coil Currents [MA]:")
coil_currents, _ = mygs.get_coil_currents()
for key, current in coil_currents.items():
    print('  {0:10} {1:10.2F}'.format(key+":",current/1.E6))

Equilibrium Statistics:
  Topology                =   Diverted
  Toroidal Current [A]    =    1.5600E+07
  Current Centroid [m]    =    6.203  0.530
  Magnetic Axis [m]       =    6.362  0.533
  Elongation              =    1.875 (U:  1.763, L:  1.987)
  Triangularity           =    0.477 (U:  0.405, L:  0.550)
  Plasma Volume [m^3]     =   820.079
  q_0, q_95               =    0.823  2.760
  Plasma Pressure [Pa]    =   Axis:  6.1923E+05, Peak:  6.1923E+05
  Stored Energy [J]       =    2.4299E+08
  <Beta_pol> [%]          =   42.6829
  <Beta_tor> [%]          =    1.7801
  <Beta_n>   [%]          =    1.1953
  Diamagnetic flux [Wb]   =    1.5403E+00
  Toroidal flux [Wb]      =    1.2187E+02
  l_i                     =    1.1597

Coil Currents [MA]:
  CS3U:           13.74
  CS2U:            9.74
  CS1U:           -9.06
  CS1L:           -8.81
  CS2L:           13.40
  CS3L:           17.88
  PF1:            12.94
  PF2:            -2.20
  PF3:            -5.95
  PF4:            -5.14
  PF5:            -5.18
  PF6:            19.93
  VS:              0.15

Compute Forward Equilibrium

We will now demonstrate how to compute a forward equilibrium with known coil currents. Note that this method generally requires a known equilibrium, either from experimental observations or another code, as the Grad-Shafranov equation represents force-balance and as such solutions only exist for appropriate combinations of the plasma current and pressure and the supporting coil currents.

Remove shape constraints from above

As we do not want to adjust the coil currents in the forward case, we must remove the shape targets from above, which can be done by passing None to both methods.

Note: by removing the shape targets the coil currents will not be adjusted and as a consequence the coil regularization defined above is no longer used.

mygs.set_isoflux(None)

mygs.set_saddles(None)

Set coil currents

We now set coil currents in TokaMaker to those from a known equilibrium solution.

eq_currents = {
    'CS3U': 5180.432355*553,
    'CS2U': -16660.61401*553,
    'CS1U': -36367.32465*553,
    'CS1L': -36367.32465*553,
    'CS2L': -16472.05814*553,
    'CS3L': 8671.71887*553,
    'PF1': 21725.94264*248.6,
    'PF2': -22213.70001*115.2,
    'PF3': -31877.98793*185.9,
    'PF4': -26721.1652*169.9,
    'PF5': -38450.07502*216.8,
    'PF6': 39603.47595*459.4,
    'VS': 0.0
}
mygs.set_coil_currents(eq_currents)

Set global targets

Again we define a target for the plasma current, but instead of a target pressure we constrain the radial and vertical position of the magnetic axis to match the equilibrium solution corresponding to the specified coil currents. During the solve the radial position is used to control the plasma pressure and VSC virtual coil is used to match the vertical position.

R0_target = 6.37
Z0_target = 0.51
mygs.set_targets(Ip=Ip_target, R0=R0_target, V0=Z0_target)

Compute equilibrium

mygs.solve()

Starting non-linear GS solver
     1  2.5039E+01  5.5296E-02  7.2031E-01  6.3624E+00  5.3283E-01 -4.0808E+04
     2  2.4452E+01  5.1835E-02  1.4442E-01  6.3634E+00  5.2979E-01 -7.3974E+04
     3  2.4210E+01  5.1078E-02  2.8927E-02  6.3646E+00  5.2614E-01 -5.9414E+04
     4  2.4083E+01  5.1086E-02  6.0418E-03  6.3659E+00  5.2236E-01 -5.2362E+04
     5  2.4011E+01  5.1312E-02  1.8537E-03  6.3671E+00  5.1856E-01 -4.7695E+04
     6  2.3967E+01  5.1591E-02  1.3838E-03  6.3684E+00  5.1476E-01 -4.3368E+04
     7  2.3939E+01  5.1874E-02  1.3593E-03  6.3697E+00  5.1095E-01 -3.9046E+04
     8  2.3974E+01  5.1787E-02  5.6303E-03  6.3699E+00  5.1019E-01 -4.6924E+02
     9  2.3965E+01  5.1722E-02  3.2792E-04  6.3700E+00  5.1004E-01 -1.0098E+04
    10  2.3963E+01  5.1701E-02  4.4648E-04  6.3700E+00  5.1001E-01 -1.2884E+04
    11  2.3962E+01  5.1695E-02  1.4514E-04  6.3700E+00  5.1000E-01 -1.3277E+04
    12  2.3963E+01  5.1693E-02  3.0525E-05  6.3700E+00  5.1000E-01 -1.3283E+04
    13  2.3963E+01  5.1693E-02  4.8811E-06  6.3700E+00  5.1000E-01 -1.3265E+04
    14  2.3963E+01  5.1692E-02  1.9505E-06  6.3700E+00  5.1000E-01 -1.3256E+04
    15  2.3963E+01  5.1692E-02  1.1889E-06  6.3700E+00  5.1000E-01 -1.3253E+04
    16  2.3963E+01  5.1692E-02  6.6367E-07  6.3700E+00  5.1000E-01 -1.3252E+04
 Timing:  0.12454300001263618     
   Source:     6.3579999899957329E-002
   Solve:      4.2287000105716288E-002
   Boundary:   3.8619999540969729E-003
   Other:      1.4814000052865595E-002

Print information and plot equilibrium

mygs.print_info()
#
fig, ax = plt.subplots(1,1)
mygs.plot_machine(fig,ax,coil_colormap='seismic',coil_symmap=True,coil_scale=1.E-6,coil_clabel=r'$I_C$ [MA]')
mygs.plot_psi(fig,ax,xpoint_color=None,vacuum_nlevels=4)

Equilibrium Statistics:
  Topology                =   Diverted
  Toroidal Current [A]    =    1.5600E+07
  Current Centroid [m]    =    6.201  0.507
  Magnetic Axis [m]       =    6.370  0.510
  Elongation              =    1.755 (U:  1.642, L:  1.868)
  Triangularity           =    0.452 (U:  0.376, L:  0.528)
  Plasma Volume [m^3]     =   862.814
  q_0, q_95               =    0.828  2.780
  Plasma Pressure [Pa]    =   Axis:  4.9294E+05, Peak:  4.9294E+05
  Stored Energy [J]       =    2.0343E+08
  <Beta_pol> [%]          =   34.4872
  <Beta_tor> [%]          =    1.4168
  <Beta_n>   [%]          =    1.0036
  Diamagnetic flux [Wb]   =    1.7430E+00
  Toroidal flux [Wb]      =    1.2828E+02
  l_i                     =    1.1825