Differentiable Refactor Plan

Purpose

This page defines the planned research-grade differentiable architecture for SPECTRAX-GK. The goal is not to rewrite the solver for aesthetics. The goal is to make the code easier to use, easier to test, easier to extend, and safer to differentiate through while preserving the validated physics behavior.

Authority note: Architecture Refactor Plan is the authoritative plan for future package layout, file naming, migration order, and conflict resolution. This page is a technical appendix for differentiability contracts, active manifest rows, completed split inventory, and AD/physics/performance gate requirements. If a layout or naming detail here conflicts with Architecture Refactor Plan, the architecture plan wins and this appendix or tools/differentiable_refactor_manifest.toml should be updated.

The active branch for this work starts with planning infrastructure only: tools/differentiable_refactor_manifest.toml and tools/release/check_package_architecture_manifest.py differentiable-refactor. Large code moves should land in later PRs only after the relevant manifest row, public facade, fast tests, parity gates, and documentation have been updated.

Design Principles

  1. Preserve public behavior first. Existing user-facing imports such as spectraxgk.linear, spectraxgk.nonlinear, spectraxgk.runtime, and spectraxgk.geometry.differentiable remain public facades while internals are split.

  2. Keep differentiable paths pure. Objective functions must take explicit PyTree parameters and arrays and return arrays or typed reports. File I/O, plotting, subprocesses, and progress printing stay outside differentiated functions.

  3. Make static and dynamic data explicit. JIT-compiled kernels should separate static model choices from dynamic array values so scans and optimizations do not recompile unnecessarily.

  4. Prefer small kernel modules over broad scripts. Physics terms, basis recurrences, field solves, brackets, integrators, diagnostics, and artifact writing each get their own testable module.

  5. Use custom derivatives only when justified. Default to native JAX autodiff. Add custom_jvp or custom_vjp only for solver, eigenvalue, fixed-point, or adjoint paths that have finite-difference and tangent tests.

  6. Treat validation adapters as boundary code. Reference-code comparison helpers remain in validation or adapter modules. Solver kernels should not import them.

  7. Make extension points deliberate. New collision operators, closure models, field solvers, polynomial bases, geometry providers, and transport models should be registered through narrow protocols instead of edited into monolithic files.

  8. Document physics and numerics at the public boundary. Public functions need docstrings describing equations, normalization, array shapes, differentiability status, static arguments, and relevant tests.

JAX and Scientific-Python Guidance

The architecture should follow current JAX ecosystem practice:

  • JAX pytrees are the natural container model for solver state, geometry, parameters, and objective reports.

  • JAX custom derivative rules should be reserved for transformable functions whose default derivative is unstable, too expensive, or mathematically wrong for the promoted objective.

  • JAX structured control flow makes adaptive controllers a solver-policy decision, not an implementation detail: scan and static-trip fori_loop are the preferred reverse-mode paths, while dynamic while_loop control is not a general reverse-mode route.

  • Diffrax already uses PyTree states, vmappable solves, and multiple adjoint methods; its adjoint guidance is a useful model for separating forward-mode, reverse-mode, direct, and checkpointed differentiation policies.

  • Equinox provides a minimal PyTree module pattern and filtered transforms without imposing a framework; use it selectively where callable parameter containers improve clarity.

  • Optax is appropriate for composable gradient transformations in noisy or staged objectives.

  • JAXopt is appropriate for differentiable deterministic solvers, implicit differentiation, and tree-structured parameters.

  • Lineax is a good model for linear operator abstractions and differentiable linear solves.

  • Orthax is relevant for future interchangeable orthogonal-polynomial basis families in JAX.

  • PEP 257 and the Google Python style guide are the baseline for docstring and API documentation conventions.

Differentiation Method Ladder

The Python research API should choose the cheapest mathematically correct derivative route for each observable rather than differentiating through every implementation detail.

  • Use native JAX grad/jvp/vjp/scan for smooth fixed-step algebraic, diagnostic, and reduced-window objectives.

  • Use implicit left/right eigenpair differentiation for isolated linear growth and frequency branches.

  • Use implicit root or fixed-point differentiation for equilibria, optimizer inner solves, and converged steady/window conditions when the defining residual is explicit.

  • Use matrix-free linear-solve adjoints for Krylov, preconditioned, and least-squares sensitivities so transpose and tolerance assumptions are inspectable.

  • Use checkpointed unrolled solves only when the transient trajectory itself is the observable and memory/runtime gates pass.

  • Treat adaptive-step derivatives as a promoted feature only after an explicit policy is recorded: fixed-grid replay of accepted steps, forward-mode through bounded controller logic for low-dimensional directions, or a custom/implicit adjoint for the accepted trajectory. The executable can use a faster non-differentiable adaptive controller, but promoted Python objectives must report the active adaptive derivative policy.

  • Treat noisy nonlinear turbulent flux objectives as statistical optimization targets first. Common-random-number finite differences, SPSA/CMA/Bayesian outer loops, late-window ensemble statistics, and transfer gates are required before claiming differentiable nonlinear turbulent-flux optimization.

Technical Layer Map

The detailed target source layout is maintained in Architecture Refactor Plan. The following layer map is retained here only to connect differentiability and validation contracts to broad responsibility groups:

src/spectraxgk/
  core/              # pytrees, dtypes, normalization, static/dynamic helpers
  basis/             # basis protocols, Hermite-Laguerre, Orthax adapters
  grids/             # spectral and field-line grids
  geometry/          # analytic, VMEC, vmec_jax, booz_xform_jax contracts
  physics/           # species, closures, fields, collisions, drives
  operators/         # pure linear/nonlinear kernels and field solves
  solvers/           # linear, nonlinear, adjoints, time steppers
  diagnostics/       # growth, frequency, transport, spectra, UQ, windows
  validation/        # literature gates, reference adapters, benchmark families
  workflows/         # runtime, scans, optimization, plotting, provenance
  io/                # TOML, NetCDF, restart, artifact schemas
  cli/               # executable commands and default demo

Existing public modules remain as facades until the planned API cleanup. New implementation code should be placed under the domain packages named in Architecture Refactor Plan, not added as new root-level prefix modules. Explicit fixed-step diagnostic integration follows this rule: public imports remain on solvers.time.explicit while solvers/time/explicit_diagnostics.py now owns method/time-policy validation, JIT stepper construction, energy/transport sampling, progress rendering, and SimulationDiagnostics assembly as named stages. Term-wise RHS assembly now follows this rule with a small public terms.assembly facade and focused cached-RHS, diagnostic-RHS, field-solve, and helper-policy owner modules. The diagnostic-RHS owner now keeps explicit private stages for state/species normalization, field/Hamiltonian construction, drift/drive/dissipation contribution assembly, fixed-order summation, and species-axis restoration so term-level parity audits exercise named physics pieces instead of one monolithic debug function. Diffrax time integration follows the same pattern: public imports stay on solvers.time while optional dependency/policy helpers, linear save paths, streaming fits, and nonlinear paths live in focused owner modules. Growth diagnostics follow the same rule: diagnostics.growth_rates remains the stable public facade for examples and runtime workflows, while diagnostics.growth_fit and diagnostics.growth_windows own reusable least-squares fitting and automatic fit-window selection; resolved mode-series diagnostics stay with their public diagnostics.growth_rates owner. Validation gates now live in one physical diagnostics owner: diagnostics.validation_gates owns frozen metric/result containers, scalar tolerance evaluation, JSON serialization, and physics/numerics report builders. Upper-limit gate semantics for convergence, mismatch, deficit, and branch-jump thresholds are centralized there instead of repeated in each report builder.

High-Risk Module Split Plan

benchmarks.py

Closed. The installed module is a compact facade for reviewed reference data, normalization constants, and comparison-only branch policies. Canonical TOML inputs plus run_runtime_linear, run_runtime_scan, and run_runtime_parameter_scan own all promoted execution. Root benchmarks/ drivers own case-specific reproduction policy. Required gates remain Cyclone, ETG, KBM, TEM, W7-X/HSX where applicable, plus branch continuity and eigenfunction comparisons; these gates no longer require a duplicate solver engine or module-level monkeypatch seams.

geometry/differentiable.py

Split backend discovery, geometry contracts, VMEC-JAX bridge, Boozer bridge, equal-arc mapping, sensitivity reports, and parity reports. Required gates: optional-backend import behavior, same-WOUT provenance, geometry parity, JVP/VJP/finite-difference agreement, and conditioning diagnostics. The internal file-backed VMEC imported-geometry backend now lives in the geometry.imported_vmec owner module with explicit helper seams for discovery, numerics, radial spline construction, field-line assembly, remap, IO, and pipeline orchestration. The in-memory VMEC/Boozer equal-arc core bridge keeps its public facade stable while optional backend execution and Boozer radial-grid validation are isolated behind private helper seams.

operators/nonlinear/parallel.py

Split domain plans, spectral communication, device-z pencil route, observable reductions, and profiling. Required gates: serial-vs-decomposed RHS identity, physical transport-window identity, profiler artifact schema, and no speedup claim without matched CPU/GPU artifacts. Domain, spectral, and strategy report contracts now live in focused parallel_contracts_* modules and are re-exported only through the public operators/nonlinear/parallel.py surface. Local spectral-state construction, chunk/layout utilities, communication/work models, pencil FFT/bracket primitives, RHS micro-routes, and tolerance helpers now live in focused operators/nonlinear/spectral_* modules behind the unchanged spectraxgk.operators.nonlinear.parallel facade. Logical spectral communication, RHS, and fixed-window integrator identity gates now live directly in focused report, RHS-routing, and fixed-window integrator identity modules. The local domain prototype gates and device-z shard-map route now live in operators/nonlinear/domain_decomposition.py and operators/nonlinear/device_z.py. The facade remains the public import surface for examples. Private FFT, bracket, sharding, and observable helpers are not re-exported: developer tests and profilers import their focused spectral_core or device_z owner directly.

solver_objective_gradients.py

Split eigenvalue objectives, linear-growth objectives, quasilinear flux objectives, nonlinear-window objectives, VMEC/Boozer objective plumbing, and gradient gates. Required gates: branch-locality, spectral-gap guards, finite-difference/JVP/VJP checks, UQ covariance, and objective conditioning. The dominant-growth implicit eigenpair VJP and branch-locality report now live in objectives/eigen.py and remain re-exported by the public solver-objective facade. Solver-objective sampling axes, physical-ky grid construction, and aggregate weights now live in objectives/sampling.py behind the same facade. Core linear/quasilinear objective constants, scalar selectors, operator materialization, growth-rate, and objective-vector evaluators now live in objectives/core.py. Solver-ready geometry objective gates, solver geometry-gradient reports, mode-21 VMEC/Boozer state-gradient reports, reduced nonlinear-window estimator metrics, VMEC/Boozer state coefficient helpers, VMEC/Boozer objective-table plumbing, and VMEC/Boozer finite-difference/line-search gates now live in objectives/geometry.py, objectives/gradient_gates.py, objectives/vmec_boozer_gradients.py, objectives/vmec_boozer_context.py, objectives/vmec_boozer.py, objectives/vmec_boozer_fd.py, objectives/vmec_boozer_line_search.py, objectives/solver_vmec.py, and objectives/solver_gradient_reports.py while solver_objective_gradients.py remains the higher-level public objective surface. Scalar and aggregate VMEC/Boozer finite-difference reports share one settings validator, VMEC-state coefficient context, perturbation helper, and three-point response/curvature diagnostic. Scalar and aggregate VMEC/Boozer line-search reports share one private curvature-gated one-parameter search loop so training, holdout, and finite-difference gates use the same accept/reject policy. Solver-ready gradient gates share one normalized heat/particle transport observable helper so growth-rate, quasilinear, and particle-flux finite-difference checks use the same field-line quadrature path.

nonlinear.py

Split RHS kernels, integrator policies, nonlinear diagnostics, and IMEX paths. Required gates: RHS identity, transport windows, spectral diagnostics, and parity-preserving output schemas. The nonlinear RHS linear-path routing and electromagnetic bracket composition now live in operators/nonlinear/rhs.py while spectraxgk.nonlinear remains the public facade for public imports, monkeypatch-based diagnostics, and runtime workflows. Duplicated explicit and IMEX state-to-diagnostic tuple assembly now lives in operators/nonlinear/diagnostic_state.py with facade-injected diagnostic kernels so existing debug seams remain intact. That module owns separate helpers for field defaulting, growth/frequency mode extraction, scalar diagnostics, and resolved spectra/channel packing, plus the shared state-to-diagnostic closure factory used by explicit and IMEX scans. Shared sampled-scan interval routing, diagnostic-stride selection, progress callback routing, scan-output sampling/finalization, resolved diagnostic packing, and SimulationDiagnostics construction now live in spectraxgk.operators.nonlinear.diagnostics. Shared diagnostic cache, quadrature-weight, omega-mask, z-index, state-projection setup, reusable IMEX operator setup, collision-split policy construction, and fixed/adaptive nonlinear time-step policy now live in spectraxgk.operators.nonlinear.policies with injected public-facade seams. Explicit RK/SSP/K10 one-step policy, cached explicit scan dispatch, explicit diagnostic step construction, and diagnostic scan-selection policy now live in solvers/nonlinear/explicit.py. Public cached RHS/state integration now lives in spectraxgk.solvers.nonlinear.state_integration and public diagnostic entry points live in spectraxgk.solvers.nonlinear.diagnostic_integration; spectraxgk.nonlinear is a small facade that re-exports the validated surface. Explicit diagnostic integration orchestration lives in solvers/nonlinear/diagnostics.py with owner-module-injected geometry, cache, RHS, diagnostic, time-step, progress, and sampled-scan seams. IMEX diagnostic integration orchestration and the fixed diagnostic scan step live in solvers/nonlinear/imex_diagnostics.py. The shared IMEX nonlinear-term closure, GMRES solve-step closure, cached scan policy, fixed-point/GMRES solve, and SSPX3 stage-composition policies live in solvers/nonlinear/imex.py. The public nonlinear facade is now below the active line-count target; repeated explicit/IMEX diagnostic option forwarding is centralized in a small policy table so future solver options are added in one place.

runtime.py and cli.py

Split executable commands, runtime workflows, scan dispatch, progress/ETA, plotting, and artifact handoff. Runtime scan orchestration and combined-ky scan batching now live in spectraxgk.workflows.runtime.orchestration_scan; progress/ETA formatting lives in spectraxgk.workflows.runtime.chunks; and nonlinear artifact/restart handoff lives in spectraxgk.workflows.runtime.orchestration_artifacts behind focused owner modules wired by the public spectraxgk.runtime facade. Runtime nonlinear diagnostics keyword assembly now lives in spectraxgk.workflows.runtime.policies so fixed-window and adaptive diagnostic branches share one policy. Generic runtime linear fit/eigenfunction extraction now lives in spectraxgk.workflows.runtime.diagnostics so the public runtime facade only wires analysis callables, normalization, quasilinear post-processing dependencies, and result construction. Runtime linear/nonlinear dispatch dependency assembly now lives in spectraxgk.workflows.runtime.execution and reads a patchable runtime facade scope at call time, so monkeypatch seams remain explicit without keeping dependency-builder bodies in spectraxgk.runtime. Runtime TOML case wrappers now delegate to spectraxgk.workflows.cases through dependency-injected facades so spectraxgk.runtime remains the public import and monkeypatch surface. Runtime restart-state loading keeps NetCDF/raw loader dispatch in the public facade while sharing one shape-keyword payload between both patchable paths. The public runtime facade also shares one private runtime linear time/fit option bundle across run_runtime_linear, run_runtime_scan, and the combined-ky batch wrapper, so method, timestep, sample-stride, fit-window, mode-method, and fit-signal forwarding cannot drift. The full-GK linear runtime workflow now delegates to spectraxgk.workflows.linear through a dependency-injected facade, including context preparation, time/Krylov dispatch, auto-fallback, fit wiring, and quasilinear post-processing. The workflow body is split into private stages so future differentiable time-policy work can test context, integration, fit, and finalization behavior independently. The full-GK nonlinear runtime workflow now delegates to spectraxgk.workflows.nonlinear through the same facade pattern, including runtime context setup, diagnostics routing, adaptive chunks, fixed-mode/source policy, final-state integration, and result assembly. The default no-input educational demo now delegates to spectraxgk.workflows.demo and uses the same runtime configuration and solver path as every other executable run. Runtime linear, scan, nonlinear, and saved-output plotting executable command bodies delegate to spectraxgk.workflows.runtime.commands so parser dispatch stays separate from simulation, plotting, path override, and artifact side effects. Non-promoted reduced-model runtime paths have been retired from main, so this facade now covers the maintained full-GK executable workflows only. Required gates: default-run behavior, --plot behavior, TOML provenance, restart/output schema, and public import contracts.

workflows/runtime/artifacts.py

Split artifact schema, NetCDF persistence, restart append, and provenance. Generic artifact I/O, linear/quasilinear writers, nonlinear table writers, nonlinear diagnostic reload helpers, and finite-value artifact validation now live under spectraxgk.artifacts. The root workflows/runtime/artifacts.py module remains the public dispatcher and monkeypatch-compatible executable seam. The nonlinear NetCDF diagnostics writer keeps the schema in one owner module but stages scalar Phi2 output, base species histories, split electromagnetic histories, zonal fields, resolved spectra, and turbulent-heating channels so optional diagnostic additions stay localized. Required gates: round-trip persistence, restart append normalization, and plot reload contracts.

linear.py

Split linear RHS, field solves, integrators, and diagnostics. Required gates: field solve identity, late-time growth/frequency metrics, eigenfunction overlap, branch continuity, and JAX transform consistency. Fixed-step diagnostic sampling now keeps the public integrate_linear_diagnostics facade while solvers/linear/integrator_diagnostics.py owns explicit stages for sample validation, cache/state setup, damping assembly, explicit/IMEX/RK step policy, density and Hermite-Laguerre observables, progress callbacks, and every-step versus strided scans.

Execution Phases

Phase 0: freeze contracts

Land this plan, the manifest checker, and a CI-safe test. Record current public imports, coverage targets, benchmark gates, and docs entry points.

Phase 1: introduce protocols and containers

Add small protocol/dataclass modules for basis families, geometry providers, collision operators, field solvers, RHS assembly, diagnostics, objective reports, and artifact schemas. Avoid behavior changes. The first Phase-1 tranche now lives in spectraxgk.core.contracts and spectraxgk.core.extension_points. The first behavior-preserving benchmark split also lives in this phase. Benchmark-family case presets now live directly in spectraxgk.config so the stable public import location is also the physical owner for those dataclasses. spectraxgk.benchmarks owns benchmark initial-condition construction and spectraxgk.benchmarks owns reference containers and CSV loaders. spectraxgk.benchmarks owns benchmark species-to-LinearParams construction and reference hypercollision policy, spectraxgk.diagnostics.growth_rates owns fit-signal and diagnostic normalization policies, spectraxgk.benchmarks owns scan batching, streaming windows, scan-window policy, and spectraxgk.benchmarks owns Krylov defaults plus branch-selection policies. spectraxgk.benchmarks and spectraxgk.benchmarks remain temporary Cyclone validation owners, while spectraxgk.benchmarks owns the kinetic-electron, ETG, KBM, and TEM benchmark implementations directly as the public benchmark entry point. The old benchmark helper bridge has been removed; runners and tests import focused benchmark modules directly. spectraxgk.diagnostics.quasilinear_transport owns the core linear-state quasilinear transport weights and differentiable saturation objectives while spectraxgk.quasilinear remains the stable public facade. spectraxgk.diagnostics.quasilinear_model_selection owns the scoped model-selection claim ledger, with candidate-skill gates, absolute-flux overclaim guardrails, optional optimized-equilibrium audit gates, and final ledger assembly staged separately from the input-normalization helpers in spectraxgk.diagnostics.quasilinear_model_selection. spectraxgk.diagnostics.transport_windows owns late-window nonlinear transport convergence reports, staged as validated window selection, finite-sample counts, drift/terminal-window statistics, block/bootstrap uncertainty, and final gate-report assembly. The first differentiable-geometry support split also lives in this phase: spectraxgk.geometry.backend_discovery owns optional vmec_jax / booz_xform_jax path discovery and spectraxgk.geometry.autodiff_checks owns finite-difference Jacobians, AD/FD reports, conditioning metadata, and strict JSON sanitation. Its public observable-gradient report is now an orchestration layer over parameter validation, observable flattening, AD/FD Jacobian construction, tangent checks, conditioning gates, failure reasons, and JSON report assembly. spectraxgk.geometry.flux_tube_contract owns solver-ready in-memory flux-tube mapping validation and geometry-observable reductions. spectraxgk.geometry.sensitivity owns geometry sensitivity, inverse-design, conditioning, and local UQ reports for solver-ready mappings. spectraxgk.geometry.booz_xform_bridge owns bounded VMEC boundary and Boozer-spectrum bridge checks, Boozer field-line |B| evaluation, and Boozer-to-flux-tube sensitivity diagnostics. spectraxgk.geometry.vmec_state_sensitivity owns optional-backend VMECState sensitivity reports for VMEC-to-Boozer, VMEC metric tensor, and VMEC field-line tensor AD/FD gates, including shared VMEC example loading, coefficient-index validation, coefficient perturbation policy, and the common tensor-observable AD/finite-difference payload builder used by the metric and field-line gates. Metric and field-line tensor gates share one VMEC geometry context/index helper before their observable-specific tensor sampling paths. spectraxgk.geometry.vmec_boozer_core owns the vmec_jax state to booz_xform_jax equal-arc core-profile bridge and solver-facing core arrays, including Boozer radial-profile interpolation and equal-arc remapping policy. Its public bridge is now an orchestration layer over radial Boozer-profile interpolation, equal-arc field-line construction, zero-beta metric/drift profile assembly, and final solver-ready mapping assembly. The profile assembly shares one dtype-aware numerical floor across field, parallel-gradient, Jacobian, metric, curvature, and safety-factor denominators. Boozer metric-gradient terms use a separate float32-safe toroidal-flux denominator floor before grad(theta), grad(phi), and grad(alpha) divisions. spectraxgk.geometry.vmec_boozer_constants owns Boozer constant preparation and equal-arc cache prewarm helpers. spectraxgk.geometry.vmec_tensor_mapping owns direct vmec_jax tensor sampling and conversion into the solver-ready flux-tube mapping contract. It now stages the bridge 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 without adding another source file. spectraxgk.geometry.vmec_flux_tube_reports owns VMEC flux-tube sensitivity and array-parity report orchestration that combines the direct tensor, Boozer equal-arc, and imported-geometry comparison paths. It consumes the same shared VMEC-state loading, coefficient-index validation, and perturbation policy as the Boozer, metric-tensor, and field-line sensitivity reports. 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, while the public report schema remains stable. The imported VMEC field-line backend keeps its public facade stable while its implementation is staged as backend fallback, scalar VMEC profile sampling, Boozer field-line state assembly, Hegna-Nakajima mode corrections, metric/drift coefficient assembly, and final normalized flux-tube packaging. spectraxgk.geometry.vmec_boundary_chain owns VMEC boundary-gradient probe classification, collection row assembly, and projected-transport line-search admission summaries. spectraxgk.geometry.autodiff_checks keeps raw per-entry AD/FD relative errors, but its summary max_rel_ad_fd_error is gate-facing: entries that already satisfy the absolute tolerance do not dominate the summary solely because the finite-difference reference is numerically zero. spectraxgk.geometry.numerics owns pure parity metrics, interpolation, radial derivative, Boozer half-mesh, Fourier field-line, and periodic sampling helpers. spectraxgk.geometry.imported_miller now owns the complete internal Miller backend, including numerics, surface/theta-grid construction, profile assembly, request-to-EIK orchestration, and NetCDF writeout. Shared JAX finite-difference kernels live in spectraxgk.geometry.kernels. spectraxgk.geometry.differentiable retains object-identical re-exports for pure helpers and thin wrappers for optional-backend bridge functions whose tests patch facade-level backend discovery.

