Ensemble Kalman Processes

Iterative ensemble methods for derivative-free parameter estimation. All methods share the same outer loop — the caller supplies forward-model evaluations G(θ⁽ʲ⁾) each step, the process supplies the ensemble update rule. The Layer-2 wrappers (filterax.EKI, filterax.EKS, filterax.TEKI, filterax.UKI) handle the loop for you; the artificial-time step Δtₙ is controlled by a scheduler, and the iteration states (ProcessState, UKIState, ProcessConfig) are documented with the Protocols & Types.

Picking a process

Method	Output	UQ	Forward evals / step	Use when
filterax.EKI	Collapsing point estimate	None	J	Default calibration; underdetermined OK (J ≥ p + 1)
filterax.EKS	Posterior samples	Full ensemble	J	Posterior uncertainty matters; longer runs; pass prior_mean / prior_cov for a Bayesian prior
filterax.TEKI	MAP	None (spread is a convergence diagnostic)	J	Point estimate under a Gaussian prior, without gradients
filterax.UKI	Parametric (μ, Σ)	Calibrated Σ	2 Nₚ + 1 (fixed)	Moderate Nₚ ≲ 100, deterministic, reproducible
filterax.processes.ETKI	MLE	None	J	Same as EKI — semantic alias; output-scalable variant for N_d ≫ J (current impl delegates to EKI's gaussx-Woodbury solve)
filterax.processes.GNKI	MAP + ensemble spread	Via spread	J	Well-conditioned, fast convergence, requires J > Nₚ
filterax.processes.SparseInversion	Sparse point estimate	None	J	Variable selection, sparse parameter recovery

What each model targets

For a Bayesian inverse problem with prior \(\mathcal{N}(m_0, \Sigma_0)\) and noise \(\Gamma\), the models converge to different things:

• EKI: a least-squares fit to \(y\) within the span of the initial ensemble. The prior enters only through that ensemble.
• TEKI: the MAP of \(\tfrac{1}{2}\|y - G(\theta)\|^2_{\Gamma^{-1}} + \tfrac{1}{2}\|\theta - m_0\|^2_{\Sigma_0^{-1}}\).
• EKS with prior_mean / prior_cov: samples of the posterior \(p(\theta \mid y)\). Without them, samples of the likelihood (a flat prior). Use a small fixed step (FixedScheduler(dt=0.02) or so) for sampling; EKSStableScheduler can pick steps too large for the data (#109).
• UKI: a Gaussian approximation, with the initial \((\mu_0, \Sigma_0)\) acting as the prior. Exact for a linear \(G\).

Decision tree

1. Need posterior samples? → EKS
2. Need calibrated parametric covariance? → UKI (if Nₚ ≲ 100)
3. Sparse parameters? → SparseInversion
4. Ill-posed / need prior regularisation? → TEKI
5. Well-conditioned, want fast convergence? → GNKI (if J > Nₚ)
6. Otherwise → EKI (simple, robust default)

Layer-2 API (run-loop wrappers)

These are the canonical entry points. Each composes the forward model, a Layer-1 process, and a scheduler into a single .run() call returning a ProcessResult.

filterax — differentiable ensemble Kalman filters and processes for JAX.

EKI

Bases: Module

Ensemble Kalman Inversion — full init → update run loop.

Composes a :class:forward_fn (the simulator) with a Layer-1 :class:~filterax._processes._processes.EKI and a scheduler. Returns the final parameter ensemble plus per-iteration history.

Attributes:

Name	Type	Description
forward_fn	ForwardFn	θ → G(θ) mapping. Applied with :func:jax.vmap over the ensemble.
obs	Float[Array, ' N_d']	Observation vector y.
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.
scheduler	AbstractScheduler	Step-size strategy. Defaults to :class:DataMisfitController which terminates when algo_time → 1. Ignored when config is provided — see config below.
config	ProcessConfig | None	Optional :class:ProcessConfig. When provided, both its scheduler and n_iterations fields take precedence over the module-level scheduler and the 50-iteration default. The scheduler attribute on the L2 module then acts as a fallback used only when config is None.

Source code in src/filterax/_processes/_process_models.py

class EKI(eqx.Module, strict=True):
    r"""Ensemble Kalman Inversion — full ``init → update`` run loop.

    Composes a :class:`forward_fn` (the simulator) with a Layer-1
    :class:`~filterax._processes._processes.EKI` and a scheduler. Returns the
    final parameter ensemble plus per-iteration history.

    Attributes:
        forward_fn: ``θ → G(θ)`` mapping. Applied with :func:`jax.vmap`
            over the ensemble.
        obs: Observation vector ``y``.
        noise_cov: Observation noise covariance ``Γ``.
        scheduler: Step-size strategy. Defaults to
            :class:`DataMisfitController` which terminates when
            ``algo_time → 1``. Ignored when ``config`` is provided —
            see ``config`` below.
        config: Optional :class:`ProcessConfig`. When provided, **both**
            its ``scheduler`` and ``n_iterations`` fields take precedence
            over the module-level ``scheduler`` and the ``50``-iteration
            default. The ``scheduler`` attribute on the L2 module then
            acts as a fallback used only when ``config`` is ``None``.
    """

    forward_fn: ForwardFn
    obs: Float[Array, " N_d"]
    noise_cov: lx.AbstractLinearOperator
    scheduler: AbstractScheduler = eqx.field(default_factory=DataMisfitController)
    config: ProcessConfig | None = None

    def run(self, init_particles: Float[Array, "J N_p"]) -> ProcessResult:
        sched = _resolve_scheduler(self.scheduler, self.config)
        n_iter = _resolve_n_iterations(self.config, default=50)
        return _run_ensemble(
            _EKI(scheduler=sched),
            self.forward_fn,
            init_particles,
            self.obs,
            self.noise_cov,
            n_iterations=n_iter,
            stop_at_algo_time_one=True,
        )

EKS

Bases: Module

Ensemble Kalman Sampler — full run loop.

Like :class:EKI but produces approximate posterior samples (the ensemble does not collapse). Uses :class:~filterax._processes._processes.EKS_Process with :class:EKSStableScheduler by default; runs the full n_iterations (no algo-time termination — burn-in + sampling).

Attributes:

Name	Type	Description
forward_fn	ForwardFn	θ → G(θ).
obs	Float[Array, ' N_d']	Observation vector y.
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.
scheduler	AbstractScheduler	Step-size strategy. Use :class:EKSStableScheduler. Ignored when config is provided.
seed	int	PRNG seed for the Brownian draws.
config	ProcessConfig | None	Optional :class:ProcessConfig. When provided, both its scheduler and n_iterations fields take precedence over the module-level scheduler and the 500-iteration default.
prior_mean	Float[Array, ' N_p'] | None	Prior mean m₀. With prior_cov, the ensemble samples the posterior p(θ | y) ∝ p(y | θ) 𝒩(θ; m₀, Σ₀); without them, the likelihood alone (a flat prior).
prior_cov	AbstractLinearOperator | None	Prior covariance Σ₀.

Source code in src/filterax/_processes/_process_models.py

class EKS(eqx.Module, strict=True):
    r"""Ensemble Kalman Sampler — full run loop.

    Like :class:`EKI` but produces approximate posterior samples (the
    ensemble does *not* collapse). Uses
    :class:`~filterax._processes._processes.EKS_Process` with
    :class:`EKSStableScheduler` by default; runs the full
    ``n_iterations`` (no algo-time termination — burn-in + sampling).

    Attributes:
        forward_fn: ``θ → G(θ)``.
        obs: Observation vector ``y``.
        noise_cov: Observation noise covariance ``Γ``.
        scheduler: Step-size strategy. Use :class:`EKSStableScheduler`.
            Ignored when ``config`` is provided.
        seed: PRNG seed for the Brownian draws.
        config: Optional :class:`ProcessConfig`. When provided, **both**
            its ``scheduler`` and ``n_iterations`` fields take precedence
            over the module-level ``scheduler`` and the ``500``-iteration
            default.
        prior_mean: Prior mean ``m₀``. With ``prior_cov``, the ensemble
            samples the posterior ``p(θ | y) ∝ p(y | θ) 𝒩(θ; m₀, Σ₀)``;
            without them, the likelihood alone (a flat prior).
        prior_cov: Prior covariance ``Σ₀``.
    """

    forward_fn: ForwardFn
    obs: Float[Array, " N_d"]
    noise_cov: lx.AbstractLinearOperator
    scheduler: AbstractScheduler = eqx.field(default_factory=EKSStableScheduler)
    seed: int = eqx.field(static=True, default=0)
    config: ProcessConfig | None = None
    prior_mean: Float[Array, " N_p"] | None = None
    prior_cov: lx.AbstractLinearOperator | None = None

    def run(self, init_particles: Float[Array, "J N_p"]) -> ProcessResult:
        sched = _resolve_scheduler(self.scheduler, self.config)
        n_iter = _resolve_n_iterations(self.config, default=500)
        return _run_ensemble(
            _EKS(
                scheduler=sched,
                key=self.seed,
                prior_mean=self.prior_mean,
                prior_cov=self.prior_cov,
            ),
            self.forward_fn,
            init_particles,
            self.obs,
            self.noise_cov,
            n_iterations=n_iter,
            stop_at_algo_time_one=False,
        )

TEKI

Bases: Module

Tikhonov Ensemble Kalman Inversion — full run loop to the MAP.

EKI on the augmented system \(\tilde G(\theta) = (G(\theta), \theta)\), \(\tilde y = (y, m_0)\), \(\tilde\Gamma = \mathrm{blockdiag}(\Gamma, \Sigma_0)\) (Chada, Stuart & Tong 2020). The prior enters as data, so the iteration targets the maximum a posteriori estimate

\[ \theta^* = \arg\min_\theta\; \tfrac{1}{2} \|y - G(\theta)\|^2_{\Gamma^{-1}} + \tfrac{1}{2} \|\theta - m_0\|^2_{\Sigma_0^{-1}}, \]

whereas :class:EKI converges toward a least-squares fit in the span of the initial ensemble. The MAP is a fixed point of every step, so the run does not stop at algo_time = 1; it runs n_iterations steps while the deterministic ensemble contracts onto it. The final ensemble spread is a convergence diagnostic, not a posterior uncertainty; use :class:EKS with the same prior for samples.

Attributes:

Name	Type	Description
forward_fn	ForwardFn	θ → G(θ).
obs	Float[Array, ' N_d']	Observation vector y.
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.
prior_mean	Float[Array, ' N_p']	Prior mean m₀.
prior_cov	AbstractLinearOperator	Prior covariance Σ₀.
scheduler	AbstractScheduler	Step-size strategy (default FixedScheduler(dt=1.0), one augmented Kalman update per iteration). Ignored when config is provided.
config	ProcessConfig | None	Optional :class:ProcessConfig. When provided, both its scheduler and n_iterations fields take precedence over the module-level scheduler and the 200-iteration default.

Examples:

>>> import jax, jax.numpy as jnp, lineax as lx
>>> import filterax as flx
>>> F = jnp.array([[1.0, 0.5], [-0.3, 1.2], [0.7, 0.1]])
>>> y = jnp.array([0.9, -1.4, 0.6])
>>> teki = flx.TEKI(
...     forward_fn=lambda th: F @ th,
...     obs=y,
...     noise_cov=lx.DiagonalLinearOperator(0.1 * jnp.ones(3)),
...     prior_mean=jnp.zeros(2),
...     prior_cov=lx.DiagonalLinearOperator(jnp.ones(2)),
... )
>>> init = jax.random.normal(jax.random.key(0), (50, 2))
>>> teki.run(init).mean.shape
(2,)

Source code in src/filterax/_processes/_process_models.py

class TEKI(eqx.Module, strict=True):
    r"""Tikhonov Ensemble Kalman Inversion — full run loop to the MAP.

    EKI on the augmented system $\tilde G(\theta) = (G(\theta), \theta)$,
    $\tilde y = (y, m_0)$, $\tilde\Gamma = \mathrm{blockdiag}(\Gamma, \Sigma_0)$
    (Chada, Stuart & Tong 2020). The prior enters as data, so the
    iteration targets the maximum a posteriori estimate

    $$
    \theta^* = \arg\min_\theta\;
        \tfrac{1}{2} \|y - G(\theta)\|^2_{\Gamma^{-1}}
        + \tfrac{1}{2} \|\theta - m_0\|^2_{\Sigma_0^{-1}},
    $$

    whereas :class:`EKI` converges toward a least-squares fit in the span
    of the initial ensemble. The MAP is a fixed point of every step, so
    the run does not stop at ``algo_time = 1``; it runs ``n_iterations``
    steps while the deterministic ensemble contracts onto it. The final
    ensemble spread is a convergence diagnostic, not a posterior
    uncertainty; use :class:`EKS` with the same prior for samples.

    Attributes:
        forward_fn: ``θ → G(θ)``.
        obs: Observation vector ``y``.
        noise_cov: Observation noise covariance ``Γ``.
        prior_mean: Prior mean ``m₀``.
        prior_cov: Prior covariance ``Σ₀``.
        scheduler: Step-size strategy (default ``FixedScheduler(dt=1.0)``,
            one augmented Kalman update per iteration). Ignored when
            ``config`` is provided.
        config: Optional :class:`ProcessConfig`. When provided, **both**
            its ``scheduler`` and ``n_iterations`` fields take precedence
            over the module-level ``scheduler`` and the ``200``-iteration
            default.

    Examples:
        >>> import jax, jax.numpy as jnp, lineax as lx
        >>> import filterax as flx
        >>> F = jnp.array([[1.0, 0.5], [-0.3, 1.2], [0.7, 0.1]])
        >>> y = jnp.array([0.9, -1.4, 0.6])
        >>> teki = flx.TEKI(
        ...     forward_fn=lambda th: F @ th,
        ...     obs=y,
        ...     noise_cov=lx.DiagonalLinearOperator(0.1 * jnp.ones(3)),
        ...     prior_mean=jnp.zeros(2),
        ...     prior_cov=lx.DiagonalLinearOperator(jnp.ones(2)),
        ... )
        >>> init = jax.random.normal(jax.random.key(0), (50, 2))
        >>> teki.run(init).mean.shape
        (2,)
    """

    forward_fn: ForwardFn
    obs: Float[Array, " N_d"]
    noise_cov: lx.AbstractLinearOperator
    prior_mean: Float[Array, " N_p"]
    prior_cov: lx.AbstractLinearOperator
    scheduler: AbstractScheduler = eqx.field(
        default_factory=lambda: FixedScheduler(dt=1.0)
    )
    config: ProcessConfig | None = None

    def run(self, init_particles: Float[Array, "J N_p"]) -> ProcessResult:
        sched = _resolve_scheduler(self.scheduler, self.config)
        n_iter = _resolve_n_iterations(self.config, default=200)
        return _run_ensemble(
            _TEKI(
                scheduler=sched, prior_mean=self.prior_mean, prior_cov=self.prior_cov
            ),
            self.forward_fn,
            init_particles,
            self.obs,
            self.noise_cov,
            n_iterations=n_iter,
            stop_at_algo_time_one=False,
        )

UKI

Bases: Module

Unscented Kalman Inversion — full parametric run loop.

Maintains μₙ and Σₙ rather than a random ensemble; uses 2 Nₚ + 1 deterministic sigma points per step.

Attributes:

Name	Type	Description
forward_fn	ForwardFn	θ → G(θ).
obs	Float[Array, ' N_d']	Observation vector y.
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.
scheduler	AbstractScheduler	Step-size strategy. Ignored when config is provided.
alpha,	(beta, kappa)	Unscented tuning parameters; see :func:~filterax._processes._processes.sigma_points.
config	ProcessConfig | None	Optional :class:ProcessConfig. When provided, both its scheduler and n_iterations fields take precedence over the module-level scheduler and the 50-iteration default.

Source code in src/filterax/_processes/_process_models.py

class UKI(eqx.Module, strict=True):
    r"""Unscented Kalman Inversion — full parametric run loop.

    Maintains ``μₙ`` and ``Σₙ`` rather than a random ensemble; uses
    ``2 Nₚ + 1`` deterministic sigma points per step.

    Attributes:
        forward_fn: ``θ → G(θ)``.
        obs: Observation vector ``y``.
        noise_cov: Observation noise covariance ``Γ``.
        scheduler: Step-size strategy. Ignored when ``config`` is
            provided.
        alpha, beta, kappa: Unscented tuning parameters; see
            :func:`~filterax._processes._processes.sigma_points`.
        config: Optional :class:`ProcessConfig`. When provided, **both**
            its ``scheduler`` and ``n_iterations`` fields take precedence
            over the module-level ``scheduler`` and the ``50``-iteration
            default.
    """

    forward_fn: ForwardFn
    obs: Float[Array, " N_d"]
    noise_cov: lx.AbstractLinearOperator
    scheduler: AbstractScheduler = eqx.field(default_factory=DataMisfitController)
    alpha: float = eqx.field(static=True, default=1.0)
    beta: float = eqx.field(static=True, default=2.0)
    kappa: float = eqx.field(static=True, default=0.0)
    config: ProcessConfig | None = None

    def run(
        self,
        init_mean: Float[Array, " N_p"],
        init_cov: lx.AbstractLinearOperator,
    ) -> ProcessResult:
        sched = _resolve_scheduler(self.scheduler, self.config)
        n_iter = _resolve_n_iterations(self.config, default=50)
        process = _UKI(
            scheduler=sched,
            alpha=self.alpha,
            beta=self.beta,
            kappa=self.kappa,
        )
        state = process.init_parametric(init_mean, init_cov, self.obs, self.noise_cov)
        # Drive the loop in parametric form for clarity; sample-point
        # evaluation happens once per iteration.
        from filterax._processes._processes import sigma_points

        algo_time = jnp.asarray(0.0)
        algo_times: list[Float[Array, ""]] = []
        converged = False
        actual = 0
        for _ in range(n_iter):
            points, _wm, _wc = sigma_points(
                state.mean,
                state.covariance,
                alpha=self.alpha,
                beta=self.beta,
                kappa=self.kappa,
            )
            evals = _vmap_forward(self.forward_fn, points)
            # Determine dt via the scheduler — wrap the parametric state
            # in a ProcessState shim so DataMisfitController can see the
            # forward evals.
            shim = ProcessState(
                particles=points,
                forward_evals=evals,
                obs=self.obs,
                noise_cov=self.noise_cov,
                step=state.step,
                algo_time=algo_time,
            )
            dt = sched.get_dt(shim)
            state = process.update_parametric(
                state, evals, obs=self.obs, noise_cov=self.noise_cov, dt=dt
            )
            algo_time = algo_time + dt
            algo_times.append(algo_time)
            actual += 1
            if float(algo_time) >= 1.0 - 1e-9:
                converged = True
                break

        return ProcessResult(
            particles=None,
            mean=state.mean,
            covariance=state.covariance,
            history_particles=None,
            history_algo_time=jnp.stack(algo_times) if algo_times else jnp.zeros(0),
            n_iterations=actual,
            converged=converged,
        )

ProcessResult

Bases: Module

Output of a full EKP run.

Attributes:

Name	Type	Description
particles	Float[Array, 'J N_p'] | None	Final ensemble (None for UKI — use mean / covariance instead).
mean	Float[Array, ' N_p']	Final posterior-mean estimate.
covariance	AbstractLinearOperator	Final posterior covariance (UKI's parametric state; the ensemble sample covariance for EKI / EKS).
history_particles	Float[Array, 'Tp1 J N_p'] | None	Particles at every iteration, shape (T+1, J, Nₚ) (initial + each post-update state).
history_algo_time	Float[Array, ' T']	algo_time after each iteration, shape (T,).
n_iterations	int	Actual number of iterations run.
converged	bool	True when the scheduler reached algo_time ≥ 1 before n_iterations was exhausted.

Source code in src/filterax/_processes/_process_models.py

class ProcessResult(eqx.Module, strict=True):
    r"""Output of a full EKP run.

    Attributes:
        particles: Final ensemble (``None`` for UKI — use ``mean`` /
            ``covariance`` instead).
        mean: Final posterior-mean estimate.
        covariance: Final posterior covariance (UKI's parametric state;
            the ensemble sample covariance for EKI / EKS).
        history_particles: Particles at every iteration, shape
            ``(T+1, J, Nₚ)`` (initial + each post-update state).
        history_algo_time: ``algo_time`` after each iteration, shape
            ``(T,)``.
        n_iterations: Actual number of iterations run.
        converged: ``True`` when the scheduler reached ``algo_time ≥ 1``
            before ``n_iterations`` was exhausted.
    """

    particles: Float[Array, "J N_p"] | None
    mean: Float[Array, " N_p"]
    covariance: lx.AbstractLinearOperator
    history_particles: Float[Array, "Tp1 J N_p"] | None
    history_algo_time: Float[Array, " T"]
    n_iterations: int = eqx.field(static=True)
    converged: bool = eqx.field(static=True)

Layer-1 components (advanced)

Drop down to these when you need per-step control of the loop (custom convergence checks, JAX scan integration, etc.). All implement AbstractProcess.

Layer-1 ensemble Kalman process components.

Re-exports the L1 process classes from :mod:filterax._processes._processes so the canonical filterax.processes.EKI() access path matches the design docs. The Layer-2 filterax.EKI / filterax.EKS / filterax.UKI wrappers live at the top level and own the full init → forward eval → update loop.

EKI

Bases: AbstractProcess

Ensemble Kalman Inversion (Iglesias, Law & Stuart 2013).

Iterative ensemble method that drives θ⁽ʲ⁾ toward a MAP / regularised-least-squares solution. The update for each member:

\[ \theta^{(j)}_{n+1} = \theta^{(j)}_n + C^{\theta G}_n \big(C^{G G}_n + \Delta t_n^{-1} \Gamma\big)^{-1} \big(y - G(\theta^{(j)}_n)\big) \]

The ensemble collapses to a point as algo_time → 1: spread → 0, no posterior uncertainty (see :class:EKS_Process if you need samples). Works in the underdetermined regime J < Nₚ.

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size strategy. Use :class:~filterax.DataMisfitController for the standard adaptive recipe.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import FixedScheduler
>>> from filterax.processes import EKI
>>> G = jnp.array([[1.0, 0.5], [-0.3, 1.2]])
>>> obs = jnp.array([0.5, 0.9])
>>> noise = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
>>> eki = EKI(scheduler=FixedScheduler(dt=1.0))
>>> particles = jax.random.normal(jax.random.key(0), (4, 2))
>>> state = eki.init(particles, obs, noise)
>>> evals = jax.vmap(lambda theta: G @ theta)(state.particles)
>>> state = eki.update(state, evals)
>>> state.particles.shape, int(state.step)
((4, 2), 1)

Source code in src/filterax/_processes/_processes.py

class EKI(AbstractProcess, strict=True):
    r"""Ensemble Kalman Inversion (Iglesias, Law & Stuart 2013).

    Iterative ensemble method that drives ``θ⁽ʲ⁾`` toward a MAP /
    regularised-least-squares solution. The update for each member:

    $$
    \theta^{(j)}_{n+1} = \theta^{(j)}_n + C^{\theta G}_n
        \big(C^{G G}_n + \Delta t_n^{-1} \Gamma\big)^{-1}
        \big(y - G(\theta^{(j)}_n)\big)
    $$

    The ensemble collapses to a point as ``algo_time → 1``: spread → 0,
    no posterior uncertainty (see :class:`EKS_Process` if you need
    samples). Works in the underdetermined regime ``J < Nₚ``.

    Attributes:
        scheduler: Step-size strategy. Use
            :class:`~filterax.DataMisfitController` for the standard
            adaptive recipe.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import FixedScheduler
        >>> from filterax.processes import EKI
        >>> G = jnp.array([[1.0, 0.5], [-0.3, 1.2]])
        >>> obs = jnp.array([0.5, 0.9])
        >>> noise = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> eki = EKI(scheduler=FixedScheduler(dt=1.0))
        >>> particles = jax.random.normal(jax.random.key(0), (4, 2))
        >>> state = eki.init(particles, obs, noise)
        >>> evals = jax.vmap(lambda theta: G @ theta)(state.particles)
        >>> state = eki.update(state, evals)
        >>> state.particles.shape, int(state.step)
        ((4, 2), 1)
    """

    scheduler: AbstractScheduler

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        check_ensemble_size(particles.shape[0])
        N_d = obs.shape[0]
        # Forward evaluations are filled in on the first update(); seed
        # with zeros so the state shape is JIT-stable.
        zero_evals = jnp.zeros((particles.shape[0], N_d), dtype=particles.dtype)
        return ProcessState(
            particles=particles,
            forward_evals=zero_evals,
            obs=obs,
            noise_cov=noise_cov,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **_: Any,
    ) -> ProcessState:
        # Drop the fresh evals into state so the scheduler's misfit
        # computation sees them.
        state_now = _with_evals(state, forward_evals)
        dt = self.scheduler.get_dt(state_now)
        delta = _eki_delta(
            state.particles, forward_evals, state.obs, state.noise_cov, dt
        )
        particles_new = state.particles + delta
        return ProcessState(
            particles=particles_new,
            forward_evals=forward_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

EKS_Process

Bases: AbstractProcess

Ensemble Kalman Sampler / ALDI (Garbuno-Inigo et al. 2020).

Interacting-Langevin ensemble that does not collapse — at stationarity the particles are approximate samples from the posterior p(θ | y). The per-step update is the EKI drift plus a finite-sample mean correction and an ensemble-preconditioned Brownian noise:

\[ \mathrm{d}\theta^{(j)} = C^{\theta G} \Gamma^{-1} \big(y - G^{(j)}\big)\, \mathrm{d}t - C^{\theta\theta} \Sigma_0^{-1} \big(\theta^{(j)} - m_0\big)\, \mathrm{d}t - (N_p + 1)\, J^{-1} \big(\theta^{(j)} - \bar{\theta}\big)\, \mathrm{d}t + \sqrt{2\, C^{\theta\theta}}\; \mathrm{d}W^{(j)} \]

The prior term is present only when prior_mean and prior_cov are given; without them the stationary law is the likelihood alone (a flat prior), which is the behaviour before the prior was added. The prior term is taken implicitly in θ, as in Garbuno-Inigo et al. (2020, §4), so a prior much tighter than the ensemble contracts the particles toward m₀ for any Δt instead of amplifying them.

The discrete-time recipe used here matches Garbuno-Inigo §4: an EKI-style mean update (with Δt⁻¹ Γ tempering), the (Nₚ + 1)/J drift term, and a noise term scaled by the ensemble-space sqrt of Cᶿᶿ (which avoids the dense Nₚ × Nₚ Cholesky — see comments below).

Use :class:~filterax.EKSStableScheduler to keep Δt inside the stability region.

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size controller. Prefer EKSStableScheduler.
key	PRNGKeyArray	Base PRNG key for the Brownian draws. Each step uses jr.fold_in(key, state.step) so successive updates are independent without manual key splitting.
prior_mean	Float[Array, ' N_p'] | None	Prior mean m₀ ∈ ℝ^{Nₚ}, or None for a flat prior. Give it together with prior_cov.
prior_cov	AbstractLinearOperator | None	Prior covariance Σ₀, or None for a flat prior. Σ₀⁻¹ is applied through :func:gaussx.solve_rows, never formed densely.

Source code in src/filterax/_processes/_processes.py

class EKS_Process(AbstractProcess, strict=True):
    r"""Ensemble Kalman Sampler / ALDI (Garbuno-Inigo et al. 2020).

    Interacting-Langevin ensemble that does **not** collapse — at
    stationarity the particles are approximate samples from the
    posterior ``p(θ | y)``. The per-step update is the EKI drift plus a
    finite-sample mean correction and an ensemble-preconditioned
    Brownian noise:

    $$
    \mathrm{d}\theta^{(j)} =
        C^{\theta G} \Gamma^{-1} \big(y - G^{(j)}\big)\, \mathrm{d}t
        - C^{\theta\theta} \Sigma_0^{-1} \big(\theta^{(j)} - m_0\big)\,
          \mathrm{d}t
        - (N_p + 1)\, J^{-1} \big(\theta^{(j)} - \bar{\theta}\big)\,
          \mathrm{d}t
        + \sqrt{2\, C^{\theta\theta}}\; \mathrm{d}W^{(j)}
    $$

    The prior term is present only when ``prior_mean`` and ``prior_cov``
    are given; without them the stationary law is the likelihood alone
    (a flat prior), which is the behaviour before the prior was added.
    The prior term is taken implicitly in θ, as in Garbuno-Inigo et al.
    (2020, §4), so a prior much tighter than the ensemble contracts the
    particles toward ``m₀`` for any ``Δt`` instead of amplifying them.

    The discrete-time recipe used here matches Garbuno-Inigo §4: an
    EKI-style mean update (with ``Δt⁻¹ Γ`` tempering), the
    ``(Nₚ + 1)/J`` drift term, and a noise term scaled by the
    *ensemble-space* sqrt of ``Cᶿᶿ`` (which avoids the dense ``Nₚ × Nₚ``
    Cholesky — see comments below).

    Use :class:`~filterax.EKSStableScheduler` to keep ``Δt`` inside the
    stability region.

    Attributes:
        scheduler: Step-size controller. Prefer EKSStableScheduler.
        key: Base PRNG key for the Brownian draws. Each step uses
            ``jr.fold_in(key, state.step)`` so successive updates are
            independent without manual key splitting.
        prior_mean: Prior mean ``m₀ ∈ ℝ^{Nₚ}``, or ``None`` for a flat
            prior. Give it together with ``prior_cov``.
        prior_cov: Prior covariance ``Σ₀``, or ``None`` for a flat prior.
            ``Σ₀⁻¹`` is applied through :func:`gaussx.solve_rows`, never
            formed densely.
    """

    scheduler: AbstractScheduler
    key: PRNGKeyArray
    prior_mean: Float[Array, " N_p"] | None
    prior_cov: lx.AbstractLinearOperator | None

    def __init__(
        self,
        scheduler: AbstractScheduler,
        key: PRNGKeyArray | int = 0,
        prior_mean: Float[Array, " N_p"] | None = None,
        prior_cov: lx.AbstractLinearOperator | None = None,
    ):
        if (prior_mean is None) != (prior_cov is None):
            raise ValueError(
                "EKS_Process needs both prior_mean and prior_cov, or neither."
            )
        self.scheduler = scheduler
        self.key = jr.PRNGKey(key) if isinstance(key, int) else key
        self.prior_mean = prior_mean
        self.prior_cov = prior_cov

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        check_ensemble_size(particles.shape[0])
        N_d = obs.shape[0]
        zero_evals = jnp.zeros((particles.shape[0], N_d), dtype=particles.dtype)
        return ProcessState(
            particles=particles,
            forward_evals=zero_evals,
            obs=obs,
            noise_cov=noise_cov,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **_: Any,
    ) -> ProcessState:
        state_now = _with_evals(state, forward_evals)
        dt = self.scheduler.get_dt(state_now)
        J, N_p = state.particles.shape

        # 1. EKI-style drift toward the data.
        drift_data = _eki_delta(
            state.particles, forward_evals, state.obs, state.noise_cov, dt
        )

        # 2. Mean-pull term ``− (Nₚ + 1) / J (θ − θ̄) Δt``.
        anom = ensemble_anomalies(state.particles)
        drift_pull = -((N_p + 1) / J) * dt * anom

        # 3. Ensemble-preconditioned noise:
        #     η⁽ʲ⁾ = √(2 Δt) · sqrt(Cᶿᶿ) · ξ⁽ʲ⁾,    ξ ~ 𝒩(0, I_{Nₚ}).
        # Using the ensemble sqrt ``sqrt(Cᶿᶿ) = X′ᵀ / √(J − 1)`` (i.e.
        # the LowRankUpdate's ``U`` factor) lets us draw the noise in
        # ensemble space (J-dimensional) and lift it: cost ``O(J²·Nₚ)``
        # instead of a ``Nₚ × Nₚ`` dense Cholesky.
        step_key = jr.fold_in(self.key, state.step)
        xi = jr.normal(step_key, (J, J), dtype=state.particles.dtype)
        # The ensemble sqrt has shape (Nₚ, J): U = X′ᵀ / √(J−1).
        # So η = √(2 Δt) · ξ U ᵀ, applied row-wise: η⁽ʲ⁾ = √(2 Δt) Σₖ ξⱼₖ Uᵀₖ
        # which collapses to ``ξ @ X′ / √(J−1)``.
        noise = (
            jnp.sqrt(2.0 * dt) * einx.dot("j k, k p -> j p", xi, anom) / jnp.sqrt(J - 1)
        )

        particles_new = state.particles + drift_data + drift_pull + noise

        # 4. Prior drift ``− Δt Cᶿᶿ Σ₀⁻¹ (θ⁽ʲ⁾ − m₀)``, taken implicitly in
        #    θ (Garbuno-Inigo et al. 2020, §4) so a prior tighter than the
        #    ensemble cannot destabilise the step:
        #        θ ← θ* − (I + Δt C Σ₀⁻¹)⁻¹ Δt C Σ₀⁻¹ (θ* − m₀),
        #    θ* being the explicit update above. With C = U Uᵀ,
        #    U = X′ᵀ / √(J − 1), the push-through identity gives
        #        (I + Δt U Uᵀ Σ₀⁻¹)⁻¹ Δt U Uᵀ Σ₀⁻¹
        #            = Δt U (I_J + Δt Uᵀ Σ₀⁻¹ U)⁻¹ Uᵀ Σ₀⁻¹,
        #    so only a J × J system is solved and no Nₚ × Nₚ matrix is formed.
        if self.prior_mean is not None and self.prior_cov is not None:
            offset = einx.subtract(
                "j p, p -> j p", particles_new, self.prior_mean
            )  # (J, Nₚ)
            precision_offset = gaussx.solve_rows(self.prior_cov, offset)  # (J, Nₚ)
            precision_anom = gaussx.solve_rows(self.prior_cov, anom)  # (J, Nₚ)
            # Uᵀ Σ₀⁻¹ U = X′ Σ₀⁻¹ X′ᵀ / (J − 1)
            gram = einx.dot("k p, l p -> k l", anom, precision_anom) / (J - 1)
            system = jnp.eye(J, dtype=gram.dtype) + dt * gram  # (J, J)
            # Uᵀ Σ₀⁻¹ (θ*ⱼ − m₀), one column per member
            rhs = einx.dot("k p, j p -> k j", anom, precision_offset) / jnp.sqrt(J - 1)
            coeff = jnp.linalg.solve(system, rhs)  # (J, J)
            correction = dt * einx.dot("k j, k p -> j p", coeff, anom) / jnp.sqrt(J - 1)
            particles_new = particles_new - correction
        return ProcessState(
            particles=particles_new,
            forward_evals=forward_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

UKI

Bases: AbstractProcess

Unscented Kalman Inversion (Huang, Schneider & Stuart 2022).

Parametric inversion that maintains a Gaussian belief θ ~ 𝒩(μₙ, Σₙ) instead of an explicit ensemble. Each step generates 2 Nₚ + 1 sigma points, evaluates the forward model on them, and applies a Kalman update on the unscented statistics:

\[ \begin{aligned} \mu_{n+1} &= \mu_n + C^{\theta y} S_n^{-1} \big(y - \hat{y}_0\big) \\ \Sigma_{n+1} &= \Sigma_n - C^{\theta y} S_n^{-1} C^{\theta y \top}, \qquad S_n = C^{yy} + \Delta t^{-1} \Gamma \end{aligned} \]

Deterministic (no sampling noise), no ensemble collapse, and the covariance Σₙ gives a calibrated posterior in the linear-Gaussian case.

The L1 protocol uses :class:ProcessState for compatibility with DataMisfitController — the field particles carries the sigma points (last 2 Nₚ + 1 evaluations); the parametric state lives in :class:UKIState and is returned by :meth:init_parametric / :meth:update_parametric for callers that prefer the explicit form.

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size strategy.
alpha	float	Unscented spread parameter; usually small (e.g. 1e-3) for highly nonlinear G, 1.0 for benign problems.
beta	float	Prior moment weighting; 2.0 is optimal for Gaussian priors.
kappa	float	Secondary scaling, usually 0 or 3 − Nₚ.

Source code in src/filterax/_processes/_processes.py

class UKI(AbstractProcess, strict=True):
    r"""Unscented Kalman Inversion (Huang, Schneider & Stuart 2022).

    Parametric inversion that maintains a Gaussian belief
    ``θ ~ 𝒩(μₙ, Σₙ)`` instead of an explicit ensemble. Each step
    generates ``2 Nₚ + 1`` sigma points, evaluates the forward model on
    them, and applies a Kalman update on the unscented statistics:

    $$
    \begin{aligned}
    \mu_{n+1} &= \mu_n
        + C^{\theta y} S_n^{-1} \big(y - \hat{y}_0\big) \\
    \Sigma_{n+1} &= \Sigma_n
        - C^{\theta y} S_n^{-1} C^{\theta y \top},
        \qquad S_n = C^{yy} + \Delta t^{-1} \Gamma
    \end{aligned}
    $$

    Deterministic (no sampling noise), no ensemble collapse, and the
    covariance ``Σₙ`` gives a calibrated posterior in the linear-Gaussian
    case.

    The L1 protocol uses :class:`ProcessState` for compatibility with
    ``DataMisfitController`` — the field ``particles`` carries the sigma
    points (last 2 Nₚ + 1 evaluations); the parametric state lives in
    :class:`UKIState` and is returned by :meth:`init_parametric` /
    :meth:`update_parametric` for callers that prefer the explicit form.

    Attributes:
        scheduler: Step-size strategy.
        alpha: Unscented spread parameter; usually small (e.g. ``1e-3``)
            for highly nonlinear ``G``, ``1.0`` for benign problems.
        beta: Prior moment weighting; ``2.0`` is optimal for Gaussian
            priors.
        kappa: Secondary scaling, usually ``0`` or ``3 − Nₚ``.
    """

    scheduler: AbstractScheduler
    alpha: float = eqx.field(static=True, default=1.0)
    beta: float = eqx.field(static=True, default=2.0)
    kappa: float = eqx.field(static=True, default=0.0)

    # --- parametric API ---------------------------------------------------

    def init_parametric(
        self,
        mean: Float[Array, " N_p"],
        covariance: lx.AbstractLinearOperator,
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> UKIState:
        """Initialise the parametric ``(μ, Σ)`` state."""
        return UKIState(
            mean=mean,
            covariance=covariance,
            step=jnp.asarray(0, dtype=jnp.int32),
        )

    def update_parametric(
        self,
        state: UKIState,
        forward_evals: Float[Array, "twoNp1 N_d"],
        *,
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
        dt: Float[Array, ""],
    ) -> UKIState:
        r"""One UKI step on the parametric state.

        Args:
            state: Current ``(μₙ, Σₙ)`` belief.
            forward_evals: Forward-model evaluations at the
                ``2 Nₚ + 1`` sigma points (in the order returned by
                :func:`sigma_points`).
            obs: Observation vector.
            noise_cov: Observation noise covariance ``Γ``.
            dt: Step size for this iteration.

        Returns:
            Updated :class:`UKIState`.
        """
        points, w_mean, w_cov = sigma_points(
            state.mean,
            state.covariance,
            alpha=self.alpha,
            beta=self.beta,
            kappa=self.kappa,
        )
        # Weighted obs-space mean.
        y_mean = einx.dot("k, k d -> d", w_mean, forward_evals)
        # Weighted cross-covariance Cᶿʸ = Σₖ Wcₖ (χₖ − μ) (ŷₖ − ȳ)ᵀ.
        theta_anom = einx.subtract("k p, p -> k p", points, state.mean)
        y_anom = einx.subtract("k d, d -> k d", forward_evals, y_mean)
        weighted_theta = einx.multiply("k p, k -> k p", theta_anom, w_cov)
        C_theta_y = einx.dot("k p, k d -> p d", weighted_theta, y_anom)
        # Sᵧᵧ = Σₖ Wcₖ (ŷₖ − ȳ)(ŷₖ − ȳ)ᵀ + Δt⁻¹ Γ.
        weighted_y = einx.multiply("k d, k -> k d", y_anom, w_cov)
        S_dense = einx.dot("k a, k b -> a b", weighted_y, y_anom)
        # Tempered Kalman solve.
        S_tempered = S_dense + _safe_inv_dt(dt) * noise_cov.as_matrix()
        innovation = obs - y_mean
        K = jnp.linalg.solve(S_tempered.T, C_theta_y.T).T  # (Nₚ, N_d)

        mean_new = state.mean + einx.dot("p d, d -> p", K, innovation)
        # Σ_{n+1} = Σ̂ − K S_tempered Kᵀ; Δt enters once, through Δt⁻¹ Γ in S.
        Sigma_dense = state.covariance.as_matrix()
        K_S = einx.dot("p a, a b -> p b", K, S_tempered)  # (Nₚ, N_d)
        Sigma_new = Sigma_dense - einx.dot("p d, q d -> p q", K_S, K)
        Sigma_new = 0.5 * (Sigma_new + Sigma_new.T)
        return UKIState(
            mean=mean_new,
            covariance=lx.MatrixLinearOperator(
                Sigma_new, tags=lx.positive_semidefinite_tag
            ),
            step=state.step + 1,
        )

    # --- AbstractProcess (ensemble-typed) API ---------------------------------

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        r"""Initialise from an ensemble proxy for the parametric belief.

        Computes the empirical mean and covariance of ``particles`` and
        stores them in the leading sigma point block. Provided so UKI
        slots into the same ``init / update`` loop as the
        ensemble-typed processes; for the cleaner parametric API use
        :meth:`init_parametric` / :meth:`update_parametric`.
        """
        check_ensemble_size(particles.shape[0])
        mean = ensemble_mean(particles)
        cov_op = gaussx.ensemble_covariance(particles, bessel=True)
        cov = lx.MatrixLinearOperator(
            cov_op.as_matrix(), tags=lx.positive_semidefinite_tag
        )
        points, _wm0, _wc0 = sigma_points(
            mean, cov, alpha=self.alpha, beta=self.beta, kappa=self.kappa
        )
        N_d = obs.shape[0]
        zero_evals = jnp.zeros((points.shape[0], N_d), dtype=particles.dtype)
        return ProcessState(
            particles=points,
            forward_evals=zero_evals,
            obs=obs,
            noise_cov=noise_cov,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "twoNp1 N_d"],
        **_: Any,
    ) -> ProcessState:
        state_now = _with_evals(state, forward_evals)
        dt = self.scheduler.get_dt(state_now)
        # Reconstruct the parametric state from the sigma points stored
        # on ``state.particles``: the first point is the mean, and the
        # sample covariance of the remaining 2 Nₚ points recovers Σ /
        # (N_p + λ).
        points = state.particles
        N_p = (points.shape[0] - 1) // 2
        mean = points[0]
        # Recover Σ from the ``Nₚ`` positive-offset sigma points alone:
        # by construction (see :func:`sigma_points`) we have
        #   χ⁺ⱼ − μ = c · [√Σ]_{:,j},   c = √(Nₚ + λ),
        # so  Σ_k (χ⁺ⱼ − μ)(χ⁺ⱼ − μ)ᵀ = c² Σ = (Nₚ + λ) Σ.
        # Dividing by ``(Nₚ + λ)`` therefore returns Σ exactly. The
        # negative-offset block is the reflection and contributes
        # nothing new, so we skip it.
        lam = self.alpha**2 * (N_p + self.kappa) - N_p
        scale = N_p + lam
        offsets = points[1 : 1 + N_p] - mean[None, :]  # (Nₚ, Nₚ)
        Sigma_mat = einx.dot("k p, k q -> p q", offsets, offsets) / scale
        Sigma_mat = 0.5 * (Sigma_mat + Sigma_mat.T)
        cov_op = lx.MatrixLinearOperator(Sigma_mat, tags=lx.positive_semidefinite_tag)
        parametric = UKIState(
            mean=mean,
            covariance=cov_op,
            step=state.step,
        )
        new_parametric = self.update_parametric(
            parametric, forward_evals, obs=state.obs, noise_cov=state.noise_cov, dt=dt
        )
        new_points, _wm, _wc = sigma_points(
            new_parametric.mean,
            new_parametric.covariance,
            alpha=self.alpha,
            beta=self.beta,
            kappa=self.kappa,
        )
        return ProcessState(
            particles=new_points,
            forward_evals=forward_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

init_parametric(mean: Float[Array, ' N_p'], covariance: lx.AbstractLinearOperator, obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator) -> UKIState

Initialise the parametric (μ, Σ) state.

Source code in src/filterax/_processes/_processes.py

def init_parametric(
    self,
    mean: Float[Array, " N_p"],
    covariance: lx.AbstractLinearOperator,
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
) -> UKIState:
    """Initialise the parametric ``(μ, Σ)`` state."""
    return UKIState(
        mean=mean,
        covariance=covariance,
        step=jnp.asarray(0, dtype=jnp.int32),
    )

update_parametric(state: UKIState, forward_evals: Float[Array, 'twoNp1 N_d'], *, obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator, dt: Float[Array, '']) -> UKIState

One UKI step on the parametric state.

Parameters:

Name	Type	Description	Default
state	UKIState	Current (μₙ, Σₙ) belief.	required
forward_evals	Float[Array, 'twoNp1 N_d']	Forward-model evaluations at the 2 Nₚ + 1 sigma points (in the order returned by :func:sigma_points).	required
obs	Float[Array, ' N_d']	Observation vector.	required
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.	required
dt	Float[Array, '']	Step size for this iteration.	required

Returns:

Name	Type	Description
Updated	UKIState	class:UKIState.

Source code in src/filterax/_processes/_processes.py

def update_parametric(
    self,
    state: UKIState,
    forward_evals: Float[Array, "twoNp1 N_d"],
    *,
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
    dt: Float[Array, ""],
) -> UKIState:
    r"""One UKI step on the parametric state.

    Args:
        state: Current ``(μₙ, Σₙ)`` belief.
        forward_evals: Forward-model evaluations at the
            ``2 Nₚ + 1`` sigma points (in the order returned by
            :func:`sigma_points`).
        obs: Observation vector.
        noise_cov: Observation noise covariance ``Γ``.
        dt: Step size for this iteration.

    Returns:
        Updated :class:`UKIState`.
    """
    points, w_mean, w_cov = sigma_points(
        state.mean,
        state.covariance,
        alpha=self.alpha,
        beta=self.beta,
        kappa=self.kappa,
    )
    # Weighted obs-space mean.
    y_mean = einx.dot("k, k d -> d", w_mean, forward_evals)
    # Weighted cross-covariance Cᶿʸ = Σₖ Wcₖ (χₖ − μ) (ŷₖ − ȳ)ᵀ.
    theta_anom = einx.subtract("k p, p -> k p", points, state.mean)
    y_anom = einx.subtract("k d, d -> k d", forward_evals, y_mean)
    weighted_theta = einx.multiply("k p, k -> k p", theta_anom, w_cov)
    C_theta_y = einx.dot("k p, k d -> p d", weighted_theta, y_anom)
    # Sᵧᵧ = Σₖ Wcₖ (ŷₖ − ȳ)(ŷₖ − ȳ)ᵀ + Δt⁻¹ Γ.
    weighted_y = einx.multiply("k d, k -> k d", y_anom, w_cov)
    S_dense = einx.dot("k a, k b -> a b", weighted_y, y_anom)
    # Tempered Kalman solve.
    S_tempered = S_dense + _safe_inv_dt(dt) * noise_cov.as_matrix()
    innovation = obs - y_mean
    K = jnp.linalg.solve(S_tempered.T, C_theta_y.T).T  # (Nₚ, N_d)

    mean_new = state.mean + einx.dot("p d, d -> p", K, innovation)
    # Σ_{n+1} = Σ̂ − K S_tempered Kᵀ; Δt enters once, through Δt⁻¹ Γ in S.
    Sigma_dense = state.covariance.as_matrix()
    K_S = einx.dot("p a, a b -> p b", K, S_tempered)  # (Nₚ, N_d)
    Sigma_new = Sigma_dense - einx.dot("p d, q d -> p q", K_S, K)
    Sigma_new = 0.5 * (Sigma_new + Sigma_new.T)
    return UKIState(
        mean=mean_new,
        covariance=lx.MatrixLinearOperator(
            Sigma_new, tags=lx.positive_semidefinite_tag
        ),
        step=state.step + 1,
    )

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

Initialise from an ensemble proxy for the parametric belief.

Computes the empirical mean and covariance of particles and stores them in the leading sigma point block. Provided so UKI slots into the same init / update loop as the ensemble-typed processes; for the cleaner parametric API use :meth:init_parametric / :meth:update_parametric.

Source code in src/filterax/_processes/_processes.py

def init(
    self,
    particles: Float[Array, "J N_p"],
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
) -> ProcessState:
    r"""Initialise from an ensemble proxy for the parametric belief.

    Computes the empirical mean and covariance of ``particles`` and
    stores them in the leading sigma point block. Provided so UKI
    slots into the same ``init / update`` loop as the
    ensemble-typed processes; for the cleaner parametric API use
    :meth:`init_parametric` / :meth:`update_parametric`.
    """
    check_ensemble_size(particles.shape[0])
    mean = ensemble_mean(particles)
    cov_op = gaussx.ensemble_covariance(particles, bessel=True)
    cov = lx.MatrixLinearOperator(
        cov_op.as_matrix(), tags=lx.positive_semidefinite_tag
    )
    points, _wm0, _wc0 = sigma_points(
        mean, cov, alpha=self.alpha, beta=self.beta, kappa=self.kappa
    )
    N_d = obs.shape[0]
    zero_evals = jnp.zeros((points.shape[0], N_d), dtype=particles.dtype)
    return ProcessState(
        particles=points,
        forward_evals=zero_evals,
        obs=obs,
        noise_cov=noise_cov,
        step=jnp.asarray(0, dtype=jnp.int32),
        algo_time=jnp.asarray(0.0, dtype=particles.dtype),
    )

ETKI

Bases: AbstractProcess

Ensemble Transform Kalman Inversion.

The "transform" variant of EKI. Conceptually identical to :class:EKI on the analysis ensemble (in batch form both write the same linear-Gaussian posterior); the intended distinction is that ETKI rewrites the per-step solve in the J × J ensemble subspace, yielding O(J² N_d) cost instead of O(N_d³) when N_d ≫ J.

The current implementation defers to the same Bessel-corrected EKI delta as :class:EKI — gaussx routes the inner solve through the Woodbury identity, so the dense (N_d, N_d) block is not formed even though the outer call signature is identical. A future optimisation can specialise the transform path explicitly when N_d is extreme; for now the two classes coexist so users can select the conceptually-appropriate one without paying a different cost.

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size strategy.

Source code in src/filterax/_processes/_processes.py

class ETKI(AbstractProcess, strict=True):
    r"""Ensemble Transform Kalman Inversion.

    The "transform" variant of EKI. Conceptually identical to
    :class:`EKI` on the analysis ensemble (in batch form both write the
    same linear-Gaussian posterior); the *intended* distinction is that
    ETKI rewrites the per-step solve in the ``J × J`` ensemble subspace,
    yielding ``O(J² N_d)`` cost instead of ``O(N_d³)`` when ``N_d ≫ J``.

    The current implementation defers to the same Bessel-corrected EKI
    delta as :class:`EKI` — gaussx routes the inner solve through the
    Woodbury identity, so the dense ``(N_d, N_d)`` block is *not* formed
    even though the outer call signature is identical. A future
    optimisation can specialise the transform path explicitly when
    ``N_d`` is extreme; for now the two classes coexist so users can
    select the conceptually-appropriate one without paying a different
    cost.

    Attributes:
        scheduler: Step-size strategy.
    """

    scheduler: AbstractScheduler

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        check_ensemble_size(particles.shape[0])
        N_d = obs.shape[0]
        zero_evals = jnp.zeros((particles.shape[0], N_d), dtype=particles.dtype)
        return ProcessState(
            particles=particles,
            forward_evals=zero_evals,
            obs=obs,
            noise_cov=noise_cov,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **_: Any,
    ) -> ProcessState:
        state_now = _with_evals(state, forward_evals)
        dt = self.scheduler.get_dt(state_now)
        delta = _eki_delta(
            state.particles, forward_evals, state.obs, state.noise_cov, dt
        )
        return ProcessState(
            particles=state.particles + delta,
            forward_evals=forward_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

GNKI

Bases: AbstractProcess

Gauss-Newton Kalman Inversion.

Estimates the Jacobian explicitly from ensemble perturbations (J̃ ≈ Cᶿᴳ (Cᶿᶿ)⁻¹) and applies a Gauss-Newton update with prior pull-back:

\[ \delta\theta^{(j)} = K \big(y - G^{(j)}\big) + \big(I - K \tilde{J}\big)\big(m_0 - \theta^{(j)}\big), \qquad K = \big(\tilde{J}^{\top} \Gamma^{-1} \tilde{J} + \Sigma_0^{-1}\big)^{-1} \tilde{J}^{\top} \Gamma^{-1}. \]

Faster convergence than :class:EKI for well-conditioned problems and requires J > Nₚ so Cᶿᶿ is invertible. In the linear-Gaussian limit GNKI recovers the exact posterior mean and covariance.

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size strategy.
prior_mean	Float[Array, ' N_p']	Prior mean m₀ ∈ ℝ^{Nₚ}.
prior_cov	AbstractLinearOperator	Prior covariance Σ₀.
jitter	float	Tikhonov regularisation added to Cᶿᶿ before inversion. Needed when J is barely larger than Nₚ.

Source code in src/filterax/_processes/_processes.py

class GNKI(AbstractProcess, strict=True):
    r"""Gauss-Newton Kalman Inversion.

    Estimates the Jacobian explicitly from ensemble perturbations
    (``J̃ ≈ Cᶿᴳ (Cᶿᶿ)⁻¹``) and applies a Gauss-Newton update with prior
    pull-back:

    $$
    \delta\theta^{(j)} = K \big(y - G^{(j)}\big)
        + \big(I - K \tilde{J}\big)\big(m_0 - \theta^{(j)}\big), \qquad
    K = \big(\tilde{J}^{\top} \Gamma^{-1} \tilde{J}
        + \Sigma_0^{-1}\big)^{-1} \tilde{J}^{\top} \Gamma^{-1}.
    $$

    Faster convergence
    than :class:`EKI` for well-conditioned problems and requires
    ``J > Nₚ`` so ``Cᶿᶿ`` is invertible. In the linear-Gaussian limit
    GNKI recovers the exact posterior mean *and* covariance.

    Attributes:
        scheduler: Step-size strategy.
        prior_mean: Prior mean ``m₀ ∈ ℝ^{Nₚ}``.
        prior_cov: Prior covariance ``Σ₀``.
        jitter: Tikhonov regularisation added to ``Cᶿᶿ`` before
            inversion. Needed when ``J`` is barely larger than ``Nₚ``.
    """

    scheduler: AbstractScheduler
    prior_mean: Float[Array, " N_p"]
    prior_cov: lx.AbstractLinearOperator
    jitter: float = eqx.field(static=True, default=1e-6)

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        J, N_p = particles.shape
        check_ensemble_size(J)
        # GNKI inverts the sample parameter covariance Cᶿᶿ ∈ ℝ^{Nₚ×Nₚ}.
        # With J ≤ Nₚ that matrix has rank ≤ J − 1 < Nₚ and is singular
        # — the jitter regulariser would dominate and silently corrupt
        # the Gauss-Newton step. Fail fast so users know to switch to
        # EKI (which works in the underdetermined regime).
        if J <= N_p:  # noqa: SIM300 — phrasing matches "J > Nₚ" requirement
            raise ValueError(
                "GNKI requires J > Nₚ so the sample parameter covariance "
                f"Cᶿᶿ is invertible; got J={J} and Nₚ={N_p}. Use "
                "filterax.EKI for the underdetermined regime."
            )
        N_d = obs.shape[0]
        zero_evals = jnp.zeros((J, N_d), dtype=particles.dtype)
        return ProcessState(
            particles=particles,
            forward_evals=zero_evals,
            obs=obs,
            noise_cov=noise_cov,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **_: Any,
    ) -> ProcessState:
        state_now = _with_evals(state, forward_evals)
        dt = self.scheduler.get_dt(state_now)
        N_p = state.particles.shape[1]

        # Ensemble Jacobian via cross-covariance solve.
        C_theta_theta = gaussx.ensemble_covariance(
            state.particles, bessel=True
        ).as_matrix() + self.jitter * jnp.eye(N_p)
        C_theta_G = gaussx.ensemble_cross_covariance(
            state.particles, forward_evals, bessel=True
        )  # (Nₚ, N_d)
        J_tilde = jnp.linalg.solve(C_theta_theta, C_theta_G)  # (Nₚ, N_d)

        # Gain K = (J̃ᵀ Γ⁻¹ J̃ + Σ₀⁻¹)⁻¹ J̃ᵀ Γ⁻¹.
        # solve_rows treats each row of J_tilde as an N_d vector and
        # returns (Γ⁻¹ J̃)_k row-by-row → shape (Nₚ, N_d).
        Gamma_inv_J = gaussx.solve_rows(state.noise_cov, J_tilde)  # (Nₚ, N_d)
        precision = einx.dot("p d, q d -> p q", J_tilde, Gamma_inv_J) + jnp.linalg.inv(
            self.prior_cov.as_matrix()
        )
        precision = 0.5 * (precision + precision.T)
        K_gain = jnp.linalg.solve(precision, Gamma_inv_J)  # (Nₚ, N_d)

        # Per-member update.
        residuals = einx.subtract("d, j d -> j d", state.obs, forward_evals)
        prior_pull = einx.subtract("p, j p -> j p", self.prior_mean, state.particles)
        data_step = einx.dot("p d, j d -> j p", K_gain, residuals)
        # (I − K J̃) m₀-pull term, expanded as m₀-pull − K J̃ (m₀-pull).
        J_pull = einx.dot("p d, j p -> j d", J_tilde, prior_pull)
        prior_step = prior_pull - einx.dot("p d, j d -> j p", K_gain, J_pull)
        delta = dt * (data_step + prior_step)
        particles_new = state.particles + delta
        return ProcessState(
            particles=particles_new,
            forward_evals=forward_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

SparseInversion

Bases: AbstractProcess

Sparse Ensemble Kalman Inversion (Schneider, Stuart & Wu 2022).

Standard EKI step followed by an L¹ proximal soft-threshold:

\[ \mathrm{prox}_{\lambda \|\cdot\|_1}(z)_i = \mathrm{sign}(z_i) \cdot \max\big(|z_i| - \lambda,\, 0\big) \]

drives inactive parameters exactly to zero. Useful for variable selection / sparse physics discovery / sensor placement.

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size strategy.
penalty_weight	float	λ > 0. Larger λ → more sparsity.

Source code in src/filterax/_processes/_processes.py

class SparseInversion(AbstractProcess, strict=True):
    r"""Sparse Ensemble Kalman Inversion (Schneider, Stuart & Wu 2022).

    Standard EKI step followed by an L¹ proximal soft-threshold:

    $$
    \mathrm{prox}_{\lambda \|\cdot\|_1}(z)_i =
        \mathrm{sign}(z_i) \cdot \max\big(|z_i| - \lambda,\, 0\big)
    $$

    drives inactive parameters exactly to zero. Useful for variable
    selection / sparse physics discovery / sensor placement.

    Attributes:
        scheduler: Step-size strategy.
        penalty_weight: ``λ > 0``. Larger ``λ`` → more sparsity.
    """

    scheduler: AbstractScheduler
    penalty_weight: float = eqx.field(static=True, default=0.1)

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        check_ensemble_size(particles.shape[0])
        N_d = obs.shape[0]
        zero_evals = jnp.zeros((particles.shape[0], N_d), dtype=particles.dtype)
        return ProcessState(
            particles=particles,
            forward_evals=zero_evals,
            obs=obs,
            noise_cov=noise_cov,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **_: Any,
    ) -> ProcessState:
        state_now = _with_evals(state, forward_evals)
        dt = self.scheduler.get_dt(state_now)
        delta = _eki_delta(
            state.particles, forward_evals, state.obs, state.noise_cov, dt
        )
        # Soft-threshold proximal operator scaled by Δt · λ.
        lam = self.penalty_weight * dt
        z = state.particles + delta
        particles_new = jnp.sign(z) * jnp.maximum(jnp.abs(z) - lam, 0.0)
        return ProcessState(
            particles=particles_new,
            forward_evals=forward_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

TEKI

Bases: AbstractProcess

Tikhonov-regularised Ensemble Kalman Inversion (Chada et al. 2019).

Standard EKI on an augmented system that includes the parameters in observation space:

\[ \tilde{G}(\theta) = \big(G(\theta),\, \theta\big)^{\top}, \qquad \tilde{y} = (y,\, m_0)^{\top}, \qquad \tilde{\Gamma} = \mathrm{blockdiag}(\Gamma,\, \Sigma_0) \]

The augmented identity block in Cᶿᴳ̃ pulls particles toward the prior mean m₀ and prevents the ensemble from drifting arbitrarily far from the prior — useful for ill-posed problems where vanilla EKI diverges. Converges to the MAP estimate (not MLE).

Attributes:

Name	Type	Description
scheduler	AbstractScheduler	Step-size strategy.
prior_mean	Float[Array, ' N_p']	Prior mean m₀ ∈ ℝ^{Nₚ}.
prior_cov	AbstractLinearOperator	Prior covariance Σ₀.

Source code in src/filterax/_processes/_processes.py

class TEKI(AbstractProcess, strict=True):
    r"""Tikhonov-regularised Ensemble Kalman Inversion (Chada et al. 2019).

    Standard EKI on an augmented system that includes the parameters in
    observation space:

    $$
    \tilde{G}(\theta) = \big(G(\theta),\, \theta\big)^{\top}, \qquad
    \tilde{y} = (y,\, m_0)^{\top}, \qquad
    \tilde{\Gamma} = \mathrm{blockdiag}(\Gamma,\, \Sigma_0)
    $$

    The augmented identity block in ``Cᶿᴳ̃`` pulls particles toward the
    prior mean ``m₀`` and prevents the ensemble from drifting arbitrarily
    far from the prior — useful for ill-posed problems where vanilla EKI
    diverges. Converges to the MAP estimate (not MLE).

    Attributes:
        scheduler: Step-size strategy.
        prior_mean: Prior mean ``m₀ ∈ ℝ^{Nₚ}``.
        prior_cov: Prior covariance ``Σ₀``.
    """

    scheduler: AbstractScheduler
    prior_mean: Float[Array, " N_p"]
    prior_cov: lx.AbstractLinearOperator

    def init(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
    ) -> ProcessState:
        check_ensemble_size(particles.shape[0])
        N_d = obs.shape[0]
        N_p = particles.shape[1]
        zero_evals = jnp.zeros((particles.shape[0], N_d + N_p), dtype=particles.dtype)
        # Store the augmented obs ỹ = (y, m₀) and augmented Γ̃ = blockdiag.
        aug_obs = jnp.concatenate([obs, self.prior_mean])
        aug_noise = (
            lx.BlockDiagonalLinearOperator([noise_cov, self.prior_cov])
            if hasattr(lx, "BlockDiagonalLinearOperator")
            else _block_diag(noise_cov, self.prior_cov)
        )
        return ProcessState(
            particles=particles,
            forward_evals=zero_evals,
            obs=aug_obs,
            noise_cov=aug_noise,
            step=jnp.asarray(0, dtype=jnp.int32),
            algo_time=jnp.asarray(0.0, dtype=particles.dtype),
        )

    def update(
        self,
        state: ProcessState,
        forward_evals: Float[Array, "J N_d"],
        **_: Any,
    ) -> ProcessState:
        # Augment evals with the particles themselves: G̃(θ⁽ʲ⁾) = (G⁽ʲ⁾, θ⁽ʲ⁾).
        aug_evals = jnp.concatenate([forward_evals, state.particles], axis=1)
        state_now = _with_evals(state, aug_evals)
        dt = self.scheduler.get_dt(state_now)
        delta = _eki_delta(state.particles, aug_evals, state.obs, state.noise_cov, dt)
        particles_new = state.particles + delta
        return ProcessState(
            particles=particles_new,
            forward_evals=aug_evals,
            obs=state.obs,
            noise_cov=state.noise_cov,
            step=state.step + 1,
            algo_time=state.algo_time + dt,
        )

Sigma-point utilities

The unscented quadrature used by UKI is also exposed for direct use.

Layer-1 ensemble Kalman process components.

Re-exports the L1 process classes from :mod:filterax._processes._processes so the canonical filterax.processes.EKI() access path matches the design docs. The Layer-2 filterax.EKI / filterax.EKS / filterax.UKI wrappers live at the top level and own the full init → forward eval → update loop.

sigma_points(mean: Float[Array, ' N_p'], cov: lx.AbstractLinearOperator, *, alpha: float = 1.0, beta: float = 2.0, kappa: float = 0.0) -> tuple[Float[Array, 'twoNp1 N_p'], Float[Array, ' twoNp1'], Float[Array, ' twoNp1']]

Generate 2 Nₚ + 1 deterministic sigma points (Wan & van der Merwe 2000).

For mean μ and covariance Σ:

\[ \lambda = \alpha^2 (N_p + \kappa) - N_p, \qquad c = \sqrt{N_p + \lambda} \]

\[ \chi^0 = \mu, \qquad \chi^{+}_j = \mu + c \, [\sqrt{\Sigma}]_{:,j}, \qquad \chi^{-}_j = \mu - c \, [\sqrt{\Sigma}]_{:,j} \]

with weights for the mean and covariance:

\[ \begin{aligned} W_m^0 &= \frac{\lambda}{N_p + \lambda}, & W_c^0 &= W_m^0 + (1 - \alpha^2 + \beta), \\ W_m^i &= W_c^i = \frac{1}{2 (N_p + \lambda)} & &\text{for } i \neq 0. \end{aligned} \]

√Σ comes from :func:gaussx.root_decomposition so structured Σ (diagonal, low-rank, …) skips the dense Cholesky.

Returns:

Type	Description
Float[Array, 'twoNp1 N_p']	(points, mean_weights, cov_weights) with shapes
Float[Array, ' twoNp1']	(2 Nₚ + 1, Nₚ), (2 Nₚ + 1,), (2 Nₚ + 1,).

Source code in src/filterax/_processes/_processes.py

def sigma_points(
    mean: Float[Array, " N_p"],
    cov: lx.AbstractLinearOperator,
    *,
    alpha: float = 1.0,
    beta: float = 2.0,
    kappa: float = 0.0,
) -> tuple[
    Float[Array, "twoNp1 N_p"],
    Float[Array, " twoNp1"],
    Float[Array, " twoNp1"],
]:
    r"""Generate ``2 Nₚ + 1`` deterministic sigma points (Wan & van der Merwe 2000).

    For mean ``μ`` and covariance ``Σ``:

    $$
    \lambda = \alpha^2 (N_p + \kappa) - N_p, \qquad
    c = \sqrt{N_p + \lambda}
    $$

    $$
    \chi^0 = \mu, \qquad
    \chi^{+}_j = \mu + c \, [\sqrt{\Sigma}]_{:,j}, \qquad
    \chi^{-}_j = \mu - c \, [\sqrt{\Sigma}]_{:,j}
    $$

    with weights for the mean and covariance:

    $$
    \begin{aligned}
    W_m^0 &= \frac{\lambda}{N_p + \lambda}, &
    W_c^0 &= W_m^0 + (1 - \alpha^2 + \beta), \\
    W_m^i &= W_c^i = \frac{1}{2 (N_p + \lambda)} & &\text{for } i \neq 0.
    \end{aligned}
    $$

    ``√Σ`` comes from :func:`gaussx.root_decomposition` so structured
    ``Σ`` (diagonal, low-rank, …) skips the dense Cholesky.

    Returns:
        ``(points, mean_weights, cov_weights)`` with shapes
        ``(2 Nₚ + 1, Nₚ)``, ``(2 Nₚ + 1,)``, ``(2 Nₚ + 1,)``.
    """
    N_p = mean.shape[0]
    lam = alpha**2 * (N_p + kappa) - N_p
    c = jnp.sqrt(N_p + lam)

    # Sigma offsets — columns of c · √Σ. UKI is parametric so Σ is
    # always a (Nₚ, Nₚ) dense or diagonal block, and Cholesky is both
    # cheap and exact. (The Lanczos default in gaussx is meant for huge
    # operators and produces NaN here on small exactly-isotropic Σ.)
    root = gaussx.root_decomposition(cov, method="cholesky").root
    offsets = c * root  # (Nₚ, Nₚ); column j is c · [√Σ]_{:,j}.

    # χ⁺ⱼ = μ + [c√Σ]_{:,j} — transpose so each row is one sigma point.
    plus = offsets.T + mean[None, :]
    minus = -offsets.T + mean[None, :]
    points = jnp.concatenate([mean[None, :], plus, minus], axis=0)

    w0_mean = lam / (N_p + lam)
    w0_cov = w0_mean + (1.0 - alpha**2 + beta)
    w_other = 1.0 / (2.0 * (N_p + lam))
    weights_mean = jnp.concatenate([jnp.asarray([w0_mean]), jnp.full(2 * N_p, w_other)])
    weights_cov = jnp.concatenate([jnp.asarray([w0_cov]), jnp.full(2 * N_p, w_other)])
    return points, weights_mean, weights_cov