pipekit integration

pipekit-cycle orchestrates data-assimilation cycles through three runtime-checkable Protocols — ForwardModel, ObservationOperator, and AnalysisStep — and algorithm libraries plug in structurally, by matching the protocol signatures. filterax never imports pipekit: the adapters in filterax.pipekit simply lift filterax components into the right shapes, so conformance is checked by pipekit's isinstance against its runtime-checkable Protocols, with no dependency in either direction.

Two impedance mismatches are handled at the boundary. pipekit's EnsembleDACycle represents the ensemble as a Python list of member states, while filterax stacks the ensemble along axis 0 of a single array — FilterAnalysisStep converts between the two and returns the same container kind it received. And pipekit never passes a PRNG key to the analysis step, so stochastic filters (e.g. StochasticEnKF) draw from the key they were constructed with; prefer deterministic filters (ETKF, EnSRF) when cycling through pipekit, or reconstruct the filter per cycle.

Wiring an EnsembleDACycle

import jax.numpy as jnp
from pipekit_cycle import DAState, EnsembleDACycle

from filterax.filters import ETKF
from filterax.pipekit import (
    DynamicsForwardModel,
    FilterAnalysisStep,
    LinearizableObsOperator,
)

cycle = EnsembleDACycle(
    forward_model=DynamicsForwardModel(lorenz96_step, dt=0.05),
    obs_op=LinearizableObsOperator(lambda x: x[::2]),  # observe every 2nd var
    analysis_step=FilterAnalysisStep(ETKF()),
    obs_source=obs_source,  # pipekit Operator: cycle index -> observation
    n_steps=100,
    n_members=20,
)

members = [x0 + 0.1 * eps for eps in perturbations]  # pipekit's list carrier
state = DAState(obs_err_cov=0.5 * jnp.eye(20))       # R rides on the DA state
members, state = cycle(members, state)

The observation-error covariance travels on the pipekit DAState (obs_err_cov) and may be a dense (N_y, N_y) array or a lineax operator; FilterAnalysisStep coerces it to the tagged operator filterax expects.

Adapters

pipekit-cycle protocol adapters.

Re-exports the adapters from filterax._integrations._pipekit so filterax.pipekit.FilterAnalysisStep(...) is the canonical access path. filterax never imports pipekit — conformance is structural, via pipekit's runtime-checkable Protocols.

FilterAnalysisStep

Bases: Module

Expose a sequential filter as a pipekit AnalysisStep.

Satisfies pipekit_cycle.protocols.AnalysisStep structurally: __call__(forecast, obs, *, obs_op, obs_err_cov). The forecast ensemble may arrive either as pipekit's list of member states or as a stacked (Nₑ, Nₓ) array; the analysis ensemble is returned in the same container kind.

Attributes:

Name	Type	Description
filter	AbstractSequentialFilter	Any :class:filterax.AbstractSequentialFilter (ETKF, EnSRF, LETKF, ...).
analysis_kwargs	dict[str, Any]	Extra keyword arguments forwarded to filter.analysis on every cycle — e.g. state_coords / obs_coords for :class:filterax.filters.LETKF.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax.filters import ETKF
>>> from filterax.pipekit import FilterAnalysisStep
>>> step = FilterAnalysisStep(ETKF())
>>> members = [jnp.array([0.0, 0.0]), jnp.array([1.0, 1.0])]
>>> analysed = step(
...     members,
...     jnp.array([0.5, 0.5]),
...     obs_op=lambda x: x,
...     obs_err_cov=0.1 * jnp.eye(2),
... )
>>> len(analysed), analysed[0].shape
(2, (2,))

Source code in src/filterax/_integrations/_pipekit.py

