Capability index: GeoNumerics

Layers: GeoStack · GeoNumerics · GeoML · GeoModels.

Every public name in the libraries geostack builds on, each listed once under the shortest module path that exports it, with the first sentence of its docstring. Search this page before writing a loader, a model, a training step or a helper in a script or a notebook: if what you need is here, use it; if it is almost here, extend it where it lives. Regenerate with make capabilities after changing a public API (tests/test_capabilities.py checks it is current).

This page is the GeoNumerics layer. It imports no other layer.

Listed at gaussx 0.6.5, finitevolx 0.0.47, spectraldiffx 0.1.1.

Upstream: `gaussx`¶

Structured linear algebra and Gaussian primitives on lineax.

`gaussx`¶

Name	Kind	What it does
`AbstractIntegrator`	class	Protocol for Gaussian integral approximation.
`AbstractLikelihood`	class	Base class for likelihood functions with optional analytical ELL.
`AbstractLogdetStrategy`	class	Protocol for log-determinant strategies.
`AbstractMultivariateNormal`	class	Common surface of `MultivariateNormal` and `MultivariateNormalPrecision`.
`AbstractPreconditioner`	class	Protocol for preconditioners producing an approximate inverse.
`AbstractSketch`	class	A random sketching matrix $S \in \mathbb{R}^{d \times m}$ , sampled once.
`AbstractSolveStrategy`	class	Protocol for linear solve strategies.
`AbstractSolverStrategy`	class	Protocol for solver strategies that pair solve + logdet.
`AnalyticalPsiStatistics`	class	Protocol for kernels with closed-form Ψ statistics.
`AssumedDensityFilter`	class	KL-optimal Gaussian projection via moment matching.
`AutoSolver`	class	Automatic solver selection based on operator type and size.
`BBMMSolver`	class	CG solve and SLQ logdet with BBMM’s defaults (Gardner et al. 2018).
`BYM2GMRF`	class	The BYM2 pair $(b, u^\ast)$ (Riebler et al., 2016), with its exact density.
`BernoulliLikelihood`	class	Bernoulli likelihood with logit link.
`BinomialLikelihood`	class	Binomial likelihood with logit link.
`BlockDiag`	class	Block diagonal operator `diag(A₁, A₂, …, Aₖ)`.
`BlockTriDiag`	class	Block-tridiagonal operator with a symmetric off-diagonal band.
`CGSolver`	class	Iterative CG solver with stochastic log-determinant.
`CUR`	class	CUR decomposition $A \approx C\,U\,R$ , $C = A_{:,J}$ , $R = A_{I,:}$ .
`CapacitanceSolver`	class	Solve a base system subject to homogeneous point constraints.
`Chandrupatla`	class	Chandrupatla’s hybrid inverse-quadratic / bisection root finder.
`ColumnID`	class	Column interpolative decomposition $A \approx A_{:,J}\,X$ .
`ComposedSolver`	class	Mix-and-match solve and logdet from different strategies.
`ConstantSDE`	class	State-space representation of a constant kernel.
`ConstrainedGMRF`	class	A `GaussianMRF` conditioned on hard linear constraints $A_cx = e$ .
`CosineSDE`	class	State-space representation of the cosine kernel.
`CubatureIntegrator`	class	Third-order spherical-radial cubature.
`DAREResult`	class	Result of DARE solver.
`DenseFallbackWarning`	class	Warning emitted when a structured primitive materialises an operator.
`DenseLogdet`	class	Dense log-determinant via gaussx structural dispatch.
`DenseSolver`	class	Dense solver strategy using gaussx structural dispatch.
`DiagonalizedOperator`	class	Linear operator `A x = V⁻¹(Λ ⊙ V x)` given by a transform pair.
`EigenDecomposition`	class	Precomputed eigendecomposition `A = V diag(λ) V⁻¹` of a square matrix.
`EmissionModel`	class	Observation (emission) model wrapping a linear observation matrix.
`FifthOrderCubatureIntegrator`	class	Fifth-order fully symmetric cubature (McNamee & Stenger).
`FilterState`	class	Output of `kalman_filter`.
`GaussHermiteIntegrator`	class	Gauss-Hermite quadrature integrator.
`GaussianExpFam`	class	Gaussian in natural (exponential family) parameters.
`GaussianLikelihood`	class	Gaussian likelihood `log N(y \| f, noise_var * I)`.
`GaussianMRF`	class	Gaussian Markov random field $x \sim \mathcal N(\mu, Q^{-1})$ .
`GaussianSites`	class	Time-varying Gaussian likelihood sites in natural parameterization.
`GaussianSketch`	class	Gaussian sketch: i.i.d. entries $S_{ij} \sim \mathcal{N}(0, 1/d)$ .
`GaussianState`	class	Gaussian distribution as (mean, covariance operator) pair.
`HeteroscedasticGaussianLikelihood`	class	Heteroscedastic Gaussian likelihood with input-dependent noise.
`IndefiniteSLQLogdet`	class	Stochastic `log\|det(A)\|` for symmetric (possibly indefinite) operators.
`IntegratedWienerSDE`	class	State-space representation of an integrated Wiener process.
`InterpolatedOperator`	class	Structured Kernel Interpolation: `K \approx W K_{uu} W^T`.
`IntrinsicGMRF`	class	Improper $\mathcal N(\mu, (\tau R)^+)$ on the complement of $\ker R$ .
`JacobiPreconditioner`	class	Diagonal preconditioner `M^{-1} = diag(1 / \|diag(A)\|)`.
`KeyedSolver`	class	Bind a PRNG key to a strategy’s stochastic log-determinant (gh-384).
`Kronecker`	class	Kronecker product operator `A₁ ⊗ A₂ ⊗ … ⊗ Aₖ`.
`KroneckerSum`	class	Kronecker sum `A \oplus B = A \otimes I_b + I_a \otimes B`.
`KroneckerSumSqrt`	class	Symmetric square root of `A \oplus B` via per-factor eigenvectors.
`LGSSM`	class	Linear-Gaussian state-space model as a density over `(T, M)`.
`LGSSMFactory`	class	Callable `obs_mask -> MaskedLGSSM`, for conditional use.
`LOOResult`	class	Result of leave-one-out cross-validation.
`LOVECache`	class	Cached rank-`k` Lanczos approximation of `K^{-1}`.
`LSMRSolver`	class	LSMR iterative least-squares solver (Fong & Saunders 2011).
`LaplaceResult`	class	The Gaussian approximation $\mathcal N(\hat x, H^{-1})$ at the mode.
`LinearizedSDE`	class	Linear surrogate $dx = (A_t x + b_t)\,dt + Q_t^{1/2}\,dW$ of an SDE.
`LineaxSolver`	class	A lineax solver as a gaussx solve strategy.
`LowRankUpdate`	class	Low-rank update operator `L + U diag(d) Vᵀ`.
`LowerBlockTriDiag`	class	Lower triangular block-bidiagonal Cholesky factor.
`MINRESSolver`	class	MINRES solver for symmetric (possibly indefinite) systems.
`MarkovGaussian`	class	Chain-Markov Gaussian over $x_{0:T-1}$ as a density over `(T, d)`.
`MaskedLGSSM`	class	`LGSSM` whose `log_prob` marginalises unobserved channels exactly.
`MaskedOperator`	class	Row/column-masked view of a base operator.
`MaternSDE`	class	State-space representation of the Matern kernel.
`MomentMatchResult`	class	Tilted log-normaliser and its derivatives w.r.t. the cavity mean.
`MonteCarloIntegrator`	class	Monte Carlo moment matching: sample, propagate, compute moments.
`MultivariateNormal`	class	Multivariate normal parameterized by a lineax linear operator.
`MultivariateNormalPrecision`	class	Multivariate normal parameterized by a precision (inverse covariance) operator.
`NegativeBinomialLikelihood`	class	Negative-binomial likelihood with log link (NB2, a gamma-Poisson mixture).
`NystromLogdet`	class	Nyström-preconditioned log-determinant of $A + \mu I$ (Wenger et al., 2022).
`NystromPreconditioner`	class	Randomized Nyström preconditioner for `A + μ I`.
`OperatorPreconditioner`	class	Use an externally supplied approximate inverse as a preconditioner.
`OrthonormalSketch`	class	Gaussian sketch with orthonormal rows, $S S^\top = I_d$ .
`PartialCholeskyPreconditioner`	class	Preconditioner from a pivoted partial Cholesky factor.
`PeriodicSDE`	class	State-space representation of the periodic (MacKay) kernel.
`PoissonLikelihood`	class	Poisson likelihood with log link.
`PreconditionedCGSolver`	class	CG solver with pivoted partial Cholesky preconditioner.
`PredictionCache`	class	Cached training solve (and factor) for amortized predictions.
`ProductSDE`	class	Product of two SDE kernels via Kronecker composition.
`PropagationResult`	class	Output of uncertainty propagation through a nonlinear function.
`QuasiPeriodicSDE`	class	Quasi-periodic kernel: product of Matern and Periodic SDE.
`RootDecomposition`	class	Tall factor `R` with `R Rᵀ ≈ A` semantics.
`RowSamplingSketch`	class	Weighted row-sampling sketch.
`SDEKernel`	class	Abstract base class for state-space kernel representations.
`SDEParams`	class	Continuous-time SDE parameters for a linear SDE.
`SLQLogdet`	class	Stochastic log-determinant via Lanczos quadrature (SLQ).
`SLRResult`	class	Linear-Gaussian surrogate produced by statistical linear regression.
`SRHTSketch`	class	Subsampled randomized Hadamard transform sketch.
`SVDLowRankUpdate`	class	Deprecated subclass of `LowRankUpdate` with `orthonormal=True`.
`SketchAndPrecondLSMR`	class	Sketch-and-precondition LSMR for tall least squares (Rokhlin & Tygert, 2008).
`SoftmaxLikelihood`	class	Softmax (categorical) likelihood for multi-class classification.
`SparseCholeskyFactor`	class	Sparse Cholesky factor `P Q Pᵀ = L Lᵀ` on a static symbolic pattern.
`SparseCholeskySolver`	class	Exact `solve` / `logdet` / `diag_inv` through a sparse Cholesky factor.
`SparseOperator`	class	Sparse matrix with a static `SparsityPattern` and traced values.
`SparseSignSketch`	class	Sparse sign sketch (sparse Johnson-Lindenstrauss transform).
`SparsityPattern`	class	Static, hashable sparsity pattern of a `(m, n)` matrix.
`SpectralFunction`	class	`f(B)` for a symmetric `B` with a factored eigenbasis.
`StudentTLikelihood`	class	Student-t likelihood for robust regression.
`SumKronecker`	class	Deprecated alias for `SumOfKroneckers`.
`SumKroneckerSqrt`	class	Deprecated alias for `SumOfKroneckersSqrt` (gh-297).
`SumOfKroneckers`	class	Sum of Kronecker products `Σ_k A_k \otimes B_k`.
`SumOfKroneckersSqrt`	class	Lazy Lanczos square-root operator for `SumOfKroneckers` covariances.
`SumSDE`	class	Sum of SDE kernels via block-diagonal composition.
`SymbolicCholesky`	class	Symbolic Cholesky factor of a symmetric sparsity pattern.
`TaylorIntegrator`	class	1st or 2nd order Taylor expansion for uncertainty propagation.
`Toeplitz`	class	Symmetric Toeplitz matrix from its first column.
`ToeplitzCholesky`	class	Circulant-embedding sample factor for a symmetric positive Toeplitz matrix.
`UDLDecomposition`	class	Banded factorisation $\Lambda = U \tilde{D} U^{\top}$ of a `BlockTriDiag`.
`UnscentedIntegrator`	class	Unscented transform: deterministic sigma points.
`UpperBlockTriDiag`	class	Upper triangular block-bidiagonal (transpose of LowerBlockTriDiag).
`add_jitter`	function	Add diagonal jitter for numerical stability: `A + eps * I`.
`ar1_precision`	function	Precision of a stationary AR(1) process with marginal precision `τ`.
`as_diagonalized`	function	Express `operator` as one `DiagonalizedOperator`, if possible.
`as_linear_operator`	function	Wrap a raw `matvec` callable as a tagged lineax operator.
`base_conditional`	function	Deprecated: use `sparse_conditional` (note `K_xz = K_mn.T`).
`besag_structure`	function	Validate a graph Laplacian as a Besag (ICAR) structure matrix.
`block_diagonal_tag`	constant	`block_diagonal_tag`
`block_tridiagonal_tag`	constant	`block_tridiagonal_tag`
`blr_diag_update`	function	Diagonal natural parameter BLR update step.
`blr_full_update`	function	Full-rank natural parameter BLR update step.
`build_prediction_cache`	function	Solve `K_y alpha = y` and cache the result.
`bym2_precision`	function	Joint precision of the BYM2 pair `(b, u*)` (Riebler et al., 2016).
`cavity_distribution`	function	Compute EP cavity distribution by removing a site.
`cavity_from_marginal`	function	Compute cavity distribution by removing a site from the marginal.
`cholesky`	function	Compute Cholesky factor L such that A = L L^T.
`cholesky_logdet`	function	Compute log\|A\| from Cholesky factor L where A = L Lᵀ.
`circulant`	function	(Block-)circulant operator from its first column.
`circulant_from_symbol`	function	(Block-)circulant operator from its DFT symbol.
`collapsed_elbo`	function	Collapsed ELBO (Titsias bound) for sparse GP regression.
`column_id`	function	Randomized column interpolative decomposition (Voronin & Martinsson, 2017).
`compute_psi_statistics`	function	Compute Ψ statistics, dispatching to analytical or numerical.
`conditional`	function	Compute `p(x_A \| x_B = b)` from a joint Gaussian `p(x_A, x_B)`.
`conditional_interpolate`	function	Two-filter fusion at `t` from a filtered state and a backward message.
`conditional_variance`	function	Predictive variance: Schur complement diagonal plus optional variational correction.
`cost_expectation`	function	Compute `E[Cost(f(x), target)]` where `x ~ N(mu, Sigma)`.
`cov_transform`	function	Covariance propagation through a linear map: `J @ Sigma @ J^T`.
`create_grid`	function	Create a regular grid from per-dimension sizes and bounds.
`cubature_points`	function	Spherical-radial cubature points and weights.
`cubic_interpolation_weights`	function	Compute cubic interpolation indices and weights for SKI.
`cur`	function	Randomized CUR decomposition (Voronin & Martinsson, 2017, §4).
`cvi_update_sites`	function	Natural gradient update for CVI sites.
`damped_natural_update`	function	Damped update in natural parameter space.
`dare`	function	Steady-state Kalman filter covariance and gain (the filtering DARE).
`diag`	function	Extract the diagonal of an operator as a 1D array.
`diag_conditional_variance`	function	Diagonal of Schur complement: `diag(K_XX - A K_ZX)`.
`diag_inv`	function	Compute the diagonal of the inverse of a linear operator.
`diagonal_tag`	constant	`diagonal_tag`
`discrepancy_step_size`	function	Adaptive EKI tempering step: the data misfit controller.
`discrete_lyapunov_solve`	function	Solve the discrete Lyapunov equation `P - G P G^T = Q`.
`discretize_mfd`	function	Discretise a linear SDE by matrix-fraction decomposition.
`discretize_mfd_sequence`	function	Vectorised `gaussx.discretize_mfd` over a vector of time steps.
`eig`	function	Compute eigenvalues and eigenvectors.
`eigh_generalized`	function	Solve $A v = \lambda B v$ for symmetric $A$ and symmetric PSD $B$ .
`eigvals`	function	Compute eigenvalues only.
`eki_step`	function	One ensemble Kalman inversion (EKI) update.
`elbo`	function	Evidence lower bound (ELBO).
`enkf_analysis`	function	Stochastic (perturbed-observation) ensemble Kalman analysis step.
`ensemble_covariance`	function	Empirical covariance from an ensemble as a low-rank operator.
`ensemble_cross_covariance`	function	Cross-covariance between two ensemble sets.
`ensemble_kalman_gain`	function	Kalman gain from an ensemble and its image in observation space.
`ep_tilted_moments`	function	Moments of the scalar EP tilted distribution, elementwise over a batch.
`estimate_spectral_bounds`	function	Estimate $(\lambda_{\min}, \lambda_{\max})$ of a symmetric operator.
`etkf_transform`	function	Ensemble Transform Kalman Filter (ETKF) analysis weights.
`euclidean_distance`	function	Pairwise Euclidean distances `\|\|a_i - b_j\|\|`.
`expectation_to_mean_chol`	function	Convert expectation parameters to mean/variance (Cholesky).
`expectation_to_natural`	function	Convert expectation parameters to natural parameters.
`expectations_to_ssm`	function	Convert expectation parameters back to SSM marginals.
`expected_log_likelihood`	function	Unified expected log-likelihood with analytical dispatch.
`fem_matrices`	function	Lumped mass `C̃` and stiffness `G` of P1 elements on a triangle mesh.
`fem_projector`	function	Observation matrix `A` of P1 interpolation: `(A w)_k = Σ_i ψ_i(s_k) w_i`.
`fifth_order_cubature_points`	function	Fifth-order fully symmetric cubature points and weights.
`fisher_info`	function	Fisher information with respect to the mean, $F_\mu = \Sigma^{-1}$ .
`frobenius_norm`	function	Compute the Frobenius norm `\|\|A\|\|_F` with structural dispatch.
`gaspari_cohn`	function	Gaspari-Cohn (1999) fifth-order compactly-supported taper.
`gauss_hermite_points`	function	Gauss-Hermite quadrature points and weights.
`gauss_kl`	function	KL(first ‖ second): `KL[q(u) \|\| p(u)]` with `p(u) = N(0, K)`.
`gauss_newton_precision`	function	Gauss-Newton precision matrix `J^T J`, optionally plus a prior.
`gaussian_entropy`	function	Entropy of a multivariate normal `N(mu, Sigma)`.
`gaussian_expected_log_lik`	function	Expected log-likelihood `E_q[log N(y \| f, R)]`.
`gaussian_kl`	function	KL(first ‖ second): `KL(p \|\| q)` between two multivariate normals.
`gaussian_log_prob`	function	Multivariate normal log-probability.
`generalized_variance_scale`	function	Generalized variance of an intrinsic GMRF (Sørbye & Rue, 2014).
`ggn_diagonal`	function	Generalized Gauss-Newton diagonal approximation.
`gradient_expectation`	function	Compute `E[nabla f(x)]` via Stein’s lemma.
`grid_coupling_indices`	function	Flat indices of masked-out grid cells adjacent to the masked-in region.
`grid_data`	function	Expand a grid to the full Cartesian product of grid points.
`hadamard_transform`	function	Unnormalized Walsh-Hadamard transform along the last axis.
`haversine_distance`	function	Pairwise great-circle (haversine) distances on a sphere.
`hutchinson_hessian_diag`	function	Stochastic Hessian diagonal via Hutchinson with Rademacher probes.
`iid_precision`	function	Precision `τ I` of `n` independent effects with variance `1/τ`.
`infinite_horizon_filter`	function	Infinite-horizon Kalman filter with fixed steady-state gain.
`infinite_horizon_smoother`	function	Infinite-horizon RTS smoother with fixed steady-state gain.
`inflate_multiplicative`	function	Multiplicative ensemble inflation about the mean.
`inflate_rtpp`	function	Relaxation to prior perturbations (RTPP; Zhang et al. 2004).
`inflate_rtps`	function	Relaxation to prior spread (RTPS; Whitaker & Hamill 2012).
`inv`	function	Return a lazy inverse operator A^{-1}.
`inv_quad_logdet`	function	Joint inverse-quadratic and log-determinant via shared BBMM CG.
`is_block_diagonal`	function	Check whether operator carries the block-diagonal tag.
`is_block_tridiagonal`	function	Check whether operator carries the block-tridiagonal tag.
`is_diagonal`	function	Re-exported from `lineax`.
`is_eigen_reducible`	function	Whether a sum of Kronecker products has the exact two-term reduction.
`is_kronecker`	function	Check whether operator carries the Kronecker tag.
`is_kronecker_sum`	function	Check whether operator carries the Kronecker sum tag.
`is_low_rank`	function	Check whether operator carries the low-rank tag.
`is_lower_triangular`	function	Re-exported from `lineax`.
`is_negative_semidefinite`	function	Re-exported from `lineax`.
`is_positive_semidefinite`	function	Re-exported from `lineax`.
`is_symmetric`	function	Re-exported from `lineax`.
`is_upper_triangular`	function	Re-exported from `lineax`.
`joseph_update`	function	Numerically stable Joseph-form covariance update.
`kalman_filter`	function	Kalman filter forward pass via `jax.lax.scan`.
`kalman_gain`	function	Compute Kalman gain `K = P @ H^T @ (H @ P @ H^T + R)^{-1}`.
`kernel_expectations`	function	Compute kernel expectations Psi_0, Psi_1, Psi_2 for uncertain inputs.
`kl_divergence`	function	KL(first ‖ second): `KL(q \|\| p)` via the Bregman form on natural parameters.
`kl_standard_normal`	function	KL(first ‖ second): `KL(N(m, S) \|\| N(0, I))`.
`kronecker_mll`	function	Exact marginal log-likelihood for a Kronecker-structured GP.
`kronecker_posterior_predictive`	function	Posterior mean and variance for a Kronecker GP at test points.
`kronecker_sum_sample`	function	Sample from `𝒩(0, A ⊕ B)` using per-factor eigendecompositions.
`kronecker_sum_solve`	function	Solve a shifted Kronecker-sum system in tensor (unvectorized) form.
`kronecker_sum_tag`	constant	`kronecker_sum_tag`
`kronecker_tag`	constant	`kronecker_tag`
`laplace_mode`	function	Mode, Hessian and log-marginal of the Laplace approximation, in precision form.
`leave_one_out_cv`	function	LOO-CV via the bordered-system identity.
`linear_solve`	function	Solve `A x = b` through the unified front door.
`linearize_sde`	function	Optimal linear drift of a nonlinear SDE along a Gaussian path.
`localization_matrix`	function	Pairwise Gaspari-Cohn taper `rho(dist(a_i, b_j); c)`.
`localized_kalman_gain`	function	Ensemble Kalman gain with Hadamard (Schur-product) localization.
`log_likelihood_expectation`	function	Compute `E[log p(y_{obs} \| f(x))]` where `x ~ N(mu, Sigma)`.
`log_marginal_likelihood`	function	GP log marginal likelihood.
`log_partition`	function	Log-partition function `A(eta)`.
`logdet`	function	Compute log \|det(A)\| with structural dispatch.
`love_cache`	function	Precompute a Lanczos approximation of `K^{-1}` for fast variance.
`love_residual`	function	Relative solve residual of a LOVE cache for one cross-covariance.
`love_variance`	function	Approximate `k_^\top K^{-1} k_` from a LOVE cache.
`love_variance_error_bound`	function	Upper bound on the variance over-estimate of a LOVE cache.
`low_rank_plus_diag`	function	Construct `diag(diag) + U diag(d) Vᵀ`.
`low_rank_plus_identity`	function	Construct `scale * I + U diag(d) Vᵀ`.
`low_rank_tag`	constant	`low_rank_tag`
`lower_triangular_tag`	constant	`lower_triangular_tag`
`masked_moment_inputs`	function	Make masked observation channels inert in a moment-matched update.
`matern_spde_params`	function	SPDE parameters `(κ, τ, α)` of a Matérn field.
`matheron_update`	function	Posterior samples via Matheron’s-rule correction.
`mean_chol_to_expectation`	function	Convert mean/variance (Cholesky) to expectation parameters.
`mean_chol_to_natural`	function	Convert mean/variance (Cholesky) to natural parameters.
`mean_cov_to_natural`	function	Convert mean/covariance to natural parameters (operator form).
`mean_expectation`	function	Compute `E[f(x)]` where `x ~ N(mu, Sigma)`.
`meanfield_kalman_filter`	function	Mean-field Kalman filter over `L = D // block_size` state blocks.
`meanfield_rts_smoother`	function	Mean-field RTS smoother over `L = D // block_size` state blocks.
`mixture_quantile`	function	Invert a CDF at one or more quantile levels by bracketed root finding.
`mixture_quantile_gaussian_approx`	function	Quantiles of the Gaussian matched to an ensemble’s first two moments.
`moment_match`	function	EP moment matching: tilted log-normaliser and its first two derivatives.
`moment_transform`	function	Moment triple of `(x, fn(x))` for `x` Gaussian.
`natural_to_expectation`	function	Convert natural parameters to expectation parameters.
`natural_to_mean_chol`	function	Convert natural parameters to mean/variance (Cholesky).
`natural_to_mean_cov`	function	Convert natural parameters to mean/covariance (operator form).
`naturals_to_ssm`	function	Convert natural parameters back to SSM parameters.
`negative_semidefinite_tag`	constant	`negative_semidefinite_tag`
`newton_update`	function	Convert a Newton step to natural pseudo-likelihood parameters.
`nonlinear_kalman_filter`	function	Moment-matched nonlinear Kalman filter.
`nonlinear_kalman_predict`	function	One moment-matched predict step.
`nonlinear_kalman_update`	function	One moment-matched update step.
`nonlinear_rts_smoother`	function	Moment-matched nonlinear Rauch-Tung-Striebel smoother.
`nonlinear_rts_step`	function	One moment-matched RTS backward step.
`oilmm_back_project`	function	Back-project latent GP predictions to the observation space.
`oilmm_project`	function	Project multi-output data to independent latent GPs via OILMM.
`pairwise_marginals`	function	Joint p(x_k, x_{k+1}) for each consecutive pair.
`parallel_kalman_filter`	function	Parallel Kalman filter via `jax.lax.associative_scan`.
`parallel_rts_smoother`	function	Parallel RTS smoother via reverse `jax.lax.associative_scan`.
`positive_semidefinite_tag`	constant	`positive_semidefinite_tag`
`predict_mean`	function	Predictive mean: `mu* = K_*f @ alpha`.
`predict_variance`	function	Predictive variance `sigma^2_* = k_** - diag(K_f K_y^{-1} K_f)`.
`process_noise_covariance`	function	Compute process noise from stationary covariance.
`product_operator`	function	Lazy matmul `(A B) v = A (B v)`.
`project`	function	Compute A_X = K_XZ @ K_ZZ^{-1} via Cholesky solve.
`pseudo_logdet`	function	Log pseudo-determinant: the log of the product of the non-zero eigenvalues.
`qb`	function	Randomized QB factorisation $A \approx QB$ with $B = Q^\top A$ .
`quadratic_form`	function	Compute `x^T A^{-1} x` via a single solve.
`randomized_eigh`	function	Partial eigendecomposition of a symmetric operator by Rayleigh-Ritz.
`randomized_nystrom`	function	Randomized Nyström approximation $\hat A = U\hat\Lambda U^\top$ of a PSD $A$ .
`randomized_svd`	function	Truncated SVD $A \approx U \operatorname{diag}(s) V^\top$ by randomized QB.
`range_finder`	function	Orthonormal basis $Q$ for the dominant range of $A$ .
`riemannian_psd_correction`	function	Riemannian gradient correction for PSD precision updates.
`root_decomposition`	function	Compute a tall factor `R` such that `R Rᵀ ≈ A`.
`root_inv_decomposition`	function	Compute a tall factor `R⁻` such that `R⁻ (R⁻)ᵀ ≈ A⁻¹`.
`rp_cholesky`	function	Randomly pivoted partial Cholesky of a PSD matrix `A`.
`rts_interpolate`	function	Smoothed marginal at an unobserved `t` from filtered and smoothed states.
`rts_smoother`	function	Rauch-Tung-Striebel backward smoother.
`rw1_structure`	function	Structure matrix `R = D₁ᵀ W D₁` of a first-order random walk.
`rw2_structure`	function	Structure matrix `R = D₂ᵀ D₂` of a second-order random walk.
`safe_cholesky`	function	Cholesky decomposition with adaptive jitter for near-singular matrices.
`sample_joint_conditional`	function	Joint draws of a partitioned Gaussian, and conditional draws from them.
`sample_mvn`	function	Draw samples from $\mathcal{N}(\mu, K)$ , dispatching on $K$ ’s structure.
`sandwich`	function	Return `A @ P @ A.T` exploiting compatible operator structure.
`scaled_operator`	function	Lazy scalar multiply `(c A) v = c (A v)`.
`schur_complement`	function	Schur complement: `K_XX - K_XZ @ K_ZZ^{-1} @ K_ZX`.
`sde_autocovariance`	function	Compute the stationary autocovariance of an SDE kernel.
`sde_kl_divergence`	function	Drift path-KL between a linear and a nonlinear SDE along a Gaussian path.
`selected_inverse`	function	The band of `Σ = Q⁻¹` for a symmetric positive-definite `BlockTriDiag`.
`sigma_points`	function	Unscented transform sigma points and weights.
`site_mean_var_from_natural`	function	Convert per-site natural parameters to mean/variance.
`site_natural_from_tilted`	function	Compute site natural parameters from tilted and cavity moments.
`sites_to_precision`	function	Convert Gaussian sites to a block-tridiagonal precision.
`sketch_and_solve`	function	Approximate least squares from a sketch.
`solve`	function	Solve `A x = b` with structural dispatch.
`solve_columns`	function	Solve A X = B column-by-column via vmap.
`solve_matrix`	function	Solve `A X = B` with a single factorization on the matrix RHS.
`solve_rows`	function	Solve A x = bᵢ for each row bᵢ of a matrix via vmap.
`sparse_cholesky`	function	Sparse Cholesky factorisation of a symmetric positive-definite operator.
`sparse_conditional`	function	Sparse-GP conditional `q(f_x) = \int p(f_x \| u) q(u) du`.
`spde_precision`	function	SPDE precision `Q_α = τ² K (C̃⁻¹K)^{α−1}` with `K = κ²C̃ + G`.
`spde_precision_grid`	function	SPDE precision on a regular grid: a function of a Kronecker sum.
`spingp_log_likelihood`	function	Log marginal likelihood via sparse inverse GP formulation.
`spingp_posterior`	function	Posterior mean and precision via SpInGP.
`sqrt`	function	Compute matrix square root S such that S @ S = A.
`sqrt_inv_matmul`	function	Compute $A^{-1/2} b$ via contour-integral quadrature.
`sqrt_matmul`	function	Compute $A^{1/2} b$ via contour-integral quadrature.
`ssm_to_expectations`	function	Convert SSM marginals to expectation parameters.
`ssm_to_naturals`	function	Convert SSM parameters to natural parameters.
`stable_squared_distances`	function	Squared Euclidean distances with mixed-precision stability.
`statistical_linear_regression`	function	Moment-match a nonlinear conditional to a linear-Gaussian surrogate.
`submatrix`	function	Extract `A[row_idx, col_idx]` without forming the full matrix.
`sufficient_stats`	function	Compute sufficient statistics `T(x) = [x, x x^T]`.
`sum_of_kroneckers_sample`	function	Sample from `𝒩(0, op)` with matrix-free Lanczos square roots.
`sum_operator`	function	Lazy sum `(A + B + …) v = A v + B v + …`.
`sumkronecker_sample`	function	Deprecated alias for `sum_of_kroneckers_sample` (gh-297).
`svd`	function	Compute the singular value decomposition `A = U diag(s) V^T`.
`svd_low_rank_plus_diag`	function	Construct `diag(diag) + U diag(S) Vᵀ` from a truncated SVD.
`svgp_variance_adjustment`	function	Compute the SVGP variance adjustment operator.
`symbolic_cholesky`	function	Symbolic analysis of a sparse Cholesky factorisation (host, cached).
`symmetric_tag`	constant	`symmetric_tag`
`symmetrize`	function	Symmetrize the trailing two axes: `0.5 * (X + X^T)`.
`theta_design`	function	Integration points and log-weights over hyperparameters $\theta$ .
`tikhonov_augment`	function	Observation augmentation for Tikhonov-regularised EKI (TEKI).
`to_expectation`	function	Deprecated alias of `to_mean_cov`.
`to_mean_cov`	function	Convert natural parameters to the mean and covariance.
`to_natural`	function	Deprecated alias of `gaussx.mean_cov_to_natural`.
`toeplitz_sample`	function	Sample from `𝒩(0, Toeplitz(column))` via FFT circulant embedding.
`trace`	function	Compute the trace of an operator.
`trace_and_diag`	function	Jointly estimate the trace and diagonal from one probe pass.
`trace_correction`	function	Trace term in Titsias collapsed ELBO.
`trace_product`	function	Trace of a matrix product: `tr(A @ B)` with structural dispatch.
`tridiagonal_solve`	function	Solve a tridiagonal system `A x = d`.
`tridiagonal_solve_batched`	function	Solve independent tridiagonal systems over leading batch dimensions.
`tridiagonal_tag`	constant	`tridiagonal_tag`
`udl_decomposition`	function	Factorise a symmetric block-tridiagonal precision as $U \tilde{D} U^{\top}$ .
`udl_from_ssm_params`	function	Build the `UDLDecomposition` of a Gauss-Markov chain’s precision.
`udl_to_ssm_params`	function	Read the equivalent Gauss-Markov chain off a `UDLDecomposition`.
`uncertain_bgplvm_predict`	function	Multi-output uncertain GP prediction for BGPLVM.
`uncertain_gp_predict`	function	Predictive mean and variance for GP with uncertain inputs.
`uncertain_gp_predict_mc`	function	Monte Carlo GP prediction with uncertain inputs.
`uncertain_svgp_predict`	function	Predictive mean and variance for SVGP with uncertain inputs.
`uncertain_vgp_predict`	function	Predictive mean and variance for dense VGP with uncertain inputs.
`unit_diagonal_tag`	constant	`unit_diagonal_tag`
`unwhiten`	function	Unwhiten variational mean: `m = L @ m_tilde`.
`unwhiten_covariance`	function	Unwhiten variational covariance: S = L S̃ Lᵀ.
`upper_triangular_tag`	constant	`upper_triangular_tag`
`variational_elbo_gaussian`	function	Uncollapsed (SVGP) ELBO for a Gaussian likelihood.
`variational_elbo_mc`	function	Monte Carlo ELBO for non-conjugate likelihoods.
`vb_mean_correction`	function	Low-rank VB correction of the Laplace mean (Van Niekerk & Rue, 2024).
`whitened_svgp_predict`	function	Whitened SVGP prediction: mean and variance at test points.
`woodbury_solve`	function	Standalone Woodbury identity solve: `(L + U diag(D) U^T)^{-1} b`.

