Deflation Methods for Equilibrium and Optimization
This notebook introduces deflation methods, and shows how to use them in two examples for equilibrium solving and stellarator optimization.
Background
Nonlinear system rootfinding \(\mathbf{F}(\mathbf{x})=0\) can exhibit multiple roots (simplest case being polynomial rootfinding with scalar \(x\)). Newton’s method is often used to rootfind, but which root is found depends on the initial guess for the method. Deflation was developed to try to allow convergence to multiple roots from the same initial guess for the rootfinding method. This is done by finding a first root, then “deflating away” that root by multiplying the original system by a deflation operator and then rootfinding on that augmented system. Concretely, let
\[\begin{equation}
\mathbf{F}(\mathbf{x}) = 0
\end{equation}\]
where \(\mathbf{F}\) is the nonlinear system of equations being solved and \(\mathbf{x}\) is the state vector of unknowns. The deflation algorithm works by attempting to find a set of distinct solutions \(\{\mathbf{x}_1^*, \mathbf{x}_2^*, ...\}\). Suppose a solution is found \(\mathbf{x}_1^*\) to the above equation (by application of Newton’s method, or some other way such as minimization of the residual). To find more solutions, one may apply a deflation operator to modify the system being solved, instead solving
\[\begin{equation}
M(\mathbf{x};\mathbf{x}_1^*)\mathbf{F}(\mathbf{x})=0
\end{equation}\]
where
\[\begin{equation}
M(\mathbf{x};\mathbf{x}_1^*) = \left(\frac{1}{| \mathbf{x}-\mathbf{x}_1^* |^p_2}\right) + \sigma
\end{equation}\]
is the deflation operator with power and shift parameters \(p>0\) and \(\sigma>0\), respectively, and \(| \cdot |_2\) is the Euclidean norm. Solving the deflated system encourages the nonlinear solver to not return the known solution \(\mathbf{x}_1^*\) (from the same initial guess) but rather a distinct solution \(\mathbf{x}_2^*\) (if the nonlinear solver converges, which is not always guaranteed). To find further solutions, the system can then be deflated again using known solutions \(\{\mathbf{x}_1^*,...,\mathbf{x}_N^*\}\) by simply applying the deflation operator repeatedly:
\[\begin{equation}
\prod_{i=1}^N M(\mathbf{x};\mathbf{x}_i^*)\mathbf{F}(\mathbf{x}) = 0
\end{equation}\]
With this deflation method, one may find multiple solutions to a given system of equations. In DESC, it can be applied to find multiple solutions to various equilibrium problems , where the nonlinear system being solved is the ideal magneto-hydrodynamic (MHD) equilibrium force balance, and minimization techniques are applied to find the solutions to the unmodified and modified system.
This can also be applied to nonconvex optimization (which many stellarator optimization problems are). One can either use the above form to modify the objective function being minimized to with the deflation operator, or one can apply a nonlinear constraint to enforce the deflation (either as an additional term to be penalized or as a nonlinear constraint through a constrained optimizer like lsq-auglag). The constraint would be of the form
\[\begin{equation}
M(\mathbf{x};\mathbf{x_1}) \leq r
\end{equation}\]
Where \(r>0\) is some positive constant. This constraint then would encourage the optimizer to not find minima similar to already-found minima.
Near Axis Constrained Equilibria
Near-axis-constrained equilibrium is an under-constrained problem which has multiple solutions. We can apply deflation to find these solutions. In this section, we will apply this to a QA NAE solution to find three different global equilibria which match the same near-axis properties
[1]:
# uncomment these if you have a GPU for a speedup!
# from desc import set_device
# set_device("gpu")
[2]:
from desc.objectives import (
DeflationOperator,
ForceBalance,
get_NAE_constraints,
get_fixed_boundary_constraints,
ObjectiveFunction,
)
from desc.io import load
# must have installed pyQsc with `pip install qsc` in order to run this notebook!
from qsc import Qsc
import numpy as np
import matplotlib.pyplot as plt
import os
from desc.equilibrium import Equilibrium
from desc.plotting import (
plot_comparison,
plot_fsa,
plot_section,
plot_surfaces,
plot_qs_error,
plot_1d,
plot_boundaries,
plot_boozer_surface,
)
plt.rcParams.update({"font.size": 20})
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]:
qsc_eq = Qsc.from_paper("precise QA", order=1)
ntheta = 12 # only using M=5 for our equilibrium, so we don't need a large ntheta
r = 0.25
eq_fit = Equilibrium.from_near_axis(
qsc_eq, # the Qsc equilibrium object
r=r, # the finite radius (m) at which to evaluate the Qsc surface to use as the DESC boundary
L=5, # DESC radial resolution
M=5, # DESC poloidal resolution
N=6, # DESC toroidal resolution
ntheta=ntheta,
)
[5]:
def solve_deflated(eq_fit, eqs_to_deflate=[], power=2, sigma=1):
# Helper Function for using deflation to find multiple NAE-constrained equilibria
eq = eq_fit.copy()
constraints = get_NAE_constraints(
eq, qsc_eq, order=1
) # get the first-order NAE constraints
# create the force balance objective, which we will use both in the deflation and to solve the NAE-constrained and final fixed-boundary solve
force_obj_bare = ForceBalance(eq)
force_obj = ObjectiveFunction(force_obj_bare)
maxiter_deflation = 15 if eqs_to_deflate[0] is not None else 50
xtol_deflation = 1e-6 if eqs_to_deflate[0] is not None else 1e-7
if eqs_to_deflate is not None:
# solve deflated
deflated_force_obj = ObjectiveFunction(
DeflationOperator(
thing=eq,
things_to_deflate=eqs_to_deflate,
objective=force_obj_bare,
power=power,
sigma=sigma,
# params_to_deflate is a dictionary
params_to_deflate_with={"Rb_lmn": True, "Zb_lmn": True},
)
)
eq.solve(
verbose=1,
ftol=1e-3,
objective=deflated_force_obj,
# only do a few iterations, to just get into a different basin of attraction for another NAE-constrained equilibrium solution
maxiter=maxiter_deflation,
xtol=xtol_deflation,
constraints=constraints,
)
# then solve just NAE-constrained equilibrium solve so it falls in new local minimum
eq.solve(
verbose=1,
ftol=1e-3,
objective=force_obj,
maxiter=50,
xtol=1e-7,
constraints=constraints,
)
# then solve fixed boundary to confirm is indeed equilibrium in usual sense
eq.solve(ftol=1e-3, maxiter=50, objective=force_obj)
return eq
With the above helper function, we can call it in a loop as many times as we’d like to get any number of NAE-constrained equilibrium solutions. For brevity’s sake this notebook will only run it three times.
[6]:
# this can be increased if more local minima are desired, though
# there is no guarantee that every solution found is "good",
# e.g. convergence to a new good solution is not guaranteed
N_deflation = 3
# we can either append to the list of equilibria we are deflating with as we go, or
# we can make the full list length we expect and fill it with None to start. The
# latter is advantageous as it will avoid recompilation of our deflation objective
# in subsequent iterations, so we will use that approach.
eqs = [None] * N_deflation
ts = []
for i in range(N_deflation):
print("#" * 20)
print(f"{i=}")
print("#" * 20)
eqs[i] = solve_deflated(eq_fit, eqs, sigma=500).copy()
####################
i=0
####################
Building objective: Deflated force
Precomputing transforms
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet current
Building objective: Fix Near Axis R Behavior
Building objective: Fix Near Axis Z Behavior
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
Number of parameters: 369
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
Warning: Maximum number of iterations has been exceeded.
Current function value: 1.734e-08
Total delta_x: 2.382e-01
Iterations: 50
Function evaluations: 56
Jacobian evaluations: 51
==============================================================================================================
Start --> End
Total (sum of squares): 1.857e+00 --> 1.734e-08,
Maximum absolute Deflated Force error: 1.797e+06 --> 8.502e+01 (N)
Minimum absolute Deflated Force error: 5.434e+00 --> 2.819e-03 (N)
Average absolute Deflated Force error: 3.535e+04 --> 4.810e+00 (N)
Maximum absolute Deflated Force error: 9.697e-01 --> 4.587e-05 (normalized)
Minimum absolute Deflated Force error: 2.932e-06 --> 1.521e-09 (normalized)
Average absolute Deflated Force error: 1.907e-02 --> 2.595e-06 (normalized)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
Fixed-Near-Axis R Behavior error: 3.993e-07 --> 1.929e-15 (m)
Fixed-Near-Axis Z Behavior error: 3.459e-07 --> 1.099e-15 (m)
==============================================================================================================
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
Number of parameters: 369
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
Warning: Maximum number of iterations has been exceeded.
Current function value: 5.840e-09
Total delta_x: 6.564e-02
Iterations: 50
Function evaluations: 63
Jacobian evaluations: 51
==============================================================================================================
Start --> End
Total (sum of squares): 1.734e-08 --> 5.840e-09,
Maximum absolute Force error: 8.502e+01 --> 4.639e+01 (N)
Minimum absolute Force error: 2.819e-03 --> 6.048e-03 (N)
Average absolute Force error: 4.810e+00 --> 2.830e+00 (N)
Maximum absolute Force error: 4.587e-05 --> 2.503e-05 (normalized)
Minimum absolute Force error: 1.521e-09 --> 3.263e-09 (normalized)
Average absolute Force error: 2.595e-06 --> 1.527e-06 (normalized)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
Fixed-Near-Axis R Behavior error: 1.929e-15 --> 1.894e-15 (m)
Fixed-Near-Axis Z Behavior error: 1.099e-15 --> 1.080e-15 (m)
==============================================================================================================
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet 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
Number of parameters: 265
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
`gtol` condition satisfied. (gtol=1.00e-08)
Current function value: 4.966e-09
Total delta_x: 1.955e-04
Iterations: 2
Function evaluations: 3
Jacobian evaluations: 3
==============================================================================================================
Start --> End
Total (sum of squares): 5.840e-09 --> 4.966e-09,
Maximum absolute Force error: 4.639e+01 --> 4.140e+01 (N)
Minimum absolute Force error: 6.048e-03 --> 3.649e-03 (N)
Average absolute Force error: 2.830e+00 --> 2.616e+00 (N)
Maximum absolute Force error: 2.503e-05 --> 2.234e-05 (normalized)
Minimum absolute Force error: 3.263e-09 --> 1.969e-09 (normalized)
Average absolute Force error: 1.527e-06 --> 1.411e-06 (normalized)
R boundary error: 0.000e+00 --> 8.050e-18 (m)
Z boundary error: 0.000e+00 --> 2.824e-17 (m)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
==============================================================================================================
####################
i=1
####################
Building objective: Deflated force
Precomputing transforms
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet current
Building objective: Fix Near Axis R Behavior
Building objective: Fix Near Axis Z Behavior
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
Number of parameters: 369
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
Warning: Maximum number of iterations has been exceeded.
Current function value: 2.277e+00
Total delta_x: 3.025e-01
Iterations: 15
Function evaluations: 21
Jacobian evaluations: 16
==============================================================================================================
Start --> End
Total (sum of squares): 8.403e+05 --> 2.277e+00,
Maximum absolute Deflated Force error: 1.209e+09 --> 9.204e+05 (N)
Minimum absolute Deflated Force error: 3.655e+03 --> 3.984e+01 (N)
Average absolute Deflated Force error: 2.378e+07 --> 5.489e+04 (N)
Maximum absolute Deflated Force error: 6.524e+02 --> 4.966e-01 (normalized)
Minimum absolute Deflated Force error: 1.972e-03 --> 2.149e-05 (normalized)
Average absolute Deflated Force error: 1.283e+01 --> 2.962e-02 (normalized)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
Fixed-Near-Axis R Behavior error: 3.993e-07 --> 1.927e-15 (m)
Fixed-Near-Axis Z Behavior error: 3.459e-07 --> 1.080e-15 (m)
==============================================================================================================
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
Number of parameters: 369
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
Warning: Maximum number of iterations has been exceeded.
Current function value: 3.340e-08
Total delta_x: 2.064e-01
Iterations: 50
Function evaluations: 53
Jacobian evaluations: 51
==============================================================================================================
Start --> End
Total (sum of squares): 5.055e-06 --> 3.340e-08,
Maximum absolute Force error: 1.372e+03 --> 1.228e+02 (N)
Minimum absolute Force error: 5.936e-02 --> 7.894e-04 (N)
Average absolute Force error: 8.180e+01 --> 6.657e+00 (N)
Maximum absolute Force error: 7.400e-04 --> 6.625e-05 (normalized)
Minimum absolute Force error: 3.203e-08 --> 4.259e-10 (normalized)
Average absolute Force error: 4.413e-05 --> 3.592e-06 (normalized)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
Fixed-Near-Axis R Behavior error: 1.927e-15 --> 1.918e-15 (m)
Fixed-Near-Axis Z Behavior error: 1.080e-15 --> 1.076e-15 (m)
==============================================================================================================
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet 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
Number of parameters: 265
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
`gtol` condition satisfied. (gtol=1.00e-08)
Current function value: 3.193e-08
Total delta_x: 5.335e-04
Iterations: 2
Function evaluations: 3
Jacobian evaluations: 3
==============================================================================================================
Start --> End
Total (sum of squares): 3.340e-08 --> 3.193e-08,
Maximum absolute Force error: 1.228e+02 --> 1.108e+02 (N)
Minimum absolute Force error: 7.894e-04 --> 3.648e-03 (N)
Average absolute Force error: 6.657e+00 --> 6.531e+00 (N)
Maximum absolute Force error: 6.625e-05 --> 5.979e-05 (normalized)
Minimum absolute Force error: 4.259e-10 --> 1.968e-09 (normalized)
Average absolute Force error: 3.592e-06 --> 3.524e-06 (normalized)
R boundary error: 0.000e+00 --> 1.238e-18 (m)
Z boundary error: 0.000e+00 --> 3.603e-18 (m)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
==============================================================================================================
####################
i=2
####################
Building objective: Deflated force
Precomputing transforms
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet current
Building objective: Fix Near Axis R Behavior
Building objective: Fix Near Axis Z Behavior
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
Number of parameters: 369
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
Warning: Maximum number of iterations has been exceeded.
Current function value: 4.990e+06
Total delta_x: 2.626e-01
Iterations: 15
Function evaluations: 23
Jacobian evaluations: 16
==============================================================================================================
Start --> End
Total (sum of squares): 3.956e+11 --> 4.990e+06,
Maximum absolute Deflated Force error: 8.296e+11 --> 9.106e+08 (N)
Minimum absolute Deflated Force error: 2.508e+06 --> 9.037e+04 (N)
Average absolute Deflated Force error: 1.632e+10 --> 7.452e+07 (N)
Maximum absolute Deflated Force error: 4.476e+05 --> 4.913e+02 (normalized)
Minimum absolute Deflated Force error: 1.353e+00 --> 4.876e-02 (normalized)
Average absolute Deflated Force error: 8.803e+03 --> 4.021e+01 (normalized)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
Fixed-Near-Axis R Behavior error: 3.993e-07 --> 1.910e-15 (m)
Fixed-Near-Axis Z Behavior error: 3.459e-07 --> 1.068e-15 (m)
==============================================================================================================
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
Number of parameters: 369
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
Warning: Maximum number of iterations has been exceeded.
Current function value: 3.608e-08
Total delta_x: 2.472e-01
Iterations: 50
Function evaluations: 59
Jacobian evaluations: 51
==============================================================================================================
Start --> End
Total (sum of squares): 4.386e-05 --> 3.608e-08,
Maximum absolute Force error: 2.700e+03 --> 1.462e+02 (N)
Minimum absolute Force error: 2.679e-01 --> 1.083e-03 (N)
Average absolute Force error: 2.209e+02 --> 6.846e+00 (N)
Maximum absolute Force error: 1.457e-03 --> 7.886e-05 (normalized)
Minimum absolute Force error: 1.446e-07 --> 5.843e-10 (normalized)
Average absolute Force error: 1.192e-04 --> 3.694e-06 (normalized)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
Fixed-Near-Axis R Behavior error: 1.910e-15 --> 1.900e-15 (m)
Fixed-Near-Axis Z Behavior error: 1.068e-15 --> 1.076e-15 (m)
==============================================================================================================
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet 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
Number of parameters: 265
Number of objectives: 1800
Starting optimization
Using method: lsq-exact
`gtol` condition satisfied. (gtol=1.00e-08)
Current function value: 3.331e-08
Total delta_x: 5.478e-04
Iterations: 2
Function evaluations: 3
Jacobian evaluations: 3
==============================================================================================================
Start --> End
Total (sum of squares): 3.608e-08 --> 3.331e-08,
Maximum absolute Force error: 1.462e+02 --> 1.299e+02 (N)
Minimum absolute Force error: 1.083e-03 --> 9.343e-04 (N)
Average absolute Force error: 6.846e+00 --> 6.447e+00 (N)
Maximum absolute Force error: 7.886e-05 --> 7.009e-05 (normalized)
Minimum absolute Force error: 5.843e-10 --> 5.041e-10 (normalized)
Average absolute Force error: 3.694e-06 --> 3.478e-06 (normalized)
R boundary error: 0.000e+00 --> 3.880e-18 (m)
Z boundary error: 0.000e+00 --> 1.434e-17 (m)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
==============================================================================================================
Now that we have our two solutions, we can check them and see how they differ by comparing their boundaries:
We see from their force errors that each are good equilibria with normalized force errors <1%. They also have different rotational transform profiles, though notably have the same near-axis properties (as expected given that all are NAE-constrained equilibria constrained by the same NAE solution):
[7]:
cs = ["tab:blue", "tab:orange", "tab:green", "c"]
plt.rcParams.update({"font.size": 18})
ax = None
ax2 = None
for i in range(len(eqs)):
fig, ax = plot_fsa(eqs[i], "|F|_normalized", log=True, ax=ax, color=cs[i])
fig2, ax2 = plot_1d(eqs[i], "iota", ax=ax2, color=cs[i])
ax2.axhline(qsc_eq.iota, c="k", linestyle="--", label="NAE", lw=2)
ax2.set_xlim([0, 1.0])
ax2.legend()
[7]:
<matplotlib.legend.Legend at 0x33c678830>
[8]:
eq = eq_fit
fig, axes = plt.subplots(1, 2, figsize=(12, 8))
phis = np.linspace(0, np.pi / eq.NFP, 2, endpoint=True)
fig, ax = plot_boundaries(eqs, lw=2, color=cs, legend=False, phi=phis[0], ax=axes[0])
fig, ax = plot_boundaries(eqs, lw=2, color=cs, legend=True, phi=phis[1], ax=axes[1])
for i, axx in enumerate(axes.flatten()):
axx.set_aspect("equal")
axx.set_title(r"$\phi\cdot N_{FP}/2\pi=$" + rf"${phis[i]*eq.NFP/2/np.pi:1.3f}$")
[9]:
ax = None
for i in range(len(eqs)):
fig, ax, dataQS = plot_qs_error(
eqs[i],
labels=[f"{i}"] * 3,
fC=False,
fT=False,
fB=True,
log=True,
color=[cs[i]] * 3,
ax=ax,
legend=False,
marker=[None] * 3,
return_data=True,
lw=[3] * 3,
)
ax.set_xscale("log")
ax.set_yscale("log")
ax.plot(
dataQS["rho"],
dataQS["rho"] ** 2 * (np.max(dataQS["f_B"]) * 0.3),
"k--",
lw=3,
)
ax.text(0.2, 2e-4, r"$\sim \rho^2$", size=18)
ax.legend()
ax.set_ylabel(r"$\hat{f}_B$", fontsize=16);
QH Optimization
Now, we will perform a stage-one QH optimization three times: the first time will result in some QH solution, while the second and third times we will use the DeflationOperator as a nonlinear constraint to penalize the problem from finding the same solutions again. The constraint will be of the form
\begin{equation} M(\mathbf{x};\{\mathbf{x}_i^*\})\leq r \end{equation}
where \(r\) is the upper bound on the deflation metric cost (the lower the deflation metric cost, the farther the solution is away from the already-found solutions as measured by the deflation operator detailed above) and \(\{\mathbf{x}_i^*\}\) is the set of already-found local minima which we wish to avoid.
[10]:
from desc.continuation import solve_continuation_automatic
from desc.equilibrium import EquilibriaFamily, Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.grid import LinearGrid
from desc.objectives import (
AspectRatio,
FixBoundaryR,
FixBoundaryZ,
FixCurrent,
FixPressure,
FixPsi,
ForceBalance,
ObjectiveFunction,
QuasisymmetryTwoTerm,
RotationalTransform,
)
from desc.optimize import Optimizer
from desc.compat import rotate_zeta
# create initial equilibrium, just a circular torus
surf = FourierRZToroidalSurface(
R_lmn=[1, 0.125, 0.0],
Z_lmn=[-0.125, -0.0],
modes_R=[[0, 0], [1, 0], [0, 1]],
modes_Z=[[-1, 0], [0, -1]],
NFP=4,
)
eq = Equilibrium(L=4, M=3, N=4, Psi=0.04, surface=surf)
[11]:
# grid to use to evaluate QH two-term errors at
grid = LinearGrid(M=eq.M, N=eq.N, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True)
def run_and_deflate(eq, eqs_deflate, r=5):
# Helper Function for using deflation to find multiple QH equilibria
this_eq = eq.copy()
objs = (
QuasisymmetryTwoTerm(
eq=this_eq,
helicity=(1, eq.NFP),
grid=grid,
),
)
constraints = (
ForceBalance(eq=this_eq), # J x B - grad(p) = 0
FixBoundaryR(
eq=this_eq, modes=[[0, 0, 0]]
), # we fix the R_000 mode corresponding to the major radius
FixPressure(eq=this_eq),
FixCurrent(eq=this_eq),
FixPsi(eq=this_eq),
AspectRatio(eq=this_eq, target=8),
RotationalTransform(this_eq, bounds=(0.7, 3), loss_function="mean"),
)
if len(eqs_deflate) > 0:
constraints += (
DeflationOperator(
this_eq,
eqs_deflate,
sigma=0,
bounds=(0, r),
power=2,
),
) # by default, the part of the state used in the deflation is the entire state i.e. Rlmn Zlmn Llmn etc,
# every part of the eq.params_dict is included in the state.
objective = ObjectiveFunction(objs)
optimizer = Optimizer("lsq-auglag")
this_eq.optimize(
objective=objective,
constraints=constraints,
optimizer=optimizer,
maxiter=300,
ftol=1e-2,
verbose=1,
x_scale="ess",
)
this_eq.solve(ftol=5e-3, verbose=0)
return this_eq
[12]:
# this can be increased if more local minima are desired, though
# there is no guarantee that every solution found is "good",
# e.g. convergence to a new good solution is not guaranteed
N_deflation = 3
eqs = (
[None] * N_deflation * 2
) # we will have 2*N_deflation total eqs to deflate since we are
# including the potential degenerate rotations as well
for i in range(0, N_deflation):
print("#" * 20)
print(f"{i=}")
print("#" * 20)
eqs[i] = run_and_deflate(eq, eqs).copy()
# also add to the list of deflated solutions the other possible parametrization
# of the same physical equilibrium in stellarator symmetry, which is the pi/NFP
# toroidally rotated equilibrium
# to avoid the exact same equilibrium being found
# This is unneeded for the NAE-constrained equilibria in the prior section
# because the NAE constraints on the axis and first-order behavior remove this possible
# parametrization degeneracy
eqs[-(i + 1)] = rotate_zeta(eqs[i].copy(), angle=np.pi / eq.NFP)
####################
i=0
####################
Building objective: QS two-term
Precomputing transforms
Building objective: lcfs R
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
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
Building objective: force
Precomputing transforms
Building objective: aspect ratio
Precomputing transforms
Building objective: rotational transform
Precomputing transforms
Building objective: Deflation
Number of parameters: 173
Number of objectives: 108
Number of equality constraints: 715
Number of inequality constraints: 2
Starting optimization
Using method: lsq-auglag
Warning: Maximum number of iterations has been exceeded.
Current function value: 2.338e-06
Constraint violation: 2.572e-04
Total delta_x: 5.764e-01
Iterations: 300
Function evaluations: 559
Jacobian evaluations: 300
==============================================================================================================
Start --> End
Total (sum of squares): 1.124e-27 --> 2.338e-06,
Maximum absolute Quasi-symmetry (1,4) two-term error: 2.576e-15 --> 4.689e-04 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error: 8.061e-18 --> 7.896e-07 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error: 8.535e-16 --> 4.739e-05 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error: 2.438e-15 --> 4.437e-04 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error: 7.628e-18 --> 7.472e-07 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error: 8.076e-16 --> 4.485e-05 (normalized)
Maximum absolute Force error: 7.928e+03 --> 3.515e+02 (N)
Minimum absolute Force error: 7.050e-13 --> 1.198e-01 (N)
Average absolute Force error: 2.645e+03 --> 3.255e+01 (N)
Maximum absolute Force error: 7.783e-03 --> 3.451e-04 (normalized)
Minimum absolute Force error: 6.922e-19 --> 1.176e-07 (normalized)
Average absolute Force error: 2.597e-03 --> 3.196e-05 (normalized)
R boundary error: 0.000e+00 --> 0.000e+00 (m)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Aspect ratio: 8.000e+00 --> 8.000e+00 (dimensionless)
Maximum Rotational transform: 1.928e-33 --> 1.225e+00 (dimensionless)
Minimum Rotational transform: 1.928e-33 --> 1.225e+00 (dimensionless)
Average Rotational transform: 1.928e-33 --> 1.225e+00 (dimensionless)
Maximum Deflation error: 0.000e+00 --> 0.000e+00 ~
Minimum Deflation error: 0.000e+00 --> 0.000e+00 ~
Average Deflation error: 0.000e+00 --> 0.000e+00 ~
Maximum Deflation error: 0.000e+00 --> 0.000e+00 (normalized)
Minimum Deflation error: 0.000e+00 --> 0.000e+00 (normalized)
Average Deflation error: 0.000e+00 --> 0.000e+00 (normalized)
==============================================================================================================
####################
i=1
####################
Building objective: QS two-term
Precomputing transforms
Building objective: lcfs R
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
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
Building objective: force
Precomputing transforms
Building objective: aspect ratio
Precomputing transforms
Building objective: rotational transform
Precomputing transforms
Building objective: Deflation
Number of parameters: 173
Number of objectives: 108
Number of equality constraints: 715
Number of inequality constraints: 2
Starting optimization
Using method: lsq-auglag
Warning: Maximum number of iterations has been exceeded.
Current function value: 9.709e-05
Constraint violation: 1.869e-03
Total delta_x: 7.126e-01
Iterations: 300
Function evaluations: 401
Jacobian evaluations: 300
==============================================================================================================
Start --> End
Total (sum of squares): 1.124e-27 --> 9.709e-05,
Maximum absolute Quasi-symmetry (1,4) two-term error: 2.576e-15 --> 3.093e-03 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error: 8.061e-18 --> 2.325e-07 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error: 8.535e-16 --> 2.830e-04 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error: 2.438e-15 --> 2.927e-03 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error: 7.628e-18 --> 2.200e-07 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error: 8.076e-16 --> 2.677e-04 (normalized)
Maximum absolute Force error: 7.928e+03 --> 2.554e+03 (N)
Minimum absolute Force error: 7.050e-13 --> 9.609e-02 (N)
Average absolute Force error: 2.645e+03 --> 2.176e+02 (N)
Maximum absolute Force error: 7.783e-03 --> 2.508e-03 (normalized)
Minimum absolute Force error: 6.922e-19 --> 9.434e-08 (normalized)
Average absolute Force error: 2.597e-03 --> 2.137e-04 (normalized)
R boundary error: 0.000e+00 --> 0.000e+00 (m)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Aspect ratio: 8.000e+00 --> 8.000e+00 (dimensionless)
Maximum Rotational transform: 1.928e-33 --> 1.343e+00 (dimensionless)
Minimum Rotational transform: 1.928e-33 --> 1.343e+00 (dimensionless)
Average Rotational transform: 1.928e-33 --> 1.343e+00 (dimensionless)
Maximum Deflation error: 8.663e+00 --> 5.000e+00 ~
Minimum Deflation error: 8.663e+00 --> 5.000e+00 ~
Average Deflation error: 8.663e+00 --> 5.000e+00 ~
Maximum Deflation error: 8.663e+00 --> 5.000e+00 (normalized)
Minimum Deflation error: 8.663e+00 --> 5.000e+00 (normalized)
Average Deflation error: 8.663e+00 --> 5.000e+00 (normalized)
==============================================================================================================
####################
i=2
####################
Building objective: QS two-term
Precomputing transforms
Building objective: lcfs R
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
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
Building objective: force
Precomputing transforms
Building objective: aspect ratio
Precomputing transforms
Building objective: rotational transform
Precomputing transforms
Building objective: Deflation
Number of parameters: 173
Number of objectives: 108
Number of equality constraints: 715
Number of inequality constraints: 2
Starting optimization
Using method: lsq-auglag
Warning: Maximum number of iterations has been exceeded.
Current function value: 2.146e-05
Constraint violation: 9.590e-04
Total delta_x: 7.146e-01
Iterations: 300
Function evaluations: 454
Jacobian evaluations: 300
==============================================================================================================
Start --> End
Total (sum of squares): 1.124e-27 --> 2.146e-05,
Maximum absolute Quasi-symmetry (1,4) two-term error: 2.576e-15 --> 1.567e-03 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error: 8.061e-18 --> 1.486e-07 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error: 8.535e-16 --> 1.399e-04 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error: 2.438e-15 --> 1.482e-03 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error: 7.628e-18 --> 1.406e-07 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error: 8.076e-16 --> 1.324e-04 (normalized)
Maximum absolute Force error: 7.928e+03 --> 1.311e+03 (N)
Minimum absolute Force error: 7.050e-13 --> 1.727e-01 (N)
Average absolute Force error: 2.645e+03 --> 8.711e+01 (N)
Maximum absolute Force error: 7.783e-03 --> 1.287e-03 (normalized)
Minimum absolute Force error: 6.922e-19 --> 1.695e-07 (normalized)
Average absolute Force error: 2.597e-03 --> 8.552e-05 (normalized)
R boundary error: 0.000e+00 --> 0.000e+00 (m)
Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Aspect ratio: 8.000e+00 --> 8.000e+00 (dimensionless)
Maximum Rotational transform: 1.928e-33 --> 1.420e+00 (dimensionless)
Minimum Rotational transform: 1.928e-33 --> 1.420e+00 (dimensionless)
Average Rotational transform: 1.928e-33 --> 1.420e+00 (dimensionless)
Maximum Deflation error: 3.533e+01 --> 5.000e+00 ~
Minimum Deflation error: 3.533e+01 --> 5.000e+00 ~
Average Deflation error: 3.533e+01 --> 5.000e+00 ~
Maximum Deflation error: 3.533e+01 --> 5.000e+00 (normalized)
Minimum Deflation error: 3.533e+01 --> 5.000e+00 (normalized)
Average Deflation error: 3.533e+01 --> 5.000e+00 (normalized)
==============================================================================================================
[13]:
eqs = eqs[0:N_deflation] # only keep the physically unique equilibria here
[14]:
from desc.plotting import *
plt.rcParams.update({"font.size": 18})
ax2 = None
for i in range(len(eqs)):
fig2, ax2 = plot_1d(eqs[i], "iota", ax=ax2, color=cs[i], label=str(i))
ax2.set_xlim([0, 1.0])
ax2.legend()
[14]:
<matplotlib.legend.Legend at 0x3be13e490>
[15]:
fig, axes = plt.subplots(2, 2, figsize=(12, 12))
axes = axes.flatten()
phis = np.linspace(0, 2 * np.pi / eqs[0].NFP, 4, endpoint=False)
fig, ax = plot_boundaries(eqs, lw=2, color=cs, legend=False, phi=phis[0], ax=axes[0])
fig, ax = plot_boundaries(eqs, lw=2, color=cs, legend=True, phi=phis[1], ax=axes[1])
fig, ax = plot_boundaries(eqs, lw=2, color=cs, legend=True, phi=phis[2], ax=axes[2])
fig, ax = plot_boundaries(eqs, lw=2, color=cs, legend=True, phi=phis[3], ax=axes[3])
for i, axx in enumerate(axes):
axx.set_aspect("equal")
axx.set_title(r"$\phi\cdot N_{FP}/2\pi=$" + rf"${phis[i]*eqs[0].NFP/2/np.pi:1.3f}$")
[16]:
fig, ax = plt.subplots(1, N_deflation, figsize=(N_deflation * 5, 5))
ax = ax.flatten()
for i in range(len(eqs)):
plot_boozer_surface(eqs[i], ax=ax[i])
[ ]: