Skip to content

Protocols & Types

filterax components talk to each other through a small set of abstract interfaces (all equinox.Module subclasses, so implementations are immutable pytrees that survive jit / grad / vmap) and a matching set of typed containers. Implement the Abstract* classes to plug your own dynamics, observation operators, localizers, or inflators into the Layer-2 loops; the containers are what those loops accept and return.

Component protocols

The pieces a data-assimilation problem is assembled from: dynamics propagate a state between two times, observation operators map state space to observation space, and noise models expose a covariance operator plus a sampler.

filterax — differentiable ensemble Kalman filters and processes for JAX.

AbstractDynamics

Bases: Module

Forward model that propagates a single state from t0 to t1.

Designed for use with :func:equinox.filter_vmap to broadcast over an ensemble::

forecast = eqx.filter_vmap(dynamics)(particles, t0, t1)
Source code in src/filterax/_protocols.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class AbstractDynamics(eqx.Module, strict=True):
    r"""Forward model that propagates a single state from ``t0`` to ``t1``.

    Designed for use with :func:`equinox.filter_vmap` to broadcast over an
    ensemble::

        forecast = eqx.filter_vmap(dynamics)(particles, t0, t1)
    """

    @abc.abstractmethod
    def __call__(
        self,
        state: Float[Array, " N_x"],
        t0: Float[Array, ""],
        t1: Float[Array, ""],
    ) -> Float[Array, " N_x"]: ...

AbstractObsOperator

Bases: Module

Map a single state vector from state space to observation space.

Vectorise over an ensemble with :func:equinox.filter_vmap.

Source code in src/filterax/_protocols.py 
45
46
47
48
49
50
51
52
53
54
55
class AbstractObsOperator(eqx.Module, strict=True):
    """Map a single state vector from state space to observation space.

    Vectorise over an ensemble with :func:`equinox.filter_vmap`.
    """

    @abc.abstractmethod
    def __call__(
        self,
        state: Float[Array, " N_x"],
    ) -> Float[Array, " N_y"]: ...

AbstractNoise

Bases: Module

Noise model backed by a gaussx / lineax covariance operator.

Provides both the covariance operator (for Kalman-gain computations) and a sampling routine (for stochastic methods).

Source code in src/filterax/_protocols.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class AbstractNoise(eqx.Module, strict=True):
    """Noise model backed by a gaussx / lineax covariance operator.

    Provides both the covariance operator (for Kalman-gain computations) and
    a sampling routine (for stochastic methods).
    """

    @abc.abstractmethod
    def covariance(self) -> lx.AbstractLinearOperator:
        """Return the covariance as a linear operator."""
        ...

    @abc.abstractmethod
    def sample(
        self,
        key: PRNGKeyArray,
        shape: tuple[int, ...],
    ) -> Float[Array, "..."]:
        """Draw noise samples of the given batch shape."""
        ...

covariance() -> lx.AbstractLinearOperator abstractmethod

Return the covariance as a linear operator.

Source code in src/filterax/_protocols.py 
65
66
67
68
@abc.abstractmethod
def covariance(self) -> lx.AbstractLinearOperator:
    """Return the covariance as a linear operator."""
    ...

sample(key: PRNGKeyArray, shape: tuple[int, ...]) -> Float[Array, '...'] abstractmethod

Draw noise samples of the given batch shape.

Source code in src/filterax/_protocols.py 
70
71
72
73
74
75
76
77
@abc.abstractmethod
def sample(
    self,
    key: PRNGKeyArray,
    shape: tuple[int, ...],
) -> Float[Array, "..."]:
    """Draw noise samples of the given batch shape."""
    ...

Regularisation protocols

Localizers and inflators are the two standard fixes for small-ensemble rank deficiency; schedulers control the artificial-time step of the ensemble Kalman processes. Concrete implementations live on the Localization, Inflation, and Schedulers pages.

filterax — differentiable ensemble Kalman filters and processes for JAX.

AbstractLocalizer

Bases: Module

Covariance localization strategy applied to the Kalman gain or covariance.

Source code in src/filterax/_protocols.py 
80
81
82
83
84
85
86
87
88
class AbstractLocalizer(eqx.Module, strict=True):
    """Covariance localization strategy applied to the Kalman gain or covariance."""

    @abc.abstractmethod
    def __call__(
        self,
        cov: Float[Array, "N_x N_y"],
        coords: Any,
    ) -> Float[Array, "N_x N_y"]: ...

AbstractInflator

Bases: Module

Ensemble inflation strategy applied after analysis.

The optional forecast_particles argument lets relaxation methods (RTPS, RTPP) recover the prior spread when needed; multiplicative inflators simply ignore it. The **kwargs slot is reserved for extension parameters supplied by the L2 run loop (e.g. step= for stochastic inflators that need a fresh per-cycle PRNG key).

Source code in src/filterax/_protocols.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
class AbstractInflator(eqx.Module, strict=True):
    """Ensemble inflation strategy applied after analysis.

    The optional ``forecast_particles`` argument lets relaxation methods
    (RTPS, RTPP) recover the prior spread when needed; multiplicative
    inflators simply ignore it. The ``**kwargs`` slot is reserved for
    extension parameters supplied by the L2 run loop (e.g. ``step=``
    for stochastic inflators that need a fresh per-cycle PRNG key).
    """

    @abc.abstractmethod
    def __call__(
        self,
        particles: Float[Array, "N_e N_x"],
        forecast_particles: Float[Array, "N_e N_x"] | None = None,
        **kwargs: Any,
    ) -> Float[Array, "N_e N_x"]: ...

AbstractScheduler

Bases: Module

Step size / learning rate strategy for iterative processes.

Source code in src/filterax/_protocols.py 
110
111
112
113
114
115
116
class AbstractScheduler(eqx.Module, strict=True):
    """Step size / learning rate strategy for iterative processes."""

    @abc.abstractmethod
    def get_dt(self, state: ProcessState) -> Float[Array, ""]:
        """Return the step size for the current iteration."""
        ...

get_dt(state: ProcessState) -> Float[Array, ''] abstractmethod

Return the step size for the current iteration.

Source code in src/filterax/_protocols.py 
113
114
115
116
@abc.abstractmethod
def get_dt(self, state: ProcessState) -> Float[Array, ""]:
    """Return the step size for the current iteration."""
    ...

Algorithm protocols

The two top-level algorithm families: sequential filters expose a one-shot analysis step (see Filters), and ensemble Kalman processes expose an init / update iteration (see Processes).

filterax — differentiable ensemble Kalman filters and processes for JAX.

AbstractSequentialFilter

Bases: Module

Protocol for sequential ensemble filters.

Concrete filters implement a single analysis step; the forecast step is typically handled externally (eqx.filter_vmap of dynamics over the ensemble) and passed back in as particles.

Source code in src/filterax/_protocols.py 
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
class AbstractSequentialFilter(eqx.Module, strict=True):
    """Protocol for sequential ensemble filters.

    Concrete filters implement a single ``analysis`` step; the forecast step
    is typically handled externally (``eqx.filter_vmap`` of dynamics over the
    ensemble) and passed back in as ``particles``.
    """

    @abc.abstractmethod
    def analysis(
        self,
        particles: Float[Array, "N_e N_x"],
        obs: Float[Array, " N_y"],
        obs_op: AbstractObsOperator
        | Callable[[Float[Array, " N_x"]], Float[Array, " N_y"]],
        obs_noise: lx.AbstractLinearOperator,
        **kwargs: Any,
    ) -> AnalysisResult:
        """Assimilate an observation vector into the ensemble.

        Concrete subclasses may require additional keyword arguments
        (LETKF needs ``state_coords`` / ``obs_coords``); every concrete
        filter accepts ``**kwargs`` so the protocol signature is uniform.
        Filters that take no extra context just ignore the extras.
        """
        ...

analysis(particles: Float[Array, 'N_e N_x'], obs: Float[Array, ' N_y'], obs_op: AbstractObsOperator | Callable[[Float[Array, ' N_x']], Float[Array, ' N_y']], obs_noise: lx.AbstractLinearOperator, **kwargs: Any) -> AnalysisResult abstractmethod

Assimilate an observation vector into the ensemble.

Concrete subclasses may require additional keyword arguments (LETKF needs state_coords / obs_coords); every concrete filter accepts **kwargs so the protocol signature is uniform. Filters that take no extra context just ignore the extras.

Source code in src/filterax/_protocols.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
@abc.abstractmethod
def analysis(
    self,
    particles: Float[Array, "N_e N_x"],
    obs: Float[Array, " N_y"],
    obs_op: AbstractObsOperator
    | Callable[[Float[Array, " N_x"]], Float[Array, " N_y"]],
    obs_noise: lx.AbstractLinearOperator,
    **kwargs: Any,
) -> AnalysisResult:
    """Assimilate an observation vector into the ensemble.

    Concrete subclasses may require additional keyword arguments
    (LETKF needs ``state_coords`` / ``obs_coords``); every concrete
    filter accepts ``**kwargs`` so the protocol signature is uniform.
    Filters that take no extra context just ignore the extras.
    """
    ...

AbstractProcess

Bases: Module

Protocol for iterative ensemble Kalman processes.

Concrete processes expose init and update; the caller owns the outer iteration loop and supplies forward evaluations at each step.

