API Reference

Public API Registry

gkx.__all__ holds 16 names: __version__ and the 15 names of gkx.api.__all__.

• Workflow: load, solve, scan, plot, prepare.

• Contracts: PreparedSimulation, Case, LinearResult, NonlinearResult, ScanResult.

• VMEX turbulence coupling: flux_tube_geometry_from_mapping, solver_objective_vector_from_geometry, solver_linear_operator_matrix_from_geometry, solver_scalar_objective_from_vector, VMEXTransportObjectiveConfig.

Every name resolves lazily from gkx.api._EXPORT_TARGETS, so import gkx does not import the solver stack. The same table keeps the GKX 1.x root names: they still import by name (from gkx import KrylovConfig) but are not in gkx.__all__, dir(gkx), or wildcard imports. They may be removed no earlier than GKX 3.0; import them from the owning module listed on this page.

Contract types

Case, LinearResult, NonlinearResult, and ScanResult are aliases, not wrappers, of RuntimeConfig, RuntimeLinearResult, RuntimeNonlinearResult, and RuntimeLinearScanResult. Result arrays are never copied, and the Runtime* names remain valid imports.

Workflow

case = gkx.load("cyclone.toml")
result = gkx.solve(case)
figure, axes = gkx.plot(result)
spectrum = gkx.scan(case, [0.1, 0.2, 0.3])
simulation = gkx.prepare(case)
result = simulation.solve()

load

Reads a TOML deck through gkx.workflows.runtime.toml and resolves case-relative paths.

solve

gkx.runtime.solve: runs the nonlinear runtime when physics.nonlinear is set, otherwise the linear runtime.

scan

gkx.runtime.run_runtime_scan: a linear \(k_y\) scan.

plot

gkx.artifacts.plotting.plot: returns (Figure, axes) for a linear, scan, or nonlinear result, and neither saves nor shows the figure.

prepare

gkx.api.prepared.prepare_simulation: validates the case and returns a frozen PreparedSimulation.

When the deck has no [run] table and no Nl/Nm argument is given, (Nl, Nm) defaults to (12, 24) for a linear case and (4, 8) for a nonlinear case, the same fallback the runtime and the CLI use.

Prepared simulations

PreparedSimulation accepts linear and nonlinear cases and owns no numerics: every method delegates to the runtime path the case would take anyway.

Member	Behavior
solve(parameters=None, initial_state=None)	Nonlinear: runs the compiled scan and returns (time, diagnostics, state, fields). Linear: calls gkx.runtime.solve. parameters must be None; a linear initial_state raises NotImplementedError.
scan(parameter, values, *, parallel="auto")	Only parameter="ky"; delegates to run_runtime_scan. An integer parallel sets the worker count.
value_and_grad(objective, parameters)	jax.value_and_grad(objective)(parameters).
warmup()	Nonlinear: runs the scan once so XLA compiles it, and returns self; a second call does nothing. Linear: no-op, because the linear runtime picks its solver per call.
estimate_memory()	Bytes for the complex state and the stage copies an explicit step keeps live. A floor: diagnostic buffers are not counted.
summary(), print_summary(), compilation_metadata()	Kind, state shape, (Nl, Nm), geometry model, precision, devices, working set, persistent-cache state, and whether it was warmed. compiled_at_prepare is always False: XLA compiles on first execution.
precision, devices	"float64" or "float32" from jax_enable_x64; visible JAX devices as platform:id.

Preparing a nonlinear case builds its explicit diagnostic scan through gkx.runtime.prepare. That route requires a serial [parallel] policy and an explicit time method (IMEX methods raise ValueError). It compiles one scan of fixed length, so a deck with [time] run_to = "saturation" raises unless steps=N is passed or the deck sets run_to = "t_max".

The prepared route and integrate_nonlinear_explicit_diagnostics_state compile the same graph and return bitwise identical arrays (tests/unit/nonlinear/test_nonlinear.py). Prefer prepare when a case runs more than once: it keeps the compiled graph, while each call to the function entry point compiles its own. Solvers gives the contract and the state-dtype caveat.

On a linked (twist-shift) deck, every nonlinear state these entry points are given – the one prepare is built with and any initial_state passed to PreparedSimulation.solve – is projected onto the linked chain cover before it is integrated, exactly as the runtime projects a supplied initial_state at intake. The rows outside the chains are decoupled, and the ExB bracket would alias whatever sits there back onto the chain rows. The projection is a mask fixed by the deck’s topology, not a check on the state, so it holds under jit and under reverse-mode AD: the cotangent of a supplied state is exactly zero on those rows and unchanged elsewhere. Periodic decks and full-cover linked grids are untouched. The raw drivers underneath, integrate_nonlinear and integrate_nonlinear_cached, project a supplied state the same way. Solvers states the rule and what it costs.

Lazy public API registry for GKX.

The runtime package keeps physics, numerics, diagnostics, and workflow ownership in their domain modules. This facade provides a compact public import surface without keeping one re-export module per domain.

gkx.api.Case

alias of RuntimeConfig

gkx.api.LinearResult

alias of RuntimeLinearResult

gkx.api.NonlinearResult

alias of RuntimeNonlinearResult

class gkx.api.PreparedSimulation(case: ~typing.Any, kind: str, state_shape: tuple[int, ...], n_laguerre: int, n_hermite: int, options: dict[str, ~typing.Any] = <factory>, _backend: ~typing.Any = None, _warmed: list[bool] = <factory>)

A reusable, compiled simulation for one case.

Construct with gkx.prepare(). The object is frozen: deriving a different topology means preparing a different case, which is what keeps the compiled callables and the reported metadata describing the same run.

compilation_metadata() → dict[str, Any]

Report whether a second run of this topology can reuse a cache.

property devices: tuple[str, ...]

Return the visible JAX devices as their platform strings.

estimate_memory() → dict[str, Any]

Estimate the resident bytes this simulation’s arrays require.

The estimate covers the distribution array and the working copies an explicit stage holds live at once. It is a floor, not a ceiling: it does not model the diagnostic buffers a long sampled run accumulates, and it says so rather than pretending to a precision it lacks.

property precision: str

Return the JAX floating precision this simulation will run in.

print_summary(*, stream: Any = None) → None

Print the summary one field per line.

scan(parameter: str, values: Sequence[float], *, parallel: str = 'auto') → Any

Solve this case once per value of a named scalar parameter.

solve(parameters: Any = None, initial_state: Any = None) → Any

Run this simulation and return its typed result.

parameters is reserved for the differentiable path and must be None here; a prepared object must not silently solve a different physics problem than the one it reports.

A nonlinear initial_state on a linked (twist-shift) deck is projected onto the linked chain cover before it is integrated, which is the same rule the runtime applies at intake (queue rows Q19 and Q23): the rows outside the chains are decoupled, and the ExB bracket would alias whatever sits there back onto the chain rows. Periodic and full-cover decks are untouched.

summary() → dict[str, Any]

Return the topology, precision, and cache facts that identify this.

value_and_grad(objective: Callable[[...], Any], parameters: Any) → Any

Return (value, gradient) of objective at parameters.

warmup() → PreparedSimulation

Force compilation now, so a later timed call measures execution.

Preparing a nonlinear case builds its scan closure but does not compile it: XLA compiles when the scan first executes. Measured on the shipped five-step KBM deck, prepare took 4.3 s, the first solve 19.6 s and the second 0.000 s. This method used to return self without doing anything, so code that called it and then timed solve was timing the compile.

Forcing compilation means executing the prepared closure once, because the scan exposes no separately lowerable handle. That costs one run – cheap next to the compile it triggers, but not free, so it happens only when asked for. A second call is a no-op.

A linear case has nothing to compile ahead of time: the linear runtime chooses its solver per call, so warming one would compile a solver the next call may not use. summary()["warmed"] reports which happened.

Returns self so a caller can chain.

gkx.api.ScanResult

alias of RuntimeLinearScanResult

class gkx.api.VMEXTransportObjectiveConfig(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 GKX objective evaluation.

property gradient_scope: str

Return the differentiated part of this objective.

objective_options() → dict[str, Any]

Return GKX solver options for this objective.

gkx.api.flux_tube_geometry_from_mapping(data: Mapping[str, Any], *, source_model: str = 'vmex', 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. vmex / 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.

gkx.api.load(path: str | Path) → RuntimeConfig

Load a resolved immutable GKX case from TOML.

gkx.api.plot(result: Any) → Any

Create the standard in-memory figure for a runtime result.

gkx.api.prepare(case: Any, **options: Any) → PreparedSimulation

Build a PreparedSimulation for case.

A nonlinear case builds its scan closure here, which is what makes a later solve cheap once that closure has been compiled – XLA does that when the scan first executes, so call warmup() if a later solve is being timed. A linear case is validated and described but not built, because the linear runtime chooses its solver per call. summary() reports both facts rather than hiding them.

gkx.api.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, window_method: str = 'stationary', krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', show_progress: bool = False, workers: int = 1, parallel_executor: str = 'thread', warm_start: bool | None = None) → RuntimeLinearScanResult

Run a ky scan; warm_start overrides the case output policy.

Coordinates serial, independent-worker, or combined-ky runtime scans.

gkx.api.solve(case: RuntimeConfig, **options: Any) → RuntimeLinearResult | RuntimeNonlinearResult

Solve one case through its existing linear or nonlinear runtime owner.

gkx.api.solver_linear_operator_matrix_from_geometry(geom: Any, *, spectral_grid: Any | None = None, selected_ky_index: int = 1, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, lx: float = 6.0, ly: float = 12.0, params_linear: LinearParams | None = None, terms: LinearTerms | None = None) → Array

Materialize the complex linear-RHS operator for one solver geometry.

This helper exposes the exact matrix whose dominant eigenvalue is used by solver_growth_rate_from_geometry(). It is intended for branch locality and AD/finite-difference admission gates; production time integration should continue to call the RHS directly.

gkx.api.solver_objective_vector_from_geometry(geom: Any, *, spectral_grid: Any | None = None, selected_ky_index: int = 1, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, lx: float = 6.0, ly: float = 12.0, params_linear: LinearParams | None = None, terms: LinearTerms | None = None, eigensolver: Literal['dense', 'adaptive-propagator'] = 'dense', adaptive_config: AdaptiveLinearEigensolverConfig | None = None) → Array

Evaluate dominant linear/quasilinear observables from geometry.

This is a reusable value-level objective builder for optimization drivers and examples. It builds the production linear RHS on the supplied solver-ready flux-tube geometry, selects the maximum-growth eigenbranch, and returns the ordered SOLVER_OBJECTIVE_NAMES vector.

eigensolver="adaptive-propagator" avoids materializing the dense operator, certifies the continuous residual, and differentiates the eigenpair implicitly. It remains explicit opt-in until branch-continuity gates admit a production optimization. The default dense path preserves established results. Pass a runtime spectral_grid that has already selected one ky to preserve linked stellarator boundary conditions; otherwise a periodic single-ky grid is built from the scalar grid options. One-dimensional species parameters add the leading multi-species state axis automatically.

gkx.api.solver_scalar_objective_from_vector(objective_vector: Array | ndarray, objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth') → Array

Select one scalar objective from SOLVER_OBJECTIVE_NAMES.

This tiny selector keeps optimizer code honest about which scalar is being minimized. It also centralizes aliases used by the examples: growth -> gamma, frequency -> omega, and quasilinear_flux -> mixing_length_heat_flux_proxy.

Collision Operator Interface

Shared interface for gyrokinetic collision operators.

class gkx.operators.collision.CollisionContext(distribution: Any, hamiltonian: Any, fields: Any, cache: Any, parameters: Any)

Distribution, Hamiltonian, fields, cache, and parameters seen by collisions.

cache: Any

Alias for field number 3

distribution: Any

Alias for field number 0

fields: Any

Alias for field number 2

hamiltonian: Any

Alias for field number 1

parameters: Any

Alias for field number 4

class gkx.operators.collision.CollisionOperator(*args, **kwargs)

JAX-compatible collision model returning a state-shaped RHS term.

apply(context: CollisionContext) → Any

Return the unit-weight collision contribution.

class gkx.operators.collision.SplitCollisionOperator(*args, **kwargs)

Collision model with a mathematically valid finite-time update.

split_step(context: CollisionContext, dt: Any) → Any

Advance the unit-weight collision model by dt.

Velocity-Space Core

Velocity-space basis and gyroaverage utilities.

gkx.core_velocity.J_l_all(b: Array, l_max: int) → Array

Gyroaveraging coefficients matching the Laguerre-Hermite quadrature convention.

gkx.core_velocity.associated_bessel_laguerre_coefficients(bessel_argument: Array, bessel_order: int, n_max: int) → Array

Return coefficients of the associated-Laguerre expansion of J_m.

For m = bessel_order, the returned coefficients are

\[A_n^m(b) = \frac{n!}{(n+m)!}\left(\frac{b}{2}\right)^m K_n(b),\]

so that \(J_m(B\sqrt{x}) = x^{m/2}\sum_n A_n^m L_n^m(x)\), with \(B=k_\perp v_{\mathrm{th}}/\Omega\). The leading axis indexes n=0, ..., n_max.

gkx.core_velocity.bessel_j0(x: Array, *, xp: Any = <module 'jax.numpy' from '/home/docs/checkouts/gkx/envs/latest/lib/python3.11/site-packages/jax/numpy/__init__.py'>) → Array

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

xp selects the array namespace. The default stages the evaluation onto the device; passing numpy evaluates the identical expression on the host, which one-shot cache construction uses to avoid compiling a kernel per polynomial term. Every operation here is arithmetic, sqrt, cos or sin, which both libraries evaluate identically in float64.

gkx.core_velocity.bessel_j1(x: Array, *, xp: Any = <module 'jax.numpy' from '/home/docs/checkouts/gkx/envs/latest/lib/python3.11/site-packages/jax/numpy/__init__.py'>) → Array

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

xp selects the array namespace, as in bessel_j0().

gkx.core_velocity.bessel_laguerre_kernels(bessel_argument: Array, n_max: int) → Array

Return finite-Larmor Bessel–Laguerre kernels through order n_max.

The coefficients

\[K_n(b) = \exp(-b^2/4)\frac{(b^2/4)^n}{n!}\]

expand \(J_0(B\sqrt{x})\) in ordinary Laguerre polynomials, where \(B=k_\perp v_{\mathrm{th}}/\Omega\). A recurrence avoids factorial overflow and preserves the exact b=0 limit. The leading axis indexes n=0, ..., n_max.

References

Frei et al., Journal of Plasma Physics 87, 905870501 (2021), Eq. (2.13).

gkx.core_velocity.gamma0(b: Array) → Array

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

gkx.core_velocity.hermite_ladder_coeffs(n_max: int) → tuple[Array, Array]

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

gkx.core_velocity.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}

gkx.core_velocity.laguerre_gyroaverage_neighbors(coefficients: Array, b: Array, *, axis: int) → tuple[Array, Array]

Return the lower and upper Laguerre neighbors of gyroaverage coefficients.

The upper neighbor at the truncation boundary is known analytically even though the corresponding distribution moment is not retained. For \(\mathcal J_\ell=(-1)^\ell e^{-b/2}(b/2)^\ell/\ell!\), it is \(\mathcal J_{L}=-\mathcal J_{L-1}(b/2)/L\). Zero-padding that value drops a physical term from the highest retained diamagnetic-drive equation.

gkx.core_velocity.laguerre_quadrature_count(nl: int) → int

Default number of Laguerre quadrature points.

gkx.core_velocity.laguerre_scaled_functions(x: ndarray, n_max: int) → ndarray

Scaled Laguerre functions \(\tilde L_n(x)=e^{-x/2}L_n(x)\).

The exponential rides along in the seed of the standard three-term recurrence instead of being applied afterwards, so the bare \(L_n(x)\) is never formed – it reaches \(10^{18}\) at the outermost Gauss node of even a modest grid. Szego’s bound gives \(|\tilde L_n(x)|\le 1\) for \(x\ge 0\), so every intermediate here is O(1):

\[\begin{split}\tilde L_0 = e^{-x/2},\qquad \tilde L_1 = (1-x)\,e^{-x/2},\\ (n+1)\,\tilde L_{n+1} = (2n+1-x)\,\tilde L_n - n\,\tilde L_{n-1}.\end{split}\]

Returns an array of shape (n_max + 1, x.size).

gkx.core_velocity.laguerre_transform(nl: int) → tuple[ndarray, ndarray, ndarray]

Return Laguerre transform matrices and roots.

The Gauss-Laguerre weight \(e^{-x}\) is split evenly between the two matrices, which are therefore built from the scaled functions \(\tilde L_\ell=e^{-x/2}L_\ell\) and the rescaled weights \(\tilde W_j=w_j e^{x_j}\):

\[\mathrm{to\_grid}[\ell,j] = (-1)^\ell\,\tilde L_\ell(x_j),\qquad \mathrm{to\_spectral}[j,\ell] = (-1)^\ell\,\tilde L_\ell(x_j)\,\tilde W_j.\]

to_grid @ to_spectral is the identity, exactly as it is for the unsplit pair \(L_\ell(x_j)\) / \(L_\ell(x_j)w_j\), and so is every to_spectral @ diag(f(x)) @ to_grid round trip through a pointwise factor on the quadrature grid: the two \(e^{\mp x/2}\) cancel. Since each nonlinear bracket carries exactly one factor of the transformed distribution against gyroaverage fields built directly from roots, the split is invisible to the physics. What it buys is conditioning – cond(to_grid) stays under 20 instead of hitting 9e18 at nl=20, where the unsplit form has no correct digits left.

The intermediate grid values are consequently \(e^{-x_j/2}g(x_j)\) rather than \(g(x_j)\), so anything comparing raw quadrature-grid values against an implementation that carries the weight separately has to undo that factor first.

The \((-1)^\ell\) sign is the established GKX runtime convention that the shipped collision tables are keyed to (the laguerre_convention fields under gkx/data), and is preserved here unchanged.

gkx.core_velocity.mu_grid_gyro_factors(vperp_max: float, b: Any, bmag: Any, nmu: int, real_dtype: Any, xp: Any) → tuple[Any, ...]

Gyro factors of the nodal mu grid, in the slots of the Laguerre ones.

The state is the node amplitude scaled by \(a_j=\sqrt{(B/B_0)\,w_j e^{\hat\mu_j}}\), so the Laguerre field sums apply unchanged with \(J_j=a_j e^{-\mu_j B/2} J_0(\alpha_j)\) and \(J^B_j=2\mu_j B\,a_j e^{-\mu_j B/2} J_1(\alpha_j)/\alpha_j\), \(\alpha_j^2=2\mu_j B\,b\). Returns (Jl, JlB, to_grid, to_spectral, roots, j0, j1_over_alpha); the transforms are identities.

gkx.core_velocity.mu_grid_laguerre_modes(nodes: Any, scaled_weights: Any, bmag: Any, n_l: int) → Any

Node values a_j e^{-x_j/2} (-1)^l L_l(x_j) of Laguerre moment l, shape (l, nmu, Nz).

Their sums against J_j give the Laguerre moments of a nodal state (P_l = sum_j mode_l,j G_j) and, as a seed, a nodal state that carries Laguerre moment l (orthonormal up to the vperp_max truncation).

gkx.core_velocity.mu_grid_lb_rate(nodes: Any, scaled_weights: Any, bmag: Any) → Any

Perpendicular Lenard-Bernstein rate R(z) on the nodes, shape (nmu, nmu, Nz).

With beta = B/B0 and u = e^{beta mu/2} G the operator is d_t u = (2/beta) e^{beta mu} d_mu[mu e^{-beta mu} d_mu u] (eigenvalues -2 l on u = L_l(beta mu)). Discretized by summation by parts, M d_t u = -D^T diag(kappa) D u with the Lagrange derivative matrix D, M_j = w_j e^{mu_j - beta mu_j} and kappa_j = (2/beta) mu_j M_j, it conserves sum_j M_j u_j exactly, annihilates u = 1 and has u^T M R u >= 0. Returned for the scaled amplitude a G (a_j = sqrt(beta w_j e^{mu_j})).

gkx.core_velocity.mu_grid_nodes(nmu: int, vperp_max: float) → tuple[ndarray, ndarray]

Nodes \(\hat\mu_j\) and scaled weights \(w_j e^{\hat\mu_j}\) of the mu grid.

The rule integrates against \(e^{-\hat\mu}\) exactly for polynomials of degree \(2 n_\mu - 1\). A finite vperp_max truncates it to \([0, v_{\perp,\max}^2/2]\); inf gives the standard Gauss-Laguerre rule, whose weights are returned in the scaled form \(w_j e^{\hat\mu_j}\) because \(w_j\) itself underflows.

Geometry

Analytic, slab, imported, and sampled flux-tube geometry contracts.

class gkx.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, cylindrical_R_profile: ndarray | None = None, cylindrical_Z_profile: ndarray | None = None, toroidal_angle_profile: ndarray | None = None)

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

trim_terminal_theta_point() → FluxTubeGeometryData

Return a copy without the terminal theta sample.

Imported *.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 gkx.geometry.GeometryConfig(model: str = 's-alpha', geometry_backend: str = 'auto', geometry_file: str | None = None, vmec_file: str | None = None, geometry_helper_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, geometry_helper_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 gkx.geometry.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, ky_layout: str = 'half')

Spectral grid configuration in a flux-tube.

ky_layout: str = 'half'

Layout of the ky axis of every spectral array built from this grid (gkx.core_ky_layout, plan 5.3 N3). "half" – the default – stores the Nyc = 1 + Ny // 2 non-negative rows and lets the reality condition hold by construction; "full" stores the two-sided fftfreq axis of length Ny and rebuilds the negative rows in the bracket and after every stage. Ny means the same thing in both: the length of the physical y axis, and therefore the resolution of the run. The physics, the resolution and the published output do not depend on the choice: nonlinear diagnostics agree to round-off (2e-16 in x64, 2e-7 in float32), the window gradient is bitwise equal, and "half" is faster (an RK3 step 0.54-0.60x of the two-sided time on CPU, 0.65x on an RTX A4000; the checkpointed window gradient 0.68-0.72x) with 12-25% less memory.

class gkx.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.

gradpar() → float

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

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.

class gkx.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)

Reference slab geometry contract.

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

Apply imported-geometry grid defaults implied by the selected geometry.

gkx.geometry.apply_imported_geometry_grid_defaults(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: GridConfig) → GridConfig

Apply imported-geometry grid defaults implied by the selected geometry.

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

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

gkx.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.

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

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

gkx.geometry.from_vmex(state: Any, runtime: Any, *, surface_index: int | None = None, alpha: float = 0.0, zeta0: float = 0.0, ntheta: int = 32, equal_arc: bool = True, arc_oversample: int = 4, validate_finite: bool = True) → Any

Return generic GKX geometry for a solved VMEX toroidal state.

gkx.geometry.from_vmex_mirror(state: Any, discretization: Any, axis: Any, *, validate_finite: bool = True, **geometry_kwargs: Any) → Any

Return generic GKX geometry for a periodic VMEX mirror field line.

gkx.geometry.from_vmex_wout(wout: Any, **geometry_kwargs: Any) → Any

Return generic GKX geometry from a VMEC-compatible WOUT via VMEX.

gkx.geometry.load_imported_geometry_netcdf(path: str | Path) → FluxTubeGeometryData

Load sampled flux-tube geometry from an imported NetCDF/eik file.

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

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

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

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

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

Return the effective zero-shear state.

Geometry Core

Flux-tube geometry builders and grid-default policy.

class gkx.geometry.core.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, cylindrical_R_profile: ndarray | None = None, cylindrical_Z_profile: ndarray | None = None, toroidal_angle_profile: ndarray | None = None)

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

trim_terminal_theta_point() → FluxTubeGeometryData

Return a copy without the terminal theta sample.

Imported *.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 gkx.geometry.core.GeometryConfig(model: str = 's-alpha', geometry_backend: str = 'auto', geometry_file: str | None = None, vmec_file: str | None = None, geometry_helper_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, geometry_helper_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 gkx.geometry.core.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, ky_layout: str = 'half')

Spectral grid configuration in a flux-tube.

ky_layout: str = 'half'

Layout of the ky axis of every spectral array built from this grid (gkx.core_ky_layout, plan 5.3 N3). "half" – the default – stores the Nyc = 1 + Ny // 2 non-negative rows and lets the reality condition hold by construction; "full" stores the two-sided fftfreq axis of length Ny and rebuilds the negative rows in the bracket and after every stage. Ny means the same thing in both: the length of the physical y axis, and therefore the resolution of the run. The physics, the resolution and the published output do not depend on the choice: nonlinear diagnostics agree to round-off (2e-16 in x64, 2e-7 in float32), the window gradient is bitwise equal, and "half" is faster (an RK3 step 0.54-0.60x of the two-sided time on CPU, 0.65x on an RTX A4000; the checkpointed window gradient 0.68-0.72x) with 12-25% less memory.

class gkx.geometry.core.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.

gradpar() → float

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

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.

class gkx.geometry.core.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)

Reference slab geometry contract.

gkx.geometry.core._bgrad_from_bmag(theta: ndarray, bmag: ndarray, gradpar_val: float, *, closed: bool) → ndarray

Reconstruct the mirror term from bmag on the solver theta grid.

gkx.geometry.core._periodic_spectral_derivative(values: ndarray, spacing: float) → ndarray

Return the periodic spectral derivative of a uniform 1D profile.

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

Apply imported-geometry grid defaults implied by the selected geometry.

gkx.geometry.core.apply_imported_geometry_grid_defaults(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: GridConfig) → GridConfig

Apply imported-geometry grid defaults implied by the selected geometry.

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

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

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

Return the effective boundary after zero-shear promotion.

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

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

gkx.geometry.core.load_imported_geometry_netcdf(path: str | Path) → FluxTubeGeometryData

Load sampled flux-tube geometry from an imported NetCDF/eik file.

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

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

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

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

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

Return the effective zero-shear state.

Analytic Geometry

Analytic flux-tube geometry models.

This module contains closed-form solver-ready geometry models. Sampled VMEC, Boozer, and imported field-line data live in gkx.geometry.flux_tube.

class gkx.geometry.analytic.MillerCoreParams(ntgrid: int, nperiod: int, rhoc: float, qinp: float, shat: float, rmaj: float, r_geo: float, shift: float, akappa: float, tri: float, akappri: float, tripri: float, betaprim: float, delrho: float = 0.001)

Core Miller parameters needed by the low-level Miller geometry formulas.

class gkx.geometry.analytic.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.

gradpar() → float

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

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.

class gkx.geometry.analytic.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)

Reference slab geometry contract.

class gkx.geometry.analytic._SpectralGeometry(*args, **kwargs)

Methods and coefficients shared by analytic and sampled geometries.

class gkx.geometry.analytic._SpectralGeometryMixin

Shared Fourier-space geometry contractions.

gkx.geometry.analytic.build_collocation_surfaces(params: MillerCoreParams) → dict[str, ndarray | float]

Construct the Miller surface on the collocation grid.

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

Return the effective boundary after zero-shear promotion.

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

Return the effective zero-shear state.

Flux-Tube Geometry

Sampled and imported flux-tube geometry contracts.

FluxTubeGeometryData is the solver-facing contract shared by analytic geometry, imported NetCDF/eik files, and differentiable VMEC/Boozer bridges.

class gkx.geometry.flux_tube.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, cylindrical_R_profile: ndarray | None = None, cylindrical_Z_profile: ndarray | None = None, toroidal_angle_profile: ndarray | None = None)

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

trim_terminal_theta_point() → FluxTubeGeometryData

Return a copy without the terminal theta sample.

Imported *.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 gkx.geometry.flux_tube._ImportedGeometrySelection(geom_vars: 'Any', theta: 'np.ndarray', theta_closed_interval: 'bool', is_grouped_output: 'bool')

gkx.geometry.flux_tube._bgrad_from_bmag(theta: ndarray, bmag: ndarray, gradpar_val: float, *, closed: bool) → ndarray

Reconstruct the mirror term from bmag on the solver theta grid.

gkx.geometry.flux_tube._infer_root_theta_closed_interval(theta: ndarray, variables: Any) → bool

Infer whether a root-level *.eik.nc file includes terminal theta.

gkx.geometry.flux_tube._periodic_spectral_derivative(values: ndarray, spacing: float) → ndarray

Return the periodic spectral derivative of a uniform 1D profile.

gkx.geometry.flux_tube._select_imported_geometry_variables(root: Any) → _ImportedGeometrySelection

Select grouped-output or root-level imported-geometry variables.

gkx.geometry.flux_tube.load_imported_geometry_netcdf(path: str | Path) → FluxTubeGeometryData

Load sampled flux-tube geometry from an imported NetCDF/eik file.

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

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

Miller EIK Generation

Miller to *.eiknc.nc generation helpers.

class gkx.geometry.miller_eik.MillerGeometryRequest(ntheta: int, nperiod: int, boundary: str, y0: float, rhoc: float, q: float, s_hat: float, Rmaj: float, R_geo: float, shift: float, akappa: float, akappri: float, tri: float, tripri: float, betaprim: float)

Miller geometry-generation contract.

gkx.geometry.miller_eik.build_miller_geometry_request(cfg: RuntimeConfig) → MillerGeometryRequest

Build a Miller generation request from a runtime config.

gkx.geometry.miller_eik.default_miller_eik_output_path(request: MillerGeometryRequest) → Path

Return a stable cache path for a Miller-generated *.eiknc.nc file.

gkx.geometry.miller_eik.generate_runtime_miller_eik(cfg: RuntimeConfig, *, output_path: str | Path | None = None, force: bool = False) → Path

Generate or reuse an internal-backend Miller *.eiknc.nc file from a runtime config.

VMEC EIK Generation

VMEC to *.eik.nc generation helpers.

class gkx.geometry.vmec_eik.VmecGeometryRequest(vmec_file: str, ntheta: int, boundary: str, y0: float, x0: float | None, jtwist: int | None, beta: float, alpha: float, torflux: float, npol: float, npol_min: float | None, isaxisym: bool, which_crossing: int | None, include_shear_variation: bool, include_pressure_variation: bool, betaprim: float | None, z: tuple[float, ...], mass: tuple[float, ...], dens: tuple[float, ...], temp: tuple[float, ...], tprim: tuple[float, ...], fprim: tuple[float, ...], vnewk: tuple[float, ...], species_type: tuple[str, ...])

VMEC geometry-generation contract.

gkx.geometry.vmec_eik._resolve_runtime_vmec_file(vmec_file: str) → Path

Resolve a runtime VMEC path with env/user expansion.

gkx.geometry.vmec_eik.build_vmec_geometry_request(cfg: RuntimeConfig) → VmecGeometryRequest

Build a VMEC generation request from a runtime config.

gkx.geometry.vmec_eik.default_vmec_eik_output_path(request: VmecGeometryRequest) → Path

Return a stable cache path for a VMEC-generated *.eik.nc file.

gkx.geometry.vmec_eik.generate_runtime_vmec_eik(cfg: RuntimeConfig, *, output_path: str | Path | None = None, force: bool = False) → Path

Generate or reuse an internal-backend *.eik.nc file from a runtime config.

Differentiable Geometry

Differentiable geometry bridge contracts for VMEC/JAX pipelines.

gkx.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.

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

Discover the optional vmex and booz_xform_jax bridge APIs.

vmex is normally the installed package; the filesystem fallbacks keep working for local checkouts (the historical vmex repository IS the vmex source tree). The report KEY NAMES intentionally keep the legacy vmex_* spelling: they are string-coupled to frozen validation artifacts, and the coordinated identifier rename happens in a later phase.

gkx.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 GKX geometry contract.

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

Central finite-difference Jacobian for small validation problems.

gkx.geometry.differentiable.flux_tube_geometry_from_mapping(data: Mapping[str, Any], *, source_model: str = 'vmex', 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. vmex / 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.

gkx.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 vmex / booz_xform_jax bridge before any turbulence observable is promoted into an optimization claim.

gkx.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 = 'vmex: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 vmex / 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.

gkx.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 = 'vmex: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.

gkx.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, jacobian_chunk_size: int | str | None = None, jacobian_mode: str = 'auto', 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. jacobian_chunk_size bounds the number of simultaneous forward-mode directions; use "auto" for SOLVAX’s device-aware policy, an integer for a fixed memory budget, or None for one full vmap. jacobian_mode="auto" chooses forward mode for few parameters (or when chunking is requested) and reverse mode for few observables.

gkx.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 vmex 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 GKX can discover a vmex checkout and differentiate through its JAX-native boundary data structures before higher-cost optimization workflows are promoted.

Differentiable Geometry Backend Discovery

Optional differentiable-geometry backend discovery helpers.

gkx.geometry.backend_discovery._candidate_paths(env_names: Sequence[str], defaults: Sequence[Path]) → list[Path]

Return existing backend roots from environment variables and defaults.

gkx.geometry.backend_discovery._find_importable_module(name: str, paths: Sequence[Path]) → Any | None

Import a backend module, preferring explicitly configured checkout paths.

gkx.geometry.backend_discovery._is_traced(value: Any) → bool

Return true when host NumPy validation would break JAX tracing.

gkx.geometry.backend_discovery._jax_float_dtype() → Any

Return the active JAX floating dtype for small validation arrays.

gkx.geometry.backend_discovery.discover_differentiable_geometry_backends() → dict[str, object]

Discover the optional vmex and booz_xform_jax bridge APIs.

vmex is normally the installed package; the filesystem fallbacks keep working for local checkouts (the historical vmex repository IS the vmex source tree). The report KEY NAMES intentionally keep the legacy vmex_* spelling: they are string-coupled to frozen validation artifacts, and the coordinated identifier rename happens in a later phase.

gkx.geometry.backend_discovery.geometry_precision() → Iterator[None]

Build geometry in float64, then hand back the caller’s JAX precision.

vmex and booz_xform_jax set jax_enable_x64 globally on import. Unguarded, the first VMEC geometry of a process left the gyrokinetic solve after it in float64 (2.2x slower, twice the memory), and a later process that hit the cached geometry ran in float32. Inside this scope geometry is always float64, as those packages require; outside it nothing changes.

gkx.geometry.backend_discovery.internal_vmec_backend_available() → bool

Return True when the internal VMEC backend dependencies are present.

Differentiable Flux-Tube Contract

Solver-ready flux-tube geometry contract for differentiable backends.

gkx.geometry.flux_tube_contract.flux_tube_geometry_from_mapping(data: Mapping[str, Any], *, source_model: str = 'vmex', 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. vmex / 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.

gkx.geometry.flux_tube_contract.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 vmex / booz_xform_jax bridge before any turbulence observable is promoted into an optimization claim.

gkx.geometry.flux_tube_contract.geometry_observable_names() → tuple[str, ...]

Return the ordered geometry observables used by bridge AD checks.

gkx.geometry.flux_tube_contract.vmec_field_line_tensor_observable_names() → tuple[str, ...]

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

gkx.geometry.flux_tube_contract.vmec_metric_tensor_observable_names() → tuple[str, ...]

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

Differentiable Geometry AD Checks

Autodiff and finite-difference validation helpers for geometry workflows.

gkx.geometry.autodiff_checks._json_ready(value: Any) → Any

Return a strict JSON-compatible copy, replacing nonfinite floats by null.

gkx.geometry.autodiff_checks._sensitivity_conditioning_metadata(jacobian_ad: Any, jacobian_fd: Any, params: Any, *, fd_step: float, observable_names: Sequence[str] | None = None, param_names: Sequence[str] | None = None, relative_floor: float = 1e-12) → dict[str, object]

Return JSON-friendly conditioning metadata for AD/FD Jacobian gates.

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

Central finite-difference Jacobian for small validation problems.

gkx.geometry.autodiff_checks.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, jacobian_chunk_size: int | str | None = None, jacobian_mode: str = 'auto', 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. jacobian_chunk_size bounds the number of simultaneous forward-mode directions; use "auto" for SOLVAX’s device-aware policy, an integer for a fixed memory budget, or None for one full vmap. jacobian_mode="auto" chooses forward mode for few parameters (or when chunking is requested) and reverse mode for few observables.

Differentiable VMEC Boozer Core

vmex (VMEC-in-JAX) to Boozer equal-arc core-profile bridge.

gkx.geometry.vmec_boozer_core._cached_booz_xform_constants(*, nfp: int, ntheta1: int, nzeta: int, mboz: int, nboz: int, asym: bool) → tuple[Any, Any]

Prepare Boozer constants outside traced vmex residual callbacks.

The mode table matches the single grid-representable mode set emitted by vmex.core.boozer_tables.boozer_input_tables (xn carries nfp; xm_nyq/xn_nyq equal xm/xn because the traceable tables project |B| and R/Z/lambda on one mode set).

gkx.geometry.vmec_boozer_core.flux_tube_geometry_from_vmec_boozer_state(state: Any, runtime: Any, inp: 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 an in-memory vmex/Boozer state.

This production bridge stays inside JAX-compatible objects: SpectralState -> Boozer tables -> booz_xform_jax -> FluxTubeGeometryData. Runtime file generation can still use NetCDF/EIK, but differentiable optimization should use this function or a higher-level objective so its gradient never passes through filesystem artifacts.

gkx.geometry.vmec_boozer_core.load_solved_vmex_case(case_name: str) → tuple[Any, Any, Any, Any]

Solve a named VMEC case with vmex; return (inp, state, runtime, wout).

The canonical loader for in-memory differentiable VMEC/Boozer geometry: the parsed vmex.VmecInput, the converged spectral state, the matching solver runtime, and the host-NumPy wout dataset.

gkx.geometry.vmec_boozer_core.resolve_vmex_case_input_path(case_name: str) → Path

Resolve a case name (path, vmex packaged resource, or tracked example).

Accepts an existing filesystem path, a vmex.resources/input.<name> resource (e.g. nfp4_QH_warm_start), or examples/vmec/input.<name>.

gkx.geometry.vmec_boozer_core.vmex_boozer_equal_arc_core_profiles_from_state(state: Any, runtime: Any, inp: 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 solved vmex state.

state/runtime/inp/wout come from load_solved_vmex_case() (or any vmex.optimize.Equilibrium). The bridge keeps the imported VMEC/EIK conventions for the scalar/core field-line quantities and the Boozer metric/drift terms from booz_xform_jax output, including the finite-beta pressure gradient read from wout; the Hegna-Nakajima local-equilibrium drift gates remain a separate promotion step.

Differentiable VMEC Boozer Drifts

Boozer field-line derivative algebra for differentiable geometry bridges.

class gkx.geometry.vmec_boozer_derivatives.BoozerCartesianDerivatives(d_x_d_theta: Array, d_x_d_phi: Array, d_x_d_s: Array, d_y_d_theta: Array, d_y_d_phi: Array, d_y_d_s: Array)

Cylindrical x/y derivatives derived from Boozer R and nu harmonics.

class gkx.geometry.vmec_boozer_derivatives.BoozerCoordinateGradients(grad_psi_x: Array, grad_psi_y: Array, grad_psi_z: Array, grad_theta_x: Array, grad_theta_y: Array, grad_theta_z: Array, grad_phi_x: Array, grad_phi_y: Array, grad_phi_z: Array)

Contravariant gradients for psi, theta, and phi coordinates.

class gkx.geometry.vmec_boozer_derivatives.BoozerFieldLineDerivatives(phi_b: Array, r_b: Array, d_mod_b_d_s: Array, d_mod_b_d_theta: Array, d_mod_b_d_phi: Array, d_r_b_d_s: Array, d_r_b_d_theta: Array, d_r_b_d_phi: Array, d_z_b_d_s: Array, d_z_b_d_theta: Array, d_z_b_d_phi: Array, nu_b: Array, d_nu_b_d_s: Array, d_nu_b_d_theta: Array, d_nu_b_d_phi: Array)

Spectral sums and first derivatives on one Boozer field line.

class gkx.geometry.vmec_boozer_derivatives.RawDriftProfiles(cvdrift: Array, cvdrift0: Array, gbdrift: Array, gbdrift0: Array)

Drift coefficients on the Boozer theta grid, before the equal-arc remap.

gkx.geometry.vmec_boozer_derivatives.boozer_cartesian_derivatives(spectral: BoozerFieldLineDerivatives) → BoozerCartesianDerivatives

Convert Boozer spectral R/nu derivatives to cylindrical x/y derivatives.

gkx.geometry.vmec_boozer_derivatives.boozer_coordinate_gradients(*, spectral: BoozerFieldLineDerivatives, cartesian: BoozerCartesianDerivatives, sqrt_g_booz: Array, etf_safe: Array) → BoozerCoordinateGradients

Return contravariant gradients for psi, theta, and phi coordinates.

gkx.geometry.vmec_boozer_derivatives.boozer_pressure_gradient(wout: Any, *, s_value: float, dtype: Any) → Array

Return dp/ds at s_value from the VMEC half-mesh pressure profile.

The half mesh is rebuilt from the wout pres length rather than from the Boozer output grid, because a stencil-restricted Boozer transform carries only a few surfaces while the pressure profile is always full-radius. This mirrors the wout runtime path, which splines pres[1:] over the same half mesh and differentiates the spline.

A wout without a pressure profile is a vacuum field; the returned zero then reproduces the zero-beta drifts exactly.

The profile stays a JAX array rather than being pulled to host NumPy. It is a constant today – VMEC pressure is prescribed input, not solved state, so no boundary-coefficient gradient flows through it – but keeping it on the tape is what would let a pressure-parameter gradient work later without revisiting this.

gkx.geometry.vmec_boozer_derivatives.evaluate_boozer_field_line_derivatives(out: dict[str, Any], *, theta_closed: Array, alpha: float, iota_safe: Array, base_dtype: Any, bmnc_b: Array, d_bmnc_b_d_s: Array, rmnc_b: Array, d_rmnc_b_d_s: Array, zmns_b: Array, d_zmns_b_d_s: Array, numns_b: Array, d_numns_b_d_s: Array) → BoozerFieldLineDerivatives

Evaluate Boozer spectral sums and first derivatives on one field line.

gkx.geometry.vmec_boozer_derivatives.raw_drift_profiles(request: Any, scales: Any, profiles: Any, equal_arc: Any, state: Any) → RawDriftProfiles

Compute curvature and grad-B drift coefficients before the equal-arc remap.

profiles.d_pressure_ds enters twice and with opposite effect: it is the finite-beta part of kappa_n that makes cvdrift the curvature drift, and it is subtracted back out to leave gbdrift the pure grad-B drift. At zero beta the two coincide, which is the aliasing this replaced.

Differentiable VMEC Tensor Mapping

Thin GKX conversions for VMEX-owned toroidal, WOUT, and mirror mappings.

gkx.geometry.vmec_tensor_mapping.from_vmex(state: Any, runtime: Any, *, surface_index: int | None = None, alpha: float = 0.0, zeta0: float = 0.0, ntheta: int = 32, equal_arc: bool = True, arc_oversample: int = 4, validate_finite: bool = True) → Any

Return generic GKX geometry for a solved VMEX toroidal state.

gkx.geometry.vmec_tensor_mapping.from_vmex_mirror(state: Any, discretization: Any, axis: Any, *, validate_finite: bool = True, **geometry_kwargs: Any) → Any

Return generic GKX geometry for a periodic VMEX mirror field line.

gkx.geometry.vmec_tensor_mapping.from_vmex_wout(wout: Any, **geometry_kwargs: Any) → Any

Return generic GKX geometry from a VMEC-compatible WOUT via VMEX.

gkx.geometry.vmec_tensor_mapping.vmex_flux_tube_mapping_from_state(state: Any, runtime: Any, *, surface_index: int | None = None, alpha: float = 0.0, zeta0: float = 0.0, ntheta: int = 32, equal_arc: bool = True, arc_oversample: int = 4) → dict[str, Any]

Build VMEX’s differentiable equal-arc PEST mapping.

Differentiable Geometry Numerics

Pure numerical helpers for differentiable geometry bridge routines.

gkx.geometry.numerics._boozer_half_mesh_s_grid(raw_jlist: Any | None, *, ns_b: int, ns_b_full: int, dtype: Any) → Array

Return normalized Boozer half-mesh coordinates from API surface indices.

booz_xform_jax exposes jlist using VMEC/Fortran half-mesh indexing where the first interior half mesh is j=2. The corresponding normalized radial coordinate is therefore (j - 1.5) / ns_b_full, matching the imported VMEC/EIK half mesh 0.5 * (s_full[:-1] + s_full[1:]).

gkx.geometry.numerics._interp_equal_arc_profile(theta_uniform_closed: Array, theta_equal_arc_closed: Array, values_closed: Array) → Array

Interpolate onto the equal-arc grid through the coordinate map.

The Boozer transform feeding this remap must use safe divisions in inactive Fourier branches; otherwise nonfinite upstream cotangents contaminate the moving equal-arc coordinate sensitivity. Sparse boundary finite-difference gates remain the promotion criterion for VMEC-JAX transport gradients.

Grids

Spectral grid utilities for flux-tube geometry.

class gkx.core_grid.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, ny_full: 'int | None' = None, ky_cut: 'float | None' = None)

ky_cut: float | None = None

|ky| of the parent grid’s last dealiased row, (Ny - 1) // 3, carried by a row selection that no longer contains it. Hyperdiffusion is normalized by the full grid’s k_perp corner, so a one-row ky scan has to keep the parent’s extent rather than its own ky.

property ky_layout: Literal['full', 'half']

Return whether this grid’s ky rows are the full axis or its half.

Reads ny_full, which is the only thing that can tell them apart; a grid that does not carry one is a selection of modes and is reported as FULL.

ny_full: int | None = None

Length of the two-sided ky axis this grid’s rows were taken from, or None when the rows are not a complete axis. A half-spectrum block cannot say how long its own full axis is (gkx.core_ky_layout), and the reduction weights of an even grid’s Nyquist row depend on the answer, so the grid carries it.

gkx.core_grid.build_spectral_grid(cfg: GridConfig, *, ky_layout: Literal['full', 'half'] | None = None) → SpectralGrid

Return the spectral grid of cfg with its ky axis in ky_layout.

FULL is the two-sided fftfreq axis of length Ny. HALF keeps the Nyc = 1 + Ny // 2 non-negative rows, which is the layout a half-spectrum evolved state uses (plan 5.3 N3); the grid still records the full length in SpectralGrid.ny_full, because Nyc alone cannot supply it.

ky_layout=None – the default – reads gkx.config.GridConfig.ky_layout, so the layout travels with the configuration rather than with whichever call site happens to build the grid. That is what lets one deck key keep the solver, the diagnostics, the restart file and the NetCDF writer on the same axis: they each build their own grid from the same cfg. Passing the argument explicitly overrides the config, which is what the profiling and evidence tools do to lower one deck onto both layouts in one process.

gkx.core_grid.real_fft_mesh(kx_grid: Array, ky_grid: Array) → tuple[Array, Array, Array, Array]

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

gkx.core_grid.real_fft_ordered_kx(kx: Array) → Array

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

gkx.core_grid.real_fft_unique_ky(ky: Array) → Array

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

gkx.core_grid.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.

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

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

gkx.core_grid.twothirds_mask(Ny: int, Nx: int, *, ky_layout: Literal['full', 'half'] = 'full') → Array

2/3 dealiasing mask for 2D Fourier grids.

Ny is always the length of the two-sided axis. ky_layout selects how many rows are returned: the full axis, or the ky >= 0 block that a half-spectrum state carries.

The ky axis layout contract for spectral arrays.

GKX represents a real perpendicular field by its Fourier coefficients on a (..., ky, kx, z) array whose ky axis is KY_AXIS. Two layouts of that axis exist in the code base, and every conversion between them goes through this module.

FULL – Nky == Ny

The two-sided FFT order [0, 1, ..., Ny//2 - 1, -Ny//2, ..., -1] that jnp.fft.fftfreq produces. Half of it is redundant: a real field satisfies the reality condition F[-ky, -kx, z] = conj(F[ky, kx, z]). This is the evolved state’s default layout (gkx.config.GridConfig.ky_layout).

HALF – Nky == Nyc == 1 + Ny // 2

The non-negative rows ky_j = j * 2*pi/Ly for j = 0 .. Ny // 2, in jnp.fft.rfftfreq order. The reality condition holds by construction because the negative rows are not stored. The evolved state takes it on request, [grid] ky_layout = "half". GKX’s restart files, NetCDF output and real-space snapshots have always used it, as do the flux-tube codes the plan compares against (docs numerics, plan 2.4).

Three facts about the contract are easy to get wrong and are therefore stated here rather than rederived at each call site.

Nyc does not determine Ny. 1 + Ny // 2 maps both Ny = 2*(Nyc - 1) and Ny = 2*Nyc - 1 to the same Nyc, so a half-spectrum array cannot say how long its own full axis is. Every function here that widens an axis takes ny_full explicitly, and ny_full_candidates() shows a caller that believes otherwise both answers.

The self-conjugate rows are not made real by the layout. Row j = 0 always, and row j = Ny//2 when Ny is even, are their own conjugate partners. On them the reality condition collapses to a constraint within the row, F[j, kx] = conj(F[j, -kx]), which storing only ky >= 0 does not enforce. A real inverse transform imposes it silently by discarding the anti-symmetric part; an operator that writes those rows directly must call symmetrize_self_conjugate_rows() instead.

Reductions need row weights, not a factor of two. A quantity that obeys Q[-ky] = Q[ky] – abs(F)**2, Re(F * conj(G)), every flux – sums over the full axis to the ky_row_weights()-weighted sum over the half axis. The weight is 2 on the paired rows and 1 on the self-conjugate rows, so a blanket factor of two over-counts the Nyquist row of an even-Ny grid.

gkx.core_ky_layout.FULL: Literal['full', 'half'] = 'full'

Two-sided fftfreq order, length Ny.

gkx.core_ky_layout.HALF: Literal['full', 'half'] = 'half'

Non-negative rfftfreq order, length Nyc = 1 + Ny // 2.

gkx.core_ky_layout.KY_AXIS: int = -3

Axis of the ky rows in a (..., ky, kx, z) spectral array.

gkx.core_ky_layout.conjugate_kx_order(nx: int) → ndarray

Return the kx -> -kx permutation [0, nx-1, ..., 1].

Kept on the host as a numpy array. Projectors built from it are cached and reused across traces, and a device constant materialized here would belong to whichever trace happened to build it first. The ordering is an exact index reversal, so numpy holds it without loss and each trace materializes its own constant at use.

gkx.core_ky_layout.describe(ny_full: int) → dict[str, Any]

Return the contract’s concrete numbers for one Ny.

gkx.core_ky_layout.half_dealias_mask(mask: Any, *, ny_full: int | None = None) → Any

Return the ky >= 0 rows of a (ky, kx) dealias mask.

gkx.core_ky_layout.half_ky_values(ky_full: Any) → Any

Return the non-negative ky magnitudes of a two-sided axis.

gkx.core_ky_layout.half_twothirds_mask(ny_full: int, nx: int) → Any

Return the two-thirds dealias mask on the ky >= 0 rows.

Built rather than sliced. The two-sided mask’s leading Nyc rows are the same booleans, but only because fftfreq puts |ky| in the first Nyc entries in increasing order; a half-spectrum grid states its own rows instead of inheriting that coincidence, which is the difference between a contract and a slice.

gkx.core_ky_layout.hermitian_mode_weights(ky: Any, nx: int, *, ny_full: int | None = None, dealias_mask: Any | None = None) → Any

Return the (ky, kx) reduction weight for a Hermitian-pair quantity.

This is ky_row_weights() as a traced kernel, for the quadrature, energy and spectral reductions of quantities obeying Q[-ky] = Q[ky]. On a two-sided axis every row of the pair is stored and each is counted once. On a half axis a paired row stands for itself and its unstored partner and is counted twice, while a self-conjugate row – ky = 0, and an even grid’s Nyquist row – stands only for itself and is counted once.

ny_full is what makes the Nyquist row findable on a half axis; without it the weight is the pre-contract one, 2 on every ky > 0.

gkx.core_ky_layout.is_half(rows: int, ny_full: int | None) → bool

Return whether an axis of rows rows is stored in HALF.

gkx.core_ky_layout.ky_layout_of(rows: int, ny_full: int | None) → Literal['full', 'half']

Return the layout of a ky axis of rows rows taken from Ny.

ny_full is required because Nyc alone cannot say how long its own full axis is (see the module docstring), and None – a grid whose rows are a selection of modes rather than a complete axis – is reported as FULL, which is what every consumer assumed before the half state existed and is correct for any selection that carries its own wavenumbers.

Ny = 1 and Ny = 2 are deliberately reported as FULL even though Nyc equals Ny there: the two layouts coincide, and the widening and the conjugate restore are both the identity, so calling the axis full keeps the untouched code path.

gkx.core_ky_layout.ky_row_weights(ny_full: int, *, dtype: Any = <class 'float'>) → ndarray

Return the half-spectrum reduction weights, length Nyc.

For any Q with Q[-ky] = Q[ky], sum_full(Q) equals sum(ky_row_weights(Ny) * Q_half). The weight is 1 on the self-conjugate rows and 2 elsewhere.

gkx.core_ky_layout.negative_ky_block(positive_ky: Any, *, ny_full: int, nx: int | None = None, conjugate_kx: Any | None = None) → Any

Return the negative-ky rows implied by positive_ky.

This is the reality condition written out: conj of the paired rows, reversed in ky and permuted in kx. The Nyquist row of an even grid is excluded because it has no distinct partner.

gkx.core_ky_layout.ny_full_candidates(nyc: int) → tuple[int, int]

Return the two full lengths that share this Nyc, as (even, odd).

gkx.core_ky_layout.nyc_from_ny(ny_full: int) → int

Return the half-spectrum row count Nyc = 1 + Ny // 2.

gkx.core_ky_layout.nyquist_row(ny_full: int) → int | None

Return the Nyquist row index, or None when the grid has none.

Only an even Ny >= 2 has one, at Nyc - 1. An odd grid’s highest stored row still has a distinct negative partner.

gkx.core_ky_layout.rows_for_layout(ny_full: int, layout: Literal['full', 'half']) → int

Return the stored ky row count of layout on a grid of Ny.

gkx.core_ky_layout.self_conjugate_row_mask(ky: Any, *, ny_full: int | None = None) → Any

Return a boolean mask of the ky rows that are their own partner.

Row ky = 0 always is. The Nyquist row of an even grid also is, and it cannot be recognized from a non-negative axis alone (see the module docstring on why Nyc does not determine Ny), so ny_full has to be supplied for it to be found. Without it the mask names only the zonal row, which is what every caller did before this function existed.

gkx.core_ky_layout.self_conjugate_rows(ny_full: int) → tuple[int, ...]

Return the half-spectrum rows that are their own conjugate partner.

gkx.core_ky_layout.source_ky_layout(source: Any) → Literal['full', 'half']

Return the layout of a grid’s or cache’s ky axis.

gkx.core_ky_layout.source_ny_full(source: Any) → int

Return the two-sided ky length behind a grid’s or cache’s rows.

Reads source.ny_full and falls back to the stored row count, which is the right answer for a source that predates the attribute and for any selection of modes that carries its own wavenumbers. The shape is taken off source.ky rather than a host copy of it, because a cache built inside a trace holds tracers whose values cannot be inspected at all.

gkx.core_ky_layout.symmetrize_self_conjugate_rows(state: Any, *, ny_full: int) → Any

Impose F[j, kx] = conj(F[j, -kx]) on the self-conjugate rows.

Applies to either layout. Averaging the pair rather than keeping one side keeps the operation self-adjoint and idempotent. It is also the published fix for a packed-transform gyroaverage, which without it left a ky = 0 asymmetry of up to 14 per cent (plan 5.3 N4 and its cited source).

gkx.core_ky_layout.to_full(positive_ky: Any, *, ny_full: int, nx: int | None = None, conjugate_kx: Any | None = None) → Any

Widen a ky >= 0 array to the two-sided axis by the reality condition.

ny_full is required: see the module docstring on why Nyc alone cannot supply it. conjugate_kx overrides conjugate_kx_order() for callers that cache the permutation.

gkx.core_ky_layout.to_half(state: Any, *, ny_full: int | None = None) → Any

Return the ky >= 0 rows of a spectral array.

An array already in HALF is returned unchanged, so the call is idempotent at a boundary that may be handed either layout.

gkx.core_ky_layout.transport_mode_weights(ky: Any, nx: int, *, ny_full: int | None = None, dealias_mask: Any | None = None) → Any

Return the (ky, kx) representative weight used by the flux kernels.

The flux kernels carry the Hermitian pair factor of two in the kernel expression, so this weight selects one representative per conjugate pair rather than weighting a reduction, and drops the zonal row, which carries no flux because the kernel’s i*ky factor vanishes on it.

On the two-sided axis the representatives are the ky > 0 rows, each standing for itself and its stored negative partner. On a half axis every row is a representative, and a self-conjugate row stands only for itself, so it takes half the weight and the kernel’s factor of two restores it to one.

The Nyquist row is a representative in both layouts (plan 5.3 N3, which Q24 left this decision to). Before the state switch the two-sided rule selected the ky > 0 rows, which drops an even grid’s Nyquist row – stored once, as -Ny/2 – from the flux altogether, while the half axis stores the same row as +Ny/2 and counts it. The two layouts therefore named different sums. They are reconciled here in favour of counting it, with the self-conjugate weight, because the flux may not depend on which sign of |ky| = Ny/2 a layout happens to store: the flux kernel carries an explicit factor i * ky, so that row’s contribution is odd under the sign choice, and a physical quantity cannot be.

What makes that safe rather than a change of answer is that the row’s contribution is identically zero on any state that represents a real field. On a self-conjugate row the reality condition reads F[kx] = conj(F[-kx]), under which the summand Im(conj(phi) * moment) is odd in kx and cancels pairwise across the row. Nothing in a run reaches it in any case: the two-thirds mask zeroes every row at or above Ny/3, and Ny/2 is always above it, so the weight moves only for an explicitly undealiased reduction.

Species and Linear Parameters

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

class gkx.operators.linear.params.Species(charge: float, mass: float, density: float, temperature: float, tprim: float, fprim: float, nu: float = 0.0, closure: float = 0.0)

Dimensionless kinetic-species inputs used to build solver parameters.

tprim and fprim are \(a/L_T\) and \(a/L_n\), copied through to LinearParams unchanged; see its docstring for the conversion to \(R/L\).

gkx.operators.linear.params.build_linear_params(species: Iterable[Species], **overrides: Any) → LinearParams

Build multi-species LinearParams from physical inputs.

Note that kpar_scale is not inferred here either; see linear_params_for_geometry().

Operators

Public operator kernels with lazy domain imports.

Linear Operators

Linear operator facade with cycle-free lazy imports.

Linear Linked Boundaries

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

gkx.operators.linear.linked._build_linked_end_damping_profile(*, linked_indices: tuple[ndarray, ...], ny: int, nx: int, nz: int, widthfrac: float, ky_mode: ndarray | None = None, ny_full: int | None = None) → ndarray

Construct the linked-boundary damping profile on the state’s FFT grid.

ny is the stored row count, which is both the profile’s own ky extent and the modulus the flat chain index is decoded with. On a two-sided grid each visited mode’s conjugate partner is written as well, so the profile is symmetric under ky -> -ky; on a half-spectrum grid there is no partner row to write and the mirror is skipped.

gkx.operators.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, ny_full: int | None = None) → tuple[tuple[ndarray, ...], tuple[ndarray, ...]]

Construct linked-chain FFT index maps for the parallel derivative.

The maps stay on the host. They are pure chain topology – which (ky, kx) modes a parallel chain visits and in what order – so the caller has to be able to read them back as integers to derive the gather metadata and the end-damping profile, which a cache built inside a trace could not do if these were staged out as device constants.

gkx.operators.linear.linked.linked_cover_mask(*, gather_mask: Any | None, use_gather: bool, full_cover: bool, ny: int, nx: int, ny_full: int | None = None) → Any | None

Return the (ky, kx) modes a linked chain couples, or None for all.

Only the linked parallel derivative, its kz hypercollisions and the chain end damping couple modes, and they act on chain members (and, on a two-sided ky grid, their conjugate mirrors); every other linear term is local in (ky, kx). A mode outside that set – a kx row outside the dealiased 1 + 2 * ((Nx - 1) // 3) set – has no coupling to the chains either way.

gather_mask is the cache’s flat chain membership, indexed ky + ny * kx where ny is the number of rows the state stores; the returned mask keeps its namespace, so a host array gives a host mask and a traced one gives a traced mask. Periodic (use_gather false) and full-cover grids return None, so every consumer keeps its present behaviour on them.

ny_full is the length of the two-sided axis those rows came from, which is what distinguishes a half-spectrum grid from a full one (Nyc alone cannot, gkx.core_ky_layout). On a half grid the conjugate mirror below is skipped: every stored row is its own representative, the chains visit the dealiased ky >= 0 rows directly, and (-j) % Nyc names an unrelated positive row rather than a partner.

gkx.operators.linear.linked.linked_cover_mask_from_cache(cache: Any) → Any | None

Return the chain cover of a built cache, or None for all modes.

linked_cover_mask() holds the topology; this reads it off the cache the operator will actually be applied with, so a caller that already has one – every nonlinear entry point does – pays nothing to resolve the twist-shift policy a second time. Periodic and full-cover caches return None.

gkx.operators.linear.linked.mask_supplied_state(state: Any, cache: Any) → Any

Apply the supplied-state contract for cache: zero the off-chain rows.

This is the library-level form of the rule the runtime applies at intake (queue row Q19): a state the caller supplies is projected onto the modes the linked chains reach, because the ExB bracket takes the whole state to real space and dealiases only its output, so off-chain content aliases back onto the chain rows. The projection is a where against a mask fixed by the deck’s topology, never a value-dependent branch, so it traces and differentiates: the returned state’s off-chain entries carry no cotangent back to the supplied array, which is the same statement as their carrying no physics forward. Periodic and full-cover decks return the same object, so they trace exactly as before.

gkx.operators.linear.linked.project_to_linked_cover(v: Any, mask: Any | None) → Any

Zero a (..., ky, kx, z) state outside mask; None is the identity.

Linear Cache

Linear-operator cache data model.

class gkx.operators.linear.cache_model.HermiteWindow(index: Array, total: int)

Global Hermite coordinates of a locally held slab of moments.

A Hermite-sharded route hands the serial kernels a slab that is a window onto the global m axis, optionally widened by ghost rows. Every operator whose coefficient is a function of the global moment index – the field drives at m = 0..3, the closure at the last moment, the hypercollision normalization – needs the global coordinates rather than the slab’s own arange. index may be a tracer (the offset is a per-shard quantity), total is static because it selects Python branches.

class gkx.operators.linear.cache_model.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, ny_full: int | None = None, ky_cut: Array | None = None, ntft_m0: Array | None = None, landau_pairs: Any = None, mu_grid: bool = False, mu_lb: Any = None, groups: tuple[~gkx.operators.linear.cache_model.LinearCache, ...] | None=None)

Precomputed arrays for the linear operator.

groups: tuple[LinearCache, ...] | None = None

one exact-size cache per run of species sharing a (Nl, Nm), in species order; None when every species keeps the state’s. The RHS runs each group at its own size.

Type:

Per-species moment counts (LinearParams.moments)

ky_cut: Array | None = None

The parent grid’s hyperdiffusion ky cutoff in cache units, or None when ky holds that row itself (see gkx.core_grid.SpectralGrid.ky_cut).

landau_pairs: Any = None

Block-sparse unit-frequency linearized Landau factors (LandauBlocks); None when the run uses the Dougherty kernel.

mu_grid: bool = False

True when the Laguerre axis holds nodes of the mu grid and laguerre_roots their mu_hat values (LinearParams.mu_vperp_max).

mu_lb: Any = None

Nodal Lenard-Bernstein perpendicular rate (Nmu, Nmu, Nz) (mu grid only).

ntft_m0: Array | None = None

row kx at z holds ballooning mode kx + m0. None off that geometry.

Type:

Non-twisting flux tube radial index shift m0(ky, z)

ny_full: int | None = None

Length of the two-sided ky axis behind ky, or None when the cached rows are not a complete axis. Carried from the grid because a half-spectrum block cannot say how long its own full axis is, and the reduction weight of an even grid’s Nyquist row depends on the answer (gkx.core_ky_layout).

gkx.operators.linear.cache_model.hermite_index_of(window: HermiteWindow | None, local_count: int) → Array

Return the global Hermite index of every row of a local slab.

gkx.operators.linear.cache_model.hermite_total_of(window: HermiteWindow | None, local_count: int) → int

Return the global Hermite extent behind a local slab.

Array factories and damping factors for the linear cache.

gkx.operators.linear.cache_arrays._build_end_damping_profile_array(Nz: int, widthfrac: float, boundary: str, real_dtype: dtype) → ndarray

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

The result stays on the host; _pack_linear_cache() moves the finished cache to the device in a single pass.

gkx.operators.linear.cache_arrays._build_gyroaverage_cache_arrays(b: Array, Nl: int, real_dtype: dtype) → tuple[Array, Array]

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

gkx.operators.linear.cache_arrays._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.

gkx.operators.linear.cache_arrays.check_mu_grid_params(params: LinearParams) → None

Refuse the terms the nodal mu grid does not implement (traced values pass).

gkx.operators.linear.cache_arrays.collision_damping(cache: LinearCache, params: LinearParams, real_dtype: jnp.dtype, *, squeeze_species: bool = False) → jnp.ndarray

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). Direct unit tests may also provide a pre-expanded array to exercise the damping assembly policy.

gkx.operators.linear.cache_arrays.hypercollision_damping(cache: LinearCache, params: LinearParams, real_dtype: jnp.dtype) → jnp.ndarray

Assemble benchmark-compatible hypercollision damping factors.

gkx.operators.linear.cache_arrays.hypercollision_kz_coefficient(cache: LinearCache, params: LinearParams, real_dtype: jnp.dtype) → jnp.ndarray

Coefficient of the Fourier multiplier abs(kz), not a spatial diagonal.

gkx.operators.linear.cache_arrays.keep_moments(state: Any, moments: tuple | None) → Any

Zero the moments each species does not keep in a (Ns, Nl, Nm, ...) state.

gkx.operators.linear.cache_arrays.with_mu_grid(cache: Any, params: LinearParams) → Any

Mark cache as nodal and add the nodal Lenard-Bernstein rate (identity otherwise).

gkx.operators.linear.cache_arrays.with_species_groups(cache: Any, params: LinearParams, build: Callable[[...], Any], *args: Any) → Any

Attach one exact-size cache per run of species keeping the same (nl, nm).

build(*args, params_g, nl, nm) makes a group’s cache. The field solve runs on the full state, so the gyro factors above a species’ nl are zeroed in cache itself: its polarization and gyroaverages then carry that species’ own Laguerre terms. Identity when params.moments is None or every species keeps the state’s.

Geometry-dependent construction of LinearCache.

gkx.operators.linear.cache_builder.build_linear_cache(grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, Nl: int, Nm: int) → LinearCache

Build reusable arrays for the linear RHS.

gkx.operators.linear.cache_builder.linked_chain_cover_mask(grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams) → ndarray | None

Return the (ky, kx) modes the linked chains couple for this deck.

This is the chain topology of build_linear_cache() without its velocity-space arrays: it resolves the twist-shift policy and the linked FFT maps only, which cost O(Nz + Nky * Nkx). Callers that only need to know which rows the linked operator can reach – state intake, for instance – use it instead of building a cache they would otherwise throw away. None means every row is reachable (periodic boundaries, full-cover linked grids), so the caller has nothing to do.

gkx.operators.linear.cache_builder.mask_off_chain_rows(state: Any, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams) → Any

Zero the rows a linked deck’s chains never reach; identity otherwise.

The linked operator maps off-chain rows to off-chain rows and never writes a chain row from one, so content there is a neutral wave that no growth-rate fit reads while free-energy and spectrum sums still include it. Runtime initial conditions never put anything there, but a supplied state or a restart written elsewhere can, so intake zeroes it once instead of the time loop zeroing it every step.

gkx.operators.linear.cache_builder.update_linear_cache_for_sheared_kx(cache: LinearCache, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, effective_kx_grid: Array) → LinearCache

Rebuild every continuously sheared kx-dependent cache array.

effective_kx_grid uses the same normalized units and (ky, kx) layout as cache.kx_grid. Periodic and linked standard flux tubes are supported. A flow-shear displacement is constant along each fixed-ky linked chain, so its precomputed twist-shift maps remain valid. Non-twist flux tubes use a separate, z-dependent radial representation and fail closed here.

Linear Collisions

Linearized collision matrices, interpolation, and runtime application.

class gkx.operators.linear.collisions.DriftKineticMomentCollisionOperator(matrix: Array)

Dense drift-kinetic collision matrix acting on gyrocenter moments.

In the zero-Larmor-radius limit the particle perturbation satisfies \(f\simeq g\); the Hermite–Laguerre matrix therefore acts on the evolved gyrocenter distribution, not on the post-field Hamiltonian \(H=g+qF_M\phi/T\). See Frei, Ernst & Ricci (2022), equation (73).

apply(context: CollisionContext) → jnp.ndarray

Apply the drift-kinetic matrix to evolved gyrocenter moments.

classmethod from_improved_species(density: Array, mass: Array, temperature: Array) → DriftKineticMomentCollisionOperator

Build the lowest-order improved-Sugama matrix for all species.

classmethod from_species(density: Array, mass: Array, temperature: Array) → DriftKineticMomentCollisionOperator

Build the ordered-pair matrix from physical species parameters.

class gkx.operators.linear.collisions.EqualSpeciesFiniteWavelengthCoulombOperator(bessel_argument_grid: Array, pair_frequency: Array, test_table: Array, field_table: Array, test_phi1: Array, field_phi1: Array, test_phi2: Array, field_phi2: Array)

Finite-wavelength Coulomb tables for one equal-species plasma.

Like-species collisions have the same target and source Bessel argument at every spatial point. Storing only that diagonal avoids the quadratic target/source wavelength table while retaining the complete test, field, and polarization terms. The interpolation remains differentiable in JAX.

apply(context: CollisionContext) → jnp.ndarray

Interpolate the diagonal table and apply it to one species.

class gkx.operators.linear.collisions.FiniteWavelengthCoulombOperator(bessel_argument_grid: Array, pair_frequency: Array, test_table: Array, field_table: Array, test_phi1: Array, field_phi1: Array, test_phi2: Array, field_phi2: Array)

Tabulated finite-wavelength Coulomb test, field, and polarization blocks.

Pair tables have independent target/source Bessel-argument axes B = kperp*v_thermal/Omega. Since the runtime cache stores b = kperp**2*T*m/(q*B_ref)**2, the interpolation coordinate is B = sqrt(2*b). Their matrices act on gyrocenter moments G; polarization vectors supply the particle-to-gyrocenter pullback terms from Frei et al. (2021), equation (3.50), without double-counting build_H.

apply(context: CollisionContext) → jnp.ndarray

Apply the resolved operator using the solved electrostatic field.

class gkx.operators.linear.collisions.LandauBlocks(transform_rows: tuple, transform_cols: tuple, transform: tuple, block_index: tuple, block_test: tuple, block_field: tuple, test_diagonal: Any, field_diagonal: Any, sonine_count: int)

Block-sparse factors of the linearized Landau operator, T B T^T.

T (Hermite-Laguerre to Legendre-Sonine) is block diagonal in the degree p + 2j and B in the Legendre index l, so the dense product is replaced by small batched ones: 4.1x fewer flops at (Nl, Nm) = (16, 48). Blocks of similar size share one padded batch; padded indices point at one extra zero row. *_diagonal hold the per-pair diagonal of the dense operator for preconditioners.

apply(hamiltonian: Array, **frequencies: Any) → Array

Apply the operator to the (s, l, m, ...) response H.

diagonal(**frequencies: Any) → Array

Return the (s, n) diagonal of the assembled operator.

gkx.operators.linear.collisions.apply_collision_moment_matrix(state: Array, matrix: Array, *, nu: Array, weight: Array = Array(1., dtype=float32, weak_type=True)) → Array

Apply a dense drift-kinetic matrix in Hermite-major moment ordering.

gkx.operators.linear.collisions.apply_finite_wavelength_coulomb_moment_operator(distribution: Array, test_matrix: Array, field_matrix: Array, test_phi1: Array, field_phi1: Array, test_phi2: Array, field_phi2: Array, *, phi: Array, pair_frequency: Array, charge_over_temperature: Array, weight: Array = Array(1., dtype=float32, weak_type=True)) → Array

Apply equations (3.47)–(3.50) on resolved moment coefficients.

Matrix axes are (target, source, output moment, input moment, ...); vector axes are (target, source, output moment, ...). Optional trailing axes must equal the state’s (ky, kx, z) shape. pair_frequency[a,b] multiplies all test, field, and polarization contributions for the ordered collision pair. Moment ordering is Hermite-major, matching the generated tables.

gkx.operators.linear.collisions.apply_linearized_landau(hamiltonian: Array, blocks: LandauBlocks, **frequencies: Any) → Array

Apply the linearized Landau operator (see LandauBlocks).

gkx.operators.linear.collisions.apply_multispecies_collision_moment_matrix(state: Array, matrix: Array, *, weight: Array = Array(1., dtype=float32, weak_type=True)) → Array

Apply a target/source-species collision matrix and sum source species.

gkx.operators.linear.collisions.assemble_drift_kinetic_improved_sugama_matrix(density: Array, mass: Array, temperature: Array) → Array

Assemble the normalized lowest-order improved-Sugama species matrix.

gkx.operators.linear.collisions.assemble_drift_kinetic_sugama_matrix(density: Array, mass: Array, temperature: Array) → Array

Assemble the normalized original-Sugama operator for all species.

The returned axes are (target species, source species, target moment, source moment). Collision frequencies use the dimensionless scaling nu_ab = n_b / (sqrt(m_a) T_a**(3/2)); callers remain responsible for any common dimensional prefactor.

gkx.operators.linear.collisions.drift_kinetic_improved_sugama_pair_matrices(mass_ratio: Array, temperature_ratio: Array) → tuple[Array, Array]

Return the lowest-order improved-Sugama test and field matrices.

The correction is Frei, Ernst & Ricci (2022), Appendix C, equations (101)–(102), added to the original-Sugama ordered-pair matrices. As for drift_kinetic_sugama_pair_matrices(), the directed frequency is not included and the result uses the code’s signed-Laguerre convention.

gkx.operators.linear.collisions.drift_kinetic_sugama_pair_matrices(mass_ratio: Array, temperature_ratio: Array) → tuple[Array, Array]

Return normalized test/field matrices for one ordered species pair.

This is Frei, Ernst & Ricci (2022), Appendix C, equations (C4)–(C5), with mass_ratio = m_a / m_b and temperature_ratio = T_a / T_b. Multiplication by the directed collision frequency nu_ab is left to the caller. The matrices use the code’s Hermite-major, signed-Laguerre ordering.

gkx.operators.linear.collisions.interpolate_collision_diagonal_table(bessel_argument_grid: Array, table: Array, bessel_argument: Array) → Array

Interpolate one equal-species vector or matrix table along its diagonal.

table has shape (B, moment) or (B, output_moment, input_moment). Coefficient axes are moved before the spatial axes in the returned array. Values outside the tabulated interval use the nearest endpoint.

gkx.operators.linear.collisions.interpolate_collision_pair_table(kperp_grid: Array, tables: Array, kperp: Array) → Array

Bilinearly interpolate target/source finite-wavelength pair tables.

tables has shape (target, source, target_k, source_k, coefficients...) with one vector or two matrix coefficient axes. kperp has a leading species axis followed by the simulation spatial axes. The result moves the coefficient axes ahead of the spatial axes for direct moment contraction.

gkx.operators.linear.collisions.load_collision_moment_matrix(model: str) → ndarray

Load a provenance-checked drift-kinetic collision moment matrix.

gkx.operators.linear.collisions.parallel_electric_field_source(maximum_hermite_order: int, maximum_laguerre_order: int, normalized_field: Array) → Array

Return the linearized parallel-electric-field gyro-moment source.

Frei, Ernst & Ricci (2022), equation (81), gives the force ladder \(E\sqrt{2p}N^{p-1,j}\). Linearizing about a Maxwellian leaves only the (p, j) = (1, 0) source in dN/dt = C N + s. The normalized field is \(eE/(v_{Te}m_e)\) in the equation’s convention, including its sign; for electrons in a positive physical field it is positive and the source is negative.

gkx.operators.linear.collisions.solve_driven_collision_response(collision_matrix: Array, source: Array, *, active_modes: tuple[int, ...]) → Array

Solve the constrained steady response C N + source = 0.

active_modes removes exact collision invariants and any intentionally truncated moments before the dense solve. The returned full vector is zero outside that subspace. The solve remains in JAX so currents and transport coefficients can be differentiated with respect to collision coefficients, species parameters, and the applied drive.

Tabulated collision matrices and differentiable wavelength interpolation.

class gkx.operators.linear.collision_tables.EqualSpeciesFiniteWavelengthSugamaOperator(bessel_argument_grid: Array, pair_frequency: Array, test_table: Array, field_table: Array)

Finite-wavelength original/improved-Sugama tables for one species.

The tabulated test and field matrices implement Frei et al. (2021), equations (3.72) and (3.79), optionally augmented by the improved field correction of Frei, Ernst & Ricci (2022), on the post-field nonadiabatic response \(H\). Unlike the Coulomb table, these models have no separately tabulated electrostatic polarization vectors.

apply(context: CollisionContext) → Array

Interpolate and apply the equal-species original-Sugama matrix.

class gkx.operators.linear.collision_tables.TabulatedMultispeciesCollisionOperator(kperp_grid: Array, matrices: Array)

Finite-wavelength target/source collision matrices on a kperp grid.

The table contains fully assembled collision-frequency-weighted blocks with shape (target, source, kperp, moment, moment). The runtime derives each target species’ normalized kperp from sqrt(cache.b) and keeps the interpolation and matrix application inside JAX.

apply(context: CollisionContext) → Array

Interpolate and apply the table to the post-field Hamiltonian.

gkx.operators.linear.collision_tables.build_finite_wavelength_coulomb_operator(density: Array, mass: Array, temperature: Array, moments: int = 8) → EqualSpeciesFiniteWavelengthCoulombOperator

Build the gyrokinetic (finite-Larmor) Coulomb operator for one species.

This is the full linearized Coulomb operator retaining finite perpendicular wavelength, from Frei, Ball, Hoffmann, Jorge, Ricci & Stenger (2021), arXiv:2104.11480, equations (3.47)–(3.50). The shipped tables are indexed by the Bessel argument \(B=k_\perp v_{\mathrm{th}}/\Omega\); the runtime interpolates them at \(B=\sqrt{2b}\) from the cached b, so the same operator covers every perpendicular wavenumber in the simulation.

As \(k_\perp\to0\) the tables reduce to the drift-kinetic Coulomb operator, which is enforced by a physics gate.

moments selects the shipped Hermite-Laguerre resolution and must equal the run’s Nl*Nm; see FINITE_WAVELENGTH_MOMENT_COUNTS.

The shipped tables are generated at mass_ratio = temperature_ratio = 1, so they are valid for like-species collisions; a multi-species request is refused rather than silently extrapolated.

gkx.operators.linear.collision_tables.collision_operator_from_config(name: str, *, density: Array, mass: Array, temperature: Array, nu: Array | float = 1.0, moments: int = 8) → CollisionOperator | None

Resolve a TOML collision_operator name to a solver collision operator.

The built-in names ("none", "lenard_bernstein", "landau" and its alias "coulomb") return None: the linear RHS applies them from the cache. "sugama" and "improved_sugama" build the dense drift-kinetic Hermite-Laguerre moment operator (Frei, Ernst & Ricci 2022) that replaces the built-in term. density/mass/temperature are the per-species normalizations (length n_species).

The assembled matrices carry only the dimensionless pair scaling n_b / (sqrt(m_a) T_a**1.5); the common collisionality prefactor is the caller’s responsibility, so nu is applied here. It is the same nu that sets the strength of the built-in Lenard-Bernstein term, which keeps every model on one collisionality axis.

gkx.operators.linear.collision_tables.finite_wavelength_coulomb_metadata(moments: int = 8) → dict[str, Any]

Return the provenance metadata of one shipped Coulomb table.

gkx.operators.linear.collision_tables.interpolate_collision_moment_matrix(kperp_grid: Array, matrices: Array, kperp: Array) → Array

Interpolate collision matrices onto a scalar or spatial kperp field.

matrices may contain one shared table, one table per species, or one table per ordered target/source species pair. Their respective shapes are (kperp, modes, modes), (species, kperp, modes, modes), and (target, source, kperp, modes, modes). Values outside the tabulated interval use the nearest endpoint. The coefficient grid is validated on the host; interpolation and its derivative with respect to kperp remain in JAX.

gkx.operators.linear.collision_tables.landau_blocks_for(params: Any, nl: int, nm: int, dtype: Any)

Return the run’s LandauBlocks in dtype, or None when unused.

They are built unless the model is not landau or every nu is a concrete zero.

gkx.operators.linear.collision_tables.linearized_landau_blocks(nl: int, nm: int, mass: Any, temperature: Any)

Build LandauBlocks for a species set.

gkx.operators.linear.collision_tables.linearized_landau_factors(nl: int, nm: int, mass_ratio: float = 1.0, temperature_ratio: float = 1.0) → tuple[ndarray, ndarray, ndarray]

Return (T, test, field): the pair operator is T @ block @ T.T.

T maps the Hermite-Laguerre box onto Legendre-Sonine functions; it is block diagonal in the degree p + 2j = l + 2k because each Hermite- Laguerre function is orthogonal to every polynomial of lower degree. The Legendre-Sonine blocks are diagonal in l.

C_ab[f_a, f_b] = C(f_a, F_b) + C(F_a, f_b) with mass_ratio = m_a/m_b and temperature_ratio = T_a/T_b, linearized about Maxwellians. Rows are species-a moments, columns the moments of the perturbed species (a for the test part, b for the field part), in the runtime’s Hermite-major p*Nl + j order and signed-Laguerre convention. Particles, momentum and energy are conserved exactly by the pair sum; the like-species matrix is symmetric negative semidefinite (H-theorem). With T_a != T_b the operator still conserves but is no longer self-adjoint.

gkx.operators.linear.collision_tables.linearized_landau_pair_matrices(nl: int, nm: int, mass_ratio: float = 1.0, temperature_ratio: float = 1.0) → tuple[ndarray, ndarray]

Return the dense (test, field) Hermite-Laguerre matrices of one pair.

gkx.operators.linear.collision_tables.linearized_landau_species_matrices(nl: int, nm: int, mass: Any, temperature: Any) → ndarray

Stack the dense unit-frequency (test, field) pairs as (2, s, s, n, n).

Linear Moments

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

gkx.operators.linear.moments.apply_hermite_v(G: Array) → Array

Multiply Hermite coefficients by v_parallel (ladder form).

gkx.operators.linear.moments.apply_hermite_v2(G: Array) → Array

Multiply Hermite coefficients by v_parallel^2.

gkx.operators.linear.moments.apply_laguerre_x(G: Array) → Array

Multiply Laguerre coefficients by the perpendicular energy variable.

This helper acts in the plain Laguerre basis, where x L_l = (2l+1) L_l - (l+1) L_{l+1} - l L_{l-1} gives negative off-diagonal couplings. The runtime state stores Laguerre coefficients in the alternating (-1)**l basis (the convention of gkx.core_velocity.J_l_all), which flips the sign of both off-diagonal terms; composing with runtime states requires the alternating-basis counterparts used by the production RHS (see the positive-neighbor couplings in gkx.terms.linear_terms, e.g. curvature_gradb_contribution). It is not used by the production RHS.

gkx.operators.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, hermite_window: HermiteWindow | None = None) → Array

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

The moment-space field transform adds electrostatic and compressional magnetic terms to m=0 and the parallel-vector-potential term to m=1. The streaming term applies its own pre-derivative field contributions.

gkx.operators.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.

gkx.operators.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)).

Built from energy_operator and therefore expressed in the plain Laguerre basis, whose off-diagonal l +/- 1 signs differ from the alternating (-1)**l basis carried by the runtime state (see gkx.core_velocity.J_l_all). Applying these coefficients to runtime states requires the alternating-basis diamagnetic-drive path in gkx.terms.linear_terms instead; this helper is not used by the production RHS.

gkx.operators.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.

The perpendicular part uses apply_laguerre_x and therefore assumes the plain Laguerre basis with negative l +/- 1 couplings, while the runtime state uses the alternating (-1)**l basis of gkx.core_velocity.J_l_all. Composing with runtime states requires the alternating-basis counterparts in gkx.terms.linear_terms (sign-flipped l +/- 1 couplings); this helper is not used by the production RHS.

gkx.operators.linear.moments.grad_z_periodic(f: Array, dz: float | Array | None = None, kz: Array | None = None, *, ny_full: int | None = None) → Array

Spectral periodic derivative along the last axis.

ny_full marks a complete real-FFT layout: see _real_fft_nyquist_derivative().

gkx.operators.linear.moments.hermite_streaming(G: Array, kpar: Array, vth: float) → Array

Parallel streaming operator acting on the Hermite index.

gkx.operators.linear.moments.lenard_bernstein_eigenvalues(Nl: int, Nm: int, nu_hermite: float, nu_laguerre: float) → Array

Diagonal Lenard-Bernstein rates in Hermite-Laguerre space.

gkx.operators.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.

gkx.operators.linear.moments.shift_axis(arr: Array, offset: int, axis: int) → Array

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

gkx.operators.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.

gkx.operators.linear.params.BUILTIN_COLLISION_OPERATORS: tuple[str, ...] = ('none', 'lenard_bernstein', 'landau', 'coulomb')

Operators applied inside the linear RHS from the cache (every solver path); the others are moment operators handed to the fixed-step integrator. coulomb is a compatibility alias of landau, which reproduces the retired eight-moment table of Frei, Ernst & Ricci (2022) at half its nu.

gkx.operators.linear.params.DAMP_ENDS_LAYER_TRANSITS = 30.0

e-folds per thermal transit of the fastest kinetic species across one damping layer (LinearParams.end_damping_rate()).

Type:

Default end-absorber strength

gkx.operators.linear.params.DEPRECATED_LINEAR_PARAM_ALIASES: dict[str, str] = {'R_over_LTe': 'tprim_e', 'R_over_LTi': 'tprim', 'R_over_Ln': 'fprim'}

Retired LinearParams field names, mapped to what replaced them. The old names said R/L; the values they held were always a/L. Kept working – as constructor keywords and as attributes, both warning – because LinearParams is public API.

class gkx.operators.linear.params.LinearParams(*args, **kwargs)

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

tprim and fprim are the gradient drives, in the units the operator actually consumes: \(a/L_T\) and \(a/L_n\) with \(L_{ref}=a\), which is what the TOML tprim/fprim keys carry and what build_linear_params() copies across unchanged. They are not \(R/L_T\) and \(R/L_n\); multiply by \(R/a\) to quote a result in those. The defaults are the Cyclone base case, matching benchmarks/cases/cyclone_nonlinear_t400.toml and gkx.config.ModelConfig, so \(R/L_T = 2.49 \times 2.77778 = 6.92\).

tprim is the per-species array the RHS diamagnetic drive reads. tprim_e does not enter the RHS at all – it only widens the frequency scale used for eigenvalue branch selection, so a kinetic-electron case must put the electron drive in tprim, not here.

closure: float | Array = 0.0

Per species, 1 closes the Hermite hierarchy with the reflectionless sink on the last moment (hermite_closure_profile()) and drops that species’ |k_z| hypercollisions; 0 keeps the hypercollisions.

collision_model: str = 'landau'

"landau" (linearized Landau, drift-kinetic limit; the default, as in the runtime) or "dougherty" (conserving Lenard-Bernstein form).

Type:

Built-in collision kernel

const_branch_nu_hyper_m()

Return the Hermite hypercollision rate the constant branch uses.

nu_hyper_m drives both hypercollision branches, so a run that wants the shipped |k_z| Hermite model and a constant-branch Laguerre sink would otherwise damp Hermite twice. nu_hyper_m_const splits the two: None keeps the historical behaviour (the constant branch reuses nu_hyper_m), and 0.0 gives a pure Laguerre sink beside an unchanged |k_z| branch. This is the only knob the Laguerre regularization contract (plan 0.5) needs that the operator could not express.

end_damping_rate(z_extent, dtype)

Return the end-absorber rate \(\nu_{end}\), the same on every route.

An explicit damp_ends_rate wins. Otherwise \(\nu_{end} = C\,\max_s v_{ts}\,k_{\parallel,s}/L_w\), with \(k_{\parallel,s}\) = kpar_scale (the geometry’s gradpar) and \(L_w\) = damp_ends_widthfrac times z_extent: C e-folds per thermal transit of the fastest kinetic species across the layer (DAMP_ENDS_LAYER_TRANSITS). No timestep enters.

moments: tuple[tuple[int | None, int | None, float], ...] | None = None

Moments each species keeps, one (Nl, Nm, p_hyper_m) per species, at most the state’s (Nl, Nm); None in a slot keeps the state’s there and None for the field gives every species the state’s. The RHS runs each run of equal species at its own size, so the moments above it stay zero and its hypercollisions are normalized to its own count (LinearCache.groups).

mu_vperp_max: float | None = None

Perpendicular velocity representation. None keeps the Laguerre expansion in mu B. A number selects the nodal magnetic-moment grid (gkx.core_velocity.mu_grid_nodes()) with that vperp_max (inf: Gauss-Laguerre); the Laguerre axis then holds the Nl nodes.

species_slice(a: int, b: int, **overrides: Any) → LinearParams

The parameters of species a:b of a moments run, with overrides.

class gkx.operators.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).

class gkx.operators.linear.params.Species(charge: float, mass: float, density: float, temperature: float, tprim: float, fprim: float, nu: float = 0.0, closure: float = 0.0)

Dimensionless kinetic-species inputs used to build solver parameters.

tprim and fprim are \(a/L_T\) and \(a/L_n\), copied through to LinearParams unchanged; see its docstring for the conversion to \(R/L\).

gkx.operators.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.

gkx.operators.linear.params.build_linear_params(species: Iterable[Species], **overrides: Any) → LinearParams

Build multi-species LinearParams from physical inputs.

Note that kpar_scale is not inferred here either; see linear_params_for_geometry().

gkx.operators.linear.params.linear_params_for_geometry(geometry, **overrides) → LinearParams

Build LinearParams whose parallel scale matches the geometry.

kpar_scale multiplies the parallel derivative, so it has to carry the geometry’s gradpar \(=1/(qR_0)\) for s-alpha. It defaults to 1.0 because LinearParams knows nothing about geometry, which means a hand-built LinearParams() silently streams at the wrong rate – by a factor of 8.4 for a typical Cyclone-like case. The runtime path sets it from the geometry; this is the equivalent for the Python API.

Pass kpar_scale explicitly to override, for instance when reproducing a published normalization that folds \(qR_0\) in elsewhere.

gkx.operators.linear.params.linear_terms_to_term_config(terms: LinearTerms | None, *, nonlinear: float = 0.0) → TermConfig

Convert LinearTerms into the modular TermConfig.

gkx.operators.linear.params.term_config_to_linear_terms(term_cfg: TermConfig | None) → LinearTerms

Convert modular TermConfig into linear-only term weights.

Linear RHS

Linear gyrokinetic RHS assembly entry points.

gkx.operators.linear.rhs.linear_rhs(G: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, terms: LinearTerms | None = None, *, collision_operator: CollisionOperator | None = None) → tuple[Array, Array]

Compute the linear RHS and electrostatic potential from grid/geometry inputs.

gkx.operators.linear.rhs.linear_rhs_cached(G: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | None = None, *, use_jit: bool = True, use_custom_vjp: bool = True, force_electrostatic_fields: bool = False, collision_operator: CollisionOperator | None = None) → tuple[Array, Array]

Compute the linear RHS using precomputed geometry/cache arrays.

Linear Dissipation

Linear collisional, hypercollisional, and damping term contributions.

class gkx.operators.linear.dissipation.CollisionInvariantRates(density: jnp.ndarray, parallel_momentum: jnp.ndarray, thermal_energy: jnp.ndarray)

Long-wavelength collisional rates of the conserved fluid moments.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new CollisionInvariantRates object from a sequence or iterable

_replace(**kwds)

Return a new CollisionInvariantRates object replacing specified fields with new values

density: Array

Alias for field number 0

parallel_momentum: Array

Alias for field number 1

thermal_energy: Array

Alias for field number 2

gkx.operators.linear.dissipation.collision_invariant_rates(contribution: Array) → CollisionInvariantRates

Return density, parallel-momentum, and thermal-energy collision rates.

The Hermite–Laguerre state must have (ell, m) axes first, optionally preceded by species. These are the discrete long-wavelength invariants; finite-Larmor-radius operators require their full gyroaveraged moments.

gkx.operators.linear.dissipation.collision_quadratic_rate(state: Array, contribution: Array, *, weights: Array | None = None) → Array

Return Re <state, C[state]> in the discrete moment norm.

A dissipative collision model has a non-positive rate. Optional broadcastable weights can supply species and spatial quadrature factors.

gkx.operators.linear.dissipation.collisions_contribution(H: Array, *, G: Array | None = None, Jl: Array | None = None, JlB: Array | None = None, b: Array | None = None, nu: Array, collision_lam: Array | None = None, lb_lam: Array | None = None, weight: Array, mu_grid: tuple[Array, Array] | None = None) → Array

mu_grid = (rate, x) selects the nodal perpendicular operator and energy weights.

gkx.operators.linear.dissipation.drift_kinetic_sugama_six_moment_contribution(state: Array, *, nu: Array, weight: Array = Array(1., dtype=float32, weak_type=True)) → Array

Apply Frei, Ernst & Ricci (2022), equations (C6a)–(C6f).

gkx.operators.linear.dissipation.hyperdiffusion_contribution(G: Array, *, kx: Array, ky: Array, dealias_mask: Array, D_hyper: Array, p_hyper_kperp: Array, weight: Array, ny_full: int | None = None, ky_cut: Array | None = None) → Array

Hyperdiffusion in k_perp following Laguerre-Hermite conventions.

ny_full is the length of the two-sided ky axis, which is what the normalizing cutoff kperp2_max is a property of: the largest dealiased row is at index (Ny - 1) // 3 in both layouts, and deriving it from the stored row count would halve the cutoff on a half-spectrum grid and inflate Dfac by 4 ** p_hyper_kperp with no shape error to show for it. A row selection does not contain that row at all, so it passes the parent’s |ky| there as ky_cut; normalizing by the selected row’s own ky would hand every scanned mode the full D_hyper.

Linear Implicit Solvers

Implicit linear solve policies for cache-backed gyrokinetic operators.

class gkx.solvers_linear_implicit.ImplicitSolveStats(max_relative_residual: jax.Array, max_iterations: jax.Array, solves: jax.Array, unconverged_solves: jax.Array)

Convergence summary of every implicit GMRES solve in one run.

Each field is a scalar array, so the summary rides through lax.scan and fori_loop carries without a host callback. max_relative_residual is the largest SOLVAX true residual ||b - A x|| / ||b||; with right preconditioning this is the physical norm. unconverged_solves counts solves whose SOLVAX converged flag was false at rtol=implicit_tol.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new ImplicitSolveStats object from a sequence or iterable

_replace(**kwds)

Return a new ImplicitSolveStats object replacing specified fields with new values

max_iterations: Array

Alias for field number 1

max_relative_residual: Array

Alias for field number 0

solves: Array

Alias for field number 2

unconverged_solves: Array

Alias for field number 3

class gkx.solvers_linear_implicit.ImplicitSolveSummary(max_relative_residual: float, max_iterations: int, solves: int, unconverged_solves: int)

Host-side ImplicitSolveStats of one run, as Python scalars.

gkx.solvers_linear_implicit._build_field_corrected_shifted_preconditioner(G0: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | TermConfig, sigma: Array, *, coarse: bool = False) → Callable[[Array], Array]

Add the exact low-moment field coupling to the Hermite-line inverse.

For fixed geometry the field solve is a linear moment map F: G -> f. Linear gyrokinetic forcing by those fields is another map R: f -> dG, so the shifted operator separates as

A - sigma I = (A_0 - sigma I) + R F.

F has at most three active fields per spatial spectral point, independent of velocity resolution. Sequentially probing R and applying F avoids materializing either the kinetic operator or its transpose. The Hermite-line inverse approximates A_0 - sigma I while a velocity-independent capacitance solve represents the field coupling.

gkx.solvers_linear_implicit._build_shifted_hermite_preconditioner(G0: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | TermConfig, sigma: Array, *, coarse: bool = False) → Callable[[Array], Array]

Approximate (A - sigma I)^-1 with the FFT/Hermite line solver.

The implicit preconditioner approximates (I - dt A)^-1. Setting dt = 1 / sigma and multiplying by -1 / sigma gives the corresponding inverse for A - sigma I while reusing exactly the same physics-aware streaming, drift, and damping factors.

gkx.solvers_linear_implicit._integrate_linear_implicit_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, *, terms: LinearTerms | 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_preconditioner: Callable[[Array], Array] | str | None = None, checkpoint: bool = False, sample_stride: int = 1, return_solve_stats: bool = False) → tuple[Any, ...]

Implicit linear integrator using GMRES with a diagonal preconditioner.

Returns (G_out, phi_t), or (G_out, phi_t, stats) with return_solve_stats=True. The traced scan cannot raise on an unconverged solve, so stats (ImplicitSolveStats) is the convergence channel; the runtime checks it with require_converged_implicit_solves().

gkx.solvers_linear_implicit.implicit_solve_payload(summary: ImplicitSolveSummary | None) → dict[str, Any]

Flatten an implicit-solve summary into scalar keys; None when none ran.

This is the one spelling of those keys. A runtime result’s summary() and the saved *.summary.json both read it, so a run held in memory and the same run read back from disk report the same status under the same names.

gkx.solvers_linear_implicit.require_converged_implicit_solves(stats: ImplicitSolveStats, *, label: str) → ImplicitSolveSummary

Fail closed at the host boundary when any implicit solve did not converge.

The traced scans cannot raise, so they carry ImplicitSolveStats and the runtime checks it here, as the eigen gates do for their residuals.

Linear Integrators

Linear time-integration and diagnostic sampling policies.

gkx.solvers_linear_integrators._integrate_linear_cached_impl(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, method: str = 'rk4', checkpoint: bool = False, terms: LinearTerms | None = None, sample_stride: int = 1, show_progress: bool = False, parallel: Any | None = None, force_electrostatic_fields: bool = False, collision_operator: CollisionOperator | None = None) → tuple[Array, Array]

Time integrate the linear system using cached geometry arrays.

gkx.solvers_linear_integrators.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_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, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...]

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

return_solve_stats=True appends the convergence summary of the implicit GMRES solves (ImplicitSolveStats) for method="implicit" and None for methods without an implicit solve.

gkx.solvers_linear_integrators.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, collision_operator: Any | 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_preconditioner: Callable[[Array], Array] | str | None = None, return_solve_stats: bool = False) → tuple[Any, ...]

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

record_hl_energy appends hl_t. return_solve_stats=True appends the implicit GMRES convergence summary as the last element (ImplicitSolveStats, or None for methods without an implicit solve).

Linear Parallel RHS

Velocity-parallel linear RHS helpers.

gkx.solvers_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.

gkx.solvers_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.

gkx.solvers_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.

gkx.solvers_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.

gkx.solvers_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.

gkx.solvers_linear_parallel.linear_rhs_electrostatic_species_hermite_sharded(G: Array, cache: LinearCache, params: LinearParams, *, terms: LinearTerms | None = None, species_chunks: int = 2, hermite_chunks: int = 2, devices: Any | None = None) → tuple[Array, Array]

Compute the electrostatic RHS on a species–Hermite mesh.

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

Compute electrostatic linear terms with one species shard per device.

gkx.solvers_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) → 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.

gkx.solvers_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 benchmark-compatible electrostatic streaming field term. It is limited to periodic field-line grids and excludes electromagnetic fields by construction.

gkx.solvers_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.

gkx.solvers_linear_parallel.prepare_electrostatic_species_hermite_state(state: Array, *, species_chunks: int = 2, hermite_chunks: int = 2, devices: Any | None = None) → Array

Place a mixed-route state once before entering a time-integration scan.

gkx.solvers_linear_parallel.prepare_electrostatic_species_inputs(G: Array, cache: LinearCache, params: LinearParams, *, num_devices: int | None = None, devices: Any | None = None, replicate_cache: bool = True) → tuple[Array, LinearCache, LinearParams]

Place concrete inputs from host memory; preserve traced inputs for AD/JIT.

Linear Parallel Policy

Shared policies for opt-in velocity-parallel linear RHS routes.

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

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

Linear Parallel Streaming

Velocity-sharded streaming routes for the linear RHS.

gkx.solvers_linear_parallel_streaming._electrostatic_streaming_field_rhs(G6: Array, *, phi: Array, Jl: Array, tz: Array, vth: Array) → Array

Build the pre-derivative electrostatic streaming field term.

gkx.solvers_linear_parallel_streaming._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.

gkx.solvers_linear_parallel_streaming.linear_rhs_electrostatic_species_hermite_sharded(G: Array, cache: LinearCache, params: LinearParams, *, terms: LinearTerms | None = None, species_chunks: int = 2, hermite_chunks: int = 2, devices: Any | None = None) → tuple[Array, Array]

Compute the electrostatic RHS on a species–Hermite mesh.

gkx.solvers_linear_parallel_streaming.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 benchmark-compatible electrostatic streaming field term. It is limited to periodic field-line grids and excludes electromagnetic fields by construction.

gkx.solvers_linear_parallel_streaming.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.

gkx.solvers_linear_parallel_streaming.prepare_electrostatic_species_hermite_state(state: Array, *, species_chunks: int = 2, hermite_chunks: int = 2, devices: Any | None = None) → Array

Place a mixed-route state once before entering a time-integration scan.

Linear Parallel Electrostatic Slices

Velocity-sharded electrostatic slice routes for the linear RHS.

gkx.solvers_linear_parallel_electrostatic._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.

gkx.solvers_linear_parallel_electrostatic.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.

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

Compute electrostatic linear terms with one species shard per device.

gkx.solvers_linear_parallel_electrostatic.prepare_electrostatic_species_inputs(G: Array, cache: LinearCache, params: LinearParams, *, num_devices: int | None = None, devices: Any | None = None, replicate_cache: bool = True) → tuple[Array, LinearCache, LinearParams]

Place concrete inputs from host memory; preserve traced inputs for AD/JIT.

Linear Krylov Solvers

Public facade preserving the API above independently tested eigenmode kernels.

class gkx.solvers_linear_krylov.KrylovConfig(krylov_dim: int = 24, restarts: int = 2, omega_min_factor: float = 0.0, omega_target_factor: float = 0.0, omega_cap_factor: float = 2.0, omega_sign: int = 0, method: str = 'auto', power_iters: int = 40, power_dt: float = 0.01, propagator_steps: int = 1, shift: complex | None = None, shift_source: str = 'propagator', shift_tol: float = 0.0001, shift_maxiter: int = 50, shift_restart: int = 20, shift_solve_method: str = 'batched', shift_preconditioner: str | None = 'auto', shift_precond_alpha: float | complex | None = None, shift_precond_block_solve: str = 'auto', shift_selection: str = 'targeted', shift_outer_residual_tol: float = 1e-06, mode_family: str = 'auto', fallback_method: str = 'propagator', fallback_real_floor: float = -1e-06, continuation: bool = False, continuation_selection: str = 'overlap', certify: bool = True)

Controls for the Krylov-based eigen solver.

The default method="auto" runs the Hermite-banded direct shift-invert where its blocks fit in memory and the residual-certified adaptive eigensolve otherwise or when the banded pair is rejected. Every method returns only a pair that passes an original-operator residual gate unless certify=False.

gkx.solvers_linear_krylov._mode_family_sign(mode_family: str) → int

Map named branch families to the sign convention for physical frequency.

gkx.solvers_linear_krylov._omega_scale(cache: LinearCache, params: LinearParams) → Array

Return the frequency scale used for branch and target selection.

gkx.solvers_linear_krylov._physical_omega(imag_part: Array) → Array

Map eigenvalue imaginary part to reported physical frequency.

gkx.solvers_linear_krylov._select_by_overlap(eigvecs: Array, V: Array, v_ref: Array, mask: Array, fallback_idx: Array) → Array

Select the eigenpair with maximal overlap to v_ref within mask.

gkx.solvers_linear_krylov._select_by_target(real_part: Array, imag_part: Array, mask: Array, omega_scale: Array, omega_target_factor: float, omega_sign: int, fallback_idx: Array) → Array

Select the branch nearest the requested physical-frequency target.

gkx.solvers_linear_krylov.adaptive_propagator_eigenpair(v0: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | None = None, *, krylov_dim: int = 16, max_restarts: int = 4, tol: float = 1e-10, chunk_horizon: float = 30.0, stability_dimension: int = 12, stability_probe_count: int = 2, stability_safety: float = 0.9, max_stability_retries: int = 2, restart_krylov_dim: int | None = None, candidate_count: int = 1, continuation_vector: Array | None = None, continuation_covector: Array | None = None, continuation_overlap_floor: float = 0.0, continuation_spectral_gap_floor: float = 0.0, exponential_krylov_dim: int | None = None, exponential_horizon: float = 5.0) → AdaptivePropagatorSolution

Adapt RK4 stability, horizon, and corrective-subspace cost.

A supplied continuation covector selects the certified candidate maximizing |w_previous**H v_candidate|. A right reference is the normalized-overlap fallback when a left vector is unavailable. This follows one physical mode through a growth-rate crossing instead of silently switching to the newly dominant branch.

tol gates a relative residual built from operator applications in the dtype of v0, so it is raised to that dtype’s noise floor. The default is a strict float64 gate; the same call certifies in float32 at the tightest residual single precision can actually deliver.

gkx.solvers_linear_krylov.build_shift_invert_preconditioner(v: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, sigma: Array, mode: str | None, factors: Pr3Factors | None = None) → tuple[Array | None, Callable[[Array], Array] | None]

Build the preconditioner used inside shift-invert Krylov GMRES solves.

None/"none" are the explicit opt-outs; any other unrecognized name raises. This used to return (None, None) for those too, so a typo disabled preconditioning while the solve reported itself as preconditioned. "auto" must not arrive here – it is a real policy resolved one level up by gkx.solvers_linear_krylov._automatic_shift_preconditioner(), and mapping it to "damping" would hand a damping diagonal to a caller that asked for a physics-aware line solve.

pr3-cm is the one mode with host-side setup: its z-local blocks cost Nl * Nm operator probes and a factorization, built once per shift by gkx.solvers_linear_precond_pr3.build_pr3_factors() and passed in as factors rather than traced. Asking for it without them raises.

gkx.solvers_linear_krylov.certifiable_residual_tolerance(tol: float, dtype: str | type[Any] | dtype | SupportsDType) → float

Raise a residual gate to the noise floor of the arithmetic feeding it.

A relative residual bottoms out a few eps above zero at every precision, so an absolute gate that is comfortable in float64 can sit below anything float32 can ever produce. Such a gate certifies nothing and only burns restarts on an answer that is already as good as the dtype allows. A request the working precision can reach is honored unchanged.

gkx.solvers_linear_krylov.dominant_eigenpair(v0: Array, cache: LinearCache, params: LinearParams, terms: LinearTerms | None = None, *, v_ref: Array | None = None, select_overlap: bool = False, krylov_dim: int = 24, restarts: int = 2, omega_min_factor: float = 0.0, omega_target_factor: float = 0.0, omega_cap_factor: float = 2.0, omega_sign: int = 0, method: str = 'auto', power_iters: int = 40, power_dt: float = 0.01, propagator_steps: int = 1, shift: complex | None = None, shift_source: str = 'propagator', shift_tol: float = 0.0001, shift_maxiter: int = 50, shift_restart: int = 20, shift_solve_method: str = 'batched', shift_preconditioner: str | None = 'auto', shift_precond_alpha: float | complex | None = None, shift_precond_block_solve: str = 'auto', shift_selection: str = 'targeted', shift_outer_residual_tol: float = 1e-06, mode_family: str = 'auto', fallback_method: str = 'propagator', fallback_real_floor: float = -1e-06, certify: bool = True, status_callback: Callable[[str], None] | None = None, return_status: bool = False) → tuple[Any, ...]

Python wrapper for cached matrix-free and sparse eigen solvers.

No method returns a pair above its original-operator residual gate: the adaptive, shift-invert and sparse routes raise on rejection, and the raw power, propagator and arnoldi routes raise above certifiable_residual_tolerance(shift_outer_residual_tol, dtype) unless certify=False explicitly accepts an uncertified pair.

Returns (eigenvalue, eigenvector); return_status=True appends the pair’s EigenSolveStatus (residual, gate, route, inner solves).

gkx.solvers_linear_krylov.dominant_eigenpair_cached(v0: Array, v_ref: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, krylov_dim: int, restarts: int, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, select_overlap: bool) → tuple[Array, Array]

Approximate the dominant eigenvalue (max real part) with restarted Arnoldi.

gkx.solvers_linear_krylov.dominant_eigenpair_power(v0: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, iterations: int, dt: float) → tuple[Array, Array]

Power iteration on an explicit-Euler propagator to target the rightmost mode.

gkx.solvers_linear_krylov.dominant_eigenpair_propagator_cached(v0: Array, v_ref: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, krylov_dim: int, restarts: int, dt: float, propagator_steps: int, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, select_overlap: bool) → tuple[Array, Array]

Filter by finite-horizon growth and return a physical Rayleigh pair.

Linear Eigenmode Solver Internals

gkx.solvers_linear_krylov owns KrylovConfig and the public status-reporting wrapper; gkx.solvers_linear_krylov_algorithms owns operator application, branch selection, shift-invert preconditioning, and the Arnoldi iterations.

Matrix-free operators and compiled Krylov eigenmode algorithms.

gkx.solvers_linear_krylov_algorithms._advance_rk4(state: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, dt: Array) → Array

Advance with a full-operator polynomial that preserves eigenvectors.

gkx.solvers_linear_krylov_algorithms._mode_family_sign(mode_family: str) → int

Map named branch families to the sign convention for physical frequency.

gkx.solvers_linear_krylov_algorithms._omega_scale(cache: LinearCache, params: LinearParams) → Array

Return the frequency scale used for branch and target selection.

gkx.solvers_linear_krylov_algorithms._physical_omega(imag_part: Array) → Array

Map eigenvalue imaginary part to reported physical frequency.

gkx.solvers_linear_krylov_algorithms._rayleigh_quotient(vector: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig) → Array

Return the least-residual eigenvalue for a fixed physical Ritz vector.

gkx.solvers_linear_krylov_algorithms._select_by_overlap(eigvecs: Array, V: Array, v_ref: Array, mask: Array, fallback_idx: Array) → Array

Select the eigenpair with maximal overlap to v_ref within mask.

gkx.solvers_linear_krylov_algorithms._select_by_target(real_part: Array, imag_part: Array, mask: Array, omega_scale: Array, omega_target_factor: float, omega_sign: int, fallback_idx: Array) → Array

Select the branch nearest the requested physical-frequency target.

gkx.solvers_linear_krylov_algorithms.build_shift_invert_preconditioner(v: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, sigma: Array, mode: str | None, factors: Pr3Factors | None = None) → tuple[Array | None, Callable[[Array], Array] | None]

Build the preconditioner used inside shift-invert Krylov GMRES solves.

None/"none" are the explicit opt-outs; any other unrecognized name raises. This used to return (None, None) for those too, so a typo disabled preconditioning while the solve reported itself as preconditioned. "auto" must not arrive here – it is a real policy resolved one level up by gkx.solvers_linear_krylov._automatic_shift_preconditioner(), and mapping it to "damping" would hand a damping diagonal to a caller that asked for a physics-aware line solve.

pr3-cm is the one mode with host-side setup: its z-local blocks cost Nl * Nm operator probes and a factorization, built once per shift by gkx.solvers_linear_precond_pr3.build_pr3_factors() and passed in as factors rather than traced. Asking for it without them raises.

gkx.solvers_linear_krylov_algorithms.dominant_eigenpair_cached(v0: Array, v_ref: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, krylov_dim: int, restarts: int, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, select_overlap: bool) → tuple[Array, Array]

Approximate the dominant eigenvalue (max real part) with restarted Arnoldi.

gkx.solvers_linear_krylov_algorithms.dominant_eigenpair_power(v0: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, iterations: int, dt: float) → tuple[Array, Array]

Power iteration on an explicit-Euler propagator to target the rightmost mode.

gkx.solvers_linear_krylov_algorithms.dominant_eigenpair_propagator_cached(v0: Array, v_ref: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, krylov_dim: int, restarts: int, dt: float, propagator_steps: int, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, select_overlap: bool) → tuple[Array, Array]

Filter by finite-horizon growth and return a physical Rayleigh pair.

Nonlinear Diagnostic State

Diagnostic tuple assembly for nonlinear time integration.

This module owns the pure state-to-diagnostics assembly logic, with diagnostic kernels injected by gkx.solvers_nonlinear_diagnostics so tests and interactive debugging can still monkeypatch the module-level functions.

class gkx.operators.nonlinear.diagnostic_state.NonlinearDiagnosticKernels(instantaneous_growth_rate_step: Callable[[...], Any], phi2_resolved: Callable[[...], Any], zonal_phi_mode_kxt: Callable[[...], Any], zonal_phi_line_kxt: Callable[[...], Any], distribution_free_energy: Callable[[...], Any], distribution_free_energy_resolved: Callable[[...], Any], electrostatic_field_energy: Callable[[...], Any], electrostatic_field_energy_resolved: Callable[[...], Any], magnetic_vector_potential_energy: Callable[[...], Any], magnetic_vector_potential_energy_resolved: Callable[[...], Any], heat_flux_species: Callable[[...], Any], heat_flux_channel_resolved_species: Callable[[...], Any], particle_flux_species: Callable[[...], Any], particle_flux_channel_resolved_species: Callable[[...], Any], turbulent_heating_species: Callable[[...], Any], turbulent_heating_resolved_species: Callable[[...], Any])

Facade-injected diagnostic callables used by nonlinear integrators.

gkx.operators.nonlinear.diagnostic_state.compute_nonlinear_diagnostic_tuple(G_state: Array, fields_state: FieldState, G_prev_step: Array, fields_prev_step: FieldState, dt_step: Array, *, grid: SpectralGrid, cache: LinearCache, params: LinearParams, vol_fac: Array, flux_fac: Array, mask: Array, z_idx: int, use_dealias: bool, real_dtype: Any, omega_ky_index: int | None, omega_kx_index: int | None, flux_scale: float, wphi_scale: float, resolved_diagnostics: bool, kernels: NonlinearDiagnosticKernels) → tuple[Any, ...]

Build the nonlinear scan diagnostic tuple for one state.

gkx.operators.nonlinear.diagnostic_state.make_nonlinear_diagnostic_tuple_fn(*, grid: SpectralGrid, cache: LinearCache, params: LinearParams, vol_fac: Array, flux_fac: Array, mask: Array, z_idx: int, use_dealias: bool, real_dtype: Any, omega_ky_index: int | None, omega_kx_index: int | None, flux_scale: float, wphi_scale: float, resolved_diagnostics: bool, kernels: NonlinearDiagnosticKernels) → Callable[[Array, FieldState, Array, FieldState, Array], tuple[Any, ...]]

Return a reusable state-to-diagnostic tuple closure for scan policies.

Nonlinear Collision Split Helpers

Collision and hypercollision split helpers for nonlinear integrations.

class gkx.operators.nonlinear.collisions.NonlinearCollisionSplitPolicy(active: bool, rhs_terms: TermConfig, damping: Array | None)

Collision split settings shared by explicit and IMEX diagnostics.

gkx.operators.nonlinear.collisions._apply_collision_split(G: Array, damping: Array, dt_local: Array, scheme: str) → Array

Apply a diagonal collision/hypercollision split update.

gkx.operators.nonlinear.collisions._collision_damping(cache: LinearCache, params: LinearParams, term_cfg: TermConfig, real_dtype: dtype, *, squeeze_species: bool) → Array

Assemble the diagonal hypercollision damping safe to split.

The conserving collision operator includes non-diagonal field-particle corrections and must remain in the RHS unless an operator supplies its own mathematically valid split update.

gkx.operators.nonlinear.collisions.build_nonlinear_collision_split_policy(cache: ~gkx.operators.linear.cache_model.LinearCache, params: ~gkx.operators.linear.params.LinearParams, term_cfg: ~gkx.terms.config.TermConfig, real_dtype: ~numpy.dtype, *, squeeze_species: bool, collision_split: bool, collision_damping_fn: ~typing.Callable[[...], ~jax.jaxlib._jax.Array] = <function _collision_damping>) → NonlinearCollisionSplitPolicy

Build collision splitting weights and RHS terms for nonlinear scans.

Nonlinear Helpers

Helper policies and operators for nonlinear gyrokinetic drivers.

class gkx.operators.nonlinear.policies.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.

class gkx.operators.nonlinear.policies.NonlinearCollisionSplitPolicy(active: bool, rhs_terms: TermConfig, damping: Array | None)

Collision split settings shared by explicit and IMEX diagnostics.

class gkx.operators.nonlinear.policies.NonlinearDiagnosticSetup(geom: Any, cache: LinearCache, vol_fac: Array, flux_fac: Array, mask: Array, z_idx: int, use_dealias: bool, project_state: Callable[[Array], Array])

Shared cache, weights, masks, and projection policy for diagnostics.

class gkx.operators.nonlinear.policies.NonlinearTimeStepPolicy(dt_init: Array, progress_total: Array, update_dt: Callable[[FieldState, Array], Array], cfl_scales: Array | None = None)

Initial step, progress horizon, and adaptive update callable.

class gkx.operators.nonlinear.policies.ShearingCoordinateUpdate(state: jnp.ndarray, effective_kx: jnp.ndarray, phase: jnp.ndarray, cumulative_mode_shift: jnp.ndarray, incremental_mode_shift: jnp.ndarray)

State and spectral coordinates after one equilibrium-flow-shear update.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new ShearingCoordinateUpdate object from a sequence or iterable

_replace(**kwds)

Return a new ShearingCoordinateUpdate object replacing specified fields with new values

cumulative_mode_shift: Array

Alias for field number 3

effective_kx: Array

Alias for field number 1

incremental_mode_shift: Array

Alias for field number 4

phase: Array

Alias for field number 2

state: Array

Alias for field number 0

gkx.operators.nonlinear.policies.advance_shearing_coordinates(state: Array, *, kx: Array, ky: Array, x0: Array | float, shear_rate: Array | float, previous_time: Array | float, time: Array | float, dealias_mask: Array | None = None) → ShearingCoordinateUpdate

Advance a Fourier state in continuously shearing coordinates.

For equilibrium \(E\times B\) shear, each shearing wave follows

\[k_x^*(t) = k_x(0) - k_y \gamma_E t.\]

The integer part of this displacement remaps Fourier amplitudes to the nearest radial mode. The sub-grid remainder is returned as the real-space phase exp(1j * delta_kx * x) and in effective_kx. Integer remap decisions are treated as locally constant under autodiff, while the continuous wavenumber and phase retain their exact tangent away from the measure-zero crossing events.

state uses (..., ky, kx, z) ordering. Modes shifted beyond the supplied two-thirds mask are discarded rather than wrapped into the resolved band.

gkx.operators.nonlinear.policies.build_nonlinear_collision_split_policy(cache: ~gkx.operators.linear.cache_model.LinearCache, params: ~gkx.operators.linear.params.LinearParams, term_cfg: ~gkx.terms.config.TermConfig, real_dtype: ~numpy.dtype, *, squeeze_species: bool, collision_split: bool, collision_damping_fn: ~typing.Callable[[...], ~jax.jaxlib._jax.Array] = <function _collision_damping>) → NonlinearCollisionSplitPolicy

Build collision splitting weights and RHS terms for nonlinear scans.

gkx.operators.nonlinear.policies.build_nonlinear_diagnostic_setup(G0: Array, grid: SpectralGrid, geom: Any, params: LinearParams, *, cache: LinearCache | None, use_dealias_mask: bool, z_index: int | None, compressed_real_fft: bool, fixed_mode_ky_index: int | None, fixed_mode_kx_index: int | None, ensure_geometry_fn: Callable[[...], Any], build_cache_fn: Callable[[...], LinearCache], quadrature_weights_fn: Callable[[...], tuple[Array, Array]], omega_mask_fn: Callable[[...], Array], midplane_index_fn: Callable[[int], int]) → NonlinearDiagnosticSetup

Build the shared diagnostic setup used by explicit and IMEX scans.

gkx.operators.nonlinear.policies.build_nonlinear_imex_operator(G0: Array, cache: LinearCache, params: LinearParams, dt: float, *, terms: TermConfig | None = None, implicit_preconditioner: str | None = None, compressed_real_fft: bool = True, build_implicit_operator_fn: Callable[[...], tuple[Any, ...]] | None = None) → IMEXLinearOperator

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

gkx.operators.nonlinear.policies.build_nonlinear_time_step_policy(grid: SpectralGrid, geom: Any, params: LinearParams, cache: LinearCache, *, method: str, dt: float, steps: int, fixed_dt: bool, dt_min: float, dt_max: float | None, cfl: float, cfl_fac: float | None, compressed_real_fft: bool, real_dtype: Any, resolve_cfl_fac_fn: Callable[[str, float | None], float], linear_frequency_bound_fn: Callable[[...], Any], laguerre_velocity_max_fn: Callable[[...], float], cfl_frequency_components_fn: Callable[[...], tuple[Array, Array]]) → NonlinearTimeStepPolicy

Build the fixed/adaptive nonlinear time-step update policy.

Nonlinear Projection Helpers

State projection helpers for nonlinear spectral integrations.

class gkx.operators.nonlinear.projection.ShearingCoordinateUpdate(state: jnp.ndarray, effective_kx: jnp.ndarray, phase: jnp.ndarray, cumulative_mode_shift: jnp.ndarray, incremental_mode_shift: jnp.ndarray)

State and spectral coordinates after one equilibrium-flow-shear update.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new ShearingCoordinateUpdate object from a sequence or iterable

_replace(**kwds)

Return a new ShearingCoordinateUpdate object replacing specified fields with new values

cumulative_mode_shift: Array

Alias for field number 3

effective_kx: Array

Alias for field number 1

incremental_mode_shift: Array

Alias for field number 4

phase: Array

Alias for field number 2

state: Array

Alias for field number 0

gkx.operators.nonlinear.projection._make_compressed_real_fft_projector(*, ny_full: int, nx: int, rows: int | None = None) → Callable[[Array], Array]

Return the Hermitian projector for a compressed real-FFT ky axis.

ny_full is the length of the two-sided axis and rows how many of them the state stores; rows is None means the full axis, which is what every caller meant before the half layout existed.

On a two-sided state the projector is to_full(to_half(G)): the bracket transforms the first Nyc rows and rebuilds the rest, so the state has to be put back on the reality-condition subspace after every stage. On a half-spectrum state (plan 5.3 N3) the reality condition holds by construction and the projector is the identity – not an approximation of the old one but exactly it, because to_full(to_half(.)) restores the unstored rows and imposes nothing on the stored ones. The self-conjugate rows’ internal constraint was never part of this projector in either layout; symmetrize_self_conjugate_rows is what imposes that, and the bracket’s irfft2 already discards the anti-symmetric part it would remove.

The layout is read from the arguments rather than from the cache’s ky, which is rho_star * grid.ky and therefore a tracer whenever the cache is built inside a trace; asking it for its sign pattern refused every compressed nonlinear gradient, for a fact that is fixed by the grid and carries no derivative either way.

gkx.operators.nonlinear.projection._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.

gkx.operators.nonlinear.projection._make_hermitian_projector(ky_vals: ndarray, nx: int) → Callable[[Array], Array]

Return a stable projector for one full-ky grid signature.

gkx.operators.nonlinear.projection._make_nonlinear_state_projector(fixed_state: Array | None, *, ky_vals: ndarray, nx: int, compressed_real_fft: bool, fixed_mode_ky_index: int | None, fixed_mode_kx_index: int | None) → Callable[[Array], Array]

Compose fixed-mode and Hermitian projections for nonlinear state scans.

gkx.operators.nonlinear.projection.advance_shearing_coordinates(state: Array, *, kx: Array, ky: Array, x0: Array | float, shear_rate: Array | float, previous_time: Array | float, time: Array | float, dealias_mask: Array | None = None) → ShearingCoordinateUpdate

Advance a Fourier state in continuously shearing coordinates.

For equilibrium \(E\times B\) shear, each shearing wave follows

\[k_x^*(t) = k_x(0) - k_y \gamma_E t.\]

The integer part of this displacement remaps Fourier amplitudes to the nearest radial mode. The sub-grid remainder is returned as the real-space phase exp(1j * delta_kx * x) and in effective_kx. Integer remap decisions are treated as locally constant under autodiff, while the continuous wavenumber and phase retain their exact tangent away from the measure-zero crossing events.

state uses (..., ky, kx, z) ordering. Modes shifted beyond the supplied two-thirds mask are discarded rather than wrapped into the resolved band.

gkx.operators.nonlinear.projection.hermitian_projector_for_signature(signature: tuple[int, bool, int]) → Callable[[Array], Array]

Return the projector of a hermitian_projector_signature() tuple.

gkx.operators.nonlinear.projection.hermitian_projector_signature(ky_vals: ndarray, nx: int) → tuple[int, bool, int]

Return the hashable grid signature _make_hermitian_projector() keys on.

The projector depends on the ky axis layout and the kx count, not on a wavenumber, so a caller that has to read the axis on the host anyway can read it once and then carry this tuple. It is hashable, which a grid is not, so it can be a jax.jit static argument and let a compiled window rebuild the same projector without touching a traced grid.

Nonlinear RHS

Nonlinear RHS assembly helpers.

This module owns the performance-critical RHS composition path. The nonlinear integration solvers in gkx.solvers_nonlinear_diagnostics pass their module-level callables into these helpers so existing tests and downstream monkeypatch/debug workflows keep the same seams while the implementation stays isolated and easier to profile.

gkx.operators.nonlinear.rhs.linear_rhs_jit_for_terms_impl(term_cfg: ~gkx.terms.config.TermConfig, *, electrostatic_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~gkx.terms.config.FieldState]] = <PjitFunction of <function assemble_rhs_cached_electrostatic_jit>>, full_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~gkx.terms.config.FieldState]] = <function assemble_rhs_cached_jit>, is_static_zero_fn: ~typing.Callable[[object], bool] = <function _is_static_zero>) → Callable[[...], tuple[Array, FieldState]]

Return the narrowest compiled linear RHS route for term_cfg.

gkx.operators.nonlinear.rhs.nonlinear_em_term_cached_impl(G: ~jax.jaxlib._jax.Array, cache: ~gkx.operators.linear.cache_model.LinearCache, params: ~gkx.operators.linear.params.LinearParams, term_cfg: ~gkx.terms.config.TermConfig, *, real_dtype: ~typing.Any | None = None, external_phi: ~jax.jaxlib._jax.Array | float | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', radial_phase: ~jax.jaxlib._jax.Array | None = None, fields_fn: ~typing.Callable[[...], ~gkx.terms.config.FieldState], nonlinear_contribution_fn: ~typing.Callable[[...], ~jax.jaxlib._jax.Array] = <function nonlinear_em_contribution>) → Array

Return the explicit electromagnetic nonlinear term used by IMEX paths.

gkx.operators.nonlinear.rhs.nonlinear_rhs_cached_impl(G: ~jax.jaxlib._jax.Array, cache: ~gkx.operators.linear.cache_model.LinearCache, params: ~gkx.operators.linear.params.LinearParams, terms: ~gkx.terms.config.TermConfig | None = None, *, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: ~jax.jaxlib._jax.Array | float | None = None, collision_operator: ~gkx.operators.collision.CollisionOperator | None = None, radial_phase: ~jax.jaxlib._jax.Array | None = None, differentiable: bool = False, electrostatic_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~gkx.terms.config.FieldState]] = <PjitFunction of <function assemble_rhs_cached_electrostatic_jit>>, full_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~gkx.terms.config.FieldState]] = <function assemble_rhs_cached_jit>, differentiable_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~gkx.terms.config.FieldState]] = <function assemble_rhs_cached>, is_static_zero_fn: ~typing.Callable[[object], bool] = <function _is_static_zero>, nonlinear_contribution_fn: ~typing.Callable[[...], ~jax.jaxlib._jax.Array] = <function nonlinear_em_contribution>) → tuple[Array, FieldState]

Compute the assembled nonlinear RHS and electromagnetic field state.

Nonlinear Bracket Kernels

Pseudo-spectral bracket kernels for nonlinear gyrokinetic terms.

gkx.operators.nonlinear.brackets._complete_hermitian_ky(positive_ky: Array, ny_full: int, nx: int, conjugate_kx: Any | None = None) → Array

Widen a ky >= 0 block onto the two-sided axis.

The rule lives in gkx.core_ky_layout; this wrapper keeps the positional signature the nonlinear kernels and their tests already use.

gkx.operators.nonlinear.brackets._fft2_xy_remove_radial_phase(x: Array, radial_phase: Array) → Array

Return a physical field to its nearest-cell shearing-wave coefficients.

gkx.operators.nonlinear.brackets._fft_scales(ky_grid: Array, *, real_dtype: dtype, fft_norm: float | None = None, ny_full: int | None = None) → tuple[Array, Array]

Return the (ifft, fft) normalizations for an Ny x Nx transform.

The pair is fixed by the physical grid, Ny * Nx, not by how many ky rows the operand stores: a half-spectrum operand is transformed with irfft2/rfft2 over the same Ny real-space rows.

gkx.operators.nonlinear.brackets._ifft2_xy_with_radial_phase(x: Array, radial_phase: Array) → Array

Transform a shearing wave after applying its fractional radial phase.

gkx.operators.nonlinear.brackets._spectral_bracket_full_core(G_hat: Array, chi_hat: Array, *, kx_grid: Array, ky_grid: Array, dealias_mask: Array, kxfac: Array, fft_norm: float | None = None, radial_phase: Array | None = None, ny_full: int | None = None, multiple_fields: bool, combine: Callable[[Array], Array] | None = None) → Array

Return the bracket by full complex transforms on a two-sided ky axis.

This route multiplies the whole two-sided spectrum in real space, so it has no half-spectrum form: a state carrying only ky >= 0 is refused here rather than transformed as if its rows were the full axis, which would silently compute the bracket of a different field.

gkx.operators.nonlinear.brackets._spectral_bracket_half_core(G_hat: Array, chi_hat: Array, *, kx_grid: Array, ky_grid: Array, dealias_mask: Array, fft_norm: float | None = None, radial_phase: Array | None = None, ny_full: int | None = None, multiple_fields: bool, combine: Callable[[Array], Array] | None = None) → tuple[Array, int, int]

Return the bracket on the ky >= 0 rows, with (ny_full, nx).

This is the whole pseudo-spectral kernel: both operands are read on their non-negative rows, transformed with irfft2, multiplied in real space, and transformed back with rfft2, which lands the product back on Nyc = 1 + Ny // 2 rows. Nothing in it needs the negative half, so it is the primitive the ky >= 0 state layout (plan 5.3 N3) calls directly; _spectral_bracket_real_fft_core() is this function plus the widening that a two-sided state still requires.

ny_full is the length of the two-sided ky axis and must be supplied whenever the operands are already in the half layout, because Nyc alone cannot say how long its own full axis is (gkx.core_ky_layout). Omitting it on a half operand is not a shape error – the slice, the irfft2 length and the mask all still compose – it just computes the bracket of a different grid, which is why the argument is explicit rather than inferred.

gkx.operators.nonlinear.brackets._spectral_bracket_multi_real_fft(G_hat: Array, chi_hat_stack: Array, combine: Callable[[Array], Array] | None = None, **kwargs) → Array

Bracket of G with each field of the stack.

combine reduces the field axis of the real-space products (linear, so one forward FFT serves every field); without it each field keeps its own.

gkx.operators.nonlinear.brackets._spectral_bracket_real_fft_core(G_hat: Array, chi_hat: Array, *, kxfac: Array, **kwargs: Any) → Array

Return the bracket in the ky layout its operands arrived in.

The kernel is the same either way. A two-sided operand pays the widening back onto Ny rows; a half-spectrum one does not, which is the whole of plan 5.3 N3 in the bracket – the reality condition is a property of what is stored rather than something restored after the fact.

Nonlinear Term Assembly

Pseudo-spectral nonlinear E×B and electromagnetic bracket terms.

class gkx.terms.nonlinear._LaguerreGridContext(to_grid: 'jnp.ndarray', to_spectral: 'jnp.ndarray', roots: 'jnp.ndarray', j0: 'jnp.ndarray | None', j1_over_alpha: 'jnp.ndarray | None', b: 'jnp.ndarray')

class gkx.terms.nonlinear._NonlinearBracketContext(tz: 'jnp.ndarray', vth: 'jnp.ndarray', sqrt_m: 'jnp.ndarray', sqrt_m_p1: 'jnp.ndarray', kx_grid: 'jnp.ndarray', ky_grid: 'jnp.ndarray', dealias_mask: 'jnp.ndarray', kxfac: 'jnp.ndarray', weight: 'jnp.ndarray', apar_weight: 'float', bpar_weight: 'float', compressed_real_fft: 'bool', ny_full: 'int | None', radial_phase: 'jnp.ndarray | None')

class gkx.terms.nonlinear._PreparedNonlinearInputs(G: 'jnp.ndarray', phi: 'jnp.ndarray', apar: 'jnp.ndarray | None', bpar: 'jnp.ndarray | None', Jl: 'jnp.ndarray', JlB: 'jnp.ndarray', squeeze_species: 'bool')

class gkx.terms.nonlinear._PreparedNonlinearPath(prep: '_PreparedNonlinearInputs', laguerre: '_LaguerreGridContext | None', electrostatic_only: 'bool')

gkx.terms.nonlinear._field_sum(c: _NonlinearBracketContext, nm: int, idx_bpar: int | None, idx_apar: int | None) → Callable[[Array], Array]

E x B + B_par + flutter, summed on the real-space products.

The Hermite shift of the flutter term and the field sum are linear, so they commute with the forward FFT: one transform instead of one per field.

gkx.terms.nonlinear._prepare_nonlinear_path(G: Array, *, phi: Array, apar: Array | None, bpar: Array | None, Jl: Array, JlB: Array, dealias_mask: Array, apar_weight: float, bpar_weight: float, laguerre_to_grid: Array | None, laguerre_to_spectral: Array | None, laguerre_roots: Array | None, laguerre_j0: Array | None, laguerre_j1_over_alpha: Array | None, b: Array | None, laguerre_mode: str) → _PreparedNonlinearPath

Prepare common nonlinear bracket routing decisions.

gkx.terms.nonlinear.nonlinear_em_contribution(G: Array, *, phi: Array, apar: Array | None, bpar: Array | None, Jl: Array, JlB: Array, tz: Array, vth: Array, sqrt_m: Array, sqrt_m_p1: Array, kx_grid: Array, ky_grid: Array, dealias_mask: Array, kxfac: Array, weight: Array, apar_weight: float, bpar_weight: float, laguerre_to_grid: Array | None = None, laguerre_to_spectral: Array | None = None, laguerre_roots: Array | None = None, laguerre_j0: Array | None = None, laguerre_j1_over_alpha: Array | None = None, b: Array | None = None, compressed_real_fft: bool = True, ny_full: int | None = None, laguerre_mode: str = 'grid', radial_phase: Array | None = None) → Array

Nonlinear E×B + flutter contribution using Laguerre gyroaveraging.

apar_weight and bpar_weight are used as on/off toggles (nonzero enables the term); the fields themselves already include any scaling.

Nonlinear Explicit Step

Explicit nonlinear time-step policies.

This module keeps the RK/SSP one-step formulas outside the public nonlinear facade. The integrator still injects the RHS and projection functions so the step policy remains pure, small, and directly testable.

gkx.solvers_nonlinear_explicit.advance_explicit_nonlinear_state(G: Array, dG: Array, dt_local: Array, *, method: str, rhs_fn: Callable[[...], tuple[Array, object]], project_state: Callable[[Array], Array], state_dtype: dtype) → Array

Advance one explicit nonlinear step with a static method string.

Barriers materialize every stage and stage derivative once. Unbarriered, XLA:CPU fuses the whole linear RHS into the transpose feeding the bracket’s irfft2 and recomputes it ~40 times per stage whenever the projector is the identity (ky >= 0 layout): 30x on an RK3 step. jaxlib < 0.11 drops the barrier before CPU fusion.

gkx.solvers_nonlinear_explicit.advance_rk4_pi_state(G: Array, dG: Array, dt_local: Array, *, rhs_fn: Callable[[...], tuple[Array, object]], project_state: Callable[[Array], Array], state_dtype: dtype) → tuple[Array, Array, Array]

One rk4-pi step: returns (G_new, dt_used, dt_next).

gkx.solvers_nonlinear_explicit.block_checkpoint_plan(steps: int, carry_bytes: int, residual_bytes: int, memory_budget_bytes: int) → int | None

Block length of the block-only schedule, or None if over budget.

It stores steps/B carries plus B step residuals, minimized by B = sqrt(steps * carry / residual).

gkx.solvers_nonlinear_explicit.checkpointed_explicit_scan(step: Callable[[...], Any], initial_carry: Any, indices: Array, *, checkpoint: bool, memory_budget_bytes: int | None = None) → tuple[Any, Any]

Scan with a two-level discrete-adjoint checkpoint schedule.

A rematerialized step alone still leaves every scan carry on the reverse tape. Grouping steps into rematerialized blocks retains only block boundaries, while the inner rematerialized scan retains the states needed to reverse one block. A block length near sqrt(steps) therefore lowers state storage from O(steps) to O(sqrt(steps)) without changing the discrete time integrator. Requested scan outputs retain their inherent O(steps) result storage.

That nested schedule runs each forward step three times. Under memory_budget_bytes the per-step remat is dropped when the block-only schedule fits (block_checkpoint_plan(), residuals sized from the abstract pullback, which errs high), so it runs twice. Same discrete map.

gkx.solvers_nonlinear_explicit.integrate_nonlinear_scan(rhs_fn: Callable[[...], tuple[Array, object]], G0: Array, dt: float, steps: int, *, method: str = 'rk4', checkpoint: bool = False, project_state: Callable[[Array], Array] | None = None, show_progress: bool = False, return_fields: bool = True, rhs_args: tuple[Any, ...] = (), rhs_static_args: tuple[Any, ...] = ()) → tuple[Array, Any] | Array

Integrate a cached nonlinear RHS using the explicit solver scan policy.

gkx.solvers_nonlinear_explicit.make_explicit_diagnostic_step(*, rhs_fn: Callable[[...], tuple[Array, object]], method: str, project_state: Callable[[Array], Array], state_dtype: dtype, real_dtype: dtype, time_step_policy: Any, compute_fields_fn: Callable[[...], Any], cache: Any, params: Any, term_cfg: Any, external_phi: Array | float | None, compute_diag_from_state: Callable[[...], Any], diagnostics_stride: int, select_diagnostics_fn: Callable[[...], Any], show_progress: bool, steps: int, emit_progress_fn: Callable[[...], Array], use_collision_split: bool = False, damping: Any | None = None, collision_scheme: str = 'implicit', apply_collision_split_fn: Callable[[Array, Any, Array, str], Array] | None = None, time_horizon: float | None = None) → Callable[[tuple[Any, Any, Any, Any, Any, Any], Any], tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]]

Build one explicit diagnostic scan step with injected runtime seams.

gkx.solvers_nonlinear_explicit.run_explicit_diagnostic_scan(step_fn: Callable[[tuple[Any, Any, Any, Any, Any, Any], Any], tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]], initial_carry: tuple[Any, Any, Any, Any, Any, Any], *, steps: int, stride: int, sampled_scan: bool, checkpoint: bool, sampled_scan_fn: Callable[[...], tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]]) → tuple[Array, tuple[Any, Any, Any]]

Run the explicit diagnostic scan using sampled or dense retention.

Nonlinear State Integration

Nonlinear integration: cached RHS, state integrators and diagnostic drivers.

The explicit and IMEX diagnostic drivers call the kernels as module-level names, so tests patch them on this module.

class gkx.solvers_nonlinear_diagnostics.PreparedExplicitNonlinearDiagnostics(initial_state: Array, geometry: Any, cache: LinearCache, params: LinearParams, _run_raw: Callable[[...], tuple[Any, Any, Any]], _run_dynamic_raw: Callable[[Array, Any, LinearCache, LinearParams], tuple[Any, Any, Any]], _finalize: Callable[[...], SimulationDiagnostics], stride: int, sampled_scan: bool, resolved_diagnostics: bool, fixed_dt: bool, cfl_scales: Any = None, cover_mask: Any = None)

Reusable compiled explicit nonlinear diagnostic simulation.

Geometry, field operators, and static numerical policy are prepared once. Calls to run() may supply a new initial state with the same shape and dtype without rebuilding the scan closure. Fixed-step sensitivity studies may instead pass matched geometry, cache, and parameter PyTrees.

This object and integrate_nonlinear_explicit_diagnostics_state() compile one graph and return bitwise identical arrays for the same inputs (plan §5.3, queue row Q18). The one difference that remains is the state dtype: this object freezes the dtype of the G0 it was prepared with and casts every later state to it, while the function entry point takes the dtype of the G0 it is called with. Feeding a promoted state to one and not the other is a precision change, not a route difference.

Supplied states (queue row Q23). On a linked (twist-shift) deck the kx rows outside the linked chains are decoupled, and the runtime’s own initial conditions never populate them. Every state this object is given – the G0 it was prepared with and any initial_state passed to run() or run_arrays() – is projected onto the chain cover before it is integrated, because the ExB bracket takes the whole state to real space and dealiases only its output, so off-chain content aliases back onto the chain rows. The projection is a mask fixed by the deck’s topology, not a check on the state’s values, so it holds under jit and under reverse-mode AD; gradients with respect to a supplied state are exactly zero on the off-chain rows and unchanged elsewhere. Periodic and full-cover decks are untouched and trace exactly as before.

property chain: Any

The chain implicit operator the scan takes as an operand, if any.

run(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | None = None, time_horizon: float | None = None) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Advance one state through the prepared compiled simulation.

run_arrays(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | None = None, time_horizon: float | None = None) → tuple[Array, tuple[Any, Any, Any], FieldState]

Run the compiled scan without host conversion or artifact assembly.

This method is the differentiable Python boundary. The initial state is dynamic. Fixed-step runs may also receive a matched cache/params pair for parameter differentiation. A changed geometry requires that pair. Grid layout and numerical policy remain fixed by prepare_nonlinear_explicit_diagnostics(); adaptive runs currently support state changes but reject traced model overrides.

A supplied initial_state is projected onto the linked chain cover first, as the class docstring describes. The prepared state was projected once at preparation, so passing None runs exactly the graph it ran before.

time_horizon (adaptive runs) stops stepping once the run time reaches it. It is a traced operand, so changing it between calls reuses the compiled scan; it overrides any prepared horizon.

with_chain(chain: Any) → PreparedExplicitNonlinearDiagnostics

Same compiled scan with another operator of the same structure.

class gkx.solvers_nonlinear_diagnostics.ShearedTransportTrace(final_state: jnp.ndarray, time: jnp.ndarray, heat_flux: jnp.ndarray, solve_stats: ImplicitSolveStats | None = None)

Final state and compact heat-flux history from a sheared run.

solve_stats is the convergence summary of every implicit GMRES solve of a method="imex" run, and is populated only when the caller asked for it with return_solve_stats=True; it stays None for explicit methods and for runs that did not ask, so a graph that does not read it is unchanged.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new ShearedTransportTrace object from a sequence or iterable

_replace(**kwds)

Return a new ShearedTransportTrace object replacing specified fields with new values

final_state: Array

Alias for field number 0

heat_flux: Array

Alias for field number 2

solve_stats: ImplicitSolveStats | None

Alias for field number 3

time: Array

Alias for field number 1

gkx.solvers_nonlinear_diagnostics._integrate_nonlinear_explicit_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Integrate an explicit nonlinear run and return diagnostics plus final state.

gkx.solvers_nonlinear_diagnostics._linear_rhs_jit_for_terms(term_cfg: TermConfig)

Return the narrowest compiled linear RHS path compatible with term_cfg.

gkx.solvers_nonlinear_diagnostics._nonlinear_diagnostic_kernels() → NonlinearDiagnosticKernels

Return the diagnostic kernels passed to the diagnostic tuple builder.

gkx.solvers_nonlinear_diagnostics._sample_indices_with_final(length: int, stride: int) → slice | ndarray

Return strided sample indices while always retaining the final step.

gkx.solvers_nonlinear_diagnostics.build_nonlinear_simulation_diagnostics(diag: tuple[Any, ...], t: Any, dt_series: Any, *, resolved_diagnostics: bool, sample_indices: slice | ndarray | None = None, resolved_to_numpy: bool = False, cfl_scales: Any | None = None) → SimulationDiagnostics

Build sampled nonlinear diagnostics from the raw scan output tuple.

gkx.solvers_nonlinear_diagnostics.finalize_nonlinear_scan_diagnostics(diag: tuple[Any, ...], t: Any, dt_series: Any, *, stride: int, sampled_scan: bool = False, resolved_diagnostics: bool, resolved_to_numpy: bool = False, cfl_scales: Any | None = None) → SimulationDiagnostics

Package raw nonlinear scan diagnostics after applying output sampling.

gkx.solvers_nonlinear_diagnostics.integrate_imex_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'imex', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Integrate an IMEX nonlinear run and return diagnostics.

Returns (t, diagnostics), or (t, diagnostics, stats) with return_solve_stats=True.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, method: str = 'rk4', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', show_progress: bool = False, return_fields: bool = True, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the nonlinear system using built-in cache construction.

A supplied G0 is projected onto the linked chain cover by integrate_nonlinear_cached(), which this driver builds the cache for and then delegates to, so both raw drivers take the same view of a state they did not build.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, method: str = 'rk4', terms: TermConfig | None = None, checkpoint: bool = False, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', show_progress: bool = False, return_fields: bool = True, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the nonlinear system using a cached geometry object.

return_solve_stats=True appends the IMEX implicit solve summary (None for explicit methods, which have no implicit solve).

G0 is a state this driver did not build, so on a linked (twist-shift) deck it is projected onto the chain cover before the run starts, exactly as the runtime does at intake (queue row Q19) and as the entry points above this one do (Q23). The projection is a where against a mask fixed by the deck’s topology, applied once here and outside the compiled scan, so the scan’s graph is unchanged and periodic or full-cover decks get the same array object back.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_explicit_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Integrate nonlinear system and return runtime diagnostics.

return_solve_stats=True appends the implicit solve summary of an IMEX run (None for the explicit methods, which take no implicit solve).

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_explicit_diagnostics_state(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False, compile_cache: dict[Any, Any] | None = None) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Integrate nonlinear system and return runtime diagnostics plus the final state.

This is the entry point run_runtime_nonlinear uses on its fixed-window, chunked and sharded routes. It runs the same jitted graph as prepare_nonlinear_explicit_diagnostics() and returns bitwise identical arrays for the same inputs (plan §5.3, queue row Q18). Prefer the prepared object when the same case is run more than once: it keeps the compiled graph, while each call here compiles its own.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_imex_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, *, terms: TermConfig | None = None, checkpoint: bool = False, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, implicit_operator: IMEXLinearOperator | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

IMEX integrator: implicit linear operator, explicit nonlinear term.

Returns (G_out, fields_t), or (G_out, fields_t, stats) with return_solve_stats=True; stats is the carried ImplicitSolveStats of every solve.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_imex_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'imex', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

IMEX nonlinear integrator with runtime diagnostics.

Returns (t, diagnostics), or (t, diagnostics, stats) with return_solve_stats=True. The scan cannot raise on a starved inner budget, so stats (ImplicitSolveStats) is this route’s convergence channel; a host-side caller turns it into a refusal with require_converged_implicit_solves().

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_sheared(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, shear_rate: Array | float, method: str = 'rk2', cache: LinearCache | None = None, terms: TermConfig | None = None, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, compressed_real_fft: bool = False, differentiable: bool = False, return_fields: bool = True, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the standard-flux-tube shearing-coordinate foundation.

This research path supports midpoint RK2, three-stage Heun RK3, and first-order IMEX. Stage states and derivatives are remapped to the stage coordinate basis before the RHS and back to the step basis before Runge–Kutta combinations. IMEX evaluates the explicit nonlinear term in the current basis and rebuilds the implicit linear operator in the endpoint basis. compressed_real_fft evaluates the nonlinear bracket in the equivalent canonical shearing-coordinate representation.

return_solve_stats=True appends the implicit solve summary of a method="imex" run (ImplicitSolveStats, None for the explicit methods). The traced scan cannot raise on a starved inner budget, so this is the convergence channel; a host-side caller turns it into a refusal with require_converged_implicit_solves().

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_sheared_transport(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, shear_rate: Array | float, method: str = 'rk2', cache: LinearCache | None = None, terms: TermConfig | None = None, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, compressed_real_fft: bool = False, flux_scale: float = 1.0, differentiable: bool = True, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, initial_time: Array | float = 0.0, initial_dt: Array | float | 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_preconditioner: str | None = None, return_solve_stats: bool = False) → ShearedTransportTrace

Integrate a sheared run and record canonical heat flux at every step.

With fixed_dt=False, steps is the accepted-step budget and time records the resulting nonuniform physical-time grid. initial_time and initial_dt continue a prior trace without resetting the shearing basis.

return_solve_stats=True fills the trace’s solve_stats field with the implicit solve summary of a method="imex" run; it stays None otherwise, so a caller that does not ask reads and compiles exactly what it did before.

gkx.solvers_nonlinear_diagnostics.maybe_emit_nonlinear_progress(state: Any, *, show_progress: bool, diag: tuple[Any, ...], idx: Any, steps: int, t_new: Any, progress_total: Any) → Any

Emit nonlinear progress callbacks when requested and return state.

gkx.solvers_nonlinear_diagnostics.nonlinear_heat_flux_window(saturated_state: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, terms: TermConfig | None = None, method: str = 'rk2', tail_steps: int | None = None, checkpoint: bool = True, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, divergence_knee_steps: int | None = 1024, adjoint_memory_budget_bytes: int | None = 4294967296) → Array

Mean physical heat flux over a differentiable nonlinear window.

saturated_state is deliberately detached: first run GKX to saturation, then differentiate this finite window through geom and params. Reverse-mode AD follows the exact discrete Runge–Kutta map. Block checkpointing retains \(O(\sqrt{N})\) distribution states instead of \(O(N)\) for a window of N steps.

On a linked (twist-shift) deck saturated_state is projected onto the linked chain cover before the window starts (queue row Q23): it is a state this function did not build, and the ExB bracket takes the whole state to real space while dealiasing only its output, so off-chain rows would alias back onto the chain rows and change the flux this objective returns. The projection does not change the gradient contract, which is that saturated_state is detached and only geom and params carry derivatives. Periodic and full-cover decks are untouched.

collision_operator is the same custom model integrate_nonlinear() accepts, and must be passed here too: a run saturated with a custom operator and then differentiated without one is a derivative of different physics from the trajectory it starts on.

steps above divergence_knee_steps warns. Past the measured knee the windowed adjoint grows with the leading Lyapunov exponent and stops being a useful design direction; pass None to silence the check when the knee has been remeasured for the case at hand with scripts/campaigns/nonlinear_gradient_window.py.

adjoint_memory_budget_bytes bounds the reverse pass’s checkpoint storage. When the block-only schedule fits, each step’s forward runs twice in a gradient instead of three times; otherwise, or with None, the nested \(O(\sqrt{N})\) schedule is used. It changes cost, not the discrete adjoint.

Host-side geometry/layout preparation stays here. The differentiated scan in _nonlinear_heat_flux_window_total() reuses its executable for matching shapes, topology and static options, with geometry arrays passed as operands. See docs/nonlinear_autodiff.rst for scope and evidence.

gkx.solvers_nonlinear_diagnostics.nonlinear_rhs_cached(G: Array, cache: LinearCache, params: LinearParams, terms: TermConfig | None = None, *, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, collision_operator: CollisionOperator | None = None, radial_phase: Array | None = None, differentiable: bool = False) → tuple[Array, FieldState]

Compute the assembled nonlinear RHS and electromagnetic field state.

gkx.solvers_nonlinear_diagnostics.prepare_explicit_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False) → PreparedExplicitNonlinearDiagnostics

Prepare one compile-stable explicit nonlinear diagnostic simulation.

gkx.solvers_nonlinear_diagnostics.prepare_nonlinear_explicit_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, **options: Any) → PreparedExplicitNonlinearDiagnostics

Prepare a reusable explicit diagnostic scan for repeated Python calls.

options accepts the same explicit-only keywords as integrate_nonlinear_explicit_diagnostics_state().

gkx.solvers_nonlinear_diagnostics.run_sampled_explicit_diagnostic_scan(step_fn: Any, initial_carry: tuple[Any, Any, Any, Any, Any, Any], *, steps: int, stride: int) → tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]

Run an explicit diagnostic scan only at retained sample intervals.

gkx.solvers_nonlinear_diagnostics.sampled_scan_intervals(length: int, stride: int) → ndarray

Return positive scan intervals that retain the requested final sample.

gkx.solvers_nonlinear_diagnostics.select_nonlinear_step_diagnostics(idx: Any, *, diagnostics_stride: int, diag_prev: Any, compute_diag_fn: Any, steps: int | None = None) → Any

Return a fresh or reused nonlinear step diagnostic tuple.

Off-stride steps reuse diag_prev. The last step of a steps-long scan is always computed fresh, because the output sampling keeps it as the final row even when it falls off-stride; reusing the carry there would report the previous stride sample’s diagnostics at the final time.

Nonlinear Diagnostic Drivers

Nonlinear integration: cached RHS, state integrators and diagnostic drivers.

The explicit and IMEX diagnostic drivers call the kernels as module-level names, so tests patch them on this module.

class gkx.solvers_nonlinear_diagnostics.PreparedExplicitNonlinearDiagnostics(initial_state: Array, geometry: Any, cache: LinearCache, params: LinearParams, _run_raw: Callable[[...], tuple[Any, Any, Any]], _run_dynamic_raw: Callable[[Array, Any, LinearCache, LinearParams], tuple[Any, Any, Any]], _finalize: Callable[[...], SimulationDiagnostics], stride: int, sampled_scan: bool, resolved_diagnostics: bool, fixed_dt: bool, cfl_scales: Any = None, cover_mask: Any = None)

Reusable compiled explicit nonlinear diagnostic simulation.

Geometry, field operators, and static numerical policy are prepared once. Calls to run() may supply a new initial state with the same shape and dtype without rebuilding the scan closure. Fixed-step sensitivity studies may instead pass matched geometry, cache, and parameter PyTrees.

This object and integrate_nonlinear_explicit_diagnostics_state() compile one graph and return bitwise identical arrays for the same inputs (plan §5.3, queue row Q18). The one difference that remains is the state dtype: this object freezes the dtype of the G0 it was prepared with and casts every later state to it, while the function entry point takes the dtype of the G0 it is called with. Feeding a promoted state to one and not the other is a precision change, not a route difference.

Supplied states (queue row Q23). On a linked (twist-shift) deck the kx rows outside the linked chains are decoupled, and the runtime’s own initial conditions never populate them. Every state this object is given – the G0 it was prepared with and any initial_state passed to run() or run_arrays() – is projected onto the chain cover before it is integrated, because the ExB bracket takes the whole state to real space and dealiases only its output, so off-chain content aliases back onto the chain rows. The projection is a mask fixed by the deck’s topology, not a check on the state’s values, so it holds under jit and under reverse-mode AD; gradients with respect to a supplied state are exactly zero on the off-chain rows and unchanged elsewhere. Periodic and full-cover decks are untouched and trace exactly as before.

property chain: Any

The chain implicit operator the scan takes as an operand, if any.

run(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | None = None, time_horizon: float | None = None) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Advance one state through the prepared compiled simulation.

run_arrays(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | None = None, time_horizon: float | None = None) → tuple[Array, tuple[Any, Any, Any], FieldState]

Run the compiled scan without host conversion or artifact assembly.

This method is the differentiable Python boundary. The initial state is dynamic. Fixed-step runs may also receive a matched cache/params pair for parameter differentiation. A changed geometry requires that pair. Grid layout and numerical policy remain fixed by prepare_nonlinear_explicit_diagnostics(); adaptive runs currently support state changes but reject traced model overrides.

A supplied initial_state is projected onto the linked chain cover first, as the class docstring describes. The prepared state was projected once at preparation, so passing None runs exactly the graph it ran before.

time_horizon (adaptive runs) stops stepping once the run time reaches it. It is a traced operand, so changing it between calls reuses the compiled scan; it overrides any prepared horizon.

with_chain(chain: Any) → PreparedExplicitNonlinearDiagnostics

Same compiled scan with another operator of the same structure.

class gkx.solvers_nonlinear_diagnostics.ShearedTransportTrace(final_state: jnp.ndarray, time: jnp.ndarray, heat_flux: jnp.ndarray, solve_stats: ImplicitSolveStats | None = None)

Final state and compact heat-flux history from a sheared run.

solve_stats is the convergence summary of every implicit GMRES solve of a method="imex" run, and is populated only when the caller asked for it with return_solve_stats=True; it stays None for explicit methods and for runs that did not ask, so a graph that does not read it is unchanged.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new ShearedTransportTrace object from a sequence or iterable

_replace(**kwds)

Return a new ShearedTransportTrace object replacing specified fields with new values

final_state: Array

Alias for field number 0

heat_flux: Array

Alias for field number 2

solve_stats: ImplicitSolveStats | None

Alias for field number 3

time: Array

Alias for field number 1

gkx.solvers_nonlinear_diagnostics._integrate_nonlinear_explicit_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Integrate an explicit nonlinear run and return diagnostics plus final state.

gkx.solvers_nonlinear_diagnostics._linear_rhs_jit_for_terms(term_cfg: TermConfig)

Return the narrowest compiled linear RHS path compatible with term_cfg.

gkx.solvers_nonlinear_diagnostics._nonlinear_diagnostic_kernels() → NonlinearDiagnosticKernels

Return the diagnostic kernels passed to the diagnostic tuple builder.

gkx.solvers_nonlinear_diagnostics._sample_indices_with_final(length: int, stride: int) → slice | ndarray

Return strided sample indices while always retaining the final step.

gkx.solvers_nonlinear_diagnostics.build_nonlinear_simulation_diagnostics(diag: tuple[Any, ...], t: Any, dt_series: Any, *, resolved_diagnostics: bool, sample_indices: slice | ndarray | None = None, resolved_to_numpy: bool = False, cfl_scales: Any | None = None) → SimulationDiagnostics

Build sampled nonlinear diagnostics from the raw scan output tuple.

gkx.solvers_nonlinear_diagnostics.finalize_nonlinear_scan_diagnostics(diag: tuple[Any, ...], t: Any, dt_series: Any, *, stride: int, sampled_scan: bool = False, resolved_diagnostics: bool, resolved_to_numpy: bool = False, cfl_scales: Any | None = None) → SimulationDiagnostics

Package raw nonlinear scan diagnostics after applying output sampling.

gkx.solvers_nonlinear_diagnostics.integrate_imex_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'imex', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Integrate an IMEX nonlinear run and return diagnostics.

Returns (t, diagnostics), or (t, diagnostics, stats) with return_solve_stats=True.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, method: str = 'rk4', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', show_progress: bool = False, return_fields: bool = True, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the nonlinear system using built-in cache construction.

A supplied G0 is projected onto the linked chain cover by integrate_nonlinear_cached(), which this driver builds the cache for and then delegates to, so both raw drivers take the same view of a state they did not build.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, method: str = 'rk4', terms: TermConfig | None = None, checkpoint: bool = False, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', show_progress: bool = False, return_fields: bool = True, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the nonlinear system using a cached geometry object.

return_solve_stats=True appends the IMEX implicit solve summary (None for explicit methods, which have no implicit solve).

G0 is a state this driver did not build, so on a linked (twist-shift) deck it is projected onto the chain cover before the run starts, exactly as the runtime does at intake (queue row Q19) and as the entry points above this one do (Q23). The projection is a where against a mask fixed by the deck’s topology, applied once here and outside the compiled scan, so the scan’s graph is unchanged and periodic or full-cover decks get the same array object back.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_explicit_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Integrate nonlinear system and return runtime diagnostics.

return_solve_stats=True appends the implicit solve summary of an IMEX run (None for the explicit methods, which take no implicit solve).

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_explicit_diagnostics_state(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False, compile_cache: dict[Any, Any] | None = None) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Integrate nonlinear system and return runtime diagnostics plus the final state.

This is the entry point run_runtime_nonlinear uses on its fixed-window, chunked and sharded routes. It runs the same jitted graph as prepare_nonlinear_explicit_diagnostics() and returns bitwise identical arrays for the same inputs (plan §5.3, queue row Q18). Prefer the prepared object when the same case is run more than once: it keeps the compiled graph, while each call here compiles its own.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_imex_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, *, terms: TermConfig | None = None, checkpoint: bool = False, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, implicit_operator: IMEXLinearOperator | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

IMEX integrator: implicit linear operator, explicit nonlinear term.

Returns (G_out, fields_t), or (G_out, fields_t, stats) with return_solve_stats=True; stats is the carried ImplicitSolveStats of every solve.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_imex_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'imex', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

IMEX nonlinear integrator with runtime diagnostics.

Returns (t, diagnostics), or (t, diagnostics, stats) with return_solve_stats=True. The scan cannot raise on a starved inner budget, so stats (ImplicitSolveStats) is this route’s convergence channel; a host-side caller turns it into a refusal with require_converged_implicit_solves().

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_sheared(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, shear_rate: Array | float, method: str = 'rk2', cache: LinearCache | None = None, terms: TermConfig | None = None, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, compressed_real_fft: bool = False, differentiable: bool = False, return_fields: bool = True, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the standard-flux-tube shearing-coordinate foundation.

This research path supports midpoint RK2, three-stage Heun RK3, and first-order IMEX. Stage states and derivatives are remapped to the stage coordinate basis before the RHS and back to the step basis before Runge–Kutta combinations. IMEX evaluates the explicit nonlinear term in the current basis and rebuilds the implicit linear operator in the endpoint basis. compressed_real_fft evaluates the nonlinear bracket in the equivalent canonical shearing-coordinate representation.

return_solve_stats=True appends the implicit solve summary of a method="imex" run (ImplicitSolveStats, None for the explicit methods). The traced scan cannot raise on a starved inner budget, so this is the convergence channel; a host-side caller turns it into a refusal with require_converged_implicit_solves().

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_sheared_transport(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, shear_rate: Array | float, method: str = 'rk2', cache: LinearCache | None = None, terms: TermConfig | None = None, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, compressed_real_fft: bool = False, flux_scale: float = 1.0, differentiable: bool = True, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, initial_time: Array | float = 0.0, initial_dt: Array | float | 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_preconditioner: str | None = None, return_solve_stats: bool = False) → ShearedTransportTrace

Integrate a sheared run and record canonical heat flux at every step.

With fixed_dt=False, steps is the accepted-step budget and time records the resulting nonuniform physical-time grid. initial_time and initial_dt continue a prior trace without resetting the shearing basis.

return_solve_stats=True fills the trace’s solve_stats field with the implicit solve summary of a method="imex" run; it stays None otherwise, so a caller that does not ask reads and compiles exactly what it did before.

gkx.solvers_nonlinear_diagnostics.maybe_emit_nonlinear_progress(state: Any, *, show_progress: bool, diag: tuple[Any, ...], idx: Any, steps: int, t_new: Any, progress_total: Any) → Any

Emit nonlinear progress callbacks when requested and return state.

gkx.solvers_nonlinear_diagnostics.nonlinear_heat_flux_window(saturated_state: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, terms: TermConfig | None = None, method: str = 'rk2', tail_steps: int | None = None, checkpoint: bool = True, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, divergence_knee_steps: int | None = 1024, adjoint_memory_budget_bytes: int | None = 4294967296) → Array

Mean physical heat flux over a differentiable nonlinear window.

saturated_state is deliberately detached: first run GKX to saturation, then differentiate this finite window through geom and params. Reverse-mode AD follows the exact discrete Runge–Kutta map. Block checkpointing retains \(O(\sqrt{N})\) distribution states instead of \(O(N)\) for a window of N steps.

On a linked (twist-shift) deck saturated_state is projected onto the linked chain cover before the window starts (queue row Q23): it is a state this function did not build, and the ExB bracket takes the whole state to real space while dealiasing only its output, so off-chain rows would alias back onto the chain rows and change the flux this objective returns. The projection does not change the gradient contract, which is that saturated_state is detached and only geom and params carry derivatives. Periodic and full-cover decks are untouched.

collision_operator is the same custom model integrate_nonlinear() accepts, and must be passed here too: a run saturated with a custom operator and then differentiated without one is a derivative of different physics from the trajectory it starts on.

steps above divergence_knee_steps warns. Past the measured knee the windowed adjoint grows with the leading Lyapunov exponent and stops being a useful design direction; pass None to silence the check when the knee has been remeasured for the case at hand with scripts/campaigns/nonlinear_gradient_window.py.

adjoint_memory_budget_bytes bounds the reverse pass’s checkpoint storage. When the block-only schedule fits, each step’s forward runs twice in a gradient instead of three times; otherwise, or with None, the nested \(O(\sqrt{N})\) schedule is used. It changes cost, not the discrete adjoint.

Host-side geometry/layout preparation stays here. The differentiated scan in _nonlinear_heat_flux_window_total() reuses its executable for matching shapes, topology and static options, with geometry arrays passed as operands. See docs/nonlinear_autodiff.rst for scope and evidence.

gkx.solvers_nonlinear_diagnostics.nonlinear_rhs_cached(G: Array, cache: LinearCache, params: LinearParams, terms: TermConfig | None = None, *, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, collision_operator: CollisionOperator | None = None, radial_phase: Array | None = None, differentiable: bool = False) → tuple[Array, FieldState]

Compute the assembled nonlinear RHS and electromagnetic field state.

gkx.solvers_nonlinear_diagnostics.prepare_explicit_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False) → PreparedExplicitNonlinearDiagnostics

Prepare one compile-stable explicit nonlinear diagnostic simulation.

gkx.solvers_nonlinear_diagnostics.prepare_nonlinear_explicit_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, **options: Any) → PreparedExplicitNonlinearDiagnostics

Prepare a reusable explicit diagnostic scan for repeated Python calls.

options accepts the same explicit-only keywords as integrate_nonlinear_explicit_diagnostics_state().

gkx.solvers_nonlinear_diagnostics.run_sampled_explicit_diagnostic_scan(step_fn: Any, initial_carry: tuple[Any, Any, Any, Any, Any, Any], *, steps: int, stride: int) → tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]

Run an explicit diagnostic scan only at retained sample intervals.

gkx.solvers_nonlinear_diagnostics.sampled_scan_intervals(length: int, stride: int) → ndarray

Return positive scan intervals that retain the requested final sample.

gkx.solvers_nonlinear_diagnostics.select_nonlinear_step_diagnostics(idx: Any, *, diagnostics_stride: int, diag_prev: Any, compute_diag_fn: Any, steps: int | None = None) → Any

Return a fresh or reused nonlinear step diagnostic tuple.

Off-stride steps reuse diag_prev. The last step of a steps-long scan is always computed fresh, because the output sampling keeps it as the final row even when it falls off-stride; reusing the carry there would report the previous stride sample’s diagnostics at the final time.

Nonlinear Diagnostic Integration

Nonlinear integration: cached RHS, state integrators and diagnostic drivers.

The explicit and IMEX diagnostic drivers call the kernels as module-level names, so tests patch them on this module.

class gkx.solvers_nonlinear_diagnostics.PreparedExplicitNonlinearDiagnostics(initial_state: Array, geometry: Any, cache: LinearCache, params: LinearParams, _run_raw: Callable[[...], tuple[Any, Any, Any]], _run_dynamic_raw: Callable[[Array, Any, LinearCache, LinearParams], tuple[Any, Any, Any]], _finalize: Callable[[...], SimulationDiagnostics], stride: int, sampled_scan: bool, resolved_diagnostics: bool, fixed_dt: bool, cfl_scales: Any = None, cover_mask: Any = None)

Reusable compiled explicit nonlinear diagnostic simulation.

Geometry, field operators, and static numerical policy are prepared once. Calls to run() may supply a new initial state with the same shape and dtype without rebuilding the scan closure. Fixed-step sensitivity studies may instead pass matched geometry, cache, and parameter PyTrees.

This object and integrate_nonlinear_explicit_diagnostics_state() compile one graph and return bitwise identical arrays for the same inputs (plan §5.3, queue row Q18). The one difference that remains is the state dtype: this object freezes the dtype of the G0 it was prepared with and casts every later state to it, while the function entry point takes the dtype of the G0 it is called with. Feeding a promoted state to one and not the other is a precision change, not a route difference.

Supplied states (queue row Q23). On a linked (twist-shift) deck the kx rows outside the linked chains are decoupled, and the runtime’s own initial conditions never populate them. Every state this object is given – the G0 it was prepared with and any initial_state passed to run() or run_arrays() – is projected onto the chain cover before it is integrated, because the ExB bracket takes the whole state to real space and dealiases only its output, so off-chain content aliases back onto the chain rows. The projection is a mask fixed by the deck’s topology, not a check on the state’s values, so it holds under jit and under reverse-mode AD; gradients with respect to a supplied state are exactly zero on the off-chain rows and unchanged elsewhere. Periodic and full-cover decks are untouched and trace exactly as before.

property chain: Any

The chain implicit operator the scan takes as an operand, if any.

run(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | None = None, time_horizon: float | None = None) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Advance one state through the prepared compiled simulation.

run_arrays(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | None = None, time_horizon: float | None = None) → tuple[Array, tuple[Any, Any, Any], FieldState]

Run the compiled scan without host conversion or artifact assembly.

This method is the differentiable Python boundary. The initial state is dynamic. Fixed-step runs may also receive a matched cache/params pair for parameter differentiation. A changed geometry requires that pair. Grid layout and numerical policy remain fixed by prepare_nonlinear_explicit_diagnostics(); adaptive runs currently support state changes but reject traced model overrides.

A supplied initial_state is projected onto the linked chain cover first, as the class docstring describes. The prepared state was projected once at preparation, so passing None runs exactly the graph it ran before.

time_horizon (adaptive runs) stops stepping once the run time reaches it. It is a traced operand, so changing it between calls reuses the compiled scan; it overrides any prepared horizon.

with_chain(chain: Any) → PreparedExplicitNonlinearDiagnostics

Same compiled scan with another operator of the same structure.

class gkx.solvers_nonlinear_diagnostics.ShearedTransportTrace(final_state: jnp.ndarray, time: jnp.ndarray, heat_flux: jnp.ndarray, solve_stats: ImplicitSolveStats | None = None)

Final state and compact heat-flux history from a sheared run.

solve_stats is the convergence summary of every implicit GMRES solve of a method="imex" run, and is populated only when the caller asked for it with return_solve_stats=True; it stays None for explicit methods and for runs that did not ask, so a graph that does not read it is unchanged.

_asdict()

Return a new dict which maps field names to their values.

classmethod _make(iterable)

Make a new ShearedTransportTrace object from a sequence or iterable

_replace(**kwds)

Return a new ShearedTransportTrace object replacing specified fields with new values

final_state: Array

Alias for field number 0

heat_flux: Array

Alias for field number 2

solve_stats: ImplicitSolveStats | None

Alias for field number 3

time: Array

Alias for field number 1

gkx.solvers_nonlinear_diagnostics._integrate_nonlinear_explicit_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Integrate an explicit nonlinear run and return diagnostics plus final state.

gkx.solvers_nonlinear_diagnostics._linear_rhs_jit_for_terms(term_cfg: TermConfig)

Return the narrowest compiled linear RHS path compatible with term_cfg.

gkx.solvers_nonlinear_diagnostics._nonlinear_diagnostic_kernels() → NonlinearDiagnosticKernels

Return the diagnostic kernels passed to the diagnostic tuple builder.

gkx.solvers_nonlinear_diagnostics._sample_indices_with_final(length: int, stride: int) → slice | ndarray

Return strided sample indices while always retaining the final step.

gkx.solvers_nonlinear_diagnostics.build_nonlinear_simulation_diagnostics(diag: tuple[Any, ...], t: Any, dt_series: Any, *, resolved_diagnostics: bool, sample_indices: slice | ndarray | None = None, resolved_to_numpy: bool = False, cfl_scales: Any | None = None) → SimulationDiagnostics

Build sampled nonlinear diagnostics from the raw scan output tuple.

gkx.solvers_nonlinear_diagnostics.finalize_nonlinear_scan_diagnostics(diag: tuple[Any, ...], t: Any, dt_series: Any, *, stride: int, sampled_scan: bool = False, resolved_diagnostics: bool, resolved_to_numpy: bool = False, cfl_scales: Any | None = None) → SimulationDiagnostics

Package raw nonlinear scan diagnostics after applying output sampling.

gkx.solvers_nonlinear_diagnostics.integrate_imex_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'imex', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Integrate an IMEX nonlinear run and return diagnostics.

Returns (t, diagnostics), or (t, diagnostics, stats) with return_solve_stats=True.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, method: str = 'rk4', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', show_progress: bool = False, return_fields: bool = True, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the nonlinear system using built-in cache construction.

A supplied G0 is projected onto the linked chain cover by integrate_nonlinear_cached(), which this driver builds the cache for and then delegates to, so both raw drivers take the same view of a state they did not build.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, method: str = 'rk4', terms: TermConfig | None = None, checkpoint: bool = False, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', show_progress: bool = False, return_fields: bool = True, collision_operator: CollisionOperator | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the nonlinear system using a cached geometry object.

return_solve_stats=True appends the IMEX implicit solve summary (None for explicit methods, which have no implicit solve).

G0 is a state this driver did not build, so on a linked (twist-shift) deck it is projected onto the chain cover before the run starts, exactly as the runtime does at intake (queue row Q19) and as the entry points above this one do (Q23). The projection is a where against a mask fixed by the deck’s topology, applied once here and outside the compiled scan, so the scan’s graph is unchanged and periodic or full-cover decks get the same array object back.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_explicit_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Integrate nonlinear system and return runtime diagnostics.

return_solve_stats=True appends the implicit solve summary of an IMEX run (None for the explicit methods, which take no implicit solve).

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_explicit_diagnostics_state(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False, compile_cache: dict[Any, Any] | None = None) → tuple[Array, SimulationDiagnostics, Array, FieldState]

Integrate nonlinear system and return runtime diagnostics plus the final state.

This is the entry point run_runtime_nonlinear uses on its fixed-window, chunked and sharded routes. It runs the same jitted graph as prepare_nonlinear_explicit_diagnostics() and returns bitwise identical arrays for the same inputs (plan §5.3, queue row Q18). Prefer the prepared object when the same case is run more than once: it keeps the compiled graph, while each call here compiles its own.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_imex_cached(G0: Array, cache: LinearCache, params: LinearParams, dt: float, steps: int, *, terms: TermConfig | None = None, checkpoint: bool = False, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, implicit_operator: IMEXLinearOperator | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

IMEX integrator: implicit linear operator, explicit nonlinear term.

Returns (G_out, fields_t), or (G_out, fields_t, stats) with return_solve_stats=True; stats is the carried ImplicitSolveStats of every solve.

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_imex_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'imex', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

IMEX nonlinear integrator with runtime diagnostics.

Returns (t, diagnostics), or (t, diagnostics, stats) with return_solve_stats=True. The scan cannot raise on a starved inner budget, so stats (ImplicitSolveStats) is this route’s convergence channel; a host-side caller turns it into a refusal with require_converged_implicit_solves().

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_sheared(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, shear_rate: Array | float, method: str = 'rk2', cache: LinearCache | None = None, terms: TermConfig | None = None, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, compressed_real_fft: bool = False, differentiable: bool = False, return_fields: bool = True, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, return_solve_stats: bool = False) → tuple[Any, ...] | Array

Integrate the standard-flux-tube shearing-coordinate foundation.

This research path supports midpoint RK2, three-stage Heun RK3, and first-order IMEX. Stage states and derivatives are remapped to the stage coordinate basis before the RHS and back to the step basis before Runge–Kutta combinations. IMEX evaluates the explicit nonlinear term in the current basis and rebuilds the implicit linear operator in the endpoint basis. compressed_real_fft evaluates the nonlinear bracket in the equivalent canonical shearing-coordinate representation.

return_solve_stats=True appends the implicit solve summary of a method="imex" run (ImplicitSolveStats, None for the explicit methods). The traced scan cannot raise on a starved inner budget, so this is the convergence channel; a host-side caller turns it into a refusal with require_converged_implicit_solves().

gkx.solvers_nonlinear_diagnostics.integrate_nonlinear_sheared_transport(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, shear_rate: Array | float, method: str = 'rk2', cache: LinearCache | None = None, terms: TermConfig | None = None, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, compressed_real_fft: bool = False, flux_scale: float = 1.0, differentiable: bool = True, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, initial_time: Array | float = 0.0, initial_dt: Array | float | 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_preconditioner: str | None = None, return_solve_stats: bool = False) → ShearedTransportTrace

Integrate a sheared run and record canonical heat flux at every step.

With fixed_dt=False, steps is the accepted-step budget and time records the resulting nonuniform physical-time grid. initial_time and initial_dt continue a prior trace without resetting the shearing basis.

return_solve_stats=True fills the trace’s solve_stats field with the implicit solve summary of a method="imex" run; it stays None otherwise, so a caller that does not ask reads and compiles exactly what it did before.

gkx.solvers_nonlinear_diagnostics.maybe_emit_nonlinear_progress(state: Any, *, show_progress: bool, diag: tuple[Any, ...], idx: Any, steps: int, t_new: Any, progress_total: Any) → Any

Emit nonlinear progress callbacks when requested and return state.

gkx.solvers_nonlinear_diagnostics.nonlinear_heat_flux_window(saturated_state: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, terms: TermConfig | None = None, method: str = 'rk2', tail_steps: int | None = None, checkpoint: bool = True, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', collision_operator: CollisionOperator | None = None, divergence_knee_steps: int | None = 1024, adjoint_memory_budget_bytes: int | None = 4294967296) → Array

Mean physical heat flux over a differentiable nonlinear window.

saturated_state is deliberately detached: first run GKX to saturation, then differentiate this finite window through geom and params. Reverse-mode AD follows the exact discrete Runge–Kutta map. Block checkpointing retains \(O(\sqrt{N})\) distribution states instead of \(O(N)\) for a window of N steps.

On a linked (twist-shift) deck saturated_state is projected onto the linked chain cover before the window starts (queue row Q23): it is a state this function did not build, and the ExB bracket takes the whole state to real space while dealiasing only its output, so off-chain rows would alias back onto the chain rows and change the flux this objective returns. The projection does not change the gradient contract, which is that saturated_state is detached and only geom and params carry derivatives. Periodic and full-cover decks are untouched.

collision_operator is the same custom model integrate_nonlinear() accepts, and must be passed here too: a run saturated with a custom operator and then differentiated without one is a derivative of different physics from the trajectory it starts on.

steps above divergence_knee_steps warns. Past the measured knee the windowed adjoint grows with the leading Lyapunov exponent and stops being a useful design direction; pass None to silence the check when the knee has been remeasured for the case at hand with scripts/campaigns/nonlinear_gradient_window.py.

adjoint_memory_budget_bytes bounds the reverse pass’s checkpoint storage. When the block-only schedule fits, each step’s forward runs twice in a gradient instead of three times; otherwise, or with None, the nested \(O(\sqrt{N})\) schedule is used. It changes cost, not the discrete adjoint.

Host-side geometry/layout preparation stays here. The differentiated scan in _nonlinear_heat_flux_window_total() reuses its executable for matching shapes, topology and static options, with geometry arrays passed as operands. See docs/nonlinear_autodiff.rst for scope and evidence.

gkx.solvers_nonlinear_diagnostics.nonlinear_rhs_cached(G: Array, cache: LinearCache, params: LinearParams, terms: TermConfig | None = None, *, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, collision_operator: CollisionOperator | None = None, radial_phase: Array | None = None, differentiable: bool = False) → tuple[Array, FieldState]

Compute the assembled nonlinear RHS and electromagnetic field state.

gkx.solvers_nonlinear_diagnostics.prepare_explicit_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'rk3', cache: LinearCache | None = None, terms: TermConfig | None = None, checkpoint: bool = False, sample_stride: int = 1, diagnostics_stride: int = 1, use_dealias_mask: bool = False, z_index: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', omega_ky_index: int | None = None, omega_kx_index: int | None = None, flux_scale: float = 1.0, wphi_scale: float = 1.0, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, time_horizon: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, fixed_mode_ky_index: int | None = None, fixed_mode_kx_index: int | None = None, external_phi: Array | float | None = None, resolved_diagnostics: bool = True, show_progress: bool = False) → PreparedExplicitNonlinearDiagnostics

Prepare one compile-stable explicit nonlinear diagnostic simulation.

gkx.solvers_nonlinear_diagnostics.prepare_nonlinear_explicit_diagnostics(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, **options: Any) → PreparedExplicitNonlinearDiagnostics

Prepare a reusable explicit diagnostic scan for repeated Python calls.

options accepts the same explicit-only keywords as integrate_nonlinear_explicit_diagnostics_state().

gkx.solvers_nonlinear_diagnostics.run_sampled_explicit_diagnostic_scan(step_fn: Any, initial_carry: tuple[Any, Any, Any, Any, Any, Any], *, steps: int, stride: int) → tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]

Run an explicit diagnostic scan only at retained sample intervals.

gkx.solvers_nonlinear_diagnostics.sampled_scan_intervals(length: int, stride: int) → ndarray

Return positive scan intervals that retain the requested final sample.

gkx.solvers_nonlinear_diagnostics.select_nonlinear_step_diagnostics(idx: Any, *, diagnostics_stride: int, diag_prev: Any, compute_diag_fn: Any, steps: int | None = None) → Any

Return a fresh or reused nonlinear step diagnostic tuple.

Off-stride steps reuse diag_prev. The last step of a steps-long scan is always computed fresh, because the output sampling keeps it as the final row even when it falls off-stride; reusing the carry there would report the previous stride sample’s diagnostics at the final time.

Nonlinear IMEX

IMEX nonlinear solve policies.

The public nonlinear facade builds operators and diagnostics. This module owns the small, reusable fixed-point predictor and GMRES solve step used by cached and diagnostic IMEX paths.

class gkx.solvers_nonlinear_imex.ChainImplicitLinear(groups: tuple[_ChainGroup, ...], shape: tuple[int, ...], dt: float, scheme: str, rows_m: tuple[int, ...], cols_m: tuple[int, ...], linear: Callable[[Array], Array], spec: tuple[tuple[int, ...], ...] = ((0,),), dims: tuple[tuple[int, int], ...] = (), on: Any = None, cfac: Any = None, collisions: Any = None)

Structured (I - gamma dt L)^{-1} for the stiff linear terms.

L = S + U Q per twist-shift chain: S is streaming with its dissipation (block-tridiagonal in Hermite, an N x N block per (species, Laguerre), N = links * Nz) and U Q the field response (Q z-local, rank nf N). I - gamma dt S is factored by block Thomas and the fields enter by Woodbury, so memory is O(ns Nl Nm N^2) per chain instead of O((ns Nl Nm N)^2). Species group by mirror and by moment count (spec, dims): each group is factored at its own (Nl, Nm) and the zero-padded rows above it stay the identity, so the cost follows the counts the species actually keep.

_ix(x: Array) → list[tuple]

Per species group: the index of its (s, l, m) block in x.

ars_step(G: Array, dG: Array, rhs: Callable[[Array], Array], project: Callable[[Array], Array], *, with_error: bool = False) → Any

One IMEX step; dG = rhs(G) is the full (linear + bracket) RHS.

Collisions given to the factor are not in rhs: the step ends with their sub-step (Lie splitting). with_error also returns the embedded error estimate (schemes with an embedded pair only).

collide(G: Array) → Array

One backward-Euler collision sub-step (identity when explicit).

solve(R: Array) → Array

(I - gamma dt L)^{-1} R on the chains (identity elsewhere).

With on set, one step of iterative refinement (set at build when the check residual exceeds 5e-3; one step takes 3e-2 to 1e-5 and 0.8 to 4e-3 at dt 0.2 and 0.8, 32x32, (4, 8), complex64).

class gkx.solvers_nonlinear_imex.ImexDtControl(op, cfl_rate, *, dt_max, cfl=0.9, rtol=0.001, iters=20)

Chunk-level dt for the chain IMEX step (fixed_dt = false).

Between chunks: dt = min(cfl 17 / w_nl, cfl 2.5 / rho, dt_err, dt_max) (dt_max also stays below COLLISION_DT_CAP / rho(C) when collisions are a sub-step). rho is the spectral radius of the explicit linear remainder (drifts, ion mirror, drive) by Arnoldi at startup and 2.5 the measured ARK3/ARS3 interval on it (0.82 x 3.6); w_nl the bracket CFL rate of the current fields; 17 = 10 x 1.7: Cyclone KE 32x32 stays finite and within 2 SE of rk4-pi to 34 at beta 1e-4, 1e-3 and 0.015, and steps stop shrinking above 17; dt_err the embedded ARK3(2) estimate against rtol: 1e-3 of the state per step, 2x the steps of 1e-2, which put the saturated heat flux 6% above rk4-pi (three seeds, Cyclone KE 32x32x16, m_e/m_i 1/1836); 1e-3 is within 1% on the seed mean (docs/guidance.rst). The factor is rebuilt only when dt moves by more than 20% (down) or grows past 1.2x, and a non-finite chunk is rejected and redone at half the step.

propose(op, G) → float

Next dt from the state at the end of a chunk.

retune(op, dt: float)

op at dt if it moved by more than 20%, else op.

gkx.solvers_nonlinear_imex.advance_imex_nonlinear_state(G: Array, *, dt_val: Array, method: str, nonlinear_term: Callable[[Array], Array], solve_step: Callable[[Array, Array], Array], project_state: Callable[[Array], Array]) → Array

Advance one IMEX nonlinear step with optional SSPX3 stage composition.

gkx.solvers_nonlinear_imex.advance_imex_nonlinear_state_with_stats(G: Array, stats: ImplicitSolveStats, *, dt_val: Array, method: str, nonlinear_term: Callable[[Array], Array], solve_step: Callable[[Array, Array, ImplicitSolveStats], tuple[Array, ImplicitSolveStats]], project_state: Callable[[Array], Array]) → tuple[Array, ImplicitSolveStats]

Advance one IMEX step and fold every implicit solve into stats.

gkx.solvers_nonlinear_imex.build_chain_implicit_linear(split_rhs: ~typing.Callable[[~jax.jaxlib._jax.Array, ~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array], fields: ~typing.Callable[[~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array], shape: tuple[int, ...], dt: float, *, modes: ~numpy.ndarray, scheme: str = 'imex-ars3', dtype=<class 'jax.numpy.complex64'>, check_tol: float | None = None, mirror_rhs: ~typing.Callable[[~jax.jaxlib._jax.Array, ~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array] | None = None, collision_rhs: ~typing.Callable[[~jax.jaxlib._jax.Array, ~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array] | None = None, moments: ~typing.Sequence[tuple[int, int]] | None = None) → ChainImplicitLinear

Probe the stiff linear terms per chain and factor I - gamma dt L.

fields(G) is the linear field solve, stacked (nf, Nky, Nkx, Nz); split_rhs(G, F) the implicit terms’ RHS of G with the fields F imposed. L G = split_rhs(G, fields(G)); S = split_rhs(., 0) and U = split_rhs(0, .) are probed separately, Q from unit moments. modes (Nky x Nkx bool) selects what is solved: every ky >= 0 mode, dealiased or not (leaving dealiased-out modes explicit makes their stiff terms explicit). ky < 0 rows pass through; the projector rebuilds them. mirror_rhs is the part of split_rhs that is the mirror force (species it is nonzero for get the Nl N blocks). collision_rhs(G, F) is the collision operator alone, same convention: given, every step ends with its local backward-Euler sub-step (_LocalCollisions), so nu does not bound dt; the caller then leaves collisions out of the explicit RHS. The factorization is checked by the residual of one solve. moments gives each species’ (Nl, Nm) when the state is zero-padded above them: species factor, and the collision sub-step acts, at their own counts.

gkx.solvers_nonlinear_imex.collisions_bind(collision_rhs, fields, shape, dtype, dt: float, moments=None) → bool

True if explicit collisions would bound dt: rho(C) dt past the limit.

gkx.solvers_nonlinear_imex.imex_fixed_point_guess(G_in: Array, G_rhs: Array, *, linear_rhs_fn: Callable[[...], tuple[Array, object]], cache: Any, params: Any, linear_cfg: object, external_phi: Array | float | None, dt_val: Array, implicit_iters: int, implicit_relax: float) → Array

Build the fixed-point predictor used as the GMRES initial guess.

gkx.solvers_nonlinear_imex.integrate_cached_imex_scan(G0: Array, cache: object, params: object, dt: float, steps: int, *, term_cfg: Any, linear_cfg: Any, linear_rhs_fn: Callable[[...], tuple[Array, object]], build_operator_fn: Callable[[...], Any], build_implicit_operator_fn: Callable[[...], tuple[Any, ...]] | None = None, fields_fn: Callable[[...], object], nonlinear_term_fn: Callable[[...], Array], nonlinear_contribution_fn: Callable[[...], Array], checkpoint: bool = False, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_preconditioner: str | None = None, implicit_operator: Any | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid', external_phi: Array | float | None = None, show_progress: bool = False, return_solve_stats: bool = False) → tuple[Any, ...]

Run the cached IMEX nonlinear scan.

The public facade injects field solves, operator construction, and RHS kernels so debug and monkeypatch seams stay outside this pure solver owner. Returns (G_out, fields_t), or (G_out, fields_t, stats) with return_solve_stats=True; stats summarizes every GMRES solve.

gkx.solvers_nonlinear_imex.make_imex_diagnostic_step(*, method: str, nonlinear_term: Callable[[Array], Array], solve_step: Callable[[Array, Array], Array], project_state: Callable[[Array], Array], state_dtype: dtype, real_dtype: dtype, dt_val: Array, compute_fields_fn: Callable[[...], object], cache: Any, params: Any, term_cfg: Any, external_phi: Array | float | None, compute_diag_from_state: Callable[[...], Any], diagnostics_stride: int, select_diagnostics_fn: Callable[[...], Any], show_progress: bool, steps: int, progress_total: Array, emit_progress_fn: Callable[[...], Array], use_collision_split: bool = False, damping: Any | None = None, collision_scheme: str = 'implicit', apply_collision_split_fn: Callable[[Array, Any, Array, str], Array] | None = None, solve_step_with_stats: Callable[[Array, Array, ImplicitSolveStats], tuple[Array, ImplicitSolveStats]] | None = None) → Callable[[tuple[Any, ...], Any], tuple[tuple[Any, ...], tuple[Any, Any]]]

Build one IMEX diagnostic scan step with injected runtime seams.

With solve_step_with_stats the carry grows one ImplicitSolveStats leaf and every implicit solve folds into it, which is how this route gets a convergence channel a traced step cannot raise on (queue row Q15). Without it the step keeps exactly the five-element carry, and the same solve, it had.

gkx.solvers_nonlinear_imex.make_imex_nonlinear_term(cache: object, params: object, term_cfg: object, *, real_dtype: object | None = None, external_phi: Array | float | None, compressed_real_fft: bool, laguerre_mode: str, fields_fn: Callable[[...], object], nonlinear_term_fn: Callable[[...], Array], nonlinear_contribution_fn: Callable[[...], Array] | None = None) → Callable[[Array], Array]

Return the explicit nonlinear term closure used by IMEX scans.

gkx.solvers_nonlinear_imex.make_imex_solve_step(*, linear_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, object]], cache: object, params: object, linear_cfg: object, external_phi: ~jax.jaxlib._jax.Array | float | None, dt_val: ~jax.jaxlib._jax.Array, implicit_iters: int, implicit_relax: float, matvec: ~typing.Callable[[~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array], shape: tuple[int, ...], implicit_tol: float, implicit_maxiter: int, implicit_restart: int, precond_op: ~typing.Callable[[~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array] | None, solve_step_fn: ~typing.Callable[[...], ~jax.jaxlib._jax.Array] = <function solve_imex_step>) → Callable[[Array, Array], Array]

Return the GMRES solve-step closure used by IMEX scan policies.

gkx.solvers_nonlinear_imex.make_imex_solve_step_with_stats(*, linear_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, object]], cache: object, params: object, linear_cfg: object, external_phi: ~jax.jaxlib._jax.Array | float | None, dt_val: ~jax.jaxlib._jax.Array, implicit_iters: int, implicit_relax: float, matvec: ~typing.Callable[[~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array], shape: tuple[int, ...], implicit_tol: float, implicit_maxiter: int, implicit_restart: int, precond_op: ~typing.Callable[[~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array] | None, solve_step_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~gkx.solvers_linear_implicit.ImplicitSolveStats]] = <function solve_imex_step_with_stats>) → Callable[[Array, Array, ImplicitSolveStats], tuple[Array, ImplicitSolveStats]]

Return make_imex_solve_step()’s closure, threading solve status.

The closure maps (state, rhs, stats) to (next state, stats), the same shape the implicit linear scan’s step uses, so a diagnostic scan can carry its convergence channel without a host callback inside the scan.

gkx.solvers_nonlinear_imex.refactor_chain_implicit_linear(op: ChainImplicitLinear, dt: float, *, free: bool = False) → Any

op at a new dt: refactor from the probed blocks, no new probes.

Same pytree structure, so a scan compiled for op reuses its executable. free releases op’s device factors first (two copies do not fit at production size); op is unusable afterwards.

gkx.solvers_nonlinear_imex.run_imex_diagnostic_scan(step_fn: Callable[[tuple[Any, ...], Any], tuple[tuple[Any, ...], tuple[Any, Any]]], initial_carry: tuple[Any, ...], *, steps: int, checkpoint: bool, return_solve_stats: bool = False) → tuple[Array, tuple[Any, Any]] | tuple[Any, ...]

Run the fixed-step IMEX diagnostic scan.

return_solve_stats=True appends the carry’s trailing ImplicitSolveStats, or None when the step was built without one. A caller that does not ask never reads the leaf, so XLA drops the status arithmetic and the compiled scan is the graph it was before.

gkx.solvers_nonlinear_imex.solve_imex_step(G_in: Array, G_rhs: Array, *, linear_rhs_fn: Callable[[...], tuple[Array, object]], cache: object, params: object, linear_cfg: object, external_phi: Array | float | None, dt_val: Array, implicit_iters: int, implicit_relax: float, matvec: Callable[[Array], Array], shape: tuple[int, ...], implicit_tol: float, implicit_maxiter: int, implicit_restart: int, precond_op: Callable[[Array], Array] | None = None) → Array

Solve one IMEX system with a predictor and an implicit solve VJP.

The primal and transpose solves use the same tolerance-controlled FGMRES policy. Reverse mode differentiates the converged linear system through SOLVAX rather than tracing the dynamic Krylov stopping loop. implicit_maxiter counts iterations. This form returns only the state; solve_imex_step_with_stats() also returns the convergence summary.

gkx.solvers_nonlinear_imex.solve_imex_step_with_stats(G_in: Array, G_rhs: Array, stats: ImplicitSolveStats, *, linear_rhs_fn: Callable[[...], tuple[Array, object]], cache: object, params: object, linear_cfg: object, external_phi: Array | float | None, dt_val: Array, implicit_iters: int, implicit_relax: float, matvec: Callable[[Array], Array], shape: tuple[int, ...], implicit_tol: float, implicit_maxiter: int, implicit_restart: int, precond_op: Callable[[Array], Array] | None = None) → tuple[Array, ImplicitSolveStats]

Solve one IMEX system as solve_imex_step() and fold its status.

SOLVAX’s true residual, iteration count and converged flag leave linear_solve as auxiliary data, so reverse mode still differentiates the solved system and never the stopping loop.

Explicit Time Integrators

Explicit linear time integrators implemented in JAX.

class gkx.solvers_time_explicit.ExplicitTimeConfig(t_max: float, dt: float, method: str = 'rk4', sample_stride: int = 1, fixed_dt: bool = False, use_dealias_mask: bool = False, dt_min: float = 1e-07, dt_max: float | None = None, cfl: float = 0.9, cfl_fac: float = 2.82)

Explicit time integration configuration.

dt is required here and defaulted in gkx.config.TimeConfig, which is why this struct can default fixed_dt to False: a caller who had to supply dt has chosen it, so using it as the CFL controller’s initial guess cannot silently substitute a step nobody asked for. The rk4/fixed_dt=False pairing is the cheap one – rk4’s CFL prefactor of 2.82 against rk2’s 1.0 more than pays for its four stages, measured at 29.1% fewer right-hand-side evaluations over the same horizon. See gkx.config.TimeConfig for the measurement and for why the deck surface defaults differently.

class gkx.solvers_time_explicit.ExplicitTimeConfigLike(*args, **kwargs)

Runtime fields required by explicit diagnostic integration.

gkx.solvers_time_explicit._apply_completed_step_state_mask(state: Array, cache: LinearCache) → Array

Apply the completed-step mask to a spectral state array.

gkx.solvers_time_explicit._completed_step_state_mask(cache: LinearCache) → Array

Return the completed-step state-space mask applied after each completed step.

gkx.solvers_time_explicit._growth_rate_mode_mask(ky: Array, kx: Array, dealias_mask: Array) → Array

Return the diagnostic mask used by explicit-time growth-rate extraction.

Single selected nonzonal ky slices should remain diagnosable even when the originating full nonlinear mesh would mark that representative row as dealiased away.

gkx.solvers_time_explicit._instantaneous_growth_rate_step(phi_now: Array, phi_prev: Array, dt: float | Array, *, z_index: int, mask: Array, mode_method: str = 'z_index') → tuple[Array, Array]

Instantaneous growth rates from phi ratios at the midplane.

gkx.solvers_time_explicit._laguerre_velocity_max(nl: int, *, node_max: bool = False) → float

nl-th Laguerre root; node_max: last node of the padded B_par grid.

gkx.solvers_time_explicit._linear_explicit_step(G: Array, cache: LinearCache, params: LinearParams, term_cfg, dt: float, *, method: str)

Explicit-module step seam used by tests and interactive diagnostics.

gkx.solvers_time_explicit._linear_native_step(G: Array, damping: Array, dt_val: Array, *, method_key: str, rhs: Callable[[Array], Array]) → Array

Advance one native explicit or diagonal-IMEX linear step.

gkx.solvers_time_explicit._parallel_periods_from_grid(grid: SpectralGrid) → float

Return the parallel extent of the flux tube in 2*pi periods.

Take the spacing off the host buffer rather than out of grid.z by jnp arithmetic. Inside a trace every jnp call stages out, so grid.z[1] - grid.z[0] is a tracer even when grid.z is the concrete axis built from the run’s own config, and reading it back on the host then fails. Nothing here is differentiable – the number of periods is grid topology – so this refused every adaptive-dt run under jit for a quantity that was never traced.

gkx.solvers_time_explicit.integrate_linear_explicit(G0: Array, grid: SpectralGrid, cache: LinearCache, params: LinearParams, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, time_cfg: ExplicitTimeConfig, terms: LinearTerms | None = None, *, mode_method: str = 'z_index', z_index: int | None = None, jit: bool = True, show_progress: bool = False) → tuple[ndarray, ndarray, ndarray, ndarray]

Explicit time integrator; shorten the final step to end at t_max.

gkx.solvers_time_explicit.integrate_linear_explicit_diagnostics(G0: Array, grid: SpectralGrid, cache: LinearCache, params: LinearParams, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, time_cfg: ExplicitTimeConfigLike, terms: LinearTerms | None = None, *, mode_method: str = 'z_index', z_index: int | None = None, jit: bool = True, show_progress: bool = False, linear_explicit_step_fn: Callable[[...], tuple[Array, Any]] | None = None) → tuple[ndarray, ndarray, ndarray, ndarray, SimulationDiagnostics]

Growth/energy/flux diagnostics, with the final step clipped to t_max.

gkx.solvers_time_explicit.integrate_linear_explicit_from_config(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, time_cfg: Any, *, Nl: int, Nm: int, terms: LinearTerms | None = None, z_index: int | None = None, show_progress: bool = False) → tuple[ndarray, ndarray]

Integrate with CFL control using the common TimeConfig contract.

Config-Driven Time Runners

Config-driven runners for time integration.

gkx.solvers_time_runners._check_moment_basis_matches_operator(operator, name: str, G0) → None

Reject a moment basis the tabulated collision matrix cannot act on.

The drift-kinetic Sugama matrices are the fixed truncation of Frei, Ernst & Ricci (2022), Appendix C, and the Coulomb tables ship at fixed orders, so each applies to one (Nl, Nm). The tables are Hermite-major: a transposed basis with the right moment count runs silently on the wrong moments.

gkx.solvers_time_runners._moment_count(G0) → int

Return the run’s Hermite-Laguerre moment count Nl*Nm.

gkx.solvers_time_runners._operator_moment_layout(operator) → tuple[int, int] | None

Return the (Nl, Nm) = (J+1, P+1) basis a tabulated operator acts on.

gkx.solvers_time_runners._resolve_config_collision_operator(time_cfg: TimeConfig, params: LinearParams, G0=None)

Build the moment collision operator selected by TimeConfig.

Returns None for "none"/"lenard_bernstein", which keeps the built-in diagonal Lenard-Bernstein term in the linear RHS. "sugama" and "improved_sugama" return the dense drift-kinetic Hermite-Laguerre moment operator that replaces that diagonal term.

gkx.solvers_time_runners._validate_nonlinear_config_state_sharding(spec: str | None) → None

Keep config-level nonlinear sharding on release-gated state axes.

gkx.solvers_time_runners.integrate_linear_from_config(G0, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, time_cfg: TimeConfig, *, cache: LinearCache | None = None, terms: LinearTerms | None = None, show_progress: bool | None = None, parallel: Any | None = None, return_solve_stats: bool = False) → tuple

Integrate the linear system using TimeConfig settings.

return_solve_stats=True appends integrate_linear()’s implicit solve summary (None for methods without an implicit solve).

gkx.solvers_time_runners.integrate_nonlinear_from_config(G0, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, time_cfg: TimeConfig, *, cache: LinearCache | None = None, terms: TermConfig | None = None, show_progress: bool | None = None, return_solve_stats: bool = False) → tuple

Integrate the nonlinear system using TimeConfig settings.

return_solve_stats=True appends the IMEX implicit solve summary (None for explicit and sharded runs, which have no implicit solve).

Benchmarks

Shared benchmark reference data, normalization, and scan policies.

class gkx.benchmarking_shared.CycloneReference(ky: 'np.ndarray', omega: 'np.ndarray', gamma: 'np.ndarray')

class gkx.benchmarking_shared.CycloneScanResult(ky: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray')

class gkx.benchmarking_shared.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 gkx.benchmarking_shared.LinearScanResult(ky: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray')

gkx.benchmarking_shared.load_cyclone_reference() → CycloneReference

Load Cyclone base case reference data (adiabatic electrons).

gkx.benchmarking_shared.load_cyclone_reference_kinetic() → CycloneReference

Load Cyclone base case reference data (kinetic electrons).

gkx.benchmarking_shared.load_etg_reference() → CycloneReference

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

gkx.benchmarking_shared.load_kbm_reference() → CycloneReference

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

gkx.benchmarking_shared.load_tem_reference() → CycloneReference

Load the provisional TEM reference digitized from the literature.

This lane remains an extended stress case while the literature case definition is being reconstructed.

Autodiff Validation

Finite-difference, covariance, and eigen-branch autodiff validation helpers.

gkx.objectives.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.

gkx.objectives.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.

gkx.objectives.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.

gkx.objectives.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.

gkx.objectives.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.

gkx.objectives.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>.

gkx.objectives.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

Parallel execution, decomposition, and sharding helpers.

class gkx.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: ~gkx.parallel.independent.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 gkx.parallel.IndependentMapExecutionError(index: int, executor: str, actual_workers: int, original_type: str, original_message: str)

Worker failure annotated with independent-map execution metadata.

class gkx.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 gkx.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.

class gkx.parallel.SpeciesHermiteRun(state: Array, traces: dict[str, Array], plan: Any, mesh_shape: tuple[int, int])

Result of one sharded nonlinear trajectory.

class gkx.parallel.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.

gkx.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.

gkx.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.

gkx.parallel.build_species_hermite_mesh(plan, *, devices: Iterable[Device] | None = None) → Mesh

Return the 2-D (species, hermite) device mesh a plan asks for.

Species vary slowest so that a species-only mesh keeps each species’ devices adjacent, which makes the halo-free configuration the natural one on a two-GPU box.

gkx.parallel.build_species_hermite_mesh_plan(state_shape: Sequence[int], *, num_devices: int, hermite_ghost_depth: int = 2) → VelocityShardingPlan

Factor num_devices over a species-first, Hermite-second device mesh.

Unlike the ceil-padded chunking of build_velocity_sharding_plan(), species are factored first, the Hermite remainder must divide the moment count exactly, and the local moment count must be at least the ghost depth so a width-hermite_ghost_depth halo is satisfiable from nearest neighbours alone. Anything else fails closed – a padded Hermite axis would put a shard’s ghost rows outside its neighbour, which is a wrong answer rather than a slow one.

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

Build a 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.

gkx.parallel.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.

gkx.parallel.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.

gkx.parallel.diamagnetic_drive_reference(state: Any, *, phi: Any, Jl: Any, b: 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.

gkx.parallel.diamagnetic_drive_shard_map(state: Any, plan: VelocityShardingPlan, *, phi: Any, Jl: Any, b: 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.

gkx.parallel.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, jacobian: Any | None = None, ky: Any | None = None) → Any

Return electrostatic phi from a full single- or multi-species state.

Pass jacobian and ky to include the ky = 0 adiabatic flux-surface-average term. Without them the zonal potential is over-screened; the production solve in terms/fields.py supplies both.

gkx.parallel.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', jacobian: Any | None = None, ky: Any | None = None) → Any

Solve electrostatic phi using a species- or Hermite-sharded reduction.

jacobian and ky enable the ky = 0 adiabatic flux-surface-average term, exactly as in electrostatic_phi_reference().

gkx.parallel.hermite_field_moment_head(local: Any, *, index: Any, chunks: int, axis_name: str, rows: int = 2) → Any

Return the leading global Hermite moment rows on every shard of the mesh.

The field solve reads quasineutrality from m = 0 and Ampere from m = 1, and the flux diagnostics reach m = 3; on a Hermite-sharded state those rows may live on more than one block. Selecting them by their global index and summing over the Hermite mesh axis delivers all of them to every shard in a single collective instead of a reduce followed by a broadcast. Each row’s sum has exactly one non-zero term, so the owner’s contribution survives bit-for-bit.

index gives the global Hermite index of each row of local.

gkx.parallel.hermite_halo_extend(local: Any, *, chunks: int, ghost_depth: int, axis_name: str, axis: int = 2) → Any

Widen a Hermite shard by ghost_depth rows from each neighbour.

One ppermute per direction carries the whole halo, so the m +- 1 ladder of streaming and mirror and the m +- 2 reach of curvature are paid for in a single exchange rather than one per coupling. Shards at the global Hermite boundary receive nothing, which is exactly the zero the serial shift_axis introduces there.

gkx.parallel.hermite_halo_interior(extended: Any, *, chunks: int, ghost_depth: int, axis: int = 2)

Drop the ghost rows a hermite_halo_extend() slab carries.

gkx.parallel.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.

gkx.parallel.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.

gkx.parallel.hermite_shift_reference(state: Any, *, offset: int) → Any

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

gkx.parallel.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.

gkx.parallel.hermite_streaming_ladder_reference(state: Any, *, vth: Any = 1.0) → Any

Return the full-array Hermite streaming ladder contribution.

gkx.parallel.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.

gkx.parallel.hermite_window_cache_arrays(cache: Any, indices: ndarray, *, hermite_count: int) → dict[str, Any]

Return every Hermite-indexed cache array gathered onto indices.

Ghost rows are zero-filled. That is safe rather than approximate: a ghost row’s coefficient only ever multiplies a ghost row’s output, and those are sliced away before the shard returns.

gkx.parallel.hermite_window_indices(hermite_count: int, *, chunks: int, ghost_depth: int) → ndarray

Return the global Hermite index of every extended row of every shard.

The shards’ windows are concatenated, so the result shards cleanly along the Hermite mesh axis and each device receives exactly its own extended coordinates. Indices outside [0, Nm) mark ghost rows whose outputs are discarded and whose physical boundary is zero.

gkx.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.

gkx.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.

gkx.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.

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

Resolve independent-task worker counts and normalized executor metadata.

gkx.parallel.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.

gkx.parallel.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, compressed_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; an extent the device count does not divide is zero-padded while it is carried (_resolve_state_split()). Every device evaluates each step on the whole state, so only the carried state is split: this spreads memory between steps, not work. It is a diagnostic whole-state sharding primitive for identity gates and profiler localization, not a production nonlinear domain decomposition or speedup claim. Domain-sharding identity reports are metadata gates only; they do not authorize routing through this whole-state integrator.

gkx.parallel.integrate_nonlinear_species_hermite(G0: Array, cache: LinearCache, params: LinearParams, *, dt: float, steps: int, method: str = 'rk2', terms: TermConfig | None = None, grid: Any | None = None, vol_fac: Array | None = None, flux_fac: Array | None = None, plan: Any | None = None, devices: Any | None = None, num_devices: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid') → SpeciesHermiteRun

Integrate the nonlinear system on a species x Hermite shard_map mesh.

The whole trajectory is one jitted lax.scan on the mesh, and the scalar traces are stacked from inside that scan rather than recomputed afterwards. Supplying grid with vol_fac/flux_fac turns the traces on; without them the route is compute-only and returns an empty trace map.

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

Return balanced k_y chunks for independent linear-scan execution.

gkx.parallel.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.

gkx.parallel.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.

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

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

gkx.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.

gkx.parallel.periodic_streaming_reference(state: Any, *, kz: Any, vth: Any = 1.0) → Any

Return periodic parallel streaming using full-array operations.

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

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

gkx.parallel.resolve_species_hermite_sharding(G0: Array, plan, *, devices: Iterable[Device] | None = None, mesh: Mesh | None = None) → NamedSharding

Return the NamedSharding that places a packed state on the mesh.

gkx.parallel.resolve_state_sharding(G0: Array, spec: str | None, *, axis_name: str = 'd', devices: Iterable[Device] | None = None) → NamedSharding | None

Return a NamedSharding for the packed state, or None if disabled.

Parameters:

• G0 (jnp.ndarray) – Initial state array with shape (Nl, Nm, Ny, Nx, Nz) or (Ns, Nl, Nm, Ny, Nx, Nz).

• spec (str | None) – Sharding directive. Allowed values: - None / “none” / “off”: disable sharding - “auto” or “ky”: shard along ky (recommended default) - “kx”, “z”, “l”, “m”, “species”: shard along the named axis

• axis_name (str) – Mesh axis name for the sharded dimension.

• devices (Iterable[jax.Device] | None) – Optional explicit device list (useful for tests).

gkx.parallel.species_hermite_device_counts(ns: int, nm: int, *, hermite_ghost_depth: int = 2, limit: int = 64) → tuple[int, ...]

Return the device counts a (Ns, Nm) grid can be decomposed over.

Reported verbatim in routing errors so an unroutable request names its alternatives instead of only its failure.

gkx.parallel.species_hermite_state_spec(ndim: int) → P

Return P('s', None, 'm', None, None, None) for a packed state.

(ky, kx), Laguerre and z stay replicated: every FFT in the bracket and in parallel streaming runs over one of them, so sharding any of them would put a distributed transpose inside the RHS.

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

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

gkx.parallel.stage_from_host(value: Any, sharding: Any) → Any

Place value on sharding from host memory rather than resharding.

Handing device_put an array that is already committed to one device asks the runtime to reshard it across the mesh, and on the two-GPU box that path silently produces a wrong answer: the sharded route drops to a max relative error of 1.0 against serial while a one-device mesh on the same GPU is exact. Round-tripping through the host is the reliable placement, and it happens once per run rather than once per step.

gkx.parallel.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

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

gkx.parallel.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.

Parallel Independent Tasks

Ordered Python-task parallelism for UQ and optimization ensembles.

class gkx.parallel.independent.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: ~gkx.parallel.independent.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 gkx.parallel.independent.IndependentMapExecutionError(index: int, executor: str, actual_workers: int, original_type: str, original_message: str)

Worker failure annotated with independent-map execution metadata.

class gkx.parallel.independent.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 gkx.parallel.independent.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.

gkx.parallel.independent.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.

gkx.parallel.independent.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.

gkx.parallel.independent.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.

gkx.parallel.independent.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.

gkx.parallel.independent.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.

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

Resolve independent-task worker counts and normalized executor metadata.

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

Return balanced k_y chunks for independent linear-scan execution.

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

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

gkx.parallel.independent.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.

gkx.parallel.independent.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 gkx.parallel.velocity.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.

gkx.parallel.velocity.build_species_hermite_mesh_plan(state_shape: Sequence[int], *, num_devices: int, hermite_ghost_depth: int = 2) → VelocityShardingPlan

Factor num_devices over a species-first, Hermite-second device mesh.

Unlike the ceil-padded chunking of build_velocity_sharding_plan(), species are factored first, the Hermite remainder must divide the moment count exactly, and the local moment count must be at least the ghost depth so a width-hermite_ghost_depth halo is satisfiable from nearest neighbours alone. Anything else fails closed – a padded Hermite axis would put a shard’s ghost rows outside its neighbour, which is a wrong answer rather than a slow one.

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

Build a 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.

gkx.parallel.velocity.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.

gkx.parallel.velocity.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.

gkx.parallel.velocity.diamagnetic_drive_reference(state: Any, *, phi: Any, Jl: Any, b: 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.

gkx.parallel.velocity.diamagnetic_drive_shard_map(state: Any, plan: VelocityShardingPlan, *, phi: Any, Jl: Any, b: 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.

gkx.parallel.velocity.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, jacobian: Any | None = None, ky: Any | None = None) → Any

Return electrostatic phi from a full single- or multi-species state.

Pass jacobian and ky to include the ky = 0 adiabatic flux-surface-average term. Without them the zonal potential is over-screened; the production solve in terms/fields.py supplies both.

gkx.parallel.velocity.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', jacobian: Any | None = None, ky: Any | None = None) → Any

Solve electrostatic phi using a species- or Hermite-sharded reduction.

jacobian and ky enable the ky = 0 adiabatic flux-surface-average term, exactly as in electrostatic_phi_reference().

gkx.parallel.velocity.hermite_field_moment_head(local: Any, *, index: Any, chunks: int, axis_name: str, rows: int = 2) → Any

Return the leading global Hermite moment rows on every shard of the mesh.

The field solve reads quasineutrality from m = 0 and Ampere from m = 1, and the flux diagnostics reach m = 3; on a Hermite-sharded state those rows may live on more than one block. Selecting them by their global index and summing over the Hermite mesh axis delivers all of them to every shard in a single collective instead of a reduce followed by a broadcast. Each row’s sum has exactly one non-zero term, so the owner’s contribution survives bit-for-bit.

index gives the global Hermite index of each row of local.

gkx.parallel.velocity.hermite_halo_extend(local: Any, *, chunks: int, ghost_depth: int, axis_name: str, axis: int = 2) → Any

Widen a Hermite shard by ghost_depth rows from each neighbour.

One ppermute per direction carries the whole halo, so the m +- 1 ladder of streaming and mirror and the m +- 2 reach of curvature are paid for in a single exchange rather than one per coupling. Shards at the global Hermite boundary receive nothing, which is exactly the zero the serial shift_axis introduces there.

gkx.parallel.velocity.hermite_halo_interior(extended: Any, *, chunks: int, ghost_depth: int, axis: int = 2)

Drop the ghost rows a hermite_halo_extend() slab carries.

gkx.parallel.velocity.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.

gkx.parallel.velocity.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.

gkx.parallel.velocity.hermite_shift_reference(state: Any, *, offset: int) → Any

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

gkx.parallel.velocity.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.

gkx.parallel.velocity.hermite_streaming_ladder_reference(state: Any, *, vth: Any = 1.0) → Any

Return the full-array Hermite streaming ladder contribution.

gkx.parallel.velocity.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.

gkx.parallel.velocity.hermite_window_cache_arrays(cache: Any, indices: ndarray, *, hermite_count: int) → dict[str, Any]

Return every Hermite-indexed cache array gathered onto indices.

Ghost rows are zero-filled. That is safe rather than approximate: a ghost row’s coefficient only ever multiplies a ghost row’s output, and those are sliced away before the shard returns.

gkx.parallel.velocity.hermite_window_indices(hermite_count: int, *, chunks: int, ghost_depth: int) → ndarray

Return the global Hermite index of every extended row of every shard.

The shards’ windows are concatenated, so the result shards cleanly along the Hermite mesh axis and each device receives exactly its own extended coordinates. Indices outside [0, Nm) mark ghost rows whose outputs are discarded and whose physical boundary is zero.

gkx.parallel.velocity.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.

gkx.parallel.velocity.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.

gkx.parallel.velocity.periodic_streaming_reference(state: Any, *, kz: Any, vth: Any = 1.0) → Any

Return periodic parallel streaming using full-array operations.

gkx.parallel.velocity.periodic_streaming_shard_map(state: Any, plan: VelocityShardingPlan, *, kz: Any, vth: Any = 1.0, ny_full: int | None = None, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

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

gkx.parallel.velocity.species_hermite_device_counts(ns: int, nm: int, *, hermite_ghost_depth: int = 2, limit: int = 64) → tuple[int, ...]

Return the device counts a (Ns, Nm) grid can be decomposed over.

Reported verbatim in routing errors so an unroutable request names its alternatives instead of only its failure.

gkx.parallel.velocity.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

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

gkx.parallel.velocity.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.

Velocity-space decomposition metadata and planning helpers.

class gkx.parallel.velocity_plan.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.

gkx.parallel.velocity_plan.build_species_hermite_mesh_plan(state_shape: Sequence[int], *, num_devices: int, hermite_ghost_depth: int = 2) → VelocityShardingPlan

Factor num_devices over a species-first, Hermite-second device mesh.

Unlike the ceil-padded chunking of build_velocity_sharding_plan(), species are factored first, the Hermite remainder must divide the moment count exactly, and the local moment count must be at least the ghost depth so a width-hermite_ghost_depth halo is satisfiable from nearest neighbours alone. Anything else fails closed – a padded Hermite axis would put a shard’s ghost rows outside its neighbour, which is a wrong answer rather than a slow one.

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

Build a 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.

gkx.parallel.velocity_plan.species_hermite_device_counts(ns: int, nm: int, *, hermite_ghost_depth: int = 2, limit: int = 64) → tuple[int, ...]

Return the device counts a (Ns, Nm) grid can be decomposed over.

Reported verbatim in routing errors so an unroutable request names its alternatives instead of only its failure.

Hermite-axis exchange and velocity-field reduction kernels.

gkx.parallel.velocity_hermite.hermite_cache_partition_specs(windowed: dict[str, Any], axis_name: str) → dict

Return the PartitionSpec for each windowed Hermite cache array.

gkx.parallel.velocity_hermite.hermite_field_moment_head(local: Any, *, index: Any, chunks: int, axis_name: str, rows: int = 2) → Any

Return the leading global Hermite moment rows on every shard of the mesh.

The field solve reads quasineutrality from m = 0 and Ampere from m = 1, and the flux diagnostics reach m = 3; on a Hermite-sharded state those rows may live on more than one block. Selecting them by their global index and summing over the Hermite mesh axis delivers all of them to every shard in a single collective instead of a reduce followed by a broadcast. Each row’s sum has exactly one non-zero term, so the owner’s contribution survives bit-for-bit.

index gives the global Hermite index of each row of local.

gkx.parallel.velocity_hermite.hermite_halo_extend(local: Any, *, chunks: int, ghost_depth: int, axis_name: str, axis: int = 2) → Any

Widen a Hermite shard by ghost_depth rows from each neighbour.

One ppermute per direction carries the whole halo, so the m +- 1 ladder of streaming and mirror and the m +- 2 reach of curvature are paid for in a single exchange rather than one per coupling. Shards at the global Hermite boundary receive nothing, which is exactly the zero the serial shift_axis introduces there.

gkx.parallel.velocity_hermite.hermite_halo_interior(extended: Any, *, chunks: int, ghost_depth: int, axis: int = 2)

Drop the ghost rows a hermite_halo_extend() slab carries.

gkx.parallel.velocity_hermite.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.

gkx.parallel.velocity_hermite.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.

gkx.parallel.velocity_hermite.hermite_shift_reference(state: Any, *, offset: int) → Any

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

gkx.parallel.velocity_hermite.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.

gkx.parallel.velocity_hermite.hermite_window_cache_arrays(cache: Any, indices: ndarray, *, hermite_count: int) → dict[str, Any]

Return every Hermite-indexed cache array gathered onto indices.

Ghost rows are zero-filled. That is safe rather than approximate: a ghost row’s coefficient only ever multiplies a ghost row’s output, and those are sliced away before the shard returns.

gkx.parallel.velocity_hermite.hermite_window_indices(hermite_count: int, *, chunks: int, ghost_depth: int) → ndarray

Return the global Hermite index of every extended row of every shard.

The shards’ windows are concatenated, so the result shards cleanly along the Hermite mesh axis and each device receives exactly its own extended coordinates. Indices outside [0, Nm) mark ghost rows whose outputs are discarded and whose physical boundary is zero.

gkx.parallel.velocity_hermite.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

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

gkx.parallel.velocity_hermite.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.

Electrostatic field reduction and diamagnetic-drive microkernels.

gkx.parallel.velocity_drive.diamagnetic_drive_reference(state: Any, *, phi: Any, Jl: Any, b: 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.

gkx.parallel.velocity_drive.diamagnetic_drive_shard_map(state: Any, plan: VelocityShardingPlan, *, phi: Any, Jl: Any, b: 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.

gkx.parallel.velocity_drive.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, jacobian: Any | None = None, ky: Any | None = None) → Any

Return electrostatic phi from a full single- or multi-species state.

Pass jacobian and ky to include the ky = 0 adiabatic flux-surface-average term. Without them the zonal potential is over-screened; the production solve in terms/fields.py supplies both.

gkx.parallel.velocity_drive.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', jacobian: Any | None = None, ky: Any | None = None) → Any

Solve electrostatic phi using a species- or Hermite-sharded reduction.

jacobian and ky enable the ky = 0 adiabatic flux-surface-average term, exactly as in electrostatic_phi_reference().

State Sharding Policy

Helpers for multi-device sharding of GK state arrays.

gkx.parallel.state.build_species_hermite_mesh(plan, *, devices: Iterable[Device] | None = None) → Mesh

Return the 2-D (species, hermite) device mesh a plan asks for.

Species vary slowest so that a species-only mesh keeps each species’ devices adjacent, which makes the halo-free configuration the natural one on a two-GPU box.

gkx.parallel.state.resolve_species_hermite_sharding(G0: Array, plan, *, devices: Iterable[Device] | None = None, mesh: Mesh | None = None) → NamedSharding

Return the NamedSharding that places a packed state on the mesh.

gkx.parallel.state.resolve_state_sharding(G0: Array, spec: str | None, *, axis_name: str = 'd', devices: Iterable[Device] | None = None) → NamedSharding | None

Return a NamedSharding for the packed state, or None if disabled.

Parameters:

• G0 (jnp.ndarray) – Initial state array with shape (Nl, Nm, Ny, Nx, Nz) or (Ns, Nl, Nm, Ny, Nx, Nz).

• spec (str | None) – Sharding directive. Allowed values: - None / “none” / “off”: disable sharding - “auto” or “ky”: shard along ky (recommended default) - “kx”, “z”, “l”, “m”, “species”: shard along the named axis

• axis_name (str) – Mesh axis name for the sharded dimension.

• devices (Iterable[jax.Device] | None) – Optional explicit device list (useful for tests).

gkx.parallel.state.species_hermite_state_spec(ndim: int) → P

Return P('s', None, 'm', None, None, None) for a packed state.

(ky, kx), Laguerre and z stay replicated: every FFT in the bracket and in parallel streaming runs over one of them, so sharding any of them would put a distributed transpose inside the RHS.

Sharded Integrators

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

class gkx.parallel.integrators.SpeciesHermiteRun(state: Array, traces: dict[str, Array], plan: Any, mesh_shape: tuple[int, int])

Result of one sharded nonlinear trajectory.

gkx.parallel.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.

gkx.parallel.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, compressed_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; an extent the device count does not divide is zero-padded while it is carried (_resolve_state_split()). Every device evaluates each step on the whole state, so only the carried state is split: this spreads memory between steps, not work. It is a diagnostic whole-state sharding primitive for identity gates and profiler localization, not a production nonlinear domain decomposition or speedup claim. Domain-sharding identity reports are metadata gates only; they do not authorize routing through this whole-state integrator.

gkx.parallel.integrators.integrate_nonlinear_species_hermite(G0: Array, cache: LinearCache, params: LinearParams, *, dt: float, steps: int, method: str = 'rk2', terms: TermConfig | None = None, grid: Any | None = None, vol_fac: Array | None = None, flux_fac: Array | None = None, plan: Any | None = None, devices: Any | None = None, num_devices: int | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid') → SpeciesHermiteRun

Integrate the nonlinear system on a species x Hermite shard_map mesh.

The whole trajectory is one jitted lax.scan on the mesh, and the scalar traces are stacked from inside that scan rather than recomputed afterwards. Supplying grid with vol_fac/flux_fac turns the traces on; without them the route is compute-only and returns an empty trace map.

gkx.parallel.integrators.stage_from_host(value: Any, sharding: Any) → Any

Place value on sharding from host memory rather than resharding.

Handing device_put an array that is already committed to one device asks the runtime to reshard it across the mesh, and on the two-GPU box that path silently produces a wrong answer: the sharded route drops to a max relative error of 1.0 against serial while a one-device mesh on the same GPU is exact. Round-tripping through the host is the reliable placement, and it happens once per run rather than once per step.

Analysis

Public facade for mode extraction and growth-rate diagnostics.

class gkx.diagnostics.analysis.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 gkx.diagnostics.analysis.CFLLimiterReport(omega_linear_floor: float, omega_total_initial: float, omega_total_final: float, omega_total_max: float, exb_share_initial: float, exb_share_final: float, exb_growth_ratio: float, limiting_term_initial: str, limiting_term_final: str, limiting_term_final_share: float, limiting_term_sample_fractions: dict[str, float], samples_total: int, samples_attributed: int, samples_at_dt_floor: int, samples_at_dt_ceiling: int)

How the CFL frequency split between terms over a nonlinear run.

class gkx.diagnostics.analysis.CFLScales(omega_magnetic_drift_radial: float, omega_magnetic_drift_binormal: float, omega_parallel_streaming: float, dt_cfl_numerator: float, dt_min: float, dt_max: float)

Run-constant CFL scales recorded alongside a nonlinear dt trajectory.

contributions_for(omega_total: float) → dict[str, float]

Per-term contributions implied by an inverted total frequency.

property omega_linear_floor: float

CFL frequency the linear terms impose regardless of turbulence.

gkx.diagnostics.analysis.CFL_TERM_NAMES: tuple[str, ...] = ('magnetic_drift_radial', 'magnetic_drift_binormal', 'parallel_streaming', 'exb')

Terms that can set the nonlinear CFL frequency, in tie-break order.

gkx.diagnostics.analysis.CFL_TERM_UNRESOLVED = 'unresolved'

Reported instead of a term name when no sample could be attributed.

class gkx.diagnostics.analysis.GrowthRateFitStats(gamma: float, omega: float, gamma_stderr: float, omega_stderr: float, r2_log: float, r2_phase: float, n_points: int, tmin: float, tmax: float)

Log-linear gamma/omega fit with standard errors over one window.

class gkx.diagnostics.analysis.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 gkx.diagnostics.analysis.ModeSelection(ky_index: 'int', kx_index: 'int', z_index: 'int' = 0)

class gkx.diagnostics.analysis.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, tau_ac: float = 0.0, n_eff: float = 0.0)

Post-transient heat-flux averaging convergence summary.

class gkx.diagnostics.analysis.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, heat_flux_tau_ac: float = 0.0, window_in_tau_ac: float = 0.0, heat_flux_n_eff: float = 0.0, heat_flux_stderr: float = inf)

Windowed transport/envelope metrics for a nonlinear run.

class gkx.diagnostics.analysis.ObservedOrderMetrics(step_sizes: ndarray, errors: ndarray, orders: ndarray, asymptotic_order: float)

Observed-order convergence summary from step sizes and errors.

gkx.diagnostics.analysis.cfl_limiter_report(dt: ndarray, scales: CFLScales) → CFLLimiterReport

Attribute a recorded dt trajectory to the CFL terms that produced it.

Samples clipped at dt_max are excluded from the attribution. There the solver wanted a larger step than the ceiling allows, so the inverted frequency is only an upper bound and naming a limiting term from it would invent one: a run that never leaves its ceiling is capped, not CFL-limited, and is reported as unresolved with the clip counts to say so. Samples pinned at dt_min are kept – there the inverted frequency is a lower bound, so any ExB excess it shows is real and understated, not invented.

gkx.diagnostics.analysis.cfl_limiting_term(contributions: dict[str, float]) → tuple[str, float]

Return the largest CFL contribution and its share of the total.

Ties resolve toward the earlier CFL_TERM_NAMES entry, so a run whose ExB frequency merely matches a drift is not called ExB-limited.

gkx.diagnostics.analysis.cfl_scales_from_array(scales: object | None) → CFLScales | None

Decode a recorded cfl_scales vector, or None when unusable.

gkx.diagnostics.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).

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

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

gkx.diagnostics.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).

gkx.diagnostics.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).

gkx.diagnostics.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 = 'stationary', 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.

window_method="stationary" (the default) selects the longest late-time window over which the instantaneous growth rate is stationary and falls back to the "loglinear" score search with a RuntimeWarning when no stationary window exists.

gkx.diagnostics.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 = 'stationary', 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.

gkx.diagnostics.analysis.fit_growth_rate_uncertainty(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None) → GrowthRateFitStats

Fit gamma/omega over a window and report stderr and R^2 diagnostics.

gkx.diagnostics.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).

gkx.diagnostics.analysis.instantaneous_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 instantaneous growth and frequency from complex mode ratios.

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

gkx.diagnostics.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.

gkx.diagnostics.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.

gkx.diagnostics.analysis.select_fit_window_stationary(t: ndarray, signal: ndarray, *, min_points: int = 20, k_std: float = 3.0, rel_tolerance: float = 0.02, min_growth_times: float = 2.0, smooth_points: int | None = None, end_scan_fraction: float = 0.5, max_end_candidates: int = 32) → Tuple[float, float] | None

Select the longest late-time window over which gamma(t) is stationary.

The instantaneous growth rate gamma(t) is computed from the log-amplitude derivative of the signal, smoothed with a short rolling mean, and windows are grown backward from a set of candidate late-time end points: a point joins the window while it stays within k_std standard deviations (or rel_tolerance times the mean, whichever band is wider) of the running window mean. Candidate end points earlier than the final sample let the search step over a trailing recurrence/overflow tail. A window is only accepted when it spans at least min_growth_times growth times (min_growth_times / gamma) and min_points samples; None is returned when no such window exists, so callers can fall back to another selection method.

gkx.diagnostics.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.

Mode Diagnostics

Mode-selection and eigenfunction extraction diagnostics.

class gkx.diagnostics.modes.EigenfunctionComparisonMetrics(overlap: float, relative_l2: float, phase_shift: float)

Phase-aligned eigenfunction comparison summary.

class gkx.diagnostics.modes.ModeSelection(ky_index: 'int', kx_index: 'int', z_index: 'int' = 0)

gkx.diagnostics.modes.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).

gkx.diagnostics.modes.extract_mode(phi_t: ndarray, sel: ModeSelection) → ndarray

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

gkx.diagnostics.modes.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).

gkx.diagnostics.modes.normalize_eigenfunction(eigenfunction: ndarray, z: ndarray) → ndarray

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

gkx.diagnostics.modes.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.

Growth-Rate Diagnostics

Growth-rate, frequency, least-squares, and fit-window diagnostics.

class gkx.diagnostics.growth_rates.GrowthRateFitStats(gamma: float, omega: float, gamma_stderr: float, omega_stderr: float, r2_log: float, r2_phase: float, n_points: int, tmin: float, tmax: float)

Log-linear gamma/omega fit with standard errors over one window.

gkx.diagnostics.growth_rates._log_amp_phase(signal: ndarray) → Tuple[ndarray, ndarray]

Return (log|signal|, unwrapped phase) with robust scaling.

gkx.diagnostics.growth_rates.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).

gkx.diagnostics.growth_rates.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 = 'stationary', 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.

window_method="stationary" (the default) selects the longest late-time window over which the instantaneous growth rate is stationary and falls back to the "loglinear" score search with a RuntimeWarning when no stationary window exists.

gkx.diagnostics.growth_rates.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 = 'stationary', 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.

gkx.diagnostics.growth_rates.fit_growth_rate_uncertainty(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None) → GrowthRateFitStats

Fit gamma/omega over a window and report stderr and R^2 diagnostics.

gkx.diagnostics.growth_rates.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).

gkx.diagnostics.growth_rates.instantaneous_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 instantaneous growth and frequency from complex mode ratios.

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

gkx.diagnostics.growth_rates.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.

gkx.diagnostics.growth_rates.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.

gkx.diagnostics.growth_rates.select_fit_window_stationary(t: ndarray, signal: ndarray, *, min_points: int = 20, k_std: float = 3.0, rel_tolerance: float = 0.02, min_growth_times: float = 2.0, smooth_points: int | None = None, end_scan_fraction: float = 0.5, max_end_candidates: int = 32) → Tuple[float, float] | None

Select the longest late-time window over which gamma(t) is stationary.

The instantaneous growth rate gamma(t) is computed from the log-amplitude derivative of the signal, smoothed with a short rolling mean, and windows are grown backward from a set of candidate late-time end points: a point joins the window while it stays within k_std standard deviations (or rel_tolerance times the mean, whichever band is wider) of the running window mean. Candidate end points earlier than the final sample let the search step over a trailing recurrence/overflow tail. A window is only accepted when it spans at least min_growth_times growth times (min_growth_times / gamma) and min_points samples; None is returned when no such window exists, so callers can fall back to another selection method.

Publication Plotting

Publication-ready benchmark, diagnostic, runtime, and zonal plots.

class gkx.artifacts.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)

gkx.artifacts.plotting.cyclone_comparison_figure(ref: CycloneReference, scan: CycloneScanResult, label: str = 'GKX') → Tuple[Figure, ndarray]

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

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

Create a two-panel Cyclone base case reference plot.

gkx.artifacts.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.

gkx.artifacts.plotting.linear_runtime_panel_figure(*, t: ndarray, signal: ndarray, z: ndarray, eigenfunction: ndarray, gamma: float, omega: float, title: str = 'GKX Linear Runtime') → Tuple[Figure, ndarray]

Create the default two-panel linear runtime plot.

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

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

gkx.artifacts.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 = 'GKX Nonlinear Runtime') → Tuple[Figure, ndarray]

Create the default three-panel nonlinear runtime plot.

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

Plot a saved linear or nonlinear output bundle.

A bundle written by another gyrokinetic code is accepted alongside GKX’s own, so a cross-code comparison is one command rather than a script; see gkx.artifacts.foreign_output for how they are told apart.

gkx.artifacts.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 = 'GKX', ref_label: str = 'Reference', log_x: bool = False) → Tuple[Figure, ndarray]

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

gkx.artifacts.plotting.set_plot_style() → None

Apply the shared publication style used by generated figures.

Transport Figures

Figures for saved or in-memory nonlinear transport diagnostics.

CSV sidecars carry time traces; resolved spectra require a NetCDF bundle.

gkx.artifacts.transport_figures.flux_spectra_figure(source: Any, *, ky: ndarray | None = None, kx: ndarray | None = None, window: tuple[float, float] | None = None, species_labels: list[str] | None = None, title: str = 'Heat-flux spectra', out: str | Path | None = None, axes: Any = None, panels: Any = None) → Tuple[Figure, ndarray]

Plot window-averaged Q(ky) and Q(kx).

NetCDF or resolved in-memory diagnostics are accepted; CSV is not. panels/axes support embedding selected panels.

gkx.artifacts.transport_figures.heat_flux_time_figure(source: Any, *, window: tuple[float, float] | None = None, species_labels: list[str] | None = None, title: str = 'Turbulent fluxes', out: str | Path | None = None, axes: Any = None) → Tuple[Figure, ndarray]

Plot Q(t) and Gamma(t) with an optional measured window.

Sources may be diagnostics, NetCDF, or CSV. Figure SEM treats samples as independent; published estimates need correlation-corrected uncertainty. axes= embeds the two panels in another figure.

gkx.artifacts.transport_figures.phi2_spectra_figure(source: Any, *, ky: ndarray | None = None, kx: ndarray | None = None, window: tuple[float, float] | None = None, title: str = 'Potential spectra', out: str | Path | None = None, axes: Any = None, panels: Any = None) → Tuple[Figure, ndarray]

Plot Phi^2 by ky, kx, 2-D map, and zonal share.

Sources and embedding follow flux_spectra_figure().

Automatic Run Figures

Render run figures; plotting failures warn but never fail a simulation.

gkx.artifacts.run_figures.auto_plot_saved_run(kind: str, paths: Mapping[str, str], *, window: tuple[float, float] | None = None) → list[str]

Render the figure set for one completed run of kind.

paths is the artifact mapping the run’s writer returned. Returns the figure paths that were actually written, which is empty – never an exception – when the run produced nothing plottable.

gkx.artifacts.run_figures.is_gkx_nonlinear_bundle(source: str | Path) → bool

Return whether source is a GKX nonlinear NetCDF history bundle.

gkx.artifacts.run_figures.measured_average_window(summary: Mapping[str, Any] | None) → tuple[float, float] | None

Return a recorded averaging window without inventing a default.

gkx.artifacts.run_figures.replot_nonlinear_bundle(source: str | Path) → list[Path]

Re-render the standard figure set for an already-saved nonlinear bundle.

gkx --plot renders one panel from whatever it is given, including other codes’ output. When the file is GKX’s own nonlinear bundle there is a whole set to reproduce, and the window to reproduce it over is in the summary sidecar the run wrote next to it – so a re-plot shows the same measured window the run reported rather than a fresh guess.

gkx.artifacts.run_figures.write_nonlinear_run_figures(source: str | Path, *, base: str | Path | None = None, window: tuple[float, float] | None = None, saturated: bool | None = None, label: str | None = None) → list[Path]

Write Q(t)/Gamma(t) plus the spectra a NetCDF bundle supports.

gkx.artifacts.run_figures.write_panel_run_figure(source: str | Path) → list[Path]

Write the single saved-output panel used by linear points and scans.

Whole-Run Summary Figure

The one figure a finished nonlinear run is read from.

The individual figures a run writes each answer one question. A person opening a result directory for the first time has a different question – what happened? – and answering it from six PNGs means knowing which to open in which order. This module composes the single page that answers it: the flux traces with the window the run actually measured, the ky spectra that say which scales carry the transport, the real-space potential, and a text panel naming the equilibrium, the resolution, the deck, and the saturated <Q> +/- SEM.

Nothing here re-derives a figure. The panels are drawn by the same functions that draw them standalone (gkx.artifacts.transport_figures and gkx.artifacts.run_summary), through their axes=/panels= arguments, so the summary cannot drift away from the figures it summarizes.

The one thing this module does own is reading the final field sidecar. The *.out.nc history carries spectra but no potential; the potential lives in the *.big.nc companion, stored through numpy.fft.ifft2 and therefore scaled by 1/(Ny Nx) relative to the convention the solver and gkx.artifacts.run_summary.potential_real_space() use. Undoing that here is what keeps the amplitude on the colorbar the same number the rest of the run reports.

gkx.artifacts.run_summary.FINAL_FIELD_SUFFIX = '.big.nc'

Companion file holding the final real-space fields of a nonlinear run.

class gkx.artifacts.run_summary.FinalField(phi_xyz: ndarray, extent: tuple[float, float] | None, geometry: Any, time: float | None)

Final-time potential of a nonlinear run, in solver amplitude units.

gkx.artifacts.run_summary.PHI_LABEL: str = '$(e\\phi/T_i)\\,/\\,\\rho_*$'

Shipped decks use diagnostic_norm = "rho_star"; this is the rho-star-normalized potential, not ephi/T_i.

gkx.artifacts.run_summary.draw_flux_tube_3d(ax3d: Any, phi_xyz: ndarray, geometry: Any, *, scale: float, turns: float = 1.5, elev: float = 32.0, azim: float = -60.0, radius_fraction: float = 0.85, cmap: str = 'RdBu_r', show_torus: bool = True) → None

Render physical imported coordinates, with an analytic fallback.

gkx.artifacts.run_summary.draw_metadata_panel(ax: Axes, lines: Sequence[tuple[str, str]], *, heading: str = 'run summary') → None

Render the text panel of the summary figure onto ax.

gkx.artifacts.run_summary.draw_phi_xy_cut(ax: Axes, phi_xy: ndarray, *, scale: float, extent: tuple[float, float] | None = None, phi_label: str = '$(e\\phi/T_i)\\,/\\,\\rho_*$', cmap: str = 'RdBu_r', colorbar: bool = True)

Render one perpendicular phi(x, y) cut onto an existing axes.

phi_xy is indexed (x, y). With extent = (Lx, Ly) the axes carry physical rho_i units; without it they fall back to grid indices and are labelled as such rather than pretending to a unit they do not have.

gkx.artifacts.run_summary.final_field_path(source: str | Path) → Path

Return the *.big.nc companion of any file in an output bundle.

gkx.artifacts.run_summary.flux_tube_3d_figure(phi_spectral: ndarray, geom: Any, *, grid: Any | None = None, scale: float | None = None, time: float | None = None, ny_full: int | None = None, turns: float = 1.5, elev: float = 32.0, azim: float = -60.0, phi_label: str = '$(e\\phi/T_i)\\,/\\,\\rho_*$', title: str = 'Flux tube along $\\mathbf{B}$', out: str | Path | None = None) → Tuple[Figure, Any]

Render the field-aligned potential, using host-only physical coordinates.

gkx.artifacts.run_summary.flux_tube_figure(source: str | Path, *, title: str | None = None, out: str | Path | None = None) → Tuple[Figure, Any]

3-D rendering of the final potential on the flux tube the run integrated.

This is the figure that shows what the geometry actually was, which no other output of a run does: everything else is a reduction over the field-aligned coordinate.

gkx.artifacts.run_summary.has_final_field(source: str | Path) → bool

Return whether the bundle carries the final fields the snapshots need.

gkx.artifacts.run_summary.label_amplitude_colorbar(bar: Any, scale: float, phi_label: str) → None

Put scientific notation in the label, away from the figure title.

gkx.artifacts.run_summary.load_final_field(source: str | Path) → FinalField

Read the final potential, box size, and geometry from *.big.nc.

gkx.artifacts.run_summary.nonlinear_summary_figure(source: str | Path, *, window: tuple[float, float] | None = None, saturated: bool | None = None, title: str | None = None, out: str | Path | None = None) → Tuple[Figure, dict[str, Axes]]

Compose the whole-run summary page for one nonlinear output bundle.

source names any file of the bundle. window is the averaging window the stop policy evaluated. A rejected window remains in the metadata but is not shaded or reported as an average on the time trace; spectra retain their labelled second-half diagnostic average.

gkx.artifacts.run_summary.phi_xy_figure(source: str | Path, *, title: str | None = None, out: str | Path | None = None) → Tuple[Figure, Axes]

Full-size outboard-midplane cut of the final potential from a bundle.

gkx.artifacts.run_summary.phi_xy_snapshot_figure(phi_spectral: ndarray, grid: Any | None = None, geom: Any | None = None, *, z_index: int | None = None, scale: float | None = None, time: float | None = None, ny_full: int | None = None, extent: tuple[float, float] | None = None, phi_label: str = '$(e\\phi/T_i)\\,/\\,\\rho_*$', title: str = 'Outboard-midplane cut of $\\phi$', out: str | Path | None = None) → Tuple[Figure, Axes]

Publication x-y cut of the potential at the outboard midplane.

phi_spectral is either the spectral field (ky, kx, z) (complex) or an already-transformed real-space field (x, y, z). grid supplies the perpendicular box size for physical rho_i axes; geom is accepted for API symmetry with flux_tube_3d_figure() and reserved for annotations. out optionally saves the figure via gkx.artifacts.figure_style.save_figure() (the figure stays open).

gkx.artifacts.run_summary.potential_real_space(phi_spectral: ndarray, *, ny_full: int | None = None) → ndarray

Transform a spectral potential phi(ky, kx, z) to real phi(x, y, z).

ky uses the nonlinear bracket’s compressed real-FFT layout; kx is a full complex axis. Accepts either the full Hermitian ky array or its non-negative block; pass ny_full with the compressed form.

gkx.artifacts.run_summary.run_label(source: str | Path) → str

Return a title for one output bundle that identifies the run.

The grouped run directory the equilibrium shorthand writes names every bundle inside it gkx, so the file’s own stem says nothing. Prefixing the directory turns the title back into the identity of the run.

gkx.artifacts.run_summary.summary_metadata_lines(source: str | Path, diag: Any, *, window: tuple[float, float] | None, measured_window: bool, saturated: bool | None = None) → list[tuple[str, str]]

Assemble the (label, value) rows of the metadata panel.

Every row is optional: the deck, the NetCDF root, and the diagnostics are read for what they happen to carry, because this figure has to render for a bundle produced by any of the output paths, not only the equilibrium shorthand that writes all three.

Foreign Output Dispatch

Dispatch to readers for other gyrokinetic codes’ output bundles.

gkx --plot accepts a bundle written by another code so that a cross-code comparison is one command rather than a bespoke script. The readers themselves live in their own modules, one per code; this indirection keeps their names out of the plotting path, so the core figure code never learns which codes exist.

gkx.artifacts.foreign_output.foreign_output_plotter(path: str | Path) → Callable[[...], Any] | None

Return a plotter for path when another code (GX) wrote it, else None.

gkx.artifacts.foreign_output.gx_summary_figure(path: str | Path, *, title: str | None = None, out: str | Path | None = None) → tuple[Any, Any]

Plot whatever a GX bundle carries: time traces plus spectra when present.

Panels are chosen by what the file actually holds, so a GX linear run yields the Phi^2(t) trace alone while a nonlinear run adds Q(t), Gamma(t), and the time-averaged ky spectrum over the second half of the run. The title names GX explicitly.

gkx.artifacts.foreign_output.is_gx_output(path: str | Path) → bool

Return whether path is a GX output bundle rather than a GKX one.

A file GKX wrote always says so on code_info; anything else that carries GX’s title or GX’s build provenance is treated as GX. Anything unreadable is not claimed either way.

gkx.artifacts.foreign_output.plot_gx_output(path: str | Path, *, out: str | Path | None = None) → Path

Render a GX bundle to out (default <stem>.gx_plot.png).

gkx.artifacts.foreign_output.read_gx_output(path: str | Path) → dict[str, ndarray]

Return every plottable series a GX bundle carries, skipping absentees.

Keys mirror the NetCDF variable names, except that the per-species *_st fluxes also gain summed heat_flux/particle_flux totals. A key is present only when the file carried it.

Persistent Compilation Cache

Persistent JAX compilation cache for the GKX executable.

A GKX run is compile-dominated at the sizes people actually iterate on. On an office GPU a 100-step nonlinear case spends ~22.9 s of a ~25 s integrator wall inside XLA compilation, so the accelerator only wins end to end past roughly seventy steps; on a laptop CPU the same compile costs ~14 s. That cost is paid again on every fresh process even when nothing about the problem changed, which is the common case while a user edits a TOML and re-runs.

JAX can persist compiled executables across processes, and this module turns that on by default for the executable. Two environment variables are the escape: GKX_JAX_CACHE=0 disables the cache entirely, and GKX_JAX_CACHE_DIR moves it. The default location mirrors how the repository already caches generated *.eik.nc geometry – a .cache/gkx directory beside the source tree – and falls back to ~/.cache/gkx/jax when that tree is read-only, which is what an installed wheel sees.

On correctness: the cache key is XLA’s, not ours. It is computed from the lowered HLO module, the compilation options, and the backend/device the executable was built for, so a changed grid, changed physics, or a different device is a different key rather than a stale hit. The one thing that key does not span is the toolchain itself, so entries are additionally namespaced by the installed jax version: upgrading JAX starts a new directory instead of reading executables built by a different compiler. jax_raise_persistent_cache_errors is left at its default of False so a corrupt or unwritable entry degrades to a recompile rather than failing a simulation.

The directory only grows, at roughly 1.6 MB per distinct problem shape. JAX 0.9 exposes no writable size cap, so clearing it is rm -rf on the directory this module reports; nothing is lost but the next cold compile.

gkx.compilation_cache.compilation_cache_directory(env: Mapping[str, str] | None = None) → Path

Return the directory that should hold persisted executables.

An explicit GKX_JAX_CACHE_DIR wins and is used verbatim. Otherwise the repository-local .cache/gkx/jax is preferred so the cache is discarded with the checkout, with ~/.cache/gkx/jax as the fallback for installs whose package tree is not writable.

gkx.compilation_cache.compilation_cache_enabled(env: Mapping[str, str] | None = None) → bool

Return whether the persistent compilation cache should be installed.

gkx.compilation_cache.enable_persistent_compilation_cache(env: Mapping[str, str] | None = None) → Path | None

Point JAX at a persistent compilation cache; return the directory used.

Returns None when the cache is disabled by GKX_JAX_CACHE, when JAX is not importable, or when the directory cannot be prepared. Never raises: a cache is an optimization, and losing it must not lose a run.

Config

class gkx.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, ky_layout='half'), time: TimeConfig = TimeConfig(t_max=150.0, dt=0.01, method='rk4', fixed_dt=False, sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, state_sharding=None, progress_bar=False, dt_min=1e-07, dt_max=0.05, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', collision_operator='landau', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid', run_to='saturation', saturation_rel_sem=0.05, saturation_min_window=None, damp_ends_rate=None), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, geometry_helper_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, geometry_helper_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(tprim_i=2.49, tprim_e=0.0, fprim=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), reference_aligned: bool = True)

Standard parameters for the Cyclone base case ITG benchmark.

class gkx.config.GeometryConfig(model: str = 's-alpha', geometry_backend: str = 'auto', geometry_file: str | None = None, vmec_file: str | None = None, geometry_helper_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, geometry_helper_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 gkx.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, ky_layout: str = 'half')

Spectral grid configuration in a flux-tube.

ky_layout: str = 'half'

Layout of the ky axis of every spectral array built from this grid (gkx.core_ky_layout, plan 5.3 N3). "half" – the default – stores the Nyc = 1 + Ny // 2 non-negative rows and lets the reality condition hold by construction; "full" stores the two-sided fftfreq axis of length Ny and rebuilds the negative rows in the bracket and after every stage. Ny means the same thing in both: the length of the physical y axis, and therefore the resolution of the run. The physics, the resolution and the published output do not depend on the choice: nonlinear diagnostics agree to round-off (2e-16 in x64, 2e-7 in float32), the window gradient is bitwise equal, and "half" is faster (an RK3 step 0.54-0.60x of the two-sided time on CPU, 0.65x on an RTX A4000; the checkpointed window gradient 0.68-0.72x) with 12-25% less memory.

class gkx.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 gkx.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, ky_layout='half'), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', fixed_dt=True, sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, state_sharding=None, progress_bar=False, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', collision_operator='landau', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid', run_to='saturation', saturation_rel_sem=0.05, saturation_min_window=None, damp_ends_rate=None), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, geometry_helper_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, geometry_helper_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(tprim_i=2.49, tprim_e=2.49, fprim=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 gkx.config.KineticElectronModelConfig(tprim_i: float = 2.49, tprim_e: float = 2.49, fprim: 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.

Gradients are \(a/L_T\) and \(a/L_n\), as in ModelConfig.

class gkx.config.ModelConfig(tprim_i: float = 2.49, tprim_e: float = 0.0, fprim: float = 0.8, nu_i: float = 0.0)

Dimensionless gradients for the Cyclone base case.

These are \(a/L_T\) and \(a/L_n\) – the same quantities the TOML tprim/fprim keys carry and the only ones the linear operator consumes. With the Cyclone \(R_0 = R/a = 2.77778\), tprim_i = 2.49 is the literature \(R/L_T = 6.92\).

class gkx.config.TimeConfig(t_max: float = 100.0, dt: float = None, method: str = None, fixed_dt: bool = None, 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, state_sharding: str | None = None, progress_bar: bool = False, 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', collision_operator: str = 'landau', compressed_real_fft: bool = True, nonlinear_dealias: bool = True, laguerre_nonlinear_mode: str = 'grid', run_to: str = 'saturation', saturation_rel_sem: float = 0.05, saturation_min_window: float | None = None, damp_ends_rate: float | None = None)

Time integration parameters.

One policy for the API and for a deck. A caller that chooses dt gets rk2 at that fixed step: the cheaper scheme, since a fixed step gives rk4 no step-size compensation for its four stages. A caller that chooses no dt gets the step from the CFL controller and rk4 with it, which reaches a given horizon in 29.1% fewer right-hand-side evaluations than rk2 because its CFL prefactor is 2.82 against 1.0 (gkx.workflows.runtime.toml._normalize_time_overrides carries the measurement). dt, method and fixed_dt default to None and are resolved in __post_init__; an explicit value always wins. Before this, a bare TimeConfig() was rk2 at a fixed dt = 0.1, which overflows on most decks.

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

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

gkx.config.method_error(method: object, allowed) → ValueError

Error for an unsupported method; retired names point to their replacement.

gkx.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 gkx.diagnostics.normalization.NormalizationContract(case: str, rho_star: float, omega_d_scale: float, omega_star_scale: float, diagnostic_norm_default: Literal['none', '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’, ‘rho_star’]

gkx.diagnostics.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.

gkx.diagnostics.normalization.get_normalization_contract(case: str) → NormalizationContract

Return the canonical normalization contract for case.

Moment And Energy Diagnostics

Gyrokinetic quadrature, energy, transport-channel, and spectral diagnostics.

gkx.operators.moments._cached_hermitian_mode_weight(cache: LinearCache, *, use_dealias: bool) → Array

Return Hermitian spectral weight from a linear cache.

gkx.operators.moments._fold_ky_into_kx(kykx: Array, fold: _KyPairFold | None, *, ky_axis: int) → Array

Sum a (..., ky, kx, ...) view over ky as the two-sided axis would.

See _KyPairFold. With fold=None this is jnp.sum over ky and nothing else, so the two-sided path keeps its exact arithmetic and its exact operation count.

gkx.operators.moments._hermitian_mode_weight(grid: SpectralGrid, *, use_dealias: bool) → Array

Return Hermitian spectral weight for (ky, kx) weighting.

gkx.operators.moments._jl_family(cache: LinearCache) → tuple[Array, Array, Array]

Return (Jl, JlB, Jfac) arrays in Laguerre-Hermite conventions.

gkx.operators.moments._ky_pair_fold(source: Any) → _KyPairFold | None

Return the kx fold for a half-spectrum source, or None.

None is the two-sided axis, where each row stands for itself and the plain sum is already the right one.

gkx.operators.moments._masked_abs2(value: Array, active: Array) → Array

Return abs(value)**2 after zeroing inactive spectral modes.

gkx.operators.moments._reduce_scalar_kykxz(contrib: Array, fold: _KyPairFold | None = None) → tuple[Array, Array, Array, Array, Array]

Reduce a (ky, kx, z) contribution to spectral diagnostics views.

The fold=None expressions are left exactly as they were, rather than routed through the folded path with an identity fold: summing (ky, z) in one reduction and summing z then ky are the same number in exact arithmetic and not always the same float, and the two-sided route is gated bitwise.

gkx.operators.moments._reduce_species_kykxz(contrib: Array, fold: _KyPairFold | None = None) → tuple[Array, Array, Array, Array, Array]

Reduce a (species, ky, kx, z) contribution to spectral diagnostics views.

gkx.operators.moments._transport_mode_weight(grid: SpectralGrid, *, use_dealias: bool) → Array

Return fac*mask that excludes ky=0 and uses positive-ky transport weighting.

gkx.operators.moments.distribution_free_energy(G: Array, grid: SpectralGrid, params: LinearParams, vol_fac: Array, *, use_dealias: bool = True) → Array

Distribution free-energy diagnostic (free energy in g).

Wg + Wphi is not a discretely conserved functional of the dissipation-free discretization: the curvature/grad-B measure weights drive bounded transient excursions of order 30-40 per cent with no secular production (a zero-drive configuration has exactly gamma = 0). Drift in Wg + Wphi is therefore not a solver-error meter.

gkx.operators.moments.distribution_free_energy_resolved(G: Array, grid: SpectralGrid, params: LinearParams, vol_fac: Array, *, use_dealias: bool = True) → tuple[Array, Array, Array, Array, Array, Array]

Return Resolved free-energy reductions from a single state.

gkx.operators.moments.electrostatic_field_energy(phi: Array, cache: LinearCache, params: LinearParams, vol_fac: Array, *, use_dealias: bool = True, wphi_scale: float = 1.0) → Array

Electrostatic field energy with exact Gamma_0 (the field solve truncates).

gkx.operators.moments.electrostatic_field_energy_resolved(phi: Array, cache: LinearCache, params: LinearParams, vol_fac: Array, *, use_dealias: bool = True, wphi_scale: float = 1.0) → tuple[Array, Array, Array, Array, Array]

Return Resolved electrostatic-energy reductions per species.

gkx.operators.moments.fieldline_quadrature_weights(geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, grid: SpectralGrid) → tuple[Array, Array]

Return (vol_fac, flux_fac) for field-line volume and flux-surface weighting.

gkx.operators.moments.heat_flux_channel_resolved_species(*args, use_dealias: bool = True, flux_scale: float = 1.0)

Return resolved ES, Apar, and Bpar heat-flux channels per species.

gkx.operators.moments.magnetic_vector_potential_energy(apar: Array, cache: LinearCache, vol_fac: Array, *, use_dealias: bool = True) → Array

Magnetic vector-potential field-energy diagnostic.

gkx.operators.moments.magnetic_vector_potential_energy_resolved(apar: Array, cache: LinearCache, vol_fac: Array, *, nspecies: int, use_dealias: bool = True) → tuple[Array, Array, Array, Array, Array]

Return Resolved magnetic-energy reductions per species.

gkx.operators.moments.particle_flux_channel_resolved_species(*args, use_dealias: bool = True, flux_scale: float = 1.0)

Return resolved ES, Apar, and Bpar particle-flux channels per species.

gkx.operators.moments.phi2_resolved(phi: Array, grid: SpectralGrid, vol_fac: Array, *, use_dealias: bool = True) → tuple[Array, Array, Array, Array, Array, Array, Array, Array]

Return Resolved Phi2 reductions from a single field state.

gkx.operators.moments.total_energy(Wg: Array | ndarray, Wphi: Array | ndarray, Wapar: Array | ndarray) → Array | ndarray

Return Wg + Wphi + Wapar.

This sum is not a discretely conserved invariant of the dissipation-free discretization: the curvature/grad-B measure weights produce bounded transient excursions of order 30-40 per cent with no secular production, and a zero-drive configuration has exactly gamma = 0. Do not read its drift as a solver-error meter.

gkx.operators.moments.turbulent_heating_resolved_species(G: Array, G_old: Array, phi: Array, apar: Array, bpar: Array, phi_old: Array, apar_old: Array, bpar_old: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, vol_fac: Array, dt: Array | float, *, use_dealias: bool = True) → tuple[Array, Array, Array, Array, Array]

Return Resolved turbulent-heating reductions per species.

gkx.operators.moments.zonal_phi_line_kxt(phi: Array, grid: SpectralGrid) → Array

Return the signed, unweighted line-averaged zonal potential per k_x.

gkx.operators.moments.zonal_phi_mode_kxt(phi: Array, grid: SpectralGrid, vol_fac: Array) → Array

Return the signed, volume-averaged zonal potential history per k_x.

This is the minimal complex zonal observable needed to construct Rosenbluth-Hinton / GAM response traces from nonlinear diagnostics without collapsing immediately to a positive-definite energy proxy.

Runtime Config

class gkx.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, ky_layout='half'), time: TimeConfig = TimeConfig(t_max=150.0, dt=0.01, method='rk4', fixed_dt=False, sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, state_sharding=None, progress_bar=False, dt_min=1e-07, dt_max=0.05, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', collision_operator='landau', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid', run_to='saturation', saturation_rel_sem=0.05, saturation_min_window=None, damp_ends_rate=None), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, geometry_helper_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, geometry_helper_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(tprim_i=2.49, tprim_e=0.0, fprim=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), reference_aligned: bool = True)

Standard parameters for the Cyclone base case ITG benchmark.

class gkx.config.GeometryConfig(model: str = 's-alpha', geometry_backend: str = 'auto', geometry_file: str | None = None, vmec_file: str | None = None, geometry_helper_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, geometry_helper_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 gkx.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, ky_layout: str = 'half')

Spectral grid configuration in a flux-tube.

ky_layout: str = 'half'

Layout of the ky axis of every spectral array built from this grid (gkx.core_ky_layout, plan 5.3 N3). "half" – the default – stores the Nyc = 1 + Ny // 2 non-negative rows and lets the reality condition hold by construction; "full" stores the two-sided fftfreq axis of length Ny and rebuilds the negative rows in the bracket and after every stage. Ny means the same thing in both: the length of the physical y axis, and therefore the resolution of the run. The physics, the resolution and the published output do not depend on the choice: nonlinear diagnostics agree to round-off (2e-16 in x64, 2e-7 in float32), the window gradient is bitwise equal, and "half" is faster (an RK3 step 0.54-0.60x of the two-sided time on CPU, 0.65x on an RTX A4000; the checkpointed window gradient 0.68-0.72x) with 12-25% less memory.

class gkx.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 gkx.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, ky_layout='half'), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', fixed_dt=True, sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, state_sharding=None, progress_bar=False, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', collision_operator='landau', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid', run_to='saturation', saturation_rel_sem=0.05, saturation_min_window=None, damp_ends_rate=None), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, geometry_helper_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, geometry_helper_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(tprim_i=2.49, tprim_e=2.49, fprim=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 gkx.config.KineticElectronModelConfig(tprim_i: float = 2.49, tprim_e: float = 2.49, fprim: 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.

Gradients are \(a/L_T\) and \(a/L_n\), as in ModelConfig.

class gkx.config.ModelConfig(tprim_i: float = 2.49, tprim_e: float = 0.0, fprim: float = 0.8, nu_i: float = 0.0)

Dimensionless gradients for the Cyclone base case.

These are \(a/L_T\) and \(a/L_n\) – the same quantities the TOML tprim/fprim keys carry and the only ones the linear operator consumes. With the Cyclone \(R_0 = R/a = 2.77778\), tprim_i = 2.49 is the literature \(R/L_T = 6.92\).

class gkx.config.TimeConfig(t_max: float = 100.0, dt: float = None, method: str = None, fixed_dt: bool = None, 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, state_sharding: str | None = None, progress_bar: bool = False, 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', collision_operator: str = 'landau', compressed_real_fft: bool = True, nonlinear_dealias: bool = True, laguerre_nonlinear_mode: str = 'grid', run_to: str = 'saturation', saturation_rel_sem: float = 0.05, saturation_min_window: float | None = None, damp_ends_rate: float | None = None)

Time integration parameters.

One policy for the API and for a deck. A caller that chooses dt gets rk2 at that fixed step: the cheaper scheme, since a fixed step gives rk4 no step-size compensation for its four stages. A caller that chooses no dt gets the step from the CFL controller and rk4 with it, which reaches a given horizon in 29.1% fewer right-hand-side evaluations than rk2 because its CFL prefactor is 2.82 against 1.0 (gkx.workflows.runtime.toml._normalize_time_overrides carries the measurement). dt, method and fixed_dt default to None and are resolved in __post_init__; an explicit value always wins. Before this, a bare TimeConfig() was rk2 at a fixed dt = 0.1, which overflows on most decks.

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

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

gkx.config.method_error(method: object, allowed) → ValueError

Error for an unsupported method; retired names point to their replacement.

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

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

Runtime Startup

Runtime startup and initialization helpers.

This module holds the geometry/loading/initial-condition logic used by the public runtime entry points. It is intentionally kept separate from the solver execution layer so startup behavior can be tested and refactored without touching the time-integration control flow.

gkx.workflows.runtime.startup._as_runtime_species_array(value: float | Array, nspecies: int, name: str) → ndarray

Return a length-nspecies NumPy array for startup-only algebra.

gkx.workflows.runtime.startup._build_initial_condition(grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, cfg: RuntimeConfig, *, ky_index: int, kx_index: int, Nl: int, Nm: int, nspecies: int) → Array

Build the runtime initial state for one selected mode or seed policy.

gkx.workflows.runtime.startup._build_single_phi_gaussian_profile(z: ndarray, *, kx: float, ky: float, s_hat: float, width: float, envelope_constant: float, envelope_sine: float) → ndarray

Return a single-mode Gaussian potential profile along the flux tube.

The W7-X zonal-flow benchmark prescribes a Gaussian electrostatic-potential perturbation centered in the middle of the flux tube. A multi-mode ballooning-angle Gaussian initializer is undefined for ky=0 because its center uses kx / (s_hat * ky); for the zonal case the physically stated center is therefore the tube midpoint, z=0.

gkx.workflows.runtime.startup._centered_glibc_random_pairs(seed: int, count: int) → ndarray

Return centered random pairs using glibc rand() semantics.

gkx.workflows.runtime.startup._dealiased_initial_mode_pairs(grid: SpectralGrid) → list[tuple[int, int]]

Return the dealiased startup-loop (kx, ky) pairs for multimode initial conditions.

The binormal index 0 is skipped only when it really is the zonal mode. A linear run selects one k_y before seeding, so its grid holds a single nonzero binormal entry at index 0; skipping it by position left the whole state at zero and the eigensolver was handed a null seed.

gkx.workflows.runtime.startup._density_moments_for_target_phi(phi_target: ndarray, *, cache, params: LinearParams, ky_i: int, kx_i: int, species_targets: tuple[int, ...]) → dict[int, ndarray]

Invert the electrostatic field solve for a requested initial phi mode.

The runtime evolves Hermite-Laguerre moments, not fields. For literature tests that prescribe an initial electrostatic-potential perturbation, seed only the density moment with the algebraic moment profile whose immediate quasineutrality solve returns phi_target. The adiabatic-electron zonal branch includes the same flux-surface-average correction used in gkx.terms.fields.

gkx.workflows.runtime.startup._enforce_full_ky_hermitian(arr: ndarray) → ndarray

Mirror positive-ky content into the negative branch for full FFT grids.

gkx.workflows.runtime.startup._expand_ky(arr: ndarray, *, ny_full: int) → ndarray

Widen a stored ky >= 0 init state onto this grid’s two-sided axis.

ny_full is passed in rather than inferred from Nyc: the two map many-to-one (gkx.core_ky_layout), and inferring the even branch silently produced an Ny - 1 row state for an odd-Ny grid.

gkx.workflows.runtime.startup._load_initial_state_from_file(path: Path, *, nspecies: int, Nl: int, Nm: int, ny: int, nx: int, nz: int, ky_layout: Literal['full', 'half'] = 'full') → ndarray

Load a restart or seed file onto ky_layout.

ny is the two-sided axis length in both layouts, which is what the file’s own contents are indexed by; ky_layout decides how many rows come back. Both file forms keep working on both axes: a binary file of Nyc rows is the ky >= 0 block and is widened only when the caller asked for the two-sided axis, and a two-sided file is narrowed only when the caller asked for the half one. Neither conversion loses anything – the discarded rows are the conjugates of the kept ones (plan 5.3 N3).

gkx.workflows.runtime.startup._periodic_zp_from_grid(z: ndarray) → float

Return periodic Zp from the discrete theta grid.

gkx.workflows.runtime.startup.build_runtime_geometry(cfg: RuntimeConfig) → SAlphaGeometry | SlabGeometry | FluxTubeGeometryData

Resolve runtime geometry, generating *.eik.nc geometry when needed.

gkx.workflows.runtime.startup.build_runtime_linear_params(cfg: RuntimeConfig, *, Nm: int | None = None, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData | None = None) → LinearParams

Build LinearParams from a unified runtime config.

gkx.workflows.runtime.startup.build_runtime_linear_terms(cfg: RuntimeConfig, *, custom_collisions: bool = False) → LinearTerms

Build LinearTerms from unified toggles.

All-zero species nu drop the built-in collision term; a caller passing a custom operator with its own rates sets custom_collisions (#324).

gkx.workflows.runtime.startup.build_runtime_term_config(cfg: RuntimeConfig) → TermConfig

Build nonlinear-ready TermConfig from unified toggles.

gkx.workflows.runtime.startup.end_damping_status(params: LinearParams, grid: Any, boundary: str) → str

Describe the end-damping rate a run uses, for the run’s status line.

gkx.workflows.runtime.startup.load_netcdf_restart_state(path: str | Path, *, nspecies: int, Nl: int, Nm: int, ny: int, nx: int, nz: int, ky_layout: Literal['full', 'half'] = 'full') → ndarray

Load a NetCDF restart file onto ky_layout.

ny is the length of the two-sided axis in both layouts; ky_layout says how many rows the caller wants back, Ny or Nyc = 1 + Ny // 2. The file itself is unchanged by the layout: it has always stored the dealiased ky >= 0 block, which is why a run started on either axis can continue from a file written by the other, and why a 2.1.0 restart still loads (plan 5.3 N3, queue row Q10).

gkx.workflows.runtime.startup.runtime_geometry_config_for_builder(cfg: RuntimeConfig, *, vmec_eik_builder: Callable[[RuntimeConfig], Any], miller_eik_builder: Callable[[RuntimeConfig], Any]) → Any

Return the geometry config consumed by the flux-tube builder.

gkx.workflows.runtime.startup.runtime_state_dtype() → dtype

Return the complex dtype the runtime’s own initial state is handed in.

Every seed in this module is assembled in complex64: the seed is an arbitrary small perturbation, so its own rounding is immaterial, and holding the assembly at one width keeps the ky-Hermitian and off-chain masking code free of a dtype branch. What is not immaterial is the width the seed imposes downstream. jnp.asarray does not promote, so a complex64 seed pinned the whole run to complex64 even under JAX_ENABLE_X64: the Krylov basis, the eigenvector and therefore the certification gate, which certifiable_residual_tolerance floors at 1e3 * eps(dtype) – 1.19e-4 in float32 against the 1e-9 the float64 gate applies. Decks that instruct the reader to run under JAX_ENABLE_X64 (the shipped Cyclone deck is one) silently got the float32 answer and a float32 gate.

The seed is therefore widened once, here, to the working precision. This does not change the default: without JAX_ENABLE_X64 the dtype is complex64 exactly as before. It matches what the initial_state API already promises for a caller-supplied seed, which preserves an explicit complex128 array including any prior initialization rounding.

Runtime Diagnostics

Runtime linear-fit and quasilinear diagnostic helpers.

class gkx.workflows.runtime.diagnostics.RuntimeLinearFitResult(gamma: float, omega: float, signal: ndarray, z: ndarray, eigenfunction: ndarray | None, fit_window_tmin: float | None, fit_window_tmax: float | None, fit_signal_used: str, gamma_stderr: float | None = None, omega_stderr: float | None = None, fit_r2: float | None = None, fit_settled: bool | None = None)

Linear runtime fit payload before diagnostic normalization.

gkx.workflows.runtime.diagnostics.ensure_finite_linear_history(name: str, values: ndarray | None) → None

Reject non-finite or overflow-scale histories before any growth fit.

gkx.workflows.runtime.diagnostics.finalize_runtime_linear_quasilinear(result: RuntimeLinearResult, *, enabled: bool, cfg: Any, grid: Any, geom: Any, params: Any, terms: Any, Nl: int, Nm: int, solver_name: str, species_names: tuple[str, ...], return_state_requested: bool, state_for_quasilinear: ndarray | None = None, status_callback: Any | None = None) → RuntimeLinearResult

Attach optional quasilinear diagnostics to a linear runtime result.

gkx.workflows.runtime.diagnostics.fit_runtime_linear_diagnostics(*, t: ndarray, phi_t: ndarray, density_t: ndarray | None, selection: Any, z: ndarray, fit_signal: str, mode_method: str, 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, window_method: str = 'stationary') → RuntimeLinearFitResult

Fit linear growth/frequency and extract the eigenfunction diagnostic.

gkx.workflows.runtime.diagnostics.half_horizon_settled_probe(t: ndarray, signal: ndarray, *, gamma: float, tmin: float | None, tmax: float | None, rel_tolerance: float = 0.05) → tuple[float | None, bool | None]

Re-fit the second half of the fit window and compare with the full fit.

A converged linear run has reached its dominant eigenmode, so halving the window must not move gamma. Returns (relative_drift, settled), or (None, None) when the half window has too few points to fit. A drift above rel_tolerance also emits a RuntimeWarning: the reported gamma is then still tracking a transient rather than an eigenvalue.

gkx.workflows.runtime.diagnostics.refit_runtime_linear_trajectory(result: RuntimeLinearResult, *, mode_method: str = 'project', auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.3, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 1.0, require_positive: bool = True, min_amp_fraction: float = 0.0, window_method: str = 'stationary') → RuntimeLinearResult

Refit one stored trajectory without repeating its integration.

gkx.workflows.runtime.diagnostics.warn_if_growth_unresolved(*, gamma: float, t: ndarray, fit_window_tmin: float | None, fit_window_tmax: float | None, min_total_growth_times: float = 7.0, min_window_growth_times: float = 2.0) → str | None

Warn when a linear run is too short to resolve its fitted growth rate.

The controlling parameter is the e-folding count N = gamma * t_max. Fixed-horizon linear scans under-report growth one-sidedly downward when N is small – measured at roughly 9-14% low for N ~ 2-3, 1-6% for N ~ 5-6, and under 0.5% for N >= 7 – so min_total_growth_times defaults to 7. The message reports the achieved e-folding count and the horizon that would reach the threshold.

Returns the emitted message, or None when the run is resolved. Only positive growth rates are checked: a decaying fit has no growth to resolve.

Runtime Diagnostic Arrays

Runtime diagnostic array validation, composition, and cost summaries.

These helpers operate on already-sampled diagnostic payloads. They deliberately stay host-side: runtime drivers use them to fail fast on invalid artifacts, to combine adaptive chunks without mixing that array bookkeeping into the linear-fit and quasilinear-finalization code, and to summarise what a run’s simulated time cost before the payload is handed to the artifact writers.

class gkx.workflows.runtime.diagnostic_arrays.TimestepCostReport(n_samples: int, t_start: float, t_end: float, t_span: float, steps: float, steps_are_exact: bool, steps_per_unit_time: float, wall_seconds: float | None, wall_seconds_per_unit_time: float | None, dt_initial: float, dt_min: float, dt_max: float, dt_final: float, dt_collapse_ratio: float, steps_per_unit_time_first_quarter: float, steps_per_unit_time_last_quarter: float, cost_growth_ratio: float, cfl: CFLLimiterReport | None, notes: tuple[str, ...])

Cost per unit of simulated time, and the dt trajectory behind it.

gkx.workflows.runtime.diagnostic_arrays.concat_runtime_diagnostics(diags: Sequence[SimulationDiagnostics]) → SimulationDiagnostics

Concatenate one or more diagnostic chunks.

gkx.workflows.runtime.diagnostic_arrays.slice_runtime_diagnostics(diag: SimulationDiagnostics, stop: int) → SimulationDiagnostics

Return the first stop diagnostic samples.

gkx.workflows.runtime.diagnostic_arrays.stride_runtime_diagnostics(diag: SimulationDiagnostics, *, stride: int, offset: int = 0, keep_last: bool = False) → SimulationDiagnostics

Apply the runtime output stride to chunk diagnostics.

offset selects the first sample kept, which is what lets the stride be applied per chunk instead of after concatenation. Striding a chunk that starts at global sample g with offset = (-g) % stride keeps exactly the global indices 0, stride, 2*stride, ... – the same samples the post-concatenation stride would have kept, at a fraction of the peak memory, because the discarded samples are never accumulated in the first place.

gkx.workflows.runtime.diagnostic_arrays.timestep_cost_payload(diag: SimulationDiagnostics, *, wall_seconds: float | None = None) → dict[str, Any]

Return the JSON-ready cost block for one run’s recorded diagnostics.

Returns an empty payload rather than raising when the series is unusable, so a summary artifact never fails to write because of a diagnostic field.

gkx.workflows.runtime.diagnostic_arrays.timestep_cost_report(t: ndarray, dt: ndarray, *, wall_seconds: float | None = None, scales: CFLScales | None = None) → TimestepCostReport

Summarise what one unit of simulated time cost, and why.

Nothing here gates. Within a healthy nonlinear run dt collapses by construction as the initial condition saturates, so a collapse-ratio limit would fire on every well-behaved surface; the ratio is reported and left to the reader. The one note emitted is objective rather than tuned: a step pinned at dt_min means the requested CFL condition was not honoured.

gkx.workflows.runtime.diagnostic_arrays.validate_finite_runtime_diagnostics(diag: SimulationDiagnostics, *, label: str = 'runtime') → None

Raise if a runtime diagnostic chunk contains NaN or infinite values.

Long validation runs can otherwise continue for thousands of fixed steps after the first unstable sample. This host-side guard keeps the expensive artifact path fail-fast and reports the first offending diagnostic channel.

Runtime Resolution Estimator

Zero-config gkx wout_XXX.nc equilibrium shorthand helpers.

class gkx.workflows.runtime.wout.GeometryFeatures(anisotropy: float, shat: float, q: float, nfp: int, bmag_wells: int, zp: float)

Host-side geometry scalars the estimator maps to a grid.

gkx.workflows.runtime.wout.KY_TARGETS_BY_CLASS = {'stellarator': {'cautious': 4.4, 'preview': 2.2, 'standard': 2.9}, 'tokamak': {'cautious': 2.9, 'preview': 1.5, 'standard': 2.2}}

ky_max*rho per (class, tier). At the default dky*rho = 0.071 these land on rungs 64/96/128 (tokamak) and 96/128/192 (stellarator). Measured y0=14 ladder bias per tier – tokamak: preview ~+8% above the converged 96/128 plateau, standard converged within SEM, cautious headroom; stellarator: preview ~+15-25% and standard ~+8-19% (both still upper estimates, flux falling at 128^2), cautious = published W7-X reach ~4.4 (unvalidated by this scan until the 192^2 references land).

gkx.workflows.runtime.wout.LINEAR_SCAN_DT = 0.01

Fixed step for the shorthand linear scan. Half the measured 0.019 bound at the top of DEFAULT_LINEAR_KY_VALUES, so every rung integrates.

gkx.workflows.runtime.wout.PERP_LADDER = (32, 64, 96, 128, 192, 256)

Perpendicular grid rungs, extending the scan ladder upward.

gkx.workflows.runtime.wout._apply_class_resolution_defaults(data: dict[str, Any], wout_path: Path) → None

Shipped-deck preview grid per equilibrium class (2026-08 y0=14 ladder).

Tokamaks (nfp = 1) saturated at every rung with 64^2 only ~8% above the converged 96/128 plateau, so their preview drops to 64^2. Stellarators keep the deck’s 96^2, which the ladder measured as an upper estimate (flux still falling at 128^2). Never applied to a user-supplied deck.

gkx.workflows.runtime.wout._apply_linear_scan_defaults(data: dict[str, Any]) → None

Switch a deck to linear physics with a default ky-scan list.

The step is reduced along with the physics. A nonlinear deck is written around its own low-k_y box, while this scan reaches k_y rho = 1, where the explicit bound is far tighter: on the shipped stellarator deck the CFL-stable step falls from 0.067 at the first finite k_y to 0.019 at the top of the scan, so the nonlinear deck’s 0.1 overflows every rung and the growth fit then refuses a non-finite history. The linear paths advance the whole RHS explicitly at a fixed step, so adaptivity in the deck does not rescue them.

gkx.workflows.runtime.wout._count_bmag_wells(bmag: ndarray) → int

Count |B| wells deeper than _WELL_PROMINENCE of the |B| range.

gkx.workflows.runtime.wout._dt_hint(cfg: RuntimeConfig, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, est: dict[str, Any]) → float

Initial-dt bound from the solver’s own linear CFL frequency estimator.

gkx.workflows.runtime.wout._flag_value(args: list[str], flag: str) → str | None

Return the value of flag inside pass-through parser arguments.

gkx.workflows.runtime.wout._force_vmec_geometry(data: dict[str, Any], wout_path: Path) → None

Point [geometry] at the wout file, forcing model=”vmec” when needed.

gkx.workflows.runtime.wout._pop_estimate_flag(args: list[str]) → str | None

Pop --estimate[=TARGET], returning the target-error tier if present.

gkx.workflows.runtime.wout._print_deck_header(*, wout_path: Path, deck_path: Path, resolved_path: Path, shipped: bool) → None

Name the deck the defaults came from before the run starts.

The resolved copy alone does not tell a first-time user what to edit: it is generated, it is inside the output directory, and it is overwritten by the next run. Naming the shipped deck – and the command that runs an edited copy of it – is what makes the next run theirs.

gkx.workflows.runtime.wout._print_resolution_estimate(wout_path: Path, config_arg: str | None, *, target_error: str) → None

Print the minimum-grid estimate table for one equilibrium.

gkx.workflows.runtime.wout._resolve_deck_paths(data: dict[str, Any], *, base_dir: Path) → None

Resolve path-valued deck fields so the resolved deck can relocate.

gkx.workflows.runtime.wout._resolved_output_prefix(data: dict[str, Any], extra: list[str], wout_path: Path) → tuple[Path, bool]

Return the output prefix for wout-run artifacts, and whether it was asked for.

The second element says the target came from --out or from the deck’s own [output] path; only the prefix this function invents is free to grow a suffix in _resolved_output_target().

gkx.workflows.runtime.wout._resolved_output_target(prefix: Path, *, explicit: bool, linear: bool) → Path

Return the [output] path a bare equilibrium run writes to.

A plain prefix makes the runtime write CSV/JSON sidecars, which carry time traces only – so the spectra, the potential map, and the restart file all silently do not exist, and the figure set shrinks to one panel. The default is therefore the NetCDF bundle, which is the format the rest of the result set is read from. A linear ky scan has no NetCDF form, and a target the user named is theirs; both keep the prefix as given.

gkx.workflows.runtime.wout._wout_nfp(path: Path) → int | None

Field periods from the wout header, or None when unreadable.

gkx.workflows.runtime.wout.default_wout_deck_path() → Path

Return the single-source default deck used for bare equilibrium runs.

gkx.workflows.runtime.wout.direct_config_shorthand_args(argv: ~typing.Sequence[str], *, load_toml_func: ~typing.Callable[[str | ~pathlib.Path], dict[str, ~typing.Any]] = <function load_toml>) → list[str] | None

Return parser arguments for gkx case.toml / gkx wout_XXX.nc shorthand.

Leading positionals may be a runtime TOML deck and/or a VMEC/VMEX wout equilibrium (in either order); --vmec FILE/--vmex FILE are explicit aliases for the wout positional. A wout argument routes through the equilibrium shorthand, which writes a fully-resolved deck next to the grouped outputs before dispatch.

gkx.workflows.runtime.wout.estimate_resolution(wout_path: str | Path, *, torflux: float | None = None, target_error: str = 'standard', deck_path: str | Path | None = None) → dict[str, Any]

Estimate the minimum adequate grid for one VMEC/VMEX equilibrium.

gkx.workflows.runtime.wout.extract_wout_flag_value(args: list[str]) → str | None

Pop --vmec/--vmex alias flags from args, returning the file.

gkx.workflows.runtime.wout.format_estimate_table(est: dict[str, Any]) → str

Render one resolution estimate as the table --estimate prints.

gkx.workflows.runtime.wout.geometry_class(features: GeometryFeatures) → str

Classify the equilibrium; nfp decides, anisotropy corroborates.

gkx.workflows.runtime.wout.geometry_features(geom: Any, *, zp: float = 1.0) → GeometryFeatures

Extract estimator features from sampled flux-tube geometry data.

gkx.workflows.runtime.wout.is_wout_file(path: str | Path) → bool

Return whether path is a NetCDF VMEC/VMEX wout equilibrium.

gkx.workflows.runtime.wout.ky_max_target(features: GeometryFeatures, target_error: str = 'standard') → float

Required binormal reach ky_max*rho for this class and error tier.

gkx.workflows.runtime.wout.perp_points_for(dky: float, ky_target: float) → int

Smallest ladder rung whose dealiased ky grid reaches ky_target.

gkx.workflows.runtime.wout.resolution_from_features(features: GeometryFeatures, *, dky: float, target_error: str = 'standard', nz_default: int = 48, hypercollisions: bool = True, kinetic_electrons: bool = False) → dict[str, Any]

Map geometry features to grid hints with a rationale per number.

gkx.workflows.runtime.wout.resolved_deck_text(data: dict[str, Any], *, wout_path: Path) → str

Render resolved deck data as TOML, mirroring the demo reproducer.

gkx.workflows.runtime.wout.wout_shorthand_args(wout_arg: str, config_arg: str | None, extra_args: list[str], *, load_toml_func: ~typing.Callable[[str | ~pathlib.Path], dict[str, ~typing.Any]] = <function load_toml>) → list[str]

Return parser args for equilibrium shorthand, writing the resolved deck.

Runtime Chunks

Bounded-memory adaptive runtime chunks.

class gkx.workflows.runtime.chunks.AdaptiveChunkResult(diagnostics: SimulationDiagnostics, state: Any, fields: FieldState, stop_decision: dict[str, Any] | None = None)

Concatenated result from one adaptive runtime chunk loop.

class gkx.workflows.runtime.chunks.RuntimeProgressSnapshot(progress: float, eta_seconds: float, chunk_wall_seconds: float, elapsed_seconds: float)

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

gkx.workflows.runtime.chunks._chunk_diagnostics_to_host(diag: SimulationDiagnostics) → SimulationDiagnostics

Move a completed chunk to host so device residency stays bounded.

gkx.workflows.runtime.chunks._chunk_end_time(diag: SimulationDiagnostics, *, label: str, chunk_index: int) → float

Return the last diagnostic time sample for one adaptive chunk.

gkx.workflows.runtime.chunks._effective_diagnostics_stride(diagnostics_stride: int) → int

Normalize runtime diagnostic stride with a floor at one.

gkx.workflows.runtime.chunks._load_spilled_chunk(path: Path) → SimulationDiagnostics

Read back one spilled chunk, restoring the resolved payload.

gkx.workflows.runtime.chunks._next_elapsed_time(diag: SimulationDiagnostics, *, previous_elapsed: float, label: str, chunk_index: int) → float

Validate and return the accumulated end time for one adaptive chunk.

gkx.workflows.runtime.chunks._offset_chunk_diagnostics_time(diag: SimulationDiagnostics, *, offset: float) → SimulationDiagnostics

Return a chunk diagnostic payload shifted onto the accumulated time axis.

gkx.workflows.runtime.chunks._spill_chunk(diag: SimulationDiagnostics, spill_dir: Path, chunk_index: int) → Path

Spill plain arrays so host memory also stays bounded.

gkx.workflows.runtime.chunks.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.

gkx.workflows.runtime.chunks.run_adaptive_runtime_chunk_loop(*, integrate_chunk: Callable[[bool, float], tuple[Any, SimulationDiagnostics, Any, FieldState | None]], t_max: float, chunk_steps: int, label: str, show_progress: bool = False, status_callback: Callable[[str], None] | None = None, diagnostics_stride: int = 1, max_chunks: int = 100000, spill_dir: Path | None = None, stop_condition: Callable[[ndarray, ndarray, ndarray, ndarray], dict[str, Any]] | None = None) → AdaptiveChunkResult

Run to a stop or t_max; check unstrided traces between chunks.

Runtime Results

Runtime result containers and small assembly helpers.

gkx.workflows.runtime.results.LinearResult

alias of RuntimeLinearResult

gkx.workflows.runtime.results.NonlinearResult

alias of RuntimeNonlinearResult

class gkx.workflows.runtime.results.RuntimeLinearResult(ky: float, gamma: float, omega: float, selection: ModeSelection, t: np.ndarray | None = None, signal: np.ndarray | None = None, field_history: np.ndarray | None = None, state: np.ndarray | None = None, z: np.ndarray | None = None, eigenfunction: np.ndarray | None = None, fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, fit_signal_used: str | None = None, gamma_stderr: float | None = None, omega_stderr: float | None = None, fit_r2: float | None = None, fit_settled: bool | None = None, quasilinear: dict[str, Any] | None = None, eigen_status: EigenSolveStatus | None = None, implicit_solve: ImplicitSolveSummary | None = None)

Result container for runtime linear runs.

static _artifact_writer(io_module: Any) → Any

Return the writer in gkx.artifacts.io that owns this result.

summary() → dict[str, Any]

Return the typed scalar diagnostics, including fit and solver status.

to_dataset() → dict[str, Any]

Return the trace and eigenfunction as an xarray-compatible mapping.

class gkx.workflows.runtime.results.RuntimeLinearScanResult(ky: ndarray, gamma: ndarray, omega: ndarray, quasilinear: tuple[dict[str, Any], ...] | None = None, parallel: dict[str, Any] | None = None, warm_start: dict[str, Any] | None = None)

Result container for runtime linear ky scans.

static _artifact_writer(io_module: Any) → Any

Return the writer in gkx.artifacts.io that owns this result.

summary() → dict[str, Any]

Return scan extent and the peak growth rate over the scanned k_y.

to_dataset() → dict[str, Any]

Return growth and frequency against k_y.

class gkx.workflows.runtime.results.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, wall_seconds: float | None = None, saturation: dict[str, Any] | None = None, implicit_solve: ImplicitSolveSummary | None = None)

Result container for runtime nonlinear runs.

static _artifact_writer(io_module: Any) → Any

Return the writer in gkx.artifacts.io that owns this result.

summary() → dict[str, Any]

Return scalar diagnostics and the saturation verdict.

saturated is reported as the stop policy recorded it, including False and None. A result must not present an unsaturated window as an accepted value, so the status travels with the number rather than being inferred by a reader.

to_dataset() → dict[str, Any]

Return the sampled time series as an xarray-compatible mapping.

class gkx.workflows.runtime.results.RuntimeParameterScanResult(parameter_name: str, values: ndarray, gamma: ndarray, omega: ndarray, runs: tuple[RuntimeLinearResult, ...])

Ordered linear results for a named scalar configuration parameter.

gkx.workflows.runtime.results.ScanResult

alias of RuntimeLinearScanResult

class gkx.workflows.runtime.results._ResultArtifacts

Shared save/plot/print_summary behaviour for public results.

Each result already had a writer and a figure builder; what it lacked was a way to reach them from the object itself. These delegate rather than reimplement, so the artifact bytes a result writes through save are the bytes the runtime already wrote.

static _artifact_writer(io_module: Any) → Any

Return the writer in gkx.artifacts.io that owns this result.

plot() → Any

Return the standard figure for this result.

print_summary(*, stream: Any = None) → None

Print the same scalar summary that save records.

save(path: str | Path) → dict[str, str]

Write this result’s artifact bundle and return the paths written.

summary() → dict[str, Any]

Return this result’s typed scalar diagnostics.

summary_text() → str

Return the one-line-per-field scalar summary as text.

gkx.workflows.runtime.results._dataset_payload(coords: dict[str, Any], data_vars: dict[str, Any], attrs: dict[str, Any]) → dict[str, Any]

Return an xarray-compatible mapping without importing xarray.

GKX does not depend on xarray, and the dependency review keeps it that way, so a result cannot hand back a real Dataset. It hands back exactly the three keyword arguments xarray.Dataset accepts, so a caller who has xarray writes xarray.Dataset(**result.to_dataset()) and a caller who does not still gets named arrays with their coordinates attached.

gkx.workflows.runtime.results.build_runtime_nonlinear_result(*, t: ndarray, diagnostics: SimulationDiagnostics | None, fields: FieldState | None, state: ndarray | None, ky_selected: float | None, kx_selected: float | None, summarize_fields: bool, saturation: dict[str, Any] | None = None) → RuntimeNonlinearResult

Build a runtime nonlinear result with optional final-field summary.

gkx.workflows.runtime.results.checked_solve_summary(outputs: Sequence[Any], request: dict[str, Any], *, label: str) → ImplicitSolveSummary | None

Fail closed on the stats a requested integrator appended last.

Returns None when no stats were requested, so explicit runs keep a null status rather than a fabricated converged one.

gkx.workflows.runtime.results.nonlinear_field_phi2(fields: FieldState) → ndarray

Return the mean electrostatic energy density from final fields.

gkx.workflows.runtime.results.plot(result: Any) → Any

Create the standard in-memory figure for a runtime result.

gkx.workflows.runtime.results.solve_stats_request(method: Any, *, kind: str) → dict[str, Any]

Return return_solve_stats=True when method takes implicit solves.

Runtime Orchestration

Runtime ky-scan orchestration and combined-batch scan policy.

class gkx.workflows.runtime.orchestration_scan.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.

gkx.workflows.runtime.orchestration_scan.run_runtime_parameter_scan(cfg: RuntimeConfig, parameter_values: Sequence[float], *, parameter_name: str, update_config: Callable[[RuntimeConfig, float, int], RuntimeConfig], ky_target: float = 0.3, linear_options: Mapping[str, Any] | None = None, point_options: Callable[[float, int, RuntimeLinearResult | None], Mapping[str, Any]] | None = None, candidate_options: Callable[[float, int, RuntimeLinearResult | None], Sequence[Mapping[str, Any]]] | None = None, select_candidate: Callable[[float, int, tuple[RuntimeLinearResult, ...], RuntimeLinearResult | None], int] | None = None, continuation: bool = False, warm_start: bool = False, warm_start_max_step: float = 0.25) → RuntimeParameterScanResult

Run a scalar scan and optionally continue a selected solution branch.

continuation follows one branch and requires every point to hand back a state. warm_start is the weaker, opportunistic form: the previous point’s state is reused only while the parameter moves by no more than warm_start_max_step relative to the point before it, and only when that state is finite and non-degenerate. A larger step, or an unusable state, silently restores the cold initial condition rather than seeding the solver from somewhere it has no reason to be. Both compile once and reuse the same executable for every point, because the shapes never change.

gkx.workflows.runtime.orchestration_scan.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, window_method: str = 'stationary', krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', show_progress: bool = False, workers: int = 1, parallel_executor: str = 'thread', warm_start: bool | None = None) → RuntimeLinearScanResult

Run a ky scan; warm_start overrides the case output policy.

Coordinates serial, independent-worker, or combined-ky runtime scans.

gkx.workflows.runtime.orchestration_scan.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, window_method: str = 'stationary') → RuntimeLinearScanResult

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

gkx.workflows.runtime.orchestration_scan.run_runtime_scan_ky_task(task: dict[str, Any]) → RuntimeLinearResult

Run one independent ky point for ordered scan-worker execution.

Runtime artifact writers, display, restart, and checkpoint handoff policies.

class gkx.workflows.runtime.artifacts.NonlinearArtifactPolicy(out_path: Path | None, netcdf_output_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.

gkx.workflows.runtime.artifacts.print_nonlinear_command_outputs(paths: Mapping[str, str], *, enabled: bool) → None

Print nonlinear artifact paths after diagnostics confirm a saved run.

gkx.workflows.runtime.artifacts.print_saved_paths(paths: Mapping[str, str], keys: Sequence[str]) → None

Print saved artifact paths in command-defined display order.

gkx.workflows.runtime.artifacts.resolve_nonlinear_artifact_policy(cfg: Any, *, out: str | Path | None, diagnostics: bool | None, steps: int | None, dt: float | None) → NonlinearArtifactPolicy

Resolve nonlinear output, restart, and checkpoint policy.

gkx.workflows.runtime.artifacts.run_runtime_nonlinear_with_artifacts(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) → tuple[RuntimeNonlinearResult, dict[str, str]]

Run nonlinear runtime chunks and hand results to artifact writers.

gkx.workflows.runtime.artifacts.write_command_outputs(out_path: str | Path | None, payload: Any | None, *, writer: Callable[[str | Path, Any], dict[str, str]], display_keys: Sequence[str]) → dict[str, str]

Write command artifacts when both destination and payload exist.

gkx.workflows.runtime.artifacts.write_linear_runtime_command_outputs(*, linear_out_path: str | Path | None, quasilinear_out_path: str | Path | None, result: Any, linear_writer: Callable[[str | Path, Any], dict[str, str]], quasilinear_writer: Callable[[str | Path, Any], dict[str, str]]) → dict[str, dict[str, str]]

Write all optional artifacts produced by one linear runtime command.

gkx.workflows.runtime.artifacts.write_runtime_nonlinear_artifacts(out: str | Path, result: Any, cfg: Any | None = None) → dict[str, str]

Write summary/diagnostics/state artifacts for a nonlinear runtime run.

gkx.workflows.runtime.artifacts.write_scan_runtime_command_outputs(out_path: str | Path | None, scan: Any, *, writer: Callable[[str | Path, Any], dict[str, str]]) → dict[str, str]

Write optional artifacts produced by one linear-scan runtime command.

Carry a converged state between related runs so repeated work restarts warm.

The persistent JAX compilation cache reuses the executable between runs. This module is its state-side counterpart: it reuses the answer. It applies only where work repeats – a ky scan, a parameter scan, an optimizer loop – and never to a single run, which still builds its own cold initial condition and pays exactly what it paid before.

Two invariants make the reuse safe rather than merely fast:

• Scale is removed, not carried. A converged linear state has grown by exp(gamma T) and would overflow (or, on a damped branch, underflow) if it were carried raw. Both the eigensolve and the growth-rate fit are invariant to the overall scale of the initial condition, so rescaling loses nothing. The rescaling factor is rounded to a power of two, which is exact in binary floating point, so the carried state is the converged state to the last bit.

• A carried state is validated before use. A non-finite, empty, or zero-norm state is refused and the caller falls back to its cold path, so a failed neighbour can never poison the point after it.

Warm start is opt-in, and the measurements in docs/performance.rst are why. On the certified adaptive eigensolver it is correctness-neutral (cold and warm agree to 1e-6 relative on the Cyclone deck, well inside the certified residual) but cost-neutral: that solver’s work is a fixed-size filtered Arnoldi whose cost is set by the Krylov dimension and the filter length, not by the quality of v0, and it already converges in one restart from the analytic seed. On a fixed-horizon time integration a warm seed does pay – it reaches the horizon-converged growth rate in half the horizon – but at any horizon short of convergence it reports a different number than a cold start, because it has removed a startup transient the cold run still contains. Neither is a reason to remove the machinery; both are reasons not to switch it on for somebody without being asked.

An eigensolver started inside one branch can also converge to that branch, so a case whose branch structure is not already known should be scanned once cold and compared point by point before the warm numbers are trusted.

class gkx.workflows.runtime.warm_start.SaturationPlan(seed: ndarray | None, steps: int, warm: bool, reason: str)

How the next spin-up should be seeded, and why.

class gkx.workflows.runtime.warm_start.SaturationRefreshPolicy(max_reuse: int = 4, geometry_tolerance: float = 0.05, warm_step_fraction: float = 0.25)

When a carried saturated state stops being a valid seed.

A saturated nonlinear state belongs to the geometry it was grown in. Two things invalidate it, and both are bounded here rather than left to trust:

• geometry_tolerance: the relative change in the flux-tube geometry signature above which the previous attractor is not the current one.

• max_reuse: how many consecutive warm spin-ups may follow one cold spin-up before a cold one is forced regardless of the metric, so a long run of individually small steps cannot drift arbitrarily far.

warm_step_fraction is the fraction of the cold spin-up budget a warm spin-up is given. It is a budget, not a claim of convergence: a warm seed starts at saturated amplitude and so skips the linear growth phase, but it still has to re-equilibrate, and the cold budget is restored the moment either bound above is crossed.

class gkx.workflows.runtime.warm_start.SaturationWarmStart(policy: SaturationRefreshPolicy = <factory>, state: ndarray | None = None, signature: ndarray | None = None, reuse_count: int = 0)

Decide, and record, whether an optimizer iteration may reuse saturation.

This deliberately does not change when the saturated state is refreshed. Callers refresh it at the same points they always did – accepted stages – so the objective stays a fixed function of its inputs for the whole of each stage and never becomes a function of the optimizer’s within-stage history. What changes is only the cost of producing the refreshed state.

plan(signature: Any, *, cold_steps: int) → SaturationPlan

Return the seed and step budget for the next spin-up.

record(state: Any, signature: Any, *, warm: bool) → None

Store the new saturated state and count the reuse chain.

class gkx.workflows.runtime.warm_start.WarmStartPolicy(enabled: bool = True, amplitude: float = 1.0)

Whether, and at what amplitude, a converged state is carried forward.

classmethod from_config(cfg: Any, *, override: bool | None = None) → WarmStartPolicy

Resolve [output] warm_start with an optional explicit override.

gkx.workflows.runtime.warm_start.carry_state(state: Any, *, amplitude: float = 1.0) → ndarray | None

Rescale a converged linear state for reuse, or refuse it as a guess.

The returned state differs from state by one power-of-two factor, which is exact in binary floating point: no information is lost and no rounding is introduced. None means the state is unusable and the caller must fall back to a cold start. Nonlinear saturated states carry a physical amplitude and must not pass through here – check them with state_norm and reuse them unscaled.

gkx.workflows.runtime.warm_start.flux_tube_signature(geom: Any) → ndarray

Return the flux-tube profiles a carried saturated state is judged against.

These are exactly the metric coefficients the nonlinear operator reads, so two geometries with the same signature drive the same turbulence and a state saturated in one is a valid seed in the other.

gkx.workflows.runtime.warm_start.linear_scan_warm_start_refusal(*, policy: WarmStartPolicy, solver_key: str, workers: int) → str | None

Return why a linear scan cannot warm start, or None when it can.

Refusals are policy, not failure: each names a contract that carrying state would break, and the caller runs its existing cold path unchanged.

gkx.workflows.runtime.warm_start.relative_change(previous: Any, current: Any) → float

Return the relative L2 distance between two parameter or geometry vectors.

gkx.workflows.runtime.warm_start.resolve_scan_warm_start(args: Any, scan_cfg: Mapping[str, Any]) → bool | None

Resolve warm start from the CLI flag, then [scan], else the config.

None means “not requested either way” and defers to [output] warm_start, so the three surfaces compose without any of them having to know the default.

gkx.workflows.runtime.warm_start.scan_visit_order(values: Sequence[float] | ndarray) → ndarray

Return the visit order that puts neighbouring scan points side by side.

A converged solution is only a good guess for a nearby point, so the scan is walked in monotone parameter order regardless of the order requested. Results are written back into the requested order by the caller, so the reported arrays are unchanged.

gkx.workflows.runtime.warm_start.signature_from_arrays(arrays: Iterable[Any]) → ndarray

Flatten geometry or parameter arrays into one comparable signature.

gkx.workflows.runtime.warm_start.state_norm(state: Any) → float | None

Return the L2 norm of a reusable state, or None if it is unusable.

Runtime Commands

Runtime executable command workflows.

class gkx.workflows.runtime.commands.RuntimeLinearCommandOptions(ky: float, Nl: int, Nm: int, solver: str, fit_signal: str, method: str | None, dt: float | None, steps: int | None, sample_stride: int, method_for_header: str, dt_for_header: float, steps_for_header: int, show_progress: bool)

Resolved executable options for one linear runtime command.

class gkx.workflows.runtime.commands.RuntimeNonlinearCommandOptions(ky: float, Nl: int, Nm: int, method: str, dt: float, steps: int | None, sample_stride: int, diagnostics_stride: int | None, diagnostics: bool, laguerre_mode: str | None, show_progress: bool, run_to: str | None = None)

Resolved executable options for one nonlinear runtime command.

class gkx.workflows.runtime.commands.RuntimeScanCommandOptions(ky_values: tuple[float, ...], Nl: int, Nm: int, solver: str, fit_signal: str, method: str | None, dt: float | None, steps: int | None, sample_stride: int, batch_ky: bool, show_progress: bool, workers: int, parallel_executor: str, warm_start: bool | None)

Resolved executable options for one linear scan command.

gkx.workflows.runtime.commands._arg_or_section(args: Any, section: dict[str, Any], name: str, default: Any) → Any

Return an executable flag override or a TOML section value.

gkx.workflows.runtime.commands._parse_ky_values(args: Any, scan_cfg: dict[str, Any]) → tuple[float, ...]

Resolve scan ky values from CLI or TOML, failing closed on empty scans.

gkx.workflows.runtime.commands._prepare_runtime_command_config(args: Any, *, path_overrides: bool, quasilinear_overrides: bool) → tuple[RuntimeConfig, dict[str, Any]]

Load runtime command config and apply the command-specific overrides.

gkx.workflows.runtime.commands._resolve_grid_time_options(args: Any, section: dict[str, Any], cfg: RuntimeConfig, hl_fallback: tuple[int, int] = (12, 24)) → tuple[int, int, str | None, float | None, int | None, int]

Resolve resolution (hl_fallback from startup), time controls, cadence.

gkx.workflows.runtime.commands._resolve_linear_command_options(args: Any, cfg: RuntimeConfig, run_cfg: dict[str, Any]) → RuntimeLinearCommandOptions

Resolve linear command options from flags, TOML, and config defaults.

gkx.workflows.runtime.commands._resolve_linear_fit_options(args: Any, section: dict[str, Any]) → tuple[str, str]

Resolve the linear eigensignal solver and fit signal.

gkx.workflows.runtime.commands._resolve_nonlinear_command_options(args: Any, cfg: RuntimeConfig, run_cfg: dict[str, Any]) → RuntimeNonlinearCommandOptions

Resolve nonlinear command options from flags, TOML, and config defaults.

gkx.workflows.runtime.commands._resolve_scan_command_options(args: Any, cfg: RuntimeConfig, scan_cfg: dict[str, Any]) → RuntimeScanCommandOptions

Resolve linear-scan command options from flags and TOML defaults.

gkx.workflows.runtime.commands._runtime_fit_config(data: dict[str, Any]) → dict[str, Any]

Return the deck’s [fit] options, refusing keys no fit owner reads.

gkx.workflows.runtime.commands.apply_quasilinear_overrides(cfg: RuntimeConfig, args: Any) → RuntimeConfig

Apply executable quasilinear diagnostic overrides.

gkx.workflows.runtime.commands.apply_runtime_path_overrides(cfg: RuntimeConfig, args: Any, *, resolve_runtime_path: Callable[[...], str | None]) → RuntimeConfig

Apply cwd-resolved executable path overrides for geometry and init files.

gkx.workflows.runtime.commands.attach_preloaded_runtime_config(args: Any, cfg: RuntimeConfig, data: dict[str, Any]) → None

Attach already-loaded runtime TOML data to a parser namespace.

The generic run dispatcher inspects a config to choose linear or nonlinear execution. Passing the loaded object forward avoids a second TOML parse while direct subcommands remain self-contained.

gkx.workflows.runtime.commands.auto_plot_runtime_outputs(kind: str, paths: Mapping[str, str], *, enabled: bool) → list[str]

Draw and announce the figure set for a run that has finished writing.

The whole call is contained. The simulation is already on disk by the time this runs, so a headless display, an unusable matplotlib backend, or an import that fails outright must cost the pictures and nothing else – the caller’s exit status never depends on what happens here.

gkx.workflows.runtime.commands.load_runtime_command_config(args: Any) → tuple[RuntimeConfig, dict[str, Any]]

Load runtime TOML data, reusing the generic-dispatch preload if present.

gkx.workflows.runtime.commands.plot_saved_output_command(argv: Sequence[str], *, plot_saved_output: Callable[[...], Path]) → int

Render a saved runtime artifact from the top-level --plot command.

gkx.workflows.runtime.commands.print_linear_run_header(*, label: str, config_path: str, ky: float, Nl: int, Nm: int, solver: str, method: str, dt: float, steps: int, grid_shape: tuple[int, int, int], show_progress: bool, extra: str | None = None) → None

Print the standard executable header for linear initial-value runs.

gkx.workflows.runtime.commands.print_nonlinear_run_header(*, config_path: str, runtime_config: Any, ky: float, Nl: int, Nm: int, method: str, dt: float, steps: int | None, grid_shape: tuple[int, int, int], diagnostics: bool, show_progress: bool) → None

Print the standard executable header for nonlinear initial-value runs.

gkx.workflows.runtime.commands.print_nonlinear_run_summary(result: Any) → bool

Print final nonlinear diagnostics and return whether diagnostics exist.

gkx.workflows.runtime.commands.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.

gkx.workflows.runtime.commands.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.

gkx.workflows.runtime.commands.run_runtime_linear_command(args: Any) → int

Execute the runtime-linear subcommand after parser dispatch.

gkx.workflows.runtime.commands.run_runtime_nonlinear_command(args: Any) → int

Execute the runtime-nonlinear subcommand after parser dispatch.

gkx.workflows.runtime.commands.runtime_output_path(args: Any, cfg: RuntimeConfig) → str | None

Return the executable output path override or TOML output path.

gkx.workflows.runtime.commands.scan_runtime_linear_command(args: Any) → int

Execute the runtime-linear ky-scan subcommand after parser dispatch.

gkx.workflows.runtime.commands.should_show_progress(args: Any, configured: bool) → bool

Resolve progress output from executable flags, TOML config, and TTY state.

gkx.workflows.runtime.commands.should_write_plots(args: Any, cfg: RuntimeConfig) → bool

Resolve whether a finished run draws its own figures.

--no-plots overrides the configuration, which itself defaults to on.

Runtime TOML Inputs

TOML-based input helpers for the executable and driver scripts.

gkx.workflows.runtime.toml._apply_runtime_section_overrides(cfg: RuntimeConfig, data: dict[str, Any]) → RuntimeConfig

Apply non-nested runtime config sections from TOML data.

gkx.workflows.runtime.toml._default_nonlinear_stepper(physics: Any) → str

Stepper of a nonlinear deck with no dt and no method.

Kinetic electrons at low beta: adaptive imex-ark3 (the electron streaming CFL is gone; falls back to rk4-pi where the factor check refuses it). Everything else: rk4-pi (about 2x fewer steps than rk4 + CFL).

gkx.workflows.runtime.toml._merge_dataclass(base: Any, overrides: dict | None, section: str = '[deck]') → Any

Recursively merge a dict into a dataclass, returning a new instance.

gkx.workflows.runtime.toml._normalize_geometry_overrides(overrides: dict | None) → dict | None

Return geometry overrides using the canonical runtime schema.

gkx.workflows.runtime.toml._normalize_time_overrides(overrides: Any, physics: Any = None) → dict[str, Any] | None

Pair the integrator with the step policy when the deck chose no step.

ExplicitTimeConfig.dt is a required field and TimeConfig.dt carries a default of 0.1. That is the whole of the real difference between the two time-configuration surfaces: a caller of the library struct has always chosen a step, a deck may not have. So the two fixed_dt defaults are not the same question and must not simply be made equal – fixed_dt=True is right for a step the user chose, and fourteen shipped decks and parity fixtures omit fixed_dt and depend on it.

What is wrong is applying a defaulted step as a fixed one. Measured on the shipped Cyclone deck at (Nl,Nm)=(4,8), ky=0.3, with the TimeConfig dataclass defaults, against that rung’s certified adaptive eigenpair gamma=0.10128645:

method / policy	gamma	steps	RHS evals	rel. error
rk2, fixed dt=0.1	FAILS	–	–	FloatingPointError
rk4, fixed dt=0.1	FAILS	–	–	FloatingPointError
rk2, CFL-controlled	0.10126899	7806	15612	-1.72e-04
rk3, CFL-controlled	0.10126041	4512	13536	-2.57e-04
rk4, CFL-controlled	0.10125984	2768	11072	-2.63e-04

The defaulted dt=0.1 is ~7.8x the CFL-stable 0.01281 for this deck, so both fixed-step arms overflow whatever the scheme. With the controller on, all three schemes agree to better than 3e-4 and rk4 reaches the same horizon in 29.1% fewer right-hand-side evaluations – exactly 4/(2.82*2), because rk4’s CFL prefactor buys a 2.82x longer step for twice the work per step. rk4 is therefore paired with the controller and only with it: at a fixed step rk4 has no step-size compensation and is simply twice the cost of rk2, so the dataclass default stays rk2 for a deck that chose its own dt.

A deck that sets dt keeps today’s behaviour exactly, and every shipped deck sets it, so this changes none of them. fixed_dt = false with dt takes the controller’s stepper: adaptive rk2 grows every streaming mode.

gkx.workflows.runtime.toml._replace_runtime_section(cfg: RuntimeConfig, data: dict[str, Any], key: str, constructor: Callable[[...], Any]) → RuntimeConfig

Replace one runtime section when the TOML section is present.

gkx.workflows.runtime.toml._resolve_runtime_config_paths(cfg: RuntimeConfig, *, base_dir: Path) → RuntimeConfig

Resolve every path-valued runtime field against the TOML directory.

gkx.workflows.runtime.toml._runtime_base_config(data: dict[str, Any]) → RuntimeConfig

Return a runtime config after applying common dataclass sections.

gkx.workflows.runtime.toml._runtime_run_from_toml(section: Any) → RuntimeRunConfig | None

Read the deck’s [run] selections, or None when it has none.

Only the keys the runtime itself resolves are carried. Anything else in the table belongs to a subcommand’s own parser and is left where it is.

gkx.workflows.runtime.toml._runtime_species_from_toml(species_raw: Any) → tuple[RuntimeSpeciesConfig, ...] | None

Parse optional [[species]] runtime entries.

gkx.workflows.runtime.toml._validate_removed_time_keys(data: dict[str, Any]) → None

Reject decks that still select the removed Diffrax integrator family.

gkx.workflows.runtime.toml.describe_binary_input(path: Path) → str | None

Return what path actually is when it is not text, else None.

gkx.workflows.runtime.toml.load(path: str | Path) → RuntimeConfig

Load a resolved immutable GKX case from TOML.

gkx.workflows.runtime.toml.load_runtime_from_toml(path: str | Path) → tuple[RuntimeConfig, dict]

Load unified runtime config from TOML, returning (cfg, data).

gkx.workflows.runtime.toml.load_toml(path: str | Path) → dict

Load TOML, naming common binary formats instead of reporting byte offsets.

gkx.workflows.runtime.toml.migrate_end_damping_reference(path: str | Path, *, route: str, dt_step: float = 0.0) → RuntimeConfig

Explicitly convert a legacy deck; write the returned case with to_toml().

dt_step is the actual fixed linear reference step, not an inferred input default. timestep_free identifies RHS calls without a solver timestep.

gkx.workflows.runtime.toml.reject_unknown_keys(section: str, keys: Any, allowed: Any) → None

Raise on deck keys no owner reads, naming the closest valid spelling.

gkx.workflows.runtime.toml.resolve_runtime_path(value: str | None, *, base_dir: Path) → str | None

Expand a config path and resolve it relative to base_dir.

gkx.workflows.runtime.toml.toml_shorthand_command(data: dict[str, Any]) → str

Return the executable command used for direct TOML path shorthand.

Runtime Artifact Package

Runtime artifact writers and reload helpers.

gkx.artifacts.validate_finite_array(value: Any, *, label: str) → None

Raise if an optional artifact array contains NaN or infinite values.

gkx.artifacts.validate_finite_runtime_result(result: Any, *, label: str) → None

Validate nonlinear runtime result payloads before artifact writes.

gkx.artifacts.write_quasilinear_artifacts(out: str | Path, quasilinear: dict[str, Any]) → dict[str, str]

Write quasilinear summary and species tables.

gkx.artifacts.write_runtime_linear_artifacts(out: str | Path, result: Any) → dict[str, str]

Write summary/timeseries/state artifacts for a linear runtime run.

gkx.artifacts.write_runtime_linear_scan_artifacts(out: str | Path, result: Any) → dict[str, str]

Write ky-scan growth/frequency and optional quasilinear spectra.

gkx.artifacts.write_runtime_nonlinear_table_artifacts(out: str | Path, result: Any) → dict[str, str]

Write non-NetCDF nonlinear summary/diagnostics/state artifacts.

Runtime Artifact IO

Read and write runtime summaries, diagnostics, restart state, and tables.

gkx.artifacts.io.load_netcdf_restart_state(path: str | Path, *, nspecies: int, Nl: int, Nm: int, ny: int, nx: int, nz: int, ky_layout: Literal['full', 'half'] = 'full') → ndarray

Load a NetCDF restart file onto ky_layout.

ny is the length of the two-sided axis in both layouts; ky_layout says how many rows the caller wants back, Ny or Nyc = 1 + Ny // 2. The file itself is unchanged by the layout: it has always stored the dealiased ky >= 0 block, which is why a run started on either axis can continue from a file written by the other, and why a 2.1.0 restart still loads (plan 5.3 N3, queue row Q10).

gkx.artifacts.io.validate_finite_array(value: Any, *, label: str) → None

Raise if an optional artifact array contains NaN or infinite values.

gkx.artifacts.io.validate_finite_runtime_result(result: Any, *, label: str) → None

Validate nonlinear runtime result payloads before artifact writes.

gkx.artifacts.io.write_netcdf_restart_state(path: str | Path, state: Array | ndarray | bool | number | bool | int | float | complex, *, ny_full: int | None = None) → Path

Write a restart state in flat complex64 restart layout.

The file carries no header, so the reader tells its two accepted forms apart by size alone: Ny stored ky rows is GKX’s own state order, and Nyc = 1 + Ny // 2 rows is the packed positive-ky interchange order, whose axes are permuted (Input Files and Executable). A half-spectrum state would land on the second size while holding the first order, and the reader has nothing to tell the two apart with – it would return a scrambled state of the right shape and dtype.

So a half-spectrum state is widened before it is written, which needs ny_full: the length of the two-sided axis. The widening is exact and costs only file size, because every row it adds is the conjugate of a row already there. Nothing about an existing file changes: a two-sided state is written exactly as before, byte for byte.

ny_full=None writes whatever it is handed, which is what a caller that already holds a two-sided state does.

gkx.artifacts.io.write_quasilinear_artifacts(out: str | Path, quasilinear: dict[str, Any]) → dict[str, str]

Write quasilinear summary and species tables.

gkx.artifacts.io.write_runtime_linear_artifacts(out: str | Path, result: Any) → dict[str, str]

Write summary/timeseries/state artifacts for a linear runtime run.

gkx.artifacts.io.write_runtime_linear_scan_artifacts(out: str | Path, result: Any) → dict[str, str]

Write ky-scan growth/frequency and optional quasilinear spectra.

gkx.artifacts.io.write_runtime_nonlinear_table_artifacts(out: str | Path, result: Any) → dict[str, str]

Write non-NetCDF nonlinear summary/diagnostics/state artifacts.

NetCDF Spectral Layout

NetCDF spectral-output layout helpers.

gkx.artifacts.spectral_layout.KY_WEIGHTING_PER_ROW = 'per_row'

How a ky-resolved diagnostic array counts a conjugate pair.

"per_row" – the value of that one stored row. This is what the two-sided reductions produce, because their weight is 1 on every row (gkx.core_ky_layout.hermitian_mode_weights()): Phi2, Wg, Wphi, Wapar and TurbulentHeating. On a half axis the same reduction carries the pair weight 2, so it has to be divided out before it is published or the file’s numbers would double.

"pair" – the sum over the conjugate pair, which the transport weights (gkx.core_ky_layout.transport_mode_weights()) already produce on both axes: HeatFlux* and ParticleFlux*. Nothing to undo.

The distinction is not cosmetic and cannot be inferred from the array: two arrays of the same shape holding the same modes differ by a factor of two, and only the reduction that built them knows which.

gkx.artifacts.spectral_layout._condense_kx_for_output(arr: ndarray, *, full_nx: int, active_nx: int) → ndarray

Return kx-resolved data on the dealiased output axis.

Fresh in-memory diagnostics carry the full spectral kx axis, while history loaded from an existing existing NetCDF output bundle is already condensed. External restart continuation appends both forms, so the writer must not apply the active-index selection a second time.

gkx.artifacts.spectral_layout._condense_ky(arr: ndarray, *, ny_full: int | None = None) → ndarray

Keep the dealiased ky >= 0 rows of a ky-resolved array.

ny_full is the length of the two-sided axis; without it the row count is taken from the array, which is the right answer only when the array is the two-sided axis.

gkx.artifacts.spectral_layout._condense_ky_for_output(arr: ndarray, *, full_ny: int, active_ny: int, ky_weighting: str | None = None) → ndarray

Return ky-resolved data on the dealiased positive-ky output axis.

The input may already be condensed (active_ny rows, which is how history reloaded from an existing bundle arrives), two-sided (full_ny rows), or half-spectrum (Nyc = 1 + Ny // 2 rows). All three publish the same numbers; see _half_ky_publish_factor() for why the last one needs ky_weighting.

gkx.artifacts.spectral_layout._condense_kykx(arr: ndarray, *, ny_full: int | None = None) → ndarray

Keep the dealiased ky >= 0 / kx block; see _condense_ky().

gkx.artifacts.spectral_layout._condense_kykx_for_output(arr: ndarray, *, full_ny: int, full_nx: int, active_ny: int, active_nx: int, ky_weighting: str | None = None) → ndarray

Return ky-kx-resolved data on dealiased output axes.

ky_weighting carries the same contract as in _condense_ky_for_output() and is needed for the same reason.

gkx.artifacts.spectral_layout._dealiased_ky_values(ky: ndarray, *, ny_full: int | None = None) → ndarray

Return the published ky axis: the dealiased ky >= 0 magnitudes.

ny_full is the length of the two-sided axis. It is what makes the answer the same from either layout’s grid.ky: the number of dealiased rows is a property of Ny, and a half-spectrum axis already holds Nyc rows, so taking the count from ky.shape[0] would publish 1 + (Nyc - 1) // 3 of them – about a third of the band – with an axis of the right dtype and the wrong length.

gkx.artifacts.spectral_layout._dealiased_spectral_field(field: ndarray, *, ky_axis: int = 0, kx_axis: int = 1, ny_full: int | None = None) → ndarray

Return the dealiased ky >= 0 / kx block of a spectral field.

A field is not a reduction, so no weight is involved: the retained rows are the same indices holding the same coefficients in either layout. What still has to come from ny_full is how many of them there are, which is a property of the two-sided axis.

gkx.artifacts.spectral_layout._half_ky_publish_factor(full_ny: int, *, ky_weighting: str | None, what: str) → ndarray | None

Return the per-row factor that republishes a half-spectrum reduction.

The published ky axis is the dealiased ky >= 0 block, which is the same rows in both layouts – so condensing a half-spectrum diagnostic is a row selection and nothing more, as long as the two layouts agree on what a row’s number means. They do not, for the reductions weighted by gkx.core_ky_layout.hermitian_mode_weights(): the two-sided axis stores each row once at weight 1, and the half axis carries the conjugate partner’s share at weight 2. Publishing the half value unchanged would double every paired row of Phi2, Wg, Wphi, Wapar and TurbulentHeating – and with them Phi2_t, which the writer derives from the condensed spectrum.

Returning the reciprocal weights restores the two-sided per-row value exactly: the weight is 1 or 2, and division by two is exact in binary floating point, so the published block is bitwise what the two-sided run would have written from the same state.

None means “publish as it stands”, which is what the transport reductions need: their weight is already the pair weight on both axes.

gkx.artifacts.spectral_layout._restart_to_netcdf_layout(state: ndarray, *, ny_full: int | None = None) → ndarray

Pack an evolved state into the restart file’s dealiased ky/kx block.

ny_full is the length of the two-sided ky axis. The retained rows are 0 .. 1 + (Ny - 1) // 3, which is the same index range in both layouts – the restart file has always stored the non-negative dealiased block – but the count is a property of Ny, not of how many rows the state carries. Deriving it from state.shape[3] on a half-spectrum state keeps only 1 + (Nyc - 1) // 3 rows, about a third of the band, and writes that file without any error, so the length is taken from ny_full when it is supplied.

gkx.artifacts.spectral_layout._spectral_species_to_xy(values: ndarray, *, ny_full: int | None = None) → ndarray

Real-space (s, y, x, z) field of a (s, ky, kx, z) spectrum.

The species-axis twin of _spectral_to_xy(); see its docstring for why a half-spectrum input must not go through the complex ifft2.

gkx.artifacts.spectral_layout._spectral_to_xy(field: ndarray, *, ny_full: int | None = None) → ndarray

Return the real-space (y, x, z) field of a (ky, kx, z) spectrum.

ny_full is the length of the two-sided ky axis. A two-sided field keeps the full complex ifft2, whose imaginary part is discarded; a half-spectrum field (gkx.core_ky_layout) has no negative rows to transform and goes through irfft2, which produces the same Ny real-space rows and the same values without materialising them. Taking np.real of a complex ifft2 of Nyc rows instead would silently return an Nyc-row image against an Ny-long y axis.

Nonlinear Output NetCDF

Write nonlinear NetCDF output, restart, diagnostics, and final-field bundles.

gkx.artifacts.nonlinear_netcdf._build_output_grid_and_geometry(cfg: Any) → tuple[Any, Any]

Resolve artifact output onto the same geometry-implied grid as the solver.

gkx.artifacts.nonlinear_netcdf._write_diagnostics_group(root: Any, diag: Any, cfg: Any, *, nspecies: int, full_nx: int, full_ny: int, active_nx: int, active_ny: int) → None

Write scalar, species, resolved, and split nonlinear diagnostics.

Quasilinear Transport Diagnostics

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 gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.normalize_quasilinear_channels(channels: Iterable[str] | str) → tuple[str, ...]

Normalize and validate quasilinear field channels.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

gkx.diagnostics.quasilinear_transport.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.

The (ky, kx) factor is the layout contract’s Hermitian reduction weight, owned by gkx.core_ky_layout.hermitian_mode_weights(), not a fourth copy of it.

Run-to-Saturation Stop Policy

Run-to-saturation stop policy for nonlinear transport runs.

The decision this module makes is when a run has produced a heat-flux mean worth stopping on: spin-up removed, the correlation time resolved inside the retained window, the correlated standard error below a relative threshold, and both halves of the window agreeing with each other – optionally with a guard trace, the field energy, held to the same stationarity so a flat-looking flux cannot end a run whose fields are still drifting.

It sits apart from transport_windows.py because the question is different. That module scores a window after the fact against promotion gates; this one is asked the same question repeatedly by a running solver, on a trace that grows under it, and has to answer “not yet” with a reason the runtime can record.

The Sokal estimator is shared with post-hoc analysis, so runtime and campaign uncertainties use one definition.

class gkx.diagnostics.saturation.SaturationStopConfig(rel_sem: float = 0.05, min_window: float | None = None, min_samples: int = 256)

Stop policy for run_to = "saturation" nonlinear runs.

gkx.diagnostics.saturation.saturation_stop_decision(time: Sequence[float] | ndarray, values: Sequence[float] | ndarray, *, guard: Sequence[float] | ndarray | None = None, free_energy_guard: Sequence[float] | ndarray | None = None, config: SaturationStopConfig | None = None) → dict[str, Any]

Decide whether a nonlinear trace has saturated well enough to stop.

Spin-up removal: everything before the trace first reaches its own median is discarded. During the linear-growth phase the flux sits below its saturated level, so once any plateau exists the overall median lies on the plateau and the first crossing lands at the end of growth; the overshoot decay that follows is left to the stationarity gate rather than trimmed.

Saturation requires all of: a resolved tau_ac (the autocorrelation crosses zero inside the window, and not at the first lag – a trace that decorrelates within one diagnostic sample has shown no correlation time, only noise), at least min_samples retained samples, window span at least 20 tau_ac and any larger configured min_window, IAT-corrected relative SEM at most rel_sem, first/second half-window means within twice their combined SEM, and – when guard (Wphi) or free_energy_guard (Wg) is given – the same half-window stationarity on each guard over the same window. Guards have no relative-SEM gate: they protect against a flat-looking flux while either energy still drifts. The trace is assumed finite; the runtime validates each chunk before this.

gkx.diagnostics.saturation.sokal_autocorrelation_time(signal: ndarray, dt: float) → tuple[float, int, ndarray]

Return the first-zero IAT, crossing index, and autocorrelation.

Solver Eigen Objectives

Implicit eigenvalue objectives and branch-continuity diagnostics.

The routines in this module are intentionally small and JAX-native: they expose an implicit eigenpair derivative (forward and reverse) for locally isolated dominant-growth branches, plus finite-difference diagnostics that verify the selected branch is consistent before using the derivative in optimization gates.

gkx.objectives.eigen.dominant_eigenvalue_branch_locality_report(base_matrix: Array | ndarray, plus_matrix: Array | ndarray, minus_matrix: Array | ndarray, *, step: float, gap_floor: float = 1e-08, slope_rtol: float = 0.01, slope_atol: float = 1e-08) → dict[str, object]

Report whether dominant-growth finite differences follow one branch.

The dominant-growth custom VJP assumes the max-real eigenvalue selected at the primal point is locally isolated. This diagnostic compares two central finite-difference slopes:

dominant_growth_fd_slope

Uses the max-real eigenvalue independently at the plus/minus points.

nearest_branch_growth_fd_slope

Uses the eigenvalue nearest to the base dominant eigenvalue at the plus/minus points.

When those branches differ, exact finite differences are measuring a branch switch rather than the local derivative used by the implicit eigenpair VJP.

gkx.objectives.eigen.dominant_real_eigenvalue(matrix: Array) → Array

Return the dominant growth rate with an implicit eigenpair derivative.

The derivative is Re(dlambda) of dominant_eigenpair(): it works in forward mode (jvp, jacfwd, VMEX’s implicit Jacobian) and in reverse mode. It treats the max-real branch selected at the primal point as locally isolated; callers that use it in optimization should keep finite-difference or branch-continuity checks near accepted candidates.

Solver Objective Core

Core linear and quasilinear solver-objective evaluators.

class gkx.objectives.core.AdaptiveLinearEigensolverConfig(krylov_dim: int = 24, restart_krylov_dim: int = 12, candidate_count: int = 2, max_restarts: int = 4, tolerance: float = 1e-09, chunk_horizon: float = 30.0, stability_dimension: int = 12, stability_probe_count: int = 2, stability_safety: float = 0.9, max_stability_retries: int = 2, exponential_krylov_dim: int | None = None, exponential_horizon: float = 5.0, adjoint_krylov_dim: int = 24, adjoint_max_restarts: int = 200, sensitivity_rtol: float = 1e-08, sensitivity_restart: int = 200, sensitivity_max_restarts: int = 4, sensitivity_solver: str = 'propagator', sensitivity_krylov_dim: int = 16, sensitivity_preconditioner: str = 'field-corrected', condition_limit: float = 100000000.0, branch_gap_floor: float = 1e-08)

Fail-closed settings for the differentiable matrix-free objective path.

gkx.objectives.core.solver_growth_rate_from_geometry(geom: Any, *, spectral_grid: Any | None = None, selected_ky_index: int = 1, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, lx: float = 6.0, ly: float = 12.0, params_linear: LinearParams | None = None, terms: LinearTerms | None = None, eigensolver: Literal['dense', 'sparse-direct'] = 'dense', shift: complex | None = None, candidates: int = 12) → Array

Evaluate the dominant linear growth rate, differentiable in both AD modes.

The dense default differentiates in both modes through gkx.objectives.eigen.dominant_eigenpair().

eigensolver="sparse-direct" assembles the operator’s exact sparse matrix from compressed products, factors A - shift I once on the host (solvax.sparse_eigenvalue) and returns the largest growth rate among the candidates eigenvalues nearest shift (a warm start such as the previous optimizer step’s eigenvalue). Its gradient is one operator VJP at the right eigenvector, weighted by the left eigenvector from the same factorization. It needs a SOLVAX with sparse_eigenvalue, returns NaN if either eigenpair misses a 1e-9 residual, and finds only eigenvalues near shift: the dense default searches the whole spectrum.

gkx.objectives.core.solver_linear_operator_matrix_from_geometry(geom: Any, *, spectral_grid: Any | None = None, selected_ky_index: int = 1, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, lx: float = 6.0, ly: float = 12.0, params_linear: LinearParams | None = None, terms: LinearTerms | None = None) → Array

Materialize the complex linear-RHS operator for one solver geometry.

This helper exposes the exact matrix whose dominant eigenvalue is used by solver_growth_rate_from_geometry(). It is intended for branch locality and AD/finite-difference admission gates; production time integration should continue to call the RHS directly.

gkx.objectives.core.solver_objective_vector_from_geometry(geom: Any, *, spectral_grid: Any | None = None, selected_ky_index: int = 1, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, lx: float = 6.0, ly: float = 12.0, params_linear: LinearParams | None = None, terms: LinearTerms | None = None, eigensolver: Literal['dense', 'adaptive-propagator'] = 'dense', adaptive_config: AdaptiveLinearEigensolverConfig | None = None) → Array

Evaluate dominant linear/quasilinear observables from geometry.

This is a reusable value-level objective builder for optimization drivers and examples. It builds the production linear RHS on the supplied solver-ready flux-tube geometry, selects the maximum-growth eigenbranch, and returns the ordered SOLVER_OBJECTIVE_NAMES vector.

eigensolver="adaptive-propagator" avoids materializing the dense operator, certifies the continuous residual, and differentiates the eigenpair implicitly. It remains explicit opt-in until branch-continuity gates admit a production optimization. The default dense path preserves established results. Pass a runtime spectral_grid that has already selected one ky to preserve linked stellarator boundary conditions; otherwise a periodic single-ky grid is built from the scalar grid options. One-dimensional species parameters add the leading multi-species state axis automatically.

gkx.objectives.core.solver_scalar_objective_from_vector(objective_vector: Array | ndarray, objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth') → Array

Select one scalar objective from SOLVER_OBJECTIVE_NAMES.

This tiny selector keeps optimizer code honest about which scalar is being minimized. It also centralizes aliases used by the examples: growth -> gamma, frequency -> omega, and quasilinear_flux -> mixing_length_heat_flux_proxy.

VMEC-JAX Transport Objective

Differentiable VMEC-JAX transport objectives and sample reductions.

class gkx.objectives.vmec_transport.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.

class gkx.objectives.vmec_transport.VMEXGKXTransportObjective(config: VMEXTransportObjectiveConfig = <factory>, wout_reference: Any | None = None)

Evaluate a configured transport metric from a solved VMEC-JAX state.

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

Return the scalar transport objective for VMEC-JAX callbacks.

class gkx.objectives.vmec_transport.VMEXTransportObjectiveConfig(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 GKX objective evaluation.

property gradient_scope: str

Return the differentiated part of this objective.

objective_options() → dict[str, Any]

Return GKX solver options for this objective.

gkx.objectives.vmec_transport.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 sample and objective weights; mean is unweighted, while max selects the worst objective-weighted sample.

gkx.objectives.vmec_transport.portfolio_objective_weight_vector(objective_rows: Any, *, objective_weights: Any | None = None) → Array

Return normalized objective-column weights.

gkx.objectives.vmec_transport.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).

gkx.objectives.vmec_transport.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 row/weight shapes and concrete finite non-negative weights.

gkx.objectives.vmec_transport.vmex_transport_objective_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, config: VMEXTransportObjectiveConfig | None = None) → Array

Evaluate a scalar GKX transport objective from a VMEC-JAX state.

Runtime Runner

Public runtime solve facade: one entry point per solve.

The solve owners live in gkx.workflows.linear, gkx.workflows.nonlinear and gkx.workflows.runtime.orchestration_scan; this module only re-exports them and sits above them, so nothing below imports it.

class gkx.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 gkx.runtime.RuntimeLinearResult(ky: float, gamma: float, omega: float, selection: ModeSelection, t: np.ndarray | None = None, signal: np.ndarray | None = None, field_history: np.ndarray | None = None, state: np.ndarray | None = None, z: np.ndarray | None = None, eigenfunction: np.ndarray | None = None, fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, fit_signal_used: str | None = None, gamma_stderr: float | None = None, omega_stderr: float | None = None, fit_r2: float | None = None, fit_settled: bool | None = None, quasilinear: dict[str, Any] | None = None, eigen_status: EigenSolveStatus | None = None, implicit_solve: ImplicitSolveSummary | None = None)

Result container for runtime linear runs.

summary() → dict[str, Any]

Return the typed scalar diagnostics, including fit and solver status.

to_dataset() → dict[str, Any]

Return the trace and eigenfunction as an xarray-compatible mapping.

class gkx.runtime.RuntimeLinearScanResult(ky: ndarray, gamma: ndarray, omega: ndarray, quasilinear: tuple[dict[str, Any], ...] | None = None, parallel: dict[str, Any] | None = None, warm_start: dict[str, Any] | None = None)

Result container for runtime linear ky scans.

summary() → dict[str, Any]

Return scan extent and the peak growth rate over the scanned k_y.

to_dataset() → dict[str, Any]

Return growth and frequency against k_y.

class gkx.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, wall_seconds: float | None = None, saturation: dict[str, Any] | None = None, implicit_solve: ImplicitSolveSummary | None = None)

Result container for runtime nonlinear runs.

summary() → dict[str, Any]

Return scalar diagnostics and the saturation verdict.

saturated is reported as the stop policy recorded it, including False and None. A result must not present an unsaturated window as an accepted value, so the status travels with the number rather than being inferred by a reader.

to_dataset() → dict[str, Any]

Return the sampled time series as an xarray-compatible mapping.

gkx.runtime.build_runtime_geometry(cfg: RuntimeConfig) → SAlphaGeometry | SlabGeometry | FluxTubeGeometryData

Resolve runtime geometry, generating *.eik.nc geometry when needed.

gkx.runtime.build_runtime_linear_params(cfg: RuntimeConfig, *, Nm: int | None = None, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData | None = None) → LinearParams

Build LinearParams from a unified runtime config.

gkx.runtime.build_runtime_linear_terms(cfg: RuntimeConfig, *, custom_collisions: bool = False) → LinearTerms

Build LinearTerms from unified toggles.

All-zero species nu drop the built-in collision term; a caller passing a custom operator with its own rates sets custom_collisions (#324).

gkx.runtime.build_runtime_nonlinear_diagnostics_kwargs(cfg: RuntimeConfig, *, dt: float, steps: int, method: str | None, term_config: Any, sample_stride: int, diagnostics_stride: int, laguerre_mode: str, ky_index: int, kx_index: int, fixed_dt: bool, fixed_mode_ky_index: int | None, fixed_mode_kx_index: int | None, external_phi: float | None, resolved_diagnostics: bool, show_progress: bool) → dict[str, Any]

Build keyword arguments for nonlinear diagnostic integration.

Runtime drivers use the same physics kwargs for fixed-window and adaptive chunked diagnostics. Keeping them in one policy helper prevents drift between the two branches while leaving the actual integrator call in the public runtime facade.

gkx.runtime.build_runtime_term_config(cfg: RuntimeConfig) → TermConfig

Build nonlinear-ready TermConfig from unified toggles.

gkx.runtime.prepare(case: RuntimeConfig, **options: Any) → Any

Prepare a reusable compiled nonlinear simulation; see its class docstring.

gkx.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, window_method: str = 'stationary', krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', return_state: bool = False, initial_state: Any | None = None, show_progress: bool = False, status_callback: Callable[[str], None] | None = None, kx_target: float | None = None) → RuntimeLinearResult

Run one linear point from a case-agnostic runtime config.

kx_target picks the fitted kx (nearest grid value; default kx = 0).

gkx.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, prepare_only: bool = False) → Any

Run a nonlinear point using the unified runtime config path.

prepare_only returns the compiled prepared simulation instead (the prepare() path).

gkx.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, window_method: str = 'stationary', krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', show_progress: bool = False, workers: int = 1, parallel_executor: str = 'thread', warm_start: bool | None = None) → RuntimeLinearScanResult

Run a ky scan; warm_start overrides the case output policy.

Coordinates serial, independent-worker, or combined-ky runtime scans.

gkx.runtime.solve(case: RuntimeConfig, **options: Any) → RuntimeLinearResult | RuntimeNonlinearResult

Solve one case through its existing linear or nonlinear runtime owner.