Skip to content

Covariance Localization

filter_update accepts two independent low-level hooks. cross_covariance_modifier receives one empirical state-observation cross-covariance and returns its modified value; it defaults to no_covariance_modifier. For observation-space covariance localization, one may optionally use construct_chol_innovation_covariance, which receives normalized observation deviations Y and a Cholesky factor of the observation noise covariance, chol_R, and returns a generalized Cholesky factor of the localized innovation covariance \(C_{yy} + R\). All localization happens before handling of missing data, so localization happens in the original coordinates.

The most common form of localization is covariance tapering; cuthbertlib provides gaspari_cohn and gaussian correlation functions for this purpose. The Gaspari-Cohn correlation function is more classical, but has compact support which may cause difficulties in optimizing localization hyperparameters. The gaussian correlation function has infinite support, and may therefore have better gradient properties. For obsservation-space tapering, cuthbertlib provides a convenience construct_tapered_chol_innovation_covariance, which uses a Cholesky factor of the taper to construct a factor of the tapered innovation covariance. This comes at slightly higher cost, due to a larger QR solve.

cuthbertlib.ensemble_kalman.localization

Covariance localization utilities for ensemble Kalman methods.

construct_tapered_chol_innovation_covariance(Y, chol_taper, chol_R)

Construct a tapered innovation covariance factor without forming it densely.

If taper = chol_taper @ chol_taper.T and Y denotes the normalized observation deviations, the returned generalized Cholesky factor chol_S satisfies chol_S @ chol_S.T = taper * (Y @ Y.T) + R.

Parameters:

Name Type Description Default
Y Array

Observation deviations transposed and divided by the square root of one less than the ensemble size, shape (y_dim, n_particles).

 required
chol_taper Array

Factor of a positive-semidefinite observation-space taper, shape (y_dim, y_dim).

 required
chol_R Array

Cholesky factor of the observation noise covariance, shape (y_dim, y_dim).

 required

Returns:

Type Description
Array

Generalized Cholesky factor of the complete tapered innovation covariance,
Array

shape (y_dim, y_dim).
Source code in cuthbertlib/ensemble_kalman/localization.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def construct_tapered_chol_innovation_covariance(
    Y: Array,
    chol_taper: Array,
    chol_R: Array,
) -> Array:
    """Construct a tapered innovation covariance factor without forming it densely.

    If ``taper = chol_taper @ chol_taper.T`` and ``Y`` denotes the normalized
    observation deviations, the returned generalized Cholesky factor ``chol_S``
    satisfies ``chol_S @ chol_S.T = taper * (Y @ Y.T) + R``.

    Args:
        Y: Observation deviations transposed and divided by the square root of one
            less than the ensemble size, shape (y_dim, n_particles).
        chol_taper: Factor of a positive-semidefinite observation-space taper,
            shape (y_dim, y_dim).
        chol_R: Cholesky factor of the observation noise covariance, shape
            (y_dim, y_dim).

    Returns:
        Generalized Cholesky factor of the complete tapered innovation covariance,
        shape (y_dim, y_dim).
    """
    y_dim = Y.shape[0]
    Y_tilde = (chol_taper[:, :, None] * Y[:, None, :]).reshape(y_dim, -1)
    return tria(jnp.concatenate([Y_tilde, chol_R], axis=1))

gaussian(distances, length_scale)

Evaluates a Gaussian covariance taper.

The taper is the squared-exponential correlation function \(\rho(d; \ell) = \exp(-\frac{1}{2}(d / \ell)^2)\). This taper has infinite support and is differentiable with respect to the length scale.

Parameters:

Name Type Description Default
distances ArrayLike

Distances at which to evaluate the taper.

 required
length_scale ScalarArrayLike

Positive characteristic distance of the taper.

 required

Returns:

Type Description
Array

Taper values with the broadcast shape of the inputs.
Source code in cuthbertlib/ensemble_kalman/localization.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def gaussian(
    distances: ArrayLike,
    length_scale: ScalarArrayLike,
) -> Array:
    r"""Evaluates a Gaussian covariance taper.

    The taper is the squared-exponential correlation function
    $\rho(d; \ell) = \exp(-\frac{1}{2}(d / \ell)^2)$. This taper
    has infinite support and is differentiable with respect
    to the length scale.

    Args:
        distances: Distances at which to evaluate the taper.
        length_scale: Positive characteristic distance of the taper.

    Returns:
        Taper values with the broadcast shape of the inputs.
    """
    scaled_distances = jnp.asarray(distances) / length_scale
    return jnp.exp(-0.5 * jnp.square(scaled_distances))

gaspari_cohn(distances, support_radius)

Evaluates the compactly supported fifth-order Gaspari-Cohn taper.

This implements Eq. (4.10) of Gaspari and Cohn (1999), https://doi.org/10.1002/qj.49712555417.

support_radius is the full support radius: taper values are exactly zero where the absolute distance is greater than or equal to it.

Parameters:

Name Type Description Default
distances ArrayLike

Distances at which to evaluate the taper.

 required
support_radius ScalarArrayLike

Positive distance at which the taper reaches zero.

 required

Returns:

Type Description
Array

Taper values with the broadcast shape of the inputs.
Source code in cuthbertlib/ensemble_kalman/localization.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def gaspari_cohn(
    distances: ArrayLike,
    support_radius: ScalarArrayLike,
) -> Array:
    """Evaluates the compactly supported fifth-order Gaspari-Cohn taper.

    This implements Eq. (4.10) of Gaspari and Cohn (1999),
    https://doi.org/10.1002/qj.49712555417.

    ``support_radius`` is the full support radius: taper values are exactly zero
    where the absolute distance is greater than or equal to it.

    Args:
        distances: Distances at which to evaluate the taper.
        support_radius: Positive distance at which the taper reaches zero.

    Returns:
        Taper values with the broadcast shape of the inputs.
    """
    distances = jnp.abs(distances)
    q = 2 * distances / support_radius

    inner = 1 - 5 / 3 * q**2 + 5 / 8 * q**3 + 1 / 2 * q**4 - 1 / 4 * q**5

    # Avoid NaNs in computing the outer branch of the polynomial
    safe_q = jnp.where(q > 0, q, 1)
    # Factored version of Eq. (4.10) for numerical stability
    outer = (2 - q) ** 4 * (2 * q**2 + 4 * q - 1) / (24 * safe_q)

    within_support = jnp.where(q <= 1, inner, outer)
    return jnp.where(distances < support_radius, within_support, 0)