Ballooning Optimizations

[matt-pharr]

The new ballooning feature takes the vast majority of the runtime of the example cases now. There were a few trivial optimizations to make to reduce its runtime by ~60%. The primary one was fixing some repeated allocations in the ode solve.

Speedup timings:

---

Isolated ballooning scaling on flux — current code

threads	compute_ballooning_stability!	ballooning_alpha_boundary	ballooning total	× vs develop
develop (serial)	1.58 s	44.16 s	45.74 s	1.0×
1	0.539 s	14.79 s	15.33 s	2.98×
2	0.288 s	7.75 s	8.04 s	5.69×
4	0.175 s	4.77 s	4.95 s	9.25×
8	0.112 s	2.94 s	3.05 s	15.0×
16	0.065 s	1.79 s	1.86 s	24.6×
32	0.044 s	1.37 s	1.41 s	32.4×

Full DIII-D example (end-to-end main, warm)

state (at 6 threads unless noted)	full main
develop	74.0 s
+ mem-alloc (commit #1)	52.9 s
+ SVector (commit #2)	45.5 s
+ threading compute only (original one-liner)	45.0 s
+ threading alpha_boundary :greedy (current state)	34.1 s
current branch @ 1 thread (reference)	50.2 s

---

Regression harness report:

================================================================
Case: diiid_n1 — DIII-D-like equilibrium, n=1, ideal + perturbed equilibrium
================================================================

Regression Report: diiid_n1
================================================================================================================
Ref 1: cfb7693  @ cfb76939 (2026-06-24)
Ref 2: local  @ local (2026-06-26)
----------------------------------------------------------------------------------------------------------------
Quantity                                      cfb7693          local            Diff               Status
----------------------------------------------------------------------------------------------------------------
root-area-weighted total energy Re(et[1])     1.535016e-15     1.535014e-15     2.5e-21            OK
root-area-weighted total energy Im(et[1])     -2.063679e-20    -3.238097e-20    1.2e-20            OK
root-area-weighted plasma energy Re(ep[1])    5.910651e-16     5.910593e-16     5.8e-21            OK
root-area-weighted vacuum energy Re(ev[1])    9.439628e-16     9.439622e-16     6.0e-22            OK
vacuum matrix min eigenvalue (XiNorm)         1.863758e-01     1.863758e-01     1.8e-13            OK
root-area-weighted plasma energy (all)        [35 elem]        [35 elem]        2.8e-17            OK
root-area-weighted vacuum energy (all)        [35 elem]        [35 elem]        9.3e-18            OK
root-area-weighted total energy (all)         [35 elem]        [35 elem]        4.5e-17            OK
ODE steps (saved)                             1278             1278             0.0e+00            OK
ODE steps (total)                             1966             1966             0.0e+00            OK
q0                                            1.209717e+00     1.209717e+00     0.0e+00            OK
q95                                           4.505011e+00     4.505011e+00     0.0e+00            OK
beta_t                                        1.301227e-02     1.301227e-02     0.0e+00            OK
beta_n                                        1.357940e+00     1.357940e+00     0.0e+00            OK
internal inductance li1                       1.046066e+00     1.046066e+00     0.0e+00            OK
internal inductance li2                       8.551636e-01     8.551636e-01     0.0e+00            OK
internal inductance li3                       8.879379e-01     8.879379e-01     0.0e+00            OK
poloidal beta betap1                          6.343911e-01     6.343911e-01     0.0e+00            OK
poloidal beta betap2                          5.186177e-01     5.186177e-01     0.0e+00            OK
poloidal beta betap3                          5.384938e-01     5.384938e-01     0.0e+00            OK
# singular surfaces                           5                5                0.0e+00            OK
singular psi locations                        [5 elem]         [5 elem]         0.0e+00            OK
singular q values                             [5 elem]         [5 elem]         0.0e+00            OK
current beta betaj                            4.252872e-01     4.252872e-01     0.0e+00            OK
plasma volume                                 1.945603e+01     1.945603e+01     0.0e+00            OK
plasma current                                1.205445e+00     1.205445e+00     0.0e+00            OK
mpert                                         35               35               0.0e+00            OK
npert                                         1                1                0.0e+00            OK
toroidal field bt0                            2.024475e+00     2.024475e+00     0.0e+00            OK
wall field bwall                              3.879849e-01     3.879849e-01     0.0e+00            OK
aspect ratio                                  2.702739e+00     2.702739e+00     0.0e+00            OK
elongation kappa                              1.658590e+00     1.658590e+00     0.0e+00            OK
q profile (checksum)                          dd3a12fb525e...  dd3a12fb525e...  identical          OK
pressure profile (checksum)                   e532229541b5...  e532229541b5...  identical          OK
Mercier D_I profile (checksum)                b8aa8e933417...  b8aa8e933417...  identical          OK
resistive interchange D_R profile (checksum)  b7c156d239c8...  b7c156d239c8...  identical          OK
ballooning Delta' profile (checksum)          a61852d006f7...  7e219728af0e...  0.000e+00 (0.00%)  ** CHANGED **
delta prime (BVP diagonal)                    [5 elem]         [5 elem]         9.1e-13            OK
island half-widths                            [5 elem]         [5 elem]         1.7e-16            OK
Chirikov parameter                            [5 elem]         [5 elem]         1.0e-14            OK
||resonant area-weighted field||              4.204598e-04     4.204598e-04     2.1e-18            OK
PE plasma energy                              3.721521e+00     3.721521e+00     8.9e-16            OK
PE vacuum energy                              3.459135e+00     3.459135e+00     1.3e-15            OK
PE surface energy                             5.576490e+00     5.576490e+00     4.4e-15            OK
PE toroidal torque                            4.584048e-02     4.584048e-02     6.7e-16            OK
Runtime (s)                                   166.3s           121.4s                              --
resonant area-weighted field b^r              [5 elem]         [5 elem]         1.7e-17            OK
================================================================================================================
Summary: 1 changed, 45 unchanged

[matt-pharr]

@logan-nc I posted the tables, this is done and ready for review

[logan-nc]

From my quick skim, this looks good to me.
I leave the rigorous testing to @jmlmir369

[logan-nc]

@matt-pharr the table says 1.0× for single thread, but it looks like there were a few improvements beyond simple parallelization - I would have thought even the single thread would have been faster than before?

Still seems slow IMO... I think the biggest lever should be the psi grid. Why not make a new psi grid for this that is concentrated where it needs to be? For example, if someone uses a ldp grid then we wast a lot of time doing deep core ballooning calculations on the equilibrium grid. Limits are pretty smooth in psi, so I bet a smart ~30-60 pt grid would be good enough for most applications. Thoughts @jmlmir369?

[matt-pharr]

@logan-nc sorry, the 1x means 1x single threaded with all performance improvements made. I updated the table to include the speedup from the single threaded improvements.

I will leave the psi grid up to you, but it does run very quickly now.

[jmlmir369]

Thank you for your help with the optimization. I have reviewed all the changes.

Additionally, I have implemented some further optimization by restoring a logic from the original balloo.jl.

Previously, the ODE integration over $(-\infty, 0)$ and $(0, \infty)$ used an arbitrary large value of 16.5 as the integration limit. However, the original implementation determined this limit dynamically. Based on the quadratic form $a(\theta - b)^2 + c$, where the condition $\theta \gg \sqrt{c/a}$ is key, the upper bound is set as theta_max = min(sqrt(c/a)*10, THETA_MAX_CAP). I have restored this logic.

The plot above compares the results using the constant value of 16.5 versus the dynamic theta_max calculation. Due to the divergence of delta prime in the core region, the plots are separated into core and non-core regions. As indicated by the dashed lines, the relative error between the two methods is very small.

The figure above shows how the time consumption decreases in the non-core region, where the $\sqrt{c/a}$ value actually becomes smaller than the max.