def test_2dof_surface_opt(self):
"""
Optimize the minor radius and elongation of an axisymmetric torus to
obtain a desired volume and area.
"""
for solver in solvers:
desired_volume = 0.6
desired_area = 8.0
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
surf = optimizable(SurfaceRZFourier())
# Set initial surface shape. It helps to make zs(1,0) larger
# than rc(1,0) since there are two solutions to this
# optimization problem, and for testing we want to find one
# rather than the other.
surf.set_zs(1, 0, 0.2)
# Parameters are all non-fixed by default, meaning they will be
# optimized. You can choose to exclude any subset of the variables
# from the space of independent variables by setting their 'fixed'
# property to True.
surf.set_fixed('rc(0,0)')
# Each function you want in the objective function is then
# equipped with a shift and weight, to become a term in a
# least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
prob = LeastSquaresProblem([(surf.volume, desired_volume, 1),
(surf.area, desired_area, 1)])
# Verify the state vector and names are what we expect
np.testing.assert_allclose(prob.x, [0.1, 0.2])
self.assertEqual(prob.dofs.names[0][:28],
'rc(1,0) of SurfaceRZFourier ')
self.assertEqual(prob.dofs.names[1][:28],
'zs(1,0) of SurfaceRZFourier ')
# Solve the minimization problem:
solver(prob)
# Check results
self.assertAlmostEqual(surf.get_rc(0, 0), 1.0, places=13)
self.assertAlmostEqual(surf.get_rc(1, 0),
0.10962565115956417,
places=13)
self.assertAlmostEqual(surf.get_zs(0, 0), 0.0, places=13)
self.assertAlmostEqual(surf.get_zs(1, 0),
0.27727411213693337,
places=13)
self.assertAlmostEqual(surf.volume(), desired_volume, places=8)
self.assertAlmostEqual(surf.area(), desired_area, places=8)
self.assertLess(np.abs(prob.objective()), 1.0e-15)
def test_stellopt_scenarios_1DOF_circularCrossSection_varyR0_targetVolume(
self):
"""
This script implements the "1DOF_circularCrossSection_varyR0_targetVolume"
example from
https://github.com/landreman/stellopt_scenarios
This optimization problem has one independent variable, representing
the mean major radius. The problem also has one objective: the plasma
volume. There is not actually any need to run an equilibrium code like
VMEC since the objective function can be computed directly from the
boundary shape. But this problem is a fast way to test the
optimization infrastructure with VMEC.
Details of the optimum and a plot of the objective function landscape
can be found here:
https://github.com/landreman/stellopt_scenarios/tree/master/1DOF_circularCrossSection_varyR0_targetVolume
"""
#logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger('[{}]'.format(MPI.COMM_WORLD.Get_rank()) +
__name__)
for ngroups in range(1, 1 + MPI.COMM_WORLD.Get_size()):
for grad in [False, True]:
# In the next line, we can adjust how many groups the pool of MPI
# processes is split into.
mpi = MpiPartition(ngroups=ngroups)
mpi.write()
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
equil = Vmec(mpi=mpi)
surf = equil.boundary
# Set the initial boundary shape. Here is one syntax:
surf.set('rc(0,0)', 1.0)
# Here is another syntax:
surf.set_rc(0, 1, 0.1)
surf.set_zs(0, 1, 0.1)
surf.set_rc(1, 0, 0.1)
surf.set_zs(1, 0, 0.1)
# VMEC parameters are all fixed by default, while surface parameters are all non-fixed by default.
# You can choose which parameters are optimized by setting their 'fixed' attributes.
surf.all_fixed()
surf.set_fixed('rc(0,0)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function
desired_volume = 0.15
term1 = LeastSquaresTerm(equil.volume, desired_volume, 1)
# A list of terms are combined to form a nonlinear-least-squares
# problem.
prob = LeastSquaresProblem([term1], mpi)
# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
prob.solve(grad=grad)
# Make sure all procs call VMEC:
objective = prob.objective()
if mpi.proc0_world:
print("At the optimum,")
print(" rc(m=0,n=0) = ", surf.get_rc(0, 0))
print(" volume, according to VMEC = ", equil.volume())
print(" volume, according to Surface = ", surf.volume())
print(" objective function = ", objective)
assert np.abs(surf.get_rc(0, 0) - 0.7599088773175) < 1.0e-5
assert np.abs(equil.volume() - 0.15) < 1.0e-6
assert np.abs(surf.volume() - 0.15) < 1.0e-6
assert prob.objective() < 1.0e-15
#equil.VMEC.reinit()
equil.finalize()
def test_2dof_surface_Garabedian_opt(self):
"""
Optimize the minor radius and elongation of an axisymmetric torus
to obtain a desired volume and area, optimizing in the space
of Garabedian coefficients.
"""
for solver in solvers:
desired_volume = 0.6
desired_area = 8.0
# Start with a default surface, which is axisymmetric with
# major radius 1 and minor radius 0.1. Setting mmax=2
# allows elongation to be added.
surf = SurfaceGarabedian(mmax=2)
# Set initial surface shape. It helps to make the initial
# surface shape slightly different from the default one
# since there are two solutions to this optimization
# problem, and for testing we want to find one rather than
# the other.
surf.set_Delta(2, 0, -0.1)
# Parameters are all non-fixed by default, meaning they will be
# optimized. You can choose to exclude any subset of the variables
# from the space of independent variables by setting their 'fixed'
# property to True.
surf.all_fixed()
surf.set_fixed('Delta(0,0)', False) # Minor radius
surf.set_fixed('Delta(2,0)', False) # Elongation
# Each function you want in the objective function is then
# equipped with a shift and weight, to become a term in a
# least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
prob = LeastSquaresProblem([(surf.volume, desired_volume, 1),
(surf.area, desired_area, 1)])
# Verify the state vector and names are what we expect
np.testing.assert_allclose(prob.x, [0.1, -0.1])
self.assertEqual(prob.dofs.names[0][:31],
'Delta(0,0) of SurfaceGarabedian')
self.assertEqual(prob.dofs.names[1][:31],
'Delta(2,0) of SurfaceGarabedian')
# Solve the minimization problem:
solver(prob)
# Check results
self.assertAlmostEqual(surf.get_Delta(0, 0),
0.193449881648249,
places=11)
self.assertAlmostEqual(surf.get_Delta(1, 0), 1.0, places=13)
self.assertAlmostEqual(surf.get_Delta(2, 0),
-0.083824230488685,
places=11)
self.assertAlmostEqual(surf.volume(), desired_volume, places=8)
self.assertAlmostEqual(surf.area(), desired_area, places=8)
# Convert the SurfaceGarabedian to a SurfaceRZFourier and
# make sure its properties match those of the optimization
# direction in RZFourier-space.
surfRZ = surf.to_RZFourier()
self.assertAlmostEqual(surfRZ.get_rc(0, 0), 1.0, places=13)
self.assertAlmostEqual(surfRZ.get_rc(1, 0),
0.10962565115956417,
places=11)
self.assertAlmostEqual(surfRZ.get_zs(0, 0), 0.0, places=13)
self.assertAlmostEqual(surfRZ.get_zs(1, 0),
0.27727411213693337,
places=11)
self.assertAlmostEqual(surfRZ.volume(), desired_volume, places=8)
self.assertAlmostEqual(surfRZ.area(), desired_area, places=8)
self.assertLess(np.abs(prob.objective()), 1.0e-15)
# Tell the curve object that the first Fourier mode is fixed, whereas
# all the other dofs are not.
curve.all_fixed(False)
curve.fixed[0] = True
# The length objective is a separate object rather than a function of
# Curve itself.
obj = CurveLength(curve)
print('Initial curve length: ', obj.J())
# Each target function is then equipped with a shift and weight, to
# become a term in a least-squares objective function.
# A list of terms are combined to form a nonlinear-least-squares
# problem.
prob = LeastSquaresProblem([(obj, 0.0, 1.0)])
# At the initial condition, get the Jacobian two ways: analytic
# derivatives and finite differencing. The difference should be small.
fd_jac = prob.dofs.fd_jac()
jac = prob.dofs.jac()
print('finite difference Jacobian:')
print(fd_jac)
print('Analytic Jacobian:')
print(jac)
print('Difference:')
print(fd_jac - jac)
# Solve the minimization problem:
least_squares_serial_solve(prob)
desired_area = 8.0
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
surf = SurfaceRZFourier()
# Parameters are all non-fixed by default, meaning they will be
# optimized. You can choose to exclude any subset of the variables
# from the space of independent variables by setting their 'fixed'
# property to True.
surf.set_fixed('rc(0,0)')
# Each target function is then equipped with a shift and weight, to
# become a term in a least-squares objective function
term1 = (surf.volume, desired_volume, 1)
term2 = (surf.area, desired_area, 1)
# A list of terms are combined to form a nonlinear-least-squares
# problem.
prob = LeastSquaresProblem([term1, term2])
# Solve the minimization problem:
least_squares_serial_solve(prob)
print("At the optimum,")
print(" rc(m=1,n=0) = ", surf.get_rc(1, 0))
print(" zs(m=1,n=0) = ", surf.get_zs(1, 0))
print(" volume = ", surf.volume())
print(" area = ", surf.area())
print(" objective function = ", prob.objective())
surf.all_fixed()
surf.set_fixed('rc(1,1)', False)
surf.set_fixed('zs(1,1)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
desired_volume = 0.15
volume_weight = 1
term1 = (spec.volume, desired_volume, volume_weight)
desired_iota = 0.41
iota_weight = 1
term2 = (vmec.iota_axis, desired_iota, iota_weight)
prob = LeastSquaresProblem([term1, term2])
# Solve the minimization problem:
least_squares_mpi_solve(prob, mpi=mpi, grad=True)
# Evaluate quantities on all processes, in case communication is
# required:
final_objective = prob.objective()
vmec_volume = vmec.volume()
spec_volume = spec.volume()
surf_volume = surf.volume()
vmec_iota = vmec.iota_axis()
spec_iota = spec.iota()
if mpi.proc0_world:
logging.info("At the optimum,")
https://github.com/landreman/stellopt_scenarios/tree/master/2DOF_circularCrossSection_varyAxis_targetIotaAndQuasisymmetry
See that website for a detailed description of the problem and plots
of the objective function landscape.
"""
vmec = Vmec(os.path.join(os.path.dirname(__file__), 'inputs', 'input.2DOF_circularCrossSection_varyAxis_targetIotaAndQuasisymmetry'))
# Define parameter space:
vmec.boundary.all_fixed()
vmec.boundary.set_fixed("rc(0,1)", False)
vmec.boundary.set_fixed("zs(0,1)", False)
# Define objective function:
boozer = Boozer(vmec, mpol=32, ntor=16)
qs = Quasisymmetry(boozer,
1.0, # Radius to target
1, 0, # (M, N) you want in |B|
normalization="symmetric",
weight="stellopt_ornl")
# Objective function is 100 * (iota - (-0.41))^2 + 1 * (qs - 0)^2
prob = LeastSquaresProblem([(vmec.iota_axis, -0.41, 100),
(qs, 0, 1)])
least_squares_serial_solve(prob)
print("Final values before shifting and scaling:", prob.dofs.f())
print("Final residuals:", prob.f())
print("Final state vector:", prob.x)
print("Final iota on axis:", vmec.iota_axis())
# Decide which degrees of freedom to optimize
stel.fix_all()
stel.unfix('rc(1)')
stel.unfix('zs(1)')
stel.unfix('etabar')
# Each target function is then equipped with a shift and weight, to
# become a term in a least-squares objective function
term1 = (stel.get_iota, -0.5, 1.0)
term2 = (stel.get_max_elongation, 0.0, 0.0001)
# Note the weight on elongation must be much smaller than the weight on iota!
# A list of terms are combined to form a nonlinear-least-squares
# problem.
prob = LeastSquaresProblem.from_tuples([term1, term2])
print('Before optimization:')
print(' Global state vector: ', prob.x)
print(' QSC dofs: ', stel.get_dofs())
print(' Iota: ', stel.iota)
print(' Max elongation: ', stel.max_elongation)
print(' objective function: ', prob.objective())
# Solve the minimization problem:
least_squares_serial_solve(prob)
print('At the optimum:')
print(' Global state vector: ', prob.x)
print(' QSC dofs: ', stel.get_dofs())
print(' Iota: ', stel.iota)
surf.fixed_range(mmin=0, mmax=2, nmin=-1, nmax=1, fixed=False)
surf.set_fixed("Delta(1,0)") # toroidally-averaged major radius
surf.set_fixed("Delta(0,0)") # toroidally-averaged minor radius
# Define objective function:
boozer = Boozer(vmec, mpol=32, ntor=16)
qs = Quasisymmetry(
boozer,
1.0, # Radius to target
1,
0, # (M, N) you want in |B|
normalization="symmetric",
weight="stellopt")
# Objective function is 100 * (iota - (-0.41))^2 + 1 * (qs - 0)^2
prob = LeastSquaresProblem([(vmec.iota_axis, -0.41, 100), (qs, 0, 1)])
residuals = prob.f()
vals = prob.dofs.f()
if mpi.proc0_world:
print("Initial values before shifting and scaling: ", vals[:10])
print("Initial residuals after shifting and scaling:", residuals[:10])
print("size of residuals:", len(residuals))
print("Initial objective function:", prob.objective())
print("Parameter space:")
for name in prob.dofs.names:
print(name)
print("Initial state vector:", prob.x)
print("Initial iota on axis:", vmec.iota_axis())
#exit(0)
surf.all_fixed()
surf.set_fixed('rc(1,1)', False)
surf.set_fixed('zs(1,1)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
desired_volume = 0.15
volume_weight = 1
term1 = (equil.volume, desired_volume, volume_weight)
desired_iota = 0.41
iota_weight = 1
term2 = (equil.iota_axis, desired_iota, iota_weight)
prob = LeastSquaresProblem([term1, term2])
# Solve the minimization problem:
least_squares_mpi_solve(prob, mpi, grad=True)
objective = prob.objective()
if mpi.proc0_world:
print("At the optimum,")
print(" rc(m=1,n=1) = ", surf.get_rc(1, 1))
print(" zs(m=1,n=1) = ", surf.get_zs(1, 1))
print(" volume, according to VMEC = ", equil.volume())
print(" volume, according to Surface = ", surf.volume())
print(" iota on axis = ", equil.iota_axis())
print(" objective function = ", objective)
assert np.abs(surf.get_rc(1, 1) - 0.0313066948) < 1.0e-3
# Garabedian representation:
surf = equil.boundary.to_Garabedian()
equil.boundary = surf
# VMEC parameters are all fixed by default, while surface parameters
# are all non-fixed by default. You can choose which parameters are
# optimized by setting their 'fixed' attributes.
surf.all_fixed()
surf.set_fixed('Delta(1,-1)', False)
# Each function we want in the objective function is then equipped
# with a shift and weight, to become a term in a least-squares
# objective function. A list of terms are combined to form a
# nonlinear-least-squares problem.
desired_iota = 0.41 # Sign was + for VMEC
prob = LeastSquaresProblem([(equil.iota, desired_iota, 1)])
# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
least_squares_mpi_solve(prob, mpi=mpi, grad=True)
final_objective = prob.objective()
final_iota = equil.iota()
if mpi.proc0_world:
print("At the optimum,")
print(" Delta(m=1,n=-1) = ", surf.get_Delta(1, -1))
print(" iota = ", final_iota)
print(" objective function = ", final_objective)
assert np.abs(surf.get_Delta(1, -1) - 0.08575) < 1.0e-4
surf.set_zs(0, 1, 0.1)
surf.set_rc(1, 0, 0.1)
surf.set_zs(1, 0, 0.1)
# VMEC parameters are all fixed by default, while surface parameters
# are all non-fixed by default. You can choose which parameters are
# optimized by setting their 'fixed' attributes.
surf.all_fixed()
surf.set_fixed('rc(0,0)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
desired_volume = 0.15
prob = LeastSquaresProblem([(equil.volume, desired_volume, 1)])
# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
least_squares_mpi_solve(prob, mpi, grad=True)
# Make sure all procs call VMEC:
objective = prob.objective()
if mpi.proc0_world:
print("At the optimum,")
print(" rc(m=0,n=0) = ", surf.get_rc(0, 0))
print(" volume, according to VMEC = ", equil.volume())
print(" volume, according to Surface = ", surf.volume())
print(" objective function = ", objective)
assert np.abs(surf.get_rc(0, 0) - 0.7599088773175) < 1.0e-5
# Garabedian representation:
surf = equil.boundary.to_Garabedian()
equil.boundary = surf
# VMEC parameters are all fixed by default, while surface parameters
# are all non-fixed by default. You can choose which parameters are
# optimized by setting their 'fixed' attributes.
surf.all_fixed()
surf.set_fixed('Delta(1,-1)', False)
# Each function we want in the objective function is then equipped
# with a shift and weight, to become a term in a least-squares
# objective function. A list of terms are combined to form a
# nonlinear-least-squares problem.
desired_iota = -0.41
prob = LeastSquaresProblem([(equil.iota_axis, desired_iota, 1)])
# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
least_squares_mpi_solve(prob, mpi, grad=False)
# Make sure all procs call VMEC:
objective = prob.objective()
if mpi.proc0_world:
print("At the optimum,")
print(" Delta(m=1,n=-1) = ", surf.get_Delta(1, -1))
print(" iota on axis = ", equil.iota_axis())
print(" objective function = ", objective)
assert np.abs(surf.get_Delta(1, -1) - 0.08575) < 1.0e-4
assert np.abs(equil.iota_axis() - desired_iota) < 1.0e-5
surf.all_fixed()
surf.set_fixed('rc(1,1)', False)
surf.set_fixed('zs(1,1)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
desired_volume = 0.15
volume_weight = 1
term1 = (equil.volume, desired_volume, volume_weight)
desired_iota = -0.41
iota_weight = 1
term2 = (equil.iota, desired_iota, iota_weight)
prob = LeastSquaresProblem([term1, term2])
# Solve the minimization problem:
least_squares_serial_solve(prob, grad=True)
print("At the optimum,")
print(" rc(m=1,n=1) = ", surf.get_rc(1, 1))
print(" zs(m=1,n=1) = ", surf.get_zs(1, 1))
print(" volume, according to VMEC = ", equil.volume())
print(" volume, according to Surface = ", surf.volume())
print(" iota on axis = ", equil.iota())
print(" objective function = ", prob.objective())
# The tests here are based on values from the VMEC version in
# https://github.com/landreman/stellopt_scenarios/tree/master/2DOF_vmecOnly_targetIotaAndVolume
# Due to this and the fact that we don't yet have iota on axis from SPEC, the tolerances are wide.
vmec = Vmec(os.path.join(os.path.dirname(__file__), 'inputs', 'input.nfp2_QA'),
mpi=mpi)
vmec.verbose = mpi.proc0_world
surf = vmec.boundary
# Configure quasisymmetry objective:
boozer = Boozer(vmec)
boozer.bx.verbose = mpi.proc0_world
qs = Quasisymmetry(
boozer,
0.5, # Radius to target
1,
0) # (M, N) you want in |B|
# Define objective function
prob = LeastSquaresProblem([(vmec.aspect, 6, 1), (vmec.iota_axis, 0.465, 1),
(vmec.iota_edge, 0.495, 1), (qs, 0, 1)])
# Fourier modes of the boundary with m <= max_mode and |n| <= max_mode
# will be varied in the optimization. A larger range of modes are
# included in the VMEC and booz_xform calculations.
for step in range(3):
max_mode = step + 1
# VMEC's mpol & ntor will be 3, 4, 5:
vmec.indata.mpol = 3 + step
vmec.indata.ntor = vmec.indata.mpol
# booz_xform's mpol & ntor will be 16, 24, 32:
boozer.mpol = 16 + step * 8
boozer.ntor = boozer.mpol
marker="x",
label=f"Periods = {fSol[3]:.0f}, Residue = {fSol[0]:.4f}")
for fSol in solInitial
]
plt.xlim([0.7, 1.3])
plt.ylim([-0.25, 0.25])
plt.legend()
plt.tight_layout()
plt.savefig('Results/ReimanInitialPoincare.png', dpi=500)
plt.savefig('Results/ReimanInitialPoincare.pdf')
p.plot_iota()
plt.tight_layout()
plt.savefig('Results/ReimanInitialIota.png')
# exit()
## Optimization
prob = LeastSquaresProblem([(obj.residue1, 0, 1)])
# Set degrees of freedom for the optimization
obj.all_fixed()
obj.set_fixed('epsilon', False)
# Run optimization problem
# nIterations = 10
print('Starting optimization...')
least_squares_serial_solve(
prob, bounds=[-0.037, 0.021], xtol=1e-12, ftol=1e-12,
gtol=1e-12) #, method='lm', max_nfev=nIterations)
print('Optimization finished...')
## Create final Poincare Plot
solFinal = [obj.residue()]
p = obj.poincare(nPpts=350, nPtrj=8, Rbegin=1.19, Rend=1.23)
p.plot(s=1.5)
surf.set_rc(1, 0, 0.1)
surf.set_zs(1, 0, 0.1)
print('rc:', surf.rc)
print('zs:', surf.zs)
# Surface parameters are all non-fixed by default. You can choose
# which parameters are optimized by setting their 'fixed' attributes.
surf.all_fixed()
surf.set_fixed('rc(0,0)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
desired_volume = 0.15
prob = LeastSquaresProblem([(equil.volume, desired_volume, 1)])
# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
least_squares_serial_solve(prob, grad=True)
print("At the optimum,")
print(" rc(m=0,n=0) = ", surf.get_rc(0, 0))
print(" volume, according to SPEC = ", equil.volume())
print(" volume, according to Surface = ", surf.volume())
print(" objective function = ", prob.objective())
assert np.abs(surf.get_rc(0, 0) - 0.7599088773175) < 1.0e-5
assert np.abs(equil.volume() - 0.15) < 1.0e-6
assert np.abs(surf.volume() - 0.15) < 1.0e-6
assert prob.objective() < 1.0e-15