Repository navigation
DKX 2.7.0
Faster by default and ready for optimization: a memory-aware Krylov restart
that removes the high-Nx stall, a self-compiling structured route for
callers that do not jit, and a drift-kinetic bootstrap-current objective for
VMEX optimization.
Correctness
-
Refuse failed bootstrap-current evidence (#274).
KineticBootstrapCurrent
replaced a nonfinite surface current with zero, so an unsuccessful solve
looked like an optimal one; nonfinite currents, residuals and squared terms
and a non-positive reference current are now refused, and the ambipolar
interpolation requires every scan point to be accepted and finite. The
scan-promotion audit likewise substituted zero when an output had neither
FSABjHatOverRootFSAB2norFSABjHat, passing the gate with a zero
objective; a missing or nonfinite bootstrap current is now refused. -
Convert the handedness of the Boozer route in one place
(boozer_route_psi_a_hat, and its traced form
convert_boozer_route_handedness):psiAHat = |phi_edge|/(2 pi) signgs sign(G + iota I). A VMEX equilibrium hassigngs = -1andbooz_xform
returnsG > 0, so takingpsiAHat = +|phi_edge|/(2 pi)flipped every flux
and<j.B>on that route against the VMEC-file route and Redl.
optimize_QA_bootstrap.py,optimize_QH_bootstrap.pyand the gradient hook
ofbootstrap_consistency_kinetic_loop.pyprinted the flipped sign; their
objectives were squares or magnitudes, so the optimizations were unaffected.
Performance
-
Retain only the active pitch rows in reusable coarse factors (#273). The
generated reusable preconditioner factored every subsystem through allNxi
rows, including the masked identity rows pastNxi_for_x; it now factors
max(1, Nxi_for_x)rows and applies the uncoupled tails as
r / (1 + floor)at factor precision. On the recorded NCSX layout that is
744 rows instead of 1,111 (predicted float64 LU and pivots 14.5 GiB instead of
21.7 GiB); a tracedNxi_for_xkeeps the rectangular layout. The operator,
masks, pinning and residual acceptance are unchanged. -
Memory-aware Krylov restart by default.
solve(restart=None)and
SolverOptions(restart=None), the new defaults, run five FGMRES cycles of
30 and two of 100, then continue from the iterate at the longest restart, at
most 1,000, whose basis (2 * restart * unknowns * 8bytes, since flexible
GMRES storesVandZ) fitskrylov_memory_budget_gb
(DKX_KRYLOV_MEMORY_BUDGET_GB, else a quarter of available memory). The
previous policy widened only to 100 and only while two such bases fit in
256 MiB, so decks above 167,772 unknowns never widened. On the HSX-like
Nxladder the iteration growth fromNx = 10to 16 was restart
stagnation: 2,788 iterations at restart 200 against 357 at 1,000, and 7,167
against 390 at(Nxi, Nx) = (40, 16)
(docs/experiments/2026-09-23-restart-and-direct-reach.md). On the public
tokamak full-Fokker-Planck deck at(Nx, Nxi) = (24, 40), 87,362 unknowns,
the default converges in 1,736 iterations where the previous one stalled
and failed every rung of the stall ladder (docs/performance.rst, "Krylov
restart length"). Decks that converged within 350 iterations follow the
same path as before. The stall
ladder's larger-budget rung spends its budget at the wide restart. An integer
restartfixes the cycle size as before; explicitmethod="gmres",
differentiable and traced solves, and caller-supplied preconditioners keep a
fixed size of 30 unless given one. -
dkx solve-v3 --restartwas printed but never applied; its default reads
autoand its help says so.
Optimization
dkx.bootstrap.KineticBootstrapMismatch: the drift-kinetic bootstrap
current as a traced VMEX objective term, with the interface and normalized
residual of VMEX'sRedlBootstrapMismatch. On each kinetic surface the chain
VMEX state ->boozer_input_tables->booz_xform_jax-> DKX structured
solve -><j.B>is traced, so VMEX's implicit Jacobian differentiates it;
mismatch=Falsegives a pure<j.B>row. Pitch-angle scattering is the
default collision operator (cheap, no momentum restoration).examples/optimization/QA_optimization_bootstrap_dkx.pyis VMEX's
QA_optimization_bootstrap.pywith the DKX row in place of Redl's
(BOOTSTRAP_MODEL = "dkx" | "redl" | "both"), replacing the
finite-difference version; the seed deck ships inexamples/data.- The structured direct route compiles itself for callers that do not
jit
(#279). The factorization, the refined substitution, the whole
differentiable solve (onecustom_linear_solvewith the residual guard
inside it),KineticOperator.rhsandprofile_moments_from_operatoreach
run as one cached executable per operator structure, so an eager
jax.value_and_gradof an objective no longer dispatches, linearizes and
transposes thousands of small operations one at a time. On the
16,230-unknown structured deck, on heavily loaded hosts, the eager primal
went from 5.0–5.9× to 1.18–1.30× of its compiled time and the eager
gradient from 7.3–9.6× to 1.24–1.38×; the fixed eager overhead fell from
1.33 s to 0.017 s (primal) and from 2.5 s to 0.10 s (gradient). Answers are
unchanged to round-off, and the plain and differentiable solves return the
same bits. The first call of each new operator structure pays a one-time
compile.