Skip to content

DKX 2.7.0

Choose a tag to compare

@rogeriojorge rogeriojorge released this 27 Sep 18:24
· 111 commits to main since this release
4d38f26

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
    FSABjHatOverRootFSAB2 nor FSABjHat, 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 has signgs = -1 and booz_xform
    returns G > 0, so taking psiAHat = +|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.py and the gradient hook
    of bootstrap_consistency_kinetic_loop.py printed 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 all Nxi
    rows, including the masked identity rows past Nxi_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 traced Nxi_for_x keeps 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 * 8 bytes, since flexible
    GMRES stores V and Z) fits krylov_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
    Nx ladder the iteration growth from Nx = 10 to 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
    restart fixes the cycle size as before; explicit method="gmres",
    differentiable and traced solves, and caller-supplied preconditioners keep a
    fixed size of 30 unless given one.

  • dkx solve-v3 --restart was printed but never applied; its default reads
    auto and 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's RedlBootstrapMismatch. 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=False gives a pure <j.B> row. Pitch-angle scattering is the
    default collision operator (cheap, no momentum restoration).
  • examples/optimization/QA_optimization_bootstrap_dkx.py is VMEX's
    QA_optimization_bootstrap.py with the DKX row in place of Redl's
    (BOOTSTRAP_MODEL = "dkx" | "redl" | "both"), replacing the
    finite-difference version; the seed deck ships in examples/data.
  • The structured direct route compiles itself for callers that do not jit
    (#279). The factorization, the refined substitution, the whole
    differentiable solve (one custom_linear_solve with the residual guard
    inside it), KineticOperator.rhs and profile_moments_from_operator each
    run as one cached executable per operator structure, so an eager
    jax.value_and_grad of 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.