Localization

A small ensemble (N_e ≪ N_x) produces a rank-deficient sample covariance whose long-range entries are dominated by sampling noise: two physically unrelated variables will show spurious correlation of order 1/√N_e, and the filter will happily use it to spread observation increments across the whole domain. Localization suppresses these spurious long-range correlations by damping covariance entries as a function of physical distance.

The standard mechanism is the Schur (element-wise) product P_loc = ρ ∘ P with a taper matrix ρ built from a compactly supported correlation function. By the Schur product theorem, the element-wise product of two positive semi-definite matrices is positive semi-definite — so as long as the taper is itself a valid correlation matrix (Gaspari-Cohn and SOAR are; a hard cutoff is not, in general), localization cannot break the PSD-ness of the covariance.

gaspari_cohn and localization_matrix delegate to their gaussx counterparts with support parameter c = 2 * radius (gaussx parameterises by the compact-support length, filterax by the half-width); the distance metrics (euclidean_distance, haversine_distance) are re-exported from gaussx unchanged. Wrap these primitives in an AbstractLocalizer to plug them into the Layer-2 loops as a drop-in component; the LETKF instead applies localization in observation space, per grid point, and localized_kalman_gain tapers the gain itself.

Taper functions

Map a distance array to weights in [0, 1]. Gaspari-Cohn is the production default (compactly supported, exactly zero beyond 2 * radius); Gaussian and SOAR decay smoothly without compact support; the hard cutoff is a diagnostic tool rather than a valid correlation function.

filterax — differentiable ensemble Kalman filters and processes for JAX.

gaspari_cohn(distances: Float[Array, '...'], radius: float) -> Float[Array, '...']

Gaspari-Cohn 5th-order piecewise polynomial taper.

With z = d / r:

\[ \begin{aligned} \rho(z) = \begin{cases} -\tfrac14 z^5 + \tfrac12 z^4 + \tfrac58 z^3 - \tfrac53 z^2 + 1 & 0 \le z \le 1 \\ \tfrac1{12} z^5 - \tfrac12 z^4 + \tfrac58 z^3 + \tfrac53 z^2 - 5 z + 4 - \tfrac{2}{3 z} & 1 < z \le 2 \\ 0 & z > 2. \end{cases} \end{aligned} \]

Properties:

• Compact support — exactly zero for d > 2 r.
• C² smoothness — value, first and second derivatives all match at z = 1 and z = 2.
• Positive definite — the gold standard for covariance localization in operational NWP and ocean DA.

Delegates to :func:gaussx.gaspari_cohn, which parameterises by the compact-support radius c = 2 r and guards the 2/(3 z) far-branch term so reverse-mode gradients stay finite at d = 0.

Parameters:

Name	Type	Description	Default
distances	Float[Array, '...']	Non-negative distance array, any shape.	required
radius	float	Positive localization half-width r. Compact support at 2 r.	required

Returns:

Type	Description
Float[Array, '...']	Taper weights in [0, 1] with the same shape as distances.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import gaspari_cohn
>>> d = jnp.array([0.0, 1.0, 2.0, 3.0])
>>> w = gaspari_cohn(d, radius=1.0)
>>> w[0], w[-1]  # full weight at zero distance, zero beyond 2r
(Array(1., dtype=float32), Array(0., dtype=float32))

Reference

Gaspari, G. & Cohn, S. E. (1999). Construction of correlation functions in two and three dimensions. Q. J. R. Meteorol. Soc., 125, 723–757.

Source code in src/filterax/_primitives/_localization.py

