Equilibrium and Coilset Optimisation

Here we explore how to optimise equilibria, coil currents, and coil positions.

This is an in-depth example, intended for developers and people familiar with plasma equilibrium problems, walking you through many of the objects, approaches, and optimisation problems that are often used when designing plasma equilibria and poloidal field coils.

There are many ways of optimising equilibria, and this example shows just one relatively crude approach. The choice of constraints, optimisation algorithms, and even the sequence of operations has a big influence on the outcome. It is a bit of a dark art, and over time you will hopefully find an approach that works for your problem.

Heavily constraining the plasma shape as we do here is not particularly robust, or even philosophically “right”. It’s however a common approach and comes in useful when one wants to optimise coil positions without affecting the plasma too much.

import contextlib
from copy import deepcopy

import matplotlib.pyplot as plt
import numpy as np
from IPython import get_ipython

from bluemira.base.constants import raw_uc
from bluemira.display import plot_defaults
from bluemira.display.plotter import plot_coordinates
from bluemira.equilibria.coils import Coil, CoilSet
from bluemira.equilibria.diagnostics import PicardDiagnostic, PicardDiagnosticOptions
from bluemira.equilibria.equilibrium import Equilibrium
from bluemira.equilibria.grid import Grid
from bluemira.equilibria.optimisation.constraints import (
    CoilFieldConstraints,
    CoilForceConstraints,
    FieldNullConstraint,
    IsofluxConstraint,
    MagneticConstraintSet,
    PsiConstraint,
)
from bluemira.equilibria.optimisation.problem import (
    MinimalCurrentCOP,
    PulsedNestedPositionCOP,
    TikhonovCurrentCOP,
    UnconstrainedTikhonovCurrentGradientCOP,
)
from bluemira.equilibria.profiles import CustomProfile
from bluemira.equilibria.shapes import JohnerLCFS
from bluemira.equilibria.solve import DudsonConvergence, PicardIterator
from bluemira.geometry.tools import make_polygon
from bluemira.utilities.positioning import PositionMapper, RegionInterpolator

plot_defaults()

with contextlib.suppress(AttributeError):
    get_ipython().run_line_magic("matplotlib", "qt")

First let’s create our inital coilset. This is taken from a reference EU-DEMO design but as you will see we will change this later on.

x = [5.4, 14.0, 17.75, 17.75, 14.0, 7.0, 2.77, 2.77, 2.77, 2.77, 2.77]
z = [9.26, 7.9, 2.5, -2.5, -7.9, -10.5, 7.07, 4.08, -0.4, -4.88, -7.86]
dx = [0.6, 0.7, 0.5, 0.5, 0.7, 1.0, 0.4, 0.4, 0.4, 0.4, 0.4]
dz = [0.6, 0.7, 0.5, 0.5, 0.7, 1.0, 2.99 / 2, 2.99 / 2, 5.97 / 2, 2.99 / 2, 2.99 / 2]

coils = []
j = 1
for i, (xi, zi, dxi, dzi) in enumerate(zip(x, z, dx, dz, strict=False)):
    if j > 6:  # noqa: PLR2004
        j = 1
    ctype = "PF" if i < 6 else "CS"  # noqa: PLR2004
    coil = Coil(xi, zi, current=0, dx=dxi, dz=dzi, ctype=ctype, name=f"{ctype}_{j}")
    coils.append(coil)
    j += 1

coilset = CoilSet(*coils)

Now we can also specify our coilset a little further, by assigning maximum current densities and peak fields. This can then be used in the bounds and constraints for our optimisation problems.

We also fix the sizes of our CS coils for this example, and ‘mesh’ them. Unless otherwise specified a Coil is represented by a single current filament. meshing here refers to the sub-division of this coil filament into several filaments equi-spaced around the coil cross-section.

coilset.assign_material("CS", j_max=16.5e6, b_max=12.5)
coilset.assign_material("PF", j_max=12.5e6, b_max=11.0)

# Later on, we will optimise the PF coil positions, but for the CS coils we can fix sizes
# and mesh them already.

cs = coilset.get_coiltype("CS")
cs.fix_sizes()
cs.discretisation = 0.3

Now, we set up our grid, equilibrium, and profiles.

We’ll just use a CustomProfile for now, but you can also use a BetaIpProfile with a flux function parameterisation if you want to directly constrain poloidal beta values in your equilibrium optimisation.

# Machine parameters
R_0 = 8.938
A = 3.1
B_0 = 4.8901  # T
I_p = 19.07e6  # A

grid = Grid(3.0, 13.0, -10.0, 10.0, 65, 65)

profiles = CustomProfile(
    np.array([86856, 86506, 84731, 80784, 74159, 64576, 52030, 36918, 20314, 4807, 0.0]),
    -np.array([
        0.125,
        0.124,
        0.122,
        0.116,
        0.106,
        0.093,
        0.074,
        0.053,
        0.029,
        0.007,
        0.0,
    ]),
    R_0=R_0,
    B_0=B_0,
    I_p=I_p,
)

eq = Equilibrium(coilset, grid, profiles, psi=None)

Now we need to specify some constraints on the plasma.

We’ll instantiate a parameterisation for the last closed flux surface (LCFS) which tends to do a good job at describing an EU-DEMO-like single null plasma.

We’ll use this to specify some constraints on the plasma equilibrium problem:

• An IsofluxConstraint forces the flux at a set of points to be equal

• A FieldNullConstraint forces the poloidal field at a point to be zero.

lcfs_parameterisation = JohnerLCFS({
    "r_0": {"value": R_0},
    "z_0": {"value": 0.0},
    "a": {"value": R_0 / A},
    "kappa_u": {"value": 1.6},
    "kappa_l": {"value": 1.9},
    "delta_u": {"value": 0.4},
    "delta_l": {"value": 0.4},
    "phi_u_neg": {"value": 0.0},
    "phi_u_pos": {"value": 0.0},
    "phi_l_neg": {"value": 45.0},
    "phi_l_pos": {"value": 30.0},
})

lcfs = lcfs_parameterisation.create_shape().discretise(byedges=True, ndiscr=50)

x_bdry, z_bdry = lcfs.x, lcfs.z
arg_inner = np.argmin(x_bdry)

isoflux = IsofluxConstraint(
    x_bdry,
    z_bdry,
    x_bdry[arg_inner],
    z_bdry[arg_inner],
    tolerance=0.5,  # Difficult to choose...
    constraint_value=0.0,  # Difficult to choose...
)

xp_idx = np.argmin(z_bdry)
x_point = FieldNullConstraint(
    x_bdry[xp_idx],
    z_bdry[xp_idx],
    tolerance=1e-4,  # [T]
)

It’s often very useful to solve an unconstrained optimised problem in order to get an initial guess for the equilibrium result. The initial equilibrium can be used as the “starting guess” for subsequent constrained optimisation problems.

This is done by using the magnetic constraints in a “set” for which the error is then minimised with an L2 norm and a Tikhonov regularisation on the currents.

Note that when we use equilibrium constraints in a “constraint set” as part of an optimisation objective (as is the case here) the tolerances are not used actively.

We can use this to optimise the current gradients during the solution of the equilibrium until convergence.

current_opt_problem = UnconstrainedTikhonovCurrentGradientCOP(
    eq, MagneticConstraintSet([isoflux, x_point]), gamma=1e-7
)
diagnostic_plotting = PicardDiagnosticOptions(plot=PicardDiagnostic.EQ)
program = PicardIterator(
    eq,
    current_opt_problem,
    convergence=DudsonConvergence(1e-3),
    fixed_coils=True,
    relaxation=0.2,
    diagnostic_plotting=diagnostic_plotting,
)
program()

Now say we want to use bounds on our current vector, and that we want to solve a constrained optimisation problem.

We can minimise the error on our target isoflux surface set with some bounds on the current vector, and some additional constraints.

First let’s set up some constraints on the coils (peak fields and peak vertical forces) are common constraints. We can even use our X-point constraint from earlier as a constraint in the optimiser. In other words, we don’t need to lump it together with the isoflux target minimisation objective.

We then instantiate a new optimisation problem, and use this in a Picard iteration scheme, using the previously converged equilibrium as a starting point.

Note that here we are optimising the current vector and not the current gradient vector.

field_constraints = CoilFieldConstraints(eq.coilset, eq.coilset.b_max, tolerance=1e-6)

force_constraints = CoilForceConstraints(
    eq.coilset, PF_Fz_max=450e6, CS_Fz_sum_max=300e6, CS_Fz_sep_max=250e6, tolerance=5e-5
)

current_opt_problem = TikhonovCurrentCOP(
    eq,
    targets=MagneticConstraintSet([isoflux]),
    gamma=0.0,
    opt_algorithm="SLSQP",
    opt_conditions={"max_eval": 2000, "ftol_rel": 1e-6},
    opt_parameters={},
    max_currents=eq.coilset.get_max_current(0.0),
    constraints=[x_point, field_constraints, force_constraints],
)
program = PicardIterator(
    eq,
    current_opt_problem,
    fixed_coils=True,
    convergence=DudsonConvergence(1e-4),
    relaxation=0.1,
)
program()

Now let’s say we don’t actually want to minimise the error, but we want to minimise the coil currents, and use the constraints that we specified above as actual constraints in the optimisation problem (rather than in the objective function as above)

Note that here we’ve got rather a lot of constraints, and that we need to choose the value and tolerance of the isoflux constraint (in particular) wisely.

Too strict a tolerance will likely result in an unsuccessful optimisation with the final result potentially violating the constraint and probably not being an actual optimum.

minimal_current_eq = deepcopy(eq)
minimal_current_opt_problem = MinimalCurrentCOP(
    minimal_current_eq,
    opt_algorithm="SLSQP",
    opt_conditions={"max_eval": 5000, "ftol_rel": 1e-6, "xtol_rel": 1e-6},
    max_currents=coilset.get_max_current(0.0),
    constraints=[isoflux, x_point, field_constraints, force_constraints],
)

program = PicardIterator(
    minimal_current_eq,
    minimal_current_opt_problem,
    fixed_coils=True,
    convergence=DudsonConvergence(1e-4),
    relaxation=0.1,
)
program()

control_coils = eq.coilset.get_control_coils()
min_control_coils = minimal_current_eq.coilset.get_control_coils()

print(
    "Total currents from minimal error optimisation problem:"
    f" {raw_uc(np.sum(np.abs(control_coils.current)), 'A', 'MA'):.2f} MA"
)
print(
    "Total currents from minimal current optimisation problem:"
    f" {raw_uc(np.sum(np.abs(min_control_coils.current)), 'A', 'MA'):.2f} MA"
)

Coil position optimisation

Now, say that we want to optimise the positions the PF coils, and the currents of the entire CoilSet for two different snapshots in a pulse. These snapshots are typically the start of flat-top (SOF) and end of flat-top (EOF) as they represent the most challenging conditions for the coils and plasma shape.

Here we set the flux at the desired LCFS to be 50 V.s and -150 V.s for the SOF and EOF, respectively. In this example, this is an arbitrary decision. In reality, this would relate to the desired pulse duration for a given machine.

We’re going to optimise the positions for an objective function that takes the maximum objective function value of two current sub-optimisation problems.

This is what we refer to as a nested optimisation, in other words that the positions and the currents (in two different situations) are being optimised separately.

First we specify the sub-optimisation problems (objective functions and constraints). Then we specify the position optimisation problem for a single current sub-optimisation problem.

isoflux = IsofluxConstraint(
    x_bdry, z_bdry, x_bdry[arg_inner], z_bdry[arg_inner], tolerance=1e-3
)

sof = deepcopy(eq)
sof_psi_boundary = PsiConstraint(
    x_bdry[arg_inner], z_bdry[arg_inner], target_value=50 / 2 / np.pi, tolerance=1e-3
)

eof = deepcopy(eq)
eof_psi_boundary = PsiConstraint(
    x_bdry[arg_inner], z_bdry[arg_inner], target_value=-150 / 2 / np.pi, tolerance=1e-3
)

current_opt_problem_sof = TikhonovCurrentCOP(
    sof,
    targets=MagneticConstraintSet([isoflux]),
    gamma=1e-12,
    opt_algorithm="SLSQP",
    opt_conditions={"max_eval": 2000, "ftol_rel": 1e-6, "xtol_rel": 1e-6},
    opt_parameters={},
    max_currents=sof.coilset.get_max_current(I_p),
    constraints=[sof_psi_boundary, x_point, field_constraints, force_constraints],
)

current_opt_problem_eof = TikhonovCurrentCOP(
    eof,
    targets=MagneticConstraintSet([isoflux]),
    gamma=1e-12,
    opt_algorithm="SLSQP",
    opt_conditions={"max_eval": 5000, "ftol_rel": 1e-6, "xtol_rel": 1e-6},
    opt_parameters={},
    max_currents=eof.coilset.get_max_current(I_p),
    constraints=[eof_psi_boundary, x_point, field_constraints, force_constraints],
)

We set up a position mapping of the regions in which we would like the PF coils to be. The positions themselves are bounded by the specification of the RegionInterpolators.

Finally, we specify our position optimisation problem, in this case with the two previously defined current sub-optimisation problems (but we could specify more if we wanted to).

Typically the currents can be varied linearly from the SOF to the EOF, so there isn’t really much point in doing more than two different equilibria here from an optimisation perspective.

Note that there are no constraints here, and if we wanted to add some they would have to pertain to the position vector in some form.

For each set of positions, we treat the plasma contribution as being “frozen” and optimise the coil currents (with the various constraints). This works as for relatively good starting guesses (converged equilibria) the plasma contribution to the various constraints does not change much when the equilibrium is subsequently converged.

For the sake of brevity, we set the maximum number of iterations for the position optimisation to 50. This is unlikely to find the best solution given our specified coil regions, but demonstrates the principle with acceptable run-times.

# We'll store this so that we can look at them again later
old_eq = deepcopy(eq)

region_interpolators = {}
pf_coils = eq.coilset.get_coiltype("PF")
for x, z, name in zip(pf_coils.x, pf_coils.z, pf_coils.name, strict=False):
    region = make_polygon(
        {"x": [x - 1, x + 1, x + 1, x - 1], "z": [z - 1, z - 1, z + 1, z + 1]},
        closed=True,
    )
    region_interpolators[name] = RegionInterpolator(region)

position_mapper = PositionMapper(region_interpolators)

position_opt_problem = PulsedNestedPositionCOP(
    eq.coilset,
    position_mapper,
    sub_opt_problems=[current_opt_problem_sof, current_opt_problem_eof],
    opt_algorithm="COBYLA",
    opt_conditions={"max_eval": 100, "ftol_rel": 1e-6, "xtol_rel": 1e-6},
    debug=False,
)

optimised_coilset = position_opt_problem.optimise(verbose=True).coilset

We’ve just optimised the PF coil positions using a single current filament at the centre of each PF coil. This is a reasonable approximation when performing a position optimisation, but we probably want to do better when it comes to our final equilibria. We will figure out the appropriate sizes of the PF coils, and fix them and mesh them accordingly.

We also need to remember to update the bounds of the current optimisation problem!

sof_pf_currents = sof.coilset.get_coiltype("PF").get_control_coils().current
eof_pf_currents = eof.coilset.get_coiltype("PF").get_control_coils().current
max_pf_currents = np.max(np.abs([sof_pf_currents, eof_pf_currents]), axis=0)

pf_coil_names = optimised_coilset.get_coiltype("PF").name

max_cs_currents = optimised_coilset.get_coiltype("CS").get_max_current()

max_currents = np.concatenate([max_pf_currents, max_cs_currents])

for problem in [current_opt_problem_sof, current_opt_problem_eof]:
    for pf_name, max_current in zip(pf_coil_names, max_pf_currents, strict=False):
        problem.eq.coilset[pf_name].resize(max_current)
        problem.eq.coilset[pf_name].fix_size()
        problem.eq.coilset[pf_name].discretisation = 0.3
    problem.set_current_bounds(max_currents)

Now that we’ve:

• optimised the coil positions for a fixed plasma,

• fixed our PF coil sizes,

• meshed our PF coils,

• updated the bounds of the current optimisation problems,

we can run the Grad-Shafranov solve again to converge the equilibria for the optimised coil positions at SOF and EOF.

program = PicardIterator(
    sof,
    current_opt_problem_sof,
    fixed_coils=True,
    convergence=DudsonConvergence(1e-4),
    relaxation=0.1,
    diagnostic_plotting=diagnostic_plotting,
)
program()

program = PicardIterator(
    eof,
    current_opt_problem_eof,
    fixed_coils=True,
    convergence=DudsonConvergence(1e-4),
    relaxation=0.05,
    diagnostic_plotting=diagnostic_plotting,
)
program()

Now let’s compare the old equilibrium and coilset to the one with optimised positions.

f, ax = plt.subplots()
x_old, z_old = old_eq.coilset.position
x_new, z_new = sof.coilset.position
plot_coordinates(old_eq.get_LCFS(), ax=ax, edgecolor="b", fill=False)
plot_coordinates(sof.get_LCFS(), ax=ax, edgecolor="r", fill=False)
plot_coordinates(eof.get_LCFS(), ax=ax, edgecolor="g", fill=False)
ax.plot(x_old, z_old, linewidth=0, marker="o", color="b")
ax.plot(x_new, z_new, linewidth=0, marker="+", color="r")
isoflux.plot(ax=ax)
plt.show()

Note that one could converge the Grad-Shafranov equation for each set of coil positions but this would be much slower and probably less robust. Personally, I don’t think it is worthwhile, but were it to succeed it would be fair to say it would be a better optimum.

Let’s look at the full result: reference equilibrium, SOF, and EOF

f, (ax_1, ax_2, ax_3) = plt.subplots(1, 3)

old_eq.plot(ax=ax_1)
old_eq.coilset.plot(ax=ax_1, label=True)

sof.plot(ax=ax_2)
sof.coilset.plot(ax=ax_2, label=True)
ax_2.set_title("SOF $\\Psi_{b} = $" + f"{sof.get_OX_psis()[1] * 2 * np.pi:.2f} V.s")


eof.plot(ax=ax_3)
eof.coilset.plot(ax=ax_3, label=True)
ax_3.set_title("EOF $\\Psi_{b} = $" + f"{eof.get_OX_psis()[1] * 2 * np.pi:.2f} V.s")
plt.show()