Upstream: `finitevolx`¶

Arakawa C-grid finite-volume operators.

`finitevolx`¶

Name	Kind	What it does
`AB2Solver`	class	Adams-Bashforth 2nd-order solver (equinox Module).
`AbstractForcing`	class	Base class for forcing operators.
`Advection1D`	class	1-D advection operator.
`Advection2D`	class	2-D advection operator.
`Advection3D`	class	3-D advection operator (horizontal plane per z-level).
`ArakawaCGrid1D`	class	1-D Arakawa C-grid topology (no metric).
`ArakawaCGrid2D`	class	2-D Arakawa C-grid topology (no metric).
`ArakawaCGrid3D`	class	3-D Arakawa C-grid topology (no metric).
`ArakawaJacobian2D`	class	Arakawa (1966) Jacobian J(f, g) at T-points, on the full grid.
`BiharmonicDiffusion2D`	class	Biharmonic (∇⁴) diffusion operator on a 2-D Arakawa C-grid.
`BiharmonicDiffusion3D`	class	Biharmonic (nabla^4) diffusion operator on a 3-D Arakawa C-grid.
`BoundaryCondition`	function	Re-exported from `typing`.
`BoundaryConditionSet`	class	Apply a different boundary condition on each domain face.
`CGInfo`	class	Convergence diagnostics returned by :func:`solve_cg`.
`CapacitanceSolver`	class	Re-exported from `spectraldiffx`.
`CartesianGrid1D`	class	1-D Cartesian Arakawa C-grid (uniform spacing).
`CartesianGrid2D`	class	2-D Cartesian Arakawa C-grid (uniform spacing).
`CartesianGrid3D`	class	3-D Cartesian Arakawa C-grid (uniform spacing).
`Coriolis2D`	class	Coriolis force operator for 2-D Arakawa C-grids.
`Coriolis3D`	class	Coriolis force operator for 3-D Arakawa C-grids.
`CurvilinearGrid1D`	class	1-D curvilinear C-grid with a uniform metric.
`CurvilinearGrid2D`	class	2-D curvilinear C-grid with a uniform metric.
`CurvilinearGrid3D`	class	3-D curvilinear C-grid with a uniform metric.
`Difference1D`	class	Finite-difference operators on a 1-D Arakawa C-grid.
`Difference2D`	class	Finite-difference operators on a 2-D Arakawa C-grid.
`Difference3D`	class	Finite-difference operators on a 3-D Arakawa C-grid.
`Diffusion2D`	class	Horizontal diffusion operator (flux form) on a 2-D Arakawa C-grid.
`Diffusion3D`	class	Horizontal diffusion operator (flux form) on a 3-D Arakawa C-grid.
`Dirichlet1D`	class	Apply a Dirichlet value on one domain face.
`DirichletHelmholtzSolver2D`	class	Re-exported from `spectraldiffx`.
`Divergence2D`	class	Divergence operator on a 2-D Arakawa C-grid.
`Energetics2D`	class	Energy diagnostics at T-points on a 2-D Arakawa C-grid.
`Extrapolation1D`	class	High-order polynomial extrapolation for ghost cells.
`FieldBCSet`	class	Dispatch per-face boundary conditions across multiple fields.
`ForcingProduct`	class	Modulated forcing: `F = alpha * F_inner`.
`ForcingSum`	class	Additive composition of forcings: `F = F_1 + F_2 + ... + F_n`.
`ForwardEulerDfx`	class	Forward Euler via the diffrax `AbstractSolver` interface.
`IMEX_SSP2`	class	IMEX-SSP2(2,2,2) solver using diffrax MultiTerm.
`Interpolation1D`	class	1-D averaging operators on an Arakawa C-grid.
`Interpolation2D`	class	2-D averaging operators on an Arakawa C-grid.
`Interpolation3D`	class	3-D averaging operators on an Arakawa C-grid.
`KnownValueLifting`	class	Pre/post-processing wrapper for the lifting trick.
`LeapfrogRAFSolver`	class	Leapfrog with Robert-Asselin filter (equinox Module).
`LinearDrag2D`	class	Linear drag on a 2-D Arakawa C-grid.
`LinearDrag3D`	class	Linear bottom drag on a 3-D Arakawa C-grid.
`Mask1D`	class	Unified Arakawa C-grid mask for a 1-D Cartesian domain.
`Mask2D`	class	Unified Arakawa C-grid mask for a 2-D Cartesian domain.
`Mask3D`	class	Unified Arakawa C-grid mask for a 3-D Cartesian domain.
`MixedBCHelmholtzSolver2D`	class	Re-exported from `spectraldiffx`.
`MixedBCHelmholtzSolver3D`	class	Re-exported from `spectraldiffx`.
`MomentumAdvection2D`	class	Energy-conserving momentum advection on a 2-D Arakawa C-grid.
`MomentumAdvection3D`	class	Energy-conserving momentum advection on a 3-D Arakawa C-grid.
`MultigridSolver`	class	Geometric multigrid V-cycle solver for the variable-coefficient Helmholtz equation::
`Neumann1D`	class	Apply an outward-normal gradient on one domain face.
`NeumannHelmholtzSolver2D`	class	Re-exported from `spectraldiffx`.
`Outflow1D`	class	Apply a one-face outflow boundary condition.
`Periodic1D`	class	Apply periodic wrapping on one domain face.
`QGPotentialVorticity2D`	class	Quasi-geostrophic potential vorticity at T-points.
`QuadraticDrag2D`	class	Quadratic bottom drag on a 2-D Arakawa C-grid.
`QuadraticDrag3D`	class	Quadratic bottom drag on a 3-D Arakawa C-grid.
`RK2Heun`	class	Heun’s method (RK2) via Butcher tableau.
`RK3SSP`	class	3rd-order Strong-Stability-Preserving Runge-Kutta.
`RK4Classic`	class	Classic 4th-order Runge-Kutta.
`RayleighDamping2D`	class	Rayleigh damping of a T-point field on a 2-D Arakawa C-grid.
`RayleighDamping3D`	class	Rayleigh damping of a T-point field on a 3-D Arakawa C-grid.
`Reconstruction1D`	class	1-D face-value reconstruction.
`Reconstruction2D`	class	2-D face-value reconstruction for advective fluxes.
`Reconstruction3D`	class	3-D face-value reconstruction.
`Reflective1D`	class	Apply an even-symmetry reflective boundary on one face.
`RegularNeumannHelmholtzSolver2D`	class	Re-exported from `spectraldiffx`.
`Robin1D`	class	Robin (mixed) boundary condition: α·u + β·∂u/∂n = γ.
`SSP_RK104`	class	4th-order SSP Runge-Kutta with 10 stages (Ketcheson 2008).
`SSP_RK2`	class	2nd-order SSP Runge-Kutta (same as Heun).
`SemiLagrangianSolver`	class	Semi-Lagrangian advection solver for diffrax.
`Slip1D`	class	Slip boundary condition for tangential velocity at a solid wall.
`SolveDomain`	class	Derived masks for the lifting trick.
`SpectralHelmholtzSolver1D`	class	Re-exported from `spectraldiffx`.
`SpectralHelmholtzSolver2D`	class	Re-exported from `spectraldiffx`.
`SpectralHelmholtzSolver3D`	class	Re-exported from `spectraldiffx`.
`SphericalAdvection2D`	class	2-D tracer advection on a spherical Arakawa C-grid.
`SphericalAdvection3D`	class	3-D tracer advection on a spherical Arakawa C-grid.
`SphericalBiharmonicDiffusion2D`	class	Biharmonic (spherical ∇⁴) diffusion on a 2-D spherical C-grid.
`SphericalBiharmonicDiffusion3D`	class	Biharmonic spherical diffusion on a 3-D Arakawa C-grid.
`SphericalDifference2D`	class	Spherical finite-difference operators on a 2-D Arakawa C-grid.
`SphericalDifference3D`	class	Spherical finite-difference operators on a 3-D Arakawa C-grid.
`SphericalDiffusion2D`	class	Horizontal tracer diffusion (flux form) on a 2-D spherical C-grid.
`SphericalDiffusion3D`	class	Horizontal tracer diffusion on a 3-D spherical Arakawa C-grid.
`SphericalDivergence2D`	class	Horizontal divergence on a 2-D spherical Arakawa C-grid.
`SphericalDivergence3D`	class	3-D spherical divergence (vmap over z-levels).
`SphericalGrid2D`	class	2-D Arakawa C-grid on a sphere.
`SphericalGrid3D`	class	3-D Arakawa C-grid on a sphere.
`SphericalLaplacian2D`	class	Spherical Laplacian on a 2-D Arakawa C-grid.
`SphericalLaplacian3D`	class	3-D spherical Laplacian (vmap over z-levels).
`SphericalMomentumAdvection2D`	class	Energy-conserving momentum advection on a 2-D spherical C-grid.
`SphericalMomentumAdvection3D`	class	Energy-conserving momentum advection on a 3-D spherical C-grid.
`SphericalVorticity2D`	class	Vorticity and PV-flux operators on a 2-D spherical Arakawa C-grid.
`SphericalVorticity3D`	class	3-D spherical vorticity (vmap over z-levels).
`SplitExplicitRKSolver`	class	Split-explicit barotropic/baroclinic solver.
`Sponge1D`	class	Relax one ghost face toward a background value.
`StaggeredDirichletHelmholtzSolver2D`	class	Re-exported from `spectraldiffx`.
`StencilCapability1D`	class	Directional count of contiguous wet neighbours for each 1-D grid cell.
`StencilCapability2D`	class	Directional count of contiguous wet neighbours for each 2-D grid cell.
`StencilCapability3D`	class	Directional count of contiguous wet neighbours for each 3-D grid cell.
`Strain2D`	class	Strain diagnostics on a 2-D Arakawa C-grid.
`TimeVaryingForcing`	class	Bundle a time-agnostic forcing with time-dependent input fields.
`Vorticity2D`	class	2-D vorticity and PV-flux operators.
`Vorticity3D`	class	3-D vorticity operators (horizontal plane per z-level).
`WindStress2D`	class	Surface wind-stress forcing on a 2-D Arakawa C-grid.
`WindStress3D`	class	Surface wind-stress forcing on a 3-D Arakawa C-grid.
`ab2_step`	function	2nd-order Adams-Bashforth.
`ab3_step`	function	3rd-order Adams-Bashforth.
`arakawa_jacobian`	function	Arakawa (1966) discretization of J(f, g).
`area_mean`	function	Area-weighted mean over physical interior T-cells.
`area_sum`	function	Area-weighted sum over physical interior T-cells.
`area_weights`	function	Dispatch to :func:`cartesian_area_weights` or :func:`spherical_area_weights`.
`available_potential_energy`	function	Available potential energy at T-points.
`avg_x_bwd`	function	Backward average in x (east face → centre).
`avg_x_bwd_1d`	function	Backward average in x (east face → centre), 1-D.
`avg_x_bwd_3d`	function	Backward average in x over all z-levels (east face → centre).
`avg_x_fwd`	function	Forward average in x (centre → east face).
`avg_x_fwd_1d`	function	Forward average in x (centre → east face), 1-D.
`avg_x_fwd_3d`	function	Forward average in x over all z-levels (centre → east face).
`avg_xbwd_yfwd`	function	Bilinear average: U-points → V-points (backward-x, forward-y).
`avg_xfwd_ybwd`	function	Bilinear average: V-points → U-points (forward-x, backward-y).
`avg_xy_bwd`	function	Bilinear average from corners to centre (backward in both x and y).
`avg_xy_fwd`	function	Bilinear average to NE corner (forward in both x and y).
`avg_y_bwd`	function	Backward average in y (north face → centre).
`avg_y_bwd_3d`	function	Backward average in y over all z-levels (north face → centre).
`avg_y_fwd`	function	Forward average in y (centre → north face).
`avg_y_fwd_3d`	function	Forward average in y over all z-levels (centre → north face).
`barotropic_filter`	function	Implicit external-gravity-wave (barotropic) mode filter.
`bernoulli_potential`	function	Bernoulli potential at T-points on an Arakawa C-grid.
`beta_param`	function	Meridional gradient of the Coriolis parameter from latitude.
`boundary_ring`	function	Wet cells adjacent to at least one dry cell (the inner ring).
`build_capacitance_solver`	function	Pre-compute the capacitance matrix and return a ready-to-use solver.
`build_coupling_matrix`	function	Build the tridiagonal vertical coupling (stratification) matrix A.
`build_multigrid_solver`	function	Build a multigrid solver with precomputed level hierarchies.
`buoyancy`	function	Compute buoyancy from density.
`cartesian_area_mean`	function	Cartesian area-weighted mean.
`cartesian_area_sum`	function	Cartesian area-weighted sum (:func:`area_sum` with Cartesian grid).
`cartesian_area_weights`	function	T-cell areas on a Cartesian grid.
`cartesian_volume_mean`	function	Cartesian volume-weighted mean.
`cartesian_volume_sum`	function	Cartesian volume-weighted sum.
`cartesian_volume_weights`	function	T-cell volumes on a Cartesian 3-D grid: `dx · dy · dz`.
`coriolis_fn`	function	Beta-plane Coriolis field.
`coriolis_param`	function	Reference Coriolis parameter from latitude.
`dct`	function	Re-exported from `spectraldiffx`.
`dct1_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dct1_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dct2_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dct2_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dct3_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dct3_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dct4_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dct4_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dctn`	function	Re-exported from `spectraldiffx`.
`decompose_vertical_modes`	function	Eigendecompose the coupling matrix A to get Rossby radii and transforms.
`diff_x_bwd`	function	Backward difference in x (east face → centre).
`diff_x_bwd_1d`	function	Backward difference in x (east face → centre), 1-D.
`diff_x_bwd_3d`	function	Backward difference in x over all z-levels (east face → centre).
`diff_x_fwd`	function	Forward difference in x (centre → east face).
`diff_x_fwd_1d`	function	Forward difference in x (centre → east face), 1-D.
`diff_x_fwd_3d`	function	Forward difference in x over all z-levels (centre → east face).
`diff_y_bwd`	function	Backward difference in y (north face → centre).
`diff_y_bwd_3d`	function	Backward difference in y over all z-levels (north face → centre).
`diff_y_fwd`	function	Forward difference in y (centre → north face).
`diff_y_fwd_3d`	function	Forward difference in y over all z-levels (centre → north face).
`diffusion_2d`	function	Horizontal tracer diffusion tendency at T-points (flux form).
`divergence_2d`	function	Compute ∇·(u, v) at T-points on the 2-D Arakawa C-grid.
`dst`	function	Re-exported from `spectraldiffx`.
`dst1_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dst1_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dst2_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dst2_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dst3_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dst3_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dst4_eigenvalues`	function	Re-exported from `spectraldiffx`.
`dst4_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`dstn`	function	Re-exported from `spectraldiffx`.
`enforce_periodic`	function	Fill the ghost-cell ring with periodic boundary conditions.
`enstrophy`	function	Enstrophy (pointwise).
`euler_step`	function	Forward Euler: y_{n+1} = y_n + dt * f(y_n).
`fft_eigenvalues`	function	Re-exported from `spectraldiffx`.
`fft_eigenvalues_ps`	function	Re-exported from `spectraldiffx`.
`fix_boundary_corners`	function	Average corner ghost cells from their two face-adjacent ghost neighbours.
`free_slip_boundaries`	function	Fill ghost ring with symmetric reflection (free-slip wall).
`heun_step`	function	Heun (RK2) predictor-corrector.
`idct`	function	Re-exported from `spectraldiffx`.
`idctn`	function	Re-exported from `spectraldiffx`.
`idst`	function	Re-exported from `spectraldiffx`.
`idstn`	function	Re-exported from `spectraldiffx`.
`imex_ssp2_step`	function	IMEX-SSP2(2,2,2) time step.
`interior`	function	Create an array shaped like `like`, zero everywhere except the interior.
`kinetic_energy`	function	Kinetic energy at T-points (cell centers) on an Arakawa C-grid.
`layer_to_mode`	function	Transform a field from layer space to mode space.
`leapfrog_raf_step`	function	Leapfrog with Robert-Asselin filter.
`linear_2pts`	function	2-point centred reconstruction at face i+1/2.
`linear_3pts_left`	function	3-point left-biased reconstruction at face i+1/2.
`linear_3pts_right`	function	3-point right-biased reconstruction at face i+1/2.
`linear_4pts`	function	4-point centred reconstruction at face i+1/2.
`linear_5pts_left`	function	5-point left-biased reconstruction at face i+1/2.
`linear_5pts_right`	function	5-point right-biased reconstruction at face i+1/2.
`linear_6pts`	function	6-point centred reconstruction at face i+1/2.
`linear_density`	function	Compute density from temperature and salinity using a linear EOS.
`linear_density_anomaly`	function	Compute density anomaly ρ’ = ρ − ρ₀ using a linear EOS.
`linear_drag`	function	Rayleigh (linear) bottom-drag tendency :math:`-\text{coef}\cdot u`.
`linear_drag_tendency`	function	Linear (Rayleigh-type) drag tendency on one velocity component.
`linear_drho_dS`	function	Partial derivative ∂ρ/∂S for the linear EOS (constant).
`linear_drho_dT`	function	Partial derivative ∂ρ/∂T for the linear EOS (constant).
`make_multigrid_preconditioner`	function	Return a preconditioner closure that applies a single multigrid V-cycle.
`make_nystrom_preconditioner`	function	Build a randomized Nyström preconditioner for a symmetric operator.
`make_preconditioner`	function	Convenience factory that builds a preconditioner by name.
`make_spectral_preconditioner`	function	Return a spectral-solve preconditioner for PCG.
`masked_laplacian`	function	Apply the masked discrete Helmholtz operator (∇² − λ)·ψ.
`masked_mean`	function	Mean over `axis`, excluding dry cells.
`masked_moments`	function	`(mean, std)` in one call — see :func:`masked_mean`, :func:`masked_std`.
`masked_std`	function	Standard deviation over `axis`, excluding dry cells, floored at `eps`.
`mc`	function	Monotonized Central (MC) flux limiter.
`minmod`	function	Minmod flux limiter.
`mode_to_layer`	function	Transform a field from mode space to layer space.
`modify_rhs_1d`	function	Re-exported from `spectraldiffx`.
`modify_rhs_2d`	function	Re-exported from `spectraldiffx`.
`modify_rhs_3d`	function	Re-exported from `spectraldiffx`.
`multilayer`	function	Lift a 2D horizontal operator to multilayer/multimode fields.
`no_flux_boundaries`	function	Fill ghost ring with zeros for normal-flux variables.
`no_slip_boundaries`	function	Fill ghost ring with sign-flipped interior values (no-slip wall).
`okubo_weiss`	function	Okubo-Weiss parameter (pointwise).
`pad_interior`	function	Extract physical interior and re-pad to original size.
`plusminus`	function	Split velocity into positive and negative parts.
`potential_enstrophy`	function	Potential enstrophy (pointwise).
`potential_vorticity`	function	Potential vorticity (pointwise).
`potential_vorticity_multilayer`	function	Multi-layer QG potential vorticity.
`pressure_from_divergence`	function	Solve ∇²p = ∇·u for the pressure correction.
`pv_flux_arakawa_lamb`	function	Arakawa--Lamb (1981) PV flux as a grid-free free function.
`pv_inversion`	function	QG potential-vorticity inversion: solve (∇² − λ)ψ = q.
`qg_potential_vorticity`	function	QG potential vorticity at T-points (single layer).
`quadratic_drag_tendency`	function	Quadratic bottom-drag tendency on one velocity component.
`rayleigh_relaxation`	function	Sponge / relaxation tendency :math:`-\text{coef}\cdot W\cdot(x - x_\text{ref})`.
`rayleigh_tendency`	function	Rayleigh damping tendency toward a reference state.
`reconstruct`	function	Compute upwind face flux along dim using the given reconstruction method.
`reduced_gravity`	function	Compute reduced gravity between two layers.
`relative_vorticity_cgrid`	function	Relative vorticity at X-points (corners) on an Arakawa C-grid.
`rk3_ssp_step`	function	3rd-order Strong-Stability-Preserving Runge-Kutta (Shu-Osher form).
`rk4_step`	function	Classic 4th-order Runge-Kutta.
`rusanov_flux`	function	Local Lax--Friedrichs (Rusanov) numerical flux at faces along `axis`.
`semi_lagrangian_step`	function	Advect a 2D scalar field using semi-Lagrangian backtracking.
`shear_strain`	function	Shear strain at X-points (corners) on an Arakawa C-grid.
`smooth_abs`	function	Smooth approximation of `\|x\|` with a nonzero gradient at the origin.
`smooth_clamp`	function	Smooth approximation of `max(x, x_min)` (a soft lower clamp).
`smooth_max`	function	Smooth approximation of the elementwise `max(x, y)`.
`solve_cg`	function	Preconditioned Conjugate Gradient solver for symmetric linear operators.
`solve_helmholtz_2d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_3d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dct`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dct1`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dct1_1d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dct1_3d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dct2_1d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dct2_3d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dst`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dst1_1d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dst1_3d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dst2`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dst2_1d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_dst2_3d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_fft`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_fft_1d`	function	Re-exported from `spectraldiffx`.
`solve_helmholtz_fft_3d`	function	Re-exported from `spectraldiffx`.
`solve_ocean_pde`	function	Integrate an ocean PDE using diffrax.
`solve_poisson_2d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_3d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dct`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dct1`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dct1_1d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dct1_3d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dct2_1d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dct2_3d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dst`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dst1_1d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dst1_3d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dst2`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dst2_1d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_dst2_3d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_fft`	function	Re-exported from `spectraldiffx`.
`solve_poisson_fft_1d`	function	Re-exported from `spectraldiffx`.
`solve_poisson_fft_3d`	function	Re-exported from `spectraldiffx`.
`spherical_area_mean`	function	Spherical area-weighted mean.
`spherical_area_sum`	function	Spherical area-weighted sum.
`spherical_area_weights`	function	T-cell areas on a spherical grid.
`spherical_volume_mean`	function	Spherical volume-weighted mean.
`spherical_volume_sum`	function	Spherical volume-weighted sum.
`spherical_volume_weights`	function	T-cell volumes on a spherical 3-D grid.
`split_explicit_step`	function	Split-explicit barotropic/baroclinic time step.
`ssh_to_streamfn`	function	Convert sea surface height to streamfunction.
`strain_magnitude_squared`	function	Squared strain magnitude (pointwise).
`streamfn_to_ssh`	function	Convert streamfunction to sea surface height.
`streamfunction_from_vorticity`	function	Invert the vorticity–streamfunction relation ∇²ψ − λψ = ζ.
`stretching_term`	function	Cross-layer stretching term for multi-layer QG models.
`superbee`	function	Superbee flux limiter.
`sw_potential_vorticity`	function	Shallow-water potential vorticity at X-points on an Arakawa C-grid.
`sw_potential_vorticity_multilayer`	function	Multi-layer shallow-water potential vorticity at X-points.
`tensor_strain`	function	Tensor (normal) strain at T-points on an Arakawa C-grid.
`total_energy`	function	Domain-integrated total energy.
`total_enstrophy`	function	Domain-integrated enstrophy.
`upwind_1pt`	function	1st-order upwind face value at interior faces along the last axis.
`upwind_3pt`	function	3-point upwind face value at i+1/2 along the last axis.
`upwind_5pt`	function	5-point upwind face value at i+1/2 along the last axis.
`upwind_flux`	function	Compute upwind flux with automatic stencil fallback near boundaries.
`uv_center_flux`	function	Compute raw face fluxes for a cell-centred scalar on a C-grid.
`uv_node_flux`	function	Compute raw face fluxes for a node-centred tracer on a C-grid.
`van_leer`	function	Van Leer flux limiter.
`vertical_velocity`	function	Vertical velocity from the continuity equation.
`volume_mean`	function	Volume-weighted mean over physical interior T-cells (3-D).
`volume_sum`	function	Volume-weighted sum over physical interior T-cells (3-D).
`volume_weights`	function	Dispatch to Cartesian/spherical volume weights based on grid type.
`vorticity_flux_upwind`	function	Dissipative upwind vortex-force fluxes on a C-grid.
`wall_boundaries`	function	Apply wall boundary conditions appropriate for a C-grid staggering.
`weno_3pts`	function	3-points non-linear left-biased stencil reconstruction:
`weno_3pts_improved`	function	3-points non-linear left-biased stencil reconstruction:
`weno_3pts_improved_right`	function	3-points non-linear right-biased WENO-Z stencil reconstruction at face q0+1/2:
`weno_3pts_right`	function	3-points non-linear right-biased stencil reconstruction at face q0+1/2:
`weno_5pts`	function	5-points non-linear left-biased stencil reconstruction
`weno_5pts_improved`	function	5-points non-linear left-biased stencil reconstruction
`weno_5pts_improved_right`	function	5-points non-linear right-biased WENO-Z stencil reconstruction at face q0+1/2:
`weno_5pts_right`	function	5-points non-linear right-biased stencil reconstruction at face q0+1/2:
`weno_7pts`	function	7-points non-linear left-biased stencil reconstruction
`weno_7pts_right`	function	7-points non-linear right-biased stencil reconstruction at face q0+1/2:
`weno_9pts`	function	9-points non-linear left-biased stencil reconstruction
`weno_9pts_right`	function	9-points non-linear right-biased stencil reconstruction at face q0+1/2:
`wind_stress_tendency`	function	Momentum tendency from a surface wind stress.
`zero_boundaries`	function	Fill ghost ring with zeros (homogeneous Dirichlet).
`zero_gradient_boundaries`	function	Fill ghost ring by copying nearest interior value (zero gradient / Neumann).