class FilterAnalysisStep(eqx.Module, strict=True):
    r"""Expose a sequential filter as a pipekit ``AnalysisStep``.

    Satisfies ``pipekit_cycle.protocols.AnalysisStep`` structurally:
    ``__call__(forecast, obs, *, obs_op, obs_err_cov)``. The forecast
    ensemble may arrive either as pipekit's list of member states or as
    a stacked ``(Nₑ, Nₓ)`` array; the analysis ensemble is returned in
    the same container kind.

    Attributes:
        filter: Any :class:`filterax.AbstractSequentialFilter`
            (ETKF, EnSRF, LETKF, ...).
        analysis_kwargs: Extra keyword arguments forwarded to
            ``filter.analysis`` on every cycle — e.g. ``state_coords`` /
            ``obs_coords`` for :class:`filterax.filters.LETKF`.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax.filters import ETKF
        >>> from filterax.pipekit import FilterAnalysisStep
        >>> step = FilterAnalysisStep(ETKF())
        >>> members = [jnp.array([0.0, 0.0]), jnp.array([1.0, 1.0])]
        >>> analysed = step(
        ...     members,
        ...     jnp.array([0.5, 0.5]),
        ...     obs_op=lambda x: x,
        ...     obs_err_cov=0.1 * jnp.eye(2),
        ... )
        >>> len(analysed), analysed[0].shape
        (2, (2,))
    """

    filter: AbstractSequentialFilter
    analysis_kwargs: dict[str, Any] = eqx.field(default_factory=dict)

    def __call__(
        self,
        forecast: Any,
        obs: Float[Array, " N_y"],
        *,
        obs_op: Callable[[Float[Array, " N_x"]], Float[Array, " N_y"]],
        obs_err_cov: Any,
    ) -> Any:
        as_members = isinstance(forecast, list | tuple)
        particles = jnp.stack(list(forecast)) if as_members else forecast
        result = self.filter.analysis(
            particles,
            obs,
            obs_op,
            _as_noise_operator(obs_err_cov),
            **self.analysis_kwargs,
        )
        return list(result.particles) if as_members else result.particles

DynamicsForwardModel

Bases: Module

Expose autonomous dynamics as a pipekit ForwardModel.

Satisfies pipekit_cycle.protocols.ForwardModel structurally: step(state, dt) plus dt and state_signature attributes. pipekit does not thread absolute time into step, so the wrapped dynamics are treated as autonomous and integrated over [0, dt].

Attributes:

Name	Type	Description
dynamics	AbstractDynamics | Callable[[Float[Array, ' N_x'], Float[Array, ''], Float[Array, '']], Float[Array, ' N_x']]	A :class:filterax.AbstractDynamics (or plain (state, t0, t1) -> state callable).
dt	float	Default integration step advertised to pipekit.

Examples:

>>> import jax.numpy as jnp
>>> from filterax.pipekit import DynamicsForwardModel
>>> fwd = DynamicsForwardModel(lambda x, t0, t1: x + (t1 - t0), dt=0.5)
>>> fwd.step(jnp.array([1.0]), fwd.dt)
Array([1.5], dtype=float32)

Source code in src/filterax/_integrations/_pipekit.py

class DynamicsForwardModel(eqx.Module, strict=True):
    r"""Expose autonomous dynamics as a pipekit ``ForwardModel``.

    Satisfies ``pipekit_cycle.protocols.ForwardModel`` structurally:
    ``step(state, dt)`` plus ``dt`` and ``state_signature`` attributes.
    pipekit does not thread absolute time into ``step``, so the wrapped
    dynamics are treated as autonomous and integrated over ``[0, dt]``.

    Attributes:
        dynamics: A :class:`filterax.AbstractDynamics` (or plain
            ``(state, t0, t1) -> state`` callable).
        dt: Default integration step advertised to pipekit.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax.pipekit import DynamicsForwardModel
        >>> fwd = DynamicsForwardModel(lambda x, t0, t1: x + (t1 - t0), dt=0.5)
        >>> fwd.step(jnp.array([1.0]), fwd.dt)
        Array([1.5], dtype=float32)
    """

    dynamics: (
        AbstractDynamics
        | Callable[
            [Float[Array, " N_x"], Float[Array, ""], Float[Array, ""]],
            Float[Array, " N_x"],
        ]
    )
    dt: float = 1.0

    def step(self, state: Float[Array, " N_x"], dt: float) -> Float[Array, " N_x"]:
        """Advance one member state by ``dt`` (autonomous: ``t0 = 0``)."""
        return self.dynamics(state, jnp.asarray(0.0), jnp.asarray(dt))

    @property
    def state_signature(self) -> None:
        """No shape contract advertised; pipekit treats ``None`` as opt-out."""
        return None

