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 :obj:`simsopt.mhd.QuasisymmetryRatioResidual` or the coil optimization residual :obj:`simsopt.geo.SquaredFlux` individually, they are both summed together to optimize a single objective function (cost function) given by
J = \text{Quasisymmetry} + \text{coils_objective_weight}*\frac{1}{2} \int |\vec{B} \cdot \vec{n}|^2 ds \\+ \text{coil regularization terms}.
where the coil regularization terms are a sum of :obj:`simsopt.geo.CurveLength` for the length of the coils, :obj:`simsopt.geo.CurveCurveDistance` for the distance between coils, :obj:`simsopt.geo.LpCurveCurvature` for the coil curvature, :obj:`simsopt.geo.MeanSquaredCurvature` for the coil mean squared curvature and :obj:`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.
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.
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.