Example: Reading and writing an equilibrium to IMAS (via IMAS-Python)

In this example, we take a MAST-U-like equilibrium and show how to save it to, and load it back from, the IMAS equilibrium IDS using the IMAS-Python package. IDSs (Interface Data Structures) are the standardised data structures IMAS uses to link different modelling codes together; here we serialise the IDS to a netCDF file.

As in the previous example, we run our standard equilibrium solve - making it lower single null by pushing the P6 current.

import matplotlib.pyplot as plt
import numpy as np

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

# initialise equilibrium object
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)
)

# initialise profile object
from freegsnke.jtor_update import ConstrainPaxisIp
profiles = ConstrainPaxisIp(
    eq=eq,
    paxis=8.1e3,
    Ip=6.2e5,
    fvac=0.5,
    alpha_m=1.8,
    alpha_n=1.2
)

# initialise solver
from freegsnke import GSstaticsolver
GSStaticSolver = GSstaticsolver.NKGSsolver(eq)

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

for key in current_values.keys():
    eq.tokamak[key].current = current_values[key]

# change this to shift plasma vertically up (-) or down (+)
eq.tokamak["P6"].current += 500

# carry out forward solve
GSStaticSolver.solve(eq=eq,
                     profiles=profiles,
                     constrain=None,
                     target_relative_tolerance=1e-9)

# plot the resulting 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()
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.
Forward static solve SUCCESS. Tolerance 4.04e-10 (vs. requested 1.00e-09) reached in 26/100 iterations.
../_images/23fac55e7e0cfbff1e9655368bd941a7fcbcdb20eb5a5a75c765fde4b6bd8d0c.png

Writing to the equilibrium IDS

FreeGSNKE provides a helper module, freegsnke.imas_read_write, that fills in the relevant fields of an IMAS equilibrium IDS from a solved eq/profiles pair: high-level ids_properties/code metadata, vacuum_toroidal_field, and a single time_slice containing the plasma boundary, global quantities, 1D flux-surface profiles, and three 2D poloidal flux maps on FreeGSNKE’s rectangular (R, Z) grid - total (also carrying j_phi), plasma-only, and vacuum (coil-only) - tagged using the standard equilibrium_profiles_2d_identifier codes.

See the equilibrium IDS data dictionary for the full structure.

from freegsnke import imas_read_write

equilibrium_ids = imas_read_write.write_equilibrium_to_ids(eq, profiles)
14:44:57 INFO     Parsing data dictionary version 4.1.1 @dd_zip.py:89

Now that it has been built, we can inspect the overall structure using IMAS-Python’s own utilities.

import imas

# inspect at a high level
imas.util.inspect(equilibrium_ids, hide_empty_nodes=True)
╭────────────────────────── IDS: equilibrium (DD version 4.1.1) ──────────────────────────╮
│ Description of a 2D, axi-symmetric, tokamak equilibrium; result of an equilibrium code. │
│ ╭──────────────────────────────────── Attributes ─────────────────────────────────────╮ │
│ │ has_value = True                                                                    │ │
│ │  metadata = <IDSMetadata for 'equilibrium'>                                         │ │
│ ╰─────────────────────────────────────────────────────────────────────────────────────╯ │
│ ╭──────────────────────────────────── Child nodes ────────────────────────────────────╮ │
│ │                  code = <IDSStructure (IDS:equilibrium, code)>                      │ │
│ │        ids_properties = <IDSStructure (IDS:equilibrium, ids_properties)>            │ │
│ │                  time = <IDSNumericArray (IDS:equilibrium, time, FLT_1D)>           │ │
│ │                         numpy.ndarray([0.])                                         │ │
│ │            time_slice = <IDSStructArray (IDS:equilibrium, time_slice with 1 items)> │ │
│ │ vacuum_toroidal_field = <IDSStructure (IDS:equilibrium, vacuum_toroidal_field)>     │ │
│ ╰─────────────────────────────────────────────────────────────────────────────────────╯ │
╰─────────────────────────────────────────────────────────────────────────────────────────╯

We can inspect at lower levels of the structure too.

imas.util.inspect(equilibrium_ids.time_slice[0].global_quantities, hide_empty_nodes=True)
╭─────────────────── IDS structure: time_slice[0]/global_quantities (DD version 4.1.1) ────────────────────╮
│ 0D parameters of the equilibrium                                                                         │
│ ╭───────────────────────────────────────────── Attributes ─────────────────────────────────────────────╮ │
│ │ has_value = True                                                                                     │ │
│ │  metadata = <IDSMetadata for 'global_quantities'>                                                    │ │
│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────╯ │
│ ╭──────────────────────────────────────────── Child nodes ─────────────────────────────────────────────╮ │
│ │            ip = <IDSFloat0D (IDS:equilibrium, time_slice[0]/global_quantities/ip, FLT_0D)>           │ │
│ │                 float(619999.9999999999)                                                             │ │
│ │ magnetic_axis = <IDSStructure (IDS:equilibrium, time_slice[0]/global_quantities/magnetic_axis)>      │ │
│ │      psi_axis = <IDSFloat0D (IDS:equilibrium, time_slice[0]/global_quantities/psi_axis, FLT_0D)>     │ │
│ │                 float(0.6311984978267434)                                                            │ │
│ │  psi_boundary = <IDSFloat0D (IDS:equilibrium, time_slice[0]/global_quantities/psi_boundary, FLT_0D)> │ │
│ │                 float(0.19775416162598092)                                                           │ │
│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────╯ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯

If unsure of where exactly a quantity is located in the structure, you can search for certain strings. Note that this returns all possible locations in the data dictionary, without revealing whether that location is actually filled.

imas.util.find_paths(equilibrium_ids, "psi_boundary")
['time_slice/global_quantities/psi_boundary',
 'time_slice/global_quantities/psi_boundary_error_upper',
 'time_slice/global_quantities/psi_boundary_error_lower']

Writing to netCDF

IMAS-Python IDSs are written to disk via an imas.DBEntry. imas_read_write.save_equilibrium_ids wraps this for the netCDF backend - see here for more details.

imas_read_write.save_equilibrium_ids(equilibrium_ids, "MASTU_LSN_equilibrium.nc")

Reading from netCDF

Now we can read the equilibrium IDS back from the netCDF file.

An equilibrium IDS doesn’t store individual coil currents (only the boundary, global quantities, and profile data), so - just as with the geqdsk example - we use the inverse solver to re-build the equilibrium from what’s actually saved.

We build the constraints entirely from what the IDS gives us:

  • an isoflux constraint from the plasma boundary outline (time_slice.boundary.outline) - all these points should sit on the same flux contour, without us needing to know its absolute value;

  • null-point constraints at the magnetic axis (time_slice.global_quantities.magnetic_axis);

  • a psi_vals constraint built from the entire saved 2D psi map (time_slice.profiles_2d[0].psi) - unlike the constraints above, this one is given the actual flux values, not just which points should share a value, so it’s what lets the optimiser pin down the absolute flux level (and, with it, the individual coil currents) rather than just the plasma shape. Passing the grid’s own R/Z axes (rather than a flattened list of points) triggers Inverse_optimizer’s “full grid” fast path, which reuses the equilibrium’s own cached Greens functions instead of recomputing them pointwise.

To check that the IDS faithfully captures the equilibrium, we’ll reconstruct it purely from what’s in the file: the grid extent, the plasma current, and the dpressure_dpsi/f_df_dpsi profile data.

FreeGSNKE’s GeneralPprimeFFprime profile class accepts exactly this kind of tabulated profile data directly, rather than a parametric form like ConstrainPaxisIp above.

We reset the coil currents on a fresh, independent tokamak to zero first, to make sure we aren’t just reusing the answer.

loaded_ids = imas_read_write.load_equilibrium_ids("MASTU_LSN_equilibrium.nc")
time_slice = loaded_ids.time_slice[0]
# grid data (IDS stores psi in Wb, i.e. with the 2*pi factor included - FreeGSNKE works in Wb/rad)
R_read = np.array(time_slice.profiles_2d[0].grid.dim1)
Z_read = np.array(time_slice.profiles_2d[0].grid.dim2)
psi_read = np.array(time_slice.profiles_2d[0].psi) / (2 * np.pi)

Rmin, Rmax = R_read[0], R_read[-1]
Zmin, Zmax = Z_read[0], Z_read[-1]
nx_read, ny_read = len(R_read), len(Z_read)

# profile data
psi_n_read = np.array(time_slice.profiles_1d.psi_norm)
pprime_read = np.array(time_slice.profiles_1d.dpressure_dpsi)
ffprime_read = np.array(time_slice.profiles_1d.f_df_dpsi)
Ip_read = time_slice.global_quantities.ip
fvac_read = loaded_ids.vacuum_toroidal_field.b0[0] * loaded_ids.vacuum_toroidal_field.r0

