def test_parallel_optimization_with_grad(self):
"""
Test a full least-squares optimization.
"""
for ngroups in range(1, 4):
for abs_step in [0, 1.0e-7]:
# Only try rel_step=0 if abs_step is positive:
rel_steps = [0, 1.0e-7]
if abs_step == 0:
rel_steps = [1.0e-7]
for rel_step in rel_steps:
for diff_method in ["forward", "centered"]:
logger.debug(f'ngroups={ngroups} abs_step={abs_step} ' \
f'rel_step={rel_step} diff_method={diff_method}')
mpi = MpiPartition(ngroups=ngroups)
o = TestFunction3(mpi.comm_groups)
term1 = (o.f0, 0, 1)
term2 = (o.f1, 0, 1)
prob = LeastSquaresProblem([term1, term2], diff_method=diff_method,
abs_step=abs_step, rel_step=rel_step)
# Set initial condition different from 0,
# because otherwise abs_step=0 causes step
# size to be 0.
prob.x = [-0.1, 0.2]
least_squares_mpi_solve(prob, mpi)
self.assertAlmostEqual(prob.x[0], 1)
self.assertAlmostEqual(prob.x[1], 1)
def test_parallel_optimization_without_grad(self):
"""
Test a full least-squares optimization.
"""
for ngroups in range(1, 4):
mpi = MpiPartition(ngroups=ngroups)
o = TestFunction3(mpi.comm_groups)
term1 = (o.f0, 0, 1)
term2 = (o.f1, 0, 1)
prob = LeastSquaresProblem([term1, term2])
least_squares_mpi_solve(prob, mpi, grad=False)
self.assertAlmostEqual(prob.x[0], 1)
self.assertAlmostEqual(prob.x[1], 1)
residue2 = Residue(spec, p, q, theta=np.pi)
if mpi.group == 0:
r1 = residue1.J()
r2 = residue2.J()
if mpi.proc0_world:
print("Initial residues:", r1, r2)
#exit(0)
# Define objective function
prob = LeastSquaresProblem([(vmec.aspect, 6, 1), (vmec.iota_axis, 0.385, 1),
(vmec.iota_edge, 0.415, 1), (qs, 0, 1),
(residue1, 0, 2), (residue2, 0, 2)])
# Check whether we're in the CI. If so, just do a single function
# evaluation rather than a real optimization.
ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true']
if ci:
obj = prob.objective()
else:
# Remove the max_nfev=1 in the next line to do a serious optimization:
least_squares_mpi_solve(prob, mpi=mpi, grad=True, max_nfev=1)
if mpi.group == 0:
r1 = residue1.J()
r2 = residue2.J()
if mpi.proc0_world:
print("Final residues:", r1, r2)
print("Good bye")
def mpi_solve_1group(prob, **kwargs):
least_squares_mpi_solve(prob, MpiPartition(ngroups=1), **kwargs)
# 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,")
logging.info(f" objective function = {final_objective}")
logging.info(f" rc(m=1,n=1) = {surf.get_rc(1, 1)}")
logging.info(f" zs(m=1,n=1) = {surf.get_zs(1, 1)}")
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__)
logger = logging.getLogger(__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
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=mpi, 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.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
assert prob.objective() < 1.0e-15