opifex.uncertainty API Reference

The opifex.uncertainty package provides predictive-distribution containers, loss-component aggregation, pure-JAX Bayesian kernels, trainable Bayesian layers, structural protocols for UQ-aware modules, a capability registry, and adapter/backend interfaces for downstream inference engines.

Module layout

Module	Purpose
opifex.uncertainty.types	Predictive value objects: PredictiveDistribution, PredictionInterval, PredictionSet, PredictiveMode.
opifex.uncertainty.objectives	ObjectiveConfig (loss weights, dataset metadata) and UQLossComponents (optimizer-facing loss decomposition); scale_kl(...) helper.
opifex.uncertainty.kernels.bayesian	Pure JAX Bayesian helpers: diagonal_gaussian_kl, sample_diagonal_gaussian.
opifex.uncertainty.layers.bayesian	Trainable NNX Bayesian layers: BayesianLinear, BayesianSpectralConvolution.
opifex.uncertainty.protocols	Structural protocols: UncertaintyAwareModule, VariationalModule, Calibrator, Conformalizer, UncertaintyEstimator.
opifex.uncertainty.registry	Capability metadata: UQCapability, DefaultStrategy, UQRegistry, register_uq_capability.
opifex.uncertainty.inference_backends	Backend protocol + base result/spec/diagnostics containers (InferenceBackendProtocol, BackendResult, BackendDiagnostics, InferenceBackendSpec, UnsupportedBackendError), BlackJAXBackend for HMC / NUTS / MALA posterior sampling, BackendRouter for Artifex-first family routing, and OptionalBackendSpec declarations for TFP / bijx / FlowJAX / Bayeux / NumPyro / GPJax / sbiax / flowMC / oryx / traceax / matfree / kfac-jax.
opifex.uncertainty.distributions	ArtifexDistributionAdapter, DistrAxAdapter, and the from_distribution(...) dispatcher that wraps a backend distribution into a PredictiveDistribution.
opifex.uncertainty.adapters	Distribution / model-uncertainty adapter protocols + concrete adapters. Protocols: DistributionAdapterProtocol, ModelUncertaintyAdapterProtocol. Concrete adapters: ModelUncertaintyAdapter (deterministic), DeepEnsembleAdapter, SnapshotEnsembleAdapter, SWAGAdapter, BatchEnsembleAdapter, MCDropoutAdapter, TestTimeAugmentationAdapter (forward a deterministic model over a fixed tuple of deterministic input augmentations and aggregate the cross-augmentation mean/variance — an ensemble over augmentations; regression mean+across-augmentation-variance form with aleatoric = 0 so total_uncertainty == epistemic; the classification analogue is the mutual-information form MI = H(mean) - mean(H)), BayesianLastLayerAdapter (closed-form analytic last-layer predictive), VBLLAdapter (closed-form analytic Variational-Bayesian-last-layer regression predictive; epistemic via the Cholesky L-form sum((phi @ L)**2); regression only), DUEAdapter (Deterministic Uncertainty Estimation, van Amersfoort et al. 2021 — wraps a FITTED spectral-normalized / bi-Lipschitz deep-kernel feature extractor + inducing-point SVGP and REUSES opifex.uncertainty.gp.predict_svgp for the single-forward-pass predictive over the features; re-tags the SVGP predictive's method to due and records num_inducing; the feature/bi-Lipschitz training is upstream in opifex.neural.operators.specialized.spectral_normalization), SNGPAdapter (Spectral-normalized Neural Gaussian Process last layer, Liu et al. 2020 — wraps a FITTED fixed random-Fourier-feature map phi(x) + last-layer weights + Laplace precision matrix and emits the closed-form regression predictive mean = phi @ output_weights with epistemic variance from the edward2 Cholesky-solve diagonal ridge * sum((cholesky(precision)^-1 phi^T)**2); REUSES the shared Gaussian-linear-head assembly; the Laplace precision is built by fit_sngp_precision, a faithful edward2 port; classification mean-field-logit adjustment is exposed as sngp_mean_field_logits but not applied in the regression predict path). Fitted-state pytrees: DeepEnsembleState, SnapshotEnsembleState, SWAGState, BatchEnsembleState, MCDropoutState, TestTimeAugmentationState (static model_fn + augmentations tuple), BayesianLastLayerState, VBLLState, DUEState (static feature_fn + nested fitted SVGPState), SNGPState (static feature_fn + ridge_penalty + observation_noise_variance; leaves output_weights + precision_matrix). Capability spec dataclasses: DistributionAdapterSpec, OperatorAdapterSpec, the FNO operator adapter family (FNOConformalAdapterSpec, FNODeepEnsembleAdapterSpec, FNOMCDropoutAdapterSpec), the DeepONet operator adapter family (DeepONetConformalAdapterSpec, DeepONetDeepEnsembleAdapterSpec, DeepONetMCDropoutAdapterSpec), The concrete diagonal-Laplace adapter (LaplaceAdapterSpec + LaplaceState) now lives in opifex.uncertainty.curvature alongside the curvature kernels it consumes.
opifex.uncertainty.curvature	Curvature primitives for second-order UQ: hessian_vector_product, ggn_vector_product, empirical_fisher_diagonal, diagonal_laplace_posterior + DiagonalLaplacePosterior. Diagonal-Laplace adapter: LaplaceAdapterSpec + LaplaceState. Linearised Neural Operator predictive: linearized_neural_operator_posterior(model_fn=, laplace_posterior=, x=) returning a PredictiveDistribution whose marginal variance is diag(J Σ Jᵀ) for Σ = diag(1 / precision_diagonal) (Magnani et al. 2024, arXiv:2406.04317). Pure-JAX; passes jit / vmap smokes.
opifex.uncertainty.tempering	Diffusion tempering schedule (Beck & Tronarp+ 2024, arXiv:2402.12231): DiffusionTemperingSchedule (Pattern-A frozen-slotted-kw-only dataclass) + TemperingScheduleType enum. Anneals physics_weight per training epoch via objective_config_at(epoch=, base_config=) which rebuilds the ObjectiveConfig with physics_weight = base / σ(epoch)². Three shapes: linear, exponential (log-linear), cosine. Pure-JAX; integrates into RobustPINNOptimizer.compute_loss_components with no new trainer class.
opifex.uncertainty.gp	Exact conjugate-Gaussian GP regression (Rasmussen & Williams 2006 Algorithm 2.1): rbf_kernel(x1, x2, *, lengthscale, output_scale) (squared-exponential), fit_exact_gp(*, x_train, y_train, lengthscale, output_scale, noise_std, kernel_fn) returning ExactGPState (Cholesky factor + pre-solved α), and predict_exact_gp(*, state, x_test) returning a PredictiveDistribution whose marginal variance follows the Algorithm-2.1 Cholesky-back-substitution recipe var = K(X*, X*) - vᵀv with v = L \ K(X, X*). Non-conjugate Laplace (RW06 §3.4 Alg. 3.1): generic fit_laplace_gp(*, log_likelihood_components_fn, …) returning LaplaceGPState plus predict_laplace_latent_moments(*, state, x_test). Any factorising (log_lik, ∇log_lik, W, √W) quadruple plugs in; shipped wrappers cover Bernoulli (binary classification, MacKay probit), Poisson (counts, exp link, log-normal predictive), Student-t (robust regression, Fisher-info W), and Beta (proportions, logit link, Fisher-info W). Higher-order OAK (Lu+ 2022 ICML): orthogonal_additive_kernel(*, base_kernel_fns, max_order, order_variances) evaluates the full ANOVA decomposition K = Σ_ℓ σ²_ℓ e_ℓ(k̃_1,…,k̃_D) in O(D · max_order) via Newton-Girard recursion, with constrained_rbf_kernel(*, input_mean, input_std) providing the canonical Gaussian-measure-orthogonal base. Celerite family (Foreman-Mackey+ 2017 / Foreman-Mackey 2018): celerite_complex_kernel(*, sine_amplitude, frequency) for the building-block term k(τ) = exp(-c τ)[a cos(d τ) + b sin(d τ)]; the Real term k(τ) = a exp(-c τ) is matern12_kernel; arbitrary sums realised via kernel_sum(*, kernel_fns). The underdamped SHO damped_oscillator_kernel(*, quality_factor) is one specific Complex-term reparametrisation. Scalable SHO via state-space Kalman: fit_quasisep_sho_gp(*, x_train, y_train, …, quality_factor) returns QuasisepGPState; predict_quasisep_sho_gp(*, state, x_test) uses the 2-state SHO SDE + opifex's tested kalman_filter / kalman_smoother for O(n) exact fit + predict — coincident with the direct damped_oscillator_kernel exact GP up to numerical precision. Stochastic SVGP for non-Gaussian likelihoods (Hensman+ 2013/2015): init_stochastic_svgp_state(*, x_inducing, lengthscale, output_scale, log_likelihood_fn, …) returns StochasticSVGPState in the whitened (μ_w, L_w) parametrisation; stochastic_svgp_elbo(*, state, x_batch, y_batch, dataset_size, num_quadrature_points) runs the minibatched ELBO with a single chol(K_zz) per call and a closed-form KL with zero K_zz operations (vs GPJax's two Cholesky factorisations + per-point vmap); predict_stochastic_svgp(*, state, x_test) returns the closed-form whitened predictive marginals. Built-in bernoulli_log_likelihood + poisson_log_likelihood per-observation helpers; users can plug in any factorising likelihood via the log_likelihood_fn callable. Natural-gradient updates (Salimbeni+ 2018): natural_gradient_step(*, state, x_batch, y_batch, dataset_size, learning_rate, num_quadrature_points) runs one Fisher-preconditioned step in the whitened-natural-parameter space — 10x-100x faster convergence than Adam on the variational params per Salimbeni+ 2018 §4. Sparse multi-output SVGP (Bonilla+ 2007 / Alvarez & Lawrence 2009): fit_svgp(*, …, kernel_fn=multi_output_icm_kernel(…)) or fit_svgp(*, …, kernel_fn=multi_output_lcm_kernel(components=…)) plus predict_svgp(*, state, x_test) work with the existing Titsias collapsed surface — the per-point K_{xx} diagonal is computed via jax.vmap so ICM/LCM task-dependent diagonals are handled transparently. Pure-JAX; passes jit smoke. Companion to the user-installed GPJaxAdapterSpec / TinygpAdapterSpec / MarkovflowAdapterSpec / BayesnewtonAdapterSpec / KalmanJaxAdapterSpec advertisements in opifex.uncertainty.adapters.gp.
opifex.uncertainty.likelihoods	Backend-neutral log-likelihoods: Gaussian, heteroscedastic Gaussian, Laplace, Student-t, mixture.
opifex.uncertainty.priors	Diagonal-Gaussian log prior.
opifex.uncertainty.losses	Uncertainty-aware loss functions: PointwiseQuantileLoss (pinball loss for UQNO residual-quantile training, mirrors neuralop.losses.data_losses.PointwiseQuantileLoss).
opifex.uncertainty.metrics	Ensemble + interval scoring metrics with no canonical CalibraX home: predictive_entropy(ensemble_probabilities) (Gal & Ghahramani 2016), mutual_information(ensemble_probabilities) (BALD epistemic decomposition, Houlsby et al. 2011), interval_score/winkler_score(lower, upper, targets, alpha) (Gneiting & Raftery 2007 strictly proper rule). Pure jax.Array kernels; no forward re-exports — Brier/ECE/pinball/Gaussian-NLL stay in opifex.uncertainty.calibration.
opifex.uncertainty.forecasting_metrics	Probabilistic forecast metrics: empirical crps (CalibraX-canonical formula), fair_crps (Ferro 2014; WeatherBenchX CRPSEnsemble(fair=True) reference), energy_score (Gneiting & Raftery 2007 §4.2), rank_histogram (Hamill 2001), spread_skill_ratio (Fortin et al. 2014), pit_histogram (Diebold-Gunther-Tay 1998), ranked_probability_score (Epstein 1969), event_reliability (Murphy 1973 Brier decomposition).
opifex.uncertainty.scientific.domain_metrics	Domain reliability metrics with the DomainMetricSummary value object: PINN (physics_residual_coverage, boundary_condition_coverage), neural operator (spectral_coverage), quantum chemistry (chemical_accuracy_coverage — caller-supplied tolerance, no hard-coded default), optimization / L2O (regret_interval_summary, feasibility_coverage), assimilation (sensor_reliability_summary — reduced χ²), parameter inference (parameter_credible_interval_coverage). It also provides reliability summaries for the three formerly-deferred surfaces — likelihood-free (likelihood_free_rank_calibration, the simulation-based-calibration expected-coverage error over SBC ranks), active learning (active_learning_acquisition_reliability, the rank correlation between acquisition score and realized error, reusing calibrax), and PAC-Bayes (pac_bayes_bound_validity, certified-bound validity and tightness against a held-out risk) — exposed as the LIKELIHOOD_FREE_RELIABILITY, ACTIVE_LEARNING_RELIABILITY, PAC_BAYES_RELIABILITY UQCapability constants. The underlying methods ship in the opifex.uncertainty.sbi, opifex.uncertainty.active and opifex.uncertainty.pac_bayes subpackages (documented below).
opifex.uncertainty.scientific.quantum	Quantum-chemistry UQ surfaces. Three ensemble-aggregating callables — EnergyUncertainty(*, method, units="hartree") (total / per-state energies), DensityUncertainty(*, grid_axes) (per-grid-point electron density), ExchangeCorrelationUncertainty(*, functional_family) (XC-functional outputs) — each take an ensemble of predictions stacked along axis 0 (shape (num_members, ...)) and return a PredictiveDistribution with mean = ensemble.mean(0) and across-member variance reported as both the marginal variance and the epistemic / total_uncertainty decomposition (aleatoric = 0), mirroring the deep-ensemble member aggregation via the shared opifex.uncertainty._predictive.ensemble_predictive factory (Rule 1 — DRY); metadata records method/units/num_members (energy), grid_axes (density), functional_family (XC). ChemicalAccuracyCoverage(*, tolerance_hartree) returns the empirical fraction mean(|prediction - reference| <= tolerance_hartree); the module constant CHEMICAL_ACCURACY_HARTREE ≈ 0.0015936 is 1 kcal/mol in Hartree — the conventional quantum-chemistry "chemical accuracy" benchmark (Pople, Rev. Mod. Phys. 71, 1267 (1999)) — and the ChemicalAccuracyCoverage.chemical_accuracy() classmethod pins that tolerance. Pure-JAX; jit / grad / vmap-compatible.
opifex.uncertainty.scientific.fields	Canonical FieldMetadata (Pattern A) schema plus field/function-space UQ metrics: spatial_calibration_error, function_space_l2_coverage, conservation_law_residual_summary, residual_uncertainty_alignment. FieldMetadata is the canonical home referenced by opifex.uncertainty.conformal.fields.FieldSplitConformalRegressor.
opifex.uncertainty.scientific.equation_discovery	Bayesian SINDy: regularized-horseshoe posterior over the SINDy candidate-library coefficients (Hirsh, Barajas-Solano & Kutz 2021 arXiv:2107.02107; Piironen & Vehtari 2017). BayesianSINDy(library, *, tau0=0.1, slab_nu=4, slab_s=2, noise_lambda=1.0, num_warmup=200, num_samples=400) ports pysindy's SBR numpyro model onto a BlackJAX-NUTS log-density sampled in unconstrained (log) space with change-of-variables Jacobians; fit(x, x_dot, *, rngs) returns PosteriorOverTerms (β / τ / σ draws + static feature_names). coefficients() is the posterior-mean β; term_inclusion_probabilities() returns per-term P(|βⱼ| > threshold) with a scale-relative threshold (10% of the largest per-target posterior-mean magnitude, floored at 1e-3) since the horseshoe is a continuous shrinkage prior with no exact zeros; coefficient_posterior_intervals(level=0.95) returns a PredictionInterval from posterior quantiles. Scope: the error model is on the derivatives x_dot (no ODE integration), matching pysindy's documented SBR simplification — σ is a derivative-noise scale; the integrated Hirsh eq. 2.4 model is a known upstream TODO. Reuses opifex.discovery.sindy.library.CandidateLibrary for Θ.
opifex.uncertainty.reports	UQReliabilityReport — Pattern-B aggregated report carrying optional metric leaves from every UQ subsystem (calibration, conformal, forecasting, OOD, selective). validate() requires at least one populated metric; to_dict() returns a deterministic JSON-compatible payload. Failed exchangeability / shift / OOD warnings propagate verbatim into metadata["assumption_warnings"] and metadata["assumption_status"]. Data class + serializer only — never an evaluator.
opifex.uncertainty.monitoring	MonitoringInputs (Pattern A provenance container) + build_reliability_report(inputs=, ...metrics) builder that assembles a validated UQReliabilityReport from already-computed metrics. Validates the payload before returning so half-populated reports never propagate silently.
opifex.uncertainty.selective	Selective prediction: risk_coverage_curve(confidences, errors) returns the per-threshold (coverages, risks) arrays; area_under_risk_coverage(confidences, errors) returns the AURC scalar; abstention_decision(confidences, threshold) returns an AbstentionDecision value object with named accepted_mask / rejected_mask + metadata (Geifman & El-Yaniv 2017, arXiv:1705.08500).
opifex.uncertainty.ood	Out-of-distribution detection scores and residual-shift diagnostics: max_softmax_probability (Hendrycks & Gimpel 2017, arXiv:1610.02136), fpr95, ShiftReport + residual_shift_diagnostic (reuses opifex.uncertainty.conformal.ks_two_sample_pvalue). AUROC / AUPRC are NOT re-exported — import directly from calibrax.metrics.functional.classification.{roc_auc, average_precision}. Predictive entropy / mutual information OOD signals live in opifex.uncertainty.metrics.
opifex.uncertainty.calibration	Calibration metrics and calibrators. Metrics: gaussian_nll, picp, mpiw, regression_calibration_error (Opifex-local), plus thin wrappers around CalibraX functionals: brier_score, expected_calibration_error, pinball_loss. Calibrator: TemperatureScaling + TemperatureScalingState (Guo et al. 2017 temperature scaling for multiclass logits; single-scalar L-BFGS fit on validation NLL).
opifex.uncertainty.conformal	Conformal prediction subsystem. Score helpers: absolute_residual_score (Lei et al. 2018), cqr_score (Romano, Patterson, Candes 2019, arXiv:1905.03222), conformal_quantile (finite-sample ceil((n+1)(1-α))/n rank). Regression calibrators: SplitConformalRegressor + SplitConformalState, ConformalizedQuantileRegressor + CQRState, GroupedSplitConformalRegressor + GroupedSplitConformalState (return PredictionInterval). Classification scores and calibrator: lac_score (Sadinle et al. 2019), aps_score (Romano et al. 2020), raps_score (Angelopoulos et al. 2021, arXiv:2009.14193), aps_prediction_set, LACConformalClassifier + LACConformalState (return PredictionSet). Cross-conformal / weighted: jackknife_plus_intervals, cv_plus_intervals (Barber et al. 2021 arXiv:1905.02928), weighted_conformal_quantile, weighted_split_conformal_intervals (Tibshirani et al. 2019 arXiv:1904.06019). Time-series: EnbPIState + enbpi_update/enbpi_predict (Xu & Xie 2021 arXiv:2010.09107), AdaptiveConformalState + aci_update/aci_metadata (Gibbs & Candès 2021 arXiv:2106.00170). Field/function-space: field_l2_score, field_linf_score, field_h1_score, FieldSplitConformalRegressor + FieldSplitConformalState with explicit norm over spatial axes. Risk control: RiskControlConfig (Pattern A), RiskControllerState (Pattern B), hoeffding_upper_bound, rcps_threshold_kernel, select_threshold_rcps, bootstrap_threshold_ci — RCPS / Learn-then-Test threshold selection (Bates et al. 2021 arXiv:2101.02703, Angelopoulos et al. 2022 arXiv:2110.01052) with calibrax.statistics.analyzer.StatisticalAnalyzer.bootstrap_ci reused for CI reporting. Exchangeability diagnostic: check_exchangeability, ks_two_sample_pvalue, ExchangeabilityReport (Vovk et al. 2005 §2.4). Numerical core for split/CQR/LAC/APS/RAPS aligned with the canonical aangelopoulos/conformal-prediction reference notebooks ('higher' interpolation rule in conformal_quantile).

Trainable Bayesian Layers

BayesianLinear

Variational diagonal-Gaussian dense layer; weights and bias each carry a (mean, log-variance) posterior. Sampling uses the reparameterization trick and requires caller-owned RNG ownership at call time.

import jax.numpy as jnp
from flax import nnx
from opifex.uncertainty import BayesianLinear

layer = BayesianLinear(in_features=4, out_features=3, prior_std=1.0, rngs=nnx.Rngs(0))

# Posterior-mean mode (no sampling). Mode resolution follows the
# canonical ``nnx.Dropout`` convention: per-call ``deterministic``
# overrides ``self.deterministic`` set by the NNX inference toggle.
x = jnp.ones((2, 4))
mean_output = layer(x, deterministic=True)

# Stochastic mode — caller owns the RNG
rngs = nnx.Rngs(posterior=42)
sample = layer(x, deterministic=False, rngs=rngs)
also_sample = layer(x, deterministic=False, rngs=rngs)  # different (stream advanced)

# Or pass an explicit JAX key
import jax
key = jax.random.PRNGKey(7)
deterministic_a = layer(x, deterministic=False, rngs=key)
deterministic_b = layer(x, deterministic=False, rngs=key)  # same (same key)

# KL divergence — used by the shared `loss_components` / `negative_elbo`
# objectives on modules built atop these layers (e.g. ProbabilisticPINN,
# UncertaintyQuantificationNeuralOperator). Prefer those built-in
# objectives over assembling ``data_loss + kl_weight * kl`` by hand.
kl = layer.kl_divergence()  # scalar; sums weight + bias KLs

BayesianSpectralConvolution

Variational Fourier-spectral convolution with complex Gaussian weights (1D and 2D modes). Implements the canonical Zongyi Li FNO spectral block (Li et al. 2021, arXiv:2010.08895) with diagonal-Gaussian posteriors over the complex weights:

