Changed the EnKF distribution to be low rank

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(

dynestyx/inference/configs/discretizer.py

@@ -4,29 +4,14 @@
 
 import abc
 import dataclasses
-import math
 
 import diffrax as dfx
 
 from dynestyx.inference.configs.simulator import (
     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

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:

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}")

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
 

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
 

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}")