API Reference

Public API Registry

The top-level spectraxgk package remains the stable user-facing facade. Implementation ownership lives in the domain modules below; the compact API registry lazily re-exports promoted symbols without keeping one re-export file per domain. Advanced users should import from the owning domain modules when they need implementation-specific extension points.

Lazy public API registry for SPECTRAX-GK.

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.

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

Continuity summary for a scanned linear branch.

class spectraxgk.api.CycloneBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=24, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=20.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=150.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=True, diffrax_rtol=1e-06, diffrax_atol=1e-08, diffrax_max_steps=200000, state_sharding=None, progress_bar=False, fixed_dt=False, dt_min=1e-07, dt_max=0.05, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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(R_over_LTi=2.49, R_over_LTe=0.0, R_over_Ln=0.8, nu_i=0.0), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False), reference_aligned: bool = True)

Standard parameters for the Cyclone base case ITG benchmark.

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

Claim-scoped shard assignment contract for a parallelization path.

property diagnostic_nonlinear_partition: bool

Whether this contract is diagnostic nonlinear-domain metadata.

property production_independent_batching: bool

Whether this contract is for production independent-work batching.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the contract.

class spectraxgk.api.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.

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

Collection of scalar gates for one validation artifact.

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

Spectral grid configuration in a flux-tube.

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

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

to_dict() → dict[str, Any]

Return a JSON-serializable provenance payload.

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

Worker failure annotated with independent-map execution metadata.

class spectraxgk.api.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.

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

Gyroaveraging coefficients matching the Laguerre-Hermite quadrature convention.

class spectraxgk.api.KBMBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=16, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=10.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Tsit5', diffrax_adaptive=True, diffrax_rtol=0.0001, diffrax_atol=1e-07, diffrax_max_steps=20000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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(R_over_LTi=2.49, R_over_LTe=2.49, R_over_Ln=0.8, Te_over_Ti=1.0, mass_ratio=3703.7037037037035, nu_i=0.0, nu_e=0.0, beta=0.015), init: InitializationConfig = InitializationConfig(init_field='all', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for an electromagnetic KBM benchmark.

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

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

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

Precomputed arrays for the linear operator.

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

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

class spectraxgk.api.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 spectraxgk.api.LinearValidationPanel(name: 'str', z: 'np.ndarray', eigenfunction: 'np.ndarray', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', x_label: 'str', x_ref: 'np.ndarray | None' = None, gamma_ref: 'np.ndarray | None' = None, omega_ref: 'np.ndarray | None' = None, ref_label: 'str' = 'Reference', log_x: 'bool' = False)

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

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

class spectraxgk.api.NonlinearDomainDecompositionPlan(state_shape: tuple[int, ...], axis: int, chunk_sizes: tuple[int, ...], halo: int = 1)

Static decomposition plan for a local nonlinear state-domain local-stencil diagnostic.

property boundary_indices: tuple[int, ...]

Return global cells that touch a decomposed halo interface.

property chunk_bounds: tuple[tuple[int, int], ...]

Return half-open (start, stop) bounds for owned chunk cells.

decomposition_metadata() → dict[str, Any]

Return derived metadata for diagnostic decomposition artifacts.

property domain_size: int

Return the global size of the decomposed axis.

property num_domains: int

Return the number of state-domain chunks.

property offsets: tuple[int, ...]

Return chunk start offsets along the decomposed axis.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the decomposition plan.

class spectraxgk.api.NonlinearDomainIdentityReport(gate_name: str, plan: NonlinearDomainDecompositionPlan, atol: float, rtol: float, max_abs_error: float, max_rel_error: float, plan_valid: bool, blocked_reasons: tuple[str, ...], identity_passed: bool, decomposed_path_enabled: bool, claim_scope: str, boundary_max_abs_error: float = 0.0, boundary_max_rel_error: float = 0.0, boundary_indices: tuple[int, ...] = ())

Numerical identity report for a decomposed nonlinear local-stencil step.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the identity report.

class spectraxgk.api.NonlinearDomainTransportWindowReport(gate_name: str, plan: NonlinearDomainDecompositionPlan, steps: int, dt: float, atol: float, rtol: float, max_abs_state_error: float, max_rel_state_error: float, max_abs_boundary_error: float, max_rel_boundary_error: float, mass_trace_max_abs_error: float, mass_trace_max_rel_error: float, free_energy_trace_max_abs_error: float, free_energy_trace_max_rel_error: float, flux_proxy_trace_max_abs_error: float, flux_proxy_trace_max_rel_error: float, serial_mass_drift: float, decomposed_mass_drift: float, serial_free_energy_drift: float, decomposed_free_energy_drift: float, plan_valid: bool, blocked_reasons: tuple[str, ...], identity_passed: bool, decomposed_path_enabled: bool, claim_scope: str, boundary_indices: tuple[int, ...] = (), serial_mass_trace: tuple[float, ...] = (), decomposed_mass_trace: tuple[float, ...] = (), serial_free_energy_trace: tuple[float, ...] = (), decomposed_free_energy_trace: tuple[float, ...] = (), serial_flux_proxy_trace: tuple[float, ...] = (), decomposed_flux_proxy_trace: tuple[float, ...] = ())

Transport-window identity report for the nonlinear domain local-stencil diagnostic.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the transport-window report.

class spectraxgk.api.NonlinearParallelStrategy(name: Literal['independent_ky_scan', 'uq_ensemble', 'whole_state_kx_ky', 'velocity_species_hermite', 'fft_axis_domain'], readiness: Literal['release_ready', 'diagnostic', 'blocked'], independent_work: bool, changes_solver_layout: bool, identity_gates: tuple[str, ...], physics_gates: tuple[str, ...], profiler_gates: tuple[str, ...], notes: str)

Readiness contract for one nonlinear parallelization candidate.

property blocked: bool

Whether this strategy is unavailable until required gates exist.

property diagnostic_only: bool

Whether this strategy is limited to correctness/profiling artifacts.

property release_ready: bool

Whether this strategy is allowed for production-facing execution.

property required_gates: tuple[str, ...]

All identity, physics, and profiler gates required by this policy.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the contract.

class spectraxgk.api.NonlinearSpectralCommunicationReport(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], atol: float, rtol: float, fft_max_abs_error: float, fft_max_rel_error: float, bracket_max_abs_error: float, bracket_max_rel_error: float, field_max_abs_error: float, field_max_rel_error: float, identity_passed: bool, decomposed_path_enabled: bool, claim_scope: str, blocked_reasons: tuple[str, ...] = (), y_offsets: tuple[int, ...] = (), x_offsets: tuple[int, ...] = ())

Numerical identity report for nonlinear spectral communication layouts.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the communication report.

class spectraxgk.api.NonlinearSpectralDevicePencilFFTBatchModel(state_shape: tuple[int, int, int, int, int], device_count: int, local_z_extent: int, max_fft_axis_extent: int, max_fft_batch_count: int, unchunked_fft_batch_count: int, suggested_z_chunk_size: int | None, effective_z_chunk_size: int, chunked_fft_batch_count: int, chunking_required: bool, chunking_active: bool, disable_gpu_preallocation_recommended: bool, profiling_candidate: bool, feasibility_blockers: tuple[str, ...], claim_scope: str)

cuFFT batch-pressure preflight model for the device-z pencil route.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the batch model.

class spectraxgk.api.NonlinearSpectralDevicePencilRHSIdentityReport(state_shape: tuple[int, int, int, int, int], sharded_axis: str, axis_name: str, requested_device_count: int, active_device_count: int, atol: float, rtol: float, rhs_max_abs_error: float, rhs_max_rel_error: float, identity_passed: bool, device_sharding_active: bool, decomposed_path_enabled: bool, claim_scope: str, blocked_reasons: tuple[str, ...] = ())

Identity report for a device-sharded fused pencil nonlinear RHS.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the identity report.

class spectraxgk.api.NonlinearSpectralDevicePencilTransportWindowReport(state_shape: tuple[int, int, int, int, int], sharded_axis: str, axis_name: str, requested_device_count: int, active_device_count: int, steps: int, dt: float, atol: float, rtol: float, final_state_max_abs_error: float, final_state_max_rel_error: float, free_energy_trace_max_abs_error: float, free_energy_trace_max_rel_error: float, field_energy_trace_max_abs_error: float, field_energy_trace_max_rel_error: float, physical_flux_trace_max_abs_error: float, physical_flux_trace_max_rel_error: float, bracket_rms_trace_max_abs_error: float, bracket_rms_trace_max_rel_error: float, serial_free_energy_drift: float, device_free_energy_drift: float, identity_passed: bool, device_sharding_active: bool, decomposed_path_enabled: bool, claim_scope: str, blocked_reasons: tuple[str, ...] = (), serial_free_energy_trace: tuple[float, ...] = (), device_free_energy_trace: tuple[float, ...] = (), serial_field_energy_trace: tuple[float, ...] = (), device_field_energy_trace: tuple[float, ...] = (), serial_physical_flux_trace: tuple[float, ...] = (), device_physical_flux_trace: tuple[float, ...] = (), serial_bracket_rms_trace: tuple[float, ...] = (), device_bracket_rms_trace: tuple[float, ...] = ())

Multi-step identity report for device-z-sharded pencil routing.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the transport report.

class spectraxgk.api.NonlinearSpectralDomainWorkModel(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], y_offsets: tuple[int, ...], x_offsets: tuple[int, ...], tile_bounds: tuple[tuple[int, int, int, int], ...], num_tiles: int, state_elements: int, field_elements: int, owned_state_elements_per_step: int, state_allgather_elements_per_step: int, bracket_allgather_elements_per_step: int, field_broadcast_elements_per_step: int, total_communication_elements_per_step: int, communication_to_owned_work_ratio: float, parallel_efficiency_ceiling: float, max_communication_to_owned_work_ratio: float, production_speedup_feasible: bool, feasibility_blockers: tuple[str, ...], claim_scope: str)

Communication/work model for the logical nonlinear spectral-domain route.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the work model.

class spectraxgk.api.NonlinearSpectralIntegratorIdentityReport(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], y_offsets: tuple[int, ...], x_offsets: tuple[int, ...], tile_bounds: tuple[tuple[int, int, int, int], ...], steps: int, dt: float, atol: float, rtol: float, final_state_max_abs_error: float, final_state_max_rel_error: float, free_energy_trace_max_abs_error: float, free_energy_trace_max_rel_error: float, field_energy_trace_max_abs_error: float, field_energy_trace_max_rel_error: float, flux_proxy_trace_max_abs_error: float, flux_proxy_trace_max_rel_error: float, serial_free_energy_drift: float, logical_free_energy_drift: float, identity_passed: bool, decomposed_path_enabled: bool, claim_scope: str, blocked_reasons: tuple[str, ...] = (), serial_free_energy_trace: tuple[float, ...] = (), logical_free_energy_trace: tuple[float, ...] = (), serial_field_energy_trace: tuple[float, ...] = (), logical_field_energy_trace: tuple[float, ...] = (), serial_flux_proxy_trace: tuple[float, ...] = (), logical_flux_proxy_trace: tuple[float, ...] = ())

Multi-step identity report for logical-shard nonlinear spectral routing.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the integrator report.

class spectraxgk.api.NonlinearSpectralPencilRHSIdentityReport(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], y_offsets: tuple[int, ...], x_offsets: tuple[int, ...], atol: float, rtol: float, field_max_abs_error: float, field_max_rel_error: float, bracket_max_abs_error: float, bracket_max_rel_error: float, rhs_max_abs_error: float, rhs_max_rel_error: float, identity_passed: bool, decomposed_path_enabled: bool, work_model: NonlinearSpectralPencilWorkModel, claim_scope: str, blocked_reasons: tuple[str, ...] = ())

Numerical identity report for the pencil-FFT nonlinear spectral RHS.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the RHS identity report.

class spectraxgk.api.NonlinearSpectralPencilTransportWindowReport(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], y_offsets: tuple[int, ...], x_offsets: tuple[int, ...], steps: int, dt: float, atol: float, rtol: float, final_state_max_abs_error: float, final_state_max_rel_error: float, free_energy_trace_max_abs_error: float, free_energy_trace_max_rel_error: float, field_energy_trace_max_abs_error: float, field_energy_trace_max_rel_error: float, physical_flux_trace_max_abs_error: float, physical_flux_trace_max_rel_error: float, bracket_rms_trace_max_abs_error: float, bracket_rms_trace_max_rel_error: float, serial_free_energy_drift: float, pencil_free_energy_drift: float, identity_passed: bool, decomposed_path_enabled: bool, work_model: NonlinearSpectralPencilWorkModel, claim_scope: str, blocked_reasons: tuple[str, ...] = (), serial_free_energy_trace: tuple[float, ...] = (), pencil_free_energy_trace: tuple[float, ...] = (), serial_field_energy_trace: tuple[float, ...] = (), pencil_field_energy_trace: tuple[float, ...] = (), serial_physical_flux_trace: tuple[float, ...] = (), pencil_physical_flux_trace: tuple[float, ...] = (), serial_bracket_rms_trace: tuple[float, ...] = (), pencil_bracket_rms_trace: tuple[float, ...] = ())

Multi-step transport-window identity report for the pencil route.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the transport report.

class spectraxgk.api.NonlinearSpectralPencilWorkModel(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], y_offsets: tuple[int, ...], x_offsets: tuple[int, ...], num_tiles: int, state_elements: int, field_elements: int, transform_payload_elements_per_step: int, pencil_transpose_elements_per_step: int, global_reconstruction_elements_per_step: int, approximate_fft_work_units_per_step: float, communication_to_fft_work_ratio: float, parallel_efficiency_ceiling: float, predicted_speedup_ceiling: float, max_communication_to_fft_work_ratio: float, min_predicted_speedup: float, production_speedup_feasible: bool, feasibility_blockers: tuple[str, ...], claim_scope: str)

Communication/work model for a pencil-FFT nonlinear bracket route.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the pencil work model.

class spectraxgk.api.NonlinearSpectralRHSIdentityReport(state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], y_offsets: tuple[int, ...], x_offsets: tuple[int, ...], tile_bounds: tuple[tuple[int, int, int, int], ...], atol: float, rtol: float, reconstruction_max_abs_error: float, reconstruction_max_rel_error: float, field_max_abs_error: float, field_max_rel_error: float, bracket_max_abs_error: float, bracket_max_rel_error: float, rhs_max_abs_error: float, rhs_max_rel_error: float, identity_passed: bool, decomposed_path_enabled: bool, claim_scope: str, blocked_reasons: tuple[str, ...] = ())

Numerical identity report for logical-shard nonlinear spectral RHS.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the RHS identity report.

class spectraxgk.api.NonlinearWindowConvergenceConfig(tmin: float | None = None, tmax: float | None = None, transient_fraction: float = 0.5, min_samples: int = 24, min_blocks: int = 4, block_size: int | None = None, bootstrap_samples: int = 256, bootstrap_seed: int = 0, max_running_mean_rel_drift: float = 0.15, terminal_fraction: float = 0.25, min_terminal_samples: int = 8, max_terminal_mean_rel_delta: float = 0.1, max_sem_rel: float = 0.25, value_floor: float = 1e-12, require_all_finite: bool = True)

Gate settings for a nonlinear post-transient transport window.

class spectraxgk.api.NonlinearWindowEnsembleConfig(min_reports: int = 2, max_mean_rel_spread: float = 0.15, max_combined_sem_rel: float = 0.25, value_floor: float = 1e-12, require_individual_passed: bool = True)

Gate settings for replicated nonlinear transport-window summaries.

class spectraxgk.api.NonlinearWindowEnsembleManifestConfig(min_replicates_per_case: int = 2, required_variant_axes: tuple[str, ...] = ('seed', 'timestep'), require_observed_windows_ready: bool = True)

Artifact requirements before a replicated nonlinear ensemble can run.

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

Windowed transport/envelope metrics for a nonlinear run.

class spectraxgk.api.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’]

class spectraxgk.api.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 spectraxgk.api.ProductionNonlinearOptimizationGuardConfig(min_replicated_ensembles: int = 2, min_reports_per_ensemble: int = 2, max_mean_rel_spread: float = 0.15, max_combined_sem_rel: float = 0.25, require_optimized_equilibrium_transport: bool = True, require_matched_optimized_transport_audit: bool = True, min_optimized_equilibrium_ensembles: int = 3, min_matched_optimized_audits: int = 3, require_seed_timestep_provenance: bool = True, min_seed_variants: int = 2, min_timestep_variants: int = 1, min_matched_optimized_relative_reduction: float = 0.05, min_matched_optimized_uncertainty_sigma: float = 1.0, value_floor: float = 1e-12)

Strict gate settings for production nonlinear optimization promotion.

validate() → None

Raise if the guard configuration is inconsistent.

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

Selection policy for projected transport line-search candidates.

to_dict() → dict[str, Any]

Return a JSON-safe policy representation.

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

Configuration for the reduced QA low-turbulence comparison.

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

JSON-ready result for one reduced QA optimization.

to_dict() → dict[str, Any]

Return a stable JSON-friendly payload.

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

One quasilinear-vs-nonlinear transport comparison point.

class spectraxgk.api.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.

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

Serial reconstruction identity report for a decomposition contract.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the report.

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

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

to_dict() → dict[str, object]

Return a JSON-friendly representation.

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

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

Collision and end-damping parameters.

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

Unified simulation config for runtime-driven GK runs.

class spectraxgk.api.RuntimeLinearResult(ky: float, gamma: float, omega: float, selection: ModeSelection, t: ndarray | None = None, signal: ndarray | None = None, field_history: ndarray | None = None, state: ndarray | None = None, z: ndarray | None = None, eigenfunction: ndarray | None = None, fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, fit_signal_used: str | None = None, quasilinear: dict[str, Any] | None = None)

Result container for runtime linear runs.

class spectraxgk.api.RuntimeLinearScanResult(ky: ndarray, gamma: ndarray, omega: ndarray, quasilinear: tuple[dict[str, Any], ...] | None = None, parallel: dict[str, Any] | None = None)

Result container for runtime linear ky scans.

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

Normalization contract selection + optional explicit overrides.

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

Artifact-output controls for runtime executable entry points.

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

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

class spectraxgk.api.RuntimeParameterScanResult(parameter_name: str, values: ndarray, gamma: ndarray, omega: ndarray, runs: tuple[RuntimeLinearResult, ...])

Ordered linear results for a named scalar configuration parameter.

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

Physics-family toggles independent from benchmark case names.

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

Quasilinear transport diagnostics computed from linear states.

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

Single species definition for runtime-configured simulations.

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

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

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

Simple s-alpha geometry with circular concentric flux surfaces.

bgrad(theta: Array) → Array

Magnetic field gradient term used in mirror force.

bmag(theta: Array) → Array

Magnetic field strength for circular s-alpha geometry.

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

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

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

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

gradpar() → float

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

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

Perpendicular wave-number squared for s-alpha geometry.

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

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

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

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

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

Magnetic drift frequency for s-alpha geometry.

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

Pass/fail result for one benchmark observable.

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

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

Linear scan plus a representative fitted eigenfunction.

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

A deterministic contiguous assignment of serial indices to one shard.

property size: int

Number of serial items assigned to this shard.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the assignment.

class spectraxgk.api.SimulationDiagnostics(t: Array | ndarray, dt_t: Array | ndarray, dt_mean: Array | ndarray, gamma_t: Array | ndarray, omega_t: Array | ndarray, Wg_t: Array | ndarray, Wphi_t: Array | ndarray, Wapar_t: Array | ndarray, heat_flux_t: Array | ndarray, particle_flux_t: Array | ndarray, energy_t: Array | ndarray, heat_flux_species_t: Array | ndarray | None = None, particle_flux_species_t: Array | ndarray | None = None, turbulent_heating_t: Array | ndarray | None = None, turbulent_heating_species_t: Array | ndarray | None = None, phi_mode_t: Array | ndarray | None = None, resolved: ResolvedDiagnostics | None = None)

Streaming diagnostics at each sample time.

class spectraxgk.api.Species(charge: float, mass: float, density: float, temperature: float, tprim: float, fprim: float, nu: float = 0.0)

Physical parameters for a kinetic species.

class spectraxgk.api.StellaratorITGOptimizationConfig(target_aspect: float = 7.0, target_iota: float = 0.41, max_mode: int = 1, aspect_weight: float = 0.25, iota_weight: float = 25.0, qa_weight: float = 5.0, turbulence_weight: float = 1.0, regularization: float = 0.002, learning_rate: float = 0.035, steps: int = 90, nonlinear_dt: float = 0.18, nonlinear_steps: int = 520, nonlinear_tail_fraction: float = 0.25, quasilinear_csat: float = 0.75, reference_density_gradient: float = 2.2, reference_temperature_gradient: float = 6.0, scan_density_gradients: tuple[float, ...] = (0.8, 1.2, 1.6, 2.2, 3.0, 3.8, 4.8), fd_step: float = 0.0001)

Configuration for the QA max-mode-1 ITG optimization examples.

with_kind_defaults(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']) → StellaratorITGOptimizationConfig

Return conservative optimizer defaults for one objective family.

class spectraxgk.api.StellaratorITGOptimizationResult(objective_kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], parameter_names: tuple[str, ...], observable_names: tuple[str, ...], initial_params: tuple[float, ...], final_params: tuple[float, ...], initial_objective: float, final_objective: float, initial_observables: tuple[float, ...], final_observables: tuple[float, ...], history: tuple[dict[str, Any], ...], gradient_gate: dict[str, Any], covariance: dict[str, Any], nonlinear_trace: dict[str, Any] | None, config: dict[str, Any], backend_info: dict[str, Any])

JSON-friendly result for one differentiable stellarator objective.

to_dict() → dict[str, Any]

Return a stable JSON-serializable representation.

class spectraxgk.api.StellaratorITGSampleSet(surfaces: tuple[float, ...] = (0.5, 0.64, 0.78), alphas: tuple[float, ...] = (0.0, 1.0471975511965976), ky_values: tuple[float, ...] = (0.1, 0.3, 0.5), surface_weights: tuple[float, ...] | None = None, alpha_weights: tuple[float, ...] | None = None, ky_weights: tuple[float, ...] | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean')

Reduced multi-surface/multi-alpha/multi-k_y ITG portfolio contract.

property n_samples: int

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

to_dict() → dict[str, Any]

Return a JSON-friendly representation.

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

Time integration parameters.

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

Policy for promoting or redesigning VMEC-JAX transport candidates.

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

to_dict() → dict[str, Any]

Return a JSON-safe representation.

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

Admission limits for launching the next nonlinear optimizer campaign.

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

to_dict() → dict[str, Any]

Return a JSON-safe representation.

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

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

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.api.VMECJAXSpectraxTransportObjective(config: VMECJAXTransportObjectiveConfig = <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 spectraxgk.api.VMECJAXTransportAdmissionPolicy(metric_keys: tuple[str, ...] = ('transport_objective_final', 'spectrax_objective_final', 'transport_metric_final', 'objective_final'), minimum_relative_improvement: float = 0.0, lower_is_better: bool = True, require_authoritative_gate: bool = True, allow_baseline_fallback: bool = True)

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

to_dict() → dict[str, Any]

Return a JSON-safe representation.

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

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

property gradient_scope: str

Return the differentiated part of this objective.

objective_options() → dict[str, Any]

Return SPECTRAX-GK solver options for this objective.

class spectraxgk.api.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.

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

Weights and floors for a minimizable zonal-flow objective.

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

objective_weights() → Array

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

to_dict() → dict[str, object]

Return a JSON-friendly representation.

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

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

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

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

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

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

Apply reporting-space normalization to growth rates/frequencies.

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

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

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

Compare JAX forward-mode derivatives against finite differences.

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

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

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

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

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

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

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

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

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

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

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

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

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

Validate a real booz_xform_jax spectral derivative when available.

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

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

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

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

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

Build a chain summary from a probe JSON payload.

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

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

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

Gate branch-continuation diagnostics for branch-followed scans.

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

Compute branch-continuity diagnostics for a linear scan.

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

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

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

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

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

Summarize several boundary-chain probes as one promotion gate.

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

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

Classify a boundary-gradient chain probe.

Parameters:

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

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

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

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

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

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

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

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

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

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

Evaluate transport residual and boundary-gradient diagnostics.

Parameters:

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

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

• label – Human-readable artifact label.

• top_n – Number of largest gradient components to keep.

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

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

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

Build a diagnostic nonlinear-domain partition contract.

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

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

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

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

Build a production independent-work decomposition contract.

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

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

Build reusable arrays for the linear RHS.

spectraxgk.api.build_linear_params(species: Iterable[Species], *, tau_e: float = 0.0, kpar_scale: float = 1.0, omega_d_scale: float = 1.0, omega_star_scale: float = 1.0, rho_star: float = 1.0, beta: float = 0.0, fapar: float = 0.0, apar_beta_scale: float = 0.5, ampere_g0_scale: float = 0.5, bpar_beta_scale: float = 0.5, nu_hyper: float = 0.0, p_hyper: float = 4.0, nu_hyper_l: float = 0.0, nu_hyper_m: float = 1.0, nu_hyper_lm: float = 0.0, p_hyper_l: float = 6.0, p_hyper_m: float = 20.0, p_hyper_lm: float = 6.0, hypercollisions_const: float = 1.0, hypercollisions_kz: float = 0.0, D_hyper: float = 0.0, p_hyper_kperp: float = 2.0) → LinearParams

Build LinearParams arrays from a list of species.

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

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

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

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

Gate the next nonlinear optimizer campaign from existing evidence.

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

spectraxgk.api.build_nonlinear_domain_decomposition_plan(state_shape: tuple[int, ...], *, axis: int = 0, num_domains: int = 2, halo: int = 1) → NonlinearDomainDecompositionPlan

Build a static chunk plan for the local state-domain identity prototype.

spectraxgk.api.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.

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

Select an uncertainty-separated nonlinear candidate from a landscape.

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

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

Combine quasilinear model-selection gates into one claim ledger.

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

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

Path-based wrapper for artifact scripts and CI checks.

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

Gate reduced transport candidates before launching nonlinear audits.

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

spectraxgk.api.build_runtime_linear_params(cfg: RuntimeConfig, *, Nm: int | None = None, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData | None = None) → LinearParams

Build runtime linear parameters using the runtime module geometry surface.

spectraxgk.api.build_runtime_linear_terms(cfg: RuntimeConfig) → LinearTerms

Build runtime linear term toggles.

spectraxgk.api.build_runtime_term_config(cfg: RuntimeConfig) → TermConfig

Build runtime nonlinear-ready term config.

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

Build a JSON-safe solved-equilibrium gate report.

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

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

Annotate and select VMEC-JAX transport candidates.

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

spectraxgk.api.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.

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

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

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

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

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

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

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

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

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

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

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

Central finite-difference Jacobian for small differentiability gates.

spectraxgk.api.classify_nonlinear_parallel_strategy(name: Literal['independent_ky_scan', 'uq_ensemble', 'whole_state_kx_ky', 'velocity_species_hermite', 'fft_axis_domain']) → Literal['release_ready', 'diagnostic', 'blocked']

Return the release-readiness classification for a strategy.

spectraxgk.api.compare_stellarator_itg_objectives(kinds: Sequence[Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']] = ('growth', 'quasilinear_flux', 'nonlinear_heat_flux'), *, initial_params: Array | Sequence[float] | None = None, config: StellaratorITGOptimizationConfig | None = None, workers: int = 1, parallel_executor: str = 'thread', finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Run the three objective reductions from a shared starting point.

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

Compute quasilinear transport weights from a linear state.

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

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

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

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

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

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

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

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

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

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

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

Create a two-panel Cyclone base case reference plot.

spectraxgk.api.default_qa_low_turbulence_initial_params() → Array

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

spectraxgk.api.default_solver_geometry_design_params() → Array

Return the small geometry-design vector used by the release gate.

spectraxgk.api.default_stellarator_initial_params() → Array

Return the shared off-optimum QA max-mode-1 starting point.

spectraxgk.api.deterministic_nonlinear_domain_state(shape: tuple[int, ...] = (6, 4)) → Array

Return a small deterministic complex state for identity gates.

spectraxgk.api.deterministic_nonlinear_spectral_state(shape: tuple[int, int, int, int, int] = (2, 3, 6, 4, 2)) → Array

Return deterministic complex spectral coefficients for communication gates.

The layout is (Nl, Nm, Ny, Nx, Nz) with the FFT axes in (Ny, Nx).

spectraxgk.api.device_z_pencil_fft_batch_pressure_model(state_shape: tuple[int, int, int, int, int], *, device_count: int, max_fft_batch_count: int = 65536, z_chunk_size: int | None = None) → NonlinearSpectralDevicePencilFFTBatchModel

Estimate cuFFT batch pressure for the device-z pencil micro-route.

The profiler traces showed that large GPU cases can fail before timing when axis-wise FFTs create too-large batched cuFFT plans. This backend-free model predicts that pressure and suggests a local z_chunk_size that keeps the largest state-gradient axis FFT batch below max_fft_batch_count. It is a profiling preflight, not a speedup or physics claim.

spectraxgk.api.device_z_pencil_nonlinear_spectral_rhs(state_hat: Array, *, devices: Sequence[Any] | None = None, axis_name: str = 'z', z_chunk_size: int | None = None, atol: float = 5e-06, rtol: float = 0.0001) → tuple[Array, NonlinearSpectralDevicePencilRHSIdentityReport]

Return the z-sharded fused pencil nonlinear RHS after identity gating.

This is the first real device-sharded nonlinear spectral route in this module. It shards over the field-line z axis so that the FFT axes remain local on every device. The function falls back to the serial RHS unless a multi-device sharding exists and the sharded fused-bracket RHS matches the serial reference within the requested tolerances.

spectraxgk.api.device_z_pencil_nonlinear_spectral_transport_window_identity_gate(state_hat: Array, *, devices: Sequence[Any] | None = None, axis_name: str = 'z', z_chunk_size: int | None = None, dt: float = 0.005, steps: int = 4, atol: float = 5e-06, rtol: float = 0.0001, observable_mode: Literal['host_gather', 'sharded_reduce'] = 'host_gather') → NonlinearSpectralDevicePencilTransportWindowReport

Validate a multi-step serial-vs-device-z-sharded nonlinear window.

The route advances the same explicit fixed-step micro-window with the serial nonlinear RHS and the shard-map z-pencil RHS. It compares the final state and physical-space scalar traces, including a density-times-radial electric-field flux proxy and bracket RMS. Passing this gate is still not a turbulent heat-flux validation; it only permits timing the decomposed device route on the same deterministic operator.

spectraxgk.api.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.

spectraxgk.api.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.

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

Discover optional vmec_jax and booz_xform_jax bridge APIs.

spectraxgk.api.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 = 'power', power_iters: int = 40, power_dt: float = 0.01, 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 = 'damping', shift_selection: str = 'targeted', shift_outer_residual_tol: float = 0.1, mode_family: str = 'auto', fallback_method: str = 'propagator', fallback_real_floor: float = -1e-06, status_callback: Callable[[str], None] | None = None) → tuple[Array, Array]

Python wrapper for the cached Krylov solver.

spectraxgk.api.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.

spectraxgk.api.dominant_real_eigenvalue(matrix: Array) → Array

Return the dominant growth rate with an implicit left/right VJP.

This helper treats the max-real eigenvalue branch selected at the primal point as locally isolated. Its reverse rule uses d lambda = w^H dA v with w^H v = 1 instead of differentiating through non-Hermitian eigenvectors. Branch isolation is still a physics gate: callers that use this in optimization should keep finite-difference or branch-continuity checks enabled near accepted candidates.

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

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

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

Gate a phase-aligned eigenfunction comparison.

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

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

Return electrostatic phi from a full single- or multi-species state.

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

Solve electrostatic phi using a species- or Hermite-sharded reduction.

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

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

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

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

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

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

Evaluate one scalar benchmark gate.

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

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

Materialize a small complex linear operator as a dense matrix.

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

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

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

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

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

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

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

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

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

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

Central finite-difference Jacobian for small validation problems.

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

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

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

Fit gamma/omega with optional auto-selected window.

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

Fit one multiplicative heat-flux scale from training points.

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

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

Build FluxTubeGeometryData from an in-memory differentiable backend.

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

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

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

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

VMECState -> BoozXformInputs -> booz_xform_jax -> FluxTubeGeometryData.

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

spectraxgk.api.flux_tube_geometry_observables(geom: FluxTubeGeometryData) → Array

Return differentiable scalar observables from solver-ready geometry.

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

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

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

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

Summarize a set of scalar gates for one artifact.

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

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

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

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

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

spectraxgk.api.geometry_observable_names() → tuple[str, ...]

Return the ordered geometry observables used by bridge AD checks.

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

Validate geometry-observable sensitivities by AD and finite differences.

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

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

Return the canonical normalization contract for case.

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

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

spectraxgk.api.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).

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

Render a growth-rate heatmap versus two gradient axes.

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

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

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

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

Exchange nearest Hermite neighbors with jax.shard_map.

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

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

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

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

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

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

Parallel streaming operator acting on the Hermite index.

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

Return the full-array Hermite streaming ladder contribution.

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

Return a shard-map Hermite streaming ladder contribution.

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

Validate implicit sensitivities of an isolated non-Hermitian eigenpair.

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

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

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

Verify independent UQ/optimization ensemble batching provenance.

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

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

Map independent Python tasks while preserving serial result ordering.

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

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

Compare independent_map against a serial list-comprehension run.

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

Resolve independent-task worker counts and normalized executor metadata.

spectraxgk.api.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).

spectraxgk.api.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) → tuple[Array, Array]

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

spectraxgk.api.integrate_linear_diffrax(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'Dopri8', cache: LinearCache | None = None, terms: LinearTerms | None = None, adaptive: bool = False, rtol: float = 1e-05, atol: float = 1e-07, max_steps: int = 4096, show_progress: bool = False, progress_bar: bool = False, checkpoint: bool = False, jit: bool | None = None, sample_stride: int = 1, return_state: bool = True, save_mode: ModeSelection | ModeSelectionBatch | None = None, mode_method: str = 'z_index', save_field: str = 'phi', density_species_index: int | None = None, state_sharding: Any | None = None, derivative_mode: str = 'reverse') → tuple[Array | None, Array]

Integrate the linear system with an explicit AD policy.

derivative_mode="reverse" preserves the custom-VJP field solve used by scalar objectives. For adaptive reverse gradients, set checkpoint=True to select Diffrax’s bounded-memory recursive checkpoint adjoint. "forward" selects native JAX rules so low-dimensional JVPs can pass through fixed or adaptive Diffrax trajectories.

spectraxgk.api.integrate_linear_diffrax_streaming(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'Dopri8', cache: LinearCache | None = None, terms: LinearTerms | None = None, adaptive: bool = False, rtol: float = 1e-05, atol: float = 1e-07, max_steps: int = 4096, show_progress: bool = False, progress_bar: bool = False, checkpoint: bool = False, jit: bool | None = None, tmin: float | None = None, tmax: float | None = None, fit_signal: str = 'density', mode_ky_indices: Sequence[int] | ndarray | Array | None = None, mode_kx_index: int = 0, mode_z_index: int = 0, mode_method: str = 'z_index', amp_floor: float = 1e-30, density_species_index: int | None = None, return_state: bool = True, state_sharding: Any | None = None) → tuple[Array | None, Array, Array]

Integrate the linear system and stream a growth-rate fit without storing time series.

spectraxgk.api.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 with growth-rate diagnostics.

spectraxgk.api.integrate_linear_explicit_diagnostics(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, SimulationDiagnostics]

Public facade for diagnostics-rich explicit linear integration.

spectraxgk.api.integrate_linear_from_config(G0, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, time_cfg: TimeConfig, *, cache: LinearCache | None = None, terms: LinearTerms | None = None, save_mode: ModeSelection | ModeSelectionBatch | None = None, mode_method: str = 'z_index', save_field: str = 'phi', density_species_index: int | None = None, show_progress: bool | None = None, parallel: Any | None = None) → tuple

Integrate the linear system using TimeConfig settings.

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

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

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

spectraxgk.api.integrate_logical_decomposed_nonlinear_spectral(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), dt: float = 0.005, steps: int = 4, atol: float = 5e-06, rtol: float = 5e-06) → tuple[Array, NonlinearSpectralIntegratorIdentityReport]

Integrate with the logical decomposed spectral route after identity gating.

This is the callable route behind the nonlinear spectral integrator identity artifact. Each step requests the logical tiled RHS; if the local RHS gate fails, that step uses the serial RHS instead. The returned report is the independent serial-vs-routed identity gate for the full fixed-step window. Passing this function’s report is required before timing the route, but it is still not a production distributed-FFT implementation or speedup claim.

spectraxgk.api.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) → tuple[Array, FieldState] | Array

Integrate the nonlinear system using built-in cache construction.

spectraxgk.api.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) → tuple[Array, FieldState] | Array

Integrate the nonlinear system using a cached geometry object.

spectraxgk.api.integrate_nonlinear_diffrax(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'KenCarp4', cache: LinearCache | None = None, terms: TermConfig | None = None, adaptive: bool = False, rtol: float = 1e-05, atol: float = 1e-07, max_steps: int = 4096, show_progress: bool = False, progress_bar: bool = False, checkpoint: bool = False, jit: bool | None = None, state_sharding: Any | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid') → tuple[Array, FieldState]

Integrate the nonlinear system with diffrax.

spectraxgk.api.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, 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]

Integrate nonlinear system and return runtime diagnostics.

spectraxgk.api.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) → tuple

Integrate the nonlinear system using TimeConfig settings.

spectraxgk.api.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. This is a diagnostic whole-state sharding primitive for identity gates and profiler localization. It is not a production nonlinear domain decomposition or speedup claim until the exact workload has communication-complete identity, conservation, transport, and profiler gates. Domain-sharding identity reports are metadata gates only; they do not authorize routing through this whole-state integrator.

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

Integrate one quasilinear spectrum column into a scalar flux estimate.

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

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

Validate AD sensitivities of an observable of one isolated eigenpair.

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

spectraxgk.api.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.

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

Return balanced k_y chunks for independent linear-scan execution.

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

Gate late-time linear growth and frequency metrics.

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

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

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

Compute the linear RHS using precomputed geometry/cache arrays.

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

Compute gated electrostatic streaming, drift, and diamagnetic slices.

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

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

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

spectraxgk.api.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.

spectraxgk.api.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.

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

Convert LinearTerms into the modular TermConfig.

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

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

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

Create summary panels with multiple external reference curves.

spectraxgk.api.load_runtime_from_toml(path: str | Path) → tuple[RuntimeConfig, dict]

Load unified runtime config from TOML, returning (cfg, data).

spectraxgk.api.local_stencil_nonlinear_domain_decomposed_step(state: Array, plan: NonlinearDomainDecompositionPlan, *, dt: float = 0.05) → Array

Apply the same local nonlinear step through static halo chunks.

spectraxgk.api.local_stencil_nonlinear_domain_serial_step(state: Array, *, axis: int = 0, dt: float = 0.05) → Array

Apply the serial local nonlinear diagnostic step along one state axis.

spectraxgk.api.logical_decomposed_nonlinear_spectral_rhs(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), atol: float = 5e-06, rtol: float = 5e-06) → tuple[Array, NonlinearSpectralRHSIdentityReport]

Return the logical-shard nonlinear spectral RHS after identity gating.

The returned RHS uses the logical decomposed route only when it is exactly equivalent to the serial reference under the provided tolerances. Otherwise the serial RHS is returned and decomposed_path_enabled is false. This is still a local diagnostic route, not a distributed runtime implementation.

spectraxgk.api.matched_nonlinear_transport_report(baseline: dict[str, Any], treatment: dict[str, Any], *, case: str = 'matched_nonlinear_transport', treatment_name: str = 'treatment', min_relative_reduction: float = 0.0, min_uncertainty_z_score: float = 0.0, value_floor: float = 1e-12) → dict[str, Any]

Compare two independently converged post-transient transport windows.

spectraxgk.api.matched_optimized_transport_report(path: str, payload: Mapping[str, Any], *, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Return whether a matched baseline-to-optimized audit promotes transport.

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

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

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

Return the mirror-drift contribution using Hermite shard exchange.

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

JAX-differentiable mixing-length squared-amplitude rule.

spectraxgk.api.nonlinear_domain_identity_report(serial_state: Array, decomposed_state: Array, plan: NonlinearDomainDecompositionPlan, *, atol: float = 1e-06, rtol: float = 1e-06) → NonlinearDomainIdentityReport

Compare decomposed and serial states and fail closed on any mismatch.

spectraxgk.api.nonlinear_domain_parallel_identity_gate(state: Array, plan: NonlinearDomainDecompositionPlan, *, dt: float = 0.05, atol: float = 1e-06, rtol: float = 1e-06) → tuple[Array, NonlinearDomainIdentityReport]

Return a fail-closed decomposed local-stencil step and its identity report.

spectraxgk.api.nonlinear_domain_transport_window_identity_gate(state: Array, plan: NonlinearDomainDecompositionPlan, *, dt: float = 0.025, steps: int = 4, atol: float = 1e-06, rtol: float = 1e-06) → NonlinearDomainTransportWindowReport

Validate a multi-step serial-vs-decomposed transport window.

The gate is deliberately stricter than a final-state check: it compares state identity, decomposed-boundary identity, and per-step scalar traces for mass, free-energy proxy, and boundary-flux proxy. The scalar drifts are compared between serial and decomposed paths; they are not claimed to be conserved by this damped diagnostic stencil.

spectraxgk.api.nonlinear_heat_flux_trace(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → tuple[Array, Array]

Return a differentiable short-window ITG heat-flux envelope trace.

The envelope evolves E with a fixed-step RK2 discretization,

dE/dt = 2 gamma E - alpha E^2, Q_env(t) = W_i E.

gamma and W_i come from the same differentiable QA/ITG feature map as the linear and quasilinear objectives. The output is therefore useful for nonlinear averaging, optimizer, and UQ gates while the full production nonlinear-GK geometry path is still being made traceable end-to-end.

spectraxgk.api.nonlinear_heat_flux_window_metrics(times: Array, heat_flux: Array, *, tail_fraction: float = 0.45, eps: float = 1e-14) → dict[str, Array]

Return mean, coefficient of variation, and trend on a late-time window.

spectraxgk.api.nonlinear_parallel_strategies() → tuple[NonlinearParallelStrategy, ...]

Return all nonlinear parallelization strategy contracts.

spectraxgk.api.nonlinear_parallel_strategy(name: Literal['independent_ky_scan', 'uq_ensemble', 'whole_state_kx_ky', 'velocity_species_hermite', 'fft_axis_domain']) → NonlinearParallelStrategy

Return the contract for a named nonlinear parallelization strategy.

spectraxgk.api.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.

spectraxgk.api.nonlinear_spectral_communication_identity_gate(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), atol: float = 5e-06, rtol: float = 5e-06) → NonlinearSpectralCommunicationReport

Validate FFT, bracket, and field layout identity under split/reassemble.

spectraxgk.api.nonlinear_spectral_communication_identity_report(serial_fft_roundtrip: Array, communicated_fft_roundtrip: Array, serial_bracket: Array, communicated_bracket: Array, serial_field: Array, communicated_field: Array, *, state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], atol: float = 5e-06, rtol: float = 5e-06) → NonlinearSpectralCommunicationReport

Compare spectral communication outputs and fail closed on mismatches.

spectraxgk.api.nonlinear_spectral_domain_work_model(state_shape: tuple[int, int, int, int, int], *, y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], max_communication_to_owned_work_ratio: float = 0.5) → NonlinearSpectralDomainWorkModel

Estimate communication pressure for the current logical spectral route.

The current diagnostic route reconstructs global spectral state/bracket arrays before returning owned output tiles. That is useful for identity gating, but it implies allgather/broadcast traffic that can dominate the owned tile work. This model is a conservative fail-closed screen for that route; it is not a performance prediction for a future distributed FFT.

spectraxgk.api.nonlinear_spectral_integrator_identity_gate(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), dt: float = 0.005, steps: int = 4, atol: float = 5e-06, rtol: float = 5e-06) → NonlinearSpectralIntegratorIdentityReport

Validate serial-vs-logical-shard nonlinear spectral integration.

This gate compares a fixed-step explicit-Euler micro-integration using the same nonlinear spectral RHS on serial and logical tiled routes. It verifies final-state identity and per-step free-energy, field-energy, and flux-proxy traces. Passing this gate is necessary but not sufficient for production nonlinear domain parallelization.

spectraxgk.api.nonlinear_spectral_pencil_rhs_identity_gate(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), atol: float = 5e-06, rtol: float = 1e-05, max_communication_to_fft_work_ratio: float = 0.35, min_predicted_speedup: float = 1.5) → NonlinearSpectralPencilRHSIdentityReport

Validate serial-vs-pencil nonlinear spectral RHS identity.

spectraxgk.api.nonlinear_spectral_pencil_transport_window_identity_gate(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), dt: float = 0.005, steps: int = 4, atol: float = 5e-06, rtol: float = 1e-05, max_communication_to_fft_work_ratio: float = 0.35, min_predicted_speedup: float = 1.5) → NonlinearSpectralPencilTransportWindowReport

Validate a serial-vs-pencil nonlinear transport window.

spectraxgk.api.nonlinear_spectral_pencil_work_model(state_shape: tuple[int, int, int, int, int], *, y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], max_communication_to_fft_work_ratio: float = 0.35, min_predicted_speedup: float = 1.5) → NonlinearSpectralPencilWorkModel

Estimate communication pressure for a pencil-FFT bracket route.

The pencil route avoids global state/bracket reconstruction. Its remaining distributed cost is the all-to-all transpose traffic needed by axis-wise 2D FFTs plus a field-reduction/broadcast. This model is intentionally simple and fail-closed: it must pass before any strong-scaling timing is treated as a meaningful candidate for production promotion.

spectraxgk.api.nonlinear_spectral_rhs_identity_gate(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), atol: float = 5e-06, rtol: float = 5e-06) → NonlinearSpectralRHSIdentityReport

Validate serial-vs-logical-shard nonlinear spectral RHS identity.

This diagnostic route owns and reassembles spectral (y, x) output tiles in row-major order. It deliberately does not install distributed FFT runtime routing or make a speedup claim.

spectraxgk.api.nonlinear_spectral_rhs_identity_report(serial_reconstruction: Array, logical_reconstruction: Array, serial_field: Array, logical_field: Array, serial_bracket: Array, logical_bracket: Array, serial_rhs: Array, logical_rhs: Array, *, state_shape: tuple[int, int, int, int, int], y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], tile_bounds: tuple[tuple[int, int, int, int], ...] | None = None, atol: float = 5e-06, rtol: float = 5e-06) → NonlinearSpectralRHSIdentityReport

Compare serial and logical-shard spectral RHS outputs fail-closed.

spectraxgk.api.nonlinear_window_convergence_from_csv(csv_path: str | Path, *, time_column: str = 't', value_column: str = 'heat_flux', case: str | None = None, config: NonlinearWindowConvergenceConfig | None = None, summary_artifact: str | None = None) → dict[str, Any]

Build a convergence report from a diagnostics CSV.

spectraxgk.api.nonlinear_window_convergence_from_summary(summary_json: str | Path, *, diagnostics_source: str = 'spectrax', time_column: str = 't', value_column: str = 'heat_flux', case: str | None = None, config: NonlinearWindowConvergenceConfig | None = None) → dict[str, Any]

Build a convergence report from a window summary and diagnostics CSV.

spectraxgk.api.nonlinear_window_convergence_report(time: Sequence[float] | ndarray, values: Sequence[float] | ndarray, *, case: str = 'nonlinear_window', observable: str = 'heat_flux', source_artifact: str | None = None, summary_artifact: str | None = None, config: NonlinearWindowConvergenceConfig | None = None) → dict[str, Any]

Return finite late-window statistics and convergence gates.

The running-mean drift compares the mean of the first and second halves of the late window, normalized by the late-window mean scale. The uncertainty gate uses the maximum of the sample SEM, contiguous-block SEM, and block-bootstrap SEM when available.

spectraxgk.api.nonlinear_window_ensemble_artifact_manifest(records: Sequence[dict[str, Any]], *, case: str = 'nonlinear_window_ensemble_artifact_manifest', config: NonlinearWindowEnsembleManifestConfig | None = None) → dict[str, Any]

Return a promotion-blocking manifest for missing ensemble artifacts.

Each record should contain a report produced by nonlinear_window_convergence_report(), plus optional variant metadata such as {"seed": 1, "timestep": 0.02}. The manifest is intentionally conservative: a production nonlinear optimization promotion needs distinct passed artifacts for every required variant axis, so a single late-window summary is recorded as useful convergence evidence but not as replicated-ensemble evidence.

spectraxgk.api.nonlinear_window_ensemble_report(reports: Sequence[dict[str, Any]], *, case: str = 'nonlinear_window_ensemble', comparison: str = 'replicate_uncertainty', config: NonlinearWindowEnsembleConfig | None = None) → dict[str, Any]

Gate repeated nonlinear-window summaries for seed/timestep robustness.

The input reports are expected to come from nonlinear_window_convergence_report(). This helper does not inspect raw time traces; it compares already-gated late-window means and their uncertainty metadata so production promotion can require seed, initial condition, or timestep robustness without rerunning simulations inside the checker.

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

Gate windowed nonlinear transport and field-energy metrics.

spectraxgk.api.nonlinear_window_stats_promotion_ready(stats: object) → tuple[bool, list[str]]

Return whether serialized nonlinear window metadata can support promotion.

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

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

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

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

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

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

Gate an observed-order convergence study.

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

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

Optimize one reduced QA low-turbulence design with Adam.

spectraxgk.api.optimize_stellarator_itg(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], initial_params: Array | Sequence[float] | None = None, config: StellaratorITGOptimizationConfig | None = None, *, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → StellaratorITGOptimizationResult

Optimize one differentiable stellarator ITG objective with Adam.

spectraxgk.api.optimized_equilibrium_transport_report(path: str, payload: Mapping[str, Any], *, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Return whether an artifact can promote optimized-equilibrium transport.

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

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

spectraxgk.api.pencil_decomposed_nonlinear_spectral_rhs(state_hat: Array, *, y_chunks: tuple[int, ...] = (3, 3), x_chunks: tuple[int, ...] = (2, 2), atol: float = 5e-06, rtol: float = 1e-05) → tuple[Array, NonlinearSpectralPencilRHSIdentityReport]

Return the pencil nonlinear spectral RHS after identity/model gating.

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

Return periodic parallel streaming using full-array operations.

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

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

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

Return the amplitude normalization used for quasilinear weights.

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

Return normalized objective-column weights.

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

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

spectraxgk.api.production_nonlinear_optimization_guard_report(*, optimization_artifact: Mapping[str, Any] | None, optimization_artifact_path: str = '', reduced_artifacts: Mapping[str, Mapping[str, Any]] | None = None, replicated_ensemble_artifacts: Mapping[str, Mapping[str, Any]] | None = None, optimized_equilibrium_artifacts: Mapping[str, Mapping[str, Any]] | None = None, matched_optimized_transport_artifacts: Mapping[str, Mapping[str, Any]] | None = None, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Build the fail-closed nonlinear turbulent-flux optimization guard.

The top-level passed field means the release is safe: reduced/startup artifacts are correctly scoped and long-window replicated holdouts are present. It does not mean production nonlinear optimization is promoted; that is reported separately by production_nonlinear_optimization_promoted.

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

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

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

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

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

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

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

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

Return the scalar reduced QA comparison objective.

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

Compare the complete controls-to-observables Jacobian with finite differences.

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

Return QA low-turbulence observables in stable order.

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

Return reduced QA constraints and ITG observables.

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

Return the stable residual names for the comparison objective.

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

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

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

Return differentiable late-window heat-flux statistics.

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

spectraxgk.api.qa_max_mode1_observables(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → dict[str, Array]

Map a QA max-mode-1 boundary/control vector to differentiable ITG observables.

The four inputs represent the active low-order controls used by the example scripts. The map is calibrated as a smooth objective-reduction gate around a QA stellarator with aspect ratio 7 and mean rotational transform 0.41. It is not a replacement for the full VMEC/Boozer flux-tube geometry contract; its purpose is to validate gradient plumbing, UQ, optimizer behavior, and figure-generation before expensive production objectives are promoted.

spectraxgk.api.qa_observable_vector(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None) → Array

Return observables in the stable order defined by OBSERVABLE_NAMES.

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

Build a JSON-friendly calibration/holdout report.

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

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

Differentiable objective from [gamma, kperp_eff2, flux_weight].

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

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

Reassemble shard values into serial index order.

spectraxgk.api.reduced_artifact_scope_report(path: str, payload: Mapping[str, Any]) → dict[str, Any]

Return whether a startup/reduced artifact is safely blocked from promotion.

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

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

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

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

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

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

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

spectraxgk.api.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) → RuntimeLinearResult

Refit one stored trajectory without repeating its integration.

spectraxgk.api.release_ready_nonlinear_parallel_strategies() → tuple[NonlinearParallelStrategy, ...]

Return production-facing strategies that do not alter solver layout.

spectraxgk.api.replicated_transport_ensemble_report(path: str, payload: Mapping[str, Any], *, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Return quality metadata for a long-window replicated transport ensemble.

spectraxgk.api.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, deps: RuntimeCaseDeps | None = None) → int

Run a linear case from a runtime TOML with optional overrides.

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

Run a deterministic pointwise linear scan over ky_values.

spectraxgk.api.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, deps: RuntimeCaseDeps | None = None) → int

Run a nonlinear case from a runtime TOML with optional overrides.

spectraxgk.api.run_runtime_linear(cfg: RuntimeConfig, *, ky_target: float = 0.3, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', return_state: bool = False, initial_state: Any | None = None, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → RuntimeLinearResult

Run one linear point from a case-agnostic runtime config.

spectraxgk.api.run_runtime_nonlinear(cfg: RuntimeConfig, *, ky_target: float = 0.3, kx_target: float | None = None, Nl: int | None = None, Nm: int | None = None, dt: float | None = None, steps: int | None = None, method: str | None = None, sample_stride: int | None = None, diagnostics_stride: int | None = None, laguerre_mode: str | None = None, diagnostics: bool | None = None, resolved_diagnostics: bool = True, return_state: bool = False, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → RuntimeNonlinearResult

Run a nonlinear point using the unified runtime config path.

spectraxgk.api.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) → RuntimeParameterScanResult

Run a scalar scan and optionally continue a selected solution branch.

spectraxgk.api.run_runtime_scan(cfg: RuntimeConfig, ky_values: Sequence[float], *, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, batch_ky: bool = False, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', show_progress: bool = False, workers: int = 1, parallel_executor: str = 'thread') → RuntimeLinearScanResult

Run a ky scan using the unified runtime config path.

The public facade keeps runtime monkeypatch seams intact while scan coordination lives in workflows/runtime/orchestration_scan.py.

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

Run a pointwise scan and extract the fastest or selected eigenmode.

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

Return a differentiable mixing-length saturated flux estimate.

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

Return the squared amplitude implied by a named saturation rule.

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

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

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

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

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

Return the promoted candidate from build_transport_admission_report().

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

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

spectraxgk.api.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.

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

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

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

Check that contract sharding reassembles exactly to serial order.

spectraxgk.api.set_plot_style() → None

Apply the shared publication style used by generated figures.

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

Differentiable shape-aware linear-weight objective.

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

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

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

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

Return values grouped according to a decomposition contract.

spectraxgk.api.solver_grid_options_from_ky_values(ky_values: float | tuple[float, ...] | list[float], *, ky_base: float | None = None, min_ny: int = 4) → dict[str, object]

Return solver grid options for physical k_y rho_i scan values.

The linear objective evaluator selects FFT row indices, while user-facing optimization and validation studies should be specified in physical k_y rho_i. This helper fixes that contract explicitly: values must be positive integer multiples of the base spacing, Ly = 2*pi/ky_base, and Ny is chosen large enough that all requested modes are represented as positive FFT rows.

spectraxgk.api.solver_growth_rate_from_geometry(geom: Any, *, 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

Evaluate the dominant linear growth rate without eigenvector AD.

spectraxgk.api.solver_linear_operator_matrix_from_geometry(geom: Any, *, 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.

spectraxgk.api.solver_objective_vector_from_geometry(geom: Any, *, 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

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.

Branch continuity and AD/finite-difference validation are deliberately handled by gate functions; this is the shared forward evaluator those production objectives can use after a branch has been selected and audited.

spectraxgk.api.solver_ready_geometry_mapping(params: Array, theta: Array) → dict[str, Any]

Map a two-parameter design vector into solver-ready flux-tube arrays.

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

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

Return a normalized sparse descent direction from a gradient report.

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

spectraxgk.api.stellarator_itg_density_gradient_scan(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradients: Sequence[float] | None = None, temperature_gradient: float | None = None) → dict[str, Any]

Return a reduced ITG response scan versus normalized density gradient.

The scan uses the same explicit drive inputs as the reduced growth, quasilinear, and nonlinear-envelope objectives. It is intended for candidate ranking and figure QA before promotion to solved VMEC/Boozer nonlinear SPECTRAX-GK scans.

spectraxgk.api.stellarator_itg_objective(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None) → Array

Return the scalar constrained QA + ITG objective for one optimization.

spectraxgk.api.stellarator_itg_objective_residual_names(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']) → tuple[str, ...]

Return stable residual names for the weighted QA + ITG objective.

spectraxgk.api.stellarator_itg_objective_residual_vector(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None) → Array

Return the weighted residual map used by the objective and covariance.

spectraxgk.api.stellarator_itg_portfolio_gate_payload(params: Array | Sequence[float] | None = None, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Return the JSON-ready reduced ITG portfolio gate artifact payload.

spectraxgk.api.stellarator_itg_portfolio_sensitivity_report(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, step: float | None = None, rtol: float | None = None, atol: float | None = None, workers: int = 1, parallel_executor: str = 'thread') → dict[str, Any]

AD/FD, conditioning, and covariance gate for the reduced ITG portfolio.

spectraxgk.api.stellarator_itg_reduced_portfolio_objective(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None) → Array

Reduce a sampled ITG growth/QL portfolio to one differentiable scalar.

spectraxgk.api.stellarator_itg_residual_sensitivity_report(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None, *, step: float | None = None, rtol: float | None = None, atol: float | None = None, min_rank: int = 4, condition_number_limit: float = 10000.0, covariance_regularization: float = 1e-08, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Check residual-Jacobian AD/FD parity and local conditioning.

spectraxgk.api.stellarator_itg_sample_objective_table(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None) → Array

Return (surface, alpha, ky, objective) reduced ITG objective rows.

This is the backend-free rehearsal of the production VMEC/Boozer sample table. It keeps the optimizer and gate semantics identical to the future real-geometry path while remaining cheap enough for CI.

spectraxgk.api.stellarator_itg_vmec_boozer_portfolio_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, **vmec_boozer_options: Any) → Array

Facade wrapper for the real VMEC/Boozer/SPECTRAX-GK portfolio scalar.

spectraxgk.api.stellarator_itg_vmec_boozer_sample_objective_table_from_state(state: Any, static: Any, indata: Any, wout: Any, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), sample_set: StellaratorITGSampleSet | None = None, **vmec_boozer_options: Any) → Array

Facade wrapper for real VMEC/Boozer/SPECTRAX-GK objective rows.

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

Convert modular TermConfig into linear-only term weights.

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

Summarize whether a transport objective has enough sample coverage.

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

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

Validate static row/weight contracts and return portfolio metadata.

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

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

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

spectraxgk.api.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.

spectraxgk.api.vmec_boozer_aggregate_line_search_holdout_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', training_weights: tuple[float, ...] | list[float] | ndarray | None = None, holdout_weights: tuple[float, ...] | list[float] | ndarray | None = None, training_surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), training_alphas: float | tuple[float, ...] | list[float] = (0.0,), training_selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), holdout_surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), holdout_alphas: float | tuple[float, ...] | list[float] = (0.0,), holdout_selected_ky_indices: int | tuple[int, ...] | list[int] = (2,), radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', initial_delta: float = 0.0, perturbation_step: float = 1e-07, update_step: float = 1e-08, max_steps: int = 3, min_improvement: float = 0.0, min_holdout_improvement: float = 0.0, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Audit a training aggregate update against held-out aggregate samples.

spectraxgk.api.vmec_boozer_aggregate_scalar_objective_finite_difference_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', weights: tuple[float, ...] | list[float] | ndarray | None = None, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), torflux_values: float | tuple[float, ...] | list[float] | None = None, alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), ky_values: float | tuple[float, ...] | list[float] | None = None, ky_base: float | None = None, radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', base_delta: float = 0.0, perturbation_step: float = 1e-07, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Finite-difference a multi-surface/multi-k_y aggregate objective.

spectraxgk.api.vmec_boozer_aggregate_scalar_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, *, objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', weights: tuple[float, ...] | list[float] | ndarray | None = None, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), **kwargs: Any) → Array

Reduce a VMEC/Boozer multi-point objective table to one scalar.

spectraxgk.api.vmec_boozer_aggregate_scalar_objective_line_search_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', weights: tuple[float, ...] | list[float] | ndarray | None = None, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', initial_delta: float = 0.0, perturbation_step: float = 1e-07, update_step: float = 1e-08, max_steps: int = 3, min_improvement: float = 0.0, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Run a curvature-gated line search for an aggregate VMEC objective.

spectraxgk.api.vmec_boozer_scalar_objective_finite_difference_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', base_delta: float = 0.0, perturbation_step: float = 1e-07, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Finite-difference a scalar objective through a VMEC state coefficient.

spectraxgk.api.vmec_boozer_scalar_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, *, objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', **kwargs: Any) → Array

Evaluate one scalar optimization objective on the VMEC/Boozer path.

spectraxgk.api.vmec_boozer_scalar_objective_line_search_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', initial_delta: float = 0.0, perturbation_step: float = 1e-07, update_step: float = 1e-08, max_steps: int = 3, min_improvement: float = 0.0, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Run a curvature-gated one-parameter VMEC/Boozer objective line search.

spectraxgk.api.vmec_boozer_solver_objective_table_from_state(state: Any, static: Any, indata: Any, wout: Any, *, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), torflux_values: float | tuple[float, ...] | list[float] | None = None, alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), ky_values: float | tuple[float, ...] | list[float] | None = None, ky_base: float | None = None, **kwargs: Any) → Array

Evaluate solver objectives over a surface/field-line/k_y table.

spectraxgk.api.vmec_boozer_solver_objective_table_with_metadata_from_state(state: Any, static: Any, indata: Any, wout: Any, *, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), torflux_values: float | tuple[float, ...] | list[float] | None = None, alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), ky_values: float | tuple[float, ...] | list[float] | None = None, ky_base: float | None = None, **kwargs: Any) → tuple[Array, list[dict[str, object]]]

Evaluate VMEC/Boozer objective rows and return sample metadata.

spectraxgk.api.vmec_boozer_solver_objective_vector_from_state(state: Any, static: Any, indata: Any, wout: Any, **kwargs: Any) → Array

Evaluate solver objectives from the in-memory VMEC/Boozer bridge.

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

Validate a real vmec_jax boundary-aspect derivative when available.

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

spectraxgk.api.vmec_field_line_tensor_observable_names() → tuple[str, ...]

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

spectraxgk.api.vmec_metric_tensor_observable_names() → tuple[str, ...]

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

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

Write a boundary-gradient diagnostic JSON artifact.

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

Write a quasilinear calibration report to JSON.

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

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

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

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

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

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

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

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

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

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

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

Gate Rosenbluth-Hinton/GAM-style response observables.

Core Refactor Contracts

Typed contracts for differentiable refactors and validation gates.

These containers are intentionally independent of solver implementation files. They make refactor ownership, JAX transformability, validation gates, and public extension points explicit before large modules are split.

class spectraxgk.core.contracts.DifferentiabilityContract(differentiable: bool, jit_safe: bool, vmap_safe: bool, static_arg_names: tuple[str, ...] = (), dynamic_arg_names: tuple[str, ...] = (), gradient_checks: tuple[str, ...] = (), custom_derivative: Literal['none', 'jvp', 'vjp', 'jvp_and_vjp'] = 'none', notes: str = '')

Declare how a callable is expected to behave under JAX transforms.

This contract separates traced array inputs from static model choices and records the gradient evidence required before a differentiable claim is promoted. It does not itself call JAX.

class spectraxgk.core.contracts.ExtensionPointContract(name: str, protocol: str, required_methods: tuple[str, ...], validation_tests: tuple[str, ...], documentation: str)

Declare a supported community extension surface.

class spectraxgk.core.contracts.ModuleRefactorContract(source_module: str, facade_module: str, target_modules: tuple[str, ...], public_api_compatible: bool, max_lines_target: int, validation_gates: tuple[~spectraxgk.core.contracts.ValidationGateContract, ...], extension_points: tuple[~spectraxgk.core.contracts.ExtensionPointContract, ...] = <factory>, differentiability: ~spectraxgk.core.contracts.DifferentiabilityContract | None = None)

Trace a large source module to its target split modules and gates.

property gate_names: tuple[str, ...]

Return the declared validation-gate names in execution order.

property target_package_names: tuple[str, ...]

Return unique parent packages for the target modules.

class spectraxgk.core.contracts.ShapeContract(axes: tuple[str, ...], description: str, dtype: str | None = None)

Array shape metadata for a public kernel, diagnostic, or artifact field.

Parameters:

• axes (tuple[str, ...]) – Ordered axis names, for example ("species", "ky", "kx", "z").

• description (str) – Human-readable role of the array in the equation or diagnostic.

• dtype (str | None) – Optional dtype policy, such as "real", "complex", or "inherits_state". This is descriptive; kernels still enforce actual dtypes locally.

class spectraxgk.core.contracts.ValidationGateContract(name: str, category: Literal['unit', 'numerics', 'physics', 'parity', 'autodiff', 'performance', 'artifact'], observable: str, evidence: tuple[str, ...], tolerance: str | None = None, literature_anchors: tuple[str, ...] = ())

Physics, numerical, autodiff, parity, or performance gate declaration.

Core Extension Points

Structural protocols for SPECTRAX-GK extension points.

The protocols are intentionally structural and dependency-light. Collision operators are active in cached RHS and explicit state-integration paths; the remaining protocols define the narrow interfaces used as other domains become pluggable.

class spectraxgk.core.extension_points.ArtifactWriter(*args, **kwargs)

Side-effectful writer kept outside differentiable objective functions.

write(payload: Any, destination: Any) → Any

Persist payload to destination and return writer metadata.

class spectraxgk.core.extension_points.BasisFamily(*args, **kwargs)

Orthogonal velocity-space basis family used by moment-space operators.

recurrence_coefficients(order: int) → Any

Return recurrence data for the requested polynomial order.

class spectraxgk.core.extension_points.CollisionContext(distribution: Any, hamiltonian: Any, fields: Any, cache: Any, parameters: Any)

Post-field state needed by a gyrokinetic collision operator.

distribution is the evolved \(G\) state, while hamiltonian is the corresponding \(H\) response after the field solve. Keeping both arrays explicit prevents finite-Larmor-radius and field-particle models from silently applying a long-wavelength approximation.

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 spectraxgk.core.extension_points.CollisionOperator(*args, **kwargs)

JAX-compatible unit-weight gyrokinetic collision model.

apply(context: CollisionContext) → Any

Return a collision RHS contribution with the same state layout.

class spectraxgk.core.extension_points.Diagnostic(*args, **kwargs)

Observable extractor used by runtime, validation, and plotting workflows.

evaluate(state: Any, fields: Any, geometry: Any, parameters: Any) → Any

Return diagnostic values for a sampled state and fields.

class spectraxgk.core.extension_points.FieldSolver(*args, **kwargs)

Field-equation solver for potential and electromagnetic response fields.

solve_fields(distribution: Any, geometry: Any, parameters: Any) → Any

Return fields generated by a distribution in a sampled geometry.

class spectraxgk.core.extension_points.GeometryProvider(*args, **kwargs)

Provider that maps equilibrium parameters to sampled flux-tube geometry.

sample_flux_tube(parameters: Any) → Any

Return a solver-ready sampled flux-tube geometry object.

class spectraxgk.core.extension_points.LinearRHS(*args, **kwargs)

Linear right-hand side callable with explicit time, state, and parameters.

class spectraxgk.core.extension_points.NonlinearRHS(*args, **kwargs)

Nonlinear right-hand side callable with explicit time, state, and parameters.

class spectraxgk.core.extension_points.Objective(*args, **kwargs)

Differentiable or derivative-free objective used by optimization workflows.

evaluate(parameters: Any) → Any

Return scalar or vector objective values for optimizer parameters.

class spectraxgk.core.extension_points.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.

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

Gyroaveraging coefficients matching the Laguerre-Hermite quadrature convention.

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

spectraxgk.core.velocity.bessel_j0(x: Array) → Array

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

spectraxgk.core.velocity.bessel_j1(x: Array) → Array

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

spectraxgk.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).

spectraxgk.core.velocity.gamma0(b: Array) → Array

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

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

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

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

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

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

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

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

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

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

spectraxgk.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}

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

spectraxgk.core.velocity.laguerre_quadrature_count(nl: int) → int

Default number of Laguerre quadrature points.

spectraxgk.core.velocity.laguerre_transform(nl: int) → tuple[ndarray, ndarray, ndarray]

Return Laguerre transform matrices and roots.

spectraxgk.core.velocity.single_precision_factorial(m: Array) → Array

Return the single-precision factorial approximation.

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

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

Geometry

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

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

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

trim_terminal_theta_point() → FluxTubeGeometryData

Return a copy without the terminal theta sample.

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 spectraxgk.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 spectraxgk.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)

Spectral grid configuration in a flux-tube.

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

Simple s-alpha geometry with circular concentric flux surfaces.

bgrad(theta: Array) → Array

Magnetic field gradient term used in mirror force.

bmag(theta: Array) → Array

Magnetic field strength for circular s-alpha geometry.

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

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

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

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

gradpar() → float

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

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

Perpendicular wave-number squared for s-alpha geometry.

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

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

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

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

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

Magnetic drift frequency for s-alpha geometry.

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

Reference slab geometry contract.

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

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

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

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

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

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

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

Return the effective boundary after zero-shear promotion.

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

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

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

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

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

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

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

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

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

Return the effective zero-shear state.

Geometry Core

Flux-tube geometry builders and grid-default policy.

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

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 spectraxgk.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 spectraxgk.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)

Spectral grid configuration in a flux-tube.

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

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

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

gradpar() → float

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

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

Perpendicular wave-number squared for s-alpha geometry.

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

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

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

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

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

Magnetic drift frequency for s-alpha geometry.

class spectraxgk.geometry.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.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

spectraxgk.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 spectraxgk.geometry.flux_tube.

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

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

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

gradpar() → float

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

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

Perpendicular wave-number squared for s-alpha geometry.

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

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

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

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

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

Magnetic drift frequency for s-alpha geometry.

class spectraxgk.geometry.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.

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

Construct the Miller surface on the collocation grid.

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

spectraxgk.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 spectraxgk.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)

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

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

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

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

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

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

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

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

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

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

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

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

Build a Miller generation request from a runtime config.

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

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

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

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

Resolve a runtime VMEC path with env/user expansion.

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

Build a VMEC generation request from a runtime config.

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

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

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

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

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

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

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

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

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

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

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

Validate a real booz_xform_jax spectral derivative when available.

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

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

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

Discover optional vmec_jax and booz_xform_jax bridge APIs.

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

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

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

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

Central finite-difference Jacobian for small validation problems.

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

Build FluxTubeGeometryData from an in-memory differentiable backend.

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

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

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

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

VMECState -> BoozXformInputs -> booz_xform_jax -> FluxTubeGeometryData.

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

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

Return differentiable scalar observables from solver-ready geometry.

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

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

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

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

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

Return the ordered geometry observables used by bridge AD checks.

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

Validate geometry-observable sensitivities by AD and finite differences.

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

spectraxgk.geometry.differentiable.observable_gradient_validation_report(observable_fn: Callable[[Array], Any], params: Array | ndarray, *, fd_step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, observable_names: Sequence[str] | None = None, param_names: Sequence[str] | None = None, tangent: Array | ndarray | None = None, relative_floor: float = 1e-12, min_rank: int | None = None, condition_number_max: float | None = 1000000000000.0, 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.

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

Validate a real vmec_jax boundary-aspect derivative when available.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Differentiable Geometry Backend Discovery

Optional differentiable-geometry backend discovery helpers.

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

Return existing backend roots from environment variables and defaults.

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

Import a backend module, preferring explicitly configured checkout paths.

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

Return true when host NumPy validation would break JAX tracing.

spectraxgk.geometry.backend_discovery._jax_float_dtype() → Any

Return the active JAX floating dtype for small validation arrays.

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

Discover optional vmec_jax and booz_xform_jax bridge APIs.

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

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

Build FluxTubeGeometryData from an in-memory differentiable backend.

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

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

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

Return the ordered geometry observables used by bridge AD checks.

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

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

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

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

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

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

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

Central finite-difference Jacobian for small validation problems.

spectraxgk.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 Geometry Sensitivity

Geometry sensitivity, inverse-design, and local UQ reports.

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

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

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

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

Validate geometry-observable sensitivities by AD and finite differences.

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

Differentiable Boozer Bridge

Bounded VMEC/Boozer differentiable bridge helpers.

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

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

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

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

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

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

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

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

Validate a real booz_xform_jax spectral derivative when available.

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

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

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

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

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

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

Validate a real vmec_jax boundary-aspect derivative when available.

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

Differentiable VMEC-State Sensitivity

VMEC-state differentiable sensitivity reports.

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

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

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

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

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

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

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

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

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

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

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

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

Differentiable VMEC Boozer Core

VMEC-JAX to Boozer equal-arc core-profile bridge.

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

Build solver-ready geometry directly from an in-memory VMEC/Boozer state.

spectraxgk.geometry.vmec_boozer_core.prewarm_vmec_boozer_equal_arc_cache(static: Any, wout: Any, *, mboz: int = 21, nboz: int = 21, asym: bool | None = None) → None

Precompute Boozer constants before VMEC-JAX jits residual callbacks.

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

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

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

Differentiable VMEC Boozer Constants

Boozer-transform constants and cache prewarm helpers for VMEC-JAX bridges.

spectraxgk.geometry.vmec_boozer_constants._cached_booz_xform_constants(*, nfp: int, mpol: int, ntor: int, ntheta: int, nzeta: int, mboz: int, nboz: int, asym: bool) → tuple[Any, Any]

Prepare Boozer constants outside traced VMEC-JAX residual callbacks.

spectraxgk.geometry.vmec_boozer_constants.prewarm_vmec_boozer_equal_arc_cache(static: Any, wout: Any, *, mboz: int = 21, nboz: int = 21, asym: bool | None = None) → None

Precompute Boozer constants before VMEC-JAX jits residual callbacks.

Differentiable VMEC Flux-Tube Reports

VMEC flux-tube sensitivity and parity reports.

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

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

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

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

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

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

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

Differentiable VMEC Tensor Mapping

Direct VMEC tensor to flux-tube mapping bridge.

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

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

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

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

Differentiable Geometry Numerics

Pure numerical helpers for differentiable geometry bridge routines.

spectraxgk.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:]).

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

spectraxgk.geometry.numerics._periodic_bilinear_sample_2d(values: Array, theta: Array, zeta: Array) → Array

Sample a uniform periodic (theta,zeta) grid with fixed bilinear weights.

Grids

Spectral grid utilities for flux-tube geometry.

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

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

spectraxgk.core.grid.real_fft_ordered_kx(kx: Array) → Array

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

spectraxgk.core.grid.real_fft_unique_ky(ky: Array) → Array

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

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

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

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

2/3 dealiasing mask for 2D Fourier grids.

Species

Species helpers for assembling linear gyrokinetic parameters.

class spectraxgk.core.species.Species(charge: float, mass: float, density: float, temperature: float, tprim: float, fprim: float, nu: float = 0.0)

Physical parameters for a kinetic species.

spectraxgk.core.species.build_linear_params(species: Iterable[Species], *, tau_e: float = 0.0, kpar_scale: float = 1.0, omega_d_scale: float = 1.0, omega_star_scale: float = 1.0, rho_star: float = 1.0, beta: float = 0.0, fapar: float = 0.0, apar_beta_scale: float = 0.5, ampere_g0_scale: float = 0.5, bpar_beta_scale: float = 0.5, nu_hyper: float = 0.0, p_hyper: float = 4.0, nu_hyper_l: float = 0.0, nu_hyper_m: float = 1.0, nu_hyper_lm: float = 0.0, p_hyper_l: float = 6.0, p_hyper_m: float = 20.0, p_hyper_lm: float = 6.0, hypercollisions_const: float = 1.0, hypercollisions_kz: float = 0.0, D_hyper: float = 0.0, p_hyper_kperp: float = 2.0) → LinearParams

Build LinearParams arrays from a list of species.

Operators

Public operator kernels with lazy domain imports.

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

Parallel streaming operator acting on the Hermite index.

Linear Operators

Linear operator facade with cycle-free lazy imports.

Linear

Public linear gyrokinetic operators, caches, and integrators.

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

Precomputed arrays for the linear operator.

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

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

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

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

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

Multiply Hermite coefficients by v_parallel (ladder form).

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

Multiply Hermite coefficients by v_parallel^2.

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

Multiply Laguerre coefficients by the perpendicular energy variable.

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

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

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.

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

Build reusable arrays for the linear RHS.

spectraxgk.linear.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.

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

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

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

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

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

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

spectraxgk.linear.grad_z_periodic(f: Array, dz: float | Array | None = None, kz: Array | None = None) → Array

Spectral periodic derivative along the last axis.

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

Parallel streaming operator acting on the Hermite index.

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

Assemble benchmark-compatible hypercollision damping factors.

spectraxgk.linear.integrate_linear(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, method: str = 'rk4', cache: LinearCache | None = None, implicit_tol: float = 1e-06, implicit_maxiter: int = 200, implicit_iters: int = 3, implicit_relax: float = 0.7, implicit_restart: int = 20, implicit_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) → tuple[Array, Array]

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

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

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

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

Diagonal Lenard-Bernstein rates in Hermite-Laguerre space.

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

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

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

Compute the linear RHS using precomputed geometry/cache arrays.

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

Compute gated electrostatic streaming, drift, and diamagnetic slices.

spectraxgk.linear.linear_rhs_electrostatic_species_hermite_sharded(G: Array, cache: LinearCache, params: LinearParams, *, terms: LinearTerms | None = None, dt: Array | float | 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.

spectraxgk.linear.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.

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

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

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

spectraxgk.linear.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.

spectraxgk.linear.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.

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

Convert LinearTerms into the modular TermConfig.

spectraxgk.linear.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.

spectraxgk.linear.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 species-dependent inputs directly from host memory before JIT.

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

Solve electrostatic quasineutrality for phi with optional adiabatic closure.

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

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

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

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

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

Convert modular TermConfig into linear-only term weights.

spectraxgk.linear.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 Linked Boundaries

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

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

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

spectraxgk.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) → tuple[tuple[Array, ...], tuple[Array, ...]]

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

Linear Cache

Linear-operator cache data model.

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

Precomputed arrays for the linear operator.

Array factories and damping factors for the linear cache.

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

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

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

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

spectraxgk.operators.linear.cache_arrays._shift_axis_for_cache(arr: Array, offset: int, axis: int) → Array

Shift an array along one axis with zeros introduced at the boundary.

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

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

Assemble benchmark-compatible hypercollision damping factors.

Geometry-dependent construction of LinearCache.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Linear Moments

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

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

Multiply Hermite coefficients by v_parallel (ladder form).

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

Multiply Hermite coefficients by v_parallel^2.

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

Multiply Laguerre coefficients by the perpendicular energy variable.

spectraxgk.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) → 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.

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

spectraxgk.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)).

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

spectraxgk.operators.linear.moments.grad_z_periodic(f: Array, dz: float | Array | None = None, kz: Array | None = None) → Array

Spectral periodic derivative along the last axis.

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

Parallel streaming operator acting on the Hermite index.

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

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

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

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

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

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

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

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

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

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

Convert LinearTerms into the modular TermConfig.

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

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

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

spectraxgk.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, dt: Array | float | None = None, 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 spectraxgk.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

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

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

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

Apply Frei, Ernst & Ricci (2022), equations (C9a)–(C9f).

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

Apply the linearized drift-kinetic Dougherty moment operator.

This is Appendix C, equation (C6), of Frei, Hoffmann & Ricci (2022), mapped to SPECTRAX-GK’s (species, ell, m, ky, kx, z) ordering and Laguerre-sign convention. The density and parallel-flow moments and the combined thermal moment sqrt(2) G[0, 2] + 2 G[1, 0] are exact null directions. Five-dimensional single-species states are also accepted.

The kernel is both an independently auditable reference and a usable long-wavelength collision operator. It is not the finite-Larmor-radius Sugama or Coulomb operator.

spectraxgk.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).

spectraxgk.operators.linear.dissipation.hyperdiffusion_contribution(G: Array, *, kx: Array, ky: Array, dealias_mask: Array, D_hyper: Array, p_hyper_kperp: Array, weight: Array) → Array

Hyperdiffusion in k_perp following Laguerre-Hermite conventions.

Linear Implicit Solvers

Implicit linear solve policies for cache-backed gyrokinetic operators.

spectraxgk.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) → tuple[Array, Array]

Implicit linear integrator using GMRES with a diagonal preconditioner.

Linear Integrators

Linear time-integration and diagnostic sampling policies.

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

spectraxgk.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) → tuple[Array, Array]

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

spectraxgk.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) → tuple[Array, Array, Array] | tuple[Array, Array, Array, Array]

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

Linear Diagnostic Integration

Diagnostic sampling integration for linear fixed-step solves.

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

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

Linear Parallel RHS

Velocity-parallel linear RHS helpers.

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

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

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

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

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

spectraxgk.solvers.linear.parallel.linear_rhs_electrostatic_species_hermite_sharded(G: Array, cache: LinearCache, params: LinearParams, *, terms: LinearTerms | None = None, dt: Array | float | 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.

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

spectraxgk.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, dt: Array | float | None = None) → tuple[Array, Array]

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

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

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

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

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

spectraxgk.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 species-dependent inputs directly from host memory before JIT.

Linear Parallel Policy

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

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

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

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

spectraxgk.solvers.linear.parallel_streaming.linear_rhs_electrostatic_species_hermite_sharded(G: Array, cache: LinearCache, params: LinearParams, *, terms: LinearTerms | None = None, dt: Array | float | 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.

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

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

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

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

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

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

spectraxgk.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 species-dependent inputs directly from host memory before JIT.

Linear Krylov Solvers

Public Krylov solver facade for linear gyrokinetic eigenmodes.

The compiled kernels live in focused eigenmode modules so that branch selection, operator application, preconditioning, and Arnoldi iterations can be tested and optimized independently. This facade keeps the documented script import path and the monkeypatch seams used by benchmark/runtime tests.

class spectraxgk.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 = 'propagator', power_iters: int = 200, power_dt: float = 0.01, 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 = 'damping', shift_selection: str = 'targeted', shift_outer_residual_tol: float = 0.1, mode_family: str = 'auto', fallback_method: str = 'propagator', fallback_real_floor: float = -1e-06, continuation: bool = False, continuation_selection: str = 'overlap')

Controls for the Krylov-based eigen solver.

spectraxgk.solvers.linear.krylov._build_shift_invert_precond(v: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, sigma: Array, mode: str | None) → tuple[Array | None, Callable[[Array], Array] | None]

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

spectraxgk.solvers.linear.krylov._mode_family_sign(mode_family: str) → int

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

spectraxgk.solvers.linear.krylov._omega_scale(cache: LinearCache, params: LinearParams) → Array

Return the frequency scale used for branch and target selection.

spectraxgk.solvers.linear.krylov._physical_omega(imag_part: Array) → Array

Map eigenvalue imaginary part to reported physical frequency.

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

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

spectraxgk.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 = 'power', power_iters: int = 40, power_dt: float = 0.01, 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 = 'damping', shift_selection: str = 'targeted', shift_outer_residual_tol: float = 0.1, mode_family: str = 'auto', fallback_method: str = 'propagator', fallback_real_floor: float = -1e-06, status_callback: Callable[[str], None] | None = None) → tuple[Array, Array]

Python wrapper for the cached Krylov solver.

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

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

spectraxgk.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, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, select_overlap: bool) → tuple[Array, Array]

Arnoldi on a stable IMEX2 propagator; eigenvalue from Rayleigh quotient.

spectraxgk.solvers.linear.krylov.dominant_eigenpair_shift_invert_cached(v0: Array, v_ref: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, krylov_dim: int, restarts: int, sigma: Array, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, gmres_tol: float, gmres_maxiter: int, gmres_restart: int, gmres_solve_method: str, shift_preconditioner: str | None, select_targeted: bool, select_growth: bool, select_overlap: bool) → tuple[Array, Array]

Restarted shift-invert Arnoldi with GMRES solves.

Linear Eigenmode Solver Internals

spectraxgk.solvers.linear.krylov owns KrylovConfig and the public status-reporting wrapper, while the focused helper modules below own operator application, branch selection, shift-invert preconditioning, and Arnoldi iterations.

Matrix-free linear operator applications used by eigenmode solvers.

Compiled Krylov algorithms for linear eigenmode extraction.

spectraxgk.solvers.linear.krylov_algorithms._build_shift_invert_precond(v: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, sigma: Array, mode: str | None) → tuple[Array | None, Callable[[Array], Array] | None]

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

spectraxgk.solvers.linear.krylov_algorithms._mode_family_sign(mode_family: str) → int

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

spectraxgk.solvers.linear.krylov_algorithms._omega_scale(cache: LinearCache, params: LinearParams) → Array

Return the frequency scale used for branch and target selection.

spectraxgk.solvers.linear.krylov_algorithms._physical_omega(imag_part: Array) → Array

Map eigenvalue imaginary part to reported physical frequency.

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

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

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

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

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

spectraxgk.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, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, select_overlap: bool) → tuple[Array, Array]

Arnoldi on a stable IMEX2 propagator; eigenvalue from Rayleigh quotient.

spectraxgk.solvers.linear.krylov_algorithms.dominant_eigenpair_shift_invert_cached(v0: Array, v_ref: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, *, krylov_dim: int, restarts: int, sigma: Array, omega_min_factor: float, omega_target_factor: float, omega_cap_factor: float, omega_sign: int, gmres_tol: float, gmres_maxiter: int, gmres_restart: int, gmres_solve_method: str, shift_preconditioner: str | None, select_targeted: bool, select_growth: bool, select_overlap: bool) → tuple[Array, Array]

Restarted shift-invert Arnoldi with GMRES solves.

Nonlinear

Public nonlinear gyrokinetic facade.

Implementation lives in focused nonlinear core, diagnostic, operator, and solver modules. This facade keeps the stable import surface small while making the owner modules explicit for tests and development.

class spectraxgk.nonlinear.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 spectraxgk.nonlinear.NonlinearCollisionSplitPolicy(active: bool, rhs_terms: TermConfig, damping: Array | None)

Collision split settings shared by explicit and IMEX diagnostics.

class spectraxgk.nonlinear.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 spectraxgk.nonlinear.NonlinearTimeStepPolicy(dt_init: Array, progress_total: Array, update_dt: Callable[[FieldState, Array], Array])

Initial step, progress horizon, and adaptive update callable.

class spectraxgk.nonlinear.PreparedExplicitNonlinearDiagnostics(initial_state: Array, geometry: Any, cache: LinearCache, params: LinearParams, _run_raw: Callable[[Array], 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)

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.

run(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | 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) → 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.

class spectraxgk.nonlinear.ShearedTransportTrace(final_state: jnp.ndarray, time: jnp.ndarray, heat_flux: jnp.ndarray)

Final state and compact heat-flux history from a sheared run.

final_state: Array

Alias for field number 0

heat_flux: Array

Alias for field number 2

time: Array

Alias for field number 1

class spectraxgk.nonlinear.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.

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

spectraxgk.nonlinear.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.

spectraxgk.nonlinear.build_nonlinear_collision_split_policy(cache: ~spectraxgk.operators.linear.cache_model.LinearCache, params: ~spectraxgk.operators.linear.params.LinearParams, term_cfg: ~spectraxgk.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.

spectraxgk.nonlinear.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.

spectraxgk.nonlinear.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.

spectraxgk.nonlinear.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) → SimulationDiagnostics

Build sampled nonlinear diagnostics from the raw scan output tuple.

spectraxgk.nonlinear.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[[int], float], cfl_frequency_components_fn: Callable[[...], tuple[Array, Array]]) → NonlinearTimeStepPolicy

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

spectraxgk.nonlinear.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) → SimulationDiagnostics

Package raw nonlinear scan diagnostics after applying output sampling.

spectraxgk.nonlinear.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) → tuple[Array, FieldState] | Array

Integrate the nonlinear system using built-in cache construction.

spectraxgk.nonlinear.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) → tuple[Array, FieldState] | Array

Integrate the nonlinear system using a cached geometry object.

spectraxgk.nonlinear.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, 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]

Integrate nonlinear system and return runtime diagnostics.

spectraxgk.nonlinear.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, 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 nonlinear system and return runtime diagnostics plus the final state.

spectraxgk.nonlinear.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) → tuple[Array, FieldState]

IMEX integrator: implicit linear operator, explicit nonlinear term.

spectraxgk.nonlinear.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) → tuple[Array, SimulationDiagnostics]

IMEX nonlinear integrator with runtime diagnostics.

spectraxgk.nonlinear.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) → tuple[Array, FieldState] | Array

Integrate the standard-flux-tube shearing-coordinate foundation.

This research path supports fixed-step Euler, 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.

spectraxgk.nonlinear.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) → 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.

spectraxgk.nonlinear.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.

spectraxgk.nonlinear.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.

spectraxgk.nonlinear.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().

spectraxgk.nonlinear.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.

spectraxgk.nonlinear.sampled_scan_intervals(length: int, stride: int) → ndarray

Return positive scan intervals that retain the requested final sample.

spectraxgk.nonlinear.select_nonlinear_step_diagnostics(idx: Any, *, diagnostics_stride: int, diag_prev: Any, compute_diag_fn: Any) → Any

Return a fresh or reused nonlinear step diagnostic tuple.

Nonlinear Diagnostics

Diagnostic packing and sampling helpers for nonlinear integrations.

spectraxgk.operators.nonlinear.diagnostics._sample_indices_with_final(length: int, stride: int) → slice | ndarray

Return strided sample indices while always retaining the final step.

spectraxgk.operators.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) → SimulationDiagnostics

Build sampled nonlinear diagnostics from the raw scan output tuple.

spectraxgk.operators.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) → SimulationDiagnostics

Package raw nonlinear scan diagnostics after applying output sampling.

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

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

spectraxgk.operators.nonlinear.diagnostics.sampled_scan_intervals(length: int, stride: int) → ndarray

Return positive scan intervals that retain the requested final sample.

spectraxgk.operators.nonlinear.diagnostics.select_nonlinear_step_diagnostics(idx: Any, *, diagnostics_stride: int, diag_prev: Any, compute_diag_fn: Any) → Any

Return a fresh or reused nonlinear step diagnostic tuple.

Nonlinear Diagnostic State

Diagnostic tuple assembly for nonlinear time integration.

The nonlinear integrators keep their public facade in spectraxgk.nonlinear. This module owns the pure state-to-diagnostics assembly logic, with diagnostic kernels injected by the facade so tests and interactive debugging can still monkeypatch the public module-level functions.

class spectraxgk.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_resolved_species: Callable[[...], Any], heat_flux_channel_resolved_species: Callable[[...], Any], particle_flux_species: Callable[[...], Any], particle_flux_resolved_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.

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

spectraxgk.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 spectraxgk.operators.nonlinear.collisions.NonlinearCollisionSplitPolicy(active: bool, rhs_terms: TermConfig, damping: Array | None)

Collision split settings shared by explicit and IMEX diagnostics.

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

Apply a diagonal collision/hypercollision split update.

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

spectraxgk.operators.nonlinear.collisions.build_nonlinear_collision_split_policy(cache: ~spectraxgk.operators.linear.cache_model.LinearCache, params: ~spectraxgk.operators.linear.params.LinearParams, term_cfg: ~spectraxgk.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 spectraxgk.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 spectraxgk.operators.nonlinear.policies.NonlinearCollisionSplitPolicy(active: bool, rhs_terms: TermConfig, damping: Array | None)

Collision split settings shared by explicit and IMEX diagnostics.

class spectraxgk.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 spectraxgk.operators.nonlinear.policies.NonlinearTimeStepPolicy(dt_init: Array, progress_total: Array, update_dt: Callable[[FieldState, Array], Array])

Initial step, progress horizon, and adaptive update callable.

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

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

spectraxgk.operators.nonlinear.policies.build_nonlinear_collision_split_policy(cache: ~spectraxgk.operators.linear.cache_model.LinearCache, params: ~spectraxgk.operators.linear.params.LinearParams, term_cfg: ~spectraxgk.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.

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

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

spectraxgk.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[[int], 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 spectraxgk.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

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

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

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

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

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

Nonlinear RHS

Nonlinear RHS assembly helpers.

This module owns the performance-critical RHS composition path. The public spectraxgk.nonlinear facade passes its 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.

spectraxgk.operators.nonlinear.rhs.linear_rhs_jit_for_terms_impl(term_cfg: ~spectraxgk.terms.config.TermConfig, *, electrostatic_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~spectraxgk.terms.config.FieldState]] = <PjitFunction of <function assemble_rhs_cached_electrostatic_jit>>, full_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~spectraxgk.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 path compatible with term_cfg.

spectraxgk.operators.nonlinear.rhs.nonlinear_em_term_cached_impl(G: ~jax.jaxlib._jax.Array, cache: ~spectraxgk.operators.linear.cache_model.LinearCache, params: ~spectraxgk.operators.linear.params.LinearParams, term_cfg: ~spectraxgk.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[[...], ~spectraxgk.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.

spectraxgk.operators.nonlinear.rhs.nonlinear_rhs_cached_impl(G: ~jax.jaxlib._jax.Array, cache: ~spectraxgk.operators.linear.cache_model.LinearCache, params: ~spectraxgk.operators.linear.params.LinearParams, terms: ~spectraxgk.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: ~spectraxgk.core.extension_points.CollisionOperator | None = None, radial_phase: ~jax.jaxlib._jax.Array | None = None, differentiable: bool = False, electrostatic_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~spectraxgk.terms.config.FieldState]] = <PjitFunction of <function assemble_rhs_cached_electrostatic_jit>>, full_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~spectraxgk.terms.config.FieldState]] = <function assemble_rhs_cached_jit>, differentiable_rhs_fn: ~typing.Callable[[...], tuple[~jax.jaxlib._jax.Array, ~spectraxgk.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.

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

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

Transform a shearing wave after applying its fractional radial phase.

Nonlinear Term Assembly

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

class spectraxgk.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 spectraxgk.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', radial_phase: 'jnp.ndarray | None')

class spectraxgk.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 spectraxgk.terms.nonlinear._PreparedNonlinearPath(prep: '_PreparedNonlinearInputs', laguerre: '_LaguerreGridContext | None', electrostatic_only: 'bool')

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

spectraxgk.terms.nonlinear.exb_nonlinear_contribution(G: Array, *, phi: Array, dealias_mask: Array, kx_grid: Array, ky_grid: Array, weight: Array, compressed_real_fft: bool = True, radial_phase: Array | None = None) → Array

Return the nonlinear E×B contribution using a pseudospectral bracket.

spectraxgk.terms.nonlinear.nonlinear_em_components(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, laguerre_mode: str = 'grid', radial_phase: Array | None = None) → dict[str, Array]

Return nonlinear E×B/flutter components for diagnostics/comparison checks.

spectraxgk.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, 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.

spectraxgk.terms.nonlinear.placeholder_nonlinear_contribution(G: Array, *, weight: Array) → Array

Return a zero contribution for shape-only tests and disabled-term paths.

Nonlinear Explicit Step

Explicit nonlinear time-step policies.

This module keeps the RK/SSP/K10 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.

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

spectraxgk.solvers.nonlinear.explicit.checkpoint_explicit_step(step: Callable[[...], object], checkpoint: bool)

Apply JAX checkpointing to an explicit scan step when requested.

spectraxgk.solvers.nonlinear.explicit.integrate_cached_explicit_scan(G0: Array, dt: float, steps: int, *, method: str, rhs_fn: Callable[[...], tuple[Array, object]], scan_fn: Callable[[...], tuple[Array, Any]], 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

Run a cached explicit nonlinear scan with injected RHS and projection.

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

spectraxgk.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) → 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.

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

Core nonlinear RHS and cached integrator drivers.

class spectraxgk.solvers.nonlinear.state_integration.ShearedTransportTrace(final_state: jnp.ndarray, time: jnp.ndarray, heat_flux: jnp.ndarray)

Final state and compact heat-flux history from a sheared run.

_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

time: Array

Alias for field number 1

spectraxgk.solvers.nonlinear.state_integration._linear_rhs_jit_for_terms(term_cfg: TermConfig)

Return the narrowest compiled linear RHS path compatible with term_cfg.

spectraxgk.solvers.nonlinear.state_integration.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) → tuple[Array, FieldState] | Array

Integrate the nonlinear system using built-in cache construction.

spectraxgk.solvers.nonlinear.state_integration.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) → tuple[Array, FieldState] | Array

Integrate the nonlinear system using a cached geometry object.

spectraxgk.solvers.nonlinear.state_integration.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) → tuple[Array, FieldState]

IMEX integrator: implicit linear operator, explicit nonlinear term.

spectraxgk.solvers.nonlinear.state_integration.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) → tuple[Array, FieldState] | Array

Integrate the standard-flux-tube shearing-coordinate foundation.

This research path supports fixed-step Euler, 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.

spectraxgk.solvers.nonlinear.state_integration.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) → 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.

spectraxgk.solvers.nonlinear.state_integration.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.

Nonlinear Diagnostic Drivers

Nonlinear diagnostic integration drivers.

This module owns the orchestration for explicit nonlinear diagnostic runs. The public spectraxgk.nonlinear facade injects the concrete kernels so tests and downstream users can still patch facade-level seams without keeping the large implementation body in the facade itself.

class spectraxgk.solvers.nonlinear.diagnostics.ExplicitNonlinearDiagnosticsDeps(ensure_geometry_fn: Callable[[...], Any], build_cache_fn: Callable[[...], Any], quadrature_weights_fn: Callable[[...], Any], omega_mask_fn: Callable[[...], Any], midplane_index_fn: Callable[[...], Any], resolve_cfl_fac_fn: Callable[[...], Any], linear_frequency_bound_fn: Callable[[...], Any], laguerre_velocity_max_fn: Callable[[...], Any], cfl_frequency_components_fn: Callable[[...], Any], collision_damping_fn: Callable[[...], Any], nonlinear_rhs_fn: Callable[[...], Any], compute_fields_fn: Callable[[...], Any], diagnostic_kernels_fn: Callable[[...], Any], build_diagnostic_setup_fn: Callable[[...], Any], build_time_step_policy_fn: Callable[[...], Any], build_collision_split_policy_fn: Callable[[...], Any], make_diagnostic_tuple_fn: Callable[[...], Any], make_explicit_step_fn: Callable[[...], Any], run_explicit_scan_fn: Callable[[...], Any], run_sampled_explicit_scan_fn: Callable[[...], Any], finalize_scan_diagnostics_fn: Callable[[...], Any], select_step_diagnostics_fn: Callable[[...], Any], emit_progress_fn: Callable[[...], Any], apply_collision_split_fn: Callable[[...], Any])

Patchable kernels used by the explicit diagnostic integrator.

class spectraxgk.solvers.nonlinear.diagnostics.IMEXNonlinearDiagnosticsDeps(ensure_geometry_fn: Callable[[...], Any], build_cache_fn: Callable[[...], Any], quadrature_weights_fn: Callable[[...], Any], omega_mask_fn: Callable[[...], Any], midplane_index_fn: Callable[[...], Any], linear_rhs_for_terms_fn: Callable[[...], Any], build_diagnostic_setup_fn: Callable[[...], Any], build_imex_operator_fn: Callable[[...], Any], build_collision_split_policy_fn: Callable[[...], Any], collision_damping_fn: Callable[[...], Any], make_imex_nonlinear_term_fn: Callable[[...], Any], make_imex_solve_step_fn: Callable[[...], Any], solve_imex_step_fn: Callable[[...], Any], make_diagnostic_tuple_fn: Callable[[...], Any], make_imex_step_fn: Callable[[...], Any], run_imex_scan_fn: Callable[[...], Any], finalize_scan_diagnostics_fn: Callable[[...], Any], select_step_diagnostics_fn: Callable[[...], Any], emit_progress_fn: Callable[[...], Any], apply_collision_split_fn: Callable[[...], Any], compute_fields_fn: Callable[[...], Any], nonlinear_term_fn: Callable[[...], Any], nonlinear_contribution_fn: Callable[[...], Any], diagnostic_kernels_fn: Callable[[...], Any])

Patchable kernels used by the IMEX diagnostic integrator.

class spectraxgk.solvers.nonlinear.diagnostics.PreparedExplicitNonlinearDiagnostics(initial_state: Array, geometry: Any, cache: LinearCache, params: LinearParams, _run_raw: Callable[[Array], 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)

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.

run(initial_state: Array | None = None, *, geometry: Any | None = None, cache: LinearCache | None = None, params: LinearParams | 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) → 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.

spectraxgk.solvers.nonlinear.diagnostics.integrate_explicit_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, deps: ExplicitNonlinearDiagnosticsDeps, 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, 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.

spectraxgk.solvers.nonlinear.diagnostics.integrate_imex_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, deps: IMEXNonlinearDiagnosticsDeps, 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) → tuple[Array, SimulationDiagnostics]

Integrate an IMEX nonlinear run and return diagnostics.

spectraxgk.solvers.nonlinear.diagnostics.prepare_explicit_nonlinear_diagnostics_impl(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, deps: ExplicitNonlinearDiagnosticsDeps, 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, 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.

Nonlinear Diagnostic Integration

Diagnostic nonlinear integration entry points and dependency wiring.

spectraxgk.solvers.nonlinear.diagnostic_integration._explicit_nonlinear_diagnostics_deps() → ExplicitNonlinearDiagnosticsDeps

Collect dependencies for explicit diagnostic integration.

spectraxgk.solvers.nonlinear.diagnostic_integration._imex_nonlinear_diagnostics_deps() → IMEXNonlinearDiagnosticsDeps

Collect dependencies for IMEX diagnostic integration.

spectraxgk.solvers.nonlinear.diagnostic_integration._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, 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 nonlinear system and return runtime diagnostics plus final state.

spectraxgk.solvers.nonlinear.diagnostic_integration._nonlinear_diagnostic_kernels() → NonlinearDiagnosticKernels

Return diagnostic kernels for dependency-injected nonlinear diagnostics.

spectraxgk.solvers.nonlinear.diagnostic_integration.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, 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]

Integrate nonlinear system and return runtime diagnostics.

spectraxgk.solvers.nonlinear.diagnostic_integration.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, 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 nonlinear system and return runtime diagnostics plus the final state.

spectraxgk.solvers.nonlinear.diagnostic_integration.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) → tuple[Array, SimulationDiagnostics]

IMEX nonlinear integrator with runtime diagnostics.

spectraxgk.solvers.nonlinear.diagnostic_integration.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().

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.

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

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

spectraxgk.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) → tuple[Array, 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.

spectraxgk.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) → Callable[[tuple[Any, Any, Any, Any, Any], Any], tuple[tuple[Any, Any, Any, Any, Any], tuple[Any, Any]]]

Build one IMEX diagnostic scan step with injected runtime seams.

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

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

spectraxgk.solvers.nonlinear.imex.run_imex_diagnostic_scan(step_fn: Callable[[tuple[Any, Any, Any, Any, Any], Any], tuple[tuple[Any, Any, Any, Any, Any], tuple[Any, Any]]], initial_carry: tuple[Any, Any, Any, Any, Any], *, steps: int, checkpoint: bool) → tuple[Array, tuple[Any, Any]]

Run the fixed-step IMEX diagnostic scan.

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

Explicit Time Integrators

Explicit linear time integrators implemented in JAX.

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

spectraxgk.solvers.time.explicit._apply_completed_step_state_mask(state: Array, cache: LinearCache) → Array

Apply the completed-step mask to a spectral state array.

spectraxgk.solvers.time.explicit._completed_step_state_mask(cache: LinearCache) → Array

Return the completed-step state-space mask applied after each completed step.

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

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

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

spectraxgk.solvers.time.explicit._rk3_heun_step(G: Array, cache: LinearCache, params: LinearParams, term_cfg, dt: float)

Single Explicit RK3/Heun step through the public explicit facade.

spectraxgk.solvers.time.explicit._rk4_step(G: Array, cache: LinearCache, params: LinearParams, term_cfg, dt: float)

Single Explicit RK4 step through the public explicit facade.

spectraxgk.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 with growth-rate diagnostics.

spectraxgk.solvers.time.explicit.integrate_linear_explicit_diagnostics(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, SimulationDiagnostics]

Public facade for diagnostics-rich explicit linear integration.

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

Explicit Step Kernels

Explicit linear step kernels and growth-rate diagnostic helpers.

spectraxgk.solvers.time.explicit_steps._apply_completed_step_state_mask(state: Array, cache: LinearCache) → Array

Apply the completed-step mask to a spectral state array.

spectraxgk.solvers.time.explicit_steps._completed_step_state_mask(cache: LinearCache) → Array

Return the completed-step state-space mask applied after each completed step.

spectraxgk.solvers.time.explicit_steps._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.

spectraxgk.solvers.time.explicit_steps._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.

spectraxgk.solvers.time.explicit_steps._linear_explicit_step(G: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, dt: float, *, method: str, assemble_rhs_cached_fn=<function assemble_rhs_cached>) → tuple[Array, FieldState]

Single explicit linear step matching explicit staged schemes.

spectraxgk.solvers.time.explicit_steps._rk3_heun_step(G: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, dt: float) → tuple[Array, FieldState]

Single Explicit RK3/Heun step for linear dynamics.

spectraxgk.solvers.time.explicit_steps._rk4_step(G: Array, cache: LinearCache, params: LinearParams, term_cfg: TermConfig, dt: float) → tuple[Array, FieldState]

Single Explicit RK4 step for linear dynamics.

Explicit Diagnostic Integrators

Diagnostics-rich explicit linear time integration.

class spectraxgk.solvers.time.explicit_diagnostics.ExplicitTimeConfigLike(*args, **kwargs)

Runtime fields required by explicit diagnostic integration.

spectraxgk.solvers.time.explicit_diagnostics.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]

Explicit time integrator with growth-rate plus energy/flux diagnostics.

Explicit CFL Policy

CFL and linear-frequency bound policy for explicit time integration.

Diffrax Time Integrators

Shared Diffrax dependency, solver-policy, and state helper routines.

Diffrax linear time integration with saved fields and mode traces.

spectraxgk.solvers.time.diffrax_linear.integrate_linear_diffrax(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'Dopri8', cache: LinearCache | None = None, terms: LinearTerms | None = None, adaptive: bool = False, rtol: float = 1e-05, atol: float = 1e-07, max_steps: int = 4096, show_progress: bool = False, progress_bar: bool = False, checkpoint: bool = False, jit: bool | None = None, sample_stride: int = 1, return_state: bool = True, save_mode: ModeSelection | ModeSelectionBatch | None = None, mode_method: str = 'z_index', save_field: str = 'phi', density_species_index: int | None = None, state_sharding: Any | None = None, derivative_mode: str = 'reverse') → tuple[Array | None, Array]

Integrate the linear system with an explicit AD policy.

derivative_mode="reverse" preserves the custom-VJP field solve used by scalar objectives. For adaptive reverse gradients, set checkpoint=True to select Diffrax’s bounded-memory recursive checkpoint adjoint. "forward" selects native JAX rules so low-dimensional JVPs can pass through fixed or adaptive Diffrax trajectories.

Streaming Diffrax linear growth/frequency estimators.

spectraxgk.solvers.time.diffrax_streaming.integrate_linear_diffrax_streaming(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'Dopri8', cache: LinearCache | None = None, terms: LinearTerms | None = None, adaptive: bool = False, rtol: float = 1e-05, atol: float = 1e-07, max_steps: int = 4096, show_progress: bool = False, progress_bar: bool = False, checkpoint: bool = False, jit: bool | None = None, tmin: float | None = None, tmax: float | None = None, fit_signal: str = 'density', mode_ky_indices: Sequence[int] | ndarray | Array | None = None, mode_kx_index: int = 0, mode_z_index: int = 0, mode_method: str = 'z_index', amp_floor: float = 1e-30, density_species_index: int | None = None, return_state: bool = True, state_sharding: Any | None = None) → tuple[Array | None, Array, Array]

Integrate the linear system and stream a growth-rate fit without storing time series.

Diffrax nonlinear time integration paths.

spectraxgk.solvers.time.diffrax_nonlinear.integrate_nonlinear_diffrax(G0: Array, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, dt: float, steps: int, *, method: str = 'KenCarp4', cache: LinearCache | None = None, terms: TermConfig | None = None, adaptive: bool = False, rtol: float = 1e-05, atol: float = 1e-07, max_steps: int = 4096, show_progress: bool = False, progress_bar: bool = False, checkpoint: bool = False, jit: bool | None = None, state_sharding: Any | None = None, compressed_real_fft: bool = True, laguerre_mode: str = 'grid') → tuple[Array, FieldState]

Integrate the nonlinear system with diffrax.

Config-Driven Time Runners

Config-driven runners for time integration.

spectraxgk.solvers.time.runners._validate_nonlinear_config_state_sharding(spec: str | None) → None

Keep config-level nonlinear sharding on release-gated state axes.

spectraxgk.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, save_mode: ModeSelection | ModeSelectionBatch | None = None, mode_method: str = 'z_index', save_field: str = 'phi', density_species_index: int | None = None, show_progress: bool | None = None, parallel: Any | None = None) → tuple

Integrate the linear system using TimeConfig settings.

spectraxgk.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) → tuple

Integrate the nonlinear system using TimeConfig settings.

Runtime Execution Dispatch

Runtime execution dispatch for linear and nonlinear configured runs.

class spectraxgk.workflows.runtime.execution.RuntimeLinearDispatchDeps(resolve_runtime_hl_dims: Callable[[...], tuple[int, int]], run_full_linear_runtime: Callable[[...], RuntimeLinearResult], full_deps: Any)

Patchable dependencies for one configured linear runtime run.

class spectraxgk.workflows.runtime.execution.RuntimeNonlinearDispatchDeps(resolve_runtime_hl_dims: Callable[[...], tuple[int, int]], run_full_nonlinear_runtime: Callable[[...], RuntimeNonlinearResult], full_deps: Any)

Patchable dependencies for one configured nonlinear runtime run.

class spectraxgk.workflows.runtime.execution._RuntimeLinearRequest(cfg: 'RuntimeConfig', ky_target: 'float', Nl: 'int | None', Nm: 'int | None', solver: 'str', 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', krylov_cfg: 'Any', mode_method: 'str', fit_signal: 'str', return_state: 'bool', initial_state: 'Any | None', show_progress: 'bool', status_callback: 'Callable[[str], None] | None', deps: 'RuntimeLinearDispatchDeps')

spectraxgk.workflows.runtime.execution.build_runtime_linear_dispatch_deps(scope: Any) → RuntimeLinearDispatchDeps

Build linear dispatch dependencies from a patchable runtime facade scope.

spectraxgk.workflows.runtime.execution.build_runtime_nonlinear_dispatch_deps(scope: Any) → RuntimeNonlinearDispatchDeps

Build nonlinear dispatch dependencies from a patchable runtime facade scope.

spectraxgk.workflows.runtime.execution.run_runtime_linear_impl(cfg: RuntimeConfig, *, ky_target: float = 0.3, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: Any = 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, deps: RuntimeLinearDispatchDeps) → RuntimeLinearResult

Run one linear point from a case-agnostic runtime config.

spectraxgk.workflows.runtime.execution.run_runtime_nonlinear_impl(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, deps: RuntimeNonlinearDispatchDeps) → RuntimeNonlinearResult

Run one nonlinear point from a case-agnostic runtime config.

Nonlinear Replicate Diagnostics

Diagnostics for replicated nonlinear transport-window spread.

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

Thresholds for classifying replicated nonlinear-window spread.

class spectraxgk.diagnostics.nonlinear_replicates.NonlinearWindowEnsembleManifestConfig(min_replicates_per_case: int = 2, required_variant_axes: tuple[str, ...] = ('seed', 'timestep'), require_observed_windows_ready: bool = True)

Artifact requirements before a replicated nonlinear ensemble can run.

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

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

Parameters:

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

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

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

spectraxgk.diagnostics.nonlinear_replicates.nonlinear_window_ensemble_artifact_manifest(records: Sequence[dict[str, Any]], *, case: str = 'nonlinear_window_ensemble_artifact_manifest', config: NonlinearWindowEnsembleManifestConfig | None = None) → dict[str, Any]

Return a promotion-blocking manifest for missing ensemble artifacts.

Each record should contain a report produced by nonlinear_window_convergence_report(), plus optional variant metadata such as {"seed": 1, "timestep": 0.02}. The manifest is intentionally conservative: a production nonlinear optimization promotion needs distinct passed artifacts for every required variant axis, so a single late-window summary is recorded as useful convergence evidence but not as replicated-ensemble evidence.

spectraxgk.diagnostics.nonlinear_replicates.summarize_window_evidence(window_artifacts: Sequence[dict[str, Any]], *, paths: Sequence[str | None] | None = None, config: NonlinearTurbulenceGradientEvidenceConfig | None = None) → dict[str, Any]

Summarize replicated long-window uncertainty evidence.

Existing nonlinear_window_ensemble_report artifacts are consumed directly. If only individual nonlinear_window_convergence_report summaries are supplied, a derived ensemble is built from those summaries using the configured uncertainty limits.

Nonlinear Replicate Follow-Up

Campaign follow-up planning for nonlinear replicate-spread blockers.

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

Options controlling the targeted replicate follow-up plan.

tools.campaigns.nonlinear_replicate_followup.compact_ensemble_payload(payload: dict[str, Any], *, output_gate_json: str, netcdf_root: str) → dict[str, Any]

Return a compact-provenance copy of a replicated ensemble payload.

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

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

Nonlinear Transport Optimization Diagnostics

Nonlinear turbulent-transport optimization promotion diagnostics.

These helpers consume already-generated nonlinear transport artifacts and keep release-scope diagnostics separate from production turbulent-flux optimization claims. They are data-only and do not launch simulations.

class spectraxgk.diagnostics.nonlinear_transport_optimization.ProductionNonlinearOptimizationGuardConfig(min_replicated_ensembles: int = 2, min_reports_per_ensemble: int = 2, max_mean_rel_spread: float = 0.15, max_combined_sem_rel: float = 0.25, require_optimized_equilibrium_transport: bool = True, require_matched_optimized_transport_audit: bool = True, min_optimized_equilibrium_ensembles: int = 3, min_matched_optimized_audits: int = 3, require_seed_timestep_provenance: bool = True, min_seed_variants: int = 2, min_timestep_variants: int = 1, min_matched_optimized_relative_reduction: float = 0.05, min_matched_optimized_uncertainty_sigma: float = 1.0, value_floor: float = 1e-12)

Strict gate settings for production nonlinear optimization promotion.

validate() → None

Raise if the guard configuration is inconsistent.

spectraxgk.diagnostics.nonlinear_transport_optimization.matched_optimized_transport_report(path: str, payload: Mapping[str, Any], *, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Return whether a matched baseline-to-optimized audit promotes transport.

spectraxgk.diagnostics.nonlinear_transport_optimization.optimization_artifact_reduction_scope(payload: Mapping[str, Any]) → dict[str, Any]

Return scope metadata for the differentiable optimization comparison.

spectraxgk.diagnostics.nonlinear_transport_optimization.optimized_equilibrium_transport_report(path: str, payload: Mapping[str, Any], *, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Return whether an artifact can promote optimized-equilibrium transport.

spectraxgk.diagnostics.nonlinear_transport_optimization.production_nonlinear_optimization_guard_report(*, optimization_artifact: Mapping[str, Any] | None, optimization_artifact_path: str = '', reduced_artifacts: Mapping[str, Mapping[str, Any]] | None = None, replicated_ensemble_artifacts: Mapping[str, Mapping[str, Any]] | None = None, optimized_equilibrium_artifacts: Mapping[str, Mapping[str, Any]] | None = None, matched_optimized_transport_artifacts: Mapping[str, Mapping[str, Any]] | None = None, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Build the fail-closed nonlinear turbulent-flux optimization guard.

The top-level passed field means the release is safe: reduced/startup artifacts are correctly scoped and long-window replicated holdouts are present. It does not mean production nonlinear optimization is promoted; that is reported separately by production_nonlinear_optimization_promoted.

spectraxgk.diagnostics.nonlinear_transport_optimization.reduced_artifact_scope_report(path: str, payload: Mapping[str, Any]) → dict[str, Any]

Return whether a startup/reduced artifact is safely blocked from promotion.

spectraxgk.diagnostics.nonlinear_transport_optimization.replicated_transport_ensemble_report(path: str, payload: Mapping[str, Any], *, config: ProductionNonlinearOptimizationGuardConfig | None = None) → dict[str, Any]

Return quality metadata for a long-window replicated transport ensemble.

Nonlinear Gradient Evidence Diagnostics

Nonlinear turbulent-gradient evidence and claim-boundary diagnostics.

This module contains pure JSON/dictionary diagnostics used to decide whether long-window nonlinear transport-gradient artifacts support production claims. It does not launch simulations. Campaign design and follow-up planning live in tools.campaigns.nonlinear_gradient_followup.

class spectraxgk.diagnostics.nonlinear_gradient_evidence.NonlinearTurbulenceGradientBracketSweepConfig(max_gradient_uncertainty_rel: float = 0.5, max_fd_asymmetry_rel: float = 0.5, max_fd_condition_number: float = 100000000.0, min_fd_response_fraction: float = 0.03, max_repeated_bracket_uncertainty_rel: float = 0.75, min_repeated_bracket_same_sign_fraction: float = 0.8, score_cap: float = 2.0, value_floor: float = 1e-12)

Decision limits for same-control perturbation-amplitude sweeps.

class spectraxgk.diagnostics.nonlinear_gradient_evidence.NonlinearTurbulenceGradientCandidateRankingConfig(max_gradient_uncertainty_rel: float = 0.5, max_fd_asymmetry_rel: float = 0.5, max_fd_condition_number: float = 100000000.0, min_fd_response_fraction: float = 0.03, score_cap: float = 2.0, value_floor: float = 1e-12, campaign_context: str = 'single_control_screen')

Scoring limits used to rank failed nonlinear-gradient control candidates.

class spectraxgk.diagnostics.nonlinear_gradient_evidence.NonlinearTurbulenceGradientEvidenceConfig(min_window_reports: int = 2, max_window_mean_rel_spread: float = 0.15, max_window_combined_sem_rel: float = 0.25, max_gradient_uncertainty_rel: float = 0.5, max_fd_asymmetry_rel: float = 0.5, max_fd_condition_number: float = 100000000.0, min_fd_response_fraction: float = 0.03, value_floor: float = 1e-12)

Acceptance limits for production nonlinear turbulence-gradient evidence.

class spectraxgk.diagnostics.nonlinear_gradient_evidence.NonlinearTurbulenceGradientFiniteDifferenceConfig(min_window_reports: int = 2, max_window_mean_rel_spread: float = 0.15, max_window_combined_sem_rel: float = 0.25, max_gradient_uncertainty_rel: float = 0.5, max_fd_asymmetry_rel: float = 0.5, max_fd_condition_number: float = 100000000.0, min_fd_response_fraction: float = 0.03, value_floor: float = 1e-12)

Acceptance limits for paired long-window finite-difference gradients.

class spectraxgk.diagnostics.nonlinear_gradient_evidence.NonlinearTurbulenceGradientGapConfig(case_slug: str = 'optimized_equilibrium_turbulence_gradient', parameter_name: str = 'vmec_state_control_or_profile_gradient', perturbation_fraction: float = 0.05, t_start: float = 0.0, analysis_tmin: float = 350.0, analysis_tmax: float = 700.0, minimum_tmax: float = 700.0, minimum_grid: str = 'n64x64x64x40x40', replicate_labels: tuple[str, ...] = ('seed31', 'seed32', 'dt0p04'))

Default campaign shape required before promoting turbulence gradients.

spectraxgk.diagnostics.nonlinear_gradient_evidence.classify_gradient_artifact(payload: dict[str, Any], *, path: str | None = None, config: NonlinearTurbulenceGradientEvidenceConfig | None = None) → dict[str, Any]

Classify a gradient/FD artifact without promoting ambiguous evidence.

spectraxgk.diagnostics.nonlinear_gradient_evidence.load_json_artifact(path: str | Path) → dict[str, Any]

Load a JSON object artifact.

spectraxgk.diagnostics.nonlinear_gradient_evidence.nonlinear_turbulence_gradient_bracket_sweep_report(artifacts: Sequence[dict[str, Any]], *, labels: Sequence[str | None] | None = None, paths: Sequence[str | None] | None = None, config: NonlinearTurbulenceGradientBracketSweepConfig | None = None) → dict[str, Any]

Summarize a same-control perturbation-amplitude sweep.

This is a planning/claim-boundary utility. It does not promote nonlinear turbulence-gradient evidence unless an input finite-difference artifact already passes the production long-window gate. Its main purpose is to decide whether the next expensive campaign should add replicas at the same bracket, change the perturbation amplitude, or move to an overdetermined profile-gradient direction.

spectraxgk.diagnostics.nonlinear_gradient_evidence.nonlinear_turbulence_gradient_candidate_ranking_report(artifacts: Sequence[dict[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, config: NonlinearTurbulenceGradientCandidateRankingConfig | None = None) → dict[str, Any]

Rank nonlinear turbulence-gradient candidates without promoting failures.

The ranking is a planning aid, not a replacement for the production gate. It scores each candidate by the weakest normalized evidence margin across response, locality, conditioning, and uncertainty. This makes the next campaign choice explicit: candidates with complementary failures should move to a profile-gradient or overdetermined least-squares design instead of repeating a single boundary coefficient indefinitely.

spectraxgk.diagnostics.nonlinear_gradient_evidence.nonlinear_turbulence_gradient_evidence_gap_report(evidence_report: dict[str, Any], *, config: NonlinearTurbulenceGradientEvidenceConfig | None = None, gap_config: NonlinearTurbulenceGradientGapConfig | None = None) → dict[str, Any]

Return the fail-closed run campaign needed to close gradient evidence.

The report is deliberately prescriptive: it requires paired plus/minus long-window nonlinear runs with the same seeds, timestep variant, grid, and post-transient analysis window before a finite-difference turbulence gradient can be promoted. It does not infer a gradient from standalone replicated transport windows.

spectraxgk.diagnostics.nonlinear_gradient_evidence.nonlinear_turbulence_gradient_evidence_report(gradient_artifact: dict[str, Any], *, window_artifacts: Sequence[dict[str, Any]] = (), gradient_path: str | None = None, window_paths: Sequence[str | None] | None = None, config: NonlinearTurbulenceGradientEvidenceConfig | None = None, gap_config: NonlinearTurbulenceGradientGapConfig | None = None) → dict[str, Any]

Return a fail-closed production nonlinear gradient evidence report.

spectraxgk.diagnostics.nonlinear_gradient_evidence.nonlinear_turbulence_gradient_finite_difference_report(*, baseline: dict[str, Any], plus: dict[str, Any], minus: dict[str, Any], delta_parameter: float, parameter_name: str, baseline_path: str | None = None, plus_path: str | None = None, minus_path: str | None = None, config: NonlinearTurbulenceGradientFiniteDifferenceConfig | None = None) → dict[str, Any]

Build a production long-window central finite-difference gradient gate.

Inputs must be replicated nonlinear_window_ensemble_report payloads for the same nonlinear case and analysis window, differing only by the perturbed parameter. The report computes the central finite-difference heat-flux gradient and checks that the response is resolved above ensemble uncertainty before allowing any turbulence-gradient claim.

spectraxgk.diagnostics.nonlinear_gradient_evidence.summarize_window_evidence(window_artifacts: Sequence[dict[str, Any]], *, paths: Sequence[str | None] | None = None, config: NonlinearTurbulenceGradientEvidenceConfig | None = None) → dict[str, Any]

Summarize replicated long-window uncertainty evidence.

Existing nonlinear_window_ensemble_report artifacts are consumed directly. If only individual nonlinear_window_convergence_report summaries are supplied, a derived ensemble is built from those summaries using the configured uncertainty limits.

Nonlinear Gradient Campaign Follow-Up Tools

Campaign follow-up planners for nonlinear turbulent-gradient evidence.

These helpers generate deterministic planning reports and launch metadata for long-window nonlinear transport-gradient campaigns. They are repository tooling, not runtime solver functionality, so they intentionally live outside the installable spectraxgk package.

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

Conditioning limits for the next nonlinear-gradient campaign design.

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

Controls for constructing the next composite nonlinear-gradient direction.

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

Acceptance limits for an independent control-mean estimate.

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

Controls for an independent control-mean campaign.

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

Acceptance and cost controls for the follow-up planner.

class tools.campaigns.nonlinear_gradient_followup.NonlinearGradientQLSeedScreenConfig(target_objectives: tuple[str, ...] = ('mixing_length_heat_flux_proxy', 'linear_heat_flux_weight', 'gamma'), primary_objective: str = 'mixing_length_heat_flux_proxy', min_distinct_controls: int = 2, min_cases_per_control: int = 2, min_sign_consistency: float = 0.75, max_objective_rel_error: float = 0.02, min_abs_sensitivity: float = 1e-12, require_artifact_passed: bool = False)

Admission limits for quasilinear-seeded nonlinear-gradient controls.

class tools.campaigns.nonlinear_gradient_followup.NonlinearGradientStateControlRunbookConfig(min_mapped_controls: int = 2, max_mapping_condition_number: float = 1000000.0, max_mapping_relative_residual: float = 0.1, default_relative_delta: float = 0.02, require_mapping_passed: bool = True)

Admission limits for mapping VMEC-state controls to launchable inputs.

class tools.campaigns.nonlinear_gradient_followup.NonlinearGradientVarianceReductionConfig(max_paired_response_uncertainty_rel: float = 0.5, max_control_variate_uncertainty_rel: float = 0.5, min_control_variate_sem_reduction: float = 0.25, require_known_control_mean: bool = True, sem_safety_factor: float = 1.1, min_common_pairs: int = 2, max_extra_paired_seeds: int = 4, value_floor: float = 1e-12)

Controls for paired-seed/control-variate nonlinear-gradient planning.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_candidate_design_report(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_candidate_design', config: NonlinearGradientCandidateDesignConfig | None = None) → dict[str, Any]

Return the next-campaign design implied by failed production gates.

The report estimates whether a failed central-FD candidate can be rescued by a larger bracket, by bounded extra replicas, or whether the next campaign should instead change the control/observable. The estimates use the usual 1/sqrt(N) SEM scaling and the local finite-difference assumption that response grows approximately linearly with bracket size before the asymmetry gate is hit.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_composite_control_report(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_composite_control_design', config: NonlinearGradientCompositeControlConfig | None = None) → dict[str, Any]

Design a normalized VMEC-boundary direction from resolved FD candidates.

This is a launch-planning gate, not nonlinear-gradient evidence. The returned controls are the steepest-descent direction in the subspace of candidates that already pass locality, response, uncertainty, coefficient, and paired-sign checks. If fewer than min_controls survive, the report fails closed and provides exact blockers instead of producing a misleading multi-coefficient launch recommendation.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_control_mean_gate(variance_report: Mapping[str, Any], *, plus_ensemble: Mapping[str, Any], minus_ensemble: Mapping[str, Any], plus_path: str | None = None, minus_path: str | None = None, case: str = 'nonlinear_turbulence_gradient_control_mean_gate', candidate_name: str | None = None, config: NonlinearGradientControlMeanGateConfig | None = None) → dict[str, Any]

Evaluate an independent control-mean estimate for a screened CV response.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_control_variate_campaign_plan(variance_report: Mapping[str, Any], *, case: str = 'nonlinear_turbulence_gradient_control_variate_campaign', candidate_name: str | None = None, config: NonlinearGradientControlVariateCampaignConfig | None = None) → dict[str, Any]

Design an independent control-mean campaign from a variance runbook.

The input is the output of nonlinear_gradient_variance_reduction_plan(). The planner is intentionally fail-closed: a sample-centered control variate can motivate new runs, but the campaign is only considered launch-ready when an independent control-mean estimate can bring the combined uncertainty under the target gate within the configured run budget.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_followup_plan(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_followup', config: NonlinearGradientFollowupConfig | None = None) → dict[str, Any]

Build a bounded, fail-closed follow-up plan from gradient artifacts.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_ql_seed_screen_report(artifacts: Sequence[Mapping[str, Any]], *, paths: Sequence[str | None] | None = None, labels: Sequence[str | None] | None = None, case: str = 'nonlinear_turbulence_gradient_ql_seed_screen', config: NonlinearGradientQLSeedScreenConfig | None = None) → dict[str, Any]

Screen QL/linear sensitivity artifacts before nonlinear-gradient runs.

The report groups full-chain VMEC/Boozer sensitivity rows by state parameter and admits a control only when the primary objective sensitivity is AD/FD-consistent, resolved, sign-consistent across enough artifacts, and tied to a distinct VMEC-state control. The output is deliberately a planning artifact: VMEC-state controls are not assumed to be patchable RBC/ZBS input-file coefficients.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_state_control_runbook_report(ql_seed_screen: Mapping[str, Any], *, mapping_artifacts: Sequence[Mapping[str, Any]] = (), case: str = 'nonlinear_gradient_state_control_runbook', config: NonlinearGradientStateControlRunbookConfig | None = None) → dict[str, Any]

Build a fail-closed launch runbook for VMEC-state nonlinear-gradient controls.

The QL seed screen operates on internal vmec_jax state coordinates. Nonlinear campaigns, however, need perturbable input directions that can be written to VMEC inputs and re-equilibrated. This report joins the admitted state controls to an explicit state-to-input mapping artifact and refuses to produce launch commands until the mapping is conditioned and complete.

tools.campaigns.nonlinear_gradient_followup.nonlinear_gradient_variance_reduction_plan(artifact: Mapping[str, Any], *, path: str | None = None, label: str | None = None, case: str = 'nonlinear_turbulence_gradient_variance_reduction_plan', config: NonlinearGradientVarianceReductionConfig | None = None) → dict[str, Any]

Plan paired-seed/control-variate follow-up for a failed central-FD artifact.

The plan uses common seed/timestep labels across plus and minus ensembles to estimate the uncertainty of paired finite-difference responses. It is a campaign-design artifact, not nonlinear-gradient evidence.

Benchmarks

Reference data and policies for documented code-comparison workflows.

Simulation execution belongs to spectraxgk.runtime. This compact facade keeps reviewed reference tables, normalization contracts, and branch selection policies together without maintaining a second solver stack.

class spectraxgk.benchmarks.CycloneBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=24, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=20.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=150.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=True, diffrax_rtol=1e-06, diffrax_atol=1e-08, diffrax_max_steps=200000, state_sharding=None, progress_bar=False, fixed_dt=False, dt_min=1e-07, dt_max=0.05, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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(R_over_LTi=2.49, R_over_LTe=0.0, R_over_Ln=0.8, nu_i=0.0), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False), reference_aligned: bool = True)

Standard parameters for the Cyclone base case ITG benchmark.

class spectraxgk.benchmarks.CycloneComparison(ky: 'float', gamma: 'float', omega: 'float', gamma_ref: 'float', omega_ref: 'float', rel_gamma: 'float', rel_omega: 'float')

class spectraxgk.benchmarks.CycloneReference(ky: 'np.ndarray', omega: 'np.ndarray', gamma: 'np.ndarray')

class spectraxgk.benchmarks.CycloneRunResult(t: 'np.ndarray', phi_t: 'np.ndarray', gamma: 'float', omega: 'float', ky: 'float', selection: 'ModeSelection')

class spectraxgk.benchmarks.CycloneScanResult(ky: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray')

class spectraxgk.benchmarks.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.

class spectraxgk.benchmarks.KBMBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=16, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=10.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Tsit5', diffrax_adaptive=True, diffrax_rtol=0.0001, diffrax_atol=1e-07, diffrax_max_steps=20000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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(R_over_LTi=2.49, R_over_LTe=2.49, R_over_Ln=0.8, Te_over_Ti=1.0, mass_ratio=3703.7037037037035, nu_i=0.0, nu_e=0.0, beta=0.015), init: InitializationConfig = InitializationConfig(init_field='all', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for an electromagnetic KBM benchmark.

class spectraxgk.benchmarks.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 = 'propagator', power_iters: int = 200, power_dt: float = 0.01, 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 = 'damping', shift_selection: str = 'targeted', shift_outer_residual_tol: float = 0.1, mode_family: str = 'auto', fallback_method: str = 'propagator', fallback_real_floor: float = -1e-06, continuation: bool = False, continuation_selection: str = 'overlap')

Controls for the Krylov-based eigen solver.

class spectraxgk.benchmarks.LinearRunResult(t: 'np.ndarray', phi_t: 'np.ndarray', gamma: 'float', omega: 'float', ky: 'float', selection: 'ModeSelection', gamma_t: 'np.ndarray | None' = None, omega_t: 'np.ndarray | None' = None)

class spectraxgk.benchmarks.LinearScanResult(ky: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray')

class spectraxgk.benchmarks.ModeSelection(ky_index: 'int', kx_index: 'int', z_index: 'int' = 0)

class spectraxgk.benchmarks.ScanFitWindowPolicy(tmin: ~typing.Any = None, tmax: ~typing.Any = None, auto_window: bool = True, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_fraction: float = 0.8, end_fraction: float = 0.9, max_amp_fraction: float = 0.9, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0, window_method: str = 'loglinear', fit_growth_rate_fn: ~collections.abc.Callable[[...], tuple[float, float]] = <function fit_growth_rate>, fit_growth_rate_auto_fn: ~collections.abc.Callable[[...], tuple[float, float, float, float]] = <function fit_growth_rate_auto>, normalize_growth_rate_fn: ~collections.abc.Callable[[float, float, ~spectraxgk.operators.linear.params.LinearParams, str], tuple[float, float]] = <function _normalize_growth_rate>)

Window-selection and normalization policy shared by benchmark scans.

fit_growth_rate_auto_fn(signal: ndarray, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0, max_amp_fraction: float = 0.9, window_method: str = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float]

Fit gamma/omega with optional auto-selected window.

fit_growth_rate_fn(signal: ndarray, tmin: float | None = None, tmax: float | None = None) → Tuple[float, float]

Fit gamma and omega from a complex signal ~ exp((gamma - i*omega) t).

fit_signal(signal: ndarray, *, idx: int, dt: float, stride: int, params: LinearParams, diagnostic_norm: str) → tuple[float, float]

Fit one scan signal and apply the configured diagnostic normalization.

spectraxgk.benchmarks.apply_auto_fit_scan_policy(fit_key: str, *, streaming_fit: bool, mode_only: bool) → tuple[bool, bool]

Disable streaming and mode-only saves when auto signal selection needs both fields.

spectraxgk.benchmarks.compare_cyclone_to_reference(result: CycloneRunResult, reference: CycloneReference) → CycloneComparison

Compare a Cyclone run result against the reference data set.

spectraxgk.benchmarks.indexed_float_value(value: Any, idx: int) → float | None

Return a scalar or indexed scan value as float for window policies.

spectraxgk.benchmarks.indexed_scan_value(value: Any, idx: int) → Any

Return a scalar or indexed scan value while preserving non-float types.

spectraxgk.benchmarks.load_cyclone_reference() → CycloneReference

Load Cyclone base case reference data (adiabatic electrons).

spectraxgk.benchmarks.load_cyclone_reference_kinetic() → CycloneReference

Load Cyclone base case reference data (kinetic electrons).

spectraxgk.benchmarks.load_etg_reference() → CycloneReference

Load ETG reference data for the tracked two-species ETG lane.

spectraxgk.benchmarks.load_kbm_reference() → CycloneReference

Load KBM reference data (finite beta, kinetic electrons).

spectraxgk.benchmarks.load_tem_reference() → CycloneReference

Load the provisional TEM reference digitized from the literature.

This lane remains an extended stress case while the literature case definition is being reconstructed.

spectraxgk.benchmarks.normalize_fit_signal(fit_signal: str) → str

Normalize and validate benchmark fit-signal selectors.

spectraxgk.benchmarks.normalize_solver_key(solver: str) → str

Normalize a benchmark solver selector to canonical SPECTRAX-GK keys.

spectraxgk.benchmarks.resolve_scan_mode_method(mode_method: str, *, mode_only: bool) → str

Use direct mode extraction when a runner saved only a mode time series.

spectraxgk.benchmarks.scan_window_valid(t: ndarray, tmin: float | None, tmax: float | None, *, min_points: int = 2) → bool

Return whether an explicit fit window contains enough sampled points.

spectraxgk.benchmarks.select_kbm_solver_auto(solver: str, *, ky_target: float, reference_aligned: bool | None = None) → str

Return deterministic KBM solver choice for auto mode.

spectraxgk.benchmarks.should_use_ky_batch(*, ky_batch: int, solver_key: str, dt: Any, steps: Any, tmin: Any, tmax: Any) → bool

Return whether a ky scan can use a fixed-shape batch path.

Validation Gates

Validation gate metrics and report builders for diagnostics artifacts.

class spectraxgk.diagnostics.validation_gates.GateReport(case: str, source: str, gates: tuple[ScalarGateResult, ...], passed: bool, max_abs_error: float, max_rel_error: float)

Collection of scalar gates for one validation artifact.

class spectraxgk.diagnostics.validation_gates.ScalarGateResult(metric: str, observed: float, reference: float, abs_error: float, rel_error: float, atol: float, rtol: float, passed: bool, units: str, notes: str)

Pass/fail result for one benchmark observable.

The tolerance convention follows numpy.isclose: a metric passes when abs_error <= atol + rtol * abs(reference). This keeps near-zero frequency and marginal-growth gates explicit through atol rather than hiding them behind unstable relative errors.

class spectraxgk.diagnostics.validation_gates.ZonalFlowResponseMetrics(initial_level: float, initial_policy: str, residual_level: float, residual_std: float, response_rms: float, gam_frequency: float, gam_damping_rate: float, damping_method: str, frequency_method: str, peak_count: int, peak_fit_count: int, tmin: float, tmax: float, fit_tmin: float, fit_tmax: float, peak_times: ndarray, peak_envelope: ndarray, max_peak_times: ndarray, max_peak_values: ndarray, min_peak_times: ndarray, min_peak_values: ndarray)

Late-time residual and GAM-envelope metrics for zonal-flow responses.

spectraxgk.diagnostics.validation_gates.branch_continuity_gate_report(metrics: BranchContinuationMetrics, *, case: str, source: str, max_rel_gamma_jump: float, max_rel_omega_jump: float, min_successive_overlap: float | None = None) → GateReport

Gate branch-continuation diagnostics for branch-followed scans.

spectraxgk.diagnostics.validation_gates.eigenfunction_gate_report(comparison: EigenfunctionComparisonMetrics, *, case: str, source: str, min_overlap: float = 0.95, max_relative_l2: float = 0.25) → GateReport

Gate a phase-aligned eigenfunction comparison.

The ideal reference is overlap equal to one and relative L2 mismatch equal to zero. min_overlap and max_relative_l2 make the acceptance policy explicit for manuscript overlays and branch-identity checks.

spectraxgk.diagnostics.validation_gates.evaluate_scalar_gate(metric: str, observed: float, reference: float, *, atol: float, rtol: float, units: str = '', notes: str = '') → ScalarGateResult

Evaluate one scalar benchmark gate.

Use this helper for publication-facing metrics such as growth rates, frequencies, windowed heat fluxes, zonal residuals, and damping rates. The explicit atol/rtol pair forces each artifact to document whether its tolerance is absolute, relative, or both.

spectraxgk.diagnostics.validation_gates.gate_report(case: str, source: str, gates: list[ScalarGateResult] | tuple[ScalarGateResult, ...]) → GateReport

Summarize a set of scalar gates for one artifact.

spectraxgk.diagnostics.validation_gates.gate_report_to_dict(report: GateReport) → dict[str, object]

Return a strict JSON-serializable representation of a gate report.

spectraxgk.diagnostics.validation_gates.linear_metrics_gate_report(observed: LateTimeLinearMetrics, reference: LateTimeLinearMetrics, *, case: str, source: str, gamma_atol: float = 0.0, gamma_rtol: float = 0.05, omega_atol: float = 0.0, omega_rtol: float = 0.05) → GateReport

Gate late-time linear growth and frequency metrics.

spectraxgk.diagnostics.validation_gates.matched_nonlinear_transport_report(baseline: dict[str, Any], treatment: dict[str, Any], *, case: str = 'matched_nonlinear_transport', treatment_name: str = 'treatment', min_relative_reduction: float = 0.0, min_uncertainty_z_score: float = 0.0, value_floor: float = 1e-12) → dict[str, Any]

Compare two independently converged post-transient transport windows.

spectraxgk.diagnostics.validation_gates.nonlinear_heat_flux_convergence_gate_report(metrics: NonlinearHeatFluxConvergenceMetrics, *, case: str, source: str, max_mean_rel_delta: float = 0.05, max_cv: float = 0.15, max_abs_trend: float = 0.1, min_samples: int = 8) → GateReport

Gate post-transient heat-flux averaging stability.

This is an internal promotion gate for nonlinear transport claims: the post-transient average must agree with its terminal subwindow, have bounded coefficient of variation, show limited normalized drift across the window, and contain enough samples to be more than a reduced-window proxy.

spectraxgk.diagnostics.validation_gates.nonlinear_window_gate_report(observed: NonlinearWindowMetrics, reference: NonlinearWindowMetrics, *, case: str, source: str, rtol: float = 0.1, atol: float = 0.0, include_envelope: bool = True) → GateReport

Gate windowed nonlinear transport and field-energy metrics.

spectraxgk.diagnostics.validation_gates.observed_order_gate_report(metrics: ObservedOrderMetrics, *, case: str, source: str, min_asymptotic_order: float, min_pairwise_order: float | None = None, max_final_error: float | None = None, order_atol: float = 1e-12) → GateReport

Gate an observed-order convergence study.

min_asymptotic_order encodes the expected method/order floor for the finest refinement pair. min_pairwise_order can additionally require the whole table to be monotone enough for publication use. max_final_error can be used when both rate and absolute accuracy matter.

spectraxgk.diagnostics.validation_gates.zonal_response_gate_report(observed: ZonalFlowResponseMetrics, reference: ZonalFlowResponseMetrics, *, case: str, source: str, residual_atol: float, residual_rtol: float = 0.0, frequency_atol: float, frequency_rtol: float = 0.0, damping_atol: float, damping_rtol: float = 0.0) → GateReport

Gate Rosenbluth-Hinton/GAM-style response observables.

Autodiff Validation

Finite-difference, covariance, and eigen-branch autodiff validation helpers.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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>.

spectraxgk.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 spectraxgk.parallel.DecompositionContract(workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble', 'diagnostic_nonlinear_domain'], claim_level: Literal['production_independent_batching', 'diagnostic_nonlinear_domain_partition'], claim_label: str, n_items: int, requested_shards: int, actual_shards: int, shards: tuple[ShardAssignment, ...], independent_work: bool, changes_solver_layout: bool, state_shape: tuple[int, ...] | None = None, axis: int | None = None)

Claim-scoped shard assignment contract for a parallelization path.

property diagnostic_nonlinear_partition: bool

Whether this contract is diagnostic nonlinear-domain metadata.

property production_independent_batching: bool

Whether this contract is for production independent-work batching.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the contract.

class spectraxgk.parallel.IndependentEnsembleProvenanceReport(kind: str, workload: str, executor: str, requested_workers: int, actual_workers: int, problem_size: int, passed: bool, identity_passed: bool, ordering_passed: bool, worker_clipping_passed: bool, reconstruction_identity_passed: bool, exception_metadata_passed: bool, serial_indices: tuple[int, ...], parallel_indices: tuple[int, ...], reconstructed_indices: tuple[int, ...], identity_report: ~spectraxgk.parallel.identity.ParallelIdentityReport, reconstruction_report: dict[str, ~typing.Any], exception_metadata: dict[str, ~typing.Any], metadata: dict[str, ~typing.Any] = <factory>)

End-to-end provenance gate for independent UQ/optimization ensembles.

to_dict() → dict[str, Any]

Return a JSON-serializable provenance payload.

exception spectraxgk.parallel.IndependentMapExecutionError(index: int, executor: str, actual_workers: int, original_type: str, original_message: str)

Worker failure annotated with independent-map execution metadata.

class spectraxgk.parallel.IndependentWorkerMetadata(requested_workers: int, actual_workers: int, problem_size: int, executor: str, parallel_enabled: bool)

Resolved worker metadata for ordered independent Python tasks.

to_dict() → dict[str, Any]

Return a JSON-serializable worker metadata payload.

class spectraxgk.parallel.ParallelIdentityReport(kind: str, backend: str, requested_workers: int, actual_workers: int, problem_size: int, identity_passed: bool, max_abs_error: float, max_rel_error: float, atol: float, rtol: float, metadata: dict[str, ~typing.Any]=<factory>)

Numerical-identity report for an independent parallel execution path.

to_dict() → dict[str, Any]

Return a JSON-serializable report for artifacts and CI gates.

class spectraxgk.parallel.ReconstructionIdentityReport(workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble', 'diagnostic_nonlinear_domain'], claim_level: Literal['production_independent_batching', 'diagnostic_nonlinear_domain_partition'], claim_label: str, n_items: int, requested_shards: int, actual_shards: int, identity_passed: bool, expected_indices: tuple[int, ...], reconstructed_indices: tuple[int, ...], missing_indices: tuple[int, ...], duplicate_indices: tuple[int, ...], out_of_range_indices: tuple[int, ...], out_of_order: bool)

Serial reconstruction identity report for a decomposition contract.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the report.

class spectraxgk.parallel.ShardAssignment(shard_id: int, start: int, stop: int, indices: tuple[int, ...], label: str)

A deterministic contiguous assignment of serial indices to one shard.

property size: int

Number of serial items assigned to this shard.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the assignment.

class spectraxgk.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.

spectraxgk.parallel.batch_map(fn: Callable[[Array], Any], values: Array | ndarray, *, batch_size: int | None = None, devices: Iterable[Device] | None = None) → Any

Map fn over independent inputs with optional multi-device batching.

This helper is intended for embarrassingly parallel physics workloads such as linear k_y scans, parameter sweeps, and UQ ensembles. It preserves numerical identity with jax.vmap(fn)(values) while allowing the leading batch axis to be distributed over available devices when more than one device is supplied.

spectraxgk.parallel.batch_map_identity_report(fn: Callable[[Array], Any], values: Array | ndarray, *, batch_size: int | None = None, devices: Iterable[Device] | None = None, atol: float = 1e-12, rtol: float = 1e-10) → ParallelIdentityReport

Compare batch_map against vmap and return a CI-ready gate report.

spectraxgk.parallel.build_diagnostic_nonlinear_domain_decomposition(state_shape: Iterable[int], *, axis: int, requested_shards: int) → DecompositionContract

Build a diagnostic nonlinear-domain partition contract.

This metadata describes split/reassemble coverage along one state axis. It is intentionally not a production nonlinear route and does not claim nonlinear speedup.

spectraxgk.parallel.build_independent_portfolio_decomposition(n_items: int, *, requested_shards: int, workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble']) → DecompositionContract

Build a production independent-work decomposition contract.

The assignment is deterministic, balanced, contiguous, and contains no empty shards. It covers release-ready independent portfolios only: independent_ky_scan, uq_ensemble, and optimization_ensemble.

spectraxgk.parallel.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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) → Any

Return electrostatic phi from a full single- or multi-species state.

spectraxgk.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') → Any

Solve electrostatic phi using a species- or Hermite-sharded reduction.

spectraxgk.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.

spectraxgk.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.

spectraxgk.parallel.hermite_shift_reference(state: Any, *, offset: int) → Any

Shift a state along the Hermite axis with zero physical boundaries.

spectraxgk.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.

spectraxgk.parallel.hermite_streaming_ladder_reference(state: Any, *, vth: Any = 1.0) → Any

Return the full-array Hermite streaming ladder contribution.

spectraxgk.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.

spectraxgk.parallel.independent_ensemble_provenance_gate(fn: Callable[[Any], Any], values: Iterable[Any], *, workers: int = 1, executor: str = 'thread', workload: str = 'uq_ensemble', atol: float = 1e-12, rtol: float = 1e-10, metadata: dict[str, Any] | None = None) → IndependentEnsembleProvenanceReport

Verify independent UQ/optimization ensemble batching provenance.

The gate intentionally runs fn serially and through independent_map. It verifies result identity, serial result ordering, worker clipping, deterministic shard reconstruction, and failure metadata for the same independent-map executor family.

spectraxgk.parallel.independent_map(fn: Callable[[Any], Any], values: Iterable[Any], *, workers: int = 1, executor: str = 'thread') → list[Any]

Map independent Python tasks while preserving serial result ordering.

batch_map handles JAX-array workloads. This helper covers file-backed calibration, finite-difference, and UQ tasks whose individual units are independent Python calls. The acceptance contract is numerical identity with [fn(value) for value in values]; timing is secondary.

spectraxgk.parallel.independent_map_identity_report(fn: Callable[[Any], Any], values: Iterable[Any], *, workers: int = 1, executor: str = 'thread', atol: float = 1e-12, rtol: float = 1e-10, metadata: dict[str, Any] | None = None) → ParallelIdentityReport

Compare independent_map against a serial list-comprehension run.

spectraxgk.parallel.independent_worker_metadata(problem_size: int, *, workers: int = 1, executor: str = 'thread') → IndependentWorkerMetadata

Resolve independent-task worker counts and normalized executor metadata.

spectraxgk.parallel.integrate_linear_sharded(G0: Array, cache: LinearCache, params: LinearParams, *, dt: float, steps: int, terms: LinearTerms | None = None, state_sharding: Any | None = None) → Array

Integrate the linear system with a pjit-sharded RK2 loop.

This is intentionally minimal: it returns the final state only and avoids saving time histories to focus on strong scaling of the RHS.

spectraxgk.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. This is a diagnostic whole-state sharding primitive for identity gates and profiler localization. It is not a production nonlinear domain decomposition or speedup claim until the exact workload has communication-complete identity, conservation, transport, and profiler gates. Domain-sharding identity reports are metadata gates only; they do not authorize routing through this whole-state integrator.

spectraxgk.parallel.ky_scan_batches(ky_values: ndarray, *, n_batches: int) → list[ndarray]

Return balanced k_y chunks for independent linear-scan execution.

spectraxgk.parallel.mirror_drift_reference(H: Any, *, vth: Any, bgrad: Any, ell: Any, sqrt_m: Any, sqrt_m_p1: Any, weight: Any = 1.0) → Any

Return the mirror-drift contribution with full-array Hermite shifts.

spectraxgk.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.

spectraxgk.parallel.pad_to_multiple(values: Array, multiple: int) → tuple[Array, int]

Pad axis zero by edge repetition so its length is divisible by multiple.

spectraxgk.parallel.parallel_identity_report(reference: Any, observed: Any, *, kind: str, problem_size: int, requested_workers: int, actual_workers: int | None = None, backend: str | None = None, atol: float = 1e-12, rtol: float = 1e-10, metadata: dict[str, Any] | None = None) → ParallelIdentityReport

Build a numerical-identity report for serial-vs-parallel outputs.

spectraxgk.parallel.periodic_streaming_reference(state: Any, *, kz: Any, vth: Any = 1.0) → Any

Return periodic parallel streaming using full-array operations.

spectraxgk.parallel.periodic_streaming_shard_map(state: Any, plan: VelocityShardingPlan, *, kz: Any, vth: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return periodic parallel streaming through the Hermite shard-map path.

spectraxgk.parallel.reconstruct_serial(contract: DecompositionContract, shard_values: Sequence[Sequence[T]]) → tuple[T, ...]

Reassemble shard values into serial index order.

spectraxgk.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).

spectraxgk.parallel.serial_reconstruction_identity_report(values: Sequence[T], contract: DecompositionContract, *, equal: Callable[[T, T], bool] | None = None) → ReconstructionIdentityReport

Check that contract sharding reassembles exactly to serial order.

spectraxgk.parallel.shard_sequence(values: Sequence[T], contract: DecompositionContract) → tuple[tuple[T, ...], ...]

Return values grouped according to a decomposition contract.

spectraxgk.parallel.split_evenly(values: ndarray, n_parts: int) → list[ndarray]

Split an array into nonempty, nearly equal chunks along axis zero.

spectraxgk.parallel.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

Return the full-array velocity-axis reduction used by field solves.

spectraxgk.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 Identity Reports

Numerical-identity reports for parallel execution paths.

class spectraxgk.parallel.identity.ParallelIdentityReport(kind: str, backend: str, requested_workers: int, actual_workers: int, problem_size: int, identity_passed: bool, max_abs_error: float, max_rel_error: float, atol: float, rtol: float, metadata: dict[str, ~typing.Any]=<factory>)

Numerical-identity report for an independent parallel execution path.

to_dict() → dict[str, Any]

Return a JSON-serializable report for artifacts and CI gates.

spectraxgk.parallel.identity.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.

Parallel Batch Mapping

JAX batch-map helpers for independent scan workloads.

spectraxgk.parallel.batch.batch_map(fn: Callable[[Array], Any], values: Array | ndarray, *, batch_size: int | None = None, devices: Iterable[Device] | None = None) → Any

Map fn over independent inputs with optional multi-device batching.

This helper is intended for embarrassingly parallel physics workloads such as linear k_y scans, parameter sweeps, and UQ ensembles. It preserves numerical identity with jax.vmap(fn)(values) while allowing the leading batch axis to be distributed over available devices when more than one device is supplied.

spectraxgk.parallel.batch.batch_map_identity_report(fn: Callable[[Array], Any], values: Array | ndarray, *, batch_size: int | None = None, devices: Iterable[Device] | None = None, atol: float = 1e-12, rtol: float = 1e-10) → ParallelIdentityReport

Compare batch_map against vmap and return a CI-ready gate report.

spectraxgk.parallel.batch.ky_scan_batches(ky_values: ndarray, *, n_batches: int) → list[ndarray]

Return balanced k_y chunks for independent linear-scan execution.

spectraxgk.parallel.batch.pad_to_multiple(values: Array, multiple: int) → tuple[Array, int]

Pad axis zero by edge repetition so its length is divisible by multiple.

spectraxgk.parallel.batch.split_evenly(values: ndarray, n_parts: int) → list[ndarray]

Split an array into nonempty, nearly equal chunks along axis zero.

Parallel Independent Tasks

Ordered Python-task parallelism for UQ and optimization ensembles.

class spectraxgk.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: ~spectraxgk.parallel.identity.ParallelIdentityReport, reconstruction_report: dict[str, ~typing.Any], exception_metadata: dict[str, ~typing.Any], metadata: dict[str, ~typing.Any] = <factory>)

End-to-end provenance gate for independent UQ/optimization ensembles.

to_dict() → dict[str, Any]

Return a JSON-serializable provenance payload.

exception spectraxgk.parallel.independent.IndependentMapExecutionError(index: int, executor: str, actual_workers: int, original_type: str, original_message: str)

Worker failure annotated with independent-map execution metadata.

class spectraxgk.parallel.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.parallel.independent.independent_worker_metadata(problem_size: int, *, workers: int = 1, executor: str = 'thread') → IndependentWorkerMetadata

Resolve independent-task worker counts and normalized executor metadata.

Nonlinear Parallel Spectral Core

Stable facade for nonlinear spectral parallelization helpers.

spectraxgk.operators.nonlinear.spectral_core._host_max_abs_rel_error(reference: Array, candidate: Array, *, atol: float) → tuple[float, float]

Return max errors after materializing arrays on the host.

spectraxgk.operators.nonlinear.spectral_core._host_staged_array_for_sharding(array: Array) → ndarray

Return a host-backed array before applying explicit device sharding.

On the CUDA stack used for the current device-z diagnostic, direct device_put from a single-device JAX array into a z-sharded NamedSharding can misplace the second z shard. Host staging keeps the identity gate about the candidate nonlinear route instead of about that source-device resharding behavior.

spectraxgk.operators.nonlinear.spectral_core._normalize_spectral_tile_bounds(tile_bounds: tuple[tuple[int, int, int, int], ...]) → tuple[tuple[int, int, int, int], ...]

Return validated fixed-width tile bounds for mypy and runtime checks.

spectraxgk.operators.nonlinear.spectral_core._pencil_fft2(arr: Array, *, y_axis: int, x_axis: int) → Array

Return a 2D forward FFT through explicit x-then-y pencil stages.

spectraxgk.operators.nonlinear.spectral_core._pencil_ifft2(arr: Array, *, y_axis: int, x_axis: int) → Array

Return a 2D inverse FFT through explicit x-then-y pencil stages.

spectraxgk.operators.nonlinear.spectral_core._pencil_spectral_bracket(state_hat: Array, phi_hat: Array) → Array

Return the pseudo-spectral bracket using pencil FFT staging.

This function is the local algorithmic route that a distributed pencil FFT implementation should follow: stack derivative operands, transform through explicit axis-transpose stages, multiply in physical space, and transform the bracket back without first reconstructing logical output tiles.

spectraxgk.operators.nonlinear.spectral_core._pencil_spectral_bracket_z_chunked(state_hat: Array, phi_hat: Array, *, z_chunk_size: int) → Array

Return the pencil bracket by processing independent z slabs.

The nonlinear bracket has no coupling along z inside this local pseudo-spectral micro-route. Chunking the local z extent therefore preserves the operator while reducing cuFFT batched-plan pressure on GPUs.

spectraxgk.operators.nonlinear.spectral_core._spectral_layout_round_trip(arr: Array, *, y_axis: int, x_axis: int, y_chunks: tuple[int, ...], x_chunks: tuple[int, ...]) → Array

Simulate the split/transposed/reassembled layout changes used by FFTs.

spectraxgk.operators.nonlinear.spectral_core._spectral_rhs_from_bracket(bracket_hat: Array) → Array

Return the ExB advection contribution used by the identity micro-route.

spectraxgk.operators.nonlinear.spectral_core._within_abs_or_rel_tolerance(max_abs_error: float, max_rel_error: float, *, atol: float, rtol: float) → bool

Return an allclose-style scalar gate for recorded max errors.

spectraxgk.operators.nonlinear.spectral_core.deterministic_nonlinear_spectral_state(shape: tuple[int, int, int, int, int] = (2, 3, 6, 4, 2)) → Array

Return deterministic complex spectral coefficients for communication gates.

The layout is (Nl, Nm, Ny, Nx, Nz) with the FFT axes in (Ny, Nx).

spectraxgk.operators.nonlinear.spectral_core.device_z_pencil_fft_batch_pressure_model(state_shape: tuple[int, int, int, int, int], *, device_count: int, max_fft_batch_count: int = 65536, z_chunk_size: int | None = None) → NonlinearSpectralDevicePencilFFTBatchModel

Estimate cuFFT batch pressure for the device-z pencil micro-route.

The profiler traces showed that large GPU cases can fail before timing when axis-wise FFTs create too-large batched cuFFT plans. This backend-free model predicts that pressure and suggests a local z_chunk_size that keeps the largest state-gradient axis FFT batch below max_fft_batch_count. It is a profiling preflight, not a speedup or physics claim.

spectraxgk.operators.nonlinear.spectral_core.nonlinear_spectral_domain_work_model(state_shape: tuple[int, int, int, int, int], *, y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], max_communication_to_owned_work_ratio: float = 0.5) → NonlinearSpectralDomainWorkModel

Estimate communication pressure for the current logical spectral route.

The current diagnostic route reconstructs global spectral state/bracket arrays before returning owned output tiles. That is useful for identity gating, but it implies allgather/broadcast traffic that can dominate the owned tile work. This model is a conservative fail-closed screen for that route; it is not a performance prediction for a future distributed FFT.

spectraxgk.operators.nonlinear.spectral_core.nonlinear_spectral_pencil_work_model(state_shape: tuple[int, int, int, int, int], *, y_chunks: tuple[int, ...], x_chunks: tuple[int, ...], max_communication_to_fft_work_ratio: float = 0.35, min_predicted_speedup: float = 1.5) → NonlinearSpectralPencilWorkModel

Estimate communication pressure for a pencil-FFT bracket route.

The pencil route avoids global state/bracket reconstruction. Its remaining distributed cost is the all-to-all transpose traffic needed by axis-wise 2D FFTs plus a field-reduction/broadcast. This model is intentionally simple and fail-closed: it must pass before any strong-scaling timing is treated as a meaningful candidate for production promotion.

Nonlinear Domain Decomposition

Local nonlinear domain local-stencil gates for parallelization validation.

These routines exercise bounded halo decomposition on a deterministic local stencil. They are diagnostic identity gates only; they do not implement a production nonlinear domain-decomposed solver route or a speedup claim.

spectraxgk.operators.nonlinear.domain_decomposition.build_nonlinear_domain_decomposition_plan(state_shape: tuple[int, ...], *, axis: int = 0, num_domains: int = 2, halo: int = 1) → NonlinearDomainDecompositionPlan

Build a static chunk plan for the local state-domain identity prototype.

spectraxgk.operators.nonlinear.domain_decomposition.deterministic_nonlinear_domain_state(shape: tuple[int, ...] = (6, 4)) → Array

Return a small deterministic complex state for identity gates.

spectraxgk.operators.nonlinear.domain_decomposition.local_stencil_nonlinear_domain_decomposed_step(state: Array, plan: NonlinearDomainDecompositionPlan, *, dt: float = 0.05) → Array

Apply the same local nonlinear step through static halo chunks.

spectraxgk.operators.nonlinear.domain_decomposition.local_stencil_nonlinear_domain_serial_step(state: Array, *, axis: int = 0, dt: float = 0.05) → Array

Apply the serial local nonlinear diagnostic step along one state axis.

spectraxgk.operators.nonlinear.domain_decomposition.nonlinear_domain_identity_report(serial_state: Array, decomposed_state: Array, plan: NonlinearDomainDecompositionPlan, *, atol: float = 1e-06, rtol: float = 1e-06) → NonlinearDomainIdentityReport

Compare decomposed and serial states and fail closed on any mismatch.

spectraxgk.operators.nonlinear.domain_decomposition.nonlinear_domain_parallel_identity_gate(state: Array, plan: NonlinearDomainDecompositionPlan, *, dt: float = 0.05, atol: float = 1e-06, rtol: float = 1e-06) → tuple[Array, NonlinearDomainIdentityReport]

Return a fail-closed decomposed local-stencil step and its identity report.

spectraxgk.operators.nonlinear.domain_decomposition.nonlinear_domain_transport_window_identity_gate(state: Array, plan: NonlinearDomainDecompositionPlan, *, dt: float = 0.025, steps: int = 4, atol: float = 1e-06, rtol: float = 1e-06) → NonlinearDomainTransportWindowReport

Validate a multi-step serial-vs-decomposed transport window.

The gate is deliberately stricter than a final-state check: it compares state identity, decomposed-boundary identity, and per-step scalar traces for mass, free-energy proxy, and boundary-flux proxy. The scalar drifts are compared between serial and decomposed paths; they are not claimed to be conserved by this damped diagnostic stencil.

Nonlinear Parallel Device-Z Core

Device-z nonlinear spectral parallel routes and identity gates.

spectraxgk.operators.nonlinear.device_z._append_device_z_transport_observables(traces: dict[str, list[float]], state_hat: Array, *, observable_mode: Literal['host_gather', 'sharded_reduce'], sharded_observables_fn: Any | None) → None

Append device-z transport observables using the selected reduction route.

spectraxgk.operators.nonlinear.device_z._append_spectral_physical_observable_vector(traces: dict[str, list[float]], values: Array) → None

Append [Wg, Wphi, Q, bracket_rms] scalar observables.

spectraxgk.operators.nonlinear.device_z._blocked_device_z_rhs_report(*, state_shape: tuple[int, int, int, int, int], axis_name: str, requested_count: int, active_count: int, atol: float, rtol: float, blocked_reasons: Sequence[str]) → NonlinearSpectralDevicePencilRHSIdentityReport

Return a fail-closed device-z RHS identity report.

spectraxgk.operators.nonlinear.device_z._blocked_device_z_transport_window_report(*, state_shape: tuple[int, int, int, int, int], axis_name: str, requested_count: int, active_count: int, steps: int, dt: float, atol: float, rtol: float, blocked_reasons: Sequence[str], serial_traces: dict[str, list[float]]) → NonlinearSpectralDevicePencilTransportWindowReport

Return a fail-closed device-z transport-window report.

spectraxgk.operators.nonlinear.device_z._device_z_pencil_shard_map_observables_fn(mesh: Any, *, axis_name: str, z_chunk_size: int | None = None) → Any

Return a shard-map scalar observable reducer for z-sharded states.

spectraxgk.operators.nonlinear.device_z._device_z_pencil_shard_map_rhs_fn(mesh: Any, *, axis_name: str, z_chunk_size: int | None = None) → Any

Return a jitted shard-map RHS for z-sharded spectral states.

spectraxgk.operators.nonlinear.device_z._device_z_rhs_identity_report(*, state_shape: tuple[int, int, int, int, int], axis_name: str, requested_count: int, active_count: int, atol: float, rtol: float, rhs_abs: float, rhs_rel: float) → NonlinearSpectralDevicePencilRHSIdentityReport

Return the passed/blocked device-z RHS identity report.

spectraxgk.operators.nonlinear.device_z._device_z_sharding_for_spectral_state(state_hat: Array, *, devices: Sequence[Any] | None, axis_name: str) → tuple[Any | None, Any | None, tuple[str, ...], int, int]

Return a z-axis sharding for the fused-bracket route, or blockers.

The nonlinear pseudo-spectral bracket transforms only the (ky, kx) axes. Sharding over z therefore keeps the FFTs local to each device and avoids the global tile reconstruction that blocked the older logical route.

spectraxgk.operators.nonlinear.device_z._device_z_transport_identity_passed(*, state_abs: float, state_rel: float, trace_errors: Mapping[str, tuple[float, float]], atol: float, rtol: float) → bool

Return whether final state and all transport traces pass identity gates.

spectraxgk.operators.nonlinear.device_z._device_z_transport_window_report(*, state_shape: tuple[int, int, int, int, int], axis_name: str, requested_count: int, active_count: int, steps: int, dt: float, atol: float, rtol: float, state_abs: float, state_rel: float, serial_trace_values: dict[str, tuple[float, ...]], device_trace_values: dict[str, tuple[float, ...]], blocked_reasons: Sequence[str]) → NonlinearSpectralDevicePencilTransportWindowReport

Return the passed/blocked device-z transport-window report.

spectraxgk.operators.nonlinear.device_z._new_transport_trace_dict() → dict[str, list[float]]

Return empty scalar traces used by nonlinear transport-window gates.

spectraxgk.operators.nonlinear.device_z._run_device_z_compute_window_states(state_hat: Array, setup: _DeviceZShardingSetup, *, axis_name: str, z_chunk_size: int | None, dt: float, steps: int) → _DeviceZComputeStates

Return final states from the compute-only serial and z-sharded routes.

The transport-window gate also collects scalar traces with host-visible instrumentation. Final-state identity must match the compute-only route that the profiler times, otherwise per-step diagnostics can create a different numerical path than the speedup candidate.

spectraxgk.operators.nonlinear.device_z._serial_transport_compute_state(state_hat: Array, *, dt: float, steps: int) → Array

Return the compute-only serial fixed-window state used for timing.

spectraxgk.operators.nonlinear.device_z._spectral_physical_transport_observable_sums(state_hat: Array, bracket_hat: Array) → Array

Return additive physical-space observable sums for local z slabs.

spectraxgk.operators.nonlinear.device_z._spectral_physical_transport_observable_vector_from_sums(sums: Array) → Array

Convert additive observable sums into [Wg, Wphi, Q, bracket_rms].

spectraxgk.operators.nonlinear.device_z._transport_trace_error_pairs(serial: dict[str, tuple[float, ...]], device: dict[str, tuple[float, ...]], *, floor: float) → dict[str, tuple[float, float]]

Return max absolute/relative errors for each transport trace.

spectraxgk.operators.nonlinear.device_z._transport_trace_tuples(traces: dict[str, list[float]]) → dict[str, tuple[float, ...]]

Freeze mutable transport traces for report assembly.

spectraxgk.operators.nonlinear.device_z.device_z_pencil_nonlinear_spectral_rhs(state_hat: Array, *, devices: Sequence[Any] | None = None, axis_name: str = 'z', z_chunk_size: int | None = None, atol: float = 5e-06, rtol: float = 0.0001) → tuple[Array, NonlinearSpectralDevicePencilRHSIdentityReport]

Return the z-sharded fused pencil nonlinear RHS after identity gating.

This is the first real device-sharded nonlinear spectral route in this module. It shards over the field-line z axis so that the FFT axes remain local on every device. The function falls back to the serial RHS unless a multi-device sharding exists and the sharded fused-bracket RHS matches the serial reference within the requested tolerances.

spectraxgk.operators.nonlinear.device_z.device_z_pencil_nonlinear_spectral_transport_window_identity_gate(state_hat: Array, *, devices: Sequence[Any] | None = None, axis_name: str = 'z', z_chunk_size: int | None = None, dt: float = 0.005, steps: int = 4, atol: float = 5e-06, rtol: float = 0.0001, observable_mode: Literal['host_gather', 'sharded_reduce'] = 'host_gather') → NonlinearSpectralDevicePencilTransportWindowReport

Validate a multi-step serial-vs-device-z-sharded nonlinear window.

The route advances the same explicit fixed-step micro-window with the serial nonlinear RHS and the shard-map z-pencil RHS. It compares the final state and physical-space scalar traces, including a density-times-radial electric-field flux proxy and bracket RMS. Passing this gate is still not a turbulent heat-flux validation; it only permits timing the decomposed device route on the same deterministic operator.

Velocity Sharding Plans

Velocity-space decomposition plans for production parallelization.

class spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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) → Any

Return electrostatic phi from a full single- or multi-species state.

spectraxgk.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') → Any

Solve electrostatic phi using a species- or Hermite-sharded reduction.

spectraxgk.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.

spectraxgk.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.

spectraxgk.parallel.velocity.hermite_shift_reference(state: Any, *, offset: int) → Any

Shift a state along the Hermite axis with zero physical boundaries.

spectraxgk.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.

spectraxgk.parallel.velocity.hermite_streaming_ladder_reference(state: Any, *, vth: Any = 1.0) → Any

Return the full-array Hermite streaming ladder contribution.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.parallel.velocity.periodic_streaming_reference(state: Any, *, kz: Any, vth: Any = 1.0) → Any

Return periodic parallel streaming using full-array operations.

spectraxgk.parallel.velocity.periodic_streaming_shard_map(state: Any, plan: VelocityShardingPlan, *, kz: Any, vth: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return periodic parallel streaming through the Hermite shard-map path.

spectraxgk.parallel.velocity.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

Return the full-array velocity-axis reduction used by field solves.

spectraxgk.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 spectraxgk.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.

spectraxgk.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.

Hermite-axis exchange and velocity-field reduction kernels.

spectraxgk.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.

spectraxgk.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.

spectraxgk.parallel.velocity_hermite.hermite_shift_reference(state: Any, *, offset: int) → Any

Shift a state along the Hermite axis with zero physical boundaries.

spectraxgk.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.

spectraxgk.parallel.velocity_hermite.velocity_field_reduce_reference(state: Any, *, axis: str = 'm') → Any

Return the full-array velocity-axis reduction used by field solves.

spectraxgk.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.

Velocity-parallel streaming and magnetic-drift microkernels.

spectraxgk.parallel.velocity_streaming.curvature_gradb_drift_reference(H: Any, *, tz: Any, omega_d_scale: Any, cv_d: Any, gb_d: Any, ell: Any, m: Any, weight_curv: Any = 1.0, weight_gradb: Any = 1.0) → Any

Return curvature and grad-B drift contributions with full-array shifts.

spectraxgk.parallel.velocity_streaming.curvature_gradb_drift_shard_map(H: Any, plan: VelocityShardingPlan, *, tz: Any, omega_d_scale: Any, cv_d: Any, gb_d: Any, ell: Any, m: Any, weight_curv: Any = 1.0, weight_gradb: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return curvature and grad-B drift contributions using Hermite exchange.

spectraxgk.parallel.velocity_streaming.hermite_streaming_ladder_reference(state: Any, *, vth: Any = 1.0) → Any

Return the full-array Hermite streaming ladder contribution.

spectraxgk.parallel.velocity_streaming.hermite_streaming_ladder_shard_map(state: Any, plan: VelocityShardingPlan, *, vth: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return a shard-map Hermite streaming ladder contribution.

spectraxgk.parallel.velocity_streaming.mirror_drift_reference(H: Any, *, vth: Any, bgrad: Any, ell: Any, sqrt_m: Any, sqrt_m_p1: Any, weight: Any = 1.0) → Any

Return the mirror-drift contribution with full-array Hermite shifts.

spectraxgk.parallel.velocity_streaming.mirror_drift_shard_map(H: Any, plan: VelocityShardingPlan, *, vth: Any, bgrad: Any, ell: Any, sqrt_m: Any, sqrt_m_p1: Any, weight: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return the mirror-drift contribution using Hermite shard exchange.

spectraxgk.parallel.velocity_streaming.periodic_streaming_reference(state: Any, *, kz: Any, vth: Any = 1.0) → Any

Return periodic parallel streaming using full-array operations.

spectraxgk.parallel.velocity_streaming.periodic_streaming_shard_map(state: Any, plan: VelocityShardingPlan, *, kz: Any, vth: Any = 1.0, devices: Sequence[Any] | None = None, axis_name: str = 'm') → Any

Return periodic parallel streaming through the Hermite shard-map path.

Electrostatic field reduction and diamagnetic-drive microkernels.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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) → Any

Return electrostatic phi from a full single- or multi-species state.

spectraxgk.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') → Any

Solve electrostatic phi using a species- or Hermite-sharded reduction.

State Sharding Policy

Helpers for multi-device sharding of GK state arrays.

spectraxgk.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).

Sharded Integrators

Sharded fixed-step integrators for multi-device scaling experiments.

spectraxgk.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.

spectraxgk.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. This is a diagnostic whole-state sharding primitive for identity gates and profiler localization. It is not a production nonlinear domain decomposition or speedup claim until the exact workload has communication-complete identity, conservation, transport, and profiler gates. Domain-sharding identity reports are metadata gates only; they do not authorize routing through this whole-state integrator.

Zonal Validation

Zonal Flow Objectives

Reduced zonal-flow objectives for differentiable stellarator optimization.

This module is intentionally backend-free. Production callers should build zonal-response metrics from VMEC/Boozer/SPECTRAX-GK rows, then pass the metric tensors here for reduction, finite-difference checks, and UQ diagnostics.

class spectraxgk.objectives.zonal.ZonalFlowObjectiveConfig(residual_weight: float = 1.0, damping_weight: float = 1.0, growth_over_residual_weight: float = 0.0, recurrence_weight: float = 0.0, residual_floor: float = 1e-06)

Weights and floors for a minimizable zonal-flow objective.

The objective rewards large residual zonal response by minimizing 1 / residual and penalizes collisionless damping, linear growth not screened by the residual, and late-time recurrence/envelope amplitude. Nonlinear heat-flux suppression remains a separate holdout gate.

objective_weights() → Array

Return the normalized objective-column weights used by the reducer.

to_dict() → dict[str, object]

Return a JSON-friendly representation.

spectraxgk.objectives.zonal.zonal_flow_objective_artifact_from_records(records: Iterable[Mapping[str, Any]], *, config: ZonalFlowObjectiveConfig | None = None, surface_keys: Sequence[str] = ('surface', 'surface_index', 'torflux'), alpha_keys: Sequence[str] = ('alpha', 'field_line_label'), kx_keys: Sequence[str] = ('kx', 'kx_target', 'kx_rhoi'), residual_keys: Sequence[str] = ('residual_level', 'spectrax_residual'), damping_keys: Sequence[str] = ('damping_rate', 'gam_damping_rate'), linear_growth_keys: Sequence[str] = ('linear_growth_rate', 'growth_rate', 'gamma'), recurrence_keys: Sequence[str] = ('recurrence_amplitude', 'tail_std_ratio', 'residual_std', 'tail_std'), missing_damping_policy: Literal['fail', 'zero'] = 'fail', claim_level: str | None = None, source_paths: Sequence[str] | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean') → dict[str, object]

Build a strict JSON-friendly zonal-flow objective artifact.

The input is a table of validated zonal-response metrics. Rows are mapped onto the shared (surface, alpha, kx) portfolio tensor used by the stellarator objective stack. Missing damping rates fail by default because a promoted zonal-flow optimization claim must know the damping convention. Diagnostic artifacts can set missing_damping_policy='zero' to produce rows while carrying an explicit promotion_ready=False flag.

spectraxgk.objectives.zonal.zonal_flow_objective_rows(*, residual_level: Any, damping_rate: Any, linear_growth_rate: Any | None = None, recurrence_amplitude: Any | None = None, config: ZonalFlowObjectiveConfig | None = None) → Array

Return objective rows with shape (surface, alpha, kx, objective).

residual_level is the late-time residual normalized to the initial zonal potential. Larger residuals reduce the first objective column. damping_rate should be positive for decaying GAM/zonal envelopes. linear_growth_rate is optional and encodes a suppression-relevance metric: high ITG growth with weak residuals is penalized. The recurrence column should be a non-negative late-envelope or moment-tail amplitude.

spectraxgk.objectives.zonal.zonal_flow_objective_sensitivity_report(metric_fn: Callable[[Array], Mapping[str, Any]], params: Any, *, config: ZonalFlowObjectiveConfig | None = None, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, min_rank: int | None = None, condition_number_limit: float = 100000000.0, covariance_regularization: float = 1e-09, workers: int = 1, parallel_executor: str = 'thread') → dict[str, object]

AD/FD, row-Jacobian, and UQ gate for a zonal-flow optimization map.

spectraxgk.objectives.zonal.zonal_flow_reduced_objective(*, residual_level: Any, damping_rate: Any, linear_growth_rate: Any | None = None, recurrence_amplitude: Any | None = None, config: ZonalFlowObjectiveConfig | None = None, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean') → Array

Reduce zonal-flow metric tensors to one differentiable scalar objective.

Analysis

Public facade for mode extraction and growth-rate diagnostics.

class spectraxgk.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 spectraxgk.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 spectraxgk.diagnostics.analysis.ModeSelection(ky_index: 'int', kx_index: 'int', z_index: 'int' = 0)

class spectraxgk.diagnostics.analysis.ModeSelectionBatch(ky_indices: 'np.ndarray', kx_index: 'int', z_index: 'int' = 0)

class spectraxgk.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)

Post-transient heat-flux averaging convergence summary.

class spectraxgk.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)

Windowed transport/envelope metrics for a nonlinear run.

class spectraxgk.diagnostics.analysis.ObservedOrderMetrics(step_sizes: ndarray, errors: ndarray, orders: ndarray, asymptotic_order: float)

Observed-order convergence summary from step sizes and errors.

spectraxgk.diagnostics.analysis.branch_continuity_metrics(ky: ndarray, gamma: ndarray, omega: ndarray, *, successive_overlap: ndarray | None = None, floor_fraction: float = 1e-08) → BranchContinuationMetrics

Compute branch-continuity diagnostics for a linear scan.

The relative jump normalization uses a local scale from adjacent values, with a floor tied to the largest value in the scan. This avoids false blow-ups near marginal points while still flagging branch jumps.

spectraxgk.diagnostics.analysis.density_moment(G: ndarray, Jl: ndarray, *, species_index: int | None = None) → ndarray

Compute the m=0 density moment for a selected species (or summed if None).

spectraxgk.diagnostics.analysis.estimate_observed_order(step_sizes: ndarray, errors: ndarray) → ObservedOrderMetrics

Estimate observed order from successive step-size refinements.

spectraxgk.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).

spectraxgk.diagnostics.analysis.extract_mode(phi_t: ndarray, sel: ModeSelection) → ndarray

Extract a complex mode time series from phi_t(t, ky, kx, z).

spectraxgk.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).

spectraxgk.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).

spectraxgk.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 = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float]

Fit gamma/omega with optional auto-selected window.

spectraxgk.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 = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float, float, float]

Fit gamma/omega and report selected window plus R^2 scores.

This wrapper intentionally calls the facade-level fit_growth_rate_with_stats() so tests and downstream users can monkeypatch the public analysis module without reaching into implementation modules.

spectraxgk.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).

spectraxgk.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).

spectraxgk.diagnostics.analysis.late_time_linear_metrics(result: object, *, tail_fraction: float = 0.5, mode_method: str = 'project') → LateTimeLinearMetrics

Return late-time growth/frequency metrics from a linear benchmark/runtime result.

spectraxgk.diagnostics.analysis.nonlinear_heat_flux_convergence_metrics(t: ndarray, heat_flux: ndarray, *, start_fraction: float = 0.5, terminal_fraction: float = 0.5, mean_floor: float = 1e-30) → NonlinearHeatFluxConvergenceMetrics

Summarize whether a post-transient heat-flux average is stable.

start_fraction discards startup samples. terminal_fraction compares the retained post-transient mean with the final subwindow of that retained region. The normalized trend is the least-squares slope multiplied by the post-transient time span and divided by the absolute post-transient mean.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.diagnostics.analysis.windowed_growth_rate_from_omega_series(gamma_t: ndarray, omega_t: ndarray, sel: ModeSelection, *, navg_fraction: float = 0.5, use_last: bool = False) → Tuple[float, float, ndarray, ndarray]

Average a resolved (time, ky, kx) growth/frequency series.

spectraxgk.diagnostics.analysis.windowed_nonlinear_metrics(result: object, *, start_fraction: float = 0.5) → NonlinearWindowMetrics

Return late-window transport and envelope metrics from a nonlinear runtime result.

Mode Diagnostics

Mode-selection and eigenfunction extraction diagnostics.

class spectraxgk.diagnostics.modes.EigenfunctionComparisonMetrics(overlap: float, relative_l2: float, phase_shift: float)

Phase-aligned eigenfunction comparison summary.

class spectraxgk.diagnostics.modes.EigenfunctionReferenceBundle(theta: ndarray, mode: ndarray, source: str, case: str, metadata: dict[str, object])

Frozen reference eigenfunction bundle for manuscript-grade overlays.

class spectraxgk.diagnostics.modes.ModeSelection(ky_index: 'int', kx_index: 'int', z_index: 'int' = 0)

class spectraxgk.diagnostics.modes.ModeSelectionBatch(ky_indices: 'np.ndarray', kx_index: 'int', z_index: 'int' = 0)

spectraxgk.diagnostics.modes.compare_eigenfunctions(eigenfunction: ndarray, reference: ndarray) → EigenfunctionComparisonMetrics

Return normalized overlap and relative L2 error after global phase alignment.

spectraxgk.diagnostics.modes.density_moment(G: ndarray, Jl: ndarray, *, species_index: int | None = None) → ndarray

Compute the m=0 density moment for a selected species (or summed if None).

spectraxgk.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).

spectraxgk.diagnostics.modes.extract_mode(phi_t: ndarray, sel: ModeSelection) → ndarray

Extract a complex mode time series from phi_t(t, ky, kx, z).

spectraxgk.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).

spectraxgk.diagnostics.modes.load_eigenfunction_reference_bundle(path: str | Path) → EigenfunctionReferenceBundle

Load a frozen reference eigenfunction bundle.

spectraxgk.diagnostics.modes.normalize_eigenfunction(eigenfunction: ndarray, z: ndarray) → ndarray

Normalize an eigenfunction by its value at theta=0 (nearest z=0).

spectraxgk.diagnostics.modes.phase_align_eigenfunction(eigenfunction: ndarray, reference: ndarray) → tuple[ndarray, float]

Phase-align eigenfunction to reference using the global complex phase.

spectraxgk.diagnostics.modes.save_eigenfunction_reference_bundle(path: str | Path, *, theta: ndarray, mode: ndarray, source: str, case: str, metadata: dict[str, object] | None = None) → Path

Write a frozen reference eigenfunction bundle as .npz.

spectraxgk.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

Public growth-rate, frequency, and fit-window diagnostics.

spectraxgk.diagnostics.growth_rates._log_amp_phase(signal: ndarray) → Tuple[ndarray, ndarray]

Return (log|signal|, unwrapped phase) with robust scaling.

spectraxgk.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).

spectraxgk.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 = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float]

Fit gamma/omega with optional auto-selected window.

spectraxgk.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 = 'loglinear', max_fraction: float = 0.8, end_fraction: float = 0.9, num_windows: int = 8, phase_weight: float = 0.2, length_weight: float = 0.05, min_r2: float = 0.0, late_penalty: float = 0.1, min_slope: float | None = None, min_slope_frac: float = 0.0, slope_var_weight: float = 0.0) → Tuple[float, float, float, float, float, float]

Fit gamma/omega and report selected window plus R^2 scores.

spectraxgk.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).

spectraxgk.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).

spectraxgk.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.

spectraxgk.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.

spectraxgk.diagnostics.growth_rates.windowed_growth_rate_from_omega_series(gamma_t: ndarray, omega_t: ndarray, sel: ModeSelection, *, navg_fraction: float = 0.5, use_last: bool = False) → Tuple[float, float, ndarray, ndarray]

Average a resolved (time, ky, kx) growth/frequency series.

Growth-Rate Fit Kernels

Least-squares growth-rate and frequency fits for complex mode traces.

spectraxgk.diagnostics.growth_fit._log_amp_phase(signal: ndarray) → Tuple[ndarray, ndarray]

Return (log|signal|, unwrapped phase) with robust scaling.

spectraxgk.diagnostics.growth_fit.fit_growth_rate(t: ndarray, signal: ndarray, tmin: float | None = None, tmax: float | None = None) → Tuple[float, float]

Fit gamma and omega from a complex signal ~ exp((gamma - i*omega) t).

spectraxgk.diagnostics.growth_fit.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).

Growth-Rate Fit Windows

Fit-window selection utilities for linear growth-rate diagnostics.

spectraxgk.diagnostics.growth_windows.late_time_window(t: ndarray, *, tail_fraction: float = 0.4) → tuple[float, float]

Return the start/end of a late-time tail window.

This is the windowing convention used for manuscript-facing eigenfunction extraction when the growth-rate fit window is not the same object as the late-time mode-shape window.

spectraxgk.diagnostics.growth_windows.select_fit_window(t: ndarray, signal: ndarray, window_fraction: float = 0.3, min_points: int = 20, start_fraction: float = 0.0, growth_weight: float = 0.0, require_positive: bool = False, min_amp_fraction: float = 0.0) → Tuple[float, float]

Pick a time window with the most exponential-like behavior.

spectraxgk.diagnostics.growth_windows.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.

Runtime Output Plots

Runtime-output plotting for saved linear and nonlinear artifacts.

spectraxgk.artifacts.runtime_plots.linear_runtime_panel_figure(*, t: ndarray, signal: ndarray, z: ndarray, eigenfunction: ndarray, gamma: float, omega: float, title: str = 'SPECTRAX-GK Linear Runtime') → Tuple[Figure, ndarray]

Create the default two-panel linear runtime plot.

spectraxgk.artifacts.runtime_plots.nonlinear_runtime_panel_figure(*, t: ndarray, phi2: ndarray | None = None, wphi: ndarray | None = None, heat_flux: ndarray | None = None, gamma: ndarray | None = None, omega: ndarray | None = None, title: str = 'SPECTRAX-GK Nonlinear Runtime') → Tuple[Figure, ndarray]

Create the default three-panel nonlinear runtime plot.

spectraxgk.artifacts.runtime_plots.plot_saved_output(path: str | Path, *, out: str | Path | None = None) → Path

Plot a saved linear or nonlinear output bundle.

Benchmark And Scan Plots

Benchmark and scan comparison plotting helpers.

class spectraxgk.artifacts.benchmark_plots.LinearValidationPanel(name: 'str', z: 'np.ndarray', eigenfunction: 'np.ndarray', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', x_label: 'str', x_ref: 'np.ndarray | None' = None, gamma_ref: 'np.ndarray | None' = None, omega_ref: 'np.ndarray | None' = None, ref_label: 'str' = 'Reference', log_x: 'bool' = False)

class spectraxgk.artifacts.benchmark_plots.MultiReferenceValidationPanel(name: 'str', z: 'np.ndarray', eigenfunction: 'np.ndarray', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', x_label: 'str', references: 'list[ReferenceSeries]', log_x: 'bool' = False)

class spectraxgk.artifacts.benchmark_plots.ReferenceSeries(label: 'str', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', color: 'str', marker: 'str' = 'o', linestyle: 'str' = '--')

spectraxgk.artifacts.benchmark_plots.cyclone_comparison_figure(ref: CycloneReference, scan: CycloneScanResult, label: str = 'SPECTRAX-GK') → Tuple[Figure, ndarray]

Create a two-panel comparison plot between reference and solver output.

spectraxgk.artifacts.benchmark_plots.cyclone_reference_figure(ref: CycloneReference) → Tuple[Figure, ndarray]

Create a two-panel Cyclone base case reference plot.

spectraxgk.artifacts.benchmark_plots.etg_trend_figure(R_over_LTe: ndarray, gamma: ndarray, omega: ndarray, ky_target: float) → Tuple[Figure, ndarray]

Create a two-panel ETG trend plot versus R/LTe.

spectraxgk.artifacts.benchmark_plots.growth_rate_heatmap(x: ndarray, y: ndarray, gamma: ndarray, title: str, x_label: str, y_label: str, cmap: str = 'jet') → Tuple[Figure, Axes]

Render a growth-rate heatmap versus two gradient axes.

spectraxgk.artifacts.benchmark_plots.linear_validation_figure(panels: list[LinearValidationPanel]) → Tuple[Figure, ndarray]

Create a multi-panel summary plot of eigenfunctions, growth rates, and frequencies.

spectraxgk.artifacts.benchmark_plots.linear_validation_multi_reference_figure(panels: list[MultiReferenceValidationPanel]) → Tuple[Figure, ndarray]

Create summary panels with multiple external reference curves.

spectraxgk.artifacts.benchmark_plots.scan_comparison_figure(x: ndarray, gamma: ndarray, omega: ndarray, x_label: str, title: str, x_ref: ndarray | None = None, gamma_ref: ndarray | None = None, omega_ref: ndarray | None = None, label: str = 'SPECTRAX-GK', ref_label: str = 'Reference', log_x: bool = False) → Tuple[Figure, ndarray]

Create a two-panel comparison plot for a generic scan.

spectraxgk.artifacts.benchmark_plots.scan_multi_reference_figure(x: ndarray, gamma: ndarray, omega: ndarray, x_label: str, title: str, references: list[ReferenceSeries], *, log_x: bool = False) → Tuple[Figure, ndarray]

Create a two-panel comparison figure against multiple reference curves.

Diagnostic Plots

Diagnostic plotting for fitted growth histories and eigenfunctions.

spectraxgk.artifacts.diagnostic_plots.eigenfunction_overlap_summary_figure(ky: ndarray, overlap: ndarray, relative_l2: ndarray, *, title: str = 'Eigenfunction overlap summary', x_label: str = '$k_y \\rho_i$', overlap_label: str = 'Normalized overlap', rel_l2_label: str = 'Relative $L^2$ error', log_x: bool = True) → Tuple[Figure, ndarray]

Render a compact two-panel eigenfunction-overlap summary.

spectraxgk.artifacts.diagnostic_plots.eigenfunction_reference_overlay_figure(theta: ndarray, eigenfunction: ndarray, theta_ref: ndarray, reference: ndarray, *, title: str = 'Eigenfunction overlay') → Tuple[Figure, ndarray]

Render a phase-aligned raw overlay against a frozen reference mode.

spectraxgk.artifacts.diagnostic_plots.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.

Zonal Response Plots

Zonal-response plotting helpers.

spectraxgk.artifacts.zonal_plots.zonal_flow_response_figure(t: ndarray, response: ndarray, *, metrics: Any = None, title: str = 'Zonal-flow response', y_label: str = 'normalized response') → tuple[Figure, ndarray]

Render a zonal-flow response trace and its envelope summary.

Plotting Facade

Public plotting facade for publication-ready figures.

The plotting implementation is split by figure family while this module keeps one stable import path for examples, tools, and external users.

class spectraxgk.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)

class spectraxgk.artifacts.plotting.MultiReferenceValidationPanel(name: 'str', z: 'np.ndarray', eigenfunction: 'np.ndarray', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', x_label: 'str', references: 'list[ReferenceSeries]', log_x: 'bool' = False)

class spectraxgk.artifacts.plotting.ReferenceSeries(label: 'str', x: 'np.ndarray', gamma: 'np.ndarray', omega: 'np.ndarray', color: 'str', marker: 'str' = 'o', linestyle: 'str' = '--')

spectraxgk.artifacts.plotting.cyclone_comparison_figure(ref: CycloneReference, scan: CycloneScanResult, label: str = 'SPECTRAX-GK') → Tuple[Figure, ndarray]

Create a two-panel comparison plot between reference and solver output.

spectraxgk.artifacts.plotting.cyclone_reference_figure(ref: CycloneReference) → Tuple[Figure, ndarray]

Create a two-panel Cyclone base case reference plot.

spectraxgk.artifacts.plotting.eigenfunction_overlap_summary_figure(ky: ndarray, overlap: ndarray, relative_l2: ndarray, *, title: str = 'Eigenfunction overlap summary', x_label: str = '$k_y \\rho_i$', overlap_label: str = 'Normalized overlap', rel_l2_label: str = 'Relative $L^2$ error', log_x: bool = True) → Tuple[Figure, ndarray]

Render a compact two-panel eigenfunction-overlap summary.

spectraxgk.artifacts.plotting.eigenfunction_reference_overlay_figure(theta: ndarray, eigenfunction: ndarray, theta_ref: ndarray, reference: ndarray, *, title: str = 'Eigenfunction overlay') → Tuple[Figure, ndarray]

Render a phase-aligned raw overlay against a frozen reference mode.

spectraxgk.artifacts.plotting.etg_trend_figure(R_over_LTe: ndarray, gamma: ndarray, omega: ndarray, ky_target: float) → Tuple[Figure, ndarray]

Create a two-panel ETG trend plot versus R/LTe.

spectraxgk.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.

spectraxgk.artifacts.plotting.growth_rate_heatmap(x: ndarray, y: ndarray, gamma: ndarray, title: str, x_label: str, y_label: str, cmap: str = 'jet') → Tuple[Figure, Axes]

Render a growth-rate heatmap versus two gradient axes.

spectraxgk.artifacts.plotting.linear_runtime_panel_figure(*, t: ndarray, signal: ndarray, z: ndarray, eigenfunction: ndarray, gamma: float, omega: float, title: str = 'SPECTRAX-GK Linear Runtime') → Tuple[Figure, ndarray]

Create the default two-panel linear runtime plot.

spectraxgk.artifacts.plotting.linear_validation_figure(panels: list[LinearValidationPanel]) → Tuple[Figure, ndarray]

Create a multi-panel summary plot of eigenfunctions, growth rates, and frequencies.

spectraxgk.artifacts.plotting.linear_validation_multi_reference_figure(panels: list[MultiReferenceValidationPanel]) → Tuple[Figure, ndarray]

Create summary panels with multiple external reference curves.

spectraxgk.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 = 'SPECTRAX-GK Nonlinear Runtime') → Tuple[Figure, ndarray]

Create the default three-panel nonlinear runtime plot.

spectraxgk.artifacts.plotting.plot_saved_output(path: str | Path, *, out: str | Path | None = None) → Path

Plot a saved linear or nonlinear output bundle.

spectraxgk.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 = 'SPECTRAX-GK', ref_label: str = 'Reference', log_x: bool = False) → Tuple[Figure, ndarray]

Create a two-panel comparison plot for a generic scan.

spectraxgk.artifacts.plotting.scan_multi_reference_figure(x: ndarray, gamma: ndarray, omega: ndarray, x_label: str, title: str, references: list[ReferenceSeries], *, log_x: bool = False) → Tuple[Figure, ndarray]

Create a two-panel comparison figure against multiple reference curves.

spectraxgk.artifacts.plotting.set_plot_style() → None

Apply the shared publication style used by generated figures.

spectraxgk.artifacts.plotting.zonal_flow_response_figure(t: ndarray, response: ndarray, *, metrics: Any = None, title: str = 'Zonal-flow response', y_label: str = 'normalized response') → tuple[Figure, ndarray]

Render a zonal-flow response trace and its envelope summary.

Config

class spectraxgk.config.CycloneBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=24, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=20.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=150.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=True, diffrax_rtol=1e-06, diffrax_atol=1e-08, diffrax_max_steps=200000, state_sharding=None, progress_bar=False, fixed_dt=False, dt_min=1e-07, dt_max=0.05, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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(R_over_LTi=2.49, R_over_LTe=0.0, R_over_Ln=0.8, nu_i=0.0), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False), reference_aligned: bool = True)

Standard parameters for the Cyclone base case ITG benchmark.

class spectraxgk.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 spectraxgk.config.GridConfig(Nx: int = 48, Ny: int = 48, Nz: int = 64, Lx: float = 62.8, Ly: float = 62.8, boundary: str = 'periodic', jtwist: int | None = None, non_twist: bool = False, kxfac: float = 1.0, z_min: float = -3.141592653589793, z_max: float = 3.141592653589793, y0: float | None = None, ntheta: int | None = None, nperiod: int | None = None, zp: int | None = None)

Spectral grid configuration in a flux-tube.

class spectraxgk.config.InitializationConfig(init_field: str = 'density', init_amp: float = 1e-05, init_single: bool = True, random_seed: int = 22, gaussian_init: bool = False, gaussian_width: float = 0.5, gaussian_envelope_constant: float = 1.0, gaussian_envelope_sine: float = 0.0, kpar_init: float = 0.0, init_file: str | None = None, init_file_scale: float = 1.0, init_file_mode: str = 'replace', init_electrons_only: bool = False)

Initialization options for linear runs.

class spectraxgk.config.KBMBaseCase(grid: GridConfig = GridConfig(Nx=1, Ny=16, Nz=96, Lx=62.8, Ly=62.8, boundary='linked', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=10.0, ntheta=32, nperiod=2, zp=None), time: TimeConfig = TimeConfig(t_max=40.0, dt=0.01, method='rk4', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Tsit5', diffrax_adaptive=True, diffrax_rtol=0.0001, diffrax_atol=1e-07, diffrax_max_steps=20000, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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(R_over_LTi=2.49, R_over_LTe=2.49, R_over_Ln=0.8, Te_over_Ti=1.0, mass_ratio=3703.7037037037035, nu_i=0.0, nu_e=0.0, beta=0.015), init: InitializationConfig = InitializationConfig(init_field='all', init_amp=1e-10, init_single=True, random_seed=22, gaussian_init=True, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False))

Parameters for an electromagnetic KBM benchmark.

class spectraxgk.config.KineticElectronModelConfig(R_over_LTi: float = 2.49, R_over_LTe: float = 2.49, R_over_Ln: float = 0.8, Te_over_Ti: float = 1.0, mass_ratio: float = 3703.7037037037035, nu_i: float = 0.0, nu_e: float = 0.0, beta: float = 1e-05)

Gradients and ratios for a kinetic-electron Cyclone-base-case setup.

class spectraxgk.config.ModelConfig(R_over_LTi: float = 2.49, R_over_LTe: float = 0.0, R_over_Ln: float = 0.8, nu_i: float = 0.0)

Dimensionless gradients for the Cyclone base case.

class spectraxgk.config.TimeConfig(t_max: float = 100.0, dt: float = 0.1, method: str = 'rk2', sample_stride: int = 1, diagnostics_stride: int = 1, diagnostics: bool = True, save_state: bool = False, checkpoint: bool = False, implicit_restart: int = 20, implicit_preconditioner: str | None = None, use_diffrax: bool = True, diffrax_solver: str = 'Dopri8', diffrax_adaptive: bool = False, diffrax_rtol: float = 1e-05, diffrax_atol: float = 1e-07, diffrax_max_steps: int = 4096, state_sharding: str | None = None, progress_bar: bool = False, fixed_dt: bool = True, dt_min: float = 1e-07, dt_max: float | None = None, cfl: float = 0.9, cfl_fac: float | None = None, nstep_restart: int | None = None, collision_split: bool = False, collision_scheme: str = 'implicit', compressed_real_fft: bool = True, nonlinear_dealias: bool = True, laguerre_nonlinear_mode: str = 'grid')

Time integration parameters.

spectraxgk.config.explicit_method_default_cfl_fac(method: str) → float

Return the reference explicit-method CFL prefactor for a given method.

spectraxgk.config.resolve_cfl_fac(method: str, cfl_fac: float | None) → float

Resolve an explicit CFL prefactor, falling back to the method default.

Normalization

Canonical normalization contracts for benchmark families.

This module centralizes the case-level calibration knobs that were previously spread across benchmark runners and tool scripts.

class spectraxgk.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’]

spectraxgk.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.

spectraxgk.diagnostics.normalization.get_normalization_contract(case: str) → NormalizationContract

Return the canonical normalization contract for case.

Energy Diagnostics

Scalar free-energy and field-energy diagnostics.

spectraxgk.diagnostics.energy._masked_abs2(value: Array, active: Array) → Array

Return abs(value)**2 after zeroing inactive spectral modes.

spectraxgk.diagnostics.energy.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).

spectraxgk.diagnostics.energy.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 diagnostic.

spectraxgk.diagnostics.energy.magnetic_vector_potential_energy(apar: Array, cache: LinearCache, vol_fac: Array, *, use_dealias: bool = True) → Array

Magnetic vector-potential field-energy diagnostic.

Transport Diagnostics

Species transport and turbulent-heating diagnostics.

spectraxgk.diagnostics.transport.heat_flux_channel_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → tuple[Array, Array, Array]

Return ES, Apar, and Bpar heat-flux channels per species.

spectraxgk.diagnostics.transport.heat_flux_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → Array

Heat-flux diagnostic per species (gyroBohm units).

spectraxgk.diagnostics.transport.heat_flux_total(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → Array

Total heat-flux diagnostic.

spectraxgk.diagnostics.transport.nonlinear_turbulence_gradient_finite_difference_report(*, baseline: dict[str, Any], plus: dict[str, Any], minus: dict[str, Any], delta_parameter: float, parameter_name: str, baseline_path: str | None = None, plus_path: str | None = None, minus_path: str | None = None, config: NonlinearTurbulenceGradientFiniteDifferenceConfig | None = None) → dict[str, Any]

Build a production long-window central finite-difference gradient gate.

Inputs must be replicated nonlinear_window_ensemble_report payloads for the same nonlinear case and analysis window, differing only by the perturbed parameter. The report computes the central finite-difference heat-flux gradient and checks that the response is resolved above ensemble uncertainty before allowing any turbulence-gradient claim.

spectraxgk.diagnostics.transport.particle_flux_channel_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → tuple[Array, Array, Array]

Return ES, Apar, and Bpar particle-flux channels per species.

spectraxgk.diagnostics.transport.particle_flux_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → Array

Particle-flux diagnostic per species.

spectraxgk.diagnostics.transport.particle_flux_total(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → Array

Total particle-flux diagnostic.

spectraxgk.diagnostics.transport.turbulent_heating_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) → Array

Turbulent-heating diagnostic per species.

spectraxgk.diagnostics.transport.turbulent_heating_total(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) → Array

Total turbulent-heating diagnostic.

Resolved Diagnostics

Resolved spectral diagnostics for nonlinear outputs.

spectraxgk.diagnostics.resolved._reduce_scalar_kykxz(contrib: Array) → tuple[Array, Array, Array, Array, Array]

Reduce a (ky, kx, z) contribution to spectral diagnostics views.

spectraxgk.diagnostics.resolved._reduce_species_kykxz(contrib: Array) → tuple[Array, Array, Array, Array, Array]

Reduce a (species, ky, kx, z) contribution to spectral diagnostics views.

spectraxgk.diagnostics.resolved.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.

spectraxgk.diagnostics.resolved.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.

spectraxgk.diagnostics.resolved.heat_flux_channel_resolved_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → tuple[tuple[Array, Array, Array, Array, Array], tuple[Array, Array, Array, Array, Array], tuple[Array, Array, Array, Array, Array]]

Return resolved ES, Apar, and Bpar heat-flux channels per species.

spectraxgk.diagnostics.resolved.heat_flux_resolved_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → tuple[Array, Array, Array, Array, Array]

Return Resolved heat-flux reductions per species.

spectraxgk.diagnostics.resolved.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.

spectraxgk.diagnostics.resolved.particle_flux_channel_resolved_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → tuple[tuple[Array, Array, Array, Array, Array], tuple[Array, Array, Array, Array, Array], tuple[Array, Array, Array, Array, Array]]

Return resolved ES, Apar, and Bpar particle-flux channels per species.

spectraxgk.diagnostics.resolved.particle_flux_resolved_species(G: Array, phi: Array, apar: Array, bpar: Array, cache: LinearCache, grid: SpectralGrid, params: LinearParams, flux_fac: Array, *, use_dealias: bool = True, flux_scale: float = 1.0) → tuple[Array, Array, Array, Array, Array]

Return Resolved particle-flux reductions per species.

spectraxgk.diagnostics.resolved.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.

spectraxgk.diagnostics.resolved.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.

spectraxgk.diagnostics.resolved.zonal_phi_line_kxt(phi: Array, grid: SpectralGrid) → Array

Return the signed, unweighted line-averaged zonal potential per k_x.

spectraxgk.diagnostics.resolved.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

Unified runtime configuration schema for linear/nonlinear GK runs.

class spectraxgk.workflows.runtime.config.RuntimeCollisionConfig(nu_hermite: float = 1.0, nu_laguerre: float = 2.0, nu_hyper: float = 0.0, p_hyper: float = 4.0, nu_hyper_l: float = 0.0, nu_hyper_m: float = 1.0, nu_hyper_lm: float = 0.0, p_hyper_l: float = 6.0, p_hyper_m: float | None = None, p_hyper_lm: float = 6.0, D_hyper: float = 0.0, p_hyper_kperp: float = 2.0, hypercollisions_const: float = 0.0, hypercollisions_kz: float = 1.0, damp_ends_amp: float = 0.1, damp_ends_widthfrac: float = 0.125, damp_ends_scale_by_dt: bool = False)

Collision and end-damping parameters.

class spectraxgk.workflows.runtime.config.RuntimeConfig(grid: GridConfig = GridConfig(Nx=48, Ny=48, Nz=64, Lx=62.8, Ly=62.8, boundary='periodic', jtwist=None, non_twist=False, kxfac=1.0, z_min=-3.141592653589793, z_max=3.141592653589793, y0=None, ntheta=None, nperiod=None, zp=None), time: TimeConfig = TimeConfig(t_max=100.0, dt=0.1, method='rk2', sample_stride=1, diagnostics_stride=1, diagnostics=True, save_state=False, checkpoint=False, implicit_restart=20, implicit_preconditioner=None, use_diffrax=True, diffrax_solver='Dopri8', diffrax_adaptive=False, diffrax_rtol=1e-05, diffrax_atol=1e-07, diffrax_max_steps=4096, state_sharding=None, progress_bar=False, fixed_dt=True, dt_min=1e-07, dt_max=None, cfl=0.9, cfl_fac=None, nstep_restart=None, collision_split=False, collision_scheme='implicit', compressed_real_fft=True, nonlinear_dealias=True, laguerre_nonlinear_mode='grid'), geometry: GeometryConfig = GeometryConfig(model='s-alpha', geometry_backend='auto', geometry_file=None, vmec_file=None, 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=1.0, B0=1.0, alpha=0.0, drift_scale=1.0, kperp2_bmag=True, bessel_bmag_power=0.0), init: InitializationConfig = InitializationConfig(init_field='density', init_amp=1e-05, init_single=True, random_seed=22, gaussian_init=False, gaussian_width=0.5, gaussian_envelope_constant=1.0, gaussian_envelope_sine=0.0, kpar_init=0.0, init_file=None, init_file_scale=1.0, init_file_mode='replace', init_electrons_only=False), species: Tuple[RuntimeSpeciesConfig, ...] = (RuntimeSpeciesConfig(name='ion', charge=1.0, mass=1.0, density=1.0, temperature=1.0, tprim=2.49, fprim=0.8, nu=0.0, kinetic=True),), physics: RuntimePhysicsConfig = RuntimePhysicsConfig(reduced_model='gyrokinetic', linear=True, nonlinear=False, electrostatic=True, electromagnetic=False, use_apar=False, use_bpar=False, adiabatic_electrons=True, adiabatic_ions=False, tau_e=1.0, tau_fac=None, z_ion=1.0, beta=0.0, collisions=True, hypercollisions=True), collisions: RuntimeCollisionConfig = RuntimeCollisionConfig(nu_hermite=1.0, nu_laguerre=2.0, nu_hyper=0.0, p_hyper=4.0, nu_hyper_l=0.0, nu_hyper_m=1.0, nu_hyper_lm=0.0, p_hyper_l=6.0, p_hyper_m=None, p_hyper_lm=6.0, D_hyper=0.0, p_hyper_kperp=2.0, hypercollisions_const=0.0, hypercollisions_kz=1.0, damp_ends_amp=0.1, damp_ends_widthfrac=0.125, damp_ends_scale_by_dt=False), normalization: RuntimeNormalizationConfig = RuntimeNormalizationConfig(contract='cyclone', rho_star=None, omega_d_scale=None, omega_star_scale=None, diagnostic_norm='rho_star', flux_scale=1.0, wphi_scale=1.0), terms: RuntimeTermsConfig = RuntimeTermsConfig(streaming=1.0, mirror=1.0, curvature=1.0, gradb=1.0, diamagnetic=1.0, collisions=1.0, hypercollisions=1.0, hyperdiffusion=0.0, end_damping=1.0, apar=1.0, bpar=1.0, nonlinear=0.0), expert: RuntimeExpertConfig = RuntimeExpertConfig(fixed_mode=False, iky_fixed=None, ikx_fixed=None, dealias_kz=False, source='default', phi_ext=0.0), output: RuntimeOutputConfig = RuntimeOutputConfig(path=None, restart=False, restart_if_exists=False, save_for_restart=True, restart_to_file=None, restart_from_file=None, restart_with_perturb=False, append_on_restart=True, resolved_diagnostics=True, restart_scale=1.0, nsave=10000), quasilinear: RuntimeQuasilinearConfig = RuntimeQuasilinearConfig(enabled=False, mode='weights', saturation_rule='none', amplitude_normalization='phi_rms', kperp_average='phi_weighted', csat=1.0, gamma_floor=0.0, include_stable_modes=False, delta_ky='auto', species='all', channels=('es',), write_spectrum=True, output_path=None), parallel: RuntimeParallelConfig = RuntimeParallelConfig(strategy='serial', axis='ky', batch_size=None, num_devices=None, strict_identity=True, profile=False, backend='auto'))

Unified simulation config for runtime-driven GK runs.

class spectraxgk.workflows.runtime.config.RuntimeExpertConfig(fixed_mode: bool = False, iky_fixed: int | None = None, ikx_fixed: int | None = None, dealias_kz: bool = False, source: str = 'default', phi_ext: float = 0.0)

Advanced runtime controls that should rarely be needed.

class spectraxgk.workflows.runtime.config.RuntimeNormalizationConfig(contract: str = 'cyclone', rho_star: float | None = None, omega_d_scale: float | None = None, omega_star_scale: float | None = None, diagnostic_norm: str = 'rho_star', flux_scale: float = 1.0, wphi_scale: float = 1.0)

Normalization contract selection + optional explicit overrides.

class spectraxgk.workflows.runtime.config.RuntimeOutputConfig(path: str | None = None, restart: bool = False, restart_if_exists: bool = False, save_for_restart: bool = True, restart_to_file: str | None = None, restart_from_file: str | None = None, restart_with_perturb: bool = False, append_on_restart: bool = True, resolved_diagnostics: bool = True, restart_scale: float = 1.0, nsave: int = 10000)

Artifact-output controls for runtime executable entry points.

class spectraxgk.workflows.runtime.config.RuntimeParallelConfig(strategy: str = 'serial', axis: str = 'ky', batch_size: int | None = None, num_devices: int | None = None, strict_identity: bool = True, profile: bool = False, backend: str = 'auto')

Parallel-execution policy for independent scans and future sharded paths.

class spectraxgk.workflows.runtime.config.RuntimePhysicsConfig(reduced_model: str = 'gyrokinetic', linear: bool = True, nonlinear: bool = False, electrostatic: bool = True, electromagnetic: bool = False, use_apar: bool = False, use_bpar: bool = False, adiabatic_electrons: bool = True, adiabatic_ions: bool = False, tau_e: float = 1.0, tau_fac: float | None = None, z_ion: float = 1.0, beta: float = 0.0, collisions: bool = True, hypercollisions: bool = True)

Physics-family toggles independent from benchmark case names.

class spectraxgk.workflows.runtime.config.RuntimeQuasilinearConfig(enabled: bool = False, mode: str = 'weights', saturation_rule: str = 'none', amplitude_normalization: str = 'phi_rms', kperp_average: str = 'phi_weighted', csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False, delta_ky: str | float = 'auto', species: str = 'all', channels: Tuple[str, ...] = ('es',), write_spectrum: bool = True, output_path: str | None = None)

Quasilinear transport diagnostics computed from linear states.

class spectraxgk.workflows.runtime.config.RuntimeSpeciesConfig(name: str = 'ion', charge: float = 1.0, mass: float = 1.0, density: float = 1.0, temperature: float = 1.0, tprim: float = 2.49, fprim: float = 0.8, nu: float = 0.0, kinetic: bool = True)

Single species definition for runtime-configured simulations.

class spectraxgk.workflows.runtime.config.RuntimeTermsConfig(streaming: float = 1.0, mirror: float = 1.0, curvature: float = 1.0, gradb: float = 1.0, diamagnetic: float = 1.0, collisions: float = 1.0, hypercollisions: float = 1.0, hyperdiffusion: float = 0.0, end_damping: float = 1.0, apar: float = 1.0, bpar: float = 1.0, nonlinear: float = 0.0)

Term toggles for assembly; applies to linear and nonlinear paths.

Runtime 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.

spectraxgk.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.

spectraxgk.workflows.runtime.startup._build_initial_condition(grid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, cfg: RuntimeConfig, *, ky_index: int, kx_index: int, Nl: int, Nm: int, nspecies: int) → Array

Build the runtime initial state using this module’s patchable params builder.

spectraxgk.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.

spectraxgk.workflows.runtime.startup._centered_glibc_random_pairs(seed: int, count: int) → ndarray

Return centered random pairs using glibc rand() semantics.

spectraxgk.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.

spectraxgk.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 spectraxgk.terms.fields.

spectraxgk.workflows.runtime.startup._enforce_full_ky_hermitian(arr: ndarray) → ndarray

Mirror positive-ky content into the negative branch for full FFT grids.

spectraxgk.workflows.runtime.startup._periodic_zp_from_grid(z: ndarray) → float

Return periodic Zp from the discrete theta grid.

spectraxgk.workflows.runtime.startup.build_runtime_geometry(cfg: RuntimeConfig) → SAlphaGeometry | SlabGeometry | FluxTubeGeometryData

Resolve runtime geometry, generating *.eik.nc geometry when needed.

spectraxgk.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.

spectraxgk.workflows.runtime.startup.build_runtime_linear_terms(cfg: RuntimeConfig) → LinearTerms

Build LinearTerms from unified toggles.

spectraxgk.workflows.runtime.startup.build_runtime_term_config(cfg: RuntimeConfig) → TermConfig

Build nonlinear-ready TermConfig from unified toggles.

spectraxgk.workflows.runtime.startup.load_netcdf_restart_state(path: str | Path, *, nspecies: int, Nl: int, Nm: int, ny: int, nx: int, nz: int) → ndarray

Load a NetCDF restart file into SPECTRAX’s full Hermitian layout.

spectraxgk.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.

Runtime Policies

Pure runtime policy helpers shared by runtime runners and tests.

class spectraxgk.workflows.runtime.policies.RuntimeIndependentParallelPlan(requested_workers: int, effective_workers: int, executor: str, strategy: str, axis: str, source: str, problem_size: int)

Resolved independent-worker policy for runtime scan workloads.

property enabled: bool

Whether the resolved plan uses more than one independent worker.

to_dict() → dict[str, Any]

Return a JSON-friendly policy payload for runtime artifacts.

spectraxgk.workflows.runtime.policies._infer_runtime_nonlinear_steps(cfg: RuntimeConfig, *, dt: float, steps: int | None) → int

Infer nonlinear explicit step counts with the same dt ceiling as the integrator.

spectraxgk.workflows.runtime.policies._parallel_requests_combined_ky_scan(cfg: RuntimeConfig) → bool

Return whether runtime parallel config requests the combined-ky scan path.

spectraxgk.workflows.runtime.policies._runtime_external_phi(cfg: RuntimeConfig) → float | None

Return a runtime external-phi source if requested.

spectraxgk.workflows.runtime.policies._runtime_independent_parallel_plan(cfg: RuntimeConfig, *, problem_size: int, workers: int, executor: str) → RuntimeIndependentParallelPlan

Resolve independent k_y worker policy from arguments and config.

spectraxgk.workflows.runtime.policies.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.

Runtime Diagnostics

Runtime linear-fit and quasilinear diagnostic helpers.

class spectraxgk.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)

Linear runtime fit payload before diagnostic normalization.

class spectraxgk.workflows.runtime.diagnostics.RuntimeQuasilinearFinalizationDeps(build_linear_cache: Any, compute_quasilinear_from_linear_state: Any, linear_terms_to_term_config: Any)

Injected dependencies for runtime quasilinear post-processing.

spectraxgk.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, deps: RuntimeQuasilinearFinalizationDeps, status_callback: Any | None = None) → RuntimeLinearResult

Attach optional quasilinear diagnostics to a linear runtime result.

spectraxgk.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, extract_mode_time_series_fn: Any = <function extract_mode_time_series>, fit_growth_rate_auto_with_stats_fn: Any = <function fit_growth_rate_auto_with_stats>, fit_growth_rate_auto_fn: Any = <function fit_growth_rate_auto>, fit_growth_rate_fn: Any = <function fit_growth_rate>, extract_eigenfunction_fn: Any = <function extract_eigenfunction>) → RuntimeLinearFitResult

Fit linear growth/frequency and extract the eigenfunction diagnostic.

spectraxgk.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) → RuntimeLinearResult

Refit one stored trajectory without repeating its integration.

Runtime Diagnostic Arrays

Runtime diagnostic array validation and composition helpers.

These helpers operate on already-sampled diagnostic payloads. They deliberately stay host-side: runtime drivers use them to fail fast on invalid artifacts and to combine adaptive chunks without mixing that array bookkeeping into the linear-fit and quasilinear-finalization code.

spectraxgk.workflows.runtime.diagnostic_arrays.concat_runtime_diagnostics(diags: Sequence[SimulationDiagnostics]) → SimulationDiagnostics

Concatenate one or more diagnostic chunks.

spectraxgk.workflows.runtime.diagnostic_arrays.slice_runtime_diagnostics(diag: SimulationDiagnostics, stop: int) → SimulationDiagnostics

Return the first stop diagnostic samples.

spectraxgk.workflows.runtime.diagnostic_arrays.stride_runtime_diagnostics(diag: SimulationDiagnostics, *, stride: int) → SimulationDiagnostics

Apply the runtime output stride after concatenating chunk diagnostics.

spectraxgk.workflows.runtime.diagnostic_arrays.truncate_runtime_diagnostics(diag: SimulationDiagnostics, *, t_max: float) → SimulationDiagnostics

Keep samples through the first entry that reaches t_max.

spectraxgk.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 Initial Conditions

Runtime initial-condition and restart-state construction.

class spectraxgk.workflows.runtime.initial_conditions._InitialConditionBuilder(grid: 'SpectralGrid', geom: 'FluxTubeGeometryLike', cfg: 'RuntimeConfig', Nl: 'int', Nm: 'int', state: 'np.ndarray', species_targets: 'tuple[int, ...]', build_runtime_linear_params_fn: 'Callable[..., LinearParams]', phi_seed_context: 'tuple[object, LinearParams] | None' = None)

spectraxgk.workflows.runtime.initial_conditions._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.

spectraxgk.workflows.runtime.initial_conditions._centered_glibc_random_pairs(seed: int, count: int) → ndarray

Return centered random pairs using glibc rand() semantics.

spectraxgk.workflows.runtime.initial_conditions._dealiased_initial_mode_pairs(grid: SpectralGrid) → list[tuple[int, int]]

Return the dealiased startup-loop (kx, ky) pairs for multimode initial conditions.

spectraxgk.workflows.runtime.initial_conditions._enforce_full_ky_hermitian(arr: ndarray) → ndarray

Mirror positive-ky content into the negative branch for full FFT grids.

spectraxgk.workflows.runtime.initial_conditions._periodic_zp_from_grid(z: ndarray) → float

Return periodic Zp from the discrete theta grid.

Runtime Phi Initializer

Runtime electrostatic-potential initial-condition inversion.

class spectraxgk.workflows.runtime.initial_phi._PhiMomentAlgebra(Jl: 'np.ndarray', jacobian: 'np.ndarray', charge: 'np.ndarray', density: 'np.ndarray', denom: 'np.ndarray', tau_e: 'float')

spectraxgk.workflows.runtime.initial_phi._as_runtime_species_array(value: float | Array, nspecies: int, name: str) → ndarray

Return a length-nspecies NumPy array for startup-only algebra.

spectraxgk.workflows.runtime.initial_phi._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 spectraxgk.terms.fields.

Runtime Chunks

Adaptive runtime chunk execution helpers.

These helpers own the repeated adaptive runtime chunk loop used by the runtime drivers. Keeping the loop outside runtime.py makes the execution layer smaller without changing the accepted diagnostics truncation/stride behavior.

class spectraxgk.workflows.runtime.chunks.AdaptiveChunkResult(diagnostics: SimulationDiagnostics, state: Any, fields: FieldState)

Concatenated result from one adaptive runtime chunk loop.

class spectraxgk.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.

spectraxgk.workflows.runtime.chunks._chunk_end_time(diag: SimulationDiagnostics, *, label: str, chunk_index: int) → float

Return the last diagnostic time sample for one adaptive chunk.

spectraxgk.workflows.runtime.chunks._effective_diagnostics_stride(diagnostics_stride: int) → int

Normalize runtime diagnostic stride with a floor at one.

spectraxgk.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.

spectraxgk.workflows.runtime.chunks._offset_chunk_diagnostics_time(diag: SimulationDiagnostics, *, offset: float) → SimulationDiagnostics

Return a chunk diagnostic payload shifted onto the accumulated time axis.

spectraxgk.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.

spectraxgk.workflows.runtime.chunks.format_duration(seconds: float) → str

Format elapsed seconds as MM:SS or H:MM:SS.

spectraxgk.workflows.runtime.chunks.run_adaptive_runtime_chunk_loop(*, integrate_chunk: Callable[[bool], 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) → AdaptiveChunkResult

Run repeated diagnostic chunks until t_max is reached.

integrate_chunk must return (t_chunk, diag_chunk, state, fields) in the same contract used by the runtime integrators.

Runtime Results

Runtime result containers and small assembly helpers.

class spectraxgk.workflows.runtime.results.RuntimeLinearResult(ky: float, gamma: float, omega: float, selection: ModeSelection, t: ndarray | None = None, signal: ndarray | None = None, field_history: ndarray | None = None, state: ndarray | None = None, z: ndarray | None = None, eigenfunction: ndarray | None = None, fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, fit_signal_used: str | None = None, quasilinear: dict[str, Any] | None = None)

Result container for runtime linear runs.

class spectraxgk.workflows.runtime.results.RuntimeLinearScanResult(ky: ndarray, gamma: ndarray, omega: ndarray, quasilinear: tuple[dict[str, Any], ...] | None = None, parallel: dict[str, Any] | None = None)

Result container for runtime linear ky scans.

class spectraxgk.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)

Result container for runtime nonlinear runs.

class spectraxgk.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.

spectraxgk.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) → RuntimeNonlinearResult

Build a runtime nonlinear result with optional final-field summary.

spectraxgk.workflows.runtime.results.nonlinear_field_phi2(fields: FieldState) → ndarray

Return the mean electrostatic energy density from final fields.

Runtime Orchestration

Runtime ky-scan orchestration and combined-batch scan policy.

class spectraxgk.workflows.runtime.orchestration_scan.RuntimeScanBatchDeps(*args, **kwargs)

Dependency surface needed by the combined-ky scan batch helper.

class spectraxgk.workflows.runtime.orchestration_scan.RuntimeScanDeps(*args, **kwargs)

Dependency surface for runtime ky-scan orchestration.

spectraxgk.workflows.runtime.orchestration_scan.build_runtime_scan_batch_deps(facade: Any) → RuntimeScanBatchDeps

Build combined-ky scan deps from the public runtime facade.

spectraxgk.workflows.runtime.orchestration_scan.build_runtime_scan_orchestration_deps(facade: Any) → RuntimeScanDeps

Build ky-scan orchestration deps from the public runtime facade.

spectraxgk.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) → RuntimeParameterScanResult

Run a scalar scan and optionally continue a selected solution branch.

spectraxgk.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, deps: RuntimeScanBatchDeps) → RuntimeLinearScanResult

Batch a ky scan using one time integration over the full grid.

spectraxgk.workflows.runtime.orchestration_scan.run_runtime_scan_ky_task(task: dict[str, Any], *, run_runtime_linear: Callable[[...], Any]) → Any

Run one independent ky point for ordered scan-worker execution.

spectraxgk.workflows.runtime.orchestration_scan.run_runtime_scan_orchestration(cfg: RuntimeConfig, ky_values: Any, *, Nl: int | None, Nm: int | None, solver: str, method: str | None, dt: float | None, steps: int | None, sample_stride: int | None, batch_ky: bool, 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, krylov_cfg: Any, mode_method: str, fit_signal: str, show_progress: bool, workers: int, parallel_executor: str, deps: RuntimeScanDeps) → RuntimeLinearScanResult

Coordinate serial, independent-worker, or combined-ky runtime scans.

Runtime artifact display, restart, and checkpoint handoff policies.

class spectraxgk.workflows.runtime.orchestration_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.

class spectraxgk.workflows.runtime.orchestration_artifacts.RuntimeArtifactHandoffDeps(is_netcdf_output_target: Callable[[Path], bool], resolve_restart_path: Callable[[str | Path, Any], Path], resolve_restart_write_path: Callable[[str | Path, Any], Path], netcdf_bundle_base: Callable[[Path], Path], load_nonlinear_netcdf_diagnostics: Callable[[str | Path], SimulationDiagnostics], condense_diagnostics_for_netcdf_output: Callable[[SimulationDiagnostics], SimulationDiagnostics], concat_runtime_diagnostics: Callable[[list[SimulationDiagnostics]], SimulationDiagnostics], validate_finite_runtime_result: Callable[[Any], None], run_runtime_nonlinear: Callable[[...], RuntimeNonlinearResult], write_runtime_nonlinear_artifacts: Callable[[str | Path, Any, Any], dict[str, str]])

Patchable functions used by nonlinear artifact handoff orchestration.

spectraxgk.workflows.runtime.orchestration_artifacts.print_nonlinear_command_outputs(paths: Mapping[str, str], *, enabled: bool) → None

Print nonlinear artifact paths after diagnostics confirm a saved run.

spectraxgk.workflows.runtime.orchestration_artifacts.print_saved_paths(paths: Mapping[str, str], keys: Sequence[str]) → None

Print saved artifact paths in command-defined display order.

spectraxgk.workflows.runtime.orchestration_artifacts.resolve_nonlinear_artifact_policy(cfg: Any, *, out: str | Path | None, diagnostics: bool | None, steps: int | None, dt: float | None, deps: RuntimeArtifactHandoffDeps) → NonlinearArtifactPolicy

Resolve nonlinear output, restart, and checkpoint policy.

spectraxgk.workflows.runtime.orchestration_artifacts.run_runtime_nonlinear_artifact_handoff(cfg: Any, *, out: str | Path | None, ky_target: float, kx_target: float | None = None, Nl: int | None = None, Nm: int | None = None, dt: float | None = None, steps: int | None = None, method: str | None = None, sample_stride: int | None = None, diagnostics_stride: int | None = None, laguerre_mode: str | None = None, diagnostics: bool | None = None, show_progress: bool = False, status_callback: Any = None, deps: RuntimeArtifactHandoffDeps) → tuple[RuntimeNonlinearResult, dict[str, str]]

Run nonlinear runtime chunks and hand results to artifact writers.

spectraxgk.workflows.runtime.orchestration_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.

spectraxgk.workflows.runtime.orchestration_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.

spectraxgk.workflows.runtime.orchestration_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.

Runtime Commands

Runtime executable command workflows.

class spectraxgk.workflows.runtime.commands.RuntimeCommandDeps(load_runtime_from_toml: Callable[[str | Path], tuple[RuntimeConfig, dict[str, Any]]], run_runtime_linear: Callable[[...], RuntimeLinearResult], run_runtime_scan: Callable[[...], Any], run_runtime_nonlinear_with_artifacts: Callable[[...], tuple[RuntimeNonlinearResult, dict[str, str]]], write_runtime_linear_artifacts: Callable[[str | Path, RuntimeLinearResult], dict[str, str]], write_runtime_linear_scan_artifacts: Callable[[str | Path, Any], dict[str, str]], write_quasilinear_artifacts: Callable[[str | Path, dict[str, Any]], dict[str, str]], resolve_runtime_path: Callable[[...], str | None])

Patchable dependencies for executable runtime subcommands.

class spectraxgk.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 spectraxgk.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)

Resolved executable options for one nonlinear runtime command.

class spectraxgk.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)

Resolved executable options for one linear scan command.

spectraxgk.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.

spectraxgk.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.

spectraxgk.workflows.runtime.commands._prepare_runtime_command_config(args: Any, *, deps: RuntimeCommandDeps, path_overrides: bool, quasilinear_overrides: bool) → tuple[RuntimeConfig, dict[str, Any]]

Load runtime command config and apply the command-specific overrides.

spectraxgk.workflows.runtime.commands._resolve_grid_time_options(args: Any, section: dict[str, Any], cfg: RuntimeConfig) → tuple[int, int, str | None, float | None, int | None, int]

Resolve resolution, optional time controls, and output cadence.

spectraxgk.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.

spectraxgk.workflows.runtime.commands._resolve_linear_fit_options(args: Any, section: dict[str, Any]) → tuple[str, str]

Resolve the linear eigensignal solver and fit signal.

spectraxgk.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.

spectraxgk.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.

spectraxgk.workflows.runtime.commands._runtime_fit_config(data: dict[str, Any]) → dict[str, Any]

Return fit options supported by runtime executable commands.

spectraxgk.workflows.runtime.commands.apply_quasilinear_overrides(cfg: RuntimeConfig, args: Any) → RuntimeConfig

Apply executable quasilinear diagnostic overrides.

spectraxgk.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.

spectraxgk.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.

spectraxgk.workflows.runtime.commands.build_runtime_command_deps(facade: Any) → RuntimeCommandDeps

Build runtime command dependencies from a patchable executable facade.

spectraxgk.workflows.runtime.commands.load_runtime_command_config(args: Any, *, deps: RuntimeCommandDeps) → tuple[RuntimeConfig, dict[str, Any]]

Load runtime TOML data, reusing the generic-dispatch preload if present.

spectraxgk.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.

spectraxgk.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.

spectraxgk.workflows.runtime.commands.print_nonlinear_run_header(*, config_path: str, 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.

spectraxgk.workflows.runtime.commands.print_nonlinear_run_summary(result: Any) → bool

Print final nonlinear diagnostics and return whether diagnostics exist.

spectraxgk.workflows.runtime.commands.run_runtime_linear_command(args: Any, *, deps: RuntimeCommandDeps) → int

Execute the runtime-linear subcommand after parser dispatch.

spectraxgk.workflows.runtime.commands.run_runtime_nonlinear_command(args: Any, *, deps: RuntimeCommandDeps) → int

Execute the runtime-nonlinear subcommand after parser dispatch.

spectraxgk.workflows.runtime.commands.runtime_output_path(args: Any, cfg: RuntimeConfig) → str | None

Return the executable output path override or TOML output path.

spectraxgk.workflows.runtime.commands.scan_runtime_linear_command(args: Any, *, deps: RuntimeCommandDeps) → int

Execute the runtime-linear ky-scan subcommand after parser dispatch.

spectraxgk.workflows.runtime.commands.should_show_progress(args: Any, configured: bool) → bool

Resolve progress output from executable flags, TOML config, and TTY state.

Runtime TOML Inputs

TOML-based input helpers for the executable and driver scripts.

spectraxgk.workflows.runtime.toml._apply_runtime_section_overrides(cfg: RuntimeConfig, data: dict[str, Any]) → RuntimeConfig

Apply non-nested runtime config sections from TOML data.

spectraxgk.workflows.runtime.toml._merge_dataclass(base: Any, overrides: dict | None) → Any

Recursively merge a dict into a dataclass, returning a new instance.

spectraxgk.workflows.runtime.toml._normalize_geometry_overrides(overrides: dict | None) → dict | None

Return geometry overrides using the canonical runtime schema.

spectraxgk.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.

spectraxgk.workflows.runtime.toml._resolve_runtime_config_paths(cfg: RuntimeConfig, *, base_dir: Path) → RuntimeConfig

Resolve every path-valued runtime field against the TOML directory.

spectraxgk.workflows.runtime.toml._runtime_base_config(data: dict[str, Any]) → RuntimeConfig

Return a runtime config after applying common dataclass sections.

spectraxgk.workflows.runtime.toml._runtime_species_from_toml(species_raw: Any) → tuple[RuntimeSpeciesConfig, ...] | None

Parse optional [[species]] runtime entries.

spectraxgk.workflows.runtime.toml.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 spectraxgk case.toml shorthand.

spectraxgk.workflows.runtime.toml.is_runtime_toml(data: dict[str, Any]) → bool

Return whether a parsed input uses the supported runtime schema.

spectraxgk.workflows.runtime.toml.load_runtime_from_toml(path: str | Path) → tuple[RuntimeConfig, dict]

Load unified runtime config from TOML, returning (cfg, data).

spectraxgk.workflows.runtime.toml.load_toml(path: str | Path) → dict

Load a TOML file into a plain dictionary.

spectraxgk.workflows.runtime.toml.resolve_runtime_path(value: str | None, *, base_dir: Path) → str | None

Expand and resolve a runtime config path.

Applies $VAR and ~ expansion, then resolves relative paths against base_dir. If an unresolved $VAR remains after expansion (env var not set), the original value is returned unchanged so downstream code can raise a clearer error. None is passed through.

Parameters:

• value (str or None) – Raw path string from a TOML config or CLI flag.

• base_dir (Path) – Directory used to resolve relative paths. Callers typically pass the config file’s parent directory (TOML values) or Path.cwd() (CLI-supplied values).

Returns:

Absolute resolved path as a string, or None when value is None.

Return type:

str or None

spectraxgk.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.

spectraxgk.artifacts.validate_finite_array(value: Any, *, label: str) → None

Raise if an optional artifact array contains NaN or infinite values.

spectraxgk.artifacts.validate_finite_runtime_result(result: Any, *, label: str) → None

Validate nonlinear runtime result payloads before artifact writes.

spectraxgk.artifacts.write_quasilinear_artifacts(out: str | Path, quasilinear: dict[str, Any]) → dict[str, str]

Write quasilinear summary and species tables.

spectraxgk.artifacts.write_runtime_linear_artifacts(out: str | Path, result: Any) → dict[str, str]

Write summary/timeseries/state artifacts for a linear runtime run.

spectraxgk.artifacts.write_runtime_linear_scan_artifacts(out: str | Path, result: Any) → dict[str, str]

Write ky-scan growth/frequency and optional quasilinear spectra.

spectraxgk.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

Generic runtime artifact path, validation, and file-writing helpers.

spectraxgk.artifacts.io.validate_finite_array(value: Any, *, label: str) → None

Raise if an optional artifact array contains NaN or infinite values.

spectraxgk.artifacts.io.validate_finite_runtime_result(result: Any, *, label: str) → None

Validate nonlinear runtime result payloads before artifact writes.

Runtime Restart Artifacts

Restart-state IO helpers.

SPECTRAX-GK uses the dealiased NetCDF restart layout so that: - runtime init_file can consume restart files directly - users can roundtrip restart state through NetCDF audit files

The file format is a raw complex64 buffer with no header. Consumers must already know the target shape from (nspecies, Nl, Nm, Ny, Nx, Nz).

spectraxgk.artifacts.restart._expand_netcdf_restart_state_full_ky(state_active: ndarray, *, nx_full: int) → ndarray

Expand a NetCDF restart that already stores the full ky axis.

spectraxgk.artifacts.restart.load_netcdf_restart_state(path: str | Path, *, nspecies: int, Nl: int, Nm: int, ny: int, nx: int, nz: int) → ndarray

Load a NetCDF restart file into SPECTRAX’s full Hermitian layout.

spectraxgk.artifacts.restart.write_netcdf_restart_state(path: str | Path, state: Array | ndarray | bool | number | bool | int | float | complex) → Path

Write a restart state in flat complex64 restart layout.

Runtime Artifact Linear Writers

Linear and quasilinear runtime artifact writers.

spectraxgk.artifacts.linear.write_quasilinear_artifacts(out: str | Path, quasilinear: dict[str, Any]) → dict[str, str]

Write quasilinear summary and species tables.

spectraxgk.artifacts.linear.write_runtime_linear_artifacts(out: str | Path, result: Any) → dict[str, str]

Write summary/timeseries/state artifacts for a linear runtime run.

spectraxgk.artifacts.linear.write_runtime_linear_scan_artifacts(out: str | Path, result: Any) → dict[str, str]

Write ky-scan growth/frequency and optional quasilinear spectra.

Runtime Artifact Nonlinear Writers

Generic nonlinear runtime artifact summaries and table writers.

spectraxgk.artifacts.nonlinear.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.

spectraxgk.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.

spectraxgk.artifacts.spectral_layout._condense_ky_for_output(arr: ndarray, *, full_ny: int, active_ny: int) → ndarray

Return ky-resolved data on the dealiased positive-ky output axis.

spectraxgk.artifacts.spectral_layout._condense_kykx_for_output(arr: ndarray, *, full_ny: int, full_nx: int, active_ny: int, active_nx: int) → ndarray

Return ky-kx-resolved data on dealiased output axes.

spectraxgk.artifacts.spectral_layout.infer_triple_dealiased_ny(nky_positive: int) → int

Infer the full Ny from the number of positive k_y points.

Reference real-FFT outputs typically store only the non-negative k_y branch. For the linked-boundary spectral grid used here, the corresponding real-space Ny follows Ny = 3 * (nky - 1) + 1.

Nonlinear Output NetCDF Geometry

Geometry and particle-moment helpers for nonlinear NetCDF output.

spectraxgk.artifacts.nonlinear_netcdf_geometry._build_output_grid_and_geometry(cfg: Any) → tuple[Any, Any]

Resolve artifact output onto the same geometry-implied grid as the solver.

Nonlinear Output NetCDF Field Writer

Final-field *.big.nc writer for nonlinear NetCDF bundles.

spectraxgk.artifacts.nonlinear_netcdf_fields._write_big_netcdf(Dataset: Any, big_path: str | Path, result: Any, cfg: Any, *, x_vals: ndarray, y_vals: ndarray, theta: ndarray, kx_vals: ndarray, ky_vals: ndarray, nspecies: int, nl: int, nm: int, time_vals: ndarray) → str | None

Write final spectral/real-space fields and moments when fields exist.

Nonlinear Output NetCDF Diagnostics Writer

Diagnostic-history writer for nonlinear NetCDF output bundles.

spectraxgk.artifacts.nonlinear_netcdf_diagnostics._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.

Nonlinear Output NetCDF Facade

Nonlinear NetCDF output schema writer.

spectraxgk.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.

spectraxgk.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.

Runtime Artifact Nonlinear Diagnostics

Nonlinear runtime diagnostic loading and restart-path helpers.

class spectraxgk.artifacts.nonlinear_diagnostics.DiagnosticTimeSeries(t: ndarray, values: ndarray, variable: str, source_path: str)

Single benchmark-facing time series loaded from an out.nc artifact.

spectraxgk.artifacts.nonlinear_diagnostics.load_diagnostic_time_series(path: str | Path, *, variable: str, diagnostics_group: str = 'Diagnostics', time_group: str = 'Grids', time_var: str = 'time', kx_index: int | None = None, component: str = 'real', align_phase: bool = False) → DiagnosticTimeSeries

Load a 1D diagnostics time series from a grouped NetCDF output artifact.

Quasilinear Transport Facade

Public quasilinear transport facade.

The implementation lives in spectraxgk.diagnostics.quasilinear_transport with the diagnostic owners. This module remains the stable public import path used by examples and downstream scripts.

class spectraxgk.quasilinear.QuasilinearTransportResult(ky: float, gamma: float, omega: float, mode: str, saturation_rule: str, amplitude_normalization: str, channels: tuple[str, ...], kperp_average: str, kperp_eff2: float, phi_norm2: float, amplitude2: float | None, heat_flux_weight_species: tuple[float, ...], particle_flux_weight_species: tuple[float, ...], saturated_heat_flux_species: tuple[float, ...] | None, saturated_particle_flux_species: tuple[float, ...] | None, species: tuple[str, ...], metadata: dict[str, Any])

JSON-friendly quasilinear diagnostic payload for one linear mode.

to_dict() → dict[str, Any]

Return a stable JSON-serializable representation.

spectraxgk.quasilinear.compute_quasilinear_from_linear_state(state: Array | ndarray, *, cache: LinearCache, grid: SpectralGrid, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData, params: LinearParams, ky: float, gamma: float, omega: float, terms: TermConfig | None = None, mode: str = 'weights', saturation_rule: str = 'none', amplitude_normalization: str = 'phi_rms', kperp_average: str = 'phi_weighted', csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False, channels: Sequence[str] | str = ('es',), species_names: Sequence[str] | None = None, use_dealias: bool = True, flux_scale: float = 1.0, metadata: dict[str, Any] | None = None) → QuasilinearTransportResult

Compute quasilinear transport weights from a linear state.

The returned heat and particle flux weights are divided by the selected mode-amplitude normalization, so they are invariant under complex phase rotations and real amplitude rescalings of the eigenstate.

spectraxgk.quasilinear.effective_kperp2(phi: Array, cache: LinearCache, vol_fac: Array, *, use_dealias: bool = True, eps: float = 1e-30) → Array

Compute <k_perp^2 |phi|^2>/<|phi|^2> for a linear mode.

spectraxgk.quasilinear.mixing_length_amplitude2_jax(gamma: Array | float, kperp_eff2_value: Array | float, *, csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False, eps: float = 1e-30) → Array

JAX-differentiable mixing-length squared-amplitude rule.

spectraxgk.quasilinear.normalize_quasilinear_channels(channels: Iterable[str] | str) → tuple[str, ...]

Normalize and validate quasilinear field channels.

spectraxgk.quasilinear.phi_norm2(phi: Array, cache: LinearCache, params: LinearParams, vol_fac: Array, *, normalization: str = 'phi_rms', use_dealias: bool = True, eps: float = 1e-30) → Array

Return the amplitude normalization used for quasilinear weights.

spectraxgk.quasilinear.quasilinear_feature_objective(features: Array | Sequence[float], *, rule: str = 'mixing_length', csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False) → Array

Differentiable objective from [gamma, kperp_eff2, flux_weight].

This helper is intentionally small: it is the reduced objective used by derivative validation tests and optimization examples once a linear scan has produced quasilinear weights.

spectraxgk.quasilinear.saturated_flux_from_linear_weight(linear_flux_weight: Array | float, gamma: Array | float, kperp_eff2_value: Array | float, *, csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False) → Array

Return a differentiable mixing-length saturated flux estimate.

spectraxgk.quasilinear.saturation_amplitude2(*, gamma: float, kperp_eff2_value: float, rule: str, csat: float = 1.0, gamma_floor: float = 0.0, include_stable_modes: bool = False) → float | None

Return the squared amplitude implied by a named saturation rule.

spectraxgk.quasilinear.shape_aware_power_law_objective(features: Array | Sequence[float], ky: Array | Sequence[float] | float, *, exponent: Array | float, csat: float = 1.0, ky_ref: float | None = None, eps: float = 1e-30) → Array

Differentiable shape-aware linear-weight objective.

features must end with [gamma, kperp_eff2, flux_weight]. The current low-dimensional shape model intentionally uses only the linear heat-flux weight and a power-law envelope in ky:

Q = C_sat * flux_weight * (ky / ky_ref)**exponent.

Growth-rate dependence is left to separately validated rules. This helper exists so the shape-aware saturation diagnostics and future optimization examples use one differentiable objective rather than plotting-only formulas.

spectraxgk.quasilinear.spectral_phi_weights(phi: Array, cache: LinearCache, vol_fac: Array, *, use_dealias: bool = True) → Array

Return (ky, kx, z) weights used for |phi|^2 averages.

Quasilinear 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 spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.diagnostics.quasilinear_transport.normalize_quasilinear_channels(channels: Iterable[str] | str) → tuple[str, ...]

Normalize and validate quasilinear field channels.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.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.

spectraxgk.diagnostics.quasilinear_transport.shape_aware_power_law_objective(features: Array | Sequence[float], ky: Array | Sequence[float] | float, *, exponent: Array | float, csat: float = 1.0, ky_ref: float | None = None, eps: float = 1e-30) → Array

Differentiable shape-aware linear-weight objective.

features must end with [gamma, kperp_eff2, flux_weight]. The current low-dimensional shape model intentionally uses only the linear heat-flux weight and a power-law envelope in ky:

Q = C_sat * flux_weight * (ky / ky_ref)**exponent.

Growth-rate dependence is left to separately validated rules. This helper exists so the shape-aware saturation diagnostics and future optimization examples use one differentiable objective rather than plotting-only formulas.

spectraxgk.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.

Quasilinear Calibration

Quasilinear calibration diagnostics and nonlinear-window ingestion.

These helpers compare linear/quasilinear flux estimates with accepted nonlinear transport windows. They are diagnostics and validation mathematics, not solver kernels: long campaign launch policy belongs in tools while public imports are re-exported through spectraxgk.api.

class spectraxgk.diagnostics.quasilinear_calibration.QuasilinearCalibrationPoint(case: str, split: str, predicted_heat_flux: float, observed_heat_flux: float, saturation_rule: str, raw_predicted_heat_flux: float | None = None, calibration_scale: float | None = None, geometry: str = 'unspecified', electron_model: str = 'unspecified', ky: float | None = None, observed_heat_flux_std: float | None = None, nonlinear_window_stats: dict[str, Any] | None = None, quasilinear_artifact: str | None = None, nonlinear_artifact: str | None = None, notes: str | None = None)

One quasilinear-vs-nonlinear transport comparison point.

spectraxgk.diagnostics.quasilinear_calibration.apply_heat_flux_scale(points: Iterable[QuasilinearCalibrationPoint | dict[str, Any]], *, scale: float, note_label: str = 'heat_flux_scale') → list[QuasilinearCalibrationPoint]

Return calibration points with heat-flux predictions multiplied by scale.

spectraxgk.diagnostics.quasilinear_calibration.calibration_point_from_nonlinear_window_summary(summary_json: str | Path, *, predicted_heat_flux: float, split: str, saturation_rule: str, diagnostics_source: str = 'spectrax', heat_flux_column: str = 'heat_flux', case: str | None = None, geometry: str = 'unspecified', electron_model: str = 'unspecified', quasilinear_artifact: str | None = None, species_index: int | None = None, window_convergence_config: NonlinearWindowConvergenceConfig | None = None, notes: str | None = None) → QuasilinearCalibrationPoint

Create a calibration point from a nonlinear window-summary JSON.

The helper reads the window bounds from a tracked nonlinear gate summary and computes the mean/std of a heat-flux column from the selected diagnostics CSV or runtime NetCDF. For NetCDF inputs, heat_flux_column='heat_flux' maps to Diagnostics/HeatFlux_st and species are summed by default.

spectraxgk.diagnostics.quasilinear_calibration.calibration_point_from_spectrum_and_nonlinear_window(spectrum_csv: str | Path, summary_json: str | Path, *, split: str, saturation_rule: str, spectrum_column: str = 'saturated_heat_flux_total', spectrum_method: str = 'sum', delta_ky: float | None = None, diagnostics_source: str = 'spectrax', heat_flux_column: str = 'heat_flux', case: str | None = None, geometry: str = 'unspecified', electron_model: str = 'unspecified', species_index: int | None = None, window_convergence_config: NonlinearWindowConvergenceConfig | None = None, notes: str | None = None) → QuasilinearCalibrationPoint

Create a calibration point from a quasilinear spectrum and nonlinear window.

spectraxgk.diagnostics.quasilinear_calibration.fit_train_heat_flux_scale(points: Iterable[QuasilinearCalibrationPoint | dict[str, Any]], *, train_split: str = 'train', prediction_floor: float = 1e-300) → dict[str, Any]

Fit one multiplicative heat-flux scale from training points.

The fit is a through-origin least-squares estimate, scale = sum(q_i Q_i) / sum(q_i^2), where q_i is the raw quasilinear heat-flux estimate and Q_i is the nonlinear window mean. This is the minimal calibration constant used by simple mixing-length models; any held-out failure after this fit is therefore a model failure, not a missing constant factor.

spectraxgk.diagnostics.quasilinear_calibration.integrated_quasilinear_flux_from_spectrum(spectrum_csv: str | Path, *, column: str = 'saturated_heat_flux_total', ky_column: str = 'ky', method: str = 'sum', delta_ky: float | None = None) → dict[str, Any]

Integrate one quasilinear spectrum column into a scalar flux estimate.

method="sum" preserves the discrete spectral-sum convention used by most runtime diagnostics. method="trapezoid" is available for smooth scan studies where the CSV is treated as a sampled function of ky.

spectraxgk.diagnostics.quasilinear_calibration.quasilinear_calibration_report(points: Iterable[QuasilinearCalibrationPoint | dict[str, Any]], *, saturation_rule: str, version: str = '0.1', holdout_mean_rel_gate: float = 0.35, observed_floor: float = 1e-12, fit_train_scale: bool = False, metadata: dict[str, Any] | None = None) → dict[str, Any]

Build a JSON-friendly calibration/holdout report.

A report is considered a calibrated absolute-flux claim only when it has at least one training point, at least one holdout point, and the holdout mean relative error passes the supplied gate.

spectraxgk.diagnostics.quasilinear_calibration.write_quasilinear_calibration_report(path: str | Path, report: dict[str, Any]) → Path

Write a quasilinear calibration report to JSON.

Quasilinear Nonlinear-Window Gates

Late-window transport diagnostics and promotion gates.

This module owns reusable nonlinear transport-window statistics that support quasilinear calibration, nonlinear transport acceptance, and release gates. It is intentionally diagnostics-focused: long-run campaign launch policy belongs in tools and tests, while solver kernels remain in operators and solvers.

class spectraxgk.diagnostics.transport_windows.NonlinearWindowConvergenceConfig(tmin: float | None = None, tmax: float | None = None, transient_fraction: float = 0.5, min_samples: int = 24, min_blocks: int = 4, block_size: int | None = None, bootstrap_samples: int = 256, bootstrap_seed: int = 0, max_running_mean_rel_drift: float = 0.15, terminal_fraction: float = 0.25, min_terminal_samples: int = 8, max_terminal_mean_rel_delta: float = 0.1, max_sem_rel: float = 0.25, value_floor: float = 1e-12, require_all_finite: bool = True)

Gate settings for a nonlinear post-transient transport window.

class spectraxgk.diagnostics.transport_windows.NonlinearWindowEnsembleConfig(min_reports: int = 2, max_mean_rel_spread: float = 0.15, max_combined_sem_rel: float = 0.25, value_floor: float = 1e-12, require_individual_passed: bool = True)

Gate settings for replicated nonlinear transport-window summaries.

spectraxgk.diagnostics.transport_windows.nonlinear_window_convergence_from_csv(csv_path: str | Path, *, time_column: str = 't', value_column: str = 'heat_flux', case: str | None = None, config: NonlinearWindowConvergenceConfig | None = None, summary_artifact: str | None = None) → dict[str, Any]

Build a convergence report from a diagnostics CSV.

spectraxgk.diagnostics.transport_windows.nonlinear_window_convergence_from_summary(summary_json: str | Path, *, diagnostics_source: str = 'spectrax', time_column: str = 't', value_column: str = 'heat_flux', case: str | None = None, config: NonlinearWindowConvergenceConfig | None = None) → dict[str, Any]

Build a convergence report from a window summary and diagnostics CSV.

spectraxgk.diagnostics.transport_windows.nonlinear_window_convergence_report(time: Sequence[float] | ndarray, values: Sequence[float] | ndarray, *, case: str = 'nonlinear_window', observable: str = 'heat_flux', source_artifact: str | None = None, summary_artifact: str | None = None, config: NonlinearWindowConvergenceConfig | None = None) → dict[str, Any]

Return finite late-window statistics and convergence gates.

The running-mean drift compares the mean of the first and second halves of the late window, normalized by the late-window mean scale. The uncertainty gate uses the maximum of the sample SEM, contiguous-block SEM, and block-bootstrap SEM when available.

spectraxgk.diagnostics.transport_windows.nonlinear_window_ensemble_report(reports: Sequence[dict[str, Any]], *, case: str = 'nonlinear_window_ensemble', comparison: str = 'replicate_uncertainty', config: NonlinearWindowEnsembleConfig | None = None) → dict[str, Any]

Gate repeated nonlinear-window summaries for seed/timestep robustness.

The input reports are expected to come from nonlinear_window_convergence_report(). This helper does not inspect raw time traces; it compares already-gated late-window means and their uncertainty metadata so production promotion can require seed, initial condition, or timestep robustness without rerunning simulations inside the checker.

spectraxgk.diagnostics.transport_windows.nonlinear_window_stats_promotion_ready(stats: object) → tuple[bool, list[str]]

Return whether serialized nonlinear window metadata can support promotion.

Quasilinear Model Selection

Quasilinear model-selection claim-boundary diagnostics.

The functions here combine dataset sufficiency, candidate-skill, calibration, and scoped optimized-equilibrium evidence into a JSON-ready model-selection status. They are diagnostics/reporting policy, not solver kernels.

spectraxgk.diagnostics.quasilinear_model_selection.build_quasilinear_model_selection_status(*, dataset_sufficiency: dict[str, Any] | str | Path, candidate_uncertainty: dict[str, Any] | str | Path, calibration_reports: Iterable[dict[str, Any] | str | Path] = (), optimized_equilibrium_nonlinear_audits: Iterable[dict[str, Any] | str | Path] = (), required_candidate: str = 'spectral_envelope_ridge', transport_gate: float | None = None, interval_coverage_gate: float | None = None, require_optimized_equilibrium_nonlinear_audit: bool = False) → dict[str, Any]

Combine quasilinear model-selection gates into one claim ledger.

The status is intentionally narrower than an absolute-flux calibration report. It passes only when the dataset-volume gate and uncertainty gate support the selected reduced candidate while all simple train/holdout calibration reports remain unpromoted. This lets documentation state a positive model-selection result without implying a runtime absolute-flux predictor. Optional optimized-equilibrium nonlinear audit artifacts can strengthen the scoped evidence ledger, but they never promote a universal absolute-flux claim.

spectraxgk.diagnostics.quasilinear_model_selection.build_quasilinear_model_selection_status_from_paths(*, dataset_sufficiency: str | Path, candidate_uncertainty: str | Path, calibration_reports: Iterable[str | Path], optimized_equilibrium_nonlinear_audits: Iterable[str | Path] = (), required_candidate: str = 'spectral_envelope_ridge', require_optimized_equilibrium_nonlinear_audit: bool = False) → dict[str, Any]

Path-based wrapper for artifact scripts and CI checks.

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 left/right eigenpair VJP for locally isolated dominant-growth branches, plus finite-difference diagnostics that verify the selected branch is consistent before using the derivative in optimization gates.

spectraxgk.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.

spectraxgk.objectives.eigen.dominant_real_eigenvalue(matrix: Array) → Array

Return the dominant growth rate with an implicit left/right VJP.

This helper treats the max-real eigenvalue branch selected at the primal point as locally isolated. Its reverse rule uses d lambda = w^H dA v with w^H v = 1 instead of differentiating through non-Hermitian eigenvectors. Branch isolation is still a physics gate: callers that use this in optimization should keep finite-difference or branch-continuity checks enabled near accepted candidates.

Solver Objective Core

Core linear and quasilinear solver-objective evaluators.

spectraxgk.objectives.core.solver_growth_rate_from_geometry(geom: Any, *, 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

Evaluate the dominant linear growth rate without eigenvector AD.

spectraxgk.objectives.core.solver_linear_operator_matrix_from_geometry(geom: Any, *, 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.

spectraxgk.objectives.core.solver_objective_vector_from_geometry(geom: Any, *, 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

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.

Branch continuity and AD/finite-difference validation are deliberately handled by gate functions; this is the shared forward evaluator those production objectives can use after a branch has been selected and audited.

spectraxgk.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.

Solver Objective Sampling

Sampling-axis and weighting helpers for solver-objective gates.

spectraxgk.objectives.sampling.solver_grid_options_from_ky_values(ky_values: float | tuple[float, ...] | list[float], *, ky_base: float | None = None, min_ny: int = 4) → dict[str, object]

Return solver grid options for physical k_y rho_i scan values.

The linear objective evaluator selects FFT row indices, while user-facing optimization and validation studies should be specified in physical k_y rho_i. This helper fixes that contract explicitly: values must be positive integer multiples of the base spacing, Ly = 2*pi/ky_base, and Ny is chosen large enough that all requested modes are represented as positive FFT rows.

Solver Geometry Objectives

Solver-ready geometry objective gates used by differentiability tests.

spectraxgk.objectives.geometry.default_solver_geometry_design_params() → Array

Return the small geometry-design vector used by the release gate.

spectraxgk.objectives.geometry.solver_ready_geometry_mapping(params: Array, theta: Array) → dict[str, Any]

Map a two-parameter design vector into solver-ready flux-tube arrays.

spectraxgk.objectives.geometry.tiny_differentiable_objective_gradient_report(params: Array | ndarray | None = None, *, fd_step: float = 0.0001, rtol: float = 0.0002, atol: float = 2e-06) → dict[str, object]

Validate a tiny differentiable objective on the solver-ready geometry map.

This is a lightweight objective-observable gate for CI and documentation. It checks the reusable AD/finite-difference report path without running the linear eigensolver or optional VMEC/Boozer backends.

Solver-Ready Gradient Gates

Solver-ready geometry-gradient gates for linear solver objectives.

spectraxgk.objectives.gradient_gates.linear_solver_geometry_gradient_report(params: Array | ndarray | None = None, *, fd_step: float = 0.001, rtol: float = 0.1, atol: float = 0.002, gap_floor: float = 1e-06) → dict[str, object]

Validate solver-objective geometry gradients on the actual linear RHS.

The report differentiates a small electrostatic Cyclone-like linear operator with respect to geometry arrays entering the production cache. It uses implicit left/right eigenpair sensitivities and compares them with nearest-branch central finite differences.

spectraxgk.objectives.gradient_gates.solver_objective_branch_gradient_report(params: Array | ndarray | None = None, *, fd_step: float = 0.001, rtol: float = 0.1, atol: float = 0.002, gap_floor: float = 1e-06, n_laguerre: int = 2, n_hermite: int = 1, _quasilinear_features_fn: Any = <function _linear_eigenpair_quasilinear_features>, _objective_vector_fn: Any = <function solver_objective_vector_from_geometry>) → dict[str, object]

Validate branch continuity and AD/FD sensitivities for solver objectives.

This gate is the lightweight counterpart of the VMEC/Boozer offline gates. It uses the solver-ready differentiable geometry contract so CI can check the objective path without optional geometry backends. The report requires the max-growth branch to stay dominant under central finite-difference perturbations and validates the objective sensitivities with the implicit left/right eigenpair method.

Solver VMEC/Boozer Gradient Gates

VMEC/Boozer full-chain gradient gates for solver objectives.

spectraxgk.objectives.vmec_boozer_gradients._mode21_vmec_boozer_linear_context(*, case_name: str, radial_index: int | None, mode_index: int, parameter_family: str, surface_index: int | None, ntheta: int, mboz: int, nboz: int, surface_stencil_width: int | None, n_laguerre: int, n_hermite: int) → dict[str, Any]

Build shared VMEC/Boozer geometry and linear-RHS closures for gates.

spectraxgk.objectives.vmec_boozer_gradients._mode21_vmec_boozer_quasilinear_features(eigenvalue: Array, eigenvector: Array, x: Array, context: dict[str, Any]) → tuple[Array, Array, Array, Array, Array]

Evaluate quasilinear observables from one isolated linear eigenpair.

spectraxgk.objectives.vmec_boozer_gradients.mode21_vmec_boozer_linear_frequency_gradient_report(*, case_name: str = 'nfp4_QH_warm_start', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', surface_index: int | None = None, fd_step: float = 1e-06, rtol: float = 0.05, atol: float = 0.02, gap_floor: float = 1e-08, ntheta: int = 4, mboz: int = 21, nboz: int = 21, surface_stencil_width: int | None = None, _linear_context_fn: Any = <function _mode21_vmec_boozer_linear_context>) → dict[str, object]

Validate a full VMEC/Boozer-state gradient of linear frequency.

This is an offline manuscript artifact gate. It perturbs one mid-surface VMEC Fourier coefficient, maps it through vmec_jax and booz_xform_jax into the mode-21 equal-arc flux-tube geometry contract, builds the SPECTRAX-GK linear RHS, and compares implicit eigenpair sensitivities against central finite differences. Quasilinear flux-weight state gradients are intentionally not promoted here because the current full-chain diagnostic is substantially heavier and remains an optimization campaign lane.

spectraxgk.objectives.vmec_boozer_gradients.mode21_vmec_boozer_nonlinear_window_gradient_report(*, case_name: str = 'nfp4_QH_warm_start', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', surface_index: int | None = None, fd_step: float = 1e-06, rtol: float = 0.075, atol: float = 0.05, gap_floor: float = 1e-08, ntheta: int = 4, mboz: int = 21, nboz: int = 21, surface_stencil_width: int | None = None, nonlinear_dt: float = 0.18, nonlinear_steps: int = 96, tail_fraction: float = 0.3, _linear_context_fn: Any = <function _mode21_vmec_boozer_linear_context>, _quasilinear_features_fn: Any = <function _mode21_vmec_boozer_quasilinear_features>, _window_metrics_fn: Any = <function _reduced_nonlinear_window_metrics_from_linear_observables>) → dict[str, object]

Validate VMEC/Boozer-state gradients of a nonlinear-window estimator.

The gate reuses the full vmec_jax state to booz_xform_jax to SPECTRAX-GK linear-RHS path from the quasilinear gradient gate, then feeds the isolated eigenpair observables into a differentiable late-time heat-flux-envelope estimator. It is a reduced nonlinear-window differentiability gate; converged nonlinear turbulence windows and optimized-equilibrium nonlinear audits remain separate promotion gates.

spectraxgk.objectives.vmec_boozer_gradients.mode21_vmec_boozer_quasilinear_gradient_report(*, case_name: str = 'nfp4_QH_warm_start', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', surface_index: int | None = None, fd_step: float = 1e-06, rtol: float = 0.02, atol: float = 0.05, gap_floor: float = 1e-08, ntheta: int = 4, mboz: int = 21, nboz: int = 21, surface_stencil_width: int | None = None, _linear_context_fn: Any = <function _mode21_vmec_boozer_linear_context>, _quasilinear_features_fn: Any = <function _mode21_vmec_boozer_quasilinear_features>) → dict[str, object]

Validate full VMEC/Boozer-state gradients of quasilinear observables.

This offline manuscript gate is the production-gradient companion to mode21_vmec_boozer_linear_frequency_gradient_report(). It uses a richer Nl=2, Nm=3 moment basis so the electrostatic heat-flux weight is nonzero, then validates implicit eigenpair sensitivities of gamma, omega, <k_perp^2>, the linear heat-flux weight, and the mixing-length heat-flux proxy against central finite differences.

Solver VMEC/Boozer Objectives

VMEC/Boozer objective-table plumbing for differentiable solver gates.

spectraxgk.objectives.vmec_boozer._split_vmec_boozer_objective_kwargs(kwargs: dict[str, Any]) → tuple[dict[str, Any], dict[str, Any]]

Split bridge options from solver-objective options and reject typos.

spectraxgk.objectives.vmec_boozer.vmec_boozer_aggregate_scalar_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, *, objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', weights: tuple[float, ...] | list[float] | ndarray | None = None, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), table_fn: Callable[[...], Array] | None = None, **kwargs: Any) → Array

Reduce a VMEC/Boozer multi-point objective table to one scalar.

spectraxgk.objectives.vmec_boozer.vmec_boozer_scalar_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, *, objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', vector_fn: Callable[[...], Array] | None = None, **kwargs: Any) → Array

Evaluate one scalar optimization objective on the VMEC/Boozer path.

spectraxgk.objectives.vmec_boozer.vmec_boozer_solver_objective_table_from_state(state: Any, static: Any, indata: Any, wout: Any, *, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), torflux_values: float | tuple[float, ...] | list[float] | None = None, alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), ky_values: float | tuple[float, ...] | list[float] | None = None, ky_base: float | None = None, table_with_metadata_fn: Callable[[...], tuple[Array, list[dict[str, object]]]] | None = None, **kwargs: Any) → Array

Evaluate solver objectives over a surface/field-line/k_y table.

spectraxgk.objectives.vmec_boozer.vmec_boozer_solver_objective_table_with_metadata_from_state(state: Any, static: Any, indata: Any, wout: Any, *, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), torflux_values: float | tuple[float, ...] | list[float] | None = None, alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), ky_values: float | tuple[float, ...] | list[float] | None = None, ky_base: float | None = None, geometry_fn: Callable[[...], Any] | None = None, objective_vector_fn: Callable[[...], Array] | None = None, **kwargs: Any) → tuple[Array, list[dict[str, object]]]

Evaluate VMEC/Boozer objective rows and return sample metadata.

spectraxgk.objectives.vmec_boozer.vmec_boozer_solver_objective_vector_from_state(state: Any, static: Any, indata: Any, wout: Any, *, geometry_fn: Callable[[...], Any] | None = None, objective_vector_fn: Callable[[...], Array] | None = None, **kwargs: Any) → Array

Evaluate solver objectives from the in-memory VMEC/Boozer bridge.

Solver VMEC/Boozer Finite-Difference Gates

Finite-difference gates for VMEC/Boozer objectives.

spectraxgk.objectives.vmec_boozer_fd._load_vmec_jax_example_state_bundle(case_name: str) → dict[str, Any]

Load a local vmec_jax example state bundle for offline gates.

spectraxgk.objectives.vmec_boozer_fd._report_float(report: dict[str, object], key: str) → float

Read a numeric finite-difference report field with mypy-safe casting.

spectraxgk.objectives.vmec_boozer_fd.vmec_boozer_aggregate_scalar_objective_finite_difference_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', weights: tuple[float, ...] | list[float] | ndarray | None = None, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), torflux_values: float | tuple[float, ...] | list[float] | None = None, alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), ky_values: float | tuple[float, ...] | list[float] | None = None, ky_base: float | None = None, radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', base_delta: float = 0.0, perturbation_step: float = 1e-07, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Finite-difference a multi-surface/multi-k_y aggregate objective.

spectraxgk.objectives.vmec_boozer_fd.vmec_boozer_scalar_objective_finite_difference_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', base_delta: float = 0.0, perturbation_step: float = 1e-07, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Finite-difference a scalar objective through a VMEC state coefficient.

This report is the safe optimization pre-step for full-chain stellarator objectives. It perturbs one VMEC state coefficient in a solved vmec_jax state, evaluates the in-memory VMEC/Boozer/SPECTRAX-GK scalar objective at x0+base_delta-h, x0+base_delta, and x0+base_delta+h, and records the central finite-difference sensitivity.

It is intentionally not an AD claim. Growth-rate objectives can later be promoted with implicit eigenpair gates; quasilinear objectives involving eigenvectors need this finite-difference/SPSA path or a custom adjoint before they are used in production optimization loops.

Solver VMEC/Boozer Line-Search Gates

Line-search and holdout gates for VMEC/Boozer objectives.

spectraxgk.objectives.vmec_boozer_line_search.vmec_boozer_aggregate_line_search_holdout_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', training_weights: tuple[float, ...] | list[float] | ndarray | None = None, holdout_weights: tuple[float, ...] | list[float] | ndarray | None = None, training_surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), training_alphas: float | tuple[float, ...] | list[float] = (0.0,), training_selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), holdout_surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), holdout_alphas: float | tuple[float, ...] | list[float] = (0.0,), holdout_selected_ky_indices: int | tuple[int, ...] | list[int] = (2,), radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', initial_delta: float = 0.0, perturbation_step: float = 1e-07, update_step: float = 1e-08, max_steps: int = 3, min_improvement: float = 0.0, min_holdout_improvement: float = 0.0, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Audit a training aggregate update against held-out aggregate samples.

A report passes only when the training line search accepts at least one curvature-gated update and the same final VMEC coefficient offset reduces the held-out aggregate objective. This is a reduced linear/quasilinear validation split, not a nonlinear transport optimization claim.

spectraxgk.objectives.vmec_boozer_line_search.vmec_boozer_aggregate_scalar_objective_line_search_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', reduction: Literal['mean', 'weighted_mean', 'max'] = 'mean', weights: tuple[float, ...] | list[float] | ndarray | None = None, surface_indices: int | None | tuple[int | None, ...] | list[int | None] = (None,), alphas: float | tuple[float, ...] | list[float] = (0.0,), selected_ky_indices: int | tuple[int, ...] | list[int] = (1,), radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', initial_delta: float = 0.0, perturbation_step: float = 1e-07, update_step: float = 1e-08, max_steps: int = 3, min_improvement: float = 0.0, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Run a curvature-gated line search for an aggregate VMEC objective.

This is the first optimizer-control gate for multi-surface, field-line, or k_y reduced objectives. It keeps the update one-dimensional so each step can be audited against the finite-difference curvature gate and rejected when the aggregate objective does not decrease.

spectraxgk.objectives.vmec_boozer_line_search.vmec_boozer_scalar_objective_line_search_report(*, case_name: str = 'nfp4_QH_warm_start', objective: Literal['growth', 'gamma', 'frequency', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'quasilinear_flux', 'mixing_length_heat_flux_proxy'] = 'growth', radial_index: int | None = None, mode_index: int = 1, parameter_family: str = 'Rcos', initial_delta: float = 0.0, perturbation_step: float = 1e-07, update_step: float = 1e-08, max_steps: int = 3, min_improvement: float = 0.0, response_atol: float = 0.0, max_curvature_ratio: float = 5.0, **kwargs: Any) → dict[str, object]

Run a curvature-gated one-parameter VMEC/Boozer objective line search.

This is the first safe optimizer scaffold for the real in-memory VMEC/Boozer/SPECTRAX-GK path. Each accepted update must pass the scalar finite-difference curvature gate, and candidate steps are accepted only when the scalar objective decreases. It is still a one-coefficient audit, not a broad stellarator-optimization claim.

Parallel Decomposition Contracts

Deterministic decomposition contracts for parallel work portfolios.

The helpers in this module describe partitioning and reconstruction contracts. They do not route solver execution, alter nonlinear state layout, or make speedup claims.

class spectraxgk.parallel.decomposition.DecompositionContract(workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble', 'diagnostic_nonlinear_domain'], claim_level: Literal['production_independent_batching', 'diagnostic_nonlinear_domain_partition'], claim_label: str, n_items: int, requested_shards: int, actual_shards: int, shards: tuple[ShardAssignment, ...], independent_work: bool, changes_solver_layout: bool, state_shape: tuple[int, ...] | None = None, axis: int | None = None)

Claim-scoped shard assignment contract for a parallelization path.

property diagnostic_nonlinear_partition: bool

Whether this contract is diagnostic nonlinear-domain metadata.

property production_independent_batching: bool

Whether this contract is for production independent-work batching.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the contract.

class spectraxgk.parallel.decomposition.ReconstructionIdentityReport(workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble', 'diagnostic_nonlinear_domain'], claim_level: Literal['production_independent_batching', 'diagnostic_nonlinear_domain_partition'], claim_label: str, n_items: int, requested_shards: int, actual_shards: int, identity_passed: bool, expected_indices: tuple[int, ...], reconstructed_indices: tuple[int, ...], missing_indices: tuple[int, ...], duplicate_indices: tuple[int, ...], out_of_range_indices: tuple[int, ...], out_of_order: bool)

Serial reconstruction identity report for a decomposition contract.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the report.

class spectraxgk.parallel.decomposition.ShardAssignment(shard_id: int, start: int, stop: int, indices: tuple[int, ...], label: str)

A deterministic contiguous assignment of serial indices to one shard.

property size: int

Number of serial items assigned to this shard.

to_dict() → dict[str, Any]

Return a JSON-friendly representation of the assignment.

spectraxgk.parallel.decomposition.build_diagnostic_nonlinear_domain_decomposition(state_shape: Iterable[int], *, axis: int, requested_shards: int) → DecompositionContract

Build a diagnostic nonlinear-domain partition contract.

This metadata describes split/reassemble coverage along one state axis. It is intentionally not a production nonlinear route and does not claim nonlinear speedup.

spectraxgk.parallel.decomposition.build_independent_portfolio_decomposition(n_items: int, *, requested_shards: int, workload: Literal['independent_ky_scan', 'uq_ensemble', 'optimization_ensemble']) → DecompositionContract

Build a production independent-work decomposition contract.

The assignment is deterministic, balanced, contiguous, and contains no empty shards. It covers release-ready independent portfolios only: independent_ky_scan, uq_ensemble, and optimization_ensemble.

spectraxgk.parallel.decomposition.reconstruct_serial(contract: DecompositionContract, shard_values: Sequence[Sequence[T]]) → tuple[T, ...]

Reassemble shard values into serial index order.

spectraxgk.parallel.decomposition.serial_reconstruction_identity_report(values: Sequence[T], contract: DecompositionContract, *, equal: Callable[[T, T], bool] | None = None) → ReconstructionIdentityReport

Check that contract sharding reassembles exactly to serial order.

spectraxgk.parallel.decomposition.shard_sequence(values: Sequence[T], contract: DecompositionContract) → tuple[tuple[T, ...], ...]

Return values grouped according to a decomposition contract.

QA Low-Turbulence Model And Gates

Reduced differentiable QA/ITG contracts, model, and objective gates.

class spectraxgk.objectives.qa_low_turbulence_model.QALowTurbulenceConfig(target_aspect: float = 6.0, min_iota: float = 0.41, iota_operating_floor: float = 0.7, max_mode: int = 1, aspect_weight: float = 8.0, iota_floor_weight: float = 160.0, iota_operating_weight: float = 70.0, qa_weight: float = 8.0, target_helical_amplitude: float = 0.16, helical_shaping_weight: float = 24.0, regularization: float = 0.002, nonlinear_weight: float = 8.0, learning_rate: float = 0.03, steps: int = 60, nonlinear_dt: float = 0.2, nonlinear_steps: int = 2000, nonlinear_tail_fraction: float = 0.5, long_window_min_time: float = 300.0, long_window_max_cv: float = 0.03, long_window_max_trend: float = 0.02, long_window_max_half_mean_rel_change: float = 0.02, fixed_density_gradient: float = 2.2, fixed_temperature_gradient: float = 6.0, scan_density_gradients: tuple[float, ...] = (0.6, 1.0, 1.4, 1.8, 2.2, 2.8, 3.4, 4.0, 4.8), fd_step: float = 0.0001, surface_ntheta: int = 72, surface_nzeta: int = 72, n_field_periods: int = 2)

Configuration for the reduced QA low-turbulence comparison.

class spectraxgk.objectives.qa_low_turbulence_model.QALowTurbulenceResult(design_name: str, includes_nonlinear_heat_flux: bool, parameter_names: tuple[str, ...], observable_names: tuple[str, ...], initial_params: tuple[float, ...], final_params: tuple[float, ...], initial_objective: float, final_objective: float, initial_observables: tuple[float, ...], final_observables: tuple[float, ...], history: tuple[dict[str, Any], ...], residual_gradient_gate: dict[str, Any], scalar_gradient_gate: dict[str, Any], observable_gradient_gate: dict[str, Any], covariance: dict[str, Any], config: dict[str, Any])

JSON-ready result for one reduced QA optimization.

to_dict() → dict[str, Any]

Return a stable JSON-friendly payload.

spectraxgk.objectives.qa_low_turbulence_model._qa_low_turbulence_core(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → dict[str, Array]

Return smooth reduced QA/ITG features for the aspect-6 comparison.

spectraxgk.objectives.qa_low_turbulence_model.default_qa_low_turbulence_initial_params() → Array

Return the shared off-optimum QA seed for the comparison.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_heat_flux_trace(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, density_gradient: float | None = None, temperature_gradient: float | None = None) → tuple[Array, Array]

Return a reduced nonlinear ITG heat-flux envelope for one gradient point.

density_gradient and temperature_gradient are normalized as a/L_n and a/L_T. The envelope is a fixed-step differentiable RK2 integration of dE/dt = 2 gamma E - alpha E^2 and Q_i = W_i E.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_objective(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, includes_nonlinear_heat_flux: bool) → Array

Return the scalar reduced QA comparison objective.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_observable_sensitivity_report(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, finite_difference_workers: int = 1) → dict[str, Any]

Compare the complete controls-to-observables Jacobian with finite differences.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_observable_vector(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → Array

Return QA low-turbulence observables in stable order.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_observables(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, density_gradient: float | None = None, temperature_gradient: float | None = None) → dict[str, Array]

Return reduced QA constraints and ITG observables.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_residual_names(*, includes_nonlinear_heat_flux: bool) → tuple[str, ...]

Return the stable residual names for the comparison objective.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_residual_vector(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None, *, includes_nonlinear_heat_flux: bool) → Array

Return weighted residuals for the aspect-6 QA low-turbulence objective.

spectraxgk.objectives.qa_low_turbulence_model.qa_low_turbulence_window_metrics(times: Array, heat_flux: Array, *, tail_fraction: float = 0.5, eps: float = 1e-12) → dict[str, Array]

Return differentiable late-window heat-flux statistics.

The standard deviation uses sqrt(var + eps) so the Jacobian remains finite when a long reduced trace has fully saturated and the late-window variance is numerically zero.

QA Low-Turbulence Optimizer

Optimizer loop for reduced QA low-turbulence objectives.

spectraxgk.objectives.qa_low_turbulence_optimizer.optimize_qa_low_turbulence(*, includes_nonlinear_heat_flux: bool, config: QALowTurbulenceConfig | None = None, initial_params: Array | Sequence[float] | None = None, finite_difference_workers: int = 1) → QALowTurbulenceResult

Optimize one reduced QA low-turbulence design with Adam.

QA Low-Turbulence Artifacts

JSON-ready payload builders for reduced QA low-turbulence comparisons.

spectraxgk.objectives.qa_low_turbulence_artifacts._long_window_convergence_gate(times: ndarray, heat_flux: ndarray, window: dict[str, Any], config: QALowTurbulenceConfig) → dict[str, Any]

Return explicit convergence checks for the fixed-gradient trace.

spectraxgk.objectives.qa_low_turbulence_artifacts.qa_low_turbulence_comparison_payload(config: QALowTurbulenceConfig | None = None, *, finite_difference_workers: int = 1) → dict[str, Any]

Build the full JSON-ready aspect-6 QA low-turbulence comparison.

spectraxgk.objectives.qa_low_turbulence_artifacts.reduced_boundary_surface(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → dict[str, Any]

Return a reduced max-mode-1 boundary surface for visualization.

spectraxgk.objectives.qa_low_turbulence_artifacts.reduced_lcfs_bmag(params: Array | Sequence[float], config: QALowTurbulenceConfig | None = None) → dict[str, Any]

Return a reduced LCFS |B| map for QA visualization.

VMEC-JAX Transport Objective

Evaluate SPECTRAX-GK transport objectives from VMEC-JAX states.

class spectraxgk.objectives.vmec_transport.VMECJAXSpectraxTransportObjective(config: VMECJAXTransportObjectiveConfig = <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 spectraxgk.objectives.vmec_transport.VMECJAXTransportObjectiveConfig(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_window_heat_flux']='nonlinear_window_heat_flux', sample_set: StellaratorITGSampleSet = <factory>, objective_weights: tuple[float, ...] | None=None, ntheta: int = 24, mboz: int = 21, nboz: int = 21, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, nonlinear_csat: float = 0.85, nonlinear_saturation_floor: float = 1e-10, reference_length: float | None = None, reference_b: float | None = None, objective_transform: Literal['raw', 'scaled', 'log1p']='raw', objective_scale: float = 1.0, surface_chunk_size: int = 0, validate_finite: bool = True)

Configuration for VMEC-JAX to SPECTRAX-GK objective evaluation.

property gradient_scope: str

Return the differentiated part of this objective.

objective_options() → dict[str, Any]

Return SPECTRAX-GK solver options for this objective.

spectraxgk.objectives.vmec_transport.vmec_jax_transport_growth_branch_locality_report_from_states(base_state: Any, plus_state: Any, minus_state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig | None = None, *, step: float, gap_floor: float = 1e-08, slope_rtol: float = 0.01, slope_atol: float = 1e-08, max_samples: int = 0) → dict[str, object]

Check dominant-growth eigenbranch locality for VMEC/Boozer samples.

The optimizer-facing transport residual can only use the implicit dominant-eigenvalue gradient when the same eigenbranch is locally selected. This report evaluates the exact SPECTRAX-GK linear operator matrix at the base, plus, and minus VMEC final states for each configured surface/alpha/k_y sample, then delegates branch classification to dominant_eigenvalue_branch_locality_report().

spectraxgk.objectives.vmec_transport.vmec_jax_transport_objective_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig | None = None) → Array

Evaluate a scalar SPECTRAX-GK transport objective from a VMEC-JAX state.

VMEC-JAX Transport Configuration

Configuration and optional-backend path policy for VMEC transport objectives.

class spectraxgk.objectives.vmec_transport_config.VMECJAXTransportObjectiveConfig(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_window_heat_flux']='nonlinear_window_heat_flux', sample_set: StellaratorITGSampleSet = <factory>, objective_weights: tuple[float, ...] | None=None, ntheta: int = 24, mboz: int = 21, nboz: int = 21, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, nonlinear_csat: float = 0.85, nonlinear_saturation_floor: float = 1e-10, reference_length: float | None = None, reference_b: float | None = None, objective_transform: Literal['raw', 'scaled', 'log1p']='raw', objective_scale: float = 1.0, surface_chunk_size: int = 0, validate_finite: bool = True)

Configuration for VMEC-JAX to SPECTRAX-GK objective evaluation.

property gradient_scope: str

Return the differentiated part of this objective.

objective_options() → dict[str, Any]

Return SPECTRAX-GK solver options for this objective.

VMEC-JAX Transport Tables

Sample-table and scalar-reduction kernels for VMEC transport objectives.

spectraxgk.objectives.vmec_transport_tables._apply_objective_transform(value: Array, config: VMECJAXTransportObjectiveConfig) → Array

Return a dimensionless transport residual with optional safe scaling.

spectraxgk.objectives.vmec_transport_tables._chunked_transport_objective_raw_value_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, cfg: VMECJAXTransportObjectiveConfig) → Array

Evaluate a weighted-mean raw objective one surface chunk at a time.

spectraxgk.objectives.vmec_transport_tables._geometry_transport_weights(geom: Any, *, selected_ky_index: int, ly: float) → tuple[Array, Array, Array]

Return trace-safe geometry weights for VMEC-JAX transport residuals.

spectraxgk.objectives.vmec_transport_tables._solver_table_to_nonlinear_window_proxy(table: Array, config: VMECJAXTransportObjectiveConfig) → Array

Map linear solver rows to a smooth reduced nonlinear heat-flux proxy.

spectraxgk.objectives.vmec_transport_tables._static_grid_options_from_ky_values(ky_values: Sequence[float], *, min_ny: int) → dict[str, Any]

Resolve physical ky samples without creating traced JAX arrays.

spectraxgk.objectives.vmec_transport_tables._surface_chunk_sample_sets(sample_set: StellaratorITGSampleSet, *, chunk_size: int) → tuple[tuple[StellaratorITGSampleSet, float], ...]

Split a sample set by surface while preserving weighted-mean algebra.

spectraxgk.objectives.vmec_transport_tables._transport_feature_table_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig, grid_options: dict[str, Any]) → Array

Return solver-objective rows with trace-safe growth-rate derivatives.

spectraxgk.objectives.vmec_transport_tables._transport_objective_raw_value_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, cfg: VMECJAXTransportObjectiveConfig) → Array

Evaluate the untransformed scalar transport objective.

spectraxgk.objectives.vmec_transport_tables.vmec_jax_transport_objective_from_state(state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig | None = None) → Array

Evaluate a scalar SPECTRAX-GK transport objective from a VMEC-JAX state.

VMEC-JAX Transport Branch Gates

Branch-locality gates for VMEC-JAX transport objectives.

spectraxgk.objectives.vmec_transport_branch.vmec_jax_transport_growth_branch_locality_report_from_states(base_state: Any, plus_state: Any, minus_state: Any, static: Any, indata: Any, wout_reference: Any, config: VMECJAXTransportObjectiveConfig | None = None, *, step: float, gap_floor: float = 1e-08, slope_rtol: float = 0.01, slope_atol: float = 1e-08, max_samples: int = 0) → dict[str, object]

Check dominant-growth eigenbranch locality for VMEC/Boozer samples.

The optimizer-facing transport residual can only use the implicit dominant-eigenvalue gradient when the same eigenbranch is locally selected. This report evaluates the exact SPECTRAX-GK linear operator matrix at the base, plus, and minus VMEC final states for each configured surface/alpha/k_y sample, then delegates branch classification to dominant_eigenvalue_branch_locality_report().

VMEC-JAX Transport Admission

VMEC-JAX transport candidate admission policies and metrics.

class spectraxgk.objectives.vmec_transport_admission.VMECJAXNonlinearAuditPolicy(minimum_relative_reduction: float = 0.02, minimum_uncertainty_z_score: float = 1.0, maximum_combined_sem_rel: float = 0.25, minimum_replicate_count: int = 3, minimum_surface_count: int = 3, minimum_alpha_count: int = 2, minimum_ky_count: int = 3, minimum_sample_count: int = 12, recommended_surfaces: tuple[float, ...] = (0.45, 0.64, 0.78), recommended_alphas: tuple[float, ...] = (0.0, 0.7853981633974483), recommended_ky_values: tuple[float, ...] = (0.1, 0.3, 0.5))

Policy for promoting or redesigning VMEC-JAX transport candidates.

Reduced growth/quasilinear/nonlinear-window objectives are useful only if they transfer to late-window nonlinear transport. This policy encodes the minimum replicated-audit evidence and sample coverage required before a candidate can be promoted beyond local reduced-metric admission.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.objectives.vmec_transport_admission.VMECJAXNonlinearCampaignPolicy(minimum_landscape_relative_reduction: float = 0.1, minimum_landscape_uncertainty_z_score: float = 3.0, maximum_landscape_sem_rel: float = 0.05, minimum_landscape_replicate_count: int = 3, require_reduced_prelaunch_passed: bool = True, require_reduced_cross_sample_gate: bool = True, require_landscape_admission_passed: bool = True)

Admission limits for launching the next nonlinear optimizer campaign.

This gate sits between a reduced candidate screen and a broader optimizer campaign. Passing it means the next campaign is worth launching; it does not promote a production nonlinear turbulent-flux optimization claim.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.objectives.vmec_transport_admission.VMECJAXReducedPrelaunchPolicy(metric_key: str = 'nonlinear_window_heat_flux', minimum_relative_reduction: float = 0.04, failed_reference_safety_factor: float = 1.5, require_sample_coverage: bool = True, maximum_cross_sample_sem_rel: float = 0.35)

Fail-closed reduced-objective gate before expensive nonlinear audits.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

class spectraxgk.objectives.vmec_transport_admission.VMECJAXTransportAdmissionPolicy(metric_keys: tuple[str, ...] = ('transport_objective_final', 'spectrax_objective_final', 'transport_metric_final', 'objective_final'), minimum_relative_improvement: float = 0.0, lower_is_better: bool = True, require_authoritative_gate: bool = True, allow_baseline_fallback: bool = True)

Fail-closed policy for selecting transport-aware VMEC candidates.

to_dict() → dict[str, Any]

Return a JSON-safe representation.

spectraxgk.objectives.vmec_transport_admission.build_transport_admission_report(summaries: Sequence[Mapping[str, Any]], *, policy: VMECJAXTransportAdmissionPolicy | None = None) → dict[str, Any]

Annotate and select VMEC-JAX transport candidates.

A transport candidate is admitted only when it passes the physical solved-WOUT gate and improves the selected transport metric relative to the admitted baseline. The baseline may be promoted only as a fallback audit target; it never counts as a transport-optimization success.

spectraxgk.objectives.vmec_transport_admission.candidate_transport_metric(candidate: Mapping[str, Any], *, metric_keys: Sequence[str] = ('transport_objective_final', 'spectrax_objective_final', 'transport_metric_final', 'objective_final')) → dict[str, Any]

Return the first finite transport metric found in a candidate summary.

spectraxgk.objectives.vmec_transport_admission.select_admitted_transport_candidate(summaries: Sequence[Mapping[str, Any]], *, policy: VMECJAXTransportAdmissionPolicy | None = None) → dict[str, Any] | None

Return the promoted candidate from build_transport_admission_report().

spectraxgk.objectives.vmec_transport_admission.transport_objective_sample_summary(sample_set: Any, *, policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Summarize whether a transport objective has enough sample coverage.

The nonlinear audit that motivated this gate was a single reduced metric: it improved locally but did not transfer to the replicated late-window heat-flux mean. Multi-surface, multi-field-line, and multi-k_y coverage is therefore treated as an admission requirement for the next candidate.

Stellarator Transport Reports

Stellarator nonlinear-transport admission and redesign reports.

spectraxgk.diagnostics.stellarator_transport_reports.build_nonlinear_audit_redesign_report(matched_comparison: Mapping[str, Any], *, objective_sample_set: Any = None, policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Decide whether a matched nonlinear audit promotes or redesigns a candidate.

This is the fail-closed bridge between reduced VMEC-JAX transport admission and expensive long-window nonlinear evidence. A candidate is promoted only if the matched replicated nonlinear comparison passes, has a positive uncertainty-separated reduction, and the reduced objective used enough surface/field-line/k_y samples to avoid a single-point overfit.

spectraxgk.diagnostics.stellarator_transport_reports.build_nonlinear_campaign_admission_report(*, reduced_prelaunch_report: Mapping[str, Any], landscape_admission_report: Mapping[str, Any], policy: VMECJAXNonlinearCampaignPolicy | None = None) → dict[str, Any]

Gate the next nonlinear optimizer campaign from existing evidence.

This report intentionally promotes only a campaign launch. It requires a reduced prelaunch pass and an uncertainty-separated replicated nonlinear landscape point. It does not convert that point into a general multi-coefficient turbulent-flux optimization result.

spectraxgk.diagnostics.stellarator_transport_reports.build_nonlinear_landscape_admission_report(baseline_ensemble: Mapping[str, Any], candidate_ensembles: Sequence[Mapping[str, Any]], *, candidate_labels: Sequence[str] | None = None, policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Select an uncertainty-separated nonlinear candidate from a landscape.

This gate is for boundary-coefficient or line-search landscapes where a small number of selected points have replicated late-window nonlinear ensembles. It does not validate multi-coefficient global optimization by itself; it only answers whether any supplied candidate has a statistically resolved lower heat flux than the supplied baseline ensemble.

spectraxgk.diagnostics.stellarator_transport_reports.build_reduced_nonlinear_audit_prelaunch_report(*, baseline_metric: float, candidate_metric: float, objective_sample_set: Any = None, baseline_sample_statistics: Mapping[str, Any] | None = None, candidate_sample_statistics: Mapping[str, Any] | None = None, failed_reference_relative_reduction: float | None = None, policy: VMECJAXReducedPrelaunchPolicy | None = None, nonlinear_policy: VMECJAXNonlinearAuditPolicy | None = None) → dict[str, Any]

Gate reduced transport candidates before launching nonlinear audits.

This is intentionally conservative. A reduced nonlinear-window improvement should exceed both an absolute release threshold and, when available, a safety factor above a known failed-transfer reference before spending another long GPU campaign.

VMEC-JAX Transport Gradient

Boundary-gradient diagnostics for VMEC-JAX transport objectives.

The helpers here intentionally avoid importing vmec_jax at module import time. They operate on the small optimizer protocol exposed by VMEC-JAX (residual_fun, objective_and_gradient_fun, and _specs), which keeps SPECTRAX-GK tests fast while letting examples evaluate the real full-chain transport sensitivity on machines with VMEC-JAX installed.

spectraxgk.objectives.vmec_transport_gradient.boundary_spec_record(spec: Any, *, fallback_index: int) → dict[str, Any]

Return a JSON-safe summary for a VMEC-JAX boundary parameter spec.

spectraxgk.objectives.vmec_transport_gradient.build_boundary_transport_gradient_report(optimizer: Any, *, params: Sequence[float] | ndarray | None = None, label: str = 'vmec_jax_transport_gradient', top_n: int = 12, sensitivity_atol: float = 1e-12, include_jacobian: bool = False) → dict[str, Any]

Evaluate transport residual and boundary-gradient diagnostics.

Parameters:

• optimizer – VMEC-JAX-like optimizer exposing residual_fun(params) and objective_and_gradient_fun(params). The optional _specs member is used only for readable boundary-coefficient labels.

• params – Active boundary-parameter vector. None means the zero-increment vector aligned with optimizer._specs.

• label – Human-readable artifact label.

• top_n – Number of largest gradient components to keep.

• sensitivity_atol – Absolute L2 threshold below which the boundary transport response is classified as locally flat for optimization purposes.

• include_jacobian – If true and the optimizer exposes jacobian_fun, include dense residual-Jacobian norms. This can be substantially more expensive than the reverse scalar-gradient path.

spectraxgk.objectives.vmec_transport_gradient.write_boundary_transport_gradient_report(report: dict[str, Any], path: str | Path) → Path

Write a boundary-gradient diagnostic JSON artifact.

VMEC-JAX Boundary Chain

Boundary-chain diagnostics for VMEC-JAX/SPECTRAX-GK gradients.

These helpers classify scalar contractions recorded by the historical boundary-chain artifacts. Those studies compared raw exact-solve finite differences, frozen-axis initial-state finite differences, and VMEC-JAX tape JVP/VJP contractions. Keeping the classification logic makes that paper-facing convention explicit and testable without retaining a tool tied to VMEC-JAX’s removed private optimizer API.

spectraxgk.geometry.vmec_boundary_chain.boundary_chain_summary_from_probe(payload: Mapping[str, Any], **kwargs: Any) → dict[str, Any]

Build a chain summary from a probe JSON payload.

spectraxgk.geometry.vmec_boundary_chain.build_boundary_chain_collection_summary(probes: Sequence[Mapping[str, Any]], *, exact_relative_tolerance: float = 0.1, internal_relative_tolerance: float = 1e-08, absolute_tolerance: float = 1e-10) → dict[str, Any]

Summarize several boundary-chain probes as one promotion gate.

A single coefficient can look well-conditioned while neighboring boundary modes still move the raw exact-solve initialization branch. The collection summary keeps the stricter manuscript/release decision explicit: frozen-axis JVP/VJP replay must be internally transposed for every component, while exact finite-difference agreement is counted separately from branch sensitivity.

spectraxgk.geometry.vmec_boundary_chain.build_boundary_chain_summary(*, exact_fd_cost_gradient: float, final_cot_dot_exact_final_fd: float, frozen_axis_replay_cost_gradient: float, frozen_axis_vjp_cost_gradient: float, frozen_axis_linear_replay_cost_gradient: float | None = None, frozen_axis_linear_vjp_cost_gradient: float | None = None, frozen_axis_initial_fd_vs_linear_abs_norm: float | None = None, frozen_axis_initial_fd_vs_linear_rel: float | None = None, raw_initial_replay_cost_gradient: float | None = None, raw_initial_fd_norm: float | None = None, frozen_axis_initial_fd_norm: float | None = None, exact_relative_tolerance: float = 0.1, internal_relative_tolerance: float = 1e-08, absolute_tolerance: float = 1e-10) → dict[str, Any]

Classify a boundary-gradient chain probe.

Parameters:

• exact_fd_cost_gradient – Central finite difference through plus/minus exact VMEC solves.

• final_cot_dot_exact_final_fd – Final-state SPECTRAX-GK cotangent dotted into the exact final-state finite-difference direction.

• frozen_axis_replay_cost_gradient – VMEC-JAX tape JVP contraction using the frozen-axis initial-state tangent used by the optimizer.

• frozen_axis_vjp_cost_gradient – VMEC-JAX tape VJP contraction projected back through the same frozen initial-state map.

• frozen_axis_linear_replay_cost_gradient – Optional contraction using VMEC-JAX’s explicit frozen-axis tangent column.

• frozen_axis_linear_vjp_cost_gradient – Optional VJP contraction using VMEC-JAX’s explicit frozen-axis tangent column.

• frozen_axis_initial_fd_vs_linear_abs_norm – Optional norm of the frozen-axis finite-difference tangent minus the explicit tangent column.

• frozen_axis_initial_fd_vs_linear_rel – Optional relative norm of the frozen-axis finite-difference tangent minus the explicit tangent column.

• raw_initial_replay_cost_gradient – Optional tape JVP contraction using raw plus/minus initial-state finite differences. This diagnoses magnetic-axis branch sensitivity, but it is not the optimizer’s advertised derivative.

VMEC-JAX Transport Line Search

Utilities for projected VMEC-JAX transport line searches.

These helpers convert a boundary-gradient diagnostic into a normalized descent direction and JSON-safe line-search manifests. The VMEC-JAX-dependent input writer lives in tools/; this module is deliberately backend-free so CI can test the admission bookkeeping without launching equilibrium solves.

class spectraxgk.objectives.vmec_transport_line_search.ProjectedLineSearchPolicy(minimum_relative_improvement: float = 0.0, lower_is_better: bool = True, require_gate_passed: bool = True)

Selection policy for projected transport line-search candidates.

to_dict() → dict[str, Any]

Return a JSON-safe policy representation.

spectraxgk.objectives.vmec_transport_line_search.boundary_chain_accepted_parameter_indices(collection: Mapping[str, Any], *, require_exact_fd: bool = True, require_growth_branch_locality: bool = False) → tuple[int, ...]

Return parameter indices admitted by a boundary-chain collection gate.

spectraxgk.objectives.vmec_transport_line_search.projected_line_search_input_manifest(report: Mapping[str, Any], *, steps: Sequence[float], top_n: int | None = None, boundary_chain_collection: Mapping[str, Any] | None = None, require_boundary_chain_exact_fd: bool = True, require_growth_branch_locality: bool = False) → dict[str, Any]

Build a JSON-safe manifest for projected line-search input generation.

spectraxgk.objectives.vmec_transport_line_search.select_projected_line_search_candidate(baseline: Mapping[str, Any], candidates: Sequence[Mapping[str, Any]], *, policy: ProjectedLineSearchPolicy | None = None) → dict[str, Any]

Select the best gate-passing projected line-search candidate.

spectraxgk.objectives.vmec_transport_line_search.sparse_descent_direction_from_gradient_report(report: Mapping[str, Any], *, parameter_count: int | None = None, top_n: int | None = None, boundary_chain_collection: Mapping[str, Any] | None = None, require_boundary_chain_exact_fd: bool = True, require_growth_branch_locality: bool = False) → ndarray

Return a normalized sparse descent direction from a gradient report.

The direction is -grad restricted to the ranked top_gradient_components. This convention makes positive line-search steps lower a lower-is-better transport objective to first order.

VMEC-JAX Candidate Gates

Acceptance gates for VMEC-JAX stellarator-optimization candidates.

spectraxgk.objectives.vmec_candidate_admission.build_authoritative_wout_candidate_gate(wout: str | Path | Mapping[str, Any], *, target_aspect: float, aspect_atol: float, min_abs_mean_iota: float, qs_residual_max: float, iota_profile_floor: float | None, helicity_m: int = 1, helicity_n: int = 0, qs_surfaces: tuple[float, ...] = (0.0, 0.1, 0.2, 0.30000000000000004, 0.4, 0.5, 0.6000000000000001, 0.7000000000000001, 0.8, 0.9, 1.0), qs_ntheta: int = 63, qs_nphi: int = 64) → dict[str, Any]

Build a solved-equilibrium gate directly from a WOUT artifact.

Use this when the deterministic replayed WOUT is the publication-facing equilibrium. It does not assert that the replayed WOUT matches an optimizer-state WOUT; that remains the role of build_wout_reproducibility_gate().

spectraxgk.objectives.vmec_candidate_admission.build_solved_vmec_candidate_gate(candidate: Any, *, target_aspect: float, aspect_atol: float, min_abs_mean_iota: float, qs_residual_max: float, iota_profile_floor: float | None, iota_profiles: tuple[ndarray, ndarray] | None = None, profile_source: str = 'provided') → dict[str, Any]

Build a JSON-safe solved-equilibrium gate report.

candidate may be a VMEC-JAX optimization result with a history property or a history mapping loaded from history.json.

spectraxgk.objectives.vmec_candidate_admission.build_wout_reproducibility_gate(reference_wout: str | Path | Mapping[str, Any], rerun_wout: str | Path | Mapping[str, Any], *, target_aspect: float, aspect_atol: float, min_abs_mean_iota: float, iota_profile_floor: float | None, mean_iota_repro_atol: float = 0.0005, aspect_repro_atol: float = 1e-06, profile_repro_atol: float = 0.0005) → dict[str, Any]

Check that a saved VMEC input reproduces the optimizer-state WOUT.

VMEC-JAX can write both an optimizer-state wout_final.nc and an input.final deck. For publication-facing transport claims, the deck must reproduce the WOUT when rerun; otherwise downstream SPECTRAX-GK metrics may be attached to a different equilibrium than the optimized state. This gate compares the original WOUT against a fresh rerun WOUT and also applies the solved-equilibrium aspect/iota/profile admission checks to the rerun.

spectraxgk.objectives.vmec_candidate_admission.final_iota_profiles_from_vmec_result(result: Any) → tuple[ndarray, ndarray] | None

Return final solved iota profiles from a VMEC-JAX result if available.

Stellarator ITG Objectives

Public facade and optimizer loop for reduced stellarator ITG objectives.

class spectraxgk.objectives.stellarator.StellaratorITGOptimizationConfig(target_aspect: float = 7.0, target_iota: float = 0.41, max_mode: int = 1, aspect_weight: float = 0.25, iota_weight: float = 25.0, qa_weight: float = 5.0, turbulence_weight: float = 1.0, regularization: float = 0.002, learning_rate: float = 0.035, steps: int = 90, nonlinear_dt: float = 0.18, nonlinear_steps: int = 520, nonlinear_tail_fraction: float = 0.25, quasilinear_csat: float = 0.75, reference_density_gradient: float = 2.2, reference_temperature_gradient: float = 6.0, scan_density_gradients: tuple[float, ...] = (0.8, 1.2, 1.6, 2.2, 3.0, 3.8, 4.8), fd_step: float = 0.0001)

Configuration for the QA max-mode-1 ITG optimization examples.

with_kind_defaults(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']) → StellaratorITGOptimizationConfig

Return conservative optimizer defaults for one objective family.

class spectraxgk.objectives.stellarator.StellaratorITGOptimizationResult(objective_kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], parameter_names: tuple[str, ...], observable_names: tuple[str, ...], initial_params: tuple[float, ...], final_params: tuple[float, ...], initial_objective: float, final_objective: float, initial_observables: tuple[float, ...], final_observables: tuple[float, ...], history: tuple[dict[str, Any], ...], gradient_gate: dict[str, Any], covariance: dict[str, Any], nonlinear_trace: dict[str, Any] | None, config: dict[str, Any], backend_info: dict[str, Any])

JSON-friendly result for one differentiable stellarator objective.

to_dict() → dict[str, Any]

Return a stable JSON-serializable representation.

class spectraxgk.objectives.stellarator.StellaratorITGSampleSet(surfaces: tuple[float, ...] = (0.5, 0.64, 0.78), alphas: tuple[float, ...] = (0.0, 1.0471975511965976), ky_values: tuple[float, ...] = (0.1, 0.3, 0.5), surface_weights: tuple[float, ...] | None = None, alpha_weights: tuple[float, ...] | None = None, ky_weights: tuple[float, ...] | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean')

Reduced multi-surface/multi-alpha/multi-k_y ITG portfolio contract.

property n_samples: int

Number of surface/alpha/ky samples in the rectangular portfolio.

to_dict() → dict[str, Any]

Return a JSON-friendly representation.

spectraxgk.objectives.stellarator.compare_stellarator_itg_objectives(kinds: Sequence[Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']] = ('growth', 'quasilinear_flux', 'nonlinear_heat_flux'), *, initial_params: Array | Sequence[float] | None = None, config: StellaratorITGOptimizationConfig | None = None, workers: int = 1, parallel_executor: str = 'thread', finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Run the three objective reductions from a shared starting point.

spectraxgk.objectives.stellarator.default_stellarator_initial_params() → Array

Return the shared off-optimum QA max-mode-1 starting point.

spectraxgk.objectives.stellarator.nonlinear_heat_flux_trace(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → tuple[Array, Array]

Return a differentiable short-window ITG heat-flux envelope trace.

The envelope evolves E with a fixed-step RK2 discretization,

dE/dt = 2 gamma E - alpha E^2, Q_env(t) = W_i E.

gamma and W_i come from the same differentiable QA/ITG feature map as the linear and quasilinear objectives. The output is therefore useful for nonlinear averaging, optimizer, and UQ gates while the full production nonlinear-GK geometry path is still being made traceable end-to-end.

spectraxgk.objectives.stellarator.nonlinear_heat_flux_window_metrics(times: Array, heat_flux: Array, *, tail_fraction: float = 0.45, eps: float = 1e-14) → dict[str, Array]

Return mean, coefficient of variation, and trend on a late-time window.

spectraxgk.objectives.stellarator.optimize_stellarator_itg(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], initial_params: Array | Sequence[float] | None = None, config: StellaratorITGOptimizationConfig | None = None, *, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → StellaratorITGOptimizationResult

Optimize one differentiable stellarator ITG objective with Adam.

spectraxgk.objectives.stellarator.qa_max_mode1_observables(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → dict[str, Array]

Map a QA max-mode-1 boundary/control vector to differentiable ITG observables.

The four inputs represent the active low-order controls used by the example scripts. The map is calibrated as a smooth objective-reduction gate around a QA stellarator with aspect ratio 7 and mean rotational transform 0.41. It is not a replacement for the full VMEC/Boozer flux-tube geometry contract; its purpose is to validate gradient plumbing, UQ, optimizer behavior, and figure-generation before expensive production objectives are promoted.

spectraxgk.objectives.stellarator.qa_observable_vector(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None) → Array

Return observables in the stable order defined by OBSERVABLE_NAMES.

spectraxgk.objectives.stellarator.smooth_positive(x: Array | float, *, beta: float = 18.0) → Array

Smooth positive part used to keep objectives differentiable near marginality.

spectraxgk.objectives.stellarator.stellarator_itg_density_gradient_scan(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradients: Sequence[float] | None = None, temperature_gradient: float | None = None) → dict[str, Any]

Return a reduced ITG response scan versus normalized density gradient.

The scan uses the same explicit drive inputs as the reduced growth, quasilinear, and nonlinear-envelope objectives. It is intended for candidate ranking and figure QA before promotion to solved VMEC/Boozer nonlinear SPECTRAX-GK scans.

spectraxgk.objectives.stellarator.stellarator_itg_objective(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None) → Array

Return the scalar constrained QA + ITG objective for one optimization.

spectraxgk.objectives.stellarator.stellarator_itg_objective_residual_names(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']) → tuple[str, ...]

Return stable residual names for the weighted QA + ITG objective.

spectraxgk.objectives.stellarator.stellarator_itg_objective_residual_vector(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None) → Array

Return the weighted residual map used by the objective and covariance.

spectraxgk.objectives.stellarator.stellarator_itg_portfolio_gate_payload(params: Array | Sequence[float] | None = None, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Return the JSON-ready reduced ITG portfolio gate artifact payload.

spectraxgk.objectives.stellarator.stellarator_itg_portfolio_sensitivity_report(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, step: float | None = None, rtol: float | None = None, atol: float | None = None, workers: int = 1, parallel_executor: str = 'thread') → dict[str, Any]

AD/FD, conditioning, and covariance gate for the reduced ITG portfolio.

spectraxgk.objectives.stellarator.stellarator_itg_reduced_portfolio_objective(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None) → Array

Reduce a sampled ITG growth/QL portfolio to one differentiable scalar.

spectraxgk.objectives.stellarator.stellarator_itg_residual_sensitivity_report(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None, *, step: float | None = None, rtol: float | None = None, atol: float | None = None, min_rank: int = 4, condition_number_limit: float = 10000.0, covariance_regularization: float = 1e-08, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Check residual-Jacobian AD/FD parity and local conditioning.

spectraxgk.objectives.stellarator.stellarator_itg_sample_objective_table(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None) → Array

Return (surface, alpha, ky, objective) reduced ITG objective rows.

This is the backend-free rehearsal of the production VMEC/Boozer sample table. It keeps the optimizer and gate semantics identical to the future real-geometry path while remaining cheap enough for CI.

spectraxgk.objectives.stellarator.stellarator_itg_vmec_boozer_portfolio_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, **vmec_boozer_options: Any) → Array

Facade wrapper for the real VMEC/Boozer/SPECTRAX-GK portfolio scalar.

spectraxgk.objectives.stellarator.stellarator_itg_vmec_boozer_sample_objective_table_from_state(state: Any, static: Any, indata: Any, wout: Any, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), sample_set: StellaratorITGSampleSet | None = None, **vmec_boozer_options: Any) → Array

Facade wrapper for real VMEC/Boozer/SPECTRAX-GK objective rows.

Stellarator ITG Objective Tables

Reduced and VMEC/Boozer ITG objective tables for stellarator studies.

spectraxgk.objectives.stellarator_tables.stellarator_itg_density_gradient_scan(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradients: Sequence[float] | None = None, temperature_gradient: float | None = None) → dict[str, Any]

Return a reduced ITG response scan versus normalized density gradient.

The scan uses the same explicit drive inputs as the reduced growth, quasilinear, and nonlinear-envelope objectives. It is intended for candidate ranking and figure QA before promotion to solved VMEC/Boozer nonlinear SPECTRAX-GK scans.

spectraxgk.objectives.stellarator_tables.stellarator_itg_portfolio_gate_payload(params: Array | Sequence[float] | None = None, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Return the JSON-ready reduced ITG portfolio gate artifact payload.

spectraxgk.objectives.stellarator_tables.stellarator_itg_portfolio_sensitivity_report(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, step: float | None = None, rtol: float | None = None, atol: float | None = None, workers: int = 1, parallel_executor: str = 'thread') → dict[str, Any]

AD/FD, conditioning, and covariance gate for the reduced ITG portfolio.

spectraxgk.objectives.stellarator_tables.stellarator_itg_reduced_portfolio_objective(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None) → Array

Reduce a sampled ITG growth/QL portfolio to one differentiable scalar.

spectraxgk.objectives.stellarator_tables.stellarator_itg_sample_objective_table(params: Array | Sequence[float], objectives: Sequence[str] = ('growth', 'quasilinear_flux'), config: StellaratorITGOptimizationConfig | None = None, sample_set: StellaratorITGSampleSet | None = None) → Array

Return (surface, alpha, ky, objective) reduced ITG objective rows.

This is the backend-free rehearsal of the production VMEC/Boozer sample table. It keeps the optimizer and gate semantics identical to the future real-geometry path while remaining cheap enough for CI.

spectraxgk.objectives.stellarator_tables.stellarator_itg_vmec_boozer_portfolio_objective_from_state(state: Any, static: Any, indata: Any, wout: Any, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), sample_set: StellaratorITGSampleSet | None = None, *, objective_weights: Sequence[float] | None = None, **vmec_boozer_options: Any) → Array

Reduce real VMEC/Boozer/SPECTRAX-GK ITG rows to one portfolio scalar.

spectraxgk.objectives.stellarator_tables.stellarator_itg_vmec_boozer_sample_objective_table_from_state(state: Any, static: Any, indata: Any, wout: Any, objectives: Sequence[str] = ('growth', 'quasilinear_flux'), sample_set: StellaratorITGSampleSet | None = None, **vmec_boozer_options: Any) → Array

Return real VMEC/Boozer/SPECTRAX-GK rows on a StellaratorITGSampleSet.

This is the production bridge counterpart to stellarator_itg_sample_objective_table(): the sample axes are physical toroidal-flux, field-line alpha, and k_y rho_i values, while the objective columns are selected from the solver objective vector.

Stellarator ITG Contracts

Contracts for reduced stellarator ITG optimization examples.

class spectraxgk.objectives.stellarator_contracts.StellaratorITGOptimizationConfig(target_aspect: float = 7.0, target_iota: float = 0.41, max_mode: int = 1, aspect_weight: float = 0.25, iota_weight: float = 25.0, qa_weight: float = 5.0, turbulence_weight: float = 1.0, regularization: float = 0.002, learning_rate: float = 0.035, steps: int = 90, nonlinear_dt: float = 0.18, nonlinear_steps: int = 520, nonlinear_tail_fraction: float = 0.25, quasilinear_csat: float = 0.75, reference_density_gradient: float = 2.2, reference_temperature_gradient: float = 6.0, scan_density_gradients: tuple[float, ...] = (0.8, 1.2, 1.6, 2.2, 3.0, 3.8, 4.8), fd_step: float = 0.0001)

Configuration for the QA max-mode-1 ITG optimization examples.

with_kind_defaults(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']) → StellaratorITGOptimizationConfig

Return conservative optimizer defaults for one objective family.

class spectraxgk.objectives.stellarator_contracts.StellaratorITGOptimizationResult(objective_kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], parameter_names: tuple[str, ...], observable_names: tuple[str, ...], initial_params: tuple[float, ...], final_params: tuple[float, ...], initial_objective: float, final_objective: float, initial_observables: tuple[float, ...], final_observables: tuple[float, ...], history: tuple[dict[str, Any], ...], gradient_gate: dict[str, Any], covariance: dict[str, Any], nonlinear_trace: dict[str, Any] | None, config: dict[str, Any], backend_info: dict[str, Any])

JSON-friendly result for one differentiable stellarator objective.

to_dict() → dict[str, Any]

Return a stable JSON-serializable representation.

class spectraxgk.objectives.stellarator_contracts.StellaratorITGSampleSet(surfaces: tuple[float, ...] = (0.5, 0.64, 0.78), alphas: tuple[float, ...] = (0.0, 1.0471975511965976), ky_values: tuple[float, ...] = (0.1, 0.3, 0.5), surface_weights: tuple[float, ...] | None = None, alpha_weights: tuple[float, ...] | None = None, ky_weights: tuple[float, ...] | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean')

Reduced multi-surface/multi-alpha/multi-k_y ITG portfolio contract.

property n_samples: int

Number of surface/alpha/ky samples in the rectangular portfolio.

to_dict() → dict[str, Any]

Return a JSON-friendly representation.

Stellarator Reduced ITG Model And Gates

Reduced QA stellarator ITG model, residuals, and sensitivity gates.

spectraxgk.objectives.stellarator_reduced._conditioning_gate_from_covariance(covariance: dict[str, Any], *, min_rank: int, condition_number_limit: float) → dict[str, Any]

Return a pass/fail gate for Gauss-Newton residual conditioning.

spectraxgk.objectives.stellarator_reduced._precision_gate_tolerances(fd_step: float) → tuple[float, float, float]

Return FD tolerances that are strict in x64 and stable in float32.

spectraxgk.objectives.stellarator_reduced._qa_core_features(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → dict[str, Array]

Return the linear/quasilinear QA-ITG features without nonlinear tracing.

spectraxgk.objectives.stellarator_reduced._qa_geometry_features(minor_shift: Array, elong_shift: Array, ripple: Array, shear_shift: Array, aspect_target: Array, iota_target: Array) → dict[str, Array]

Return reduced QA shape, iota, residual, and curvature metrics.

spectraxgk.objectives.stellarator_reduced._qa_gradient_drives(config: StellaratorITGOptimizationConfig, dtype: dtype, *, density_gradient: float | Array | None, temperature_gradient: float | Array | None) → tuple[Array, Array]

Return normalized density and temperature-gradient offsets.

spectraxgk.objectives.stellarator_reduced._qa_linear_itg_features(geometry: dict[str, Array], elong_shift: Array, ripple: Array, shear_shift: Array, iota_target: Array, density_drive: Array, temperature_drive: Array, dtype: dtype) → dict[str, Array]

Return reduced linear ITG frequency, growth, and flux-weight features.

spectraxgk.objectives.stellarator_reduced._qa_quasilinear_heat_flux(linear: dict[str, Array], config: StellaratorITGOptimizationConfig, dtype: dtype) → Array

Return the shipped reduced mixing-length heat-flux proxy.

spectraxgk.objectives.stellarator_reduced._residual_precision_gate_tolerances(fd_step: float) → tuple[float, float, float]

Return residual-Jacobian FD tolerances that remain local near zero.

spectraxgk.objectives.stellarator_reduced._sampled_qa_itg_fields(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig, sample_set: StellaratorITGSampleSet) → dict[str, Array]

Return smooth reduced ITG fields over a surface/alpha/ky sample set.

spectraxgk.objectives.stellarator_reduced.default_stellarator_initial_params() → Array

Return the shared off-optimum QA max-mode-1 starting point.

spectraxgk.objectives.stellarator_reduced.nonlinear_heat_flux_trace(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → tuple[Array, Array]

Return a differentiable short-window ITG heat-flux envelope trace.

The envelope evolves E with a fixed-step RK2 discretization,

dE/dt = 2 gamma E - alpha E^2, Q_env(t) = W_i E.

gamma and W_i come from the same differentiable QA/ITG feature map as the linear and quasilinear objectives. The output is therefore useful for nonlinear averaging, optimizer, and UQ gates while the full production nonlinear-GK geometry path is still being made traceable end-to-end.

spectraxgk.objectives.stellarator_reduced.nonlinear_heat_flux_window_metrics(times: Array, heat_flux: Array, *, tail_fraction: float = 0.45, eps: float = 1e-14) → dict[str, Array]

Return mean, coefficient of variation, and trend on a late-time window.

spectraxgk.objectives.stellarator_reduced.qa_max_mode1_observables(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None, *, density_gradient: float | Array | None = None, temperature_gradient: float | Array | None = None) → dict[str, Array]

Map a QA max-mode-1 boundary/control vector to differentiable ITG observables.

The four inputs represent the active low-order controls used by the example scripts. The map is calibrated as a smooth objective-reduction gate around a QA stellarator with aspect ratio 7 and mean rotational transform 0.41. It is not a replacement for the full VMEC/Boozer flux-tube geometry contract; its purpose is to validate gradient plumbing, UQ, optimizer behavior, and figure-generation before expensive production objectives are promoted.

spectraxgk.objectives.stellarator_reduced.qa_observable_vector(params: Array | Sequence[float], config: StellaratorITGOptimizationConfig | None = None) → Array

Return observables in the stable order defined by OBSERVABLE_NAMES.

spectraxgk.objectives.stellarator_reduced.smooth_positive(x: Array | float, *, beta: float = 18.0) → Array

Smooth positive part used to keep objectives differentiable near marginality.

spectraxgk.objectives.stellarator_reduced.stellarator_itg_objective(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None) → Array

Return the scalar constrained QA + ITG objective for one optimization.

spectraxgk.objectives.stellarator_reduced.stellarator_itg_objective_residual_names(kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux']) → tuple[str, ...]

Return stable residual names for the weighted QA + ITG objective.

spectraxgk.objectives.stellarator_reduced.stellarator_itg_objective_residual_vector(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None) → Array

Return the weighted residual map used by the objective and covariance.

spectraxgk.objectives.stellarator_reduced.stellarator_itg_residual_sensitivity_report(params: Array | Sequence[float], kind: Literal['growth', 'quasilinear_flux', 'nonlinear_heat_flux'], config: StellaratorITGOptimizationConfig | None = None, *, step: float | None = None, rtol: float | None = None, atol: float | None = None, min_rank: int = 4, condition_number_limit: float = 10000.0, covariance_regularization: float = 1e-08, finite_difference_workers: int = 1, finite_difference_executor: str = 'thread') → dict[str, Any]

Check residual-Jacobian AD/FD parity and local conditioning.

Stellarator Objective Portfolio Contracts

Backend-free objective portfolio shape, weight, and reduction contracts.

class spectraxgk.objectives.portfolio_contracts.StellaratorObjectivePortfolioContract(n_surfaces: int, n_alphas: int, n_ky: int, n_objectives: int, reduction: Literal['weighted_mean', 'mean', 'max'], uses_sample_weights: bool, uses_separable_sample_weights: bool, uses_objective_weights: bool)

Static shape/weight contract for a reduced objective portfolio.

property n_samples: int

Number of surface/alpha/ky samples in the portfolio.

property row_shape: tuple[int, int, int, int]

Expected objective-table shape (surface, alpha, ky, objective).

property sample_shape: tuple[int, int, int]

Expected sample-weight shape (surface, alpha, ky).

to_dict() → dict[str, object]

Return a JSON-friendly representation.

spectraxgk.objectives.portfolio_contracts.aggregate_objective_portfolio(objective_rows: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, objective_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean', validate: bool = True) → Array

Reduce a (surface, alpha, ky, objective) table to one scalar.

weighted_mean normalizes both sample and objective weights to unit sum, making the scalar invariant to the caller’s absolute weight scale. mean is the unweighted mean over every table entry. max returns the worst-case objective-weighted sample and is intended for diagnostics rather than smooth gradient-based optimization.

spectraxgk.objectives.portfolio_contracts.portfolio_objective_weight_vector(objective_rows: Any, *, objective_weights: Any | None = None) → Array

Return normalized objective-column weights.

spectraxgk.objectives.portfolio_contracts.portfolio_sample_weight_tensor(objective_rows: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None) → Array

Return normalized sample weights with shape (surface, alpha, ky).

spectraxgk.objectives.portfolio_contracts.validate_objective_portfolio_contract(objective_rows: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, objective_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean') → StellaratorObjectivePortfolioContract

Validate static row/weight contracts and return portfolio metadata.

Concrete weights must be finite, non-negative, and have positive sum. Under JAX tracing, value-level weight checks are deferred to the caller, but shape contracts remain enforced from static array shapes.

Stellarator Objective Portfolio Sensitivity

Autodiff, finite-difference, conditioning, and covariance gates for portfolios.

spectraxgk.objectives.portfolio_sensitivity.objective_portfolio_sensitivity_report(objective_row_fn: Callable[[Array], Any], params: Any, *, sample_weights: Any | None = None, surface_weights: Any | None = None, alpha_weights: Any | None = None, ky_weights: Any | None = None, objective_weights: Any | None = None, reduction: Literal['weighted_mean', 'mean', 'max'] = 'weighted_mean', step: float = 0.0001, rtol: float = 0.0001, atol: float = 1e-06, min_rank: int | None = None, condition_number_limit: float = 100000000.0, covariance_regularization: float = 1e-09, workers: int = 1, parallel_executor: str = 'thread') → dict[str, object]

AD/FD and conditioning report for a reduced objective-row portfolio.

objective_row_fn is the backend boundary: production callers can wire a VMEC/Boozer/quasilinear row builder into this gate while tests can use a cheap fixture. The report checks both the final scalar reduction and the unreduced row sensitivity map so a passing scalar gradient cannot hide a rank-deficient or badly conditioned objective table.

Stellarator Objective Portfolio Artifact Guards

Promotion guards for real VMEC/Boozer reduced objective-portfolio artifacts.

class spectraxgk.objectives.portfolio_artifacts.ReducedPortfolioArtifactGuardConfig(min_alphas: int = 2, min_ky: int = 2, min_objectives: int = 1, min_boozer_mode: int = 21, require_growth_objective: bool = True, require_quasilinear_objective: bool = True, require_vmec_paths: bool = True, value_rtol: float = 1e-08, value_atol: float = 1e-08)

Requirements for promoting real VMEC/Boozer reduced-portfolio rows.

to_dict() → dict[str, object]

Return a JSON-friendly representation.

spectraxgk.objectives.portfolio_artifacts.reduced_portfolio_artifact_guard_report(row_artifact: dict[str, Any], *, gradient_artifacts: list[dict[str, Any]] | tuple[dict[str, Any], ...] = (), config: ReducedPortfolioArtifactGuardConfig | None = None) → dict[str, object]

Validate a real VMEC/Boozer reduced-portfolio artifact before promotion.

The guard is backend-free: it consumes already-generated JSON payloads, rebuilds a (surface, alpha, ky, objective) reducer table from real VMEC/Boozer sample rows, and checks that provenance, coverage, FD/AD diagnostics, and nonlinear-claim boundaries are explicit.

Runtime Runner

Unified runtime-configured linear driver (case-agnostic core path).

class spectraxgk.runtime.RuntimeIndependentParallelPlan(requested_workers: int, effective_workers: int, executor: str, strategy: str, axis: str, source: str, problem_size: int)

Resolved independent-worker policy for runtime scan workloads.

property enabled: bool

Whether the resolved plan uses more than one independent worker.

to_dict() → dict[str, Any]

Return a JSON-friendly policy payload for runtime artifacts.

class spectraxgk.runtime.RuntimeLinearResult(ky: float, gamma: float, omega: float, selection: ModeSelection, t: ndarray | None = None, signal: ndarray | None = None, field_history: ndarray | None = None, state: ndarray | None = None, z: ndarray | None = None, eigenfunction: ndarray | None = None, fit_window_tmin: float | None = None, fit_window_tmax: float | None = None, fit_signal_used: str | None = None, quasilinear: dict[str, Any] | None = None)

Result container for runtime linear runs.

class spectraxgk.runtime.RuntimeLinearScanResult(ky: ndarray, gamma: ndarray, omega: ndarray, quasilinear: tuple[dict[str, Any], ...] | None = None, parallel: dict[str, Any] | None = None)

Result container for runtime linear ky scans.

class spectraxgk.runtime.RuntimeNonlinearResult(t: ndarray, diagnostics: SimulationDiagnostics | None, phi2: ndarray | None = None, fields: FieldState | None = None, state: ndarray | None = None, ky_selected: float | None = None, kx_selected: float | None = None)

Result container for runtime nonlinear runs.

spectraxgk.runtime.build_runtime_geometry(cfg: RuntimeConfig) → SAlphaGeometry | SlabGeometry | FluxTubeGeometryData

Resolve runtime geometry while preserving the runtime module patch surface.

spectraxgk.runtime.build_runtime_linear_params(cfg: RuntimeConfig, *, Nm: int | None = None, geom: SAlphaGeometry | SlabGeometry | FluxTubeGeometryData | None = None) → LinearParams

Build runtime linear parameters using the runtime module geometry surface.

spectraxgk.runtime.build_runtime_linear_terms(cfg: RuntimeConfig) → LinearTerms

Build runtime linear term toggles.

spectraxgk.runtime.build_runtime_term_config(cfg: RuntimeConfig) → TermConfig

Build runtime nonlinear-ready term config.

spectraxgk.runtime.run_runtime_linear(cfg: RuntimeConfig, *, ky_target: float = 0.3, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', return_state: bool = False, initial_state: Any | None = None, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → RuntimeLinearResult

Run one linear point from a case-agnostic runtime config.

spectraxgk.runtime.run_runtime_nonlinear(cfg: RuntimeConfig, *, ky_target: float = 0.3, kx_target: float | None = None, Nl: int | None = None, Nm: int | None = None, dt: float | None = None, steps: int | None = None, method: str | None = None, sample_stride: int | None = None, diagnostics_stride: int | None = None, laguerre_mode: str | None = None, diagnostics: bool | None = None, resolved_diagnostics: bool = True, return_state: bool = False, show_progress: bool = False, status_callback: Callable[[str], None] | None = None) → RuntimeNonlinearResult

Run a nonlinear point using the unified runtime config path.

spectraxgk.runtime.run_runtime_scan(cfg: RuntimeConfig, ky_values: Sequence[float], *, Nl: int | None = None, Nm: int | None = None, solver: str = 'auto', method: str | None = None, dt: float | None = None, steps: int | None = None, sample_stride: int | None = None, batch_ky: bool = False, auto_window: bool = True, tmin: float | None = None, tmax: float | None = None, window_fraction: float = 0.4, min_points: int = 40, start_fraction: float = 0.2, growth_weight: float = 0.2, require_positive: bool = True, min_amp_fraction: float = 0.0, krylov_cfg: KrylovConfig | None = None, mode_method: str = 'project', fit_signal: str = 'auto', show_progress: bool = False, workers: int = 1, parallel_executor: str = 'thread') → RuntimeLinearScanResult

Run a ky scan using the unified runtime config path.

The public facade keeps runtime monkeypatch seams intact while scan coordination lives in workflows/runtime/orchestration_scan.py.