Changed the EnKF distribution to be low rank

[MatthieuDarcy]

Summary

Adresses #355.

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.

Bug fixes

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.

Verification

I'm still double checking that this yields the same results as before but I locally ran the following tests:

1. Equivalence on a single distribution for the same ensemble at various dimensionality and ensemble size ($d= 8, 16, 64, 256, N = 8, 16, 32$).
2. End-to-end filtering: same marginal_loglik as before, distributions agree when both work, new low rank works when the distributioj is degenerate.
3. Speed test: $d = 1024, 4096, 10 000$, measures the speedup and memory saving (between 7x-608x improvement).

[mattlevine22]

looks quite reasonable to me!

Question: Do we want the new jitter to be used in regular MVN stuff too (currently it is)?

My recommendation would be:

• only use this jitter for the low rank MVN
• give it a > 0 default (so it protects from the .log_prob nan issue)
• rename / redoc this in the EnKF config to be clear that it is only for low rank and only for the recorded dists (I suggested a name above)

Certainly open to the current choice if you think its worth it.

[mattlevine22]

inference/configs/discretizer.py already has _validate_jitter_scale and _validate_covariance_jitter.

Can we use those?

[mattlevine22]

I'd worry a user thinks this will appear in the EnKF algorithm (which it doesn't).

How about a name like recorded_filtered_states_[lowrank]_cov_jitter?

Do we want this jitter usable even for full rank MVN? or just to deal w LowRank stuff

[mattlevine22]

Maybe recorded_filtered_states_cov_jitter?

[mattlevine22]

so this uses the jitter in the full-rank MVN...do we want to even allow that?

[mattlevine22]

Since this test uses default ...jitter=0.0, a d.log_prob(...) call would fail / NaN here right?

May want a regression test specifically that jitter protects you in that case (its the failure mode that made you build it).