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_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_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_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 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 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