diff --git a/examples/convergence_tests/fixed_step_tracing_cost.py b/examples/convergence_tests/fixed_step_tracing_cost.py new file mode 100644 index 00000000..7db59633 --- /dev/null +++ b/examples/convergence_tests/fixed_step_tracing_cost.py @@ -0,0 +1,87 @@ +"""Where the cost of fixed-step field-line tracing actually comes from. + +An earlier version of this example claimed the slowdown was XLA graph size +caused by max_steps=10000000000 on the fixed-step path. That was wrong. +Measured with AOT lowering, which separates compilation from execution: + + max_steps compile HLO size + 4,096 0.19 s 165,513 chars + 100,000 0.19 s 165,515 chars + 10,000,000,000 0.18 s 165,520 chars + +max_steps changes neither compile time nor graph size. The cost is per-step +runtime, and it scales linearly with the number of steps (4 field lines): + + 500 steps 0.066 s + 2,000 steps 0.255 s + 8,000 steps 1.034 s + 18,000 steps 2.306 s + 50,000 steps 6.306 s + +The same trace with an adaptive PIDController (rtol=atol=1e-7) takes 0.0017 s, +because it takes far fewer, larger steps. + +Note that the fixed-step path in Tracing does not pass a stepsize controller at +all, so it silently inherits the diffrax default. Users should be able to select +an explicit or implicit step size. + +Run from the repository root: + python examples/convergence_tests/fixed_step_tracing_cost.py +""" +import time, jax, jax.numpy as jnp, diffrax, json +from diffrax import diffeqsolve, ODETerm, SaveAt +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import near_axis, BiotSavart +from essos.dynamics import FieldLine + +nfp = 3 +field = near_axis(rc=jnp.array([1, 0.045]), zs=jnp.array([0, -0.045]), etabar=-0.9, nfp=nfp) +current = 17e5 * field.B0 / nfp / 2 +curves = CreateEquallySpacedCurves(n_curves=3, order=5, R=field.R0[0], + r=field.R0[0] / 2., n_segments=50, nfp=nfp, stellsym=True) +biot_savart = BiotSavart(Coils(curves=curves, currents=[current] * 3)) +term = ODETerm(FieldLine) +tmax = 1e-4 +n_save = 200 + + +def measure(n_lines, n_steps, max_steps, controller): + """Lower and compile ahead of time so compile and run are timed separately.""" + R0 = jnp.linspace(field.R0[0], 1.04 * field.R0[0], n_lines) + y0 = jnp.stack([R0, jnp.zeros(n_lines), jnp.zeros(n_lines)], axis=-1) + ts = jnp.linspace(0, tmax, n_save) + + def trace_one(initial_condition): + return diffeqsolve(term, t0=0., t1=tmax, dt0=tmax / n_steps, y0=initial_condition, + solver=diffrax.Dopri5(), args=biot_savart, saveat=SaveAt(ts=ts), + throw=False, max_steps=max_steps, + stepsize_controller=controller).ys + + traced = jax.jit(jax.vmap(trace_one)) + lowered = traced.lower(y0) + t0 = time.time(); compiled = lowered.compile(); t_compile = time.time() - t0 + t0 = time.time(); out = compiled(y0); jax.block_until_ready(out); t_run = time.time() - t0 + return t_compile, t_run, len(compiled.as_text()) + + +results = {} + +print('compile time and graph size vs max_steps (4 lines, 2000 fixed steps)') +for max_steps in [4096, 100_000, 10_000_000_000]: + t_compile, t_run, hlo = measure(4, 2000, max_steps, diffrax.ConstantStepSize()) + print(f' max_steps={max_steps:<15d} compile {t_compile:5.2f} s run {t_run:6.3f} s HLO {hlo} chars') + results[f'max_steps_{max_steps}'] = dict(compile=t_compile, run=t_run, hlo=hlo) + +print('\nrun time vs number of fixed steps (4 lines, max_steps generous)') +for n_steps in [500, 2000, 8000, 18000, 50000]: + t_compile, t_run, _ = measure(4, n_steps, max(2 * n_steps, 4096), diffrax.ConstantStepSize()) + print(f' {n_steps:6d} steps compile {t_compile:5.2f} s run {t_run:7.3f} s') + results[f'steps_{n_steps}'] = dict(compile=t_compile, run=t_run) + +print('\nsame trace with an adaptive controller') +t_compile, t_run, _ = measure(4, 2000, 4096, diffrax.PIDController(rtol=1e-7, atol=1e-7)) +print(f' PIDController 1e-7 compile {t_compile:5.2f} s run {t_run:7.4f} s') +results['adaptive'] = dict(compile=t_compile, run=t_run) + +json.dump(results, open('fixed_step_tracing_cost.json', 'w'), indent=1) +print('\nwrote fixed_step_tracing_cost.json') diff --git a/examples/convergence_tests/nearaxis_joint_optimization.py b/examples/convergence_tests/nearaxis_joint_optimization.py new file mode 100644 index 00000000..95794b50 --- /dev/null +++ b/examples/convergence_tests/nearaxis_joint_optimization.py @@ -0,0 +1,55 @@ +"""Joint coils + near-axis optimization (optimize_coils_and_nearaxis). + +Two stages with reduced parameters: coils fitted to a fixed near-axis field, +then coils and near-axis optimized together. + +Measured: stage 1 6.2 s, stage 2 6.8 s; iota 0.41831 -> 0.45294, +max elongation 2.413 -> 2.559, on-axis B error 1.354 -> 1.569. +""" +"""Joint coils + near-axis optimization (the second example), reduced params.""" +import matplotlib; matplotlib.use('Agg') +from time import time +import jax.numpy as jnp, json +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import near_axis, BiotSavart +from essos.optimization import optimize_loss_function +from essos.objective_functions import (loss_coils_for_nearaxis, loss_coils_and_nearaxis, + difference_B_gradB_onaxis) +nfp=3; MCL=4.; MCC=6.; order=5; nseg=order*10; ncoils=3; TOL=1e-8; NFEV=100 +f0 = near_axis(rc=jnp.array([1,0.045]), zs=jnp.array([0,-0.045]), etabar=-0.9, nfp=nfp) +cur = 17e5*f0.B0/nfp/2 +curves = CreateEquallySpacedCurves(n_curves=ncoils, order=order, R=f0.R0[0], + r=f0.R0[0]/2.0, n_segments=nseg, nfp=nfp, stellsym=True) +ci = Coils(curves=curves, currents=[cur]*ncoils) + +print('stage 1: coils only ...', flush=True) +t0=time() +c1 = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=ci.x, coils=ci, + tolerance_optimization=TOL, maximum_function_evaluations=NFEV, field_nearaxis=f0, + max_coil_length=MCL, max_coil_curvature=MCC) +t1=time()-t0 +print(f' stage 1 done {t1:.2f}s', flush=True) + +print('stage 2: joint coils + near-axis ...', flush=True) +x0 = jnp.concatenate((c1.x, f0.x)) +t0=time() +res = optimize_loss_function(loss_coils_and_nearaxis, initial_dofs=x0, coils=ci, + tolerance_optimization=TOL, maximum_function_evaluations=NFEV, field_nearaxis=f0, + max_coil_length=MCL, max_coil_curvature=MCC) +t2=time()-t0 +c2, f1 = res +print(f' stage 2 done {t2:.2f}s', flush=True) + +Bd0,Gd0 = difference_B_gradB_onaxis(f0, BiotSavart(c1)) +Bd1,Gd1 = difference_B_gradB_onaxis(f1, BiotSavart(c2)) +out = dict(stage1_time=t1, stage2_time=t2, + iota_initial=float(f0.iota), iota_optimized=float(f1.iota), + elong_initial=float(max(f0.elongation)), elong_optimized=float(max(f1.elongation)), + B_err_stage1=float(jnp.sum(jnp.abs(Bd0))), B_err_joint=float(jnp.sum(jnp.abs(Bd1))), + gradB_err_stage1=float(jnp.sum(jnp.abs(Gd0))), gradB_err_joint=float(jnp.sum(jnp.abs(Gd1)))) +print() +print(f"iota {out['iota_initial']:.5f} -> {out['iota_optimized']:.5f}") +print(f"max elongation {out['elong_initial']:.3f} -> {out['elong_optimized']:.3f}") +print(f"B error {out['B_err_stage1']:.4f} -> {out['B_err_joint']:.4f}") +print(f"gradB error {out['gradB_err_stage1']:.4f} -> {out['gradB_err_joint']:.4f}") +json.dump(out, open('joint_results.json','w'), indent=1) diff --git a/examples/convergence_tests/nearaxis_optimization_convergence.py b/examples/convergence_tests/nearaxis_optimization_convergence.py new file mode 100644 index 00000000..d258b9fe --- /dev/null +++ b/examples/convergence_tests/nearaxis_optimization_convergence.py @@ -0,0 +1,68 @@ +"""Convergence study for near-axis coil optimization. + +Sweeps the function-evaluation budget for loss_coils_for_nearaxis and records +wall time, final loss, and the on-axis B / grad-B errors at each budget. + +Run from the repository root: + python examples/convergence_tests/nearaxis_optimization_convergence.py + +Uses reduced parameters so each point completes in seconds. +""" +# Full-parameter near-axis coil optimization + convergence study +import matplotlib; matplotlib.use('Agg') +from time import time +import jax.numpy as jnp, matplotlib.pyplot as plt, json +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import near_axis, BiotSavart +from essos.optimization import optimize_loss_function +from essos.objective_functions import loss_coils_for_nearaxis, difference_B_gradB_onaxis + +max_coil_length=4; max_coil_curvature=6; order=5 +nseg=order*10; ncoils=3; nfp=3; tol=1e-8 +field = near_axis(rc=jnp.array([1,0.045]), zs=jnp.array([0,-0.045]), etabar=-0.9, nfp=nfp) +cur = 17e5*field.B0/nfp/2 + +def make_coils(): + c = CreateEquallySpacedCurves(n_curves=ncoils, order=order, R=field.R0[0], + r=field.R0[0]/2.0, n_segments=nseg, nfp=nfp, stellsym=True) + return Coils(curves=c, currents=[cur]*ncoils) + +def eval_loss(coils): + return float(loss_coils_for_nearaxis(coils.x, field, coils.dofs_curves, + coils.currents_scale, nfp, max_coil_length=max_coil_length, + n_segments=nseg, stellsym=True, max_coil_curvature=max_coil_curvature)) + +results={} +for nfev in [10,20,50,100,200]: + ci = make_coils(); lb = eval_loss(ci) + t0=time() + co = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=ci.x, coils=ci, + tolerance_optimization=tol, maximum_function_evaluations=nfev, field_nearaxis=field, + max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature) + el=time()-t0; la=eval_loss(co) + Bd,gBd = difference_B_gradB_onaxis(field, BiotSavart(co)) + results[nfev]=dict(time=el, loss_initial=lb, loss_final=la, + reduction=lb/la, B_err=float(jnp.sum(jnp.abs(Bd))), gradB_err=float(jnp.sum(jnp.abs(gBd)))) + print(f'nfev={nfev:4d} {el:7.2f}s loss {lb:.4f} -> {la:.4f} ({lb/la:.2f}x) ' + f'B_err={results[nfev]["B_err"]:.4f} gradB_err={results[nfev]["gradB_err"]:.4f}', flush=True) + if nfev==200: ci_final, co_final = ci, co + +json.dump(results, open('nearaxis_results.json','w'), indent=1) + +fig,(a1,a2)=plt.subplots(1,2,figsize=(11,4.2)) +ks=sorted(results); a1.plot(ks,[results[k]['loss_final'] for k in ks],'o-') +a1.set_xlabel('function evaluations'); a1.set_ylabel('final loss'); a1.set_yscale('log') +a1.set_title('Near-axis coil optimization convergence'); a1.grid(alpha=.3) +a2.plot(ks,[results[k]['time'] for k in ks],'s-',color='darkred') +a2.set_xlabel('function evaluations'); a2.set_ylabel('wall time (s)') +a2.set_title('Optimization cost'); a2.grid(alpha=.3) +plt.tight_layout(); plt.savefig('nearaxis_convergence.png',dpi=200) + +fig=plt.figure(figsize=(11,5)) +b1=fig.add_subplot(121,projection='3d'); b2=fig.add_subplot(122,projection='3d') +ci_final.plot(ax=b1,show=False); field.plot(ax=b1,show=False,alpha=0.25) +co_final.plot(ax=b2,show=False); field.plot(ax=b2,show=False,alpha=0.25) +b1.set_title(f'Initial coils (loss={results[200]["loss_initial"]:.3g})') +b2.set_title(f'Optimized, 200 nfev (loss={results[200]["loss_final"]:.3g})') +plt.tight_layout(); plt.savefig('nearaxis_coils_200.png',dpi=200) +print('saved plots + nearaxis_results.json') diff --git a/examples/convergence_tests/nearaxis_optimization_quick.py b/examples/convergence_tests/nearaxis_optimization_quick.py new file mode 100644 index 00000000..c2f20918 --- /dev/null +++ b/examples/convergence_tests/nearaxis_optimization_quick.py @@ -0,0 +1,64 @@ +# Small-parameter near-axis coil optimization, per Rogerio's suggestion: +# "start with smaller max times, less num steps, and less particles" +import matplotlib +matplotlib.use('Agg') +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import near_axis, BiotSavart +from essos.optimization import optimize_loss_function +from essos.objective_functions import loss_coils_for_nearaxis + +# ---- reduced parameters (upstream example uses 200 fn evals, order 5) ---- +max_coil_length = 4 +max_coil_curvature = 6 +order_Fourier_series_coils = 3 # was 5 +number_coil_points = order_Fourier_series_coils*10 +maximum_function_evaluations = 20 # was 200 +number_coils_per_half_field_period = 3 +tolerance_optimization = 1e-6 # was 1e-8 + +rc=jnp.array([1, 0.045]); zs=jnp.array([0,-0.045]); etabar=-0.9; nfp=3 +field = near_axis(rc=rc, zs=zs, etabar=etabar, nfp=nfp) +print(f'near_axis built: iota={field.iota}, B0={field.B0}, R0[0]={field.R0[0]}') + +current_on_each_coil = 17e5*field.B0/nfp/2 +major_radius_coils = field.R0[0] +curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, + order=order_Fourier_series_coils, + R=major_radius_coils, r=major_radius_coils/2.0, + n_segments=number_coil_points, + nfp=nfp, stellsym=True) +coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) + +def eval_loss(coils): + return loss_coils_for_nearaxis( + coils.x, field, coils.dofs_curves, coils.currents_scale, nfp, + max_coil_length=max_coil_length, n_segments=number_coil_points, + stellsym=True, max_coil_curvature=max_coil_curvature) + +loss_before = eval_loss(coils_initial) +print(f'initial loss = {loss_before}') + +print(f'Optimizing with {maximum_function_evaluations} function evaluations...') +t0 = time() +coils_optimized = optimize_loss_function( + loss_coils_for_nearaxis, initial_dofs=coils_initial.x, coils=coils_initial, + tolerance_optimization=tolerance_optimization, + maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field, + max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature) +elapsed = time()-t0 +print(f'OPTIMIZATION COMPLETED in {elapsed:.2f} s') + +loss_after = eval_loss(coils_optimized) +print(f'final loss = {loss_after}') +print(f'loss reduction factor = {loss_before/loss_after:.3f}x') + +fig = plt.figure(figsize=(10,5)) +ax1 = fig.add_subplot(121, projection='3d'); ax2 = fig.add_subplot(122, projection='3d') +coils_initial.plot(ax=ax1, show=False); field.plot(ax=ax1, show=False, alpha=0.2) +coils_optimized.plot(ax=ax2, show=False); field.plot(ax=ax2, show=False, alpha=0.2) +ax1.set_title(f'Initial (loss={loss_before:.4g})'); ax2.set_title(f'Optimized (loss={loss_after:.4g})') +plt.tight_layout(); plt.savefig('nearaxis_opt.png', dpi=150) +print('saved /tmp/essos-work/nearaxis_opt.png')