Skip to content

DKX 2.10.0

Choose a tag to compare

@rogeriojorge rogeriojorge released this 20 Oct 13:13
· 22 commits to main since this release
ee8cc41

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 eager lax.map over 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, and dkx run no longer calls jax.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 of scan-er and scan-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.scan sweeps eagerly over fresh closures, so every solve (and every E_r field of a scan) retraced and compiled one executable per Nxi_for_x group (8 on the tokamak decks). Now _truncated_sweeps is a module-level jit taking the operator as an argument. tokamak_1species_PASCollisions_noEr through dkx 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_budget reads available memory, so on a loaded host (office, 0 GB free) six test_speed_triangle_precond and test_tier2_memory_guard tests failed on a correct tree; with the budget pinned all pass at 0 GB free.

  • dkx sfincs scan-er no longer loses a scan to one bad point. A point that raises is recorded in scan_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 upstream sfincsScan with scanType = 5 (radius × E_r from a profile file and the deck's !ss lines). A malformed !ss or 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 coupled preconditioner, 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. auto now starts with coarse_triangle when 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_triangle rung of the escalation always failed with UnexpectedTracerError. SOLVAX's compile reuse turns arrays closed over by a jitted callback into jit arguments, including those of the nested host callback, so upper and 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 one auto had used.

  • The last escalation rung runs SuperLU above max_dense_size when 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.py gains a resolution ladder. Figures are palette-compressed (docs/figures/publish_tutorial_figures.py, 244 KB for 13) and registered in docs/figure_provenance.json.

  • Fix: find_ambipolar_er kept every field's full solve state, preconditioner included, in its J_r memo, 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 when precond_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.scan over fresh closures on every application), and dkx run clears 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.py keeps the SFINCS-compat commands and main (1,096 lines), cli_user.py the user commands, cli_parsers.py the argparse builders, cli_front.py the argv front door. dkx.cli re-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 to representative_plots.py (649 lines; representative.py is 1,559). dkx.representative re-exports them; no behaviour change.

  • dkx/workflows/optimization.py (2,124 lines) is split by owner: optimization.py keeps objectives and scan promotion (1,078 lines), optimization_evidence.py, optimization_comparison.py and optimization_ladder.py hold the evidence plans, payload comparison and resolution ladder. dkx.workflows.optimization re-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_* into test_plotting, test_audit_suite_output_keys into test_audit_suite (was test_audit_suite_runtime_drift), test_repo_size_policy into test_release_size_contracts, test_cached_grids_readonly into test_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.py keeps the docstring and the preconditioner builders (586 lines). dkx.multigrid re-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.py keeps the Legendre pitch machinery, Grids and make_grids (668 lines). dkx.phase_space re-exports every name; no behaviour change.

  • Monoenergetic solves stay lean. When the truncated block-Thomas solve misses tol by rounding (low collisionality), auto now 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 .bc surface, 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 .bc surfaces 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: KineticBootstrapCurrent read aHat with netCDF4 only and fell back to aHat = 1 on any error. A LIBSTELL text wout (which geometryScheme = 5 accepts) or a missing Aminor_p therefore scaled every d/drHat gradient in the deck by aHat without warning. It now reads through read_vmec_wout and raises when Aminor_p is missing or not positive.

  • The persistent compilation cache keys on the XLA:CPU thread count: entries live
    under threads<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.py is 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) to dkx/operator_build.py. dkx.drift_kinetic still exports kinetic_operator_from_namelist, kinetic_operator_build_from_namelist and KineticOperatorBuild; no behaviour change.

  • The namelist-to-operator builder moved from dkx/drift_kinetic.py (2,697 to 1,890 lines) to dkx/operator_build.py. dkx.drift_kinetic still exports kinetic_operator_from_namelist, kinetic_operator_build_from_namelist and KineticOperatorBuild; no behaviour change.

  • dkx/magnetic_geometry.py (2,373 lines) keeps geometry construction (1,471 lines); the Boozer .bc reader moved to boozer_bc.py and the VMEC wout reader to vmec_wout.py. dkx.magnetic_geometry re-exports both; no behaviour change.

  • Ten tests/test_output_h5_*_parity.py modules (1,232 lines) are one parametrized tests/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.py keeps routing and solve() (986 lines), solve_structured.py the structured direct route, solve_krylov.py recycled Krylov, solve_direct.py sparse direct, solve_common.py the result types and residual guards. dkx.solve re-exports its public names; no behaviour change.

  • dkx/solve.py (3,641 lines) is split by route: solve.py keeps routing and solve() (965 lines), solve_structured.py the structured direct route, solve_krylov.py recycled Krylov, solve_direct.py sparse direct, solve_common.py the result types and residual guards. dkx.solve re-exports its public names; no behaviour change.

  • dkx/solve.py (3,641 lines) is split by route: solve.py keeps routing and solve() (966 lines), solve_structured.py the structured direct route, solve_krylov.py recycled Krylov, solve_direct.py sparse direct, solve_common.py the result types and residual guards. dkx.solve re-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=1 runs any of them in seconds. The heavy VMEX scripts moved to examples/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.md gives the reasons for the velocity discretization (Legendre modes, polynomial speed grid) and for the native defaults (Fokker–Planck collisions, DKES trajectories).

  • docs/design_decisions.md compares 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_profiles brackets only fields whose solve passes the original-equation check (algebraic_converged, with the full state retained) and lists the rest in uncertified_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 silent except Exception: break fallbacks are gone. Records carry their resolution.

  • find_ambipolar_er no 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_window opts into reuse near the build field, as bounded GMRES (about 200 iterations, under auto too) with a fresh-build retry. See docs/numerics/factor_reuse.md.

  • find_ambipolar_er solves each field once: J_r is memoized, the cold duplicate solve at the root is gone (every accepted state already passes the original-equation check), and all_roots reuses Brent's fields and reports the Brent root itself for its own sign change, so the enumerated and selected roots agree exactly. all_roots is now opt-in. wall_time_s and n_solves report the whole call.

  • current_tol is relative to the largest |J_r| seen while bracketing (current_scale), default 1e-6; it was an absolute 1e-10 in normalized units. dkx sfincs ambipolar --current-tolerance and run_ambipolar_brent follow.

  • A field whose solves all fail is retried cold, then under auto with a 100x tighter kinetic tolerance; if that fails too the result is status="solve_failure" with the attempts in failures, instead of an exception.