Skip to content

Observation Operators

Observation operators map model state to observation space. All of them satisfy the pipekit-cycle ObservationOperator Protocol — __call__(state) produces the predicted observation, and linearize(state) returns the tangent-linear (Jacobian) operator that the incremental and posterior-covariance machinery needs. Because conformance is structural, custom operators plug in by matching that signature; nothing here needs to be subclassed. See Observation Operators in the Mathematical Reference for the modelling background.

Core operators

LinearObs wraps an explicit observation matrix; CallableObs turns any differentiable JAX function (a model run, a simulator, an emulator) into an operator, with its tangent-linear map from autodiff; MaskedIdentity handles the ubiquitous "observe a subset of grid points" case (satellite tracks, sparse sensor networks); AveragingKernel implements the smoothing kernels of retrieval products such as atmospheric-composition column averages.

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

LinearObs

Bases: Module

\(H(x) = H x\) for a linear operator \(H\) (operator).

Examples:

>>> import jax.numpy as jnp, lineax as lx
>>> import vardax as vdx
>>> H = lx.MatrixLinearOperator(jnp.eye(3)[:2])
>>> obs = vdx.LinearObs(operator=H)
>>> obs(jnp.array([1.0, 2.0, 3.0]))
Array([1., 2.], dtype=float32)
>>> obs.linearize(jnp.zeros(3)) is H
True

A structured operator is applied without a dense matrix:

>>> diag = vdx.LinearObs(
...     operator=lx.DiagonalLinearOperator(jnp.array([2.0, 3.0]))
... )
>>> diag(jnp.ones(2))
Array([2., 3.], dtype=float32)

The tangent-linear operator of a linear map does not depend on the linearisation point, so the cost of linearising is zero:

>>> diag.linearize(jnp.ones(2)) is diag.linearize(jnp.zeros(2))
True
Source code in src/vardax/_src/obs_operators/linear.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class LinearObs(eqx.Module):
    r"""$H(x) = H x$ for a linear operator $H$ (``operator``).

    Examples:
        >>> import jax.numpy as jnp, lineax as lx
        >>> import vardax as vdx
        >>> H = lx.MatrixLinearOperator(jnp.eye(3)[:2])
        >>> obs = vdx.LinearObs(operator=H)
        >>> obs(jnp.array([1.0, 2.0, 3.0]))
        Array([1., 2.], dtype=float32)
        >>> obs.linearize(jnp.zeros(3)) is H
        True

        A structured operator is applied without a dense matrix:

        >>> diag = vdx.LinearObs(
        ...     operator=lx.DiagonalLinearOperator(jnp.array([2.0, 3.0]))
        ... )
        >>> diag(jnp.ones(2))
        Array([2., 3.], dtype=float32)

        The tangent-linear operator of a linear map does not depend on the
        linearisation point, so the cost of linearising is zero:

        >>> diag.linearize(jnp.ones(2)) is diag.linearize(jnp.zeros(2))
        True
    """

    operator: lx.AbstractLinearOperator

    def __call__(self, x: Float[Array, ...]) -> Float[Array, ...]:
        return self.operator.mv(x)

    def linearize(self, x: Float[Array, ...]) -> lx.AbstractLinearOperator:
        """Tangent-linear operator at ``x``: just ``operator`` (linear).

        Args:
            x: State at which to linearise; unused.

        Returns:
            The stored operator ``operator``.
        """
        return self.operator

linearize 

linearize(x: Float[Array, ...]) -> AbstractLinearOperator

Tangent-linear operator at x: just operator (linear).

Parameters:

Name Type Description Default
x Float[Array, ...]

State at which to linearise; unused.

 required

Returns:

Type Description
AbstractLinearOperator

The stored operator operator.
Source code in src/vardax/_src/obs_operators/linear.py 
51
52
53
54
55
56
57
58
59
60
def linearize(self, x: Float[Array, ...]) -> lx.AbstractLinearOperator:
    """Tangent-linear operator at ``x``: just ``operator`` (linear).

    Args:
        x: State at which to linearise; unused.

    Returns:
        The stored operator ``operator``.
    """
    return self.operator

CallableObs

Bases: Module

\(H(x) = f(x)\) for any differentiable JAX function \(f\).

linearize(x) is a matrix-free lineax.JacobianLinearOperator at x: its mv is a JVP of \(f\) and its transpose().mv a VJP, so nothing of size \(\dim y \times \dim x\) is ever formed. The operator is an Equinox module, so it works under eqx.filter_jit and jax.vmap; when fn is itself an Equinox module (an eqx.nn.MLP emulator, say), its parameters are pytree leaves and eqx.filter_grad differentiates through them.

Prefer LinearObs when you already hold an explicit or structured matrix you want to reuse; use this for a model run, a simulator or an emulator.

Pass jac="bwd" for a model that defines only a reverse-mode rule (a jax.custom_vjp function): forward mode raises on it as soon as an analysis or posterior linearises the operator. Matrix-free, lineax then applies \(H'v\) by transposing the VJP, which fails when the backward pass contains a loop JAX cannot transpose, as a diffrax solve's does (RecursiveCheckpointAdjoint runs a while_loop). For those, also pass materialise=True: linearize then forms the Jacobian once, with one VJP per observation (jax.jacrev), and both \(H'\) and \(H'^\top\) are dense products. That suits a small control, the usual case for source and parameter estimation.

Attributes:

Name Type Description
fn Callable[[Float[Array, ...]], Float[Array, ...]]

The observation map \(f\), from a state of any shape to an observation of any shape.
jac Literal['fwd', 'bwd'] | None

How linearize differentiates fn: None (lineax's default, forward mode), "fwd" or "bwd" (reverse mode).
materialise bool

Form the Jacobian in linearize instead of keeping it matrix-free, with jax.jacrev for jac="bwd" and jax.jacfwd otherwise. Default False.

Examples:

A linear \(f(x) = F x\) on a (1, N) state with \(M \ne N\) observations: OptimalInterpolation with \(B = R = I\) gives the closed-form BLUE, \(x^* = (I + F^\top F)^{-1} F^\top y\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax as vdx
>>> F = jnp.array([[1.0, 0.0], [0.0, 2.0], [1.0, 1.0]])
>>> obs = vdx.CallableObs(fn=lambda x: x @ F.T)
>>> obs(jnp.ones((1, 2))).shape
(1, 3)
>>> eye_x = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 2), jnp.float32))
>>> eye_y = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 3), jnp.float32))
>>> oi = vdx.OptimalInterpolation(
...     obs_op=obs,
...     prior_mean=jnp.zeros((1, 2)),
...     prior_cov_op=eye_x,
...     obs_cov_op=eye_y,
... )
>>> y = jnp.array([1.0, 2.0, 3.0])
>>> batch = vdx.Batch1D(input=y[None, None], mask=jnp.ones((1, 1, 3)))
>>> blue = jnp.linalg.solve(jnp.eye(2) + F.T @ F, F.T @ y)
>>> bool(jnp.allclose(oi(batch)[0, 0], blue, atol=1e-4))
True

A nonlinear map is linearised by autodiff, and materialise=True forms the dense Jacobian:

>>> sq = vdx.CallableObs(fn=lambda x: x**2, materialise=True)
>>> jac = sq.linearize(jnp.array([1.0, 2.0, 3.0])).as_matrix()
>>> bool(jnp.allclose(jac, jnp.diag(jnp.array([2.0, 4.0, 6.0]))))
True
Source code in src/vardax/_src/obs_operators/callable.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
class CallableObs(eqx.Module):
    r"""$H(x) = f(x)$ for any differentiable JAX function $f$.

    ``linearize(x)`` is a matrix-free ``lineax.JacobianLinearOperator`` at
    ``x``: its ``mv`` is a JVP of $f$ and its ``transpose().mv`` a VJP, so
    nothing of size $\dim y \times \dim x$ is ever formed. The operator is
    an Equinox module, so it works under ``eqx.filter_jit`` and
    ``jax.vmap``; when ``fn`` is itself an Equinox module (an
    ``eqx.nn.MLP`` emulator, say), its parameters are pytree leaves and
    ``eqx.filter_grad`` differentiates through them.

    Prefer [`LinearObs`][vardax.LinearObs] when you already hold an explicit
    or structured matrix you want to reuse; use this for a model run, a
    simulator or an emulator.

    Pass ``jac="bwd"`` for a model that defines only a reverse-mode rule
    (a ``jax.custom_vjp`` function): forward mode raises on it as soon as
    an analysis or posterior linearises the operator. Matrix-free, lineax
    then applies $H'v$ by transposing the VJP, which fails when the
    backward pass contains a loop JAX cannot transpose, as a ``diffrax``
    solve's does (``RecursiveCheckpointAdjoint`` runs a ``while_loop``).
    For those, also pass ``materialise=True``: ``linearize`` then forms the
    Jacobian once, with one VJP per observation (``jax.jacrev``), and both
    $H'$ and $H'^\top$ are dense products. That suits a small control, the
    usual case for source and parameter estimation.

    Attributes:
        fn: The observation map $f$, from a state of any shape to an
            observation of any shape.
        jac: How ``linearize`` differentiates ``fn``: ``None`` (lineax's
            default, forward mode), ``"fwd"`` or ``"bwd"`` (reverse mode).
        materialise: Form the Jacobian in ``linearize`` instead of keeping
            it matrix-free, with ``jax.jacrev`` for ``jac="bwd"`` and
            ``jax.jacfwd`` otherwise. Default ``False``.

    Examples:
        A linear $f(x) = F x$ on a ``(1, N)`` state with $M \ne N$
        observations: ``OptimalInterpolation`` with $B = R = I$ gives the
        closed-form BLUE, $x^* = (I + F^\top F)^{-1} F^\top y$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax as vdx
        >>> F = jnp.array([[1.0, 0.0], [0.0, 2.0], [1.0, 1.0]])
        >>> obs = vdx.CallableObs(fn=lambda x: x @ F.T)
        >>> obs(jnp.ones((1, 2))).shape
        (1, 3)
        >>> eye_x = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 2), jnp.float32))
        >>> eye_y = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 3), jnp.float32))
        >>> oi = vdx.OptimalInterpolation(
        ...     obs_op=obs,
        ...     prior_mean=jnp.zeros((1, 2)),
        ...     prior_cov_op=eye_x,
        ...     obs_cov_op=eye_y,
        ... )
        >>> y = jnp.array([1.0, 2.0, 3.0])
        >>> batch = vdx.Batch1D(input=y[None, None], mask=jnp.ones((1, 1, 3)))
        >>> blue = jnp.linalg.solve(jnp.eye(2) + F.T @ F, F.T @ y)
        >>> bool(jnp.allclose(oi(batch)[0, 0], blue, atol=1e-4))
        True

        A nonlinear map is linearised by autodiff, and ``materialise=True``
        forms the dense Jacobian:

        >>> sq = vdx.CallableObs(fn=lambda x: x**2, materialise=True)
        >>> jac = sq.linearize(jnp.array([1.0, 2.0, 3.0])).as_matrix()
        >>> bool(jnp.allclose(jac, jnp.diag(jnp.array([2.0, 4.0, 6.0]))))
        True
    """

    fn: Callable[[Float[Array, ...]], Float[Array, ...]]
    jac: Literal["fwd", "bwd"] | None = eqx.field(static=True, default=None)
    materialise: bool = eqx.field(static=True, default=False)

    def __call__(self, x: Float[Array, ...]) -> Float[Array, ...]:
        """Forward map $f(x)$.

        Args:
            x: State.

        Returns:
            The observation prediction $f(x)$.
        """
        return self.fn(x)

    def linearize(self, x: Float[Array, ...]) -> lx.AbstractLinearOperator:
        r"""Tangent-linear operator $H'(x) = \partial f / \partial x$ at ``x``.

        Args:
            x: State at which to linearise.

        Returns:
            A matrix-free ``JacobianLinearOperator``, or its dense
            materialisation when ``materialise=True``.
        """
        tangent = lx.JacobianLinearOperator(lambda z, _: self.fn(z), x, jac=self.jac)
        if self.materialise:
            # dense ∂f/∂x: one VJP per output (jac="bwd") or JVP per input
            return lx.materialise(tangent)
        return tangent

linearize 

linearize(x: Float[Array, ...]) -> AbstractLinearOperator

Tangent-linear operator \(H'(x) = \partial f / \partial x\) at x.

Parameters:

Name Type Description Default
x Float[Array, ...]

State at which to linearise.

 required

Returns:

Type Description
AbstractLinearOperator

A matrix-free JacobianLinearOperator, or its dense
AbstractLinearOperator

materialisation when materialise=True.
Source code in src/vardax/_src/obs_operators/callable.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def linearize(self, x: Float[Array, ...]) -> lx.AbstractLinearOperator:
    r"""Tangent-linear operator $H'(x) = \partial f / \partial x$ at ``x``.

    Args:
        x: State at which to linearise.

    Returns:
        A matrix-free ``JacobianLinearOperator``, or its dense
        materialisation when ``materialise=True``.
    """
    tangent = lx.JacobianLinearOperator(lambda z, _: self.fn(z), x, jac=self.jac)
    if self.materialise:
        # dense ∂f/∂x: one VJP per output (jac="bwd") or JVP per input
        return lx.materialise(tangent)
    return tangent

MaskedIdentity

Bases: Module

\(H(x) = m \odot x\) with optional element-wise mask.

Stateless: no parameters. The mask can be supplied at call time (per-batch) rather than baked in.

Examples:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> op = vdx.MaskedIdentity()
>>> x = jnp.array([1.0, 2.0, 3.0])
>>> op(x, mask=jnp.array([1.0, 0.0, 1.0]))
Array([1., 0., 3.], dtype=float32)
>>> type(op.linearize(x)).__name__
'JacobianLinearOperator'

Without a mask the operator is the identity:

>>> bool(jnp.all(op(x) == x))
True

A per-batch mask broadcasts over the batch axis:

>>> xb = jnp.ones((2, 3))
>>> op(xb, mask=jnp.array([1.0, 0.0, 1.0])).sum(axis=-1)
Array([2., 2.], dtype=float32)
Source code in src/vardax/_src/obs_operators/masked.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class MaskedIdentity(eqx.Module):
    r"""$H(x) = m \odot x$ with optional element-wise mask.

    Stateless: no parameters. The mask can be supplied at call time
    (per-batch) rather than baked in.

    Examples:
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> op = vdx.MaskedIdentity()
        >>> x = jnp.array([1.0, 2.0, 3.0])
        >>> op(x, mask=jnp.array([1.0, 0.0, 1.0]))
        Array([1., 0., 3.], dtype=float32)
        >>> type(op.linearize(x)).__name__
        'JacobianLinearOperator'

        Without a mask the operator is the identity:

        >>> bool(jnp.all(op(x) == x))
        True

        A per-batch mask broadcasts over the batch axis:

        >>> xb = jnp.ones((2, 3))
        >>> op(xb, mask=jnp.array([1.0, 0.0, 1.0])).sum(axis=-1)
        Array([2., 2.], dtype=float32)
    """

    def __call__(
        self,
        x: Float[Array, ...],
        mask: Float[Array, ...] | None = None,
    ) -> Float[Array, ...]:
        if mask is None:
            return x
        return x * mask

    def linearize(self, x: Float[Array, ...]) -> lx.AbstractLinearOperator:
        """Tangent-linear operator at ``x``: identity (since $H$ is linear).

        Args:
            x: State at which to linearise.

        Returns:
            An identity-like Jacobian operator on the shape of ``x``.
        """
        return lx.JacobianLinearOperator(lambda y, _args=None: y, x)

linearize 

linearize(x: Float[Array, ...]) -> AbstractLinearOperator

Tangent-linear operator at x: identity (since \(H\) is linear).

Parameters:

Name Type Description Default
x Float[Array, ...]

State at which to linearise.

 required

Returns:

Type Description
AbstractLinearOperator

An identity-like Jacobian operator on the shape of x.
Source code in src/vardax/_src/obs_operators/masked.py 
53
54
55
56
57
58
59
60
61
62
def linearize(self, x: Float[Array, ...]) -> lx.AbstractLinearOperator:
    """Tangent-linear operator at ``x``: identity (since $H$ is linear).

    Args:
        x: State at which to linearise.

    Returns:
        An identity-like Jacobian operator on the shape of ``x``.
    """
    return lx.JacobianLinearOperator(lambda y, _args=None: y, x)

AveragingKernel

Bases: Module

RTM L2-style averaging-kernel obs operator.

\[H(x) = A \, (h \odot x + (1 - h) \odot x_a)\]
Algorithm
  1. Blend the state with the retrieval prior: \(z = h \odot x + (1 - h) \odot x_a\).
  2. Apply the kernel: \(\hat y = A z\).
  3. linearize returns \(A \, \mathrm{diag}(h)\) without evaluating anything, because the map is affine in \(x\).

Attributes:

Name Type Description
kernel_matrix AbstractLinearOperator

Averaging kernel matrix, stored as an lineax.AbstractLinearOperator. May be MatrixLinearOperator for dense kernels or a structured operator (Kronecker for separable kernels, LowRank for retrievals with low DOF).
prior_profile Float[Array, ' N']

Retrieval prior, same shape as the model state.
pressure_weights Float[Array, ' N']

Weighting vector (often pressure-weighted).

Examples:

With \(h = 1\) the prior drops out and the operator is just \(A x\):

>>> import jax.numpy as jnp, lineax as lx
>>> import vardax as vdx
>>> A = jnp.array([[0.5, 0.5], [0.0, 1.0]])
>>> kernel = lx.MatrixLinearOperator(A)
>>> op = vdx.AveragingKernel(
...     kernel_matrix=kernel,
...     prior_profile=jnp.zeros(2),
...     pressure_weights=jnp.ones(2),
... )
>>> bool(jnp.allclose(op(jnp.array([1.0, 3.0])), A @ jnp.array([1.0, 3.0])))
True

Where \(h = 0\) the output falls back to the retrieval prior \(x_a\):

>>> op0 = vdx.AveragingKernel(
...     kernel_matrix=kernel,
...     prior_profile=jnp.array([2.0, 2.0]),
...     pressure_weights=jnp.zeros(2),
... )
>>> bool(jnp.allclose(op0(jnp.array([9.0, 9.0])), A @ jnp.array([2.0, 2.0])))
True

The tangent-linear operator is \(A \, \mathrm{diag}(h)\):

>>> bool(jnp.allclose(op.linearize(jnp.zeros(2)).as_matrix(), A))
True
Source code in src/vardax/_src/obs_operators/averaging_kernel.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class AveragingKernel(eqx.Module):
    r"""RTM L2-style averaging-kernel obs operator.

    $$H(x) = A \, (h \odot x + (1 - h) \odot x_a)$$

    Algorithm:
        1. Blend the state with the retrieval prior:
           $z = h \odot x + (1 - h) \odot x_a$.
        2. Apply the kernel: $\hat y = A z$.
        3. ``linearize`` returns $A \, \mathrm{diag}(h)$ without evaluating
           anything, because the map is affine in $x$.

    Attributes:
        kernel_matrix: Averaging kernel matrix, stored as an
            ``lineax.AbstractLinearOperator``. May be ``MatrixLinearOperator``
            for dense kernels or a structured operator (Kronecker for
            separable kernels, LowRank for retrievals with low DOF).
        prior_profile: Retrieval prior, same shape as the model state.
        pressure_weights: Weighting vector (often pressure-weighted).

    Examples:
        With $h = 1$ the prior drops out and the operator is just $A x$:

        >>> import jax.numpy as jnp, lineax as lx
        >>> import vardax as vdx
        >>> A = jnp.array([[0.5, 0.5], [0.0, 1.0]])
        >>> kernel = lx.MatrixLinearOperator(A)
        >>> op = vdx.AveragingKernel(
        ...     kernel_matrix=kernel,
        ...     prior_profile=jnp.zeros(2),
        ...     pressure_weights=jnp.ones(2),
        ... )
        >>> bool(jnp.allclose(op(jnp.array([1.0, 3.0])), A @ jnp.array([1.0, 3.0])))
        True

        Where $h = 0$ the output falls back to the retrieval prior $x_a$:

        >>> op0 = vdx.AveragingKernel(
        ...     kernel_matrix=kernel,
        ...     prior_profile=jnp.array([2.0, 2.0]),
        ...     pressure_weights=jnp.zeros(2),
        ... )
        >>> bool(jnp.allclose(op0(jnp.array([9.0, 9.0])), A @ jnp.array([2.0, 2.0])))
        True

        The tangent-linear operator is $A \, \mathrm{diag}(h)$:

        >>> bool(jnp.allclose(op.linearize(jnp.zeros(2)).as_matrix(), A))
        True
    """

    kernel_matrix: lx.AbstractLinearOperator
    prior_profile: Float[Array, " N"]
    pressure_weights: Float[Array, " N"]

    def __call__(self, x: Float[Array, " N"]) -> Float[Array, " N"]:
        w = self.pressure_weights
        inner = w * x + (1.0 - w) * self.prior_profile
        return self.kernel_matrix.mv(inner)

    def linearize(self, x: Float[Array, " N"]) -> lx.AbstractLinearOperator:
        r"""Tangent-linear operator at ``x``: $A \cdot \mathrm{diag}(h)$.

        Args:
            x: State at which to linearise; unused because the map is
                affine.

        Returns:
            The operator $A \, \mathrm{diag}(h)$.
        """
        diag_w = lx.DiagonalLinearOperator(self.pressure_weights)
        return self.kernel_matrix @ diag_w

linearize 

linearize(x: Float[Array, ' N']) -> AbstractLinearOperator

Tangent-linear operator at x: \(A \cdot \mathrm{diag}(h)\).

Parameters:

Name Type Description Default
x Float[Array, ' N']

State at which to linearise; unused because the map is affine.

 required

Returns:

Type Description
AbstractLinearOperator

The operator \(A \, \mathrm{diag}(h)\).
Source code in src/vardax/_src/obs_operators/averaging_kernel.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def linearize(self, x: Float[Array, " N"]) -> lx.AbstractLinearOperator:
    r"""Tangent-linear operator at ``x``: $A \cdot \mathrm{diag}(h)$.

    Args:
        x: State at which to linearise; unused because the map is
            affine.

    Returns:
        The operator $A \, \mathrm{diag}(h)$.
    """
    diag_w = lx.DiagonalLinearOperator(self.pressure_weights)
    return self.kernel_matrix @ diag_w

InterpObs

Bases: Module

Grid → scattered-obs sparse bilinear interpolation.

\[(Hx)_k = \sum_{i \in \mathrm{stencil}(p_k)} w_{k,i} \, x_i\]

Build a precomputed sparse (n_obs, 4) weight stencil from a 2-D grid and an (n_obs, 2) array of obs coordinates. The forward indexing is pure JAX (so jit / grad / vmap all work); linearize wraps the same forward as a lineax.FunctionLinearOperator (its transpose comes from jax.vjp automatically).

Algorithm
  1. Flatten the grid state.
  2. Gather the four stencil corners of each observation with indices.
  3. Take the weighted sum of the corners with weights.
  4. Multiply by mask if one is given.

Attributes:

Name Type Description
indices Int[Array, 'n_obs 4']

Integer flat grid indices of each observation's 2x2 stencil.
weights Float[Array, 'n_obs 4']

Bilinear weights of the four corners, summing to 1 for each observation.
grid_shape tuple[int, int]

Static (ny, nx) extent of the input grid.

Examples:

Interpolate a linear field at the centre of a unit cell:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> axis = jnp.array([0.0, 1.0])
>>> op = vdx.interp_obs_from_coords((axis, axis), jnp.array([[0.5, 0.5]]))
>>> field = jnp.array([[0.0, 1.0], [2.0, 3.0]])
>>> bool(jnp.allclose(op(field), 1.5))
True

A mask zeroes out dropped observations:

>>> pts = jnp.array([[0.0, 0.0], [1.0, 1.0]])
>>> op2 = vdx.interp_obs_from_coords((axis, axis), pts)
>>> op2(field, mask=jnp.array([1.0, 0.0]))
Array([0., 0.], dtype=float32)

The operator is linear, so its linearisation applies the same map:

>>> bool(jnp.allclose(op2.linearize(field).mv(field), op2(field)))
True
Source code in src/vardax/_src/obs_operators/interp_obs.py 
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
class InterpObs(eqx.Module):
    r"""Grid → scattered-obs sparse bilinear interpolation.

    $$(Hx)_k = \sum_{i \in \mathrm{stencil}(p_k)} w_{k,i} \, x_i$$

    Build a precomputed sparse ``(n_obs, 4)`` weight stencil from a 2-D
    grid and an ``(n_obs, 2)`` array of obs coordinates. The forward
    indexing is pure JAX (so ``jit`` / ``grad`` / ``vmap`` all work);
    ``linearize`` wraps the same forward as a
    ``lineax.FunctionLinearOperator`` (its transpose comes from
    ``jax.vjp`` automatically).

    Algorithm:
        1. Flatten the grid state.
        2. Gather the four stencil corners of each observation with
           ``indices``.
        3. Take the weighted sum of the corners with ``weights``.
        4. Multiply by ``mask`` if one is given.

    Attributes:
        indices: Integer flat grid indices of each observation's 2x2
            stencil.
        weights: Bilinear weights of the four corners, summing to 1 for
            each observation.
        grid_shape: Static ``(ny, nx)`` extent of the input grid.

    Examples:
        Interpolate a linear field at the centre of a unit cell:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> axis = jnp.array([0.0, 1.0])
        >>> op = vdx.interp_obs_from_coords((axis, axis), jnp.array([[0.5, 0.5]]))
        >>> field = jnp.array([[0.0, 1.0], [2.0, 3.0]])
        >>> bool(jnp.allclose(op(field), 1.5))
        True

        A mask zeroes out dropped observations:

        >>> pts = jnp.array([[0.0, 0.0], [1.0, 1.0]])
        >>> op2 = vdx.interp_obs_from_coords((axis, axis), pts)
        >>> op2(field, mask=jnp.array([1.0, 0.0]))
        Array([0., 0.], dtype=float32)

        The operator is linear, so its linearisation applies the same map:

        >>> bool(jnp.allclose(op2.linearize(field).mv(field), op2(field)))
        True
    """

    indices: Int[Array, "n_obs 4"]
    weights: Float[Array, "n_obs 4"]
    grid_shape: tuple[int, int] = eqx.field(static=True)

    def __call__(
        self,
        state: Float[Array, "*grid"],
        mask: Float[Array, " n_obs"] | None = None,
    ) -> Float[Array, " n_obs"]:
        """Apply ``H`` (optionally masked) to a grid state.

        Args:
            state: Grid state, either ``(ny, nx)`` or flattened.
            mask: Optional observation mask multiplying the predictions.

        Returns:
            The interpolated observations, one per location.
        """
        flat = flatten(state)
        corners = flat[self.indices]
        # yₒ = Σₖ wₒₖ x[iₒₖ]: bilinear weights over the four corners;
        # jnp, not einx: a time step may have no observations
        predicted_obs = jnp.sum(corners * self.weights, axis=-1)
        if mask is not None:
            predicted_obs = predicted_obs * mask
        return predicted_obs

    def linearize(
        self,
        x: Float[Array, ...],
    ) -> lx.AbstractLinearOperator:
        """Tangent-linear at ``x``: ``H`` itself (linear operator).

        Wraps ``__call__`` as a ``lineax.FunctionLinearOperator`` so it
        plugs straight into ``IncrementalFourDVar``'s GN Hessian
        assembly and ``LaplaceCovariance``. The transpose ``H^T`` is
        derived via ``jax.vjp`` (cached numpy adjoint is no longer
        exposed; use ``op.transpose().mv(y)`` instead).

        The input structure is taken from ``x.shape`` so the operator
        composes with both 2-D ``(Ny, Nx)`` and flat ``(Ny*Nx,)`` state
        layouts — variational models that carry flat state vectors then
        get a flat-shaped ``H`` instead of being forced through the
        grid shape.

        Args:
            x: State whose shape and dtype define the operator input.

        Returns:
            The operator $H$, which does not depend on ``x`` otherwise.
        """
        in_struct = jax.ShapeDtypeStruct(x.shape, x.dtype)
        return lx.FunctionLinearOperator(self.__call__, in_struct)

linearize 

linearize(x: Float[Array, ...]) -> AbstractLinearOperator

Tangent-linear at x: H itself (linear operator).

Wraps __call__ as a lineax.FunctionLinearOperator so it plugs straight into IncrementalFourDVar's GN Hessian assembly and LaplaceCovariance. The transpose H^T is derived via jax.vjp (cached numpy adjoint is no longer exposed; use op.transpose().mv(y) instead).

The input structure is taken from x.shape so the operator composes with both 2-D (Ny, Nx) and flat (Ny*Nx,) state layouts — variational models that carry flat state vectors then get a flat-shaped H instead of being forced through the grid shape.

Parameters:

Name Type Description Default
x Float[Array, ...]

State whose shape and dtype define the operator input.

 required

Returns:

Type Description
AbstractLinearOperator

The operator \(H\), which does not depend on x otherwise.
Source code in src/vardax/_src/obs_operators/interp_obs.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def linearize(
    self,
    x: Float[Array, ...],
) -> lx.AbstractLinearOperator:
    """Tangent-linear at ``x``: ``H`` itself (linear operator).

    Wraps ``__call__`` as a ``lineax.FunctionLinearOperator`` so it
    plugs straight into ``IncrementalFourDVar``'s GN Hessian
    assembly and ``LaplaceCovariance``. The transpose ``H^T`` is
    derived via ``jax.vjp`` (cached numpy adjoint is no longer
    exposed; use ``op.transpose().mv(y)`` instead).

    The input structure is taken from ``x.shape`` so the operator
    composes with both 2-D ``(Ny, Nx)`` and flat ``(Ny*Nx,)`` state
    layouts — variational models that carry flat state vectors then
    get a flat-shaped ``H`` instead of being forced through the
    grid shape.

    Args:
        x: State whose shape and dtype define the operator input.

    Returns:
        The operator $H$, which does not depend on ``x`` otherwise.
    """
    in_struct = jax.ShapeDtypeStruct(x.shape, x.dtype)
    return lx.FunctionLinearOperator(self.__call__, in_struct)

interp_obs_from_coords 

interp_obs_from_coords(
    grid_coords: tuple[
        Float[ArrayLike, " Ny"], Float[ArrayLike, " Nx"]
    ],
    obs_coords: Float[ArrayLike, "n_obs 2"],
    *,
    order: int = 1,
) -> InterpObs

Build an InterpObs from grid axis arrays + obs locations.

Locates each observation in its grid cell and caches the four corner indices and bilinear weights. Points outside the grid are clamped to the edge.

Parameters:

Name Type Description Default
grid_coords tuple[Float[ArrayLike, ' Ny'], Float[ArrayLike, ' Nx']]

(y_axis, x_axis) tuple of 1-D axis arrays.

 required
obs_coords Float[ArrayLike, 'n_obs 2']

(y, x) observation locations, one row each.

 required
order int

interpolation order. Only order=1 (bilinear) is currently supported.

 1

Returns:

Type Description
InterpObs

Configured InterpObs instance with precomputed cached
InterpObs

(indices, weights).

Raises:

Type Description
GridError

If grid_coords does not hold two axes, or an axis is not strictly increasing with at least two points.
ShapeError

If obs_coords is not an (n_obs, 2) array.
NotImplementedError

If order is not 1.

Examples:

Cache the stencil for two observations on a 3 x 3 grid:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> axis = jnp.array([0.0, 1.0, 2.0])
>>> obs = jnp.array([[0.5, 1.5], [2.0, 0.0]])
>>> op = vdx.interp_obs_from_coords((axis, axis), obs)
>>> op.indices.shape, op.weights.shape
((2, 4), (2, 4))

Predict the observations from a state whose value is 3 * y + x:

>>> op(jnp.arange(9.0).reshape(3, 3))
Array([3., 6.], dtype=float32)

A descending axis is rejected:

>>> vdx.interp_obs_from_coords(
...     (axis[::-1], axis), obs
... )
Traceback (most recent call last):
...
GridError: y_axis must be strictly increasing ...
Source code in src/vardax/_src/obs_operators/interp_obs.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def interp_obs_from_coords(
    grid_coords: tuple[Float[ArrayLike, " Ny"], Float[ArrayLike, " Nx"]],
    obs_coords: Float[ArrayLike, "n_obs 2"],
    *,
    order: int = 1,
) -> InterpObs:
    """Build an `InterpObs` from grid axis arrays + obs locations.

    Locates each observation in its grid cell and caches the four corner
    indices and bilinear weights. Points outside the grid are clamped to
    the edge.

    Args:
        grid_coords: ``(y_axis, x_axis)`` tuple of 1-D axis arrays.
        obs_coords: ``(y, x)`` observation locations, one row each.
        order: interpolation order. Only ``order=1`` (bilinear) is
            currently supported.

    Returns:
        Configured ``InterpObs`` instance with precomputed cached
        ``(indices, weights)``.

    Raises:
        GridError: If ``grid_coords`` does not hold two axes, or an axis
            is not strictly increasing with at least two points.
        ShapeError: If ``obs_coords`` is not an ``(n_obs, 2)`` array.
        NotImplementedError: If ``order`` is not 1.

    Examples:
        Cache the stencil for two observations on a 3 x 3 grid:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> axis = jnp.array([0.0, 1.0, 2.0])
        >>> obs = jnp.array([[0.5, 1.5], [2.0, 0.0]])
        >>> op = vdx.interp_obs_from_coords((axis, axis), obs)
        >>> op.indices.shape, op.weights.shape
        ((2, 4), (2, 4))

        Predict the observations from a state whose value is 3 * y + x:

        >>> op(jnp.arange(9.0).reshape(3, 3))
        Array([3., 6.], dtype=float32)

        A descending axis is rejected:

        >>> vdx.interp_obs_from_coords(
        ...     (axis[::-1], axis), obs
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        GridError: y_axis must be strictly increasing ...
    """
    if order != 1:
        raise NotImplementedError(
            f"interp_obs_from_coords supports order=1 (bilinear), got {order}."
        )
    if len(grid_coords) != 2:
        raise GridError(
            "interp_obs_from_coords expects grid_coords=(y_axis, x_axis); "
            f"got len={len(grid_coords)}."
        )

    y_axis = _np.asarray(grid_coords[0], dtype=float)
    x_axis = _np.asarray(grid_coords[1], dtype=float)
    for name, axis in (("y_axis", y_axis), ("x_axis", x_axis)):
        if axis.size < 2 or _np.any(_np.diff(axis) <= 0):
            raise GridError(
                f"{name} must be strictly increasing with at least 2 points; "
                "reverse a descending axis before calling interp_obs_from_coords."
            )
    points = _np.asarray(obs_coords, dtype=float)
    if points.ndim != 2 or points.shape[1] != 2:
        raise ShapeError(f"obs_coords must be (n_obs, 2); got {points.shape}.")

    n_y, n_x = y_axis.size, x_axis.size
    y_lower_index, y_upper_weight = _bracket(y_axis, points[:, 0])
    x_lower_index, x_upper_weight = _bracket(x_axis, points[:, 1])
    y_upper_index, x_upper_index = y_lower_index + 1, x_lower_index + 1
    y_lower_weight, x_lower_weight = 1.0 - y_upper_weight, 1.0 - x_upper_weight

    # flat index of the 2x2 stencil: (lo, lo), (lo, hi), (hi, lo), (hi, hi)
    indices = _np.stack(
        [
            y_lower_index * n_x + x_lower_index,
            y_lower_index * n_x + x_upper_index,
            y_upper_index * n_x + x_lower_index,
            y_upper_index * n_x + x_upper_index,
        ],
        axis=1,
    )
    # bilinear weight of each corner, w_y w_x, in the same order
    weights = _np.stack(
        [
            y_lower_weight * x_lower_weight,
            y_lower_weight * x_upper_weight,
            y_upper_weight * x_lower_weight,
            y_upper_weight * x_upper_weight,
        ],
        axis=1,
    )

    return InterpObs(
        indices=jnp.asarray(indices, dtype=jnp.int32),
        # JAX's default float dtype, so the weights are float64 under x64
        weights=jnp.asarray(weights),
        grid_shape=(int(n_y), int(n_x)),
    )

Scattered point observations

Observations at arbitrary coordinates and times, carried by a PointObsBatch with per-observation \(\sigma_i\) and a valid mask for padding. GridPointObs evaluates a gridded state or trajectory at the points (linear / bilinear in space, at each observation's time_index); FieldPointObs evaluates a callable field; point_obs_cost is the masked, \(\sigma\)-weighted cost, exactly zero in value and gradient on padding rows. The classical models dispatch on the batch type.

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

GridPointObs

Bases: Module

\(\mathcal H_t(x_t)\) at scattered points for a gridded state.

The state carries a leading time axis, (n_times, *grid) (a 4D-Var trajectory, or a ThreeDVar state (T, N)), with grid (N,) (1-D) or (Ny, Nx) (2-D) on the axes grid_axes. For observation \(i\) the operator takes the state at time_index[i] (\(0\) when the batch has no time_index) and interpolates it linearly in each spatial dimension at the first len(grid_axes) columns of coords[i], clamping points outside the grid to its edge — the order=1 scheme of interp_obs_from_coords, with the stencil computed at call time so the coordinates can be traced. A time_index outside [0, n_times - 1] on a valid row is an error when the indices are concrete (the models check it); traced indices are clamped to that range. Any further coordinate columns (a time) are ignored: the time enters through time_index.

Algorithm
  1. Replace padding rows by finite placeholders.
  2. Pick each observation's frame state[time_index].
  3. Bracket the coordinates in the grid axes, clamping at the edges.
  4. Combine the stencil corners with linear (1-D) or bilinear (2-D) weights.

Attributes:

Name Type Description
grid_axes tuple[Float[Array, ' _'], ...]

One strictly increasing 1-D axis per spatial dimension: (x_axis,) or (y_axis, x_axis).

Raises:

Type Description
GridError

At construction, if grid_axes holds neither one nor two axes, or an axis is not 1-D, strictly increasing and at least two points long.

Examples:

A 1-D grid at \(0, 1, 2, 3\): linear interpolation between nodes.

>>> import jax.numpy as jnp, vardax
>>> H = vardax.GridPointObs((jnp.arange(4.0),))
>>> state = jnp.array([[0.0, 10.0, 20.0, 30.0]])  # (n_times=1, N=4)
>>> batch = vardax.PointObsBatch(
...     coords=jnp.array([[0.5], [2.25]]),
...     values=jnp.zeros(2),
...     sigma=jnp.ones(2),
...     valid=jnp.ones(2, bool),
... )
>>> H(state, batch)
Array([ 5. , 22.5], dtype=float32)

Points outside the grid are clamped to the edge value:

>>> far = vardax.PointObsBatch(
...     coords=jnp.array([[-5.0], [9.0]]),
...     values=jnp.zeros(2),
...     sigma=jnp.ones(2),
...     valid=jnp.ones(2, bool),
... )
>>> H(state, far)
Array([ 0., 30.], dtype=float32)

A non-increasing axis is rejected:

>>> vardax.GridPointObs(
...     (jnp.array([1.0, 0.0]),)
... )
Traceback (most recent call last):
...
GridError: grid_axes[0] must be strictly increasing.
Source code in src/vardax/_src/obs_operators/point.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
class GridPointObs(eqx.Module):
    r"""$\mathcal H_t(x_t)$ at scattered points for a gridded state.

    The state carries a leading time axis, ``(n_times, *grid)`` (a 4D-Var
    trajectory, or a ``ThreeDVar`` state ``(T, N)``), with ``grid``
    ``(N,)`` (1-D) or ``(Ny, Nx)`` (2-D) on the axes ``grid_axes``. For
    observation $i$ the operator takes the state at ``time_index[i]``
    ($0$ when the batch has no ``time_index``) and interpolates it
    linearly in each spatial dimension at the first ``len(grid_axes)``
    columns of ``coords[i]``, clamping points outside the grid to its
    edge — the ``order=1`` scheme of `interp_obs_from_coords`, with the
    stencil computed at call time so the coordinates can be traced. A
    ``time_index`` outside ``[0, n_times - 1]`` on a valid row is an error
    when the indices are concrete (the models check it); traced indices
    are clamped to that range. Any
    further coordinate columns (a time) are ignored: the time enters
    through ``time_index``.

    Algorithm:
        1. Replace padding rows by finite placeholders.
        2. Pick each observation's frame ``state[time_index]``.
        3. Bracket the coordinates in the grid axes, clamping at the edges.
        4. Combine the stencil corners with linear (1-D) or bilinear (2-D)
           weights.

    Attributes:
        grid_axes: One strictly increasing 1-D axis per spatial
            dimension: ``(x_axis,)`` or ``(y_axis, x_axis)``.

    Raises:
        GridError: At construction, if ``grid_axes`` holds neither one nor
            two axes, or an axis is not 1-D, strictly increasing and at
            least two points long.

    Examples:
        A 1-D grid at $0, 1, 2, 3$: linear interpolation between nodes.

        >>> import jax.numpy as jnp, vardax
        >>> H = vardax.GridPointObs((jnp.arange(4.0),))
        >>> state = jnp.array([[0.0, 10.0, 20.0, 30.0]])  # (n_times=1, N=4)
        >>> batch = vardax.PointObsBatch(
        ...     coords=jnp.array([[0.5], [2.25]]),
        ...     values=jnp.zeros(2),
        ...     sigma=jnp.ones(2),
        ...     valid=jnp.ones(2, bool),
        ... )
        >>> H(state, batch)
        Array([ 5. , 22.5], dtype=float32)

        Points outside the grid are clamped to the edge value:

        >>> far = vardax.PointObsBatch(
        ...     coords=jnp.array([[-5.0], [9.0]]),
        ...     values=jnp.zeros(2),
        ...     sigma=jnp.ones(2),
        ...     valid=jnp.ones(2, bool),
        ... )
        >>> H(state, far)
        Array([ 0., 30.], dtype=float32)

        A non-increasing axis is rejected:

        >>> vardax.GridPointObs(
        ...     (jnp.array([1.0, 0.0]),)
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        GridError: grid_axes[0] must be strictly increasing.
    """

    grid_axes: tuple[Float[Array, " _"], ...]

    def __init__(self, grid_axes: tuple[Float[Array, " _"], ...]) -> None:
        axes = tuple(jnp.asarray(axis_values) for axis_values in grid_axes)
        if len(axes) not in (1, 2):
            raise GridError(f"grid_axes must hold 1 or 2 axes; got {len(axes)}.")
        for axis_number, axis in enumerate(axes):
            if axis.ndim != 1 or axis.shape[0] < 2:
                raise GridError(
                    f"grid_axes[{axis_number}] must be 1-D with at least 2 points; "
                    f"got shape {axis.shape}."
                )
            if not bool(jnp.all(jnp.diff(axis) > 0)):
                raise GridError(
                    f"grid_axes[{axis_number}] must be strictly increasing."
                )
        self.grid_axes = axes

    def __call__(
        self, state: Float[Array, "n_times ..."], batch: PointObsBatch
    ) -> Float[Array, " N_obs"]:
        r"""Interpolate ``state[time_index]`` at the observation coordinates.

        Args:
            state: Gridded state with a leading time axis.
            batch: The point observations.

        Returns:
            The interpolated prediction $\hat y_i$ for every row,
            padding included.
        """
        safe = _safe(batch)
        n_obs = safe.values.shape[0]
        grid_shape = tuple(axis.shape[0] for axis in self.grid_axes)
        flat_frames = einx.id("t s... -> t (s...)", state)
        frames = einx.id("t (g...) -> t g...", flat_frames, g=grid_shape)
        time_indices = (
            jnp.zeros(n_obs, dtype=jnp.int32)
            if safe.time_index is None
            else jnp.clip(safe.time_index, 0, state.shape[0] - 1)
        )
        if len(self.grid_axes) == 1:
            left_index, right_weight = _bracket(self.grid_axes[0], safe.coords[:, 0])
            left = frames[time_indices, left_index]
            right = frames[time_indices, left_index + 1]
            # ŷ = (1 − w) f_left + w f_right
            return (1.0 - right_weight) * left + right_weight * right
        y_index, y_weight = _bracket(self.grid_axes[0], safe.coords[:, 0])
        x_index, x_weight = _bracket(self.grid_axes[1], safe.coords[:, 1])
        corner_00 = frames[time_indices, y_index, x_index]
        corner_01 = frames[time_indices, y_index, x_index + 1]
        corner_10 = frames[time_indices, y_index + 1, x_index]
        corner_11 = frames[time_indices, y_index + 1, x_index + 1]
        # ŷ = (1 − w_y) [(1 − w_x) f₀₀ + w_x f₀₁] + w_y [(1 − w_x) f₁₀ + w_x f₁₁]
        return (1.0 - y_weight) * (
            (1.0 - x_weight) * corner_00 + x_weight * corner_01
        ) + y_weight * ((1.0 - x_weight) * corner_10 + x_weight * corner_11)

FieldPointObs

Bases: Module

\(\mathcal H(f)\) for a state that is a callable field.

Evaluates the single-example field \(f(s) \mapsto\) (1,) (or a scalar) at every observation coordinate, \(\hat y_i = f(s_i)\); all D coordinate columns (including a time column) are passed to the field. Padding rows are evaluated at the first valid row's coordinates, so the field is never queried where no observation is.

Algorithm
  1. Replace padding rows by finite placeholders.
  2. vmap the field over the coordinate rows.
  3. Reduce each result to one value per observation.

Examples:

A field of two coordinates:

>>> import jax.numpy as jnp, vardax
>>> H = vardax.FieldPointObs()
>>> batch = vardax.PointObsBatch(
...     coords=jnp.array([[1.0, 2.0], [3.0, 4.0]]),
...     values=jnp.zeros(2),
...     sigma=jnp.ones(2),
...     valid=jnp.ones(2, bool),
... )
>>> H(lambda s: jnp.sum(s, keepdims=True), batch)
Array([3., 7.], dtype=float32)

A scalar-returning field works too, and a padding row reuses the first valid coordinates:

>>> padded = vardax.PointObsBatch(
...     coords=jnp.array([[1.0, 2.0], [9.0, 9.0]]),
...     values=jnp.zeros(2),
...     sigma=jnp.ones(2),
...     valid=jnp.array([True, False]),
... )
>>> H(lambda s: s[0] * s[1], padded)
Array([2., 2.], dtype=float32)
Source code in src/vardax/_src/obs_operators/point.py 
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
class FieldPointObs(eqx.Module):
    r"""$\mathcal H(f)$ for a state that is a callable field.

    Evaluates the single-example field $f(s) \mapsto$ ``(1,)`` (or a
    scalar) at every observation coordinate, $\hat y_i = f(s_i)$; all
    ``D`` coordinate columns (including a time column) are passed to the
    field. Padding rows are evaluated at the first valid row's
    coordinates, so the field is never queried where no observation is.

    Algorithm:
        1. Replace padding rows by finite placeholders.
        2. ``vmap`` the field over the coordinate rows.
        3. Reduce each result to one value per observation.

    Examples:
        A field of two coordinates:

        >>> import jax.numpy as jnp, vardax
        >>> H = vardax.FieldPointObs()
        >>> batch = vardax.PointObsBatch(
        ...     coords=jnp.array([[1.0, 2.0], [3.0, 4.0]]),
        ...     values=jnp.zeros(2),
        ...     sigma=jnp.ones(2),
        ...     valid=jnp.ones(2, bool),
        ... )
        >>> H(lambda s: jnp.sum(s, keepdims=True), batch)
        Array([3., 7.], dtype=float32)

        A scalar-returning field works too, and a padding row reuses the
        first valid coordinates:

        >>> padded = vardax.PointObsBatch(
        ...     coords=jnp.array([[1.0, 2.0], [9.0, 9.0]]),
        ...     values=jnp.zeros(2),
        ...     sigma=jnp.ones(2),
        ...     valid=jnp.array([True, False]),
        ... )
        >>> H(lambda s: s[0] * s[1], padded)
        Array([2., 2.], dtype=float32)
    """

    def __call__(
        self, field: Callable[[Array], Array], batch: PointObsBatch
    ) -> Float[Array, " N_obs"]:
        r"""Evaluate ``field`` at each coordinate row.

        Args:
            field: Callable mapping one coordinate row to a value.
            batch: The point observations.

        Returns:
            The field value $\hat y_i = f(s_i)$ for every row.
        """
        safe = _safe(batch)
        values = jax.vmap(field)(safe.coords)
        # a scalar field may return (1,) per point; one value per observation
        return as_shape(values, (values.shape[0],))

point_obs_cost 

point_obs_cost(
    predictions: Float[Array, " N_obs"],
    batch: PointObsBatch,
) -> Float[Array, ""]

Masked, \(\sigma\)-weighted observation cost of a point batch.

\[J_o = \tfrac12 \sum_i m_i (y_i - \hat y_i)^2 / \sigma_i^2\]

The whitened residual is set to zero by jnp.where on padding rows before squaring, so padding contributes zero value and zero gradient.

Algorithm
  1. Replace padding rows by finite placeholders.
  2. Form the whitened residual \((y_i - \hat y_i) / \sigma_i\) and set it to zero where \(m_i\) is false.
  3. Return half the sum of squares.

Parameters:

Name Type Description Default
predictions Float[Array, ' N_obs']

Predicted observations \(\hat y_i = \mathcal H(x)(s_i, t_i)\).

 required
batch PointObsBatch

The observations (one sample, no batch axis).

 required

Returns:

Type Description
Float[Array, '']

The scalar cost.

Examples:

A batch whose third row is padding (NaN value, zero sigma):

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.PointObsBatch(
...     coords=jnp.zeros((3, 1)),
...     values=jnp.array([1.0, 2.0, jnp.nan]),
...     sigma=jnp.array([1.0, 2.0, 0.0]),
...     valid=jnp.array([True, True, False]),
... )
>>> float(vardax.point_obs_cost(jnp.zeros(3), batch))  # ½ (1 + 1)
1.0

Padding contributes zero gradient, not NaN:

>>> import jax
>>> grad = jax.grad(lambda p: vardax.point_obs_cost(p, batch))(jnp.zeros(3))
>>> bool(jnp.all(jnp.isfinite(grad)))
True

Composed with a point operator on a gridded state:

>>> H = vardax.GridPointObs((jnp.arange(3.0),))
>>> b = vardax.PointObsBatch(
...     coords=jnp.array([[0.5]]),
...     values=jnp.array([1.0]),
...     sigma=jnp.array([1.0]),
...     valid=jnp.array([True]),
... )
>>> float(vardax.point_obs_cost(H(jnp.zeros((1, 3)), b), b))
0.5
Source code in src/vardax/_src/obs_operators/point.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def point_obs_cost(
    predictions: Float[Array, " N_obs"], batch: PointObsBatch
) -> Float[Array, ""]:
    r"""Masked, $\sigma$-weighted observation cost of a point batch.

    $$J_o = \tfrac12 \sum_i m_i (y_i - \hat y_i)^2 / \sigma_i^2$$

    The whitened residual is set to zero by ``jnp.where`` on padding rows
    before squaring, so padding contributes zero value and zero gradient.

    Algorithm:
        1. Replace padding rows by finite placeholders.
        2. Form the whitened residual $(y_i - \hat y_i) / \sigma_i$ and set
           it to zero where $m_i$ is false.
        3. Return half the sum of squares.

    Args:
        predictions: Predicted observations
            $\hat y_i = \mathcal H(x)(s_i, t_i)$.
        batch: The observations (one sample, no batch axis).

    Returns:
        The scalar cost.

    Examples:
        A batch whose third row is padding (``NaN`` value, zero sigma):

        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.PointObsBatch(
        ...     coords=jnp.zeros((3, 1)),
        ...     values=jnp.array([1.0, 2.0, jnp.nan]),
        ...     sigma=jnp.array([1.0, 2.0, 0.0]),
        ...     valid=jnp.array([True, True, False]),
        ... )
        >>> float(vardax.point_obs_cost(jnp.zeros(3), batch))  # ½ (1 + 1)
        1.0

        Padding contributes zero gradient, not ``NaN``:

        >>> import jax
        >>> grad = jax.grad(lambda p: vardax.point_obs_cost(p, batch))(jnp.zeros(3))
        >>> bool(jnp.all(jnp.isfinite(grad)))
        True

        Composed with a point operator on a gridded state:

        >>> H = vardax.GridPointObs((jnp.arange(3.0),))
        >>> b = vardax.PointObsBatch(
        ...     coords=jnp.array([[0.5]]),
        ...     values=jnp.array([1.0]),
        ...     sigma=jnp.array([1.0]),
        ...     valid=jnp.array([True]),
        ... )
        >>> float(vardax.point_obs_cost(H(jnp.zeros((1, 3)), b), b))
        0.5
    """
    return 0.5 * jnp.sum(jnp.square(whitened_residual(predictions, batch)))

Multi-instrument fusion

Assimilating several instruments at once: each instrument is described by an InstrumentSpec, registered in an InstrumentRegistry, and MultiInstrumentFusion stacks the per-instrument operators into a single observation operator over the concatenated observation vector.

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

InstrumentSpec

Bases: Module

Per-instrument (obs_op, mask, obs_cov_op, id) bundle.

A plain container: it only groups one instrument's operator, quality mask, error covariance and name.

Attributes:

Name Type Description
obs_op Any

ObservationOperator-conforming operator (typically AveragingKernel).
mask Float[Array, ...]

Quality mask of shape compatible with the instrument's observation space. 1 for valid pixels, 0 for flagged/dropped.
obs_cov_op AbstractLinearOperator

Observation-error covariance as a lineax.AbstractLinearOperator. Often a DiagonalLinearOperator keyed on the per-pixel retrieval uncertainty.
instrument_id str

Identifier (e.g. "TROPOMI", "EMIT", "GHGSat").

Examples:

Bundle a masked-identity instrument with unit error covariance:

>>> import jax, jax.numpy as jnp, lineax as lx
>>> import vardax as vdx
>>> spec = vdx.InstrumentSpec(
...     obs_op=vdx.MaskedIdentity(),
...     mask=jnp.ones(3),
...     obs_cov_op=lx.IdentityLinearOperator(
...         jax.ShapeDtypeStruct((3,), jnp.float32)
...     ),
...     instrument_id="TROPOMI",
... )
>>> spec.instrument_id
'TROPOMI'
Source code in src/vardax/_src/obs_operators/multi_instrument.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class InstrumentSpec(eqx.Module):
    """Per-instrument ``(obs_op, mask, obs_cov_op, id)`` bundle.

    A plain container: it only groups one instrument's operator, quality
    mask, error covariance and name.

    Attributes:
        obs_op: ``ObservationOperator``-conforming operator (typically
            [`AveragingKernel`][vardax.AveragingKernel]).
        mask: Quality mask of shape compatible with the instrument's
            observation space. ``1`` for valid pixels, ``0`` for
            flagged/dropped.
        obs_cov_op: Observation-error covariance as a
            ``lineax.AbstractLinearOperator``. Often a
            ``DiagonalLinearOperator`` keyed on the per-pixel
            retrieval uncertainty.
        instrument_id: Identifier (e.g. ``"TROPOMI"``, ``"EMIT"``,
            ``"GHGSat"``).

    Examples:
        Bundle a masked-identity instrument with unit error covariance:

        >>> import jax, jax.numpy as jnp, lineax as lx
        >>> import vardax as vdx
        >>> spec = vdx.InstrumentSpec(
        ...     obs_op=vdx.MaskedIdentity(),
        ...     mask=jnp.ones(3),
        ...     obs_cov_op=lx.IdentityLinearOperator(
        ...         jax.ShapeDtypeStruct((3,), jnp.float32)
        ...     ),
        ...     instrument_id="TROPOMI",
        ... )
        >>> spec.instrument_id
        'TROPOMI'
    """

    obs_op: Any  # ObservationOperator (we don't import the Protocol to avoid a cycle)
    mask: Float[Array, ...]
    obs_cov_op: lx.AbstractLinearOperator
    instrument_id: str = eqx.field(static=True)

InstrumentRegistry

Bases: Module

Keyed lookup of InstrumentSpec by instrument_id.

Attributes:

Name Type Description
entries dict[str, InstrumentSpec]

dict[instrument_id, InstrumentSpec].

Examples:

An empty registry by default, then one keyed by instrument id:

>>> import jax, jax.numpy as jnp, lineax as lx
>>> import vardax as vdx
>>> vdx.InstrumentRegistry().entries
{}
>>> spec = vdx.InstrumentSpec(
...     obs_op=vdx.MaskedIdentity(),
...     mask=jnp.ones(2),
...     obs_cov_op=lx.IdentityLinearOperator(
...         jax.ShapeDtypeStruct((2,), jnp.float32)
...     ),
...     instrument_id="EMIT",
... )
>>> list(vdx.InstrumentRegistry(entries={"EMIT": spec}).entries)
['EMIT']
Source code in src/vardax/_src/obs_operators/multi_instrument.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
class InstrumentRegistry(eqx.Module):
    """Keyed lookup of ``InstrumentSpec`` by ``instrument_id``.

    Attributes:
        entries: ``dict[instrument_id, InstrumentSpec]``.

    Examples:
        An empty registry by default, then one keyed by instrument id:

        >>> import jax, jax.numpy as jnp, lineax as lx
        >>> import vardax as vdx
        >>> vdx.InstrumentRegistry().entries
        {}
        >>> spec = vdx.InstrumentSpec(
        ...     obs_op=vdx.MaskedIdentity(),
        ...     mask=jnp.ones(2),
        ...     obs_cov_op=lx.IdentityLinearOperator(
        ...         jax.ShapeDtypeStruct((2,), jnp.float32)
        ...     ),
        ...     instrument_id="EMIT",
        ... )
        >>> list(vdx.InstrumentRegistry(entries={"EMIT": spec}).entries)
        ['EMIT']
    """

    entries: dict[str, InstrumentSpec] = eqx.field(default_factory=dict)

MultiInstrumentFusion

Bases: Module

Compose per-instrument operators at the likelihood level.

__call__ returns dict[instrument_id, predicted_obs] — one array per instrument. The cost function consumes the dict and sums per-instrument terms with their respective \(R_i^{-1}\). There is no shared coordinate system; each instrument keeps its native footprint and resolution.

For strict pipekit_cycle.ObservationOperator contexts where a single observation vector + single linear operator are required, call .to_observation_operator() for a flattened wrapper.

Attributes:

Name Type Description
registry InstrumentRegistry

Per-instrument InstrumentRegistry.
weights dict[str, float] | None

Optional {instrument_id: alpha} mapping. None ⇒ uniform \(\alpha_i = 1\).

Examples:

Two instruments observing the same state, one an identity and one a linear projection:

>>> import jax, jax.numpy as jnp, lineax as lx
>>> import vardax as vdx
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((2,), jnp.float32))
>>> a = vdx.InstrumentSpec(vdx.MaskedIdentity(), jnp.ones(2), eye, "a")
>>> H = lx.MatrixLinearOperator(jnp.array([[1.0, 1.0]]))
>>> b = vdx.InstrumentSpec(vdx.LinearObs(operator=H), jnp.ones(1), eye, "b")
>>> fusion = vdx.MultiInstrumentFusion(
...     vdx.InstrumentRegistry(entries={"a": a, "b": b})
... )
>>> preds = fusion(jnp.array([1.0, 2.0]))
>>> {k: v.shape for k, v in preds.items()}
{'a': (2,), 'b': (1,)}

The flattened adapter concatenates the outputs in sorted id order:

>>> flat = fusion.to_observation_operator()
>>> flat(jnp.array([1.0, 2.0])).shape
(3,)
Source code in src/vardax/_src/obs_operators/multi_instrument.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
class MultiInstrumentFusion(eqx.Module):
    r"""Compose per-instrument operators at the likelihood level.

    ``__call__`` returns ``dict[instrument_id, predicted_obs]`` — one
    array per instrument. The cost function consumes the dict and
    sums per-instrument terms with their respective $R_i^{-1}$.
    There is no shared coordinate system; each instrument keeps its
    native footprint and resolution.

    For strict ``pipekit_cycle.ObservationOperator`` contexts where a
    single observation vector + single linear operator are required,
    call ``.to_observation_operator()`` for a flattened wrapper.

    Attributes:
        registry: Per-instrument
            [`InstrumentRegistry`][vardax.InstrumentRegistry].
        weights: Optional ``{instrument_id: alpha}`` mapping. ``None``
            ⇒ uniform $\alpha_i = 1$.

    Examples:
        Two instruments observing the same state, one an identity and one a
        linear projection:

        >>> import jax, jax.numpy as jnp, lineax as lx
        >>> import vardax as vdx
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((2,), jnp.float32))
        >>> a = vdx.InstrumentSpec(vdx.MaskedIdentity(), jnp.ones(2), eye, "a")
        >>> H = lx.MatrixLinearOperator(jnp.array([[1.0, 1.0]]))
        >>> b = vdx.InstrumentSpec(vdx.LinearObs(operator=H), jnp.ones(1), eye, "b")
        >>> fusion = vdx.MultiInstrumentFusion(
        ...     vdx.InstrumentRegistry(entries={"a": a, "b": b})
        ... )
        >>> preds = fusion(jnp.array([1.0, 2.0]))
        >>> {k: v.shape for k, v in preds.items()}
        {'a': (2,), 'b': (1,)}

        The flattened adapter concatenates the outputs in sorted id order:

        >>> flat = fusion.to_observation_operator()
        >>> flat(jnp.array([1.0, 2.0])).shape
        (3,)
    """

    registry: InstrumentRegistry
    weights: dict[str, float] | None = None

    def __call__(self, x: Float[Array, ...]) -> dict[str, Float[Array, ...]]:
        """Apply every instrument operator to the same state.

        Args:
            x: State shared by all instruments.

        Returns:
            Predicted observations keyed by instrument id.
        """
        return {
            inst_id: spec.obs_op(x) for inst_id, spec in self.registry.entries.items()
        }

    def linearize(self, x: Float[Array, ...]) -> dict[str, lx.AbstractLinearOperator]:
        """Per-instrument tangent-linear operators.

        Returns ``{instrument_id: H_i'(x)}``. The fused tangent linear
        is the block-diagonal stack of these — assembled lazily by
        the cost function or via ``to_observation_operator()``.

        Args:
            x: State at which to linearise.
        """
        return {
            inst_id: spec.obs_op.linearize(x)
            for inst_id, spec in self.registry.entries.items()
        }

    def to_observation_operator(self) -> _FlattenedMultiInstrument:
        """Adapt to the strict ``pipekit_cycle.ObservationOperator`` protocol.

        Returns a wrapper that concatenates per-instrument outputs and
        exposes a block-diagonal linear operator. Use this when the
        consumer requires a single ``(state) -> Array`` signature.
        """
        return _FlattenedMultiInstrument(fusion=self)

linearize 

linearize(
    x: Float[Array, ...],
) -> dict[str, AbstractLinearOperator]

Per-instrument tangent-linear operators.

Returns {instrument_id: H_i'(x)}. The fused tangent linear is the block-diagonal stack of these — assembled lazily by the cost function or via to_observation_operator().

Parameters:

Name Type Description Default
x Float[Array, ...]

State at which to linearise.

 required
Source code in src/vardax/_src/obs_operators/multi_instrument.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def linearize(self, x: Float[Array, ...]) -> dict[str, lx.AbstractLinearOperator]:
    """Per-instrument tangent-linear operators.

    Returns ``{instrument_id: H_i'(x)}``. The fused tangent linear
    is the block-diagonal stack of these — assembled lazily by
    the cost function or via ``to_observation_operator()``.

    Args:
        x: State at which to linearise.
    """
    return {
        inst_id: spec.obs_op.linearize(x)
        for inst_id, spec in self.registry.entries.items()
    }

to_observation_operator 

to_observation_operator() -> _FlattenedMultiInstrument

Adapt to the strict pipekit_cycle.ObservationOperator protocol.

Returns a wrapper that concatenates per-instrument outputs and exposes a block-diagonal linear operator. Use this when the consumer requires a single (state) -> Array signature.

Source code in src/vardax/_src/obs_operators/multi_instrument.py 
185
186
187
188
189
190
191
192
def to_observation_operator(self) -> _FlattenedMultiInstrument:
    """Adapt to the strict ``pipekit_cycle.ObservationOperator`` protocol.

    Returns a wrapper that concatenates per-instrument outputs and
    exposes a block-diagonal linear operator. Use this when the
    consumer requires a single ``(state) -> Array`` signature.
    """
    return _FlattenedMultiInstrument(fusion=self)