# boundary and magnetic axis
rbdry, zbdry = np.array(time_slice.boundary.outline.r), np.array(time_slice.boundary.outline.z)
rmagx, zmagx = time_slice.global_quantities.magnetic_axis.r, time_slice.global_quantities.magnetic_axis.z

Now let’s set up and run the inverse solver, using those quantities plus the constraints described above to reconstruct the coil currents.

from freegsnke.jtor_update import GeneralPprimeFFprime
from freegsnke.inverse import Inverse_optimizer

# coil names
active_coil_names = tokamak.coils_list[: tokamak.n_active_coils]

# fresh tokamak so we can set coil currents to zero and compare to old eq object
tokamak_inv = 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",
)

# fresh equilibrium object
eq_inv = equilibrium_update.Equilibrium(
    tokamak=tokamak_inv,
    Rmin=Rmin, Rmax=Rmax,
    Zmin=Zmin, Zmax=Zmax,
    nx=nx_read,
    ny=ny_read,
)

# fresh profiles object, built from the tabulated data saved to the IDS
profiles_inv = GeneralPprimeFFprime(
    eq=eq_inv,
    Ip=Ip_read,
    fvac=fvac_read,
    psi_n=psi_n_read,
    pprime_data=pprime_read,
    ffprime_data=ffprime_read,
)

# isoflux constraint: subsample the boundary outline stored in the IDS
isoflux_set = [[rbdry[::5], zbdry[::5]]]

# null-point constraint: magnetic axis (from the IDS)
null_points = [
    [rmagx],
    [zmagx],
]

# psi_vals constraint built from the *entire* saved 2D psi map (full-grid fast path)
psi_vals = [eq_inv.R, eq_inv.Z, psi_read]

# build constraints
constrain = Inverse_optimizer(
    isoflux_set=isoflux_set,
    null_points=null_points,
    psi_vals=psi_vals,
)

# initialise solver object
GSStaticSolver_inv = GSstaticsolver.NKGSsolver(eq_inv)
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.
# set coils to zero current
for label in active_coil_names:
    eq_inv.tokamak.set_coil_current(label, 0.0)

# default plasma flux guess
eq_inv.set_plasma_psi(eq_inv.create_psi_plasma_default())

