# The DESC Stellarator Code Suite Part II: Perturbation and continuation
methods

Rory Conlin 1<EMAIL_ADDRESS>Daniel W. Dudt 1  Dario Panici 1  Egemen
Kolemen 1<EMAIL_ADDRESS>1Princeton University, Princeton, New Jersey
08544

###### Abstract

A new perturbation and continuation method is presented for computing and
analyzing stellarator equilibria. The method is formally derived from a series
expansion about the equilibrium condition
$\bm{F}\equiv\bm{J}\times\bm{B}-\nabla p=0$, and an efficient algorithm for
computing solutions to 2nd and 3rd order perturbations is developed. The
method has been implemented in the DESC stellarator equilibrium code, using
automatic differentiation to compute the required derivatives. Examples are
shown demonstrating its use for computing complicated equilibria, perturbing a
tokamak into a stellarator, and performing parameter scans in pressure,
rotational transform and boundary shape in a fraction of the time required for
a full solution.

##  1 Introduction

In the search for controlled nuclear fusion, 3D magnetic confinement devices
such as stellarators have been shown to have several advantages over 2D
magnetic geometries such as tokamaks, such as lower risk of disruption
(Helander et al., 2012) and current-free steady state operation (Helander,
2014). However, achieving good performance in a stellarator often requires
significant optimization of the plasma equilibrium. An additional advantage is
that because of the generally lower plasma current and consequently fewer
instabilities, more of the design and optimization can be done
computationally, without requiring building and testing full scale devices
(Boozer, 2015). Despite this, designing and optimizing 3D magnetic equilibria
that have good properties is still a computationally intensive task, for which
a number of codes and software packages have been developed (Dudt & Kolemen,
2020; Hirshman & Whitson, 1983a; Landreman et al., 2021; Hudson et al., 2012;
Spong et al., 1998; Lazerson et al., 2020; Drevlak et al., 2019)

Perturbation methods have been used heavily in tokamak plasma physics,
primarily to analyze the stability of axisymmetric MHD equilibria by searching
for a perturbation that minimizes the energy of the plasma, as in Bernstein’s
energy principle (Bernstein et al., 1958). This has been extended to a wide
range of of codes for analyzing fusion devices under small perturbations to an
MHD equilibrium, such as the GPEC suite of codes (Glasser, 2016; Glasser et
al., 2018; Park et al., 2007, 2009) for analyzing tokamak configurations under
3D perturbations.

Continuation methods have received less attention in the fusion community,
though they have seen extensive use in other fields such as(Howell, 2009;
Richter & DeCarlo, 1983). For the purposes of the present work, continuation
methods can be used to solve parameterized equations of the form
$F(x,\eta)=0$, where we identify $x$ as the solution vector, and $\eta$ as a
continuation parameter. When $\eta$ is varied we find a family of solutions
connected in parameter space, as well as possible branches and bifurcations of
this family, indicated by points where the Jacobian of $F$ is singular.
Starting from the initial value $(x_{0},\eta_{0})$ where $F$ is zero, we
continuously vary $\eta$ while simultaneously varying $x$ such that the
equation $F(x,\eta)=0$ is satisfied. Various methods exist for finding the
solution curve (the locus of points $(x,\eta)$ where $F(x,\eta)=0$ is
satisfied) such as piecewise linear (simplex) continuation and psuedo-
arclength methods (Allgower & Georg, 1990). In an abstract sense, the standard
multigrid method for solving partial differential equations can be though of
as a discrete continuation method, where the continuation parameter $\eta$
governs the level of numerical resolution (ie grid spacing, number of basis
functions etc.)

In general, existing codes for computing stellarator equilibria (Hirshman &
Whitson, 1983a; Hudson et al., 2012) only find discrete equilibria, and
solving for a new equilibrium requires running the code from scratch (possibly
with a different starting guess). Several questions that may then be asked
are:

  1. 1.

Can these equilibria be computed more efficiently?

  2. 2.

Given a single equilibrium solution, can we find other similar solutions?

  3. 3.

What does the full phase space of 3D MHD equilibria look like?

This is part II of a three-part series of papers on the DESC stellarator
optimization code suite. Part I details the DESC equilibrium solver in
comparison with the VMEC (Hirshman & Whitson, 1983b) 3D MHD equilibrium code.
Computing 3D MHD equilibria is the preliminary step in stellarator analysis,
so computing these equilibria quickly and accurately is important for further
studies. In this paper, we describe a new continuation method for computing
complicated stellarator equilibria using perturbations that attempts to
resolve these questions. In section 2 we describe the mathematical background
to the perturbation method, while in section 3 we describe the implementation
of the method in the DESC code. In section 4 demonstrate how to use these
perturbations in a continuation method for computing stellarator equilibria.
In section 5 we demonstrate other applications of the perturbation method for
computing and analyzing stellarator equilibria. Part III (Dudt et al., 2022)
presents DESC’s unique stellarator optimization capabilities made possible by
the efficient equilibrium solver and the perturbation method, resulting in
orders of magnitude speed-up in optimization. These advantages are shown in
the context of quasi-symmetry optimization, where results are compared to
conventional tools (Spong et al., 1998). Three different quasi-symmetry
objective formulations are also shown, with the relative advantages of each
compared, highlighting the flexibility of DESC as an optimization code.

##  2 Perturbations

