Omnigenity Optimization
This tutorial demonstrates how to optimize for omnigenity in DESC. It will go through an example using omnigenity with poloidally closed contours of magnetic field strength (OP), but the method is capable of optimizing for any general omnigenous magnetic fields as explained in Dudt et al. (2024).
[1]:
import sys
import os
sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../../../"))
If you have access to a GPU, uncomment the following two lines.
[2]:
# from desc import set_device
# set_device("gpu")
[3]:
import matplotlib.pyplot as plt
import numpy as np
from desc.backend import jnp
from desc.equilibrium import Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.grid import LinearGrid
from desc.magnetic_fields import OmnigenousField
from desc.objectives import (
AspectRatio,
CurrentDensity,
FixCurrent,
FixOmniBmax,
FixOmniMap,
FixPressure,
FixPsi,
GenericObjective,
LinearObjectiveFromUser,
ObjectiveFunction,
Omnigenity,
)
from desc.optimize import Optimizer
from desc.plotting import plot_boozer_surface, plot_boundaries
plt.rcParams["font.size"] = 14
DESC version 0.10.4+1069.g4a9334fbc.dirty,using JAX backend, jax version=0.4.11, jaxlib version=0.4.11, dtype=float64
Using device: CPU, with 11.03 GB available memory
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/jaxtyping/__init__.py:221: UserWarning: jaxtyping version >=0.2.23 should be used with Equinox version >=0.11.1
warnings.warn(
As an initial guess for the optimization, we will start with a boundary shape generated by an analytic model for (very approximate) quasi-poloidal symmetry (QP). In this example, we will seek a vacuum magnetic field with two field periods (\(N_{FP}=2\)), an aspect ratio of \(R_0/a\leq10\), and a mirror ratio on axis of \(\Delta=0.2\).
[4]:
surf = FourierRZToroidalSurface.from_qp_model(
major_radius=1,
aspect_ratio=10,
elongation=3,
mirror_ratio=0.2,
torsion=0.1,
NFP=2,
sym=True,
)
# this value of Psi gives an average |B| on axis of about 1 T
# the Equilibrium class defaults to vacuum pressure and current profiles
eq = Equilibrium(Psi=3e-2, M=4, N=4, surface=surf)
Now that the equilibrium is initialized, we need to solve the fixed-boundary vacuum equilibrium:
[5]:
eq, _ = eq.solve(objective="force", verbose=3)
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 1.78 sec
Timer: Objective build = 2.99 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed-Psi
Building objective: fixed-pressure
Building objective: fixed-current
Building objective: self_consistency R
Building objective: self_consistency Z
Building objective: lambda gauge
Building objective: axis R self consistency
Building objective: axis Z self consistency
Timer: Objective build = 3.65 sec
Timer: Linear constraint projection build = 1.61 sec
Number of parameters: 120
Number of objectives: 850
Starting optimization
Using method: lsq-exact
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 1.043e-01 4.108e+01
1 2 2.850e-03 1.014e-01 1.372e-01 1.723e+01
2 3 5.358e-05 2.797e-03 6.424e-02 1.375e+00
3 4 1.835e-05 3.524e-05 5.304e-02 1.697e-01
4 5 1.632e-05 2.021e-06 1.809e-02 2.100e-02
5 6 1.627e-05 5.057e-08 1.944e-02 5.184e-02
Optimization terminated successfully.
`ftol` condition satisfied.
Current function value: 1.627e-05
Total delta_x: 1.421e-01
Iterations: 5
Function evaluations: 6
Jacobian evaluations: 6
Timer: Solution time = 6.82 sec
Timer: Avg time per step = 1.13 sec
Start of solver
Total (sum of squares): 1.043e-01,
Maximum absolute Force error: 4.672e+04 (N)
Minimum absolute Force error: 3.799e+00 (N)
Average absolute Force error: 4.703e+03 (N)
Maximum absolute Force error: 6.138e-02 (normalized)
Minimum absolute Force error: 4.991e-06 (normalized)
Average absolute Force error: 6.178e-03 (normalized)
R boundary error: 0.000e+00 (m)
Z boundary error: 0.000e+00 (m)
Fixed-Psi error: 0.000e+00 (Wb)
Fixed-pressure profile error: 0.000e+00 (Pa)
Fixed-current profile error: 0.000e+00 (A)
End of solver
Total (sum of squares): 1.627e-05,
Maximum absolute Force error: 6.159e+02 (N)
Minimum absolute Force error: 2.030e-01 (N)
Average absolute Force error: 5.960e+01 (N)
Maximum absolute Force error: 8.092e-04 (normalized)
Minimum absolute Force error: 2.667e-07 (normalized)
Average absolute Force error: 7.831e-05 (normalized)
R boundary error: 6.939e-18 (m)
Z boundary error: 0.000e+00 (m)
Fixed-Psi error: 0.000e+00 (Wb)
Fixed-pressure profile error: 0.000e+00 (Pa)
Fixed-current profile error: 0.000e+00 (A)
Let us make a copy of this initial equilibrium so that we can compare our final solution to it later and see how well the omnigenity optimization worked. By plotting the \(|B|\) contours in Boozer coordinates, we can see that this equilibrium is already somewhat close to being omnigenous with poloidal contours, but far from perfect.
[6]:
eq0 = eq.copy()
plot_boozer_surface(eq0, fieldlines=4) # defaults to the rho=1 surface
[6]:
(<Figure size 576.078x576.078 with 2 Axes>,
<AxesSubplot: title={'center': '$|\\mathbf{B}|~(T)$'}, xlabel='$\\zeta_{Boozer}$', ylabel='$\\theta_{Boozer}$'>)
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/IPython/core/pylabtools.py:152: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
fig.canvas.print_figure(bytes_io, **kw)
In order to optimize the equilibrium for omnigenity, we need to create a target omnigenous magnetic field. The OmnigenousField class has two attributes that represent parameters in the omnigenous magnetic field model: - B_lm specifies the shape of the “magnetic well” on each flux surface. - x_lmn specifies how the well shape varies along different field lines.
The helicity is given by the tuple of integers \((M, N)\), and is set to \((0, N_{FP})\) for omnigenity with poloidal contours in this example. The typical case for helical contours would be \((1, N_{FP})\), and for toroidal contours would be \((1, 0)\).
We need to specify the resolution of the parameter space, given by the following integers: - L_B is the maximum power of \(\rho\) used in the Chebyshev polynomial expansion for B_lm. - M_B is the number of spline knots used on each surface for B_lm. - L_x is the maximum power of \(\rho\) used in the Chebyshev polynomial expansion for x_lmn. - M_x is the maximum mode number used in the cosine series expansion in \(\eta\) for x_lmn. - N_x is the
maximum mode number used in the Fourier series expansion in \(\alpha\) for x_lmn. Quasi-symmetry corresponds to N_x=0.
We provide initial values for the well shape parameters B_lm so that we can set the mirror ratio. The total number of parameters is B_lm.size = M_B * (L_B + 1). In this example, the well shape on each surface is represented by three spline knots, and there is \(\mathcal{O}(\rho)\) variation across the flux surfaces. Here we set only the constant terms of the Chebyshev polynomials such that the initial target field has the same magnetic well from
\(B_{\mathrm{min}}=0.8\mathrm{~T}\) to \(B_{\mathrm{max}}=1.2\mathrm{~T}\) on each surface (corresponding to a mirror ratio of \(\Delta=0.2\)).
The x_lmn parameters are left to their default values of 0, which corresponds to a quasi-poloidally symmetric (QP) initial target field. These parameters will not be fixed during the optimization, so the final result will not be constrained to QP symmetry.
[7]:
field = OmnigenousField(
L_B=1, # radial resolution of B_lm parameters
M_B=3, # number of spline knots on each flux surface
L_x=1, # radial resolution of x_lmn parameters
M_x=1, # eta resolution of x_lmn parameters
N_x=1, # alpha resolution of x_lmn parameters
NFP=eq.NFP, # number of field periods; should always be equal to Equilibrium.NFP
helicity=(0, eq.NFP), # helicity for poloidally closed |B| contours
B_lm=np.array( # magnetic well shape parameters
[
[0.8, 1.0, 1.2], # the first M_well coefficients are the L_B=0 spline knots
[0, 0, 0],
] # the next M_well coefficients are the L_B=1 spline knots, etc.
).flatten(),
)
We can use the field.compute function to visualize what the current target well shape for this field is, with the minimum and maximum as we have prescribed. This magnetic well will be allowed to change during the optimization process according to how the B_lm coefficients change.
[8]:
# plot initial target well |B|
grid_well = LinearGrid(rho=[0.0], M=50)
data_initial = field.compute("|B|", grid=grid_well)
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
ax.plot(data_initial["eta"], data_initial["|B|"], c="b", lw=2)
ax.scatter(
[0, np.pi / 4, np.pi / 2],
field.B_lm[0 : field.M_B] - field.B_lm[field.M_B :],
c="k",
marker=".",
label="Spline Knots",
s=120,
)
ax.set_xticks([-np.pi / 2, 0, np.pi / 2])
ax.set_xticklabels(["$-\pi/2$", "0", "$\pi/2$"])
ax.set_xlabel(r"$\eta$")
ax.set_ylabel("|B|")
ax.legend(loc="upper center")
[8]:
<matplotlib.legend.Legend at 0x7f2bc0733430>
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/IPython/core/pylabtools.py:152: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
fig.canvas.print_figure(bytes_io, **kw)
Next we create the objective function for the optimization. We include one objective to keep the major radius at \(R_0=1~m\), and another to keep the aspect ratio at \(R_0/a\leq10\). The elongation will be unconstrained, but that is another common objective we could choose to include.
We will also target omnigenity on two flux surfaces: \(\rho=0.5\) and \(\rho=1\). The Omnigenity objective class requires two different computational grids: - eq_grid is the grid used to compute the Boozer transform. - field_grid is the grid corresponding to \((\rho,\eta,\alpha)\) coordinates where the omnigenity residuals are minimized.
A separate Omnigenity objective is required for each flux surface, but they all reference the same Equilibrium and OmnigenousField. Make sure both grids for each objective are at the desired surface and have sym=False!
[9]:
eq_half_grid = LinearGrid(rho=0.5, M=4 * eq.M, N=4 * eq.N, NFP=eq.NFP, sym=False)
eq_lcfs_grid = LinearGrid(rho=1.0, M=4 * eq.M, N=4 * eq.N, NFP=eq.NFP, sym=False)
field_half_grid = LinearGrid(rho=0.5, theta=16, zeta=8, NFP=field.NFP, sym=False)
field_lcfs_grid = LinearGrid(rho=1.0, theta=16, zeta=8, NFP=field.NFP, sym=False)
objective = ObjectiveFunction(
(
# target major radius of R0=1 m
GenericObjective("R0", eq=eq, target=1.0, name="major radius"),
# target aspect ratio R0/a<=10
AspectRatio(eq=eq, bounds=(0, 10)),
# omnigenity on the rho=0.5 surface
Omnigenity(
eq=eq,
field=field,
eq_grid=eq_half_grid,
field_grid=field_half_grid,
eta_weight=1,
),
# omnigenity on the rho=1.0 surface
Omnigenity(
eq=eq,
field=field,
eq_grid=eq_lcfs_grid,
field_grid=field_lcfs_grid,
eta_weight=2,
),
)
)
Next we set the optimization constraints. The CurrentDensity, FixPressure, FixCurrent, and FixPsi objectives ensure that we maintain a good vacuum equilibrium during the optimization.
We also include three additional constraints that are unique to the omnigenity optimization: - A perfect omnigenous magnetic field must have a straight \(B_{\mathrm{max}}\) contour in Boozer coordinates, and this is accomplished with the FixOmniBmax objective. - The FixOmniMap objective is used to fix values of the field.x_lmn parameters. In this OP example our \(B_{\mathrm{max}}\) contour is located at \(\zeta_B=0\), and this constraint is used to ensure that the
\(B_{\mathrm{min}}\) contour is located at \(\zeta_B=\pi/N_{FP}\) on average. This constraint should always be true for stellarator symmetry. - The LinearObjectiveFromUser objective is used to fix the sum of values of the field.B_lm parameters. Here we use it to fix the values of \(B_{\mathrm{min}}\) and \(B_{\mathrm{max}}\) on axis to contrain the mirror ratio. The shape of the magnetic well will still have one degree of freedom on the magnetic axis, and the mirror ratio
is unconstrained on other flux surfaces.
[10]:
def mirrorRatio(params):
"""Custom linear function to constrain the mirror ratio on axis."""
B_lm = params["B_lm"]
f = jnp.array(
[
B_lm[0] - B_lm[field.M_B], # B_min on axis
B_lm[field.M_B - 1] - B_lm[-1], # B_max on axis
]
)
return f
constraints = (
CurrentDensity(eq=eq), # vacuum equilibrium force balance
FixPressure(eq=eq), # fix vacuum pressure profile
FixCurrent(eq=eq), # fix vacuum current profile
FixPsi(eq=eq), # fix total toroidal magnetic flux
# ensure the B_max contour is straight in Boozer coordinates
FixOmniBmax(field=field),
# ensure the average B_min contour is at zeta_B=pi/NFP
FixOmniMap(field=field, indices=np.where(field.x_basis.modes[:, 1] == 0)[0]),
# fix the mirror ratio on the magnetic axis
LinearObjectiveFromUser(mirrorRatio, field, target=[0.8, 1.2]),
)
Finally we are ready to run the optimization! We use a least-squares augmented Lagrangian optimizer, but the “proximal” least-squares optimizer would also work. Note that because we are optimizing multiple “things” (the equilibrium eq and the omnigenous field field) we must use optimizer.optimize() instead of Equilibrium.optimize().
[11]:
optimizer = Optimizer("lsq-auglag")
(eq, field), _ = optimizer.optimize(
(eq, field), objective, constraints, maxiter=100, verbose=3
)
Building objective: major radius
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 59.7 ms
Building objective: omnigenity
Precomputing transforms
Timer: Precomputing transforms = 567 ms
Building objective: omnigenity
Precomputing transforms
Timer: Precomputing transforms = 66.2 ms
Timer: Objective build = 795 ms
Building objective: fixed-pressure
Building objective: fixed-current
Building objective: fixed-Psi
Building objective: fixed omnigenity B_max
Building objective: fixed omnigenity map
Building objective: custom linear
Building objective: self_consistency R
Building objective: self_consistency Z
Building objective: lambda gauge
Building objective: axis R self consistency
Building objective: axis Z self consistency
Timer: Objective build = 779 ms
Building objective: current density
Precomputing transforms
Timer: Precomputing transforms = 84.0 ms
Timer: Objective build = 261 ms
Timer: Linear constraint projection build = 1.65 sec
Timer: Linear constraint projection build = 88.5 ms
Number of parameters: 211
Number of objectives: 258
Number of equality constraints: 1275
Number of inequality constraints: 0
Starting optimization
Using method: lsq-auglag
Iteration Total nfev Cost Cost reduction Step norm Optimality Constr viol. Penalty param max(|mltplr|)
0 1 3.256e+01 1.475e+04 2.143e-02 1.000e+01 0.000e+00
1 2 3.153e+01 1.037e+00 9.907e-04 1.382e+04 2.287e-02 1.000e+01 0.000e+00
2 3 2.779e+01 3.741e+00 4.386e-03 1.061e+04 5.549e-02 1.000e+01 0.000e+00
3 4 1.810e+01 9.683e+00 2.541e-02 4.472e+03 1.385e-01 1.000e+01 0.000e+00
4 5 6.789e+00 1.131e+01 1.356e-01 1.434e+03 1.582e-01 1.000e+01 0.000e+00
5 6 1.679e+00 5.110e+00 1.541e-01 9.284e+02 9.832e-02 1.000e+01 0.000e+00
6 8 1.141e+00 5.381e-01 1.281e-01 1.429e+03 6.245e-02 1.000e+01 0.000e+00
7 10 7.641e-01 3.769e-01 8.749e-02 4.986e+02 6.297e-02 1.000e+01 0.000e+00
8 12 6.396e-01 1.244e-01 9.038e-02 1.397e+03 5.265e-02 1.000e+01 0.000e+00
9 13 5.337e-01 1.060e-01 7.602e-02 3.731e+02 3.530e-02 1.000e+01 0.000e+00
10 15 4.125e-01 1.211e-01 7.919e-02 2.077e+02 2.877e-02 1.000e+01 0.000e+00
11 17 3.494e-01 6.312e-02 5.712e-02 2.011e+02 2.656e-02 1.000e+01 0.000e+00
12 18 3.368e-01 1.258e-02 1.154e-01 1.634e+03 2.887e-02 1.000e+01 0.000e+00
13 19 1.547e-01 1.822e-01 3.043e-02 2.252e+02 2.152e-02 1.000e+01 0.000e+00
14 21 1.352e-01 1.949e-02 2.488e-02 7.064e+01 2.007e-02 1.000e+01 0.000e+00
15 23 1.192e-01 1.602e-02 2.433e-02 7.707e+01 1.928e-02 1.000e+01 0.000e+00
16 25 1.054e-01 1.373e-02 2.254e-02 6.183e+01 1.819e-02 1.000e+01 0.000e+00
17 27 9.370e-02 1.173e-02 2.376e-02 5.963e+01 1.719e-02 1.000e+01 0.000e+00
18 29 8.336e-02 1.034e-02 2.283e-02 5.346e+01 1.609e-02 1.000e+01 0.000e+00
19 31 7.452e-02 8.845e-03 2.290e-02 4.734e+01 1.489e-02 1.000e+01 0.000e+00
20 33 6.683e-02 7.693e-03 2.266e-02 4.097e+01 1.378e-02 1.000e+01 0.000e+00
21 35 6.018e-02 6.647e-03 2.300e-02 3.964e+01 1.329e-02 1.000e+01 0.000e+00
22 36 5.599e-02 4.191e-03 8.023e-02 3.972e+02 1.615e-02 1.000e+01 0.000e+00
23 37 3.813e-02 1.786e-02 2.208e-02 5.597e+01 1.227e-02 1.000e+01 0.000e+00
24 38 3.391e-02 4.221e-03 6.744e-02 3.674e+02 1.556e-02 1.000e+01 0.000e+00
25 39 2.600e-02 7.906e-03 1.665e-02 5.092e+01 1.079e-02 1.000e+01 0.000e+00
26 41 2.431e-02 1.696e-03 1.696e-02 2.782e+01 1.071e-02 1.000e+01 0.000e+00
27 43 2.268e-02 1.622e-03 1.488e-02 2.990e+01 1.059e-02 1.000e+01 0.000e+00
28 45 2.125e-02 1.433e-03 1.448e-02 2.618e+01 1.045e-02 1.000e+01 0.000e+00
29 47 1.997e-02 1.283e-03 1.393e-02 2.335e+01 1.030e-02 1.000e+01 0.000e+00
30 49 1.883e-02 1.137e-03 1.374e-02 2.334e+01 1.014e-02 1.000e+01 0.000e+00
31 51 1.782e-02 1.009e-03 1.362e-02 2.766e+01 9.978e-03 1.000e+01 0.000e+00
32 53 1.693e-02 8.936e-04 1.363e-02 3.202e+01 9.808e-03 1.000e+01 0.000e+00
33 55 1.613e-02 7.929e-04 1.368e-02 3.600e+01 9.634e-03 1.000e+01 0.000e+00
34 57 1.543e-02 7.073e-04 1.376e-02 3.923e+01 9.458e-03 1.000e+01 0.000e+00
35 59 1.479e-02 6.369e-04 1.384e-02 4.139e+01 9.280e-03 1.000e+01 0.000e+00
36 61 1.421e-02 5.805e-04 1.392e-02 4.382e+01 9.097e-03 1.000e+01 0.000e+00
37 63 1.367e-02 5.359e-04 1.400e-02 4.557e+01 8.910e-03 1.000e+01 0.000e+00
38 65 1.317e-02 5.009e-04 1.411e-02 4.634e+01 8.715e-03 1.000e+01 0.000e+00
39 67 1.270e-02 4.736e-04 1.426e-02 4.610e+01 8.508e-03 1.000e+01 0.000e+00
40 69 1.225e-02 4.519e-04 1.447e-02 4.481e+01 8.287e-03 1.000e+01 0.000e+00
41 71 1.181e-02 4.330e-04 1.473e-02 4.247e+01 8.044e-03 1.000e+01 0.000e+00
42 73 1.140e-02 4.132e-04 1.502e-02 3.920e+01 8.130e-03 1.000e+01 0.000e+00
43 75 1.101e-02 3.886e-04 1.531e-02 3.518e+01 8.206e-03 1.000e+01 0.000e+00
44 77 1.066e-02 3.568e-04 1.548e-02 3.059e+01 8.226e-03 1.000e+01 0.000e+00
45 79 1.034e-02 3.195e-04 1.543e-02 2.550e+01 8.186e-03 1.000e+01 0.000e+00
46 81 1.005e-02 2.823e-04 1.518e-02 1.995e+01 8.087e-03 1.000e+01 0.000e+00
47 83 9.802e-03 2.521e-04 1.494e-02 2.009e+01 7.938e-03 1.000e+01 0.000e+00
48 85 9.571e-03 2.310e-04 1.485e-02 2.268e+01 7.753e-03 1.000e+01 0.000e+00
49 87 9.359e-03 2.119e-04 1.484e-02 2.708e+01 7.555e-03 1.000e+01 0.000e+00
50 89 9.181e-03 1.783e-04 1.474e-02 3.621e+01 7.366e-03 1.000e+01 0.000e+00
51 90 9.044e-03 1.367e-04 1.456e-02 4.504e+01 7.202e-03 1.000e+01 0.000e+00
52 91 8.897e-03 1.474e-04 1.445e-02 4.873e+01 7.046e-03 1.000e+01 0.000e+00
53 92 8.669e-03 2.284e-04 1.442e-02 4.589e+01 6.846e-03 1.000e+01 0.000e+00
54 93 8.394e-03 2.747e-04 1.436e-02 4.209e+01 6.580e-03 1.000e+01 0.000e+00
55 94 8.148e-03 2.455e-04 1.423e-02 4.347e+01 6.292e-03 1.000e+01 0.000e+00
56 95 7.915e-03 2.338e-04 1.395e-02 4.362e+01 6.378e-03 1.000e+01 0.000e+00
57 96 7.678e-03 2.368e-04 1.333e-02 4.033e+01 6.491e-03 1.000e+01 0.000e+00
58 98 7.456e-03 2.215e-04 1.241e-02 3.512e+01 6.541e-03 1.000e+01 0.000e+00
59 100 7.263e-03 1.932e-04 1.148e-02 2.985e+01 6.525e-03 1.000e+01 0.000e+00
60 102 7.095e-03 1.685e-04 1.072e-02 2.553e+01 6.476e-03 1.000e+01 0.000e+00
61 104 6.941e-03 1.535e-04 1.013e-02 2.211e+01 6.420e-03 1.000e+01 0.000e+00
62 106 6.797e-03 1.442e-04 9.663e-03 1.929e+01 6.362e-03 1.000e+01 0.000e+00
63 108 6.660e-03 1.371e-04 9.265e-03 1.676e+01 6.301e-03 1.000e+01 0.000e+00
64 110 6.529e-03 1.310e-04 8.918e-03 1.426e+01 6.238e-03 1.000e+01 0.000e+00
65 112 6.403e-03 1.257e-04 8.609e-03 1.174e+01 6.176e-03 1.000e+01 0.000e+00
66 114 6.282e-03 1.209e-04 8.332e-03 9.232e+00 6.115e-03 1.000e+01 0.000e+00
67 116 6.166e-03 1.167e-04 8.083e-03 6.917e+00 6.057e-03 1.000e+01 0.000e+00
68 117 6.063e-03 1.023e-04 2.519e-02 4.373e+01 5.776e-03 1.000e+01 0.000e+00
69 118 5.622e-03 4.410e-04 7.956e-03 6.986e+00 5.777e-03 1.000e+01 0.000e+00
70 120 5.517e-03 1.053e-04 8.047e-03 7.558e+00 5.758e-03 1.000e+01 0.000e+00
71 122 5.429e-03 8.832e-05 7.594e-03 6.763e+00 5.717e-03 1.000e+01 0.000e+00
72 124 5.335e-03 9.378e-05 7.595e-03 6.796e+00 5.676e-03 1.000e+01 0.000e+00
73 126 5.248e-03 8.635e-05 7.315e-03 7.292e+00 5.622e-03 1.000e+01 0.000e+00
74 128 5.165e-03 8.371e-05 7.059e-03 7.938e+00 5.566e-03 1.000e+01 0.000e+00
75 130 5.086e-03 7.884e-05 6.788e-03 8.856e+00 5.500e-03 1.000e+01 0.000e+00
76 132 5.011e-03 7.480e-05 6.625e-03 1.025e+01 5.430e-03 1.000e+01 0.000e+00
77 134 4.940e-03 7.151e-05 6.566e-03 1.124e+01 5.354e-03 1.000e+01 0.000e+00
78 136 4.871e-03 6.829e-05 6.604e-03 1.211e+01 5.276e-03 1.000e+01 0.000e+00
79 138 4.806e-03 6.583e-05 6.699e-03 1.253e+01 5.194e-03 1.000e+01 0.000e+00
80 140 4.742e-03 6.325e-05 6.818e-03 1.272e+01 5.110e-03 1.000e+01 0.000e+00
81 142 4.681e-03 6.098e-05 6.941e-03 1.338e+01 5.023e-03 1.000e+01 0.000e+00
82 144 4.623e-03 5.867e-05 7.061e-03 1.416e+01 4.934e-03 1.000e+01 0.000e+00
83 146 4.566e-03 5.650e-05 7.185e-03 1.509e+01 4.842e-03 1.000e+01 0.000e+00
84 148 4.512e-03 5.439e-05 7.332e-03 1.631e+01 4.746e-03 1.000e+01 0.000e+00
85 150 4.459e-03 5.231e-05 7.526e-03 1.806e+01 4.646e-03 1.000e+01 0.000e+00
86 151 4.409e-03 4.996e-05 7.807e-03 2.071e+01 4.540e-03 1.000e+01 0.000e+00
87 152 4.363e-03 4.650e-05 8.246e-03 2.500e+01 4.426e-03 1.000e+01 0.000e+00
88 153 4.323e-03 3.949e-05 9.002e-03 3.217e+01 4.297e-03 1.000e+01 0.000e+00
89 154 4.230e-03 9.383e-05 4.685e-03 6.640e+00 4.242e-03 1.000e+01 0.000e+00
90 156 4.206e-03 2.331e-05 5.468e-03 7.071e+00 4.292e-03 1.000e+01 0.000e+00
91 158 4.179e-03 2.714e-05 6.067e-03 7.628e+00 4.358e-03 1.000e+01 0.000e+00
92 160 4.153e-03 2.614e-05 6.467e-03 7.587e+00 4.440e-03 1.000e+01 0.000e+00
93 162 4.128e-03 2.524e-05 6.371e-03 7.179e+00 4.512e-03 1.000e+01 0.000e+00
94 164 4.104e-03 2.337e-05 5.866e-03 6.305e+00 4.553e-03 1.000e+01 0.000e+00
95 166 4.082e-03 2.237e-05 5.291e-03 5.277e+00 4.564e-03 1.000e+01 0.000e+00
96 168 4.060e-03 2.234e-05 4.842e-03 4.324e+00 4.550e-03 1.000e+01 0.000e+00
97 170 4.037e-03 2.262e-05 4.549e-03 3.549e+00 4.517e-03 1.000e+01 0.000e+00
98 172 4.014e-03 2.290e-05 4.377e-03 2.928e+00 4.474e-03 1.000e+01 0.000e+00
99 174 3.991e-03 2.320e-05 4.293e-03 2.437e+00 4.424e-03 1.000e+01 0.000e+00
100 176 3.967e-03 2.359e-05 4.274e-03 2.053e+00 4.369e-03 1.000e+01 0.000e+00
Warning: Maximum number of iterations has been exceeded.
Current function value: 3.967e-03
Constraint violation: 4.369e-03
Total delta_x: 6.815e-01
Iterations: 100
Function evaluations: 176
Jacobian evaluations: 100
Timer: Solution time = 1.47 min
Timer: Avg time per step = 878 ms
Start of solver
Total (sum of squares): 3.256e+01,
GenericObjective value: 9.999e-01 (m)
GenericObjective value: -1.493e-04 (normalized error)
Aspect ratio: 9.973e+00 (dimensionless)
Maximum absolute Omnigenity error: 2.009e-01 (T)
Minimum absolute Omnigenity error: 3.144e-03 (T)
Average absolute Omnigenity error: 8.289e-02 (T)
Maximum absolute Omnigenity error: 2.009e-01 (normalized)
Minimum absolute Omnigenity error: 3.144e-03 (normalized)
Average absolute Omnigenity error: 8.289e-02 (normalized)
Maximum absolute Omnigenity error: 4.613e-01 (T)
Minimum absolute Omnigenity error: 1.207e-02 (T)
Average absolute Omnigenity error: 1.420e-01 (T)
Maximum absolute Omnigenity error: 4.613e-01 (normalized)
Minimum absolute Omnigenity error: 1.207e-02 (normalized)
Average absolute Omnigenity error: 1.420e-01 (normalized)
Maximum absolute Current density: 1.149e+04 (A*m)
Minimum absolute Current density: 3.072e+00 (A*m)
Average absolute Current density: 9.327e+02 (A*m)
Maximum absolute Current density: 6.593e-03 (normalized)
Minimum absolute Current density: 1.763e-06 (normalized)
Average absolute Current density: 5.354e-04 (normalized)
Fixed-pressure profile error: 0.000e+00 (Pa)
Fixed-current profile error: 0.000e+00 (A)
Fixed-Psi error: 0.000e+00 (Wb)
Fixed omnigenity B_max error: 0.000e+00 (rad)
Fixed omnigenity map error: 0.000e+00 (rad)
Custom linear Objective value: 0.000e+00(Unknown)
End of solver
Total (sum of squares): 3.967e-03,
GenericObjective value: 1.023e+00 (m)
GenericObjective value: 2.266e-02 (normalized error)
Aspect ratio: 9.779e+00 (dimensionless)
Maximum absolute Omnigenity error: 4.628e-03 (T)
Minimum absolute Omnigenity error: 6.075e-06 (T)
Average absolute Omnigenity error: 1.614e-03 (T)
Maximum absolute Omnigenity error: 4.628e-03 (normalized)
Minimum absolute Omnigenity error: 6.075e-06 (normalized)
Average absolute Omnigenity error: 1.614e-03 (normalized)
Maximum absolute Omnigenity error: 3.188e-03 (T)
Minimum absolute Omnigenity error: 2.929e-05 (T)
Average absolute Omnigenity error: 7.248e-04 (T)
Maximum absolute Omnigenity error: 3.188e-03 (normalized)
Minimum absolute Omnigenity error: 2.929e-05 (normalized)
Average absolute Omnigenity error: 7.248e-04 (normalized)
Maximum absolute Current density: 1.203e+04 (A*m)
Minimum absolute Current density: 5.637e-01 (A*m)
Average absolute Current density: 4.930e+02 (A*m)
Maximum absolute Current density: 6.907e-03 (normalized)
Minimum absolute Current density: 3.236e-07 (normalized)
Average absolute Current density: 2.830e-04 (normalized)
Fixed-pressure profile error: 0.000e+00 (Pa)
Fixed-current profile error: 0.000e+00 (A)
Fixed-Psi error: 0.000e+00 (Wb)
Fixed omnigenity B_max error: 0.000e+00 (rad)
Fixed omnigenity map error: 0.000e+00 (rad)
Custom linear Objective value: 3.140e-16(Unknown)
Since we used an augmented Lagrangian optimizer, the nonlinear equilibrium constraint is not guaranteed to be satisfied. It is typically smart to re-solve the fixed-boundary equilibrium after optimization to ensure we have low force balance residuals.
[12]:
eq, _ = eq.solve(objective="force", verbose=3)
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 115 ms
Timer: Objective build = 304 ms
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed-Psi
Building objective: fixed-pressure
Building objective: fixed-current
Building objective: self_consistency R
Building objective: self_consistency Z
Building objective: lambda gauge
Building objective: axis R self consistency
Building objective: axis Z self consistency
Timer: Objective build = 898 ms
Timer: Linear constraint projection build = 1.06 sec
Number of parameters: 120
Number of objectives: 850
Starting optimization
Using method: lsq-exact
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 2.032e-05 2.827e-01
1 2 8.473e-06 1.185e-05 3.102e-02 1.069e-01
2 3 8.278e-06 1.949e-07 3.186e-03 2.764e-03
3 4 8.278e-06 3.635e-10 9.147e-04 1.453e-04
Optimization terminated successfully.
`ftol` condition satisfied.
Current function value: 8.278e-06
Total delta_x: 3.211e-02
Iterations: 3
Function evaluations: 4
Jacobian evaluations: 4
Timer: Solution time = 6.02 sec
Timer: Avg time per step = 1.50 sec
Start of solver
Total (sum of squares): 2.032e-05,
Maximum absolute Force error: 1.187e+03 (N)
Minimum absolute Force error: 4.845e-02 (N)
Average absolute Force error: 6.262e+01 (N)
Maximum absolute Force error: 1.514e-03 (normalized)
Minimum absolute Force error: 6.177e-08 (normalized)
Average absolute Force error: 7.984e-05 (normalized)
R boundary error: 0.000e+00 (m)
Z boundary error: 0.000e+00 (m)
Fixed-Psi error: 0.000e+00 (Wb)
Fixed-pressure profile error: 0.000e+00 (Pa)
Fixed-current profile error: 0.000e+00 (A)
End of solver
Total (sum of squares): 8.278e-06,
Maximum absolute Force error: 7.519e+02 (N)
Minimum absolute Force error: 4.210e-02 (N)
Average absolute Force error: 4.333e+01 (N)
Maximum absolute Force error: 9.586e-04 (normalized)
Minimum absolute Force error: 5.368e-08 (normalized)
Average absolute Force error: 5.524e-05 (normalized)
R boundary error: 0.000e+00 (m)
Z boundary error: 0.000e+00 (m)
Fixed-Psi error: 0.000e+00 (Wb)
Fixed-pressure profile error: 0.000e+00 (Pa)
Fixed-current profile error: 0.000e+00 (A)
Let us again plot the \(|B|\) contours in Boozer coordinates to get a qualitative picture of the solution. Although still not perfect, the optimized equilibrium is clearly more omnigenous compared to the initial one we plotted above. Now almost all of the contours are closed poloidally, except for a few “puddles” near the minimum of the field strength. The omnigenity could probably be further improved by using higher resolutions and running for more iterations.
[13]:
plot_boozer_surface(eq, fieldlines=4) # defaults to the rho=1 surface
[13]:
(<Figure size 576.078x576.078 with 2 Axes>,
<AxesSubplot: title={'center': '$|\\mathbf{B}|~(T)$'}, xlabel='$\\zeta_{Boozer}$', ylabel='$\\theta_{Boozer}$'>)
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/IPython/core/pylabtools.py:152: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
fig.canvas.print_figure(bytes_io, **kw)
We can also plot the omnigenous field that was used as a target during the optimization, as shown below. This is a perfectly omnigenous magnetic field and is physically unrealistic to achieve by an equilibrium, but it represents the “closest” omnigenous field to the optimized solution. In the limit of lower omnigenity errors, this plot and the one from the equilibrium plotted above should approach becoming identical.
Plotting the omnigenous field in Boozer coordinates requires a value for the rotational transform, so we use \(\iota\) from the optimized equilibrium.
[14]:
# compute the rotational transform at rho=1
grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)
iota = eq.compute("iota", grid=grid)["iota"][0]
plot_boozer_surface(field, iota=iota, fieldlines=4)
[14]:
(<Figure size 576.078x576.078 with 2 Axes>,
<AxesSubplot: title={'center': '$|\\mathbf{B}|~(T)$'}, xlabel='$\\zeta_{Boozer}$', ylabel='$\\theta_{Boozer}$'>)
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/IPython/core/pylabtools.py:152: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
fig.canvas.print_figure(bytes_io, **kw)
It is also useful to compare the boundaries of the initial and optimized equilibria, as shown in this plot. This reveals that the optimization improved the omnigenity by adding some torsion and elongation.
[15]:
plot_boundaries((eq0, eq), labels=["Initial", "Optimized"], phi=3, lw=2)
[15]:
(<Figure size 576.078x576.078 with 1 Axes>,
<AxesSubplot: xlabel='$R ~(\\mathrm{m})$', ylabel='$Z ~(\\mathrm{m})$'>)
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/IPython/core/pylabtools.py:152: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
fig.canvas.print_figure(bytes_io, **kw)
We can plot the magnetic well again to see how its shape changed during the optimization.
[16]:
# plot final target well |B|
data_optimal = field.compute("|B|", grid=grid_well)
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
ax.plot(data_initial["eta"], data_initial["|B|"], c="b", lw=2, label="Initial")
ax.plot(data_optimal["eta"], data_optimal["|B|"], c="r", lw=2, label="Optimized")
ax.scatter(
[0, np.pi / 4, np.pi / 2],
field.B_lm[0 : field.M_B] - field.B_lm[field.M_B :],
c="k",
marker=".",
label="Spline Knots",
s=120,
)
ax.set_xticks([-np.pi / 2, 0, np.pi / 2])
ax.set_xticklabels(["$-\pi/2$", "0", "$\pi/2$"])
ax.set_xlabel(r"$\eta$")
ax.set_ylabel("|B|")
ax.legend(loc="upper center")
[16]:
<matplotlib.legend.Legend at 0x7f29938d1960>
/home/dudt/anaconda3/envs/desc/lib/python3.10/site-packages/IPython/core/pylabtools.py:152: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
fig.canvas.print_figure(bytes_io, **kw)
[ ]: