Add Localization for the EnKF

[DanWaxman]

This PR adds localization via covariance tapering, namely through the Gaspari-Cohn and Gaussian correlation functions.

• The proposed API adds a GetCovarianceTapers callback to the filtering interface, which specifies a cross-covariance taper and, optionally, a marginal covariance taper.
• Custom tapers are possible, but we also provide Gaspari-Cohn and Gaussian in cuthbertlib/ensemble_kalman/localization.py.
  • Gaspari-Cohn is more standard, but Gaussian has infinite support, and therefore has better gradient properties if one wants to perform gradient-based optimization of localization hyperparameters.
• We add a tutorial using Lorenz-96.
• We have end-to-end tests to show (i) localization doesn't crash things, basically and (ii) gradients w.r.t. lengthscale of a Gaussian taper are well-behaved, matching closely to a finite difference approximation.

We purposefully avoid localization for smoothing, as there lacks a straightforward "correct" way to do so, as far as I can tell.

[AdrienCorenflos]

Aouch, a cholesky! Can we discuss this before it goes in? Can you write the actual equations so we see if we can avoid the cholesky

[DanWaxman]

I know I know I was sad I couldn't figure it out :(. The actual equations are:

$$ C_{yy} = \rho_{yy} \odot \sum_n (\hat{y}^{(n)}_{t | t-1} - \bar{\hat{y}}_{t | t-1}) (\hat{y}^{(n)}_{t | t-1} - \bar{\hat{y}}_{t | t-1})^\top / N-1 ,$$

where $\bar{\hat{y}}_{t | t-1}$ is the mean over the ensemble. The ensemble Kalman gain used later is then

$$ K = C_{xy} \underbrace{(C_{yy} + R)}_{S}{}^{-1}, $$

which is used in the update

$$ x^{(n)}_{t|t} = x^{(n)}_{t|t-1} + K (y_t - \hat{y}_{t|t-1}). $$

The issue is then that we must get a factor of $S.$ I didn't see an obvious way to get a factor of $S = C_{yy} + R$ unless the taper $\rho_{yy}$ comes factored, i.e., $\rho_{yy} = L_\rho L_\rho^\top$.

It's worth noting that I don't anticipate hitting this branch very often. It only happens if a marginal covariance taper is specified. To the best of my knowledge, it is more common to just supply a cross covariance taper, which avoids a Cholesky call.

[AdrienCorenflos]

@jax.jit
def compute_enkf_cholesky_factor_qr(Y, L_rho, L_R):
    """
    Computes the lower-triangular Cholesky factor of S = rho_yy * (Y @ Y.T) + R
    using an augmented matrix QR formulation to avoid explicit covariance construction.
    
    Parameters
    ----------
    Y : jax.numpy.ndarray
        Normalized ensemble perturbations array of shape (dy, N).
    L_rho : jax.numpy.ndarray
        Lower-triangular Cholesky factor of the tapering matrix, shape (dy, dy).
    L_R : jax.numpy.ndarray
        Lower-triangular Cholesky factor of the observation noise matrix, shape (dy, dy).
        
    Returns
    -------
    L_S : jax.numpy.ndarray
        Lower-triangular Cholesky factor of S, shape (dy, dy).
    """
    dy, N = Y.shape
    
    # Construct $\tilde{Y} = [\text{diag}(l_1)Y \mid \dots \mid \text{diag}(l_{d_y})Y]$ 
    # L_rho[:, :, None] -> shape (dy, dy, 1)
    # Y[:, None, :]     -> shape (dy, 1, N)
    # Z                 -> shape (dy, dy, N)
    # Y_tilde           -> shape (dy, dy * N)
    Z = L_rho[:, :, None] * Y[:, None, :]
    Y_tilde = Z.reshape((dy, dy * N))

    # Construct the augmented matrix A = [L_R | Y_tilde]
    A = jnp.hstack([L_R, Y_tilde])

    # Compute QR decomposition of A^T. A^T is of shape (dy + N*dy, dy).
    R_qr = jnp.linalg.qr(A.T, mode='r')
    L_S_qr = R_qr.T

    # Enforce positive diagonal elements for uniqueness
    sign_diag = jnp.sign(jnp.diagonal(L_S_qr))
    # Handle zeros to avoid zeroing out entire columns
    zero_flag = (sign_diag < 0.5) & (sign_diag > -0.5) 
    sign_diag = jnp.where(zero_flag, jnp.ones_like(sign_diag), sign_diag)
    L_S_qr = L_S_qr * sign_diag[None, :]
    
    # Clean up tiny floating point noise above the diagonal
    L_S_qr = jnp.tril(L_S_qr)
    
    return L_S_qr

[AdrienCorenflos]

@DanWaxman that should do the trick, I tested on some cases, trusting you to check properly :D

[AdrienCorenflos]

Here

$$ Y = (\hat{y}^{(n)}_{t | t-1} - \bar{\hat{y}}_{t | t-1}) /\sqrt{N-1} $$

Btw, if this works can you add me as a co-author in the commit so I appear on the blame?

Co-authored-by: AdrienCorenflos adrien.corenflos@gmail.com

same as in https://docs.github.com/en/pull-requests/how-tos/commit-changes/creating-a-commit-with-multiple-authors

[DanWaxman]

Thanks Adrien! I agree that this works. I guess two things from there:

1. I was trying to avoid assuming that $\rho_{yy}$ had a computed Cholesky factor. We can compute it, but this additionally assumes that the taper is PSD/well-conditioned. This is probably a generally reasonable assumption, though I can imagine places where this runs into numerical trouble.
2. Sam had suggested moving to a more general marginal_covariance_modifier(C_yy) API. I quite like this, as it is more general (and there are some things you may want to implement that are only possible this way; for example, shrinkage-based estimators). Using the above form would necessitate moving back to the taper-only form (or a more complicated API for the user). I don't have very strong feelings here, but it does seem relevant to talk through.

[SamDuffield]

I'm happy with requiring $\rho_{yy}$ in Cholesky form, if it becomes a problem later we can address it.

Regarding 2, yeah I think the functional form would be good. Maybe we can expand the input function to be construct_chol_cross_covariance(Y, L_R) for ensemble perturbations $Y$ and Cholesky factor of R in Adrien's notation. Then the default can be Adrien's function with L_rho=jnp.eye. What do you think?

[DanWaxman]

Then the default can be Adrien's function with L_rho=jnp.eye. What do you think?

Hmm, I think the drawback here is requiring a much larger tria solve if the user isn't actually doing tapering. I would probably prefer keeping a None path that reflects the current code structure, just using tria instead of chol for the cross-covariance-modified branch?

[SamDuffield]

oh yes good point, for sure!

[DanWaxman]

Updated this!

This commit adds use of generalized Cholesky factors under observation-space localization. This works by creating a general hook for creating an innovation covaraince factor, of which, tapering is a special case. The tapering special case is given as an augmented tria call in construct_tapered_chol_innovation_covariance, assuming that the taper matrix has a Cholesky factor.

The form of the tapered solve is due to Adrien Corenflos (see #273 (comment)).

The branching logic in the EnKF remains unchanged, i.e., allows a branch that skips localization, as (i) observation-space tapering here implies a larger-dimensional QR solve and (ii) missingness requires an unnecessary Cholesky pivot if a default is provided.

[Sahel13]

Any reason docs/api_cuthbertlib/ensemble_kalman.md was deleted? There a couple of links in zensical.toml and docs/api_cuthbertlib/index.md still pointing to this.

[Sahel13]

The docs can be cleaned up later as well, there are some other minor issues that also need to be fixed unrelated to the ensemble Kalman

[DanWaxman]

Thanks! I updated the docs (I got a bit confused, was using the structure for docs/cuthbert/ensemble_kalman from #271 that I don't think made it over fully to cuthbertlib). Should be fixed now. I'll take a look at Adrien's comment now!

[Sahel13]

I think it's best to wait to get @SamDuffield's opinion on this

[SamDuffield]

Do we need args and kwargs here?

[DanWaxman]

It's for typing, otherwise it doesn't conform to the ModifyCrossCovariance Protocol.

[SamDuffield]

Really? I'm only seeing CrossCovarianceModifier = Callable[[Array], Array] which would actually require no args or kwargs. Am I missing something?

[SamDuffield]

Oh I see it's with the cuthbert interface, hmm we should definitely be able to do this without needing args, kwargs

[SamDuffield]

For example I'd be happy to have two default functions for no_covariance_modifier in both cuthbert and cuthbertlib. They can even be called _no_covariance_modifier or similar so they don't need docstrings. I just don't like the args, kwargs 😄

[SamDuffield]

I would maybe rename this construct_chol_innovation_covariance since there is no requirement that is has to be localised, this is only an obvious use case, but there may be others

[DanWaxman]

Sure, done.

[SamDuffield]

Would state clearly in the docs here what the input and output are for construct_localized_chol_innovation_covariance or if you like consider a Protocol rather than Callable to make it clear

[DanWaxman]

I documented this in the docstring; I thought about more protocols, but they differ between cuthbert and cuthbertlib (the former takes in model inputs), and this seemed more easily confused.

[SamDuffield]

Happy with this, thanks a lot @DanWaxman! As always, if in future we find it's not suitably general for some use cases we can reevaluate

Think we also need @Sahel13 to approve as he is currently on Changes requested