A general fixed-boundary equilibrium problem can be described by a set of
parameters $\bm{c}=\\{R_{b},Z_{b},p,\iota,\Psi\\}$ where $R_{b},Z_{b}$ are the
$R,Z$ coordinates of the boundary surface, $p$ is the pressure profile,
$\iota$ the rotational transform, and $\Psi$ the total toroidal flux through
the torus. In many spectral equilibrium codes (Dudt & Kolemen, 2020; Hirshman
& Whitson, 1983a), the independent variables that define the equilibrium can
be given by $\bm{x}=[R_{lmn},Z_{lmn},\lambda_{lmn}]$ where $R_{lmn}$,
$Z_{lmn}$ and $\lambda_{lmn}$ are spectral coefficients of the flux surface
positions and the poloidal stream function (indexed by $l$ in the radial
direction, $m$ in the poloidal direction, and $n$ in the toroidal direction).
The condition of MHD equilibrium can then be written as a (possibly vector
valued) nonlinear algebraic equation involving the fixed parameters and
independent variables, $\bm{f}(\bm{x},\bm{c})=0$. The function $f$ is the
discretized form of the general MHD force balance $\bm{J}\times\bm{B}-\nabla
p=0$(Dudt & Kolemen, 2020), or a condition on the gradient of the energy
functional $W=\int_{V}(B^{2}/2\mu_{0}+p)dV$ (Hirshman & Whitson, 1983a).

Given a set of parameters $\bm{c}$ and a solution vector $\bm{x}$ which
satisfy $\bm{f}(\bm{x},\bm{c})=0$, we wish to find how the equilibrium would
change if the parameters $\bm{c}$ are perturbed to $\bm{c}+\Delta\bm{c}$. For
instance, we may have a solution for a vacuum equilibrium and want to see how
adding finite pressure changes it, or we may start with a 2D tokamak solution
and add a 3D perturbation to the boundary shape to form a stellarator.

We assume that the new equilibrium is given by $\bm{x}+\Delta\bm{x}$, and
expand $f$ in a Taylor series ***Some care must be taken here, as we are
implicitly assuming that the function $f$ is at least $C^{2}$ continuous. The
DESC code assumes the existence of nested flux surfaces so that there is an
analytic mapping between real space and magnetic coordinates. In cases where
the assumption of nested surfaces is violated, this mapping may not exist.
However, we note that the force balance equations are still analytic functions
of the dependent variables (it can be shown that they form a high order
polynomial), though their physical meaning will be somewhat unclear in regions
where nested surfaces don’t exist :

$\bm{f}(\bm{x}+\Delta\bm{x},\bm{c}+\Delta\bm{c})=\bm{f}(\bm{x},\bm{c})+\frac{\partial\bm{f}}{\partial\bm{x}}\Delta\bm{x}+\frac{\partial\bm{f}}{\partial\bm{c}}\Delta\bm{c}+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{x}^{2}}\Delta\bm{x}\Delta\bm{x}\\\
+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{c}^{2}}\Delta\bm{c}\Delta\bm{c}+\frac{\partial^{2}\bm{f}}{\partial\bm{x}\partial\bm{c}}\Delta\bm{x}\Delta\bm{c}+O(\Delta
x^{3})$ (1)

We wish to solve this equation for $\Delta\bm{x}$ such that
$\bm{f}\bm{x}+\Delta\bm{x},\bm{c}+\Delta\bm{c})=0$. At first order this is a
straightforward algebraic equation, but at higher orders it becomes a tensor
polynomial equation which can be difficult or impossible to solve efficiently
(for example, if $\bm{f}$ is a vector valued function, just storing the 2nd
derivative tensor in memory could require upwards of 100 GB). Instead of
seeking a direct solution, we can try a perturbative approach, were we
introduce an arbitrary small parameter $\epsilon$ and further expand
$\Delta\bm{x}$ and $\Delta\bm{c}$ in a perturbation series in powers of
$\epsilon$ (we assume that $\Delta\bm{c}$ is known a-priori and so only a
first order term is required).

$\displaystyle\Delta\bm{x}$ $\displaystyle=$
$\displaystyle\epsilon\bm{x}_{1}+\epsilon^{2}\bm{x}_{2}+...$ (2)
$\displaystyle\Delta\bm{c}$ $\displaystyle=$ $\displaystyle\epsilon\bm{c}_{1}$
(3)

Plugging this into Equation 1 (and setting
$\bm{f}(\bm{x},\bm{c})=\bm{f}(\bm{x}+\Delta\bm{x},\bm{c}+\Delta\bm{c})=0$) we
get:

$0=\frac{\partial\bm{f}}{\partial\bm{x}}(\epsilon\bm{x}_{1}+\epsilon^{2}\bm{x}_{2})+\frac{\partial\bm{f}}{\partial\bm{c}}\epsilon\bm{c}_{1}+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{x}^{2}}(\epsilon\bm{x}_{1}+\epsilon^{2}\bm{x}_{2})(\epsilon\bm{x}_{1}+\epsilon^{2}\bm{x}_{2})\\\
+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{c}^{2}}\epsilon\bm{c}_{1}\epsilon\bm{c}_{1}+\frac{\partial^{2}\bm{f}}{\partial\bm{x}\partial\bm{c}}(\epsilon\bm{x}_{1}+\epsilon^{2}\bm{x}_{2})\epsilon\bm{c}_{1}+O(\Delta
x^{3})\\\ $ (5)

We can then collect powers of $\epsilon$ and set each order of $\epsilon$ to
zero in turn. The first order equation gives:

$\displaystyle{0}$ $\displaystyle=$
$\displaystyle\frac{\partial\bm{f}}{\partial\bm{x}}\epsilon\bm{x}_{1}+\frac{\partial\bm{f}}{\partial\bm{c}}\epsilon\bm{c}_{1}$
(6) $\displaystyle x_{1}$ $\displaystyle=$
$\displaystyle-\Bigg{(}\frac{\partial\bm{f}}{\partial\bm{x}}\Bigg{)}^{-1}\Bigg{(}\frac{\partial\bm{f}}{\partial\bm{c}}c_{1}\Bigg{)}$
(7)

The second order term gives:

