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.

[SamDuffield]

I'm wondering if the tapering API could be a function that modifies the cross or marginal covariance rather than an array itself (the current CovarianceTapers object). This would be more flexible to the user and allow tapering that doesn't require instantiating an extra d x d matrix. The only issue I see is I'm not sure it will play nice with the argsort/nan code above

[SamDuffield]

I.e. modified_cross, modified_marginal = covariance_modify(cross, marginal, model_inputs)

what do you think?

[DanWaxman]

Sure, I like this idea! One thing that I'm a bit unclear on is how to do this for cross vs. marginal. In particular, I think the marginal taper requires a direct Cholesky decomposition (unless Adrien has any ideas), but cross-covariance-only tapering can avoid this.

One option is to just write modify_cross_covariance and modify_marginal_covariance separately?

Another option would be to use a ModifiedCovariance return type that stores modified_cross_cov : Array and modified_marginal_cov : Array | None. I think I like the former better, but can also see merits in the latter.

[SamDuffield]

Former is better: more functional and more transparent!

[DanWaxman]

Updated! Relevant L96 example code now looks like (streamlined a little bit from the example notebook, which is slightly more general):

def periodic_distance(left, right):
    """Return pairwise shortest distances on the Lorenz–96 ring."""
    direct = jnp.abs(left[:, None] - right[None, :])
    return jnp.minimum(direct, state_dim - direct)

cross_distances = periodic_distance(state_locations, observation_locations)
marginal_distances = periodic_distance(observation_locations, observation_locations)

def modify_cross_covariance(C_xy, model_inputs):
    return gaspari_cohn(cross_distances, support_radius) * C_xy

def modify_marginal_covariance(C_yy, model_inputs):
    return gaspari_cohn(marginal_distances, support_radius) * C_yy

filter_obj = ensemble_kalman_filter.build_filter(
        init_sample=init_sample,
        get_dynamics=get_dynamics,
        get_observations=get_observations,
        n_particles=n_members,
        inflation=inflation,
        perturbed_obs=True,
        modify_cross_covariance=modify_cross_covariance,
        modify_marginal_covariance=modify_marginal_covariance,
    )

Passing in modify_marginal_covariance=None maintains the square root update. Modifiers are applied before pivoting from NaNs, so they should retain the correct indexing.

[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.