Source code in src/filterax/_protocols.py 
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
class AbstractProcess(eqx.Module, strict=True):
    """Protocol for iterative ensemble Kalman processes.

    Concrete processes expose ``init`` and ``update``; the caller owns the
    outer iteration loop and supplies forward evaluations at each step.
    """

    @abc.abstractmethod
    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        """Initialise process state from a prior ensemble."""
        ...

    @abc.abstractmethod
    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **kwargs: Any,
    ) -> ProcessState:
        """Apply a single update step given forward evaluations.

        Concrete processes may extend with extra keyword arguments
        (per-step PRNG keys, hyperparameters, …); the ``**kwargs`` slot
        keeps the protocol open without forcing every process to thread
        every extension parameter.
        """
        ...

init(particles: Float[Array, 'J N_p'], obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator) -> ProcessState abstractmethod

Initialise process state from a prior ensemble.

Source code in src/filterax/_protocols.py 
154
155
156
157
158
159
160
161
162
@abc.abstractmethod
def init(
    self,
    particles: Float[Array, "J N_p"],
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
) -> ProcessState:
    """Initialise process state from a prior ensemble."""
    ...

update(state: ProcessState, forward_evals: Float[Array, 'J N_d'], **kwargs: Any) -> ProcessState abstractmethod

Apply a single update step given forward evaluations.

Concrete processes may extend with extra keyword arguments (per-step PRNG keys, hyperparameters, …); the **kwargs slot keeps the protocol open without forcing every process to thread every extension parameter.

Source code in src/filterax/_protocols.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@abc.abstractmethod
def update(
    self,
    state: ProcessState,
    forward_evals: Float[Array, "J N_d"],
    **kwargs: Any,
) -> ProcessState:
    """Apply a single update step given forward evaluations.

    Concrete processes may extend with extra keyword arguments
    (per-step PRNG keys, hyperparameters, …); the ``**kwargs`` slot
    keeps the protocol open without forcing every process to thread
    every extension parameter.
    """
    ...

State & result containers

What flows through and out of the algorithms: per-cycle filter state, the one-shot AnalysisResult, the full-trajectory AssimilationResult (and its latent-space variant), the smoothers' SmoothingResult, and the process iteration states.

filterax — differentiable ensemble Kalman filters and processes for JAX.

FilterState

Bases: Module

Running state for sequential filters.

Source code in src/filterax/_types.py 
19
20
21
22
23
class FilterState(eqx.Module, strict=True):
    """Running state for sequential filters."""

    particles: Float[Array, "N_e N_x"]
    step: Int[Array, ""]

AnalysisResult

Bases: Module

Output of any analysis / update step.

Source code in src/filterax/_types.py 
48
49
50
51
52
53
class AnalysisResult(eqx.Module, strict=True):
    """Output of any analysis / update step."""

    particles: Float[Array, "N_e N_x"]
    log_likelihood: Float[Array, ""] | None = None
    diagnostics: dict | None = None

AssimilationResult

Bases: Module

Output of a full forecast-analysis-inflate loop.

The history arrays are stacked along a leading time axis so callers can index into per-window analyses without materialising a Python list.

Source code in src/filterax/_types.py 
56
57
58
59
60
61
62
63
64
65
66
class AssimilationResult(eqx.Module, strict=True):
    """Output of a full forecast-analysis-inflate loop.

    The history arrays are stacked along a leading time axis so callers
    can index into per-window analyses without materialising a Python list.
    """

    particles: Float[Array, "N_e N_x"]
    forecast_history: Float[Array, "T N_e N_x"]
    analysis_history: Float[Array, "T N_e N_x"]
    log_likelihoods: Float[Array, " T"] | None = None

LatentAssimilationResult

Bases: Module

Output of a latent-space L2 assimilation loop.

Mirrors :class:AssimilationResult for the x-space view (particles / forecast_history / analysis_history) while carrying the latent ensembles alongside. The _z fields are the primary representation — the x-space fields are decoded views surfaced for ergonomics and for chaining into the existing smoothers and diagnostics.

Identical-axis convention to :class:AssimilationResult: histories stack along a leading T axis, particles is the terminal ensemble for follow-up forecasts.