$\displaystyle{0}$ $\displaystyle=$
$\displaystyle\frac{\partial\bm{f}}{\partial\bm{x}}\epsilon^{2}\bm{x}_{2}+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{x}^{2}}\epsilon^{2}\bm{x}_{1}\bm{x}_{1}+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{c}^{2}}\epsilon^{2}\bm{c}_{1}\bm{c}_{1}+\frac{\partial^{2}\bm{f}}{\partial\bm{x}\partial\bm{c}}\epsilon^{2}\bm{x}_{1}\bm{c}_{1}$
$\displaystyle\bm{x}_{2}$ $\displaystyle=$
$\displaystyle-\Bigg{(}\frac{\partial\bm{f}}{\partial\bm{x}}\Bigg{)}^{-1}\Bigg{(}\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{x}^{2}}\bm{x}_{1}\bm{x}_{1}+\frac{1}{2}\frac{\partial^{2}\bm{f}}{\partial\bm{c}^{2}}\bm{c}_{1}\bm{c}_{1}+\frac{\partial^{2}\bm{f}}{\partial\bm{x}\partial\bm{c}}\bm{x}_{1}\bm{c}_{1}\Bigg{)}$

In general, the second derivative terms will be large, dense, rank 3 tensors
(recall that
$\bm{$}canbeavectorvaluedfunction),whichmaybeextremelyexpensivetocomputeandmaynotevenfitinmemoryforhighresolutioncases.However,itisimportanttonotethatthefullsecondderivativesareneverneeded(orindeedfullfirstderivativesapartfromtheJacobianmatrix$∂f∂x$).AllthatisneededaredirectionalderivativesorJacobianvectorproducts.IntheDESCcode(\autoref{sec:code_improvements})thesearecalculatedusingforwardmodeautomaticdifferentiationatacostafewtimesgreaterthanthecostofasingleevaluationof$f$.Thesetermscouldalsobeapproximatedusingfinitedifferencesatasimilarcost.Thismeansthatateachorderthemostexpensiveoperationissolvingalinearsystemoftheform\begin{equation}\bm{J}\bm{x}_{i}=\bm{b}\end{equation}where$J
≡∂f /
∂x$istheJacobianmatrixandisthesameateachorder,soonlyneedstobedecomposedonceand$i$istheorderofthetermintheperturbationseries.Becauseofthesecomputationandmemorysavings,themethodhasbeenextendedto3rdorderintheDESCcode(see\autoref{sec:code_improvements}),thoughinmostsituations1stor2ndorderperturbationsgivesacceptableresults(see\autoref{fig:dshape}).\par\par
Itisinstructivetonotethatthefirstordertermintheperturbationseriesfor$Δx$iseffectivelytheNewtonstep,whilethesecondordertermiscommonlyreferredtoastheHalleystepinoptimizationliterature\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{gander1985halley,gundersen2010}{\@@citephrase{,
}}{})}.WhileNewton^{\prime}smethodconvergesquadratically\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{nocedal2006numerical}{\@@citephrase{,
}}{})}whenstartedsufficientlyclosetoasolution,ifstartedfarawayitmaydivergeandsoinpracticethemethodmustbeglobalizedusingeitheralinesearchoratrustregionframework.Althoughweassumethatwebeginnearoratanequilibriumsothat$f(x,c)
=
0$,dependingonthesizeoftheperturbationthesolutionlandscapemaychangesuchthatthisisnolongerthecase,andwefindinpracticeinseveralcasestheunconstrainedNewtonstepcausesthesolutiontodiverge.Thereisalsoanadditionalconstraintinherentintheapproachtakenthatthesecondordertermshouldbesmallerthanthefirstordertermbyafactor$ϵ<<
1$.\par
Thissuggeststhatanaturalwaytoenforcetheseconditionsandensurereasonableconvergencepropertiesistouseatrustregionmethod.WhenexpandingtheobjectivefunctioninaTaylorseriesweareeffectivelyapproximatingitbyalinearorquadraticfunction.Insteadofseekingtheglobalminimizertothesemodelfunctions,thetrustregionapproachinsteadrecognizesthattheTaylorapproximationisonlyvalidinsomesmallneighborhoodandrestrictsthestepsizeaccordingly.Mathematically,insteadoffindingtheexactsolution$x_i^*$tothelinearsystem\begin{equation}\bm{J}\bm{x}_{i}=\bm{b}\end{equation}Weinsteadseekasolutiontothefollowingoptimizationproblem:\begin{equation}\min_{\bm{x}_{i}}||\bm{J}\bm{x}_{i}-\bm{b}||^{2}\quad\quad
s.t.\quad||\bm{x}_{i}||\leq
r\end{equation}where$r$istheradiusofthetrustregion.Thissubproblemhastwopossiblesolutions\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{nocedal2006numerical}{\@@citephrase{,
}}{})}:eitherthetruesolution$x_i^*$lieswithinthetrustregion,orelsethatthereisascalar$α>
0$suchthat\begin{equation}(\bm{J}+\alpha
I)\bm{x}_{i}=\bm{b},\quad\quad||\bm{x}_{i}||=r\end{equation}Theconstrainedvaluefor$x_i$isbeefficientlyfoundusinga1Drootfindingmethodin$α$,intheDESCcodeweuseabracketedNewtonmethod.ThecostofsolvingthissubproblemisroughlythecostofonesingularvaluedecompositionoftheJacobian$J$,ortwotothreeCholeskyfactorizationsoftheapproximateHessianmatrix$B
≡J^T J$.\par
Selectionofthetrustregionradius$r$istraditionallydoneadaptivelywithinanoptimizationloop,bycomparingtheactualreductionintheresidualateachstepwiththatpredictedbythemodelfunction.Inaperturbation,onlyasingle"step"istaken,andwehavefoundempiricallythatsettingthetrustregionatasmallfractionof$||x||$generallygivesgoodresults,typically$r=0.1
||x||$.Forsecondandhigherorderperturbations,suchalargetrustregioncanresultleadtoinaccurateresults.Thiscanbeunderstoodbyconsideringtheexpansionmadein\autoref{eq:expansion},whereitwasimplicitlyassumedthatthe2ndordercorrectionwasafactorof$ϵ$smallerthanthefirstorderterm.Usingalargetrustregionforthefirstorderterm($r=0.1
||x||$)andasmallertrustregionforthehigherordertermscorrectlyenforcesthisassumedscaling.Inpracticeweobtainedgoodperformanceusing$r_i=0.1
||x_i||$,sothesecondordertermisanorderofmagnitudelessthanthefirstordertermandsoonforhigherorders.\par$

##  3 The DESC code

DESC is a recently developed (Dudt & Kolemen, 2020) pseudo-spectral code for
computing 3D MHD equilibria. DESC computes 3D MHD equilibia by solving the
force balance equations $\bm{J}\times\bm{B}-\nabla p=0$ as opposed to the more
common variational method which minimizes the MHD energy
$\int_{V}(B^{2}/2\mu_{0}+p)\ dV$. The independent variables in the equation
are the positions of the flux surfaces $(R,Z)$ as well as the stream function
$\lambda$, defined as the difference between the boundary poloidal angle
$\theta$ and the straight field line poloidal angle:
$\vartheta=\theta+\lambda$. These quantities are discretized in a Fourier-
Zernike basis, using a Fourier series in the toroidal direction and Zernike
polynomials in the poloidal/radial directions. After discretization, the
equilibrium equation is expressed as a set of nonlinear algebraic equations
$\bm{f}(\bm{x},\bm{c})=0$ where $\bm{x}$ is a vector containing the spectral
coefficients of $R$, $Z$, and $\lambda$, while $\bm{c}$ contains fixed
parameters that define the equilibrium problem, such as the pressure and
rotational transform profiles and the fixed boundary shape.

Since it’s original publication (Dudt & Kolemen, 2020), DESC has undergone a
major upgrade that involved porting it from MATLAB to Python in order to take
advantage of the JAX library (Bradbury et al., 2018) for automatic
differentiation (AD) and just-in-time (JIT) compilation. Initially developed
for machine learning applications, JAX provides an NumPy(Van Der Walt et al.,
2011; Harris et al., 2020)-like API for common mathematical operations and
allows arbitrary functions to be differentiated using forward or reverse mode
AD. This allows the calculation of exact derivatives of the objective function
automatically, rather than having to code them by hand which is time-consuming
and error prone, or using finite differences which are computationally
expensive and can be inaccurate. It is also much more flexible as new
objective functions can be added and optimized by defining only the forward
pass, which is of great use in stellarator optimization where new objectives
may be added in the future. JAX also allows for JIT compiling of code to both
CPUs and GPUs which significantly speeds up calculation, approaching speeds of
traditional compiled languages, avoiding one of the primary limitations of
Python for scientific computing. Additionally, given that the vast majority of
new supercomputers heavily leverage GPUs, and this trend is likely to
continue, using JAX allows DESC to take full advantage of all the compute
capability available, rather than being limited to CPU-only parallelization
like many legacy codes. This allows a "best of both worlds" approach where the
code is easy to use, maintain, adapt, and upgrade, while still being fast
enough for production applications.

Several aspects also make DESC the ideal code for implementing the
perturbation method outlined in section 2:

  1. 1.

Using a pseudo-spectral discretization significantly reduces the number of
independent variables, resulting in smaller Jacobian matrices.

  2. 2.

Using JAX allows fast and accurate computation of the required derivatives and
Jacobian-vector products.

  3. 3.

Formulating the equilibrium problem as a system of nonlinear equations means
the first order perturbation term spans the entire tangent space at the
initial point (assuming the Jacobian is full rank), rather than just the
single line along the gradient of a scalar objective function.

In addition, formulating the problem as a system of nonlinear equations and
solving them in a least squares sense also effectively gives an extra order of
derivative for free. This can be seen by considering the least squares problem

$\min_{\bm{x}}y(\bm{x})\equiv\frac{1}{2}\bm{f}(\bm{x})^{T}\bm{f}(\bm{x})$ (12)

where $\bm{f}$ is a vector valued function, and $y$ is the sum of squares of
the residuals of $\bm{f}$. The gradient of $y$ is given by

$\frac{dy}{d\bm{x}}=\bm{f}^{T}\frac{d\bm{f}}{d\bm{x}}$ (13)

and its Hessian (matrix of second partial derivatives) is given by

$\frac{d^{2}y}{d\bm{x}^{2}}=\frac{d\bm{f}}{d\bm{x}}^{T}\frac{d\bm{f}}{d\bm{x}}+\bm{f}^{T}\frac{d^{2}\bm{f}}{d\bm{x}^{2}}$
(14)

Given that $\bm{f}$ is the function we are trying to minimize in the least
squares sense, we can generally assume that the 2nd term in the Hessian is
negligible compared to the first, the so-called "small residual approximation"
(Nocedal & Wright, 2006). This gives an approximate Hessian
$\frac{d^{2}y}{d\bm{x}^{2}}\sim\frac{d\bm{f}}{d\bm{x}}^{T}\frac{d\bm{f}}{d\bm{x}}$,
meaning that using only first derivative information about $\bm{f}$ gives us
both first and second derivative information about $y$. In addition, the
approximate Hessian this gives is always positive semi-definite, leading to a
convex subproblem which is easy to solve.

##  4 Continuation Method

