Optimizing for quasisymmetry
In this tutorial it is shown how to optimize the boundary shape of a
VMEC configuration to achieve quasisymmetry. Several variants of this
problem are shown here, including both quasi-helical symmetry and
quasi-axisymmetry. There are two objective functions available in
simsopt to achieve quasisymmetry, and both are demonstrated here. One
objective function is
simsopt.mhd.QuasisymmetryRatioResidual
, based
on the fact that \((\vec{B}\times\nabla B
\cdot\nabla\psi)/(\vec{B}\cdot\nabla B)\) is a flux function in
quasisymmetry, demonstrated in the paper arXiv:2108.03711. The other objective function,
using simsopt.mhd.Quasisymmetry
, is a more traditional
one based on the Fourier modes in Boozer coordinates \(B_{m,n}\)
that break the symmetry. The recommended approach is to use
QuasisymmetryRatioResidual
, since
it has the modest advantage of not requiring a transformation to
Boozer coordinates, although both methods are effective. The first two
examples on this page do not require the booz_xform
package,
whereas the third example does.
For these optimization examples, you probably want to write a script that is submitted to the queue on a computing cluster, where they are likely to take minutes or tens of minutes on one node.
Fixed resolution
This example is also available as QH_fixed_resolution.py
in the
examples/2_Intermediate
directory. As usual, a driver script begins with
imports of the classes and functions we will need:
import numpy as np
from simsopt.util import MpiPartition
from simsopt.mhd import Vmec, QuasisymmetryRatioResidual
from simsopt.objectives import LeastSquaresProblem
from simsopt.solve import least_squares_mpi_solve
For this problem we will want MPI for parallelized finite difference
gradients. As explained below, this particular problem has 24
independent variables, so we can take advantage of 24 + 1 = 25 concurrent
function evaluations for one-sided differences. It is therefore
optimal to divide the MPI processes into 25 worker groups. For more
information about MPI and worker groups, see MPI partitions and worker groups or
MpiPartition
. In our script, we therefore
use:
mpi = MpiPartition(25)
It is not necessary to use 25 groups - the code will function properly for any choice. If you select more than 25 groups, the groups beyond the 25th will sit idle. If you select fewer than 25 groups, the 25 function evaluations needed for a finite difference gradient will be distributed among the groups that are available. There is no need to make the number of groups a multiple of the number of available MPI processes, although there cannot be more groups than processes.
We next initialize a VMEC object from an input file:
vmec = Vmec("input.nfp4_QH_warm_start", mpi=mpi)
This file can be found in the examples/2_Intermediate/inputs
directory. The file describes a configuration that has already been
partially optimized for quasi-helical symmetry in a very small
parameter space, keeping poloidal and toroidal mode numbers (the
latter divided by the number of field periods) only up through 1:
In the present example we will refine the configuration by optimizing
in a larger parameter space, with poloidal and toroidal mode numbers
up through 2. To define this parameter space, we set the fixed
property of the boundary’s Fourier modes as follows:
# Define parameter space:
surf = vmec.boundary
surf.fix_all()
max_mode = 2
surf.fixed_range(mmin=0, mmax=max_mode,
nmin=-max_mode, nmax=max_mode, fixed=False)
surf.fix("rc(0,0)") # Major radius
The above code first fixes all modes of the boundary, since we want the mode numbers greater than 2 to all be fixed. Then the desired range of modes is set to be not fixed. This range includes the m=n=0 mode which is essentially the mean major radius. We don’t need or want to vary the overall scale of the configuration, so it is convenient to remove this mode from the parameter space. To confirm which modes will be varied, you can print out the names of the free degrees of freedom (dofs):
print('Parameter space:', surf.dof_names)
The result is
Parameter space: ['SurfaceRZFourier1:rc(0,1)', 'SurfaceRZFourier1:rc(0,2)',
'SurfaceRZFourier1:rc(1,-2)', 'SurfaceRZFourier1:rc(1,-1)',
'SurfaceRZFourier1:rc(1,0)', 'SurfaceRZFourier1:rc(1,1)',
'SurfaceRZFourier1:rc(1,2)', 'SurfaceRZFourier1:rc(2,-2)',
'SurfaceRZFourier1:rc(2,-1)', 'SurfaceRZFourier1:rc(2,0)',
'SurfaceRZFourier1:rc(2,1)', 'SurfaceRZFourier1:rc(2,2)',
'SurfaceRZFourier1:zs(0,1)', 'SurfaceRZFourier1:zs(0,2)',
'SurfaceRZFourier1:zs(1,-2)', 'SurfaceRZFourier1:zs(1,-1)',
'SurfaceRZFourier1:zs(1,0)', 'SurfaceRZFourier1:zs(1,1)',
'SurfaceRZFourier1:zs(1,2)', 'SurfaceRZFourier1:zs(2,-2)',
'SurfaceRZFourier1:zs(2,-1)', 'SurfaceRZFourier1:zs(2,0)',
'SurfaceRZFourier1:zs(2,1)', 'SurfaceRZFourier1:zs(2,2)']
Next, we need to configure a term in the objective function to represent the departure from quasisymmetry. This can be done as follows:
# Configure quasisymmetry objective:
qs = QuasisymmetryRatioResidual(vmec,
np.arange(0, 1.01, 0.1), # Radii to target
helicity_m=1, helicity_n=-1) # (M, N) you want in |B|
There are several adjustable options, the details of which can be
found in the API documentation for
QuasisymmetryRatioResidual
.
There you can also find the mathematical expression for the objective
function. The second argument to
QuasisymmetryRatioResidual
above
sets the quasisymmetry objective to be evaluated at a uniform grid of
11 surfaces [0, 0.1, 0.2, ..., 1]
in the normalized toroidal flux
\(s\), with the result that quasisymmetry is targeted throughout
the volume. You are free to provide different values, or a single
float if you only want to target quasisymmetry on a single
surface. There is also an optional argument weights
if you wish to
have different weights in the objective function for quasisymmetry on
different surfaces. The helicity_n
argument can also be set to
+1
rather than -1
for quasi-helical symmetry, amounting to a
mirror-reversal, though the initial configuration used here is
consistent with the -1
handedness.
We are now ready to define the total objective function. Here we will include quasisymmetry and aspect ratio. Aspect ratio must be included because otherwise quasisymmetry can be made arbitrarily good by increasing the aspect ratio to infinity. The simsopt objective function is defined as follows:
# Define objective function
prob = LeastSquaresProblem.from_tuples([(vmec.aspect, 7, 1),
(qs.residuals, 0, 1)])
It can be seen that we are targeting an aspect ratio of 7. This
objective function will be a sum of 44,353 least-squares terms, 44,352
of which correspond to the quasisymmetry residual on 63x64 grid points
on the 11 flux surfaces targeted, plus one additional term
(vmec.aspect - 7) ** 2
. (The 63x64 resolution is a default in
QuasisymmetryRatioResidual
.) This
large number of residual terms is no problem - it introduces
negligible computational cost compared to the cost of the equilibrium
calculations, so we may as well use this high resolution.
You can check the value of the objective functions before the
optimization. Rather than print each residual term, the scalar total
for the quasisymmetry term can be obtained with the .total()
method.
print("Quasisymmetry objective before optimization:", qs.total())
print("Total objective before optimization:", prob.objective())
The results are both 0.304, since the aspect ratio term is negligible.
Finally, we solve the optimization problem:
least_squares_mpi_solve(prob, mpi, grad=True)
Suppose you have written the above commands in a file named
simsopt_driver
. Depending on your computing system, the script
can be run using a command like srun python simsopt_driver
(for
SLURM systems) or mpirun -n 25 simsopt_driver
.
Since this objective function has multiple local minima, the final result of the optimization can be sensitive to small changes in simsopt, VMEC, or the packages they depend on. Therefore you will not necessarily obtain exactly the result shown here. But one result produced by this optimization script is the following configuration:
This last figure shows that reasonably good quasisymmetry has been
achieved. The quality of quasisymmetry can be improved significantly
by further refining the configuration using one or more rounds of
optimization with more Fourier modes in the parameter space. Printing
qs.total()
or prob.objective()
at the end of the optimization,
it can be seen that both have been reduced significantly, to 0.00794
for the result shown here.
Dynamic resolution
Since simsopt optimization problems are defined using a python script,
you are free to add other scripting in your problem definition. Here
we show how this capability can be used to increase the numerical
resolution of codes such as VMEC during the optimization. At the same
time, we will increase the number of Fourier modes in the parameter
space during the optimization. This example can also be found in the
examples/2_Intermediate
directory as
resolution_increase.py
. This example is very similar to the
quasi-axisymmetry optimization shown in arXiv:2108.03711.
As usual, we begin with the necessary imports:
import numpy as np
from simsopt.util import MpiPartition
from simsopt.mhd import Vmec, QuasisymmetryRatioResidual
from simsopt.objectives import LeastSquaresProblem
from simsopt.solve import least_squares_mpi_solve
We again split the pool of MPI processes into worker groups. Here, for simplicity, we make each process its own worker group, by omitting the argument:
mpi = MpiPartition()
We initialize a VMEC configuration from an input file. This starting configuration is axisymmetric with a circular cross-section, so we are starting “from scratch”:
vmec = Vmec("input.nfp2_QA", mpi=mpi)
This input file can be found in the examples/2_Intermediate/inputs
directory. We define the quasisymmetry objective as in the previous
section, except that we specify a helicity of (1,0) instead of (1,1)
or (1,-1) to get quasi-axisymmetry instead of quasi-helical symmetry:
# Configure quasisymmetry objective:
qs = QuasisymmetryRatioResidual(vmec,
np.arange(0, 1.01, 0.1), # Radii to target
helicity_m=1, helicity_n=0) # (M, N) you want in |B|
We now define the total objective function. For this example, it is necessary to include a nonzero target value for the rotational transform in the objective, to prevent the optimum from being truly axisymmetric:
# Define objective function
prob = LeastSquaresProblem.from_tuples([(vmec.aspect, 6, 1),
(vmec.mean_iota, 0.42, 1),
(qs, 0, 1)])
It can be seen here that we are seeking a configuration with aspect
ratio 6, and average iota of 0.42, slightly above the resonance at 2 /
5 = 0.4. The function simsopt.mhd.Vmec.mean_iota()
used
here returns \(\int_0^1 \iota\, ds\) where \(s\) is the
toroidal flux normalized by its value at the VMEC boundary.
Now, we set up a loop over several optimization steps. At each step,
the resolution parameters mpol
and ntor
for VMEC increase. At
the same time, in each optimization step a larger range of poloidal
and toroidal mode numbers are set to be varied in the optimization:
for step in range(4):
max_mode = step + 1
# VMEC's mpol & ntor will be 3, 4, 5, 6:
vmec.indata.mpol = 3 + step
vmec.indata.ntor = vmec.indata.mpol
if mpi.proc0_world:
print("Beginning optimization with max_mode =", max_mode, \
", vmec mpol=ntor=", vmec.indata.mpol, \
". Previous vmec iteration = ", vmec.iter)
# Define parameter space:
surf.fix_all()
surf.fixed_range(mmin=0, mmax=max_mode,
nmin=-max_mode, nmax=max_mode, fixed=False)
surf.fix("rc(0,0)") # Major radius
# Carry out the optimization for this step:
least_squares_mpi_solve(prob, mpi, grad=True)
if mpi.proc0_world:
print("Done optimization with max_mode =", max_mode, \
". Final vmec iteration = ", vmec.iter)
If you like, other parameters could be adjusted at each step too, such as the radial resolution or number of iterations in VMEC, the solver tolerances, or the maximum number of iteration of the optimization algorithm.
As in the previous section, the final result of this optimization can be sensitive to small changes in simsopt, VMEC, or the packages they depend on. Therefore you will not necessarily obtain exactly the result shown here. But one result produced by this optimization script is the following configuration:
Bmn objective
Here we show an alternative method of quasisymmetry optimization using
a different objective function,
simsopt.mhd.Quasisymmetry
, based on the
symmetry-breaking Fourier mode aplitudes \(B_{m,n}\) in Boozer
coordinates. This example can also be found in the
examples/2_Intermediate
directory as
resolution_increase_boozer.py
.
In this case, the imports needed are:
from simsopt.util import MpiPartition
from simsopt.mhd import Vmec, Boozer, Quasisymmetry
from simsopt.objectives import LeastSquaresProblem
from simsopt.solve import least_squares_mpi_solve
We again split the pool of MPI processes into worker groups and
initialize a Vmec
object as in the previous example:
mpi = MpiPartition()
vmec = Vmec("input.nfp2_QA", mpi=mpi)
This input file, corresponding to an axisymmetric torus with circular
cross-section, can be found in the examples/2_Intermediate/inputs
directory. Next, this alternative quasisymmetry objective can be
created as follows:
# Configure quasisymmetry objective:
boozer = Boozer(vmec)
qs = Quasisymmetry(boozer,
0.5, # Radius to target
1, 0) # (M, N) you want in |B|
There are several adjustable options, the details of which can be
found in the API documentation for Boozer
and Quasisymmetry
. The numerical resolution
of the Boozer-coordinate transformation can be adjusted by passing
parameters to the Boozer
constructor, as in
Boozer(vmec, mpol=64, ntor=32)
. The second argument to
Quasisymmetry
above sets the quasisymmetry objective to be
evaluated at normalized toroidal flux of 0.5, but you are free to
provide different values. Or, a list of values can be provided to
target quasisymmetry on multiple surfaces. The
Quasisymmetry
also has optional arguments
to adjust the normalization and weighting of different Fourier modes.
We now define the total objective function. As with the previous quasi-axisymmetry example, it is necessary to include a nonzero target value for the rotational transform in the objective, to prevent the optimum from being truly axisymmetric. Here we will constrain iota at the edge and magnetic axis, in order to prescribe the magnetic shear:
# Define objective function
prob = LeastSquaresProblem.from_tuples([(vmec.aspect, 6, 1),
(vmec.iota_axis, 0.465, 1),
(vmec.iota_edge, 0.495, 1),
(qs, 0, 1)])
It can be seen here that we are seeking a configuration with aspect ratio 6, and iota slightly below 0.5.
Now, we set up a loop over several optimization steps. At each step,
the resolution parameters mpol
and ntor
for VMEC increase, as
do the the Fourier resolution parameters for booz_xform
. At the
same time, in each optimization step a larger range of poloidal and
toroidal mode numbers are set to be varied in the optimization:
for step in range(4):
max_mode = step + 1
# VMEC's mpol & ntor will be 3, 4, 5, 6:
vmec.indata.mpol = 3 + step
vmec.indata.ntor = vmec.indata.mpol
# booz_xform's mpol & ntor will be 16, 24, 32, 40:
boozer.mpol = 16 + step * 8
boozer.ntor = boozer.mpol
if mpi.proc0_world:
print("Beginning optimization with max_mode =", max_mode, \
", vmec mpol=ntor=", vmec.indata.mpol, \
", boozer mpol=ntor=", boozer.mpol, \
". Previous vmec iteration = ", vmec.iter)
# Define parameter space:
surf.fix_all()
surf.fixed_range(mmin=0, mmax=max_mode,
nmin=-max_mode, nmax=max_mode, fixed=False)
surf.fix("rc(0,0)") # Major radius
# Carry out the optimization for this step:
least_squares_mpi_solve(prob, mpi, grad=True)
if mpi.proc0_world:
print("Done optimization with max_mode =", max_mode, \
". Final vmec iteration = ", vmec.iter)
If you like, other parameters could be adjusted at each step too, such as the radial resolution or number of iterations in VMEC, the solver tolerances, or the maximum number of iteration of the optimization algorithm.
As with the previous examples, the final result of this optimization can be sensitive to small changes in simsopt, VMEC, or the packages they depend on. Therefore you will not necessarily obtain exactly the result shown here. But one result produced by this optimization script is the following configuration: