API Reference

Basis

Hermite and Laguerre basis utilities.

spectraxgk.basis.hermite_ladder_coeffs(n_max: int) → tuple[Array, Array]

Return sqrt(n+1) and sqrt(n) arrays for Hermite ladder operators.

spectraxgk.basis.hermite_normed(x: Array, n_max: int) → Array

Normalized Hermite functions with weight exp(-x**2).

psi_n = H_n(x) / sqrt(2**n * n! * sqrt(pi))

spectraxgk.basis.hermite_physicists(x: Array, n_max: int) → Array

Physicists’ Hermite polynomials H_n(x) for n=0..n_max.

Weight: exp(-x**2). Recurrence:

H_0 = 1 H_1 = 2x H_{n+1} = 2x H_n - 2n H_{n-1}

spectraxgk.basis.laguerre(x: Array, l_max: int) → Array

Laguerre polynomials L_l(x) for l=0..l_max.

Weight: exp(-x). Recurrence:

L_0 = 1 L_1 = 1 - x (l+1) L_{l+1} = (2l+1-x) L_l - l L_{l-1}

Gyroaverage

Gyroaveraging coefficients for Laguerre velocity space.

spectraxgk.gyroaverage.J_l_all(b: Array, l_max: int) → Array

Gyroaveraging coefficients matching the GX Laguerre-Hermite convention.

spectraxgk.gyroaverage.bessel_j0(x: Array) → Array

Return J0(x) using a Cephes-style approximation (GX-compatible).

spectraxgk.gyroaverage.bessel_j1(x: Array) → Array

Return J1(x) using a Cephes-style approximation (GX-compatible).

spectraxgk.gyroaverage.gamma0(b: Array) → Array

Compute Gamma_0(b) = exp(-b) I_0(b) using i0e for stability.

spectraxgk.gyroaverage.gx_factorial(m: Array) → Array

Return GX’s single-precision factorial approximation.

spectraxgk.gyroaverage.gx_laguerre_nj(nl: int) → int

GX default for number of Laguerre quadrature points.

spectraxgk.gyroaverage.gx_laguerre_transform(nl: int) → tuple[ndarray, ndarray, ndarray]

Return GX-style Laguerre transform matrices and roots.

spectraxgk.gyroaverage.sum_Jl2(b: Array, l_max: int) → Array

Truncated sum of J_l(b)^2, useful for Gamma_0 convergence checks.

Geometry

Analytic flux-tube geometry for the Cyclone base case.

class spectraxgk.geometry.FluxTubeGeometryData(theta: Array, gradpar_value: float, bmag_profile: Array, bgrad_profile: Array, gds2_profile: Array, gds21_profile: Array, gds22_profile: Array, cv_profile: Array, gb_profile: Array, cv0_profile: Array, gb0_profile: Array, jacobian_profile: Array, grho_profile: Array, q: float, s_hat: float, epsilon: float, R0: float, B0: float = 1.0, alpha: float = 0.0, drift_scale: float = 1.0, kxfac: float = 1.0, theta_scale: float = 1.0, nfp: int = 1, kperp2_bmag: bool = True, bessel_bmag_power: float = 0.0, source_model: str = 'sampled', theta_closed_interval: bool = False)

Sampled flux-tube geometry contract for solver-ready metric profiles.

trim_terminal_theta_point() → FluxTubeGeometryData

Return a copy without the terminal theta sample.

GX *.eik.nc files commonly store a closed theta interval, while the spectral solver uses the matching open interval with the terminal point excluded. Trimming keeps the imported coefficients aligned with the runtime grid without changing the physical extent.

class spectraxgk.geometry.SAlphaGeometry(q: float, s_hat: float, epsilon: float, R0: float = 1.0, B0: float = 1.0, alpha: float = 0.0, drift_scale: float = 1.0, kperp2_bmag: bool = True, bessel_bmag_power: float = 0.0)

Simple s-alpha geometry with circular concentric flux surfaces.

bgrad(theta: Array) → Array

Magnetic field gradient term used in mirror force.

bmag(theta: Array) → Array

Magnetic field strength for circular s-alpha geometry.

drift_coeffs(theta: Array) → tuple[Array, Array, Array, Array]

Curvature and grad-B drift coefficients for s-alpha geometry.

drift_components(kx: Array, ky: Array, theta: Array) → tuple[Array, Array]

Return cv_d and gb_d drift components in (ky, kx, theta).

gradpar() → float

Parallel gradient factor for s-alpha geometry (constant for equal-arc).

k_perp2(kx0: Array, ky: Array, theta: Array) → Array

Perpendicular wave-number squared for s-alpha geometry.

kx_effective(kx0: Array, ky: Array, theta: Array) → Array

Field-aligned kx(theta) with s-alpha shear shift.

metric_coeffs(theta: Array) → tuple[Array, Array, Array]

Metric coefficients (gds2, gds21, gds22) for s-alpha geometry.

omega_d(kx: Array, ky: Array, theta: Array) → Array

Magnetic drift frequency for s-alpha geometry.

class spectraxgk.geometry.SlabGeometry(s_hat: float = 0.0, z0: float | None = None, q: float = 1.0, epsilon: float = 0.0, R0: float = 1.0, B0: float = 1.0, alpha: float = 0.0, drift_scale: float = 0.0, kperp2_bmag: bool = True, bessel_bmag_power: float = 0.0, zero_shat: bool = False)

GX slab geometry contract.

spectraxgk.geometry.apply_geometry_grid_defaults(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: GridConfig) → GridConfig

Apply GX-aligned grid defaults implied by the selected geometry.

spectraxgk.geometry.apply_gx_geometry_grid_defaults(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: GridConfig) → GridConfig

Apply GX-aligned grid defaults implied by the selected geometry.

spectraxgk.geometry.build_flux_tube_geometry(cfg: GeometryConfig) → SAlphaGeometry | SlabGeometry | FluxTubeGeometryData

Build an analytic or imported flux-tube geometry from config.

spectraxgk.geometry.effective_boundary(boundary: str, *, s_hat: float, zero_shat: bool = False, threshold: float = 1e-05) → str

Return the effective boundary after zero-shear promotion.

spectraxgk.geometry.ensure_flux_tube_geometry_data(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, theta: Array) → FluxTubeGeometryData

Return sampled geometry data for analytic or pre-sampled inputs.

spectraxgk.geometry.gx_effective_boundary(boundary: str, *, s_hat: float, zero_shat: bool = False, threshold: float = 1e-05) → str

Return the effective boundary after zero-shear promotion.

spectraxgk.geometry.gx_twist_shift_params(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: GridConfig) → tuple[int, float]

Return (jtwist, x0) for twist-and-shift grid defaults.

spectraxgk.geometry.gx_zero_shat_enabled(s_hat: float, *, zero_shat: bool = False, threshold: float = 1e-05) → bool

Return the effective zero-shear state.

spectraxgk.geometry.load_gx_geometry_netcdf(path: str | Path) → FluxTubeGeometryData

Load sampled geometry from a GX-style NetCDF file.

spectraxgk.geometry.sample_flux_tube_geometry(geom: SAlphaGeometry | SlabGeometry, theta: Array) → FluxTubeGeometryData

Sample an analytic geometry model onto a flux-tube theta grid.

spectraxgk.geometry.twist_shift_params(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: GridConfig) → tuple[int, float]

Return (jtwist, x0) for twist-and-shift grid defaults.

spectraxgk.geometry.zero_shear_enabled(s_hat: float, *, zero_shat: bool = False, threshold: float = 1e-05) → bool

Return the effective zero-shear state.

Differentiable Geometry

Differentiable geometry bridge contracts for VMEC/JAX pipelines.

spectraxgk.geometry.differentiable.booz_xform_flux_tube_mapping_from_inputs(inputs: Any, *, mboz: int = 2, nboz: int = 1, ntheta: int = 96, alpha: float = 0.0, surface_index: int = 0, magnetic_shear: float = 0.35, R0: float = 1.0, B0: float = 1.0, drift_scale: float = 1.0, jit: bool = False) → dict[str, Any]

Build a solver-ready flux-tube mapping from booz_xform_jax output.

This is the first bounded production bridge step between JAX-native Boozer coordinates and SPECTRAX-GK. It uses the real Boozer magnetic-field spectrum for bmag/bgrad and supplies smooth metric/drift profiles with the same solver-ready names accepted by flux_tube_geometry_from_mapping().

Full VMEC/Boozer metric parity remains a separate promotion gate: a high-fidelity backend must replace the smooth metric/drift closure here with sampled VMEC/Boozer metric tensors before nonlinear optimization claims are made.

spectraxgk.geometry.differentiable.booz_xform_flux_tube_sensitivity_report(*, params: Array | None = None, fd_step: float = 2e-05, mboz: int = 2, nboz: int = 1, ntheta: int = 64) → dict[str, object]

AD/FD-check a Boozer-spectrum-to-flux-tube geometry bridge.

params = [axisymmetric_B_ripple, helical_B_ripple] perturbs a tiny one-surface VMEC-to-Boozer input bundle. The real booz_xform_jax transform is run for each parameter vector; its Boozer |B| spectrum is sampled on a field line and converted into FluxTubeGeometryData.

spectraxgk.geometry.differentiable.booz_xform_spectral_sensitivity_report(*, ripple: float = 0.05, fd_step: float = 2e-05, mboz: int = 2, nboz: int = 0) → dict[str, object]

Validate a real booz_xform_jax spectral derivative when available.

This is a deliberately tiny Boozer-transform gate. It constructs an axisymmetric one-surface VMEC-to-Boozer input bundle, runs the real booz_xform_jax functional API, and checks the derivative of a Boozer magnetic-spectrum norm with respect to a magnetic-ripple coefficient against central finite differences.

The gate strengthens the bridge beyond import discovery while remaining bounded enough for examples and optional local validation. It is not a full VMEC-state-to-flux-tube parity claim; that requires an equilibrium solve, field-line sampling, and comparison against the production imported-VMEC geometry path.

spectraxgk.geometry.differentiable.discover_differentiable_geometry_backends() → dict[str, object]

Discover optional vmec_jax and booz_xform_jax bridge APIs.

spectraxgk.geometry.differentiable.evaluate_boozer_bmag_on_field_line(theta: Array, *, bmnc_b: Array, ixm_b: Array, ixn_b: Array, iota: Array | float, alpha: float = 0.0) → tuple[Array, Array]

Evaluate a Boozer |B| spectrum and theta derivative on a field line.

The field-line label convention is \(\alpha = \theta - \iota\zeta\). This helper is intentionally small and JAX-native so that the booz_xform_jax spectral output can be differentiated all the way into the sampled SPECTRAX-GK geometry contract.

spectraxgk.geometry.differentiable.finite_difference_jacobian(fn: Any, params: Array, *, step: float = 0.0001) → Array

Central finite-difference Jacobian for small validation problems.

spectraxgk.geometry.differentiable.flux_tube_geometry_from_mapping(data: Mapping[str, Any], *, source_model: str = 'vmec_jax', validate_finite: bool = True) → FluxTubeGeometryData

Build FluxTubeGeometryData from an in-memory differentiable backend.

The input is intentionally the solver-ready flux-tube contract, not a fake equilibrium. vmec_jax / booz_xform_jax pipelines should first produce the sampled field-line arrays named here, then this function validates shapes/finite values and hands them to the existing solver.

spectraxgk.geometry.differentiable.flux_tube_geometry_from_vmec_boozer_state(state: Any, static: Any, indata: Any, wout: Any, *, surface_index: int | None = None, torflux: float | None = None, alpha: float = 0.0, ntheta: int = 32, mboz: int = 21, nboz: int = 21, jit: bool = False, surface_stencil_width: int | None = None, reference_length: float | None = None, reference_b: float | None = None, source_model: str = 'mode21_vmec_boozer_state', validate_finite: bool = True) → FluxTubeGeometryData

Build solver-ready geometry directly from a solved vmec_jax state.

This is the production-facing in-memory bridge for differentiable optimization workflows. It keeps the path inside JAX-compatible objects:

VMECState -> BoozXformInputs -> booz_xform_jax -> FluxTubeGeometryData.

Runtime VMEC file generation can still use the NetCDF/EIK route, but differentiable stellarator optimization should call this function or a higher-level objective wrapper around it so gradients never pass through filesystem artifacts.

spectraxgk.geometry.differentiable.flux_tube_geometry_observables(geom: FluxTubeGeometryData) → Array

Return differentiable scalar observables from solver-ready geometry.

The observables are intentionally geometry-level quantities: mean field strength, relative ripple, metric norm, drift norm, mean Jacobian, and mean parallel-gradient factor. They are used to validate the differentiable vmec_jax / booz_xform_jax bridge before any turbulence observable is promoted into an optimization claim.

spectraxgk.geometry.differentiable.geometry_inverse_design_report(mapping_fn: Any, initial_params: Array, target_observables: Array, *, observable_indices: Sequence[int] | None = None, max_steps: int = 8, damping: float = 1e-08, fd_step: float = 0.0001, regularization: float = 1e-08, source_model: str = 'vmec_jax:in-memory') → dict[str, object]

Run a small Gauss-Newton geometry inverse-design validation.

mapping_fn(params) must be the same solver-ready field-line mapping accepted by flux_tube_geometry_from_mapping(). The routine is meant for differentiable vmec_jax / booz_xform_jax workflows: it keeps the optimization, sensitivity check, and local UQ covariance in one JSON-friendly report so examples can validate the full AD contract without depending on a long equilibrium solve in CI.

spectraxgk.geometry.differentiable.geometry_observable_names() → tuple[str, ...]

Return the ordered geometry observables used by bridge AD checks.

spectraxgk.geometry.differentiable.geometry_sensitivity_report(mapping_fn: Any, params: Array, *, fd_step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, source_model: str = 'vmec_jax:in-memory') → dict[str, object]

Validate geometry-observable sensitivities by AD and finite differences.

mapping_fn(params) must return the solver-ready field-line mapping accepted by flux_tube_geometry_from_mapping(). The report is strict JSON friendly so examples and CI gates can preserve the derivative contract without depending on large VMEC solves.

spectraxgk.geometry.differentiable.observable_gradient_validation_report(observable_fn: Callable[[Array], Any], params: Array | ndarray, *, fd_step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, observable_names: Sequence[str] | None = None, param_names: Sequence[str] | None = None, tangent: Array | ndarray | None = None, relative_floor: float = 1e-12, min_rank: int | None = None, condition_number_max: float | None = 1000000000000.0, report_kind: str = 'observable_gradient_validation') → dict[str, object]

Validate observable gradients by AD, finite differences, and conditioning.

observable_fn(params) may return any array-like observable vector. The returned report is strict JSON-compatible: nonfinite diagnostic numbers are represented as None while finite flags and failure reasons preserve why the gate failed.

spectraxgk.geometry.differentiable.vmec_boundary_aspect_sensitivity_report(params: Array, *, fd_step: float = 2e-05, mpol: int = 2, ntor: int = 0, ntheta: int = 96, nphi: int = 1, nfp: int = 1) → dict[str, object]

Validate a real vmec_jax boundary-aspect derivative when available.

The check intentionally stops at the boundary Fourier API. Full VMEC solves are too expensive and environment-sensitive for the default package tests, but the boundary-aspect path verifies that SPECTRAX-GK can discover a vmec_jax checkout and differentiate through its JAX-native boundary data structures before higher-cost optimization workflows are promoted.

spectraxgk.geometry.differentiable.vmec_field_line_tensor_observable_names() → tuple[str, ...]

Return the ordered observables used by the VMEC field-line tensor gate.

spectraxgk.geometry.differentiable.vmec_jax_boozer_equal_arc_core_profiles_from_state(state: Any, static: Any, indata: Any, wout: Any, *, surface_index: int | None = None, torflux: float | None = None, alpha: float = 0.0, ntheta: int = 32, mboz: int = 21, nboz: int = 21, jit: bool = False, surface_stencil_width: int | None = None, reference_length: float | None = None, reference_b: float | None = None) → dict[str, Any]

Return Boozer equal-arc core profiles from a real vmec_jax state.

This bridge follows the same high-level convention as the imported VMEC/EIK runtime path for scalar/core field-line quantities and the zero-beta Boozer metric/drift terms that can be reconstructed directly from booz_xform_jax output: Boozer |B|, equal-arc constant gradpar, q, magnetic shear, solver Jacobian normalization, gds*/grho, and loaded-convention cvdrift/gbdrift coefficients. General finite-beta pressure corrections and broader-equilibrium drift gates remain separate promotion steps.

spectraxgk.geometry.differentiable.vmec_jax_boozer_flux_tube_sensitivity_report(*, params: Array | None = None, case_name: str = 'circular_tokamak', radial_index: int | None = None, mode_index: int = 1, surface_index: int | None = None, fd_step: float = 1e-05, mboz: int = 2, nboz: int = 0, ntheta: int = 32) → dict[str, object]

AD/FD-check vmec_jax state coefficients through the Boozer bridge.

This is the first end-to-end optional-backend gate that starts from a real vmec_jax VMECState instead of a hand-built Boozer input bundle. It loads a small bundled VMEC example, perturbs two VMEC Fourier coefficients [Rcos(radial_index, mode_index), Zsin(radial_index, mode_index)], converts the perturbed state to booz_xform_jax inputs, samples the resulting Boozer |B| spectrum on a field line, and checks SPECTRAX-GK geometry-observable derivatives against central finite differences.

The current metric/drift closure is still intentionally smooth and local to SPECTRAX-GK. Full production promotion requires replacing it with sampled VMEC/Boozer metric tensors and parity-checking those arrays against the imported VMEC/EIK path.

spectraxgk.geometry.differentiable.vmec_jax_field_line_tensor_sensitivity_report(*, params: Array | None = None, case_name: str = 'nfp4_QH_warm_start', radial_index: int | None = None, mode_index: int = 1, surface_index: int | None = None, alpha: float = 0.0, ntheta: int = 32, fd_step: float = 1e-06, b2_floor: float = 1e-24, rms_epsilon: float = 1e-24) → dict[str, object]

AD/FD-check VMEC field-line metric and |B| tensors from vmec_jax.

This optional-backend gate is deliberately upstream of the production SPECTRAX-GK metric/drift closure. It loads a real stellarator vmec_jax example state, perturbs two VMEC Fourier coefficients, evaluates vmec_jax.geom.eval_geom and vmec_jax.vmec_bcovar, samples raw metric/|B| tensors on a fixed VMEC field-line convention, and checks those observable derivatives against central finite differences.

The gate proves differentiability from VMECState through real VMEC metric and magnetic-field tensors. The later production gate must still convert those tensors into the exact SPECTRAX-GK gds*/drift contract and compare against the imported VMEC/EIK path.

spectraxgk.geometry.differentiable.vmec_jax_flux_tube_array_parity_report(*, case_name: str = 'nfp4_QH_warm_start', surface_index: int | None = None, alpha: float = 0.0, ntheta: int = 16, mboz: int = 21, nboz: int = 21, boundary: str = 'none', include_shear_variation: bool = True, include_pressure_variation: bool = True, core_tolerance: float = 0.05, scalar_tolerance: float = 0.005, equal_arc_core_tolerance: float = 0.01, equal_arc_derivative_tolerance: float = 0.03, equal_arc_metric_tolerance: float = 0.08, equal_arc_drift_tolerance: float = 0.08) → dict[str, object]

Compare the direct vmec_jax flux-tube arrays to imported VMEC/EIK.

This is a diagnostic promotion gate, not a differentiability check. It starts from the same real vmec_jax example state used by vmec_jax_flux_tube_sensitivity_report(), builds the direct VMEC-tensor-derived flux-tube mapping, then generates the existing imported VMEC/EIK geometry on the same surface and compares solver-facing arrays.

The expected current result is that q and magnetic shear are close while metric/drift arrays remain open because the direct path still uses a VMEC-coordinate/equal-theta convention and a local grad-\(B\) closure instead of the production Boozer equal-arc/Hegna-Nakajima convention.

spectraxgk.geometry.differentiable.vmec_jax_flux_tube_mapping_from_state(state: Any, static: Any, wout: Any, *, surface_index: int | None = None, alpha: float = 0.0, ntheta: int = 32, b2_floor: float = 1e-24, metric_floor: float = 1e-24, reference_length: float | None = None, reference_b: float | None = None, drift_scale: float = 1.0) → dict[str, Any]

Build a solver-ready flux-tube mapping directly from vmec_jax tensors.

This is the VMEC-native bridge step: it evaluates vmec_jax.geom and vmec_jax.vmec_bcovar, samples the covariant metric and |B| on a fixed field line, inverts the sampled metric to construct perpendicular flux-tube metric coefficients, and emits the flux_tube_geometry_from_mapping() contract.

The metric and magnetic-field derivatives are differentiable with respect to the VMEC state. The drift coefficients use a local grad-\(B\) projection closure so that downstream solver contracts remain populated; the full Hegna-Nakajima/imported-VMEC drift parity gate remains a separate production promotion step.

spectraxgk.geometry.differentiable.vmec_jax_flux_tube_sensitivity_report(*, params: Array | None = None, case_name: str = 'nfp4_QH_warm_start', radial_index: int | None = None, mode_index: int = 1, surface_index: int | None = None, alpha: float = 0.0, ntheta: int = 24, fd_step: float = 2e-06) → dict[str, object]

AD/FD-check VMEC-state derivatives through a solver-ready flux tube.

Unlike the Boozer-only bridge, this report starts from a real vmec_jax state, evaluates VMEC metric and magnetic-field tensors, emits the SPECTRAX-GK FluxTubeGeometryData mapping, and differentiates geometry observables through that path.

spectraxgk.geometry.differentiable.vmec_jax_metric_tensor_sensitivity_report(*, params: Array | None = None, case_name: str = 'circular_tokamak', radial_index: int | None = None, mode_index: int = 1, surface_index: int | None = None, fd_step: float = 1e-05, rms_epsilon: float = 1e-24) → dict[str, object]

AD/FD-check real vmec_jax metric tensors from a VMECState.

The Boozer bridge validates the straight-field-line |B| spectrum, but SPECTRAX-GK’s production geometry contract also needs sampled metric and drift tensors. This gate stays upstream of any reduced closure: it loads a real vmec_jax example state, perturbs two VMEC Fourier coefficients, evaluates vmec_jax.geom.eval_geom, and checks metric-tensor observable derivatives against central finite differences.

This is a prerequisite for replacing the smooth metric/drift closure in booz_xform_flux_tube_mapping_from_inputs(); it is not by itself the final Boozer-field-line metric parity gate.

spectraxgk.geometry.differentiable.vmec_metric_tensor_observable_names() → tuple[str, ...]

Return the ordered observables used by the VMEC metric-tensor gate.

Grids

Spectral grid utilities for flux-tube geometry.

class spectraxgk.grids.SpectralGrid(kx: 'jnp.ndarray', ky: 'jnp.ndarray', z: 'jnp.ndarray', kx_grid: 'jnp.ndarray', ky_grid: 'jnp.ndarray', dealias_mask: 'jnp.ndarray', y0: 'float', x0: 'float', boundary: 'str', jtwist: 'int | None', non_twist: 'bool', kxfac: 'float', ky_mode: 'jnp.ndarray | None' = None)

spectraxgk.grids.gx_real_fft_kx(kx: Array) → Array

Return the kx ordering used with real-FFT nonlinear kernels.

spectraxgk.grids.gx_real_fft_ky(ky: Array) → Array

Return the compressed non-negative ky block for a real FFT.

spectraxgk.grids.gx_real_fft_mesh(kx_grid: Array, ky_grid: Array) → tuple[Array, Array, Array, Array]

Return compressed real-FFT (kx, ky) multipliers and meshgrids.

spectraxgk.grids.real_fft_mesh(kx_grid: Array, ky_grid: Array) → tuple[Array, Array, Array, Array]

Return compressed real-FFT (kx, ky) multipliers and meshgrids.

spectraxgk.grids.real_fft_ordered_kx(kx: Array) → Array

Return the kx ordering used with real-FFT nonlinear kernels.

spectraxgk.grids.real_fft_unique_ky(ky: Array) → Array

Return the compressed non-negative ky block for a real FFT.

spectraxgk.grids.select_gx_real_fft_ky_grid(grid: SpectralGrid, ky_values: Array | ndarray | Sequence[float]) → SpectralGrid

Return a positive-ky real-FFT view of grid.

spectraxgk.grids.select_ky_grid(grid: SpectralGrid, ky_index: int | Array | ndarray | Sequence[int]) → SpectralGrid

Return a linear-solver grid sliced down to one or more ky indices.

The parent grid’s two-thirds mask belongs to nonlinear FFT products. A linear ky scan must not zero a selected high-ky mode just because that row would be dealiased in a nonlinear convolution, so sliced linear grids carry an all-true mask.

spectraxgk.grids.select_real_fft_ky_grid(grid: SpectralGrid, ky_values: Array | ndarray | Sequence[float]) → SpectralGrid

Return a positive-ky real-FFT view of grid.

spectraxgk.grids.twothirds_mask(Ny: int, Nx: int) → Array

2/3 dealiasing mask for 2D Fourier grids.

Operators

Linear operators used in the Hermite-Laguerre gyrokinetic system.

spectraxgk.operators.hermite_streaming(G: Array, kpar: Array, vth: float) → Array

Parallel streaming operator acting on the Hermite index.

Parameters:

• G (jnp.ndarray) – Array with Hermite index on the last axis.

• kpar (jnp.ndarray) – Parallel wave number (broadcastable to G without Hermite axis).

• vth (float) – Thermal speed scaling.

Linear

Linear electrostatic gyrokinetic building blocks (Hermite-Laguerre).

spectraxgk.linear.integrate_linear(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, method: str = 'rk4', cache: LinearCache | None = None, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_solve_method: str = 'batched', implicit_preconditioner: Callable[[Array], Array] | str | None = None, terms: LinearTerms | None = None, checkpoint: bool = False, sample_stride: int = 1, donate: bool = False, show_progress: bool = False, parallel: Any | None = None) → tuple[Array, Array]

Time integrate the linear system using a fixed-step scheme.

spectraxgk.linear.integrate_linear_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk4', cache: LinearCache | None = None, terms: LinearTerms | None = None, sample_stride: int = 1, species_index: int | None = 0, record_hl_energy: bool = False, show_progress: bool = False) → tuple[Array, Array, Array] | tuple[Array, Array, Array, Array]

Integrate and return (G_out, phi_t, density_t) for diagnostics.

spectraxgk.linear.linear_rhs(G: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, terms: LinearTerms | None = None, *, dt: Array | float | None = None) → tuple[Array, Array]

Compute the linear RHS and electrostatic potential.

Parameters:

• G (jnp.ndarray) – Laguerre-Hermite moments with shape (Nl, Nm, Ny, Nx, Nz).

• grid (SpectralGrid) – Flux-tube spectral grid.

• geom (SAlphaGeometry) – Analytic s-alpha geometry.

• params (LinearParams) – Physical and normalization parameters.

spectraxgk.linear.linear_rhs_cached(G: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | None = None, *, use_jit: bool = True, use_custom_vjp: bool = True, dt: Array | float | None = None, force_electrostatic_fields: bool = False) → tuple[Array, Array]

Compute the linear RHS using precomputed geometry arrays.

Linear Linked Boundaries

Linked-boundary FFT maps and damping profiles for linear operators.

spectraxgk.linear_linked._build_linked_end_damping_profile(*, linked_indices: tuple[Array, ...], ny: int, nx: int, nz: int, widthfrac: float, ky_mode: ndarray | None = None) → ndarray

Construct the GX linked-boundary damping profile on the full FFT grid.

spectraxgk.linear_linked._build_linked_fft_maps(kx: ndarray, ky: ndarray, y0: float, jtwist: int, dz: float, nz: int, real_dtype: dtype, ky_mode: ndarray | None = None) → tuple[tuple[Array, ...], tuple[Array, ...]]

Construct GX-style linked FFT index maps for the parallel derivative.

Linear Cache

Linear operator cache construction and damping helpers.

class spectraxgk.linear_cache.LinearCache(Jl: Array, b: Array, kperp2: Array, kperp2_bmag: bool, bmag: Array, omega_d: Array, cv_d: Array, gb_d: Array, bgrad: Array, jacobian: Array, mask0: Array, dz: Array, kz: Array, ky: Array, kx: Array, kx_grid: Array, ky_grid: Array, dealias_mask: Array, kxfac: Array, lb_lam: Array, collision_lam: Array, hyper_ratio: Array, ratio_l: Array, ratio_m: Array, ratio_lm: Array, mask_const: Array, mask_kz: Array, m_pow: Array, m_norm_kz_factor: Array, damp_profile: Array, linked_damp_profile: Array, l: Array, m: Array, l4: Array, sqrt_m: Array, sqrt_m_p1: Array, sqrt_p: Array, sqrt_m_ladder: Array, JlB: Array, laguerre_to_grid: Array, laguerre_to_spectral: Array, laguerre_roots: Array, laguerre_j0: Array, laguerre_j1_over_alpha: Array, kx_link_plus: Array, kx_link_minus: Array, kx_link_mask_plus: Array, kx_link_mask_minus: Array, linked_inverse_permutation: Array = <factory>, linked_gather_map: Array = <factory>, linked_gather_mask: Array = <factory>, linked_full_cover: bool = False, linked_use_gather: bool = False, linked_indices: tuple[~jax.jaxlib._jax.Array, ...]=(), linked_kz: tuple[~jax.jaxlib._jax.Array, ...]=(), use_twist_shift: bool = False, jtwist: int = 0)

Precomputed arrays for the linear operator.

spectraxgk.linear_cache._build_end_damping_profile_array(Nz: int, widthfrac: float, boundary: str, real_dtype: dtype) → Array

Build the one-dimensional end-damping profile as one host array.

spectraxgk.linear_cache._build_gyroaverage_cache_arrays(b: Array, Nl: int, real_dtype: dtype) → tuple[Array, Array]

Build species-major gyroaverage factors without a Python-level vmap.

spectraxgk.linear_cache._build_low_rank_moment_cache_arrays(Nl: int, Nm: int, params: LinearParams, real_dtype: dtype) → dict[str, Array]

Build small moment-space cache arrays without many eager JAX dispatches.

spectraxgk.linear_cache.build_linear_cache(grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, Nl: int, Nm: int) → LinearCache

Build reusable arrays for the linear RHS.

spectraxgk.linear_cache.collision_damping(cache: LinearCache, params: LinearParams, real_dtype: dtype, *, squeeze_species: bool = False) → Array

Assemble collision damping from cached low-rank factors.

Runtime caches store lb_lam as the Hermite-Laguerre Lenard-Bernstein diagonal only, with shape (Nl, Nm). Older tests may still provide a pre-expanded array; support that for compatibility.

spectraxgk.linear_cache.hypercollision_damping(cache: LinearCache, params: LinearParams, real_dtype: dtype) → Array

Assemble GX-style hypercollision damping factors.

Linear Moments

Hermite-Laguerre moment and field-coupling primitives for linear operators.

spectraxgk.linear_moments.apply_hermite_v(G: Array) → Array

Multiply Hermite coefficients by v_parallel (ladder form).

spectraxgk.linear_moments.apply_hermite_v2(G: Array) → Array

Multiply Hermite coefficients by v_parallel^2.

spectraxgk.linear_moments.apply_laguerre_x(G: Array) → Array

Multiply Laguerre coefficients by the perpendicular energy variable.

spectraxgk.linear_moments.build_H(G: Array, Jl: Array, phi: Array, tz: Array, apar: Array | None = None, vth: Array | None = None, bpar: Array | None = None, JlB: Array | None = None) → Array

Map G -> H for mirror/curvature/grad-B/collision terms.

GX builds H by adding the field terms for m=0 (phi, Bpar) and the A_parallel term for m=1, while the streaming term applies its own pre-derivative field contributions. We mirror that behavior here.

spectraxgk.linear_moments.compute_b(grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, rho: float) → Array

Compute b = rho^2 * k_perp^2(kx, ky, theta) for s-alpha geometry.

spectraxgk.linear_moments.diamagnetic_drive_coeffs(Nl: int, Nm: int, eta_i: Array, coeff_const: float, coeff_par: float, coeff_perp: float) → Array

Return velocity-space coefficients for (1 + eta_i(E - 3/2)).

spectraxgk.linear_moments.energy_operator(G: Array, coeff_const: float, coeff_par: float, coeff_perp: float) → Array

Apply the energy operator (1 + v_par^2 + mu) in Hermite-Laguerre space.

spectraxgk.linear_moments.grad_z_periodic(f: Array, dz: float | Array) → Array

Spectral periodic derivative along the last axis.

spectraxgk.linear_moments.lenard_bernstein_eigenvalues(Nl: int, Nm: int, nu_hermite: float, nu_laguerre: float) → Array

Diagonal Lenard-Bernstein rates in Hermite-Laguerre space.

spectraxgk.linear_moments.quasineutrality_phi(G: Array, Jl: Array, tau_e: float | Array, charge: Array, density: Array, tz: Array) → Array

Solve electrostatic quasineutrality for phi with optional adiabatic closure.

spectraxgk.linear_moments.shift_axis(arr: Array, offset: int, axis: int) → Array

Shift an array along an axis with zero padding (non-periodic).

spectraxgk.linear_moments.streaming_term(H: Array, dz: float | Array, vth: float | Array) → Array

Streaming term using Hermite ladder and real-space z derivative.

Linear Parameters

Linear parameter, term-toggle, and validation policy helpers.

class spectraxgk.linear_params.LinearParams(charge_sign: float | Array = 1.0, density: float | Array = 1.0, mass: float | Array = 1.0, temp: float | Array = 1.0, tau_e: float = 1.0, vth: float | Array = 1.0, rho: float | Array = 1.0, kpar_scale: float = 1.0, R_over_Ln: float | Array = 2.2, R_over_LTi: float | Array = 6.9, R_over_LTe: float | Array = 0.0, omega_d_scale: float = 1.0, omega_star_scale: float = 1.0, energy_const: float = 0.0, energy_par_coef: float = 0.5, energy_perp_coef: float = 1.0, nu: float | Array = 0.0, nu_hermite: float = 1.0, nu_laguerre: float = 2.0, nu_hyper: float = 0.0, p_hyper: float = 4.0, nu_hyper_l: float = 0.0, nu_hyper_m: float = 1.0, nu_hyper_lm: float = 0.0, p_hyper_l: float = 6.0, p_hyper_m: float = 20.0, p_hyper_lm: float = 6.0, hypercollisions_const: float = 1.0, hypercollisions_kz: float = 0.0, D_hyper: float = 0.0, p_hyper_kperp: float = 2.0, damp_ends_widthfrac: float | Array = 0.125, damp_ends_amp: float | Array = 0.1, tz: float | Array = 1.0, rho_star: float = 1.0, beta: float = 0.0, fapar: float = 0.0, apar_beta_scale: float = 0.5, ampere_g0_scale: float = 0.5, bpar_beta_scale: float = 0.5)

Parameters for the linear gyrokinetic operator (supports multi-species arrays).

class spectraxgk.linear_params.LinearTerms(streaming: float = 1.0, mirror: float = 1.0, curvature: float = 1.0, gradb: float = 1.0, diamagnetic: float = 1.0, collisions: float = 1.0, hypercollisions: float = 1.0, hyperdiffusion: float = 0.0, end_damping: float = 1.0, apar: float = 1.0, bpar: float = 1.0)

Switches for linear-operator components (1.0 = on, 0.0 = off).

spectraxgk.linear_params._as_species_array(value: float | Array, ns: int, name: str) → Array

Ensure a parameter is a 1D array of length ns for multi-species handling.

spectraxgk.linear_params.linear_terms_to_term_config(terms: LinearTerms | None, *, nonlinear: float = 0.0) → TermConfig

Convert LinearTerms into the modular TermConfig.

spectraxgk.linear_params.term_config_to_linear_terms(term_cfg: TermConfig | None) → LinearTerms

Convert modular TermConfig into linear-only term weights.

Linear Parallel RHS

Velocity-parallel linear RHS helpers.

spectraxgk.linear_parallel._electrostatic_streaming_field_rhs(G6: Array, *, phi: Array, Jl: Array, tz: Array, vth: Array) → Array

Build the pre-derivative electrostatic streaming field term.

spectraxgk.linear_parallel._linear_rhs_electrostatic_slices_velocity_sharded_fused(arr: Array, cache: LinearCache, params: LinearParams, term_weights: LinearTerms, *, plan: Any, devices: Any, axis_name: str = 'm') → tuple[Array, Array]

Fuse the current single-species periodic electrostatic shard-map route.

spectraxgk.linear_parallel._resolve_parallel_devices(*, num_devices: int | None = None, devices: Any | None = None) → list[Any]

Return an explicit device list for opt-in parallel diagnostics.

spectraxgk.linear_parallel._streaming_electrostatic_from_phi_velocity_sharded(arr: Array, cache: LinearCache, params: LinearParams, *, phi: Array, plan: Any, devices: Any) → Array

Apply electrostatic streaming with a precomputed electrostatic field.

spectraxgk.linear_parallel.linear_rhs_electrostatic_slices_velocity_sharded(G: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | None = None, *, num_devices: int | None = None, devices: Any | None = None) → tuple[Array, Array]

Compute gated electrostatic streaming, drift, and diamagnetic slices.

spectraxgk.linear_parallel.linear_rhs_parallel_cached(G: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | None = None, *, parallel: Any | None = None, use_jit: bool = True, use_custom_vjp: bool = True, dt: Array | float | None = None) → tuple[Array, Array]

Compute linear RHS with an explicit, disabled-by-default parallel route.

parallel=None and parallel.strategy="serial" are exact aliases for linear_rhs_cached(). The non-serial velocity routes are opt-in, Hermite-axis-only identity gates. backend="auto" selects the most complete currently gated electrostatic route when the term set is eligible; otherwise callers must request a narrower explicit backend.

spectraxgk.linear_parallel.linear_rhs_streaming_electrostatic_velocity_sharded(G: Array, cache: LinearCache, params: LinearParams, *, num_devices: int | None = None, devices: Any | None = None, use_custom_vjp: bool = True) → tuple[Array, Array]

Compute electrostatic streaming RHS with Hermite-sharded particle streaming.

This route solves phi with the production electrostatic field solve, applies the Hermite velocity-sharded particle-streaming operator, and adds the GX-style electrostatic streaming field term. It is limited to periodic field-line grids and excludes electromagnetic fields by construction.

spectraxgk.linear_parallel.linear_rhs_streaming_velocity_sharded(G: Array, cache: LinearCache, params: LinearParams, *, num_devices: int | None = None, devices: Any | None = None) → tuple[Array, Array]

Compute the streaming-only linear RHS with the Hermite shard-map path.

This diagnostic route is intentionally narrower than linear_rhs_cached(): it covers the velocity-space streaming operator only and returns a zero electrostatic potential. It is used to gate the future production velocity decomposition before field solves, drifts, collisions, and nonlinear terms are exposed through the runtime path.

Nonlinear Diagnostics

Diagnostic packing and sampling helpers for nonlinear integrations.

spectraxgk.nonlinear_diagnostics._sample_indices_with_final(length: int, stride: int) → slice | ndarray

Return strided sample indices while always retaining the final step.

Nonlinear Helpers

Helper policies and operators for nonlinear gyrokinetic drivers.

class spectraxgk.nonlinear_helpers.IMEXLinearOperator(state_dtype: dtype, shape: tuple[int, ...], dt_val: Array, precond_op: Callable[[Array], Array] | None, matvec: Callable[[Array], Array], squeeze_species: bool)

Reusable matrix-free linear operator for nonlinear IMEX solves.

spectraxgk.nonlinear_helpers._apply_collision_split(G: Array, damping: Array, dt_local: Array, scheme: str) → Array

Apply a diagonal collision/hypercollision split update.

spectraxgk.nonlinear_helpers._collision_damping(cache: LinearCache, params: LinearParams, term_cfg: TermConfig, real_dtype: dtype, *, squeeze_species: bool) → Array

Assemble collision + hypercollision damping for operator splitting.

spectraxgk.nonlinear_helpers._gx_nonlinear_omega_components(fields: FieldState, grid: SpectralGrid, cache: LinearCache, *, gx_real_fft: bool, kx_max: float, ky_max: float, kxfac: float, vpar_max: float, muB_max: float) → tuple[Array, Array]

GX-style nonlinear x/y CFL frequency components from grad(phi,apar,bpar).

spectraxgk.nonlinear_helpers._gx_omega_mode_mask(grid: SpectralGrid, cache: LinearCache, *, gx_real_fft: bool) → Array

Mask used to reduce mode-wise GX omega/gamma diagnostics.

spectraxgk.nonlinear_helpers._make_fixed_mode_projector(fixed_state: Array | None, *, ky_index: int | None, kx_index: int | None) → Callable[[Array], Array] | None

Return a projector that keeps one Fourier mode equal to fixed_state.

spectraxgk.nonlinear_helpers._make_hermitian_projector(ky_vals: ndarray, nx: int) → Callable[[Array], Array]

Project full-ky states onto the GX real-FFT Hermitian manifold.

spectraxgk.nonlinear_helpers.build_nonlinear_imex_operator(G0: Array, cache: LinearCache, params: LinearParams, dt: float, *, terms: TermConfig | None = None, implicit_preconditioner: str | None = None, gx_real_fft: bool = True) → IMEXLinearOperator

Build and cache the matrix-free linear operator used by nonlinear IMEX.

Nonlinear Replicate Diagnostics

Diagnostics for replicated nonlinear transport-window spread.

The routines here are intentionally data-only. They consume already-generated ensemble JSON artifacts and classify why a replicated nonlinear window passed or failed before more GPU time is spent on follow-up runs.

class spectraxgk.nonlinear_replicate_diagnostics.NonlinearReplicateSpreadConfig(max_mean_rel_spread: float = 0.15, value_floor: float = 1e-12)

Thresholds for classifying replicated nonlinear-window spread.

spectraxgk.nonlinear_replicate_diagnostics.nonlinear_replicate_spread_report(ensembles: Sequence[Mapping[str, Any]], *, case: str = 'nonlinear_replicate_spread_diagnostic', config: NonlinearReplicateSpreadConfig | None = None) → dict[str, Any]

Classify which replicate/state drives nonlinear-window ensemble spread.

Parameters:

• ensembles – Sequence of ensemble JSON payloads, typically produced by tools/check_nonlinear_window_ensemble.py.

• case – Human-readable label for the diagnostic artifact.

• config – Spread threshold and numerical floor used for relative deviations.

Nonlinear Replicate Follow-Up

Follow-up planning for nonlinear replicate-spread blockers.

The planner consumes the replicate-spread diagnostic and already-known variant metadata, then selects the smallest cross checks needed to separate seed variability from timestep sensitivity.

class spectraxgk.nonlinear_replicate_followup.NonlinearReplicateFollowupConfig(include_extra_nominal_seed: bool = True, extra_seed_increment: int = 1, max_runs_per_state: int = 3)

Options controlling the targeted replicate follow-up plan.

spectraxgk.nonlinear_replicate_followup.nonlinear_replicate_followup_plan(spread_report: Mapping[str, Any], *, variant_metadata: Sequence[Mapping[str, Any]], case: str = 'nonlinear_replicate_followup_plan', config: NonlinearReplicateFollowupConfig | None = None) → dict[str, Any]

Return targeted cross-run follow-ups for failed replicate-spread states.

Nonlinear Gradient Follow-Up

Targeted follow-up plans for nonlinear turbulence-gradient audits.

This module turns failed long-window central finite-difference artifacts into a bounded run prescription. It is deliberately conservative: extra replicas are recommended only when the finite-difference response is resolved and local, but the propagated gradient uncertainty is slightly too large.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientCandidateDesignConfig(max_gradient_uncertainty_rel: float = 0.5, max_fd_asymmetry_rel: float = 0.5, max_window_mean_rel_spread: float = 0.15, max_window_sem_rel: float = 0.25, min_fd_response_fraction: float = 0.03, sem_safety_factor: float = 1.1, max_extra_replicates_per_state: int = 4, max_checked_bracket_scale: float = 1.5, locality_safety_factor: float = 0.95, value_floor: float = 1e-12)

Conditioning limits for the next nonlinear-gradient campaign design.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientCompositeControlConfig(max_gradient_uncertainty_rel: float = 1.0, max_fd_asymmetry_rel: float = 0.5, min_fd_response_fraction: float = 0.03, min_same_sign_fraction: float = 0.8, min_controls: int = 2, default_relative_delta: float = 0.02, max_weight_abs: float = 1.0, value_floor: float = 1e-12)

Controls for constructing the next composite nonlinear-gradient direction.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientControlMeanGateConfig(target_response_uncertainty_rel: float = 0.5, min_control_mean_pairs: int = 4, require_state_ensembles_passed: bool = True, value_floor: float = 1e-12)

Acceptance limits for an independent control-mean estimate.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientControlVariateCampaignConfig(target_response_uncertainty_rel: float = 0.5, sem_safety_factor: float = 1.1, min_control_mean_pairs: int = 4, max_control_mean_pairs: int = 32, first_new_seed: int = 34, value_floor: float = 1e-12)

Controls for an independent control-mean campaign.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientFollowupConfig(max_gradient_uncertainty_rel: float = 0.5, max_fd_asymmetry_rel: float = 0.5, min_fd_response_fraction: float = 0.03, sem_safety_factor: float = 1.1, max_extra_replicates_per_state: int = 4, default_nominal_timestep: float = 0.05, value_floor: float = 1e-12)

Acceptance and cost controls for the follow-up planner.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientQLSeedScreenConfig(target_objectives: tuple[str, ...] = ('mixing_length_heat_flux_proxy', 'linear_heat_flux_weight', 'gamma'), primary_objective: str = 'mixing_length_heat_flux_proxy', min_distinct_controls: int = 2, min_cases_per_control: int = 2, min_sign_consistency: float = 0.75, max_objective_rel_error: float = 0.02, min_abs_sensitivity: float = 1e-12, require_artifact_passed: bool = False)

Admission limits for quasilinear-seeded nonlinear-gradient controls.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientStateControlRunbookConfig(min_mapped_controls: int = 2, max_mapping_condition_number: float = 1000000.0, max_mapping_relative_residual: float = 0.1, default_relative_delta: float = 0.02, require_mapping_passed: bool = True)

Admission limits for mapping VMEC-state controls to launchable inputs.

class spectraxgk.nonlinear_gradient_followup.NonlinearGradientVarianceReductionConfig(max_paired_response_uncertainty_rel: float = 0.5, max_control_variate_uncertainty_rel: float = 0.5, min_control_variate_sem_reduction: float = 0.25, require_known_control_mean: bool = True, sem_safety_factor: float = 1.1, min_common_pairs: int = 2, max_extra_paired_seeds: int = 4, value_floor: float = 1e-12)