Many MHD equilibrium codes (Hirshman & Whitson, 1983a) (Hudson et al., 2012)
use a multigrid approach, where an initial guess is specified on a coarse
grid, and the error is minimized on that grid before being interpolated to a
finer grid and re-solved, continuing until the finest resolution level has
been reached. This is done both to speed computation by doing more
calculations on coarser grids, and also to make it more robust to poor initial
guesses and avoid additional saddle points that can appear in high dimensional
spaces (Dauphin et al., 2014). DESC offers this as well, allowing the
resolution to be varied in the radial, poloidal, and toroidal directions
independently.

In addition to the standard multigrid approach, DESC has also implemented a
new continuation method using the perturbation techniques described in section
2. In general, continuation methods seek to find how the solution to an
equation varies as parameters are changed. Using the notation of section 2 we
can view $\bm{x}$ to be an implicit function of $\bm{c}$, related by the
constraint equation $\bm{f}(\bm{x},\bm{c})=0$, and try to find the map that
relates $\bm{x}$ and $\bm{c}$ along lines of $\bm{f}(\bm{x},\bm{c})=0$.

In this work, we note two properties that make continuation methods especially
relevant:

  1. 1.

2D (axisymmetric) equilibria are much easier to compute than 3D equilibria,
due to the guaranteed existence of flux surfaces in axisymmetric equilibria,
and the significant reduction in the number of variables needed to represent
the solution.

  2. 2.

Vacuum (vanishing beta) equilibria are easier to compute than finite pressure
equilibria, due to the lack of Shafranov shift.

We can take advantage of these properties to compute complicated stellarator
equilibria by first solving an "easy" problem, and then introduce a
continuation parameter that transforms our initial solution into the solution
to a "hard" problem.

To do this, we introduce 2 scalar continuation parameters: $\eta_{b}$, a
multiplier on all of the non-axisymmetric boundary modes, and $\eta_{p}$, a
scaling factor for the pressure profile. Setting $\eta_{b}=0$ would give a 2D
axisymmetric boundary that is "close" to the desired 3D equilibrium, while
setting $\eta_{p}=0$ would give the vacuum equilibrium with the desired
boundary shape. By varying these two parameters, we find a "family" of
solutions that are connected continuously.

To vary the parameters, we apply perturbations as outlined in section 2. A
standard method for solving continuation problems, known as "natural parameter
continuation" would be analagous to a 0th order perturbation followed by
several Newton iterations, while a first order perturbation would be similar
to Gauss-Newton continuation. The higher order perturbations discussed above
do not seem to have been explored in the more general continuation method
literature, but can be seen as a version of Halley’s method applied to the
combined system.

In most traditional continuation methods, the step size in the continuation
parameter is determined adaptively based on local error estimates and the
ratio of predicted to achieved error reduction. In practice we have found that
this is often not necessary, and fixing the step sizes a-priori provides
sufficiently accurate results. In most cases going from a vacuum equilibrium
to a moderate beta of $\sim 3\%$ can be done in a single step, and boundary
perturbations in anywhere from 1-4 steps, depending on the desired accuracy
and complexity of the boundary. This does require the user to specify the
desired perturbation steps explicitly, thought a future upgrade to the code is
planned to allow adaptive perturbations where the user need only supply an
initial guess and the final desired parameters.

As a first example (Figure 1), we solve for a vacuum heliotron by first
solving for a simple circular tokamak, then applying a 3D perturbation to the
boundary. After the perturbation, a small number (usually 5-10) of regular
Newton iterations are applied to ensure convergence.

Figure 1: Comparison of perturbed tokamak equilibrium vs true solution for a
heliotron like stellarator. The initial tokamak (grey) is axisymmetric. By
applying a perturbation (1st order in red, 2nd order in green) to the non-
axisymmetric boundary modes, we obtain an approximation to the true 3D
solution (blue), obtained by solving from the start with the full 3D boundary
shape. The first order perturbation captures the majority of the differences
between the initial and true solutions. The second order effects bring the
perturbed solution even closer to the true one.

Similarly, when computing finite pressure equilibria it can be difficult to
estimate the Shafranov shift a-priori for generating a good initial guess.
DESC avoids this by first computing a vacuum solution, for which a good
initial guess can generally be found by simply scaling the boundary surface.
The finite pressure is then added back in as a perturbation, which then often
only requires a small number of further iterations to reach convergence as
shown in Figure 2.

Figure 2: Comparison of perturbed equilibrium vs true solution for a change in pressure from $\beta=0\%$ to $\beta=3\%$. The first order perturbation (red) captures the majority of the differences between the initial (grey) and true (blue) solutions. Second and third order effects (not shown for clarity) are visually very similar but significantly reduce the resulting force balance error as shown in Table 1. Perturbation order | Normalized force error ($|F|/|\nabla p|)$
---|---
1st order | $87.6\%$
2nd order | $29.6\%$
3rd order | $9.7\%$
Table 1: Normalized force balance error vs perturbation order for a pressure
increase from $\beta=0\%$ to $\beta=3\%$.

While in general the two parameters could be varied in any order, or
simultaneously when solving for a 3D finite-beta equilibrium, we have found
that first varying $\eta_{p}$, followed by $\eta_{b}$ to be more efficient.
Varying $\eta_{p}$ while holding $\eta_{b}$ fixed at $0$ allows the pressure
perturbations to be done on a 2D axisymmetric configuration, which reduces the
computational cost and ensures the existence of good flux surfaces. After
reaching a high resolution finite-beta axisymmetric equilibrium, 3D modes are
added to the basis functions and the boundary is perturbed to give the final
desired 3D finite-beta equilibrium.

This procedure is demonstrated in Figure 3, where the initial solution is an
axisymmetric vacuum tokamak. The pressure is then increased, as demonstrated
by the Shafranov shift, followed by 3D deformation of the boundary shape to
arrive at a W7X like configuration. After each large perturbation, a small
number of regular Newton iterations are performed to ensure that an
equilibrium has been reached (recall that in the derivation of the
perturbations, it was assumed that the initial state before perturbing was an
equilibrium). This method was also used to compute the solutions shown in Part
I (Panici et al., 2022), where a detailed analysis of the force error is
given.

