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

cuthbertlib/linalg/symmetric_inv_sqrt.py

@@ -0,0 +1,126 @@
+"""Implements inverse square root of a symmetric matrix."""
+
+import jax.numpy as jnp
+
+from cuthbertlib.linalg.tria import tria
+from cuthbertlib.types import Array, ArrayLike
+
+
+def symmetric_inv_sqrt(
+    A: ArrayLike,
+    rtol: float | ArrayLike | None = None,
+    ignore_nan_dims: bool = False,
+) -> Array:
+    r"""Computes the inverse square root of a symmetric matrix.
+
+    I.e., a lower triangular matrix $L$ such that $L L^{\top} = A^{-1}$ (for positive definite
+    $A$). Note that this is not unique and will generally not match the Cholesky factor
+    of $A^{-1}$.
+
+    For singular matrices, small singular values will be cut off reminiscent of
+    the Moore-Penrose pseudoinverse - https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.pinv.html.
+
+    In the case of singular or indefinite $A$, the output will be an approximation
+    and $L L^{\top} = A^{-1}$ will not hold in general.
+
+    Args:
+        A: A symmetric matrix.
+        rtol: The relative tolerance for the singular values.
+            Cutoff for small singular values; singular values smaller than
+            `rtol * largest_singular_value` are treated as zero.
+            See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.pinv.html.
+        ignore_nan_dims: Whether to treat dimensions with NaN on the diagonal as missing
+            and ignore all rows and columns associated with them (with result in those
+            dimensions being NaN on the diagonal and zero off-diagonal).
+
+    Returns:
+        A lower triangular matrix $L$ such that $L L^{\top} = A^{-1}$ (for valid dimensions).
+    """
+    arr = jnp.asarray(A)
+
+    # Check for NaNs on the diagonal (missing dimensions)
+    diag_vals = jnp.diag(arr)
+    nan_diag_mask = jnp.isnan(diag_vals) * ignore_nan_dims
+
+    # Check for dimensions whose row and column are all 0
+    zero_mask = jnp.all(arr == 0.0, axis=0) & jnp.all(arr == 0.0, axis=1)
+
+    nan_mask = nan_diag_mask | zero_mask
+
+    # Sort to group valid dimensions first (needed for SVD to work correctly)
+    argsort = jnp.argsort(nan_mask, stable=True)
+    arr_sorted = arr[argsort[:, None], argsort]
+    nan_mask_sorted = nan_mask[argsort]
+
+    # Zero out invalid dimensions before computation
+    invalid_mask_2d = ((nan_mask_sorted[:, None]) | (nan_mask_sorted[None, :])) & (
+        ignore_nan_dims
+    )
+    arr_sorted = jnp.where(invalid_mask_2d, 0.0, arr_sorted)
+
+    # Compute inverse square root on sorted, masked matrix
+    L_sorted = _symmetric_inv_sqrt(arr_sorted, rtol)
+
+    # Post-process: zero out invalid rows/cols, set NaN on invalid diagonal
+    L_sorted = jnp.where(invalid_mask_2d, 0.0, L_sorted)
+    diag_L = jnp.where(nan_mask_sorted, jnp.nan, jnp.diag(L_sorted))
+    L_sorted = L_sorted.at[jnp.diag_indices_from(L_sorted)].set(diag_L)
+
+    # Un-sort to restore original order
+    inv_argsort = jnp.argsort(argsort)
+    L = L_sorted[inv_argsort[:, None], inv_argsort]
+
+    return L
+
+
+def _symmetric_inv_sqrt(A: ArrayLike, rtol: float | ArrayLike | None = None) -> Array:
+    """Implementation of symmetric inverse square root without NaN handling."""
+    arr = jnp.asarray(A)
+
+    # From https://github.com/jax-ml/jax/blob/75d8702023fca6fe4a223bf1e08545c1c80581c0/jax/_src/numpy/linalg.py#L972
+    if rtol is None:
+        max_rows_cols = max(arr.shape[-2:])
+        rtol = jnp.asarray(10.0 * max_rows_cols * jnp.finfo(arr.dtype).eps)
+    u, s, _ = jnp.linalg.svd(arr, full_matrices=False, hermitian=True)
+    cutoff = rtol * s[0]
+    # Use 0 for invalid singular values to avoid inf/NaN propagation in tria
+    valid_mask = s > cutoff
+    inv_sqrt_s = jnp.where(valid_mask, 1.0 / jnp.sqrt(s), 0.0).astype(u.dtype)
+    B = u * inv_sqrt_s  # Square root but not lower triangular
+    L = tria(B)  # Make lower triangular
+    # Mark dimensions with all 0 rows and columns as NaN
+    zero_dims_mask = jnp.all(L == 0.0, axis=0) & jnp.all(L == 0.0, axis=1)
+    L = jnp.where(zero_dims_mask[:, None] | zero_dims_mask[None, :], jnp.nan, L)
+    return L
+
+
+def chol_cov_with_nans_to_cov(chol_cov: ArrayLike) -> Array:
+    """Converts a Cholesky factor to a covariance matrix.
+
+    NaNs on the diagonal specify dimensions to be ignored.
+
+    Args:
+        chol_cov: A Cholesky factor of a covariance matrix with NaNs on the diagonal
+            specifying dimensions to be ignored.
+
+    Returns:
+        A covariance matrix equivalent to chol_cov @ chol_cov.T in dimensions where
+            the Cholesky factor is valid and for invalid dimensions (ones with NaN on the
+            diagonal in chol_cov) with NaN on the diagonal and zero off-diagonal.
+    """
+    chol_cov = jnp.asarray(chol_cov)
+
+    nan_mask = jnp.isnan(jnp.diag(chol_cov))
+
+    # Set all rows and columns with invalid diagonal to zero
+    chol_cov = jnp.where(nan_mask[:, None] | nan_mask[None, :], 0, chol_cov)
+
+    # Calculate the covariance matrix
+    cov = chol_cov @ chol_cov.T
+
+    # Set the diagonal to NaN
+    cov = cov.at[jnp.diag_indices_from(cov)].set(
+        jnp.where(nan_mask, jnp.nan, jnp.diag(cov))
+    )
+
+    return cov

