From 3c96e38f0e359e12b9584c27aeff2fe5521057bc Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Thu, 3 Sep 2026 17:43:38 -0400
Subject: [PATCH 1/7] Changed the EnKF distribution to be low rank

Return LowRankMultivariateNormal for rank-deficient EnKF distributions (for both `Filter` and `Smooth`)

Previously the distributions returned for the EnKF were always
MultivariateNormal, with the covariance formed densely from the QR
factors. When the ensemble covariance was low rank (n_particles - 1 <
state_dim) this had two consequences:

1. The returned distributions were singular and yielded NaN when `.sample` or `.log_prob` were used. NumPyro
   accepts a singular covariance matrix, which then gives NaN downstream
   when it is factorized.

2. It was computationally expensive for large state_dim.

The solution is to check whether the covariance is singular (a simple
comparison of n_particles against state_dim) and return a
LowRankMultivariateNormal in that case, which can be sampled. The same
is done in filter_state_dist, used by DiscreteControlLoopSimulator. If in the non degenerate case, we use the same `MultivariateNormal` to ensure existence of densities.

Note that this affects only the returned distributions. The filtered/smoothed
states are unchanged and behavior is otherwise the same as before (`marginal_loglik` is identical).

One design choice I am unsure of: I added a jitter parameter to the EnKF
config (filtered_covariance_jitter, default 0.0). LowRankMultivariateNormal
has no valid log_prob unless a jitter is added, so this lets the user
obtain a distribution with a valid density when needed. The jitter does
not affect the ability to sample, but it does change the resulting
distribution.

Claude claims to have solved 2 bugs:

1. In DiscreteControlLoopSimulator, filter_state_dist passed the rectangular ensemble factor as scale_tril, so the belief reported event_shape == (n_particles,). .mean and .log_prob would then raise an error.
2. Posterior rollout (predict_times with a filtered_result) returned all-NaN states: it grafts the filtered distributions in as the forecast initial condition and samples them, and those samples were NaN because the dense covariance was singular.
---
 .../control/discrete_controller_simulators.py | 35 +++++--
 dynestyx/inference/configs/filter.py          | 20 ++++
 dynestyx/inference/filters.py                 |  1 +
 .../integrations/cuthbert/discrete_filter.py  |  1 +
 .../cuthbert/discrete_smoother.py             |  1 +
 dynestyx/inference/smoothers.py               |  1 +
 .../inference/utils/distribution_utils.py     | 98 ++++++++++++++++++-
 tests/test_distribution_utils.py              | 79 +++++++++++++++
 tests/test_filters.py                         | 35 +++++++
 9 files changed, 261 insertions(+), 10 deletions(-)

diff --git a/dynestyx/control/discrete_controller_simulators.py b/dynestyx/control/discrete_controller_simulators.py
index ceb895d8..4a9dee56 100644
--- a/dynestyx/control/discrete_controller_simulators.py
+++ b/dynestyx/control/discrete_controller_simulators.py
@@ -17,6 +17,10 @@
     compute_cuthbert_filter_update,
 )
 from dynestyx.inference.integrations.utils import WeightedParticles
+from dynestyx.inference.utils.distribution_utils import (
+    _check_if_ensemble_low_rank,
+    _ensemble_sequence_to_low_rank_gaussian_dists,
+)
 from dynestyx.models import DynamicalModel
 from dynestyx.simulation.base import BaseSimulator
 from dynestyx.simulation.utils import _ensure_trailing_dim, _tile_times
@@ -45,14 +49,27 @@ def filter_state_mean(state: Any) -> Real[Array, "..."]:
 def filter_state_dist(state: Any) -> Distribution:
     """Full-belief NumPyro distribution for a cuthbert filter state, any family.
 
-    Kalman-family states (`KFConfig`, `EKFConfig`, `EnKFConfig`) expose
-    `.mean`/`.chol_cov`, giving an exact `MultivariateNormal`. `PFConfig`
-    states have no such property -- their belief is a weighted particle
-    cloud (`.particles`, `.log_weights`), represented via `WeightedParticles`
-    (dynestyx's own `Distribution`; NumPyro has no built-in equivalent).
-    Unlike `filter_state_mean`, this does not broadcast over a leading
-    time/batch axis -- call it once per (unbatched) state.
+    Kalman-family states (`KFConfig`, `EKFConfig`) expose `.mean`/`.chol_cov`
+    with a square Cholesky factor, giving an exact `MultivariateNormal`.
+    `PFConfig` states have no such property -- their belief is a weighted
+    particle cloud (`.particles`, `.log_weights`), represented via
+    `WeightedParticles` (dynestyx's own `Distribution`; NumPyro has no built-in
+    equivalent). Unlike `filter_state_mean`, this does not broadcast over a
+    leading time/batch axis -- call it once per (unbatched) state.
+
+    Ensemble states (`EnKFConfig`) have covariances that are at most rank `n_particles - 1`. 
+    When that is below `state_dim`, we use a `LowRankMultivariateNormal` built
+    from the ensemble factor. 
+    Otherwise the exact `MultivariateNormal` is kept.
+
+    A singular belief has the exact ensemble mean and covariance, but no density -- its `log_prob` is `nan`. 
+    Raise `n_particles` above `state_dim`, or set `EnKFConfig.filtered_covariance_jitter` on the
+    inference path, if you need one.
     """
+    if hasattr(state, "ensemble") and _check_if_ensemble_low_rank(state.ensemble):
+        return _ensemble_sequence_to_low_rank_gaussian_dists(state.ensemble[None, ...])[
+            0
+        ]
     if hasattr(state, "chol_cov"):
         return dist.MultivariateNormal(state.mean, scale_tril=state.chol_cov)
     if hasattr(state, "particles") and hasattr(state, "log_weights"):
@@ -70,7 +87,9 @@ class PolicyCallable(Protocol):
     $$u_k, s_{k+1} = \pi(\hat x_{k|k}, t_k, t_{k+1}, s_k)$$
 
     `x_hat` is a NumPyro `Distribution` -- `MultivariateNormal` for
-    `KFConfig`/`EKFConfig`/`EnKFConfig`, `WeightedParticles` for `PFConfig`
+    `KFConfig`/`EKFConfig`, `WeightedParticles` for `PFConfig`, and for
+    `EnKFConfig` either of `MultivariateNormal` or, once the ensemble is rank
+    deficient (`n_particles - 1 < state_dim`), `LowRankMultivariateNormal`
     (see `filter_state_dist`); use `x_hat.mean` for a
     family-agnostic point estimate, or the distribution itself for
     uncertainty-aware planning. `t_now`/`t_next` are the current and next
diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index b0139035..40886341 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -22,6 +22,17 @@
 CuthbertOrCDDynamaxFilterSource = CuthbertOnlyFilterSource | CDDynamaxOnlyFilterSource
 
 
