Efficient EnKF analysis step (including N << d_obs)

[MatthieuDarcy]

Addresses issue #277.

What was implemented

0. Efficiency gain when N < d_x

The current implementation computes the Kalman gain matrix

$$K = C_{XY}S^{-1}$$

Assuming $S$ is already factored, computing this from left to right incurs a cost of $\mathcal{O}(d_xd_y^2 + Nd_yd_x)$, whereas from right to left the cost is $\mathcal{O}(Nd_y^2 + Nd_yd_x)$. I added a simple dimension check to use the more efficient path.

1. Efficient implementation when N << d_obs

For the math, see #277. I provide a new implementation ("new branch") as an alternative to the previous implementation ("previous branch"); the new branch is more efficient when N << d_obs. The goal is to preserve the previous branch's behavior exactly.

Choices that were made:

1. The new implementation is used through an optional keyword ensemble_subspace (defaults to False) passed to build_filter and to update. The default preserves the previous branch, but I believe the new branch could be the default in some cases.
2. Covariance tapering is not supported: passing ensemble_subspace=True together with modify_cross_covariance != no_covariance_modifier or construct_chol_innovation_covariance != None raises an error (both at filter construction and in the filter update).
3. The new branch supports a Cholesky factor for R that is either a scalar or a 1D vector (implicitly assuming the Cholesky factor is diagonal, with either that scalar or that vector on the diagonal), as well as the full Cholesky factor. The reason for this choice is that computing $R^{-1}y$ costs $\mathcal{O}(d_y^2)$ even when $R$ is already factored; allowing a diagonal Cholesky factor reduces this to $\mathcal{O}(d_y)$. The docs reflect this. I did not change the previous implementation in this respect; it might be worth considering doing this for the previous branch too, but it would most likely bring little or no benefit.
4. The expression $d^\intercal R^{-1}d - d^\intercal R^{-1}YC^{-1}Y^{\intercal}R^{-1}d$ (see Improving EnKF computational efficiency when N << d_obs #277) is implemented as a least-squares problem, which avoids subtracting two large values. Claude claims this is more stable at a small extra cost. See below for the mathematical details. A test in the test suite ensures that both formulations match.
5. collect_nans_chol still costs $\mathcal{O}(d_y^3)$ because it refactorizes the Cholesky factor, even when there are no NaNs. In both branches, I added a check that skips this step when no NaNs are present. This function also converted square roots into lower-triangular factors; is this dangerous?

2. Tests

tests/cuthbertlib/ensemble_kalman/test_filtering.py

1. test_update_ensemble_subspace_matches_dense: basic consistency test. The previous branch and the new branch give the same answer (they differ only in computational complexity).
2. test_update_ensemble_subspace_chol_R_forms: scalar, 1D and 2D factors of the same R reproduce the previous branch.
3. test_quadratic_form_residual_matches_difference: _quadratic_form_residual (the least-squares implementation) matches the naive difference formula.
4. test_update_ensemble_subspace_missing_observations: with partially missing observations, the new branch matches the previous branch, for both a dense and a 1D chol_R.
5. test_update_ensemble_subspace_rejects_invalid_arguments: update rejects ensemble_subspace=True combined with cross_covariance_modifier or construct_chol_innovation_covariance.

tests/cuthbert/ensemble_kalman/test_enkf.py

1. test_build_filter_ensemble_subspace_matches_default: checks that a 10-step filter run with build_filter(..., ensemble_subspace=True) matches build_filter(..., ensemble_subspace=False) in ensemble and log normalizing constant, testing end-to-end behavior.
2. test_build_filter_ensemble_subspace_rejects_localization: build_filter raises at construction time when ensemble_subspace=True is combined with either localization callback.

3. Sharp edges and limitations

1. In the new branch, passing a non-triangular square root of R as chol_R (instead of a lower-triangular Cholesky factor) can silently produce wrong results. This is reflected in the documentation and, as far as I can tell, chol_R was never intended to be anything other than a lower-triangular factor, but it might be worth double-checking that this is acceptable. In the previous branch this is not an issue: results remain correct, and only the realization of the perturbed observations changes.
2. Passing a dense Cholesky factor erases much of the gain of the new branch; this is reflected in the documentation.
3. With a dense Cholesky factor, every step with missing data incurs an $\mathcal{O}(d_y^3)$ cost (in both branches). Scalar and 1D factors handle missing data in $\mathcal{O}(d_y)$.
4. The new branch is always a manual choice.

4. Mathematical details of the least-squares formulation

$$q = d^\top R^{-1}d - d^\top R^{-1}Y C^{-1}Y^\top R^{-1}d = \min_{w} |d - Yw|_{R^{-1}}^2 + |w|^2 = |d - Yw^\ast|_{R^{-1}}^2 + |w^\ast|^2$$

where $w^\ast = C^{-1}g$ and $g := Y^\top R^{-1}d$.

Proof:

$$ \begin{aligned} |d - Yw|_{R^{-1}}^2 + |w|^2 &= d^\top R^{-1}d - 2w^\top Y^\top R^{-1}d + w^\top Y^\top R^{-1}Yw + w^\top w \\ &= d^\top R^{-1}d - 2w^\top g + w^\top\big(I_N + Y^\top R^{-1}Y\big)w \\ &= d^\top R^{-1}d - 2w^\top g + w^\top Cw \end{aligned} $$

Taking the gradient and setting it to zero yields the optimum $w^\ast = C^{-1}g$. We now compute

$$ \begin{aligned} |d - Yw^\ast|_{R^{-1}}^2 + |w^\ast|^2 &= d^\top R^{-1}d - 2w^{\ast\top} g + w^{\ast\top} Cw^\ast \\ &= d^\top R^{-1}d - 2g^\top C^{-1}g + g^\top C^{-1}CC^{-1}g \\ &= d^\top R^{-1}d - g^\top C^{-1}g \\ &= d^\top R^{-1}d - d^\top R^{-1}YC^{-1}Y^\top R^{-1}d \end{aligned} $$

[SamDuffield]

Why does chol_R have to be lower-triangular rather than generalised Cholesky factor for this?

[SamDuffield]

Actually, I believe the "generalised Cholesky factor" is lower triangular, just not necessarily the unique output from jnp.linalg.cholesky. See

cuthbert/cuthbertlib/linalg/tria.py

Line 15 in 689dd02

def tria(A: Array) -> Array:

- so not sure the "lower-triangular" needs clarifying or could be replaced with "generalized" ?

[MatthieuDarcy]

Yes, "generalized Cholesky factor" is fine and I should've worded it that way.

Claude was worried that people could pass general square root factors (i.e. not necessarily lower triangular), hence the slightly strange wording.

Another thing: the factors cannot be singular (i.e. $R$ should be invertible) for the new branch.

[SamDuffield]

Cool cool yeah I think let's use "generalized Cholesky factor" here as it matches our terminology. Although we do need to state clearly what that means in the docs xD - this is issue #59

[MatthieuDarcy]

Done

[SamDuffield]

This looks good, great stuff! And as always we can come back and refactor the interface in future if we see good reason