• 1D — one weight tensor of shape (out, in, modes[0]).
• 2D — two weight tensors of shape (out, in, modes[0], modes[1] // 2 + 1) covering positive and negative H-axis low-frequency bands respectively.

Same RNG-safety contract as BayesianLinear.

import jax.numpy as jnp
from flax import nnx
from opifex.uncertainty import BayesianSpectralConvolution

# 2D spectral conv (Fourier neural operator-style block)
layer = BayesianSpectralConvolution(
    in_channels=2, out_channels=3, modes=(8, 8), prior_std=1.0, rngs=nnx.Rngs(0)
)

# Input: (batch, in_channels, height, width)
x = jnp.ones((1, 2, 32, 32))
mean_output = layer(x, sample=False)         # (1, 3, 32, 32)
sample = layer(x, sample=True, rngs=nnx.Rngs(posterior=0))

# KL over real + imaginary Fourier weights
kl = layer.kl_divergence()

Inference Backends

BackendRouter

Selects the highest-priority available backend for a given family ("flow" / "sampler" / "distribution") using Artifex-first resolution order. Optional backends remain available by name but raise ImportError on instantiation when their package is not installed.

from opifex.uncertainty import BackendRouter

router = BackendRouter()
sampler_spec = router.resolve("sampler")               # default = BlackJAX
flow_spec = router.resolve("flow")                     # default = Artifex RealNVP
dist_spec = router.resolve("distribution")             # default = Artifex Distribution
print(sampler_spec.name, flow_spec.name, dist_spec.name)

Listing all registered specs for a family:

from opifex.uncertainty import BackendRouter

router = BackendRouter()
for spec in router.available("sampler"):
    status = "installed" if spec.probe() else "needs install"
    print(f"  {spec.name} ({spec.source_package}) — {status}")

BlackJAXBackend

Thin adapter over Artifex's BlackJAX HMC / NUTS / MALA samplers conforming to InferenceBackendProtocol. Supports method in {"hmc", "nuts", "mala"}; other sampler families raise UnsupportedBackendError.

import jax.numpy as jnp
from flax import nnx
from opifex.uncertainty import BlackJAXBackend


def log_density(theta):
    return -0.5 * jnp.sum(theta * theta)


backend = BlackJAXBackend(
    target_log_prob=log_density,
    init_state=jnp.zeros(10),
    n_samples=1000,
    n_burnin=500,
    method="nuts",
    step_size=0.1,
)
result = backend.fit(log_density, rngs=nnx.Rngs(sample=0))
samples = result.sampler_state             # (n_samples, ...)
diagnostics = result.diagnostics           # BackendDiagnostics

# Convert posterior samples into a predictive distribution.
predictive = backend.predict_distribution(jnp.zeros((4, 10)), rngs=nnx.Rngs(sample=1))

Distribution Adapters

from opifex.uncertainty import from_distribution
from artifex.generative_models.core.distributions.continuous import Normal
from flax import nnx

# Wrap an Artifex distribution.
dist = Normal(loc=jnp.zeros(3), scale=jnp.ones(3), rngs=nnx.Rngs(0))
predictive = from_distribution(dist)       # PredictiveDistribution

# Distrax-like objects (anything exposing sample/log_prob/mean/variance) are
# accepted as a secondary fallback.

from_distribution resolves Artifex Distribution first, then falls back to Distrax-like objects. Unsupported objects raise TypeError.

Predictive Value Objects

import jax.numpy as jnp
from opifex.uncertainty import (
    PredictiveDistribution,
    PredictionInterval,
    PredictionSet,
    PredictiveMode,
)

mean = jnp.zeros((4,))
variance = jnp.ones((4,))
pd = PredictiveDistribution(mean=mean, variance=variance)
std = pd.std()  # sqrt(variance)

PredictiveDistribution is a JAX pytree (round-trips through jit / vmap) with optional samples, covariance, epistemic, aleatoric, total_uncertainty, quantiles, interval, and prediction_set fields, plus tuple-of-pairs metadata.

SolutionDistribution (solver-side, per-field)

import jax.numpy as jnp
from opifex.uncertainty.scientific import SolutionDistribution

sd = SolutionDistribution(
    mean={"u": jnp.zeros((64,)), "p": jnp.ones((64,))},
    epistemic={"u": jnp.full((64,), 0.1), "p": jnp.full((64,), 0.2)},
    aleatoric={"u": jnp.full((64,), 0.3), "p": jnp.full((64,), 0.4)},
    total_uncertainty={"u": jnp.full((64,), 0.4), "p": jnp.full((64,), 0.6)},
    metadata=(
        ("uncertainty_sources", ("numerical", "parameter")),
        ("spatial_axes", (0,)),
        ("function_space_norm", "L2"),
        ("covariance_form", "diag"),
    ),
)
sd.validate()  # public preflight; never called during pytree unflatten

# Project a single field back onto the canonical PredictiveDistribution
# contract so downstream calibration / conformal code consumes either
# container without solver-specific awareness.
pd_u = sd.as_predictive_distribution("u")

SolutionDistribution is the multi-field solver counterpart of PredictiveDistribution. Every uncertainty leaf (samples, variance, covariance, epistemic, aleatoric, total_uncertainty) is a dict[str, jax.Array] keyed by field name; metadata is static aux_data and must declare a tuple-of-strings uncertainty_sources drawn from {"numerical", "parameter", "observation", "model_discrepancy", "ensemble", "calibration"}. The same _VARIANCE_RTOL / _VARIANCE_ATOL tolerances as PredictiveDistribution apply to per-field variance additivity, so a SolutionDistribution round-tripped through as_predictive_distribution(field) does not flap downstream additivity checks.

aggregate_solver_solutions (solver-side Monte-Carlo / ensemble aggregation)

from opifex.uncertainty.scientific import aggregate_solver_solutions

# Caller writes the explicit replay or ensemble loop.
replays = [solver.solve(problem, rngs=nnx.Rngs(seed)) for seed in range(num_samples)]
out = aggregate_solver_solutions(
    replays,
    quantiles=(0.05, 0.95),
    metadata=(("method", "monte_carlo_empirical"),),
)
mean_field_u = out.fields["u"]
band_lower = out.auxiliary_data["uq"]["quantiles"][0.05]["u"]
band_upper = out.auxiliary_data["uq"]["quantiles"][0.95]["u"]

Stacks per-field arrays across the sequence, reports the mean in Solution.fields, and stores the full UQ payload (samples, ddof=1 variance, optional quantile bands, metadata) under auxiliary_data["uq"]. Replaces four previous wrapper classes (BayesianWrapper / ConformalWrapper / EnsembleWrapper / GenerativeWrapper); deep ensemble, stochastic-replay, and empirical-interval flows all reduce to this single call.

summarize_stacked_sample_solution (generative-sampler aggregation)

from opifex.uncertainty.scientific import summarize_stacked_sample_solution

# Generative base solver returns one Solution whose fields are stacked
# (num_samples, *field_shape) arrays — typically a Diffusion or
# Flow-Matching model from Artifex.
raw = generative_solver.solve(problem)
out = summarize_stacked_sample_solution(raw, quantiles=(0.05, 0.95))

Aggregates along the sample axis without overwriting the underlying batch. Scalar fields pass through unchanged; non-scalar fields land in auxiliary_data["uq"]["samples"] and a mean lands in Solution.fields[key]. Replaces the GenerativeWrapper class.

Objectives

import jax.numpy as jnp
from flax import nnx
from opifex.uncertainty import (
    BayesianLinear,
    ObjectiveConfig,
    UQLossComponents,
)

config = ObjectiveConfig(
    kl_weight=1.0,
    dataset_size=1000,
    physics_weight=1.0,
    data_weight=1.0,
    boundary_weight=1.0,
    initial_condition_weight=1.0,
    regularization_weight=1.0,
    calibration_weight=1.0,
    conformal_weight=1.0,
    pac_bayes_weight=1.0,
)

# Stand-in loss tensors from your training step.
data_loss = jnp.array(0.5)
physics_loss = jnp.array(0.3)
layer = BayesianLinear(in_features=4, out_features=3, rngs=nnx.Rngs(0))

components = UQLossComponents.from_components(
    config=config,
    data=data_loss,
    physics_residual=physics_loss,
    kl=layer.kl_divergence(),
)

scalar_for_optimizer = components.total

Canonical KL helper

import jax.numpy as jnp
from opifex.uncertainty.kernels.bayesian import diagonal_gaussian_kl

mean = jnp.zeros(4)
logvar = jnp.zeros(4)

# Standard N(0, 1) prior delegates to Artifex gaussian_kl_divergence:
kl = diagonal_gaussian_kl(mean, logvar, prior_mean=0.0, prior_std=1.0)

# Parametric (prior_mean, prior_std) extension is provided by Opifex:
kl_parametric = diagonal_gaussian_kl(mean, logvar, prior_mean=2.0, prior_std=3.0)

Capability Registry

from opifex.uncertainty import (
    UQCapability,
    UQRegistry,
    DefaultStrategy,
    register_uq_capability,
)

cap = UQCapability(
    native_bayesian=True,
    native_nnx_module=True,
    supports_function_space=True,
    default_strategy=DefaultStrategy.BAYESIAN,
    source_package="opifex",
)

@register_uq_capability("my_model", cap)
class MyModel:
    pass

registry = UQRegistry()
assert registry.get("my_model") is cap

Protocols

The structural-typing protocols in opifex.uncertainty.protocols describe what a model surface must offer to interoperate with the rest of the UQ platform. None of them require inheritance — any class that exposes the named methods passes isinstance(model, Protocol) via @runtime_checkable.

from opifex.uncertainty.protocols import (
    UncertaintyAwareModule,
    VariationalModule,
    UncertaintyEstimator,
    Calibrator,
    Conformalizer,
)

• UncertaintyAwareModule — exposes predict_distribution(x, *, rngs) -> PredictiveDistribution; the minimum every UQ-aware model must implement so calibration / conformal code can drive it generically.
• VariationalModule — extends UncertaintyAwareModule with loss_components(batch, *, rngs, objective) -> UQLossComponents, negative_elbo(batch, *, rngs, objective) -> UQLossComponents and kl_divergence() -> jax.Array. Used by Bayesian layers and ProbabilisticPINN.
• UncertaintyEstimator — produces uncertainty arrays from raw predictions / ensemble outputs.
• Calibrator — exposes the fit(...) / predict(...) / with_state(...) triple used by calibration tools (TemperatureScaling, …).
• Conformalizer — exposes the same triple shape but the predict output is a PredictionInterval or PredictionSet. Implemented by every class under opifex.uncertainty.conformal.

Inference-Backend and Distribution-Adapter Specs

from opifex.uncertainty.inference_backends import (
    InferenceBackendProtocol,
    InferenceBackendSpec,
    BackendResult,
    BackendDiagnostics,
    UnsupportedBackendError,
)
from opifex.uncertainty.adapters import (
    DistributionAdapterProtocol,
    DistributionAdapterSpec,
    ModelUncertaintyAdapterProtocol,
)

• InferenceBackendProtocol — fit(...) -> BackendResult plus posterior_predictive(...) -> PredictiveDistribution and predict_distribution(...) -> PredictiveDistribution. Implemented by BlackJAXBackend; optional NumPyro / Bayeux / FlowJAX backends follow the same shape.
• BackendResult carries the raw sampler / fitted-state object unchanged (e.g. BlackJAXSamplerState) so callers can inspect it.
• BackendDiagnostics — Pattern-B flax.struct.dataclass(slots=True, kw_only=True) with ess, r_hat, acceptance_rate, divergences, step_size, tree_depth leaves (all optional). Survives flax.struct.replace() and pytree round-tripping.
• InferenceBackendSpec — frozen capability declaration with sampler_names: tuple[str, ...]. The router walks the tuple left-to-right and selects the first installed backend.
• DistributionAdapterProtocol and DistributionAdapterSpec — wrap backend distribution objects (Artifex Distribution, Distrax-like) into the canonical PredictiveDistribution value object via from_distribution(...).
• ModelUncertaintyAdapterProtocol — wraps deterministic / ensemble / dropout / Laplace-style models. The concrete adapters live in opifex.uncertainty.adapters.{model,ensemble,operators} and opifex.uncertainty.curvature (concrete LaplaceAdapterSpec); the remaining deferred operator spec classes (FNOConformalAdapterSpec, etc.) raise actionable NotImplementedError from .wrap() until the backend lands.

Model UQ Adapters

from opifex.uncertainty.adapters import (
    ModelUncertaintyAdapter,
    DeepEnsembleAdapter,
    SnapshotEnsembleAdapter,
    SWAGAdapter,
    BatchEnsembleAdapter,
    MCDropoutAdapter,
    DeepEnsembleState,
    SnapshotEnsembleState,
    SWAGState,
    BatchEnsembleState,
    MCDropoutState,
)

• ModelUncertaintyAdapter — wraps a single deterministic callable. Rejects any non-DefaultStrategy.DETERMINISTIC capability claim so a raw point estimator cannot silently advertise epistemic uncertainty.
• DeepEnsembleAdapter — aggregates mean / variance over a fixed member tuple (DeepEnsembleState).
• SnapshotEnsembleAdapter — averages cyclic-LR snapshots of a single training run (SnapshotEnsembleState; Huang et al. ICLR 2017).
• SWAGAdapter — samples weights from the low-rank-plus-diagonal SWAG Gaussian N(θ_SWA, ½(Σ_diag + Σ_lowrank)), forwards each draw, and aggregates predictive mean / variance (SWAGState; caller-owned nnx.Rngs at predict time; Maddox et al. NeurIPS 2019).
• BatchEnsembleAdapter — forwards over M rank-1 fast-weight members y_m = ((x ∘ r_m) W) ∘ s_m of a shared kernel (BatchEnsembleState; Wen et al. ICLR 2020).
• MCDropoutAdapter — caller-owned nnx.Rngs at predict time; deterministic at mode="deterministic", stochastic at mode="bayesian".

The neural-operator family ships parallel specs in opifex.uncertainty.adapters.operators (OperatorAdapterSpec, FNOConformalAdapterSpec, FNODeepEnsembleAdapterSpec, FNOMCDropoutAdapterSpec, DeepONetConformalAdapterSpec, …) that declare operator-family axes (spatial_axes, spectral_axes), supported metrics, and required capabilities.

Likelihoods

import jax.numpy as jnp
from opifex.uncertainty.likelihoods import (
    gaussian_log_likelihood,
    heteroscedastic_gaussian_log_likelihood,
    laplace_log_likelihood,
    student_t_log_likelihood,
    mixture_log_likelihood,
    LikelihoodSpec,
)

ll = gaussian_log_likelihood(
    predictions=jnp.zeros((10,)),
    targets=jnp.zeros((10,)),
    scale=1.0,
)

Each helper returns the per-sample log-likelihood as a jax.Array. The heteroscedastic Gaussian takes a per-sample scale array; Student-t takes a df parameter; the mixture form takes (weights, means, scales) triplets and returns the log-sum-exp likelihood. LikelihoodSpec is the Pattern-A frozen-dataclass capability container that registers a likelihood family with the UQ registry.

Priors

from opifex.uncertainty.priors import diagonal_gaussian_log_prior, PriorSpec

log_prior = diagonal_gaussian_log_prior(
    params=jnp.zeros((4,)),
    prior_mean=0.0,
    prior_std=1.0,
)

diagonal_gaussian_log_prior returns the sum-over-features log-prior under a diagonal-Gaussian N(prior_mean, prior_std² I). PriorSpec records the prior family (name, family, parameter_names) for the registry, mirroring LikelihoodSpec.

Physics-Informed Priors

opifex.uncertainty.priors_physics exposes five NNX-based prior modules for physics-aware Bayesian modelling:

• PhysicsInformedPriors — hard-constraint projector with conservation-law and boundary-condition application.
• ConservationLawPriors — uncertainty modifier that inflates uncertainty based on per-law violation magnitude.
• DomainSpecificPriors — domain-tailored parameter ranges and distribution families (quantum chemistry, fluid dynamics, …).
• HierarchicalBayesianFramework — multi-level uncertainty propagation across hierarchy levels (multiplicative or additive modes).
• PhysicsAwareUncertaintyPropagation — confidence adjustment under physics-constraint violations.

Each module follows the canonical NNX convention: parameters declared as nnx.Param, indexed with [...] (not the deprecated .value property), pure-array methods trace under jax.jit / jax.grad.

Probabilistic-numerics and linear algebra

Matrix-free estimators and decompositions backing scalable Bayesian inference.

Matrix-free linear algebra primitives for uncertainty quantification.

Vendored JAX-native implementations of Krylov decompositions, stochastic trace and diagonal estimators, low-rank approximations, randomized SVD, matrix functions, log-determinant integrands, differentiable least-squares, and higher-moment trace UQ.

Pure JAX; no NNX imports anywhere in this subpackage. Each algorithm cites its canonical reference (paper + sibling-repo path) in the module docstring.

The sibling repositories /mnt/ssd2/Works/{matfree,traceax,cola} are reference implementations only — opifex never carries them as runtime dependencies. Algorithms are implemented natively in JAX and cite the sibling source line-by-line.

References

• matfree — Krämer arXiv:2405.17277 (differentiable Lanczos/Arnoldi).
• traceax — Nahid et al. (XTrace, XNysTrace, Hutch++).
• cola — Potapczynski et al. arXiv:2309.03060 (structured operators; vendored under :mod:opifex.uncertainty.curvature.structured).

eig_partial

eig_partial(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int) -> tuple[Array, Array]

Partial eigendecomposition of a general (non-symmetric) operator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	matvec callable for a (possibly non-symmetric) operator.	required
init_vec	Array	starting vector for the Krylov sequence.	required
num_matvecs	int	number of Arnoldi iterations.	required

Returns:

Type	Description
Array	(eigenvalues, eigenvectors) where eigenvalues has shape
Array	(num_matvecs,) and may be complex; eigenvectors has shape
tuple[Array, Array]	(dim, num_matvecs) with each column lifted from the Hessenberg
tuple[Array, Array]	spectrum into the full space.

eigh_partial

eigh_partial(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int) -> tuple[Array, Array]

Partial eigendecomposition of a symmetric matrix-free operator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	symmetric matvec callable.	required
init_vec	Array	starting vector for the Krylov sequence.	required
num_matvecs	int	number of Lanczos iterations.	required

Returns:

Type	Description
Array	(eigenvalues, eigenvectors) where eigenvalues has shape
Array	(num_matvecs,) and eigenvectors has shape
tuple[Array, Array]	(dim, num_matvecs) with each column a unit eigenvector of the
tuple[Array, Array]	projected tridiagonal lifted into the full space.

svd_partial

svd_partial(*, matvec: Callable[[Array], Array], matvec_transpose: Callable[[Array], Array], init_vec: Array, num_matvecs: int) -> tuple[Array, Array, Array]

Partial SVD of a matrix-free operator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable computing A @ v.	required
matvec_transpose	Callable[[Array], Array]	callable computing A.T @ v.	required
init_vec	Array	starting vector for the Golub-Kahan sequence.	required
num_matvecs	int	number of bidiagonalisation iterations.	required

Returns:

Type	Description
Array	(left_singvecs, singvals, right_singvecs) where singvals has
Array	shape (num_matvecs,) and the singular-vector matrices have
Array	shape (dim, num_matvecs).

dense_funm_sym_eigh

dense_funm_sym_eigh(*, matrix: Array, matfun: Callable[[Array], Array]) -> Array

Apply matfun to a dense symmetric matrix via eigendecomposition.

Computes U @ diag(matfun(eigvals)) @ U.T where eigvals, U are the eigenvalues and orthonormal eigenvectors of matrix. matfun must be jittable and broadcast over the eigenvalue array.

Sibling reference: matfree/matfree/funm.py:300 dense_funm_sym_eigh.

funm_arnoldi

funm_arnoldi(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int, matfun: Callable[[Array], Array]) -> Array

Compute matfun(A) @ init_vec via Arnoldi for general square A.

Pipeline: Arnoldi → Hessenberg H → dense f(H) via eigendecomposition → lift. Eigenvalues of H may be complex; the dense f is taken via jnp.linalg.eig and reconstruction uses the right-eigenvector matrix.

Sibling reference: matfree/matfree/funm.py:147 funm_arnoldi.

funm_chebyshev

funm_chebyshev(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int, matfun: Callable[[Array], Array]) -> Array

Compute matfun(A) @ init_vec via Chebyshev polynomial expansion.

Assumes the spectrum of A is contained in (-1, 1) and matfun is analytic on that interval. Faster than Lanczos / Arnoldi when both conditions hold because no Krylov basis is built.

Sibling reference (port of the Clenshaw-style recurrence): matfree/matfree/funm.py:44 funm_chebyshev.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping (dim,) to (dim,).	required
init_vec	Array	vector v such that the function returns f(A) @ v.	required
num_matvecs	int	degree of the Chebyshev expansion (static).	required
matfun	Callable[[Array], Array]	scalar function broadcasting over the Chebyshev nodes.	required

Returns:

Type	Description
Array	Approximation of matfun(A) @ init_vec.

funm_lanczos_sym

funm_lanczos_sym(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int, matfun: Callable[[Array], Array]) -> Array

Compute matfun(A) @ init_vec via Lanczos for symmetric A.

Pipeline: Lanczos → tridiagonal T → dense f(T) → lift.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	symmetric matvec callable.	required
init_vec	Array	vector v such that the function returns f(A) @ v.	required
num_matvecs	int	number of Lanczos iterations. Static.	required
matfun	Callable[[Array], Array]	scalar function applied to eigenvalues of the tridiagonal projection. Must broadcast over a 1-D array.	required

Returns:

Type	Description
Array	The approximation of matfun(A) @ init_vec as a (dim,) array.

Sibling reference: matfree/matfree/funm.py:116 funm_lanczos_sym.

arnoldi_hessenberg

arnoldi_hessenberg(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int) -> tuple[Array, Array]

Arnoldi process — non-symmetric Krylov decomposition.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping (dim,) to (dim,).	required
init_vec	Array	starting vector of shape (dim,). Normalised internally.	required
num_matvecs	int	number of Arnoldi iterations. Static.	required

Returns:

Type	Description
Array	(basis, hessenberg) where basis is (dim, num_matvecs)
Array	orthonormal and hessenberg is (num_matvecs, num_matvecs)
tuple[Array, Array]	upper Hessenberg.

golub_kahan_bidiag

golub_kahan_bidiag(*, matvec: Callable[[Array], Array], matvec_transpose: Callable[[Array], Array], init_vec: Array, num_matvecs: int) -> tuple[Array, Array, Array]

Golub-Kahan bidiagonalisation of a general (possibly non-symmetric) operator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping (dim,) to (dim,) (A @ v).	required
matvec_transpose	Callable[[Array], Array]	callable mapping (dim,) to (dim,) (A.T @ v). Required because bidiagonalisation alternates between primal and adjoint applications.	required
init_vec	Array	starting vector of shape (dim,). Normalised internally.	required
num_matvecs	int	number of bidiagonal iterations. Static.	required

Returns:

Type	Description
Array	(left_basis, right_basis, bidiag) where left_basis and
Array	right_basis are (dim, num_matvecs) orthonormal and
Array	bidiag is (num_matvecs, num_matvecs) upper bidiagonal.
tuple[Array, Array, Array]	The singular values of bidiag match those of A.

lanczos_tridiag

lanczos_tridiag(*, matvec: Callable[[Array], Array], init_vec: Array, num_matvecs: int) -> tuple[Array, Array, Array]

Lanczos tridiagonalisation of a symmetric matrix-free operator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping (dim,) to (dim,). Must be symmetric for the tridiagonal output to be accurate; non-symmetric inputs yield a still-orthogonal basis but the projection is no longer tridiagonal.	required
init_vec	Array	starting vector of shape (dim,). Normalised internally.	required
num_matvecs	int	number of Lanczos iterations. Static.	required

Returns:

Type	Description
Array	(basis, diag, off_diag) where basis is (dim, num_matvecs)
Array	orthonormal, diag is (num_matvecs,) containing the main
Array	diagonal of the tridiagonal projection, and off_diag is
tuple[Array, Array, Array]	(num_matvecs - 1,) containing the sub/super-diagonal.

slq_logdet

slq_logdet(*, matvec: Callable[[Array], Array], dim: int, num_samples: int, num_matvecs: int, key: Array) -> Array

Estimate log det(A) for symmetric positive-definite A.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	SPD matvec callable.	required
dim	int	matrix dimension n (static).	required
num_samples	int	number of Rademacher probes (static). Variance of the estimator scales as O(1 / num_samples).	required
num_matvecs	int	Lanczos depth per probe (static). Higher depth improves the per-probe quadratic-form approximation.	required
key	Array	PRNG key, split into one substream per probe.	required

Returns:

Type	Description
Array	Scalar JAX array containing the SLQ estimate of log det(A).

cholesky_greedy

cholesky_greedy(*, matrix: Array, rank: int) -> Array

Compute a partial Cholesky factorisation with greedy pivoting.

At each step the algorithm selects the diagonal entry of the residual with the largest current value, then performs the corresponding Cholesky-style column update on the running factor.

Parameters:

Name	Type	Description	Default
matrix	Array	dense (dim, dim) symmetric positive-semidefinite matrix.	required
rank	int	number of Cholesky iterations (static); rank <= dim.	required

Returns:

Type	Description
Array	Factor L of shape (dim, rank). L @ L.T is the
Array	rank-rank greedy partial Cholesky approximation to
Array	matrix.

rp_cholesky

rp_cholesky(*, matrix: Array, rank: int, key: Array) -> Array

Compute a partial Cholesky factorisation with randomly-pivoted Cholesky.

Pivots are sampled at each step from a categorical distribution whose probability mass is proportional to the current residual diagonal. Per Chen+ arXiv:2207.06503, this scheme yields strong expected-error guarantees for low-rank approximation of PSD matrices (in particular, optimal expected Frobenius error after the rank-r iteration for a rank-r ground truth).

Parameters:

Name	Type	Description	Default
matrix	Array	dense (dim, dim) symmetric positive-semidefinite matrix.	required
rank	int	number of Cholesky iterations (static); rank <= dim.	required
key	Array	PRNG key, split into one substream per pivot draw.	required

Returns:

Type	Description
Array	Factor L of shape (dim, rank). L @ L.T is the
Array	rank-rank randomly-pivoted Cholesky approximation.

lsmr

lsmr(*, matvec: Callable[[Array], Array], matvec_transpose: Callable[[Array], Array], rhs: Array, dim_cols: int, num_matvecs: int, damping: float = 0.0) -> Array

Solve min ||A x - rhs||^2 + damping^2 ||x||^2 via bidiagonalised reduction.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable computing A @ v for v of length dim_cols.	required
matvec_transpose	Callable[[Array], Array]	callable computing A.T @ u for u of length len(rhs).	required
rhs	Array	right-hand-side vector b.	required
dim_cols	int	number of columns of A (static); x has this length.	required
num_matvecs	int	number of bidiagonalisation iterations (static). Larger values yield more accurate truncated solutions; for full-rank A with num_matvecs >= rank, the routine recovers the exact least-squares solution.	required
damping	float	Tikhonov damping coefficient λ (default 0 for plain least squares). Regularises ill-conditioned systems.	0.0

Returns:

Type	Description
Array	Approximate solution x of shape (dim_cols,).

trace_moments

trace_moments(*, matvec: Callable[[Array], Array], dim: int, num_samples: int, powers: tuple[int, ...], key: Array) -> tuple[Array, ...]

Estimate raw moments of the per-probe quadratic form v.T A v.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	symmetric matvec callable.	required
dim	int	matrix dimension n (static).	required
num_samples	int	number of Rademacher probes (static).	required
powers	tuple[int, ...]	tuple of integer powers — one entry per moment to compute. (1,) returns the Hutchinson trace estimate. (1, 2) adds the second raw moment for variance recovery via Var = E[X^2] - E[X]^2.	required
key	Array	PRNG key.	required

Returns:

Type	Description
Array	Tuple of scalar arrays, one per power in powers. The m-th
...	entry equals (1/K) * sum_k (v_k.T A v_k)**powers[m] over
tuple[Array, ...]	Rademacher probes v_k.

randomized_svd

randomized_svd(*, matvec: Callable[[Array], Array], matvec_transpose: Callable[[Array], Array], dim_rows: int, dim_cols: int, rank: int, oversampling: int, num_iterations: int, key: Array) -> tuple[Array, Array, Array]

Approximate the top-rank SVD of A via HMT randomised SVD.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable computing A @ v for v of shape (dim_cols,).	required
matvec_transpose	Callable[[Array], Array]	callable computing A.T @ u for u of shape (dim_rows,).	required
dim_rows	int	number of rows of A (static).	required
dim_cols	int	number of columns of A (static).	required
rank	int	target rank (static); rank <= min(dim_rows, dim_cols).	required
oversampling	int	extra sketch columns. Typical values 5–10 trade additional matvec cost for robust spectral capture (Halko+ §4.4).	required
num_iterations	int	number of subspace iterations (static). 0 recovers the basic projection-based estimate; >= 1 amplifies leading directions when the spectrum decays slowly.	required
key	Array	PRNG key, split into one substream for the initial sketch and one per subspace-iteration step.	required

Returns:

Type	Description
Array	(left_vectors, singvals, right_vectors) where left_vectors
Array	has shape (dim_rows, rank), singvals has shape
Array	(rank,), and right_vectors has shape (dim_cols, rank).

hutch_plus_plus_trace

hutch_plus_plus_trace(*, matvec: Callable[[Array], Array], dim: int, num_samples: int, key: Array) -> Array

Estimate trace(A) via Hutch++ (Meyer et al. arXiv:2010.09649).

The total query budget num_samples is split evenly across three stages — sketch, exact reduction, and Hutchinson residual — so the estimator uses 2 * (num_samples // 3) matvecs to build a leading subspace and num_samples // 3 Rademacher probes for the orthogonal residual.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping a (dim,) vector to a (dim,) vector.	required
dim	int	the dimension n of the operator (static).	required
num_samples	int	total matvec budget. Static; must be divisible by 3 to allocate equal sub-budgets cleanly.	required
key	Array	PRNG key; split into sketch and hutchinson substreams.	required

Returns:

Type	Description
Array	A scalar JAX array — the Hutch++ trace estimate. Captures the
Array	contribution of the top num_samples // 3 directions exactly and
Array	adds an unbiased Hutchinson term for the orthogonal residual.

hutchinson_trace

hutchinson_trace(*, matvec: Callable[[Array], Array], dim: int, num_samples: int, key: Array) -> Array

Estimate trace(A) from a matrix-free matvec using Rademacher probes.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping a (dim,) vector to a (dim,) vector. Conceptually computes A @ v for the matrix whose trace is being estimated.	required
dim	int	the dimension n of the operator. Static — used to allocate the probe array.	required
num_samples	int	number of Rademacher probes. Static — controls variance.	required
key	Array	a PRNG key.	required

Returns:

Type	Description
Array	A scalar JAX array — the empirical mean of v^T A v over the
Array	Rademacher probes. Unbiased for symmetric A; variance scales as
Array	O(1 / num_samples).

xnys_trace

xnys_trace(*, matvec: Callable[[Array], Array], dim: int, num_samples: int, key: Array) -> Array

Estimate trace(A) via XNysTrace for a PSD operator A.

Builds a Nyström approximation from num_samples matvecs and adds an exchangeable residual correction. On a PSD matrix whose rank is at most num_samples, XNysTrace recovers the trace exactly up to numerical roundoff.

Sibling reference (line-by-line port): traceax/src/traceax/_estimators.py:241 XNysTraceEstimator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping (dim,) to (dim,). Must be a positive-semi-definite operator; the algorithm forms a symmetric Cholesky factor of Ω^T A Ω internally.	required
dim	int	matrix dimension n (static).	required
num_samples	int	total matvec budget. Static.	required
key	Array	PRNG key.	required

Returns:

Type	Description
Array	Scalar JAX array — the mean of the per-probe XNysTrace estimates.

xtrace

xtrace(*, matvec: Callable[[Array], Array], dim: int, num_samples: int, key: Array) -> Array

Estimate trace(A) via XTrace (Epperly+ arXiv:2301.07825).

XTrace enforces exchangeability across the m = num_samples // 2 random probes by reusing each probe in both the sketch and the residual stages — implemented efficiently through leave-one-out algebra so the total cost is num_samples matvecs (m for the sketch, m for the basis image Z = A Q). The improved scaling uses traceax's _get_scale factor that orthogonalises probes against the low-rank approximation.

Sibling reference (line-by-line port): traceax/src/traceax/_estimators.py:170 XTraceEstimator.

Parameters:

Name	Type	Description	Default
matvec	Callable[[Array], Array]	callable mapping (dim,) to (dim,).	required
dim	int	matrix dimension n (static).	required
num_samples	int	total matvec budget. Static. num_samples // 2 sketch matvecs plus the same number of basis-image matvecs.	required
key	Array	PRNG key.	required

Returns:

Type	Description
Array	Scalar JAX array — the mean of the per-probe XTrace estimates.

Bayesian-quadrature primitives.

Functional building blocks for integrating a function f against a probability measure π with calibrated uncertainty in the integral. The initial slice provides only the no-GP baseline (:func:bayesian_monte_carlo) plus a typed result container (:class:IntegralEstimate); advanced methods (WSABI-L, vanilla BQ, SOBER, FFBQ) will land as additional modules and re-export from this package's namespace.

References

• Rasmussen, C. E. & Ghahramani, Z. 2003 — Bayesian Monte Carlo, NeurIPS 16.
• Briol, F.-X. et al. 2019 — Probabilistic Integration: A Role in Statistical Computation?, Statistical Science 34(1).

EmukitQuadratureAdapterSpec dataclass

EmukitQuadratureAdapterSpec(*, default_strategy: DefaultStrategy = BAYESIAN_QUADRATURE, source_package: str = 'emukit', required_capabilities: tuple[str, ...] = (), family_tags: tuple[str, ...] = ('reference_baseline', 'numpy_only'), notes: str = 'Metadata-only — emukit is the NumPy reference implementation. Vendored under opifex.uncertainty.quadrature; not a runtime dependency. See ../emukit/emukit/quadrature/* for the source.')

Bases: _BQAdapterSpecBase

Read-only adapter pointing at the emukit (NumPy) BQ baselines.

emukit is vendored as a reference implementation, not installed as a runtime dependency. Use this spec for benchmarking; concrete integration drives through opifex's vanilla / WSABI-L JAX-native implementations.

FFBQAdapterSpec dataclass

FFBQAdapterSpec(*, default_strategy: DefaultStrategy = BAYESIAN_QUADRATURE, source_package: str = 'opifex', required_capabilities: tuple[str, ...] = (), family_tags: tuple[str, ...] = ('frank_wolfe', 'convex_iterates', 'point_set'), notes: str = 'FFBQ — Frank-Wolfe Bayesian Quadrature (Briol+ NeurIPS 2015, arXiv:1506.02681). FW-Vanilla iterate with mass-redistribution step alpha_n = 1/(n+1); MMD = O(1/n). Vendored in frank_wolfe_bq.py per the SOBER <-> FFBQ design split (fix #190).')

Bases: _BQAdapterSpecBase

Frank-Wolfe Bayesian quadrature (Briol+ 2015 NeurIPS).

Despite the legacy FFBQ mnemonic the algorithm is Frank-Wolfe BQ (FW-Vanilla, Algorithm 1 in the paper) — not a Fourier-features method. The design notes (fix #190) pin this spec to Briol+ 2015.

wrap

wrap(model: Any, capability: UQCapability) -> Any

Return the JAX-native Frank-Wolfe BQ callable.

SOBERAdapterSpec dataclass

SOBERAdapterSpec(*, default_strategy: DefaultStrategy = BAYESIAN_QUADRATURE, source_package: str = 'opifex', required_capabilities: tuple[str, ...] = (), family_tags: tuple[str, ...] = ('point_set', 'kernel_recombination', 'discrete_mixed'), notes: str = 'SOBER (Adachi+ 2022 NeurIPS arXiv:2206.04734 + 2023 TMLR arXiv:2301.11832). Kernel-recombination via Tchernychova-Lyons CAR + Nyström low-rank approximation; vendored in sober.py per the SOBER ↔ FFBQ design split.')

Bases: _BQAdapterSpecBase

Kernel-recombination Bayesian quadrature (point-set method).

wrap

wrap(model: Any, capability: UQCapability) -> Any

Return the JAX-native SOBER kernel-recombination callable.

VanillaBayesianQuadratureAdapterSpec dataclass

VanillaBayesianQuadratureAdapterSpec(*, default_strategy: DefaultStrategy = BAYESIAN_QUADRATURE, source_package: str = 'opifex', required_capabilities: tuple[str, ...] = (), family_tags: tuple[str, ...] = ('gp_prior', 'closed_form_posterior'), notes: str = "Vanilla BQ — uses a Gaussian-process prior over the integrand and closed-form formulas for the posterior integral mean and variance. Coexists with WSABI-L in opifex's bayesian_quadrature.py module.")

Bases: _BQAdapterSpecBase

GP-prior Bayesian quadrature with closed-form mean and variance.

wrap

wrap(model: Any, capability: UQCapability) -> Any

Return the JAX-native closed-form Vanilla BQ callable.

WSABILAdapterSpec dataclass

WSABILAdapterSpec(*, default_strategy: DefaultStrategy = BAYESIAN_QUADRATURE, source_package: str = 'opifex', required_capabilities: tuple[str, ...] = (), family_tags: tuple[str, ...] = ('warped_gp_prior', 'linear_warp', 'sequential_active'), notes: str = "WSABI-L (Gunter+ 2014). Linear approximation of the warped-GP posterior moments for non-negative integrands. Coexists with VanillaBayesianQuadratureAdapterSpec in opifex's bayesian_quadrature.py module.")

Bases: _BQAdapterSpecBase

Warped Sequential Active Bayesian Integration — linear approximation.

wrap

wrap(model: Any, capability: UQCapability) -> Any

Return the JAX-native WSABI-L bounded-BQ mean callable.

IntegralEstimate

Point estimate + standard-error variance of an integral.

Registered as a JAX pytree (flax.struct.dataclass) so functions returning this type can be jit-compiled.

Attributes:

Name	Type	Description
mean	Array	Scalar estimate of ∫ f dπ.
variance	Array	Standard-error variance of mean (i.e., Var(f(X)) / N for N samples).
num_samples	int	Number of samples that produced the estimate (carried as static pytree aux-data, not as a leaf).

GaussianMeasure dataclass

GaussianMeasure(*, mean: Array, variance: Array)

Diagonal-covariance Gaussian integration measure.

Density:

.. math::

p(x) = (2\pi)^{-d/2} \prod_i \sigma_i^{-1}
       \exp\Bigl(-\tfrac12 \sum_i (x_i - \mu_i)^2 / \sigma_i^2\Bigr).

Attributes:

Name	Type	Description
mean	Array	Per-dimension mean μ of shape (d,).
variance	Array	Per-dimension diagonal variance σ² of shape (d,). Must be entry-wise strictly positive.

input_dim property

input_dim: int

Dimensionality of the support.

sample

sample(*, num_samples: int, key: Array) -> Array

Draw num_samples IID samples from N(μ, diag(σ²)).

LebesgueMeasure dataclass

LebesgueMeasure(*, lower: Array, upper: Array)

Uniform Lebesgue measure over a hyperrectangle.

Attributes:

Name	Type	Description
lower	Array	Per-dimension lower bound of shape (d,).
upper	Array	Per-dimension upper bound of shape (d,). Must be entry-wise strictly greater than lower.

input_dim property

input_dim: int

Dimensionality of the support.

volume property

volume: Array

Volume ∏_i (upper_i - lower_i) of the hyperrectangle.

sample

sample(*, num_samples: int, key: Array) -> Array

Draw num_samples IID uniform samples on [lower, upper].

Gaussian processes and state-space models

Exact and approximate Gaussian-Process regression for opifex.

Phase 11 Task 11.1 ships the foundational exact conjugate-Gaussian GP regression surface on top of the Task 6.3 adapter specs at :mod:opifex.uncertainty.adapters.gp. Subsequent slices will add deep kernels, multi-output ICM/LCM, RFF approximations, natural-gradient SVGP, heteroscedastic likelihoods, OAK kernels, and CARMA/Celerite/SHO state-space kernels per the Phase-11 plan.

References

• Rasmussen, C. E., Williams, C. K. I. 2006 — Gaussian Processes for Machine Learning, MIT Press; Algorithm 2.1 §2.2 (PRIMARY for exact).

ExactGPState dataclass

ExactGPState(*, x_train: Array, y_train: Array, cholesky: Array, alpha: Array, lengthscale: float, output_scale: float, noise_std: float, kernel_fn: Callable[..., Array] = rbf_kernel)

Fitted state for an exact conjugate-Gaussian GP.

Attributes:

Name	Type	Description
x_train	Array	Training inputs (n, d).
y_train	Array	Training targets (n,).
cholesky	Array	Lower-triangular Cholesky factor of K + σ² I.
alpha	Array	α = (K + σ² I)^{-1} y — pre-solved coefficient vector used at predict time.
lengthscale	float	Kernel length-scale used at fit time (forwarded into predict_exact_gp for kernel evaluations at test points).
output_scale	float	Kernel output-scale used at fit time.
noise_std	float	Observation noise scale σ used at fit time.
kernel_fn	Callable[..., Array]	Callable kernel used during fit. Defaults to :func:rbf_kernel.

LaplaceGPState dataclass

LaplaceGPState(*, x_train: Array, y_train: Array, f_mode: Array, sqrt_w: Array, cholesky_b: Array, gradient_log_likelihood_at_mode: Array, log_marginal_likelihood: Array, lengthscale: float, output_scale: float, kernel_fn: Callable[..., Array] = rbf_kernel)

Laplace-approximated GP posterior state, likelihood-agnostic.

Attributes:

Name	Type	Description
x_train	Array	(n, d) training inputs.
y_train	Array	(n,) training targets in the likelihood's support ({-1, +1} for Bernoulli, non-negative integers for Poisson, reals for Student-t, (0, 1) for Beta, ...).
f_mode	Array	Posterior mode f̂ of shape (n,).
sqrt_w	Array	Diagonal of W^{1/2} at f̂ (shape (n,)).
cholesky_b	Array	Lower-triangular Cholesky factor of B = I + W^{1/2} K W^{1/2}, shape (n, n).
gradient_log_likelihood_at_mode	Array	∇ log p(y | f̂), cached for use in the predictive-mean computation.
log_marginal_likelihood	Array	Laplace-approximated log marginal (scalar).
lengthscale	float	Kernel length-scale used at fit time.
output_scale	float	Kernel output-scale used at fit time.
kernel_fn	Callable[..., Array]	Kernel callable.

QuasisepGPState dataclass

QuasisepGPState(*, x_train: Array, y_train: Array, log_marginal_likelihood: Array, lengthscale: float, output_scale: float, noise_std: float, quality_factor: float)

Fitted state for the scalable SHO Gaussian process.

Attributes:

Name	Type	Description
x_train	Array	(n,) or (n, 1) strictly-increasing training times.
y_train	Array	(n,) training observations.
log_marginal_likelihood	Array	Exact log marginal evaluated in O(n) via :func:kalman_log_likelihood.
lengthscale	float	ℓ = 1/ω — kernel length-scale.
output_scale	float	σ_f — kernel output-scale.
noise_std	float	σ — observation noise scale.
quality_factor	float	Q > 1/2 — SHO quality factor.

QuasisepCarmaGPState dataclass

QuasisepCarmaGPState(*, x_train: Array, y_train: Array, log_marginal_likelihood: Array, ar_coefficients: Array, ma_coefficients: Array, lengthscale: float, output_scale: float | Array, noise_std: float)

Fitted state for the scalable CARMA Gaussian process.

Attributes:

Name	Type	Description
x_train	Array	(n,) or (n, 1) strictly-increasing training times.
y_train	Array	(n,) training observations.
log_marginal_likelihood	Array	Exact log marginal evaluated in O(n) via :func:kalman_log_likelihood.
ar_coefficients	Array	(p,) autoregressive coefficients α_0, …, α_{p-1} (implicit α_p = 1).
ma_coefficients	Array	(q + 1,) moving-average coefficients β_0, …, β_q with q + 1 ≤ p.
lengthscale	float	ℓ — rescales the lag (τ → τ / ℓ).
output_scale	float | Array	σ_f — kernel output-scale (k → σ_f² k).
noise_std	float	σ — observation noise scale.

RFFGPState dataclass

RFFGPState(*, features_train: Array, omega: Array, feature_scale: Array, gram: Array, gram_cholesky: Array, alpha: Array, lengthscale: float, output_scale: float, noise_std: float)

Fitted state for an RFF-approximate conjugate GP.

Attributes:

Name	Type	Description
features_train	Array	Feature matrix Φ(X) ∈ R^{n × D}.
omega	Array	Spectral samples (D/2, d) used to build the feature map; needed at predict time so the same map is applied to test inputs.
feature_scale	Array	Scalar prefactor σ_f sqrt(2 / D) applied to the cos/sin features.
gram	Array	(D, D) ridge-Gram Φ^T Φ + σ² I.
gram_cholesky	Array	Lower-triangular Cholesky factor of gram.
alpha	Array	Pre-solved (Φ^T Φ + σ² I)^{-1} Φ^T y of shape (D,).
lengthscale	float	RBF length-scale used at fit time.
output_scale	float	RBF output-scale used at fit time.
noise_std	float	Observation noise scale σ.

SpatioTemporalGPState dataclass

SpatioTemporalGPState(*, smoothed_means: Array, smoothed_covs: Array, measurement: Array, noise_std: Array, num_space: int)

Fitted state of a conjugate spatio-temporal variational GP.

Attributes:

Name	Type	Description
smoothed_means	Array	RTS-smoothed latent state means per time, shape (T, M * state_dim_t) (spatial index outer, temporal SDE state inner — the I_M \\otimes · Kronecker layout).
smoothed_covs	Array	RTS-smoothed latent state covariances per time, shape (T, M * state_dim_t, M * state_dim_t).
measurement	Array	Lifted measurement matrix H = I_M \\otimes H_t of shape (M, M * state_dim_t).
noise_std	Array	Observation noise scale σ used at fit time.
num_space	int	Number of spatial points M.

SVGPState dataclass

SVGPState(*, x_inducing: Array, cholesky_kmm: Array, cholesky_b: Array, scaled_alpha: Array, lengthscale: float, output_scale: float, noise_std: float, jitter: float, kernel_fn: Callable[..., Array] = rbf_kernel, cached_y_squared_norm: Array = (lambda: asarray(0.0))(), cached_a_y_inside_norm: Array = (lambda: asarray(0.0))(), cached_trace_aat: Array = (lambda: asarray(0.0))(), cached_kxx_diag_sum: Array = (lambda: asarray(0.0))(), cached_log_det_b: Array = (lambda: asarray(0.0))(), cached_n: int = 0)

Fitted state for the Titsias-collapsed sparse-variational GP.

Attributes:

Name	Type	Description
x_inducing	Array	(m, d) inducing inputs Z.
cholesky_kmm	Array	Lower-triangular L_z = chol(K_zz + jitter I).
cholesky_b	Array	Lower-triangular L_B = chol(I + A A^T) where A = L_z^{-1} K_zx / σ.
scaled_alpha	Array	L_B^{-1} (A y / σ) of shape (m,) — the pre-solved coefficient vector used at predict time.
lengthscale	float	Kernel length-scale.
output_scale	float	Kernel output-scale.
noise_std	float	Observation noise scale σ.
jitter	float	Numerical jitter added to K_zz for PD stability.
kernel_fn	Callable[..., Array]	Kernel callable.
cached_y_squared_norm	Array	||y||² cached for the collapsed-ELBO quadratic term.
cached_a_y_inside_norm	Array	||L_B^{-1} A y / σ||² cached for the collapsed-ELBO quadratic term.
cached_trace_aat	Array	tr(A A^T) cached for the collapsed-ELBO trace term.
cached_kxx_diag_sum	Array	Σ_i K(x_i, x_i) cached for the collapsed-ELBO trace term.
cached_log_det_b	Array	log|B| = 2 Σ_i log L_B[i, i] cached for the collapsed-ELBO log-det term.
cached_n	int	Training-set size n (static int).

StochasticSVGPState dataclass

StochasticSVGPState(*, x_inducing: Array, whitened_mean: Array, whitened_root_cov: Array, lengthscale: float, output_scale: float, kernel_fn: Callable[..., Array] = rbf_kernel, log_likelihood_fn: LogLikelihoodFn = lambda f, y: -(y - f) ** 2, jitter: float = 1e-06)

Whitened-parametrisation stochastic SVGP state.

Attributes:

Name	Type	Description
x_inducing	Array	(m, d) inducing inputs Z.
whitened_mean	Array	(m,) μ_w — variational mean in the whitened basis (μ = L_z μ_w in the unwhitened basis).
whitened_root_cov	Array	(m, m) lower-triangular L_w such that the whitened variational covariance is L_w L_w^T (unwhitened: L_z L_w L_w^T L_z^T).
lengthscale	float	Kernel length-scale.
output_scale	float	Kernel output-scale.
kernel_fn	Callable[..., Array]	Kernel callable.
log_likelihood_fn	LogLikelihoodFn	Per-observation log p(y_i | f_i).
jitter	float	Numerical jitter on K_zz for PD stability.

exact_gp_loocv_log_predictive

exact_gp_loocv_log_predictive(*, state: ExactGPState) -> Array

Leave-one-out predictive log-likelihood for a fitted exact GP.

Rasmussen & Williams 2006 §5.4.2 derives the closed form

.. math::

\mu_{-i} &= y_{i} - \frac{\alpha_{i}}{[K^{-1}]_{ii}}, \\
\sigma_{-i}^{2} &= \frac{1}{[K^{-1}]_{ii}}, \\
\log p(y_{i} \mid X, y_{-i})
&= -\tfrac{1}{2}\!\left[\log(2\pi)
    + \log \sigma_{-i}^{2}
    + \frac{(y_{i} - \mu_{-i})^{2}}{\sigma_{-i}^{2}}\right],

so all n leave-one-out terms are recovered from a single Cholesky factorisation (no refitting required).

Parameters:

Name	Type	Description	Default
state	ExactGPState	Fitted :class:ExactGPState.	required

Returns:

Type	Description
Array	Scalar LOOCV log-predictive Σ_i log p(y_i | X, y_{-i}).

fit_exact_gp

fit_exact_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, noise_std: float, kernel_fn: Callable[..., Array] = rbf_kernel) -> ExactGPState

Fit an exact conjugate-Gaussian GP (RW06 Algorithm 2.1 prelude).

Computes the Cholesky factor of K + σ² I and the pre-solved α = (K + σ² I)^{-1} y for downstream :func:predict_exact_gp calls.

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) training targets.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
noise_std	float	Observation noise scale σ (strictly positive — also acts as the jitter that keeps K + σ² I positive definite).	required
kernel_fn	Callable[..., Array]	Optional kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Name	Type	Description
An	ExactGPState	class:ExactGPState carrying the Cholesky factor + α.

Raises:

Type	Description
ValueError	If noise_std is non-positive.

fit_heteroscedastic_exact_gp

fit_heteroscedastic_exact_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, noise_std: Array, kernel_fn: Callable[..., Array] = rbf_kernel) -> ExactGPState

Fit an exact GP with per-observation Gaussian noise.

Generalises :func:fit_exact_gp from a scalar noise level to a per-observation vector noise_std[i] = σ_i. The Gram matrix becomes K + diag(σ_i²); the rest of the algorithm is unchanged so the returned :class:ExactGPState is consumable by the existing :func:predict_exact_gp driver. Setting every σ_i = σ reproduces :func:fit_exact_gp exactly (verified by the test suite).

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) training targets.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
noise_std	Array	(n,) per-observation noise scales (every entry must be strictly positive).	required
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Type	Description
ExactGPState	class:ExactGPState carrying the Cholesky factor + α and
ExactGPState	a noise_std field set to -1.0 so callers can detect
ExactGPState	that the fitted state used the heteroscedastic path (the per-
ExactGPState	observation noises live inside the Cholesky factor by then).

Raises:

Type	Description
ValueError	If noise_std has a non-positive entry or a length not matching y_train.

predict_exact_gp

predict_exact_gp(*, state: ExactGPState, x_test: Array) -> PredictiveDistribution

Predict at x_test using a fitted :class:ExactGPState.

Implements the Algorithm-2.1 Cholesky-back-substitution recipe:

.. code-block:: text

v    = L \ K(X, X*)
mean = K(X*, X) @ α
var  = K(X*, X*) - v^T v   (per test point)

Parameters:

Name	Type	Description	Default
state	ExactGPState	Fitted :class:ExactGPState.	required
x_test	Array	(m, d) test inputs.	required

Returns:

Name	Type	Description
A	PredictiveDistribution	class:PredictiveDistribution carrying the predictive
	PredictiveDistribution	mean (m,) and the marginal variance (m,) plus
	PredictiveDistribution	metadata advertising method=gaussian_process and the
	PredictiveDistribution	opifex source-package tag.

rbf_kernel

rbf_kernel(x1: Array, x2: Array, *, lengthscale: float, output_scale: float) -> Array

Squared-exponential (RBF) kernel.

.. math::

k(x, x') = \sigma_{f}^{2} \exp\!\left(-\tfrac{1}{2}\,
    \frac{\lVert x - x' \rVert_{2}^{2}}{\ell^{2}}\right).

Parameters:

Name	Type	Description	Default
x1	Array	Inputs of shape (n, d).	required
x2	Array	Inputs of shape (m, d).	required
lengthscale	float	ℓ (must be strictly positive).	required
output_scale	float	σ_f (must be strictly positive).	required

Returns:

Type	Description
Array	(n, m) kernel matrix.

Raises:

Type	Description
ValueError	If lengthscale or output_scale is non-positive.

additive_kernel

additive_kernel(*, component_kernel_fns: tuple[Callable[..., Array], ...]) -> Callable[..., Array]

Additive (OAK-base) kernel — sum of per-dimension univariate kernels.

.. math::

k_{\text{add}}(x, x') = \sum_{d=1}^{D} k_{d}(x_{d}, x'_{d}).

Each component_kernel_fns[d] is evaluated on the d-th input dimension only; the returned kernel has the standard (x1, x2, *, lengthscale, output_scale) -> Gram signature. This is the first-order (max_order = 1) case of the Orthogonal Additive Kernel (Lu, Boukouvalas, Hensman 2022 ICML — Additive Gaussian Processes Revisited). Higher-order interactions via Newton-Girard recursion + the Gaussian-measure orthogonality constraint are deferred to a follow-up slice; the first-order form already enables ANOVA-style interpretable GPs.

Parameters:

Name	Type	Description	Default
component_kernel_fns	tuple[Callable[..., Array], ...]	One kernel callable per input dimension. All callables share the standard kernel signature.	required

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature.

Raises:

Type	Description
ValueError	If component_kernel_fns is empty.

carma_kernel

carma_kernel(*, alpha: Array, beta: Array) -> Callable[..., Array]

Direct-evaluation CARMA(p, q) kernel (Kelly+ 2014 arXiv:1402.5978).

Continuous-time autoregressive moving-average process. The kernel is a closed-form sum of real-and-complex exponential terms:

.. math::

k(\tau) = \sigma_{\text{out}}^{2}\,
    \mathrm{Re}\!\left[
        \sum_{j=1}^{p} \mathrm{acf}_{j}\,
        \exp(\rho_{j}\,|\tau|\,/\,\ell)
    \right],

where ρ_j are the AR roots and acf_j the ACVF coefficients derived from (α, β) per Kelly+ 2014 eq. 4. Reduces to known cases at low orders:

• CARMA(1, 0) ≡ Matern-½ with ℓ_eff = 1/α_0.
• CARMA(2, 1) ≡ Celerite Complex term (Foreman-Mackey+ 2017).

The opifex implementation ships the direct O(n²) Gram. The scalable O(n) state-space quasiseparable port is filed as a deferred follow-up in memory-bank/.../notes/06-task-11.1-deferred-items.md (D7).

Parameters:

Name	Type	Description	Default
alpha	Array	(p,) array of AR coefficients α_0, …, α_{p-1}. The implicit highest-order coefficient α_p = 1 is appended internally. For stationarity all roots of the AR polynomial must have negative real parts (the factory does not enforce this — NaN propagates if violated).	required
beta	Array	(q + 1,) array of MA coefficients β_0, …, β_q with q + 1 ≤ p. β_0 is absorbed as the overall amplitude factor.	required

Returns:

Type	Description
Callable[..., Array]	Standard kernel callable
Callable[..., Array]	(x1, x2, *, lengthscale, output_scale) -> (n, m) where
Callable[..., Array]	lengthscale rescales the lag (τ → τ / ℓ) and
Callable[..., Array]	output_scale² multiplies the kernel amplitude.

Raises:

Type	Description
ValueError	If len(beta) > len(alpha) (MA order exceeds AR order, violating Kelly+ 2014 Eq. 1).

Reference implementation consulted (READ-ONLY): ../tinygp/src/tinygp/kernels/quasisep.py:CARMA (lines 672-885) + the carma_roots / carma_acvf helpers.

celerite_complex_kernel

celerite_complex_kernel(*, sine_amplitude: float, frequency: float) -> Callable[..., Array]

Celerite Complex term — building block of the celerite kernel family.

Foreman-Mackey, Agol, Ambikasaran, Angus 2017 (AJ, arXiv:1703.09710) and Foreman-Mackey 2018 (RNAAS, Scalable backpropagation for Gaussian processes using celerite) define the Complex term as

.. math::

k(\tau) = \exp(-c\,|\tau|)\,\bigl[a\,\cos(d\,|\tau|)
                                  + b\,\sin(d\,|\tau|)\bigr],

parametrised by (a, b, c, d) with a, c > 0 and the positive-definiteness constraint a c - b d > 0. The opifex parametrisation maps a = output_scale², c = 1/lengthscale (matching :func:matern12_kernel for the cosine envelope), and closes over b = sine_amplitude, d = frequency via the factory.

The Complex term reduces to the celerite Real term when sine_amplitude = 0, frequency = 0 — i.e. it collapses to :func:matern12_kernel. Sums of Complex terms via :func:kernel_sum realise the full celerite kernel family (Foreman-Mackey 2018 §2); the underdamped SHO (:func:damped_oscillator_kernel) is one specific Complex-term reparametrisation.

Reference implementation consulted (READ-ONLY): ../tinygp/src/tinygp/kernels/quasisep.py:Celerite.

Parameters:

Name	Type	Description	Default
sine_amplitude	float	b — sine-component amplitude. Constraint: ``output_scale² / lengthscale - sine_amplitude * frequency 0`` (positive-definiteness; NaN propagates if violated).	required
frequency	float	d — oscillation angular frequency.	required

Returns:

Type	Description
Callable[..., Array]	Standard kernel callable
Callable[..., Array]	(x1, x2, *, lengthscale, output_scale) -> (n, m).

constrained_rbf_kernel

constrained_rbf_kernel(*, input_mean: float = 0.0, input_std: float = 1.0) -> Callable[..., Array]

RBF kernel projected to be orthogonal to constants under N(μ, ζ²).

Lu, Boukouvalas, Hensman 2022 (ICML, Additive Gaussian Processes Revisited) eq. 10 gives the closed-form Gaussian-measure projection of the squared-exponential kernel onto the orthogonal complement of the constants. For p(x) = N(μ, ζ²):

.. math::

\tilde{k}(x, x')
    &= k(x, x') - \frac{\mu_{k}(x)\,\mu_{k}(x')}{m_{k}}, \\
\mu_{k}(x) &= \int k(x, y)\,p(y)\,\mathrm{d}y
    = \frac{\sigma_{f}^{2}\,\ell}{\sqrt{\ell^{2} + \zeta^{2}}}\,
      \exp\!\left(-\frac{(x - \mu)^{2}}{2(\ell^{2} + \zeta^{2})}\right), \\
m_{k}      &= \int \mu_{k}(x)\,p(x)\,\mathrm{d}x
    = \frac{\sigma_{f}^{2}\,\ell}{\sqrt{\ell^{2} + 2\zeta^{2}}},

yielding the projection

.. math::

\tilde{k}(x, x') = k(x, x')
    - \frac{\sigma_{f}^{2}\,\ell\,\sqrt{\ell^{2} + 2\zeta^{2}}}
           {\ell^{2} + \zeta^{2}}\,
      \exp\!\left(-\frac{(x - \mu)^{2} + (x' - \mu)^{2}}
                         {2(\ell^{2} + \zeta^{2})}\right).

This is the canonical Gaussian-measure-orthogonal base kernel for the Orthogonal Additive Kernel (:func:orthogonal_additive_kernel). For μ = 0, ζ = 1 the formula matches GPJax's additive/oak.py:_constrained_se_kernel line-for-line.

Parameters:

Name	Type	Description	Default
input_mean	float	Gaussian input-measure mean μ. Defaults to 0.	0.0
input_std	float	Gaussian input-measure std ζ > 0. Defaults to 1.	1.0

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature.

damped_oscillator_kernel

damped_oscillator_kernel(*, quality_factor: float) -> Callable[..., Array]

Underdamped simple-harmonic-oscillator (SHO) kernel.

Foreman-Mackey, Agol, Ambikasaran, Angus 2017 (AJ, arXiv:1703.09710) introduce the celerite SHO kernel — a damped oscillator with closed-form covariance ideal for time-series GPs. For Q > 1/2 (underdamped) the covariance at lag τ = x_1 - x_2 is

.. math::

k(\tau) = \sigma^{2}\,\exp\!\left(-\frac{\omega\,|\tau|}{2 Q}\right)
    \left[\cos\!\left(\frac{g\,\omega\,|\tau|}{2 Q}\right)
          + \frac{1}{g}\sin\!\left(\frac{g\,\omega\,|\tau|}{2 Q}\right)\right]

where g = √(4 Q² − 1) and ω is the natural angular frequency. The opifex implementation maps the standard kernel hyperparameters to output_scale = σ and lengthscale = 1/ω; the quality factor is closed over by this kernel factory.

Parameters:

Name	Type	Description	Default
quality_factor	float	Q > 1/2. Higher Q produces a sharper oscillator (slower envelope decay, longer oscillation train).	required

Returns:

Type	Description
Callable[..., Array]	Standard kernel callable
Callable[..., Array]	(x1, x2, *, lengthscale, output_scale) -> Gram.

Raises:

Type	Description
ValueError	If quality_factor <= 0.5.

Notes

Only the underdamped regime (Q > 1/2) is implemented here. The critically-damped (Q = 1/2) and overdamped (Q < 1/2) forms are documented in the docstring of the tinygp reference ../tinygp/src/tinygp/kernels/quasisep.py:SHO but are deferred to a follow-up slice. The direct-evaluation form here is O(n²) per Gram; the celerite scalable state-space form (O(n) Kalman-style) is a separate Phase-11 deliverable.

deep_kernel

deep_kernel(*, feature_map: Callable[[Array], Array], base_kernel_fn: Callable[..., Array]) -> Callable[..., Array]

Deep-kernel composition k_deep(x, x') = k_base(φ(x), φ(x')).

Wilson, Hu, Salakhutdinov, Xing 2016 (arXiv:1511.02222) introduced Deep Kernel Learning: combine a learnable neural feature map φ_θ with any base kernel k_base so the resulting kernel inherits the base hyperparameters while the NN provides the representation.

The opifex composition is a thin wrapper over the existing kernel_fn API: any callable mapping (n, d) -> (n, d') qualifies as the feature map — Python lambdas, flax.nnx modules, nnx.Sequential extractors, etc. No equinox dependency: opifex uses flax.nnx for the NN component (matching the rest of the opifex neural backbone).

Parameters:

Name	Type	Description	Default
feature_map	Callable[[Array], Array]	Callable that lifts (n, d) -> (n, d').	required
base_kernel_fn	Callable[..., Array]	Standard kernel callable (x1, x2, *, lengthscale, output_scale) -> (n, m).	required

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature.

graph_diffusion_kernel

graph_diffusion_kernel(*, laplacian_eigenvalues: Array, laplacian_eigenvectors: Array) -> Callable[..., Array]

Heat-kernel on a graph via the Laplacian spectral decomposition.

For a graph with normalised Laplacian L = V Λ V^T (orthonormal eigenvectors V, eigenvalues Λ = diag(λ_k) ≥ 0), the diffusion (heat) kernel between nodes i and j is

.. math::

K_{\text{heat}}(i, j) = \sigma_{f}^{2}\,
    \sum_{k} e^{-\lambda_{k}\,\ell^{2}/2}\,
    v_{k}[i]\,v_{k}[j].

The lengthscale plays the role of diffusion time: t = ℓ²/2. Larger ℓ → more diffusion → predictions spread further across the graph (matches the standard "larger lengthscale → broader correlations" GP convention).

The construction is the Gaussian / squared-exponential entry point of the Matern-on-graphs family (Borovitskiy, Mostowsky, Lindgren, Hensman 2020, arXiv:2006.10160); replacing exp(-λ/ℓ²) with the spectral Matern density (2ν/ℓ² + λ)^{-ν - d/2} recovers the full Matern variant.

The returned callable consumes integer node indices in (n, 1) or (n,) shape; values must be in {0, …, num_nodes - 1}.

Parameters:

Name	Type	Description	Default
laplacian_eigenvalues	Array	(N,) array of Laplacian eigenvalues (typically obtained via jnp.linalg.eigh on the normalised Laplacian).	required
laplacian_eigenvectors	Array	(N, N) matrix of orthonormal eigenvectors arranged column-wise.	required

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature
Callable[..., Array]	(x1, x2, *, lengthscale, output_scale) -> Gram.

Raises:

Type	Description
ValueError	If the eigenvalue / eigenvector shapes are inconsistent.

kernel_sum

kernel_sum(*, kernel_fns: tuple[Callable[..., Array], ...]) -> Callable[..., Array]

Pointwise sum of kernel callables k = Σ_q k_q.

A common building block in stationary-kernel composition: sums of Matern12 + celerite-Complex terms reproduce the full celerite family (Foreman-Mackey+ 2017 / Foreman-Mackey 2018); sums of RBF + periodic + linear realise the standard interpretable-GP decomposition (Duvenaud 2014 thesis).

Each component kernel sees the same (x1, x2, lengthscale, output_scale) arguments and contributes its full Gram matrix to the sum. Use a closure to fix per-component parameters that differ from the global (lengthscale, output_scale) (e.g. celerite_complex_kernel already does this for the Complex-term (b, d) parameters).

Parameters:

Name	Type	Description	Default
kernel_fns	tuple[Callable[..., Array], ...]	Tuple of kernel callables, each matching the standard signature.	required

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature.

Raises:

Type	Description
ValueError	If kernel_fns is empty.

matern12_kernel

matern12_kernel(x1: Array, x2: Array, *, lengthscale: float, output_scale: float) -> Array

Matern-½ (exponential / Ornstein-Uhlenbeck) kernel.

.. math::

k(r) = \sigma_{f}^{2}\,\exp\!\left(-\frac{r}{\ell}\right),
\quad r = \lVert x_{1} - x_{2} \rVert_{2}.

matern32_kernel

matern32_kernel(x1: Array, x2: Array, *, lengthscale: float, output_scale: float) -> Array

Matern-3/2 kernel.

.. math::

k(r) = \sigma_{f}^{2}\,
    \left(1 + \frac{\sqrt{3}\,r}{\ell}\right)
    \exp\!\left(-\frac{\sqrt{3}\,r}{\ell}\right).

matern52_kernel

matern52_kernel(x1: Array, x2: Array, *, lengthscale: float, output_scale: float) -> Array

Matern-5/2 kernel.

.. math::

k(r) = \sigma_{f}^{2}\,
    \left(1 + \frac{\sqrt{5}\,r}{\ell}
          + \frac{5\,r^{2}}{3\,\ell^{2}}\right)
    \exp\!\left(-\frac{\sqrt{5}\,r}{\ell}\right).

multi_output_icm_kernel

multi_output_icm_kernel(*, base_kernel_fn: Callable[..., Array], coregionalisation: Array) -> Callable[..., Array]

Intrinsic Coregionalisation Model (ICM) multi-output kernel.

Wraps a scalar base kernel and a coregionalisation matrix B ∈ R^{T × T} into a multi-output kernel of the form k_ICM((x, i), (x', j)) = k_base(x, x') · B[i, j] (Alvarez, Rosasco, Lawrence 2012, §3.1).

The inputs to the returned callable are arrays whose last column is an integer-valued task index in {0, …, T-1} and whose preceding columns are the spatial / feature coordinates that base_kernel_fn consumes. coregionalisation must be a square (T, T) array.

Parameters:

Name	Type	Description	Default
base_kernel_fn	Callable[..., Array]	Scalar kernel callable (x1_features, x2_features, *, lengthscale, output_scale) -> (n, m).	required
coregionalisation	Array	B matrix of shape (T, T) (typically B = W W^T + diag(κ) for low-rank plus diagonal structure).	required

Returns:

Type	Description
Callable[..., Array]	A callable with the same signature as base_kernel_fn but
Callable[..., Array]	consuming task-indexed inputs.

Raises:

Type	Description
ValueError	If coregionalisation is not square.

multi_output_lcm_kernel

multi_output_lcm_kernel(*, components: tuple[tuple[Callable[..., Array], Array], ...]) -> Callable[..., Array]

Linear Coregionalisation Model (LCM) multi-output kernel.

Generalises ICM to a sum of base-kernel / coregionalisation pairs (Alvarez, Rosasco, Lawrence 2012 §3.2):

.. math::

k_{LCM}((x, i), (x', j)) = \sum_{q} k_{q}(x, x')\,B_{q}[i, j].

Each B_q may be low-rank (rank-1 components reproduce the classical co-kriging linear model) or full-rank. With Q = 1 the LCM collapses to a single ICM block.

Parameters:

Name	Type	Description	Default
components	tuple[tuple[Callable[..., Array], Array], ...]	A tuple of (base_kernel_fn, coregionalisation) pairs. Each coregionalisation must be a square (T, T) array; all components must share the same T.	required

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature
Callable[..., Array]	(x1, x2, *, lengthscale, output_scale) -> Gram.

Raises:

Type	Description
ValueError	If components is empty, if any coregionalisation is non-square, or if the components disagree on T.

orthogonal_additive_kernel

orthogonal_additive_kernel(*, base_kernel_fns: tuple[Callable[..., Array], ...], max_order: int, order_variances: Array) -> Callable[..., Array]

Higher-order Orthogonal Additive Kernel (Lu+ 2022 ICML).

Extends :func:additive_kernel (max_order = 1) to the full OAK construction:

.. math::

k_{\text{OAK}}(\mathbf{x}, \mathbf{x}')
  = \sum_{\ell = 0}^{D_{\text{tilde}}}\sigma_{\ell}^{2}\,
        e_{\ell}\!\bigl(
            \tilde{k}_{1}(x_{1}, x'_{1}),\dots,\tilde{k}_{D}(x_{D}, x'_{D})
        \bigr),

where D_{tilde} = max_order, e_ℓ is the ℓ-th elementary symmetric polynomial of the D per-dimension constrained base kernels (orthogonal to constants under the chosen input measure), and σ²_ℓ ≥ 0 are the per-order variances. e_ℓ is evaluated in O(D · max_order) via the Newton-Girard recursion e_n = (1/n) Σ_{k=1}^{n} (-1)^{k-1} e_{n-k} s_k with s_k = Σ_d k̃_d^k.

Use :func:constrained_rbf_kernel for the canonical Gaussian- measure base; any callable matching the standard kernel signature and integrating to zero against the chosen input measure works.

Parameters:

Name	Type	Description	Default
base_kernel_fns	tuple[Callable[..., Array], ...]	Tuple of D per-dimension base kernels. Each must satisfy the standard kernel signature (x1, x2, *, lengthscale, output_scale) -> (n, m) and be orthogonal to constants under the desired input measure (use :func:constrained_rbf_kernel for the Gaussian-measure default).	required
max_order	int	Maximum interaction order 0 ≤ max_order ≤ D. max_order = 1 recovers :func:additive_kernel; max_order = D enables full ANOVA decomposition.	required
order_variances	Array	(max_order + 1,) non-negative variances σ²_0, σ²_1, …, σ²_{max_order}. σ²_0 is the constant-offset variance; σ²_1 weights first-order terms; σ²_D weights the full-product interaction.	required

Returns:

Type	Description
Callable[..., Array]	A callable matching the standard kernel signature.

Raises:

Type	Description
ValueError	If base_kernel_fns is empty, max_order is outside [0, D], or order_variances has the wrong shape.

fit_laplace_gp

fit_laplace_gp(*, log_likelihood_components_fn: LikelihoodComponentsFn, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, num_newton_iterations: int = 50, kernel_fn: Callable[..., Array] = rbf_kernel) -> LaplaceGPState

Fit a Laplace-approximated GP for any factorising likelihood.

Runs num_newton_iterations of RW06 Algorithm 3.1's Newton update (factorised through the B = I + W^{1/2} K W^{1/2} Cholesky for PSD stability) and returns a :class:LaplaceGPState consumable by :func:predict_laplace_latent_moments.

Parameters:

Name	Type	Description	Default
log_likelihood_components_fn	LikelihoodComponentsFn	Callable (f, y) -> (log_lik, ∇log_lik, W, √W) for the chosen likelihood. log_lik is a scalar; the rest are (n,) arrays. W must be non-negative element-wise; each wrapper clips internally before returning.	required
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) training targets in the likelihood's support.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
num_newton_iterations	int	Fixed Newton-loop count. Static under jax.jit so the jax.lax.scan shape is known at trace time. Defaults to 50.	50
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Type	Description
LaplaceGPState	class:LaplaceGPState carrying f_mode, the Cholesky
LaplaceGPState	factor of B at the mode, and the Laplace-approximated log
LaplaceGPState	marginal likelihood (RW06 eq. 3.32).

predict_laplace_latent_moments

predict_laplace_latent_moments(*, state: LaplaceGPState, x_test: Array) -> tuple[Array, Array]

Compute latent Gaussian-posterior moments at x_test.

RW06 Algorithm 3.2:

.. math::

\bar{f}_*       &= k_*^{T}\,\nabla \log p(y \mid \hat{f}), \\
\mathrm{Var}(f_*) &= k(x_*, x_*)
    - (W^{1/2} k_*)^{T}\,B^{-1}\,(W^{1/2} k_*).

The per-likelihood wrappers feed these moments through the appropriate response-mean / response-variance map (MacKay probit for Bernoulli, exp(μ + V/2) for Poisson with log link, ...).

Parameters:

Name	Type	Description	Default
state	LaplaceGPState	Fitted :class:LaplaceGPState.	required
x_test	Array	(m, d) test inputs.	required

Returns:

Type	Description
Array	(latent_mean, latent_variance) each of shape (m,); the
Array	variance is clipped to ≥ 1e-12 for numerical safety.

fit_bernoulli_laplace_gp

fit_bernoulli_laplace_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, num_newton_iterations: int = 50, kernel_fn: Callable[..., Array] = rbf_kernel) -> LaplaceGPState

Fit a binary GP classifier via the Laplace approximation.

Thin Bernoulli-specific wrapper around :func:opifex.uncertainty.gp.laplace.fit_laplace_gp that supplies the (log_lik, ∇log_lik, W = σ(f)(1 - σ(f))) triple from RW06 §3.4.

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) targets in {-1, +1}.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
num_newton_iterations	int	Fixed Newton-loop count. Defaults to 50. Static under jax.jit.	50
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Type	Description
LaplaceGPState	class:LaplaceGPState carrying f_mode, the Cholesky
LaplaceGPState	factor of B at the mode, and the Laplace-approximated log
LaplaceGPState	marginal likelihood.

predict_bernoulli_laplace_gp

predict_bernoulli_laplace_gp(*, state: LaplaceGPState, x_test: Array) -> PredictiveDistribution

Predict class probabilities at x_test via MacKay's approximation.

RW06 Algorithm 3.2 gives the latent predictive moments

.. math::

\bar{f}_*       &= k_*^{T}\,\nabla \log p(y \mid \hat{f}), \\
\mathrm{Var}(f_*) &= k(x_*, x_*)
    - (W^{1/2} k_*)^{T}\,B^{-1}\,(W^{1/2} k_*),

obtained via :func:predict_laplace_latent_moments. The class probability uses MacKay's probit approximation (RW06 eq. 3.25):

.. math::

p(y_* = +1 \mid x_*) \approx \sigma\!\left(
    \frac{\bar{f}_*}{\sqrt{1 + \pi\,\mathrm{Var}(f_*) / 8}}
\right).

Parameters:

Name	Type	Description	Default
state	LaplaceGPState	Fitted :class:LaplaceGPState.	required
x_test	Array	(m, d) test inputs.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean is
PredictiveDistribution	p(y_*=+1 | x_*) (in [0, 1]) and whose variance is
PredictiveDistribution	the latent Var(f_*) (a useful epistemic-uncertainty
PredictiveDistribution	proxy). Metadata records the source paper.

fit_beta_laplace_gp

fit_beta_laplace_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, scale: float = 10.0, num_newton_iterations: int = 50, kernel_fn: Callable[..., Array] = rbf_kernel) -> LaplaceGPState

Fit a Beta GP regressor for proportion data (logit link).

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) observations in (0, 1).	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
scale	float	Beta precision s = α + β > 0 (gpflow / bayesnewton convention). Larger s ⇒ narrower per-point Beta. Defaults to 10.0.	10.0
num_newton_iterations	int	Fixed Newton-loop count. Defaults to 50. Static under jax.jit.	50
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Type	Description
LaplaceGPState	class:LaplaceGPState ready for
LaplaceGPState	func:predict_beta_laplace_gp.

fit_poisson_laplace_gp

fit_poisson_laplace_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, num_newton_iterations: int = 50, kernel_fn: Callable[..., Array] = rbf_kernel) -> LaplaceGPState

Fit a Poisson GP regressor (count data, exp link).

Thin wrapper around :func:opifex.uncertainty.gp.laplace.fit_laplace_gp that supplies the Poisson (log_lik, ∇log_lik, W = exp(f)) triple.

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) non-negative count observations.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
num_newton_iterations	int	Fixed Newton-loop count. Defaults to 50. Static under jax.jit.	50
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Type	Description
LaplaceGPState	class:LaplaceGPState ready for
LaplaceGPState	func:predict_poisson_laplace_gp.

fit_studentst_laplace_gp

fit_studentst_laplace_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, df: float = 4.0, scale: float = 1.0, num_newton_iterations: int = 50, kernel_fn: Callable[..., Array] = rbf_kernel) -> LaplaceGPState

Fit a Student-t robust GP regressor (identity link).

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) real-valued observations.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
df	float	Student-t degrees of freedom ν > 0. Smaller values (e.g. ν = 3 - 5) yield heavier tails and more outlier robustness. Defaults to 4.0.	4.0
scale	float	Student-t scale σ > 0. Defaults to 1.0.	1.0
num_newton_iterations	int	Fixed Newton-loop count. Defaults to 50. Static under jax.jit.	50
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel

Returns:

Type	Description
LaplaceGPState	class:LaplaceGPState ready for
LaplaceGPState	func:predict_studentst_laplace_gp.

predict_beta_laplace_gp

predict_beta_laplace_gp(*, state: LaplaceGPState, x_test: Array, scale: float = 10.0) -> PredictiveDistribution

Predict Beta response moments at x_test (logit link).

Approximates E[σ(f*)] by MacKay's probit-style sigmoid approximation:

.. math::

\mathbb{E}[σ(f_*)] \approx σ\!\left(
    \frac{\mu}{\sqrt{1 + π\,V / 8}}
\right),

then returns the implied Beta variance Var[y*] = m̂ (1 - m̂) / (s + 1) at that mean.

Parameters:

Name	Type	Description	Default
state	LaplaceGPState	Fitted :class:LaplaceGPState.	required
x_test	Array	(m, d) test inputs.	required
scale	float	Beta precision used at fit time. Defaults to 10.0.	10.0

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean is
PredictiveDistribution	E[y* | x*] ∈ (0, 1) and whose variance is the implied
PredictiveDistribution	Beta variance. epistemic carries the latent Var f_*.

predict_poisson_laplace_gp

predict_poisson_laplace_gp(*, state: LaplaceGPState, x_test: Array) -> PredictiveDistribution

Predict Poisson intensity moments at x_test.

Under the latent Gaussian posterior f* ~ N(μ, V) and the exp link, the predicted rate λ* = exp(f*) is log-normal:

.. math::

\mathbb{E}[λ_*] &= \exp\!\left(\mu + \tfrac{1}{2} V\right), \\
\mathrm{Var}[λ_*] &= \mathbb{E}[λ_*]^{2}\,
    \bigl(\mathrm{e}^{V} - 1\bigr).

Parameters:

Name	Type	Description	Default
state	LaplaceGPState	Fitted :class:LaplaceGPState.	required
x_test	Array	(m, d) test inputs.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean is the
PredictiveDistribution	predicted Poisson intensity E[λ*] and whose variance
PredictiveDistribution	is the rate variance. epistemic carries the latent
PredictiveDistribution	Var f_*.

predict_studentst_laplace_gp

predict_studentst_laplace_gp(*, state: LaplaceGPState, x_test: Array, df: float = 4.0, scale: float = 1.0) -> PredictiveDistribution

Predict Student-t response moments at x_test.

Under the latent Gaussian f* ~ N(μ, V) and a Student-t likelihood with location f* and scale σ, the conditional mean is f* (for ν > 1) and the conditional variance is σ² ν / (ν - 2) (for ν > 2). By the law of total variance:

.. math::

\mathbb{E}[y_*]   &= \mu, \\
\mathrm{Var}[y_*] &= V + \frac{\sigma^{2}\,\nu}{\nu - 2}.

Parameters:

Name	Type	Description	Default
state	LaplaceGPState	Fitted :class:LaplaceGPState.	required
x_test	Array	(m, d) test inputs.	required
df	float	Student-t degrees of freedom used at fit time (ν > 2 required for finite predictive variance). Defaults to 4.0.	4.0
scale	float	Student-t scale used at fit time. Defaults to 1.0.	1.0

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean is the
PredictiveDistribution	predictive E[y*] and whose variance is the total
PredictiveDistribution	predictive variance. epistemic carries the latent
PredictiveDistribution	Var f_*.

fit_quasisep_sho_gp

fit_quasisep_sho_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, noise_std: float, quality_factor: float) -> QuasisepGPState

Fit the scalable SHO state-space GP in O(n).

Maps the SHO hyperparameters into a 2-state linear-Gaussian SDE and evaluates the exact log marginal likelihood via opifex's Kalman primitives.

Parameters:

Name	Type	Description	Default
x_train	Array	(n,) or (n, 1) strictly-increasing training times.	required
y_train	Array	(n,) training observations.	required
lengthscale	float	ℓ > 0 — kernel length-scale.	required
output_scale	float	σ_f > 0 — kernel output-scale.	required
noise_std	float	σ > 0 — observation noise scale.	required
quality_factor	float	Q > 1/2 — SHO quality factor (underdamped regime only in this slice).	required

Returns:

Type	Description
QuasisepGPState	class:QuasisepGPState carrying the data, hyperparameters
QuasisepGPState	and exact log marginal likelihood.

Raises:

Type	Description
ValueError	If hyperparameters are non-positive or quality_factor ≤ 1/2, or if x_train is not strictly increasing (checked at trace-time only).

predict_quasisep_sho_gp

predict_quasisep_sho_gp(*, state: QuasisepGPState, x_test: Array) -> PredictiveDistribution

Posterior moments at x_test via augmented filter + smoother.

Augments x_train ∪ x_test into a single sorted sequence, runs the forward Kalman filter and backward RTS smoother on it, and extracts the smoothed first-state moments at the test positions. Test entries are marked as "missing" via a large observation covariance so the Kalman gain at those points is numerically zero — equivalent to running the joint GP posterior conditioned on the training data only.

Parameters:

Name	Type	Description	Default
state	QuasisepGPState	Fitted :class:QuasisepGPState.	required
x_test	Array	(m,) or (m, 1) test times (any order).	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean and
PredictiveDistribution	variance carry the latent E[f(x*)] and Var[f(x*)]
PredictiveDistribution	for each test point. Metadata records the source paper.

fit_quasisep_carma_gp

fit_quasisep_carma_gp(*, x_train: Array, y_train: Array, ar_coefficients: Array, ma_coefficients: Array, lengthscale: float, output_scale: float | Array, noise_std: float) -> QuasisepCarmaGPState

Fit the scalable CARMA state-space GP in O(n).

Maps the CARMA(p, q) coefficients into the celerite real linear-Gaussian SDE and evaluates the exact log marginal likelihood via opifex's Kalman primitives.

Parameters:

Name	Type	Description	Default
x_train	Array	(n,) or (n, 1) strictly-increasing training times.	required
y_train	Array	(n,) training observations.	required
ar_coefficients	Array	(p,) autoregressive coefficients α_0, …, α_{p-1} (implicit α_p = 1). For stationarity all AR roots must have negative real parts (NaN propagates otherwise).	required
ma_coefficients	Array	(q + 1,) moving-average coefficients β_0, …, β_q with q + 1 ≤ p.	required
lengthscale	float	ℓ > 0 — rescales the lag (τ → τ / ℓ).	required
output_scale	float | Array	σ_f > 0 — kernel output-scale.	required
noise_std	float	σ > 0 — observation noise scale.	required

Returns:

Type	Description
QuasisepCarmaGPState	class:QuasisepCarmaGPState carrying the data, coefficients,
QuasisepCarmaGPState	hyperparameters and exact log marginal likelihood.

Raises:

Type	Description
ValueError	If q + 1 > p, if hyperparameters are non-positive, or if x_train is not strictly increasing (checked at trace-time only).

predict_quasisep_carma_gp

predict_quasisep_carma_gp(*, state: QuasisepCarmaGPState, x_test: Array) -> PredictiveDistribution

Posterior moments at x_test via augmented filter + smoother.

Augments x_train ∪ x_test into a single sorted sequence, runs the forward Kalman filter and backward RTS smoother on it, and extracts the smoothed observed-channel moments at the test positions. Test entries are marked "missing" via a large observation covariance so the Kalman gain at those points is numerically zero — equivalent to the joint GP posterior conditioned on the training data only.

Parameters:

Name	Type	Description	Default
state	QuasisepCarmaGPState	Fitted :class:QuasisepCarmaGPState.	required
x_test	Array	(m,) or (m, 1) test times (any order).	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean and variance
PredictiveDistribution	carry E[f(x*)] and Var[f(x*)] for each test point.

fit_rff_gp

fit_rff_gp(*, x_train: Array, y_train: Array, lengthscale: float, output_scale: float, noise_std: float, num_features: int, rngs: Rngs | Array) -> RFFGPState

Fit the RFF-approximate conjugate GP (Rahimi & Recht 2007).

Solves the ridge-regression normal equations (Φ^T Φ + σ² I) α = Φ^T y via Cholesky.

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) targets.	required
lengthscale	float	RBF length-scale.	required
output_scale	float	RBF output-scale.	required
noise_std	float	Observation noise scale σ (also acts as the ridge jitter).	required
num_features	int	Feature count D (even, strictly positive).	required
rngs	Rngs | Array	Caller-owned RNG for the spectral sample.	required

Returns:

Type	Description
RFFGPState	class:RFFGPState carrying the feature map data + Cholesky
RFFGPState	factor + pre-solved α.

Raises:

Type	Description
ValueError	If noise_std is non-positive.

predict_rff_gp

predict_rff_gp(*, state: RFFGPState, x_test: Array) -> PredictiveDistribution

Predict at x_test using a fitted :class:RFFGPState.

The lifted predictive moments

.. math::

\mu(x^{*})  &= \phi(x^{*})^{T}\,\alpha, \\
\mathrm{Var}(x^{*}) &= \sigma^{2}\,
    \phi(x^{*})^{T}\,(\Phi^{T}\Phi + \sigma^{2} I)^{-1}\,\phi(x^{*}),

are evaluated via the Cholesky factor of the ridge-Gram.

Parameters:

Name	Type	Description	Default
state	RFFGPState	Fitted :class:RFFGPState.	required
x_test	Array	(m, d) test inputs.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution with mean (m,),
PredictiveDistribution	variance (m,), epistemic == variance, and metadata
PredictiveDistribution	advertising estimator=rff_gp plus the Rahimi & Recht
PredictiveDistribution	citation.

rbf_random_fourier_features

rbf_random_fourier_features(*, x: Array, lengthscale: float, output_scale: float, num_features: int, rngs: Rngs | Array) -> Array

Construct the RFF feature map for the RBF kernel.

Parameters:

Name	Type	Description	Default
x	Array	(n, d) inputs.	required
lengthscale	float	RBF length-scale ℓ.	required
output_scale	float	RBF output-scale σ_f.	required
num_features	int	D — total feature count. Must be a strictly positive even integer (cos and sin contribute a pair per ω).	required
rngs	Rngs | Array	Caller-owned nnx.Rngs or raw key.	required

Returns:

Type	Description
Array	Feature matrix (n, num_features) such that
Array	φ(x)^T φ(x') ≈ k_RBF(x, x') in expectation.

Raises:

Type	Description
ValueError	If num_features is non-positive or odd.

fit_spatiotemporal_vgp

fit_spatiotemporal_vgp(*, times: Array, space: Array, observations: Array, temporal_kernel: StateSpaceKernel, spatial_lengthscale: float, spatial_output_scale: float, noise_std: Array | float) -> SpatioTemporalGPState

Fit the conjugate spatio-temporal variational GP on gridded data.

Runs the Kronecker-lifted Kalman filter and RTS smoother over the T time steps; the observation at each step is the M-vector of spatial measurements with diagonal Gaussian noise σ^2 I_M (bayesnewton MarkovGaussianProcess.update_posterior line 692, specialised to the conjugate Gaussian case where the pseudo likelihood equals the true likelihood).

Parameters:

Name	Type	Description	Default
times	Array	Time stamps of shape (T, 1) (or (T,)), assumed sorted ascending.	required
space	Array	Spatial locations of shape (M, d).	required
observations	Array	Field observations of shape (T, M).	required
temporal_kernel	StateSpaceKernel	A :class:StateSpaceKernel temporal prior.	required
spatial_lengthscale	float	RBF spatial length-scale (strictly positive).	required
spatial_output_scale	float	RBF spatial output-scale (strictly positive).	required
noise_std	Array | float	Observation noise scale σ (strictly positive).	required

Returns:

Name	Type	Description
A	SpatioTemporalGPState	class:SpatioTemporalGPState carrying the smoothed latent
	SpatioTemporalGPState	posterior and the lifted measurement model.

Raises:

Type	Description
ValueError	If a static (non-traced) hyperparameter is non-positive, or observations is not (T, M).

predict_spatiotemporal_vgp

predict_spatiotemporal_vgp(*, state: SpatioTemporalGPState, include_observation_noise: bool = False) -> PredictiveDistribution

Predict the field at the training space-time grid from a fitted state.

Projects the smoothed latent state to function space at every time via the measurement model f(t) = H x(t) and extracts the per-point marginal variance \mathrm{diag}(H P(t) H^\top) (bayesnewton MarkovGaussianProcess.predict line 800, spatio-temporal branch: test_var = diag(W P W^\top) with W = H for inducing-at-data).

Parameters:

Name	Type	Description	Default
state	SpatioTemporalGPState	A fitted :class:SpatioTemporalGPState.	required
include_observation_noise	bool	When True, adds σ^2 to each marginal variance to form the observation predictive (p(y_*)); when False (default) returns the latent predictive (p(f_*)).	False

Returns:

Name	Type	Description
A	PredictiveDistribution	class:PredictiveDistribution whose mean and variance
	PredictiveDistribution	are shaped (T, M).

separable_spatiotemporal_kernel

separable_spatiotemporal_kernel(*, times: Array, space: Array, temporal_kernel: StateSpaceKernel, spatial_lengthscale: float, spatial_output_scale: float) -> Array

Dense separable space-time covariance on the flattened (T, M) grid.

Returns the (T M, T M) Gram matrix of the separable kernel k((t, R), (t', R')) = k_t(t, t') k_s(R, R') evaluated on the Cartesian product of times and space, ordered with the spatial index varying fastest (time-major flattening). On a regular grid this equals the Kronecker product K_t \otimes K_s (bayesnewton SpatioTemporalKernel.K line 462).

The temporal Gram matrix is reconstructed from the state-space kernel via its stationary covariance and closed-form state transition: k_t(t, t') = H P_\infty A(|t - t'|)^\top H^\top (Sarkka 2013 §6; bayesnewton Kernel.K for SDE kernels).

Parameters:

Name	Type	Description	Default
times	Array	Time stamps of shape (T, 1) (or (T,)).	required
space	Array	Spatial locations of shape (M, d).	required
temporal_kernel	StateSpaceKernel	A :class:StateSpaceKernel temporal prior.	required
spatial_lengthscale	float	RBF spatial length-scale (strictly positive).	required
spatial_output_scale	float	RBF spatial output-scale (strictly positive).	required

Returns:

Type	Description
Array	(T M, T M) separable covariance matrix.

spatiotemporal_vgp_log_marginal

spatiotemporal_vgp_log_marginal(*, times: Array, space: Array, observations: Array, temporal_kernel: StateSpaceKernel, spatial_lengthscale: float, spatial_output_scale: float, noise_std: Array | float) -> Array

Marginal log-likelihood of the gridded observations under the ST-VGP.

Equals the Kalman innovation log-likelihood of the Kronecker-lifted state-space model (bayesnewton compute_log_lik line 729 — the conjugate-Gaussian case where the pseudo likelihood is exact). Used for hyperparameter learning via jax.grad.

Parameters:

Name	Type	Description	Default
times	Array	Time stamps of shape (T, 1) (or (T,)).	required
space	Array	Spatial locations of shape (M, d).	required
observations	Array	Field observations of shape (T, M).	required
temporal_kernel	StateSpaceKernel	A :class:StateSpaceKernel temporal prior.	required
spatial_lengthscale	float	RBF spatial length-scale (strictly positive).	required
spatial_output_scale	float	RBF spatial output-scale (strictly positive).	required
noise_std	Array | float	Observation noise scale σ (strictly positive).	required

Returns:

Type	Description
Array	Scalar marginal log-likelihood log p(Y).

fit_svgp

fit_svgp(*, x_train: Array, y_train: Array, x_inducing: Array, lengthscale: float, output_scale: float, noise_std: float, kernel_fn: Callable[..., Array] = rbf_kernel, jitter: float = 1e-06) -> SVGPState

Fit the Titsias-collapsed sparse-variational GP.

Parameters:

Name	Type	Description	Default
x_train	Array	(n, d) training inputs.	required
y_train	Array	(n,) training targets.	required
x_inducing	Array	(m, d) inducing inputs Z. Choose m << n for the scalability benefit; m == n recovers the exact GP up to numerical jitter.	required
lengthscale	float	Kernel length-scale.	required
output_scale	float	Kernel output-scale.	required
noise_std	float	Observation noise scale σ (strictly positive).	required
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel
jitter	float	Numerical jitter added to K_zz for PD stability.	1e-06

Returns:

Type	Description
SVGPState	class:SVGPState carrying the L_z / L_B Cholesky
SVGPState	factors + pre-solved coefficient vector + cached ELBO terms.

Raises:

Type	Description
ValueError	If noise_std is not strictly positive.

predict_svgp

predict_svgp(*, state: SVGPState, x_test: Array) -> PredictiveDistribution

Predictive distribution at x_test (Titsias 2009 SOR/DTC form).

Parameters:

Name	Type	Description	Default
state	SVGPState	Fitted :class:SVGPState.	required
x_test	Array	(t, d) test inputs.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution with mean (t,),
PredictiveDistribution	variance (t,), epistemic == variance, and metadata
PredictiveDistribution	advertising estimator=titsias_collapsed_svgp.

svgp_collapsed_elbo

svgp_collapsed_elbo(*, state: SVGPState) -> Array

Collapsed ELBO for a fitted :class:SVGPState (scalar, jit-safe).

Computes Titsias 2009 eq. 9 / GPJax collapsed_elbo exactly via the cached A-Cholesky factorisation. Useful as the objective for hyperparameter optimisation when wrapped under jax.grad.

bernoulli_log_likelihood

bernoulli_log_likelihood(f: Array, y: Array) -> Array

log p(y_i | f_i) = log σ(y_i f_i) for y ∈ {-1, +1} (RW06 §3.4).

Uses :func:jax.nn.log_sigmoid for overflow / underflow safety.

init_stochastic_svgp_state

init_stochastic_svgp_state(*, x_inducing: Array, lengthscale: float, output_scale: float, log_likelihood_fn: LogLikelihoodFn, kernel_fn: Callable[..., Array] = rbf_kernel, jitter: float = 1e-06) -> StochasticSVGPState

Initialise q(u) = N(0, K_zz) (μ_w = 0, L_w = I).

Parameters:

Name	Type	Description	Default
x_inducing	Array	(m, d) inducing inputs.	required
lengthscale	float	Kernel length-scale (strictly positive).	required
output_scale	float	Kernel output-scale (strictly positive).	required
log_likelihood_fn	LogLikelihoodFn	Per-observation log-likelihood callable (f, y) -> log p(y | f) evaluated element-wise. Use :func:bernoulli_log_likelihood / :func:poisson_log_likelihood for the canonical cases or close over additional hyperparameters via a factory.	required
kernel_fn	Callable[..., Array]	Kernel callable. Defaults to :func:rbf_kernel.	rbf_kernel
jitter	float	Numerical jitter for K_zz factorisation.	1e-06

Returns:

Type	Description
StochasticSVGPState	class:StochasticSVGPState ready for gradient-based ELBO
StochasticSVGPState	optimisation.

Raises:

Type	Description
ValueError	If lengthscale or output_scale is non- positive or x_inducing is not 2-D.

natural_gradient_step

natural_gradient_step(*, state: StochasticSVGPState, x_batch: Array, y_batch: Array, dataset_size: int, learning_rate: float, num_quadrature_points: int = 20) -> StochasticSVGPState

One Salimbeni+ 2018 natural-gradient update of (μ_w, L_w).

For the whitened variational Gaussian q(u_w) = N(μ_w, S_w) with S_w = L_w L_w^{T}, the natural parameters are

.. math::

\theta_{1} = S_{w}^{-1}\,\mu_{w},
\qquad \theta_{2} = -\tfrac{1}{2}\,S_{w}^{-1},

and the expectation parameters are η_1 = μ_w, η_2 = S_w + μ_w μ_w^{T}. Salimbeni, Eleftheriadis, Hensman 2018 (AISTATS — Natural Gradients in Practice) shows that the natural-gradient step on θ equals the regular gradient on η:

.. math::

\theta_{1}^{\text{new}} &= \theta_{1} + \rho\,\bigl(
    \partial \mathcal{L}/\partial \mu_{w}
    - 2\,(\partial \mathcal{L}/\partial S_{w})\,\mu_{w}\bigr),\\
\theta_{2}^{\text{new}} &= \theta_{2} + \rho\,
    \partial \mathcal{L}/\partial S_{w}.

After the update we recover (μ_w, L_w) via S_w = (-2 \theta_{2})^{-1}, μ_w = S_w \theta_{1}, L_w = chol(S_w). The natural-gradient direction is the Fisher-information-preconditioned direction; Salimbeni+ 2018 §4 reports 10x-100x faster convergence than Adam on the variational parameters for non-Gaussian likelihoods.

Performance: each step costs one regular jax.grad pass plus three O(m^{3}) linear-algebra operations (S_w inverse, S_w_new recovery, Cholesky). For m up to a few hundred, this is dominated by the gradient pass.

Parameters:

Name	Type	Description	Default
state	StochasticSVGPState	Current variational state.	required
x_batch	Array	(b, d) minibatch inputs.	required
y_batch	Array	(b,) minibatch targets.	required
dataset_size	int	Full-dataset size N for the ELBO scaling.	required
learning_rate	float	Natural-gradient step size ρ. For the Gaussian-likelihood + conjugate case, ρ = 1 reaches the variational optimum in a single step (Salimbeni+ 2018 Algorithm 1). For non-Gaussian likelihoods use ρ ≈ 0.1 - 1.0 and decay over iterations.	required
num_quadrature_points	int	Gauss-Hermite quadrature nodes for the ELBO data-fit term. Defaults to 20.	20

Returns:

Type	Description
StochasticSVGPState	class:StochasticSVGPState with updated (μ_w, L_w);
StochasticSVGPState	every other field is unchanged.

poisson_log_likelihood

poisson_log_likelihood(f: Array, y: Array) -> Array

log p(y_i | f_i) = y_i f_i - exp(f_i) - log Γ(y_i + 1) (exp link).

Matches the Poisson likelihood used by D5 (:func:opifex.uncertainty.gp.fit_poisson_laplace_gp).

predict_stochastic_svgp

predict_stochastic_svgp(*, state: StochasticSVGPState, x_test: Array) -> PredictiveDistribution

Closed-form latent predictive q(f(x_*)) (Hensman+ 2015 eq. 8-9).

Reuses :func:_latent_marginal_moments so the predictive moments are computed in O(m^2 t) after a single chol(K_zz).

Parameters:

Name	Type	Description	Default
state	StochasticSVGPState	Trained :class:StochasticSVGPState.	required
x_test	Array	(t, d) test inputs.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean and
PredictiveDistribution	variance carry the latent marginal moments. Map through
PredictiveDistribution	the appropriate response link (MacKay probit for classification,
PredictiveDistribution	exp for Poisson, identity for regression) downstream.

stochastic_svgp_elbo

stochastic_svgp_elbo(*, state: StochasticSVGPState, x_batch: Array, y_batch: Array, dataset_size: int, num_quadrature_points: int = 20) -> Array

Hensman+ 2013 minibatched ELBO (scalar, jit-safe).

ELBO = (N / |B|) Σ_{i ∈ B} E_q[log p(y_i | f_i)] - KL[q(u) || p(u)] with the whitened KL and a Gauss-Hermite-quadrature data-fit. Exactly one chol(K_zz) per call.

Parameters:

Name	Type	Description	Default
state	StochasticSVGPState	Variational state to be differentiated through.	required
x_batch	Array	(b, d) minibatch inputs.	required
y_batch	Array	(b,) minibatch targets.	required
dataset_size	int	N — full-dataset size. The data-fit term is scaled by N / b so the ELBO is an unbiased estimate of the full-data ELBO.	required
num_quadrature_points	int	Q — Gauss-Hermite quadrature nodes (static; defaults to 20). Static under jax.jit.	20

Returns:

Type	Description
Array	Scalar ELBO array.

State-space math primitives for opifex Kalman filtering and smoothing.

Provides the math layer for opifex.uncertainty.assimilation (Task 6.7 applied data-assimilation layer). Pure JAX; no NNX imports.

The sibling repositories /mnt/ssd2/Works/{bayesnewton,kalman-jax,markovflow, ComputationAwareKalman.jl} serve as reference implementations only — opifex never carries them as runtime dependencies. Algorithms here are JAX-native and cite the sibling repo line-by-line.

References

• Kalman 1960 — A New Approach to Linear Filtering and Prediction Problems.
• Rauch, Tung, Striebel 1965 — Maximum Likelihood Estimates of Linear Dynamic Systems, AIAA J.
• Särkkä 2013 — Bayesian Filtering and Smoothing.
• Pförtner, Wenger, Cockayne, Hennig arXiv:2405.08971 — Compute-Aware Kalman Filtering and Smoothing (primary CAKF/CAKS reference).

LowRankDowndatedMatrix dataclass

LowRankDowndatedMatrix(*, dense: Array, left: Array, right: Array)

Implicit A - U V^T matrix for low-rank-corrected covariances.

Attributes:

Name	Type	Description
dense	Array	The base dense matrix A of shape (n, n) (typically the prior marginal covariance Σ).
left	Array	Left low-rank factor U of shape (n, r).
right	Array	Right low-rank factor V of shape (n, r).

StateSpaceKernel dataclass

StateSpaceKernel(*, feedback: Array, noise_effect: Array, diffusion: Array, measurement: Array, stationary_cov: Array, state_transition: Callable[[Array], Array])

Continuous-time SDE representation of a stationary temporal GP kernel.

Attributes:

Name	Type	Description
feedback	Array	Drift matrix F (shape (n, n)).
noise_effect	Array	Dispersion matrix L (shape (n, k)).
diffusion	Array	Wiener diffusion Q_c (shape (k, k)).
measurement	Array	Observation matrix H (shape (1, n)).
stationary_cov	Array	Stationary covariance P_inf (shape (n, n)), satisfying F P_inf + P_inf F^T + L Q_c L^T = 0.
state_transition	Callable[[Array], Array]	Closed-form A(dt) = exp(F dt).

state_dim property

state_dim: int

Dimension of the state-space representation.

cakf_predict

cakf_predict(*, mean: Array, factor: Array, transition: Array) -> tuple[Array, Array]

Propagate the CAKF state through a transition matrix.

Ports ComputationAwareKalman.jl/src/filter/predict.jl. The prior marginal covariance Σ_{k+1} is determined by the Gauss-Markov chain externally — only the mean and the low-rank correction factor flow through the transition.

Parameters:

Name	Type	Description	Default
mean	Array	Filter mean at time k with shape (n,).	required
factor	Array	Low-rank correction factor M with shape (n, r).	required
transition	Array	Transition matrix A_k with shape (n, n).	required

Returns:

Type	Description
tuple[Array, Array]	Predicted (mean, factor) for time k + 1.

cakf_smooth

cakf_smooth(*, filter_means: Array, filter_covs: Array, transitions: Array, process_noises: Array) -> tuple[Array, Array]

CAKS Rauch-Tung-Striebel backward smoother (Pförtner+ 2024).

Ports ../ComputationAwareKalman.jl/src/smoother/loop.jl. Runs the standard RTS recursion

.. math::

\hat{\mu}_t &= \mu_t + G_t (\hat{\mu}_{t+1} - A_t \mu_t),\\
\hat{\Sigma}_t &= \Sigma_t + G_t (\hat{\Sigma}_{t+1}
                                  - A_t \Sigma_t A_t^T - Q_t) G_t^T,\\
G_t &= \Sigma_t A_t^T (A_t \Sigma_t A_t^T + Q_t)^{-1},

over the forward CAKF-filter outputs. The CAKF filter pass produces posterior moments that already reflect computation-aware uncertainty (the implicit Σ_t = Σ_prior_t - M_t M_t^T is the posterior cov returned in dense form to this smoother). Future slices may add a low-rank smoothed-factor variant that keeps the cov in :class:LowRankDowndatedMatrix form throughout.

Parameters:

Name	Type	Description	Default
filter_means	Array	Forward-pass posterior means, shape (num_steps, state_dim).	required
filter_covs	Array	Forward-pass posterior covariances, shape (num_steps, state_dim, state_dim).	required
transitions	Array	Per-step transition matrices, shape (num_steps, state_dim, state_dim).	required
process_noises	Array	Per-step process-noise covariances, shape (num_steps, state_dim, state_dim).	required

Returns:

Type	Description
Array	(smoothed_means, smoothed_covs) matching the
Array	func:opifex.uncertainty.statespace.kalman_smoother shape.

cakf_step

cakf_step(*, mean: Array, factor: Array, transition: Array, prior_cov: Array, observation: Array, observation_matrix: Array, observation_cov: Array, max_iter: int) -> tuple[Array, Array]

Fused CAKF predict + update step (one full filter iteration).

Composes :func:cakf_predict and :func:cakf_update in the order used by ../ComputationAwareKalman.jl/src/filter/loop.jl. Useful in :func:jax.lax.scan-driven filter loops where the per-step operation is a single (mean, factor) -> (mean, factor) map.

Parameters:

Name	Type	Description	Default
mean	Array	Filter mean at time k with shape (n,).	required
factor	Array	Low-rank correction factor M with shape (n, r0).	required
transition	Array	Discrete transition matrix A_k, shape (n, n).	required
prior_cov	Array	Marginal prior covariance Σ_{k+1} after the transition (provided externally by the Gauss-Markov chain), shape (n, n).	required
observation	Array	Observed value at time k + 1, shape (p,).	required
observation_matrix	Array	Observation operator H, shape (p, n).	required
observation_cov	Array	Observation noise covariance Λ, shape (p, p).	required
max_iter	int	Maximum number of CG iterations for the update step (static int).	required

Returns:

Type	Description
tuple[Array, Array]	Posterior (mean, factor) at time k + 1.

cakf_update

cakf_update(*, mean: Array, prior_cov: Array, factor: Array, observation: Array, observation_matrix: Array, observation_cov: Array, max_iter: int, policy: CAKFPolicy = CG, key: Array | None = None) -> tuple[Array, Array]

CAKF update step using the CG search-direction policy.

Iterates up to max_iter times. Each iteration:

1. Uses the current residual r as the search direction (CGPolicy).
2. Conjugates r against previously selected directions via d = r - U (U^T S r) where S is the symmetric operator H Σ H^T - H M M^T H^T + Λ and U is the running Gram-Schmidt factor.
3. Normalises d by η = r^T S d and appends sqrt(1/η) d to U.
4. Updates the action u so that H^T u is the cumulative posterior-mean correction.

Parameters:

Name	Type	Description	Default
mean	Array	Prior mean with shape (n,).	required
prior_cov	Array	Prior marginal covariance Σ of shape (n, n).	required
factor	Array	Existing low-rank correction factor M of shape (n, r0). May be empty (r0 = 0).	required
observation	Array	Observed value with shape (k,).	required
observation_matrix	Array	Observation matrix H of shape (k, n).	required
observation_cov	Array	Observation noise Λ of shape (k, k).	required
max_iter	int	Maximum number of CG iterations (must be a static Python int so the shapes of U and M_new are known at trace time).	required

Returns:

Type	Description
Array	Posterior (mean, factor). The returned factor has shape
Array	(n, r0 + max_iter).

diagonal_ek1_step

diagonal_ek1_step(*, mean: Array, cov_sqrt: Array, transition: Array, process_sqrt: Array, drift: Callable[[Array], Array], jacobian_diagonal: Callable[[Array], Array]) -> tuple[Array, Array, Array, Array]

One unit step of the diagonal Extended Kalman (EK1) ODE solver.

Parameters:

Name	Type	Description	Default
mean	Array	State mean of shape (num_derivatives + 1, ode_dim). Row 0 holds the function value, row 1 the first derivative, and so on.	required
cov_sqrt	Array	Per-dimension lower-triangular Cholesky factor of the marginal covariance, shape (ode_dim, num_derivatives + 1, num_derivatives + 1) with P_d = L_d L_d^T.	required
transition	Array	Discrete-time transition matrix Φ of shape (num_derivatives + 1, num_derivatives + 1) (typically the IBM(n) transition over dt).	required
process_sqrt	Array	Cholesky factor of the discrete process noise, shape (num_derivatives + 1, num_derivatives + 1).	required
drift	Callable[[Array], Array]	ODE right-hand side f(x) mapping (ode_dim,) to (ode_dim,). Time is assumed absorbed into drift.	required
jacobian_diagonal	Callable[[Array], Array]	Diagonal of df/dx mapping (ode_dim,) to (ode_dim,).	required

Returns:

Type	Description
Array	(new_mean, new_cov_sqrt, error_estimate, sigma):
Array	new_mean with shape (num_derivatives + 1, ode_dim);
Array	new_cov_sqrt with shape (ode_dim, n + 1, n + 1);
Array	error_estimate/sigma: per-dimension (ode_dim,) arrays for step-size control.

kalman_filter

kalman_filter(*, transitions: Array, process_noises: Array, observations: Array, observation_matrix: Array, observation_covs: Array, initial_mean: Array, initial_cov: Array) -> tuple[Array, Array]

Run a sequential Kalman filter over a fixed-length sequence.

Sibling reference: bayesnewton/bayesnewton/ops.py:154 _sequential_kf.

Parameters:

Name	Type	Description	Default
transitions	Array	shape (num_steps, state_dim, state_dim).	required
process_noises	Array	shape (num_steps, state_dim, state_dim).	required
observations	Array	shape (num_steps, obs_dim).	required
observation_matrix	Array	shape (obs_dim, state_dim) (time-invariant).	required
observation_covs	Array	shape (num_steps, obs_dim, obs_dim).	required
initial_mean	Array	shape (state_dim,).	required
initial_cov	Array	shape (state_dim, state_dim).	required

Returns:

Type	Description
Array	(means, covs) of shapes (num_steps, state_dim) and
Array	(num_steps, state_dim, state_dim) containing the posterior
tuple[Array, Array]	mean and covariance at each step after the update.

kalman_log_likelihood

kalman_log_likelihood(*, transitions: Array, process_noises: Array, observations: Array, observation_matrix: Array, observation_covs: Array, initial_mean: Array, initial_cov: Array) -> Array

Marginal log-likelihood of the observation sequence.

Returns sum_k log p(y_k | y_{0:k-1}) evaluated at each step via the innovation Gaussian whose covariance is H P_- H^T + R. Used for hyperparameter learning via jax.grad.

Sibling reference: bayesnewton/bayesnewton/ops.py:170-180 (the mvn_logpdf accumulator inside _sequential_kf).

kalman_predict

kalman_predict(*, mean: Array, cov: Array, transition: Array, process_noise: Array) -> tuple[Array, Array]

Linear-Gaussian Kalman prediction step.

Returns (transition @ mean, transition @ cov @ transition.T + process_noise).

Parameters:

Name	Type	Description	Default
mean	Array	prior mean of shape (state_dim,).	required
cov	Array	prior covariance of shape (state_dim, state_dim).	required
transition	Array	state-transition matrix A of shape (state_dim, state_dim).	required
process_noise	Array	process-noise covariance Q of shape (state_dim, state_dim).	required

Returns:

Type	Description
tuple[Array, Array]	Predicted (mean, cov).

kalman_smoother

kalman_smoother(*, filter_means: Array, filter_covs: Array, transitions: Array, process_noises: Array) -> tuple[Array, Array]

Rauch-Tung-Striebel backward smoother.

Sibling reference: bayesnewton/bayesnewton/ops.py:288 _sequential_rts.

Parameters:

Name	Type	Description	Default
filter_means	Array	forward-pass posterior means (num_steps, state_dim).	required
filter_covs	Array	forward-pass posterior covariances (num_steps, state_dim, state_dim).	required
transitions	Array	same shape as in kalman_filter.	required
process_noises	Array	same shape as in kalman_filter.	required

Returns:

Type	Description
Array	(smoothed_means, smoothed_covs) — posteriors conditioned on
Array	the full observation sequence.

kalman_update

kalman_update(*, mean: Array, cov: Array, observation: Array, observation_matrix: Array, observation_cov: Array) -> tuple[Array, Array]

Linear-Gaussian Kalman update (innovation + Joseph-form covariance).

Parameters:

Name	Type	Description	Default
mean	Array	predicted mean m_- of shape (state_dim,).	required
cov	Array	predicted covariance P_- of shape (state_dim, state_dim).	required
observation	Array	observed measurement y of shape (obs_dim,).	required
observation_matrix	Array	linear observation operator H of shape (obs_dim, state_dim).	required
observation_cov	Array	observation-noise covariance R of shape (obs_dim, obs_dim).	required

Returns:

Type	Description
tuple[Array, Array]	Posterior (mean, cov) after conditioning on observation.

cosine_kernel

cosine_kernel(*, frequency: float) -> StateSpaceKernel

Cosine kernel as SDE. Ports bayesnewton Cosine (line 770).

State dim 2; transition is the 2-D rotation by angle frequency dt.

matern12_kernel

matern12_kernel(*, variance: float, lengthscale: float) -> StateSpaceKernel

Matern-½ (exponential) kernel in SDE form.

State dimension 1. F = [[-1/ell]], Q_c = 2 sigma^2 / ell, A(dt) = exp(-dt/ell). Ports bayesnewton Matern12 (line 141).

matern32_kernel

matern32_kernel(*, variance: float, lengthscale: float) -> StateSpaceKernel

Matern-3/2 kernel in SDE form. Ports bayesnewton Matern32 (line 200).

matern52_kernel

matern52_kernel(*, variance: float, lengthscale: float) -> StateSpaceKernel

Matern-5/2 kernel in SDE form. Ports bayesnewton Matern52 (line 253).

matern72_kernel

matern72_kernel(*, variance: float, lengthscale: float) -> StateSpaceKernel

Matern-7/2 kernel in SDE form. Ports bayesnewton Matern72 (line 321).

periodic_kernel

periodic_kernel(*, variance: float, lengthscale: float, period: float, order: int = 6) -> StateSpaceKernel

Periodic kernel via Bessel-weighted sum of harmonic rotations.

Ports bayesnewton Periodic (line 802). State dim 2(order + 1); transition is block-diagonal of harmonic rotation matrices. The Bessel-weighted spectrum uses the exponentially-scaled modified Bessel functions i0e (analogous to bessel_ive in bayesnewton).

quasi_periodic_matern12_kernel

quasi_periodic_matern12_kernel(*, variance: float, lengthscale_periodic: float, period: float, lengthscale_matern: float, order: int = 6) -> StateSpaceKernel

Quasi-periodic Matern-½ kernel: product of Periodic and Matern-½.

Ports bayesnewton QuasiPeriodicMatern12 (line 882). Constructed as the Kronecker product of the Matern-½ SDE with the Periodic SDE.

discretize_lti_sde

discretize_lti_sde(*, drift_matrix: Array, dispersion_matrix: Array, dt: Array, diffusion: Array | None = None) -> tuple[Array, Array]

Discretize a continuous-time linear time-invariant SDE.

Computes the discrete-time state transition matrix :math:A = \exp(F\, \Delta t) and the integrated process noise covariance via Van Loan's matrix-fraction decomposition.

Parameters:

Name	Type	Description	Default
drift_matrix	Array	Continuous-time drift :math:F of shape (n, n).	required
dispersion_matrix	Array	Dispersion :math:L of shape (n, k).	required
dt	Array	Time step :math:\Delta t (scalar).	required
diffusion	Array | None	Wiener-process diffusion :math:Q_c of shape (k, k). Defaults to the identity matrix if None.	None

Returns:

Type	Description
Array	(transition, process_noise) of shapes (n, n) and
Array	(n, n).

process_noise_covariance

process_noise_covariance(*, drift_matrix: Array, dispersion_matrix: Array, dt: Array, diffusion: Array | None = None) -> Array

Return only Q (Van Loan process-noise) from the LTI-SDE discretisation.

Thin convenience wrapper around :func:discretize_lti_sde for the common case where only the process-noise covariance is needed.

state_transition_matrix

state_transition_matrix(*, drift_matrix: Array, dispersion_matrix: Array, dt: Array, diffusion: Array | None = None) -> Array

Return only A = exp(F dt) from the LTI-SDE discretisation.

Thin convenience wrapper around :func:discretize_lti_sde for the common case where only the transition matrix is needed (e.g. independent discretisation of a kernel's continuous-time SDE).

kalman_filter_parallel

kalman_filter_parallel(*, transitions: Array, process_noises: Array, observations: Array, observation_matrix: Array, observation_covs: Array, initial_mean: Array, initial_cov: Array) -> tuple[Array, Array]

Parallel-scan Kalman filter via associative composition.

Identical output to :func:kalman_filter but runs in :math:O(\log N) parallel depth via :func:jax.lax.associative_scan.

The bayesnewton reference uses an "observe at step t then transition to step t+1" indexing, whereas opifex uses "predict via A_t then update with y_t" (i.e., observation y_t is observed at time t+1 relative to the initial prior). To bridge the two, the first transition is folded into an effective prior N(A_0 m_0, A_0 P_0 A_0^T + Q_0) and the first transition slot in the parallel pipeline is replaced with the identity.

Parameters:

Name	Type	Description	Default
transitions	Array	Per-step transition matrices, shape (N, n, n).	required
process_noises	Array	Per-step process-noise covariances, (N, n, n).	required
observations	Array	Observation sequence, (N, k).	required
observation_matrix	Array	Time-invariant observation matrix, (k, n).	required
observation_covs	Array	Per-step observation noise covariances, (N, k, k).	required
initial_mean	Array	Prior mean m_0, (n,).	required
initial_cov	Array	Prior covariance P_0, (n, n).	required

Returns:

Type	Description
tuple[Array, Array]	Filter means (N, n) and filter covariances (N, n, n).

kalman_smoother_parallel

kalman_smoother_parallel(*, filter_means: Array, filter_covs: Array, transitions: Array, process_noises: Array) -> tuple[Array, Array]

Parallel-scan Rauch-Tung-Striebel smoother.

Identical output to :func:kalman_smoother but runs in :math:O(\log N) parallel depth via :func:jax.lax.associative_scan with reverse=True.

Parameters:

Name	Type	Description	Default
filter_means	Array	Sequential filter means, (N, n).	required
filter_covs	Array	Sequential filter covariances, (N, n, n).	required
transitions	Array	Per-step transition matrices, (N, n, n). The element at index t is the transition from t to t+1.	required
process_noises	Array	Per-step process noise covariances, (N, n, n).	required

Returns:

Type	Description
tuple[Array, Array]	Smoothed means (N, n) and smoothed covariances (N, n, n).

sqrt_kalman_predict

sqrt_kalman_predict(*, mean: Array, cov_sqrt: Array, transition: Array, process_noise_sqrt: Array) -> tuple[Array, Array]

Square-root predict step.

Computes (A m, L') where L' L'^T = A L L^T A^T + Q_sqrt Q_sqrt^T.

The new factor is obtained by stacking the right square roots [(A L)^T; Q_sqrt^T] and reading off the upper-triangular R of a thin QR factorisation: R^T R = (A L)(A L)^T + Q_sqrt Q_sqrt^T. The returned left factor is L' = R^T.

Parameters:

Name	Type	Description	Default
mean	Array	Prior mean m with shape (n,).	required
cov_sqrt	Array	Prior left square root L with shape (n, n) such that P = L @ L.T.	required
transition	Array	Transition matrix A with shape (n, n).	required
process_noise_sqrt	Array	Left square root of the process-noise covariance, shape (n, n).	required

Returns:

Type	Description
Array	Predicted mean A m and the lower-triangular factor L' of
Array	the predicted covariance.

sqrt_kalman_update

sqrt_kalman_update(*, mean: Array, cov_sqrt: Array, observation: Array, observation_matrix: Array, observation_cov_sqrt: Array) -> tuple[Array, Array]

Square-root update step via the joint-QR revert formula.

Given prior N(m, L L^T) and observation model y | x ~ N(H x, R_sqrt R_sqrt^T), computes the posterior N(m', L'_post L'_post^T) using the QR-revert identity (probdiffeq revert_conditional line 51):

Stack ::

R_block = [[R_YX,    0  ],
           [R_X_F,  R_X ]]

where R_X_F = (H L)^T (shape (n, k)), R_X = L^T (shape (n, n)) and R_YX = R_sqrt^T (shape (k, k)). Computing R_full = qr_r(R_block) and partitioning yields r_obs (innovation sqrt), r_cor (posterior sqrt) and the Kalman gain G = solve_triangular(r_obs, R12, lower=False).T.

Parameters:

Name	Type	Description	Default
mean	Array	Prior mean with shape (n,).	required
cov_sqrt	Array	Prior left square root L with shape (n, n).	required
observation	Array	Observed value with shape (k,).	required
observation_matrix	Array	Observation matrix H with shape (k, n).	required
observation_cov_sqrt	Array	Left square root of observation noise, shape (k, k).	required

Returns:

Type	Description
tuple[Array, Array]	Posterior mean and lower-triangular posterior covariance factor.

Non-Gaussian inference on Markov GPs (Task 11.2).

Builds on the linear-Gaussian state-space layer in :mod:opifex.uncertainty.statespace and the per-observation LikelihoodComponentsFn interface introduced for non-conjugate GP inference in :mod:opifex.uncertainty.gp.laplace (Task 11.1 D5).

References

• Sarkka 2013 — Bayesian Filtering and Smoothing, CUP (state-space GPs + iterated extended Kalman smoothing).
• Wilkinson, Solin, Adam 2020+ — bayesnewton (primary reference for the inference algorithm catalogue: PEP / VI / Laplace / Posterior Linearisation / Newton).
• Solin, Hensman, Turner 2018 — Infinite-horizon Gaussian processes, NeurIPS (steady-state Kalman variants — deferred).

MarkovLaplaceGPState dataclass

MarkovLaplaceGPState(*, times: Array, observations: Array, smoothed_means: Array, smoothed_variances: Array, smoothed_state_means: Array, smoothed_state_covariances: Array, log_marginal_likelihood: Array, state_space_kernel: StateSpaceKernel, log_likelihood_components_fn: LikelihoodComponentsFn)

Fitted state for the Markov-Laplace non-conjugate GP.

Attributes:

Name	Type	Description
times	Array	(n,) strictly-increasing training time stamps.
observations	Array	(n,) training observations.
smoothed_means	Array	(n,) posterior mean of the latent f(t) at training times.
smoothed_variances	Array	(n,) posterior variance of the latent (marginal of the smoothed state via the kernel's measurement operator).
smoothed_state_means	Array	(n, state_dim) smoothed full state trajectory (used at predict time to propagate to held-out times).
smoothed_state_covariances	Array	(n, state_dim, state_dim) full smoothed state covariances.
log_marginal_likelihood	Array	Laplace-approximated log marginal (scalar).
state_space_kernel	StateSpaceKernel	The :class:StateSpaceKernel used at fit time (closed over for the predict path).
log_likelihood_components_fn	LikelihoodComponentsFn	The per-observation likelihood quadruple (log_lik, ∇log_lik, W, √W) callable used at fit time. Kept on the state so that :func:predict_markov_laplace_gp can re-evaluate W / grad at the smoothed mode for the response-distribution map.

MarkovPEPGPState dataclass

MarkovPEPGPState(*, times: Array, observations: Array, smoothed_means: Array, smoothed_variances: Array, smoothed_state_means: Array, smoothed_state_covariances: Array, site_eta_1: Array, site_eta_2: Array, log_marginal_likelihood: Array, state_space_kernel: StateSpaceKernel, log_likelihood_fn: PerObservationLogLikelihoodFn, power: float)

Fitted state for Power EP on a Markov-GP prior.

Carries the smoothed posterior moments, the converged per-site natural parameters, and the approximate log marginal likelihood.

MarkovPLGPState dataclass

MarkovPLGPState(*, times: Array, observations: Array, smoothed_means: Array, smoothed_variances: Array, smoothed_state_means: Array, smoothed_state_covariances: Array, state_space_kernel: StateSpaceKernel, conditional_moments_fn: ConditionalMomentsFn)

Fitted state for Posterior Linearisation on a Markov-GP prior.

MarkovVIGPState dataclass

MarkovVIGPState(*, times: Array, observations: Array, smoothed_means: Array, smoothed_variances: Array, smoothed_state_means: Array, smoothed_state_covariances: Array, evidence_lower_bound: Array, state_space_kernel: StateSpaceKernel, log_likelihood_components_fn: LikelihoodComponentsFn)

Fitted state for VI on a Markov-GP prior.

Attributes are the Markov-Laplace counterparts plus the variational evidence_lower_bound (ELBO) — the scalar objective conjugate-computation VI maximises.

fit_markov_laplace_gp

fit_markov_laplace_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, log_likelihood_components_fn: LikelihoodComponentsFn, num_iterations: int = 25) -> MarkovLaplaceGPState

Fit the Laplace approximation for a Markov-GP non-conjugate likelihood.

Runs num_iterations Newton iterations through the pseudo-Gaussian-observation linearisation; each iteration is a single Kalman filter + smoother pass over the time grid (linear in n).

Parameters:

Name	Type	Description	Default
times	Array	(n,) strictly-increasing training time stamps.	required
observations	Array	(n,) training observations in the likelihood's support.	required
state_space_kernel	StateSpaceKernel	SDE-form GP prior. Any :class:opifex.uncertainty.statespace.StateSpaceKernel (matern12 / 32 / 52 / 72 / cosine / periodic / quasi-periodic) is supported.	required
log_likelihood_components_fn	LikelihoodComponentsFn	Callable (f, y) -> (log_lik_total, ∇log_lik, W, √W) returning a scalar log_lik_total and per-observation (n,) gradient + curvature + sqrt-curvature arrays. The same interface as :data:opifex.uncertainty.gp.laplace.LikelihoodComponentsFn.	required
num_iterations	int	Newton-loop count (static under jax.jit).	25

Returns:

Type	Description
MarkovLaplaceGPState	class:MarkovLaplaceGPState carrying smoothed posterior
MarkovLaplaceGPState	moments, full smoothed state, and the Laplace-approximated
MarkovLaplaceGPState	log marginal.

predict_markov_laplace_gp

predict_markov_laplace_gp(*, state: MarkovLaplaceGPState, times_test: Array) -> PredictiveDistribution

Posterior latent moments at times_test via state-space interpolation.

For each test time t*, locate the latest training time t_k ≤ t* and propagate the smoothed state at t_k forward by Δt = t* − t_k using the SDE transition matrix. For test times before the first training time, propagate from the stationary prior (zero mean, stationary covariance).

Parameters:

Name	Type	Description	Default
state	MarkovLaplaceGPState	Fitted :class:MarkovLaplaceGPState.	required
times_test	Array	(m,) test time stamps (any order).	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution whose mean and
PredictiveDistribution	variance carry the latent f(t*) marginal moments at
PredictiveDistribution	each test time. Map through the per-likelihood response
PredictiveDistribution	link (MacKay probit for Bernoulli, exp for Poisson, ...)
PredictiveDistribution	downstream.

fit_bernoulli_markov_laplace_gp

fit_bernoulli_markov_laplace_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, num_iterations: int = 25) -> MarkovLaplaceGPState

Bernoulli classification on a Markov-GP prior (slice-25 wrapper).

fit_beta_markov_laplace_gp

fit_beta_markov_laplace_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, scale: float = 10.0, num_iterations: int = 25) -> MarkovLaplaceGPState

Beta proportion regression on a Markov-GP prior (logit link, Fisher W).

fit_gaussian_markov_laplace_gp

fit_gaussian_markov_laplace_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, noise_std: float = 0.1, num_iterations: int = 1) -> MarkovLaplaceGPState