+def _validate_filtered_covariance_jitter(filtered_covariance_jitter: float) -> None:
+    if (
+        not math.isfinite(filtered_covariance_jitter)
+        or filtered_covariance_jitter < 0.0
+    ):
+        raise ValueError(
+            "filtered_covariance_jitter must be a finite, nonnegative float, "
+            f"got {filtered_covariance_jitter!r}."
+        )
+
+
 @dataclasses.dataclass
 class BaseFilterConfig(abc.ABC):
     r"""Shared configuration options inherited by all filter configs.
@@ -134,6 +145,11 @@ class EnKFConfig(BaseFilterConfig):
         inflation_delta (float | None): Scale ensemble anomalies by
             \(\sqrt{1 + \delta}\) before the update to prevent collapse.
             `None` disables inflation.
+        filtered_covariance_jitter (float): Nonnegative \(\epsilon\) added to the
+            reported filtered-state covariance as \(\epsilon I\). Defaults to `0.0`.
+            When the filtered covariance is singular (i.e., when the ensemble size is smaller than the state dimension),
+            this jitter ensures that the resulting distribution has a well-defined density.
+            It never affects the filter recursion or the marginal likelihood.
         filter_source (FilterSource): Backend. Defaults to `"cuthbert"`.
 
     ??? note "Algorithm Reference"
@@ -185,8 +201,12 @@ class EnKFConfig(BaseFilterConfig):
     )
     perturb_measurements: bool | None = None
     inflation_delta: float | None = None
+    filtered_covariance_jitter: float = 0.0
     filter_source: CuthbertOnlyFilterSource = "cuthbert"
 
+    def __post_init__(self) -> None:
+        _validate_filtered_covariance_jitter(self.filtered_covariance_jitter)
+
 
 @dataclasses.dataclass
 class PFResamplingConfig:
diff --git a/dynestyx/inference/filters.py b/dynestyx/inference/filters.py
index 1b2d3c08..4603106e 100644
--- a/dynestyx/inference/filters.py
+++ b/dynestyx/inference/filters.py
@@ -716,6 +716,7 @@ def compute_output_member(dyn, ot, ov, ovf, om, ct, cv, k, *idxs):
                 states,
                 particle_mode=isinstance(config, PFConfig),
                 plate_shapes=plate_shapes,
+                covariance_jitter=getattr(config, "filtered_covariance_jitter", 0.0),
             )
 
         raise ValueError(f"Unsupported batched output kind: {output_kind}")
diff --git a/dynestyx/inference/integrations/cuthbert/discrete_filter.py b/dynestyx/inference/integrations/cuthbert/discrete_filter.py
index 36142350..69dc8b93 100644
--- a/dynestyx/inference/integrations/cuthbert/discrete_filter.py
+++ b/dynestyx/inference/integrations/cuthbert/discrete_filter.py
@@ -436,6 +436,7 @@ def run_discrete_filter(
     filtered_dists = _cholesky_state_sequence_to_dists(
         states,
         particle_mode=isinstance(filter_config, PFConfig),
+        covariance_jitter=getattr(filter_config, "filtered_covariance_jitter", 0.0),
     )
     return marginal_loglik, states, filtered_dists
 
diff --git a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
index c7087b63..fb03b7ee 100644
--- a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
+++ b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
@@ -289,6 +289,7 @@ def run_discrete_smoother(
     smoothed_dists = _cholesky_state_sequence_to_dists(
         states,
         particle_mode=isinstance(smoother_config, PFSmootherConfig),
+        covariance_jitter=getattr(smoother_config, "filtered_covariance_jitter", 0.0),
     )
     return marginal_loglik, states, smoothed_dists
 
diff --git a/dynestyx/inference/smoothers.py b/dynestyx/inference/smoothers.py
index 06f1f3c7..d0832b31 100644
--- a/dynestyx/inference/smoothers.py
+++ b/dynestyx/inference/smoothers.py
@@ -582,6 +582,7 @@ def compute_output_member(dyn, ot, ov, ct, cv, k, *idxs):
                 states,
                 particle_mode=isinstance(config, PFSmootherConfig),
                 plate_shapes=plate_shapes,
+                covariance_jitter=getattr(config, "filtered_covariance_jitter", 0.0),
             )
 
         raise ValueError(f"Unsupported batched output kind: {output_kind}")
diff --git a/dynestyx/inference/utils/distribution_utils.py b/dynestyx/inference/utils/distribution_utils.py
index 0b1c3616..bf4c21d6 100644
--- a/dynestyx/inference/utils/distribution_utils.py
+++ b/dynestyx/inference/utils/distribution_utils.py
@@ -120,6 +120,76 @@ def _gaussian_sequence_to_dists(
     ]
 
 
+def _check_if_ensemble_low_rank(
+    ensemble: Real[Array, "... n_particles state_dim"],
+) -> bool:
+    r"""Whether the ensemble sample covariance is rank deficient.
+
+    $P = X'X'^{\top}$ has rank at most $N-1$, so it is
+    singular exactly when ``n_particles - 1 < state_dim``.
+    Used to determine whether to use a low-rank representation of the covariance
+    or to expand it to a dense matrix for `MultivariateNormal`.
+    """
+    n_particles, state_dim = ensemble.shape[-2], ensemble.shape[-1]
+    return n_particles - 1 < state_dim
+
+
+def _ensemble_sequence_to_low_rank_gaussian_dists(
+    ensemble: Real[Array, "*plate time n_particles state_dim"],
+    *,
+    covariance_jitter: float = 0.0,
+    plate_shapes: tuple[int, ...] = (),
+) -> list[dist.Distribution]:
+    r"""Convert an ensemble to per-time Gaussians with low-rank covariance representation.
+
+    The ensemble sample covariance is a low-rank object:
+    $$
+    P_t = \frac{1}{N-1}\sum_{i}\left(x_t^{(i)}-\bar x_t\right)
+                              \left(x_t^{(i)}-\bar x_t\right)^{\top}
+        = X'_t X_t'^{\top},
+    \qquad \operatorname{rank} P_t \le N-1,
+    $$
+
+    The low-rank representation is a $(\text{state\_dim}, N)$ factor $X'_t$
+    (not expanded into a dense $(\text{state\_dim}, \text{state\_dim})$ matrix).
+
+    ``covariance_jitter`` is the $\epsilon$ of $P_t + \epsilon I$.
+    Default of ``0.0``, the distributions have the exact covariance, but no Lebesgue density.
+
+    Note:
+        `LowRankMultivariateNormal.log_prob` will yield ``nan`` unless a positive
+        `covariance_jitter` is provided. Hence only use this for genuinely rank deficient matrices;
+        otherwise see `_cholesky_state_sequence_to_dists`.
+
+    Args:
+        ensemble: Ensemble states, ``(*plate, time, n_particles, state_dim)``.
+        covariance_jitter: Nonnegative $\epsilon$ added to the covariance as
+            $\epsilon I$.
+        plate_shapes: Leading plate dimensions, as elsewhere in this module.
+
+    Returns:
+        One `numpyro.distributions.LowRankMultivariateNormal` per time index.
+    """
+    n_particles = ensemble.shape[-2]
+    state_dim = ensemble.shape[-1]
+    mean = jnp.mean(ensemble, axis=-2)
+    # (..., time, state_dim, n_particles): the factor X', not the product X' X'^T.
+    cov_factor = jnp.swapaxes(ensemble - mean[..., None, :], -1, -2) / jnp.sqrt(
+        jnp.asarray(n_particles - 1, dtype=ensemble.dtype)
+    )
+    cov_diag = jnp.full((state_dim,), covariance_jitter, dtype=ensemble.dtype)
+
+    t_len = _time_len_from_array(mean, plate_shapes)
+    return [
+        dist.LowRankMultivariateNormal(
+            _slice_time_axis(mean, t, plate_shapes),
+            _slice_time_axis(cov_factor, t, plate_shapes),
+            cov_diag,
+        )
+        for t in range(t_len)
+    ]
+
+
 def _particle_sequence_to_dists(
     particles: Real[Array, "*plate time n_particles state_dim"]
     | Real[Array, "*plate time n_particles"],
@@ -174,8 +244,20 @@ def _cholesky_state_sequence_to_dists(
     *,
     particle_mode: bool,
     plate_shapes: tuple[int, ...] = (),
+    covariance_jitter: float = 0.0,
 ) -> list[dist.Distribution]:
-    """Convert cuthbert state objects to per-time distributions."""
+    r"""Convert cuthbert state objects to per-time distributions.
+
+    Three state families are handled, dispatched structurally:
+
+    - particle states (`.particles`, `.log_weights`) become `WeightedParticles`;
+    - ensemble states (`.ensemble`, i.e. `EnKFState` and `EnRTSState`) become low-rank Gaussians,
+    when the ensemble is rank-deficient (``n_particles - 1 < state_dim``);
+    - everything else becomes a dense `MultivariateNormal`.
+
+    ``covariance_jitter`` is applied in both Gaussian branches; see
+    `_ensemble_sequence_to_low_rank_gaussian_dists`.
+    """
     if particle_mode:
         return _particle_sequence_to_dists(
             states.particles,
@@ -183,9 +265,21 @@ def _cholesky_state_sequence_to_dists(
             plate_shapes=plate_shapes,
         )
 
+    if hasattr(states, "ensemble") and _check_if_ensemble_low_rank(states.ensemble):
+        return _ensemble_sequence_to_low_rank_gaussian_dists(
+            states.ensemble,
+            covariance_jitter=covariance_jitter,
+            plate_shapes=plate_shapes,
+        )
+
+    covariances = covariance_from_cholesky(states.chol_cov)
+    if covariance_jitter:
+        covariances = covariances + covariance_jitter * jnp.eye(
+            covariances.shape[-1], dtype=covariances.dtype
+        )
     return _gaussian_sequence_to_dists(
         states.mean,
-        covariance_from_cholesky(states.chol_cov),
+        covariances,
         plate_shapes=plate_shapes,
     )
 
diff --git a/tests/test_distribution_utils.py b/tests/test_distribution_utils.py
index c5dcabaf..0595b012 100644
--- a/tests/test_distribution_utils.py
+++ b/tests/test_distribution_utils.py
@@ -2,6 +2,8 @@
 
 import jax
 import jax.numpy as jnp
+import jax.random as jr
+import numpyro.distributions as dist
 
 from dynestyx.inference.utils.distribution_utils import (
     _categorical_log_probs_to_dists,
@@ -65,6 +67,83 @@ def test_cholesky_state_sequence_to_dists_gaussian():
     assert jnp.allclose(dists[0].covariance_matrix, 4.0 * jnp.eye(2))
 
 
+def test_cholesky_state_sequence_to_dists_ensemble_is_low_rank():
+    """Ensemble states keep the rank-(N-1) factor instead of a dense covariance.
+
+    Expanding it would be both quadratic in ``state_dim`` and singular, and
+    ``MultivariateNormal`` takes its Cholesky eagerly, so the dense path yields
+    ``nan`` for ``sample`` and ``log_prob``. ``state_dim`` is deliberately larger
+    than ``n_particles`` here: that is the regime the dense path gets wrong, and
+    the one every other EnKF test misses by using ``state_dim=2``.
+    """
+    t_len, n_particles, state_dim = 3, 4, 16
+    ensemble = jr.normal(jr.PRNGKey(0), (t_len, n_particles, state_dim))
+    states = SimpleNamespace(ensemble=ensemble)
+
+    dists = _cholesky_state_sequence_to_dists(states, particle_mode=False)
+
+    assert len(dists) == t_len
+    for t, d in enumerate(dists):
+        assert isinstance(d, dist.LowRankMultivariateNormal)
+        assert d.batch_shape == ()
+        assert d.event_shape == (state_dim,)
+        assert d.cov_factor.shape == (state_dim, n_particles)
+
+        members = ensemble[t]
+        assert jnp.allclose(d.mean, members.mean(axis=0), atol=1e-5)
+
+        deviations = members - members.mean(axis=0)
+        expected_cov = deviations.T @ deviations / (n_particles - 1)
+        assert jnp.allclose(d.covariance_matrix, expected_cov, atol=1e-5)
+
+        # The point of the change: the dense path samples nan here.
+        assert jnp.isfinite(d.sample(jr.PRNGKey(t))).all()
+
+
+def test_cholesky_state_sequence_to_dists_full_rank_ensemble_stays_dense():
+    """A full-rank ensemble keeps the dense `MultivariateNormal`, and its `log_prob`.
+
+    `LowRankMultivariateNormal` divides by `cov_diag` to form its Woodbury
+    capacitance factor, so with the default zero jitter its `log_prob` is `nan`
+    even when the factor has full rank. Switching representation there would be a
+    silent regression for every model with `n_particles > state_dim` and buys
+    nothing: the dense covariance is well posed and at most `n_particles` wide.
+    """
+    n_particles, state_dim = 16, 4  # n_particles - 1 >= state_dim: full rank
+    ensemble = jr.normal(jr.PRNGKey(0), (2, n_particles, state_dim))
+    deviations = ensemble - ensemble.mean(axis=-2, keepdims=True)
+    chol_cov = jnp.linalg.cholesky(
+        jnp.einsum("tni,tnj->tij", deviations, deviations) / (n_particles - 1)
+    )
+    states = SimpleNamespace(
+        ensemble=ensemble, mean=ensemble.mean(axis=-2), chol_cov=chol_cov
+    )
+
+    dists = _cholesky_state_sequence_to_dists(states, particle_mode=False)
+
+    assert isinstance(dists[0], dist.MultivariateNormal)
+    assert jnp.isfinite(dists[0].log_prob(dists[0].mean))
+
+
+def test_cholesky_state_sequence_to_dists_ensemble_jitter_enables_log_prob():
+    """The ensemble covariance is singular, so a density needs explicit jitter."""
+    ensemble = jr.normal(jr.PRNGKey(0), (2, 4, 16))
+    states = SimpleNamespace(ensemble=ensemble)
+
+    without = _cholesky_state_sequence_to_dists(states, particle_mode=False)[0]
+    assert jnp.isnan(without.log_prob(without.mean))
+
+    with_jitter = _cholesky_state_sequence_to_dists(
+        states, particle_mode=False, covariance_jitter=1e-2
+    )[0]
+    assert jnp.isfinite(with_jitter.log_prob(with_jitter.mean))
+    assert jnp.allclose(
+        with_jitter.covariance_matrix,
+        without.covariance_matrix + 1e-2 * jnp.eye(16),
+        atol=1e-5,
+    )
+
+
 def test_categorical_log_probs_to_dists_plate_batched():
     logits = jnp.arange(24.0).reshape(2, 3, 4)
     log_probs = jax.nn.log_softmax(logits, axis=-1)
diff --git a/tests/test_filters.py b/tests/test_filters.py
index 4818acbe..bfb38a5d 100644
--- a/tests/test_filters.py
+++ b/tests/test_filters.py
@@ -669,6 +669,41 @@ def test_cuthbert_enkf_sparse_h_matches_dense_h():
     assert jnp.allclose(means_dense, means_sparse)
 
 
+def test_cuthbert_enkf_filtered_dists_are_low_rank_and_samplable():
+    """EnKF filtered distributions keep the ensemble factor and stay samplable.
+
+    The ensemble covariance has rank at most `n_particles - 1`, so expanding it into
+    a dense `MultivariateNormal` gives a singular matrix whose eager Cholesky is
+    `nan` -- which silently propagated into posterior rollout, since the rollout
+    grafts these distributions in as the forecast initial condition and samples
+    them. `n_particles=3` against a 3-state model puts us in that rank-deficient
+    regime on purpose.
+    """
+    dynamics = _sparse_h_test_dynamics(jnp.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]))
+    obs_times = jnp.arange(6.0)
+    ground_truth = dsx.simulate(
+        dynamics, rng_key=jr.PRNGKey(0), predict_times=obs_times, n_simulations=1
+    )
+    obs_values = jnp.asarray(ground_truth.observations)[0]
+
+    n_particles = 3
+    with Filter(
+        filter_config=EnKFConfig(n_particles=n_particles, crn_seed=jr.PRNGKey(42))
+    ):
+        result = dsx.condition(
+            "f", dynamics, obs_times=obs_times, obs_values=obs_values
+        )
+
+    ensemble = result.states.ensemble
+    assert len(result.dists) == len(obs_times)
+    for t, d in enumerate(result.dists):
+        assert isinstance(d, dist.LowRankMultivariateNormal)
+        assert d.event_shape == (dynamics.state_dim,)
+        assert d.cov_factor.shape == (dynamics.state_dim, n_particles)
+        assert jnp.allclose(d.mean, ensemble[t].mean(axis=0), atol=1e-5)
+        assert jnp.isfinite(d.sample(jr.PRNGKey(t))).all()
+
+
 def test_kf_raises_on_sparse_observation_matrix_cuthbert():
     """KF's cuthbert backend does not support a sparse H: cuthbert's Kalman internals
     cannot handle it (vmap's sparse tracing breaks inside a jnp.block call), so dynestyx

From 76ccf557a02fdfc10ae6f67580760def48e93976 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Thu, 3 Sep 2026 17:53:38 -0400
Subject: [PATCH 2/7] ruff fix

---
 dynestyx/control/discrete_controller_simulators.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/dynestyx/control/discrete_controller_simulators.py b/dynestyx/control/discrete_controller_simulators.py
index 4a9dee56..6399f20f 100644
--- a/dynestyx/control/discrete_controller_simulators.py
+++ b/dynestyx/control/discrete_controller_simulators.py
@@ -57,12 +57,12 @@ def filter_state_dist(state: Any) -> Distribution:
     equivalent). Unlike `filter_state_mean`, this does not broadcast over a
     leading time/batch axis -- call it once per (unbatched) state.
 