Figure 3: Continuation solution for W7X like equilibrium. Starting from a
vacuum 2D equilibrium, the pressure is increased to $\beta\sim 3\%$ in a
series of two steps. Then a 3D perturbation is applied to the boundary, broken
up into 4 steps (2 shown) to arrive at the final finite $\beta$ 3D solution.

Further insight can be gained by looking at the toroidal current profile of
the W7X like equilibrium as the boundary ratio is varied in Figure 4. At each
step we keep the rotational transform profile fixed, and so as expected the
axisymmetric case ($\eta_{b}=0$) requires a large toroidal current to generate
the poloidal field. As we increase the 3D shaping, more of the rotational
transform is generated by axis torsion and boundary shaping, reducing the
required plasma current to near zero.

Figure 4: Toroidal current and rotational transform for W7X like equilibrium
with varying boundary ratios. When $\eta_{b}=0$, corresponding to an
axisymmetric solution, a large toroidal current is required to generate the
rotational transform (which is held fixed as the boundary is varied).
Increasing the 3D shaping allows more of the rotational transform to come from
external fields, reducing the plasma current.

An important feature to note about the aforementioned continuation method is
that the solution at each step is in fact the "exact" solution to the
equilibrium problem with perturbed parameters, and is hence still a valid
equilibrium that satisfies MHD force balance and may have desirable physics
properties worth studying that may be absent in the final solution. By
applying different perturbations and varying different parameters, one can
explore whole families of solutions that are "nearby" to a starting
equilibrium.

##  5 Other Applications

###  5.1 Parameter Scans

Another important feature of continuation methods is revealing how solutions
change as parameters of the problem are varied. In the previous section we
used this to find single equilibria, but the method can also be used to
explore families of equilibria related by a parameter or group of parameters.
In this, we find equilibria "nearby" to an equilibrium already found, such as
examining the same boundary shape at different values of $\beta$ or boundary
perturbations of varying magnitude and shape such as resonant magnetic
perturbations (RMP) in a tokamak. Traditionally this requires re-solving the
equilibrium for each new value of the parameter. An alternative approach is to
apply a perturbation to an initial equilibrium and find the corresponding
changes in the flux surfaces. After applying the first perturbation, we re-
linearize about the new state and perturb again, and so step through different
equilibria for the cost of a single linear system solve at each step. This
represents a significant computational savings compared to solving the full
nonlinear problem for each value of the parameter.

Figure 5: Flux surfaces on the $\zeta=0$ and $\zeta=\pi$ planes for a helical
stellarator. The initial equilibrium had $\beta=0\%$, and a series of 2nd
order perturbations were applied to increase the pressure rather than solving
a new equilibrium from scratch each time.

Figure 5 shows the flux surfaces on the $\zeta=0$ and $\zeta=\pi$ plane for a
helical stellarator for the same pressure profile scaled to different values
of $\beta$. The initial solution was at $\beta=0\%$, and 2nd order
perturbations were applied sequentially to step the pressure up to
$\beta=10\%$. The flux surfaces are indistinguishable from those obtained by
solving from scratch at each value of $\beta$, and the computational cost is
significantly reduced, as shown in Table 2. Performing such a scan requires
only a few lines of code, as shown in 1.

⬇

1import numpy as np

2import desc.io

3from desc.equilibrium import EquilibriaFamily

4eq0 = desc.io.load("heliotron_vacuum_solution.h5")[-1]

5eqf = EquilibriaFamily(eq0)

6# this corresponds to a change in beta of ~1%

7dp = np.array([ 1800., 0., -3600., 0., 1800.])

8for i in range(10):

9 eqf.append(eqf[-1].perturb(dp=dp, order=1))

10 # polish off solution

11 eqf[-1].solve(maxiter=5)

Listing 1: Code example for $\beta$ scan $\beta$ | Time (s) | Time (s)
---|---|---
| without perturbations | with perturbations
1 | 164 | 78.0
2 | 174 | 16.3 s
3 | 164 | 14.4 s
4 | 122 | 12.6 s
5 | 127 | 17.8 s
6 | 122 | 19.1 s
7 | 121 | 15.2 s
8 | 134 | 13.6 s
9 | 140 | 16.0 s
10 | 138 | 18.3 s
Total | 23.4 min | 3.68 min
Table 2: Computation times for different values of $\beta$ in the pressure
scan shown in Figure 5, both with and without perturbations, starting from a
vacuum solution. The first perturbation takes longer as the time includes JIT
compilation for the specific problem being solved, resulting in significantly
faster times for the subsequent perturbation steps which reuse the same
compiled code. All computation done on an AMD Ryzen 7 PRO 4750U with 32 GB
memory.

As another example (Figure 6), a D-shaped tokamak was used as the starting
point for variations in the rotational transform profile. The initial
equilibrium had a rotational transform on axis of $\iota_{0}=1$, and
perturbations were applied to reduce this down to $\iota_{0}=0.34$. As the
rotational transform is decreased, we see increased Shafranov shift and
deformation of the flux surfaces as the reduced poloidal field struggles to
contain the pressure ($\beta\sim 3\%$ in this example).

Figure 6: Top: Flux surfaces on the $\zeta=0$ plane for a D-shaped tokamak
with varying axis rotational transform. The initial equilibrium (blue)
corresponds to $\iota_{0}=1$, and perturbations were used to reduce the axis
rotational transform without recomputing the entire equilibrium solution.
Bottom: Rotational transform profiles corresponding to flux surfaces shown
above.

As in the pressure scan example, the solution obtained by perturbation is
indistinguishable by eye from the solution obtained by solving the full
equilibrium problem, and significantly faster. Once an initial equilibrium is
solved, parameter scans can be performed an order of magnitude faster using
perturbations than other approaches which require a full solution at each
step. The code to perform such a scan is shown in 2.

⬇

1import numpy as np

2import desc.io

3eq0 = desc.io.load("DSHAPE.h5")[-1]

4# perturbing a 2nd order polynomial basis for iota

5di = np.array([-0.33, 0, 0.33])

6eq1 = eq0.perturb(di=di, order=2)

7eq1.solve(maxiter=5)

8eq2 = eq1.perturb(di=di, order=2)

9eq2.solve(maxiter=5)

Listing 2: Code example for perturbing iota profile

###  5.2 Optimization

The perturbation techniques described previously can also be adapted for
optimization, where instead of choosing a particular change in the parameters
$\Delta\bm{c}$ we instead let $\Delta\bm{c}$ be a free parameter that is
chosen to minimize a cost function $\bm{g}(\bm{x},\bm{c})$, such as
quasisymmetry error (Boozer, 1983; Helander, 2014), coil complexity (Zhu et
al., 2018b, a; McGreivy et al., 2020), or fast particle confinement (Nemov et
al., 2008; Velasco et al., 2021):

$\Delta\bm{c}^{*}=\operatorname*{arg\,min}_{\Delta\bm{c}}\bm{g}(\bm{x}+\Delta\bm{x},\bm{c}+\Delta\bm{c})\quad
s.t.\quad\bm{f}(\bm{x}+\Delta\bm{x},\bm{c}+\Delta\bm{c})=0$ (15)

Where as before $\Delta\bm{x}$ is an implicit function of $\Delta\bm{c}$.

This finds the perturbation that most decreases the cost function $\bm{g}$
while maintaining force balance, without the need for additional equilibrium
solves, and is analogous to a Newton method applied to the objective
$\bm{g}(\bm{x},\bm{c})$ and its constraints $\bm{f}(\bm{x},\bm{c})=0$. In the
optimization literature, this is part of a more general class of methods for
constrained optimization. Like the perturbations described in section 2, this
method can also be extended to higher order. This extension and further
details on this method of optimization and applications to quasisymmetry will
be given in Part III (Dudt et al., 2022).

##  6 Conclusions

We have demonstrated a new technique for computing stellarator equilibria, and
for exploring how those equilibria change as parameters are varied. The
methods are computationally efficient, and offer significant speedups compared
to existing techniques, and in many cases offer possibilities that have not
existed before. An important future application of these methods is in
exploring the connections between different classes of equilibria, such as how
tokamaks bifurcate into stellarators, and how different classes of
quasisymmetric stellarators may be related.

This work was supported by the U.S. Department of Energy under contract
numbers DE-AC02-09CH11466, DE- SC0022005 and Field Work Proposal No. 1019. The
United States Government retains a non-exclusive, paid-up, irrevocable, world-
wide license to publish or reproduce the published form of this manuscript, or
allow others to do so, for United States Government purposes.