Gaussian regression on a Markov-GP prior.

For Gaussian likelihood the log-likelihood is quadratic in f, so Newton converges in one step (the algorithm collapses to the exact conjugate Kalman smoother). Additional iterations are fixed-point no-ops.

fit_poisson_markov_laplace_gp

fit_poisson_markov_laplace_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, num_iterations: int = 25) -> MarkovLaplaceGPState

Poisson count regression on a Markov-GP prior (exp link).

fit_studentst_markov_laplace_gp

fit_studentst_markov_laplace_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, df: float = 4.0, scale: float = 1.0, num_iterations: int = 25) -> MarkovLaplaceGPState

Student-t robust regression on a Markov-GP prior (Fisher-info W).

predict_bernoulli_markov_laplace_gp

predict_bernoulli_markov_laplace_gp(*, state: MarkovLaplaceGPState, times_test: Array) -> PredictiveDistribution

Predict p(y_* = +1 | t_*) via the MacKay probit collapse.

predict_beta_markov_laplace_gp

predict_beta_markov_laplace_gp(*, state: MarkovLaplaceGPState, times_test: Array, scale: float = 10.0) -> PredictiveDistribution

Predict Beta response: mean σ(κ μ), variance m̂(1-m̂)/(s+1).

predict_gaussian_markov_laplace_gp

predict_gaussian_markov_laplace_gp(*, state: MarkovLaplaceGPState, times_test: Array, noise_std: float = 0.1) -> PredictiveDistribution

Predict y* | t*: mean μ, variance V + σ² (latent + noise).

predict_poisson_markov_laplace_gp

predict_poisson_markov_laplace_gp(*, state: MarkovLaplaceGPState, times_test: Array) -> PredictiveDistribution

Predict Poisson intensity E[λ*] = exp(μ + ½ V) under log-normal moments.

predict_studentst_markov_laplace_gp

predict_studentst_markov_laplace_gp(*, state: MarkovLaplaceGPState, times_test: Array, df: float = 4.0, scale: float = 1.0) -> PredictiveDistribution

Student-t predictive: mean μ, variance V + σ² ν / (ν − 2).

fit_markov_pep_gp

fit_markov_pep_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, log_likelihood_fn: PerObservationLogLikelihoodFn, power: float = 0.5, num_iterations: int = 25, learning_rate: float = 0.5, num_quadrature_points: int = 20, log_partition_fn: LogZAndDerivativesFn | None = None) -> MarkovPEPGPState

Fit Power EP on a Markov-GP prior.

Parameters:

Name	Type	Description	Default
times	Array	(n,) strictly-increasing training times.	required
observations	Array	(n,) training observations.	required
state_space_kernel	StateSpaceKernel	SDE-form GP prior.	required
log_likelihood_fn	PerObservationLogLikelihoodFn	(f, y) -> log p(y|f) per-observation.	required
power	float	EP power α ∈ (0, 1]. Defaults to 0.5 (Power EP).	0.5
num_iterations	int	Number of full Kalman-and-site sweeps.	25
learning_rate	float	Damping for the site update (1.0 = no damping).	0.5
num_quadrature_points	int	GH nodes for cubature.	20
log_partition_fn	LogZAndDerivativesFn | None	Optional closed-form log Z and its first two derivatives w.r.t. cavity mean. When provided, cubature is bypassed — used by the Gaussian wrapper, which has the exact analytical form.	None

Returns:

Type	Description
MarkovPEPGPState	class:MarkovPEPGPState carrying smoothed posterior moments,
MarkovPEPGPState	converged sites, and approximate log marginal likelihood.

predict_markov_pep_gp

predict_markov_pep_gp(*, state: MarkovPEPGPState, times_test: Array) -> PredictiveDistribution

Posterior latent moments at times_test via state-space interpolation.

Identical predict path as :func:predict_markov_laplace_gp and :func:predict_markov_vi_gp — the inference algorithm differs but the smoothed-state interpolation is shared.

fit_bernoulli_markov_pep_gp

fit_bernoulli_markov_pep_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, power: float = 0.5, num_iterations: int = 25, learning_rate: float = 0.5, num_quadrature_points: int = 20) -> MarkovPEPGPState

Bernoulli classification on a Markov-GP prior via Power EP.

fit_beta_markov_pep_gp

fit_beta_markov_pep_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, scale: float = 10.0, power: float = 0.5, num_iterations: int = 25, learning_rate: float = 0.5, num_quadrature_points: int = 20) -> MarkovPEPGPState

Beta regression on a Markov-GP prior via Power EP with logit link.

fit_gaussian_markov_pep_gp

fit_gaussian_markov_pep_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, noise_std: float = 0.1, power: float = 1.0, num_iterations: int = 5, learning_rate: float = 1.0, num_quadrature_points: int = 20) -> MarkovPEPGPState

Gaussian regression on a Markov-GP prior via Power EP.

Defaults to classical EP (power = 1, learning_rate = 1). For Gaussian likelihood the tilted distribution is exactly Gaussian, so we use the closed-form partition function and its derivatives (_gaussian_log_partition_factory) — bypassing Gauss-Hermite cubature, which would break down for narrow noise_std relative to the cavity scale (a known limitation of cavity-centered GH; see e.g. Hernández-Lobato et al. on BB-α). One EP sweep then exactly recovers the conjugate Kalman solution.

fit_poisson_markov_pep_gp

fit_poisson_markov_pep_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, power: float = 0.5, num_iterations: int = 25, learning_rate: float = 0.5, num_quadrature_points: int = 20) -> MarkovPEPGPState

Poisson count regression on a Markov-GP prior via Power EP.

fit_studentst_markov_pep_gp

fit_studentst_markov_pep_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, df: float = 4.0, scale: float = 1.0, power: float = 0.5, num_iterations: int = 25, learning_rate: float = 0.5, num_quadrature_points: int = 20) -> MarkovPEPGPState

Robust regression on a Markov-GP prior via Power EP with Student-t likelihood.

predict_bernoulli_markov_pep_gp

predict_bernoulli_markov_pep_gp(*, state: MarkovPEPGPState, times_test: Array) -> PredictiveDistribution

Predict p(y_* = +1) via the MacKay probit collapse.

predict_beta_markov_pep_gp

predict_beta_markov_pep_gp(*, state: MarkovPEPGPState, times_test: Array, scale: float = 10.0) -> PredictiveDistribution

Predict Beta response mean E[y*] = sigmoid(latent_mean).

Predictive variance under the unit-interval link uses the standard Beta(α, β) marginal variance mean (1 - mean) / (scale + 1) at mean = sigmoid(latent_mean).

predict_gaussian_markov_pep_gp

predict_gaussian_markov_pep_gp(*, state: MarkovPEPGPState, times_test: Array, noise_std: float = 0.1) -> PredictiveDistribution

Predict y*: mean μ, variance V + σ² (latent + obs noise).

predict_poisson_markov_pep_gp

predict_poisson_markov_pep_gp(*, state: MarkovPEPGPState, times_test: Array) -> PredictiveDistribution

Predict Poisson intensity E[λ] = exp(μ + ½ V) under log-normal moments.

predict_studentst_markov_pep_gp

predict_studentst_markov_pep_gp(*, state: MarkovPEPGPState, times_test: Array, df: float = 4.0, scale: float = 1.0) -> PredictiveDistribution

Predict y* under the Student-t response.

Latent mean μ; response variance is the latent variance plus the Student-t marginal variance scale² ν / (ν - 2) (finite for ν > 2).

fit_markov_pl_gp

fit_markov_pl_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, conditional_moments_fn: ConditionalMomentsFn, num_iterations: int = 20, num_quadrature_points: int = 20) -> MarkovPLGPState

Fit Posterior Linearisation on a Markov-GP prior.

Parameters:

Name	Type	Description	Default
times	Array	(n,) strictly-increasing training times.	required
observations	Array	(n,) training observations.	required
state_space_kernel	StateSpaceKernel	SDE-form GP prior.	required
conditional_moments_fn	ConditionalMomentsFn	Per-likelihood f -> (E[y|f], Var[y|f]).	required
num_iterations	int	Number of full re-linearise-and-Kalman sweeps.	20
num_quadrature_points	int	GH nodes for the cubature in SLR.	20

Returns:

Type	Description
MarkovPLGPState	class:MarkovPLGPState with the converged smoothed posterior
MarkovPLGPState	moments and the full smoothed state trajectory.

predict_markov_pl_gp

predict_markov_pl_gp(*, state: MarkovPLGPState, times_test: Array) -> PredictiveDistribution

Posterior latent moments at times_test via state-space interpolation.

Identical predict path as the Laplace / VI / PEP routes — only the inference algorithm differs.

fit_bernoulli_markov_pl_gp

fit_bernoulli_markov_pl_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, num_iterations: int = 20, num_quadrature_points: int = 20) -> MarkovPLGPState

Bernoulli classification on a Markov-GP prior via Posterior Linearisation.

fit_beta_markov_pl_gp

fit_beta_markov_pl_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, scale: float = 10.0, num_iterations: int = 20, num_quadrature_points: int = 20) -> MarkovPLGPState

Beta regression on a Markov-GP prior via Posterior Linearisation.

fit_gaussian_markov_pl_gp

fit_gaussian_markov_pl_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, noise_std: float = 0.1, num_iterations: int = 3, num_quadrature_points: int = 20) -> MarkovPLGPState

Gaussian regression via Posterior Linearisation.

For Gaussian likelihood the SLR linearisation is exact (the conditional mean is the identity in f, so A = 1, b = 0, omega = noise_std**2), and one iteration suffices to recover the conjugate Kalman solution.

fit_poisson_markov_pl_gp

fit_poisson_markov_pl_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, num_iterations: int = 20, num_quadrature_points: int = 20) -> MarkovPLGPState

Poisson count regression on a Markov-GP prior via Posterior Linearisation.

fit_studentst_markov_pl_gp

fit_studentst_markov_pl_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, df: float = 4.0, scale: float = 1.0, num_iterations: int = 20, num_quadrature_points: int = 20) -> MarkovPLGPState

Robust regression on a Markov-GP prior via Posterior Linearisation.

For a location-scale Student-t with df > 2, E[y|f] = f and the response variance is constant, so the SLR linearisation collapses to a Gaussian observation model with that constant variance.

predict_bernoulli_markov_pl_gp

predict_bernoulli_markov_pl_gp(*, state: MarkovPLGPState, times_test: Array) -> PredictiveDistribution

Predict p(y_* = +1) via the MacKay probit collapse.

predict_beta_markov_pl_gp

predict_beta_markov_pl_gp(*, state: MarkovPLGPState, times_test: Array, scale: float = 10.0) -> PredictiveDistribution

Predict Beta response mean sigmoid(latent_mean) + Beta marginal variance.

predict_gaussian_markov_pl_gp

predict_gaussian_markov_pl_gp(*, state: MarkovPLGPState, times_test: Array, noise_std: float = 0.1) -> PredictiveDistribution

Predict y*: mean mu, variance V + sigma**2 (latent + obs noise).

predict_poisson_markov_pl_gp

predict_poisson_markov_pl_gp(*, state: MarkovPLGPState, times_test: Array) -> PredictiveDistribution

Predict Poisson intensity E[lambda] = exp(mu + 0.5 V).

predict_studentst_markov_pl_gp

predict_studentst_markov_pl_gp(*, state: MarkovPLGPState, times_test: Array, df: float = 4.0, scale: float = 1.0) -> PredictiveDistribution

Predict y* under the Student-t response (df > 2 finite variance).

fit_markov_vi_gp

fit_markov_vi_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, log_likelihood_components_fn: LikelihoodComponentsFn, num_iterations: int = 25, num_quadrature_points: int = 20) -> MarkovVIGPState

Fit conjugate-computation VI on a Markov-GP prior.

Parameters:

Name	Type	Description	Default
times	Array	(n,) strictly-increasing training times.	required
observations	Array	(n,) training observations.	required
state_space_kernel	StateSpaceKernel	SDE-form GP prior.	required
log_likelihood_components_fn	LikelihoodComponentsFn	D5 LikelihoodComponentsFn — the same callable used by :func:fit_markov_laplace_gp.	required
num_iterations	int	VI iteration count (static under jax.jit).	25
num_quadrature_points	int	Gauss-Hermite nodes for the expected components (static; defaults to 20).	20

Returns:

Type	Description
MarkovVIGPState	class:MarkovVIGPState with the smoothed posterior moments,
MarkovVIGPState	full smoothed state, and the final ELBO value.

predict_markov_vi_gp

predict_markov_vi_gp(*, state: MarkovVIGPState, times_test: Array) -> PredictiveDistribution

Posterior latent moments at times_test via state-space interpolation.

Mirrors :func:opifex.uncertainty.markov.predict_markov_laplace_gp — the state-space interpolation step is identical between the Laplace and VI posteriors once the smoothed state trajectory has been computed at the training grid.

fit_bernoulli_markov_vi_gp

fit_bernoulli_markov_vi_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, num_iterations: int = 25, num_quadrature_points: int = 20) -> MarkovVIGPState

Bernoulli classification on a Markov-GP prior via conjugate VI.

fit_beta_markov_vi_gp

fit_beta_markov_vi_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, scale: float = 10.0, num_iterations: int = 25, num_quadrature_points: int = 20) -> MarkovVIGPState

Beta proportion regression on a Markov-GP prior via conjugate VI.

fit_gaussian_markov_vi_gp

fit_gaussian_markov_vi_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, noise_std: float = 0.1, num_iterations: int = 1, num_quadrature_points: int = 20) -> MarkovVIGPState

Gaussian regression on a Markov-GP prior.

For Gaussian likelihood the log-likelihood gradient is linear in f and the curvature is constant, so the expected components (over any q(f)) coincide with the mode-evaluated components. VI converges in one iteration to the exact conjugate Kalman solution — identical to :func:opifex.uncertainty.markov.fit_gaussian_markov_laplace_gp.

fit_poisson_markov_vi_gp

fit_poisson_markov_vi_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, num_iterations: int = 25, num_quadrature_points: int = 20) -> MarkovVIGPState

Poisson count regression on a Markov-GP prior via conjugate VI.

fit_studentst_markov_vi_gp

fit_studentst_markov_vi_gp(*, times: Array, observations: Array, state_space_kernel: StateSpaceKernel, df: float = 4.0, scale: float = 1.0, num_iterations: int = 25, num_quadrature_points: int = 20) -> MarkovVIGPState

Student-t robust regression on a Markov-GP prior via conjugate VI.

predict_bernoulli_markov_vi_gp

predict_bernoulli_markov_vi_gp(*, state: MarkovVIGPState, times_test: Array) -> PredictiveDistribution

Predict p(y_* = +1) via the MacKay probit collapse.

predict_beta_markov_vi_gp

predict_beta_markov_vi_gp(*, state: MarkovVIGPState, times_test: Array, scale: float = 10.0) -> PredictiveDistribution

Predict Beta response: E[y] = σ(κ μ), variance m̂(1-m̂)/(s+1).

predict_gaussian_markov_vi_gp

predict_gaussian_markov_vi_gp(*, state: MarkovVIGPState, times_test: Array, noise_std: float = 0.1) -> PredictiveDistribution

Predict y*: mean μ, variance V + σ² (latent + obs noise).

predict_poisson_markov_vi_gp

predict_poisson_markov_vi_gp(*, state: MarkovVIGPState, times_test: Array) -> PredictiveDistribution

Predict Poisson intensity E[λ] = exp(μ + ½ V) under log-normal moments.

predict_studentst_markov_vi_gp

predict_studentst_markov_vi_gp(*, state: MarkovVIGPState, times_test: Array, df: float = 4.0, scale: float = 1.0) -> PredictiveDistribution

Predict Student-t response: mean μ, variance V + σ²ν/(ν−2).

Data-assimilation / digital-twin state utilities (Task 6.7).

A thin orchestration layer over :mod:opifex.uncertainty.statespace: this package supplies the digital-twin-aware containers (:class:AssimilationState), sensor-mask helpers, and a jax.lax.scan-driven sequential update loop. The actual Kalman math is re-exported from the canonical state-space module so callers have a one-stop import surface, but no formula bodies are re-implemented here (Phase 9 Task 9.3 enforces this).

AssimilationState

Digital-twin state estimate + uncertainty bookkeeping.

Attributes:

Name	Type	Description
mean	Array	Physical-state mean, shape (state_dim,).
covariance	Array	State covariance, shape (state_dim, state_dim).
time	Array	Scalar timestamp.
metadata	tuple[tuple[str, Any], ...]	Static, hashable annotations distinguishing physical state, observation uncertainty, model discrepancy, numerical uncertainty, and calibration uncertainty. Stored as tuple[tuple[str, Any], ...] so it survives jax.tree.flatten as static aux data.

metadata_dict

metadata_dict() -> dict[str, Any]

Return metadata as a regular dict for ergonomic read access.

validate

validate() -> None

Public validation — must be called by callers, not the tree machinery.

Raises:

Type	Description
ValueError	If mean is not 1-D, covariance is not square / matching mean size, or required metadata fields are missing.

build_default_metadata

build_default_metadata(*, physical_state: str, observation_uncertainty: float, model_discrepancy: float, numerical_uncertainty: float, calibration_uncertainty: float, extra: dict[str, Any] | None = None) -> tuple[tuple[str, Any], ...]

Helper that builds a metadata tuple satisfying validate().

Designed so the canonical five uncertainty-source slots are named consistently across callers; extra allows experiment-specific annotations without touching the required slots.

annotate_metadata

annotate_metadata(state: AssimilationState, **extra: Any) -> AssimilationState

Return a copy of state with extra merged into metadata.

Useful when each assimilation step records additional bookkeeping (e.g. sensor name, calibration status). Preserves the required canonical keys.

observation_matrix_from_mask

observation_matrix_from_mask(mask: Array, state_dim: int) -> Array

Build H that selects only observed dimensions of the state vector.

Parameters:

Name	Type	Description	Default
mask	Array	1-D boolean / 0-1 array, length state_dim. True entries are observed.	required
state_dim	int	Full state dimension.	required

Returns:

Type	Description
Array	H of shape (num_observed, state_dim) whose rows are the
Array	canonical basis vectors of the observed dimensions.

predict

predict(state: AssimilationState, *, transition: Array, process_noise: Array, new_time: Array | None = None) -> AssimilationState

Advance state one step using the Kalman prediction primitive.

Math is delegated to opifex.uncertainty.statespace.kalman_predict.

Parameters:

Name	Type	Description	Default
state	AssimilationState	Current AssimilationState.	required
transition	Array	A matrix, shape (state_dim, state_dim).	required
process_noise	Array	Q matrix, shape (state_dim, state_dim).	required
new_time	Array | None	Optional new timestamp; defaults to state.time + 1.	None

Returns:

Type	Description
AssimilationState	Updated AssimilationState with carried-over metadata.

sequential_update

sequential_update(state: AssimilationState, *, transitions: Array, process_noises: Array, observations: Array, observation_matrices: Array, observation_covs: Array) -> tuple[AssimilationState, AssimilationState]

Roll state forward through T predict/update steps via lax.scan.

All time-varying inputs have a leading T axis. The state and the per-step posterior history are returned so the caller can recover smoothed estimates downstream.

Parameters:

Name	Type	Description	Default
state	AssimilationState	Initial AssimilationState.	required
transitions	Array	(T, state_dim, state_dim).	required
process_noises	Array	(T, state_dim, state_dim).	required
observations	Array	(T, obs_dim).	required
observation_matrices	Array	(T, obs_dim, state_dim).	required
observation_covs	Array	(T, obs_dim, obs_dim).	required

Returns:

Type	Description
AssimilationState	(final_state, history_state). history_state has each
AssimilationState	leaf prefixed with the time axis, so e.g.
tuple[AssimilationState, AssimilationState]	history_state.mean.shape == (T, state_dim).

update

update(state: AssimilationState, *, observation: Array, observation_matrix: Array, observation_cov: Array) -> AssimilationState

Condition state on a (possibly sparse) sensor observation.

Math is delegated to opifex.uncertainty.statespace.kalman_update. Sparse / partial observations are handled by passing a thinned observation_matrix that selects only the observed dimensions (or a boolean-mask-built matrix on the caller side).

Parameters:

Name	Type	Description	Default
state	AssimilationState	Predicted AssimilationState.	required
observation	Array	Sensor reading y of shape (obs_dim,).	required
observation_matrix	Array	H of shape (obs_dim, state_dim).	required
observation_cov	Array	R of shape (obs_dim, obs_dim).	required

Returns:

Type	Description
AssimilationState	Updated AssimilationState with carried-over metadata.

kalman_predict

kalman_predict(*, mean: Array, cov: Array, transition: Array, process_noise: Array) -> tuple[Array, Array]

Linear-Gaussian Kalman prediction step.

Returns (transition @ mean, transition @ cov @ transition.T + process_noise).

Parameters:

Name	Type	Description	Default
mean	Array	prior mean of shape (state_dim,).	required
cov	Array	prior covariance of shape (state_dim, state_dim).	required
transition	Array	state-transition matrix A of shape (state_dim, state_dim).	required
process_noise	Array	process-noise covariance Q of shape (state_dim, state_dim).	required

Returns:

Type	Description
tuple[Array, Array]	Predicted (mean, cov).

kalman_update

kalman_update(*, mean: Array, cov: Array, observation: Array, observation_matrix: Array, observation_cov: Array) -> tuple[Array, Array]

Linear-Gaussian Kalman update (innovation + Joseph-form covariance).

Parameters:

Name	Type	Description	Default
mean	Array	predicted mean m_- of shape (state_dim,).	required
cov	Array	predicted covariance P_- of shape (state_dim, state_dim).	required
observation	Array	observed measurement y of shape (obs_dim,).	required
observation_matrix	Array	linear observation operator H of shape (obs_dim, state_dim).	required
observation_cov	Array	observation-noise covariance R of shape (obs_dim, obs_dim).	required

Returns:

Type	Description
tuple[Array, Array]	Posterior (mean, cov) after conditioning on observation.

Multi-fidelity emulation (Task 11.3).

Two canonical multi-fidelity surrogate families plus the multi-fidelity acquisition pillar (MUMBO):

• :mod:opifex.uncertainty.multi_fidelity.linear — Kennedy & O'Hagan 2000 AR(1) linear multi-fidelity GP (slice 33). Models each fidelity as f_i(x) = rho_i f_{i-1}(x) + delta_i(x) where each delta_i is an independent GP.
• (Planned) :mod:opifex.uncertainty.multi_fidelity.nonlinear — Perdikaris, Raissi, Damianou, Lawrence, Karniadakis 2017 NARGP non-linear multi-fidelity emulator.
• (Planned) :mod:opifex.uncertainty.multi_fidelity.acquisition — Moss, Leslie, Rayson 2020 MUMBO max-value entropy acquisition for multi-fidelity Bayesian optimisation.

References

• Kennedy, O'Hagan 2000 — Predicting the output from a complex computer code when fast approximations are available, Biometrika.
• Perdikaris, Raissi, Damianou, Lawrence, Karniadakis 2017 — Nonlinear information fusion algorithms for data-efficient multi-fidelity modelling, Proc. R. Soc. A.
• Moss, Leslie, Rayson 2020 — MUMBO: MUlti-task Max-value Bayesian Optimisation, ECML-PKDD.

LinearMultiFidelityGPState dataclass

LinearMultiFidelityGPState(*, x_augmented: Array, y_train: Array, cholesky: Array, alpha: Array, lengthscales: tuple[float, ...], output_scales: tuple[float, ...], scaling_factors: tuple[float, ...], noise_std: float, base_kernel_fn: KernelFn = rbf_kernel)

Fitted state for a linear (K&O AR(1)) multi-fidelity GP.

Attributes:

Name	Type	Description
x_augmented	Array	(n_total, d+1) augmented training inputs including the fidelity-level column.
y_train	Array	(n_total,) training targets (concatenated across levels in increasing fidelity order).
cholesky	Array	Lower-triangular Cholesky factor of the joint K + noise_std**2 I Gram matrix.
alpha	Array	Pre-solved (K + sigma**2 I)^{-1} y.
lengthscales	tuple[float, ...]	Per-level length-scales used at fit time.
output_scales	tuple[float, ...]	Per-level output-scales used at fit time.
scaling_factors	tuple[float, ...]	AR(1) coupling factors used at fit time.
noise_std	float	Observation noise scale.
base_kernel_fn	KernelFn	Kernel used internally (defaults to RBF).

NonLinearMultiFidelityGPState dataclass

NonLinearMultiFidelityGPState(*, level_states: tuple[ExactGPState, ...], lengthscales: tuple[float, ...], output_scales: tuple[float, ...], noise_std: float)

Fitted state for a NARGP non-linear multi-fidelity GP.

Stores one :class:ExactGPState per fidelity level. Levels >= 1 are fitted on the augmented design [x, f_{i-1}(x)].

Attributes:

Name	Type	Description
level_states	tuple[ExactGPState, ...]	One :class:ExactGPState per fidelity level.
lengthscales	tuple[float, ...]	Per-level length-scale used at fit time.
output_scales	tuple[float, ...]	Per-level output-scale used at fit time.
noise_std	float	Observation noise scale used across levels.

fit_linear_multi_fidelity_gp

fit_linear_multi_fidelity_gp(*, x_train_per_level: Sequence[Array], y_train_per_level: Sequence[Array], lengthscales: Sequence[float], output_scales: Sequence[float], scaling_factors: Sequence[float], noise_std: float, base_kernel_fn: KernelFn = rbf_kernel) -> LinearMultiFidelityGPState

Fit a linear (K&O AR(1)) multi-fidelity GP.

Parameters:

Name	Type	Description	Default
x_train_per_level	Sequence[Array]	Per-fidelity training inputs (each (n_i, d)).	required
y_train_per_level	Sequence[Array]	Per-fidelity training targets (each (n_i,)).	required
lengthscales	Sequence[float]	(L,) per-level length-scales.	required
output_scales	Sequence[float]	(L,) per-level output-scales.	required
scaling_factors	Sequence[float]	(L-1,) AR(1) coupling factors.	required
noise_std	float	Observation noise scale (also acts as the joint jitter for Cholesky stability).	required
base_kernel_fn	KernelFn	Optional base kernel; defaults to :func:opifex.uncertainty.gp.rbf_kernel.	rbf_kernel

Returns:

Type	Description
LinearMultiFidelityGPState	class:LinearMultiFidelityGPState with Cholesky + alpha.

linear_multi_fidelity_kernel

linear_multi_fidelity_kernel(x1: Array, x2: Array, *, lengthscales: Sequence[float], output_scales: Sequence[float], scaling_factors: Sequence[float], base_kernel_fn: KernelFn = rbf_kernel) -> Array

Kennedy & O'Hagan AR(1) joint kernel block matrix.

Parameters:

Name	Type	Description	Default
x1	Array	(n1, d+1) augmented inputs whose final column is the integer fidelity level in {0, 1, ..., L-1}.	required
x2	Array	(n2, d+1) augmented inputs.	required
lengthscales	Sequence[float]	One length-scale per fidelity level.	required
output_scales	Sequence[float]	One output-scale per fidelity level.	required
scaling_factors	Sequence[float]	(L-1,) AR(1) coupling factors rho_1, ..., rho_{L-1}. Empty tuple for a single fidelity.	required
base_kernel_fn	KernelFn	Base kernel callable; defaults to :func:opifex.uncertainty.gp.rbf_kernel.	rbf_kernel

Returns:

Type	Description
Array	(n1, n2) joint Gram matrix.

predict_linear_multi_fidelity_gp

predict_linear_multi_fidelity_gp(*, state: LinearMultiFidelityGPState, x_test: Array, target_level: int) -> PredictiveDistribution

Predict the GP posterior at x_test for the given fidelity level.

Parameters:

Name	Type	Description	Default
state	LinearMultiFidelityGPState	Fitted linear-MF state.	required
x_test	Array	(m, d) test inputs (without level column).	required
target_level	int	Integer fidelity level at which to predict.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution with mean and variance.

mumbo_acquisition

mumbo_acquisition(*, state: LinearMultiFidelityGPState, x_candidates: Array, candidate_levels: Array, target_level: int, rng_key: Array, grid_size: int = 1000, num_gumbel_samples: int = 10, num_quadrature_points: int = 1000) -> Array

MUMBO multi-fidelity acquisition score per candidate.

Parameters:

Name	Type	Description	Default
state	LinearMultiFidelityGPState	Fitted linear-MF GP state.	required
x_candidates	Array	(m, d) candidate inputs (without level column).	required
candidate_levels	Array	(m,) integer fidelity levels.	required
target_level	int	Highest-fidelity (objective) level index.	required
rng_key	Array	JAX PRNG key for Gumbel sampling.	required
grid_size	int	Random grid size used to fit the Gumbel approximation to the target-level GP marginal.	1000
num_gumbel_samples	int	Number of Monte-Carlo samples drawn from the fitted Gumbel for the outer expectation.	10
num_quadrature_points	int	Number of Simpson-rule points used for the ESG entropy integral.	1000

Returns:

Type	Description
Array	(m,) acquisition scores (information gain in nats).
Array	Higher = more informative candidate.

fit_nonlinear_multi_fidelity_gp

fit_nonlinear_multi_fidelity_gp(*, x_train_per_level: Sequence[Array], y_train_per_level: Sequence[Array], lengthscales: Sequence[float], output_scales: Sequence[float], noise_std: float) -> NonLinearMultiFidelityGPState

Greedy recursive fit of a NARGP non-linear multi-fidelity GP.

Parameters:

Name	Type	Description	Default
x_train_per_level	Sequence[Array]	Per-fidelity training inputs (each (n_i, d)).	required
y_train_per_level	Sequence[Array]	Per-fidelity training targets (each (n_i,)).	required
lengthscales	Sequence[float]	(L,) per-level length-scales for the independent ExactGPs.	required
output_scales	Sequence[float]	(L,) per-level output-scales.	required
noise_std	float	Observation noise scale (shared across levels — also acts as the joint jitter for Cholesky stability).	required

Returns:

Type	Description
NonLinearMultiFidelityGPState	class:NonLinearMultiFidelityGPState with one fitted
NonLinearMultiFidelityGPState	class:ExactGPState per fidelity level.

predict_nonlinear_multi_fidelity_gp

predict_nonlinear_multi_fidelity_gp(*, state: NonLinearMultiFidelityGPState, x_test: Array, target_level: int, num_samples: int = 64, rng_key: Array) -> PredictiveDistribution

Predict at x_test for the given fidelity level via MC propagation.

Uncertainty propagates through the chain by Monte-Carlo: at each level above 0, num_samples posterior samples of the previous level are fed as the (d+1)-th input column, and the resulting per-sample predictive means / variances are aggregated via the standard mixture moments.

Parameters:

Name	Type	Description	Default
state	NonLinearMultiFidelityGPState	Fitted :class:NonLinearMultiFidelityGPState.	required
x_test	Array	(m, d) test inputs (without level column).	required
target_level	int	Integer fidelity level at which to predict.	required
num_samples	int	Monte-Carlo samples used to propagate previous-level uncertainty.	64
rng_key	Array	JAX PRNG key for sampling.	required

Returns:

Type	Description
PredictiveDistribution	class:PredictiveDistribution with mean / variance.

Inference, sensitivity, and decision-making

Simulation-Based Inference subsystem (Task 8.2).

Subpackage modules:

• :mod:opifex.uncertainty.sbi.simulators — :class:Simulator static container + :func:sample_joint joint-sampling helper.
• :mod:opifex.uncertainty.sbi.posterior_estimation — :class:NeuralPosteriorEstimator (NPE) fitting q(theta | x).
• :mod:opifex.uncertainty.sbi.likelihood_estimation — :class:NeuralLikelihoodEstimator (NLE) fitting q(x | theta) + posterior MCMC via the BlackJAX backend.
• :mod:opifex.uncertainty.sbi.ratio_estimation — :class:NeuralRatioEstimator (NRE) fitting log r(theta, x) + posterior MCMC via the BlackJAX backend.
• :mod:opifex.uncertainty.sbi.diagnostics — Simulation-Based Calibration (SBC) + expected posterior contraction.

Default density-estimator backend wraps Artifex's NNX-native flows (ConditionalRealNVP); optional backends (bijx / sbiax / flowMC) raise :class:ImportError with the canonical install hint when not present. MCMC sampling for NLE / NRE routes through :class:opifex.uncertainty.inference_backends.BlackJAXBackend.

References (read-only):

• Greenberg, Nonnenmacher, Macke (2019) — APT/NPE, arXiv:1905.07488.
• Papamakarios, Sterratt, Murray (2019) — NLE, arXiv:1805.07226.
• Hermans, Begy, Louppe (2020) — NRE, arXiv:1903.04057.
• Talts, Betancourt, Simpson, Vehtari, Gelman (2018) — SBC, arXiv:1804.06788.

PosteriorContractionResult

Typed result of an :func:expected_posterior_contraction run.

Fields:

• contraction — scalar mean contraction across observations and parameter dimensions (1 - var_post / var_prior averaged).
• per_dim — (theta_dim,) per-dimension mean contraction.

metadata_dict

metadata_dict() -> dict[str, Any]

Return a fresh dict view of the immutable metadata.

validate

validate() -> None

Eager-validate shape invariants.

SBCResult

Typed result of a Simulation-Based Calibration run (pattern (B)).

Fields:

• ranks — (num_runs, theta_dim) per-dimension rank of the ground-truth theta_star among the posterior samples.
• ks_statistic — (theta_dim,) per-dimension Kolmogorov-Smirnov statistic against Uniform(0, 1) on the normalised ranks.
• ks_pvalue — (theta_dim,) per-dimension KS p-value (large values support uniformity, i.e., calibration).

metadata_dict

metadata_dict() -> dict[str, Any]

Return a fresh dict view of the immutable metadata.

validate

validate() -> None

Eager-validate field-shape invariants. Not called from pytree path.

Raises:

Type	Description
ValueError	When ranks/statistic/p-value shapes are inconsistent.

NeuralLikelihoodEstimator dataclass

NeuralLikelihoodEstimator(theta_dim: int, x_dim: int, backend: str = _DEFAULT_BACKEND_NAME, num_steps: int = 100, learning_rate: float = 0.001, hidden_dim: int = 32, num_coupling_layers: int = 4, mcmc_method: str = 'nuts', mcmc_samples: int = 200, mcmc_burnin: int = 50, mcmc_step_size: float = 0.1)

Neural Likelihood Estimator with a conditional-flow likelihood and BlackJAX MCMC.

Parameters:

Name	Type	Description	Default
theta_dim	int	Parameter-space dimension.	required
x_dim	int	Observation-space dimension.	required
backend	str	Density-estimator backend name (default "RealNVP").	_DEFAULT_BACKEND_NAME
num_steps	int	Number of training steps (full-batch Adam).	100
learning_rate	float	Adam learning rate.	0.001
hidden_dim	int	Width of the coupling MLP hidden layers.	32
num_coupling_layers	int	Coupling-layer depth.	4
mcmc_method	str	BlackJAX sampler family ("nuts" / "hmc" / "mala").	'nuts'
mcmc_samples	int	Number of posterior MCMC samples to draw at predict time.	200
mcmc_burnin	int	Number of warmup MCMC samples to discard.	50
mcmc_step_size	float	BlackJAX step size; tuned to a small value for low-dim toys where NUTS adaptation can still mix well.	0.1

fit

fit(simulator: Simulator, num_simulations: int, *, rngs: Rngs) -> NeuralLikelihoodEstimator

Fit q(x | theta) on num_simulations joint draws.

Raises:

Type	Description
ImportError	When the requested optional backend is not installed.

predict_distribution

predict_distribution(observation: Array, *, rngs: Rngs, num_samples: int, log_prior: Callable[[Array], Array]) -> PredictiveDistribution

Run posterior MCMC at observation and return a :class:PredictiveDistribution.

NLEState

Bases: _SBIFittedState

Fitted-state container for :class:NeuralLikelihoodEstimator (pattern (B)).

NeuralPosteriorEstimator dataclass

NeuralPosteriorEstimator(theta_dim: int, x_dim: int, backend: str = _DEFAULT_BACKEND_NAME, num_steps: int = 100, learning_rate: float = 0.001, hidden_dim: int = 32, num_coupling_layers: int = 4)

Neural Posterior Estimator with a conditional-flow density estimator.

Parameters:

Name	Type	Description	Default
theta_dim	int	Parameter-space dimension.	required
x_dim	int	Observation-space dimension (post-summary if a summary_fn is configured on the simulator).	required
backend	str	Density-estimator backend name; must be one of the names registered in :class:BackendRouter "flow" family. Default "RealNVP" uses Artifex ConditionalRealNVP.	_DEFAULT_BACKEND_NAME
num_steps	int	Number of training steps (full-batch Adam).	100
learning_rate	float	Adam learning rate.	0.001
hidden_dim	int	Width of the coupling MLP hidden layers.	32
num_coupling_layers	int	Coupling-layer depth in the conditional flow.	4

Public state:

``state`` — fitted :class:`NPEState` after ``fit``; ``None`` otherwise.

fit

fit(simulator: Simulator, num_simulations: int, *, rngs: Rngs) -> NeuralPosteriorEstimator

Fit the conditional flow on num_simulations (theta, x) pairs.

Deterministic given a fixed rngs. Returns self with self.state populated.

Raises:

Type	Description
ImportError	When the requested optional backend is not installed.

predict_distribution

predict_distribution(observation: Array, *, rngs: Rngs, num_samples: int) -> PredictiveDistribution

Sample the fitted posterior q(theta | observation).

refine_round

refine_round(simulator: Simulator, *, observation: Array, num_simulations: int, num_steps: int, rngs: Rngs) -> NeuralPosteriorEstimator

Run one sequential SBI round centred on observation.

Sequential NPE narrows the proposal toward the current posterior: we draw the new round's parameter proposals from q(theta | observation) (the current flow) and train a fresh flow on the resulting (theta, x) pairs. With informative observations the posterior tightens across rounds.

NPEState

Bases: _SBIFittedState

Fitted-state container for :class:NeuralPosteriorEstimator (pattern (B)).

Shares its body and validate() with the other SBI estimators via :class:opifex.uncertainty.sbi._base._SBIFittedState; only the divergence-diagnostic label differs.

NeuralRatioEstimator dataclass

NeuralRatioEstimator(theta_dim: int, x_dim: int, backend: str = _DEFAULT_BACKEND_NAME, num_steps: int = 100, learning_rate: float = 0.001, hidden_dim: int = 32, mcmc_method: str = 'nuts', mcmc_samples: int = 200, mcmc_burnin: int = 50, mcmc_step_size: float = 0.1)

Neural Ratio Estimator with an MLP classifier and BlackJAX MCMC.

Parameters:

Name	Type	Description	Default
theta_dim	int	Parameter-space dimension.	required
x_dim	int	Observation-space dimension.	required
backend	str	Density-estimator backend name ("RealNVP" by default mirrors NPE/NLE for consistent backend-selection semantics).	_DEFAULT_BACKEND_NAME
num_steps	int	Number of training steps.	100
learning_rate	float	Adam learning rate.	0.001
hidden_dim	int	Width of the classifier MLP.	32
mcmc_method	str	BlackJAX sampler family.	'nuts'
mcmc_samples	int	Number of posterior samples to draw.	200
mcmc_burnin	int	MCMC warmup samples to discard.	50
mcmc_step_size	float	BlackJAX step size.	0.1

fit

fit(simulator: Simulator, num_simulations: int, *, rngs: Rngs) -> NeuralRatioEstimator

Fit the ratio classifier on positives (theta, x) vs marginals.

Raises:

Type	Description
ImportError	When the requested optional backend is not installed.

predict_distribution

predict_distribution(observation: Array, *, rngs: Rngs, num_samples: int, log_prior: Callable[[Array], Array]) -> PredictiveDistribution

Run MCMC over log r(theta, x_obs) + log prior(theta).

NREState

Bases: _SBIFittedState

Fitted-state container for :class:NeuralRatioEstimator (pattern (B)).

Simulator dataclass

Simulator(*, prior_sampler: PriorSampler, simulate_fn: SimulateFn, summary_fn: SummaryFn | None = None, metadata: MetadataItems = ())

Static SBI simulator description (pattern (A)).

Parameters:

Name	Type	Description	Default
prior_sampler	PriorSampler	(rng_key, num_samples) -> theta returning an array of shape (num_samples, *theta_event_shape).	required
simulate_fn	SimulateFn	(rng_key, theta) -> x returning observations of shape (num_samples, *x_event_shape) whose leading axis matches the leading axis of theta.	required
summary_fn	SummaryFn | None	Optional x -> s(x) mapping applied per-element after simulation. Used for summary-statistic compression (e.g., neural compression to a low-dim s(x)).	None
metadata	MetadataItems	Immutable tuple[tuple[str, Any], ...] of static metadata. Hashable so callers can pass the simulator through jit's static-argnums without losing the cache.	()

The dataclass is intentionally minimal — array state belongs in the estimator state container (pattern (B)), not here.

metadata_dict

metadata_dict() -> dict[str, Any]

Return a mutable dict view of the immutable metadata.

validate

validate() -> None

Eager-validate the simulator's contract.

Public method — never called from __post_init__ or the pytree path (GUIDE_ALIGNMENT item 7).

Raises:

Type	Description
TypeError	When prior_sampler / simulate_fn / summary_fn are not callable, or metadata is not a tuple-of-pairs.

expected_posterior_contraction

expected_posterior_contraction(estimator: NeuralPosteriorEstimator, simulator: Simulator, *, rngs: Rngs, num_observations: int = 16, num_posterior_samples: int = 200) -> PosteriorContractionResult

Compute mean 1 - var_post / var_prior across observations.

Positive on informative observations, ~0 on uninformative ones, and can be slightly negative under finite-sample noise (no clipping here — leave the raw signal for the caller).

simulation_based_calibration

simulation_based_calibration(estimator: NeuralPosteriorEstimator, simulator: Simulator, *, rngs: Rngs, num_runs: int, num_posterior_samples: int = 100) -> SBCResult

Run SBC on estimator using fresh draws from simulator.

Parameters:

Name	Type	Description	Default
estimator	NeuralPosteriorEstimator	Fitted NPE.	required
simulator	Simulator	The same simulator used to generate the fit data.	required
rngs	Rngs	Caller-owned :class:nnx.Rngs carrying sbi_simulate and sbi_sample streams.	required
num_runs	int	Number of ground-truth draws used to populate ranks.	required
num_posterior_samples	int	Posterior samples drawn per run for the rank computation.	100

Returns:

Type	Description
SBCResult	class:SBCResult with rank statistics + per-dim KS statistics.

sample_joint

sample_joint(simulator: Simulator, *, num_simulations: int, rngs: Rngs) -> Batch

Sample num_simulations joint draws from the prior and simulator.

Returns a Datarax Batch where every element pairs theta with its observation x (and optionally a compressed summary).

The RNG resolution uses Artifex's :func:extract_rng_key with the canonical SBI stream name "sbi_simulate".

Parameters:

Name	Type	Description	Default
simulator	Simulator	Static simulator description.	required
num_simulations	int	Number of (theta, x) pairs to draw.	required
rngs	Rngs	Caller-owned nnx.Rngs carrying the sbi_simulate named stream.	required

Raises:

Type	Description
ValueError	When num_simulations is non-positive, when the prior sampler returns a malformed leading axis, or when the simulator output's leading axis disagrees with the parameters.

Global-sensitivity utilities (Task 6.4).

Two complementary techniques for model-input sensitivity:

• :func:sobol_indices — variance-based first-order and total-order indices via the Saltelli (2002) pick-freeze scheme.
• :func:morris_screening — elementary-effects screening for cheap ranking of influential inputs (Morris 1991 / Campolongo+ 2007).

Both utilities are pure JAX and accept arbitrary scalar-valued model: Callable[[jax.Array], jax.Array] callables; neither uses SALib (Task 6.4 forbids it as a hard dependency).

MorrisResult dataclass

MorrisResult(*, mu_star: Array, mu: Array, sigma: Array, num_trajectories: int, num_levels: int)

Container for Morris screening statistics.

Attributes:

Name	Type	Description
mu_star	Array	Mean of absolute elementary effects, shape (d,). Monotone with overall input influence.
mu	Array	Mean of signed elementary effects, shape (d,). Sign indicates effect direction.
sigma	Array	Standard deviation of elementary effects, shape (d,). Large sigma relative to mu_star flags non-linearity or interaction with other inputs.
num_trajectories	int	r, the number of independent trajectories.
num_levels	int	p, the number of grid levels per dimension.

SobolResult dataclass

SobolResult(*, first_order: Array, total_order: Array, variance: Array, num_samples: int)

Container for first-order + total-order Sobol indices.

Attributes:

Name	Type	Description
first_order	Array	S_i = V[E[Y|X_i]] / V[Y], shape (d,).
total_order	Array	T_i = E[V[Y|X_{~i}]] / V[Y], shape (d,).
variance	Array	V[Y] estimate over the combined sample.
num_samples	int	Base sample count N (total f evaluations = N * (d + 2)).

morris_screening

morris_screening(model: Callable[[Array], Array], *, num_trajectories: int, num_levels: int, lower: Array, upper: Array, rng_key: Array) -> MorrisResult

Run Morris elementary-effects screening on model.

Parameters:

Name	Type	Description	Default
model	Callable[[Array], Array]	Scalar-valued model f(x) accepting input arrays of shape (..., d) and returning (...,). JAX-traceable.	required
num_trajectories	int	r, the number of independent trajectories.	required
num_levels	int	p, the number of grid levels per dimension. Must be >= 2.	required
lower	Array	Lower box bounds, shape (d,).	required
upper	Array	Upper box bounds, shape (d,).	required
rng_key	Array	Caller-owned JAX PRNG key.	required

Returns:

Name	Type	Description
A	MorrisResult	class:MorrisResult with mu_star, mu, and sigma
	MorrisResult	statistics over r trajectories.

Raises:

Type	Description
ValueError	If lower / upper shapes mismatch, num_trajectories <= 0, or num_levels < 2.

sobol_indices

sobol_indices(model: Callable[[Array], Array], *, num_samples: int, lower: Array, upper: Array, rng_key: Array) -> SobolResult

Estimate first- + total-order Sobol indices via Saltelli (2002).

Parameters:

Name	Type	Description	Default
model	Callable[[Array], Array]	Scalar-valued model f(x) accepting input arrays of shape (..., d) and returning shape (...,). The function must be JAX-traceable.	required
num_samples	int	Base sample count N (total evaluations = N * (d + 2)).	required
lower	Array	Lower box bounds, shape (d,).	required
upper	Array	Upper box bounds, shape (d,).	required
rng_key	Array	Caller-owned JAX PRNG key. Two independent N x d sample matrices are drawn from rng_key.	required

Returns:

Name	Type	Description
A	SobolResult	class:SobolResult with first-order, total-order, and
	SobolResult	variance estimates.

Raises:

Type	Description
ValueError	If lower and upper have different shapes, or num_samples <= 0.

Active Learning and Bayesian Experimental Design subsystem (Task 8.3).

Subpackage modules:

• :mod:opifex.uncertainty.active.acquisition — pure JAX BALD / EI / Log-EI / UCB / LCB / PI single-point acquisition kernels operating on :class:~opifex.uncertainty.types.PredictiveDistribution objects and a named-strategy :func:acquire dispatcher.
• :mod:opifex.uncertainty.active.batch_active — BatchBALD greedy joint-MI maximisation, reparameterised Monte-Carlo batch EI, and the q-EHVI multi-objective acquisition.
• :mod:opifex.uncertainty.active.experimental_design — :func:expected_information_gain (linear-Gaussian closed form plus Monte-Carlo nested-sampling fallback) and the :func:bayesian_experimental_design_loop BO loop driver.
• :mod:opifex.uncertainty.active.pinn_acquisition — PINN residual-norm acquisition that ranks candidates by PDE residual magnitude and exposes residual + uncertainty metadata.

All RNG-dependent acquisitions route through :func:artifex.generative_models.core.rng.extract_rng_key with named streams "active_acquire", "active_bald", "active_eig".

Container patterns (GUIDE_ALIGNMENT §5a):

• :class:AcquisitionStrategy — :class:~enum.StrEnum of strategy names.
• :class:ActiveLearningConfig — pattern (A): frozen, slotted, keyword-only dataclass.
• :class:AcquiredBatch — pattern (B): flax.struct.dataclass that carries jax.Array indices + scores through the batch loop.

Primary acquisition-function reference: trieste (TensorFlow original ported to JAX). Each kernel docstring cites the trieste source line it was ported from.

AcquiredBatch

Result of one acquisition round.

GUIDE_ALIGNMENT pattern (B): flax.struct.dataclass carrying jax.Array payloads. strategy and metadata are marked pytree_node=False so they participate in the JIT cache key rather than the leaf list. :meth:validate is public and is never called from the unflatten path.

metadata_dict

metadata_dict() -> dict[str, Any]

Return a fresh dict view of the immutable metadata tuple.

validate

validate() -> None

Eager-validate shapes; not called from __post_init__.

AcquisitionStrategy

Bases: StrEnum

Named acquisition strategies dispatched by :func:acquire.

ActiveLearningConfig dataclass

ActiveLearningConfig(*, strategy: AcquisitionStrategy, batch_size: int, acquisition_family: str = 'single_point', extra_streams: tuple[str, ...] = ())

Configuration container for an active-learning round.

GUIDE_ALIGNMENT pattern (A): plain @dataclass(frozen=True, slots=True, kw_only=True). Sequence fields use tuple[...] (GUIDE_ALIGNMENT item 22a). __post_init__ performs eager validation.

BayesianExperimentalDesignResult dataclass

BayesianExperimentalDesignResult(*, acquired_indices: tuple[int, ...], acquired_x: tuple[Any, ...] = tuple(), acquired_y: tuple[Any, ...] = tuple(), history_variance: tuple[float, ...] = tuple(), history_score: tuple[float, ...] = tuple(), metadata: tuple[tuple[str, Any], ...] = tuple())

Result container for :func:bayesian_experimental_design_loop.

GUIDE_ALIGNMENT pattern (A): plain frozen-slotted-kw-only dataclass. All array history is materialised as tuple[...] so the container is hashable (sequence fields use tuple, item 22a).

acquire

acquire(predictive_dist: PredictiveDistribution, *, strategy: AcquisitionStrategy | str, batch_size: int, rngs: Rngs | Array, metadata: MetadataItems = (), **kwargs: Any) -> AcquiredBatch

Named-strategy acquisition dispatcher.

The top-batch_size candidates (by descending score) are returned inside an :class:AcquiredBatch. This is the entry point invoked by the rewritten :class:opifex.training.uncertainty_trainers.ActiveUncertaintyLearner — the trainer wraps the predictive distribution and delegates here.

Note: this is a "naive top-K" greedy batch over single-point scores; for diversity-aware batch acquisition use :func:batch_bald or :func:batch_mc_expected_improvement from :mod:opifex.uncertainty.active.batch_active.

bald

bald(predictive_dist: PredictiveDistribution, *, rngs: Rngs | Array) -> Array

Bayesian Active Learning by Disagreement (regression-ensemble form).

Ported from trieste/acquisition/function/active_learning.py:418 (bayesian_active_learning_by_disagreement). The trieste original targets binary GP-classification with a Bernoulli likelihood; opifex generalises to the regression ensemble carried by :class:PredictiveDistribution.samples because that's the shape every Phase-7 Bayesian backend already produces. The mutual information is

.. math:: \mathrm{BALD}(x) = H[\hat p(y \mid x)] - \mathbb{E}_{\theta} [H[p(y \mid x, \theta)]]

where the predictive marginal entropy uses the moment-matched Gaussian approximation (same approximation as the trieste original) and the per-sample entropies are exact Gaussians with the carried aleatoric variance. The rngs argument is unused in the closed-form regression branch but is retained for API parity with sampling-based variants (it is consumed eagerly so callers see no spurious key reuse).

expected_improvement

expected_improvement(predictive_dist: PredictiveDistribution, *, best_value: float) -> Array

Single-point Expected Improvement (minimisation convention).

Ported from trieste/acquisition/function/function.py:226 (expected_improvement.__call__):

.. math:: \mathrm{EI}(x) = (\eta - \mu(x)) \Phi!\left(\tfrac{\eta - \mu(x)}{\sigma(x)}\right) + \sigma(x) \, \phi!\left(\tfrac{\eta - \mu(x)}{\sigma(x)}\right)

Substitutions vs. trieste: tfp.distributions.Normal → jax.scipy.stats.norm.

log_expected_improvement

log_expected_improvement(predictive_dist: PredictiveDistribution, *, best_value: float) -> Array

Numerically stable log-EI.

Ported from trieste/acquisition/function/function.py:269 (log_expected_improvement.__call__). Mitigates the vanishing-gradient issue of standard EI in regions where EI is tiny (Ament et al. 2023, Unexpected EI).

lower_confidence_bound

lower_confidence_bound(predictive_dist: PredictiveDistribution, *, beta: float) -> Array

LCB mu - beta * sigma (minimise for exploration).

Ported from trieste/acquisition/function/function.py:571 (lower_confidence_bound).

probability_of_improvement

probability_of_improvement(predictive_dist: PredictiveDistribution, *, best_value: float) -> Array

Probability that f(x) < best_value under the posterior.

Ported from trieste/acquisition/function/function.py:684 (probability_below_threshold.__call__). For a Gaussian posterior f(x) ~ N(mu, sigma^2) the closed form is Phi((best_value - mu) / sigma).

upper_confidence_bound

upper_confidence_bound(predictive_dist: PredictiveDistribution, *, beta: float) -> Array

UCB mu + beta * sigma (maximise for exploration).

Mirror of LCB; same trieste reference at trieste/acquisition/function/function.py:571 (NegativeLowerConfidenceBound simply negates lower_confidence_bound).

batch_bald

batch_bald(predictive_dist: PredictiveDistribution, *, batch_size: int, rngs: Rngs | Array, num_mc_samples: int = 1024) -> AcquiredBatch

Greedy BatchBALD acquisition.

Selects batch_size candidates by greedily maximising the joint mutual information

.. math:: I[{y_i}{i \in S}; \theta] = H[\hat p({y_i})] - \sum_{i \in S} \mathbb{E}_\theta [H[p(y_i \mid \theta)]]

Reference: Kirsch, van Amersfoort & Gal (2019), BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning (Algorithm 1). The joint predictive entropy H[E_theta p(y_S)] is a Gaussian-mixture entropy with no closed form; we use the Monte Carlo estimator from the paper (sample from the mixture, evaluate log-mean-likelihood). The conditional sum E_theta H[p|theta] has the trivial Gaussian closed form because the noise per candidate is conditionally independent in the ensemble model.

The MC mixture-entropy estimator correctly handles BatchBALD's key edge case: when two candidates have identical ensemble samples (perfect redundancy), the joint entropy estimator returns the marginal entropy of either point, so the MI gain for the second redundant point collapses to zero and the greedy step picks a diverse alternative.

The greedy outer structure mirrors trieste/acquisition/function/greedy_batch.py.

batch_mc_expected_improvement

batch_mc_expected_improvement(*, mean: Array, std: Array, best_value: float, num_samples: int, rngs: Rngs | Array) -> Array

Reparameterised Monte-Carlo batch Expected Improvement.

Ported from trieste/acquisition/function/function.py:1364 (batch_monte_carlo_expected_improvement.__call__):

.. code-block:: text

samples = sampler.sample(x, jitter=jitter)            # [S, B]
min_per_sample = reduce_min(samples, axis=-1)         # [S]
improvement = maximum(eta - min_per_sample, 0.0)      # [S]
return reduce_mean(improvement, axis=-1)              # scalar

Substitutions vs. trieste: the model's reparam sampler is replaced by an explicit diagonal-Gaussian reparameterisation y = mean + std * eps with eps ~ N(0, I) because the active subsystem operates on :class:PredictiveDistribution moments rather than a model object.

dominated_hypervolume

dominated_hypervolume(front: Array, reference_point: Array) -> Array

General-M dominated hypervolume via a grid box decomposition.

Computes the volume of \{z : z \le r, \exists p \in P. p \preceq z\} (minimisation) — the region dominated by the Pareto front P below the reference point r. The dominated region is tiled by an axis-aligned grid whose lines are the front's per-objective coordinates plus the reference point; a cell is summed when its centre is dominated by some front point.

This is the JAX-native analogue of trieste's Pareto.hypervolume_indicator box decomposition (trieste/acquisition/multi_objective/pareto.py) and the grid decomposition described by Daulton, Balandat & Bakshy (2020, arXiv:2006.05078). The M = 2 case is routed to the exact staircase :func:_pareto_volume_2d (the fast special-case path); higher M generalises that staircase to a hyper-rectangle grid and is unit-tested to agree with it at M = 2. Dominated (redundant) front points contribute no extra volume.

Parameters:

Name	Type	Description	Default
front	Array	Shape (P, M) Pareto front (need not be screened of dominated points).	required
reference_point	Array	Shape (M,) upper bound of the hypervolume.	required

Returns:

Type	Description
Array	Scalar dominated hypervolume.

q_expected_hypervolume_improvement

q_expected_hypervolume_improvement(*, candidate_mean: Array, candidate_std: Array, pareto_front: Array, reference_point: Array, num_samples: int, rngs: Rngs | Array) -> Array

Monte-Carlo q-EHVI multi-objective acquisition (general M).

Implements the parallel Expected Hypervolume Improvement of Daulton, Balandat & Bakshy (2020), Differentiable Expected Hypervolume Improvement for Parallel Multi-Objective Bayesian Optimization (NeurIPS, arXiv:2006.05078). Ported from trieste's batch_ehvi (trieste/acquisition/function/multi_objective.py:211). The acquisition value is

.. math:: \alpha_{\text{qEHVI}}(\mathcal{X}) = \mathbb{E}!\Big[ \sum_{\emptyset \ne J \subseteq {1,\dots,q}} (-1)^{|J|+1} \sum_{k} \prod_{m} \max!\big(u^{(k)}m - \max_m), 0\big) \Big],} \max(z^{(j)}_m, l^{(k)

where \{[l^{(k)}, u^{(k)})\}_k is a box decomposition of the non-dominated region below the reference point and the expectation is over the posterior of the batch.

The opifex port:

• Replaces trieste's model reparam sampler with an explicit diagonal-Gaussian reparameterisation y = mean + std * eps, eps ~ N(0, I) (the active subsystem operates on :class:PredictiveDistribution moments, not a model object).
• Replaces trieste's prepare_default_non_dominated_partition_bounds with a JAX-native grid box decomposition (:func:_non_dominated_partition_bounds) so the whole acquisition stays jit/grad/vmap compatible for any M.

The exact 2-D staircase (:func:_pareto_volume_2d) is retained as a fast special case for the dominated-hypervolume helper and the general path is unit-tested to agree with it at M = 2.

Parameters:

Name	Type	Description	Default
candidate_mean	Array	Shape (q, M) predictive mean per candidate and per objective.	required
candidate_std	Array	Shape (q, M) predictive std.	required
pareto_front	Array	Shape (P, M) current non-dominated set.	required
reference_point	Array	Shape (M,) reference point.	required
num_samples	int	Monte-Carlo sample count.	required
rngs	Rngs | Array	Caller-owned nnx.Rngs bundle or a raw jax.Array PRNG key used for the reparameterisation draws.	required

Returns:

Type	Description
Array	Scalar Monte-Carlo estimate of the q-EHVI (non-negative).

bayesian_experimental_design_loop

bayesian_experimental_design_loop(*, surrogate: Surrogate, candidates: Array, oracle: Callable[[Array], Array], acquisition: Callable[[PredictiveDistribution], Array], num_rounds: int, rngs: Rngs | Array) -> BayesianExperimentalDesignResult

Ask-tell Bayesian-experimental-design / BO loop.

Mirrors the structure of trieste/ask_tell_optimization.py:742 (AskTellOptimizer). The opifex variant exposes a tiny surface:

1. Predict the surrogate's :class:PredictiveDistribution over candidates.
2. Apply the acquisition callable to the predictive distribution.
3. Pick argmax of the scores; record the index and the per-candidate mean predictive variance.
4. Query the oracle at the chosen candidate.
5. Update the surrogate with (x_new, y_new).

rngs is consumed up-front to keep the loop deterministic given a fixed key (the acquisition function itself is responsible for any internal stochasticity).

expected_information_gain

expected_information_gain(*, model: Callable[[Array, Array], Array], design: Array, prior_samples: Array, noise_std: float, rngs: Rngs | Array, num_outer_samples: int | None = None) -> Array

Monte-Carlo estimator of the Expected Information Gain.

.. math:: \mathrm{EIG}(d) = \mathbb E_{\theta} !\left[ \mathbb E_{y \mid \theta} !\left[ \log p(y \mid \theta, d) - \log p(y \mid d) \right] \right]

With y = model(design, theta) + N(0, sigma^2) we compute the nested-MC estimator (Foster et al. 2019, Variational Bayesian Optimal Experimental Design) using the same prior-sample set for both outer and inner expectations (the "amortised" form). For the linear-Gaussian case y = d^T theta with prior theta ~ N(0, sigma_p^2 I) the estimator converges to 0.5 * log(1 + sigma_p^2 ||d||^2 / sigma_n^2).

The opifex implementation reuses the same prior samples as the inner and outer expectations, which is the standard amortised nested-MC estimator. num_outer_samples defaults to all prior samples.

pinn_residual_acquisition

pinn_residual_acquisition(*, model: Module | Callable[..., Array], candidates: Array, residual_fn: Callable[[Array, Array], Array], batch_size: int, rngs: Rngs | Array, uncertainty_fn: Callable[[Array], Array] | None = None) -> AcquiredBatch

Rank candidates by |residual_fn(predictions, candidates)|.

Returns the top-batch_size candidates and exposes both the per-candidate residual norm and a (placeholder) predictive uncertainty signal as metadata items so downstream consumers can inspect why each point was selected. When uncertainty_fn is supplied it is invoked on the model predictions; otherwise the metadata field carries zeros (deterministic surrogate).

The acquisition is a JAX-traceable kernel — both _invoke_model and residual_fn must operate purely on jax.Array inputs.

PAC-Bayes objective + certificate subsystem.

Modules:

• :mod:opifex.uncertainty.pac_bayes.bounds — pure JAX bound formulas (McAllester, Catoni, quadratic / kl-inversion).
• :mod:opifex.uncertainty.pac_bayes.objectives — :func:pac_bayes_kl_objective differentiable training objective.
• :mod:opifex.uncertainty.pac_bayes.certificates — :class:PACBayesCertificate typed result and the :func:pac_bayes_certificate driver.

All bound formulas delegate KL computation to :func:opifex.uncertainty.kernels.bayesian.diagonal_gaussian_kl (which itself delegates to Artifex gaussian_kl_divergence for the standard-normal prior), so the KL formula lives in exactly one place across the stack.

Canonical references (read-only):

• Dziugaite & Roy (2017) — arXiv:1703.11008.
• Pérez-Ortiz et al. (JMLR v22) — empirical PAC-Bayes for deep networks.
• Alquier (2024) — arXiv:2110.11216 (survey).

PACBayesCertificate

Typed PAC-Bayes certificate record.

Array fields (bound_value, kl, dataset_size) are pytree leaves so the container flows transparently through jax.jit / jax.grad; bound_name, delta, and metadata are static aux_data.

Fields:

• bound_value — scalar JAX array; the population-risk upper bound.
• kl — scalar JAX array; KL divergence between posterior and prior.
• dataset_size — scalar JAX array; n used to evaluate the bound.
• bound_name — one of "mcallester", "catoni", "quadratic".
• delta — confidence parameter in (0, 1).
• metadata — immutable tuple[tuple[str, Any], ...] aux_data; at minimum carries a "tightness" entry (bound_value - empirical_risk).

validate() is public and not wired into __post_init__ / tree_unflatten so JIT-traced reconstructions never trigger spurious failures.

metadata_dict

metadata_dict() -> dict[str, Any]

Return a fresh dict view of the immutable metadata tuple.

validate

validate() -> None

Eager-validate the certificate; never called from the pytree path.

Raises:

Type	Description
ValueError	If bound_name is not a known family, delta is outside (0, 1), or any scalar is non-finite.

catoni_bound

catoni_bound(empirical_risk: Array, kl: Array, dataset_size: Array | int, delta: float, *, beta: float = 1.0) -> Array

Catoni (2007) PAC-Bayes bound with temperature beta > 0.

With confidence 1 - delta:

``B = (beta * R + (KL + log(1 / delta)) / n) / (1 - exp(-beta))``.

The Catoni bound is not a probability — it is a real-valued upper bound on the population risk; the temperature beta trades empirical-risk sensitivity against the KL term.

Parameters:

Name	Type	Description	Default
empirical_risk	Array	Scalar empirical risk R.	required
kl	Array	Scalar KL divergence between posterior and prior, in nats.	required
dataset_size	Array | int	Sample size n (positive scalar).	required
delta	float	Confidence parameter in (0, 1).	required
beta	float	Strictly positive temperature.	1.0

Returns:

Type	Description
Array	Scalar Catoni upper bound.

Raises:

Type	Description
ValueError	If delta is not in (0, 1) or beta <= 0.

kl_bernoulli

kl_bernoulli(p: Array, q: Array) -> Array

Bernoulli KL kl(p || q) = p log(p/q) + (1-p) log((1-p)/(1-q)).

Inputs are clamped to (_EPS, 1 - _EPS) so the log terms stay finite at the boundary p in {0, 1}. The function is symmetric in neither argument and is convex in q for fixed p — the property used by the bisection in :func:quadratic_bound.

mcallester_bound

mcallester_bound(empirical_risk: Array, kl: Array, dataset_size: Array | int, delta: float) -> Array

Evaluate the McAllester (1999) PAC-Bayes bound.

With confidence 1 - delta:

``B = R + sqrt((KL + log(2 * sqrt(n) / delta)) / (2 * n))``.

Parameters:

Name	Type	Description	Default
empirical_risk	Array	Scalar empirical risk R in [0, 1].	required
kl	Array	Scalar KL divergence between posterior and prior, in nats.	required
dataset_size	Array | int	Sample size n (positive scalar).	required
delta	float	Confidence parameter in (0, 1).	required

Returns:

Type	Description
Array	Scalar PAC-Bayes upper bound on the population risk.

Raises:

Type	Description
ValueError	If delta is not in (0, 1).

quadratic_bound

quadratic_bound(empirical_risk: Array, kl: Array, dataset_size: Array | int, delta: float) -> Array

Quadratic / Maurer kl-inversion PAC-Bayes bound.

Returns the unique B in [R, 1) solving

``kl_bernoulli(R, B) = (KL + log(2 * sqrt(n) / delta)) / n``,

obtained by _INVERT_ITERS bisection steps. Because kl_bernoulli is strictly convex in its second argument with minimum at B = R, the root is unique on [R, 1) whenever the right-hand side is non-negative.

Parameters:

Name	Type	Description	Default
empirical_risk	Array	Scalar empirical risk in [0, 1].	required
kl	Array	Scalar KL divergence, in nats.	required
dataset_size	Array | int	Sample size n.	required
delta	float	Confidence parameter in (0, 1).	required

Returns:

Type	Description
Array	Scalar quadratic-bound upper bound on the population risk.

Raises:

Type	Description
ValueError	If delta is not in (0, 1).

pac_bayes_certificate

pac_bayes_certificate(model: Any, posterior: Any, data: dict[str, Any], *, delta: float, bound_name: str = 'mcallester', beta: float = 1.0) -> PACBayesCertificate

Evaluate a PAC-Bayes bound and return a typed certificate.

Parameters:

Name	Type	Description	Default
model	Any	Reference to the model whose posterior is being certified. The driver does not introspect the model; it is recorded in metadata for traceability.	required
posterior	Any	An object exposing kl_divergence (attribute or method) returning a scalar JAX array. Bayesian-NNX layers and VariationalResult containers both satisfy this protocol.	required
data	dict[str, Any]	Mapping with required keys "empirical_risk" (scalar) and "dataset_size" (positive scalar / int).	required
delta	float	Confidence parameter in (0, 1).	required
bound_name	str	Which bound family to evaluate; one of "mcallester", "catoni", "quadratic".	'mcallester'
beta	float	Catoni temperature; ignored for other bound families.	1.0

Returns:

Name	Type	Description
A	PACBayesCertificate	class:PACBayesCertificate whose bound_value is a scalar JAX
	PACBayesCertificate	array.

Raises:

Type	Description
ValueError	If delta is outside (0, 1) or bound_name is not a known family.
KeyError	If data lacks "empirical_risk" or "dataset_size".

pac_bayes_kl_objective

pac_bayes_kl_objective(empirical_risk: Array, kl: Array, dataset_size: Array | int, *, delta: float = 0.05) -> Array

Differentiable PAC-Bayes training objective.

Returns the McAllester bound

``R_hat + sqrt((KL + log(2 * sqrt(n) / delta)) / (2 * n))``

for delta in (0, 1). When delta == 1 the helper reduces to the Phase-1 ELBO scaling R_hat + KL / n so PAC-Bayes consumers can fall back to ELBO training by setting delta=1.

Parameters:

Name	Type	Description	Default
empirical_risk	Array	Scalar empirical risk (per-example loss after reduction).	required
kl	Array	Scalar KL divergence between posterior and prior.	required
dataset_size	Array | int	Sample size n; must be positive.	required
delta	float	Confidence parameter in (0, 1]. delta=1 selects the ELBO limit.	0.05

Returns:

Type	Description
Array	Scalar differentiable objective.

Raises:

Type	Description
ValueError	If delta is outside (0, 1].

Aggregation, reliability, and surrogates

Uncertainty quantification utilities for Bayesian neural networks.

Submodules:

• :mod:opifex.uncertainty.aggregators.types — value-object containers (UncertaintyComponents, CalibrationMetrics, UncertaintyIntegrationResults, EnhancedUncertaintyComponents).
• :mod:opifex.uncertainty.aggregators.basic — basic epistemic / aleatoric estimators and the :class:UncertaintyQuantifier integration interface.
• :mod:opifex.uncertainty.aggregators.calibration — reliability-binning helper and :class:CalibrationAssessment tooling.
• :mod:opifex.uncertainty.aggregators.enhanced — ensemble / distributional / multi-source quantifiers (:class:EnhancedUncertaintyQuantifier).

AleatoricUncertainty

Aleatoric (data) uncertainty estimation.

homoscedastic_uncertainty staticmethod

homoscedastic_uncertainty(_predictions: Float[Array, 'batch output'], log_variance: Float[Array, 'batch output']) -> Float[Array, 'batch output']

Compute homoscedastic (constant) aleatoric uncertainty.

heteroscedastic_uncertainty staticmethod

heteroscedastic_uncertainty(input_dependent_variance: Float[Array, 'batch output']) -> Float[Array, 'batch output']

Compute heteroscedastic (input-dependent) aleatoric uncertainty.

predictive_variance staticmethod

predictive_variance(predictions: Float[Array, 'samples batch output'], individual_variances: Float[Array, 'samples batch output']) -> Float[Array, 'batch output']

Compute total predictive variance including aleatoric component.

noise_estimation staticmethod

noise_estimation(residuals: Float[Array, 'batch output'], predictions: Float[Array, 'batch output']) -> Float[Array, 'batch output']

Estimate aleatoric uncertainty from residuals.

EpistemicUncertainty

Epistemic (model) uncertainty estimation.

compute_variance staticmethod

compute_variance(predictions: Float[Array, 'samples batch output']) -> Float[Array, 'batch output']

Compute epistemic uncertainty as variance across model samples.

compute_entropy staticmethod

compute_entropy(predictions: Float[Array, 'samples batch classes']) -> Float[Array, 'batch classes']

Compute predictive entropy for classification tasks.

compute_mutual_information staticmethod

compute_mutual_information(predictions: Float[Array, 'samples batch classes']) -> Float[Array, 'batch classes']

Compute mutual information between predictions and model parameters.

compute_variance_of_expected staticmethod

compute_variance_of_expected(predictions: Float[Array, 'samples batch output']) -> Float[Array, 'batch output']

Variance over the sample axis — pure epistemic uncertainty.

Equivalent to :meth:compute_variance; kept as a named alias so call sites can express intent (Var_θ[E[y|θ]]) when used alongside the aleatoric component E_θ[Var[y|θ]].

compute_ensemble_disagreement staticmethod

compute_ensemble_disagreement(ensemble_predictions: Float[Array, 'models batch output'], aggregation_method: str = 'variance') -> Float[Array, 'batch output']

Epistemic uncertainty from ensemble disagreement under multiple statistics.

aggregation_method selects the dispersion statistic — variance (same as :meth:compute_variance), std (standard deviation), range (max − min), or iqr (75^th − 25^th percentile).

compute_predictive_diversity staticmethod

compute_predictive_diversity(ensemble_predictions: Float[Array, 'models batch output'], diversity_metric: str = 'pairwise_distance') -> Float[Array, 'batch']

Average diversity across ensemble predictions.

diversity_metric selects the comparison:

• pairwise_distance — mean L2 distance between every ordered pair of ensemble members (batched along the batch axis).
• cosine_diversity — 1 − mean cos(member, mean_prediction), rewarding angular spread.

UncertaintyQuantifier

UncertaintyQuantifier(num_samples: int = 100, confidence_level: float = 0.95)

Enhanced uncertainty quantification interface with integration capabilities.

decompose_uncertainty

decompose_uncertainty(predictions: Float[Array, 'samples batch output'], aleatoric_variance: Float[Array, 'samples batch output'] | None = None) -> UncertaintyComponents

Decompose total uncertainty into epistemic and aleatoric components.

enhanced_uncertainty_decomposition

enhanced_uncertainty_decomposition(predictions: Float[Array, 'samples batch output'], true_values: Float[Array, 'batch output'] | None = None, inputs: Float[Array, 'batch input_dim'] | None = None) -> UncertaintyComponents

Enhanced uncertainty decomposition with additional context.

compute_confidence_intervals

compute_confidence_intervals(predictions: Float[Array, 'samples batch output'], confidence_level: float | None = None) -> tuple[Float[Array, 'batch output'], Float[Array, 'batch output']]

Compute confidence intervals from prediction samples.

compute_prediction_intervals

compute_prediction_intervals(mean_predictions: Float[Array, 'batch output'], total_variance: Float[Array, 'batch output'], confidence_level: float | None = None) -> tuple[Float[Array, 'batch output'], Float[Array, 'batch output']]

Compute prediction intervals using Gaussian assumption.

propagate_uncertainty

propagate_uncertainty(predictions: Float[Array, 'samples batch output'], inputs: Float[Array, 'batch input_dim'], true_values: Float[Array, 'batch output'] | None = None) -> UncertaintyIntegrationResults

Propagate uncertainty through the entire prediction pipeline.

CalibrationAssessment

Enhanced uncertainty calibration assessment tools.

expected_calibration_error staticmethod

expected_calibration_error(confidences: Float[Array, 'n_samples'], accuracies: Float[Array, 'n_samples'], n_bins: int = 10) -> float

Compute Expected Calibration Error (ECE).

maximum_calibration_error staticmethod

maximum_calibration_error(confidences: Float[Array, 'n_samples'], accuracies: Float[Array, 'n_samples'], n_bins: int = 10) -> float

Compute Maximum Calibration Error (MCE).

reliability_diagram_data staticmethod

reliability_diagram_data(confidences: Float[Array, 'n_samples'], accuracies: Float[Array, 'n_samples'], n_bins: int = 10) -> dict[str, Array]

Compute reliability diagram data for visualization.

assess_calibration

assess_calibration(confidences: Float[Array, 'n_samples'], accuracies: Float[Array, 'n_samples'], n_bins: int = 10) -> CalibrationMetrics

Assess overall calibration with multiple metrics.

DistributionalAleatoricUncertainty

Distributional modeling of aleatoric uncertainty.

sample_gaussian

sample_gaussian(mean: Float[Array, 'batch output'], log_std: Float[Array, 'batch output'], num_samples: int, *, rngs: Rngs) -> Float[Array, 'samples batch output']

Sample from Gaussian distributional output.

Parameters:

Name	Type	Description	Default
mean	Float[Array, 'batch output']	Mean predictions	required
log_std	Float[Array, 'batch output']	Log standard deviation predictions	required
num_samples	int	Number of samples to draw	required
rngs	Rngs	Caller-owned RNG bundle; advances its sample stream once per call to produce reproducible-given-seed Monte-Carlo draws.	required

Returns:

Type	Description
Float[Array, 'samples batch output']	Samples from the distributional output

compute_gaussian_uncertainty

compute_gaussian_uncertainty(mean: Float[Array, 'batch output'], log_std: Float[Array, 'batch output']) -> Float[Array, 'batch output']

Compute uncertainty from Gaussian distributional parameters.

Parameters:

Name	Type	Description	Default
mean	Float[Array, 'batch output']	Mean predictions	required
log_std	Float[Array, 'batch output']	Log standard deviation predictions	required

Returns:

Type	Description
Float[Array, 'batch output']	Aleatoric uncertainty (variance)

compute_laplace_uncertainty

compute_laplace_uncertainty(scale: Float[Array, 'batch output']) -> Float[Array, 'batch output']

Std-equivalent uncertainty from a Laplace(0, scale) likelihood.

A Laplace distribution with scale parameter b has variance 2 * b**2 and std b * sqrt(2). Returning the std lets the Laplace branch line up with :meth:compute_gaussian_uncertainty for callers that switch on the noise model.

compute_mixture_uncertainty

compute_mixture_uncertainty(mixture_weights: Float[Array, 'batch components'], means: Float[Array, 'batch components output'], log_stds: Float[Array, 'batch components output']) -> Float[Array, 'batch output']

Compute uncertainty from mixture of Gaussians.

Parameters:

Name	Type	Description	Default
mixture_weights	Float[Array, 'batch components']	Mixture component weights	required
means	Float[Array, 'batch components output']	Component means	required
log_stds	Float[Array, 'batch components output']	Component log standard deviations	required

Returns:

Type	Description
Float[Array, 'batch output']	Total uncertainty from mixture model

EnhancedUncertaintyQuantifier

EnhancedUncertaintyQuantifier(ensemble_size: int = 5, distributional_output: bool = True, multi_source_aggregation: bool = True, confidence_level: float = 0.95)

Enhanced uncertainty quantifier with multiple decomposition methods.

Parameters:

Name	Type	Description	Default
ensemble_size	int	Number of models in ensemble	5
distributional_output	bool	Whether to use distributional outputs	True
multi_source_aggregation	bool	Whether to aggregate multiple uncertainty sources	True
confidence_level	float	Confidence level for intervals	0.95

enhanced_decompose_uncertainty

enhanced_decompose_uncertainty(ensemble_predictions: Float[Array, 'models batch output'], distributional_std: Float[Array, 'batch output'] | None = None, inputs: Float[Array, 'batch input_dim'] | None = None, dropout_predictions: Float[Array, 'samples batch output'] | None = None) -> EnhancedUncertaintyComponents

Enhanced uncertainty decomposition with multiple sources.

Parameters:

Name	Type	Description	Default
ensemble_predictions	Float[Array, 'models batch output']	Predictions from ensemble models	required
distributional_std	Float[Array, 'batch output'] | None	Standard deviation from distributional output	None
inputs	Float[Array, 'batch input_dim'] | None	Input data for context-dependent uncertainty	None
dropout_predictions	Float[Array, 'samples batch output'] | None	Predictions with dropout for additional epistemic uncertainty	None

Returns:

Type	Description
EnhancedUncertaintyComponents	Enhanced uncertainty components with detailed breakdown

EnsembleEpistemicUncertainty

EnsembleEpistemicUncertainty(num_models: int)

Ensemble-based epistemic uncertainty estimation.

Parameters:

Name	Type	Description	Default
num_models	int	Number of models in the ensemble	required

add_model

add_model(model: Any) -> None

Add a model to the ensemble.

Parameters:

Name	Type	Description	Default
model	Any	Neural network model to add to ensemble	required

aggregate_predictions

aggregate_predictions(ensemble_predictions: Float[Array, 'models batch output'], method: str = 'mean') -> Float[Array, 'batch output']

Aggregate predictions from ensemble models.

Parameters:

Name	Type	Description	Default
ensemble_predictions	Float[Array, 'models batch output']	Predictions from all ensemble models	required
method	str	Aggregation method ("mean", "median", "weighted_mean")	'mean'

Returns:

Type	Description
Float[Array, 'batch output']	Aggregated predictions

compute_epistemic_uncertainty

compute_epistemic_uncertainty(ensemble_predictions: Float[Array, 'models batch output']) -> Float[Array, 'batch output']

Compute epistemic uncertainty from ensemble predictions.

Parameters:

Name	Type	Description	Default
ensemble_predictions	Float[Array, 'models batch output']	Predictions from all ensemble models	required

Returns:

Type	Description
Float[Array, 'batch output']	Epistemic uncertainty (variance across models)

compute_prediction_disagreement

compute_prediction_disagreement(ensemble_predictions: Float[Array, 'models batch output']) -> Float[Array, 'batch output']

Compute prediction disagreement metric.

Parameters:

Name	Type	Description	Default
ensemble_predictions	Float[Array, 'models batch output']	Predictions from all ensemble models	required

Returns:

Type	Description
Float[Array, 'batch output']	Disagreement metric (pairwise prediction variance)

MultiSourceUncertaintyAggregator

Aggregation of uncertainty from multiple sources.

aggregate_uncertainties

aggregate_uncertainties(epistemic_sources: list[Float[Array, 'batch output']], aleatoric_sources: list[Float[Array, 'batch output']], method: str = 'variance_sum', epistemic_weights: Array | None = None, aleatoric_weights: Array | None = None) -> Float[Array, 'batch output']

Aggregate uncertainties from multiple sources.

Parameters:

Name	Type	Description	Default
epistemic_sources	list[Float[Array, 'batch output']]	List of epistemic uncertainty estimates	required
aleatoric_sources	list[Float[Array, 'batch output']]	List of aleatoric uncertainty estimates	required
method	str	Aggregation method ("variance_sum", "weighted_sum", "max")	'variance_sum'
epistemic_weights	Array | None	Weights for epistemic sources	None
aleatoric_weights	Array | None	Weights for aleatoric sources	None

Returns:

Type	Description
Float[Array, 'batch output']	Total aggregated uncertainty

compute_uncertainty_breakdown

compute_uncertainty_breakdown(epistemic_sources: list[Float[Array, 'batch output']], aleatoric_sources: list[Float[Array, 'batch output']], source_names: list[str] | None = None) -> dict[str, Float[Array, 'batch output']]

Compute detailed uncertainty breakdown by source.

Parameters:

Name	Type	Description	Default
epistemic_sources	list[Float[Array, 'batch output']]	List of epistemic uncertainty estimates	required
aleatoric_sources	list[Float[Array, 'batch output']]	List of aleatoric uncertainty estimates	required
source_names	list[str] | None	Names for uncertainty sources	None

Returns:

Type	Description
dict[str, Float[Array, 'batch output']]	Dictionary mapping source names to uncertainty values

adaptive_weighting staticmethod

adaptive_weighting(uncertainty_sources: list[Float[Array, 'batch output']], reliability_scores: list[Float[Array, 'batch']] | None = None, adaptation_method: str = 'reliability_based') -> Float[Array, 'sources batch']

Compute per-source, per-sample weights for uncertainty fusion.

adaptation_method:

• reliability_based — weights ∝ per-source reliability scores supplied by the caller (must be non-None).
• inverse_variance — weights ∝ 1 / Var(source).
• entropy_based — weights ∝ predictive entropy of each source.
• uniform — flat weights, primarily a baseline.

assess_uncertainty_quality staticmethod

assess_uncertainty_quality(predictions: Float[Array, 'batch output'], uncertainties: Float[Array, 'batch output'], true_values: Float[Array, 'batch output'] | None = None) -> dict[str, float]

Diagnostic summary of an uncertainty estimate.

Always returns the source-side statistics (mean_uncertainty, uncertainty_std, uncertainty_range, mean_confidence). If true_values is supplied, also returns coverage_probability (under a 2σ band), mean_interval_width, and calibration_error (mean |y - μ| / σ).

CalibrationMetrics dataclass

CalibrationMetrics(expected_calibration_error: float, maximum_calibration_error: float, reliability_diagram: dict[str, Array], confidence_histogram: Array, accuracy_histogram: Array)

Uncertainty calibration assessment metrics.

EnhancedUncertaintyComponents dataclass

EnhancedUncertaintyComponents(epistemic_ensemble: Float[Array, 'batch output'], aleatoric_distributional: Float[Array, 'batch output'], total_uncertainty: Float[Array, 'batch output'], uncertainty_breakdown: dict[str, Float[Array, 'batch output']], epistemic_dropout: Float[Array, 'batch output'] | None = None)

Enhanced uncertainty components with multiple sources.

UncertaintyComponents dataclass

UncertaintyComponents(epistemic: Float[Array, ...], aleatoric: Float[Array, ...], total: Float[Array, ...])

Decomposed uncertainty components.

UncertaintyIntegrationResults dataclass

UncertaintyIntegrationResults(predictions: Float[Array, 'batch output'], uncertainty_components: UncertaintyComponents, calibration_metrics: CalibrationMetrics, confidence_intervals: tuple[Float[Array, 'batch output'], Float[Array, 'batch output']], prediction_intervals: tuple[Float[Array, 'batch output'], Float[Array, 'batch output']])

Results from uncertainty propagation through model pipeline.

Reliability-engineering utilities (Task 6.5).

Two primitives that operate on arrays — never on NNX modules — so they compose cleanly with sampling utilities and pre-tabulated indicator sequences:

• :func:failure_probability — empirical p_f + Wilson binomial confidence interval (Wilson 1927).
• :func:reliability_index — Cornell / Hasofer-Lind index beta = -Phi^{-1}(p_f) (Hasofer-Lind 1974).

Advanced FORM / SORM / subset-simulation adapters are out of scope for this slice per the Task 6.5 implementation requirements.

ReliabilityResult dataclass

ReliabilityResult(*, probability: Array, confidence_interval: tuple[Array, Array], confidence_level: float, num_samples: int, num_failures: Array)

Container for empirical failure probability + binomial CI.

Attributes:

Name	Type	Description
probability	Array	Empirical failure probability p_f = sum / N.
confidence_interval	tuple[Array, Array]	(lower, upper) Wilson binomial CI at the requested confidence level.
confidence_level	float	Confidence level of the interval (e.g. 0.95).
num_samples	int	Total sample count N.
num_failures	Array	Sum of failure indicators.

Surrogate-model uncertainty primitives (Task 6.6).

Currently a single helper:

• :func:decompose_surrogate_uncertainty — combine independent prediction / residual / calibration variance components into a single total uncertainty.

The PCE primitives live under :mod:opifex.uncertainty.scientific.polynomial_chaos (Task 6.6 explicitly forbids an intermediate surrogate/pce.py).

SurrogateUncertaintyResult dataclass

SurrogateUncertaintyResult(*, total_std: Array, prediction_std: Array, residual_std: Array, calibration_std: Array)

Container for the per-source variance decomposition.

Attributes:

Name	Type	Description
total_std	Array	sqrt(sigma_pred^2 + sigma_resid^2 + sigma_cal^2).
prediction_std	Array	Surrogate predictive standard deviation.
residual_std	Array	Held-out residual standard deviation (zero if not supplied).
calibration_std	Array	Post-hoc calibration standard deviation (zero if not supplied).

decompose_surrogate_uncertainty

decompose_surrogate_uncertainty(*, prediction_std: Array, residual_std: Array | None = None, calibration_std: Array | None = None) -> SurrogateUncertaintyResult

Combine independent variance sources into a total-uncertainty estimate.

Parameters:

Name	Type	Description	Default
prediction_std	Array	Predictive standard deviation reported by the surrogate. Shape (N,) or (N, K).	required
residual_std	Array | None	Optional held-out residual standard deviation. Defaults to zero (no residual contribution). Must broadcast to prediction_std if supplied.	None
calibration_std	Array | None	Optional calibration standard deviation (e.g. from temperature scaling cross-validation). Defaults to zero. Must broadcast to prediction_std if supplied.	None

Returns:

Type	Description
SurrogateUncertaintyResult	class:SurrogateUncertaintyResult with each component and the
SurrogateUncertaintyResult	combined total.

Raises:

Type	Description
ValueError	If any standard deviation is negative.