state_signature: None property

No shape contract advertised; pipekit treats None as opt-out.

step(state: Float[Array, ' N_x'], dt: float) -> Float[Array, ' N_x']

Advance one member state by dt (autonomous: t0 = 0).

Source code in src/filterax/_integrations/_pipekit.py

def step(self, state: Float[Array, " N_x"], dt: float) -> Float[Array, " N_x"]:
    """Advance one member state by ``dt`` (autonomous: ``t0 = 0``)."""
    return self.dynamics(state, jnp.asarray(0.0), jnp.asarray(dt))

LinearizableObsOperator

Bases: Module

Expose an observation operator as a pipekit ObservationOperator.

Satisfies pipekit_cycle.protocols.ObservationOperator structurally: __call__(state) plus linearize(state). The tangent-linear operator is derived with :func:jax.jacfwd, so any JAX-traceable H gets an exact Jacobian for free.

Attributes:

Name	Type	Description
obs_op	AbstractObsOperator | Callable[[Float[Array, ' N_x']], Float[Array, ' N_y']]	A :class:filterax.AbstractObsOperator or plain state -> obs callable.

Examples:

>>> import jax.numpy as jnp
>>> from filterax.pipekit import LinearizableObsOperator
>>> H = LinearizableObsOperator(lambda x: x[:1] ** 2)
>>> H(jnp.array([3.0, 4.0]))
Array([9.], dtype=float32)
>>> H.linearize(jnp.array([3.0, 4.0]))  # dH/dx = [2x, 0]
Array([[6., 0.]], dtype=float32)

Source code in src/filterax/_integrations/_pipekit.py

class LinearizableObsOperator(eqx.Module, strict=True):
    r"""Expose an observation operator as a pipekit ``ObservationOperator``.

    Satisfies ``pipekit_cycle.protocols.ObservationOperator``
    structurally: ``__call__(state)`` plus ``linearize(state)``. The
    tangent-linear operator is derived with :func:`jax.jacfwd`, so any
    JAX-traceable ``H`` gets an exact Jacobian for free.

    Attributes:
        obs_op: A :class:`filterax.AbstractObsOperator` or plain
            ``state -> obs`` callable.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax.pipekit import LinearizableObsOperator
        >>> H = LinearizableObsOperator(lambda x: x[:1] ** 2)
        >>> H(jnp.array([3.0, 4.0]))
        Array([9.], dtype=float32)
        >>> H.linearize(jnp.array([3.0, 4.0]))  # dH/dx = [2x, 0]
        Array([[6., 0.]], dtype=float32)
    """

    obs_op: AbstractObsOperator | Callable[[Float[Array, " N_x"]], Float[Array, " N_y"]]

    def __call__(self, state: Float[Array, " N_x"]) -> Float[Array, " N_y"]:
        return self.obs_op(state)

    def linearize(self, state: Float[Array, " N_x"]) -> Float[Array, "N_y N_x"]:
        """Jacobian ``∂H/∂x`` at ``state``, shape ``(N_y, N_x)``."""
        return jax.jacfwd(self.obs_op)(state)

linearize(state: Float[Array, ' N_x']) -> Float[Array, 'N_y N_x']

Jacobian ∂H/∂x at state, shape (N_y, N_x).

Source code in src/filterax/_integrations/_pipekit.py

def linearize(self, state: Float[Array, " N_x"]) -> Float[Array, "N_y N_x"]:
    """Jacobian ``∂H/∂x`` at ``state``, shape ``(N_y, N_x)``."""
    return jax.jacfwd(self.obs_op)(state)