Ensemble Kalman Filtering

def no_covariance_modifier(covariance: Array) -> Array:
    """Return an empirical covariance unchanged.

    The identity covariance modifier, used as the default when no modification
    (e.g. localization) is requested.

    Args:
        covariance: Empirical covariance matrix.

    Returns:
        The covariance matrix, unchanged.
    """
    return covariance

def predict(
    key: KeyArray,
    ensemble: Array,
    dynamics_fn: DynamicsFn,
    inflation: float = 0.0,
) -> Array:
    """Propagate ensemble members through an arbitrary simulator p(x_{t+1} | x_t).

    Args:
        key: JAX PRNG key.
        ensemble: Ensemble of state vectors, shape (N, x_dim).
        dynamics_fn: Dynamics function mapping (state, key) -> state.
        inflation: Multiplicative inflation factor applied to ensemble deviations.

    Returns:
        Predicted ensemble, shape (N, x_dim).
    """
    N, x_dim = ensemble.shape

    # Propagate each member through the dynamics
    keys = random.split(key, N)
    propagated = jax.vmap(dynamics_fn, (0, 0))(ensemble, keys)

    # Apply multiplicative inflation
    mean = jnp.mean(propagated, axis=0)
    propagated = mean + (1 + inflation) * (propagated - mean)

    return propagated