## References

  * Allgower & Georg (1990) Allgower, Eugene L. & Georg, Kurt 1990 Numerical Continuation Methods: An Introduction. Berlin, Heidelberg: Springer-Verlag.
  * Bernstein et al. (1958) Bernstein, I.B., Frieman, E.A., Kruskal, M. D. & Kulsrud, R. M. 1958 An energy principle for hydromagnetic stability problems. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences 244 (1236), 17–40, citation Key: Bernstein1958.
  * Boozer (1983) Boozer, Allen H. 1983 Transport and isomorphic equilibria. Physics of Fluids 26 (2), 496–499.
  * Boozer (2015) Boozer, Allen H. 2015 Stellarator design. Journal of Plasma Physics 81 (6).
  * Bradbury et al. (2018) Bradbury, James, Frostig, Roy, Hawkins, Peter, Johnson, Matthew James, Leary, Chris, Maclaurin, Dougal, Necula, George, Paszke, Adam, VanderPlas, Jake, Wanderman-Milne, Skye & Zhang, Qiao 2018 JAX: composable transformations of Python+NumPy programs.
  * Dauphin et al. (2014) Dauphin, Yann, Pascanu, Razvan, Gulcehre, Caglar, Cho, Kyunghyun, Ganguli, Surya & Bengio, Yoshua 2014 Identifying and attacking the saddle point problem in high-dimensional non-convex optimization. arXiv:1406.2572 [cs, math, stat] ArXiv: 1406.2572.
  * Drevlak et al. (2019) Drevlak, M., Beidler, C. D., Geiger, J., Helander, P. & Turkin, Y. 2019 Optimisation of stellarator equilibria with rose. Nuclear Fusion 59 (1).
  * Dudt et al. (2022) Dudt, Daniel, Conlin, Rory, Panici, Dario & Kolemen, Egemen 2022 The desc stellarator code suite part iii: Quasi-symmetry optimization .
  * Dudt & Kolemen (2020) Dudt, DW & Kolemen, E 2020 Desc: A stellarator equilibrium solver. Physics of Plasmas 27 (10), 102513.
  * Gander (1985) Gander, Walter 1985 On halley’s iteration method. The American Mathematical Monthly 92 (2), 131–134.
  * Glasser et al. (2018) Glasser, Alexander, Kolemen, Egemen & Glasser, A. H. 2018 A riccati solution for the ideal mhd plasma response with applications to real-time stability control. Physics of Plasmas 25 (3), 032507, citation Key: Glasser2018.
  * Glasser (2016) Glasser, A. H. 2016 The direct criterion of newcomb for the ideal mhd stability of an axisymmetric toroidal plasma. Physics of Plasmas 23 (7), 072505, citation Key: Glasser2016.
  * Gundersen & Steihaug (2010) Gundersen, Geir & Steihaug, Trond 2010 On large-scale unconstrained optimization problems and higher order methods. Optimization Methods and Software 25 (3), 337–358.
  * Harris et al. (2020) Harris, Charles R., Millman, K. Jarrod, van der Walt, Stéfan J., Gommers, Ralf, Virtanen, Pauli, Cournapeau, David, Wieser, Eric, Taylor, Julian, Berg, Sebastian, Smith, Nathaniel J. & et al. 2020 Array programming with numpy. Nature 585 (7825), 357–362, arXiv: 2006.10256 publisher: Springer US.
  * Helander (2014) Helander, Per 2014 Theory of plasma confinement in non-axisymmetric magnetic fields. Reports on Progress in Physics 77 (8).
  * Helander et al. (2012) Helander, P., Beidler, C. D., Bird, T. M., Drevlak, M., Feng, Y., Hatzky, R., Jenko, F., Kleiber, R., Proll, J. H.E., Turkin, Y. & et al. 2012 Stellarator and tokamak plasmas: A comparison. Plasma Physics and Controlled Fusion 54 (12), 1–33.
  * Hirshman & Whitson (1983a) Hirshman, S. P. & Whitson, J. C. 1983a Steepest-descent moment method for three-dimensional magnetohydrodynamic equilibria. Physics of Fluids 26 (12), 3553–3568.
  * Hirshman & Whitson (1983b) Hirshman, S. P. & Whitson, J. C. 1983b Steepest-descent moment method for three-dimensional magnetohydrodynamic equilibria. The Physics of Fluids 26 (12), 3553–3568, publisher: American Institute of Physics.
  * Howell (2009) Howell, Jason S. 2009 Computation of viscoelastic fluid flows using continuation methods. Journal of Computational and Applied Mathematics 225 (1), 187–201.
  * Hudson et al. (2012) Hudson, S R, Dewar, R L, Hole, M J & McGann, M 2012 Non-axisymmetric, multi-region relaxed magnetohydrodynamic equilibrium solutions. Plasma Physics and Controlled Fusion 54 (1), 014005.
  * Landreman et al. (2021) Landreman, Matt, Medasani, Bharat, Wechsung, Florian, Giuliani, Andrew, Jorge, Rogerio & Zhu, Caoxiang 2021 Simsopt: A flexible framework for stellarator optimization. Journal of Open Source Software 6 (65), 3525.
  * Lazerson et al. (2020) Lazerson, Samuel, Schmitt, John, Zhu, Caoxiang, Breslau, Joshua, STELLOPT Developers, All & of Science, USDOE Office 2020 Stellopt, version 2.7.5.
  * McGreivy et al. (2020) McGreivy, Nick, Hudson, Stuart R & Zhu, Caoxiang 2020 Optimized finite-build stellarator coils using automatic differentiation. arXiv ArXiv: 2009.00196.
  * Nemov et al. (2008) Nemov, V. V., Kasilov, S. V., Kernbichler, W. & Leitold, G. O. 2008 Poloidal motion of trapped particle orbits in real-space coordinates. Physics of Plasmas 15 (5), 052501\.
  * Nocedal & Wright (2006) Nocedal, Jorge & Wright, Stephen 2006 Numerical optimization. Springer Science & Business Media.
  * Panici et al. (2022) Panici, Dario, Conlin, Rory, Dudt, Daniel W. & Kolemen, Egemen 2022 The desc stellarator code suite part i: Quick and accurate equilibria computations .
  * Park et al. (2007) Park, Jong Kyu, Boozer, Allen H. & Glasser, Alan H. 2007 Computation of three-dimensional tokamak and spherical torus equilibria. Physics of Plasmas 14 (5), 0–9.
  * Park et al. (2009) Park, Jong-kyu, Boozer, Allen H, Menard, Jonathan E, Garofalo, Andrea M, Schaffer, Michael J, Kaye, Stanley M, Gerhardt, Stefan P & Sabbagh, Steve A 2009 Importance of plasma response to nonaxisymmetric perturbations in tokamaks. Phys. Plasmas p. 12.
  * Richter & DeCarlo (1983) Richter, Stephen L & DeCarlo, Raymond A 1983 Continuation methods: Theory and applications. IEEE Transactions on Systems, Man, and Cybernetics (4), 459–464.
  * Spong et al. (1998) Spong, D. A., Hirshman, S. P., Whitson, J. C., Batchelor, D. B., Carreras, B. A., Lynch, V. E. & Rome, J. A. 1998 J* optimization of small aspect ratio stellarator/tokamak hybrid devices. Physics of Plasmas 5 (5), 1752–1758.
  * Van Der Walt et al. (2011) Van Der Walt, Stéfan, Colbert, S. Chris & Varoquaux, Gaël 2011 The numpy array: A structure for efficient numerical computation. Computing in Science and Engineering 13 (2), 22–30, arXiv: 1102.1523.
  * Velasco et al. (2021) Velasco, J. L., Calvo, I., Mulas, S., Sanchez, E., Parra, F. I., Cappa, A & team, the W7-X. 2021 A model for the fast evaluation of prompt losses of energetic ions in stellarators. arXiv:2106.05697 [physics] ArXiv: 2106.05697.
  * Zhu et al. (2018a) Zhu, Caoxiang, Hudson, Stuart R., Lazerson, Samuel A., Song, Yuntao & Wan, Yuanxi 2018a Hessian matrix approach for determining error field sensitivity to coil deviations. Plasma Physics and Controlled Fusion 60 (5).
  * Zhu et al. (2018b) Zhu, Caoxiang, Hudson, Stuart R., Song, Yuntao & Wan, Yuanxi 2018b New method to design stellarator coils without the winding surface. Nuclear Fusion 58.