Repository navigation
DKX 2.10.0
Production scans that finish: a robust ambipolar root search, radial profiles
that bracket only certified points and continue between radii, scans that keep
going past failed points, and dkx sfincs scan-profile for radius x Er scans
from a SFINCS namelist and profile file. Large-|E_r| Fokker-Planck points that
stalled for 20 minutes now solve in about a minute. Memory: monoenergetic runs
use 1.3 GB instead of 24 GB, and multi-surface runs stay flat without clearing
the compile cache. New cross-code validation against SFINCS, an 8-page guided
tutorial, and a faster CI.
-
Fix: host memory grew with every surface of an ambipolar profile because each batched solve (
batched_solve, one per E_r evaluation) ran an eagerlax.mapover a fresh closure: one executable compiled and kept per call, 11 per surface. The compiled map is now cached on the operator it closes over and freed with it, anddkx runno longer callsjax.clear_caches()between surfaces. W7-X wout, 5 surfaces, 15x15x12x4 FP ambipolar, 8 pinned cores of the shared office host (load 10-13), cold/warm compile cache: RSS after the last surface 1.59/1.40 GB with the clear on main, 4.5/4.4 GB without it, 0.80/0.61 GB now; wall time 198/201 s on main, 170/113 s now; backend compiles 481 (main), 82 (now). Plain profile, same case: 33.6/24.3 s to 21.0/12.4 s, 461 to 97 compiles. Single-device and sharded maps now agree to the last ulp (3.5e-18 on states of 1e-3) rather than bitwise. -
Fix:
scan_failures.json(and the failed-point list ofscan-erandscan-profile) recorded the failed E_r rounded to 4 significant figures, read back from the point's directory name. It now records the float that was solved. -
The ramped truncated structured solve is one compiled executable per layout. It ran its
lax.scansweeps eagerly over fresh closures, so every solve (and every E_r field of a scan) retraced and compiled one executable perNxi_for_xgroup (8 on the tokamak decks). Now_truncated_sweepsis a module-leveljittaking the operator as an argument.tokamak_1species_PASCollisions_noErthroughdkx sfincs write-output, 2 pinned cores of the shared office host, alternating A/B/A/B: cold 14.5-15.0 s to 9.5-9.6 s (235 to 131 compiles), second solve in the same process 0.87-0.94 s to 0.13 s (8 to 2 compiles). Peak RSS +50-70 MB (465 to 515-537 MB). -
Tests: the dense-route memory budget is pinned for the suite (
tests/conftest.py)._coarse_memory_budgetreads available memory, so on a loaded host (office, 0 GB free) sixtest_speed_triangle_precondandtest_tier2_memory_guardtests failed on a correct tree; with the budget pinned all pass at 0 GB free. -
dkx sfincs scan-erno longer loses a scan to one bad point. A point that raises is recorded inscan_failures.json, the other points are written, and the command exits 1. When the shared-operator batch cannot converge a point, that point is retried alone. When the batch would be a traced GMRES (full Fokker–Planck with the E_r terms, which has no preconditioner escalation and no early exit) the scan goes point by point from the start; a batch of 9 E_r values for an HSX QHS full-FP deck at 66,004 unknowns had written no point after 50 minutes (8 pinned cores of the shared office host, load 25). -
New
dkx sfincs scan-profile: the namelist-route equivalent of upstreamsfincsScanwithscanType = 5(radius × E_r from a profile file and the deck's!sslines). A malformed!ssor profile line raises an error naming the file and line, and a radius outside the profile file's range is refused (it used to take the edge value and a zero gradient). -
Krylov escalation at large |E_r| with the xDot/xiDot terms. On an HSX QHS full-FP deck (66,004 unknowns, 8 pinned cores of the shared office host) the points at |E_r| = 15 and 20 took 19 to 24 minutes each: the
coupledpreconditioner, which factors the operator without the E_r terms, stalled after 768 to 1,262 iterations, and each escalation rung then retried at the probe restart of 30, where only the last rung (wide restart) converged.autonow starts withcoarse_trianglewhen the E_r terms are active and nonzero and its bands fit, and every preconditioner rung follows the restart policy of the first attempt: the same four points take 44 to 114 s and all nine E_r points from -20 to 20 converge (main: 16 to 1,415 s). -
Fix: the
sparse_trianglerung of the escalation always failed withUnexpectedTracerError. SOLVAX's compile reuse turns arrays closed over by a jitted callback into jit arguments, including those of the nested host callback, soupperand the operator reached it as tracers. The host state is now on one opaque object. -
Fix: the stall message named the requested preconditioner (
coarse) instead of the oneautohad used. -
The last escalation rung runs SuperLU above
max_dense_sizewhen its estimated peak fits the memory available now (1.25 x 9.9 GB x (n/66,004)^1.6; the HSX QHS full-FP deck at E_r = -20, 66,004 unknowns, converged on the sparse direct route in 47 minutes at 9.9 GB on 8 pinned cores of a loaded host). Before, any stall above 8,192 unknowns raised after the iterative rungs, and the message asked for MUMPS. When the estimate does not fit, the error states it. -
Docs: the tutorial track is rewritten as eight guided pages (first run, reading results, geometry and species, collisions and solvers, scans and resolution, ambipolar
E_r, transport matrices,Phi1/bootstrap/gradients) with the real output and figure of every step, replacing the four-part track.13_convergence.pygains a resolution ladder. Figures are palette-compressed (docs/figures/publish_tutorial_figures.py, 244 KB for 13) and registered indocs/figure_provenance.json. -
Fix:
find_ambipolar_erkept every field's full solve state, preconditioner included, in itsJ_rmemo, so RSS grew with each field (a W7-X 15x45x80x6 FP run passed 14 GB after 5 fields). The memo now keeps the warm start only, and the single reuse preconditioner is held only whenprecond_window > 0. -
CI: the coverage shards restore a per-shard JAX compilation cache (written by pushes to main, read by every run, keyed on jax version and
pyproject.toml) and the matrix drops from 13 to 9 shards. Required check names are unchanged. -
Fix: host RSS grew over an ambipolar profile because executables were compiled and kept per field and per surface. The coupled preconditioner now applies through one module-level jitted sweep (it ran an eager
lax.scanover fresh closures on every application), anddkx runclears JAX compile caches between surfaces. W7-X, 5 surfaces, 13x37x40x5 FP, 2 species, 6 cores of the shared office host: peak RSS 13.5 to 6.9 GB, 5/5 surfaces converged both times, total 1450 to 1611 s at load 9 to 22. -
dkx/cli.py(2,481 lines) is split by owner:cli.pykeeps the SFINCS-compat commands andmain(1,096 lines),cli_user.pythe user commands,cli_parsers.pythe argparse builders,cli_front.pythe argv front door.dkx.clire-exports every name tests and entry points use; no behaviour change. -
dkx/representative.py(2,164 lines) loses its figure panels and plotting entry points torepresentative_plots.py(649 lines;representative.pyis 1,559).dkx.representativere-exports them; no behaviour change. -
dkx/workflows/optimization.py(2,124 lines) is split by owner:optimization.pykeeps objectives and scan promotion (1,078 lines),optimization_evidence.py,optimization_comparison.pyandoptimization_ladder.pyhold the evidence plans, payload comparison and resolution ladder.dkx.workflows.optimizationre-exports every public name; no behaviour change. -
Tests: four small modules folded into their natural owners (four fewer files, same 89 tests collected):
test_plot_pdf_paging_*intotest_plotting,test_audit_suite_output_keysintotest_audit_suite(wastest_audit_suite_runtime_drift),test_repo_size_policyintotest_release_size_contracts,test_cached_grids_readonlyintotest_phase_space. -
dkx/multigrid.py(1,901 lines) is split by owner:multigrid_levels.py(settings, simplified operator, transfers, coarsening),multigrid_smoothers.py(level matvec, coarse solve, line and plane smoothers),multigrid_diagnostics.py(pitch-basis measurements);multigrid.pykeeps the docstring and the preconditioner builders (586 lines).dkx.multigridre-exports every name; no behaviour change. -
dkx/phase_space.py(1,877 lines) is split by owner:phase_space_angles.py(uniform and widened-upwind angle grids),phase_space_speed.py(speed grid and its derivative matrices);phase_space.pykeeps the Legendre pitch machinery,Gridsandmake_grids(668 lines).dkx.phase_spacere-exports every name; no behaviour change. -
Monoenergetic solves stay lean. When the truncated block-Thomas solve misses
tolby rounding (low collisionality),autonow applies one defect correction before falling back to dense-preconditioned Krylov. W7-X at ν*=3e-4, grid 23×55×140, office Xeon on 8 pinned cores, NTX's cross-code harness: peak RSS 23.9 GB (recorded run) to 1.26 GB, warm 101 to 85 s. Alternating A/B/A/B at 19×41×100 with a 1 GB structured budget: 6.1–7.4 to 1.09 GB, warm 47–49 to 19–21 s. D11, D31, D33 unchanged to 1e-9. -
Validation: production-resolution DKX/SFINCS cells for W7-X ambipolar roots at three radii, precise QA and QH, a Boozer
.bcsurface, the trajectory-model Er sweep and RHSMode 2/3 (validation/cross_code_gaps_v1.json); fluxes agree within 2.7e-4. -
Fix: the memoized speed grids (
make_x_grid,make_speed_grid) and parsed.bcsurfaces returned shared, writable arrays, so one caller's in-place edit changed every later result for the same arguments. Their arrays are now read-only. -
Fix:
KineticBootstrapCurrentreadaHatwith netCDF4 only and fell back toaHat = 1on any error. A LIBSTELL textwout(whichgeometryScheme = 5accepts) or a missingAminor_ptherefore scaled everyd/drHatgradient in the deck byaHatwithout warning. It now reads throughread_vmec_woutand raises whenAminor_pis missing or not positive. -
The persistent compilation cache keys on the XLA:CPU thread count: entries live
underthreads<NPROC>/, so an executable compiled under one--cores
setting is not reloaded under another. -
Smaller clone: committed PNGs are palette-quantized and recompressed (1.91 MB to 1.15 MB, same pixels to the eye, same generators), 15 evidence JSONs that nothing reads are deleted, and
tools/parity/output_parity_vs_fortran_fixture.pyis gone (the CLI driver CI runs does the same comparison). -
The namelist-to-operator builder moved from
dkx/drift_kinetic.py(2,697 to 1,877 lines) todkx/operator_build.py.dkx.drift_kineticstill exportskinetic_operator_from_namelist,kinetic_operator_build_from_namelistandKineticOperatorBuild; no behaviour change. -
The namelist-to-operator builder moved from
dkx/drift_kinetic.py(2,697 to 1,890 lines) todkx/operator_build.py.dkx.drift_kineticstill exportskinetic_operator_from_namelist,kinetic_operator_build_from_namelistandKineticOperatorBuild; no behaviour change. -
dkx/magnetic_geometry.py(2,373 lines) keeps geometry construction (1,471 lines); the Boozer.bcreader moved toboozer_bc.pyand the VMECwoutreader tovmec_wout.py.dkx.magnetic_geometryre-exports both; no behaviour change. -
Ten
tests/test_output_h5_*_parity.pymodules (1,232 lines) are one parametrizedtests/test_output_h5_parity.py(254 lines): 34 deck cases over the writer and canonical routes, every compared key and tolerance kept (checked key-for-key against the old lists). -
dkx/solve.py(3,641 lines) is split by route:solve.pykeeps routing andsolve()(986 lines),solve_structured.pythe structured direct route,solve_krylov.pyrecycled Krylov,solve_direct.pysparse direct,solve_common.pythe result types and residual guards.dkx.solvere-exports its public names; no behaviour change. -
dkx/solve.py(3,641 lines) is split by route:solve.pykeeps routing andsolve()(965 lines),solve_structured.pythe structured direct route,solve_krylov.pyrecycled Krylov,solve_direct.pysparse direct,solve_common.pythe result types and residual guards.dkx.solvere-exports its public names; no behaviour change. -
dkx/solve.py(3,641 lines) is split by route:solve.pykeeps routing andsolve()(966 lines),solve_structured.pythe structured direct route,solve_krylov.pyrecycled Krylov,solve_direct.pysparse direct,solve_common.pythe result types and residual guards.dkx.solvere-exports its public names; no behaviour change. -
examples/is a numbered tutorial ladder,examples/tutorials/01..15, from a first native run to a shape-optimization loop. Each script reads top to bottom: imports, input parameters, run with progress printed, then plot, save and print a summary.DKX_EXAMPLES_CI=1runs any of them in seconds. The heavy VMEX scripts moved toexamples/advanced/; the old numbered folders,getting_started/,autodiff/,transport/,vmex_finite_beta/and the notebooks are gone. -
The docs have a four-part tutorial track (
docs/tutorials/) that follows the scripts: first steps, building a case, transport physics, and bootstrap current, gradients and optimization. Each part gives the equations, the command and the output it printed. The README links the new example paths and the quick start runs as written. -
docs/design_decisions.mdgives the reasons for the velocity discretization (Legendre modes, polynomial speed grid) and for the native defaults (Fokker–Planck collisions, DKES trajectories). -
docs/design_decisions.mdcompares DKX's solver with DKES, MONKES, SFINCS, yancc, NEO-2 and KNOSOS and gives the route per run type; plan section 17 records the decisions and the GPU kernel measurement. -
radial_profilesbrackets only fields whose solve passes the original-equation check (algebraic_converged, with the full state retained) and lists the rest inuncertified_er. Before, a failed field entered the bracketing. Each surface also samples its neighbour's root (er_seed), extension probes merge into the scan instead of a second full scan, and a failure anywhere on a surface, including root refinement, is recorded for that surface rather than ending the profile. The two silentexcept Exception: breakfallbacks are gone. Records carry theirresolution. -
find_ambipolar_erno longer stalls on a stale preconditioner. Brent's first two fields are the bracket ends, and reusing the preconditioner built at -12 kV/m at +12 kV/m held a W7-X FP solve (r/a = 0.5, 15x31x32x6) past 2,400 s. Preconditioners are now rebuilt at every field by default and each field starts from the nearest solved state;precond_windowopts into reuse near the build field, as bounded GMRES (about 200 iterations, underautotoo) with a fresh-build retry. Seedocs/numerics/factor_reuse.md. -
find_ambipolar_ersolves each field once:J_ris memoized, the cold duplicate solve at the root is gone (every accepted state already passes the original-equation check), andall_rootsreuses Brent's fields and reports the Brent root itself for its own sign change, so the enumerated and selected roots agree exactly.all_rootsis now opt-in.wall_time_sandn_solvesreport the whole call. -
current_tolis relative to the largest|J_r|seen while bracketing (current_scale), default1e-6; it was an absolute1e-10in normalized units.dkx sfincs ambipolar --current-toleranceandrun_ambipolar_brentfollow. -
A field whose solves all fail is retried cold, then under
autowith a 100x tighter kinetic tolerance; if that fails too the result isstatus="solve_failure"with the attempts infailures, instead of an exception.