-    Ensemble states (`EnKFConfig`) have covariances that are at most rank `n_particles - 1`. 
+    Ensemble states (`EnKFConfig`) have covariances that are at most rank `n_particles - 1`.
     When that is below `state_dim`, we use a `LowRankMultivariateNormal` built
-    from the ensemble factor. 
+    from the ensemble factor.
     Otherwise the exact `MultivariateNormal` is kept.
 
-    A singular belief has the exact ensemble mean and covariance, but no density -- its `log_prob` is `nan`. 
+    A singular belief has the exact ensemble mean and covariance, but no density -- its `log_prob` is `nan`.
     Raise `n_particles` above `state_dim`, or set `EnKFConfig.filtered_covariance_jitter` on the
     inference path, if you need one.
     """

From 6399365e8ed66bb08fb4321d59c79f0cca88bb00 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Fri, 4 Sep 2026 11:45:40 -0400
Subject: [PATCH 3/7] unified jitter check

---
 dynestyx/inference/configs/discretizer.py | 13 +++----------
 dynestyx/inference/configs/filter.py      | 17 +++++------------
 dynestyx/utils.py                         | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/dynestyx/inference/configs/discretizer.py b/dynestyx/inference/configs/discretizer.py
index 9b9153d6..86b92146 100644
--- a/dynestyx/inference/configs/discretizer.py
+++ b/dynestyx/inference/configs/discretizer.py
@@ -4,7 +4,6 @@
 
 import abc
 import dataclasses
-import math
 
 import diffrax as dfx
 
@@ -12,21 +11,15 @@
     ODESimulatorConfig,
     SDESimulatorConfig,
 )
+from dynestyx.utils import _validate_nonnegative_float
 
 
 def _validate_covariance_jitter(covariance_jitter: float) -> None:
-    if not math.isfinite(covariance_jitter) or covariance_jitter < 0.0:
-        raise ValueError(
-            "covariance_jitter must be a finite, nonnegative float, "
-            f"got {covariance_jitter!r}."
-        )
+    _validate_nonnegative_float("covariance_jitter", covariance_jitter)
 
 
 def _validate_jitter_scale(jitter_scale: float) -> None:
-    if not math.isfinite(jitter_scale) or jitter_scale < 0.0:
-        raise ValueError(
-            f"jitter_scale must be a finite, nonnegative float, got {jitter_scale!r}."
-        )
+    _validate_nonnegative_float("jitter_scale", jitter_scale)
 
 
 def _default_diffrax_sde_solver() -> SDESimulatorConfig:
diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index 40886341..5f6323b7 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -8,6 +8,8 @@
 import jax.random as jr
 from jaxtyping import PRNGKeyArray
 
+from dynestyx.utils import _validate_nonnegative_float
+
 ResamplingBaseMethod = Literal["systematic", "multinomial", "stratified"]
 ResamplingDifferentiableMethod = Literal["stop_gradient", "straight_through", "soft"]
 FilterEmissionOrder = Literal["zeroth", "first", "second"]
@@ -22,17 +24,6 @@
 CuthbertOrCDDynamaxFilterSource = CuthbertOnlyFilterSource | CDDynamaxOnlyFilterSource
 
 
-def _validate_filtered_covariance_jitter(filtered_covariance_jitter: float) -> None:
-    if (
-        not math.isfinite(filtered_covariance_jitter)
-        or filtered_covariance_jitter < 0.0
-    ):
-        raise ValueError(
-            "filtered_covariance_jitter must be a finite, nonnegative float, "
-            f"got {filtered_covariance_jitter!r}."
-        )
-
-
 @dataclasses.dataclass
 class BaseFilterConfig(abc.ABC):
     r"""Shared configuration options inherited by all filter configs.
@@ -205,7 +196,9 @@ class EnKFConfig(BaseFilterConfig):
     filter_source: CuthbertOnlyFilterSource = "cuthbert"
 
     def __post_init__(self) -> None:
-        _validate_filtered_covariance_jitter(self.filtered_covariance_jitter)
+        _validate_nonnegative_float(
+            "filtered_covariance_jitter", self.filtered_covariance_jitter
+        )
 
 
 @dataclasses.dataclass
diff --git a/dynestyx/utils.py b/dynestyx/utils.py
index ff122ded..61c8ce57 100644
--- a/dynestyx/utils.py
+++ b/dynestyx/utils.py
@@ -294,6 +294,24 @@ def _should_record_field(
     return math.prod(shape) <= max_elems
 
 
+def _validate_nonnegative_float(name: str, value: float) -> None:
+    """Validate a nonnegative, finite float-valued config field.
+
+    Shared by the jitter fields on the discretizer and filter configs, which
+    all have the same admissible range. ``name`` is the field's own name, so
+    the error message points at the attribute the user actually set.
+
+    Args:
+        name: Name of the field being validated, as it appears on the config.
+        value: Value to validate.
+
+    Raises:
+        ValueError: If ``value`` is not finite or is negative.
+    """
+    if not math.isfinite(value) or value < 0.0:
+        raise ValueError(f"{name} must be a finite, nonnegative float, got {value!r}.")
+
+
 def _validate_control_dim(
     dynamics: DynamicalModel,
     ctrl_values: Real[Array, "*ctrl_value_plate ctrl_time control_dim"]

From 675ca55df56c19f43096d2d0bfd47ea891441207 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Fri, 4 Sep 2026 16:00:00 -0400
Subject: [PATCH 4/7] Rename, tests improvements, new defaults

Rename `filtered_covariance_jitter` to `recorded_filtered_states_cov_jitter`.
The old name suggested it entered the EnKF algorithm; it only perturbs the
distributions reported in `ConditionedResult.dists`, never the filter
recursion or the marginal likelihood.

Give it a positive default via an `"auto"` sentinel: resolves
once in `EnKFConfig.__post_init__` to a precision-dependent value (1e-5 in
float32, 1e-12 in float64). Compute so that variance ~ 1 should work, but may fail for larger variances.

In control, the old `filter_state_dist` now delegates to
`_cholesky_state_sequence_to_dists` with the same arguments
`run_discrete_filter` passes it, instead of duplicating the three-branch
dispatch, thus unifying with the behavior from filter.

_validate_jitter_scale and _validate_covariance_jitter were redundant. Created a new utility `_validate_nonnegative_float` into `dynestyx/utils.py` which all jitter validators can now call. Each caller passes its own field name to ensure so error messages still name the attribute the user
set.

Simplified and improved test suite:
`test_ensemble_jitter_is_off_by_default_and_enabled_only_by_config` verifies that each filter config has the right jitter value and that manually setting a jitter value gives the right answer.
`test_covariance_jitter_shifts_only_the_covariance_diagonal` checks the perturbation is exactly `eps * I` and the mean untouched, on both the dense and low-rank branches.
`test_cuthbert_enkf_filtered_dists_are_low_rank_and_samplable` rebuilds the dists as `run_discrete_filter` does and asserts they are no longer singular.
---
 .../control/discrete_controller_simulators.py | 80 +++++++++++--------
 dynestyx/inference/configs/filter.py          | 28 +++++--
 dynestyx/inference/filters.py                 |  4 +-
 .../integrations/cuthbert/discrete_filter.py  |  4 +-
 .../cuthbert/discrete_smoother.py             |  4 +-
 dynestyx/inference/smoothers.py               |  4 +-
 .../inference/utils/distribution_utils.py     | 33 +++++++-
 tests/test_discrete_control.py                | 38 +++++----
 tests/test_distribution_utils.py              | 75 ++++++++++++-----
 tests/test_filters.py                         | 46 +++++++++++
 10 files changed, 233 insertions(+), 83 deletions(-)

diff --git a/dynestyx/control/discrete_controller_simulators.py b/dynestyx/control/discrete_controller_simulators.py
index 6399f20f..87617089 100644
--- a/dynestyx/control/discrete_controller_simulators.py
+++ b/dynestyx/control/discrete_controller_simulators.py
@@ -1,25 +1,23 @@
 """Closed-loop simulation for controlled discrete-time dynamical models."""
 
+from types import SimpleNamespace
 from typing import Any, Protocol, runtime_checkable
 
 import jax
 import jax.numpy as jnp
 import jax.random as jr
-import numpyro.distributions as dist
 from jax import Array
 from jaxtyping import PRNGKeyArray, PyTree, Real
 from numpyro.distributions import Distribution
 
-from dynestyx.inference.configs.filter import BaseFilterConfig
+from dynestyx.inference.configs.filter import BaseFilterConfig, PFConfig
 from dynestyx.inference.filters import _default_filter_config
 from dynestyx.inference.integrations.cuthbert.discrete_filter import (
     build_cuthbert_filter,
     compute_cuthbert_filter_update,
 )
-from dynestyx.inference.integrations.utils import WeightedParticles
 from dynestyx.inference.utils.distribution_utils import (
-    _check_if_ensemble_low_rank,
-    _ensemble_sequence_to_low_rank_gaussian_dists,
+    _cholesky_state_sequence_to_dists,
 )
 from dynestyx.models import DynamicalModel
 from dynestyx.simulation.base import BaseSimulator
@@ -46,38 +44,49 @@ def filter_state_mean(state: Any) -> Real[Array, "..."]:
     raise TypeError(f"Cannot summarize filter state of type {type(state).__name__}")
 
 
-def filter_state_dist(state: Any) -> Distribution:
-    """Full-belief NumPyro distribution for a cuthbert filter state, any family.
+def filter_state_dist(state: Any, filter_config: BaseFilterConfig) -> Distribution:
+    """Full-belief NumPyro distribution for a filter state.
 
-    Kalman-family states (`KFConfig`, `EKFConfig`) expose `.mean`/`.chol_cov`
-    with a square Cholesky factor, giving an exact `MultivariateNormal`.
-    `PFConfig` states have no such property -- their belief is a weighted
-    particle cloud (`.particles`, `.log_weights`), represented via
-    `WeightedParticles` (dynestyx's own `Distribution`; NumPyro has no built-in
-    equivalent). Unlike `filter_state_mean`, this does not broadcast over a
-    leading time/batch axis -- call it once per (unbatched) state.
+    Only `filter_source="cuthbert"` is supported today, matching
+    `DiscreteControlLoopSimulator`'s own restriction. Converts the filter state to a NumPyro `Distribution`
+    in the same way that `ConditionedResult.dists` does for the recorded filtered states.
 
-    Ensemble states (`EnKFConfig`) have covariances that are at most rank `n_particles - 1`.
-    When that is below `state_dim`, we use a `LowRankMultivariateNormal` built
-    from the ensemble factor.
-    Otherwise the exact `MultivariateNormal` is kept.
+    The shared conversion is time-indexed, so the state is given a leading axis of
+    length one and the single distribution unwrapped.
 
-    A singular belief has the exact ensemble mean and covariance, but no density -- its `log_prob` is `nan`.
-    Raise `n_particles` above `state_dim`, or set `EnKFConfig.filtered_covariance_jitter` on the
-    inference path, if you need one.
+    Args:
+        state: A single, unbatched filter state produced by `filter_config`.
+        filter_config: The config that produced `state`. Selects the backend and
+            carries `recorded_filtered_states_cov_jitter` for ensemble states.
+
+    Returns:
+        The belief as a NumPyro `Distribution`.
+
+    Raises:
+        ValueError: If `filter_config.filter_source` is not `"cuthbert"`.
     """
-    if hasattr(state, "ensemble") and _check_if_ensemble_low_rank(state.ensemble):
-        return _ensemble_sequence_to_low_rank_gaussian_dists(state.ensemble[None, ...])[
-            0
-        ]
-    if hasattr(state, "chol_cov"):
-        return dist.MultivariateNormal(state.mean, scale_tril=state.chol_cov)
-    if hasattr(state, "particles") and hasattr(state, "log_weights"):
-        log_weights = jax.nn.log_softmax(state.log_weights, axis=-1)
-        return WeightedParticles(state.particles, log_weights)
-    raise TypeError(
-        f"Cannot build a distribution for filter state of type {type(state).__name__}"
-    )
+
+    if filter_config.filter_source == "cuthbert":
+        # Give the time-indexed conversion a leading axis of one
+        with_time_axis = SimpleNamespace(
+            **{
+                name: jnp.asarray(getattr(state, name))[None]
+                for name in ("ensemble", "mean", "chol_cov", "particles", "log_weights")
+                if hasattr(state, name)
+            }
+        )
+        return _cholesky_state_sequence_to_dists(
+            with_time_axis,
+            particle_mode=isinstance(filter_config, PFConfig),
+            covariance_jitter=getattr(
+                filter_config, "recorded_filtered_states_cov_jitter", 0.0
+            ),
+        )[0]
+    else:
+        raise ValueError(
+            "filter_state_dist currently only supports filter_source='cuthbert' only, got "
+            f"{filter_config.filter_source!r}."
+        )
 
 
 @runtime_checkable
@@ -305,7 +314,10 @@ def _step(carry, t_idx):
             t_next = times[t_idx + 1]
 
             u_k, s_next = self.control_policy(
-                filter_state_dist(x_hat_prev), t_now, t_next, s_prev
+                filter_state_dist(x_hat_prev, filter_config),
+                t_now,
+                t_next,
+                s_prev,
             )
             if isinstance(u_k, Distribution):
                 raise ValueError(
diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index 5f6323b7..a3df3db9 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -8,6 +8,10 @@
 import jax.random as jr
 from jaxtyping import PRNGKeyArray
 
+from dynestyx.inference.utils.distribution_utils import (
+    CovarianceJitter,
+    _default_covariance_jitter,
+)
 from dynestyx.utils import _validate_nonnegative_float
 
 ResamplingBaseMethod = Literal["systematic", "multinomial", "stratified"]
@@ -136,11 +140,17 @@ class EnKFConfig(BaseFilterConfig):
         inflation_delta (float | None): Scale ensemble anomalies by
             \(\sqrt{1 + \delta}\) before the update to prevent collapse.
             `None` disables inflation.
-        filtered_covariance_jitter (float): Nonnegative \(\epsilon\) added to the
-            reported filtered-state covariance as \(\epsilon I\). Defaults to `0.0`.
-            When the filtered covariance is singular (i.e., when the ensemble size is smaller than the state dimension),
-            this jitter ensures that the resulting distribution has a well-defined density.
-            It never affects the filter recursion or the marginal likelihood.
+        recorded_filtered_states_cov_jitter (float | Literal["auto"]): Nonnegative \(\epsilon\) added to
+            the **recorded** filtered-state covariance as \(\epsilon I\).
+            This only affects the covariance when converted to a `MultivariateNormal` or `LowRankMultivariateNormal`
+            distribution (notably those returned in `ConditionedResult.dists`); it never
+            enters the EnKF update, the filter recursion, or the marginal
+            likelihood.
+            When `n_particles - 1 < state_dim`, the ensemble covariance is singular,
+            this regularization is necessary to give the recorded distributions a well-defined density (sampling will work nonethelss).
+            `"auto"` (default) selects a small precision-dependent value
+            (`1e-5` in float32, `1e-12` in float64). Will work for variance around 1, but may need a bigger value
+            for larger magnitudes. Pass `0.0` for the exact, unregularised covariance.
         filter_source (FilterSource): Backend. Defaults to `"cuthbert"`.
 
     ??? note "Algorithm Reference"
@@ -192,12 +202,16 @@ class EnKFConfig(BaseFilterConfig):
     )
     perturb_measurements: bool | None = None
     inflation_delta: float | None = None
-    filtered_covariance_jitter: float = 0.0
+    recorded_filtered_states_cov_jitter: CovarianceJitter = "auto"
     filter_source: CuthbertOnlyFilterSource = "cuthbert"
 
     def __post_init__(self) -> None:
+        if self.recorded_filtered_states_cov_jitter == "auto":
+            self.recorded_filtered_states_cov_jitter = _default_covariance_jitter()
+        # Check that the jitter is nonnegative float
         _validate_nonnegative_float(
-            "filtered_covariance_jitter", self.filtered_covariance_jitter
+            "recorded_filtered_states_cov_jitter",
+            self.recorded_filtered_states_cov_jitter,
         )
 
 
diff --git a/dynestyx/inference/filters.py b/dynestyx/inference/filters.py
index 4603106e..deb4cbcc 100644
--- a/dynestyx/inference/filters.py
+++ b/dynestyx/inference/filters.py
@@ -716,7 +716,9 @@ def compute_output_member(dyn, ot, ov, ovf, om, ct, cv, k, *idxs):
                 states,
                 particle_mode=isinstance(config, PFConfig),
                 plate_shapes=plate_shapes,
-                covariance_jitter=getattr(config, "filtered_covariance_jitter", 0.0),
+                covariance_jitter=getattr(
+                    config, "recorded_filtered_states_cov_jitter", 0.0
+                ),
             )
 
         raise ValueError(f"Unsupported batched output kind: {output_kind}")
diff --git a/dynestyx/inference/integrations/cuthbert/discrete_filter.py b/dynestyx/inference/integrations/cuthbert/discrete_filter.py
index 69dc8b93..9440a0d7 100644
--- a/dynestyx/inference/integrations/cuthbert/discrete_filter.py
+++ b/dynestyx/inference/integrations/cuthbert/discrete_filter.py
@@ -436,7 +436,9 @@ def run_discrete_filter(
     filtered_dists = _cholesky_state_sequence_to_dists(
         states,
         particle_mode=isinstance(filter_config, PFConfig),
-        covariance_jitter=getattr(filter_config, "filtered_covariance_jitter", 0.0),
+        covariance_jitter=getattr(
+            filter_config, "recorded_filtered_states_cov_jitter", 0.0
+        ),
     )
     return marginal_loglik, states, filtered_dists
 
diff --git a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
index fb03b7ee..5c44a5e2 100644
--- a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
+++ b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
@@ -289,7 +289,9 @@ def run_discrete_smoother(
     smoothed_dists = _cholesky_state_sequence_to_dists(
         states,
         particle_mode=isinstance(smoother_config, PFSmootherConfig),
-        covariance_jitter=getattr(smoother_config, "filtered_covariance_jitter", 0.0),
+        covariance_jitter=getattr(
+            smoother_config, "recorded_filtered_states_cov_jitter", 0.0
+        ),
     )
     return marginal_loglik, states, smoothed_dists
 
diff --git a/dynestyx/inference/smoothers.py b/dynestyx/inference/smoothers.py
index d0832b31..b465afe2 100644
--- a/dynestyx/inference/smoothers.py
+++ b/dynestyx/inference/smoothers.py
@@ -582,7 +582,9 @@ def compute_output_member(dyn, ot, ov, ct, cv, k, *idxs):
                 states,
                 particle_mode=isinstance(config, PFSmootherConfig),
                 plate_shapes=plate_shapes,
-                covariance_jitter=getattr(config, "filtered_covariance_jitter", 0.0),
+                covariance_jitter=getattr(
+                    config, "recorded_filtered_states_cov_jitter", 0.0
+                ),
             )
 
         raise ValueError(f"Unsupported batched output kind: {output_kind}")
diff --git a/dynestyx/inference/utils/distribution_utils.py b/dynestyx/inference/utils/distribution_utils.py
index bf4c21d6..97f73f58 100644
--- a/dynestyx/inference/utils/distribution_utils.py
+++ b/dynestyx/inference/utils/distribution_utils.py
@@ -19,6 +19,10 @@
 
 MissingPolicy = Literal["raise", "empty"]
 
+# Sentinel for "pick a small value appropriate to the working precision",
+# following the `"auto"` convention used elsewhere in the package.
+CovarianceJitter = float | Literal["auto"]
+
 
 class _ForwardSimulationImproperUniform(dist.ImproperUniform):
     """An improper distribution sampled by dynamical forward simulation.
@@ -120,6 +124,26 @@ def _gaussian_sequence_to_dists(
     ]
 
 
+_DEFAULT_COV_JITTER_F32 = 1e-5
+_DEFAULT_COV_JITTER_F64 = 1e-12
+
+
+def _default_covariance_jitter() -> float:
+    r"""Default jitter for covariance regularization.
+    Chosen depending on the precision of the current JAX default float type.
+    Values were chosen empirically to give no failures at around
+    unit variance. States of much larger magnitude may still need a bigger value.
+
+    Resolved at call time rather than at import, since `jax_enable_x64` may be
+    toggled after a config is constructed.
+    """
+    return (
+        _DEFAULT_COV_JITTER_F64
+        if jnp.zeros(()).dtype == jnp.float64
+        else _DEFAULT_COV_JITTER_F32
+    )
+
+
 def _check_if_ensemble_low_rank(
     ensemble: Real[Array, "... n_particles state_dim"],
 ) -> bool:
@@ -153,8 +177,9 @@ def _ensemble_sequence_to_low_rank_gaussian_dists(
     The low-rank representation is a $(\text{state\_dim}, N)$ factor $X'_t$
     (not expanded into a dense $(\text{state\_dim}, \text{state\_dim})$ matrix).
 
-    ``covariance_jitter`` is the $\epsilon$ of $P_t + \epsilon I$.
-    Default of ``0.0``, the distributions have the exact covariance, but no Lebesgue density.
+    ``covariance_jitter`` is the $\epsilon$ of $P_t + \epsilon I$. At the default of
+    ``0.0`` the distributions carry the exact covariance but have no Lebesgue
+    density, so ``log_prob`` is ``nan``.
 
     Note:
         `LowRankMultivariateNormal.log_prob` will yield ``nan`` unless a positive
@@ -255,8 +280,8 @@ def _cholesky_state_sequence_to_dists(
     when the ensemble is rank-deficient (``n_particles - 1 < state_dim``);
     - everything else becomes a dense `MultivariateNormal`.
 
-    ``covariance_jitter`` is applied in both Gaussian branches; see
-    `_ensemble_sequence_to_low_rank_gaussian_dists`.
+    ``covariance_jitter`` is applied in both Gaussian branches. It defaults to
+    ``0.0`` -- the exact, unregularised covariance.
     """
     if particle_mode:
         return _particle_sequence_to_dists(
diff --git a/tests/test_discrete_control.py b/tests/test_discrete_control.py
index 97ba1229..70fab1a8 100644
--- a/tests/test_discrete_control.py
+++ b/tests/test_discrete_control.py
@@ -21,7 +21,13 @@
     EulerMaruyamaConfig,
     _discretize_state_evolution,
 )
-from dynestyx.inference.configs.filter import EKFConfig, EnKFConfig, KFConfig, PFConfig
+from dynestyx.inference.configs.filter import (
+    EKFConfig,
+    EnKFConfig,
+    KFConfig,
+    PFConfig,
+    UKFConfig,
+)
 from dynestyx.inference.integrations.cuthbert.discrete_filter import (
     build_cuthbert_filter,
     compute_cuthbert_filter,
@@ -135,7 +141,7 @@ class _KFLikeState:
         mean = jnp.array([1.0, 2.0])
         chol_cov = jnp.array([[1.0, 0.0], [0.5, 1.0]])
 
-    result = filter_state_dist(_KFLikeState())
+    result = filter_state_dist(_KFLikeState(), KFConfig(filter_source="cuthbert"))
     assert isinstance(result, dist.MultivariateNormal)
     assert jnp.allclose(result.mean, _KFLikeState.mean)
     assert jnp.allclose(result.scale_tril, _KFLikeState.chol_cov)
@@ -146,7 +152,7 @@ class _PFLikeState:
         particles = jnp.array([[0.0], [2.0], [4.0]])  # 3 particles, state_dim=1
         log_weights = jnp.log(jnp.array([0.25, 0.25, 0.5]))
 
-    result = filter_state_dist(_PFLikeState())
+    result = filter_state_dist(_PFLikeState(), PFConfig(n_particles=3))
     assert isinstance(result, WeightedParticles)
     assert jnp.allclose(result.particles, _PFLikeState.particles)
     assert jnp.allclose(
@@ -154,19 +160,23 @@ class _PFLikeState:
     )
 
 
-@pytest.mark.parametrize(
-    ("fn", "match"),
-    [
-        (filter_state_mean, "Cannot summarize filter state"),
-        (filter_state_dist, "Cannot build a distribution"),
-    ],
-)
-def test_unsupported_filter_state_type_raises(fn, match):
+def test_unsupported_filter_state_type_raises():
     class _Neither:
         pass
 
-    with pytest.raises(TypeError, match=match):
-        fn(_Neither())
+    with pytest.raises(TypeError, match="Cannot summarize filter state"):
+        filter_state_mean(_Neither())
+
+
+def test_filter_state_dist_rejects_non_cuthbert_backend():
+    """`filter_state_dist` dispatches on the backend, not on the state's shape.
+
+    Only cuthbert states carry a Cholesky factor; cd-dynamax reports dense
+    covariances, so it needs its own branch rather than being duck-typed into
+    this one.
+    """
+    with pytest.raises(ValueError, match="filter_source='cuthbert' only"):
+        filter_state_dist(object(), UKFConfig())
 
 
 # ---------------------------------------------------------------------------
@@ -379,7 +389,7 @@ def test_filter_state_dist_matches_family_and_agrees_with_mean(
         t=_OBS_TIMES[0],
         t_prev=_OBS_TIMES[0] - 1.0,
     )
-    result = filter_state_dist(state)
+    result = filter_state_dist(state, filter_config)
     assert isinstance(result, expected_dist_type)
 
     if isinstance(result, WeightedParticles):
diff --git a/tests/test_distribution_utils.py b/tests/test_distribution_utils.py
index 0595b012..3713f35a 100644
--- a/tests/test_distribution_utils.py
+++ b/tests/test_distribution_utils.py
@@ -101,14 +101,7 @@ def test_cholesky_state_sequence_to_dists_ensemble_is_low_rank():
 
 
 def test_cholesky_state_sequence_to_dists_full_rank_ensemble_stays_dense():
-    """A full-rank ensemble keeps the dense `MultivariateNormal`, and its `log_prob`.
-
-    `LowRankMultivariateNormal` divides by `cov_diag` to form its Woodbury
-    capacitance factor, so with the default zero jitter its `log_prob` is `nan`
-    even when the factor has full rank. Switching representation there would be a
-    silent regression for every model with `n_particles > state_dim` and buys
-    nothing: the dense covariance is well posed and at most `n_particles` wide.
-    """
+    """A full-rank ensemble keeps the dense `MultivariateNormal`, and its `log_prob`."""
     n_particles, state_dim = 16, 4  # n_particles - 1 >= state_dim: full rank
     ensemble = jr.normal(jr.PRNGKey(0), (2, n_particles, state_dim))
     deviations = ensemble - ensemble.mean(axis=-2, keepdims=True)
@@ -125,23 +118,65 @@ def test_cholesky_state_sequence_to_dists_full_rank_ensemble_stays_dense():
     assert jnp.isfinite(dists[0].log_prob(dists[0].mean))
 
 
-def test_cholesky_state_sequence_to_dists_ensemble_jitter_enables_log_prob():
-    """The ensemble covariance is singular, so a density needs explicit jitter."""
-    ensemble = jr.normal(jr.PRNGKey(0), (2, 4, 16))
-    states = SimpleNamespace(ensemble=ensemble)
+def test_covariance_jitter_shifts_only_the_covariance_diagonal():
+    """The jitter adds exactly ``eps * I`` to the covariance and nothing else.
 
-    without = _cholesky_state_sequence_to_dists(states, particle_mode=False)[0]
-    assert jnp.isnan(without.log_prob(without.mean))
+    Checked on both Gaussian branches, since they apply it by different means:
+    the dense branch adds ``eps * I`` to the covariance directly, while the
+    low-rank branch passes ``eps`` as `LowRankMultivariateNormal`'s ``cov_diag``
+    and never forms the covariance at all. The mean must be untouched either way
+    -- the jitter regularizes the reported covariance so a singular one gains a
+    density, it is not a change of location.
+    """
+    jitter = 1e-5
 
-    with_jitter = _cholesky_state_sequence_to_dists(
-        states, particle_mode=False, covariance_jitter=1e-2
+    # Dense branch: a square Cholesky factor. Covariance is 2 * I @ (2 * I).T = 4 * I.
+    dense_states = SimpleNamespace(
+        mean=jnp.array([[1.0, 2.0], [3.0, 4.0]]),
+        chol_cov=jnp.broadcast_to(2.0 * jnp.eye(2), (2, 2, 2)),
+    )
+    exact = _cholesky_state_sequence_to_dists(
+        dense_states, particle_mode=False, covariance_jitter=0.0
     )[0]
-    assert jnp.isfinite(with_jitter.log_prob(with_jitter.mean))
+    jittered = _cholesky_state_sequence_to_dists(
+        dense_states, particle_mode=False, covariance_jitter=jitter
+    )[0]
+
+    assert isinstance(exact, dist.MultivariateNormal)
+    assert jnp.array_equal(exact.covariance_matrix, 4.0 * jnp.eye(2))
+    assert jnp.allclose(
+        jittered.covariance_matrix,
+        exact.covariance_matrix + jitter * jnp.eye(2),
+        atol=1e-8,
+    )
+    assert jnp.array_equal(jittered.mean, exact.mean)
+
+    # Low-rank branch: the ensemble covariance is
+    # singular and only the jitter gives it a density.
+    state_dim = 16
+    ensemble_states = SimpleNamespace(
+        ensemble=jr.normal(jr.PRNGKey(0), (2, 4, state_dim))
+    )
+    lr_exact = _cholesky_state_sequence_to_dists(
+        ensemble_states, particle_mode=False, covariance_jitter=0.0
+    )[0]
+    lr_jittered = _cholesky_state_sequence_to_dists(
+        ensemble_states, particle_mode=False, covariance_jitter=jitter
+    )[0]
+
+    assert isinstance(lr_exact, dist.LowRankMultivariateNormal)
     assert jnp.allclose(
-        with_jitter.covariance_matrix,
-        without.covariance_matrix + 1e-2 * jnp.eye(16),
-        atol=1e-5,
+        lr_jittered.covariance_matrix,
+        lr_exact.covariance_matrix + jitter * jnp.eye(state_dim),
+        atol=1e-6,
     )
+    assert jnp.array_equal(lr_jittered.mean, lr_exact.mean)
+    # The factor itself is untouched; the jitter lives entirely in cov_diag.
+    assert jnp.array_equal(lr_jittered.cov_factor, lr_exact.cov_factor)
+    assert jnp.allclose(lr_jittered.cov_diag, jnp.full((state_dim,), jitter))
+    # Only the jittered one has a density.
+    assert jnp.isnan(lr_exact.log_prob(lr_exact.mean))
+    assert jnp.isfinite(lr_jittered.log_prob(lr_jittered.mean))
 
 
 def test_categorical_log_probs_to_dists_plate_batched():
diff --git a/tests/test_filters.py b/tests/test_filters.py
index bfb38a5d..cdc62a03 100644
--- a/tests/test_filters.py
+++ b/tests/test_filters.py
@@ -23,6 +23,10 @@
 from dynestyx.inference.integrations.cuthbert.discrete import (
     run_discrete_filter as run_cuthbert_discrete_filter,
 )
+from dynestyx.inference.utils.distribution_utils import (
+    _cholesky_state_sequence_to_dists,
+    _default_covariance_jitter,
+)
 from dynestyx.models import (
     ContinuousTimeStateEvolution,
     DynamicalModel,
@@ -669,6 +673,34 @@ def test_cuthbert_enkf_sparse_h_matches_dense_h():
     assert jnp.allclose(means_dense, means_sparse)
 
 
+def test_ensemble_jitter_is_off_by_default_and_enabled_only_by_config():
+    """Only `EnKFConfig` carries a jitter; every other cuthbert filter reads 0.0.
+
+    This is the exact expression the four conversion call sites use. The
+    `getattr` fallback is load-bearing: `KFConfig`/`EKFConfig`/`PFConfig` do not
+    declare the field at all, and they are exact filters whose reported
+    covariance is compared against analytic solutions, so they must not inherit
+    a perturbation from the shared code path.
+    """
+
+    def jitter_of(filter_config):
+        return getattr(filter_config, "recorded_filtered_states_cov_jitter", 0.0)
+
+    for filter_config in (
+        KFConfig(filter_source="cuthbert"),
+        EKFConfig(filter_source="cuthbert"),
+        PFConfig(n_particles=16, filter_source="cuthbert"),
+    ):
+        assert jitter_of(filter_config) == 0.0
+
+    # EnKF opts in by default, resolving "auto" to a precision-dependent value.
+    assert jitter_of(EnKFConfig()) == _default_covariance_jitter()
+
+    # ... and an explicit value is passed through untouched, including zero.
+    assert jitter_of(EnKFConfig(recorded_filtered_states_cov_jitter=0.0)) == 0.0
+    assert jitter_of(EnKFConfig(recorded_filtered_states_cov_jitter=3.14)) == 3.14
+
+
 def test_cuthbert_enkf_filtered_dists_are_low_rank_and_samplable():
     """EnKF filtered distributions keep the ensemble factor and stay samplable.
 
@@ -703,6 +735,20 @@ def test_cuthbert_enkf_filtered_dists_are_low_rank_and_samplable():
         assert jnp.allclose(d.mean, ensemble[t].mean(axis=0), atol=1e-5)
         assert jnp.isfinite(d.sample(jr.PRNGKey(t))).all()
 
+    # Rebuilt exactly as `run_discrete_filter` does it: with the config's own
+    # jitter the belief is no longer singular, so it has a usable density.
+    filter_config = EnKFConfig(n_particles=n_particles, crn_seed=jr.PRNGKey(42))
+    with_jitter = _cholesky_state_sequence_to_dists(
+        result.states,
+        particle_mode=isinstance(filter_config, PFConfig),
+        covariance_jitter=getattr(
+            filter_config, "recorded_filtered_states_cov_jitter", 0.0
+        ),
+    )
+    for d in with_jitter:
+        assert isinstance(d, dist.LowRankMultivariateNormal)
+        assert jnp.isfinite(d.log_prob(d.mean))
+
 
 def test_kf_raises_on_sparse_observation_matrix_cuthbert():
     """KF's cuthbert backend does not support a sparse H: cuthbert's Kalman internals

From b0a64f87ee38ba732e6158017f8d1faee9843d55 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Tue, 8 Sep 2026 15:05:50 -0400
Subject: [PATCH 5/7] jitter is now at a default of 1e-5

---
 dynestyx/inference/configs/filter.py          | 15 ++++--------
 .../inference/utils/distribution_utils.py     | 24 -------------------
 tests/test_filters.py                         |  3 +--
 3 files changed, 5 insertions(+), 37 deletions(-)

diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index a3df3db9..dcc416c3 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -8,10 +8,6 @@
 import jax.random as jr
 from jaxtyping import PRNGKeyArray
 
-from dynestyx.inference.utils.distribution_utils import (
-    CovarianceJitter,
-    _default_covariance_jitter,
-)
 from dynestyx.utils import _validate_nonnegative_float
 
 ResamplingBaseMethod = Literal["systematic", "multinomial", "stratified"]
@@ -140,7 +136,7 @@ class EnKFConfig(BaseFilterConfig):
         inflation_delta (float | None): Scale ensemble anomalies by
             \(\sqrt{1 + \delta}\) before the update to prevent collapse.
             `None` disables inflation.
-        recorded_filtered_states_cov_jitter (float | Literal["auto"]): Nonnegative \(\epsilon\) added to
+        recorded_filtered_states_cov_jitter (float): Nonnegative \(\epsilon\) added to
             the **recorded** filtered-state covariance as \(\epsilon I\).
             This only affects the covariance when converted to a `MultivariateNormal` or `LowRankMultivariateNormal`
             distribution (notably those returned in `ConditionedResult.dists`); it never
@@ -148,9 +144,8 @@ class EnKFConfig(BaseFilterConfig):
             likelihood.
             When `n_particles - 1 < state_dim`, the ensemble covariance is singular,
             this regularization is necessary to give the recorded distributions a well-defined density (sampling will work nonethelss).
-            `"auto"` (default) selects a small precision-dependent value
-            (`1e-5` in float32, `1e-12` in float64). Will work for variance around 1, but may need a bigger value
-            for larger magnitudes. Pass `0.0` for the exact, unregularised covariance.
+            Defaults to `1e-5`. Will work for variance around 1, but may need a bigger value
+            for larger magnitudes and may want to reduce when using float64. Pass `0.0` for the exact, unregularised covariance.
         filter_source (FilterSource): Backend. Defaults to `"cuthbert"`.
 
     ??? note "Algorithm Reference"
@@ -202,12 +197,10 @@ class EnKFConfig(BaseFilterConfig):
     )
     perturb_measurements: bool | None = None
     inflation_delta: float | None = None
-    recorded_filtered_states_cov_jitter: CovarianceJitter = "auto"
+    recorded_filtered_states_cov_jitter: float = 1e-5 # this is good for float32, may want to reduce for float64
     filter_source: CuthbertOnlyFilterSource = "cuthbert"
 
     def __post_init__(self) -> None:
-        if self.recorded_filtered_states_cov_jitter == "auto":
-            self.recorded_filtered_states_cov_jitter = _default_covariance_jitter()
         # Check that the jitter is nonnegative float
         _validate_nonnegative_float(
             "recorded_filtered_states_cov_jitter",
diff --git a/dynestyx/inference/utils/distribution_utils.py b/dynestyx/inference/utils/distribution_utils.py
index 97f73f58..62229bea 100644
--- a/dynestyx/inference/utils/distribution_utils.py
+++ b/dynestyx/inference/utils/distribution_utils.py
@@ -19,10 +19,6 @@
 
 MissingPolicy = Literal["raise", "empty"]
 
-# Sentinel for "pick a small value appropriate to the working precision",
-# following the `"auto"` convention used elsewhere in the package.
-CovarianceJitter = float | Literal["auto"]
-
 
 class _ForwardSimulationImproperUniform(dist.ImproperUniform):
     """An improper distribution sampled by dynamical forward simulation.
@@ -124,26 +120,6 @@ def _gaussian_sequence_to_dists(
     ]
 
 
-_DEFAULT_COV_JITTER_F32 = 1e-5
-_DEFAULT_COV_JITTER_F64 = 1e-12
-
-
-def _default_covariance_jitter() -> float:
-    r"""Default jitter for covariance regularization.
-    Chosen depending on the precision of the current JAX default float type.
-    Values were chosen empirically to give no failures at around
-    unit variance. States of much larger magnitude may still need a bigger value.
-
-    Resolved at call time rather than at import, since `jax_enable_x64` may be
-    toggled after a config is constructed.
-    """
-    return (
-        _DEFAULT_COV_JITTER_F64
-        if jnp.zeros(()).dtype == jnp.float64
-        else _DEFAULT_COV_JITTER_F32
-    )
-
-
 def _check_if_ensemble_low_rank(
     ensemble: Real[Array, "... n_particles state_dim"],
 ) -> bool:
diff --git a/tests/test_filters.py b/tests/test_filters.py
index cdc62a03..20b6be01 100644
--- a/tests/test_filters.py
+++ b/tests/test_filters.py
@@ -25,7 +25,6 @@
 )
 from dynestyx.inference.utils.distribution_utils import (
     _cholesky_state_sequence_to_dists,
-    _default_covariance_jitter,
 )
 from dynestyx.models import (
     ContinuousTimeStateEvolution,
@@ -694,7 +693,7 @@ def jitter_of(filter_config):
         assert jitter_of(filter_config) == 0.0
 
     # EnKF opts in by default, resolving "auto" to a precision-dependent value.
-    assert jitter_of(EnKFConfig()) == _default_covariance_jitter()
+    assert jitter_of(EnKFConfig()) == 1e-5
 
     # ... and an explicit value is passed through untouched, including zero.
     assert jitter_of(EnKFConfig(recorded_filtered_states_cov_jitter=0.0)) == 0.0

From 8bb69c37641b48b2be4cbe05726c313f66348630 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Tue, 8 Sep 2026 15:06:22 -0400
Subject: [PATCH 6/7] Update filter.py

---
 dynestyx/inference/configs/filter.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index dcc416c3..0ecd0040 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -197,7 +197,9 @@ class EnKFConfig(BaseFilterConfig):
     )
     perturb_measurements: bool | None = None
     inflation_delta: float | None = None
-    recorded_filtered_states_cov_jitter: float = 1e-5 # this is good for float32, may want to reduce for float64
+    recorded_filtered_states_cov_jitter: float = (
+        1e-5  # this is good for float32, may want to reduce for float64
+    )
     filter_source: CuthbertOnlyFilterSource = "cuthbert"
 
     def __post_init__(self) -> None:

From 7668258bb5b3c47565854e8d9b2fc0034895f250 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Wed, 9 Sep 2026 13:37:22 -0400
Subject: [PATCH 7/7] Simplified some tests and fixed typos

---
 dynestyx/inference/configs/discretizer.py | 18 +++++-------------
 dynestyx/inference/configs/filter.py      |  2 +-
 tests/test_filters.py                     | 16 ----------------
 3 files changed, 6 insertions(+), 30 deletions(-)

diff --git a/dynestyx/inference/configs/discretizer.py b/dynestyx/inference/configs/discretizer.py
index 86b92146..5116426f 100644
--- a/dynestyx/inference/configs/discretizer.py
+++ b/dynestyx/inference/configs/discretizer.py
@@ -14,14 +14,6 @@
 from dynestyx.utils import _validate_nonnegative_float
 
 
-def _validate_covariance_jitter(covariance_jitter: float) -> None:
-    _validate_nonnegative_float("covariance_jitter", covariance_jitter)
-
-
-def _validate_jitter_scale(jitter_scale: float) -> None:
-    _validate_nonnegative_float("jitter_scale", jitter_scale)
-
-
 def _default_diffrax_sde_solver() -> SDESimulatorConfig:
     return SDESimulatorConfig(source="diffrax", solver=dfx.Euler())
 
@@ -69,7 +61,7 @@ class ODEFlowConfig(BaseDiscretizerConfig):
     jitter_scale: float = 0.0
 
     def __post_init__(self) -> None:
-        _validate_jitter_scale(self.jitter_scale)
+        _validate_nonnegative_float("jitter_scale", self.jitter_scale)
 
 
 @dataclasses.dataclass
@@ -126,7 +118,7 @@ class EulerMaruyamaConfig(BaseDiscretizerConfig):
     covariance_jitter: float = 0.0
 
     def __post_init__(self) -> None:
-        _validate_covariance_jitter(self.covariance_jitter)
+        _validate_nonnegative_float("covariance_jitter", self.covariance_jitter)
 
 
 @dataclasses.dataclass
@@ -201,7 +193,7 @@ class ExactAffineConfig(BaseDiscretizerConfig):
     covariance_jitter: float = 0.0
 
     def __post_init__(self) -> None:
-        _validate_covariance_jitter(self.covariance_jitter)
+        _validate_nonnegative_float("covariance_jitter", self.covariance_jitter)
 
 
 @dataclasses.dataclass
@@ -267,7 +259,7 @@ class LocalLinearizationConfig(BaseDiscretizerConfig):
     covariance_jitter: float = 0.0
 
     def __post_init__(self) -> None:
-        _validate_covariance_jitter(self.covariance_jitter)
+        _validate_nonnegative_float("covariance_jitter", self.covariance_jitter)
 
 
 @dataclasses.dataclass
@@ -343,7 +335,7 @@ class MeanTrajectoryLinearizationConfig(BaseDiscretizerConfig):
     covariance_jitter: float = 0.0
 
     def __post_init__(self) -> None:
-        _validate_covariance_jitter(self.covariance_jitter)
+        _validate_nonnegative_float("covariance_jitter", self.covariance_jitter)
 
 
 @dataclasses.dataclass
diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index 0ecd0040..8c51a0f5 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -143,7 +143,7 @@ class EnKFConfig(BaseFilterConfig):
             enters the EnKF update, the filter recursion, or the marginal
             likelihood.
             When `n_particles - 1 < state_dim`, the ensemble covariance is singular,
-            this regularization is necessary to give the recorded distributions a well-defined density (sampling will work nonethelss).
+            this regularization is necessary to give the recorded distributions a well-defined density (sampling will work nonetheless).
             Defaults to `1e-5`. Will work for variance around 1, but may need a bigger value
             for larger magnitudes and may want to reduce when using float64. Pass `0.0` for the exact, unregularised covariance.
         filter_source (FilterSource): Backend. Defaults to `"cuthbert"`.
diff --git a/tests/test_filters.py b/tests/test_filters.py
index 20b6be01..ccc9372a 100644
--- a/tests/test_filters.py
+++ b/tests/test_filters.py
@@ -23,9 +23,6 @@
 from dynestyx.inference.integrations.cuthbert.discrete import (
     run_discrete_filter as run_cuthbert_discrete_filter,
 )
-from dynestyx.inference.utils.distribution_utils import (
-    _cholesky_state_sequence_to_dists,
-)
 from dynestyx.models import (
     ContinuousTimeStateEvolution,
     DynamicalModel,
@@ -733,19 +730,6 @@ def test_cuthbert_enkf_filtered_dists_are_low_rank_and_samplable():
         assert d.cov_factor.shape == (dynamics.state_dim, n_particles)
         assert jnp.allclose(d.mean, ensemble[t].mean(axis=0), atol=1e-5)
         assert jnp.isfinite(d.sample(jr.PRNGKey(t))).all()
-
-    # Rebuilt exactly as `run_discrete_filter` does it: with the config's own
-    # jitter the belief is no longer singular, so it has a usable density.
-    filter_config = EnKFConfig(n_particles=n_particles, crn_seed=jr.PRNGKey(42))
-    with_jitter = _cholesky_state_sequence_to_dists(
-        result.states,
-        particle_mode=isinstance(filter_config, PFConfig),
-        covariance_jitter=getattr(
-            filter_config, "recorded_filtered_states_cov_jitter", 0.0
-        ),
-    )
-    for d in with_jitter:
-        assert isinstance(d, dist.LowRankMultivariateNormal)
         assert jnp.isfinite(d.log_prob(d.mean))