Code Structure
Purpose
This page documents where the main physics, numerics, runtime, and artifact surfaces live in the source tree. It is meant to make refactoring safer by keeping the boundary between public APIs and internal implementation modules explicit.
The long-term consolidation target is documented in
Architecture Refactor Plan. New refactor work should reduce package
surface area by merging, deleting, or moving code out of the installable
package. Add a new module only when it replaces multiple existing files or
isolates a measured JAX compilation/performance boundary. Domain packages such
as operators, solvers, objectives, parallel, diagnostics,
workflows, and artifacts are preferred over more root-level
runtime_*, nonlinear_*, vmec_*, quasilinear_*, or
benchmark_* modules. Validation and campaign implementation should not live
in the installable package; reusable metrics belong in diagnostics and
long-run orchestration belongs in tools or root benchmarks.
Repository Roles
The repository uses four non-source-code areas with distinct jobs:
examples/contains small, copyable user workflows.benchmarks/contains reproducible benchmark drivers and compact benchmark inputs that researchers can run directly.tools/contains maintainer commands for artifacts, campaigns, comparisons, profiling, and releases.tests/contains automated gates that CI can run without relying on raw generated outputs. Consolidated test owners keep shared imports once at the file boundary; section markers describe the retained physics families but do not recreate former files through repeated imports or module aliases.
Files that do not fit one of those roles should be deleted from main or
moved to a draft experiment branch. Benchmark or comparison references to other
codes are allowed only in explicit benchmark/comparison contexts.
Artifact builders under tools/artifacts should be organized by validation
family rather than by one output file per script. For example, the VMEC/Boozer
aggregate holdout evidence now uses one
build_vmec_boozer_aggregate_holdout_gate.py command with alpha and
surface subcommands while preserving separate documented artifact names.
Public API vs Internal Modules
Public surfaces that examples, scripts, and external users are expected to rely on:
documented module pages in API Reference
spectraxgk.artifactsspectraxgk.geometryspectraxgk.clispectraxgk.runtimespectraxgk.workflows.runtime.configspectraxgk.workflows.runtime.artifactsspectraxgk.artifacts.plottingspectraxgk.parallelspectraxgk.operators.nonlinear.paralleldocumented benchmark/example scripts under
benchmarks/andexamples/documented repository-maintenance entry points under purpose-specific
tools/subfolders
Internal modules that are free to move as long as the public behavior and tests remain unchanged:
spectraxgk.terms.*spectraxgk.workflows.runtime.startupspectraxgk.workflows.runtime.diagnosticsspectraxgk.workflows.runtime.diagnostic_arraysspectraxgk.workflows.runtime.initial_phispectraxgk.workflows.runtime.chunksspectraxgk.workflows.runtime.resultsspectraxgk.workflows.runtime.commandslow-level geometry adapters and import bridges inside
spectraxgk.geometry
Large refactor status for this push: the split runtime, diagnostics, validation-gate, zonal-validation, and parallelization-policy modules are documented as behavior-preserving infrastructure. They should not be cited as new physics validation, production nonlinear optimization, or broad performance claims unless the corresponding artifact gate in Release Scope and Claim Boundaries promotes that claim separately.
Runtime Flow
The executable-facing runtime path is split conceptually into four layers:
configuration and startup -
workflows/runtime/config.py-workflows/runtime/startup.py-workflows/runtime/policies.pysolver execution -
runtime.py-linear.py-nonlinear.py-solvers/time/explicit.py-solvers/time/explicit_steps.py-solvers/time/explicit_cfl.py-solvers/time/diffrax_*owner modules exported bysolvers/time/__init__.py-solvers/time/runners.pydiagnostics and artifacts -
diagnostics/core.py-diagnostics/energy.py-diagnostics/transport.py-diagnostics/resolved.py-diagnostics/modes.py-diagnostics/growth_rates.py-diagnostics/growth_fit.py-diagnostics/growth_windows.py-workflows/runtime/diagnostics.py-workflows/runtime/diagnostic_arrays.py-workflows/runtime/results.py-workflows/runtime/orchestration_scan.pyandworkflows/runtime/orchestration_artifacts.py-workflows/runtime/artifacts.py-artifacts/-artifacts/plotting.pyexecutable workflows -
workflows/runtime/commands.py-workflows/linear.py-workflows/nonlinear.py-workflows/cases.py-workflows/demo.py-cli.pybenchmark and validation tooling -
spectraxgk.benchmarks- rootbenchmarks/drivers - purpose-specifictools/commands -tests/validationphysics and benchmark gates
Physics / Numerics / IO Map
Responsibility |
Primary files |
Typical tests |
|---|---|---|
Basis and spectral grids |
|
orthonormality, indexing, symmetry |
Geometry and imported equilibria |
|
parser, remap, normalization, geometry-response tests, Miller/VMEC finite-difference geometry and NetCDF writeout gates |
Linear operators and fields |
|
manufactured solutions, observed-order, eigenfunction and branch tests; cache-builder tests cover staged grid, geometry, twist-shift, gyro/moment, drift, and linked-boundary packing |
Solver objectives and eigen-AD gates |
top-level |
core linear/quasilinear observables, implicit eigenpair VJP, branch-locality, sampling-axis, shared solver-ready linear context for gradient gates, solver-ready and VMEC/Boozer gradient gates, VMEC/Boozer context/feature extraction, reduced nonlinear-window metrics, stellarator ITG table/residual gates, reduced QA low-turbulence model/residual/artifact gates, VMEC transport objective adapters, staged VMEC growth-branch locality setup/sample evaluation/report packing, VMEC-state coefficient helpers, and finite-difference line-search tests |
Nonlinear operators |
|
RHS routing, bracket payload, explicit stepping, explicit diagnostic orchestration, IMEX diagnostic orchestration, cached IMEX operator/state policy, diagnostic tuple assembly, fixed-mode and Hermitian projection, collision-split, staged electrostatic/electromagnetic nonlinear contribution helpers, transport-window tests |
Parallelization policy and helpers |
|
identity gates, one-device fallback, velocity-space plan/exchange/streaming/field-reduction microkernels, domain/spectral/strategy contracts, spectral state/layout/work-model/bracket/tolerance helpers, logical spectral reports/RHS/integrator gates, device-z routing gates with explicit RHS and transport-trace/report policies, diagnostic-only nonlinear sharding policy |
Runtime/executable behavior |
|
runtime contract, startup/restart, output-path, full-GK linear/nonlinear workflows, runtime TOML case dependency defaults, saved-output plot command routing, executable artifact path display and progress/summary printing, linear-fit diagnostics, electrostatic-potential initializers, quasilinear finalization, diagnostic-array validation/composition, chunking, result assembly, runtime command workflows, executable smoke tests |
Public import registry |
|
compact lazy public export registry, top-level |
Diagnostic extraction and growth-rate fitting |
|
mode selection, eigenfunction extraction, least-squares growth/frequency fitting, automatic fit-window selection, quasilinear transport weights and saturation helpers, late-time growth/frequency tests |
Artifacts and plots |
|
serialization, reload, restart append schema, dealiased-axis contracts, runtime-output plots, benchmark/scan panels, diagnostic/eigenfunction figures, zonal-response figures, plotting contract tests |
Benchmark harness |
|
late-time/windowed gate tests, eigenfunction reference/phase utilities, diagnostics time-series loading, benchmark case presets, physics metric extraction, scan/eigenmode orchestration, reference loading, fallback policy tests |
Refactor Mapping
The current modularization branch is preserving the public runtime surface while
extracting internal responsibilities out of runtime.py and other large
modules.
Completed extractions:
compact public API registry in
spectraxgk.api. The rootspectraxgkpackage andspectraxgk.apiregistry are lazy facades that derive the stable documented__all__order from one authoritative lazy target mapping while exporting directly from the owning configuration, geometry, solver, validation, parallelization, objective, and plotting modules. Pure-contract imports such asspectraxgk.parallel.decompositionmust remain dependency-light and must not import NumPy/JAX-heavy solver stacks through package initializers.zero-shear boundary promotion, analytic s-alpha/slab geometry models, and sampled/imported flux-tube geometry data/loading:
geometry/analytic.pyandgeometry/flux_tube.py. Imported NetCDF/eik loading keeps schema selection, scalar/profile reads, root-level terminal-theta inference, mirror-term reconstruction, drift/Jacobian normalization, andFluxTubeGeometryDatapacking as separate private stages so geometry-file variants can be tested without one large loader body.focused imported-geometry owners.
geometry.imported_millerowns the complete Miller imported-geometry pipeline.geometry.imported_vmecis a compact VMEC/Boozer-to-EIK orchestrator; backend loading, field-line sampling, metric derivatives, and state construction live respectively ingeometry.backend_discovery,geometry.vmec_field_line_sampling,geometry.vmec_boozer_derivatives, andgeometry.vmec_state_controls. Shared finite-difference and period-extension kernels live ingeometry.kernels. Imported Miller profile assembly keeps central-surface normalization, period extension, Bishop coefficients, metric coefficients, magnetic drifts, target-grid interpolation, ballooning conversion, and final EIK profile packing as explicit stages insidegeometry.imported_miller. Ownership tests import the VMEC orchestrator once and compare only its genuinely shared backend and field-line contracts against those focused owners; former section names are not represented by duplicate aliases of the same module.mode selection/eigenfunction extraction and late-time growth/frequency fitting:
diagnostics/modes.py,diagnostics/growth_rates.py,diagnostics/growth_fit.py, anddiagnostics/growth_windows.py. The publicdiagnostics.analysisanddiagnostics.growth_ratesmodules remain small facades over focused diagnostic owners. Fit-window selection keeps argument validation, least-squares scoring, amplitude/slope thresholds, candidate-window search, and fallback policy as named stages insidediagnostics/growth_windows.pyso benchmark auto-windowing can be tested without duplicating fit logic.scalar energy, species transport/heating, and resolved spectral diagnostics:
diagnostics/energy.py,diagnostics/transport.py, anddiagnostics/resolved.py. The publicdiagnostics.coremodule remains a small facade re-exported byspectraxgk.diagnostics.explicit linear step kernels, diagnostics-rich linear IVP integration, explicit CFL/frequency-bound policy, and progress formatting:
solvers/time/explicit_steps.py,solvers/time/explicit_diagnostics.py, andsolvers/time/explicit_cfl.py. The publicsolvers.time.explicitmodule remains the import facade for existing debug tools and tests. Its linear IVP facade now keeps method validation, adaptive CFL timing, JIT stepper construction, sample-history collection, progress emission, and array packaging as named private stages. The diagnostics owner separately stages method/time-policy validation, JIT stepper construction, energy/transport sampling, progress rendering, andSimulationDiagnosticsconstruction so saved explicit-time benchmark paths exercise named numerical pieces instead of one monolithic loop.Diffrax time-integration internals.
solvers/time/__init__.pyremains the package-level public facade while optional dependency/policy helpers, linear save-path integration, streaming growth/frequency fits, and nonlinear integration live insolvers/time/diffrax_core.py,diffrax_linear.py,diffrax_streaming.py, anddiffrax_nonlinear.py. The streaming owner stages state/cache preparation, monitored-mode extraction, optimized density-mode extraction, weighted log-derivative accumulation, Diffrax RHS construction, IMEX zero-term routing, and saved-result finalization so differentiable streaming fits can be audited without changing the publicintegrate_linear_diffrax_streamingcontract. The saved-field linear owner stages state/cache preparation, packed-state sharding, RHS construction, save-field/mode extraction, save-time policy, Diffrax solve execution, and final unpacking while keeping the publicintegrate_linear_diffraxcontract stable. The fixed-step cached linear integrator stages method validation, state/damping preparation, serial/parallel RHS routing, RK/IMEX/SSPX3 stepping, progress callbacks, and sample-stride scans separately while keeping the publicintegrate_linearand donated-buffer wrappers stable. The nonlinear explicit solver owns fixed-step RK/SSP/K10 scan policy, progress callbacks, state projection, and diagnostic explicit stepping insolvers/nonlinear/explicit.py. The nonlinear owner stages state/cache preparation, packed-state sharding, linear and nonlinear RHS construction, IMEX term routing, saved-phiextraction, solve execution, and finalFieldStatepacking while keeping the publicintegrate_nonlinear_diffraxcontract stable.term-wise RHS assembly internals.
terms/assembly.pyowns public RHS assembly, cached RHS composition, per-term diagnostic decomposition, field-only solves, external-field sources, electrostatic-field policy, and collision-skip policy. The production RHS and diagnostic decomposition share one staging layer for state/species normalization, scalar parameter expansion, field and Hamiltonian construction, drift/drive/dissipation contribution assembly, fixed-order term summation, and species-axis restoration. This keepsassemble_rhs_cachedandassemble_rhs_terms_cachednumerically aligned without duplicating RHS policy. Linear electromagnetic field solves interms/fields.pyare staged as coefficient casting, gyrocenter moments, electrostatic potential, compressionalbparcoupling, parallelaparsolve, and finalFieldStatepacking, while the public custom-VJP wrapper keeps the same differentiability boundary. Linear Streaming and diamagnetic contribution kernels keep Hermite-mode drive insertion centralized interms/linear_terms.py. Collisional, hypercollisional, hyperdiffusive, and end-damping kernels have one physical owner inoperators/linear/dissipation.py; assembly imports that owner directly rather than retaining a second implementation underterms.nonlinear public-driver internals.
nonlinear.pyremains the public facade while cached RHS/state integration lives insolvers/nonlinear/state_integration.pyand explicit/IMEX diagnostic entry points live insolvers/nonlinear/diagnostic_integration.py. Lower-level nonlinear RHS, diagnostic-state, policy, explicit-step, explicit and IMEX diagnostic scan preparation/finalization, and IMEX mechanics remain owned byoperators/nonlinear/*andsolvers/nonlinear/*modules. The explicit nonlinear diagnostic implementation keeps the broad public signature as a facade, then packs method, timestep, stride, Fourier-layout, selected-mode, collision, and output-resolution knobs into a private options object before constructing state/policy/closure scan components. This keeps runtime diagnostics patchable while making the numerical stages testable without a monolithic integration body.startup/loading/initial-condition helpers:
workflows/runtime/startup.pyruntime mode-index, nonlinear step-count, external-source, parallel-scan, and nonlinear diagnostics keyword policies:
workflows/runtime/policies.pyruntime linear fit/eigenfunction extraction and quasilinear finalization:
workflows/runtime/diagnostics.pyfinite-value checks plus runtime diagnostic slicing, truncation, striding, and concatenation:
workflows/runtime/diagnostic_arrays.py. These helpers are imported directly from this owner rather than through the runtime fit module.adaptive chunk execution used by runtime and comparison artifacts:
workflows/runtime/chunks.py. Progress/ETA formatting is imported from the canonical runtime progress formatter rather than wrapped in the chunk owner.runtime result containers and nonlinear result assembly:
workflows/runtime/results.pyruntime progress formatting, combined-
kyscan batching, serial/worker scan orchestration, progress formatting, and nonlinear artifact handoff policy:workflows/runtime/orchestration_scan.py,workflows/runtime/chunks.py, andworkflows/runtime/orchestration_artifacts.py. Scan dependency-bundle builders live with the scan owner and read the publicruntime.pyfacade only as a patchable symbol source. Scan routing keeps combined-kyeligibility, independent-worker task generation, ordered result packing, batch initial-state assembly, diagnostic extraction, and fit selection as named stages behind one private scan-options contract. Nonlinear artifact handoff keeps restart input resolution, append-history loading, checkpoint chunk sizing, and diagnostic-history merging as explicit policies in the artifact owner.saved runtime-output plotting command routing:
workflows/runtime/commands.pyplusworkflows/runtime/orchestration_artifacts.py. The publiccli.pyfacade still owns executable parser dispatch and renderer/runtime-command patch seams, while the command workflow ownsspectraxgk --plotargument validation and runtime command dependency construction. Command artifact display, executable headers, and nonlinear summaries live with the runtime artifact orchestration policy so saved-output behavior and restart/checkpoint handoff share one owner.runtime restart-state dispatch: the public
runtime.pyfacade keeps the patchable NetCDF/raw state loaders but shares one shape-keyword payload between both paths, so restart dimensions cannot drift between loader calls.runtime linear time/fit option forwarding:
run_runtime_linear,run_runtime_scan, and the combined-kybatch wrapper share one private option bundle in the publicruntime.pyfacade, keeping method, timestep, sample-stride, fit-window, mode-method, and fit-signal forwarding aligned.runtime TOML case workflows:
workflows/cases.pyowns the stablerun_linear_caseandrun_nonlinear_casesignatures and their default wiring; the package-level API exports them directly from that owner.root benchmark drivers and result pointers:
benchmarks/. This directory is the canonical lightweight benchmark entry point at repository root. It stores drivers, TOML inputs, andbenchmarks/results/manifest.tomlonly; raw solver products stay in scratch directories and promoted benchmark results are displayed from Benchmarks.full-GK executable linear runtime workflow:
workflows/linear.py. Context preparation, time integration, Krylov fallback, linear fitting, and quasilinear finalization are separate private stages so runtime dispatch can stay compact while preserving the publicspectraxgk.runtimefacade and monkeypatch seams.full-GK executable nonlinear runtime workflow:
workflows/nonlinear.pyshared plot style plus runtime-output, benchmark/scan, diagnostic, and zonal-response figure families:
artifacts/runtime_plots.py,artifacts/benchmark_plots.py,artifacts/diagnostic_plots.py, andartifacts/zonal_plots.py. The shared publication style and public imports live inartifacts/plotting.py, which remains the stable facade for examples and user scripts.validation acceptance contracts and JSON-ready gate helpers:
diagnostics/validation_gates.pyowns scalar tolerance evaluation, JSON serialization, and report builders. Metric extraction and its result contracts live indiagnostics/analysis.py; eigenfunction comparison and reference-bundle contracts live indiagnostics/modes.py. One private upper-limit scalar-gate policy is shared for convergence, mismatch, deficit, and branch-jump thresholds so tolerance semantics remain auditable.autodiff validation helpers:
objectives/autodiff_validation.pyowns finite-difference Jacobians, Gauss-Newton covariance diagnostics, dense operator materialization, and isolated eigenbranch sensitivity reports. Eigenbranch AD gates share selector/gap classification, complex-observable realification, unsupported-AD fallback packing, implicit left/right eigenpair tangent solves, split observable chain-rule assembly, and finite-difference comparison stages.benchmark-harness physics metric extraction and scan/mode orchestration:
diagnostics/analysis.pyandworkflows/runtime. Reviewed reference tables and comparison-only policies live inbenchmarking/shared.pybehind the compactspectraxgk.benchmarksfacade. Eigenfunction normalization, phase alignment, comparison metrics, and reference-bundle IO live indiagnostics/modes.py; diagnostic time-series loading lives inartifacts/nonlinear_diagnostics.py; late/leading windows and analytic signals live indiagnostics/growth_windows.py; real-FFT grid inference lives inartifacts/spectral_layout.py. Late-time linear metrics, nonlinear transport windows, heat-flux convergence, observed-order checks, and branch-continuity metrics live indiagnostics/analysis.py. Zonal-flow residual/GAM metric extraction lives indiagnostics/zonal_validation.py. Zonal-response metrics are staged as trace coercion, initial-level normalization, tail residual extraction, extrema/envelope detection, damping/frequency fits, and finalZonalFlowResponseMetricspacking, keeping Rosenbluth-Hinton/GAM conventions explicit.zonal-response reference/trace normalization helpers:
diagnostics/zonal_validation.pydominant-eigenvalue custom VJP and branch-locality diagnostics:
objectives/eigen.pycore solver-objective constants plus value-level linear/quasilinear observables:
objectives/core.pysolver-objective sampling axes, physical-
kygrid mapping, and aggregate weights:objectives/sampling.pysolver-ready geometry objective gates, reduced nonlinear-window metrics, solver-ready gradient gates, mode-21 VMEC/Boozer gradient gates, backend-free portfolio row/weight contracts, portfolio AD/FD sensitivity gates, artifact promotion guards, geometry-owned VMEC/Boozer state coefficient helpers, VMEC/Boozer objective-table plumbing, and VMEC/Boozer finite-difference/ line-search gates:
objectives/geometry.py,objectives/gradient_gates.py,objectives/vmec_boozer_gradients.py,objectives/portfolio_contracts.py,objectives/portfolio_sensitivity.py,objectives/portfolio_artifacts.py,objectives/vmec_boozer.py,objectives/vmec_boozer_fd.py,objectives/vmec_boozer_line_search.py. Scalar and aggregate VMEC/Boozer finite-difference reports share one settings validator, one VMEC-state coefficient context, one perturbation helper, and one three-point response/curvature diagnostic so scalar and aggregate sensitivity gates use the same finite-value, response-resolution, and curvature policy. Scalar and aggregate reports now keep aggregate surface/field-line/kyplanning, three-point VMEC-state evaluation, weighted-sample metadata, and payload assembly in separate helpers so optimizer-facing sensitivity reports remain auditable without changing their public JSON schema. Mode-21 VMEC/Boozer gradient reports now share context construction, observable-vector assembly, implicit sensitivity-gate execution, and common payload metadata, while the public report functions remain the physics-facing entry points for linear, quasilinear, and reduced nonlinear-window differentiability claims. Scalar and aggregate VMEC/Boozer line-search reports share one private curvature-gated one-parameter search loop plus focused scalar/aggregate probe builders, common payload assembly, and explicit held-out training/probe helpers, keeping finite-difference, training, and held-out aggregate gates on the same accept/reject policy. Reduced QA low-turbulence comparison artifacts keep optimized-design generation, per-design diagnostic payloads, gate booleans, comparison metrics, and differentiability-plumbing metadata in separate helpers so publication JSON structure and optimization diagnostics can evolve independently. Solver-ready gradient gates share one normalized heat/particle transport-weight helper for eigenmode observables so linear-growth, quasilinear, and particle-flux AD checks use the same quadrature and normalization path. Solver-ready branch-continuity and linear-RHS gradient reports now stage parameter validation, objective construction, eigensystem branch checks, implicit AD/FD gate rows, value-evaluator checks, and report packing separately. The top-levelspectraxgk.objectivesAPI re-exports the portfolio helpers directly from these owner modules.production nonlinear turbulent-flux optimization guardrails now live in
diagnostics/nonlinear_transport_optimization.py. The diagnostics owner stages optimization-scope normalization, reduced-artifact scope checks, replicated long-window transport extraction, matched baseline-to-optimized audit gates, safety gates, promotion gates, evidence-gap accounting, and summary assembly so release safety and production-claim promotion cannot be conflated. Persisted pass flags and named gate rows require explicit Boolean values, and replicate counts require finite nonnegative integers; malformed evidence therefore produces blockers rather than promotion or exceptions. Explicit finite comparison metrics take precedence over fallback statistics, including a physical zero; zero uncertainty separation can therefore never be replaced by a stale positive fallback and accidentally promote an audit. Replicate-spread diagnostics now live indiagnostics/nonlinear_replicates.pyand stage ensemble row normalization, high/low variant selection, state classification, replicate-row packing, and summary assembly. The same owner builds seed/timestep artifact-readiness manifests, whilediagnostics/transport_windows.pyowns individual-window statistics and ensemble uncertainty gates. These decisions remain testable without rerunning nonlinear simulations. Follow-up launch planning is not runtime package functionality; it lives intools/campaigns/nonlinear_replicate_followup.py, where report normalization, classification-specific cross-run selection, dedupe/limits, state-plan packing, and config serialization keep GPU follow-up campaigns deterministic and reviewable.quasilinear nonlinear-window convergence metadata is consolidated in
diagnostics/transport_windows.pyfor statistics, CSV/summary IO, promotion readiness, and ensemble uncertainty; replicate readiness belongs todiagnostics/nonlinear_replicates.py. The public API re-exports the documented transport-window helpers directly from these diagnostics owners. Persisted report, gate, and row pass flags must be explicit booleans; strings and numeric lookalikes fail closed. The statistics owner stages validated late-window selection, finite-sample counts, drift/terminal-window metrics, block/bootstrap uncertainty, and gate-report assembly in one file so nonlinear transport admission rules remain auditable. The ensemble owner stages replicate-row normalization, uncertainty statistics, gate packing, artifact grouping, missing-replicate hints, and readiness-manifest packing so seed/timestep promotion evidence can be tested without rerunning simulations.quasilinear model-selection claim boundaries live in
diagnostics/quasilinear_model_selection.py. This owner separates candidate-skill gate rows, absolute-flux overclaim guardrails, optional optimized-equilibrium audit gates, and final ledger assembly, while artifact loading and required-candidate metric normalization.nonlinear parallelization policy metadata, local domain prototypes, and spectral-core work models/RHS primitives plus device-z shard-map routes:
operators/nonlinear/parallel.py,operators/nonlinear/parallel_contracts_domain.py,operators/nonlinear/parallel_contracts_spectral.py,operators/nonlinear/parallel_contracts_strategy.py,operators/nonlinear/domain_decomposition.py,operators/nonlinear/spectral_core.py,operators/nonlinear/spectral_layout.py,operators/nonlinear/spectral_identity_reports.py,operators/nonlinear/spectral_identity_rhs.py,operators/nonlinear/spectral_identity_integrator.py,operators/nonlinear/device_z.pyThe local-domain transport gate stages trace collection, trace-error scoring, fail-closed blockers, and report packing inoperators/nonlinear/domain_decomposition.py. The spectral fixed-window and pencil-transport gates use the same structure inoperators/nonlinear/spectral_identity_integrator.pyso future profiler-backed nonlinear decomposition work can change routing without weakening the numerical-identity policy. The device-z route keeps sharding setup, fail-closed topology checks, shard-map RHS execution, transport-window sampling, compute-only final-state identity, and final report packing as separate stages insideoperators/nonlinear/device_z.pywhile report-schema construction remains inoperators/nonlinear/device_z_reports.py.velocity-space parallelization is split into decomposition metadata (
parallel/velocity_plan.py), Hermite exchange and velocity-field reductions (parallel/velocity_hermite.py), streaming/magnetic-drift microkernels (parallel/velocity_streaming.py), and electrostatic/ diamagnetic field-drive microkernels (parallel/velocity_drive.py). The publicparallel/velocity.pymodule remains the stable facade used by tools and performance gates.independent-work parallelization is split into numerical-identity reports (
parallel/identity.py), JAX-array batch mapping (parallel/batch.py), and ordered Python ensemble execution (parallel/independent.py). The publicparallel/core.pyandparallel/__init__.pymodules remain stable facades for user imports. The package facade uses a static lazy-export registry rather than reading sibling source files at import time. This keeps imports dependency-light and compatible with installed wheels and non-filesystem Python loaders; a unit test requires the registry to match each owning module’s public exports.parallel/independent.pykeeps indexed payload collection, reconstruction contracts, identity report construction, exception provenance, ordered executor routing, and metadata packing as separate helpers so UQ and optimization ensembles stay auditable without changing solver layout.nonlinear RHS composition and state-to-diagnostic tuple assembly:
operators/nonlinear/rhs.pyandoperators/nonlinear/diagnostic_state.py. The diagnostic-state owner separates field defaulting, growth/frequency mode extraction, scalar diagnostics, resolved field/transport group evaluation, and resolved spectra/channel schema packing while preserving the explicit/IMEX scan tuple schema. The old root nonlinear helper shims were removed; normal users should usespectraxgk.nonlinearand developer helpers should import fromspectraxgk.operators.nonlinear.full-GK nonlinear executable orchestration lives in
workflows/nonlinear.pybehind the publicspectraxgk.runtimefacade. The owner separates runtime context construction, fixed-mode/source policy, diagnostic keyword forwarding, adaptive/fixed diagnostic execution, final-state integration, and result assembly so runtime branch tests can exercise each policy without duplicating executable wiring.fixed-step linear integration keeps public dispatch in
solvers/linear/integrators.pywhile diagnostic sampling lives insolvers/linear/integrator_diagnostics.py. The diagnostic owner now keeps sampling validation, cache/state setup, damping assembly, explicit/IMEX/RK step policy, density/Hermite-Laguerre observables, progress callbacks, and every-step versus strided scans in named stages. The facade preserves the publicintegrate_linear_diagnosticsimport and test monkeypatch seams.velocity-parallel linear RHS routing is split into common eligibility/device policy (
solvers/linear/parallel_common.py), Hermite streaming routes (solvers/linear/parallel_streaming.py), and electrostatic slice/fused shard-map routes (solvers/linear/parallel_electrostatic.py). The publicsolvers/linear/parallel.pymodule remains the stable dispatcher. The electrostatic fused route is staged as route validation, sharding specs, closure constants, Hermite exchange, phi solve, streaming/mirror/drift/drive term builders, and JIT cache lookup; the serial single-device path has a separate helper so identity tests can exercise both routes.explicit RK/SSP/K10 one-step policy, cached explicit scan policy, explicit diagnostic step and scan-selection policy, explicit diagnostic integration orchestration, IMEX diagnostic integration orchestration, cached IMEX scan policy, IMEX diagnostic step and scan-execution policy, and IMEX fixed-point/GMRES solve policy:
solvers/nonlinear/explicit.py,solvers/nonlinear/diagnostics.py,solvers/nonlinear/imex_diagnostics.py, andsolvers/nonlinear/imex.py. Developer helpers should import fromspectraxgk.solvers.nonlinear.linear cache, linked-boundary maps, Hermite-Laguerre moments, parameter pytrees, cache-backed RHS assembly, implicit linear GMRES/preconditioner policy, fixed-step/diagnostic integration policy, eigenmode policy/operator/ branch-selection/preconditioner/Krylov algorithms, and velocity-parallel RHS dispatch live under
operators/linear/andsolvers/linear/. The cache path is split intocache_model.pyfor the JAX pytree,cache_arrays.pyfor moment/damping/gyroaverage array factories,cache_builder.pyfor geometry-dependent construction, andcollisions.pyfor collision-matrix assembly, table interpolation, and runtime application. This keeps the collision equations independent of cache allocation without introducing a compatibility facade. The builder itself now has explicit private stages for twist-shift policy, perpendicular wavenumber/drift arrays, Laguerre gyroaverage construction, and linked-boundary metadata so extension work can test one numerical policy at a time;operators.linearremains the package-level public cache facade. The public Krylov import path remainssolvers/linear/krylov.py; that facade now keeps option normalization, user-facing progress messages, shift-invert seed selection, shift-selection flags, and fallback policy as explicit private stages while delegating compiled kernels to the focused owner modules;solvers/linear/implicit.pykeeps implicit state normalization, damping/ drift diagonal assembly, linked Hermite-line solves, coarse kx projection, preconditioner selection, and matrix-free matvec construction as separate private stages; focused developer helpers live ineigen_operator.pyandkrylov_algorithms.py. The old rootlinear_*helper shims were removed; normal users should usespectraxgk.linearfor the public linear API or import focused developer helpers from the domain packages.nonlinear turbulence-gradient follow-up shared configs, JSON parsing, candidate design, composite-control, matched-replicate, QL-seed, state-runbook, and variance-reduction/control-variate report helpers now live in
tools/campaigns/nonlinear_gradient_followup.py. They are campaign planning tools, not runtime package functionality. Variance-reduction, control-mean campaign, and control-mean gate reports share one control-variate candidate parsing/ranking policy and stage paired-label extraction, control-candidate construction, uncertainty propagation, campaign sizing, independent-control pairing, and report packing so noisy follow-up campaign decisions stay deterministic. QL-seed screening, state-control runbooks, matched-replicate follow-up planning, and candidate campaign design are all staged inside that single tools owner so they do not re-enter the installable package as validation-campaign code.nonlinear turbulence-gradient evidence scope markers, acceptance config dataclasses, JSON-safe parsing, finite-difference conditioning gates, and artifact classification live in
diagnostics/metadata.py. Replicated window summaries live indiagnostics/nonlinear_replicates.pyand central finite-difference transport response/uncertainty lives indiagnostics/transport.py. The compactdiagnostics/nonlinear_gradient_evidence.pyfacade owns bracket and candidate report orchestration plus production evidence-gap reports. It has one final export contract;diagnostics/metadata.pyanddiagnostics/transport.pylikewise each have one complete owner contract covering their shared data, transport, and evidence interfaces. Intermediate section exports and alias-based ownership tests were removed so each helper is tested at its real owner. Candidate ranking and evidence-gap reporting remain fail-closed so production nonlinear-gradient promotion cannot be inferred from startup, pilot, reduced, or single-window artifacts.runtime artifact read/write, generic I/O and finite-value validation helpers, linear/quasilinear artifact writers, generic nonlinear table writers, dealiased-axis layout, NetCDF output-bundle orchestration, NetCDF diagnostic-history schema writing with staged Phi2, base species-history, split electromagnetic, zonal-field, resolved species-spectra, and turbulent-heating helpers, NetCDF output geometry, restart-file writing, final-field big-file writing, nonlinear diagnostic reload helpers, and restart-append schema coverage:
workflows/runtime/artifacts.py,artifacts/io.py,artifacts/linear.py,artifacts/nonlinear.py,artifacts/spectral_layout.py,artifacts/nonlinear_netcdf.py,artifacts/nonlinear_netcdf_diagnostics.py,artifacts/nonlinear_netcdf_geometry.py,artifacts/nonlinear_netcdf_fields.py,artifacts/nonlinear_diagnostics.py. The old rootruntime_artifact_*helper modules were removed; import implementation helpers fromspectraxgk.artifactsinstead.
Next planned extractions:
benchmark-family runners and fit-policy helpers
remaining linear result-assembly helpers
runtime output/artifact handoff helpers
Traceability For Refactors
Refactor work is tracked in tools/validation_coverage_manifest.toml. The
manifest is checked by tools/release/check_validation_coverage_manifest.py and
requires every high-priority module to name its source path, owning lane,
reference anchors, physics and numerics contracts, fast tests, artifacts, and
next coverage tests. Update it whenever a source extraction changes module
ownership or validation responsibility.
The manifest also owns the package inventory. Direct rows cover the public or
high-risk refactor modules. Smaller implementation modules are listed in an
owner row’s owned_modules field when their behavior is validated by the
same fast tests and artifacts. Package plumbing such as __init__.py and
version metadata is the only normal exclusion. Adding a new
src/spectraxgk/*.py file without one of those declarations should fail the
manifest checker.
Use this rule of thumb when changing ownership:
add a direct row when a module has public imports, independent physics or numerics contracts, separate artifact traceability, or high refactor risk;
add it to
owned_moduleswhen it is a narrow helper split whose contract is still fully exercised by the owning row’s tests;update the owning row’s
next_testswhen the helper split creates coverage debt that is not closed in the same change.
The refactor manifest tests also cross-check public and size-sensitive
surfaces. Every .. automodule:: spectraxgk.* entry in API Reference must
resolve to a source file or package __init__ that is accounted for by the
manifest, and public package __init__ entries must be explicit exceptions
rather than generic plumbing. The same test requires any non-__init__
source module with at least 2,000 non-comment source lines to have a direct
manifest row, so large modules cannot be hidden under another owner’s
owned_modules list.
The authoritative target package layout, naming policy, and conflict-resolution
rules live in Architecture Refactor Plan. The executable migration ledger
is tools/differentiable_refactor_manifest.toml, checked by
tools/release/check_package_architecture_manifest.py differentiable-refactor. See
Differentiable Refactor Plan for differentiability contracts, extension
points, active manifest rows, and physics/autodiff/parity gates. If this page or
the manifest conflicts with Architecture Refactor Plan, update the
current-tree documentation or manifest rather than adding another root-level
prefix module.
The first behavior-preserving contract modules are spectraxgk.core.contracts
and spectraxgk.core.extension_points. They introduce typed refactor,
validation-gate, differentiability, and extension-point protocols without
moving solver kernels or changing public numerical behavior.
Runtime command dispatch now keeps parser construction in spectraxgk.cli,
with parser registration split by command family, and moves runtime linear,
runtime scan, and runtime nonlinear command execution into
spectraxgk.workflows.runtime.commands. The generic run executable path
attaches the already-loaded runtime config/data to the parser namespace before
dispatch, so command execution does not parse the same TOML twice.
Linear, scan, and nonlinear command flags are resolved once into typed
command-option records before solver calls are made, which keeps precedence
between CLI flags, TOML sections, and runtime defaults inspectable without
spreading executable policy across command bodies. Saved-artifact display order
and optional artifact writing for linear, scan, quasilinear, and nonlinear
commands also live in focused command-output helpers so user-facing executable
messages can be tested without launching solver runs.
spectraxgk.workflows.cases remains the TOML case-workflow owner and
re-exports the command helpers only for public executable dispatch; path override, progress,
quasilinear override, and preload-reuse policies have one canonical workflow
owner.
Runtime initial-condition construction is staged inside
spectraxgk.workflows.runtime.initial_conditions. Validation of initializer
options, restart-state scaling, kinetic-species targeting, single-mode profile
assembly, random/Gaussian multimode seeding, electrostatic-potential seeding,
Hermitian full-ky completion, and restart merge policy now have named helper
stages. The startup facade still exposes the same _build_initial_condition
entry point for runtime workflows and tests, but the implementation no longer
mixes restart I/O, moment normalization policy, and phi-inversion cache setup in
one long branch.
The benchmark stack has two explicit roles:
spectraxgk.benchmarking.sharedowns compact reference containers, CSV loaders, normalization constants, and comparison-only defaults.Canonical TOML inputs plus
run_runtime_linearandrun_runtime_scanown promoted solver execution. ETG and Cyclone use only this path; artifact figures consume reviewed tables rather than launching a second hidden solve.
spectraxgk.benchmarks is now a small facade for reviewed reference data and
comparison policies. KBM time histories and fixed-beta k_y scans use
generic runtime orchestration. The former kbm_beta_scan.py was removed
after audit showed that it incorrectly interpreted a k_y reference table
as beta values. TEM and kinetic-electron execution also use canonical runtime
inputs and scan paths. New examples must not add named-case execution APIs.
External-code names and raw reference handling stay in root benchmarks/ or
tools/comparison; reusable source modules use physics and numerical names.
Ordered continuation scans use run_runtime_parameter_scan. The caller
supplies a named scalar axis and a pure RuntimeConfig update callback; the
runtime records one RuntimeLinearResult per point and can pass the previous
state into the next solve. Case-specific branch targets and acceptance rules
remain callback policy rather than hidden runtime behavior. Continuation scans
are intentionally sequential because adjacent points depend on one another;
independent parameter ensembles should instead use the parallel batch APIs.
For spectra with competing branches, candidate_options declares the
per-point solver targets and select_candidate chooses the retained result;
only that result is continued.
Time-domain linear results retain the selected-k_y field history.
refit_runtime_linear_trajectory applies another fit window or mode
extractor to that stored trajectory without repeating the simulation. This is
the preferred path for branch audits: integration cost is paid once, while
z_index, max, project, and svd diagnostics remain independently
reviewable.
Quasilinear calibration now lives in
spectraxgk.diagnostics.quasilinear_calibration. It owns calibration-point
schemas, spectrum integration, train/holdout scale fitting, nonlinear-window
CSV/NetCDF ingestion, and report writing behind one diagnostics owner. Each
nonlinear trace is loaded once and reused for both the selected window and its
convergence report; persisted ensemble pass flags must be explicit booleans.
Late-window transport gates live in spectraxgk.diagnostics.transport_windows.
The public validation API re-exports user-facing helpers while
campaign launch and artifact-building policy stays in tools. Model-selection status construction keeps scoped candidate-skill gates,
input normalization, optimized-equilibrium audit summaries, and absolute-flux
claim guardrails in diagnostics/quasilinear_model_selection.py.
VMEC-JAX candidate and transport admission gates now have explicit owners.
spectraxgk.objectives.vmec_candidate_admission owns solved-equilibrium,
authoritative-WOUT, and WOUT-reproducibility candidate gates. It keeps aspect,
iota, iota-profile, quasisymmetry, and pass/fail helpers together so optimizer
state and WOUT gates share one JSON schema and threshold semantics.
spectraxgk.objectives.vmec_transport_admission owns transport-admission
policy dataclasses, reduced transport metric selection, multi-surface/
field-line/k_y sample coverage, and promoted transport-candidate selection.
The public spectraxgk.objectives.vmec_transport module contains only the
optimizer callback and supported objective exports. Optional-backend path
policy lives in vmec_transport_config, differentiable sample-table and
reduction kernels live in vmec_transport_tables, and eigenbranch-locality
gates live in vmec_transport_branch. Tests and developer extensions inject
dependencies at those owners; the public facade does not synchronize or mutate
implementation-module globals at runtime.
spectraxgk.diagnostics.stellarator_transport_reports owns report-style
nonlinear transport diagnostics: landscape admission, reduced prelaunch gates,
next-campaign admission, and matched nonlinear audit redesign. Persisted gate
flags must be explicit booleans and replicate counts must be finite,
nonnegative integers; malformed values fail closed into report blockers rather
than becoming truthy or raising during report construction. The public
spectraxgk.api re-exports user-facing admission helpers directly
from these owners, while installable validation-campaign subpackages have
been removed.
The first differentiable-geometry split keeps
spectraxgk.geometry.differentiable as the public facade while
moving optional backend lookup and strict AD/finite-difference gate utilities
into spectraxgk.geometry.backend_discovery and
spectraxgk.geometry.autodiff_checks. The solver-ready in-memory flux-tube
mapping and geometry-observable contract lives in
spectraxgk.geometry.flux_tube_contract. Pure parity metrics,
interpolation, radial derivative, Boozer half-mesh, Fourier field-line, and
periodic sampling helpers live in spectraxgk.geometry.numerics. This
separates import-side effects, validation-report plumbing, public contract
validation, and small numerical kernels from the VMEC/Boozer field-line bridge.
The AD/FD validation owner stages parameter validation, observable flattening,
Jacobian construction, tangent checks, conditioning gates, failure reasons, and
strict JSON report assembly so differentiability tests can target each
research-grade gate directly.
Imported VMEC/Boozer geometry generation enters through
spectraxgk.geometry.imported_vmec. Its focused owners separate optional
backend discovery, radial spline and Boozer-mode sampling, Hegna-Nakajima and
metric derivatives, and VMEC field-line state construction. The orchestrator
retains only flux-tube cutting, equal-arc remapping, atomic EIK output, and the
high-level request path. This separation keeps the formulas discoverable and
lets backend, sampling, derivative, and orchestration tests fail independently.
Zero-shear boundary policy and analytic s-alpha/slab geometry models live in
spectraxgk.geometry.analytic. Sampled solver-ready geometry data, analytic
sampling, imported-NetCDF loading, and periodic mirror-term reconstruction live
in spectraxgk.geometry.flux_tube. Twist-shift and grid-default geometry
policy live in spectraxgk.geometry.core. The
spectraxgk.geometry package remains a thin public facade that
re-exports the same classes and functions for existing user code.
Geometry sensitivity, inverse-design, and local UQ reports live in
spectraxgk.geometry.sensitivity so backend bridge modules can depend on
the report contract without importing the public facade. Bounded VMEC
boundary and Boozer-spectrum bridge checks, Boozer |B| field-line
evaluation, and Boozer-to-flux-tube sensitivity diagnostics live in
spectraxgk.geometry.booz_xform_bridge. Pure helper imports retain object
identity; backend-discovery-dependent bridge functions use thin facade wrappers
so existing monkeypatch-based optional-backend tests still target
spectraxgk.geometry.differentiable. VMEC-state-to-Boozer, VMEC metric
tensor, and VMEC field-line tensor AD/FD sensitivity reports live in
spectraxgk.geometry.vmec_state_sensitivity; the public facade uses the
same hook-preserving wrapper pattern, while the implementation owns shared
VMEC example loading, coefficient-index validation, perturbation policy, and
the common tensor-observable AD/finite-difference payload builder used by
metric and field-line gates. Metric and field-line tensor reports also share
one VMEC geometry context/index helper before diverging into their
observable-specific sampling paths.
Boozer constant preparation and equal-arc cache prewarm helpers live in
spectraxgk.geometry.vmec_boozer_constants. Core Boozer equal-arc profile
construction, radial Boozer-profile interpolation, and equal-arc remapping
live in spectraxgk.geometry.vmec_boozer_core. That owner now stages the
bridge as radial Boozer-profile interpolation, equal-arc field-line
construction, zero-beta metric/drift profile assembly, and final solver-ready
mapping assembly; optional backend execution and Boozer radial-grid validation
remain private helper seams so the public state-to-profile bridge stays
focused on physics profile assembly. Metric/drift profile assembly is further
split into differential-geometry evaluation, raw metric coefficients, raw
curvature-drift coefficients, open equal-arc remapping, and final
_MetricDriftProfiles packing so the tensor algebra and solver-normalized
coefficients can be reviewed independently. The core bridge delegates Boozer
field-line spectral sums, cylindrical derivatives, and coordinate-gradient
algebra to
spectraxgk.geometry.vmec_boozer_derivatives so the differentiable geometry
path has a small, unit-testable tensor-algebra owner rather than one long
VMEC/Boozer orchestration function. The core
profile assembly shares one dtype-aware numerical floor across |B|,
gradpar, Jacobian, metric, curvature, and q denominators so
regularization policy is visible and consistent. Boozer metric-gradient terms
use a separate float32-safe toroidal-flux denominator floor before
grad(theta), grad(phi), and grad(alpha) divisions. Direct
vmec_jax tensor sampling and conversion into the solver-ready flux-tube
mapping contract lives in spectraxgk.geometry.vmec_tensor_mapping. That
bridge is staged as surface/reference-scale validation, shared VMEC field-line
coordinate construction, raw tensor loading, periodic line sampling,
perpendicular metric assembly, local grad-B drift closure, and final mapping
packaging. VMEC flux-tube sensitivity and array-parity report orchestration
lives in
spectraxgk.geometry.vmec_flux_tube_reports; it reuses the shared
VMEC-state example loading, coefficient-index validation, and perturbation
policy from spectraxgk.geometry.vmec_state_sensitivity so the flux-tube,
Boozer, metric-tensor, and field-line AD/FD gates stay on one setup contract.
Direct-array parity, imported-EIK loading, optional Boozer equal-arc parity,
production parity metrics, and final JSON packing now have separate private
stages in that owner. The public report functions remain schema-preserving
facades over those stages, which keeps the VMEC/Boozer differentiability and
parity gates easier to review without adding another module.
VMEC boundary-gradient
probe classification, collection row assembly, and projected-transport
line-search admission summaries live in
spectraxgk.geometry.vmec_boundary_chain. Boundary-chain scalar error
construction and pass/fail policy helpers are kept explicit in that owner so
VMEC/Boozer finite-difference, JVP, and VJP convention gates can be tested
without launching expensive VMEC solves.
Release-scope synchronization for refactors is tracked separately in
Release Scope and Claim Boundaries. In particular, the current restartable NetCDF append
contract normalizes diagnostics loaded from *.out.nc to the persisted
schema before concatenation; transient in-memory traces that are not stored in
the NetCDF artifact are not treated as release data on continuation.
Testing Taxonomy
The source tree should be validated through five distinct test classes:
unit tests - cheap, deterministic, local behavior
numerical verification tests - observed-order, manufactured solutions, invariants
benchmark/validation tests - growth rates, frequencies, eigenfunctions, transport windows
autodiff tests - finite-difference and complex-step gradient checks, tangent/adjoint consistency
regression tests - runtime contracts, saved artifacts, reference-lane preservation
Every future source extraction should update this page if module ownership changes materially.
Repository Artifact Hygiene
The repository should keep reproducible source inputs, tests, small gate reports, and lightweight documentation previews in Git. Heavy runtime outputs, profiler traces, raw NetCDF comparisons, and high-resolution publication exports should live in release artifacts with checksums and replay commands.
Use the non-destructive audit helper before release cleanups:
python tools/release/check_repository_size_manifest.py audit --top 30
python tools/release/check_repository_size_manifest.py
python tools/release/check_repository_size_manifest.py release-artifacts
The report separates tracked file size from ignored local artifact roots such
as tools_out/, docs/_build/, dist/, virtual environments, and caches.
The checked manifest is tools/repository_size_manifest.toml. It defines the
tracked-size budget, the maximum size of any unlisted tracked file, and the
temporary whitelist for any intentionally retained large files.
The release migration manifest is tools/release_artifact_manifest.toml. It
records checksums, replay commands, and planned destinations for high-resolution
panels and other large assets. The checker validates provenance only; it does
not upload or delete artifacts. move_to_release entries may be absent from
Git after migration, but only when the manifest records the immutable
release_tag and release_url together with the original size and SHA-256
checksum.
Documentation figures should use lightweight checked-in previews, with high-resolution publication exports regenerated from the replay commands or hosted as release assets. The reproducible preview-compression command is:
python tools/release/check_repository_size_manifest.py compress-previews --mode release --max-width 2200 --colors 192
python tools/release/check_repository_size_manifest.py compress-previews --mode docs --min-bytes 300000 --max-width 1800 --colors 192
The first command only touches release-manifest previews, so update
tools/release_artifact_manifest.toml with the new sizes and checksums after
running it. The second command skips release-manifest paths by default and is
intended for ordinary checked-in documentation previews. Rerun both manifest
checkers after either cleanup.
History rewrites are not part of routine development; they require a coordinated maintenance window because every collaborator must reclone or reset local branches after a force push.