Upstream: `spectraldiffx`¶

Pseudospectral differentiation, filters and PDE solvers.

`spectraldiffx`¶

Name	Kind	What it does
`ChebyshevDerivative1D`	class	1D Chebyshev derivative operator using the precomputed differentiation matrix.
`ChebyshevDerivative2D`	class	2D Chebyshev derivative operators on [−Lx, Lx] × [−Ly, Ly].
`ChebyshevDerivative3D`	class	3D Chebyshev derivative operators on [−Lz, Lz] × [−Ly, Ly] × [−Lx, Lx].
`ChebyshevFilter1D`	class	1D Chebyshev spectral filter for smoothing and numerical stabilisation.
`ChebyshevFilter2D`	class	2D Chebyshev spectral filter on [−Lx, Lx] × [−Ly, Ly].
`ChebyshevGrid1D`	class	1D Chebyshev grid on the domain [-L, L].
`ChebyshevGrid2D`	class	2D Chebyshev grid on [-Lx, Lx] × [-Ly, Ly].
`ChebyshevGrid3D`	class	3D Chebyshev grid on [-Lz, Lz] × [-Ly, Ly] × [-Lx, Lx].
`ChebyshevHelmholtzSolver1D`	class	1D Chebyshev-collocation Helmholtz/Poisson solver with Dirichlet or Neumann BCs.
`ChebyshevHelmholtzSolver2D`	class	2D Chebyshev-collocation Helmholtz/Poisson solver with Dirichlet BCs.
`ChebyshevPoissonSolver1D`	class	1D Chebyshev Poisson solver: d²u/dx² = f on [−L, L].
`ChebyshevPoissonSolver2D`	class	2D Chebyshev Poisson solver: ∇²u = f with Dirichlet BCs.
`ChebyshevTransform1D`	class	Forward/inverse 1D Chebyshev transform.
`ChebyshevTransform2D`	class	Forward/inverse 2D Chebyshev transform (tensor product of 1D).
`FourierGrid1D`	class	Fourier grid setup for periodic domain.
`FourierGrid2D`	class	2D Fourier grid for doubly periodic domain.
`FourierGrid3D`	class	3D Fourier grid for triply periodic domain.
`SpectralDerivative1D`	class	1D Spectral derivative operator using the Fast Fourier Transform (FFT).
`SpectralDerivative2D`	class	2D Spectral derivative operators for doubly periodic rectangular domains.
`SpectralDerivative3D`	class	3D Spectral derivative operators for triply periodic domains.
`SpectralFilter1D`	class	1D Spectral filter for smoothing and numerical stabilization.
`SpectralFilter2D`	class	2D Spectral filter for doubly periodic domains.
`SpectralFilter3D`	class	3D Spectral filter for triply periodic domains.
`SphericalDerivative1D`	class	1D spectral derivative operator on Gauss-Legendre latitude grid.
`SphericalDerivative2D`	class	2D pseudo-spectral derivative operators on the sphere.
`SphericalDivergenceInversionSolver`	class	Divergence-inversion solver on the sphere.
`SphericalFilter1D`	class	1D spectral filter on the Gauss–Legendre latitude grid.
`SphericalFilter2D`	class	2D spectral filter on the full sphere lat-lon grid.
`SphericalGrid1D`	class	Gauss-Legendre grid for the latitude (colatitude) direction.
`SphericalGrid2D`	class	Full sphere latitude-longitude pseudo-spectral grid.
`SphericalHarmonicTransform`	class	Full 2D Spherical Harmonic Transform (SHT).
`SphericalHelmholtzDecomposition`	class	Helmholtz decomposition of a horizontal vector field on the sphere.
`SphericalHelmholtzSolver`	class	Spectral Helmholtz solver on the sphere: (∇² − α) φ = f.
`SphericalPoissonSolver`	class	Spectral Poisson solver on the sphere: ∇²φ = f.
`SphericalVorticityInversionSolver`	class	Vorticity-inversion solver on the sphere.
`build_capacitance_solver`	function	Pre-compute the capacitance factorisation and return a solver.
`cheb_dealias_product`	function	Compute the dealiased pointwise product a·b on a Chebyshev grid.
`chebyshev_antiderivative_coeffs`	function	Indefinite integral of a Chebyshev series, vanishing at x = −L.
`chebyshev_derivative_coeffs`	function	Differentiate a Chebyshev series in coefficient space (last axis).
`chebyshev_integral_coeffs`	function	Definite integral ∫_{−L}^{L} u(x) dx from Chebyshev coefficients.
`clenshaw_curtis_integrate_1d`	function	Integrate a 1D nodal field over [−L, L] using Clenshaw–Curtis.
`clenshaw_curtis_integrate_2d`	function	Integrate a 2D nodal field over [−Lx, Lx] × [−Ly, Ly].
`clenshaw_curtis_weights`	function	Clenshaw–Curtis quadrature weights on Gauss–Lobatto nodes of [−L, L].

Capability index: GeoNumerics

Upstream: gaussx¶

gaussx¶

Upstream: finitevolx¶

finitevolx¶

Upstream: spectraldiffx¶