cuthbertlib/linalg/tria.py

@@ -0,0 +1,21 @@
+"""Implements triangularization operator a matrix via QR decomposition."""
+
+import jax
+
+from cuthbertlib.types import Array
+
+
+def tria(A: Array) -> Array:
+    r"""A triangularization operator using QR decomposition.
+
+    Args:
+        A: The matrix to triangularize.
+
+    Returns:
+        A lower triangular matrix $R$ such that $R R^\top = A A^\top$.
+
+    Reference:
+        [Arasaratnam and Haykin (2008)](https://ieeexplore.ieee.org/document/4524036): Square-Root Quadrature Kalman Filtering
+    """
+    _, R = jax.scipy.linalg.qr(A.T, mode="economic")
+    return R.T

cuthbertlib/linearize/__init__.py

@@ -0,0 +1,7 @@
+from cuthbertlib.linalg import symmetric_inv_sqrt
+from cuthbertlib.linearize.log_density import (
+    linearize_log_density,
+    linearize_log_density_given_chol_cov,
+)
+from cuthbertlib.linearize.moments import linearize_moments
+from cuthbertlib.linearize.taylor import linearize_taylor

cuthbertlib/linearize/log_density.py

@@ -0,0 +1,175 @@
+"""Implements linearization of conditional log densities."""
+
+from typing import overload
+
+import jax.numpy as jnp
+from jax import grad, hessian, jacobian
+
+from cuthbertlib.linalg import chol_cov_with_nans_to_cov, symmetric_inv_sqrt
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    LogConditionalDensity,
+    LogConditionalDensityAux,
+)
+
+
+@overload
+def linearize_log_density(
+    log_density: LogConditionalDensity,
+    x: ArrayLike,
+    y: ArrayLike,
+    has_aux: bool = False,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array, Array]: ...
+@overload
+def linearize_log_density(
+    log_density: LogConditionalDensityAux,
+    x: ArrayLike,
+    y: ArrayLike,
+    has_aux: bool = True,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array, Array, ArrayTree]: ...
+
+
+def linearize_log_density(
+    log_density: LogConditionalDensity | LogConditionalDensityAux,
+    x: ArrayLike,
+    y: ArrayLike,
+    has_aux: bool = False,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array, Array] | tuple[Array, Array, Array, ArrayTree]:
+    r"""Linearizes a conditional log density around given points.
+
+    The linearization is exact in the case of a linear-Gaussian `log_density`, i.e., it returns
+    $(H, d, L)$ if `log_density` is of the form
+
+    $$
+    \log p(y \mid x) = -\frac{1}{2}(y - H x - d)^\top (LL^\top)^{-1} (y - H x - d) + \textrm{const}.
+    $$
+
+    The Cholesky factor of the covariance is calculated using the negative Hessian
+    of `log_density` with respect to `y` as the precision matrix.
+    `symmetric_inv_sqrt` is used to calculate the inverse square root by
+    ignoring any singular values that are sufficiently close to zero
+    (this is a projection in the case the Hessian is not positive definite).
+
+    Alternatively, the Cholesky factor can be provided directly
+    in `linearize_log_density_given_chol_cov`.
+
+    Args:
+        log_density: A conditional log density of y given x. Returns a scalar.
+        x: The input points.
+        y: The output points.
+        has_aux: Whether `log_density` returns an auxiliary value.
+        rtol: The relative tolerance for the singular values of the precision matrix
+            when passed to `symmetric_inv_sqrt`.
+            Cutoff for small singular values; singular values smaller than
+            `rtol * largest_singular_value` are treated as zero.
+            The default is determined based on the floating point precision of the dtype.
+            See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.pinv.html.
+        ignore_nan_dims: Whether to treat dimensions with NaN on the diagonal of the
+            precision matrix as missing and ignore all rows and columns associated with
+            them.
+
+    Returns:
+        Linearized matrix, shift, and Cholesky factor of the covariance matrix.
+            The auxiliary value is also returned if `has_aux` is `True`.
+    """
+    prec_and_maybe_aux = hessian(log_density, 1, has_aux=has_aux)(x, y)
+    prec = -prec_and_maybe_aux[0] if has_aux else -prec_and_maybe_aux
+    if ignore_nan_dims:
+        prec_diag = jnp.diag(prec)
+        nan_mask = jnp.isnan(y) | jnp.isnan(prec_diag)
+        prec = prec.at[jnp.diag_indices_from(prec)].set(
+            jnp.where(nan_mask, jnp.nan, prec_diag)
+        )
+
+    chol_cov = symmetric_inv_sqrt(prec, rtol=rtol, ignore_nan_dims=ignore_nan_dims)
+    mat, shift, *extra = linearize_log_density_given_chol_cov(
+        log_density, x, y, chol_cov, has_aux=has_aux, ignore_nan_dims=ignore_nan_dims
+    )
+    return mat, shift, chol_cov, *extra
+
+
+@overload
+def linearize_log_density_given_chol_cov(
+    log_density: LogConditionalDensity,
+    x: ArrayLike,
+    y: ArrayLike,
+    chol_cov: ArrayLike,
+    has_aux: bool = False,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array]: ...
+@overload
+def linearize_log_density_given_chol_cov(
+    log_density: LogConditionalDensityAux,
+    x: ArrayLike,
+    y: ArrayLike,
+    chol_cov: ArrayLike,
+    has_aux: bool = True,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array, ArrayTree]: ...
+
+
+def linearize_log_density_given_chol_cov(
+    log_density: LogConditionalDensity | LogConditionalDensityAux,
+    x: ArrayLike,
+    y: ArrayLike,
+    chol_cov: ArrayLike,
+    has_aux: bool = False,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array] | tuple[Array, Array, ArrayTree]:
+    r"""Linearizes a conditional log density around given points.
+
+    The linearization is exact in the case of a linear-Gaussian `log_density`, i.e., it returns
+    $(H, d)$ if `log_density` is of the form
+
+    $$
+    \log p(y \mid x) = -\frac{1}{2}(y - H x - d)^\top (LL^\top)^{-1} (y - H x - d) + \textrm{const},
+    $$
+
+    where $L$ is the argument `chol_cov`.
+
+    Args:
+        log_density: A conditional log density of y given x. Returns a scalar.
+        x: The input points.
+        y: The output points.
+        chol_cov: The Cholesky factor of the covariance matrix of the Gaussian.
+        has_aux: Whether `log_density` returns an auxiliary value.
+        ignore_nan_dims: Whether to ignore dimensions with NaN on the diagonal of the
+            precision matrix or in y.
+
+    Returns:
+        Linearized matrix and shift. The auxiliary value is also returned if `has_aux` is `True`.
+    """
+    chol_cov = jnp.asarray(chol_cov)
+
+    cov = (
+        chol_cov_with_nans_to_cov(chol_cov)
+        if ignore_nan_dims
+        else chol_cov @ chol_cov.T
+    )
+
+    if has_aux:
+
+        def grad_log_density_wrapper_aux(x, y):
+            g, aux = grad(log_density, 1, has_aux=True)(x, y)
+            return g, (g, aux)
+
+        jac, (g, *extra) = jacobian(grad_log_density_wrapper_aux, 0, has_aux=True)(x, y)
+    else:
+
+        def grad_log_density_wrapper(x, y):
+            g = grad(log_density, 1)(x, y)
+            return g, (g,)
+
+        jac, (g, *extra) = jacobian(grad_log_density_wrapper, 0, has_aux=True)(x, y)
+
+    mat = cov @ jac
+    shift = y - mat @ x + cov @ g
+    return mat, shift, *extra

cuthbertlib/linearize/moments.py

@@ -0,0 +1,94 @@
+"""Implements moment-based linearization."""
+
+from typing import Callable, cast, overload
+
+import jax
+from jax.typing import ArrayLike
+
+from cuthbertlib.types import Array, ArrayTree
+
+MeanAndCholCovFunc = Callable[[ArrayLike], tuple[Array, Array]]
+MeanAndCholCovFuncAux = Callable[[ArrayLike], tuple[Array, Array, ArrayTree]]
+
+
+@overload
+def linearize_moments(
+    mean_and_chol_cov_function: MeanAndCholCovFunc,
+    x: ArrayLike,
+    has_aux: bool = False,
+) -> tuple[Array, Array, Array]: ...
+@overload
+def linearize_moments(
+    mean_and_chol_cov_function: MeanAndCholCovFuncAux,
+    x: ArrayLike,
+    has_aux: bool = True,
+) -> tuple[Array, Array, Array, ArrayTree]: ...
+
+
+def linearize_moments(
+    mean_and_chol_cov_function: MeanAndCholCovFunc | MeanAndCholCovFuncAux,
+    x: ArrayLike,
+    has_aux: bool = False,
+) -> tuple[Array, Array, Array] | tuple[Array, Array, Array, ArrayTree]:
+    r"""Linearizes conditional mean and chol_cov functions into a linear-Gaussian form.
+
+    Takes a function `mean_and_chol_cov_function(x)` that returns the
+    conditional mean and Cholesky factor of the covariance matrix of the distribution
+    $p(y \mid x)$ for a given input `x`.
+
+    Returns $(H, d, L)$ defining a linear-Gaussian approximation to the conditional
+    distribution $p(y \mid x) \approx N(y \mid H x + d, L L^\top)$.
+
+    `mean_and_chol_cov_function` has the following signature with `has_aux` = False:
+    ```
+    m, chol = mean_and_chol_cov_function(x)
+    ```
+    or with `has_aux` = True:
+    ```
+    m, chol, aux = mean_and_chol_cov_function(x)
+    ```
+
+    Args:
+        mean_and_chol_cov_function: A callable that returns the conditional mean and
+            Cholesky factor of the covariance matrix of the distribution for a given
+            input.
+        x: The point to linearize around.
+        has_aux: Whether `mean_and_chol_cov_function` returns an auxiliary value.
+
+    Returns:
+        Linearized matrix, shift, and Cholesky factor of the covariance matrix.
+            The auxiliary value is also returned if `has_aux` is `True`.
+
+    References:
+        - [sqrt-parallel-smoothers](https://github.com/EEA-sensors/sqrt-parallel-smoothers/blob/main/parsmooth/linearization/_extended.py)
+    """
+    if has_aux:
+        mean_and_chol_cov_function = cast(
+            MeanAndCholCovFuncAux, mean_and_chol_cov_function
+        )
+
+        def mean_and_chol_cov_function_wrapper_aux(
+            x: ArrayLike,
+        ) -> tuple[Array, tuple[Array, Array, ArrayTree]]:
+            mean, chol_cov, aux = mean_and_chol_cov_function(x)
+            return mean, (mean, chol_cov, aux)
+
+        F, (m, *extra) = jax.jacfwd(
+            mean_and_chol_cov_function_wrapper_aux, has_aux=True
+        )(x)
+
+    else:
+        mean_and_chol_cov_function = cast(
+            MeanAndCholCovFunc, mean_and_chol_cov_function
+        )
+
+        def mean_and_chol_cov_function_wrapper(
+            x: ArrayLike,
+        ) -> tuple[Array, tuple[Array, Array]]:
+            mean, chol_cov = mean_and_chol_cov_function(x)
+            return mean, (mean, chol_cov)
+
+        F, (m, *extra) = jax.jacfwd(mean_and_chol_cov_function_wrapper, has_aux=True)(x)
+
+    b = m - F @ x
+    return F, b, *extra

cuthbertlib/linearize/taylor.py

@@ -0,0 +1,83 @@
+"""Implements Taylor-like linearization."""
+
+from typing import Callable, overload
+
+import jax
+from jax import numpy as jnp
+from jax.typing import ArrayLike
+
+from cuthbertlib.linalg import symmetric_inv_sqrt
+from cuthbertlib.types import Array, ArrayTree
+
+
+@overload
+def linearize_taylor(
+    log_potential: Callable[[ArrayLike], Array],
+    x: ArrayLike,
+    has_aux: bool = False,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array]: ...
+@overload
+def linearize_taylor(
+    log_potential: Callable[[ArrayLike], tuple[Array, ArrayTree]],
+    x: ArrayLike,
+    has_aux: bool = True,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array, ArrayTree]: ...
+
+
+def linearize_taylor(
+    log_potential: Callable[[ArrayLike], Array]
+    | Callable[[ArrayLike], tuple[Array, ArrayTree]],
+    x: ArrayLike,
+    has_aux: bool = False,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array] | tuple[Array, Array, ArrayTree]:
+    r"""Linearizes a log potential function around a given point using Taylor expansion.
+
+    Unlike the other linearization methods, this applies to a potential function
+    with no required notion of observation $y$ or conditional dependence.
+
+    Instead we have the linearization
+
+    $$
+    \log G(x) = -\frac{1}{2} (x - m)^\top (L L^\top)^{-1} (x - m).
+    $$
+
+    Args:
+        log_potential: A callable that returns a non-negative scalar. Does not need
+            to be a normalized probability density in its input.
+        x: The point to linearize around.
+        has_aux: Whether `log_potential` returns an auxiliary value.
+        rtol: The relative tolerance for the singular values of the precision matrix
+            when passed to `symmetric_inv_sqrt`.
+            Cutoff for small singular values; singular values smaller than
+            `rtol * largest_singular_value` are treated as zero.
+            The default is determined based on the floating point precision of the dtype.
+            See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.pinv.html.
+        ignore_nan_dims: Whether to treat dimensions with NaN on the diagonal of the
+            precision matrix as missing and ignore all rows and columns associated with
+            them.
+
+    Returns:
+        Linearized mean and Cholesky factor of the covariance matrix.
+            The auxiliary value is also returned if `has_aux` is `True`.
+    """
+    g_and_maybe_aux = jax.grad(log_potential, has_aux=has_aux)(x)
+    prec_and_maybe_aux = jax.hessian(log_potential, has_aux=has_aux)(x)
+
+    g, aux = g_and_maybe_aux if has_aux else (g_and_maybe_aux, None)
+    prec = -prec_and_maybe_aux[0] if has_aux else -prec_and_maybe_aux
+
+    L = symmetric_inv_sqrt(prec, rtol=rtol, ignore_nan_dims=ignore_nan_dims)
+
+    # Change nans on diag to zeros for L @ L.T @ g, still retain nans on diag for L for bookkeeping
+    # If ignore_nan_dims, change all rows and columns with nans on the diagonal to 0
+    L_diag = jnp.diag(L)
+    nan_mask = jnp.isnan(L_diag) * ignore_nan_dims
+    L_temp = jnp.where(nan_mask[:, None] | nan_mask[None, :], 0.0, L)
+    m = x + L_temp @ L_temp.T @ g
+    return (m, L, aux) if has_aux else (m, L)

