SPECTRAX-GK
SPECTRAX-GK is a JAX-native gyrokinetic solver using Hermite-Laguerre velocity space, Fourier perpendicular coordinates, and field-aligned flux-tube geometry.
Install with pip install spectraxgk and run spectraxgk in a terminal to
launch the default linear Cyclone demo. The same executable also accepts
checked-in TOMLs directly and can plot saved runtime outputs with
spectraxgk --plot <artifact>.
Documentation map
The documentation is organized so that the core physics, equations, numerical methods, and model-specific paths are easy to locate:
Theory summarizes the gyrokinetic ordering, field equations, and flux-tube assumptions.
Linear Model And Derivations gives the operator-level derivation tied directly to the implemented linear equations and diagnostics.
Operators And Terms lists every implemented term, collision model, hyperdiffusion/hypercollision control, and the runtime parameters that select them.
Numerics documents discretization, time integration, FFT brackets, solver contracts, and JAX parallelization.
Geometry and Input Files and Executable cover the supported model paths (analytic s-alpha, Miller, VMEC/imported geometry, slab) together with the TOML schema used by the executable and Python drivers.
For a first technical read, start with Theory, then move to Operators And Terms and Numerics. That path gives the governing equations, the term-by-term implementation contract, and the numerical approximations in the same order they appear in the code.
Contents
- Quickstart
- Input Files and Executable
- Outputs and Restart Artifacts
- Theory
- Linear Model And Derivations
- Operators And Terms
- Numerics
- Spectral discretization
- Algorithm mapping (numerics → code)
- JAX execution model
- Time integration algorithms
- Diffrax integration
- Nonlinear FFT bracket
- De-aliasing and hyperdiffusion
- Performance tuning
- Automatic solver + fit-signal selection
- Gyroaverage and polarization
- Parallel streaming
- Curvature, grad-B, and mirror couplings
- Field solve and electromagnetic coupling
- Normalization control
- Diamagnetic drive
- Time integration
- Boundary damping
- Dealiasing
- Nonlinear Electromagnetic Terms
- Parallelization policy
- Normalization
- Quasilinear Transport
- Current validated scope
- Literature anchors and claim policy
- Executable usage
- Model details
- Implementation map
- Algorithmic workflow
- Numerics and differentiability
- Validation gates
- Calibration reports
- Train/holdout transfer
- Non-axisymmetric HSX holdout
- W7-X NetCDF nonlinear-window path
- Saturation-rule sweep
- Shape-aware saturation diagnostic
- Candidate uncertainty gate
- Dataset-sufficiency gate
- Model-selection status
- Holdout-gap report
- External-VMEC next-holdout runbook
- VMEC equilibrium portfolio for future holdouts
- Axisymmetric spectrum-shape gates
- Differentiable Stellarator Optimization
- Purpose
- Source Map
- VMEC-JAX-Style QA Transport Scripts
- Full Max-Mode-5 Optimizer Sweeps
- Configurable Solved-Boundary Driver
- Development-Only Reduced Diagnostics
- Aspect-6 QA Low-Turbulence Comparison
- VMEC-JAX Geometry Examples
- Production VMEC-JAX Optimization Plan
- Development Portfolio Gate
- Objective
- Numerics and Differentiation
- Objective-portfolio reducer gate
- Development Diagnostic Results
- Zonal-flow Objective Contract
- Connection to Literature
- Solver-objective Geometry Gradients
- Promotion Gates for Full VMEC/Boozer/GK Optimization
- Geometry
- Algorithms
- Solvers
- Architecture
- Benchmarks
- Examples
- Config-backed runtime cases
- Runtime TOML entry points
- Scaling utilities
- Plotting outputs
- Geometry examples
- Nonlinear restart and continuation
- Geometry helper workflows
- Benchmark and scan helpers
- Foundational demos
- Differentiable optimization examples
- Secondary slab workflow
- Reduced-model runtime
- Full-GK ETG nonlinear pilot
- Related Codes
- Performance
- JAX performance model
- Cache profiling
- Nonlinear profiling
- JAX/XProf operational notes
- Recent nonlinear profiling (Cyclone, benchmark-locked config)
- Nonlinear benchmark harness
- RHS kernel profile (nonlinear Cyclone)
- Benchmark-size Cyclone Miller RHS profile
- Runtime-mode stellarator RHS smoke profile
- Linear RHS term profile
- Full fused linear RHS trace
- Parallelization scaling guardrail
- Fixed-step nonlinear state sharding
- Spectral nonlinear mode (gated fast toggle)
- Runtime and memory comparison workflow
- Final runtime/memory figure
- Interpretation of short nonlinear GPU rows
- Startup phase profiler
- Cached basis indices
- GMRES preconditioner iterations
- JIT considerations
- Planned optimizations
- Linear-to-nonlinear optimization roadmap
- Testing
- Testing philosophy
- Current testing target
- Test categories
- Unit tests (numerical invariants)
- Physics regression tests
- Benchmark regression tests
- Literature-anchored response and spectrum tests
- Diffrax and nonlinear smoke tests
- Parallelization identity gates
- Nonlinear parity snapshots
- Linear physics checks
- Running tests
- Benchmark reproducibility stack
- Stress-matrix parity gates
- CI split: fast PR vs manual full
- Core modular coverage gate
- Core solver coverage gates
- Validation And Coverage Strategy
- Verification Matrix
- Release Scope and Claim Boundaries
- Code Structure
- Manuscript Figures
- References
- Roadmap
- API Reference
- Basis
- Gyroaverage
- Geometry
- Differentiable Geometry
- Grids
- Operators
- Linear
- Linear Linked Boundaries
- Linear Cache
- Linear Moments
- Linear Parameters
- Linear Parallel RHS
- Nonlinear Diagnostics
- Nonlinear Helpers
- Nonlinear Replicate Diagnostics
- Nonlinear Replicate Follow-Up
- Nonlinear Gradient Follow-Up
- Benchmarks
- Benchmark Defaults
- Benchmark Scan Policies
- Benchmarking
- Validation Gates
- Autodiff Validation
- Parallelization
- Velocity Sharding Plans
- Sharded Integrators
- Zonal Validation
- Zonal Flow Objectives
- Analysis
- Plotting
- Config
- Normalization
- Runtime Config
- Runtime Policies
- Runtime Orchestration
- Quasilinear Transport
- Quasilinear Calibration
- Quasilinear Model Selection
- External-VMEC Holdout Planning
- Parallel Decomposition Contracts
- QA Low-Turbulence Optimization
- VMEC-JAX Transport Objective
- VMEC-JAX Transport Admission
- VMEC-JAX Transport Gradient
- VMEC-JAX Boundary Chain
- VMEC-JAX Transport Line Search
- VMEC-JAX Candidate Gates
- Stellarator Objective Portfolios
- Runtime Runner