Source code in src/filterax/_types.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class LatentAssimilationResult(eqx.Module, strict=True):
    """Output of a latent-space L2 assimilation loop.

    Mirrors :class:`AssimilationResult` for the ``x``-space view
    (``particles`` / ``forecast_history`` / ``analysis_history``)
    while carrying the latent ensembles alongside. The ``_z`` fields
    are the *primary* representation — the ``x``-space fields are
    decoded views surfaced for ergonomics and for chaining into the
    existing smoothers and diagnostics.

    Identical-axis convention to :class:`AssimilationResult`:
    histories stack along a leading ``T`` axis, ``particles`` is the
    terminal ensemble for follow-up forecasts.
    """

    particles: Float[Array, "N_e N_x"]
    forecast_history: Float[Array, "T N_e N_x"]
    analysis_history: Float[Array, "T N_e N_x"]
    log_likelihoods: Float[Array, " T"] | None
    particles_z: Float[Array, "N_e N_z"]
    forecast_history_z: Float[Array, "T N_e N_z"]
    analysis_history_z: Float[Array, "T N_e N_z"]

SmoothingResult

Bases: Module

Output of a backward-pass ensemble smoother.

smoothed_history shares the time axis of the filter history that produced it. particles is the smoothed ensemble at the most recent time (smoothed_history[-1]) — matching the :class:AssimilationResult convention so callers can chain a smoother into a follow-up forecast without re-indexing. (For pure backward smoothers the terminal smoothed ensemble equals the filter's terminal analysis by construction.)

Source code in src/filterax/_types.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
class SmoothingResult(eqx.Module, strict=True):
    """Output of a backward-pass ensemble smoother.

    ``smoothed_history`` shares the time axis of the filter history that
    produced it. ``particles`` is the smoothed ensemble at the *most
    recent* time (``smoothed_history[-1]``) — matching the
    :class:`AssimilationResult` convention so callers can chain a
    smoother into a follow-up forecast without re-indexing. (For pure
    backward smoothers the terminal smoothed ensemble equals the
    filter's terminal analysis by construction.)
    """

    particles: Float[Array, "N_e N_x"]
    smoothed_history: Float[Array, "T N_e N_x"]

ProcessState

Bases: Module

Running state for iterative EKP processes.

Source code in src/filterax/_types.py 
26
27
28
29
30
31
32
33
34
class ProcessState(eqx.Module, strict=True):
    """Running state for iterative EKP processes."""

    particles: Float[Array, "J N_p"]
    forward_evals: Float[Array, "J N_d"]
    obs: Float[Array, " N_d"]
    noise_cov: lx.AbstractLinearOperator
    step: Int[Array, ""]
    algo_time: Float[Array, ""]

UKIState

Bases: Module

Parametric state for Unscented Kalman Inversion.

UKI tracks an explicit mean and covariance rather than particles.

Source code in src/filterax/_types.py 
37
38
39
40
41
42
43
44
45
class UKIState(eqx.Module, strict=True):
    """Parametric state for Unscented Kalman Inversion.

    UKI tracks an explicit mean and covariance rather than particles.
    """

    mean: Float[Array, " N_p"]
    covariance: lx.AbstractLinearOperator
    step: Int[Array, ""]

Configuration

filterax — differentiable ensemble Kalman filters and processes for JAX.

FilterConfig

Bases: Module

Static configuration for sequential filters.

localizer and inflator are forward-referenced through their abstract base classes to avoid a circular import with :mod:filterax._protocols.

Source code in src/filterax/_types.py 
109
110
111
112
113
114
115
116
117
118
119
class FilterConfig(eqx.Module, strict=True):
    """Static configuration for sequential filters.

    ``localizer`` and ``inflator`` are forward-referenced through their
    abstract base classes to avoid a circular import with
    :mod:`filterax._protocols`.
    """

    n_ensemble: int = eqx.field(static=True)
    localizer: eqx.Module | None = None
    inflator: eqx.Module | None = None

ProcessConfig

Bases: Module

Static configuration for EKP processes.

Holds the step-size strategy and iteration cap shared by the Layer-2 :class:filterax.EKI / :class:filterax.EKS / :class:filterax.UKI wrappers. When a ProcessConfig is passed to one of those models, both fields take precedence over the module-level scheduler and the model's default iteration count.

The scheduler field is forward-referenced through :class:equinox.Module to avoid a circular import with the L1 protocols; concrete subclasses live in :mod:filterax._processes._schedulers.

Source code in src/filterax/_types.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
class ProcessConfig(eqx.Module, strict=True):
    """Static configuration for EKP processes.

    Holds the step-size strategy and iteration cap shared by the
    Layer-2 :class:`filterax.EKI` / :class:`filterax.EKS` /
    :class:`filterax.UKI` wrappers. When a ``ProcessConfig`` is passed
    to one of those models, **both** fields take precedence over the
    module-level ``scheduler`` and the model's default iteration count.

    The ``scheduler`` field is forward-referenced through
    :class:`equinox.Module` to avoid a circular import with the L1
    protocols; concrete subclasses live in
    :mod:`filterax._processes._schedulers`.
    """

    scheduler: eqx.Module
    n_iterations: int = eqx.field(static=True)