def test_solve_quadratic_fixed_supplying_objects(self):
"""
Same as test_solve_quadratic_fixed, except supplying objects
rather than functions as targets.
"""
for solver in solvers:
iden1 = Identity()
iden2 = Identity()
iden3 = Identity()
iden1.x = 4
iden2.x = 5
iden3.x = 6
iden1.names = ['x1']
iden2.names = ['x2']
iden3.names = ['x3']
iden1.fixed = [True]
iden3.fixed = [True]
term1 = [iden1, 1, 1]
term2 = [iden2, 2, 1 / 4.]
term3 = [iden3, 3, 1 / 9.]
prob = LeastSquaresProblem([term1, term2, term3])
solver(prob)
self.assertAlmostEqual(prob.objective(), 10)
self.assertAlmostEqual(iden1.x, 4)
self.assertAlmostEqual(iden2.x, 2)
self.assertAlmostEqual(iden3.x, 6)
def test_solve_quadratic_fixed_supplying_attributes(self):
"""
Same as test_solve_quadratic_fixed, except supplying attributes
rather than functions as targets.
"""
for solver in solvers:
iden1 = Identity()
iden2 = Identity()
iden3 = Identity()
iden1.x = 4
iden2.x = 5
iden3.x = 6
iden1.names = ['x1']
iden2.names = ['x2']
iden3.names = ['x3']
iden1.fixed = [True]
iden3.fixed = [True]
# Try a mix of explicit LeastSquaresTerms and tuples
term1 = LeastSquaresTerm(Target(iden1, 'x'), 1, 1)
term2 = (iden2, 'x', 2, 1 / 4.)
term3 = (iden3, 'x', 3, 1 / 9.)
prob = LeastSquaresProblem([term1, term2, term3])
solver(prob)
self.assertAlmostEqual(prob.objective(), 10)
self.assertAlmostEqual(iden1.x, 4)
self.assertAlmostEqual(iden2.x, 2)
self.assertAlmostEqual(iden3.x, 6)
def test_solve_quadratic_fixed(self):
"""
Same as test_solve_quadratic, except with different weights and x
and z are fixed, so only y is optimized.
"""
for solver in solvers:
iden1 = Identity()
iden2 = Identity()
iden3 = Identity()
iden1.x = 4
iden2.x = 5
iden3.x = 6
iden1.names = ['x1']
iden2.names = ['x2']
iden3.names = ['x3']
iden1.fixed = [True]
iden3.fixed = [True]
term1 = (iden1.J, 1, 1)
term2 = (iden2.J, 2, 1 / 4.)
term3 = (iden3.J, 3, 1 / 9.)
prob = LeastSquaresProblem([term1, term2, term3])
solver(prob)
self.assertAlmostEqual(prob.objective(), 10)
self.assertAlmostEqual(iden1.x, 4)
self.assertAlmostEqual(iden2.x, 2)
self.assertAlmostEqual(iden3.x, 6)
def test_solve_quadratic_fixed_supplying_properties(self):
"""
Same as test_solve_quadratic_fixed, except supplying @properties
rather than functions as targets.
"""
for solver in solvers:
iden1 = Identity()
iden2 = Identity()
iden3 = Identity()
iden1.x = 4
iden2.x = 5
iden3.x = 6
iden1.names = ['x1']
iden2.names = ['x2']
iden3.names = ['x3']
iden1.fixed = [True]
iden3.fixed = [True]
# Try a mix of explicit LeastSquaresTerms and lists
term1 = [iden1, 'f', 1, 1]
term2 = [iden2, 'f', 2, 1 / 4.]
term3 = LeastSquaresTerm.from_sigma(Target(iden3, 'f'), 3, sigma=3)
prob = LeastSquaresProblem([term1, term2, term3])
solver(prob)
self.assertAlmostEqual(prob.objective(), 10)
self.assertAlmostEqual(iden1.x, 4)
self.assertAlmostEqual(iden2.x, 2)
self.assertAlmostEqual(iden3.x, 6)
def test_integrated_stellopt_scenarios_1dof_Garabedian(self):
"""
This script implements the "1DOF_circularCrossSection_varyAxis_targetIota"
example from
https://github.com/landreman/stellopt_scenarios
This example demonstrates optimizing a surface shape using the
Garabedian representation instead of VMEC's RBC/ZBS representation.
This optimization problem has one independent variable, the Garabedian
Delta_{m=1, n=-1} coefficient, representing the helical excursion of
the magnetic axis. The objective function is (iota - iota_target)^2,
where iota is measured on the magnetic axis.
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_varyAxis_targetIota
"""
filename = os.path.join(TEST_DIR, '1DOF_Garabedian.sp')
for mpol_ntor in [2, 4]:
# Start with a default surface.
equil = Spec(filename)
equil.inputlist.mpol = mpol_ntor
equil.inputlist.ntor = mpol_ntor
# We will optimize in the space of Garabedian coefficients
# rather than RBC/ZBS coefficients. To do this, we convert the
# boundary to the Garabedian representation:
surf = equil.boundary.to_Garabedian()
equil.boundary = surf
# SPEC 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)
# Use low resolution, for speed:
equil.inputlist.lrad[0] = 4
equil.inputlist.nppts = 100
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function
desired_iota = 0.41 # Sign was + for VMEC
prob = LeastSquaresProblem([(equil.iota, desired_iota, 1)])
# Check that the problem was set up correctly:
self.assertEqual(len(prob.dofs.names), 1)
self.assertEqual(prob.dofs.names[0][:11], 'Delta(1,-1)')
np.testing.assert_allclose(prob.x, [0.1])
self.assertEqual(prob.dofs.all_owners, [equil, surf])
self.assertEqual(prob.dofs.dof_owners, [surf])
# Solve the minimization problem:
least_squares_serial_solve(prob)
self.assertAlmostEqual(surf.get_Delta(1, -1), 0.08575, places=4)
self.assertAlmostEqual(equil.iota(), desired_iota, places=5)
self.assertLess(np.abs(prob.objective()), 1.0e-15)
def subtest_curve_length_optimisation(self, rotated):
nquadrature = 100
nfourier = 4
nfp = 5
curve = CurveRZFourier(nquadrature, nfourier, nfp, True)
if rotated:
curve = RotatedCurve(curve, 0.5, flip=False)
# Initialize the Fourier amplitudes to some random values
x0 = np.random.rand(curve.num_dofs()) - 0.5
x0[0] = 3.0
curve.set_dofs(x0)
print('Initial curve dofs: ', curve.get_dofs())
# 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
# Presently in simsgeo, the length objective is a separate object
# rather than a function of Curve itself.
obj = make_optimizable(CurveLength(curve))
# For now, we need to add this attribute to CurveLength. Eventually
# this would hopefully be done in simsgeo, but for now I'll put it here.
obj.depends_on = ['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)
assert np.allclose(fd_jac, jac, rtol=1e-4, atol=1e-4)
# Solve the minimization problem:
least_squares_serial_solve(prob, ftol=1e-10, xtol=1e-10, gtol=1e-10)
print('At the optimum, x: ', prob.x)
print(' Final curve dofs: ', curve.get_dofs())
print(' Final curve length: ', obj.J())
print(' Expected final length: ', 2 * np.pi * x0[0])
print(' objective function: ', prob.objective())
assert abs(obj.J() - 2 * np.pi * x0[0]) < 1e-8
def test_supply_tuples(self):
"""
Test basic usage
"""
# Objective function f(x) = ((x - 3) / 2) ** 2
iden1 = Identity()
term1 = (iden1.J, 3, 0.25)
prob = LeastSquaresProblem([term1])
self.assertAlmostEqual(prob.objective(), 2.25)
self.assertAlmostEqual(prob.objective(), sum(t.f_out() for t in prob.terms))
self.assertEqual(len(prob.dofs.f()), 1)
self.assertAlmostEqual(prob.dofs.f()[0], 0)
self.assertEqual(len(prob.f()), 1)
self.assertAlmostEqual(prob.f()[0], -1.5)
self.assertAlmostEqual(prob.objective_from_shifted_f(prob.f()), 2.25)
self.assertAlmostEqual(prob.objective_from_unshifted_f(prob.dofs.f()), 2.25)
iden1.set_dofs([10])
self.assertAlmostEqual(prob.objective(), 12.25)
self.assertAlmostEqual(prob.objective(), sum(t.f_out() for t in prob.terms))
self.assertAlmostEqual(prob.objective_from_shifted_f(prob.f()), 12.25)
self.assertAlmostEqual(prob.objective_from_unshifted_f(prob.dofs.f()), 12.25)
self.assertAlmostEqual(prob.objective([0]), 2.25)
self.assertAlmostEqual(prob.objective([10]), 12.25)
self.assertEqual(prob.dofs.all_owners, [iden1])
self.assertEqual(prob.dofs.dof_owners, [iden1])
# Objective function
# f(x,y) = ((x - 3) / 2) ** 2 + ((y + 4) / 5) ** 2
iden2 = Identity()
term2 = (iden2.J, -4, 0.04)
prob = LeastSquaresProblem([term1, term2])
self.assertAlmostEqual(prob.objective(), 12.89)
self.assertAlmostEqual(prob.objective(), sum(t.f_out() for t in prob.terms))
self.assertEqual(len(prob.f()), 2)
self.assertAlmostEqual(prob.objective_from_shifted_f(prob.f()), 12.89)
self.assertAlmostEqual(prob.objective_from_unshifted_f(prob.dofs.f()), 12.89)
iden1.set_dofs([5])
iden2.set_dofs([-7])
self.assertAlmostEqual(prob.objective(), 1.36)
self.assertAlmostEqual(prob.objective(), sum(t.f_out() for t in prob.terms))
self.assertEqual(len(prob.f()), 2)
self.assertAlmostEqual(prob.objective_from_shifted_f(prob.f()), 1.36)
self.assertAlmostEqual(prob.objective_from_unshifted_f(prob.dofs.f()), 1.36)
self.assertAlmostEqual(prob.objective([10, 0]), 12.89)
self.assertAlmostEqual(prob.objective([5, -7]), 1.36)
self.assertEqual(prob.dofs.dof_owners, [iden1, iden2])
self.assertEqual(prob.dofs.all_owners, [iden1, iden2])
def test_solve_rosenbrock_using_vector(self):
"""
Minimize the Rosenbrock function using a single vector-valued
least-squares term.
"""
for solver in solvers:
for grad in [True, False]:
r = Rosenbrock()
prob = LeastSquaresProblem([(r.terms, 0, 1)])
if solver == serial_solve:
if grad == True:
continue
else:
solver(prob, tol=1e-12)
else:
solver(prob, grad=grad)
self.assertAlmostEqual(prob.objective(), 0)
v = r.get_dofs()
self.assertAlmostEqual(v[0], 1)
self.assertAlmostEqual(v[1], 1)
def test_solve_quadratic(self):
"""
Minimize f(x,y,z) = 1 * (x - 1) ^ 2 + 2 * (y - 2) ^ 2 + 3 * (z - 3) ^ 2.
The optimum is at (x,y,z)=(1,2,3), and f=0 at this point.
"""
for solver in solvers:
iden1 = Identity()
iden2 = Identity()
iden3 = Identity()
term1 = (iden1.J, 1, 1)
term2 = (iden2.J, 2, 2)
term3 = (iden3.J, 3, 3)
prob = LeastSquaresProblem([term1, term2, term3])
if solver == serial_solve:
solver(prob, tol=1e-12)
else:
solver(prob)
self.assertAlmostEqual(prob.objective(), 0)
self.assertAlmostEqual(iden1.x, 1)
self.assertAlmostEqual(iden2.x, 2)
self.assertAlmostEqual(iden3.x, 3)
def test_solve_rosenbrock_using_scalars(self):
"""
Minimize the Rosenbrock function using two separate least-squares
terms.
"""
for solver in solvers:
for grad in [True, False]:
r = Rosenbrock()
term1 = (r.term1, 0, 1)
term2 = (r.term2, 0, 1)
prob = LeastSquaresProblem((term1, term2))
if solver == serial_solve:
if grad == True:
continue
else:
solver(prob, tol=1e-12)
else:
solver(prob, grad=grad)
self.assertAlmostEqual(prob.objective(), 0)
v = r.get_dofs()
self.assertAlmostEqual(v[0], 1)
self.assertAlmostEqual(v[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 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
def test_integrated_stellopt_scenarios_1dof(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
SPEC 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 SPEC.
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
"""
for grad in [True, False]:
# Start with a default surface.
equil = Spec()
surf = equil.boundary
# Set the initial boundary shape. Here is one way to do it:
surf.set('rc(0,0)', 1.0)
# Here is another syntax that works:
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)
surf.set_rc(1, 1, 0)
surf.set_zs(1, 1, 0)
# SPEC 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)
# Turn off Poincare plots and use low resolution, for speed:
equil.inputlist.nptrj[0] = 0
equil.inputlist.lrad[0] = 2
# 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 = (equil.volume, desired_volume, 1)
# A list of terms are combined to form a nonlinear-least-squares
# problem.
prob = LeastSquaresProblem([term1])
# Check that the problem was set up correctly:
self.assertEqual(len(prob.dofs.names), 1)
self.assertEqual(prob.dofs.names[0][:7], 'rc(0,0)')
np.testing.assert_allclose(prob.x, [1.0])
self.assertEqual(prob.dofs.all_owners, [equil, surf])
self.assertEqual(prob.dofs.dof_owners, [surf])
# Solve the minimization problem:
least_squares_serial_solve(prob, grad=grad)
self.assertAlmostEqual(surf.get_rc(0, 0), 0.7599088773175, places=5)
self.assertAlmostEqual(equil.volume(), 0.15, places=6)
self.assertAlmostEqual(surf.volume(), 0.15, places=6)
self.assertLess(np.abs(prob.objective()), 1.0e-15)
def test_integrated_stellopt_scenarios_2dof(self):
"""
This script implements the "2DOF_vmecOnly_targetIotaAndVolume" example from
https://github.com/landreman/stellopt_scenarios
This optimization problem has two independent variables, representing
the helical shape of the magnetic axis. The problem also has two
objectives: the plasma volume and the rotational transform on the
magnetic axis.
The resolution in this example (i.e. ns, mpol, and ntor) is somewhat
lower than in the stellopt_scenarios version of the example, just so
this example runs fast.
Details of the optimum and a plot of the objective function landscape
can be found here:
https://github.com/landreman/stellopt_scenarios/tree/master/2DOF_vmecOnly_targetIotaAndVolume
"""
filename = os.path.join(TEST_DIR, '2DOF_targetIotaAndVolume.sp')
# Initialize SPEC from an input file
equil = Spec(filename)
surf = equil.boundary
# 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(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)
# 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.
"""
assert np.abs(surf.get_rc(1, 1) - 0.0313066948) < 0.001
assert np.abs(surf.get_zs(1, 1) - (-0.031232391)) < 0.001
assert np.abs(equil.volume() - 0.178091) < 0.001
assert np.abs(surf.volume() - 0.178091) < 0.001
assert np.abs(equil.iota() - (-0.4114567)) < 0.001
assert (prob.objective() - 7.912501330E-04) < 0.2e-4
"""
self.assertAlmostEqual(surf.get_rc(1, 1), 0.0313066948, places=3)
self.assertAlmostEqual(surf.get_zs(1, 1), -0.031232391, places=3)
self.assertAlmostEqual(equil.volume(), 0.178091, places=3)
self.assertAlmostEqual(surf.volume(), 0.178091, places=3)
self.assertAlmostEqual(equil.iota(), -0.4114567, places=3)
self.assertAlmostEqual(prob.objective(), 7.912501330E-04, places=3)
# Define some Target objects that depend on Parameter objects. In the
# future these functions would involve codes like VMEC, but for now we
# just use the functions f(x) = x.
iden1 = Identity()
iden2 = Identity()
iden3 = Identity()
# Parameters are all not fixed by default, meaning they will not be
# optimized. You can choose to exclude any subset of the parameters
# from the space of independent variables by setting their 'fixed'
# property to True.
#iden1.fixed[0] = True
#iden2.fixed[0] = True
#iden3.fixed[0] = True
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function
term1 = (iden1, 1, 1)
term2 = (iden2, 2, 2)
term3 = (iden3, 3, 3)
# A list of terms are combined to form a nonlinear-least-squares problem.
prob = LeastSquaresProblem([term1, term2, term3])
# Solve the minimization problem:
least_squares_serial_solve(prob)
print("An optimum was found at x=", iden1.x, ", y=", iden2.x, \
", z=", iden3.x)
print("The minimum value of the objective function is ", prob.objective())