Near-axis optimization convergence study + fixed-step tracing compile fix - #49
Open
Agastyarathee wants to merge 3 commits into
Open
Near-axis optimization convergence study + fixed-step tracing compile fix#49Agastyarathee wants to merge 3 commits into
Agastyarathee wants to merge 3 commits into
Conversation
Adds an isolated example that sweeps the function-evaluation budget for loss_coils_for_nearaxis and reports wall time, final loss, and on-axis B / grad-B errors, so the convergence behaviour can be reproduced locally. Results on essos 0.17.dev86, jax 0.11.0, diffrax 0.7.2: nfev=10 0.68s loss 143.07 -> 27.90 (5.13x) nfev=20 1.02s loss 143.07 -> 10.38 (13.79x) nfev=50 2.72s loss 143.07 -> 4.49 (31.86x) nfev=100 5.65s loss 143.07 -> 3.63 (39.46x) nfev=200 11.66s loss 143.07 -> 3.31 (43.21x) grad-B error falls 12.6x across the sweep while B error falls 1.7x, indicating the on-axis field gradient is the binding constraint.
… example The fixed-step field-line path (model='FieldLine') sets max_steps=1e10 and relies on an implicit stepsize controller, which makes XLA build a graph so large that a 2-step warm-up projected ~43 h to compile. Passing diffrax.ConstantStepSize() explicitly with a realistic max_steps compiles in seconds: 10 field lines, 18,000 fixed steps: 5.60 s compile, 5.17 s run 10 field lines, 2,000 fixed steps: 1.00 s compile, 0.58 s run Also adds an example for the joint coils + near-axis optimization (optimize_coils_and_nearaxis) with reduced parameters: stage 1 (coils only) 6.2 s stage 2 (joint) 6.8 s iota 0.41831 -> 0.45294, max elongation 2.413 -> 2.559 on-axis B error 1.354 -> 1.569 (expected trade for higher iota)
My earlier claim that the fixed-step slowdown was XLA graph size caused by max_steps=10000000000 was wrong. Measured with AOT lowering, which separates compilation from execution, max_steps changes neither compile time nor graph size: 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 The cost is per-step runtime and scales linearly with step count (4 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) runs in 0.0017 s, taking far fewer and larger steps. Renames fixed_step_tracing_fix.py -> fixed_step_tracing_cost.py, since the example characterises the cost rather than fixing it. Also notes that the fixed-step path in Tracing passes no stepsize controller and so inherits the diffrax default; selecting explicit vs implicit stepping should be a user choice.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Isolated PR for the near-axis optimization convergence test, as requested. Adds example scripts under
examples/convergence_tests/— no existing files modified.Convergence sweep (
nearaxis_optimization_convergence.py) —loss_coils_for_nearaxisacross nfev 10→200:∇B error falls 12.6× while B error falls only 1.7×.
Fixed-step tracing cost (
fixed_step_tracing_cost.py) — correcting an earlier claim of mine. I had said the fixed-step slowdown was XLA graph size frommax_steps=10000000000. That was wrong. Measured with AOT lowering, which separates compilation fromexecution,
max_stepschanges neither compile time nor graph size:The cost is per-step runtime, scaling linearly with step count (4 field lines): 500 steps 0.066 s, 2,000 0.255 s, 8,000 1.034 s, 18,000 2.306 s, 50,000 6.306 s. The same trace with an adaptive
PIDController(rtol=atol=1e-7) runs in 0.0017 s.Worth noting the fixed-step path passes no stepsize controller, so it inherits the diffrax default — selecting explicit vs implicit stepping should probably be a user choice.
Joint optimization (
nearaxis_joint_optimization.py) —optimize_coils_and_nearaxis: iota 0.4183 → 0.4529, max elongation 2.41 → 2.56, on-axis B error 1.354 → 1.569 (expected trade), 6.8 s.Run on essos 0.17.dev86, jax 0.11.0, diffrax 0.7.2.