# solve
GSStaticSolver_inv.solve(
    eq=eq_inv,
    profiles=profiles_inv,
    constrain=constrain,
    target_relative_tolerance=1e-4,
    target_relative_psit_update=1e-3,
    max_solving_iterations=50,
    l2_reg=1e-9,
    verbose=True,
    Picard_handover=1e-5,
    full_jacobian_handover=[2e-2,1e-2],
)
-----
Inverse static solve starting...
Successfully computed GS residual using initial plasma_psi and tokamak_psi guesses.
Initial relative error = 1.00e+00
-----
Iteration: 0
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['-1.92e+05', '1.49e+04', '-1.76e+04', '-2.63e+04', '-1.41e+04', '-9.81e+03', '7.02e+03', '-1.01e+04', '2.83e+03', '-5.34e+03', '-3.12e+03', '-7.44e+02']
Constraint losses = 6.15e+00
Relative update of tokamak psi (in plasma core): 4.01e+07
Handing off to forward solve (with updated currents).
Relative error =  9.40e-01
-----
Iteration: 1
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['6.08e+04', '-4.59e+03', '7.55e+03', '9.42e+03', '5.24e+03', '2.10e+03', '-2.06e+03', '2.97e+03', '-5.94e+02', '3.08e+02', '-2.20e+02', '1.19e+02']
Constraint losses = 3.40e+00
Relative update of tokamak psi (in plasma core): 1.85e-01
Handing off to forward solve (with updated currents).
Relative error =  8.65e-01
-----
Iteration: 2
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['4.47e+04', '-3.36e+03', '5.56e+03', '6.93e+03', '3.86e+03', '1.55e+03', '-1.52e+03', '2.19e+03', '-4.55e+02', '2.50e+02', '-1.34e+02', '1.09e+02']
Constraint losses = 2.37e+00
Relative update of tokamak psi (in plasma core): 1.86e-01
Handing off to forward solve (with updated currents).
Relative error =  7.63e-01
-----
Iteration: 3
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['3.29e+04', '-2.48e+03', '4.09e+03', '5.10e+03', '2.84e+03', '1.15e+03', '-1.12e+03', '1.62e+03', '-3.56e+02', '2.07e+02', '-6.05e+01', '1.00e+02']
Constraint losses = 1.62e+00
Relative update of tokamak psi (in plasma core): 1.80e-01
Handing off to forward solve (with updated currents).
Relative error =  6.30e-01
-----
Iteration: 4
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['2.42e+04', '-1.86e+03', '3.01e+03', '3.74e+03', '2.08e+03', '8.48e+02', '-8.27e+02', '1.21e+03', '-2.87e+02', '1.79e+02', '1.14e+01', '9.39e+01']
Constraint losses = 1.09e+00
Relative update of tokamak psi (in plasma core): 1.78e-01
Handing off to forward solve (with updated currents).
Relative error =  4.55e-01
-----
Iteration: 5
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['1.73e+04', '-9.80e+02', '2.17e+03', '2.70e+03', '1.50e+03', '6.54e+02', '-6.01e+02', '8.81e+02', '-2.31e+02', '1.84e+02', '-1.84e+00', '1.02e+02']
Constraint losses = 7.67e-01
Relative update of tokamak psi (in plasma core): 1.79e-01
Handing off to forward solve (with updated currents).
Relative error =  2.19e-01
-----
Iteration: 6
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['1.21e+04', '-1.10e+02', '1.55e+03', '1.92e+03', '1.05e+03', '5.32e+02', '-4.32e+02', '6.38e+02', '-1.95e+02', '2.09e+02', '-4.38e+01', '1.28e+02']
Constraint losses = 6.06e-01
Relative update of tokamak psi (in plasma core): 1.77e-01
Handing off to forward solve (with updated currents).
Relative error =  3.45e-02
-----
Iteration: 7
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['3.41e+03', '7.91e+02', '5.11e+02', '4.84e+02', '1.84e+02', '5.06e+02', '-2.61e+02', '2.81e+02', '-1.32e+02', '2.29e+02', '-1.02e+02', '1.61e+02']
Constraint losses = 5.53e-01
Relative update of tokamak psi (in plasma core): 7.26e-02
Handing off to forward solve (with updated currents).
Relative error =  9.80e-03
-----
Iteration: 8
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['1.26e+03', '5.88e+02', '1.67e+02', '1.88e+02', '1.28e+02', '6.09e+01', '3.55e+01', '3.91e+00', '-5.52e+01', '1.84e+02', '-3.48e+02', '5.79e+01']
Constraint losses = 5.76e-01
Relative update of tokamak psi (in plasma core): 1.87e-02
Handing off to forward solve (with updated currents).
Relative error =  1.04e-02
-----
Iteration: 9
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['1.02e+02', '4.79e+02', '1.20e+02', '1.11e+02', '8.88e+01', '-9.84e+00', '5.68e+00', '3.74e+01', '-4.49e+01', '3.11e+01', '2.51e+02', '6.68e+01']
Constraint losses = 5.75e-01
Relative update of tokamak psi (in plasma core): 2.28e-02
Handing off to forward solve (with updated currents).
Relative error =  1.17e-02
-----
Iteration: 10
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['4.24e+02', '-4.23e+01', '-1.32e+02', '1.65e+01', '1.33e+01', '-7.17e+01', '1.55e+01', '-3.82e+01', '6.07e+01', '-3.56e+01', '-3.56e+02', '8.56e+01']
Constraint losses = 6.04e-01
Relative update of tokamak psi (in plasma core): 2.51e-02
Handing off to forward solve (with updated currents).
Relative error =  1.22e-02
-----
Iteration: 11
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['-3.26e+02', '3.46e+02', '1.15e+02', '5.92e+01', '2.65e+01', '5.02e+01', '-1.14e+01', '5.25e+01', '-7.41e+01', '5.26e+01', '4.02e+02', '3.50e+01']
Constraint losses = 6.29e-01
Relative update of tokamak psi (in plasma core): 2.59e-02
Handing off to forward solve (with updated currents).
Relative error =  1.42e-02
-----
Iteration: 12
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['4.22e+02', '-2.20e+02', '-1.76e+02', '-2.30e+01', '-1.75e+01', '-7.06e+01', '1.54e+01', '-5.53e+01', '8.06e+01', '-5.15e+01', '-4.67e+02', '6.10e+01']
Constraint losses = 6.58e-01
Relative update of tokamak psi (in plasma core): 3.73e-02
Handing off to forward solve (with updated currents).
Relative error =  1.48e-02
-----
Iteration: 13
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['-4.31e+02', '3.33e+02', '1.50e+02', '5.16e+01', '2.16e+01', '6.92e+01', '-1.53e+01', '6.04e+01', '-8.73e+01', '5.95e+01', '4.87e+02', '8.38e+00']
Constraint losses = 6.74e-01
Relative update of tokamak psi (in plasma core): 2.93e-02
Handing off to forward solve (with updated currents).
Relative error =  1.60e-02
-----
Iteration: 14
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['2.42e+02', '-1.51e+02', '-9.92e+01', '-1.96e+01', '-1.24e+01', '-4.04e+01', '8.85e+00', '-3.34e+01', '4.83e+01', '-3.14e+01', '-2.77e+02', '2.44e+01']
Constraint losses = 6.80e-01
Relative update of tokamak psi (in plasma core): 2.26e-02
Handing off to forward solve (with updated currents).
Relative error =  5.80e-03
-----
Iteration: 15
Using simplified Green's Jacobian (of constraints wrt coil currents) to optimise the currents.
Change in coil currents (being controlled): ['4.11e+01', '-3.84e+01', '-1.98e+01', '1.03e+01', '3.10e+00', '-1.19e+01', '2.15e+00', '-6.77e+00', '1.31e+01', '-1.18e+01', '-3.68e+01', '4.05e+01']
Constraint losses = 5.40e-01
Relative update of tokamak psi (in plasma core): 2.79e-03
Handing off to forward solve (with updated currents).
Relative error =  4.40e-03
-----
Iteration: 16
Using full Jacobian (of constraints wrt coil currents) to optimsise currents.
 - calculating derivatives for coil 1/12
 - calculating derivatives for coil 2/12
 - calculating derivatives for coil 3/12
 - calculating derivatives for coil 4/12
 - calculating derivatives for coil 5/12
 - calculating derivatives for coil 6/12
 - calculating derivatives for coil 7/12
 - calculating derivatives for coil 8/12
 - calculating derivatives for coil 9/12
 - calculating derivatives for coil 10/12
 - calculating derivatives for coil 11/12
 - calculating derivatives for coil 12/12