cuthbertlib/quadrature/__init__.py

@@ -0,0 +1,4 @@
+from cuthbertlib.quadrature import cubature, gauss_hermite, unscented
+from cuthbertlib.quadrature.common import Quadrature, SigmaPoints
+from cuthbertlib.quadrature.linearize import conditional_moments, functional
+from cuthbertlib.quadrature.utils import cholesky_update_many

cuthbertlib/quadrature/common.py

@@ -0,0 +1,102 @@
+"""Common types and protocols for quadrature."""
+
+from typing import NamedTuple, Protocol, Self, runtime_checkable
+
+import jax.numpy as jnp
+
+from cuthbertlib.linalg import tria
+from cuthbertlib.types import Array, ArrayLike
+
+__all__ = ["SigmaPoints", "Quadrature"]
+
+
+class SigmaPoints(NamedTuple):
+    """Represents integration (quadrature) sigma points as a collection of points.
+
+    Weights correspond to mean and covariance calculations.
+
+    Attributes:
+        points: The sigma points.
+        wm: The mean weights.
+        wc: The covariance weights.
+
+    Methods:
+        mean: Computes the mean of the sigma points.
+        covariance: Computes the covariance between the sigma points and the other
+            sigma points (or itself).
+        sqrt: Computes a square root of the covariance matrix of the sigma points.
+
+    References:
+        Simo Särkkä, Lennard Svensson. *Bayesian Filtering and Smoothing.*
+            In: Cambridge University Press 2023.
+    """
+
+    points: Array
+    wm: Array
+    wc: Array
+
+    @property
+    def mean(self) -> Array:
+        """Computes the mean of the sigma points.
+
+        Returns:
+            The mean of the sigma points.
+        """
+        return jnp.dot(self.wm, self.points)
+
+    # Should this be property too?
+    def covariance(self, other: Self | None = None) -> Array:
+        """Computes the covariance between the sigma points and the other sigma points.
+
+        Args:
+            other: The optional other sigma points.
+
+        Returns:
+            The covariance matrix.
+        """
+        mean = self.mean
+        if other is None:
+            return _cov(self.wc, self.points, mean, self.points, mean)
+
+        other_mean = other.mean
+        return _cov(self.wc, self.points, mean, other.points, other_mean)
+
+    @property
+    def sqrt(self) -> Array:
+        """Computes the square root of the covariance matrix of the sigma points.
+
+        Returns:
+            The square root of the covariance matrix.
+        """
+        sqrt = jnp.sqrt(self.wc[:, None]) * (self.points - self.mean[None, :])
+        sqrt = tria(sqrt.T)
+        return sqrt
+
+
+@runtime_checkable
+class Quadrature(Protocol):
+    """Protocol for quadrature methods."""
+
+    def get_sigma_points(self, m: ArrayLike, chol: ArrayLike) -> SigmaPoints:
+        """Get the sigma points.
+
+        Args:
+            m: The mean.
+            chol: The Cholesky factor of the covariance.
+
+        Returns:
+            SigmaPoints: The sigma points.
+        """
+        ...
+
+
+def _cov(
+    wc: Array,
+    x_pts: Array,
+    x_mean: Array,
+    y_points: Array,
+    y_mean: Array,
+) -> Array:
+    one = (x_pts - x_mean[None, :]).T * wc[None, :]
+    two = y_points - y_mean[None, :]
+    return jnp.dot(one, two)

