forked from rogeriojorge/single_stage_optimization
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
executable file
·434 lines (421 loc) · 26.2 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
#!/usr/bin/env python3
"""
Single stage stellarator optimization (3SO) code
Repository: https://github.com/rogeriojorge/single_stage_optimization
TEMPORARY NOTE: use single_stage branch from SIMSOPT while it is not merged to main branch
"""
import os
import sys
import glob
import json
import time
import shutil
import logging
import argparse
import numpy as np
import pandas as pd
from pathlib import Path
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from simsopt.mhd import Vmec, Boozer
from simsopt.util import MpiPartition
from simsopt._core.derivative import Derivative
from simsopt.solve import least_squares_mpi_solve
from simsopt._core.finite_difference import MPIFiniteDifference
from simsopt.geo import curves_to_vtk
from src.vmecPlot2 import main as vmecPlot2_main
from src.field_from_coils import main as field_from_coils_main
from simsopt.util import proc0_print as pprint
from src.initialization_functions import (recalculate_inputs,
create_results_folders, create_initial_coils)
from src.stage_1 import form_stage_1_objective_function
from src.stage_2 import form_stage_2_objective_function, inner_coil_loop
# Start the timer
start = time.time()
# Start mpi process
mpi = MpiPartition()
# Set the logger to print debugs
logger = logging.getLogger(__name__)
# from simsopt.util.mpi import log
# log(level=logging.DEBUG)
# log(level=logging.INFO)
# Check if user selected QA or QH when launching main.py
QAQHQICNTselected=False
if len(sys.argv) > 1:
if sys.argv[1]=='QA' or sys.argv[1]=='QH' or sys.argv[1]=='QI' or sys.argv[1]=='CNT':
QAorQHorQIorCNT = sys.argv[1]
QAQHQICNTselected=True
if not QAQHQICNTselected:
pprint('First line argument (QA, QH, QI or CNT) not selected. Defaulting to QA.')
QAorQHorQIorCNT = 'QA'
###############
# Parse the command line arguments and overwrite inputs.py if needed
parser = argparse.ArgumentParser()
inputs = recalculate_inputs(parser, QAQHQICNTselected, QAorQHorQIorCNT, sys.argv)
# Create results folders
parent_path = str(Path(__file__).parent.resolve())
current_path = os.path.join(parent_path, 'results', f'{inputs.name}')
if mpi.proc0_world and inputs.remove_previous_results and os.path.isdir(current_path):
shutil.copytree(current_path, current_path + '_backup', dirs_exist_ok=True)
shutil.rmtree(current_path)
if mpi.proc0_world: Path(current_path).mkdir(parents=True, exist_ok=True)
current_path = str(Path(current_path).resolve())
if mpi.proc0_world: shutil.copy(os.path.join(parent_path,'src','inputs.py'), os.path.join(current_path,f'{inputs.name}.py'))
time.sleep(0.5)
try: os.chdir(current_path)
except:
time.sleep(3)
os.chdir(current_path)
coils_results_path, vmec_results_path, figures_results_path = create_results_folders(inputs)
# Write inputs to file f'inputs_{inputs.name}.json'
inputs_dict = dict([(att, getattr(inputs,att)) for att in dir(inputs) if '__' not in att])
with open(os.path.join(current_path, f'inputs_{inputs.name}.json'), 'w', encoding='utf-8') as f:
json.dump(inputs_dict, f, ensure_ascii=False, indent=4)
# If the user did not specify the stage, do single stage by default
if (not inputs.stage_1) and (not inputs.stage_2) and (not inputs.single_stage):
inputs.single_stage = True
######################################################
pprint("============================================")
pprint("Starting single stage optimization")
pprint("============================================")
######################################################
# Check if a previous optimization has already taken place and load it if exists
vmec_files_list = os.listdir(vmec_results_path)
if os.path.exists(os.path.join(current_path, f'input.final')):
vmec_input_filename = os.path.join(current_path, f'input.final')
elif len(vmec_files_list)==0:
vmec_input_filename = os.path.join(parent_path, 'vmec_inputs', inputs.vmec_input_start)
else:
vmec_input_files = [file for file in vmec_files_list if 'input.' in file]
vmec_input_files.sort(key=lambda item: (len(item), item), reverse=False)
vmec_input_filename = vmec_input_files[-1]
pprint(f' Using vmec input file {os.path.join(vmec_results_path,vmec_input_filename)}')
if inputs.use_half_period:
vmec = Vmec(os.path.join(vmec_results_path,vmec_input_filename), mpi=mpi, verbose=inputs.vmec_verbose, nphi=inputs.nphi, ntheta=inputs.ntheta, range_surface='half period')
vmec_full_boundary = Vmec(os.path.join(vmec_results_path,vmec_input_filename), mpi=mpi, verbose=inputs.vmec_verbose, nphi=inputs.nphi, ntheta=inputs.ntheta)
else:
vmec = Vmec(os.path.join(vmec_results_path,vmec_input_filename), mpi=mpi, verbose=inputs.vmec_verbose, nphi=inputs.nphi, ntheta=inputs.ntheta)
vmec_full_boundary = vmec
surf = vmec.boundary
surf_full_boundary = vmec_full_boundary.boundary
bs, coils, curves, base_curves = create_initial_coils(vmec, parent_path, coils_results_path, inputs, surf_full_boundary, mpi)
################################################################
######### DEFINE OBJECTIVE FUNCTION AND ITS DERIVATIVES ########
################################################################
class single_stage_obj_and_der():
def __init__(self) -> None:
pass
def fun(self, dofs, prob_jacobian=None, info={'Nfeval':0}, max_mode=1, oustr_dict=[]):
logger.info('Entering fun')
info['Nfeval'] += 1
JF.x = dofs[:-number_vmec_dofs]
prob.x = dofs[-number_vmec_dofs:]
bs.set_points(surf.gamma().reshape((-1, 3)))
os.chdir(vmec_results_path)
J_stage_1 = prob.objective()
J_stage_2 = inputs.coils_objective_weight * JF.J()
J = J_stage_1 + J_stage_2
if J > inputs.JACOBIAN_THRESHOLD or np.isnan(J):
logger.info(f"Exception caught during function evaluation with J={J}. Returning J={inputs.JACOBIAN_THRESHOLD}")
J = inputs.JACOBIAN_THRESHOLD
jf = Jf.J()
BdotN = np.mean(np.abs(np.sum(bs.B().reshape((inputs.nphi, inputs.ntheta, 3)) * surf.unitnormal(), axis=2)))
outstr = f"\n\nfun#{info['Nfeval']} - J={J:.1e}, Jf={jf:.1e}, ⟨B·n⟩={BdotN:.1e}"
dict1 = {}
dict1.update({
'Nfeval': info['Nfeval'], 'J':float(J), 'Jf': float(jf),'J_length':float(J_LENGTH.J()),
'J_CC':float(J_CC.J()),'J_CURVATURE':float(J_CURVATURE.J()), 'J_LENGTH_PENALTY': float(J_LENGTH_PENALTY.J()),
# ,'J_CS':float(J_CS.J())
'J_MSC':float(J_MSC.J()), 'J_ALS':float(J_ALS.J()), 'Lengths':float(sum(j.J() for j in Jls)),
'curvatures':float(np.sum([np.max(c.kappa()) for c in base_curves])),'msc':float(np.sum([j.J() for j in Jmscs])),
'B.n':float(BdotN),
# 'gradJcoils':float(np.linalg.norm(JF.dJ())),
'C-C-Sep':float(Jccdist.shortest_distance())
#, 'C-S-Sep':float(Jcsdist.shortest_distance())
})
if inputs.debug_coils_outputtxt:
# outstr += f", ║∇J coils║={np.linalg.norm(JF.dJ()):.1e}"
outstr += f", , C-C-Sep={Jccdist.shortest_distance():.2f}"#, C-S-Sep={Jcsdist.shortest_distance():.2f}"
outstr += f"\nJf={jf:.1e}, J_length={J_LENGTH.J():.1e}, J_CC={(J_CC.J()):.1e}, J_CURVATURE={J_CURVATURE.J():.1e}, J_MSC={J_MSC.J():.1e}, J_ALS={J_ALS.J():.1e}, J_LENGTH_PENALTY={J_LENGTH_PENALTY.J():.1e}"#, J_CS={J_CS.J():.1e}
cl_string = ", ".join([f"{j.J():.1f}" for j in Jls])
kap_string = ", ".join(f"{np.max(c.kappa()):.1f}" for c in base_curves)
msc_string = ", ".join(f"{j.J():.1f}" for j in Jmscs)
outstr += f"\n Coil lengths=sum([{cl_string}])={sum(j.J() for j in Jls):.1f}, curvature=[{kap_string}], mean squared curvature=[{msc_string}]"
try:
outstr += f"\n surface dofs="+", ".join([f"{pr}" for pr in dofs[-number_vmec_dofs:]])
if J<inputs.JACOBIAN_THRESHOLD:
dict1.update({'Jquasisymmetry':float(qs.total()), 'Jiota':float((vmec.mean_iota()-inputs.iota_target)**2), 'Jaspect':float((vmec.aspect()-inputs.aspect_ratio_target)**2)})
if not QAorQHorQIorCNT=='QI': outstr += f"\n Quasisymmetry objective={qs.total()}"
outstr += f"\n aspect={vmec.aspect()}"
outstr += f"\n mean iota={vmec.mean_iota()}"
else:
dict1.update({'Jquasisymmetry':0, 'Jiota':0,'Jaspect':0})
except Exception as e:
pprint(e)
if J<inputs.JACOBIAN_THRESHOLD:
logger.info(f'Objective function {J} is smaller than the threshold {inputs.JACOBIAN_THRESHOLD}')
## Finite differences for the first-stage objective function
prob_dJ = prob_jacobian.jac(prob.x)
if not inputs.finite_beta:
## Finite differences for the second-stage objective function
coils_dJ = JF.dJ()
## Mixed term - derivative of squared flux with respect to the surface shape
n = surf.normal()
absn = np.linalg.norm(n, axis=2)
B = bs.B().reshape((inputs.nphi, inputs.ntheta, 3))
dB_by_dX = bs.dB_by_dX().reshape((inputs.nphi, inputs.ntheta, 3, 3))
Bcoil = bs.B().reshape(n.shape)
unitn = n * (1./absn)[:, :, None]
Bcoil_n = np.sum(Bcoil*unitn, axis=2)
mod_Bcoil = np.linalg.norm(Bcoil, axis=2)
B_n = Bcoil_n
B_diff = Bcoil
B_N = np.sum(Bcoil * n, axis=2)
assert Jf.definition == "local"
dJdx = (B_n/mod_Bcoil**2)[:, :, None] * (np.sum(dB_by_dX*(n-B*(B_N/mod_Bcoil**2)[:, :, None])[:, :, None, :], axis=3))
dJdN = (B_n/mod_Bcoil**2)[:, :, None] * B_diff - 0.5 * (B_N**2/absn**3/mod_Bcoil**2)[:, :, None] * n
deriv = surf.dnormal_by_dcoeff_vjp(dJdN/(inputs.nphi*inputs.ntheta)) + surf.dgamma_by_dcoeff_vjp(dJdx/(inputs.nphi*inputs.ntheta))
mixed_dJ = Derivative({surf: deriv})(surf)
## Put both gradients together
grad_with_respect_to_coils = inputs.coils_objective_weight * coils_dJ
grad_with_respect_to_surface = np.ravel(prob_dJ) + inputs.coils_objective_weight * mixed_dJ
else:
pprint('Finite beta needs to be implemented')
exit()
grad = np.concatenate((grad_with_respect_to_coils, grad_with_respect_to_surface))
else:
logger.info(f'Objective function {J} is greater than the threshold {inputs.JACOBIAN_THRESHOLD}')
grad = [0] * len(dofs)
os.chdir(current_path)
with open(inputs.debug_output_file, "a") as myfile:
myfile.write(outstr)
oustr_dict.append(dict1)
if np.mod(info['Nfeval'],inputs.output_interval)==0:
pointData = {"B_N": np.sum(bs.B().reshape((inputs.nphi, inputs.ntheta, 3)) * surf.unitnormal(), axis=2)[:, :, None]}
surf_full_boundary.x = surf.x
surf_full_boundary.to_vtk(os.path.join(coils_results_path,f"surf_intermediate_max_mode_{max_mode}_{info['Nfeval']}"), extra_data=pointData)
curves_to_vtk(curves, os.path.join(coils_results_path,f"curves_intermediate_max_mode_{max_mode}_{info['Nfeval']}"))
return J, grad
# Stage 2 objective function
JF_simple, JF, Jls, Jmscs, Jccdist, Jcsdist, Jf, J_LENGTH, J_CC, J_CS, J_CURVATURE, J_MSC, J_ALS, J_LENGTH_PENALTY = form_stage_2_objective_function(surf, bs, base_curves, curves, inputs)
if inputs.stage_1 or inputs.stage_2 or inputs.single_stage:
oustr_dict_outer=[]
ran_stage1 = False # Only running stage1 optimization once
previous_max_mode=0
# Loop over the number of predefined maximum poloidal/toroidal modes
for max_mode in inputs.max_modes:
pprint(f' Optimization with max_mode={max_mode}')
if max_mode != previous_max_mode: oustr_dict_inner=[]
vmec, vmec_full_boundary, surf, surf_full_boundary, qs, qi, number_vmec_dofs, prob = form_stage_1_objective_function(vmec, vmec_full_boundary, surf, surf_full_boundary, max_mode, inputs)
dofs = np.concatenate((JF.x, vmec.x))
# Stage 1 Optimization
if inputs.stage_1 and not ran_stage1:
os.chdir(vmec_results_path)
pprint(f' Performing Stage 1 optimization with {inputs.MAXITER_stage_1} iterations')
least_squares_mpi_solve(prob, mpi, grad=True, rel_step=inputs.finite_difference_rel_step, abs_step=inputs.finite_difference_abs_step, max_nfev=inputs.MAXITER_stage_1, ftol=inputs.ftol)
os.chdir(current_path)
vmec.write_input(os.path.join(current_path, f'input.stage1'))
if mpi.proc0_world:
with open(inputs.debug_output_file, "a") as myfile:
try:
myfile.write(f"\nAfter Stage 1 optimization with {inputs.MAXITER_stage_1} iterations")
myfile.write(f"\nAspect ratio at max_mode {max_mode}: {vmec.aspect()}")
myfile.write(f"\nMean iota at {max_mode}: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': myfile.write(f"\nQuasisymmetry objective at max_mode {max_mode}: {qs.total()}")
myfile.write(f"\nSquared flux at max_mode {max_mode}: {Jf.J()}")
except Exception as e:
myfile.write(e)
ran_stage1 = True
# Stage 2 Optimization
if not previous_max_mode==max_mode and ((inputs.stage_2 and inputs.stage_1) or previous_max_mode==0):
pprint(f' Performing Stage 2 optimization with {inputs.MAXITER_stage_2+inputs.MAXITER_stage_2_simple} iterations')
surf = vmec.boundary
bs.set_points(surf.gamma().reshape((-1, 3)))
if mpi.proc0_world:
dofs, bs, JF = inner_coil_loop(mpi, JF_simple, JF, Jls, Jmscs, Jccdist, Jcsdist, Jf, J_LENGTH, J_CC, J_CS, J_CURVATURE, J_MSC, J_ALS, J_LENGTH_PENALTY, vmec, curves, base_curves, surf, coils_results_path, number_vmec_dofs, bs, max_mode, inputs, figures_results_path, surf_full_boundary)
with open(inputs.debug_output_file, "a") as myfile:
try:
myfile.write(f"\nAfter Stage 2 optimization with {inputs.MAXITER_stage_2+inputs.MAXITER_stage_2_simple} iterations")
myfile.write(f"\nAspect ratio at max_mode {max_mode}: {vmec.aspect()}")
myfile.write(f"\nMean iota at {max_mode}: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': myfile.write(f"\nQuasisymmetry objective at max_mode {max_mode}: {qs.total()}")
myfile.write(f"\nSquared flux at max_mode {max_mode}: {Jf.J()}")
except Exception as e:
myfile.write(e)
try:
pprint(f"\nAspect ratio at max_mode {max_mode} after stage 2: {vmec.aspect()}")
pprint(f"\nMean iota at {max_mode} after stage 2: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': pprint(f"\nQuasisymmetry objective at max_mode {max_mode} after stage 2: {qs.total()}")
pprint(f"\nSquared flux at max_mode {max_mode} after stage 2: {Jf.J()}")
except Exception as e:
pprint(e)
mpi.comm_world.Bcast(dofs, root=0)
# Single stage Optimization
if inputs.single_stage:
pprint(f' Performing Single Stage optimization with {inputs.MAXITER_single_stage} iterations')
x0 = np.copy(np.concatenate((JF.x, vmec.x)))
dofs = np.concatenate((JF.x, vmec.x))
obj_and_der = single_stage_obj_and_der()
with MPIFiniteDifference(prob.objective, mpi, rel_step=inputs.finite_difference_rel_step, abs_step=inputs.finite_difference_abs_step, diff_method=inputs.diff_method) as prob_jacobian:
if mpi.proc0_world:
res = minimize(obj_and_der.fun, dofs, args=(prob_jacobian,{'Nfeval':0},max_mode,oustr_dict_inner), jac=True, method='BFGS', options={'maxiter': inputs.MAXITER_single_stage}, tol=1e-15)
with open(inputs.debug_output_file, "a") as myfile:
try:
myfile.write(f'\nAfter Single Stage optimization with {inputs.MAXITER_single_stage} iterations')
myfile.write(f"\nAspect ratio at max_mode {max_mode}: {vmec.aspect()}")
myfile.write(f"\nMean iota at {max_mode}: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': myfile.write(f"\nQuasisymmetry objective at max_mode {max_mode}: {qs.total()}")
myfile.write(f"\nSquared flux at max_mode {max_mode}: {Jf.J()}")
except Exception as e:
myfile.write(e)
mpi.comm_world.Bcast(dofs, root=0)
# Stage 2 Optimization after single_stage
if (inputs.stage_2 and inputs.single_stage) or (inputs.stage_2 and not inputs.single_stage and not inputs.stage_1 and not previous_max_mode==0):
pprint(f' Performing Stage 2 optimization with {inputs.MAXITER_stage_2+inputs.MAXITER_stage_2_simple} iterations')
surf = vmec.boundary
bs.set_points(surf.gamma().reshape((-1, 3)))
if mpi.proc0_world:
dofs, bs, JF = inner_coil_loop(mpi, JF_simple, JF, Jls, Jmscs, Jccdist, Jcsdist, Jf, J_LENGTH, J_CC, J_CS, J_CURVATURE, J_MSC, J_ALS, J_LENGTH_PENALTY, vmec, curves, base_curves, surf, coils_results_path, number_vmec_dofs, bs, max_mode, inputs, figures_results_path, surf_full_boundary, write_coils=False)
with open(inputs.debug_output_file, "a") as myfile:
try:
myfile.write(f'\nAfter last Second Stage optimization with {inputs.MAXITER_stage_2+inputs.MAXITER_stage_2_simple} iterations')
myfile.write(f"\nAspect ratio at max_mode {max_mode}: {vmec.aspect()}")
myfile.write(f"\nMean iota at {max_mode}: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': myfile.write(f"\nQuasisymmetry objective at max_mode {max_mode}: {qs.total()}")
myfile.write(f"\nSquared flux at max_mode {max_mode}: {Jf.J()}")
except Exception as e:
myfile.write(e)
mpi.comm_world.Bcast(dofs, root=0)
if mpi.proc0_world:
pointData = {"B_N": np.sum(bs.B().reshape((inputs.nphi, inputs.ntheta, 3)) * surf.unitnormal(), axis=2)[:, :, None]}
surf_full_boundary.x = surf.x
surf_full_boundary.to_vtk(os.path.join(coils_results_path,inputs.resulting_surface+'max_mode_'+str(max_mode)), extra_data=pointData)
curves_to_vtk(curves, os.path.join(coils_results_path,inputs.resulting_coils+'max_mode_'+str(max_mode)))
bs.save(os.path.join(coils_results_path,inputs.resulting_field_json+'max_mode_'+str(max_mode)+'.json'))
vmec.write_input(os.path.join(vmec_results_path, f'input.{inputs.name}_maxmode{max_mode}'))
os.chdir(vmec_results_path)
try:
pprint(f"\nAspect ratio at max_mode {max_mode}: {vmec.aspect()}")
pprint(f"\nMean iota at {max_mode}: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': pprint(f"\nQuasisymmetry objective at max_mode {max_mode}: {qs.total()}")
pprint(f"\nSquared flux at max_mode {max_mode}: {Jf.J()}")
except Exception as e:
pprint(e)
os.chdir(current_path)
if inputs.single_stage and mpi.proc0_world:
try:
df = pd.DataFrame(oustr_dict_inner)
df.to_csv(os.path.join(current_path, f'output_max_mode_{max_mode}.csv'), index_label='index')
ax=df.plot(
kind='line',
logy=True,
y=['J','Jf','J_length','J_CC','J_CURVATURE','J_MSC','J_ALS','J_LENGTH_PENALTY','Jquasisymmetry','Jiota','Jaspect'],#,'C-C-Sep','C-S-Sep'],
linewidth=0.8)
ax.set_ylim(bottom=1e-9, top=None)
ax.set_xlabel('Number of function evaluations')
ax.set_ylabel('Objective function')
plt.legend(loc=3, prop={'size': 6})
plt.tight_layout()
plt.savefig(os.path.join(figures_results_path, f'optimization_stage3_max_mode_{max_mode}.pdf'), bbox_inches = 'tight', pad_inches = 0)
plt.close()
except Exception as e:
pprint(e)
previous_max_mode = max_mode
oustr_dict_outer.append(oustr_dict_inner)
if mpi.proc0_world:
pointData = {"B_N": np.sum(bs.B().reshape((inputs.nphi, inputs.ntheta, 3)) * surf.unitnormal(), axis=2)[:, :, None]}
surf_full_boundary.x = surf.x
surf_full_boundary.to_vtk(os.path.join(coils_results_path,inputs.resulting_surface), extra_data=pointData)
curves_to_vtk(curves, os.path.join(coils_results_path,inputs.resulting_coils))
bs.save(os.path.join(coils_results_path,inputs.resulting_field_json))
vmec.write_input(os.path.join(current_path, f'input.final'))
if inputs.single_stage:
try:
df = pd.DataFrame(oustr_dict_outer)
df.to_csv(os.path.join(current_path, f'output_final.csv'), index_label='index')
ax=df.plot(kind='line',
logy=True,
y=['J','Jf','J_length','J_CC','J_CURVATURE','J_MSC','J_ALS','J_LENGTH_PENALTY','Jquasisymmetry', 'Jiota','Jaspect'],#,'C-C-Sep','C-S-Sep'],
linewidth=0.8)
ax.set_ylim(bottom=1e-9, top=None)
plt.legend(loc=3, prop={'size': 6})
ax.set_xlabel('Number of function evaluations')
ax.set_ylabel('Objective function')
plt.tight_layout()
plt.savefig(os.path.join(figures_results_path, f'optimization_stage3_final.pdf'), bbox_inches = 'tight', pad_inches = 0)
plt.close()
except Exception as e:
pprint(e)
os.chdir(vmec_results_path)
try:
pprint(f"Aspect ratio after optimization: {vmec.aspect()}")
pprint(f"Mean iota after optimization: {vmec.mean_iota()}")
if not QAorQHorQIorCNT=='QI': pprint(f"Quasisymmetry objective after optimization: {qs.total()}")
pprint(f"Squared flux after optimization: {Jf.J()}")
except Exception as e:
pprint(e)
os.chdir(current_path)
if inputs.create_wout_final:
try:
vmec_final = Vmec(os.path.join(current_path, f'input.final'))
vmec_final.indata.ns_array[:3] = [ 16, 51, 101]
vmec_final.indata.niter_array[:3] = [ 2000, 3000, 20000]
vmec_final.indata.ftol_array[:3] = [1e-14, 1e-14, 1e-14]
vmec_final.run()
if mpi.proc0_world:
shutil.move(os.path.join(current_path, f"wout_final_000_000000.nc"), os.path.join(current_path, f"wout_final.nc"))
os.remove(os.path.join(current_path, f'input.final_000_000000'))
except Exception as e:
pprint('Exception when creating final vmec file:')
pprint(e)
## Create results figures
if os.path.isfile(os.path.join(current_path, f"wout_final.nc")):
pprint('Found final vmec file')
if mpi.proc0_world:
if inputs.vmec_plot_result:
pprint("Plot VMEC result")
vmecPlot2_main(file=os.path.join(current_path, f"wout_final.nc"), name=inputs.name, figures_folder=figures_results_path, coils_curves=curves)
if inputs.booz_xform_plot_result:
pprint('Creating Boozer class for vmec_final')
b1 = Boozer(vmec_final, mpol=64, ntor=64)
pprint('Defining surfaces where to compute Boozer coordinates')
booz_surfaces = np.linspace(0,1,inputs.boozxform_nsurfaces,endpoint=False)
pprint(f' booz_surfaces={booz_surfaces}')
b1.register(booz_surfaces)
pprint('Running BOOZ_XFORM')
b1.run()
if mpi.proc0_world:
# b1.bx.write_boozmn(os.path.join(vmec_results_path,"boozmn_"+inputs.name+".nc"))
pprint("Plot BOOZ_XFORM")
import booz_xform as bx
fig = plt.figure(); bx.surfplot(b1.bx, js=1, fill=False, ncontours=35)
plt.savefig(os.path.join(figures_results_path, "Boozxform_surfplot_1_"+inputs.name+'.pdf'), bbox_inches = 'tight', pad_inches = 0); plt.close()
fig = plt.figure(); bx.surfplot(b1.bx, js=int(inputs.boozxform_nsurfaces/2), fill=False, ncontours=35)
plt.savefig(os.path.join(figures_results_path, "Boozxform_surfplot_2_"+inputs.name+'.pdf'), bbox_inches = 'tight', pad_inches = 0); plt.close()
fig = plt.figure(); bx.surfplot(b1.bx, js=inputs.boozxform_nsurfaces-1, fill=False, ncontours=35)
plt.savefig(os.path.join(figures_results_path, "Boozxform_surfplot_3_"+inputs.name+'.pdf'), bbox_inches = 'tight', pad_inches = 0); plt.close()
if inputs.name[0:2] == 'QH':
helical_detail = True
else:
helical_detail = False
fig = plt.figure(); bx.symplot(b1.bx, helical_detail = helical_detail, sqrts=True)
plt.savefig(os.path.join(figures_results_path, "Boozxform_symplot_"+inputs.name+'.pdf'), bbox_inches = 'tight', pad_inches = 0); plt.close()
fig = plt.figure(); bx.modeplot(b1.bx, sqrts=True); plt.xlabel(r'$s=\psi/\psi_b$')
plt.savefig(os.path.join(figures_results_path, "Boozxform_modeplot_"+inputs.name+'.pdf'), bbox_inches = 'tight', pad_inches = 0); plt.close()
os.chdir(parent_path)
# Stop the timer
stop = time.time()
# Remove spurious files
if mpi.proc0_world:
for objective_file in glob.glob(os.path.join(current_path,f"jac_*")): os.remove(objective_file)
for objective_file in glob.glob(os.path.join(vmec_results_path,f"jac_*")): os.remove(objective_file)
for objective_file in glob.glob(os.path.join(vmec_results_path,f"objective_*")): os.remove(objective_file)
for objective_file in glob.glob(os.path.join(vmec_results_path,f"residuals_*")): os.remove(objective_file)
for objective_file in glob.glob(os.path.join(vmec_results_path,f"*000_*")): os.remove(objective_file)
for objective_file in glob.glob(os.path.join(vmec_results_path,f"parvmec*")): os.remove(objective_file)
for objective_file in glob.glob(os.path.join(vmec_results_path,f"threed*")): os.remove(objective_file)
# Finish
pprint("============================================")
pprint("End of single stage optimization")
pprint(f"Took {stop-start} seconds")
pprint("============================================")