cuthbert 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

cuthbertlib/discrete/smoothing.py

@@ -0,0 +1,35 @@
+"""Implements the discrete HMM smoothing associative operator."""
+
+import jax.numpy as jnp
+
+from cuthbertlib.types import Array, ArrayLike
+
+
+def get_reverse_kernel(x_t_dist: ArrayLike, trans_matrix: ArrayLike) -> Array:
+    r"""Computes reverse transition probabilities $p(x_{t-1} \mid x_{t}, \dots)$ for a discrete HMM.
+
+    Args:
+        x_t_dist: Array of shape (N,) where `x_t_dist[i]` = $p(x_{t} = i \mid \dots)$.
+        trans_matrix: Array of shape (N, N) where
+            `trans_matrix[i, j]` = $p(x_{t} = j \mid x_{t-1} = i)$.
+
+    Returns:
+        An (N, N) matrix `x_tm1_dist[i, j]` = $p(x_{t-1} = j \mid x_{t} = i, \dots)$.
+    """
+    x_t_dist, trans_matrix = jnp.asarray(x_t_dist), jnp.asarray(trans_matrix)
+    pred = jnp.dot(trans_matrix.T, x_t_dist)
+    x_tm1_dist = trans_matrix.T * x_t_dist[None, :] / pred[:, None]
+    return x_tm1_dist
+
+
+def smoothing_operator(elem_ij: Array, elem_jk: Array) -> Array:
+    """Binary associative operator for smoothing in discrete HMMs.
+
+    Args:
+        elem_ij: Smoothing scan element.
+        elem_jk: Smoothing scan element.
+
+    Returns:
+        The output of the associative operator applied to the input elements.
+    """
+    return elem_jk @ elem_ij

cuthbertlib/kalman/__init__.py

@@ -0,0 +1,4 @@
+from cuthbertlib.kalman import generate
+from cuthbertlib.kalman.filtering import predict
+from cuthbertlib.kalman.filtering import update as filter_update
+from cuthbertlib.kalman.smoothing import update as smoother_update

cuthbertlib/kalman/filtering.py