Change in coil currents (being controlled): ['-1.63e+01', '2.54e+01', '1.82e+01', '-4.15e+00', '6.98e+00', '2.34e+01', '-3.16e+00', '1.62e+00', '-8.51e+00', '1.57e+01', '4.85e+01', '4.46e+01']
Constraint losses = 1.71e+00
Relative update of tokamak psi (in plasma core): 3.94e-03
Handing off to forward solve (with updated currents).
Relative error =  2.82e-03
-----
Iteration: 17
Using full Jacobian (of constraints wrt coil currents) to optimsise currents.
 - calculating derivatives for coil 1/12
 - calculating derivatives for coil 2/12
 - calculating derivatives for coil 3/12
 - calculating derivatives for coil 4/12
 - calculating derivatives for coil 5/12
 - calculating derivatives for coil 6/12
 - calculating derivatives for coil 7/12
 - calculating derivatives for coil 8/12
 - calculating derivatives for coil 9/12
 - calculating derivatives for coil 10/12
 - calculating derivatives for coil 11/12
 - calculating derivatives for coil 12/12
Change in coil currents (being controlled): ['-1.26e+01', '8.50e+00', '6.30e+00', '6.52e+00', '1.02e+00', '-1.51e+00', '9.93e-01', '8.25e-01', '-3.76e+00', '-8.51e-02', '1.42e+01', '1.55e+00']
Constraint losses = 1.71e+00
Relative update of tokamak psi (in plasma core): 8.64e-04
Handing off to forward solve (with updated currents).
Relative error =  5.88e-04
-----
Iteration: 18
Using full Jacobian (of constraints wrt coil currents) to optimsise currents.
 - calculating derivatives for coil 1/12
 - calculating derivatives for coil 2/12
 - calculating derivatives for coil 3/12
 - calculating derivatives for coil 4/12
 - calculating derivatives for coil 5/12
 - calculating derivatives for coil 6/12
 - calculating derivatives for coil 7/12
 - calculating derivatives for coil 8/12
 - calculating derivatives for coil 9/12
 - calculating derivatives for coil 10/12
 - calculating derivatives for coil 11/12
 - calculating derivatives for coil 12/12
