Single stage optimization
In this tutorial it is shown how to optimize the boundary shape of a
VMEC configuration to achieve quasisymmetry at the same time that the
shape of coils are optimized to reproduce the VMEC field. The approach
employed here follows the work in the paper arXiv:2302.10622 and includes both a vacuum and a
finite beta implementation. This generalizes the “stage 1” and “stage 2”
approaches in previous tutorials to a “single stage” approach. In this
case, instead of optimizing the simsopt.mhd.QuasisymmetryRatioResidual
or the coil optimization residual simsopt.geo.SquaredFlux individually,
they are both summed together to optimize a single objective function
(cost function) given by
where the coil regularization terms are a sum of simsopt.geo.CurveLength for
the length of the coils, simsopt.geo.CurveCurveDistance for the distance
between coils, simsopt.geo.LpCurveCurvature for the coil curvature,
simsopt.geo.MeanSquaredCurvature for the coil mean squared curvature and
simsopt.geo.ArclengthVariation for the arclength variation of the coils.
In order to minimize the number of single stage iterations, in this script,
an initial stage 2 optimization is performed.
The finite beta optimization script requires the Virtual Casing module
from the hiddenSymmetries repository.
Vacuum optimization
This example is available as single_stage_optimization.py in the
examples/3_Advanced directory. As usual, a driver script begins with
imports of the classes and functions we will need
import os
import numpy as np
from mpi4py import MPI
from pathlib import Path
from scipy.optimize import minimize
from simsopt.util import MpiPartition
from simsopt._core.derivative import Derivative
from simsopt.mhd import Vmec, QuasisymmetryRatioResidual
from simsopt._core.finite_difference import MPIFiniteDifference
from simsopt.field import BiotSavart, Current, coils_via_symmetries
from simsopt.objectives import SquaredFlux, QuadraticPenalty, LeastSquaresProblem
from simsopt.geo import (CurveLength, CurveCurveDistance, MeanSquaredCurvature,
LpCurveCurvature, ArclengthVariation, curves_to_vtk, create_equally_spaced_curves)
For a description of each class and function imported, please see the stage 1
script QH_fixed_resolution.py in the examples/2_Intermediate directory
and the stage 2 script stage_two_optimization_minimal.py in the examples/1_Simple
directory. For convenience, a print function is defined that only prints once even when
the code is ran in parallel:
def pprint(*args, **kwargs):
if comm.rank == 0:
print(*args, **kwargs)
The input parameters are given by
MAXITER_stage_2 = 10, which selects the maximum number of iterations for the initial stage 2 optimization.MAXITER_single_stage = 10, which selects the maximum number of iterations for the single stage optimization.max_mode = 1, which selects the maximum poloidal and toroidal modes on the surface being optimized with VMEC.vmec_input_filename = os.path.join(parent_path, 'inputs', 'input.nfp4_QH_warm_start'), which represents the initial VMEC input file.ncoils = 3: the number of coils per field period.aspect_ratio_target = 7.0: target aspect ratio for the VMEC surface.coils_objective_weight = 1e+3: the weight given to the coils objective function with respect to the stage 1 optimization.
The remaining input parameters follow the convention of the stage 2 optimization script.
Then, the results directory are created to hold the VMEC configurations and the coils
directory = 'optimization_QH'
vmec_verbose = False
# Create output directories
this_path = os.path.join(parent_path, directory)
os.makedirs(this_path, exist_ok=True)
os.chdir(this_path)
vmec_results_path = os.path.join(this_path, "vmec")
coils_results_path = os.path.join(this_path, "coils")
if comm.rank == 0:
os.makedirs(vmec_results_path, exist_ok=True)
os.makedirs(coils_results_path, exist_ok=True)
The function fun_coils returns the objective function and gradients
used in the initial stage 2 optimization, while the fun function
returns the objective function and gradients used in the single stage
optimization. In this function, the derivatives with respect to the coils
and to the surface are computed separately. The derivatives with respect
to the coils are analytical, while the derivatives with respect to the surface
are a mix of analytical (defined as mixed_dJ) and finite-diference
derivatives
def fun(dofs, prob_jacobian=None, info={'Nfeval': 0}):
info['Nfeval'] += 1
JF.x = dofs[:-number_vmec_dofs]
prob.x = dofs[-number_vmec_dofs:]
bs.set_points(surf.gamma().reshape((-1, 3)))
os.chdir(vmec_results_path)
J_stage_1 = prob.objective()
J_stage_2 = coils_objective_weight * JF.J()
J = J_stage_1 + J_stage_2
if J > JACOBIAN_THRESHOLD or np.isnan(J):
pprint(f"Exception caught during function evaluation with J={J}."
f" Returning J={JACOBIAN_THRESHOLD}")
J = JACOBIAN_THRESHOLD
grad_with_respect_to_surface = [0] * number_vmec_dofs
grad_with_respect_to_coils = [0] * len(JF.x)
else:
pprint(f"fun#{info['Nfeval']}: Objective function = {J:.4f}")
prob_dJ = prob_jacobian.jac(prob.x)
## Finite differences for the second-stage objective function
coils_dJ = JF.dJ()
## Mixed term - derivative of squared flux with respect to the surface shape
n = surf.normal()
absn = np.linalg.norm(n, axis=2)
B = bs.B().reshape((nphi_VMEC, ntheta_VMEC, 3))
dB_by_dX = bs.dB_by_dX().reshape((nphi_VMEC, ntheta_VMEC, 3, 3))
Bcoil = bs.B().reshape(n.shape)
unitn = n * (1./absn)[:, :, None]
Bcoil_n = np.sum(Bcoil*unitn, axis=2)
mod_Bcoil = np.linalg.norm(Bcoil, axis=2)
B_n = Bcoil_n
B_diff = Bcoil
B_N = np.sum(Bcoil * n, axis=2)
assert Jf.definition == "local"
dJdx = (B_n/mod_Bcoil**2)[:, :, None] * (np.sum(dB_by_dX*(n-B*(B_N/mod_Bcoil**2)[:, :, None])[:, :, None, :], axis=3))
dJdN = (B_n/mod_Bcoil**2)[:, :, None] * B_diff - 0.5 * (B_N**2/absn**3/mod_Bcoil**2)[:, :, None] * n
deriv = surf.dnormal_by_dcoeff_vjp(dJdN/(nphi_VMEC*ntheta_VMEC)) + surf.dgamma_by_dcoeff_vjp(dJdx/(nphi_VMEC*ntheta_VMEC))
mixed_dJ = Derivative({surf: deriv})(surf)
## Put both gradients together
grad_with_respect_to_coils = coils_objective_weight * coils_dJ
grad_with_respect_to_surface = np.ravel(prob_dJ) + coils_objective_weight * mixed_dJ
grad = np.concatenate((grad_with_respect_to_coils, grad_with_respect_to_surface))
return J, grad
The initial stage 2 optimization is then performed at the line
res = minimize(fun_coils, dofs[:-number_vmec_dofs], jac=True,
args=({'Nfeval': 0}), method='L-BFGS-B',
options={'maxiter': MAXITER_stage_2, 'maxcor': 300},
tol=1e-12)
while the single stage optimization is performed at
with MPIFiniteDifference(prob.objective, mpi,
diff_method=diff_method,
abs_step=finite_difference_abs_step,
rel_step=finite_difference_rel_step) as prob_jacobian:
if mpi.proc0_world:
res = minimize(fun, dofs,
args=(prob_jacobian, {'Nfeval': 0}),
jac=True, method='BFGS',
options={'maxiter': MAXITER_single_stage},
tol=1e-15)
The results are then printed and stored in files.
Finite beta optimization
The finite beta generalization example is available as
single_stage_optimization_finite_beta.py in the
examples/3_Advanced directory. In addition to the parameters in the
previous example, the finite beta script uses the Virtual Casing principle
to decouple the plasma magnetic field from the coil magnetic field.
The VirtualCasing module is imported in
from simsopt.mhd import Vmec, QuasisymmetryRatioResidual, VirtualCasing
and its resolution is set in vc_src_nphi = ntheta_VMEC.
The initialization of the VirtualCasing is performed at the line
vc = VirtualCasing.from_vmec(
vmec, src_nphi=vc_src_nphi, trgt_nphi=nphi_VMEC,
trgt_ntheta=ntheta_VMEC)
total_current_vmec = vmec.external_current() / (2 * surf.nfp)
Now the gradients of the objective function are computed using
finite differences instead of a mix of analytical and finite difference
derivatives. The objective function is then wrapped in the fun_J function
def fun_J(prob, coils_prob):
global previous_surf_dofs
J_stage_1 = prob.objective()
if np.any(previous_surf_dofs != prob.x): # Only run virtual casing if surface dofs have changed
previous_surf_dofs = prob.x
try:
vc = VirtualCasing.from_vmec(
vmec, src_nphi=vc_src_nphi, trgt_nphi=nphi_VMEC,
trgt_ntheta=ntheta_VMEC)
Jf.target = vc.B_external_normal
except ObjectiveFailure as e:
pass
bs.set_points(surf.gamma().reshape((-1, 3)))
J_stage_2 = coils_objective_weight * JF.J()
J = J_stage_1 + J_stage_2
return J
And the resulting objective function and gradients are computed using
the fun function
def fun(dofss, prob_jacobian, info={'Nfeval': 0}):
info['Nfeval'] += 1
os.chdir(vmec_results_path)
prob.x = dofss[-number_vmec_dofs:]
coil_dofs = dofss[:-number_vmec_dofs]
# Un-fix the desired coil dofs so they can be updated:
JF.full_unfix(free_coil_dofs)
JF.x = coil_dofs
J = fun_J(prob, JF)
if J > JACOBIAN_THRESHOLD or isnan(J):
pprint(f"fun#{info['Nfeval']}: Exception caught during function evaluation with J={J}. Returning J={JACOBIAN_THRESHOLD}")
J = JACOBIAN_THRESHOLD
grad_with_respect_to_surface = [0] * number_vmec_dofs
grad_with_respect_to_coils = [0] * len(coil_dofs)
else:
pprint(f"fun#{info['Nfeval']}: Objective function = {J:.4f}")
coils_dJ = JF.dJ()
grad_with_respect_to_coils = coils_objective_weight * coils_dJ
JF.fix_all() # Must re-fix the coil dofs before beginning the finite differencing.
grad_with_respect_to_surface = prob_jacobian.jac(prob.x)[0]
JF.fix_all()
grad = np.concatenate((grad_with_respect_to_coils,
grad_with_respect_to_surface))
return J, grad
The initial stage 2 optimization and single stage optimization follow the previous vacuum case, with the exception of the lines
dofs[:-number_vmec_dofs] = res.x
JF.x = dofs[:-number_vmec_dofs]
mpi.comm_world.Bcast(dofs, root=0)
opt = make_optimizable(fun_J, prob, JF)
free_coil_dofs = JF.dofs_free_status
JF.fix_all()
where the coils and surface degrees of freedom are defined and MPI broadcasted and
JF.full_unfix(free_coil_dofs) # Needed to evaluate JF.dJ
where the coils degrees of freedom are unfixed to evaluate their Jacobian.