@@ -0,0 +1,213 @@
+"""Implements the square root parallel Kalman filter and associative variant."""
+
+from typing import NamedTuple
+
+import jax.numpy as jnp
+from jax.scipy.linalg import cho_solve, solve_triangular
+
+from cuthbertlib.linalg import tria
+from cuthbertlib.stats import multivariate_normal
+from cuthbertlib.stats.multivariate_normal import collect_nans_chol
+from cuthbertlib.types import Array, ArrayLike, ScalarArray
+
+
+class FilterScanElement(NamedTuple):
+    """Arrays carried through the Kalman filter scan."""
+
+    A: Array
+    b: Array
+    U: Array
+    eta: Array
+    Z: Array
+    ell: ScalarArray
+
+
+def predict(
+    m: ArrayLike,
+    chol_P: ArrayLike,
+    F: ArrayLike,
+    c: ArrayLike,
+    chol_Q: ArrayLike,
+) -> tuple[Array, Array]:
+    """Propagate the mean and square root covariance through linear Gaussian dynamics.
+
+    Args:
+        m: Mean of the state.
+        chol_P: Generalized Cholesky factor of the covariance of the state.
+        F: Transition matrix.
+        c: Transition shift.
+        chol_Q: Generalized Cholesky factor of the transition noise covariance.
+
+    Returns:
+        Propagated mean and square root covariance.
+
+    References:
+        Paper: G. J. Bierman, Factorization Methods for Discrete Sequential Estimation,
+        Code: https://github.com/EEA-sensors/sqrt-parallel-smoothers/tree/main/parsmooth/sequential
+    """
+    m, chol_P = jnp.asarray(m), jnp.asarray(chol_P)
+    F, c, chol_Q = jnp.asarray(F), jnp.asarray(c), jnp.asarray(chol_Q)
+    m1 = F @ m + c
+    A = jnp.concatenate([F @ chol_P, chol_Q], axis=1)
+    chol_P1 = tria(A)
+    return m1, chol_P1
+
+
+# Note that `update` is aliased as `kalman.filter_update` in `kalman.__init__.py`
+def update(
+    m: ArrayLike,
+    chol_P: ArrayLike,
+    H: ArrayLike,
+    d: ArrayLike,
+    chol_R: ArrayLike,
+    y: ArrayLike,
+    log_normalizing_constant: ArrayLike = 0.0,
+) -> tuple[tuple[Array, Array], Array]:
+    """Update the mean and square root covariance with a linear Gaussian observation.
+
+    Args:
+        m: Mean of the state.
+        chol_P: Generalized Cholesky factor of the covariance of the state.
+        H: Observation matrix.
+        d: Observation shift.
+        chol_R: Generalized Cholesky factor of the observation noise covariance.
+        y: Observation.
+        log_normalizing_constant: Optional input of log normalizing constant to be added to
+            log normalizing constant of the Bayesian update.
+
+    Returns:
+        Updated mean and square root covariance as well as the log marginal likelihood.
+
+    References:
+        Paper: G. J. Bierman, Factorization Methods for Discrete Sequential Estimation,
+        Code: https://github.com/EEA-sensors/sqrt-parallel-smoothers/tree/main/parsmooth/sequential
+    """
+    # Handle case where there is no observation
+    flag = jnp.isnan(y)
+    flag, chol_R, H, d, y = collect_nans_chol(flag, chol_R, H, d, y)
+
+    m, chol_P = jnp.asarray(m), jnp.asarray(chol_P)
+    H, d, chol_R = jnp.asarray(H), jnp.asarray(d), jnp.asarray(chol_R)
+    y = jnp.asarray(y)
+
+    n_y, n_x = H.shape
+
+    y_hat = H @ m + d
+    y_diff = y - y_hat
+
+    M = jnp.block(
+        [
+            [H @ chol_P, chol_R],
+            [chol_P, jnp.zeros((n_x, n_y), dtype=chol_P.dtype)],
+        ]
+    )
+    chol_S = tria(M)
+    chol_Py = chol_S[n_y:, n_y:]
+
+    Gmat = chol_S[n_y:, :n_y]
+    Imat = chol_S[:n_y, :n_y]
+
+    my = m + Gmat @ solve_triangular(Imat, y_diff, lower=True)
+
+    ell = multivariate_normal.logpdf(y, y_hat, Imat, nan_support=False)
+    return (my, chol_Py), jnp.asarray(ell + log_normalizing_constant)
+
+
+def associative_params_single(
+    F: Array, c: Array, chol_Q: Array, H: Array, d: Array, chol_R: Array, y: Array
+) -> FilterScanElement:
+    """Single time step for scan element for square root parallel Kalman filter.
+
+    Args:
+        F: State transition matrix.
+        c: State transition shift vector.
+        chol_Q: Generalized Cholesky factor of the state transition noise covariance.
+        H: Observation matrix.
+        d: Observation shift.
+        chol_R: Generalized Cholesky factor of the observation noise covariance.
+        y: Observation.
+
+    Returns:
+        Prepared scan element for the square root parallel Kalman filter.
+    """
+    # Handle case where there is no observation
+    flag = jnp.isnan(y)
+    flag, chol_R, H, d, y = collect_nans_chol(flag, chol_R, H, d, y)
+
+    ny, nx = H.shape
+
+    # joint over the predictive and the observation
+    Psi_ = jnp.block([[H @ chol_Q, chol_R], [chol_Q, jnp.zeros((nx, ny))]])
+
+    Tria_Psi_ = tria(Psi_)
+
+    Psi11 = Tria_Psi_[:ny, :ny]
+    Psi21 = Tria_Psi_[ny : ny + nx, :ny]
+    U = Tria_Psi_[ny : ny + nx, ny:]
+
+    # pre-compute inverse of Psi11: we apply it to matrices and vectors alike.
+    Psi11_inv = solve_triangular(Psi11, jnp.eye(ny), lower=True)
+
+    # predictive model given one observation
+    K = Psi21 @ Psi11_inv  # local Kalman gain
+    HF = H @ F  # temporary variable
+    A = F - K @ HF  # corrected transition matrix
+
+    b = c + K @ (y - H @ c - d)  # corrected transition offset
+
+    # information filter
+    Z = HF.T @ Psi11_inv.T
+    eta = Psi11_inv @ (y - H @ c - d)
+    eta = Z @ eta
+
+    if nx > ny:
+        Z = jnp.concatenate([Z, jnp.zeros((nx, nx - ny))], axis=1)
+    else:
+        Z = tria(Z)
+
+    # local log marginal likelihood
+    ell = jnp.asarray(
+        multivariate_normal.logpdf(y, H @ c + d, Psi11, nan_support=False)
+    )
+
+    return FilterScanElement(A, b, U, eta, Z, ell)
+
+
+def filtering_operator(
+    elem_i: FilterScanElement, elem_j: FilterScanElement
+) -> FilterScanElement:
+    """Binary associative operator for the square root Kalman filter.
+
+    Args:
+        elem_i: Filter scan element for the previous time step.
+        elem_j: Filter scan element for the current time step.
+
+    Returns:
+        FilterScanElement: The output of the associative operator applied to the input elements.
+    """
+    A1, b1, U1, eta1, Z1, ell1 = elem_i
+    A2, b2, U2, eta2, Z2, ell2 = elem_j
+
+    nx = Z2.shape[0]
+
+    Xi = jnp.block([[U1.T @ Z2, jnp.eye(nx)], [Z2, jnp.zeros_like(A1)]])
+    tria_xi = tria(Xi)
+    Xi11 = tria_xi[:nx, :nx]
+    Xi21 = tria_xi[nx : nx + nx, :nx]
+    Xi22 = tria_xi[nx : nx + nx, nx:]
+
+    tmp_1 = solve_triangular(Xi11, U1.T, lower=True).T
+    D_inv = jnp.eye(nx) - tmp_1 @ Xi21.T
+    tmp_2 = D_inv @ (b1 + U1 @ (U1.T @ eta2))
+
+    A = A2 @ D_inv @ A1
+    b = A2 @ tmp_2 + b2
+    U = tria(jnp.concatenate([A2 @ tmp_1, U2], axis=1))
+    eta = A1.T @ (D_inv.T @ (eta2 - Z2 @ (Z2.T @ b1))) + eta1
+    Z = tria(jnp.concatenate([A1.T @ Xi22, Z1], axis=1))
+
+    mu = cho_solve((U1, True), b1)
+    t1 = b1 @ mu - (eta2 + mu) @ tmp_2
+    ell = ell1 + ell2 - 0.5 * t1 + 0.5 * jnp.linalg.slogdet(D_inv)[1]
+
+    return FilterScanElement(A, b, U, eta, Z, ell)

cuthbertlib/kalman/generate.py

@@ -0,0 +1,85 @@
+"""Utilities to generate linear-Gaussian state-space models (LGSSMs)."""
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax import Array, random
+
+from cuthbertlib.types import KeyArray
+
+
+@partial(jax.jit, static_argnames=("x_dim", "y_dim", "num_time_steps"))
+def generate_lgssm(seed: int, x_dim: int, y_dim: int, num_time_steps: int):
+    """Generates an LGSSM along with a set of observations."""
+    key = random.key(seed)
+
+    key, init_key, sample_key, obs_model_key, obs_key = random.split(key, 5)
+    m0, chol_P0 = generate_init_model(init_key, x_dim)
+    x0 = m0 + chol_P0 @ random.normal(sample_key, (x_dim,))
+
+    # Generate an observation for time 0
+    H0, d0, chol_R0 = generate_obs_model(obs_model_key, x_dim, y_dim)
+    obs_noise = chol_R0 @ random.normal(obs_key, (y_dim,))
+    y0 = H0 @ x0 + d0 + obs_noise
+
+    def body(_x, _key):
+        trans_model_key, trans_key, obs_model_key, obs_key = random.split(_key, 4)
+
+        F, c, chol_Q = generate_trans_model(trans_model_key, x_dim)
+        state_noise = chol_Q @ random.normal(trans_key, (x_dim,))
+        x = F @ _x + c + state_noise
+
+        H, d, chol_R = generate_obs_model(obs_model_key, x_dim, y_dim)
+        obs_noise = chol_R @ random.normal(obs_key, (y_dim,))
+        y = H @ x + d + obs_noise
+
+        return x, (F, c, chol_Q, H, d, chol_R, y)
+
+    scan_keys = random.split(key, num_time_steps)
+    _, (Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys) = jax.lax.scan(body, x0, scan_keys)
+
+    # Prepend the observation model parameters and observation for t=0
+    Hs, ds, chol_Rs, ys = jax.tree.map(
+        lambda x, xs: jnp.concatenate([x[None], xs]),
+        (H0, d0, chol_R0, y0),
+        (Hs, ds, chol_Rs, ys),
+    )
+
+    return m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys
+
+
+def generate_cholesky_factor(key: KeyArray, dim: int) -> Array:
+    """Generates a random Cholesky factor (lower-triangular matrix)."""
+    chol_A = random.uniform(key, (dim, dim))
+    chol_A = chol_A.at[jnp.triu_indices(dim, 1)].set(0.0)
+    return chol_A
+
+
+def generate_init_model(key: KeyArray, x_dim: int) -> tuple[Array, Array]:
+    """Generates a random initial state for an LGSSM."""
+    keys = random.split(key)
+    m0 = random.normal(keys[0], (x_dim,))
+    chol_P0 = generate_cholesky_factor(keys[1], x_dim)
+    return m0, chol_P0
+
+
+def generate_trans_model(key: KeyArray, x_dim: int) -> tuple[Array, Array, Array]:
+    """Generates a random transition model for an LGSSM."""
+    keys = random.split(key, 3)
+    exp_eig_max = 0.75  # Chosen less than one to stop exploding states (in expectation)
+    F = exp_eig_max * random.normal(keys[0], (x_dim, x_dim)) / jnp.sqrt(x_dim)
+    c = 0.1 * random.normal(keys[1], (x_dim,))
+    chol_Q = generate_cholesky_factor(keys[2], x_dim)
+    return F, c, chol_Q
+
+
+def generate_obs_model(
+    key: KeyArray, x_dim: int, y_dim: int
+) -> tuple[Array, Array, Array]:
+    """Generates a random observation model for an LGSSM."""
+    keys = random.split(key, 3)
+    H = random.normal(keys[0], (y_dim, x_dim))
+    d = random.normal(keys[1], (y_dim,))
+    chol_R = generate_cholesky_factor(keys[2], y_dim)
+    return H, d, chol_R