Controls for paired-seed/control-variate nonlinear-gradient planning.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_candidate_design_report(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_candidate_design', config: NonlinearGradientCandidateDesignConfig | None = None) → dict[str, Any]

Return the next-campaign design implied by failed production gates.

The report estimates whether a failed central-FD candidate can be rescued by a larger bracket, by bounded extra replicas, or whether the next campaign should instead change the control/observable. The estimates use the usual 1/sqrt(N) SEM scaling and the local finite-difference assumption that response grows approximately linearly with bracket size before the asymmetry gate is hit.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_composite_control_report(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_composite_control_design', config: NonlinearGradientCompositeControlConfig | None = None) → dict[str, Any]

Design a normalized VMEC-boundary direction from resolved FD candidates.

This is a launch-planning gate, not nonlinear-gradient evidence. The returned controls are the steepest-descent direction in the subspace of candidates that already pass locality, response, uncertainty, coefficient, and paired-sign checks. If fewer than min_controls survive, the report fails closed and provides exact blockers instead of producing a misleading multi-coefficient launch recommendation.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_control_mean_gate(variance_report: Mapping[str, Any], *, plus_ensemble: Mapping[str, Any], minus_ensemble: Mapping[str, Any], plus_path: str | None = None, minus_path: str | None = None, case: str = 'nonlinear_turbulence_gradient_control_mean_gate', candidate_name: str | None = None, config: NonlinearGradientControlMeanGateConfig | None = None) → dict[str, Any]

Evaluate an independent control-mean estimate for a screened CV response.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_control_variate_campaign_plan(variance_report: Mapping[str, Any], *, case: str = 'nonlinear_turbulence_gradient_control_variate_campaign', candidate_name: str | None = None, config: NonlinearGradientControlVariateCampaignConfig | None = None) → dict[str, Any]

Design an independent control-mean campaign from a variance runbook.

The input is the output of nonlinear_gradient_variance_reduction_plan(). The planner is intentionally fail-closed: a sample-centered control variate can motivate new runs, but the campaign is only considered launch-ready when an independent control-mean estimate can bring the combined uncertainty under the target gate within the configured run budget.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_followup_plan(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_followup', config: NonlinearGradientFollowupConfig | None = None) → dict[str, Any]

Build a bounded, fail-closed follow-up plan from gradient artifacts.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_ql_seed_screen_report(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_ql_seed_screen', config: NonlinearGradientQLSeedScreenConfig | None = None) → dict[str, Any]

Screen QL/linear sensitivity artifacts before nonlinear-gradient runs.

The report groups full-chain VMEC/Boozer sensitivity rows by state parameter and admits a control only when the primary objective sensitivity is AD/FD-consistent, resolved, sign-consistent across enough artifacts, and tied to a distinct VMEC-state control. The output is deliberately a planning artifact: VMEC-state controls are not assumed to be patchable RBC/ZBS input-file coefficients.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_state_control_runbook_report(ql_seed_screen: Mapping[str, Any], *, mapping_artifacts: Sequence[Mapping[str, Any]] = (), case: str = 'nonlinear_gradient_state_control_runbook', config: NonlinearGradientStateControlRunbookConfig | None = None) → dict[str, Any]

Build a fail-closed launch runbook for VMEC-state nonlinear-gradient controls.

The QL seed screen operates on internal vmec_jax state coordinates. Nonlinear campaigns, however, need perturbable input directions that can be written to VMEC inputs and re-equilibrated. This report joins the admitted state controls to an explicit state-to-input mapping artifact and refuses to produce launch commands until the mapping is conditioned and complete.

spectraxgk.nonlinear_gradient_followup.nonlinear_gradient_variance_reduction_plan(artifact: Mapping[str, Any], *, path: str | None = None, label: str | None = None, case: str = 'nonlinear_turbulence_gradient_variance_reduction_plan', config: NonlinearGradientVarianceReductionConfig | None = None) → dict[str, Any]

Plan paired-seed/control-variate follow-up for a failed central-FD artifact.

The plan uses common seed/timestep labels across plus and minus ensembles to estimate the uncertainty of paired finite-difference responses. It is a campaign-design artifact, not nonlinear-gradient evidence.

Benchmarks

Benchmark utilities for linear Cyclone base case comparisons.

class spectraxgk.benchmarks.CycloneComparison(ky: 'float', gamma: 'float', omega: 'float', gamma_ref: 'float', omega_ref: 'float', rel_gamma: 'float', rel_omega: 'float')

class spectraxgk.benchmarks.CycloneReference(ky: 'np.ndarray', omega: 'np.ndarray', gamma: 'np.ndarray')

class spectraxgk.benchmarks.CycloneRunResult(t: 'np.ndarray', phi_t: 'np.ndarray', gamma: 'float', omega: 'float', ky: 'float', selection: 'ModeSelection')

class spectraxgk.benchmarks.CycloneScanResult(ky: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray')

class spectraxgk.benchmarks.LinearRunResult(t: 'np.ndarray', phi_t: 'np.ndarray', gamma: 'float', omega: 'float', ky: 'float', selection: 'ModeSelection', gamma_t: 'np.ndarray | None' = None, omega_t: 'np.ndarray | None' = None)

class spectraxgk.benchmarks.LinearScanResult(ky: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray')

spectraxgk.benchmarks.compare_cyclone_to_reference(result: CycloneRunResult, reference: CycloneReference) → CycloneComparison

Compare a Cyclone run result against the reference data set.

spectraxgk.benchmarks.load_cyclone_reference() → CycloneReference

Load Cyclone base case reference data (adiabatic electrons).

spectraxgk.benchmarks.load_cyclone_reference_kinetic() → CycloneReference

Load Cyclone base case reference data (kinetic electrons).

spectraxgk.benchmarks.load_etg_reference() → CycloneReference

Load GX-backed ETG reference data for the tracked two-species ETG lane.

spectraxgk.benchmarks.load_kbm_reference() → CycloneReference

Load KBM reference data (finite beta, kinetic electrons).

spectraxgk.benchmarks.load_tem_reference() → CycloneReference

Load the provisional TEM reference digitized from the literature.

This lane is not backed by a GX reference dump. It remains an extended stress case while the literature case definition is being reconstructed.

spectraxgk.benchmarks.run_cyclone_linear(ky_target: float = 0.3, Nl: int = 6, Nm: int = 12, dt: float = 0.01, steps: int = 800, method: str = 'rk4', params: LinearParams | None = None, cfg: CycloneBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, max_fraction: float = 0.8, end_fraction: float = 1.0, max_amp_fraction: float = 1.0, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.0, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0, window_method: str = 'loglinear', mode_method: str = 'project', terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', init_cfg: InitializationConfig | None = None, diagnostic_norm: str = 'none', use_jit: bool = True, gx_reference: bool | None = None, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → CycloneRunResult

Run the linear Cyclone benchmark and extract growth rate.

spectraxgk.benchmarks.run_cyclone_scan(ky_values: ndarray, Nl: int = 6, Nm: int = 12, dt: float | ndarray = 0.01, steps: int | ndarray = 800, method: str = 'imex2', params: LinearParams | None = None, cfg: CycloneBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, max_fraction: float = 0.8, end_fraction: float = 1.0, max_amp_fraction: float = 1.0, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.0, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0, window_method: str = 'loglinear', mode_method: str = 'project', mode_only: bool = True, terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', diagnostic_norm: str = 'none', use_jit: bool = True, ky_batch: int = 4, fixed_batch_shape: bool = True, streaming_fit: bool = True, streaming_amp_floor: float = 1e-30, mode_follow: bool = True, gx_reference: bool | None = None, show_progress: bool = False) → CycloneScanResult

Run the linear Cyclone benchmark for a list of ky values.

If time_cfg is provided, its dt and t_max override dt/steps.

spectraxgk.benchmarks.run_etg_linear(ky_target: float = 3.0, Nl: int = 6, Nm: int = 12, dt: float = 0.01, steps: int = 800, method: str = 'rk4', params: LinearParams | None = None, cfg: ETGBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', streaming_fit: bool = False, streaming_amp_floor: float = 1e-30, gx_growth: bool = False, gx_navg_fraction: float = 0.5, diagnostic_norm: str = 'none', show_progress: bool = False) → LinearRunResult

Run an ETG linear benchmark and extract growth rate.

spectraxgk.benchmarks.run_etg_scan(ky_values: ndarray, Nl: int = 6, Nm: int = 12, dt: float | ndarray = 0.01, steps: int | ndarray = 800, method: str = 'imex2', params: LinearParams | None = None, cfg: ETGBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, max_fraction: float = 0.8, end_fraction: float = 1.0, max_amp_fraction: float = 1.0, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.0, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0, window_method: str = 'loglinear', mode_method: str = 'project', mode_only: bool = True, terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', ky_batch: int = 4, fixed_batch_shape: bool = True, streaming_fit: bool = True, streaming_amp_floor: float = 1e-30, gx_growth: bool = False, gx_navg_fraction: float = 0.5, diagnostic_norm: str = 'none', show_progress: bool = False) → LinearScanResult

Run an ETG linear benchmark for a list of ky values.

If time_cfg is provided, its dt and t_max override dt/steps.

spectraxgk.benchmarks.run_kbm_beta_scan(betas: ndarray, ky_target: float = 0.3, Nl: int = 6, Nm: int = 12, dt: float | ndarray = 0.01, steps: int | ndarray = 800, method: str = 'imex2', cfg: KBMBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, kbm_target_factors: Sequence[float] | None = (0.7, 1.5), kbm_beta_transition: float | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', mode_only: bool = True, terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', ky_batch: int = 4, fixed_batch_shape: bool = True, streaming_fit: bool = True, streaming_amp_floor: float = 1e-30, init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', fapar_override: float | None = None, apar_beta_scale: float | None = None, ampere_g0_scale: float | None = None, bpar_beta_scale: float | None = None, gx_reference: bool | None = True) → LinearScanResult

Run a KBM beta scan at fixed ky.

If time_cfg is provided, its dt and t_max override dt/steps.

spectraxgk.benchmarks.run_kbm_linear(ky_target: float = 0.3, *, beta_value: float | None = None, Nl: int = 6, Nm: int = 12, dt: float = 0.01, steps: int = 800, method: str = 'imex2', params: LinearParams | None = None, cfg: KBMBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, kbm_target_factors: Sequence[float] | None = (0.7, 1.5), kbm_beta_transition: float | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', streaming_fit: bool = False, init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', fapar_override: float | None = None, apar_beta_scale: float | None = None, ampere_g0_scale: float | None = None, bpar_beta_scale: float | None = None, gx_reference: bool | None = True, show_progress: bool = False) → LinearRunResult

Run a single linear KBM point and return the stored field history.

spectraxgk.benchmarks.run_kbm_scan(ky_values: ndarray, *, beta_value: float | None = None, Nl: int = 6, Nm: int = 12, dt: float | ndarray = 0.01, steps: int | ndarray = 800, method: str = 'imex2', cfg: KBMBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, kbm_target_factors: Sequence[float] | None = (0.7, 1.5), kbm_beta_transition: float | None = None, tmin: float | ndarray | None = None, tmax: float | ndarray | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', mode_only: bool = True, terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'auto', ky_batch: int = 4, fixed_batch_shape: bool = True, streaming_fit: bool = True, streaming_amp_floor: float = 1e-30, init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', fapar_override: float | None = None, apar_beta_scale: float | None = None, ampere_g0_scale: float | None = None, bpar_beta_scale: float | None = None, gx_reference: bool | None = True) → LinearScanResult

Run a KBM ky scan at fixed beta.

This is a thin wrapper over run_kbm_beta_scan() used for GX-reference workflows where the GX benchmark is a ky scan at fixed beta.

spectraxgk.benchmarks.run_kinetic_linear(ky_target: float = 0.3, Nl: int = 6, Nm: int = 12, dt: float = 0.01, steps: int = 800, method: str = 'rk4', params: LinearParams | None = None, cfg: KineticElectronBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'krylov', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'density', init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', gx_reference: bool | None = True, show_progress: bool = False) → LinearRunResult

Run a kinetic-electron ITG/TEM benchmark and extract growth rate.

spectraxgk.benchmarks.run_kinetic_scan(ky_values: ndarray, Nl: int = 6, Nm: int = 12, dt: float | ndarray = 0.01, steps: int | ndarray = 800, method: str = 'imex2', params: LinearParams | None = None, cfg: KineticElectronBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', mode_only: bool = True, terms: LinearTerms | None = None, sample_stride: int | None = None, fit_signal: str = 'density', ky_batch: int = 4, fixed_batch_shape: bool = True, streaming_fit: bool = True, streaming_amp_floor: float = 1e-30, init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', gx_reference: bool | None = True, show_progress: bool = False) → LinearScanResult

Run a kinetic-electron ITG/TEM benchmark for a list of ky values.

If time_cfg is provided, its dt and t_max override dt/steps.

spectraxgk.benchmarks.run_tem_linear(ky_target: float = 0.3, Nl: int = 6, Nm: int = 12, dt: float = 0.01, steps: int = 800, method: str = 'rk4', params: LinearParams | None = None, cfg: TEMBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'krylov', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', fit_signal: str = 'phi', terms: LinearTerms | None = None, sample_stride: int | None = None, init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', show_progress: bool = False) → LinearRunResult

Run the TEM benchmark and extract growth rate.

spectraxgk.benchmarks.run_tem_scan(ky_values: ndarray, Nl: int = 6, Nm: int = 12, dt: float | ndarray = 0.01, steps: int | ndarray = 800, method: str = 'imex2', params: LinearParams | None = None, cfg: TEMBaseCase | None = None, time_cfg: TimeConfig | None = None, solver: str = 'auto', krylov_cfg: KrylovConfig | None = None, tmin: float | None = None, tmax: float | None = None, auto_window: bool = True, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, mode_method: str = 'project', mode_only: bool = True, terms: LinearTerms | None = None, sample_stride: int | None = None, ky_batch: int = 4, fixed_batch_shape: bool = True, streaming_fit: bool = True, streaming_amp_floor: float = 1e-30, init_species_index: int = 1, density_species_index: int = 1, diagnostic_norm: str = 'none', show_progress: bool = False) → LinearScanResult

Run the TEM benchmark for a list of ky values.

If time_cfg is provided, its dt and t_max override dt/steps.

spectraxgk.benchmarks.select_kbm_solver_auto(solver: str, *, ky_target: float, gx_reference: bool) → str

Return deterministic KBM solver choice for auto mode.

Benchmark Defaults

Default normalization and Krylov policies for shipped benchmark lanes.

The runner functions in spectraxgk.benchmarks keep these names available for compatibility, but the constants live here so benchmark policy is separated from the long-running scan orchestration.

Benchmark Scan Policies

Shared scan policies for benchmark runners.

The public benchmark API stays in spectraxgk.benchmarks; this module keeps small, deterministic scan decisions separate from solver orchestration.

class spectraxgk.benchmark_scan.ScanFitWindowPolicy(tmin: ~typing.Any = None, tmax: ~typing.Any = None, auto_window: bool = True, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_fraction: float = 0.8, end_fraction: float = 0.9, max_amp_fraction: float = 0.9, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0, window_method: str = 'loglinear', fit_growth_rate_fn: ~collections.abc.Callable[[...], tuple[float, float]] = <function fit_growth_rate>, fit_growth_rate_auto_fn: ~collections.abc.Callable[[...], tuple[float, float, float, float]] = <function fit_growth_rate_auto>, normalize_growth_rate_fn: ~collections.abc.Callable[[float, float, ~spectraxgk.linear_params.LinearParams, str], tuple[float, float]] = <function _normalize_growth_rate>)

Window-selection and normalization policy shared by benchmark scans.

fit_growth_rate_auto_fn(signal: ndarray, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_amp_fraction: float = 0.9, window_method: str = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float]

Fit gamma/omega with optional auto-selected window.

fit_growth_rate_fn(signal: ndarray, tmin: float | None = None, tmax: float | None = None) → Tuple[float, float]

Fit gamma and omega from a complex signal ~ exp((gamma - i*omega) t).

fit_signal(signal: ndarray, *, idx: int, dt: float, stride: int, params: LinearParams, diagnostic_norm: str) → tuple[float, float]

Fit one scan signal and apply the configured diagnostic normalization.

spectraxgk.benchmark_scan.apply_auto_fit_scan_policy(fit_key: str, *, streaming_fit: bool, mode_only: bool) → tuple[bool, bool]

Disable streaming and mode-only saves when auto signal selection needs both fields.

spectraxgk.benchmark_scan.indexed_float_value(value: Any, idx: int) → float | None

Return a scalar or indexed scan value as float for window policies.

spectraxgk.benchmark_scan.indexed_scan_value(value: Any, idx: int) → Any

Return a scalar or indexed scan value while preserving non-float types.

spectraxgk.benchmark_scan.normalize_fit_signal(fit_signal: str) → str

Normalize and validate benchmark fit-signal selectors.

spectraxgk.benchmark_scan.normalize_solver_key(solver: str) → str

Normalize a benchmark solver selector without changing its semantics.

spectraxgk.benchmark_scan.resolve_scan_mode_method(mode_method: str, *, mode_only: bool) → str

Use direct mode extraction when a runner saved only a mode time series.

spectraxgk.benchmark_scan.scan_window_valid(t: ndarray, tmin: float | None, tmax: float | None, *, min_points: int = 2) → bool

Return whether an explicit fit window contains enough sampled points.

spectraxgk.benchmark_scan.should_use_ky_batch(*, ky_batch: int, solver_key: str, dt: Any, steps: Any, tmin: Any, tmax: Any) → bool

Return whether a ky scan can use a fixed-shape batch path.

Benchmarking

High-level benchmark helpers for scans and eigenfunction extraction.

class spectraxgk.benchmarking.ScanAndModeResult(scan: 'LinearScanResult', eigenfunction: 'np.ndarray', grid: 'SpectralGrid', ky_selected: 'float', tmin: 'float | None', tmax: 'float | None')

spectraxgk.benchmarking.branch_continuity_metrics(ky: ndarray, gamma: ndarray, omega: ndarray, *, successive_overlap: ndarray | None = None, floor_fraction: float = 1e-08) → BranchContinuationMetrics

Compute branch-continuity diagnostics for a linear scan.

The relative jump normalization uses a local scale from adjacent values, with a floor tied to the largest value in the scan. This avoids false blow-ups near marginal points while still flagging branch jumps.

spectraxgk.benchmarking.compare_eigenfunctions(eigenfunction: ndarray, reference: ndarray) → EigenfunctionComparisonMetrics

Return normalized overlap and relative L2 error after global phase alignment.

spectraxgk.benchmarking.estimate_observed_order(step_sizes: ndarray, errors: ndarray) → ObservedOrderMetrics

Estimate observed order from successive step-size refinements.

spectraxgk.benchmarking.infer_triple_dealiased_ny(nky_positive: int) → int

Infer the full Ny from the number of positive k_y points.

GX-style reference outputs typically store only the non-negative k_y branch. For the linked-boundary spectral grid used here, the corresponding real-space Ny follows Ny = 3 * (nky - 1) + 1.

spectraxgk.benchmarking.late_time_linear_metrics(result: object, *, tail_fraction: float = 0.5, mode_method: str = 'project') → LateTimeLinearMetrics

Return late-time growth/frequency metrics from a linear benchmark/runtime result.

spectraxgk.benchmarking.late_time_window(t: ndarray, *, tail_fraction: float = 0.4) → tuple[float, float]

Return the start/end of a late-time tail window.

This is the windowing convention used for manuscript-facing eigenfunction extraction when the growth-rate fit window is not the same object as the late-time mode-shape window.

spectraxgk.benchmarking.load_diagnostic_time_series(path: str | Path, *, variable: str, diagnostics_group: str = 'Diagnostics', time_group: str = 'Grids', time_var: str = 'time', kx_index: int | None = None, component: str = 'real', align_phase: bool = False) → DiagnosticTimeSeries

Load a 1D diagnostics time series from a GX-style out.nc artifact.

spectraxgk.benchmarking.load_eigenfunction_reference_bundle(path: str | Path) → EigenfunctionReferenceBundle

Load a frozen reference eigenfunction bundle.

spectraxgk.benchmarking.nonlinear_heat_flux_convergence_metrics(t: ndarray, heat_flux: ndarray, *, start_fraction: float = 0.5, terminal_fraction: float = 0.5, mean_floor: float = 1e-30) → NonlinearHeatFluxConvergenceMetrics

Summarize whether a post-transient heat-flux average is stable.

start_fraction discards startup samples. terminal_fraction compares the retained post-transient mean with the final subwindow of that retained region. The normalized trend is the least-squares slope multiplied by the post-transient time span and divided by the absolute post-transient mean.

spectraxgk.benchmarking.normalize_eigenfunction(eigenfunction: ndarray, z: ndarray) → ndarray

Normalize an eigenfunction by its value at theta=0 (nearest z=0).

spectraxgk.benchmarking.phase_align_eigenfunction(eigenfunction: ndarray, reference: ndarray) → tuple[ndarray, float]

Phase-align eigenfunction to reference using the global complex phase.

spectraxgk.benchmarking.run_linear_scan(*, ky_values: ndarray, run_linear_fn, cfg, Nl: int, Nm: int, dt: float | ndarray, steps: int | ndarray, method: str, solver: str, krylov_cfg, window_kw: dict, tmin: float | ndarray | None = None, tmax: float | ndarray | None = None, auto_window: bool = True, run_kwargs: dict | None = None, resolution_policy: Callable[[float], tuple[int, int]] | None = None, krylov_policy: Callable[[float], object] | None = None) → LinearScanResult

Run a linear scan over ky values.

spectraxgk.benchmarking.run_scan_and_mode(*, ky_values: ndarray, scan_fn, linear_fn, cfg, Nl: int, Nm: int, dt: float | ndarray, steps: int | ndarray, method: str, solver: str, mode_solver: str, krylov_cfg, window_kw: dict, tmin: float | ndarray | None = None, tmax: float | ndarray | None = None, auto_window: bool = True, run_kwargs: dict | None = None, mode_kwargs: dict | None = None, resolution_policy: Callable[[float], tuple[int, int]] | None = None, krylov_policy: Callable[[float], object] | None = None, select_ky: Callable[[LinearScanResult], float] | None = None) → ScanAndModeResult

Run a scan and extract a representative eigenfunction.

spectraxgk.benchmarking.save_eigenfunction_reference_bundle(path: str | Path, *, theta: ndarray, mode: ndarray, source: str, case: str, metadata: dict[str, object] | None = None) → Path

Write a frozen reference eigenfunction bundle as .npz.

spectraxgk.benchmarking.windowed_nonlinear_metrics(result: object, *, start_fraction: float = 0.5) → NonlinearWindowMetrics

Return late-window transport and envelope metrics from a nonlinear runtime result.

spectraxgk.benchmarking.zonal_flow_response_metrics(t: ndarray, response: ndarray, *, tail_fraction: float = 0.3, initial_fraction: float = 0.1, initial_policy: str = 'window_abs_mean', initial_level_override: float | None = None, peak_fit_max_peaks: int | None = None, damping_fit_mode: str = 'combined_envelope', frequency_fit_mode: str = 'peak_spacing', fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, hilbert_trim_fraction: float = 0.2) → ZonalFlowResponseMetrics

Estimate residual level and GAM envelope metrics from a zonal response.

The input response should be a scalar zonal observable such as zonal potential or a normalized zonal-energy proxy on a uniform time trace. initial_policy="first_abs" follows Rosenbluth-Hinton/GAM convention by normalizing to the initial potential magnitude; "window_abs_mean" keeps the older robust behavior for generic noisy traces. initial_level_override supports benchmarks whose published normalization is an external initial amplitude, for example a Gaussian potential maximum rather than the first line-averaged sample.

Validation Gates

Validation gate primitives for benchmark and manuscript artifacts.

class spectraxgk.validation_gates.BranchContinuationMetrics(ky: ndarray, gamma: ndarray, omega: ndarray, rel_gamma_jumps: ndarray, rel_omega_jumps: ndarray, max_rel_gamma_jump: float, max_rel_omega_jump: float, min_successive_overlap: float | None)

Continuity summary for a scanned linear branch.

class spectraxgk.validation_gates.DiagnosticTimeSeries(t: ndarray, values: ndarray, variable: str, source_path: str)

Single benchmark-facing time series loaded from an out.nc artifact.

class spectraxgk.validation_gates.EigenfunctionComparisonMetrics(overlap: float, relative_l2: float, phase_shift: float)

Phase-aligned eigenfunction comparison summary.

class spectraxgk.validation_gates.EigenfunctionReferenceBundle(theta: ndarray, mode: ndarray, source: str, case: str, metadata: dict[str, object])

Frozen reference eigenfunction bundle for manuscript-grade overlays.

class spectraxgk.validation_gates.GateReport(case: str, source: str, gates: tuple[ScalarGateResult, ...], passed: bool, max_abs_error: float, max_rel_error: float)

Collection of scalar gates for one validation artifact.

class spectraxgk.validation_gates.LateTimeLinearMetrics(gamma_fit: float, omega_fit: float, gamma_tail_mean: float, omega_tail_mean: float, gamma_tail_std: float, omega_tail_std: float, tmin: float | None, tmax: float | None, nsamples: int, signal_source: str)

Late-time growth/frequency metrics for a linear run.

class spectraxgk.validation_gates.NonlinearHeatFluxConvergenceMetrics(tmin: float, tmax: float, nsamples: int, heat_flux_mean: float, heat_flux_std: float, heat_flux_cv: float, heat_flux_rms: float, terminal_tmin: float, terminal_tmax: float, terminal_nsamples: int, terminal_heat_flux_mean: float, mean_rel_delta: float, trend: float, abs_trend: float, start_fraction: float, terminal_fraction: float)

Post-transient heat-flux averaging convergence summary.

class spectraxgk.validation_gates.NonlinearWindowMetrics(tmin: float, tmax: float, nsamples: int, heat_flux_mean: float, heat_flux_std: float, heat_flux_rms: float, wphi_mean: float, wphi_std: float, wg_mean: float, wg_std: float, phi_mode_envelope_mean: float | None, phi_mode_envelope_std: float | None, phi_mode_envelope_max: float | None)

Windowed transport/envelope metrics for a nonlinear run.

class spectraxgk.validation_gates.ObservedOrderMetrics(step_sizes: ndarray, errors: ndarray, orders: ndarray, asymptotic_order: float)

Observed-order convergence summary from step sizes and errors.

class spectraxgk.validation_gates.ScalarGateResult(metric: str, observed: float, reference: float, abs_error: float, rel_error: float, atol: float, rtol: float, passed: bool, units: str, notes: str)

Pass/fail result for one benchmark observable.

The tolerance convention follows numpy.isclose: a metric passes when abs_error <= atol + rtol * abs(reference). This keeps near-zero frequency and marginal-growth gates explicit through atol rather than hiding them behind unstable relative errors.

class spectraxgk.validation_gates.ZonalFlowResponseMetrics(initial_level: float, initial_policy: str, residual_level: float, residual_std: float, response_rms: float, gam_frequency: float, gam_damping_rate: float, damping_method: str, frequency_method: str, peak_count: int, peak_fit_count: int, tmin: float, tmax: float, fit_tmin: float, fit_tmax: float, peak_times: ndarray, peak_envelope: ndarray, max_peak_times: ndarray, max_peak_values: ndarray, min_peak_times: ndarray, min_peak_values: ndarray)

Late-time residual and GAM-envelope metrics for zonal-flow responses.

spectraxgk.validation_gates.branch_continuity_gate_report(metrics: BranchContinuationMetrics, *, case: str, source: str, max_rel_gamma_jump: float, max_rel_omega_jump: float, min_successive_overlap: float | None = None) → GateReport

Gate branch-continuation diagnostics for branch-followed scans.

spectraxgk.validation_gates.eigenfunction_gate_report(comparison: EigenfunctionComparisonMetrics, *, case: str, source: str, min_overlap: float = 0.95, max_relative_l2: float = 0.25) → GateReport

Gate a phase-aligned eigenfunction comparison.

The ideal reference is overlap equal to one and relative L2 mismatch equal to zero. min_overlap and max_relative_l2 make the acceptance policy explicit for manuscript overlays and branch-identity checks.

spectraxgk.validation_gates.evaluate_scalar_gate(metric: str, observed: float, reference: float, *, atol: float, rtol: float, units: str = '', notes: str = '') → ScalarGateResult

Evaluate one scalar benchmark gate.

Use this helper for publication-facing metrics such as growth rates, frequencies, windowed heat fluxes, zonal residuals, and damping rates. The explicit atol/rtol pair forces each artifact to document whether its tolerance is absolute, relative, or both.

spectraxgk.validation_gates.gate_report(case: str, source: str, gates: list[ScalarGateResult] | tuple[ScalarGateResult, ...]) → GateReport

Summarize a set of scalar gates for one artifact.

spectraxgk.validation_gates.gate_report_to_dict(report: GateReport) → dict[str, object]

Return a strict JSON-serializable representation of a gate report.

spectraxgk.validation_gates.linear_metrics_gate_report(observed: LateTimeLinearMetrics, reference: LateTimeLinearMetrics, *, case: str, source: str, gamma_atol: float = 0.0, gamma_rtol: float = 0.05, omega_atol: float = 0.0, omega_rtol: float = 0.05) → GateReport

Gate late-time linear growth and frequency metrics.

spectraxgk.validation_gates.nonlinear_heat_flux_convergence_gate_report(metrics: NonlinearHeatFluxConvergenceMetrics, *, case: str, source: str, max_mean_rel_delta: float = 0.05, max_cv: float = 0.15, max_abs_trend: float = 0.1, min_samples: int = 8) → GateReport

Gate post-transient heat-flux averaging stability.

This is an internal promotion gate for nonlinear transport claims: the post-transient average must agree with its terminal subwindow, have bounded coefficient of variation, show limited normalized drift across the window, and contain enough samples to be more than a reduced-window proxy.

spectraxgk.validation_gates.nonlinear_window_gate_report(observed: NonlinearWindowMetrics, reference: NonlinearWindowMetrics, *, case: str, source: str, rtol: float = 0.1, atol: float = 0.0, include_envelope: bool = True) → GateReport

Gate windowed nonlinear transport and field-energy metrics.

spectraxgk.validation_gates.observed_order_gate_report(metrics: ObservedOrderMetrics, *, case: str, source: str, min_asymptotic_order: float, min_pairwise_order: float | None = None, max_final_error: float | None = None, order_atol: float = 1e-12) → GateReport

Gate an observed-order convergence study.

min_asymptotic_order encodes the expected method/order floor for the finest refinement pair. min_pairwise_order can additionally require the whole table to be monotone enough for publication use. max_final_error can be used when both rate and absolute accuracy matter.

spectraxgk.validation_gates.zonal_response_gate_report(observed: ZonalFlowResponseMetrics, reference: ZonalFlowResponseMetrics, *, case: str, source: str, residual_atol: float, residual_rtol: float = 0.0, frequency_atol: float, frequency_rtol: float = 0.0, damping_atol: float, damping_rtol: float = 0.0) → GateReport

Gate Rosenbluth-Hinton/GAM-style response observables.

Autodiff Validation

Autodiff validation helpers for inverse and UQ examples.

spectraxgk.autodiff_validation.autodiff_finite_difference_report(fn: Callable[[Array], Any], params: Array | ndarray, *, step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, direction: Array | ndarray | None = None, workers: int = 1, parallel_executor: str = 'thread') → dict[str, object]

Compare JAX forward-mode derivatives against finite differences.

spectraxgk.autodiff_validation.central_finite_difference_jacobian(fn: Callable[[Array], Any], params: Array | ndarray, *, step: float = 0.0001, workers: int = 1, parallel_executor: str = 'thread') → Array

Central finite-difference Jacobian for small differentiability gates.

spectraxgk.autodiff_validation.covariance_diagnostics(jacobian: ndarray, residual: ndarray, *, regularization: float = 1e-09) → dict[str, object]

Return covariance and conditioning diagnostics for a least-squares inverse.

The covariance uses the local Gauss-Newton approximation sigma^2 (J^T J + lambda I)^-1. The returned dictionary is strict-JSON friendly and records enough metadata to decide whether an inverse demo is identifiable, merely locally sensitive, or ill-conditioned.

spectraxgk.autodiff_validation.explicit_complex_operator_matrix(operator: Callable[[Array], Any], state_shape: tuple[int, ...], *, dtype: Any | None = None) → Array

Materialize a small complex linear operator as a dense matrix.

This helper is intended for validation fixtures, not production solves. It applies operator to each basis vector of state_shape and returns a matrix whose columns are the flattened outputs. Small dense matrices make eigenvalue AD-vs-finite-difference gates easy to express while keeping the production code matrix-free.

spectraxgk.autodiff_validation.implicit_eigenpair_observable_sensitivity_report(matrix_fn: Callable[[Array], Any], observable_fn: Callable[[Array, Array, Array], Any], params: Array | ndarray, *, selector: str = 'max_real', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, gap_floor: float = 1e-08) → dict[str, object]

Validate implicit sensitivities of an isolated non-Hermitian eigenpair.

JAX currently supports first derivatives of non-Hermitian eigenvalues but not eigenvectors. This helper avoids differentiating through jnp.linalg.eig. It differentiates the matrix entries with JAX, solves the left/right eigenvector perturbation equation for each parameter, and compares the resulting observable Jacobian against nearest-branch central finite differences.

The observable should be phase-invariant under v -> exp(i alpha) v. The implicit solve fixes the gauge with w^H dv = 0 using the left eigenvector w normalized by w^H v = 1.

spectraxgk.autodiff_validation.isolated_eigenpair_observable_sensitivity_report(matrix_fn: Callable[[Array], Any], observable_fn: Callable[[Array, Array, Array], Any], params: Array | ndarray, *, selector: str = 'max_real', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, gap_floor: float = 1e-08) → dict[str, object]

Validate AD sensitivities of an observable of one isolated eigenpair.

observable_fn receives (eigenvalue, eigenvector, params) for the branch selected at the base point. The selected index is held fixed during finite differences, so this gate is appropriate for branch-continuous, phase-invariant quantities such as gamma / <k_perp^2>.

spectraxgk.autodiff_validation.isolated_eigenvalue_sensitivity_report(matrix_fn: Callable[[Array], Any], params: Array | ndarray, *, selector: str = 'max_real', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, gap_floor: float = 1e-08) → dict[str, object]

Validate AD sensitivities of one isolated eigenvalue branch.

The branch index is selected at the base point and then held fixed during the finite-difference comparison. This mirrors the branch-continuity assumption used for linear growth/frequency sensitivities.

Parallelization

Production parallelization helpers for independent scan and ensemble work.

class spectraxgk.parallel.IndependentEnsembleProvenanceReport(kind: str, workload: str, executor: str, requested_workers: int, actual_workers: int, problem_size: int, passed: bool, identity_passed: bool, ordering_passed: bool, worker_clipping_passed: bool, reconstruction_identity_passed: bool, exception_metadata_passed: bool, serial_indices: tuple[int, ...], parallel_indices: tuple[int, ...], reconstructed_indices: tuple[int, ...], identity_report: ~spectraxgk.parallel.ParallelIdentityReport, reconstruction_report: dict[str, ~typing.Any], exception_metadata: dict[str, ~typing.Any], metadata: dict[str, ~typing.Any] = <factory>)

End-to-end provenance gate for independent UQ/optimization ensembles.

to_dict() → dict[str, Any]

Return a JSON-serializable provenance payload.

exception spectraxgk.parallel.IndependentMapExecutionError(index: int, executor: str, actual_workers: int, original_type: str, original_message: str)

Worker failure annotated with independent-map execution metadata.

class spectraxgk.parallel.IndependentWorkerMetadata(requested_workers: int, actual_workers: int, problem_size: int, executor: str, parallel_enabled: bool)

Resolved worker metadata for ordered independent Python tasks.

to_dict() → dict[str, Any]

Return a JSON-serializable worker metadata payload.

class spectraxgk.parallel.ParallelIdentityReport(kind: str, backend: str, requested_workers: int, actual_workers: int, problem_size: int, identity_passed: bool, max_abs_error: float, max_rel_error: float, atol: float, rtol: float, metadata: dict[str, ~typing.Any]=<factory>)

Numerical-identity report for an independent parallel execution path.

to_dict() → dict[str, Any]

Return a JSON-serializable report for artifacts and CI gates.

spectraxgk.parallel.batch_map(fn: Callable[[Array], Any], values: Array | ndarray, *, batch_size: int | None = None, devices: Iterable[Device] | None = None) → Any

Map fn over independent inputs with optional multi-device batching.

This helper is intended for embarrassingly parallel physics workloads such as linear k_y scans, parameter sweeps, and UQ ensembles. It preserves numerical identity with jax.vmap(fn)(values) while allowing the leading batch axis to be distributed over available devices when more than one device is supplied.

spectraxgk.parallel.batch_map_identity_report(fn: Callable[[Array], Any], values: Array | ndarray, *, batch_size: int | None = None, devices: Iterable[Device] | None = None, atol: float = 1e-12, rtol: float = 1e-10) → ParallelIdentityReport

Compare batch_map against vmap and return a CI-ready gate report.

spectraxgk.parallel.independent_ensemble_provenance_gate(fn: Callable[[Any], Any], values: Iterable[Any], *, workers: int = 1, executor: str = 'thread', workload: str = 'uq_ensemble', atol: float = 1e-12, rtol: float = 1e-10, metadata: dict[str, Any] | None = None) → IndependentEnsembleProvenanceReport

Verify independent UQ/optimization ensemble batching provenance.

The gate intentionally runs fn serially and through independent_map. It verifies result identity, serial result ordering, worker clipping, deterministic shard reconstruction, and failure metadata for the same independent-map executor family.

spectraxgk.parallel.independent_map(fn: Callable[[Any], Any], values: Iterable[Any], *, workers: int = 1, executor: str = 'thread') → list[Any]

Map independent Python tasks while preserving serial result ordering.

batch_map handles JAX-array workloads. This helper covers file-backed calibration, finite-difference, and UQ tasks whose individual units are independent Python calls. The acceptance contract is numerical identity with [fn(value) for value in values]; timing is secondary.

spectraxgk.parallel.independent_map_identity_report(fn: Callable[[Any], Any], values: Iterable[Any], *, workers: int = 1, executor: str = 'thread', atol: float = 1e-12, rtol: float = 1e-10, metadata: dict[str, Any] | None = None) → ParallelIdentityReport

Compare independent_map against a serial list-comprehension run.

spectraxgk.parallel.independent_worker_metadata(problem_size: int, *, workers: int = 1, executor: str = 'thread') → IndependentWorkerMetadata

Resolve independent-task worker counts and normalized executor metadata.

spectraxgk.parallel.ky_scan_batches(ky_values: ndarray, *, n_batches: int) → list[ndarray]

Return balanced k_y chunks for independent linear-scan execution.

spectraxgk.parallel.pad_to_multiple(values: Array, multiple: int) → tuple[Array, int]

Pad axis zero by edge repetition so its length is divisible by multiple.

spectraxgk.parallel.parallel_identity_report(reference: Any, observed: Any, *, kind: str, problem_size: int, requested_workers: int, actual_workers: int | None = None, backend: str | None = None, atol: float = 1e-12, rtol: float = 1e-10, metadata: dict[str, Any] | None = None) → ParallelIdentityReport

Build a numerical-identity report for serial-vs-parallel outputs.

spectraxgk.parallel.split_evenly(values: ndarray, n_parts: int) → list[ndarray]

Split an array into nonempty, nearly equal chunks along axis zero.

Velocity Sharding Plans

Velocity-space decomposition plans for production parallelization.

class spectraxgk.velocity_sharding.VelocityShardingPlan(state_shape: tuple[int, ...], dims: tuple[str, ...], num_devices: int, chunks: dict[str, int], shard_shape: tuple[int, ...], active_axes: tuple[str, ...], hermite_ghost_depth: int, needs_hermite_exchange: bool, needs_field_reduction: bool, field_reduction_axes: tuple[str, ...], communication_pattern: str, load_balance: float)

JSON-friendly plan for decomposing a packed GK state over devices.

spectraxgk.velocity_sharding.build_velocity_sharding_plan(state_shape: Sequence[int], *, num_devices: int, axes: Sequence[str] | None = None, hermite_ghost_depth: int = 1) → VelocityShardingPlan

Build a GX-inspired species/Hermite velocity-space decomposition plan.

The plan is metadata only. It does not move arrays or claim speedup. It records which axes should be split, where Hermite ghost exchange is needed, and which velocity axes require field-solve reductions/broadcasts before a production shard_map implementation is allowed to use the layout.

spectraxgk.velocity_sharding.curvature_gradb_drift_reference(H: Any, *, tz: Any, omega_d_scale: Any, cv_d: Any, gb_d: Any, ell: Any, m: Any, weight_curv: Any = 1.0, weight_gradb: Any = 1.0) → Any

Return curvature and grad-B drift contributions with full-array shifts.

spectraxgk.velocity_sharding.curvature_gradb_drift_shard_map(H: Any, plan: VelocityShardingPlan, *, tz: Any, omega_d_scale: Any, cv_d: Any, gb_d: Any, ell: Any, m: Any, weight_curv: Any = 1.0, weight_gradb: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return curvature and grad-B drift contributions using Hermite exchange.

spectraxgk.velocity_sharding.diamagnetic_drive_reference(state: Any, *, phi: Any, Jl: Any, l4: Any, tprim: Any, fprim: Any, omega_star_scale: Any, ky: Any, weight: Any = 1.0) → Any

Return the single-species electrostatic diamagnetic drive.

spectraxgk.velocity_sharding.diamagnetic_drive_shard_map(state: Any, plan: VelocityShardingPlan, *, phi: Any, Jl: Any, l4: Any, tprim: Any, fprim: Any, omega_star_scale: Any, ky: Any, weight: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return the diamagnetic drive through a Hermite-sharded local map.

spectraxgk.velocity_sharding.electrostatic_phi_reference(state: Any, *, Jl: Any, tau_e: Any, charge: Any = 1.0, density: Any = 1.0, tz: Any = 1.0, mask0: Any | None = None) → Any

Return single-species electrostatic phi from a full 5D state.

spectraxgk.velocity_sharding.electrostatic_phi_shard_map(state: Any, plan: VelocityShardingPlan, *, Jl: Any, tau_e: Any, charge: Any = 1.0, density: Any = 1.0, tz: Any = 1.0, mask0: Any | None = None, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Solve electrostatic phi using a Hermite-sharded density reduction.

spectraxgk.velocity_sharding.hermite_neighbor_reference(state: Any) → tuple[Any, Any]

Return full-array lower/upper Hermite-neighbor states.

The Hermite streaming ladder couples moment m to m-1 and m+1. Physical boundaries outside [0, Nm-1] are zeros. The returned arrays have the same shape as state and provide the lower and upper neighbor values for every Hermite index.

spectraxgk.velocity_sharding.hermite_neighbor_shard_map(state: Any, plan: VelocityShardingPlan, *, devices: Sequence[Any] | None = None, axis_name: str = 'm') → tuple[Any, Any]

Exchange nearest Hermite neighbors with jax.shard_map.

This is a communication-kernel identity primitive, not a production nonlinear solver path. It currently supports one-dimensional Hermite decomposition plans. More complex species-Hermite meshes should first add a separate field-reduction and broadcast gate.

spectraxgk.velocity_sharding.hermite_shift_reference(state: Any, *, offset: int) → Any

Shift a state along the Hermite axis with zero physical boundaries.

spectraxgk.velocity_sharding.hermite_shift_shard_map(state: Any, plan: VelocityShardingPlan, *, offset: int, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Shift a Hermite-sharded state by offset moments with shard exchange.

spectraxgk.velocity_sharding.hermite_streaming_ladder_reference(state: Any, *, vth: Any = 1.0) → Any

Return the full-array Hermite streaming ladder contribution.

spectraxgk.velocity_sharding.hermite_streaming_ladder_shard_map(state: Any, plan: VelocityShardingPlan, *, vth: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return a shard-map Hermite streaming ladder contribution.

spectraxgk.velocity_sharding.mirror_drift_reference(H: Any, *, vth: Any, bgrad: Any, ell: Any, sqrt_m: Any, sqrt_m_p1: Any, weight: Any = 1.0) → Any

Return the mirror-drift contribution with full-array Hermite shifts.

spectraxgk.velocity_sharding.mirror_drift_shard_map(H: Any, plan: VelocityShardingPlan, *, vth: Any, bgrad: Any, ell: Any, sqrt_m: Any, sqrt_m_p1: Any, weight: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return the mirror-drift contribution using Hermite shard exchange.

spectraxgk.velocity_sharding.periodic_streaming_reference(state: Any, *, kz: Any, vth: Any = 1.0) → Any

Return periodic parallel streaming using full-array operations.

spectraxgk.velocity_sharding.periodic_streaming_shard_map(state: Any, plan: VelocityShardingPlan, *, kz: Any, vth: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return periodic parallel streaming through the Hermite shard-map path.

spectraxgk.velocity_sharding.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

Return the full-array velocity-axis reduction used by field solves.

spectraxgk.velocity_sharding.velocity_field_reduce_shard_map(state: Any, plan: VelocityShardingPlan, *, axis: str = 'm', devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Reduce one velocity axis across a shard-map mesh and broadcast it.

Sharded Integrators

Sharded fixed-step integrators for multi-device scaling experiments.

spectraxgk.sharded_integrators.integrate_linear_sharded(G0: Array, cache: LinearCache, params: LinearParams, *, dt: float, steps: int, terms: LinearTerms | None = None, state_sharding: Any | None = None) → Array

Integrate the linear system with a pjit-sharded RK2 loop.

This is intentionally minimal: it returns the final state only and avoids saving time histories to focus on strong scaling of the RHS.

spectraxgk.sharded_integrators.integrate_nonlinear_sharded(G0: Array, cache: LinearCache, params: LinearParams, *, dt: float, steps: int, method: str = 'rk2', terms: TermConfig | None = None, state_sharding: Any | None = None, gx_real_fft: bool = True, laguerre_mode: str = 'grid', return_fields: bool = True) → tuple[Array, FieldState] | Array

Integrate the nonlinear system with an explicit pjit-sharded scan.

The state array can be partitioned along a resolve_state_sharding axis such as ky or kx. This is a diagnostic whole-state sharding primitive for identity gates and profiler localization. It is not a production nonlinear domain decomposition or speedup claim until the exact workload has communication-complete identity, conservation, transport, and profiler gates. Domain-sharding identity reports are metadata gates only; they do not authorize routing through this whole-state integrator.

Zonal Validation

Zonal Flow Objectives

Reduced zonal-flow objectives for differentiable stellarator optimization.

This module is intentionally backend-free. Production callers should build zonal-response metrics from VMEC/Boozer/SPECTRAX-GK rows, then pass the metric tensors here for reduction, finite-difference checks, and UQ diagnostics.

class spectraxgk.zonal_objective.ZonalFlowObjectiveConfig(residual_weight: float = 1.0, damping_weight: float = 1.0, growth_over_residual_weight: float = 0.0, recurrence_weight: float = 0.0, residual_floor: float = 1e-06)

Weights and floors for a minimizable zonal-flow objective.

The objective rewards large residual zonal response by minimizing 1 / residual and penalizes collisionless damping, linear growth not screened by the residual, and late-time recurrence/envelope amplitude. Nonlinear heat-flux suppression remains a separate holdout gate.

objective_weights() → Array

Return the normalized objective-column weights used by the reducer.

to_dict() → dict[str, object]

Return a JSON-friendly representation.

spectraxgk.zonal_objective.zonal_flow_objective_artifact_from_records(records: Iterable[Mapping[str, Any]], *, config: ZonalFlowObjectiveConfig | None = None, surface_keys: Sequence[str] = ('surface', 'surface_index', 'torflux'), alpha_keys: Sequence[str] = ('alpha', 'field_line_label'), kx_keys: Sequence[str] = ('kx', 'kx_target', 'kx_rhoi'), residual_keys: Sequence[str] = ('residual_level', 'spectrax_residual'), damping_keys: Sequence[str] = ('damping_rate', 'gam_damping_rate'), linear_growth_keys: Sequence[str] = ('linear_growth_rate', 'growth_rate', 'gamma'), recurrence_keys: Sequence[str] = ('recurrence_amplitude', 'tail_std_ratio', 'residual_std', 'tail_std'), missing_damping_policy: Literal['fail', 'zero'] = 'fail', claim_level: str | None = None, source_paths: Sequence[str] | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean') → dict[str, object]

Build a strict JSON-friendly zonal-flow objective artifact.

The input is a table of validated zonal-response metrics. Rows are mapped onto the shared (surface, alpha, kx) portfolio tensor used by the stellarator objective stack. Missing damping rates fail by default because a promoted zonal-flow optimization claim must know the damping convention. Diagnostic artifacts can set missing_damping_policy='zero' to produce rows while carrying an explicit promotion_ready=False flag.

spectraxgk.zonal_objective.zonal_flow_objective_rows(*, residual_level: Any, damping_rate: Any, linear_growth_rate: Any | None = None, recurrence_amplitude: Any | None = None, config: ZonalFlowObjectiveConfig | None = None) → Array

Return objective rows with shape (surface, alpha, kx, objective).

residual_level is the late-time residual normalized to the initial zonal potential. Larger residuals reduce the first objective column. damping_rate should be positive for decaying GAM/zonal envelopes. linear_growth_rate is optional and encodes a suppression-relevance metric: high ITG growth with weak residuals is penalized. The recurrence column should be a non-negative late-envelope or moment-tail amplitude.

spectraxgk.zonal_objective.zonal_flow_objective_sensitivity_report(metric_fn: Callable[[Array], Mapping[str, Any]], params: Any, *, config: ZonalFlowObjectiveConfig | None = None, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, min_rank: int | None = None, condition_number_limit: float = 100000000.0, covariance_regularization: float = 1e-09, workers: int = 1, parallel_executor: str = 'thread') → dict[str, object]

AD/FD, row-Jacobian, and UQ gate for a zonal-flow optimization map.

spectraxgk.zonal_objective.zonal_flow_reduced_objective(*, residual_level: Any, damping_rate: Any, linear_growth_rate: Any | None = None, recurrence_amplitude: Any | None = None, config: ZonalFlowObjectiveConfig | None = None, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean') → Array

Reduce zonal-flow metric tensors to one differentiable scalar objective.

Analysis

Analysis helpers for extracting growth rates and mode signals.

class spectraxgk.analysis.ModeSelection(ky_index: 'int', kx_index: 'int', z_index: 'int' = 0)

class spectraxgk.analysis.ModeSelectionBatch(ky_indices: 'np.ndarray', kx_index: 'int', z_index: 'int' = 0)

spectraxgk.analysis.density_moment(G: ndarray, Jl: ndarray, *, species_index: int | None = None) → ndarray

Compute the m=0 density moment for a selected species (or summed if None).

spectraxgk.analysis.extract_eigenfunction(phi_t: ndarray, t: ndarray, sel: ModeSelection, z: ndarray | None = None, method: str = 'svd', tmin: float | None = None, tmax: float | None = None) → ndarray

Extract a normalized eigenfunction in z from phi_t(t, ky, kx, z).

spectraxgk.analysis.extract_mode(phi_t: ndarray, sel: ModeSelection) → ndarray

Extract a complex mode time series from phi_t(t, ky, kx, z).

spectraxgk.analysis.extract_mode_time_series(phi_t: ndarray, sel: ModeSelection, method: str = 'z_index') → ndarray

Extract a complex mode time series from phi_t(t, ky, kx, z).

spectraxgk.analysis.fit_growth_rate(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None) → Tuple[float, float]

Fit gamma and omega from a complex signal ~ exp((gamma - i*omega) t).

spectraxgk.analysis.fit_growth_rate_auto(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_amp_fraction: float = 0.9, window_method: str = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float]

Fit gamma/omega with optional auto-selected window.

spectraxgk.analysis.fit_growth_rate_auto_with_stats(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_amp_fraction: float = 0.9, window_method: str = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float, float, float]

Fit gamma/omega and report selected window plus R^2 scores.

spectraxgk.analysis.fit_growth_rate_with_stats(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None) → Tuple[float, float, float, float]

Fit gamma/omega and return (gamma, omega, r2_log_amp, r2_phase).

spectraxgk.analysis.gx_growth_rate_from_omega_series(gamma_t: ndarray, omega_t: ndarray, sel: ModeSelection, *, navg_fraction: float = 0.5, use_last: bool = False) → Tuple[float, float, ndarray, ndarray]

Compute GX-style averaged growth rates from precomputed omega_kxky(t).

Parameters:

• gamma_t – Arrays with shape (t, ky, kx).

• omega_t – Arrays with shape (t, ky, kx).

• sel – Mode selection used to choose the (ky, kx) series.

• navg_fraction – Fractional start index for late-time averaging (GX-style).

• use_last – If true, use the last finite sample instead of late-time average.

spectraxgk.analysis.gx_growth_rate_from_phi(phi_t: ndarray, t: ndarray | None, sel: ModeSelection, *, navg_fraction: float = 0.5, use_last: bool = False, mode_method: str = 'z_index') → Tuple[float, float, ndarray, ndarray, ndarray]

Compute GX-style instantaneous growth rates from phi ratios.

Returns (gamma_avg, omega_avg, gamma_t, omega_t, t_mid).

spectraxgk.analysis.select_fit_window(t: ndarray, signal: ndarray, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0) → Tuple[float, float]

Pick a time window with the most exponential-like behavior.

spectraxgk.analysis.select_fit_window_loglinear(t: ndarray, signal: ndarray, min_points: int = 20, start_fraction: float = 0.0, max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_amp_fraction: float = 0.9, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float]

Select a window where log-amplitude is closest to linear.

spectraxgk.analysis.select_ky_index(ky: ndarray, ky_target: float) → int

Return the best ky index for a requested target.

For nonzero requests, prefer a nonzonal mode with the closest absolute magnitude, then prefer a sign match when one exists. This avoids collapsing sparse signed grids such as [0, -k] onto the zonal row when the user requests +k.

Plotting

Plotting utilities for publication-ready figures.

class spectraxgk.plotting.LinearValidationPanel(name: 'str', z: 'np.ndarray', eigenfunction: 'np.ndarray', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', x_label: 'str', x_ref: 'np.ndarray | None' = None, gamma_ref: 'np.ndarray | None' = None, omega_ref: 'np.ndarray | None' = None, ref_label: 'str' = 'Reference', log_x: 'bool' = False)

class spectraxgk.plotting.MultiReferenceValidationPanel(name: 'str', z: 'np.ndarray', eigenfunction: 'np.ndarray', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', x_label: 'str', references: 'list[ReferenceSeries]', log_x: 'bool' = False)

class spectraxgk.plotting.ReferenceSeries(label: 'str', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', color: 'str', marker: 'str' = 'o', linestyle: 'str' = '--')

spectraxgk.plotting.cyclone_comparison_figure(ref: CycloneReference, scan: CycloneScanResult, label: str = 'SPECTRAX-GK') → Tuple[Figure, ndarray]

Create a two-panel comparison plot between reference and solver output.

spectraxgk.plotting.cyclone_reference_figure(ref: CycloneReference) → Tuple[Figure, ndarray]

Create a two-panel Cyclone base case reference plot.

spectraxgk.plotting.eigenfunction_overlap_summary_figure(ky: ndarray, overlap: ndarray, relative_l2: ndarray, *, title: str = 'Eigenfunction overlap summary', x_label: str = '$k_y \\rho_i$', overlap_label: str = 'Normalized overlap', rel_l2_label: str = 'Relative $L^2$ error', log_x: bool = True) → Tuple[Figure, ndarray]

Render a compact two-panel eigenfunction-overlap summary.

spectraxgk.plotting.eigenfunction_reference_overlay_figure(theta: ndarray, eigenfunction: ndarray, theta_ref: ndarray, reference: ndarray, *, title: str = 'Eigenfunction overlay') → Tuple[Figure, ndarray]

Render a phase-aligned raw overlay against a frozen reference mode.

spectraxgk.plotting.etg_trend_figure(R_over_LTe: ndarray, gamma: ndarray, omega: ndarray, ky_target: float) → Tuple[Figure, ndarray]

Create a two-panel ETG trend plot versus R/LTe.

spectraxgk.plotting.growth_fit_figure(t: ndarray, signal: ndarray, *, tmin: float | None = None, tmax: float | None = None, title: str = 'Growth-fit window') → Tuple[Figure, ndarray]

Plot \(|s|^2\) and \(\log |s|^2\) with an optional fit window.

spectraxgk.plotting.growth_rate_heatmap(x: ndarray, y: ndarray, gamma: ndarray, title: str, x_label: str, y_label: str, cmap: str = 'jet') → Tuple[Figure, Axes]

Render a growth-rate heatmap versus two gradient axes.

spectraxgk.plotting.linear_runtime_panel_figure(*, t: ndarray, signal: ndarray, z: ndarray, eigenfunction: ndarray, gamma: float, omega: float, title: str = 'SPECTRAX-GK Linear Runtime') → Tuple[Figure, ndarray]

Create the default two-panel linear runtime plot.

spectraxgk.plotting.linear_validation_figure(panels: list[LinearValidationPanel]) → Tuple[Figure, ndarray]

Create a multi-panel summary plot of eigenfunctions, growth rates, and frequencies.

spectraxgk.plotting.linear_validation_multi_reference_figure(panels: list[MultiReferenceValidationPanel]) → Tuple[Figure, ndarray]

Create summary panels with multiple external reference curves.

spectraxgk.plotting.nonlinear_runtime_panel_figure(*, t: ndarray, phi2: ndarray | None = None, wphi: ndarray | None = None, heat_flux: ndarray | None = None, gamma: ndarray | None = None, omega: ndarray | None = None, title: str = 'SPECTRAX-GK Nonlinear Runtime') → Tuple[Figure, ndarray]

Create the default three-panel nonlinear runtime plot.

spectraxgk.plotting.plot_saved_output(path: str | Path, *, out: str | Path | None = None) → Path

Plot a saved linear or nonlinear runtime artifact bundle.

spectraxgk.plotting.scan_comparison_figure(x: ndarray, gamma: ndarray, omega: ndarray, x_label: str, title: str, x_ref: ndarray | None = None, gamma_ref: ndarray | None = None, omega_ref: ndarray | None = None, label: str = 'SPECTRAX-GK', ref_label: str = 'Reference', log_x: bool = False) → Tuple[Figure, ndarray]

Create a two-panel comparison plot for a generic scan.

spectraxgk.plotting.scan_multi_reference_figure(x: ndarray, gamma: ndarray, omega: ndarray, x_label: str, title: str, references: list[ReferenceSeries], *, log_x: bool = False) → Tuple[Figure, ndarray]

Create a two-panel comparison figure against multiple reference curves.

spectraxgk.plotting.set_plot_style() → None

Apply a consistent plotting style suitable for publications.

spectraxgk.plotting.zonal_flow_response_figure(t: ndarray, response: ndarray, *, metrics=None, title: str = 'Zonal-flow response', y_label: str = 'normalized response') → Tuple[Figure, ndarray]

Render a zonal-flow response trace and its envelope summary.

Config

class spectraxgk.config.CycloneBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=24, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=20.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=150.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, implicit_solve_method='batched', use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=True, diffrax_rtol=1e-06, diffrax_atol=1e-08, diffrax_max_steps=200000, state_sharding=None, progress_bar=False, fixed_dt=False, dt_min=1e-07, dt_max=0.05, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', gx_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, gx_python=None, rhoc=0.5, R_geo=None, shift=0.0, akappa=1.0, akappri=0.0, tri=0.0, tripri=0.0, torflux=None, npol=None, npol_min=None, isaxisym=False, which_crossing=None, include_shear_variation=False, include_pressure_variation=False, betaprim=None, gx_repo=None, q=1.4, s_hat=0.8, z0=None, zero_shat=False, epsilon=0.18, R0=2.77778, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), model: ModelConfig = ModelConfig(R_over_LTi=2.49, R_over_LTe=0.0, R_over_Ln=0.8, nu_i=0.0), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False), gx_reference: bool = True)

Standard parameters for the Cyclone base case ITG benchmark.

class spectraxgk.config.ETGBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=24, Nz=96, Lx=6.28, Ly=6.28, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=0.2, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=10.0, dt=0.05, method='rk2', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, implicit_solve_method='batched', use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=True, diffrax_rtol=1e-05, diffrax_atol=1e-07, diffrax_max_steps=200000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', gx_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, gx_python=None, rhoc=0.5, R_geo=None, shift=0.0, akappa=1.0, akappri=0.0, tri=0.0, tripri=0.0, torflux=None, npol=None, npol_min=None, isaxisym=False, which_crossing=None, include_shear_variation=False, include_pressure_variation=False, betaprim=None, gx_repo=None, q=1.4, s_hat=0.8, z0=None, zero_shat=False, epsilon=0.18, R0=2.77778, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), model: ETGModelConfig = ETGModelConfig(R_over_LTi=2.49, R_over_LTe=2.49, R_over_Ln=0.8, R_over_Lni=None, R_over_Lne=None, Te_over_Ti=1.0, mass_ratio=3703.7037037037035, nu_i=0.0, nu_e=0.0, beta=1e-05, adiabatic_ions=True), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for a reduced ETG linear benchmark.

class spectraxgk.config.ETGModelConfig(R_over_LTi: float = 2.49, R_over_LTe: float = 2.49, R_over_Ln: float = 0.8, R_over_Lni: float | None = None, R_over_Lne: float | None = None, Te_over_Ti: float = 1.0, mass_ratio: float = 3703.7037037037035, nu_i: float = 0.0, nu_e: float = 0.0, beta: float = 1e-05, adiabatic_ions: bool = True)

Dimensionless gradients and ratios for a canonical ETG setup.

class spectraxgk.config.GeometryConfig(model: str = 's-alpha', geometry_backend: str = 'auto', geometry_file: str | None = None, vmec_file: str | None = None, gx_python: str | None = None, rhoc: float = 0.5, R_geo: float | None = None, shift: float = 0.0, akappa: float = 1.0, akappri: float = 0.0, tri: float = 0.0, tripri: float = 0.0, torflux: float | None = None, npol: float | None = None, npol_min: float | None = None, isaxisym: bool = False, which_crossing: int | None = None, include_shear_variation: bool = False, include_pressure_variation: bool = False, betaprim: float | None = None, gx_repo: str | None = None, q: float = 1.4, s_hat: float = 0.8, z0: float | None = None, zero_shat: bool = False, epsilon: float = 0.18, R0: float = 1.0, B0: float = 1.0, alpha: float = 0.0, drift_scale: float = 1.0, kperp2_bmag: bool = True, bessel_bmag_power: float = 0.0)

Flux-tube geometry parameters or imported sampled geometry settings.

class spectraxgk.config.GridConfig(Nx: int = 48, Ny: int = 48, Nz: int = 64, Lx: float = 62.8, Ly: float = 62.8, boundary: str = 'periodic', jtwist: int | None = None, non_twist: bool = False, kxfac: float = 1.0, z_min: float = -3.141592653589793, z_max: float = 3.141592653589793, y0: float | None = None, ntheta: int | None = None, nperiod: int | None = None, zp: int | None = None)

Spectral grid configuration in a flux-tube.

class spectraxgk.config.InitializationConfig(init_field: str = 'density', init_amp: float = 1e-05, init_single: bool = True, random_seed: int = 22, gaussian_init: bool = False, gaussian_width: float = 0.5, gaussian_envelope_constant: float = 1.0, gaussian_envelope_sine: float = 0.0, kpar_init: float = 0.0, init_file: str | None = None, init_file_scale: float = 1.0, init_file_mode: str = 'replace', init_electrons_only: bool = False)

Initialization options for linear runs.

class spectraxgk.config.KBMBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=16, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=10.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, implicit_solve_method='batched', use_diffrax=True, diffrax_solver='Tsit5', diffrax_adaptive=True, diffrax_rtol=0.0001, diffrax_atol=1e-07, diffrax_max_steps=20000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', gx_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, gx_python=None, rhoc=0.5, R_geo=None, shift=0.0, akappa=1.0, akappri=0.0, tri=0.0, tripri=0.0, torflux=None, npol=None, npol_min=None, isaxisym=False, which_crossing=None, include_shear_variation=False, include_pressure_variation=False, betaprim=None, gx_repo=None, q=1.4, s_hat=0.8, z0=None, zero_shat=False, epsilon=0.18, R0=2.77778, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), model: KineticElectronModelConfig = KineticElectronModelConfig(R_over_LTi=2.49, R_over_LTe=2.49, R_over_Ln=0.8, Te_over_Ti=1.0, mass_ratio=3703.7037037037035, nu_i=0.0, nu_e=0.0, beta=0.015), init: InitializationConfig = InitializationConfig(init_field='all', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for an electromagnetic KBM benchmark.

class spectraxgk.config.KineticElectronBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=16, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=10.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, implicit_solve_method='batched', use_diffrax=True, diffrax_solver='Tsit5', diffrax_adaptive=True, diffrax_rtol=0.0001, diffrax_atol=1e-07, diffrax_max_steps=20000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', gx_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, gx_python=None, rhoc=0.5, R_geo=None, shift=0.0, akappa=1.0, akappri=0.0, tri=0.0, tripri=0.0, torflux=None, npol=None, npol_min=None, isaxisym=False, which_crossing=None, include_shear_variation=False, include_pressure_variation=False, betaprim=None, gx_repo=None, q=1.4, s_hat=0.8, z0=None, zero_shat=False, epsilon=0.18, R0=2.77778, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), model: KineticElectronModelConfig = KineticElectronModelConfig(R_over_LTi=2.49, R_over_LTe=2.49, R_over_Ln=0.8, Te_over_Ti=1.0, mass_ratio=3703.7037037037035, nu_i=0.0, nu_e=0.0, beta=1e-05), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for kinetic-electron Cyclone benchmarks.

class spectraxgk.config.KineticElectronModelConfig(R_over_LTi: float = 2.49, R_over_LTe: float = 2.49, R_over_Ln: float = 0.8, Te_over_Ti: float = 1.0, mass_ratio: float = 3703.7037037037035, nu_i: float = 0.0, nu_e: float = 0.0, beta: float = 1e-05)

Gradients and ratios for a kinetic-electron Cyclone-base-case setup.

class spectraxgk.config.ModelConfig(R_over_LTi: float = 2.49, R_over_LTe: float = 0.0, R_over_Ln: float = 0.8, nu_i: float = 0.0)

Dimensionless gradients for the Cyclone base case.

class spectraxgk.config.TEMBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=24, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=20.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=8.0, dt=0.01, method='rk2', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, implicit_solve_method='batched', use_diffrax=True, diffrax_solver='Tsit5', diffrax_adaptive=True, diffrax_rtol=0.0001, diffrax_atol=1e-07, diffrax_max_steps=20000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', gx_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, gx_python=None, rhoc=0.5, R_geo=None, shift=0.0, akappa=1.0, akappri=0.0, tri=0.0, tripri=0.0, torflux=None, npol=None, npol_min=None, isaxisym=False, which_crossing=None, include_shear_variation=False, include_pressure_variation=False, betaprim=None, gx_repo=None, q=2.7, s_hat=0.5, z0=None, zero_shat=False, epsilon=0.18, R0=1.0, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), model: TEMModelConfig = TEMModelConfig(R_over_LTi=20.0, R_over_LTe=20.0, R_over_Ln=20.0, Te_over_Ti=1.0, mass_ratio=370.0, nu_i=0.0, nu_e=0.0, beta=0.0001), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for the provisional literature-backed TEM stress case.

class spectraxgk.config.TEMModelConfig(R_over_LTi: float = 20.0, R_over_LTe: float = 20.0, R_over_Ln: float = 20.0, Te_over_Ti: float = 1.0, mass_ratio: float = 370.0, nu_i: float = 0.0, nu_e: float = 0.0, beta: float = 0.0001)

Parameters for a trapped-electron-mode benchmark.

class spectraxgk.config.TimeConfig(t_max: float = 100.0, dt: float = 0.1, method: str = 'rk2', sample_stride: int = 1, diagnostics_stride: int = 1, diagnostics: bool = True, save_state: bool = False, checkpoint: bool = False, implicit_restart: int = 20, implicit_preconditioner: str | None = None, implicit_solve_method: str = 'batched', use_diffrax: bool = True, diffrax_solver: str = 'Dopri8', diffrax_adaptive: bool = False, diffrax_rtol: float = 1e-05, diffrax_atol: float = 1e-07, diffrax_max_steps: int = 4096, state_sharding: str | None = None, progress_bar: bool = False, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, nstep_restart: int | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', gx_real_fft: bool = True, nonlinear_dealias: bool = True, laguerre_nonlinear_mode: str = 'grid')

Time integration parameters.

spectraxgk.config.explicit_method_default_cfl_fac(method: str) → float

Return the reference explicit-method CFL prefactor for a given method.

spectraxgk.config.gx_default_cfl_fac(method: str) → float

Return the reference explicit-method CFL prefactor for a given method.

spectraxgk.config.resolve_cfl_fac(method: str, cfl_fac: float | None) → float

Resolve an explicit CFL prefactor, falling back to the method default.

Normalization

Canonical normalization contracts for benchmark families.

This module centralizes the case-level calibration knobs that were previously spread across benchmark runners and tool scripts.

class spectraxgk.normalization.NormalizationContract(case: str, rho_star: float, omega_d_scale: float, omega_star_scale: float, diagnostic_norm_default: Literal['none', 'gx', 'rho_star'] = 'none')

Case-level normalization parameters.

case

Canonical case key (e.g. "cyclone").

Type:

str

rho_star

Multiplier applied to kx/ky in drift and drive terms.

Type:

float

omega_d_scale

Curvature / grad-B / mirror scaling.

Type:

float

omega_star_scale

Diamagnetic drive scaling.

Type:

float

diagnostic_norm_default

Default post-processing normalization for reported (gamma, omega).

Type:

Literal[‘none’, ‘gx’, ‘rho_star’]

spectraxgk.normalization.apply_diagnostic_normalization(gamma: float, omega: float, *, rho_star: float, diagnostic_norm: str) → tuple[float, float]

Apply reporting-space normalization to growth rates/frequencies.

spectraxgk.normalization.get_normalization_contract(case: str) → NormalizationContract

Return the canonical normalization contract for case.

Runtime Config

Unified runtime configuration schema for linear/nonlinear GK runs.

class spectraxgk.runtime_config.RuntimeCollisionConfig(nu_hermite: float = 1.0, nu_laguerre: float = 2.0, nu_hyper: float = 0.0, p_hyper: float = 4.0, nu_hyper_l: float = 0.0, nu_hyper_m: float = 1.0, nu_hyper_lm: float = 0.0, p_hyper_l: float = 6.0, p_hyper_m: float | None = None, p_hyper_lm: float = 6.0, D_hyper: float = 0.0, p_hyper_kperp: float = 2.0, hypercollisions_const: float = 0.0, hypercollisions_kz: float = 1.0, damp_ends_amp: float = 0.1, damp_ends_widthfrac: float = 0.125, damp_ends_scale_by_dt: bool = False)

Collision and end-damping parameters.

class spectraxgk.runtime_config.RuntimeConfig(grid: GridConfig = GridConfig(Nx=48, Ny=48, Nz=64, Lx=62.8, Ly=62.8, boundary='periodic', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=None, ntheta=None, nperiod=None, zp=None), time: TimeConfig = TimeConfig(t_max=100.0, dt=0.1, method='rk2', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, implicit_solve_method='batched', use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=False, diffrax_rtol=1e-05, diffrax_atol=1e-07, diffrax_max_steps=4096, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', gx_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, gx_python=None, rhoc=0.5, R_geo=None, shift=0.0, akappa=1.0, akappri=0.0, tri=0.0, tripri=0.0, torflux=None, npol=None, npol_min=None, isaxisym=False, which_crossing=None, include_shear_variation=False, include_pressure_variation=False, betaprim=None, gx_repo=None, q=1.4, s_hat=0.8, z0=None, zero_shat=False, epsilon=0.18, R0=1.0, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-05, init_single=True, random_seed=22, gaussian_init=False, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False), species: Tuple[RuntimeSpeciesConfig, ...] = (RuntimeSpeciesConfig(name='ion', charge=1.0, mass=1.0, density=1.0, temperature=1.0, tprim=2.49, fprim=0.8, nu=0.0, kinetic=True),), physics: RuntimePhysicsConfig = RuntimePhysicsConfig(reduced_model='gyrokinetic', linear=True, nonlinear=False, electrostatic=True, electromagnetic=False, use_apar=False, use_bpar=False, adiabatic_electrons=True, adiabatic_ions=False, tau_e=1.0, tau_fac=None, z_ion=1.0, beta=0.0, collisions=True, hypercollisions=True), collisions: RuntimeCollisionConfig = RuntimeCollisionConfig(nu_hermite=1.0, nu_laguerre=2.0, nu_hyper=0.0, p_hyper=4.0, nu_hyper_l=0.0, nu_hyper_m=1.0, nu_hyper_lm=0.0, p_hyper_l=6.0, p_hyper_m=None, p_hyper_lm=6.0, D_hyper=0.0, p_hyper_kperp=2.0, hypercollisions_const=0.0, hypercollisions_kz=1.0, damp_ends_amp=0.1, damp_ends_widthfrac=0.125, damp_ends_scale_by_dt=False), normalization: RuntimeNormalizationConfig = RuntimeNormalizationConfig(contract='cyclone', rho_star=None, omega_d_scale=None, omega_star_scale=None, diagnostic_norm='gx', flux_scale=1.0, wphi_scale=1.0), terms: RuntimeTermsConfig = RuntimeTermsConfig(streaming=1.0, mirror=1.0, curvature=1.0, gradb=1.0, diamagnetic=1.0, collisions=1.0, hypercollisions=1.0, hyperdiffusion=0.0, end_damping=1.0, apar=1.0, bpar=1.0, nonlinear=0.0), expert: RuntimeExpertConfig = RuntimeExpertConfig(fixed_mode=False, iky_fixed=None, ikx_fixed=None, dealias_kz=False, source='default', phi_ext=0.0), output: RuntimeOutputConfig = RuntimeOutputConfig(path=None, restart=False, restart_if_exists=False, save_for_restart=True, restart_to_file=None, restart_from_file=None, restart_with_perturb=False, append_on_restart=True, restart_scale=1.0, nsave=10000), quasilinear: RuntimeQuasilinearConfig = RuntimeQuasilinearConfig(enabled=False, mode='weights', saturation_rule='none', amplitude_normalization='phi_rms', kperp_average='phi_weighted', csat=1.0, gamma_floor=0.0, include_stable_modes=False, delta_ky='auto', species='all', channels=('es',), write_spectrum=True, output_path=None), parallel: RuntimeParallelConfig = RuntimeParallelConfig(strategy='serial', axis='ky', batch_size=None, num_devices=None, strict_identity=True, profile=False, backend='auto'))

Unified simulation config for runtime-driven GK runs.

class spectraxgk.runtime_config.RuntimeExpertConfig(fixed_mode: bool = False, iky_fixed: int | None = None, ikx_fixed: int | None = None, dealias_kz: bool = False, source: str = 'default', phi_ext: float = 0.0)

Advanced runtime controls that should rarely be needed.

class spectraxgk.runtime_config.RuntimeNormalizationConfig(contract: str = 'cyclone', rho_star: float | None = None, omega_d_scale: float | None = None, omega_star_scale: float | None = None, diagnostic_norm: str = 'gx', flux_scale: float = 1.0, wphi_scale: float = 1.0)

Normalization contract selection + optional explicit overrides.

class spectraxgk.runtime_config.RuntimeOutputConfig(path: str | None = None, restart: bool = False, restart_if_exists: bool = False, save_for_restart: bool = True, restart_to_file: str | None = None, restart_from_file: str | None = None, restart_with_perturb: bool = False, append_on_restart: bool = True, restart_scale: float = 1.0, nsave: int = 10000)

Artifact-output controls for runtime executable entry points.

class spectraxgk.runtime_config.RuntimeParallelConfig(strategy: str = 'serial', axis: str = 'ky', batch_size: int | None = None, num_devices: int | None = None, strict_identity: bool = True, profile: bool = False, backend: str = 'auto')

Parallel-execution policy for independent scans and future sharded paths.

class spectraxgk.runtime_config.RuntimePhysicsConfig(reduced_model: str = 'gyrokinetic', linear: bool = True, nonlinear: bool = False, electrostatic: bool = True, electromagnetic: bool = False, use_apar: bool = False, use_bpar: bool = False, adiabatic_electrons: bool = True, adiabatic_ions: bool = False, tau_e: float = 1.0, tau_fac: float | None = None, z_ion: float = 1.0, beta: float = 0.0, collisions: bool = True, hypercollisions: bool = True)

Physics-family toggles independent from benchmark case names.

class spectraxgk.runtime_config.RuntimeQuasilinearConfig(enabled: bool = False, mode: str = 'weights', saturation_rule: str = 'none', amplitude_normalization: str = 'phi_rms', kperp_average: str = 'phi_weighted', csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False, delta_ky: str | float = 'auto', species: str = 'all', channels: Tuple[str, ...] = ('es',), write_spectrum: bool = True, output_path: str | None = None)

Quasilinear transport diagnostics computed from linear states.

class spectraxgk.runtime_config.RuntimeSpeciesConfig(name: str = 'ion', charge: float = 1.0, mass: float = 1.0, density: float = 1.0, temperature: float = 1.0, tprim: float = 2.49, fprim: float = 0.8, nu: float = 0.0, kinetic: bool = True)

Single species definition for runtime-configured simulations.

class spectraxgk.runtime_config.RuntimeTermsConfig(streaming: float = 1.0, mirror: float = 1.0, curvature: float = 1.0, gradb: float = 1.0, diamagnetic: float = 1.0, collisions: float = 1.0, hypercollisions: float = 1.0, hyperdiffusion: float = 0.0, end_damping: float = 1.0, apar: float = 1.0, bpar: float = 1.0, nonlinear: float = 0.0)

Term toggles for assembly; applies to linear and nonlinear paths.

Runtime Policies

Pure runtime policy helpers shared by runtime runners and tests.

class spectraxgk.runtime_policies.RuntimeIndependentParallelPlan(requested_workers: int, effective_workers: int, executor: str, strategy: str, axis: str, source: str, problem_size: int)

Resolved independent-worker policy for runtime scan workloads.

property enabled: bool

Whether the resolved plan uses more than one independent worker.

to_dict() → dict[str, Any]

Return a JSON-friendly policy payload for runtime artifacts.

spectraxgk.runtime_policies._infer_runtime_nonlinear_steps(cfg: RuntimeConfig, *, dt: float, steps: int | None) → int

Infer nonlinear explicit step counts with the same dt ceiling as the integrator.

spectraxgk.runtime_policies._parallel_requests_combined_ky_scan(cfg: RuntimeConfig) → bool

Return whether runtime parallel config requests the combined-ky scan path.

spectraxgk.runtime_policies._runtime_external_phi(cfg: RuntimeConfig) → float | None

Return a GX-style runtime external-phi source if requested.

spectraxgk.runtime_policies._runtime_independent_parallel_plan(cfg: RuntimeConfig, *, problem_size: int, workers: int, executor: str) → RuntimeIndependentParallelPlan

Resolve independent k_y worker policy from arguments and config.

Runtime Orchestration

Runtime orchestration helpers split from public runtime entry points.

This module owns coordination policy that is not itself a solver kernel: progress/ETA formatting, combined-ky scan batching, and nonlinear artifact restart/checkpoint handoff. Callers pass dependency tables so legacy spectraxgk.runtime and spectraxgk.runtime_artifacts monkeypatch seams remain effective.

class spectraxgk.runtime_orchestration.NonlinearArtifactPolicy(out_path: Path | None, gx_target: bool, diagnostics_on: bool, restart_from: Path | None, restart_to: Path | None, resume_requested: bool, remaining_steps: int | None, checkpoint_steps: int | None)

Resolved nonlinear artifact/restart policy for a single handoff.

class spectraxgk.runtime_orchestration.RuntimeArtifactHandoffDeps(is_gx_netcdf_target: Callable[[Path], bool], resolve_restart_path: Callable[[str | Path, Any], Path], resolve_restart_write_path: Callable[[str | Path, Any], Path], gx_bundle_base: Callable[[Path], Path], load_runtime_nonlinear_gx_diagnostics: Callable[[str | Path], SimulationDiagnostics], condense_gx_diagnostics_for_output: Callable[[SimulationDiagnostics], SimulationDiagnostics], concat_gx_diagnostics: Callable[[list[SimulationDiagnostics]], SimulationDiagnostics], validate_finite_runtime_result: Callable[[Any], None], run_runtime_nonlinear: Callable[[...], RuntimeNonlinearResult], write_runtime_nonlinear_artifacts: Callable[[str | Path, Any, Any], dict[str, str]])

Patchable functions used by nonlinear artifact handoff orchestration.

class spectraxgk.runtime_orchestration.RuntimeProgressSnapshot(progress: float, eta_seconds: float, chunk_wall_seconds: float, elapsed_seconds: float)

Computed wall-clock progress fields for a chunked runtime update.

class spectraxgk.runtime_orchestration.RuntimeScanBatchDeps(*args, **kwargs)

Dependency surface needed by the combined-ky scan batch helper.

spectraxgk.runtime_orchestration.build_runtime_progress_message(*, label: str, chunk_index: int, t_elapsed: float, t_max: float, chunk_wall_seconds: float, elapsed_seconds: float) → tuple[str, RuntimeProgressSnapshot]

Return the standard adaptive-runtime progress line and policy snapshot.

spectraxgk.runtime_orchestration.format_duration(seconds: float) → str

Format elapsed seconds as MM:SS or H:MM:SS.

spectraxgk.runtime_orchestration.resolve_nonlinear_artifact_policy(cfg: Any, *, out: str | Path | None, diagnostics: bool | None, steps: int | None, dt: float | None, deps: RuntimeArtifactHandoffDeps) → NonlinearArtifactPolicy

Resolve nonlinear output, restart, and checkpoint policy.

spectraxgk.runtime_orchestration.run_runtime_nonlinear_artifact_handoff(cfg: Any, *, out: str | Path | None, ky_target: float, kx_target: float | None = None, Nl: int | None = None, Nm: int | None = None, dt: float | None = None, steps: int | None = None, method: str | None = None, sample_stride: int | None = None, diagnostics_stride: int | None = None, laguerre_mode: str | None = None, diagnostics: bool | None = None, show_progress: bool = False, status_callback: Any = None, deps: RuntimeArtifactHandoffDeps) → tuple[RuntimeNonlinearResult, dict[str, str]]

Run nonlinear runtime chunks and hand results to artifact writers.

spectraxgk.runtime_orchestration.run_runtime_scan_batch(cfg: RuntimeConfig, ky_arr: ndarray, *, Nl: int, Nm: int, method: str | None, dt: float | None, steps: int | None, sample_stride: int | None, auto_window: bool, tmin: float | None, tmax: float | None, window_fraction: float, min_points: int, start_fraction: float, growth_weight: float, require_positive: bool, min_amp_fraction: float, mode_method: str, fit_signal: str, show_progress: bool, deps: RuntimeScanBatchDeps) → RuntimeLinearScanResult

Batch a ky scan using one time integration over the full grid.

Quasilinear Transport

Quasilinear transport diagnostics from linear gyrokinetic states.

The routines in this module compute linear heat and particle flux weights from an eigenstate or late-time linear state. Saturation rules are kept explicitly separate from the linear weights so calibration and uncertainty metadata can be audited case by case.

class spectraxgk.quasilinear.QuasilinearTransportResult(ky: float, gamma: float, omega: float, mode: str, saturation_rule: str, amplitude_normalization: str, channels: tuple[str, ...], kperp_average: str, kperp_eff2: float, phi_norm2: float, amplitude2: float | None, heat_flux_weight_species: tuple[float, ...], particle_flux_weight_species: tuple[float, ...], saturated_heat_flux_species: tuple[float, ...] | None, saturated_particle_flux_species: tuple[float, ...] | None, species: tuple[str, ...], metadata: dict[str, Any])

JSON-friendly quasilinear diagnostic payload for one linear mode.

to_dict() → dict[str, Any]

Return a stable JSON-serializable representation.

spectraxgk.quasilinear.compute_quasilinear_from_linear_state(state: Array | ndarray, *, cache: LinearCache, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, ky: float, gamma: float, omega: float, terms: TermConfig | None = None, mode: str = 'weights', saturation_rule: str = 'none', amplitude_normalization: str = 'phi_rms', kperp_average: str = 'phi_weighted', csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False, channels: Sequence[str] | str = ('es',), species_names: Sequence[str] | None = None, use_dealias: bool = True, flux_scale: float = 1.0, metadata: dict[str, Any] | None = None) → QuasilinearTransportResult

Compute quasilinear transport weights from a linear state.

The returned heat and particle flux weights are divided by the selected mode-amplitude normalization, so they are invariant under complex phase rotations and real amplitude rescalings of the eigenstate.

spectraxgk.quasilinear.effective_kperp2(phi: Array, cache: LinearCache, vol_fac: Array, *, use_dealias: bool = True, eps: float = 1e-30) → Array

Compute <k_perp^2 |phi|^2>/<|phi|^2> for a linear mode.

spectraxgk.quasilinear.mixing_length_amplitude2_jax(gamma: Array | float, kperp_eff2_value: Array | float, *, csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False, eps: float = 1e-30) → Array

JAX-differentiable mixing-length squared-amplitude rule.

spectraxgk.quasilinear.normalize_quasilinear_channels(channels: Iterable[str] | str) → tuple[str, ...]

Normalize and validate quasilinear field channels.

spectraxgk.quasilinear.phi_norm2(phi: Array, cache: LinearCache, params: LinearParams, vol_fac: Array, *, normalization: str = 'phi_rms', use_dealias: bool = True, eps: float = 1e-30) → Array

Return the amplitude normalization used for quasilinear weights.

spectraxgk.quasilinear.quasilinear_feature_objective(features: Array | Sequence[float], *, rule: str = 'mixing_length', csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False) → Array

Differentiable objective from [gamma, kperp_eff2, flux_weight].

This helper is intentionally small: it is the reduced objective used by derivative validation tests and optimization examples once a linear scan has produced quasilinear weights.

spectraxgk.quasilinear.saturated_flux_from_linear_weight(linear_flux_weight: Array | float, gamma: Array | float, kperp_eff2_value: Array | float, *, csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False) → Array

Return a differentiable mixing-length saturated flux estimate.

spectraxgk.quasilinear.saturation_amplitude2(*, gamma: float, kperp_eff2_value: float, rule: str, csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False) → float | None

Return the squared amplitude implied by a named saturation rule.

spectraxgk.quasilinear.shape_aware_power_law_objective(features: Array | Sequence[float], ky: Array | Sequence[float] | float, *, exponent: Array | float, csat: float = 1.0, ky_ref: float | None = None, eps: float = 1e-30) → Array

Differentiable shape-aware linear-weight objective.

features must end with [gamma, kperp_eff2, flux_weight]. The current low-dimensional shape model intentionally uses only the linear heat-flux weight and a power-law envelope in ky:

Q = C_sat * flux_weight * (ky / ky_ref)**exponent.

Growth-rate dependence is left to separately validated rules. This helper exists so the shape-aware saturation diagnostics and future optimization examples use one differentiable objective rather than plotting-only formulas.

spectraxgk.quasilinear.spectral_phi_weights(phi: Array, cache: LinearCache, vol_fac: Array, *, use_dealias: bool = True) → Array

Return (ky, kx, z) weights used for |phi|^2 averages.

Quasilinear Calibration

Calibration artifact helpers for quasilinear transport models.

class spectraxgk.quasilinear_calibration.QuasilinearCalibrationPoint(case: str, split: str, predicted_heat_flux: float, observed_heat_flux: float, saturation_rule: str, raw_predicted_heat_flux: float | None = None, calibration_scale: float | None = None, geometry: str = 'unspecified', electron_model: str = 'unspecified', ky: float | None = None, observed_heat_flux_std: float | None = None, nonlinear_window_stats: dict[str, Any] | None = None, quasilinear_artifact: str | None = None, nonlinear_artifact: str | None = None, notes: str | None = None)

One quasilinear-vs-nonlinear transport comparison point.

spectraxgk.quasilinear_calibration.apply_heat_flux_scale(points: Iterable[QuasilinearCalibrationPoint | dict[str, Any]], *, scale: float, note_label: str = 'heat_flux_scale') → list[QuasilinearCalibrationPoint]

Return calibration points with heat-flux predictions multiplied by scale.

spectraxgk.quasilinear_calibration.calibration_point_from_nonlinear_window_summary(summary_json: str | Path, *, predicted_heat_flux: float, split: str, saturation_rule: str, diagnostics_source: str = 'spectrax', heat_flux_column: str = 'heat_flux', case: str | None = None, geometry: str = 'unspecified', electron_model: str = 'unspecified', quasilinear_artifact: str | None = None, species_index: int | None = None, window_convergence_config: NonlinearWindowConvergenceConfig | None = None, notes: str | None = None) → QuasilinearCalibrationPoint

Create a calibration point from a nonlinear window-summary JSON.

The helper reads the window bounds from a tracked nonlinear gate summary and computes the mean/std of a heat-flux column from the selected diagnostics CSV or runtime NetCDF. For NetCDF inputs, heat_flux_column='heat_flux' maps to Diagnostics/HeatFlux_st and species are summed by default.

spectraxgk.quasilinear_calibration.calibration_point_from_spectrum_and_nonlinear_window(spectrum_csv: str | Path, summary_json: str | Path, *, split: str, saturation_rule: str, spectrum_column: str = 'saturated_heat_flux_total', spectrum_method: str = 'sum', delta_ky: float | None = None, diagnostics_source: str = 'spectrax', heat_flux_column: str = 'heat_flux', case: str | None = None, geometry: str = 'unspecified', electron_model: str = 'unspecified', species_index: int | None = None, window_convergence_config: NonlinearWindowConvergenceConfig | None = None, notes: str | None = None) → QuasilinearCalibrationPoint

Create a calibration point from a quasilinear spectrum and nonlinear window.

spectraxgk.quasilinear_calibration.fit_train_heat_flux_scale(points: Iterable[QuasilinearCalibrationPoint | dict[str, Any]], *, train_split: str = 'train', prediction_floor: float = 1e-300) → dict[str, Any]

Fit one multiplicative heat-flux scale from training points.

The fit is a through-origin least-squares estimate, scale = sum(q_i Q_i) / sum(q_i^2), where q_i is the raw quasilinear heat-flux estimate and Q_i is the nonlinear window mean. This is the minimal calibration constant used by simple mixing-length models; any held-out failure after this fit is therefore a model failure, not a missing constant factor.

spectraxgk.quasilinear_calibration.integrated_quasilinear_flux_from_spectrum(spectrum_csv: str | Path, *, column: str = 'saturated_heat_flux_total', ky_column: str = 'ky', method: str = 'sum', delta_ky: float | None = None) → dict[str, Any]

Integrate one quasilinear spectrum column into a scalar flux estimate.

method="sum" preserves the discrete spectral-sum convention used by most runtime diagnostics. method="trapezoid" is available for smooth scan studies where the CSV is treated as a sampled function of ky.

spectraxgk.quasilinear_calibration.quasilinear_calibration_report(points: Iterable[QuasilinearCalibrationPoint | dict[str, Any]], *, saturation_rule: str, version: str = '0.1', holdout_mean_rel_gate: float = 0.35, observed_floor: float = 1e-12, fit_train_scale: bool = False, metadata: dict[str, Any] | None = None) → dict[str, Any]

Build a JSON-friendly calibration/holdout report.

A report is considered a calibrated absolute-flux claim only when it has at least one training point, at least one holdout point, and the holdout mean relative error passes the supplied gate.

spectraxgk.quasilinear_calibration.write_quasilinear_calibration_report(path: str | Path, report: dict[str, Any]) → Path

Write a quasilinear calibration report to JSON.

Quasilinear Model Selection

Claim-boundary checks for quasilinear model-selection artifacts.

spectraxgk.quasilinear_model_selection.build_quasilinear_model_selection_status(*, dataset_sufficiency: dict[str, Any] | str | Path, candidate_uncertainty: dict[str, Any] | str | Path, calibration_reports: Iterable[dict[str, Any] | str | Path] = (), optimized_equilibrium_nonlinear_audits: Iterable[dict[str, Any] | str | Path] = (), required_candidate: str = 'spectral_envelope_ridge', transport_gate: float | None = None, interval_coverage_gate: float | None = None, require_optimized_equilibrium_nonlinear_audit: bool = False) → dict[str, Any]

Combine quasilinear model-selection gates into one claim ledger.

The status is intentionally narrower than an absolute-flux calibration report. It passes only when the dataset-volume gate and uncertainty gate support the selected reduced candidate while all simple train/holdout calibration reports remain unpromoted. This lets documentation state a positive model-selection result without implying a runtime absolute-flux predictor. Optional optimized-equilibrium nonlinear audit artifacts can strengthen the scoped evidence ledger, but they never promote a universal absolute-flux claim.

spectraxgk.quasilinear_model_selection.build_quasilinear_model_selection_status_from_paths(*, dataset_sufficiency: str | Path, candidate_uncertainty: str | Path, calibration_reports: Iterable[str | Path], optimized_equilibrium_nonlinear_audits: Iterable[str | Path] = (), required_candidate: str = 'spectral_envelope_ridge', require_optimized_equilibrium_nonlinear_audit: bool = False) → dict[str, Any]

Path-based wrapper for artifact scripts and CI checks.

External-VMEC Holdout Planning

Planning utilities for external-VMEC nonlinear quasilinear holdouts.

The functions in this module do not promote a quasilinear absolute-flux model and do not run simulations. They turn the tracked holdout-gap metadata and a linear candidate screen into a reproducible launch/runbook contract for the next expensive nonlinear validation campaign.

class spectraxgk.external_holdout_plan.ExternalHoldoutScreenRow(case: str, vmec_file: str, returncode: int, best_ky: float | None, best_gamma: float | None, best_omega: float | None, log: str = '')

One row from an external-VMEC linear candidate screen.

property family: str

Geometry family inferred from case and source path.

to_dict() → dict[str, object]

Return a JSON-friendly representation.

property unstable: bool

Whether the screen row is finite, successful, and linearly unstable.

spectraxgk.external_holdout_plan.build_external_holdout_runbook(*, gap_report: dict[str, Any], screen_rows: Iterable[ExternalHoldoutScreenRow], out_dir: str = 'tools_out/external_vmec_holdouts', grids: tuple[str, ...] = ('n48:48:48:32:32', 'n64:64:64:40:40'), dt: float = 0.05, min_launch_gamma: float = 0.02, max_candidates: int = 6) → dict[str, Any]

Build a JSON-ready runbook for the next nonlinear holdout campaign.

The ranking first preserves the gap-report priority, then prefers unstable families absent from the current train/holdout portfolio. The runbook is a launch plan only; promotion still requires the generated nonlinear traces to pass grid/window convergence and then enter calibration metadata as split=holdout.

spectraxgk.external_holdout_plan.external_vmec_family(case: str, source: str = '') → str

Return a stable family label for an external-VMEC candidate.

spectraxgk.external_holdout_plan.read_external_holdout_screen(path: str | Path) → list[ExternalHoldoutScreenRow]

Load a candidate-screen CSV produced by the linear external-VMEC sweep.

Parallel Decomposition Contracts

Deterministic decomposition contracts for parallel work portfolios.

The helpers in this module describe partitioning and reconstruction contracts. They do not route solver execution, alter nonlinear state layout, or make speedup claims.

class spectraxgk.parallel_decomposition.DecompositionContract(workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble', 'diagnostic_nonlinear_domain'], claim_level: Literal['production_independent_batching', 'diagnostic_nonlinear_domain_partition'], claim_label: str, n_items: int, requested_shards: int, actual_shards: int, shards: tuple[ShardAssignment, ...], independent_work: bool, changes_solver_layout: bool, state_shape: tuple[int, ...] | None = None, axis: int | None = None)

Claim-scoped shard assignment contract for a parallelization path.

property diagnostic_nonlinear_partition: bool

Whether this contract is diagnostic nonlinear-domain metadata.

property production_independent_batching: bool

Whether this contract is for production independent-work batching.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the contract.

class spectraxgk.parallel_decomposition.ReconstructionIdentityReport(workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble', 'diagnostic_nonlinear_domain'], claim_level: Literal['production_independent_batching', 'diagnostic_nonlinear_domain_partition'], claim_label: str, n_items: int, requested_shards: int, actual_shards: int, identity_passed: bool, expected_indices: tuple[int, ...], reconstructed_indices: tuple[int, ...], missing_indices: tuple[int, ...], duplicate_indices: tuple[int, ...], out_of_range_indices: tuple[int, ...], out_of_order: bool)

Serial reconstruction identity report for a decomposition contract.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the report.

class spectraxgk.parallel_decomposition.ShardAssignment(shard_id: int, start: int, stop: int, indices: tuple[int, ...], label: str)

A deterministic contiguous assignment of serial indices to one shard.

property size: int

Number of serial items assigned to this shard.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the assignment.

spectraxgk.parallel_decomposition.build_diagnostic_nonlinear_domain_decomposition(state_shape: Iterable[int], *, axis: int, requested_shards: int) → DecompositionContract

Build a diagnostic nonlinear-domain partition contract.

This metadata describes split/reassemble coverage along one state axis. It is intentionally not a production nonlinear route and does not claim nonlinear speedup.

spectraxgk.parallel_decomposition.build_independent_portfolio_decomposition(n_items: int, *, requested_shards: int, workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble']) → DecompositionContract

Build a production independent-work decomposition contract.

The assignment is deterministic, balanced, contiguous, and contains no empty shards. It covers release-ready independent portfolios only: independent_ky_scan, uq_ensemble, and optimization_ensemble.

spectraxgk.parallel_decomposition.reconstruct_serial(contract: DecompositionContract, shard_values: Sequence[Sequence[T]]) → tuple[T, ...]

Reassemble shard values into serial index order.

spectraxgk.parallel_decomposition.serial_reconstruction_identity_report(values: Sequence[T], contract: DecompositionContract, *, equal: Callable[[T, T], bool] | None = None) → ReconstructionIdentityReport

Check that contract sharding reassembles exactly to serial order.

spectraxgk.parallel_decomposition.shard_sequence(values: Sequence[T], contract: DecompositionContract) → tuple[tuple[T, ...], ...]

Return values grouped according to a decomposition contract.

QA Low-Turbulence Optimization

Reduced QA low-turbulence stellarator optimization comparison tools.

This module adds a deliberately scoped, fully JAX-differentiable comparison between two low-order quasi-axisymmetric (QA) stellarator designs:

• a control design constrained by quasisymmetry, aspect ratio, and an iota floor; and

• a transport-aware design with the same constraints plus a reduced nonlinear ITG heat-flux envelope in the objective.

The utilities are intended for optimization plumbing, sensitivity validation, and manuscript figure generation. The heat-flux trace is a smooth reduced envelope, not a turbulent nonlinear gyrokinetic time trace, and the visualized surfaces are reduced boundary models rather than solved VMEC-JAX equilibria. They are not a substitute for a production VMEC/Boozer/full-nonlinear-GK optimization loop.

class spectraxgk.qa_low_turbulence.QALowTurbulenceConfig(target_aspect: float = 6.0, min_iota: float = 0.41, iota_operating_floor: float = 0.7, max_mode: int = 1, aspect_weight: float = 8.0, iota_floor_weight: float = 160.0, iota_operating_weight: float = 70.0, qa_weight: float = 8.0, target_helical_amplitude: float = 0.16, helical_shaping_weight: float = 24.0, regularization: float = 0.002, nonlinear_weight: float = 8.0, learning_rate: float = 0.03, steps: int = 60, nonlinear_dt: float = 0.2, nonlinear_steps: int = 2000, nonlinear_tail_fraction: float = 0.5, long_window_min_time: float = 300.0, long_window_max_cv: float = 0.03, long_window_max_trend: float = 0.02, long_window_max_half_mean_rel_change: float = 0.02, fixed_density_gradient: float = 2.2, fixed_temperature_gradient: float = 6.0, scan_density_gradients: tuple[float, ...] = (0.6, 1.0, 1.4, 1.8, 2.2, 2.8, 3.4, 4.0, 4.8), fd_step: float = 0.0001, surface_ntheta: int = 72, surface_nzeta: int = 72, n_field_periods: int = 2)

Configuration for the reduced QA low-turbulence comparison.

class spectraxgk.qa_low_turbulence.QALowTurbulenceResult(design_name: str, includes_nonlinear_heat_flux: bool, parameter_names: tuple[str, ...], observable_names: tuple[str, ...], initial_params: tuple[float, ...], final_params: tuple[float, ...], initial_objective: float, final_objective: float, initial_observables: tuple[float, ...], final_observables: tuple[float, ...], history: tuple[dict[str, Any], ...], residual_gradient_gate: dict[str, Any], scalar_gradient_gate: dict[str, Any], observable_gradient_gate: dict[str, Any], covariance: dict[str, Any], config: dict[str, Any])

JSON-ready result for one reduced QA optimization.

to_dict() → dict[str, Any]

Return a stable JSON-friendly payload.

spectraxgk.qa_low_turbulence.default_qa_low_turbulence_initial_params() → Array

Return the shared off-optimum QA seed for the comparison.

spectraxgk.qa_low_turbulence.optimize_qa_low_turbulence(*, includes_nonlinear_heat_flux: bool, config: QALowTurbulenceConfig | None = None, initial_params: Array | Sequence[float] | None = None, finite_difference_workers: int = 1) → QALowTurbulenceResult

Optimize one reduced QA low-turbulence design with Adam.

spectraxgk.qa_low_turbulence.qa_low_turbulence_comparison_payload(config: QALowTurbulenceConfig | None = None, *, finite_difference_workers: int = 1) → dict[str, Any]

Build the full JSON-ready aspect-6 QA low-turbulence comparison.

spectraxgk.qa_low_turbulence.qa_low_turbulence_heat_flux_trace(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, density_gradient: float | None = None, temperature_gradient: float | None = None) → tuple[Array, Array]

Return a reduced nonlinear ITG heat-flux envelope for one gradient point.

density_gradient and temperature_gradient are normalized as a/L_n and a/L_T. The envelope is a fixed-step differentiable RK2 integration of dE/dt = 2 gamma E - alpha E^2 and Q_i = W_i E.

spectraxgk.qa_low_turbulence.qa_low_turbulence_objective(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, includes_nonlinear_heat_flux: bool) → Array

Return the scalar reduced QA comparison objective.

spectraxgk.qa_low_turbulence.qa_low_turbulence_observable_sensitivity_report(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, finite_difference_workers: int = 1) → dict[str, Any]

Check the full controls-to-observables differentiable plumbing.

This gate is stricter than the scalar objective check: it differentiates the full reduced observable vector, including the long-window nonlinear heat-flux mean, CV, and trend, and compares the JAX Jacobian against central finite differences.

spectraxgk.qa_low_turbulence.qa_low_turbulence_observable_vector(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → Array

Return QA low-turbulence observables in stable order.

spectraxgk.qa_low_turbulence.qa_low_turbulence_observables(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, density_gradient: float | None = None, temperature_gradient: float | None = None) → dict[str, Array]

Return reduced QA constraints and ITG observables.

spectraxgk.qa_low_turbulence.qa_low_turbulence_residual_names(*, includes_nonlinear_heat_flux: bool) → tuple[str, ...]

Return the stable residual names for the comparison objective.

spectraxgk.qa_low_turbulence.qa_low_turbulence_residual_vector(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, includes_nonlinear_heat_flux: bool) → Array

Return weighted residuals for the aspect-6 QA low-turbulence objective.

spectraxgk.qa_low_turbulence.qa_low_turbulence_window_metrics(times: Array, heat_flux: Array, *, tail_fraction: float = 0.5, eps: float = 1e-12) → dict[str, Array]

Return differentiable late-window heat-flux statistics.

The standard deviation uses sqrt(var + eps) so the Jacobian remains finite when a long reduced trace has fully saturated and the late-window variance is numerically zero.

spectraxgk.qa_low_turbulence.reduced_boundary_surface(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → dict[str, Any]

Return a reduced max-mode-1 boundary surface for visualization.

spectraxgk.qa_low_turbulence.reduced_lcfs_bmag(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → dict[str, Any]

Return a reduced LCFS |B| map for QA visualization.

VMEC-JAX Transport Objective

VMEC-JAX QA optimization hooks for SPECTRAX-GK transport objectives.

The classes in this module are optional glue for examples that start from a vmec_jax fixed-boundary optimization and append a SPECTRAX-GK transport term. They are intentionally small: VMEC-JAX remains responsible for aspect, iota, and quasisymmetry constraints, while SPECTRAX-GK supplies the local ITG linear/quasilinear or reduced nonlinear-window transport residual.

The gradient contract is deliberately conservative. The growth objective uses eigenvalue-only differentiation. The quasilinear and reduced nonlinear window objectives used by the VMEC-JAX optimizer differentiate the same solver growth rate and combine it with differentiable geometry-level transport weights. The eigenfunction-resolved quasilinear weights remain an audit path, not the traced VMEC-JAX optimizer residual, because nonsymmetric eigenvector AD is not supported by JAX.

class spectraxgk.vmec_jax_transport_objective.VMECJAXSpectraxTransportObjective(config: VMECJAXTransportObjectiveConfig = <factory>, wout_reference: Any | None = None, name: str = 'spectraxgk_transport')

Least-squares objective object for vmec_jax QA optimizers.

J(ctx: Any, state: Any) → Array

Return the scalar transport objective for VMEC-JAX callbacks.

to_objective_term(*, target: float | ndarray, residual_weight: float) → Any

Return a VMEC-JAX ObjectiveTerm when VMEC-JAX is installed.

class spectraxgk.vmec_jax_transport_objective.VMECJAXTransportObjectiveConfig(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_window_heat_flux']='nonlinear_window_heat_flux', sample_set: StellaratorITGSampleSet = <factory>, objective_weights: tuple[float, ...] | None=None, ntheta: int = 24, mboz: int = 21, nboz: int = 21, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, nonlinear_csat: float = 0.85, nonlinear_saturation_floor: float = 1e-10, reference_length: float | None = None, reference_b: float | None = None, objective_transform: Literal['raw', 'scaled', 'log1p']='raw', objective_scale: float = 1.0, surface_chunk_size: int = 0, validate_finite: bool = True)

Configuration for VMEC-JAX to SPECTRAX-GK objective evaluation.

property gradient_scope: str

Return the differentiated part of this objective.

objective_options() → dict[str, Any]

Return SPECTRAX-GK solver options for this objective.

spectraxgk.vmec_jax_transport_objective.spectrax_transport_objective_tuple(*, weight: float, config: VMECJAXTransportObjectiveConfig | None = None, wout_reference: Any | None = None, target: float = 0.0) → tuple[Any, float, float]

Return a tuple that can be appended to LeastSquaresProblem.from_tuples.

spectraxgk.vmec_jax_transport_objective.vmec_jax_transport_growth_branch_locality_report_from_states(base_state: Any, plus_state: Any, minus_state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig | None = None, *, step: float, gap_floor: float = 1e-08, slope_rtol: float = 0.01, slope_atol: float = 1e-08, max_samples: int = 0) → dict[str, object]

Check dominant-growth eigenbranch locality for VMEC/Boozer samples.

The optimizer-facing transport residual can only use the implicit dominant-eigenvalue gradient when the same eigenbranch is locally selected. This report evaluates the exact SPECTRAX-GK linear operator matrix at the base, plus, and minus VMEC final states for each configured surface/alpha/k_y sample, then delegates branch classification to dominant_eigenvalue_branch_locality_report().

spectraxgk.vmec_jax_transport_objective.vmec_jax_transport_objective_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig | None = None) → Array

Evaluate a scalar SPECTRAX-GK transport objective from a VMEC-JAX state.

VMEC-JAX Transport Admission

Admission policy for VMEC-JAX transport-optimization candidates.

This module is deliberately small and independent of VMEC-JAX internals. It operates on JSON-safe candidate summaries produced by the optimization and ladder tools, then answers the release-critical question: which solved WOUT, if any, is physically admissible for expensive long-window nonlinear audits?

class spectraxgk.vmec_jax_transport_admission.VMECJAXNonlinearAuditPolicy(minimum_relative_reduction: float = 0.02, minimum_uncertainty_z_score: float = 1.0, maximum_combined_sem_rel: float = 0.25, minimum_replicate_count: int = 3, minimum_surface_count: int = 3, minimum_alpha_count: int = 2, minimum_ky_count: int = 3, minimum_sample_count: int = 12, recommended_surfaces: tuple[float, ...] = (0.45, 0.64, 0.78), recommended_alphas: tuple[float, ...] = (0.0, 0.7853981633974483), recommended_ky_values: tuple[float, ...] = (0.1, 0.3, 0.5))

Policy for promoting or redesigning VMEC-JAX transport candidates.

Reduced growth/quasilinear/nonlinear-window objectives are useful only if they transfer to late-window nonlinear transport. This policy encodes the minimum replicated-audit evidence and sample coverage required before a candidate can be promoted beyond local reduced-metric admission.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.vmec_jax_transport_admission.VMECJAXNonlinearCampaignPolicy(minimum_landscape_relative_reduction: float = 0.1, minimum_landscape_uncertainty_z_score: float = 3.0, maximum_landscape_sem_rel: float = 0.05, minimum_landscape_replicate_count: int = 3, require_reduced_prelaunch_passed: bool = True, require_reduced_cross_sample_gate: bool = True, require_landscape_admission_passed: bool = True)

Admission limits for launching the next nonlinear optimizer campaign.

This gate sits between a reduced candidate screen and a broader optimizer campaign. Passing it means the next campaign is worth launching; it does not promote a production nonlinear turbulent-flux optimization claim.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.vmec_jax_transport_admission.VMECJAXReducedPrelaunchPolicy(metric_key: str = 'nonlinear_window_heat_flux', minimum_relative_reduction: float = 0.04, failed_reference_safety_factor: float = 1.5, require_sample_coverage: bool = True, maximum_cross_sample_sem_rel: float = 0.35)

Fail-closed reduced-objective gate before expensive nonlinear audits.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.vmec_jax_transport_admission.VMECJAXTransportAdmissionPolicy(metric_keys: tuple[str, ...] = ('transport_objective_final', 'spectrax_objective_final', 'transport_metric_final', 'objective_final'), minimum_relative_improvement: float = 0.0, lower_is_better: bool = True, require_authoritative_gate: bool = True, allow_baseline_fallback: bool = True)

Fail-closed policy for selecting transport-aware VMEC candidates.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

spectraxgk.vmec_jax_transport_admission.build_nonlinear_audit_redesign_report(matched_comparison: Mapping[str, Any], *, objective_sample_set: Any = None, policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Decide whether a matched nonlinear audit promotes or redesigns a candidate.

This is the fail-closed bridge between reduced VMEC-JAX transport admission and expensive long-window nonlinear evidence. A candidate is promoted only if the matched replicated nonlinear comparison passes, has a positive uncertainty-separated reduction, and the reduced objective used enough surface/field-line/k_y samples to avoid a single-point overfit.

spectraxgk.vmec_jax_transport_admission.build_nonlinear_campaign_admission_report(*, reduced_prelaunch_report: Mapping[str, Any], landscape_admission_report: Mapping[str, Any], policy: VMECJAXNonlinearCampaignPolicy | None = None) → dict[str, Any]

Gate the next nonlinear optimizer campaign from existing evidence.

This report intentionally promotes only a campaign launch. It requires a reduced prelaunch pass and an uncertainty-separated replicated nonlinear landscape point. It does not convert that point into a general multi-coefficient turbulent-flux optimization result.

spectraxgk.vmec_jax_transport_admission.build_nonlinear_landscape_admission_report(baseline_ensemble: Mapping[str, Any], candidate_ensembles: Sequence[Mapping[str, Any]], *, candidate_labels: Sequence[str] | None = None, policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Select an uncertainty-separated nonlinear candidate from a landscape.

This gate is for boundary-coefficient or line-search landscapes where a small number of selected points have replicated late-window nonlinear ensembles. It does not validate multi-coefficient global optimization by itself; it only answers whether any supplied candidate has a statistically resolved lower heat flux than the supplied baseline ensemble.

spectraxgk.vmec_jax_transport_admission.build_reduced_nonlinear_audit_prelaunch_report(*, baseline_metric: float, candidate_metric: float, objective_sample_set: Any = None, baseline_sample_statistics: Mapping[str, Any] | None = None, candidate_sample_statistics: Mapping[str, Any] | None = None, failed_reference_relative_reduction: float | None = None, policy: VMECJAXReducedPrelaunchPolicy | None = None, nonlinear_policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Gate reduced transport candidates before launching nonlinear audits.

This is intentionally conservative. A reduced nonlinear-window improvement should exceed both an absolute release threshold and, when available, a safety factor above a known failed-transfer reference before spending another long GPU campaign.

spectraxgk.vmec_jax_transport_admission.build_transport_admission_report(summaries: Sequence[Mapping[str, Any]], *, policy: VMECJAXTransportAdmissionPolicy | None = None) → dict[str, Any]

Annotate and select VMEC-JAX transport candidates.

A transport candidate is admitted only when it passes the physical solved-WOUT gate and improves the selected transport metric relative to the admitted baseline. The baseline may be promoted only as a fallback audit target; it never counts as a transport-optimization success.

spectraxgk.vmec_jax_transport_admission.candidate_transport_metric(candidate: Mapping[str, Any], *, metric_keys: Sequence[str] = ('transport_objective_final', 'spectrax_objective_final', 'transport_metric_final', 'objective_final')) → dict[str, Any]

Return the first finite transport metric found in a candidate summary.

spectraxgk.vmec_jax_transport_admission.select_admitted_transport_candidate(summaries: Sequence[Mapping[str, Any]], *, policy: VMECJAXTransportAdmissionPolicy | None = None) → dict[str, Any] | None

Return the promoted candidate from build_transport_admission_report().

spectraxgk.vmec_jax_transport_admission.transport_objective_sample_summary(sample_set: Any, *, policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Summarize whether a transport objective has enough sample coverage.

The nonlinear audit that motivated this gate was a single reduced metric: it improved locally but did not transfer to the replicated late-window heat-flux mean. Multi-surface, multi-field-line, and multi-k_y coverage is therefore treated as an admission requirement for the next candidate.

VMEC-JAX Transport Gradient

Boundary-gradient diagnostics for VMEC-JAX transport objectives.

The helpers here intentionally avoid importing vmec_jax at module import time. They operate on the small optimizer protocol exposed by VMEC-JAX (residual_fun, objective_and_gradient_fun, and _specs), which keeps SPECTRAX-GK tests fast while letting examples evaluate the real full-chain transport sensitivity on machines with VMEC-JAX installed.

spectraxgk.vmec_jax_transport_gradient.boundary_spec_record(spec: Any, *, fallback_index: int) → dict[str, Any]

Return a JSON-safe summary for a VMEC-JAX boundary parameter spec.

spectraxgk.vmec_jax_transport_gradient.build_boundary_transport_gradient_report(optimizer: Any, *, params: Sequence[float] | ndarray | None = None, label: str = 'vmec_jax_transport_gradient', top_n: int = 12, sensitivity_atol: float = 1e-12, include_jacobian: bool = False) → dict[str, Any]

Evaluate transport residual and boundary-gradient diagnostics.

Parameters:

• optimizer – VMEC-JAX-like optimizer exposing residual_fun(params) and objective_and_gradient_fun(params). The optional _specs member is used only for readable boundary-coefficient labels.

• params – Active boundary-parameter vector. None means the zero-increment vector aligned with optimizer._specs.

• label – Human-readable artifact label.

• top_n – Number of largest gradient components to keep.

• sensitivity_atol – Absolute L2 threshold below which the boundary transport response is classified as locally flat for optimization purposes.

• include_jacobian – If true and the optimizer exposes jacobian_fun, include dense residual-Jacobian norms. This can be substantially more expensive than the reverse scalar-gradient path.

spectraxgk.vmec_jax_transport_gradient.write_boundary_transport_gradient_report(report: dict[str, Any], path: str | Path) → Path

Write a boundary-gradient diagnostic JSON artifact.

VMEC-JAX Boundary Chain

Boundary-chain diagnostics for VMEC-JAX/SPECTRAX-GK gradients.

These helpers classify the scalar contractions produced by the expensive tools/probe_vmec_jax_boundary_chain.py diagnostic. The diagnostic compares raw exact-solve finite differences, frozen-axis initial-state finite differences, and VMEC-JAX exact-tape JVP/VJP contractions. Keeping the classification logic in the package makes the paper-facing convention explicit and unit-testable without launching VMEC solves.

spectraxgk.vmec_jax_boundary_chain.boundary_chain_summary_from_probe(payload: Mapping[str, Any], **kwargs: Any) → dict[str, Any]

Build a chain summary from a probe JSON payload.

spectraxgk.vmec_jax_boundary_chain.build_boundary_chain_collection_summary(probes: Sequence[Mapping[str, Any]], *, exact_relative_tolerance: float = 0.1, internal_relative_tolerance: float = 1e-08, absolute_tolerance: float = 1e-10) → dict[str, Any]

Summarize several boundary-chain probes as one promotion gate.

A single coefficient can look well-conditioned while neighboring boundary modes still move the raw exact-solve initialization branch. The collection summary keeps the stricter manuscript/release decision explicit: frozen-axis JVP/VJP replay must be internally transposed for every component, while exact finite-difference agreement is counted separately from branch sensitivity.

spectraxgk.vmec_jax_boundary_chain.build_boundary_chain_summary(*, exact_fd_cost_gradient: float, final_cot_dot_exact_final_fd: float, frozen_axis_replay_cost_gradient: float, frozen_axis_vjp_cost_gradient: float, frozen_axis_linear_replay_cost_gradient: float | None = None, frozen_axis_linear_vjp_cost_gradient: float | None = None, frozen_axis_initial_fd_vs_linear_abs_norm: float | None = None, frozen_axis_initial_fd_vs_linear_rel: float | None = None, raw_initial_replay_cost_gradient: float | None = None, raw_initial_fd_norm: float | None = None, frozen_axis_initial_fd_norm: float | None = None, exact_relative_tolerance: float = 0.1, internal_relative_tolerance: float = 1e-08, absolute_tolerance: float = 1e-10) → dict[str, Any]

Classify a boundary-gradient chain probe.

Parameters:

• exact_fd_cost_gradient – Central finite difference through plus/minus exact VMEC solves.

• final_cot_dot_exact_final_fd – Final-state SPECTRAX-GK cotangent dotted into the exact final-state finite-difference direction.

• frozen_axis_replay_cost_gradient – VMEC-JAX tape JVP contraction using the frozen-axis initial-state tangent used by the optimizer.

• frozen_axis_vjp_cost_gradient – VMEC-JAX tape VJP contraction projected back through the same frozen initial-state map.

• frozen_axis_linear_replay_cost_gradient – Optional contraction using VMEC-JAX’s explicit frozen-axis tangent column.

• frozen_axis_linear_vjp_cost_gradient – Optional VJP contraction using VMEC-JAX’s explicit frozen-axis tangent column.

• frozen_axis_initial_fd_vs_linear_abs_norm – Optional norm of the frozen-axis finite-difference tangent minus the explicit tangent column.

• frozen_axis_initial_fd_vs_linear_rel – Optional relative norm of the frozen-axis finite-difference tangent minus the explicit tangent column.

• raw_initial_replay_cost_gradient – Optional tape JVP contraction using raw plus/minus initial-state finite differences. This diagnoses magnetic-axis branch sensitivity, but it is not the optimizer’s advertised derivative.

VMEC-JAX Transport Line Search

Utilities for projected VMEC-JAX transport line searches.

These helpers convert a boundary-gradient diagnostic into a normalized descent direction and JSON-safe line-search manifests. The VMEC-JAX-dependent input writer lives in tools/; this module is deliberately backend-free so CI can test the admission bookkeeping without launching equilibrium solves.

class spectraxgk.vmec_jax_transport_line_search.ProjectedLineSearchPolicy(minimum_relative_improvement: float = 0.0, lower_is_better: bool = True, require_gate_passed: bool = True)

Selection policy for projected transport line-search candidates.

to_dict() → dict[str, Any]

Return a JSON-safe policy representation.

spectraxgk.vmec_jax_transport_line_search.boundary_chain_accepted_parameter_indices(collection: Mapping[str, Any], *, require_exact_fd: bool = True, require_growth_branch_locality: bool = False) → tuple[int, ...]

Return parameter indices admitted by a boundary-chain collection gate.

spectraxgk.vmec_jax_transport_line_search.projected_line_search_input_manifest(report: Mapping[str, Any], *, steps: Sequence[float], top_n: int | None = None, boundary_chain_collection: Mapping[str, Any] | None = None, require_boundary_chain_exact_fd: bool = True, require_growth_branch_locality: bool = False) → dict[str, Any]

Build a JSON-safe manifest for projected line-search input generation.

spectraxgk.vmec_jax_transport_line_search.select_projected_line_search_candidate(baseline: Mapping[str, Any], candidates: Sequence[Mapping[str, Any]], *, policy: ProjectedLineSearchPolicy | None = None) → dict[str, Any]

Select the best gate-passing projected line-search candidate.

spectraxgk.vmec_jax_transport_line_search.sparse_descent_direction_from_gradient_report(report: Mapping[str, Any], *, parameter_count: int | None = None, top_n: int | None = None, boundary_chain_collection: Mapping[str, Any] | None = None, require_boundary_chain_exact_fd: bool = True, require_growth_branch_locality: bool = False) → ndarray

Return a normalized sparse descent direction from a gradient report.

The direction is -grad restricted to the ranked top_gradient_components. This convention makes positive line-search steps lower a lower-is-better transport objective to first order.

VMEC-JAX Candidate Gates

Acceptance gates for VMEC-JAX stellarator-optimization candidates.

The helpers here are intentionally independent of the SPECTRAX-GK time integrator. They answer a narrower question: is a solved VMEC-JAX equilibrium candidate physically acceptable enough to spend expensive nonlinear GK audit time on it?

spectraxgk.vmec_jax_candidate_gate.build_authoritative_wout_candidate_gate(wout: str | Path | Mapping[str, Any], *, target_aspect: float, aspect_atol: float, min_abs_mean_iota: float, qs_residual_max: float, iota_profile_floor: float | None, helicity_m: int = 1, helicity_n: int = 0, qs_surfaces: tuple[float, ...] = (0.0, 0.1, 0.2, 0.30000000000000004, 0.4, 0.5, 0.6000000000000001, 0.7000000000000001, 0.8, 0.9, 1.0), qs_ntheta: int = 63, qs_nphi: int = 64) → dict[str, Any]

Build a solved-equilibrium gate directly from a WOUT artifact.

Use this when the deterministic replayed WOUT is the publication-facing equilibrium. It does not assert that the replayed WOUT matches an optimizer-state WOUT; that remains the role of build_wout_reproducibility_gate().

spectraxgk.vmec_jax_candidate_gate.build_solved_vmec_candidate_gate(candidate: Any, *, target_aspect: float, aspect_atol: float, min_abs_mean_iota: float, qs_residual_max: float, iota_profile_floor: float | None, iota_profiles: tuple[ndarray, ndarray] | None = None, profile_source: str = 'provided') → dict[str, Any]

Build a JSON-safe solved-equilibrium gate report.

candidate may be a VMEC-JAX optimization result with a history property or a history mapping loaded from history.json.

spectraxgk.vmec_jax_candidate_gate.build_wout_reproducibility_gate(reference_wout: str | Path | Mapping[str, Any], rerun_wout: str | Path | Mapping[str, Any], *, target_aspect: float, aspect_atol: float, min_abs_mean_iota: float, iota_profile_floor: float | None, mean_iota_repro_atol: float = 0.0005, aspect_repro_atol: float = 1e-06, profile_repro_atol: float = 0.0005) → dict[str, Any]

Check that a saved VMEC input reproduces the optimizer-state WOUT.

VMEC-JAX can write both an optimizer-state wout_final.nc and an input.final deck. For publication-facing transport claims, the deck must reproduce the WOUT when rerun; otherwise downstream SPECTRAX-GK metrics may be attached to a different equilibrium than the optimized state. This gate compares the original WOUT against a fresh rerun WOUT and also applies the solved-equilibrium aspect/iota/profile admission checks to the rerun.

spectraxgk.vmec_jax_candidate_gate.final_iota_profiles_from_vmec_result(result: Any) → tuple[ndarray, ndarray] | None

Return final solved iota profiles from a VMEC-JAX result if available.

Stellarator Objective Portfolios

Backend-free reduced objective portfolios for stellarator optimization.

This module only reduces already-evaluated objective rows. It intentionally does not import VMEC, Boozer, or solver backends so the same contract can be used by fast CI fixtures and by production VMEC/Boozer objective drivers after they have built a per-surface/per-alpha/per-ky objective table.

class spectraxgk.stellarator_objective_portfolio.ReducedPortfolioArtifactGuardConfig(min_alphas: int = 2, min_ky: int = 2, min_objectives: int = 1, min_boozer_mode: int = 21, require_growth_objective: bool = True, require_quasilinear_objective: bool = True, require_vmec_paths: bool = True, value_rtol: float = 1e-08, value_atol: float = 1e-08)

Requirements for promoting real VMEC/Boozer reduced-portfolio rows.

to_dict() → dict[str, object]

Return a JSON-friendly representation.

class spectraxgk.stellarator_objective_portfolio.StellaratorObjectivePortfolioContract(n_surfaces: int, n_alphas: int, n_ky: int, n_objectives: int, reduction: Literal['weighted_mean', 'mean', 'max'], uses_sample_weights: bool, uses_separable_sample_weights: bool, uses_objective_weights: bool)

Static shape/weight contract for a reduced objective portfolio.

property n_samples: int

Number of surface/alpha/ky samples in the portfolio.

property row_shape: tuple[int, int, int, int]

Expected objective-table shape (surface, alpha, ky, objective).

property sample_shape: tuple[int, int, int]

Expected sample-weight shape (surface, alpha, ky).

to_dict() → dict[str, object]

Return a JSON-friendly representation.

spectraxgk.stellarator_objective_portfolio.aggregate_objective_portfolio(objective_rows: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, objective_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean', validate: bool = True) → Array

Reduce a (surface, alpha, ky, objective) table to one scalar.

weighted_mean normalizes both sample and objective weights to unit sum, making the scalar invariant to the caller’s absolute weight scale. mean is the unweighted mean over every table entry. max returns the worst-case objective-weighted sample and is intended for diagnostics rather than smooth gradient-based optimization.

spectraxgk.stellarator_objective_portfolio.objective_portfolio_sensitivity_report(objective_row_fn: Callable[[Array], Any], params: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, objective_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, min_rank: int | None = None, condition_number_limit: float = 100000000.0, covariance_regularization: float = 1e-09, workers: int = 1, parallel_executor: str = 'thread') → dict[str, object]

AD/FD and conditioning report for a reduced objective-row portfolio.

objective_row_fn is the backend boundary: production callers can wire a VMEC/Boozer/quasilinear row builder into this gate while tests can use a cheap fixture. The report checks both the final scalar reduction and the unreduced row sensitivity map so a passing scalar gradient cannot hide a rank-deficient or badly conditioned objective table.

spectraxgk.stellarator_objective_portfolio.portfolio_objective_weight_vector(objective_rows: Any, *, objective_weights: Any | None = None) → Array

Return normalized objective-column weights.

spectraxgk.stellarator_objective_portfolio.portfolio_sample_weight_tensor(objective_rows: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None) → Array

Return normalized sample weights with shape (surface, alpha, ky).

spectraxgk.stellarator_objective_portfolio.reduced_portfolio_artifact_guard_report(row_artifact: dict[str, Any], *, gradient_artifacts: list[dict[str, Any]] | tuple[dict[str, Any], ...] = (), config: ReducedPortfolioArtifactGuardConfig | None = None) → dict[str, object]

Validate a real VMEC/Boozer reduced-portfolio artifact before promotion.

The guard is backend-free: it consumes already-generated JSON payloads, rebuilds a (surface, alpha, ky, objective) reducer table from real VMEC/Boozer sample rows, and checks that provenance, coverage, FD/AD diagnostics, and nonlinear-claim boundaries are explicit.

spectraxgk.stellarator_objective_portfolio.validate_objective_portfolio_contract(objective_rows: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, objective_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean') → StellaratorObjectivePortfolioContract

Validate static row/weight contracts and return portfolio metadata.

Concrete weights must be finite, non-negative, and have positive sum. Under JAX tracing, value-level weight checks are deferred to the caller, but shape contracts remain enforced from static array shapes.

Runtime Runner

Unified runtime-configured linear driver (case-agnostic core path).

class spectraxgk.runtime.RuntimeIndependentParallelPlan(requested_workers: int, effective_workers: int, executor: str, strategy: str, axis: str, source: str, problem_size: int)

Resolved independent-worker policy for runtime scan workloads.

property enabled: bool

Whether the resolved plan uses more than one independent worker.

to_dict() → dict[str, Any]

Return a JSON-friendly policy payload for runtime artifacts.

class spectraxgk.runtime.RuntimeLinearResult(ky: float, gamma: float, omega: float, selection: ModeSelection, t: ndarray | None = None, signal: ndarray | None = None, state: ndarray | None = None, z: ndarray | None = None, eigenfunction: ndarray | None = None, fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, fit_signal_used: str | None = None, quasilinear: dict[str, Any] | None = None)

Result container for runtime linear runs.

class spectraxgk.runtime.RuntimeLinearScanResult(ky: ndarray, gamma: ndarray, omega: ndarray, quasilinear: tuple[dict[str, Any], ...] | None = None, parallel: dict[str, Any] | None = None)

Result container for runtime linear ky scans.

class spectraxgk.runtime.RuntimeNonlinearResult(t: ndarray, diagnostics: SimulationDiagnostics | None, phi2: ndarray | None = None, fields: FieldState | None = None, state: ndarray | None = None, ky_selected: float | None = None, kx_selected: float | None = None)

Result container for runtime nonlinear runs.

spectraxgk.runtime.build_runtime_geometry(cfg: RuntimeConfig) → SAlphaGeometry | SlabGeometry | FluxTubeGeometryData

Resolve runtime geometry while preserving the runtime module patch surface.

spectraxgk.runtime.build_runtime_linear_params(cfg: RuntimeConfig, *, Nm: int | None = None, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData | None = None) → LinearParams

Build runtime linear parameters using the runtime module geometry surface.

spectraxgk.runtime.build_runtime_linear_terms(cfg: RuntimeConfig) → LinearTerms

Build runtime linear term toggles.

spectraxgk.runtime.build_runtime_term_config(cfg: RuntimeConfig) → TermConfig

Build runtime nonlinear-ready term config.

spectraxgk.runtime.run_linear_case(config_path: str | Path, *, ky: float | None = None, Nl: int | None = None, Nm: int | None = None, solver: str | None = None, method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, show_progress: bool = True) → int

Run a linear case from a runtime TOML with optional overrides.

spectraxgk.runtime.run_nonlinear_case(config_path: str | Path, *, ky: float | None = None, Nl: int | None = None, Nm: int | None = None, method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, diagnostics_stride: int | None = None, show_progress: bool = True) → int

Run a nonlinear case from a runtime TOML with optional overrides.

spectraxgk.runtime.run_runtime_linear(cfg: RuntimeConfig, *, ky_target: float = 0.3, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', return_state: bool = False, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → RuntimeLinearResult

Run one linear point from a case-agnostic runtime config.

spectraxgk.runtime.run_runtime_nonlinear(cfg: RuntimeConfig, *, ky_target: float = 0.3, kx_target: float | None = None, Nl: int | None = None, Nm: int | None = None, dt: float | None = None, steps: int | None = None, method: str | None = None, sample_stride: int | None = None, diagnostics_stride: int | None = None, laguerre_mode: str | None = None, diagnostics: bool | None = None, resolved_diagnostics: bool = True, return_state: bool = False, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → RuntimeNonlinearResult

Run a nonlinear point using the unified runtime config path.

spectraxgk.runtime.run_runtime_scan(cfg: RuntimeConfig, ky_values: Sequence[float], *, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, batch_ky: bool = False, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', show_progress: bool = False, workers: int = 1, parallel_executor: str = 'thread') → RuntimeLinearScanResult

Run a ky scan using the unified runtime config path.

When batch_ky is enabled, all ky points are integrated together using the time integrator (Krylov is not supported in this mode).