Phase 2: split pure kernels

Move basis, gyroaverage, field-solve, linear-term, nonlinear-bracket, and diagnostic kernels first. These have the clearest unit and numerical tests. The nonlinear bracket and linear dissipation slices now live under operators/nonlinear/brackets.py and operators/linear/dissipation.py; no compatibility modules remain at their former terms paths. The linear cache builder now keeps one public construction function but delegates twist-shift policy, perpendicular-wavenumber/drift arrays, Laguerre gyroaverage arrays, and linked-boundary metadata to focused private stages. Future linear-cache work should extend or test those stages before changing the final LinearCache assembly. The implicit linear solver now follows the same rule: the public operator builder delegates state normalization, preconditioner-diagonal assembly, linked Hermite-line solves, coarse kx projection, preconditioner selection, and matrix-free matvec construction to private stages. Future implicit-solver changes should add identity/finite gates for the relevant stage before changing GMRES integration behavior.

Phase 3: split differentiable geometry

Move the remaining VMEC/Boozer bridge and parity routines behind in-memory geometry contracts. Optional backend discovery and geometry AD/FD gate utilities, the in-memory flux-tube contract, sensitivity reports, and pure numerical helpers are already split into Phase-1 support modules. Keep same-WOUT and finite-difference gates mandatory.

Phase 4: split objectives and AD policies

Separate linear, quasilinear, nonlinear-window, and VMEC/Boozer objectives. Add explicit adjoint policy selection and branch-locality guards.

Phase 5: split runtime, CLI, and workflows

Move side-effectful command/runtime/plotting code out of solver kernels. Add progress and provenance tests for the executable. Default-demo side effects now have an injectable workflow seam for tests and downstream examples.

Phase 6: extension registries

Add registries for collision operators, basis families, closure models, field solvers, geometry providers, transport diagnostics, and benchmark families. Each registry entry must declare tests and docs.

Phase 7: deprecation cleanup

Keep stable public facades through the release series. Remove old helper imports once downstream examples, docs, and tests use canonical owners.

Differentiability Contract

Every promoted differentiable workflow should declare:

  • the optimized parameters and their PyTree structure;

  • which inputs are static and which are traced arrays;

  • whether gradients use forward mode, reverse mode, direct adjoint, checkpointed adjoint, implicit differentiation, or custom JVP/VJP;

  • the finite-difference step policy and accepted error window;

  • spectral-gap or branch-locality criteria for eigenvalue objectives;

  • conditioning, rank, and covariance diagnostics for inverse/UQ examples;

  • whether the observable is a production physics observable or a reduced differentiable proxy.

Parallel trajectories follow the same rule. The production two-species electrostatic fixed-step integrator now has a reverse-mode parameter-gradient gate through its enclosing pmap: a thermodynamic-drive derivative agrees with centered finite differences in float32. Input placement remains outside the differentiated function, and traced parameters bypass host conversion. A separate adaptive Diffrax gate now promotes low-dimensional forward-mode JVPs: derivative_mode="forward" uses native JAX rules through the accepted Tsit5 trajectory, agrees with centered finite differences to 1.9e-5 relative, and remains stable when the controller tolerance is tightened. The same observable now passes a checkpointed reverse gradient through Diffrax’s RecursiveCheckpointAdjoint at 1.9e-5 relative to finite differences. Direct non-checkpointed adaptive reverse mode, IMEX trajectories, electromagnetic trajectories, and mixed species–Hermite derivatives remain fail-closed until each has a bounded-memory observable-level gate.