cuthbertlib/quadrature/cubature.py

@@ -0,0 +1,73 @@
+"""Implements cubature quadrature."""
+
+from typing import NamedTuple
+
+import jax.numpy as jnp
+import numpy as np
+from jax.typing import ArrayLike
+
+from cuthbertlib.quadrature.common import Quadrature, SigmaPoints
+
+__all__ = ["weights", "CubatureQuadrature"]
+
+
+class CubatureQuadrature(NamedTuple):
+    """Cubature quadrature.
+
+    Attributes:
+        wm: The mean weights.
+        wc: The covariance weights.
+        xi: The sigma points.
+    """
+
+    wm: ArrayLike
+    wc: ArrayLike
+    xi: ArrayLike
+
+    def get_sigma_points(self, m: ArrayLike, chol: ArrayLike) -> SigmaPoints:
+        """Get the sigma points.
+
+        Args:
+            m: The mean.
+            chol: The Cholesky factor of the covariance.
+
+        Returns:
+            SigmaPoints: The sigma points.
+        """
+        return get_sigma_points(m, chol, self.xi, self.wm, self.wc)
+
+
+def get_sigma_points(
+    m: ArrayLike, chol: ArrayLike, xi: ArrayLike, wm: ArrayLike, wc: ArrayLike
+) -> SigmaPoints:
+    # TODO: Add docstring here
+    m = jnp.asarray(m)
+    chol = jnp.asarray(chol)
+    xi = jnp.asarray(xi)
+    wm = jnp.asarray(wm)
+    wc = jnp.asarray(wc)
+    sigma_points = m[None, :] + jnp.dot(chol, xi.T).T
+
+    return SigmaPoints(sigma_points, wm, wc)
+
+
+def weights(n_dim: int) -> Quadrature:
+    """Computes the weights associated with the spherical cubature method.
+
+    The number of sigma-points is 2 * n_dim.
+
+    Args:
+        n_dim: Dimensionality of the problem.
+
+    Returns:
+        The quadrature object with the weights and sigma-points.
+
+    References:
+        Simo Särkkä, Lennard Svensson. *Bayesian Filtering and Smoothing.*
+            In: Cambridge University Press 2023.
+    """
+    wm = np.ones(shape=(2 * n_dim,)) / (2 * n_dim)
+    wc = wm
+    xi = np.concatenate([np.eye(n_dim), -np.eye(n_dim)], axis=0) * np.sqrt(n_dim)
+
+    return CubatureQuadrature(wm=wm, wc=wc, xi=xi)