def gaspari_cohn(
    distances: Float[Array, "..."],
    radius: float,
) -> Float[Array, "..."]:
    r"""Gaspari-Cohn 5th-order piecewise polynomial taper.

    With ``z = d / r``:

    $$
    \begin{aligned}
    \rho(z) = \begin{cases}
      -\tfrac14 z^5 + \tfrac12 z^4 + \tfrac58 z^3 - \tfrac53 z^2 + 1
        & 0 \le z \le 1 \\
      \tfrac1{12} z^5 - \tfrac12 z^4 + \tfrac58 z^3 + \tfrac53 z^2
        - 5 z + 4 - \tfrac{2}{3 z}
        & 1 < z \le 2 \\
      0 & z > 2.
    \end{cases}
    \end{aligned}
    $$

    Properties:

    * **Compact support** — exactly zero for ``d > 2 r``.
    * **C² smoothness** — value, first and second derivatives all match
      at ``z = 1`` and ``z = 2``.
    * **Positive definite** — the gold standard for covariance
      localization in operational NWP and ocean DA.

    Delegates to :func:`gaussx.gaspari_cohn`, which parameterises by
    the compact-support radius ``c = 2 r`` and guards the ``2/(3 z)``
    far-branch term so reverse-mode gradients stay finite at ``d = 0``.

    Args:
        distances: Non-negative distance array, any shape.
        radius: Positive localization half-width ``r``. Compact support
            at ``2 r``.

    Returns:
        Taper weights in ``[0, 1]`` with the same shape as ``distances``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import gaspari_cohn
        >>> d = jnp.array([0.0, 1.0, 2.0, 3.0])
        >>> w = gaspari_cohn(d, radius=1.0)
        >>> w[0], w[-1]  # full weight at zero distance, zero beyond 2r
        (Array(1., dtype=float32), Array(0., dtype=float32))

    Reference:
        Gaspari, G. & Cohn, S. E. (1999). *Construction of correlation
        functions in two and three dimensions.* Q. J. R. Meteorol. Soc.,
        125, 723–757.
    """
    return gaussx.gaspari_cohn(distances, 2.0 * radius)

gaussian_taper(distances: Float[Array, '...'], radius: float) -> Float[Array, '...']

Gaussian taper ρ(d) = exp(−d² / (2 r²)).

Infinitely smooth (C^∞) and positive definite, but not compactly supported — decays exponentially but never reaches zero. Useful as a simpler alternative to Gaspari-Cohn when compact support is not required (e.g., adjoint sensitivity studies).

Parameters:

Name	Type	Description	Default
distances	Float[Array, '...']	Non-negative distance array, any shape.	required
radius	float	Positive characteristic length scale r.	required

Returns:

Type	Description
Float[Array, '...']	Taper weights in (0, 1] with the same shape as distances.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import gaussian_taper
>>> gaussian_taper(jnp.array([0.0, 1.0]), radius=1.0)
Array([1.        , 0.60653067], dtype=float32)

Source code in src/filterax/_primitives/_localization.py

def gaussian_taper(
    distances: Float[Array, "..."],
    radius: float,
) -> Float[Array, "..."]:
    r"""Gaussian taper ``ρ(d) = exp(−d² / (2 r²))``.

    Infinitely smooth (``C^∞``) and positive definite, but **not**
    compactly supported — decays exponentially but never reaches zero.
    Useful as a simpler alternative to Gaspari-Cohn when compact support
    is not required (e.g., adjoint sensitivity studies).

    Args:
        distances: Non-negative distance array, any shape.
        radius: Positive characteristic length scale ``r``.

    Returns:
        Taper weights in ``(0, 1]`` with the same shape as ``distances``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import gaussian_taper
        >>> gaussian_taper(jnp.array([0.0, 1.0]), radius=1.0)
        Array([1.        , 0.60653067], dtype=float32)
    """
    return jnp.exp(-(distances**2) / (2.0 * radius**2))

soar_taper(distances: Float[Array, '...'], radius: float) -> Float[Array, '...']

Second-Order Auto-Regressive taper (Thiebaux & Pedder 1987).

\[ \rho(d) = \left( 1 + \frac{d}{r} \right) e^{-d / r}. \]

Properties:

