Basic QS Optimization
[1]:
import sys
import os
sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../../../"))
This tutorial demonstrates how to perform optimization with DESC. It will go through two examples:
targeting the “triple product” quasi-symmetry objective \(f_T\) in the full plasma volume
targeting the “two-term” objective \(f_C\) for quasi-helical symmetry at the plasma boundary surface
These examples are adapted from a problem first presented in Rodriguez et al. (2022) and reproduced in Dudt et al. (2023).
If you have access to a GPU, uncomment the following two lines before any DESC or JAX related imports. You should see about an order of magnitude speed improvement with only these two lines of code!
[ ]:
# from desc import set_device
# set_device("gpu")
As mentioned in DESC Documentation on performance tips, one can use compilation cache directory to reduce the compilation overhead time. Note: One needs to create jax-caches folder manually.
[3]:
# import jax
# jax.config.update("jax_compilation_cache_dir", "../jax-caches")
# jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
# jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
[4]:
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams["font.size"] = 20
import desc.io
from desc.grid import LinearGrid, ConcentricGrid
from desc.objectives import (
ObjectiveFunction,
FixBoundaryR,
FixBoundaryZ,
FixPressure,
FixIota,
FixPsi,
AspectRatio,
ForceBalance,
QuasisymmetryBoozer,
QuasisymmetryTwoTerm,
QuasisymmetryTripleProduct,
)
from desc.optimize import Optimizer
from desc.plotting import (
plot_grid,
plot_boozer_modes,
plot_boozer_surface,
plot_qs_error,
plot_boundaries,
plot_boundary,
)
Initial guess
To save time during this tutorial, we will load an equilibrium solution to use as the initial guess for optimization. This equilibrium is already somewhat close to quasi-helical symmetry (QH) but can be improved.
[5]:
# load initial equilibrium
eq_init = desc.io.load("qs_initial_guess.h5")
[6]:
plot_boundary(eq_init, figsize=(8, 8));
There are several plotting functions available in DESC that are useful for visualizing the quasi-symmetry (QS) errors. Let us look at the initial errors before optimization:
[7]:
# plot modes of |B| in Boozer coordinates
plot_boozer_modes(eq_init, num_modes=8, rho=10)
# plot |B| contours in Boozer coordinates on a surface (default is rho=1)
plot_boozer_surface(eq_init)
# plot normalized QS metrics
plot_qs_error(eq_init, helicity=(1, eq_init.NFP), rho=10);
QS metrics
The last plot shows the three different QS errors implemented in DESC. Their definitions and the normalized scalar quantities plotted for each flux surface are shown below:
Boozer Coordinates
This is the “traditional” definition of QS. The drawbacks of this metric are that Boozer coordinates are expensive to compute, and it only provides “global” information about the error on each surface. But it has the nice property of \(\hat{f}_B \in [0,1]\).
\begin{equation} |\mathbf{B}| = B(\psi, M\vartheta_{B} - N\zeta_{B}) \end{equation}
\begin{equation} \mathbf{f}_{B} = \{ B_{mn} ~|~ m/n \neq M/N \} \end{equation}
\begin{equation} \hat{f}_{B} = \frac{|\mathbf{f}_{B}|}{\sqrt{\sum_{m,n} B_{mn}^2}} \end{equation}
Two-Term
This definition is advantageous because it does not require a transformation to Boozer coordinates, and it reveals “local” errors within a flux surface.
\begin{equation} f_{C} = \left( M \iota - N \right) \left( \mathbf{B} \times \nabla \psi \right) \cdot \nabla B - \left( M G + N I \right) \mathbf{B} \cdot \nabla B. \end{equation}
\begin{equation} \mathbf{f}_{C} = \{ f_{C}(\theta_{i},\zeta_{j}) ~|~ i \in [0,2\pi), j \in [0,2\pi/N_{FP}) \} \end{equation}
\begin{equation} \hat{f}_{C} = \frac{\langle |f_C| \rangle}{\langle B \rangle^3}. \end{equation}
Triple Product
This is also a local error metric that can be evaluated without Boozer coordinates. The potential advantages of this definition are that it does not require specifying the helicity (type of quasi-symmetry) and does not assume \(\mathbf{J}\cdot\nabla\psi=0\).
\begin{equation} f_{T} = \nabla \psi \times \nabla B \cdot \nabla \left( \mathbf{B} \cdot \nabla B \right) \end{equation}
\begin{equation} \mathbf{f}_{T} = \{ f_{T}(\theta_{i},\zeta_{j}) ~|~ i \in [0,2\pi), j \in [0,2\pi/N_{FP}) \} \end{equation}
\begin{equation} \hat{f}_{T}(\psi) = \frac{\langle R \rangle^2 \langle |f_{T}| \rangle}{\langle B \rangle^4} \end{equation}
Optimizer
To set up the optimization problem, we need to choose an optimization algorithm. "proximal-lsq-exact" is a custom least-squares routine (with proximal referring to the proximal projection done at each step to ensure force balance i.e. it perturbs the boundary in the direction to optimize the objective function, then resolves the equilibrium to ensure force balance), but there are other options available (see
documentation).
[8]:
optimizer = Optimizer("proximal-lsq-exact")
Specifying constraints
Next we need to define the optimization constraints. In DESC, we explicitly specify the constraints, and all other parameters become free variables for optimization. In this example, we would like the following 4 stellarator-symmetric boundary coefficients to be our only optimization variables:
\(R^{b}_{1,2} \cos(\theta) \cos(2N_{FP}\phi)\)
\(R^{b}_{-1,-2} \sin(\theta) \sin(2N_{FP}\phi)\)
\(Z^{b}_{-1,2} \sin(\theta) \cos(2N_{FP}\phi)\)
\(Z^{b}_{1,-2} \cos(\theta) \sin(2N_{FP}\phi)\)
We will fix all other boundary modes besides these. We will also fix the pressure profile, rotational transform profile, and total toroidal magnetic flux. Finally, we also specify equilibrium force balance as a constraint by including the ForceBalance() objective in the list of constraints.
[9]:
# indices of boundary modes we want to optimize
idx_Rcc = eq_init.surface.R_basis.get_idx(M=1, N=2)
idx_Rss = eq_init.surface.R_basis.get_idx(M=-1, N=-2)
idx_Zsc = eq_init.surface.Z_basis.get_idx(M=-1, N=2)
idx_Zcs = eq_init.surface.Z_basis.get_idx(M=1, N=-2)
print("surface.R_basis.modes is an array of [l,m,n] of the surface modes:")
print(eq_init.surface.R_basis.modes[0:10])
# boundary modes to constrain
R_modes = np.delete(eq_init.surface.R_basis.modes, [idx_Rcc, idx_Rss], axis=0)
Z_modes = np.delete(eq_init.surface.Z_basis.modes, [idx_Zsc, idx_Zcs], axis=0)
eq_qs_T = eq_init.copy() # make a copy of the original one
# constraints
constraints = (
ForceBalance(eq=eq_qs_T), # enforce JxB-grad(p)=0 during optimization
FixBoundaryR(eq=eq_qs_T, modes=R_modes), # fix specified R boundary modes
FixBoundaryZ(eq=eq_qs_T, modes=Z_modes), # fix specified Z boundary modes
FixPressure(eq=eq_qs_T), # fix pressure profile
FixIota(eq=eq_qs_T), # fix rotational transform profile
FixPsi(eq=eq_qs_T), # fix total toroidal magnetic flux
)
surface.R_basis.modes is an array of [l,m,n] of the surface modes:
[[ 0 -8 -4]
[ 0 -7 -4]
[ 0 -6 -4]
[ 0 -5 -4]
[ 0 -4 -4]
[ 0 -3 -4]
[ 0 -2 -4]
[ 0 -1 -4]
[ 0 -8 -3]
[ 0 -7 -3]]
Optimizing for Triple Product QS in Volume
First we are going to optimize for QS using the “triple product” objective, evaluated throughout the plasma volume. We start by creating the grid of collocation points where we want the \(f_T\) errors to be evaluated. We then initialize the appropriate objective function with this grid.
[10]:
# objective
grid_vol = ConcentricGrid(
L=eq_init.L_grid,
M=eq_init.M_grid,
N=eq_init.N_grid,
NFP=eq_init.NFP,
sym=eq_init.sym,
)
plot_grid(grid_vol, figsize=(8, 8))
objective_fT = ObjectiveFunction(QuasisymmetryTripleProduct(eq=eq_qs_T, grid=grid_vol))
We are now ready to run the optimization problem. The syntax is shown below with many of the optimization options that are available – you can try changing these parameters.
[11]:
eq_qs_T, result_T = eq_qs_T.optimize(
objective=objective_fT,
constraints=constraints,
optimizer=optimizer,
ftol=5e-2, # stopping tolerance on the function value
xtol=1e-6, # stopping tolerance on the step size
gtol=1e-6, # stopping tolerance on the gradient
maxiter=50, # maximum number of iterations
options={
"perturb_options": {"order": 2, "verbose": 0}, # use 2nd-order perturbations
"solve_options": {
"ftol": 5e-3,
"xtol": 1e-6,
"gtol": 1e-6,
"verbose": 0,
}, # for equilibrium subproblem
},
copy=False, # copy=False we will overwrite the eq_qs_T object with the optimized result
verbose=3,
)
Building objective: QS triple product
Precomputing transforms
Timer: Precomputing transforms = 635 ms
Timer: Objective build = 1.22 sec
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 202 ms
Timer: Objective build = 236 ms
Timer: Objective build = 1.55 ms
Timer: Eq Update LinearConstraintProjection build = 4.11 sec
Timer: Proximal projection build = 19.9 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed iota
Building objective: fixed Psi
Timer: Objective build = 894 ms
Timer: LinearConstraintProjection build = 1.79 sec
Number of parameters: 4
Number of objectives: 1377
Timer: Initializing the optimization = 22.7 sec
Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations : 251
Maximum Allowed Total Δx Norm : inf
Scaled Termination : True
Trust Region Method : qr
Initial Trust Radius : 9.423e+00
Maximum Trust Radius : inf
Minimum Trust Radius : 2.220e-16
Trust Radius Increase Ratio : 2.000e+00
Trust Radius Decrease Ratio : 2.500e-01
Trust Radius Increase Threshold : 7.500e-01
Trust Radius Decrease Threshold : 2.500e-01
------------------------------------------------------------
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 4.463e+02 2.921e+01
1 2 1.446e+02 3.017e+02 2.073e-02 1.043e+01
2 3 2.249e+01 1.221e+02 3.599e-02 1.847e+00
3 4 3.440e+00 1.905e+01 4.339e-02 3.022e-01
4 5 4.338e-01 3.006e+00 3.382e-02 5.883e-02
5 6 3.474e-02 3.990e-01 2.079e-02 1.151e-02
6 7 1.076e-02 2.397e-02 1.202e-02 1.023e-03
7 8 1.061e-02 1.515e-04 1.043e-03 2.165e-05
Optimization terminated successfully.
`ftol` condition satisfied. (ftol=5.00e-02)
Current function value: 1.061e-02
Total delta_x: 1.071e-01
Iterations: 7
Function evaluations: 8
Jacobian evaluations: 8
Timer: Solution time = 43.5 sec
Timer: Avg time per step = 5.43 sec
==============================================================================================================
Start --> End
Total (sum of squares): 4.465e+02 --> 1.061e-02,
Maximum absolute Quasi-symmetry error: 4.730e+02 --> 2.086e+00 (T^4/m^2)
Minimum absolute Quasi-symmetry error: 4.422e-04 --> 2.514e-05 (T^4/m^2)
Average absolute Quasi-symmetry error: 1.117e+01 --> 1.081e-01 (T^4/m^2)
Maximum absolute Quasi-symmetry error: 1.104e+01 --> 4.869e-02 (normalized)
Minimum absolute Quasi-symmetry error: 1.032e-05 --> 5.868e-07 (normalized)
Average absolute Quasi-symmetry error: 2.607e-01 --> 2.523e-03 (normalized)
Maximum absolute Force error: 1.516e+05 --> 1.611e+04 (N)
Minimum absolute Force error: 2.431e+00 --> 2.797e-01 (N)
Average absolute Force error: 5.922e+03 --> 7.447e+02 (N)
Maximum absolute Force error: 1.250e-02 --> 1.329e-03 (normalized)
Minimum absolute Force error: 2.005e-07 --> 2.306e-08 (normalized)
Average absolute Force error: 4.884e-04 --> 6.141e-05 (normalized)
R boundary error: 5.649e-17 --> 7.758e-18 (m)
Z boundary error: 9.866e-18 --> 6.942e-18 (m)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed iota profile error: 0.000e+00 --> 0.000e+00 (dimensionless)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
==============================================================================================================
If you would like to store the objective values before and after optimization, you can access them through result that is returned by optimization.
[12]:
# info["Objective values"] is a dictionary of the objective values
for key in result_T["Objective values"].keys():
if isinstance(result_T["Objective values"][key], dict):
for subkey in result_T["Objective values"][key].keys():
print(key, subkey, result_T["Objective values"][key][subkey])
# if there are multiple instances of the same objective type,
# there will be a list of dictionaries, each storing the values of the objective
elif isinstance(result_T["Objective values"][key], list):
for i in range(len(result_T["Objective values"][key])):
for subkey in result_T["Objective values"][key][i].keys():
print(key, subkey, result_T["Objective values"][key][i][subkey])
Total (sum of squares) f 0.01061003559091122
Total (sum of squares) f0 446.51269799795404
Quasi-symmetry error: f_max 2.0861726187482885
Quasi-symmetry error: f0_max 473.01432460839777
Quasi-symmetry error: f_min 2.514216931882426e-05
Quasi-symmetry error: f0_min 0.0004421575033381454
Quasi-symmetry error: f_mean 0.10811726291657031
Quasi-symmetry error: f0_mean 11.171841092895564
Quasi-symmetry error: f_max_norm 0.0486910939174104
Quasi-symmetry error: f0_max_norm 11.04011465628716
Quasi-symmetry error: f_min_norm 5.868161227831812e-07
Quasi-symmetry error: f0_min_norm 1.0319919036346524e-05
Quasi-symmetry error: f_mean_norm 0.0025234478467667295
Quasi-symmetry error: f0_mean_norm 0.2607498339283862
Force error: f_max 16114.099037013826
Force error: f0_max 151574.9774697879
Force error: f_min 0.279676979582177
Force error: f0_min 2.431091953493361
Force error: f_mean 744.6671866516118
Force error: f0_mean 5922.427405369982
Force error: f_max_norm 0.0013288886317115976
Force error: f0_max_norm 0.01250000164134959
Force error: f_min_norm 2.306424689735961e-08
Force error: f0_min_norm 2.0048595036074404e-07
Force error: f_mean_norm 6.141080283030013e-05
Force error: f0_mean_norm 0.000488407476772704
R boundary error: f 7.757919228897728e-18
R boundary error: f0 5.648891010403884e-17
Z boundary error: f 6.942281208919283e-18
Z boundary error: f0 9.86564403615321e-18
Fixed pressure profile error: f 0.0
Fixed pressure profile error: f0 0.0
Fixed iota profile error: f 0.0
Fixed iota profile error: f0 0.0
Fixed Psi error: f 0.0
Fixed Psi error: f0 0.0
Let us plot the QS errors again to see how well the optimization worked:
[13]:
# |B| contours at rho=1 surface
plot_boozer_surface(eq_qs_T);
[14]:
# compare f_T & f_C before (o) vs after (x) optimization
fig, ax = plot_qs_error(
eq_init, helicity=(1, eq_init.NFP), fB=False, legend=False, rho=10
)
plot_qs_error(
eq_qs_T, helicity=(1, eq_init.NFP), fB=False, ax=ax, marker=["x", "x"], rho=10
);
[15]:
# compare f_B before (o) vs after (x) optimization
fig, ax = plot_qs_error(
eq_init, helicity=(1, eq_init.NFP), fT=False, fC=False, legend=False, rho=10
)
plot_qs_error(
eq_qs_T, helicity=(1, eq_init.NFP), fT=False, fC=False, ax=ax, marker=["x"], rho=10
);
Optimizing for Two-Term QH at Boundary Surface
Next we are going to optimize for QH using the “two-term” objective, evaluated at the \(\rho=1\) flux surface. Again we need to create the grid of collocation points where we want the \(f_C\) errors to be evaluated. We then initialize the appropriate objective function with this grid.
[16]:
grid_rho1 = LinearGrid(
M=eq_init.M_grid,
N=eq_init.N_grid,
NFP=eq_init.NFP,
sym=eq_init.sym,
rho=np.array(1.0),
)
plot_grid(grid_rho1, figsize=(8, 8))
eq_qs_C = eq_init.copy()
# constraints
constraints = (
ForceBalance(eq=eq_qs_C), # enforce JxB-grad(p)=0 during optimization
FixBoundaryR(eq=eq_qs_C, modes=R_modes), # fix specified R boundary modes
FixBoundaryZ(eq=eq_qs_C, modes=Z_modes), # fix specified Z boundary modes
FixPressure(eq=eq_qs_C), # fix pressure profile
FixIota(eq=eq_qs_C), # fix rotational transform profile
FixPsi(eq=eq_qs_C), # fix total toroidal magnetic flux
)
# two-term objective
objective_fC = ObjectiveFunction(
QuasisymmetryTwoTerm(eq=eq_qs_C, grid=grid_rho1, helicity=(1, eq_init.NFP))
)
Now we can re-run the optimization from the same intial guess as before, but using this different objective function. We can reuse the same optimizer from the first run.
[17]:
eq_qs_C, result_C = eq_qs_C.optimize(
objective=objective_fC,
constraints=constraints,
optimizer=optimizer,
ftol=1e-2,
xtol=1e-6,
gtol=1e-6,
maxiter=50,
options={
"perturb_options": {"order": 2, "verbose": 0},
"solve_options": {"ftol": 1e-2, "xtol": 1e-6, "gtol": 1e-6, "verbose": 0},
},
copy=False,
verbose=3,
)
Building objective: QS two-term
Precomputing transforms
Timer: Precomputing transforms = 54.3 ms
Timer: Objective build = 107 ms
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 60.2 ms
Timer: Objective build = 69.1 ms
Timer: Objective build = 1.53 ms
Timer: Eq Update LinearConstraintProjection build = 143 ms
Timer: Proximal projection build = 1.79 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed iota
Building objective: fixed Psi
Timer: Objective build = 166 ms
Timer: LinearConstraintProjection build = 161 ms
Number of parameters: 4
Number of objectives: 289
Timer: Initializing the optimization = 2.17 sec
Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations : 251
Maximum Allowed Total Δx Norm : inf
Scaled Termination : True
Trust Region Method : qr
Initial Trust Radius : 5.985e+01
Maximum Trust Radius : inf
Minimum Trust Radius : 2.220e-16
Trust Radius Increase Ratio : 2.000e+00
Trust Radius Decrease Ratio : 2.500e-01
Trust Radius Increase Threshold : 7.500e-01
Trust Radius Decrease Threshold : 2.500e-01
------------------------------------------------------------
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 7.636e+04 3.235e+02
1 2 5.328e+04 2.308e+04 1.082e-02 2.131e+02
2 3 2.706e+04 2.622e+04 2.240e-02 9.898e+01
3 4 8.870e+03 1.818e+04 5.144e-02 2.247e+01
4 5 5.963e+03 2.907e+03 5.411e-02 5.528e+00
5 6 5.903e+03 6.015e+01 8.265e-03 9.483e-01
6 7 5.897e+03 5.393e+00 2.423e-03 1.802e-01
Optimization terminated successfully.
`ftol` condition satisfied. (ftol=1.00e-02)
Current function value: 5.897e+03
Total delta_x: 1.174e-01
Iterations: 6
Function evaluations: 7
Jacobian evaluations: 7
Timer: Solution time = 14.5 sec
Timer: Avg time per step = 2.08 sec
==============================================================================================================
Start --> End
Total (sum of squares): 7.631e+04 --> 5.897e+03,
Maximum absolute Quasi-symmetry (1,4) two-term error: 1.353e+02 --> 2.011e+01 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error: 1.562e-02 --> 5.200e-02 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error: 2.315e+01 --> 7.302e+00 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error: 4.015e+01 --> 5.968e+00 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error: 4.635e-03 --> 1.543e-02 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error: 6.870e+00 --> 2.167e+00 (normalized)
Maximum absolute Force error: 1.516e+05 --> 2.183e+04 (N)
Minimum absolute Force error: 2.431e+00 --> 8.773e-01 (N)
Average absolute Force error: 5.922e+03 --> 9.689e+02 (N)
Maximum absolute Force error: 1.250e-02 --> 1.801e-03 (normalized)
Minimum absolute Force error: 2.005e-07 --> 7.235e-08 (normalized)
Average absolute Force error: 4.884e-04 --> 7.990e-05 (normalized)
R boundary error: 5.649e-17 --> 7.758e-18 (m)
Z boundary error: 9.866e-18 --> 6.942e-18 (m)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed iota profile error: 0.000e+00 --> 0.000e+00 (dimensionless)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
==============================================================================================================
Let us plot the QS errors again for this case. Was one objective more effective than the other at improving quasi-symmetry for this example?
[18]:
# |B| contours at rho=1 surface
plot_boozer_surface(eq_qs_C);
[19]:
# compare f_T & f_C before (o) vs after (x) optimization
fig, ax = plot_qs_error(
eq_init, helicity=(1, eq_init.NFP), fB=False, legend=False, rho=10
)
plot_qs_error(
eq_qs_C, helicity=(1, eq_init.NFP), fB=False, ax=ax, marker=["x", "x"], rho=10
);
[20]:
# compare f_B before (o) vs after (x) optimization
fig, ax = plot_qs_error(
eq_init, helicity=(1, eq_init.NFP), fT=False, fC=False, legend=False, rho=10
)
plot_qs_error(
eq_qs_C, helicity=(1, eq_init.NFP), fT=False, fC=False, ax=ax, marker=["x"], rho=10
);
Combining Multiple Objectives
It is very easy in DESC to combine multiple optimization objectives with relative weights between them. Here is an example of how to create an objective optimizing for both \(f_B\) and a target aspect ratio:
[21]:
objective = ObjectiveFunction(
(
QuasisymmetryBoozer(eq=eq_init, helicity=(1, eq_init.NFP), weight=1e-2),
AspectRatio(eq=eq_init, target=6, weight=1e1),
)
)