Generic structured derivatives should use reviewed SOLVAX primitives when their contracts pass the required dtype and transformation gates. In particular, memory-chunked Jacobians are preferred when an unchunked jacfwd/jacrev materialization is the measured memory bottleneck, and implicit linear/root solve rules are preferred for converged solves whose residual equation and transpose are explicit. SPECTRAX-GK still owns branch locality, adaptive-controller assumptions, turbulent-window statistics, and finite-difference acceptance. A solver dependency cannot turn an ill-conditioned or non-converged physical observable into a valid derivative.

Before replacing a local path, require real and complex dtype coverage as appropriate, JIT/vmap/JVP/VJP checks, transpose or Hermitian-adjoint identity, finite-difference agreement, and a physical observable gate. SOLVAX 0.6.1 satisfies those contracts for complex Krylov, fixed-point, tridiagonal, and chunked-Jacobian primitives. SPECTRAX-GK uses the admitted primitives only on paths whose physical trajectory and branch gates pass; shift-invert remains on its prior backend because its outer eigenpair gate is still open.

Recommended default choices:

  • low-dimensional geometry or scalar scans: jax.jvp and finite differences;

  • many-parameter smooth deterministic objectives: reverse mode or JAXopt implicit differentiation after branch-locality gates pass;

  • time-integration objectives: Diffrax-style explicit adjoint policy selection with checkpointing documented;

  • noisy nonlinear transport windows: common-random-number replicas, robust window statistics, derivative-free outer loops when gradients are not conditioned, and no production claim until long post-transient audits pass.

Test Matrix

Unit and numerical tests

Basis orthonormality and recurrence, gyroaverage limits, field-solve consistency, nonlinear bracket symmetries, conservation/free-energy reduced limits, observed-order checks, and artifact round trips.

Autodiff tests

Finite-difference versus JVP/VJP for geometry observables, eigenvalue objectives, quasilinear flux weights, reduced nonlinear-window estimators, UQ covariance, and sensitivity maps.

Parity tests

Reference-code comparisons remain external validation inputs. Required metrics include growth rate, real frequency, eigenfunction overlap, late-window nonlinear transport, spectra, and diagnostic schema equality.

Literature-anchored physics gates

Keep the existing Cyclone, shaped tokamak, KBM, ETG, TEM, W7-X/HSX, Rosenbluth-Hinton/Merlo, and quasilinear/nonlinear holdout gates. Add new literature gates only when the observable, window, normalization, and tolerance are documented.

Performance and parallel tests

Require numerical identity before speedup. Require profiler artifacts before claims. Track compile time, warm throughput, memory, communication model, and CPU/GPU behavior separately.

Acceptance Criteria

A refactor tranche is ready only when all of the following are true:

  • public imports either remain object-compatible or have documented deprecation shims;

  • package-wide coverage remains at or above 95%;

  • each moved module has unit tests and at least one physics, numerical, autodiff, parity, or artifact gate appropriate to its responsibility;

  • no public module exceeds 800 lines unless the manifest records an exception;

  • no internal module exceeds 1200 lines unless the manifest records an exception;

  • docs list equations, normalization, shapes, differentiability status, and extension points for public APIs;

  • benchmark and reference-code adapters are isolated from solver kernels;

  • all published figures or claims have replay commands and artifact provenance.

Developer Checklist For Future PRs

  • Update tools/differentiable_refactor_manifest.toml before moving code.

  • Add or update fast tests before changing numerical behavior.

  • Keep stable public facades; remove old helper shims after examples and docs are migrated.

  • Run the manifest checker and the affected fast tests locally.

  • If a gate tolerance changes, update the validation docs and artifact ledger.

  • If a new extension point is added, include one minimal real implementation and one test that demonstrates how a contributor would add another one.

References and External Models

The plan is anchored in current JAX and gyrokinetic validation practice: