Posterior Covariance¶

A variational analysis returns a point estimate; this layer attaches the uncertainty around it. The design is deliberately lazy and matrix-free: a posterior adapter never materialises the covariance matrix. Instead it builds a lineax operator for the Hessian (or precision) at the analysis point, and inverses, diagonals (pointwise variances), and ensemble estimates are delegated to gaussx — gaussx.inv, gaussx.diag, and gaussx.ensemble_covariance — so structured and iterative solves come for free. See Posterior Covariance in the Mathematical Reference for derivations and when each approximation is trustworthy.

Adapters satisfy the PosteriorAdapter Protocol and all produce the same Posterior container, so downstream diagnostics (e.g. the posterior-agreement gate) are adapter-agnostic. LaplaceCovariance and GaussNewtonHessian both build the inverse Gauss–Newton Hessian at the mode, matrix-free; for a 4DVar model it runs through the model's tangent-linear and adjoint rollouts, so pass the model and the window's batch. GaussNewtonHessian exposes the Krylov budget. Choose EnsembleCovariance when samples are cheaper than curvature or the posterior is not Gaussian.

Container¶

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

Posterior ¶

Bases: Module

Posterior container — Decision D10.

A plain container for \(p(x \mid y) \approx \mathcal{N}(\text{mean}, \text{cov})\), plus optional samples and provenance.

Attributes:

Name	Type	Description
`mean`	`Float[Array, ...]`	Posterior mean / MAP estimate.
`cov`	`AbstractLinearOperator \| None`	Posterior covariance as `lineax.AbstractLinearOperator`, or `None` for amortized / sampling-only heads. Stored lazily — never materialised into a dense matrix unless the user asks.
`samples`	`Float[Array, ...] \| None`	Optional posterior samples (from a flow head, ensemble, or score-based diffusion). `None` for analytical adapters.
`provenance`	`dict[str, Any]`	Free-form dict carrying audit info — at minimum `{forward_model_id, n_iter, J_star, converged, vardax_version}`. Consumed by `GaussianMarkLikelihood` when serialising for downstream population models.

Examples:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> post = vdx.Posterior(mean=jnp.zeros(4))
>>> post.mean.shape
(4,)
>>> post.cov is None
True
>>> post.provenance
{}

A Gaussian posterior carries a lazy covariance operator and audit metadata:

>>> import jax, lineax as lx
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((2,), jnp.float32))
>>> gauss = vdx.Posterior(
...     mean=jnp.zeros(2), cov=eye, provenance={"adapter": "demo"}
... )
>>> gauss.cov.mv(jnp.ones(2)).shape, gauss.provenance["adapter"]
((2,), 'demo')

Source code in src/vardax/_src/posterior/container.py

class Posterior(eqx.Module):
    r"""Posterior container — Decision D10.

    A plain container for $p(x \mid y) \approx \mathcal{N}(\text{mean}, \text{cov})$,
    plus optional samples and provenance.

    Attributes:
        mean: Posterior mean / MAP estimate.
        cov: Posterior covariance as ``lineax.AbstractLinearOperator``,
            or ``None`` for amortized / sampling-only heads. Stored
            lazily — never materialised into a dense matrix unless the
            user asks.
        samples: Optional posterior samples (from a flow head, ensemble,
            or score-based diffusion). ``None`` for analytical adapters.
        provenance: Free-form dict carrying audit info — at minimum
            ``{forward_model_id, n_iter, J_star, converged,
            vardax_version}``. Consumed by
            [`GaussianMarkLikelihood`][vardax.GaussianMarkLikelihood]
            when serialising for downstream population models.

    Examples:
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> post = vdx.Posterior(mean=jnp.zeros(4))
        >>> post.mean.shape
        (4,)
        >>> post.cov is None
        True
        >>> post.provenance
        {}

        A Gaussian posterior carries a lazy covariance operator and audit
        metadata:

        >>> import jax, lineax as lx
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((2,), jnp.float32))
        >>> gauss = vdx.Posterior(
        ...     mean=jnp.zeros(2), cov=eye, provenance={"adapter": "demo"}
        ... )
        >>> gauss.cov.mv(jnp.ones(2)).shape, gauss.provenance["adapter"]
        ((2,), 'demo')
    """

    mean: Float[Array, ...]
    cov: lx.AbstractLinearOperator | None = None
    samples: Float[Array, ...] | None = None
    provenance: dict[str, Any] = eqx.field(default_factory=dict)

Adapters¶

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

LaplaceCovariance ¶

Bases: Module

Laplace approximation at MAP.

\[P^* \approx \big(B^{-1} + H'^\top R^{-1} H'\big)^{-1}\]

Algorithm

Linearise the observation operator (and, for 4DVar, the model rollout) at the analysis.
Assemble the Gauss-Newton precision \(B^{-1} + \sum_t G_t^\top R_t^{-1} G_t\) matrix-free, with \(G_t = H'_t M'_{t \leftarrow 0}\).
Return its inverse as a lazy operator applied by solver.

Attributes:

Name	Type	Description
`prior_cov_op`	`AbstractLinearOperator \| None`	Background-error covariance \(B\). Default `None`: the model's `prior_cov_op`. Unused for a linear control, whose prior is the basis's \(Q\).
`obs_cov_op`	`AbstractLinearOperator \| None`	Observation-error covariance \(R\) at one time. Default `None`: the model's `obs_cov_op`.
`solver`	`AbstractLinearSolver`	Linear solver for the lazy inverses (default CG, `atol = rtol = 1e-6`, 200 steps).

Pass \(B\) and \(R\) explicitly only to evaluate the posterior under different error statistics from the ones the analysis used.

Examples:

>>> import jax, jax.numpy as jnp, lineax as lx
>>> import vardax as vdx
>>> from types import SimpleNamespace
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((4,), jnp.float32))
>>> laplace = vdx.LaplaceCovariance(prior_cov_op=eye, obs_cov_op=eye)
>>> model = SimpleNamespace(obs_op=vdx.LinearObs(operator=eye))
>>> post = laplace(jnp.zeros(4), model, batch=None)
>>> post.mean.shape
(4,)
>>> # B = R = H = I, so P* = (I + I)^{-1} = I / 2
>>> bool(jnp.allclose(post.cov.mv(jnp.ones(4)), 0.5, atol=1e-2))
True

Weaker observations (R = 3 I) leave a larger posterior, \(P^* = (1 + 1/3)^{-1} = 3/4\):

>>> weak = jnp.float32(3.0) * eye
>>> post = vdx.LaplaceCovariance(prior_cov_op=eye, obs_cov_op=weak)(
...     jnp.zeros(4), model, batch=None
... )
>>> bool(jnp.allclose(post.cov.mv(jnp.ones(4)), 0.75, atol=1e-2))
True

A model with no obs_op is rejected:

>>> laplace(
...     jnp.zeros(4), SimpleNamespace(), batch=None
... )
Traceback (most recent call last):
...
AttributeError: ...

Source code in src/vardax/_src/posterior/laplace.py

class LaplaceCovariance(eqx.Module):
    r"""Laplace approximation at MAP.

    $$P^* \approx \big(B^{-1} + H'^\top R^{-1} H'\big)^{-1}$$

    Algorithm:
        1. Linearise the observation operator (and, for 4DVar, the model
           rollout) at the analysis.
        2. Assemble the Gauss-Newton precision
           $B^{-1} + \sum_t G_t^\top R_t^{-1} G_t$
           matrix-free, with $G_t = H'_t M'_{t \leftarrow 0}$.
        3. Return its inverse as a lazy operator applied by ``solver``.

    Attributes:
        prior_cov_op: Background-error covariance $B$. Default ``None``:
            the model's ``prior_cov_op``. Unused for a linear control,
            whose prior is the basis's $Q$.
        obs_cov_op: Observation-error covariance $R$ at one time. Default
            ``None``: the model's ``obs_cov_op``.
        solver: Linear solver for the lazy inverses (default CG,
            ``atol = rtol = 1e-6``, 200 steps).

    Pass $B$ and $R$ explicitly only to evaluate the posterior under
    different error statistics from the ones the analysis used.

    Examples:
        >>> import jax, jax.numpy as jnp, lineax as lx
        >>> import vardax as vdx
        >>> from types import SimpleNamespace
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((4,), jnp.float32))
        >>> laplace = vdx.LaplaceCovariance(prior_cov_op=eye, obs_cov_op=eye)
        >>> model = SimpleNamespace(obs_op=vdx.LinearObs(operator=eye))
        >>> post = laplace(jnp.zeros(4), model, batch=None)
        >>> post.mean.shape
        (4,)
        >>> # B = R = H = I, so P* = (I + I)^{-1} = I / 2
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones(4)), 0.5, atol=1e-2))
        True

        Weaker observations (R = 3 I) leave a larger posterior,
        $P^* = (1 + 1/3)^{-1} = 3/4$:

        >>> weak = jnp.float32(3.0) * eye
        >>> post = vdx.LaplaceCovariance(prior_cov_op=eye, obs_cov_op=weak)(
        ...     jnp.zeros(4), model, batch=None
        ... )
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones(4)), 0.75, atol=1e-2))
        True

        A model with no ``obs_op`` is rejected:

        >>> laplace(
        ...     jnp.zeros(4), SimpleNamespace(), batch=None
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        AttributeError: ...
    """

    prior_cov_op: lx.AbstractLinearOperator | None = None
    obs_cov_op: lx.AbstractLinearOperator | None = None
    solver: lx.AbstractLinearSolver = eqx.field(static=True, default=DEFAULT_SOLVER)

    def __call__(
        self,
        analysis: Float[Array, ...] | tuple[Array, Array],
        model: Any,  # AnalysisStep-compliant
        batch: Any,
    ) -> Posterior:
        r"""Build the Laplace posterior at ``analysis``.

        Args:
            analysis: MAP of one window: $x_0^*$, or $(x_0^*, \eta^*)$ for
                ``WeakFourDVar``.
            model: The analysis model or its ``.as_analysis_step()``
                wrapper. Its ``obs_op`` (and, for 4DVar, its rollout) build
                the precision.
            batch: The single window the analysis was computed from; 4DVar
                models read its length and mask. A batch with a leading
                sample axis must hold one sample.

        Returns:
            [`Posterior`][vardax.Posterior] with ``mean`` $= x_0^*$ and
            ``cov`` a lazy ``AbstractLinearOperator`` representing $P^*$.

        Raises:
            AttributeError: If the model exposes no ``obs_op``.
            UnsupportedInputError: For a point batch, a non-``Batch1D``
                batch, a ``basis=`` or ``parameterization=`` control, or
                a ``WeakFourDVar`` analysis that is not ``(x0, eta)``.
            ShapeError: If the batch holds several windows.
            MissingInputError: If neither the adapter nor the model has
                a covariance operator.
        """
        mean, covariance, structure = gauss_newton_covariance(
            analysis,
            model,
            batch,
            self.prior_cov_op,
            self.obs_cov_op,
            self.solver,
        )
        return Posterior(
            mean=mean,
            cov=covariance,
            samples=None,
            provenance={"adapter": "LaplaceCovariance", "structure": structure},
        )

GaussNewtonHessian ¶

Bases: Module

Gauss-Newton Hessian inversion at MAP.

Computes \((B^{-1} + \sum_t (H'_t M'_t)^\top R_t^{-1} (H'_t M'_t))^{-1}\) for 4DVar models and \((B^{-1} + H'^\top R^{-1} H')^{-1}\) for single-time ones, like LaplaceCovariance, with CG capped at n_krylov iterations per covariance mat-vec.

Algorithm

Linearise the observation operator (and, for 4DVar, the model rollout) at the analysis.
Assemble the precision \(B^{-1} + \sum_t G_t^\top R_t^{-1} G_t\) matrix-free, with \(G_t = H'_t M'_{t \leftarrow 0}\).
Return its inverse as a lazy operator applied by CG with at most n_krylov steps.

Attributes:

Name	Type	Description
`prior_cov_op`	`AbstractLinearOperator \| None`	\(B\). Default `None`: the model's `prior_cov_op`. Unused for a linear control, whose prior is the basis's \(Q\).
`obs_cov_op`	`AbstractLinearOperator \| None`	\(R\) at one time. Default `None`: the model's `obs_cov_op`.
`n_krylov`	`int`	Maximum CG iterations per covariance mat-vec.
`tol`	`float`	CG absolute and relative tolerance.

Examples:

Linear dynamics \(x_{t+1} = a x_t\), a fully observed two-state window, \(B = R = I\): the observations at \(t = 0, 1\) both inform \(x_0\), so \(P^* = 1 / (1 + 1 + a^2)\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> class Decay:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return 0.5 * x
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((2,), jnp.float32))
>>> model = vardax.StrongFourDVar(
...     forward=Decay(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(2),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 2, 2)), mask=jnp.ones((1, 2, 2)))
>>> adapter = vardax.GaussNewtonHessian(prior_cov_op=eye, obs_cov_op=eye)
>>> post = adapter(model(batch)[0], model, batch)
>>> bool(jnp.allclose(post.cov.mv(jnp.ones(2)), 1 / 2.25, atol=1e-5))
True

The Krylov budget is recorded in the provenance:

>>> post.provenance["n_krylov"]
50

A point batch is not supported yet:

>>> pts = vardax.PointObsBatch(
...     coords=jnp.zeros((1, 1)),
...     values=jnp.zeros(1),
...     sigma=jnp.ones(1),
...     valid=jnp.ones(1, bool),
... )
>>> adapter(model(batch)[0], model, pts)
Traceback (most recent call last):
...
UnsupportedInputError: The posterior adapters do not support ...

Source code in src/vardax/_src/posterior/gauss_newton.py

class GaussNewtonHessian(eqx.Module):
    r"""Gauss-Newton Hessian inversion at MAP.

    Computes $(B^{-1} + \sum_t (H'_t M'_t)^\top R_t^{-1} (H'_t M'_t))^{-1}$
    for 4DVar models and $(B^{-1} + H'^\top R^{-1} H')^{-1}$ for
    single-time ones, like
    [`LaplaceCovariance`][vardax.LaplaceCovariance], with CG capped at
    ``n_krylov`` iterations per covariance mat-vec.

    Algorithm:
        1. Linearise the observation operator (and, for 4DVar, the model
           rollout) at the analysis.
        2. Assemble the precision $B^{-1} + \sum_t G_t^\top R_t^{-1} G_t$
           matrix-free, with $G_t = H'_t M'_{t \leftarrow 0}$.
        3. Return its inverse as a lazy operator applied by CG with at
           most ``n_krylov`` steps.

    Attributes:
        prior_cov_op: $B$. Default ``None``: the model's ``prior_cov_op``.
            Unused for a linear control, whose prior is the basis's $Q$.
        obs_cov_op: $R$ at one time. Default ``None``: the model's
            ``obs_cov_op``.
        n_krylov: Maximum CG iterations per covariance mat-vec.
        tol: CG absolute and relative tolerance.

    Examples:
        Linear dynamics $x_{t+1} = a x_t$, a fully observed two-state
        window, $B = R = I$: the observations at $t = 0, 1$ both inform
        $x_0$, so $P^* = 1 / (1 + 1 + a^2)$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> class Decay:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return 0.5 * x
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((2,), jnp.float32))
        >>> model = vardax.StrongFourDVar(
        ...     forward=Decay(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(2),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 2, 2)), mask=jnp.ones((1, 2, 2)))
        >>> adapter = vardax.GaussNewtonHessian(prior_cov_op=eye, obs_cov_op=eye)
        >>> post = adapter(model(batch)[0], model, batch)
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones(2)), 1 / 2.25, atol=1e-5))
        True

        The Krylov budget is recorded in the provenance:

        >>> post.provenance["n_krylov"]
        50

        A point batch is not supported yet:

        >>> pts = vardax.PointObsBatch(
        ...     coords=jnp.zeros((1, 1)),
        ...     values=jnp.zeros(1),
        ...     sigma=jnp.ones(1),
        ...     valid=jnp.ones(1, bool),
        ... )
        >>> adapter(model(batch)[0], model, pts)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        UnsupportedInputError: The posterior adapters do not support ...
    """

    prior_cov_op: lx.AbstractLinearOperator | None = None
    obs_cov_op: lx.AbstractLinearOperator | None = None
    n_krylov: int = eqx.field(static=True, default=50)
    tol: float = eqx.field(static=True, default=1e-6)

    def __call__(
        self,
        analysis: Float[Array, ...] | tuple[Array, Array],
        model: Any,
        batch: Any,
    ) -> Posterior:
        r"""Build the Gauss–Newton posterior at ``analysis``.

        Args:
            analysis: MAP of one window: $x_0^*$, or $(x_0^*, \eta^*)$ for
                ``WeakFourDVar``.
            model: The analysis model or its ``.as_analysis_step()`` wrapper.
            batch: The single window the analysis was computed from.

        Returns:
            [`Posterior`][vardax.Posterior] with a lazy covariance.

        Raises:
            AttributeError: If the model exposes no ``obs_op``.
            UnsupportedInputError: For a point batch, a non-``Batch1D``
                batch, a ``basis=`` or ``parameterization=`` control, or
                a ``WeakFourDVar`` analysis that is not ``(x0, eta)``.
            ShapeError: If the batch holds several windows.
            MissingInputError: If neither the adapter nor the model has
                a covariance operator.
        """
        solver = lx.CG(atol=self.tol, rtol=self.tol, max_steps=self.n_krylov)
        mean, covariance, structure = gauss_newton_covariance(
            analysis,
            model,
            batch,
            self.prior_cov_op,
            self.obs_cov_op,
            solver,
        )
        return Posterior(
            mean=mean,
            cov=covariance,
            samples=None,
            provenance={
                "adapter": "GaussNewtonHessian",
                "n_krylov": self.n_krylov,
                "structure": structure,
            },
        )

EnsembleCovariance ¶

Bases: Module

Sample-covariance posterior from an ensemble of analyses.

\[\bar x = \frac{1}{M} \sum_m x^{(m)}, \quad P^* = \lambda \, \frac{1}{M-1} \sum_m (x^{(m)} - \bar x)(x^{(m)} - \bar x)^\top\]

Algorithm

Average the analyses over the ensemble axis to get \(\bar x\).
Flatten each member and form the Bessel-corrected sample covariance as a low-rank operator (never dense).
Scale it by the inflation factor \(\lambda\).
Package mean, covariance, the members as samples and the provenance into a Posterior.

Attributes:

Name	Type	Description
`n_members`	`int`	Expected ensemble size (used for diagnostics).
`inflation`	`float`	Multiplicative inflation factor \(\lambda\) applied to the sample covariance. Default 1.0 (no inflation).

Examples:

Three scalar-pair analyses: the mean is the member average.

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> analyses = jnp.array([[0.0, 1.0], [2.0, 3.0], [4.0, 8.0]])
>>> post = vdx.EnsembleCovariance(n_members=3)(analyses, None, None)
>>> bool(jnp.allclose(post.mean, jnp.array([2.0, 4.0])))
True
>>> post.samples.shape
(3, 2)

Inflation scales the covariance (here by 4):

>>> big = vdx.EnsembleCovariance(n_members=3, inflation=4.0)(
...     analyses, None, None
... )
>>> v = jnp.array([1.0, 1.0])
>>> bool(jnp.allclose(big.cov.mv(v), 4.0 * post.cov.mv(v)))
True

Members of any shape are flattened, so a (M, 2, 2) stack gives a 4 x 4 covariance:

>>> grid = jnp.arange(12.0).reshape(3, 2, 2) ** 2
>>> vdx.EnsembleCovariance(n_members=3)(grid, None, None).cov.mv(
...     jnp.ones(4)
... ).shape
(4,)

Source code in src/vardax/_src/posterior/ensemble.py

class EnsembleCovariance(eqx.Module):
    r"""Sample-covariance posterior from an ensemble of analyses.

    $$\bar x = \frac{1}{M} \sum_m x^{(m)}, \quad
    P^* = \lambda \, \frac{1}{M-1} \sum_m (x^{(m)} - \bar x)(x^{(m)} - \bar x)^\top$$

    Algorithm:
        1. Average the analyses over the ensemble axis to get $\bar x$.
        2. Flatten each member and form the Bessel-corrected sample
           covariance as a low-rank operator (never dense).
        3. Scale it by the inflation factor $\lambda$.
        4. Package mean, covariance, the members as samples and the
           provenance into a [`Posterior`][vardax.Posterior].

    Attributes:
        n_members: Expected ensemble size (used for diagnostics).
        inflation: Multiplicative inflation factor $\lambda$ applied
            to the sample covariance. Default 1.0 (no inflation).

    Examples:
        Three scalar-pair analyses: the mean is the member average.

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> analyses = jnp.array([[0.0, 1.0], [2.0, 3.0], [4.0, 8.0]])
        >>> post = vdx.EnsembleCovariance(n_members=3)(analyses, None, None)
        >>> bool(jnp.allclose(post.mean, jnp.array([2.0, 4.0])))
        True
        >>> post.samples.shape
        (3, 2)

        Inflation scales the covariance (here by 4):

        >>> big = vdx.EnsembleCovariance(n_members=3, inflation=4.0)(
        ...     analyses, None, None
        ... )
        >>> v = jnp.array([1.0, 1.0])
        >>> bool(jnp.allclose(big.cov.mv(v), 4.0 * post.cov.mv(v)))
        True

        Members of any shape are flattened, so a ``(M, 2, 2)`` stack gives a
        4 x 4 covariance:

        >>> grid = jnp.arange(12.0).reshape(3, 2, 2) ** 2
        >>> vdx.EnsembleCovariance(n_members=3)(grid, None, None).cov.mv(
        ...     jnp.ones(4)
        ... ).shape
        (4,)
    """

    n_members: int = eqx.field(static=True)
    inflation: float = eqx.field(static=True, default=1.0)

    def __call__(
        self,
        analyses: Float[Array, "M ..."],
        model: Any,
        batch: Any,
    ) -> Posterior:
        """Build the ensemble posterior from per-member analyses.

        Args:
            analyses: Stack of the ``M`` member analyses, member axis
                first.
            model: AnalysisStep instance (not used directly; kept for
                interface compatibility).
            batch: The batch (not used directly; kept for interface
                compatibility).

        Returns:
            [`Posterior`][vardax.Posterior] with mean = ensemble mean,
            ``cov`` = sample covariance as a dense
            ``MatrixLinearOperator`` (only assembled if you ask for
            it; safe for moderate M).
        """
        n_analyses = analyses.shape[0]
        if n_analyses != self.n_members:
            # Allow runtime mismatch; the configured ``n_members`` is
            # advisory rather than enforced.
            pass

        # x̄ = (1/m) Σᵢ xᵢ
        mean = einx.mean("[m] ...", analyses)
        # Flatten everything but the leading ensemble dim; the
        # Bessel-corrected sample covariance comes from gaussx as a
        # rank-(m-1) LowRankUpdate operator (never dense), scaled by
        # the multiplicative inflation factor.
        flat_analyses = einx.id("m s... -> m (s...)", analyses)
        # Σ = λ (1/(m−1)) Σᵢ (xᵢ − x̄)(xᵢ − x̄)ᵀ
        cov_op = self.inflation * gaussx.ensemble_covariance(flat_analyses, bessel=True)

        return Posterior(
            mean=mean,
            cov=cov_op,
            samples=analyses,
            provenance={
                "adapter": "EnsembleCovariance",
                "n_members": int(n_analyses),
                "inflation": float(self.inflation),
            },
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "n_members": self.n_members,
            "inflation": self.inflation,
        }

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> EnsembleCovariance:
        return cls(**config)

Likelihoods¶

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

GaussianMarkLikelihood ¶

Bases: Module

Serialise a Gaussian posterior into a mark-likelihood dict.

A JSON-friendly summary of a Posterior: its mean, the diagonal of its covariance, any samples, the provenance and the event metadata.

Attributes:

Name	Type	Description
`posterior`	`Posterior`	A vardax `Posterior` (with Gaussian `cov`).
`event_metadata`	`dict[str, Any]`	Free-form dict (event_id, time, geometry, instruments_used, …). Passed through verbatim into the output.

Use .to_dict() to get a JSON-friendly representation that downstream consumers (GeoCatalog, hierarchical Bayesian models, DuckDB ingest) can store as-is.

Examples:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> post = vdx.Posterior(mean=jnp.zeros(3))
>>> mark = vdx.GaussianMarkLikelihood(
...     posterior=post, event_metadata={"event_id": "e1"}
... )
>>> sorted(mark.to_dict())
['cov_diag', 'event_metadata', 'mean', 'provenance', 'samples']

Without a covariance the diagonal is None:

>>> mark.to_dict()["cov_diag"] is None
True

Composed with an ensemble posterior, the diagonal and samples are serialised as nested lists:

>>> analyses = jnp.array([[0.0, 1.0], [2.0, 3.0], [4.0, 8.0]])
>>> ens = vdx.EnsembleCovariance(n_members=3)(analyses, None, None)
>>> out = vdx.GaussianMarkLikelihood(posterior=ens).to_dict()
>>> len(out["cov_diag"]), len(out["samples"])
(2, 3)

Source code in src/vardax/_src/posterior/adapter.py

class GaussianMarkLikelihood(eqx.Module):
    """Serialise a Gaussian posterior into a mark-likelihood dict.

    A JSON-friendly summary of a [`Posterior`][vardax.Posterior]: its mean,
    the diagonal of its covariance, any samples, the provenance and the
    event metadata.

    Attributes:
        posterior: A vardax ``Posterior`` (with Gaussian ``cov``).
        event_metadata: Free-form dict (event_id, time, geometry,
            instruments_used, …). Passed through verbatim into the
            output.

    Use ``.to_dict()`` to get a JSON-friendly representation that
    downstream consumers (GeoCatalog, hierarchical Bayesian models,
    DuckDB ingest) can store as-is.

    Examples:
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> post = vdx.Posterior(mean=jnp.zeros(3))
        >>> mark = vdx.GaussianMarkLikelihood(
        ...     posterior=post, event_metadata={"event_id": "e1"}
        ... )
        >>> sorted(mark.to_dict())
        ['cov_diag', 'event_metadata', 'mean', 'provenance', 'samples']

        Without a covariance the diagonal is ``None``:

        >>> mark.to_dict()["cov_diag"] is None
        True

        Composed with an ensemble posterior, the diagonal and samples are
        serialised as nested lists:

        >>> analyses = jnp.array([[0.0, 1.0], [2.0, 3.0], [4.0, 8.0]])
        >>> ens = vdx.EnsembleCovariance(n_members=3)(analyses, None, None)
        >>> out = vdx.GaussianMarkLikelihood(posterior=ens).to_dict()
        >>> len(out["cov_diag"]), len(out["samples"])
        (2, 3)
    """

    posterior: Posterior
    event_metadata: dict[str, Any] = eqx.field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly mark-likelihood representation.

        Returns:
            A dict with keys ``mean``, ``cov_diag``, ``samples``,
            ``provenance`` and ``event_metadata``; arrays become nested
            lists and missing parts become ``None``.
        """
        out: dict[str, Any] = {
            "mean": _to_list(self.posterior.mean),
            "cov_diag": (
                _to_list(self._cov_diag()) if self.posterior.cov is not None else None
            ),
            "samples": (
                _to_list(self.posterior.samples)
                if self.posterior.samples is not None
                else None
            ),
            "provenance": dict(self.posterior.provenance),
            "event_metadata": dict(self.event_metadata),
        }
        return out

    def _cov_diag(self):
        """Extract the diagonal of the covariance operator.

        For materialised operators this is cheap; for lazy operators
        it costs one mat-vec per state dim (probe with unit vectors).
        We pick the materialised path when the operator exposes it,
        else fall back to a single ``as_matrix()`` call.
        """
        cov = self.posterior.cov
        if cov is None:
            return None
        try:
            return gaussx.diag(cov)
        except (AttributeError, NotImplementedError):
            # Lazy operators don't always materialise; just emit a
            # null marker.
            return None

to_dict ¶

to_dict() -> dict[str, Any]

Return a JSON-friendly mark-likelihood representation.

Returns:

Type	Description
`dict[str, Any]`	A dict with keys `mean`, `cov_diag`, `samples`,
`dict[str, Any]`	`provenance` and `event_metadata`; arrays become nested
`dict[str, Any]`	lists and missing parts become `None`.

Source code in src/vardax/_src/posterior/adapter.py

def to_dict(self) -> dict[str, Any]:
    """Return a JSON-friendly mark-likelihood representation.

    Returns:
        A dict with keys ``mean``, ``cov_diag``, ``samples``,
        ``provenance`` and ``event_metadata``; arrays become nested
        lists and missing parts become ``None``.
    """
    out: dict[str, Any] = {
        "mean": _to_list(self.posterior.mean),
        "cov_diag": (
            _to_list(self._cov_diag()) if self.posterior.cov is not None else None
        ),
        "samples": (
            _to_list(self.posterior.samples)
            if self.posterior.samples is not None
            else None
        ),
        "provenance": dict(self.posterior.provenance),
        "event_metadata": dict(self.event_metadata),
    }
    return out