diff --git a/dynestyx/control/discrete_controller_simulators.py b/dynestyx/control/discrete_controller_simulators.py index ceb895d8..87617089 100644 --- a/dynestyx/control/discrete_controller_simulators.py +++ b/dynestyx/control/discrete_controller_simulators.py @@ -1,22 +1,24 @@ """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 ( + _cholesky_state_sequence_to_dists, +) from dynestyx.models import DynamicalModel from dynestyx.simulation.base import BaseSimulator from dynestyx.simulation.utils import _ensure_trailing_dim, _tile_times @@ -42,25 +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. + + 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. + + The shared conversion is time-indexed, so the state is given a leading axis of + length one and the single distribution unwrapped. + + 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`. - 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. + Raises: + ValueError: If `filter_config.filter_source` is not `"cuthbert"`. """ - 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 @@ -70,7 +96,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 @@ -286,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/discretizer.py b/dynestyx/inference/configs/discretizer.py index 9b9153d6..5116426f 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,7 @@ ODESimulatorConfig, SDESimulatorConfig, ) - - -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}." - ) - - -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}." - ) +from dynestyx.utils import _validate_nonnegative_float def _default_diffrax_sde_solver() -> SDESimulatorConfig: @@ -76,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 @@ -133,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 @@ -208,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 @@ -274,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 @@ -350,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 b0139035..8c51a0f5 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"] @@ -134,6 +136,16 @@ 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): 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 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"`. ??? note "Algorithm Reference" @@ -185,8 +197,18 @@ 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 + ) filter_source: CuthbertOnlyFilterSource = "cuthbert" + def __post_init__(self) -> None: + # Check that the jitter is nonnegative float + _validate_nonnegative_float( + "recorded_filtered_states_cov_jitter", + self.recorded_filtered_states_cov_jitter, + ) + @dataclasses.dataclass class PFResamplingConfig: diff --git a/dynestyx/inference/filters.py b/dynestyx/inference/filters.py index 1b2d3c08..deb4cbcc 100644 --- a/dynestyx/inference/filters.py +++ b/dynestyx/inference/filters.py @@ -716,6 +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, "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 36142350..9440a0d7 100644 --- a/dynestyx/inference/integrations/cuthbert/discrete_filter.py +++ b/dynestyx/inference/integrations/cuthbert/discrete_filter.py @@ -436,6 +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, "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 c7087b63..5c44a5e2 100644 --- a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py +++ b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py @@ -289,6 +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, "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 06f1f3c7..b465afe2 100644 --- a/dynestyx/inference/smoothers.py +++ b/dynestyx/inference/smoothers.py @@ -582,6 +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, "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 0b1c3616..62229bea 100644 --- a/dynestyx/inference/utils/distribution_utils.py +++ b/dynestyx/inference/utils/distribution_utils.py @@ -120,6 +120,77 @@ 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$. 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 + `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 +245,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. It defaults to + ``0.0`` -- the exact, unregularised covariance. + """ if particle_mode: return _particle_sequence_to_dists( states.particles, @@ -183,9 +266,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/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"] 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 c5dcabaf..3713f35a 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,118 @@ 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`.""" + 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_covariance_jitter_shifts_only_the_covariance_diagonal(): + """The jitter adds exactly ``eps * I`` to the covariance and nothing else. + + 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 + + # 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] + 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( + 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(): 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..ccc9372a 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -669,6 +669,70 @@ 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()) == 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 + 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. + + 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() + 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 cannot handle it (vmap's sparse tracing breaks inside a jnp.block call), so dynestyx