Change in coil currents (being controlled): ['-9.21e+00', '5.03e+00', '2.71e+00', '9.69e-02', '-4.49e-01', '-2.57e+00', '-1.89e-01', '1.99e+00', '1.30e+00', '-4.31e+00', '2.07e+00', '2.79e+00']
Constraint losses = 1.71e+00
Relative update of tokamak psi (in plasma core): 1.49e-04
Handing off to forward solve (with updated currents).
Relative error =  1.65e-04
-----
Iteration: 19
Using full Jacobian (of constraints wrt coil currents) to optimsise currents.
 - calculating derivatives for coil 1/12
 - calculating derivatives for coil 2/12
 - calculating derivatives for coil 3/12
 - calculating derivatives for coil 4/12
 - calculating derivatives for coil 5/12
 - calculating derivatives for coil 6/12
 - calculating derivatives for coil 7/12
 - calculating derivatives for coil 8/12
 - calculating derivatives for coil 9/12
 - calculating derivatives for coil 10/12
 - calculating derivatives for coil 11/12
 - calculating derivatives for coil 12/12
Change in coil currents (being controlled): ['-8.68e+00', '4.76e+00', '2.54e+00', '-3.09e-01', '-3.41e-01', '-1.29e+00', '-1.24e-01', '3.46e-01', '-1.23e+00', '1.73e+00', '8.63e-01', '2.41e-01']
Constraint losses = 1.71e+00
Relative update of tokamak psi (in plasma core): 8.63e-05
Handing off to forward solve (with updated currents).
Relative error =  8.53e-05
-----
Inverse static solve SUCCESS. Tolerance 8.53e-05 (vs. requested 0.0001) reached in 20/50 iterations.

This converges properly and gives the correct total flux, coil currents close to their true values, and psi_axis/psi_bndry matching to several decimal places.

# print coil currents
reconstructed_currents = {label: eq_inv.tokamak[label].current for label in active_coil_names}
original_active_currents = {label: eq.tokamak[label].current for label in active_coil_names}

print("Coil currents [A]:")
print(f"{'coil':>10s}  {'original':>12s}  {'reconstructed':>14s}")
for label in active_coil_names:
    print(f"{label:>10s}  {original_active_currents[label]:12.1f}  {reconstructed_currents[label]:14.1f}")

# some small differences
psi_diff_inv = np.abs(eq.psi() - eq_inv.psi())
print(f"\nMagnetic axis:   original = {eq.magneticAxis()[0:2]}, reconstructed = {eq_inv.magneticAxis()[0:2]}")
print(f"psi_axis:        original = {eq.psi_axis:.6g}, reconstructed = {eq_inv.psi_axis:.6g}")
print(f"psi_boundary:    original = {eq.psi_bndry:.6g}, reconstructed = {eq_inv.psi_bndry:.6g}")
print(f"Max |Δpsi| relative to max|psi|: {np.max(psi_diff_inv) / np.max(np.abs(eq.psi())):.3e}")

# plot the two equilibria side by side
fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)
for ax, e, title in [(ax1, eq, "Original"), (ax2, eq_inv, "Reconstruction from IDS")]:
    ax.grid(True, which="both")
    e.plot(axis=ax, show=False)
    e.tokamak.plot(axis=ax, show=False)
    ax.set_xlim(0.1, 2.15)
    ax.set_ylim(-2.25, 2.25)
    ax.set_title(title)
    if ax == ax2:
        constrain.plot(axis=ax2)
plt.tight_layout()
Coil currents [A]:
      coil      original   reconstructed
  Solenoid        5000.0          5069.6
        PX        3696.8          3621.5
        D1        6941.2          6968.7
        D2        4390.9          4381.3
        D3        2924.4          2935.4
        Dp       -2488.2         -2479.8
        D5         249.2           249.3
        D6        -315.0          -320.1
        D7         511.1           513.6
        P4       -3564.4         -3565.4
        P5       -3954.2         -3951.0
        P6         500.0           497.7

Magnetic axis:   original = [ 0.98617225 -0.051466  ], reconstructed = [ 0.98567621 -0.05133485]
psi_axis:        original = 0.100458, reconstructed = 0.100493
psi_boundary:    original = 0.0314736, reconstructed = 0.0315588
Max |Δpsi| relative to max|psi|: 2.228e-03
../_images/3e85eaa8ac3842ba41b1cbc61bfac8163790ebbb2880e819b5cd25d8f3927088.png
<Figure size 640x480 with 0 Axes>

In the future, we’ll work on adding in capabilities to save the machine description and coil currents to the pf_active and pf_passive IDS formats.