def update(
    key: KeyArray,
    predicted_ensemble: Array,
    observation_fn: ObservationFn,
    chol_R: Array,
    y: Array,
    perturbed_obs: bool = True,
    cross_covariance_modifier: CrossCovarianceModifier = no_covariance_modifier,
    construct_chol_innovation_covariance: ConstructCholInnovationCovariance
    | None = None,
    ensemble_subspace: bool = False,
) -> tuple[Array, ScalarArray]:
    """Update ensemble members with an observation using the EnKF update.

    NaNs in ``y`` are treated as missing dimensions and are excluded from the
    update. When ``y`` is entirely NaN, the update is a no-op: the predicted
    ensemble is returned unchanged with zero log-likelihood contribution.

    Args:
        key: JAX PRNG key.
        predicted_ensemble: Predicted ensemble, shape (N, x_dim).
        observation_fn: Observation function mapping state -> obs.
        chol_R: Generalized Cholesky factor of the observation noise covariance,
            shape (y_dim, y_dim). Square roots that are not generalized Cholesky
            factors, such as a symmetric R ** 0.5, are not supported.
            When ``ensemble_subspace`` is True this may instead be a scalar (a multiple
            of the identity) or a 1D array of shape (y_dim,) (a diagonal factor).
            Prefer those forms when the structure allows: a 2D factor must be applied
            by triangular solve, at O(N * y_dim ** 2) instead of O(N * y_dim), and
            stored densely in y_dim ** 2 entries. On either path, a 2D factor is also
            refactored in O(y_dim ** 3) at any step where ``y`` has missing values,
            whereas scalar and 1D factors handle missingness in O(y_dim). Steps with
            nothing missing skip the refactor.
        y: Observation vector, shape (y_dim,). NaNs indicate missing dimensions.
        perturbed_obs: If True, use perturbed observations (stochastic EnKF).
            If False, use deterministic update.
        cross_covariance_modifier: Function that modifies the empirical
            state-observation cross-covariance, shape (x_dim, y_dim), and returns
            an array with the same shape. Defaults to the identity.
        construct_chol_innovation_covariance: Optional function that
            receives normalized observation deviations with shape (y_dim, N) and
            ``chol_R`` with shape (y_dim, y_dim). It must return a generalized
            Cholesky factor of the complete innovation covariance with shape
            (y_dim, y_dim). The deviations have already been divided by
            ``sqrt(N - 1)``. Both inputs use the original observation order.
            ``None`` uses the standard, unlocalized square-root construction.
        ensemble_subspace: If True, perform the analysis in the N-dimensional
            ensemble subspace. This is algebraically exact and costs
            O(N ** 2 * x_dim) in the state dimension rather than
            O(N * x_dim * y_dim), so it is preferable when ``N << y_dim``. It is
            incompatible with both localization arguments above, which it rejects.
            Defaults to False; the choice is never made automatically.

    Returns:
        Tuple of (updated_ensemble, log_likelihood).

    Raises:
        ValueError: If ``ensemble_subspace`` is combined with either localization
            argument, or if a non-2D ``chol_R`` is given without it.
    """
    if ensemble_subspace:
        if cross_covariance_modifier is not no_covariance_modifier:
            raise ValueError(
                "ensemble_subspace=True is incompatible with cross_covariance_modifier: "
                "the ensemble-subspace update never forms the state-observation "
                "cross-covariance, so there is nothing to modify."
            )
        if construct_chol_innovation_covariance is not None:
            raise ValueError(
                "ensemble_subspace=True is incompatible with "
                "construct_chol_innovation_covariance: tapering the innovation "
                "covariance destroys the rank-N structure the update relies on."
            )
        return _update_ensemble_subspace(
            key,
            predicted_ensemble,
            observation_fn,
            chol_R,
            y,
            perturbed_obs,
        )

    if jnp.ndim(chol_R) != 2:
        raise ValueError(
            "chol_R must be 2D, of shape (y_dim, y_dim). Scalar and diagonal factors "
            "are only supported with ensemble_subspace=True, because this path "
            "factorizes the dense y_dim x y_dim innovation covariance."
        )

    N, x_dim = predicted_ensemble.shape

    # Map ensemble to observation space
    y_pred = jax.vmap(observation_fn, (0,))(predicted_ensemble)
    x_mean = jnp.mean(predicted_ensemble, axis=0)
    x_dev = predicted_ensemble - x_mean

    missing = jnp.isnan(y)

    # Modify or construct covariances before reordering due to NaNs.
    argsort = jnp.argsort(missing, stable=True)
    original_y_dev = y_pred - jnp.mean(y_pred, axis=0)
    normalized_original_y_dev = original_y_dev.T / jnp.sqrt(N - 1)

    C_xy = x_dev.T @ original_y_dev / (N - 1)
    C_xy = cross_covariance_modifier(C_xy)

    if construct_chol_innovation_covariance is not None:
        original_chol_S = construct_chol_innovation_covariance(
            normalized_original_y_dev, chol_R
        )

    # Handle partially-missing observations by reordering and zeroing missing dims.
    # Use y_pred.T because y_pred is (N, y_dim) and we want to reorder along axis 0.
    # Refactoring chol_R is O(y_dim ** 3); skip it when nothing is missing, in which
    # case the reordering is the identity and the inputs are returned unchanged.
    flag, chol_R, y, y_pred = jax.lax.cond(
        jnp.any(missing),
        lambda args: collect_nans_chol(missing, *args[1:]),
        lambda args: args,
        (missing, chol_R, y, y_pred.T),
    )
    y_pred = y_pred.T
    y_dim = y.shape[0]

    y_mean = jnp.mean(y_pred, axis=0)
    y_dev = y_pred - y_mean
    C_xy = C_xy[:, argsort]
    C_xy = jnp.where(flag[None, :], 0.0, C_xy)

    if construct_chol_innovation_covariance is None:
        chol_S = tria(jnp.concatenate([y_dev.T / jnp.sqrt(N - 1), chol_R], axis=1))
    else:
        # The constructor sees the original indexing. Only collect and refactor its
        # result when dimensions are missing; otherwise preserve its returned factor.
        chol_S = jax.lax.cond(
            jnp.any(missing),
            lambda chol: collect_nans_chol(missing, chol)[1],
            lambda chol: chol,
            original_chol_S,
        )

    # Innovation per member
    if perturbed_obs:
        y_n = y[None, :] + (chol_R @ random.normal(key, (y_dim, N))).T
    else:
        y_n = jnp.broadcast_to(y[None, :], (N, y_dim))

    innovations = y_n - y_pred

    # Doing K = C_xy @ S^{-1}\delta right to left has cost O(Nd_y^2 + Nd_yd_x), left to right O(d_xd_y^2 + Nd_yd_x).
    # If N < d_x, then right to left is cheaper; otherwise left to right is cheaper.
    if N < x_dim:
        increment = cho_solve((chol_S, True), innovations.T).T @ C_xy.T
    else:
        increment = innovations @ cho_solve((chol_S, True), C_xy.T)

    # Update ensemble
    updated = predicted_ensemble + increment

    # Log-likelihood
    ll = multivariate_normal.logpdf(y, y_mean, chol_S, nan_support=False)

    return updated, jnp.asarray(ll)