• C¹ smooth — value and first derivative continuous, second is not (compare with Gaspari-Cohn's C²).
• Positive definite — Schur product preserves PSD.
• Approximate compact support — decays to < 0.01 by d ≈ 5 r; finite but exponentially small beyond that.

Often used as a correlation model for background error covariance in operational variational systems (e.g. Met Office VAR). Simpler than Gaspari-Cohn while still being positive definite — useful as a localization taper when the strict compact support of :func:gaspari_cohn is not required.

Parameters:

Name	Type	Description	Default
distances	Float[Array, '...']	Non-negative distance array, any shape.	required
radius	float	Positive characteristic length scale r.	required

Returns:

Type	Description
Float[Array, '...']	Taper weights in (0, 1] with the same shape as distances.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import soar_taper
>>> soar_taper(jnp.array([0.0, 1.0]), radius=1.0)
Array([1.       , 0.7357589], dtype=float32)

Reference

Thiebaux, H. J. & Pedder, M. A. (1987). Spatial Objective Analysis. Academic Press.

Source code in src/filterax/_primitives/_localization.py

def soar_taper(
    distances: Float[Array, "..."],
    radius: float,
) -> Float[Array, "..."]:
    r"""Second-Order Auto-Regressive taper (Thiebaux & Pedder 1987).

    $$
    \rho(d) = \left( 1 + \frac{d}{r} \right) e^{-d / r}.
    $$

    Properties:

    * **C¹ smooth** — value and first derivative continuous, second is
      not (compare with Gaspari-Cohn's C²).
    * **Positive definite** — Schur product preserves PSD.
    * **Approximate compact support** — decays to ``< 0.01`` by
      ``d ≈ 5 r``; finite but exponentially small beyond that.

    Often used as a correlation model for background error covariance
    in operational variational systems (e.g. Met Office VAR). Simpler
    than Gaspari-Cohn while still being positive definite — useful as
    a localization taper when the strict compact support of
    :func:`gaspari_cohn` is not required.

    Args:
        distances: Non-negative distance array, any shape.
        radius: Positive characteristic length scale ``r``.

    Returns:
        Taper weights in ``(0, 1]`` with the same shape as ``distances``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import soar_taper
        >>> soar_taper(jnp.array([0.0, 1.0]), radius=1.0)
        Array([1.       , 0.7357589], dtype=float32)

    Reference:
        Thiebaux, H. J. & Pedder, M. A. (1987). *Spatial Objective
        Analysis.* Academic Press.
    """
    z = distances / radius
    return (1.0 + z) * jnp.exp(-z)

hard_cutoff(distances: Float[Array, '...'], radius: float) -> Float[Array, '...']

Binary cutoff ρ(d) = 𝟙{d ≤ r}.

Discontinuous at d = r; not positive definite in general (the resulting localized matrix may lose PSD-ness, which can cause filter instability). Provided as a debugging primitive only; prefer :func:gaspari_cohn or :func:gaussian_taper in production.

Parameters:

Name	Type	Description	Default
distances	Float[Array, '...']	Non-negative distance array, any shape.	required
radius	float	Positive cutoff radius r.	required

Returns:

Type	Description
Float[Array, '...']	Indicator weights (0.0 or 1.0) with the same shape as
Float[Array, '...']	distances, same dtype.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import hard_cutoff
>>> hard_cutoff(jnp.array([0.5, 1.0, 1.5]), radius=1.0)
Array([1., 1., 0.], dtype=float32)

Source code in src/filterax/_primitives/_localization.py

def hard_cutoff(
    distances: Float[Array, "..."],
    radius: float,
) -> Float[Array, "..."]:
    r"""Binary cutoff ``ρ(d) = 𝟙{d ≤ r}``.

    Discontinuous at ``d = r``; **not** positive definite in general
    (the resulting localized matrix may lose PSD-ness, which can cause
    filter instability). Provided as a debugging primitive only; prefer
    :func:`gaspari_cohn` or :func:`gaussian_taper` in production.

    Args:
        distances: Non-negative distance array, any shape.
        radius: Positive cutoff radius ``r``.

    Returns:
        Indicator weights (``0.0`` or ``1.0``) with the same shape as
        ``distances``, same dtype.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import hard_cutoff
        >>> hard_cutoff(jnp.array([0.5, 1.0, 1.5]), radius=1.0)
        Array([1., 1., 0.], dtype=float32)
    """
    return jnp.where(distances <= radius, 1.0, 0.0).astype(distances.dtype)

Distances & taper matrices

localization_matrix fuses a pairwise distance computation with the Gaspari-Cohn taper, producing the dense ρ consumed by localize and localized_kalman_gain. Use haversine_distance as the metric for spherical (lat, lon)-in-radians grids.

filterax — differentiable ensemble Kalman filters and processes for JAX.

localization_matrix(coords_a: Float[Array, 'N_a D'], coords_b: Float[Array, 'N_b D'], radius: float, *, metric: Callable[[Float[Array, 'N_a D'], Float[Array, 'N_b D']], Float[Array, 'N_a N_b']] = euclidean_distance) -> Float[Array, 'N_a N_b']

Pairwise Gaspari-Cohn taper matrix ρᵢⱼ = ρ(d(aᵢ, bⱼ); r).

Builds the dense localization matrix consumed by :func:filterax.localize and :func:filterax.localized_kalman_gain directly from coordinates, fusing the distance computation with the taper. Delegates to :func:gaussx.localization_matrix using filterax's half-width convention (compact support at 2 r), so localization_matrix(a, b, r) matches gaspari_cohn(metric(a, b), r) entry for entry.

Parameters:

Name	Type	Description	Default
coords_a	Float[Array, 'N_a D']	First coordinate set, shape (N_a, D).	required
coords_b	Float[Array, 'N_b D']	Second coordinate set, shape (N_b, D).	required
radius	float	Positive localization half-width r; weights vanish beyond distance 2 r.	required
metric	Callable[[Float[Array, 'N_a D'], Float[Array, 'N_b D']], Float[Array, 'N_a N_b']]	Pairwise distance function (coords_a, coords_b) -> (N_a, N_b). Defaults to :func:filterax.euclidean_distance; pass :func:filterax.haversine_distance for spherical (lat, lon)-in-radians grids.	euclidean_distance

Returns:

Type	Description
Float[Array, 'N_a N_b']	Taper matrix in [0, 1] of shape (N_a, N_b).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import localization_matrix
>>> grid = jnp.arange(3.0)[:, None]  # three points on a line
>>> rho = localization_matrix(grid, grid, radius=1.0)
>>> rho.shape, float(rho[0, 0]), float(rho[0, 2])
((3, 3), 1.0, 0.0)

Source code in src/filterax/_primitives/_localization.py

def localization_matrix(
    coords_a: Float[Array, "N_a D"],
    coords_b: Float[Array, "N_b D"],
    radius: float,
    *,
    metric: Callable[
        [Float[Array, "N_a D"], Float[Array, "N_b D"]], Float[Array, "N_a N_b"]
    ] = euclidean_distance,
) -> Float[Array, "N_a N_b"]:
    r"""Pairwise Gaspari-Cohn taper matrix ``ρᵢⱼ = ρ(d(aᵢ, bⱼ); r)``.

    Builds the dense localization matrix consumed by
    :func:`filterax.localize` and :func:`filterax.localized_kalman_gain`
    directly from coordinates, fusing the distance computation with the
    taper. Delegates to :func:`gaussx.localization_matrix` using
    filterax's half-width convention (compact support at ``2 r``), so
    ``localization_matrix(a, b, r)`` matches
    ``gaspari_cohn(metric(a, b), r)`` entry for entry.

    Args:
        coords_a: First coordinate set, shape ``(N_a, D)``.
        coords_b: Second coordinate set, shape ``(N_b, D)``.
        radius: Positive localization half-width ``r``; weights vanish
            beyond distance ``2 r``.
        metric: Pairwise distance function
            ``(coords_a, coords_b) -> (N_a, N_b)``. Defaults to
            :func:`filterax.euclidean_distance`; pass
            :func:`filterax.haversine_distance` for spherical
            (lat, lon)-in-radians grids.

    Returns:
        Taper matrix in ``[0, 1]`` of shape ``(N_a, N_b)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import localization_matrix
        >>> grid = jnp.arange(3.0)[:, None]  # three points on a line
        >>> rho = localization_matrix(grid, grid, radius=1.0)
        >>> rho.shape, float(rho[0, 0]), float(rho[0, 2])
        ((3, 3), 1.0, 0.0)
    """
    return gaussx.localization_matrix(coords_a, coords_b, 2.0 * radius, metric=metric)

euclidean_distance(coords_a: Float[Array, 'Na D'], coords_b: Float[Array, 'Nb D']) -> Float[Array, 'Na Nb']

Pairwise Euclidean distances ||a_i - b_j||.

A default metric for localization_matrix. Builds on stable_squared_distances and takes a gradient-safe square root so zero distances (e.g. the diagonal of a self-distance matrix) do not produce NaN gradients.

Parameters:

Name	Type	Description	Default
coords_a	Float[Array, 'Na D']	First set of points, shape (Na, D).	required
coords_b	Float[Array, 'Nb D']	Second set of points, shape (Nb, D).	required

Returns:

Type	Description
Float[Array, 'Na Nb']	Distance matrix of shape (Na, Nb).

Source code in .venv/lib/python3.12/site-packages/gaussx/_inference/_localization.py

def euclidean_distance(
    coords_a: Float[Array, "Na D"],
    coords_b: Float[Array, "Nb D"],
) -> Float[Array, "Na Nb"]:
    """Pairwise Euclidean distances ``||a_i - b_j||``.

    A default ``metric`` for `localization_matrix`. Builds on
    `stable_squared_distances` and takes a gradient-safe square root so
    zero distances (e.g. the diagonal of a self-distance matrix) do not produce
    ``NaN`` gradients.

    Args:
        coords_a: First set of points, shape ``(Na, D)``.
        coords_b: Second set of points, shape ``(Nb, D)``.

    Returns:
        Distance matrix of shape ``(Na, Nb)``.
    """
    sq = stable_squared_distances(
        coords_a,
        coords_b,
        compute_dtype=coords_a.dtype,
        accumulate_dtype=coords_a.dtype,
    )
    sq_safe = jnp.where(sq > 0.0, sq, 1.0)
    return jnp.where(sq > 0.0, jnp.sqrt(sq_safe), 0.0)

haversine_distance(coords_a: Float[Array, 'Na 2'], coords_b: Float[Array, 'Nb 2'], radius: float = 6371000.0) -> Float[Array, 'Na Nb']

Pairwise great-circle (haversine) distances on a sphere.

A metric for localization_matrix on geophysical grids. Coordinates are (latitude, longitude) in radians.

Parameters:

Name	Type	Description	Default
coords_a	Float[Array, 'Na 2']	First set of points (lat, lon) in radians, shape (Na, 2).	required
coords_b	Float[Array, 'Nb 2']	Second set of points (lat, lon) in radians, shape (Nb, 2).	required
radius	float	Sphere radius in the units of the returned distance (default the Earth mean radius, 6.371e6 m).	6371000.0

Returns:

Type	Description
Float[Array, 'Na Nb']	Great-circle distance matrix of shape (Na, Nb).

Source code in .venv/lib/python3.12/site-packages/gaussx/_inference/_localization.py

def haversine_distance(
    coords_a: Float[Array, "Na 2"],
    coords_b: Float[Array, "Nb 2"],
    radius: float = 6.371e6,
) -> Float[Array, "Na Nb"]:
    """Pairwise great-circle (haversine) distances on a sphere.

    A ``metric`` for `localization_matrix` on geophysical grids.
    Coordinates are ``(latitude, longitude)`` in **radians**.

    Args:
        coords_a: First set of points ``(lat, lon)`` in radians, shape ``(Na, 2)``.
        coords_b: Second set of points ``(lat, lon)`` in radians, shape ``(Nb, 2)``.
        radius: Sphere radius in the units of the returned distance (default the
            Earth mean radius, ``6.371e6`` m).

    Returns:
        Great-circle distance matrix of shape ``(Na, Nb)``.
    """
    lat_a = coords_a[:, 0][:, None]
    lon_a = coords_a[:, 1][:, None]
    lat_b = coords_b[:, 0][None, :]
    lon_b = coords_b[:, 1][None, :]
    dlat = lat_b - lat_a
    dlon = lon_b - lon_a
    h = (
        jnp.sin(dlat / 2.0) ** 2
        + jnp.cos(lat_a) * jnp.cos(lat_b) * jnp.sin(dlon / 2.0) ** 2
    )
    return 2.0 * radius * jnp.arcsin(jnp.sqrt(jnp.clip(h, 0.0, 1.0)))

Applying localization

localize is the Schur product itself; adaptive_localization builds the weight matrix from the ensemble's own sampling-error statistics (Anderson 2007) instead of a prescribed distance taper — useful when no meaningful distance metric exists (e.g. parameter spaces).

filterax — differentiable ensemble Kalman filters and processes for JAX.

localize(cov: Float[Array, 'M N'], taper: Float[Array, 'M N']) -> Float[Array, 'M N']

Apply localization via Schur (element-wise) product.

Computes P_loc = ρ ∘ P where ρ ∘ P denotes elementwise multiplication. Works on covariance, gain, or any other dense (M, N) matrix; for structured operators apply localization closer to where the matrix is consumed.

Parameters:

Name	Type	Description	Default
cov	Float[Array, 'M N']	Dense covariance or gain matrix of shape (M, N).	required
taper	Float[Array, 'M N']	Taper matrix of the same shape; elementwise multipliers in [0, 1].	required

Returns:

Type	Description
Float[Array, 'M N']	Localized matrix of shape (M, N).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import localize
>>> cov = jnp.array([[1.0, 0.5], [0.5, 1.0]])
>>> taper = jnp.array([[1.0, 0.0], [0.0, 1.0]])
>>> localize(cov, taper)  # off-diagonal entries suppressed
Array([[1., 0.],
       [0., 1.]], dtype=float32)

Source code in src/filterax/_primitives/_localization.py

def localize(
    cov: Float[Array, "M N"],
    taper: Float[Array, "M N"],
) -> Float[Array, "M N"]:
    r"""Apply localization via Schur (element-wise) product.

    Computes ``P_loc = ρ ∘ P`` where ``ρ ∘ P`` denotes
    elementwise multiplication. Works on covariance, gain, or any other
    dense ``(M, N)`` matrix; for structured operators apply localization
    closer to where the matrix is consumed.

    Args:
        cov: Dense covariance or gain matrix of shape ``(M, N)``.
        taper: Taper matrix of the same shape; elementwise multipliers
            in ``[0, 1]``.

    Returns:
        Localized matrix of shape ``(M, N)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import localize
        >>> cov = jnp.array([[1.0, 0.5], [0.5, 1.0]])
        >>> taper = jnp.array([[1.0, 0.0], [0.0, 1.0]])
        >>> localize(cov, taper)  # off-diagonal entries suppressed
        Array([[1., 0.],
               [0., 1.]], dtype=float32)
    """
    return cov * taper

localized_kalman_gain(particles: Float[Array, 'N_e N_x'], obs_particles: Float[Array, 'N_e N_y'], obs_noise: lx.AbstractLinearOperator, rho_xy: Float[Array, 'N_x N_y'], rho_yy: Float[Array, 'N_y N_y'], *, solver: gaussx.AbstractSolverStrategy | None = None) -> Float[Array, 'N_x N_y']

Kalman gain with Schur-product (Hadamard) localization.

Thin wrapper over :func:gaussx.localized_kalman_gain with the EnKF Bessel convention. Computes

\[ K = (\rho^{xy} \circ C^{xH})\,(\rho^{yy} \circ C^{HH} + R)^{-1}. \]

Tapering both covariances with the Schur (element-wise) product suppresses spurious long-range correlations; by the Schur product theorem the tapered innovation covariance stays PSD when ρʸʸ is (use :func:filterax.localization_matrix / :func:filterax.gaspari_cohn to build the tapers).

Unlike :func:filterax.kalman_gain, the Hadamard product destroys the low-rank structure of Cᴴᴴ, so the innovation covariance is materialised densely — cost O(Nₑ Nₓ Nᵧ + Nᵧ³). With ρ ≡ 1 this reduces exactly to the unlocalized gain.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Prior ensemble in state space, shape (Nₑ, Nₓ).	required
obs_particles	Float[Array, 'N_e N_y']	H applied to each member, shape (Nₑ, Nᵧ).	required
obs_noise	AbstractLinearOperator	Observation error covariance R as a linear operator.	required
rho_xy	Float[Array, 'N_x N_y']	State-observation taper, shape (Nₓ, Nᵧ).	required
rho_yy	Float[Array, 'N_y N_y']	Observation-observation taper, shape (Nᵧ, Nᵧ).	required
solver	AbstractSolverStrategy | None	Optional :class:gaussx.AbstractSolverStrategy. None uses structural dispatch on the dense innovation covariance.	None

Returns:

Type	Description
Float[Array, 'N_x N_y']	Dense localized Kalman gain of shape (Nₓ, Nᵧ).

Raises:

Type	Description
ValueError	if Nₑ < 2.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import kalman_gain, localized_kalman_gain
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> obs_particles = particles[:, :1]
>>> R = lx.DiagonalLinearOperator(0.5 * jnp.ones(1))
>>> ones = jnp.ones((2, 1)), jnp.ones((1, 1))
>>> K_loc = localized_kalman_gain(particles, obs_particles, R, *ones)
>>> bool(jnp.allclose(K_loc, kalman_gain(particles, obs_particles, R)))
True

Source code in src/filterax/_primitives/_gain.py

def localized_kalman_gain(
    particles: Float[Array, "N_e N_x"],
    obs_particles: Float[Array, "N_e N_y"],
    obs_noise: lx.AbstractLinearOperator,
    rho_xy: Float[Array, "N_x N_y"],
    rho_yy: Float[Array, "N_y N_y"],
    *,
    solver: gaussx.AbstractSolverStrategy | None = None,
) -> Float[Array, "N_x N_y"]:
    r"""Kalman gain with Schur-product (Hadamard) localization.

    Thin wrapper over :func:`gaussx.localized_kalman_gain` with the EnKF
    Bessel convention. Computes

    $$
    K = (\rho^{xy} \circ C^{xH})\,(\rho^{yy} \circ C^{HH} + R)^{-1}.
    $$

    Tapering both covariances with the Schur
    (element-wise) product suppresses spurious long-range correlations;
    by the Schur product theorem the tapered innovation covariance stays
    PSD when ``ρʸʸ`` is (use :func:`filterax.localization_matrix` /
    :func:`filterax.gaspari_cohn` to build the tapers).

    Unlike :func:`filterax.kalman_gain`, the Hadamard product destroys
    the low-rank structure of ``Cᴴᴴ``, so the innovation covariance is
    materialised densely — cost ``O(Nₑ Nₓ Nᵧ + Nᵧ³)``. With
    ``ρ ≡ 1`` this reduces exactly to the unlocalized gain.

    Args:
        particles: Prior ensemble in state space, shape ``(Nₑ, Nₓ)``.
        obs_particles: ``H`` applied to each member, shape ``(Nₑ, Nᵧ)``.
        obs_noise: Observation error covariance ``R`` as a linear operator.
        rho_xy: State-observation taper, shape ``(Nₓ, Nᵧ)``.
        rho_yy: Observation-observation taper, shape ``(Nᵧ, Nᵧ)``.
        solver: Optional :class:`gaussx.AbstractSolverStrategy`. ``None``
            uses structural dispatch on the dense innovation covariance.

    Returns:
        Dense localized Kalman gain of shape ``(Nₓ, Nᵧ)``.

    Raises:
        ValueError: if ``Nₑ < 2``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import kalman_gain, localized_kalman_gain
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> obs_particles = particles[:, :1]
        >>> R = lx.DiagonalLinearOperator(0.5 * jnp.ones(1))
        >>> ones = jnp.ones((2, 1)), jnp.ones((1, 1))
        >>> K_loc = localized_kalman_gain(particles, obs_particles, R, *ones)
        >>> bool(jnp.allclose(K_loc, kalman_gain(particles, obs_particles, R)))
        True
    """
    check_ensemble_size(particles.shape[0])
    return gaussx.localized_kalman_gain(
        particles, obs_particles, obs_noise, rho_xy, rho_yy, solver=solver, bessel=True
    )

adaptive_localization(state_particles: Float[Array, 'N_e N_x'], obs_particles: Float[Array, 'N_e N_y'], *, significance: float = 1.0) -> Float[Array, 'N_x N_y']

Adaptive localization weights from ensemble correlations (Anderson 2007).

Estimates a per-pair localization weight on the (Nₓ, Nᵧ) Kalman gain by checking whether each sample correlation r_{ij} = Cˣᴴ_{ij} / (σ_xᵢ σ_yⱼ) exceeds its sampling uncertainty:

\[ \operatorname{se}(r) \approx \frac{1 - r^{2}}{\sqrt{N_{e} - 2}}. \]

Correlations smaller than significance · se(r) are zeroed. Correlations above that threshold are kept at unit weight (a hard-mask variant). The resulting weight matrix is multiplied into the Kalman gain (Schur product) to suppress spurious long-range correlations driven by sampling noise rather than physical structure.

Unlike a distance-based taper, this localizer is data-driven — no radius to tune. The trade-off: it requires enough ensemble members for the correlation noise floor to be informative (typically Nₑ ≥ 20) and an extra O(Nₑ Nₓ Nᵧ) work per cycle.

Parameters:

Name	Type	Description	Default
state_particles	Float[Array, 'N_e N_x']	Prior ensemble in state space, (Nₑ, Nₓ).	required
obs_particles	Float[Array, 'N_e N_y']	Prior ensemble in obs space (Nₑ, Nᵧ) — H applied to each member.	required
significance	float	Threshold multiplier on se(r). Larger → more aggressive zeroing.	1.0

Returns:

Type	Description
Float[Array, 'N_x N_y']	Weight matrix ρ ∈ ℝ^{Nₓ × Nᵧ} with entries in {0, 1}.

Raises:

Type	Description
ValueError	if Nₑ < 3 (the √(Nₑ − 2) noise floor is undefined for the smallest ensembles).

Examples:

>>> import jax
>>> from filterax import adaptive_localization
>>> state = jax.random.normal(jax.random.key(0), (20, 2))
>>> obs_p = state[:, :1]  # obs perfectly correlated with dim 0
>>> w = adaptive_localization(state, obs_p)
>>> w.shape
(2, 1)
>>> w[0, 0]  # significant correlation kept at unit weight
Array(1., dtype=float32)

Reference

Anderson, J. L. (2007). Exploring the need for localization in ensemble data assimilation using a hierarchical ensemble filter. Physica D, 230, 99-111.

Source code in src/filterax/_primitives/_localization.py

def adaptive_localization(
    state_particles: Float[Array, "N_e N_x"],
    obs_particles: Float[Array, "N_e N_y"],
    *,
    significance: float = 1.0,
) -> Float[Array, "N_x N_y"]:
    r"""Adaptive localization weights from ensemble correlations (Anderson 2007).

    Estimates a per-pair localization weight on the ``(Nₓ, Nᵧ)`` Kalman
    gain by checking whether each sample correlation
    ``r_{ij} = Cˣᴴ_{ij} / (σ_xᵢ σ_yⱼ)`` exceeds its sampling
    uncertainty:

    $$
    \operatorname{se}(r) \approx \frac{1 - r^{2}}{\sqrt{N_{e} - 2}}.
    $$

    Correlations smaller than ``significance · se(r)`` are zeroed.
    Correlations above that threshold are kept at unit weight (a
    hard-mask variant). The resulting weight matrix is multiplied into
    the Kalman gain (Schur product) to suppress spurious long-range
    correlations driven by sampling noise rather than physical
    structure.

    Unlike a distance-based taper, this localizer is *data-driven* —
    no radius to tune. The trade-off: it requires enough ensemble
    members for the correlation noise floor to be informative
    (typically ``Nₑ ≥ 20``) and an extra ``O(Nₑ Nₓ Nᵧ)`` work per
    cycle.

    Args:
        state_particles: Prior ensemble in state space, ``(Nₑ, Nₓ)``.
        obs_particles: Prior ensemble in obs space ``(Nₑ, Nᵧ)`` —
            ``H`` applied to each member.
        significance: Threshold multiplier on ``se(r)``. Larger →
            more aggressive zeroing.

    Returns:
        Weight matrix ``ρ ∈ ℝ^{Nₓ × Nᵧ}`` with entries in ``{0, 1}``.

    Raises:
        ValueError: if ``Nₑ < 3`` (the ``√(Nₑ − 2)`` noise floor is
            undefined for the smallest ensembles).

    Examples:
        >>> import jax
        >>> from filterax import adaptive_localization
        >>> state = jax.random.normal(jax.random.key(0), (20, 2))
        >>> obs_p = state[:, :1]  # obs perfectly correlated with dim 0
        >>> w = adaptive_localization(state, obs_p)
        >>> w.shape
        (2, 1)
        >>> w[0, 0]  # significant correlation kept at unit weight
        Array(1., dtype=float32)

    Reference:
        Anderson, J. L. (2007). *Exploring the need for localization
        in ensemble data assimilation using a hierarchical ensemble
        filter.* Physica D, 230, 99-111.
    """
    N_e = state_particles.shape[0]
    # √(N_e − 2) appears in the sampling-noise denominator; we also
    # rely on the Bessel-corrected std (N_e ≥ 2). Demand N_e ≥ 3.
    check_ensemble_size(N_e)
    if N_e < 3:
        raise ValueError(
            "adaptive_localization needs N_e ≥ 3 so the √(N_e − 2) "
            f"sampling-noise floor is defined; got N_e={N_e}."
        )
    state_mean = jnp.mean(state_particles, axis=0)
    obs_mean = jnp.mean(obs_particles, axis=0)
    state_anom = state_particles - state_mean[None, :]
    obs_anom = obs_particles - obs_mean[None, :]
    state_std = jnp.std(state_particles, axis=0, ddof=1)
    obs_std = jnp.std(obs_particles, axis=0, ddof=1)

    # Sample cross-correlation r_{ij} (Bessel-corrected).
    cross_cov = jnp.einsum("ex,ey->xy", state_anom, obs_anom) / (N_e - 1)
    denom = jnp.maximum(jnp.outer(state_std, obs_std), 1e-30)
    corr = cross_cov / denom
    # Sampling-noise floor; the (1 − r²) factor uses the absolute
    # value so the floor stays positive even for slightly noisy r.
    se = (1.0 - corr * corr) / jnp.sqrt(jnp.asarray(N_e - 2, dtype=corr.dtype))
    keep = jnp.abs(corr) > significance * jnp.maximum(se, 0.0)
    return keep.astype(state_particles.dtype)