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