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_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_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_exceptions(self):
"""
Verify that exceptions are raised when invalid inputs are
provided.
"""
# Argument must be a list in which each element is a
# LeastSquaresTerm or a 3- or 4-element tuple/list.
with self.assertRaises(TypeError):
prob = LeastSquaresProblem(7)
with self.assertRaises(ValueError):
prob = LeastSquaresProblem([])
with self.assertRaises(TypeError):
prob = LeastSquaresProblem([7, 1])
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)])
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_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))
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])
solver(prob)
self.assertAlmostEqual(prob.objective(), 0)
self.assertAlmostEqual(iden1.x, 1)
self.assertAlmostEqual(iden2.x, 2)
self.assertAlmostEqual(iden3.x, 3)
def test_parallel_optimization(self):
"""
Test a full least-squares optimization.
"""
for ngroups in range(1, 4):
for grad in [True, False]:
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=grad)
self.assertAlmostEqual(prob.x[0], 1)
self.assertAlmostEqual(prob.x[1], 1)
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)
iden1.set_dofs([10])
self.assertAlmostEqual(prob.objective(), 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)
iden1.set_dofs([5])
iden2.set_dofs([-7])
self.assertAlmostEqual(prob.objective(), 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 optaxis(stel, iotaTarget, nIterations=30, nquadrature=100):
from simsopt.core.least_squares_problem import LeastSquaresProblem
from simsopt.solve.serial_solve import least_squares_serial_solve
from glob import glob
from os import remove
from numpy import concatenate
from simsopt.geo.curverzfourier import CurveRZFourier
stel.all_fixed()
stel.set_fixed('rc(1)', False)
stel.set_fixed('zs(1)', False)
stel.set_fixed('rc(2)', False)
stel.set_fixed('zs(2)', False)
stel.set_fixed('rc(3)', False)
stel.set_fixed('zs(3)', False)
stel.set_fixed('rc(4)', False)
stel.set_fixed('zs(4)', False)
stel.set_fixed('rc(5)', False)
stel.set_fixed('zs(5)', False)
stel.set_fixed('rc(6)', False)
stel.set_fixed('zs(6)', False)
stel.set_fixed('etabar', False)
stel.set_fixed('B2c', False)
# stel.set_fixed('p2',False)
## Form a nonlinear-least-squares optimization problem
term = [
(stel, 'iota', iotaTarget, 1e4),
#(stel, 'p2', 0.0, 1e-1),
(stel, 'max_elongation', 0.0, 1e-1),
(stel, 'B20_anomaly', 0.0, 1e2),
(stel, 'X20', 0.0, 1e-3),
(stel, 'X2c', 0.0, 1e-3),
(stel, 'X2s', 0.0, 1e-3),
(stel, 'Y20', 0.0, 1e-3),
(stel, 'Y2c', 0.0, 1e-3),
(stel, 'Y2s', 0.0, 1e-3),
(stel, 'Z20', 0.0, 1e-3),
(stel, 'Z2c', 0.0, 1e-3),
(stel, 'Z2s', 0.0, 1e-3),
#(stel, 'DMerc_times_r2', 1e+0,1e-0), ### CHECK POSITIVE
#(stel, 'grad_grad_B_inverse_scale_length', 0.0,1e-1)
]
prob = LeastSquaresProblem(term)
## Print initial conditions
print('Before optimization:')
print('NFP = ', stel.nfp)
print('p2 = ', stel.p2)
print('etabar = ', stel.etabar)
print('B2c = ', stel.B2c)
print('iota = ', stel.iota)
print('rc = [', ','.join([str(elem) for elem in stel.rc]), ']')
print('zs = [', ','.join([str(elem) for elem in stel.zs]), ']')
print('DMerc = ', stel.DMerc_times_r2)
print('Max elongation = ', stel.max_elongation)
print('gradgradB inverse length: ', stel.grad_grad_B_inverse_scale_length)
print('objective function: ', prob.objective())
## Solve the minimization problem:
try:
least_squares_serial_solve(prob,
xtol=1e-15,
ftol=1e-15,
gtol=1e-15,
max_nfev=nIterations,
method='lm')
except KeyboardInterrupt:
print("Terminated optimization")
for f in glob("residuals_2021*.dat"):
remove(f)
for f in glob("simsopt_2021*.dat"):
remove(f)
## Print final conditions
print('After optimization:')
print('NFP = ', stel.nfp)
print('p2 = ', stel.p2)
print('etabar = ', stel.etabar)
print('B2c = ', stel.B2c)
print('iota = ', stel.iota)
print('rc = [', ','.join([str(elem) for elem in stel.rc]), ']')
print('zs = [', ','.join([str(elem) for elem in stel.zs]), ']')
print('DMerc = ', stel.DMerc_times_r2)
print('Max elongation = ', stel.max_elongation)
print('gradgradB inverse length: ', stel.grad_grad_B_inverse_scale_length)
print('objective function: ', prob.objective())
nN = stel.iota - stel.iotaN
if nN == 0:
print('Quasi-axisymmetric solution')
else:
print('Quasi-helically symmetric solution with N =', nN)
axis = CurveRZFourier(nquadrature, len(stel.rc) - 1, stel.nfp, True)
axis.set_dofs(concatenate((stel.rc, stel.zs[1:])))
return stel, axis
def optaxiscoil(qvfilename, stel, iotaTarget, nIterations=30, nquadrature=150):
from simsopt.core.least_squares_problem import LeastSquaresProblem
from simsopt.solve.serial_solve import least_squares_serial_solve
from simsopt.geo.curverzfourier import CurveRZFourier
from axisOptFuncs import getFourierCurve, plot_stellarator, export_coils
from numpy import sqrt, dot
from os import remove
from glob import glob
import matplotlib.pyplot as plt
outputFile = qvfilename + "_coil_coeffs.dat"
current = 3.11049660e+05
coils, _ = getFourierCurve(outputFile, current)
axis = CurveRZFourier(nquadrature, len(stel.rc) - 1, stel.nfp, True)
axis.set_dofs(concatenate((stel.rc, stel.zs[1:])))
obj = objaxiscoil(coils, current, axis, stel)
term = [(obj.iota, iotaTarget, 1e4), (obj.devBonAxis, 0, 1e0),
(obj, 'max_elongation', 0, 1e1), (obj.stdBonAxis, 0, 1e5),
(obj.residue, 0, 1e5)]
prob = LeastSquaresProblem(term)
obj.all_fixed()
# obj.set_fixed('rc(1)', False)
# obj.set_fixed('zs(1)', False)
# obj.set_fixed('rc(2)', False)
# obj.set_fixed('zs(2)', False)
obj.set_fixed('rc(3)', False)
obj.set_fixed('zs(3)', False)
obj.set_fixed('rc(4)', False)
obj.set_fixed('zs(4)', False)
obj.set_fixed('rc(5)', False)
obj.set_fixed('zs(5)', False)
#obj.set_fixed('etabar', False)
#obj.set_fixed('current', False)
#obj.fixed[obj.nAxisFour:obj.nAxisFour+(obj.ncoils+1)*obj.nCoilFour]=False
################################################
rcInit = obj.stel.rc
zsInit = obj.stel.zs
plt.figure()
plt.plot(obj.BonAxis())
plt.xlabel('phi')
plt.ylabel('B')
################ OPTIMIZE ####################
least_squares_serial_solve(prob, max_nfev=nIterations) #, method='lm')
################################################
for f in glob("residuals_2021*.dat"):
remove(f)
for f in glob("simsopt_2021*.dat"):
remove(f)
################################################
rcFinal = obj.stel.rc
zsFinal = obj.stel.zs
rcDelta = [
100 * (rcf - rcInit[count]) / rcInit[count]
for count, rcf in enumerate(rcFinal)
]
zsDelta = [
100 * (zsf - zsInit[count]) / zsInit[count]
for count, zsf in enumerate(zsFinal)
]
print('Initial rc = [',
','.join([str(elem) for elem in rcInit]), ']')
print('Final rc = [',
','.join([str(elem) for elem in rcFinal]), ']')
print('Initial zs = [',
','.join([str(elem) for elem in zsInit]), ']')
print('Final zs = [',
','.join([str(elem) for elem in zsFinal]), ']')
print('Percentage change in rc = [',
','.join([str(elem) for elem in rcDelta]), ']')
print('Percentage change in zs = [',
','.join([str(elem) for elem in zsDelta]), ']')
print('Target iota = ', iotaTarget, ', final iota = ', obj.stel.iota)
################################################
print("Plot on-axis B from these coils")
plt.plot(obj.BonAxis())
plt.savefig("optimizedBonaxis_" + qvfilename + '.pdf',
bbox_inches='tight',
pad_inches=0)
################################################
print("Plot final coils and axis")
plot_stellarator("coils_axis_Optimized_", qvfilename, obj.coils, obj.nfp,
obj.axis)
filename = "coils." + qvfilename + "Optimized"
export_coils(obj.coils, filename, obj.currents, obj.nfp)
return obj.stel, obj.axis