cuthbertlib/kalman/sampling.py

@@ -0,0 +1,68 @@
+"""Implements the square root parallel Kalman associative operator for sampling.
+
+Samples from the smoothing distribution without doing the smoothing scan for means
+and (chol) covariances.
+"""
+
+from typing import NamedTuple, Sequence
+
+import jax
+import jax.numpy as jnp
+
+from cuthbertlib.kalman.smoothing import associative_params_single
+from cuthbertlib.types import Array, ArrayLike
+
+
+class SamplerScanElement(NamedTuple):
+    """Kalman sampling scan element."""
+
+    gain: Array
+    sample: Array
+
+
+def sqrt_associative_params(
+    key: ArrayLike,
+    ms: Array,
+    chol_Ps: Array,
+    Fs: Array,
+    cs: Array,
+    chol_Qs: Array,
+    shape: Sequence[int],
+) -> SamplerScanElement:
+    """Compute the sampler scan elements."""
+    shape = tuple(shape)
+    eps = jax.random.normal(key, ms.shape[:1] + shape + ms.shape[1:])
+    interm_elems = jax.vmap(_sqrt_associative_params_interm)(
+        ms[:-1], chol_Ps[:-1], Fs, cs, chol_Qs, eps[:-1]
+    )
+    last_elem = _sqrt_associative_params_final(ms[-1], chol_Ps[-1], eps[-1])
+    return jax.tree.map(
+        lambda x, y: jnp.concatenate([x, y[None]]), interm_elems, last_elem
+    )
+
+
+def _sqrt_associative_params_interm(
+    m: Array, chol_P: Array, F: Array, c: Array, chol_Q: Array, eps: Array
+) -> SamplerScanElement:
+    inc_m, gain, L = associative_params_single(m, chol_P, F, c, chol_Q)
+    inc = inc_m + eps @ L.T
+    return SamplerScanElement(gain, inc)
+
+
+def _sqrt_associative_params_final(
+    m: Array, chol_P: Array, eps: Array
+) -> SamplerScanElement:
+    gain = jnp.zeros_like(chol_P)
+    sample = m + eps @ chol_P.T
+    return SamplerScanElement(gain, sample)
+
+
+def sampling_operator(
+    elem_i: SamplerScanElement, elem_j: SamplerScanElement
+) -> SamplerScanElement:
+    """Binary associative operator for sampling."""
+    G_i, e_i = elem_i
+    G_j, e_j = elem_j
+    G = G_j @ G_i
+    e = e_i @ G_j.T + e_j
+    return SamplerScanElement(G, e)

cuthbertlib/kalman/smoothing.py

@@ -0,0 +1,121 @@
+"""Implements the square root Rauch–Tung–Striebel (RTS) smoother and associative variant."""
+
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+from jax.scipy.linalg import solve_triangular
+
+from cuthbertlib.linalg import tria
+from cuthbertlib.types import Array, ArrayLike
+
+
+class SmootherScanElement(NamedTuple):
+    """Kalman smoother scan element."""
+
+    g: Array
+    E: Array
+    D: Array
+
+
+# Note that `update` is aliased as `kalman.smoother_update` in `kalman.__init__.py`
+def update(
+    filter_m: ArrayLike,
+    filter_chol_P: ArrayLike,
+    smoother_m: ArrayLike,
+    smoother_chol_P: ArrayLike,
+    F: ArrayLike,
+    c: ArrayLike,
+    chol_Q: ArrayLike,
+) -> tuple[tuple[Array, Array], Array]:
+    """Single step of the square root Rauch–Tung–Striebel (RTS) smoother.
+
+    Args:
+        filter_m: Mean of the filtered state.
+        filter_chol_P: Generalized Cholesky factor of the filtering covariance.
+        smoother_m: Mean of the smoother state.
+        smoother_chol_P: Generalized Cholesky factor of the smoothing covariance.
+        F: State transition matrix.
+        c: State transition shift vector.
+        chol_Q: Generalized Cholesky factor of the state transition noise covariance.
+
+    Returns:
+        A tuple `(smooth_state, info)`.
+        `smooth_state` contains the smoothed mean and square root covariance.
+        `info` contains the smoothing gain matrix.
+
+    References:
+        Paper: Park and Kailath (1994) - Square-root RTS smoothing algorithms
+        Code: https://github.com/EEA-sensors/sqrt-parallel-smoothers/tree/main/parsmooth/sequential
+    """
+    filter_m, filter_chol_P = jnp.asarray(filter_m), jnp.asarray(filter_chol_P)
+    smoother_m, smoother_chol_P = jnp.asarray(smoother_m), jnp.asarray(smoother_chol_P)
+    F, c, chol_Q = jnp.asarray(F), jnp.asarray(c), jnp.asarray(chol_Q)
+
+    nx = F.shape[0]
+    Phi = jnp.block([[F @ filter_chol_P, chol_Q], [filter_chol_P, jnp.zeros_like(F)]])
+    tria_Phi = tria(Phi)
+    Phi11 = tria_Phi[:nx, :nx]
+    Phi21 = tria_Phi[nx:, :nx]
+    Phi22 = tria_Phi[nx:, nx:]
+    gain = solve_triangular(Phi11, Phi21.T, trans=True, lower=True).T
+
+    mean_diff = smoother_m - (c + F @ filter_m)
+    mean = filter_m + gain @ mean_diff
+    chol = tria(jnp.concatenate([Phi22, gain @ smoother_chol_P], axis=1))
+    return (mean, chol), gain
+
+
+def associative_params_single(
+    m: Array,
+    chol_P: Array,
+    F: Array,
+    c: Array,
+    chol_Q: Array,
+) -> SmootherScanElement:
+    """Single time step for scan element for square root parallel Kalman smoother.
+
+    Args:
+        m: Mean of the smoother state.
+        chol_P: Generalized Cholesky factor of the smoothing covariance.
+        F: State transition matrix.
+        c: State transition shift vector.
+        chol_Q: Generalized Cholesky factor of the state transition noise covariance.
+
+    Returns:
+        SmootherScanElement: The output of the associative operator applied to the input
+            elements.
+    """
+    nx = chol_Q.shape[0]
+
+    Phi = jnp.block([[F @ chol_P, chol_Q], [chol_P, jnp.zeros_like(chol_Q)]])
+    Tria_Phi = tria(Phi)
+    Phi11 = Tria_Phi[:nx, :nx]
+    Phi21 = Tria_Phi[nx:, :nx]
+    D = Tria_Phi[nx:, nx:]
+
+    E = jax.scipy.linalg.solve_triangular(Phi11.T, Phi21.T).T
+    g = m - E @ (F @ m + c)
+    return SmootherScanElement(g, E, D)
+
+
+def smoothing_operator(
+    elem_i: SmootherScanElement, elem_j: SmootherScanElement
+) -> SmootherScanElement:
+    """Binary associative operator for the square root Kalman smoother.
+
+    Args:
+        elem_i: Smoother scan element.
+        elem_j: Smoother scan element.
+
+    Returns:
+        SmootherScanElement: The output of the associative operator applied to the input elements.
+    """
+    g_i, E_i, D_i = elem_i
+    g_j, E_j, D_j = elem_j
+
+    g = E_j @ g_i + g_j
+    E = E_j @ E_i
+    D = tria(jnp.concatenate([E_j @ D_i, D_j], axis=1))
+
+    return SmootherScanElement(g, E, D)

cuthbertlib/linalg/__init__.py

@@ -0,0 +1,7 @@
+from cuthbertlib.linalg.collect_nans_chol import collect_nans_chol
+from cuthbertlib.linalg.marginal_sqrt_cov import marginal_sqrt_cov
+from cuthbertlib.linalg.symmetric_inv_sqrt import (
+    chol_cov_with_nans_to_cov,
+    symmetric_inv_sqrt,
+)
+from cuthbertlib.linalg.tria import tria

cuthbertlib/linalg/collect_nans_chol.py

@@ -0,0 +1,90 @@
+"""Implements collection of NaNs and reordering within a Cholesky factor."""
+
+from typing import Any
+
+from jax import numpy as jnp
+from jax import tree
+
+from cuthbertlib.linalg.tria import tria
+from cuthbertlib.types import Array, ArrayLike
+
+
+def _set_to_zero(flag: ArrayLike, x: ArrayLike) -> Array:
+    x = jnp.asarray(x)
+    broadcast_flag = jnp.expand_dims(flag, list(range(1, x.ndim)))
+    return jnp.where(broadcast_flag, 0.0, x)
+
+
+def collect_nans_chol(flag: ArrayLike, chol: ArrayLike, *rest: Any) -> Any:
+    """Converts chol into an order chol factor with NaNs moved to the bottom right.
+
+    Specifically, converts a generalized Cholesky factor of a covariance matrix wit
+    NaNs into an ordered generalized Cholesky factor with NaNs rows and columns
+    moved to the end with diagonal elements set to 1.
+
+    Also reorders the rest of the arguments in the same way along the first axis
+    and sets to 0 for dimensions where flag is True.
+
+    Example behavior:
+    ```
+    flag = jnp.array([False, True, False, True])
+    new_flag, new_chol, new_mean = collect_nans_chol(flag, chol, mean)
+    ```
+
+    Args:
+        flag: Array, boolean array indicating which entries are NaN
+            True for NaN entries, False for valid
+        chol: Array, Cholesky factor of the covariance matrix
+        rest: Any, rest of the arguments to be reordered in the same way
+            along the first axis
+
+    Returns:
+        flag, chol and rest reordered so that valid entries are first and NaNs are last.
+            Diagonal elements of chol are set to 1/√2π so that normalization is correct
+    """
+    flag = jnp.asarray(flag)
+    chol = jnp.asarray(chol)
+
+    # TODO: Can we support batching? I.e. when `chol` is a batch of Cholesky factors,
+    # possibly with multiple leading dimensions
+
+    if flag.ndim > 1 or chol.ndim > 2:
+        raise ValueError("Batched flag or chol not supported yet")
+
+    if not flag.shape:
+        return (
+            flag,
+            _set_to_zero(flag, chol),
+            *tree.map(lambda x: _set_to_zero(flag, x), rest),
+        )
+
+    if chol.size == 1:
+        chol *= jnp.ones_like(flag, dtype=chol.dtype)
+
+    # group the NaN entries together
+    argsort = jnp.argsort(flag, stable=True)
+
+    if chol.ndim == 1:
+        chol = chol[argsort]
+        flag = flag[argsort]
+        chol = jnp.where(flag, 1 / jnp.sqrt(2 * jnp.pi), chol)
+
+    else:
+        chol = jnp.where(flag[:, None], 0.0, chol)
+        chol = chol[argsort]
+        # compute the tria of the covariance matrix with NaNs set to 0
+        chol = tria(chol)
+
+        flag = flag[argsort]
+
+        # set the diagonal of chol_cov to 1/√2π where nans were present so that normalization is correct
+        diag_chol = jnp.diag(chol)
+        diag_chol = jnp.where(flag, 1 / jnp.sqrt(2 * jnp.pi), diag_chol)
+        diag_indices = jnp.diag_indices_from(chol)
+        chol = chol.at[diag_indices].set(diag_chol)
+
+    # Only reorder non-scalar arrays in rest
+    rest = tree.map(lambda x: x[argsort] if jnp.asarray(x).shape else x, rest)
+    rest = tree.map(lambda x: _set_to_zero(flag, x), rest)
+
+    return flag, chol, *rest

cuthbertlib/linalg/marginal_sqrt_cov.py

@@ -0,0 +1,34 @@
+"""Extract marginal square root covariance from a joint square root covariance."""
+
+from typing import Sequence
+
+from jax import numpy as jnp
+
+from cuthbertlib.linalg.tria import tria
+from cuthbertlib.types import Array, ArrayLike
+
+
+def marginal_sqrt_cov(chol_cov: ArrayLike, start: int, end: int) -> Array:
+    """Extracts square root submatrix from a joint square root matrix.
+
+    Specifically, returns B such that
+    B @ B.T = (chol_cov @ chol_cov.T)[start:end, start:end]
+
+    Args:
+        chol_cov: Generalized Cholesky factor of the covariance matrix.
+        start: Start index of the submatrix.
+        end: End index of the submatrix.
+
+    Returns:
+        Lower triangular square root matrix of the marginal covariance matrix.
+    """
+    chol_cov = jnp.asarray(chol_cov)
+    assert chol_cov.ndim == 2, "chol_cov must be a 2D array"
+    assert chol_cov.shape[0] == chol_cov.shape[1], "chol_cov must be square"
+    assert start >= 0 and end <= chol_cov.shape[0], (
+        "start and end must be within the bounds of chol_cov"
+    )
+    assert start < end, "start must be less than end"
+
+    chol_cov_select_rows = chol_cov[start:end, :]
+    return tria(chol_cov_select_rows)