Near-axis optimization convergence study + fixed-step tracing compile fix

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')

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)

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')

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')