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

cuthbertlib/quadrature/gauss_hermite.py

@@ -0,0 +1,62 @@
+"""Implements Gauss-Hermite quadrature."""
+
+from itertools import product
+from typing import NamedTuple
+
+import numpy as np
+from jax.typing import ArrayLike
+from numpy.polynomial.hermite_e import hermegauss
+
+from cuthbertlib.quadrature import cubature
+from cuthbertlib.quadrature.common import Quadrature, SigmaPoints
+
+__all__ = ["weights", "GaussHermiteQuadrature"]
+
+
+class GaussHermiteQuadrature(NamedTuple):
+    """Gauss-Hermite 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 cubature.get_sigma_points(m, chol, self.xi, self.wm, self.wc)
+
+
+def weights(n_dim: int, order: int = 3) -> Quadrature:
+    """Computes the weights associated with the Gauss-Hermite quadrature method.
+
+    The Hermite polynomial is in the probabilist's version.
+
+    Args:
+        n_dim: Dimensionality of the problem.
+        order: The order of Hermite polynomial. Defaults to 3.
+
+    Returns:
+        The quadrature object with the weights and sigma-points.
+
+    References:
+        Simo Särkkä. *Bayesian Filtering and Smoothing.*
+            In: Cambridge University Press 2013.
+    """
+    x, w = hermegauss(order)
+    xn = np.array(list(product(*(x,) * n_dim)))
+    wn = np.prod(np.array(list(product(*(w,) * n_dim))), 1)
+    wn /= np.sqrt(2 * np.pi) ** n_dim
+    return GaussHermiteQuadrature(wm=wn, wc=wn, xi=xn)

cuthbertlib/quadrature/linearize.py

@@ -0,0 +1,143 @@
+"""Implements quadrature-based linearization of conditional moments and functional."""
+
+from typing import Callable
+
+import jax.numpy as jnp
+from jax import vmap
+from jax.scipy.linalg import cho_solve
+
+from cuthbertlib.linalg import tria
+from cuthbertlib.quadrature.common import Quadrature, SigmaPoints
+from cuthbertlib.quadrature.utils import cholesky_update_many
+from cuthbertlib.types import Array, ArrayLike
+
+__all__ = ["conditional_moments", "functional"]
+
+
+def conditional_moments(
+    mean_fn: Callable[[ArrayLike], Array],
+    cov_fn: Callable[[ArrayLike], Array],
+    m: ArrayLike,
+    cov: ArrayLike,
+    quadrature: Quadrature,
+    mode: str = "covariance",
+) -> tuple[Array, Array, Array]:
+    r"""Linearizes the conditional mean and covariance of a Gaussian distribution.
+
+    Args:
+        mean_fn: The mean function $\mathbb{E}[Y \mid x] =$ `mean_fn(x)`.
+        cov_fn: The covariance function $\mathbb{C}[Y \mid x] =$ `cov_fn(x)`.
+        m: The mean of the Gaussian distribution.
+        cov: The covariance of the Gaussian distribution.
+        quadrature: The quadrature object with the weights and sigma-points.
+        mode: The mode of the covariance. Default is 'covariance', which means that cov
+            and cov_fn are given as covariance matrices.
+            Otherwise, the Cholesky factor of the covariances are given.
+
+    Returns:
+        A, b, Q where A, b are the linearized model parameters and Q is either given as
+            a full covariance matrix or as a square root factor depending on
+            the `mode`.
+    """
+    if mode == "covariance":
+        chol = jnp.linalg.cholesky(cov)
+    else:
+        chol = cov
+    x_pts: SigmaPoints = quadrature.get_sigma_points(m, chol)
+
+    f_pts = SigmaPoints(vmap(mean_fn)(x_pts.points), x_pts.wm, x_pts.wc)
+    Psi_x = x_pts.covariance(f_pts)
+
+    A = cho_solve((chol, True), Psi_x).T
+    b = f_pts.mean - A @ m
+    if mode != "covariance":
+        # This can probably be abstracted better.
+        sqrt_Phi = f_pts.sqrt
+
+        chol_pts = vmap(cov_fn)(x_pts.points)
+        temp = jnp.sqrt(x_pts.wc[:, None, None]) * chol_pts
+
+        # concatenate the blocks properly, it's a bit urk, but what can you do...
+        temp = jnp.transpose(temp, [1, 0, 2]).reshape(temp.shape[1], -1)
+        chol_Q = tria(jnp.concatenate([sqrt_Phi, temp], axis=1))
+        chol_Q = cholesky_update_many(chol_Q, (A @ chol).T, -1.0)
+        return A, b, chol_Q
+
+    V_pts = vmap(cov_fn)(x_pts.points)
+    v_f = jnp.sum(x_pts.wc[:, None, None] * V_pts, 0)
+
+    Phi = f_pts.covariance()
+    Q = Phi + v_f - A @ cov @ A.T
+
+    return A, b, 0.5 * (Q + Q.T)
+
+
+def functional(
+    fn: Callable[[ArrayLike], Array],
+    S: ArrayLike,
+    m: ArrayLike,
+    cov: ArrayLike,
+    quadrature: Quadrature,
+    mode: str = "covariance",
+) -> tuple[Array, Array, Array]:
+    r"""Linearizes a nonlinear function of a Gaussian distribution.
+
+    For a given Gaussian distribution $p(x) = N(x \mid m, P)$,
+    and $Y = f(X) + \epsilon$, where $\epsilon$ is a zero-mean Gaussian noise
+    with covariance S, this function computes an approximation $Y = A X + b + \epsilon$
+    using the sigma points method given by get_sigma_points.
+
+    Args:
+        fn: The function $Y = f(X) + N(0, S)$.
+            Because the function is linearized, the function should be vectorized.
+        S: The covariance of the noise.
+        m: The mean of the Gaussian distribution.
+        cov: The covariance of the Gaussian distribution.
+        quadrature: The quadrature object with the weights and sigma-points.
+        mode: The mode of the covariance. Default is 'covariance', which means that cov
+            and cov_fn are given as covariance matrices. Otherwise, the Cholesky factor
+            of the covariances are given.
+
+    Returns:
+        A, b, Q: The linearized model parameters $Y = A X + b + N(0, Q)$.
+            Q is either given as a full covariance matrix or as a square root factor depending on the `mode`.
+
+    Notes:
+        We do not support non-additive noise in this method.
+        If you have a non-additive noise, you should use the `conditional_moments` or
+        the Taylor linearization method.
+        Another solution is to form the covariance function using the quadrature method
+        itself. For example, if you have a function $f(x, q)$, where $q$ is a zero-mean
+        random variable with covariance `S`,
+        you can form the mean and covariance function as follows:
+
+        ```python
+        def linearize_q_part(x):
+            n_dim = S.shape[0]
+            m_q = jnp.zeros(n_dim)
+            A, b, Q = functional(lambda x: f(x, q_sigma_points.points), 0. * S, m_q, S, quadrature, mode)
+            return A, b, Q
+
+        def cov_fn(x):
+            A, b, Q = linearize_q_part(x)
+            return Q + A @ S @ A.T
+
+        def mean_fn(x):
+            A, b, Q = linearize_q_part(x)
+            m_q = jnp.zeros(n_dim)
+            return b + f(x, m_q)
+        ```
+
+        This technique is a bit wasteful due to our current separation of duties between
+        the mean and covariance functions, but as we develop the library further, we
+        will provide a more elegant solution.
+    """
+
+    # make the equivalent conditional_moments model
+    def mean_fn(x):
+        return fn(x)
+
+    def cov_fn(x):
+        return jnp.asarray(S)
+
+    return conditional_moments(mean_fn, cov_fn, m, cov, quadrature, mode)

cuthbertlib/quadrature/unscented.py

@@ -0,0 +1,79 @@
+"""Implements unscented quadrature."""
+
+from typing import NamedTuple
+
+import jax.numpy as jnp
+
+from cuthbertlib.quadrature.common import SigmaPoints
+from cuthbertlib.types import Array, ArrayLike
+
+__all__ = ["weights", "UnscentedQuadrature"]
+
+
+class UnscentedQuadrature(NamedTuple):
+    """Unscented quadrature.
+
+    Attributes:
+        wm: The mean weights.
+        wc: The covariance weights.
+        lamda: The lambda parameter.
+    """
+
+    wm: Array
+    wc: Array
+    lamda: float
+
+    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.
+        """
+        m = jnp.asarray(m)
+        chol = jnp.asarray(chol)
+
+        n_dim = m.shape[0]
+        scaled_chol = jnp.sqrt(n_dim + self.lamda) * chol
+
+        zeros = jnp.zeros((1, n_dim))
+        sigma_points = m[None, :] + jnp.concatenate(
+            [zeros, scaled_chol.T, -scaled_chol.T], axis=0
+        )
+        return SigmaPoints(sigma_points, self.wm, self.wc)
+
+
+def weights(
+    n_dim: int, alpha: float = 0.5, beta: float = 2.0, kappa: float | None = None
+) -> UnscentedQuadrature:
+    """Computes the weights associated with the unscented cubature method.
+
+    The number of sigma-points is 2 * n_dim.
+    This method is also known as the Unscented Transform, and generalizes the
+    `cubature.py` weights: the cubature method is a special case of the unscented
+    for the parameters `alpha=1.0`, `beta=0.0`, `kappa=0.0`.
+
+    Args:
+        n_dim: Dimension of the space.
+        alpha: Parameter of the unscented transform, default is 0.5.
+        beta: Parameter of the unscented transform, default is 2.0.
+        kappa: Parameter of the unscented transform, default is 3 + n_dim.
+
+    Returns:
+        UnscentedQuadrature: The quadrature object with the weights and sigma-points.
+
+    References:
+        - https://groups.seas.harvard.edu/courses/cs281/papers/unscented.pdf
+    """
+    if kappa is None:
+        kappa = 3.0 + n_dim
+
+    lamda = alpha**2 * (n_dim + kappa) - n_dim
+    wm = jnp.full(2 * n_dim + 1, 1 / (2 * (n_dim + lamda)))
+
+    wm = wm.at[0].set(lamda / (n_dim + lamda))
+    wc = wm.at[0].set(lamda / (n_dim + lamda) + (1 - alpha**2 + beta))
+    return UnscentedQuadrature(wm=wm, wc=wc, lamda=lamda)

cuthbertlib/quadrature/utils.py

@@ -0,0 +1,109 @@
+"""Utility functions (Cholesky updating) for quadrature."""
+
+import jax
+import jax.numpy as jnp
+
+from cuthbertlib.types import Array, ArrayLike
+
+__all__ = ["cholesky_update_many"]
+
+
+def cholesky_update_many(
+    chol_init: ArrayLike, update_vectors: ArrayLike, multiplier: float
+) -> Array:
+    r"""Update the Cholesky decomposition of a matrix with multiple update vectors.
+
+    In mathematical terms, we compute :math:`A + \sum_{i=1}^{n} \alpha v_i v_i^T`
+    where :math:`A` is the original matrix, :math:`v_i` are the update vectors and
+    :math:`\alpha` is the multiplier.
+
+    Args:
+        chol_init: Initial Cholesky decomposition of the matrix, :math:`A`.
+        update_vectors: Update vectors, :math:`v_i`.
+        multiplier: The multiplier, :math:`\alpha`.
+
+    Returns:
+        The updated Cholesky decomposition of the matrix.
+
+    Notes:
+        If the updated matrix does not correspond to a positive definite matrix, the
+        function has undefined behaviour. It is the responsibility of the caller to
+        ensure that the updated matrix is positive definite as we cannot check this at runtime.
+    """
+
+    def body(chol, update_vector):
+        res = _cholesky_update(chol, update_vector, multiplier=multiplier)
+        return res, None
+
+    final_chol, _ = jax.lax.scan(body, jnp.asarray(chol_init), update_vectors)
+    return final_chol
+
+
+def _set_diagonal(x: Array, y: Array) -> Array:
+    N, _ = x.shape
+    i, j = jnp.diag_indices(N)
+    return x.at[i, j].set(y)
+
+
+def _set_triu(x: Array, val: ArrayLike) -> Array:
+    N, _ = x.shape
+    i = jnp.triu_indices(N, 1)
+    return x.at[i].set(val)
+
+
+def _cholesky_update(
+    chol: ArrayLike, update_vector: ArrayLike, multiplier: float = 1.0
+) -> Array:
+    chol = jnp.asarray(chol)
+    chol_diag = jnp.diag(chol)
+
+    # The algorithm in [1] is implemented as a double for loop. We can treat
+    # the inner loop in Algorithm 3.1 as a vector operation, and thus the
+    # whole algorithm as a single for loop, and hence can use a `tf.scan`
+    # on it.
+
+    # We use for accumulation omega and b as defined in Algorithm 3.1, since
+    # these are updated per iteration.
+
+    def scan_body(carry, inp):
+        _, _, omega, b = carry
+        index, diagonal_member, col = inp
+        omega_at_index = omega[..., index]
+
+        # Line 4
+        new_diagonal_member = jnp.sqrt(
+            jnp.square(diagonal_member) + multiplier / b * jnp.square(omega_at_index)
+        )
+        # `scaling_factor` is the same as `gamma` on Line 5.
+        scaling_factor = jnp.square(diagonal_member) * b + multiplier * jnp.square(
+            omega_at_index
+        )
+
+        # The following updates are the same as the for loop in lines 6-8.
+        omega = omega - (omega_at_index / diagonal_member)[..., None] * col
+        new_col = new_diagonal_member[..., None] * (
+            col / diagonal_member[..., None]
+            + (multiplier * omega_at_index / scaling_factor)[..., None] * omega
+        )
+        b = b + multiplier * jnp.square(omega_at_index / diagonal_member)
+        return (new_diagonal_member, new_col, omega, b), (
+            new_diagonal_member,
+            new_col,
+            omega,
+            b,
+        )
+
+    # We will scan over the columns.
+    chol = chol.T
+
+    _, (new_diag, new_chol, _, _) = jax.lax.scan(
+        scan_body,
+        (0.0, jnp.zeros_like(chol[0]), update_vector, 1.0),
+        (jnp.arange(0, chol.shape[0]), chol_diag, chol),
+    )
+
+    new_chol = new_chol.T
+    new_chol = _set_diagonal(new_chol, new_diag)
+    new_chol = _set_triu(new_chol, 0.0)
+    new_chol = jnp.where(jnp.isfinite(new_chol), new_chol, 0.0)
+    return new_chol

cuthbertlib/resampling/__init__.py

@@ -0,0 +1,3 @@
+from cuthbertlib.resampling import killing, multinomial, systematic
+from cuthbertlib.resampling.protocols import ConditionalResampling, Resampling
+from cuthbertlib.resampling.utils import inverse_cdf

cuthbertlib/resampling/killing.py

@@ -0,0 +1,79 @@
+"""Implements killing resampling."""
+
+from functools import partial
+
+from jax import numpy as jnp
+from jax import random
+from jax.scipy.special import logsumexp
+
+from cuthbertlib.resampling import multinomial
+from cuthbertlib.resampling.protocols import (
+    conditional_resampling_decorator,
+    resampling_decorator,
+)
+from cuthbertlib.types import Array, ArrayLike
+
+_DESCRIPTION = """
+The Killing resampling is a simple resampling mechanism that checks if 
+particles should be replaced or not, based on their weights. 
+If they should be replaced, they are replaced by another particle using 
+multinomial resampling on residual weights. It presents the benefit of not 
+"breaking" trajectories as much as multinomial resampling, and therefore is
+stable in contexts where the trajectories are important (typically when dealing
+with continuous-time models). 
+
+By construction, it requires the same number of sampled indices `n` as the
+number of particles `logits.shape[0]`.
+"""
+
+
+@partial(resampling_decorator, name="Killing", desc=_DESCRIPTION)
+def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
+    logits = jnp.asarray(logits)
+    key_1, key_2 = random.split(key)
+    N = logits.shape[0]
+    if n != N:
+        raise AssertionError(
+            "The number of sampled indices must be equal to the number of "
+            f"particles for `Killing` resampling. Got {n} instead of {N}."
+        )
+
+    max_logit = jnp.max(logits)
+    log_uniforms = jnp.log(random.uniform(key_1, (N,)))
+
+    survived = log_uniforms <= logits - max_logit
+    if_survived = jnp.arange(N)  # If the particle survives, it keeps its index
+    otherwise = multinomial.resampling(
+        key_2, logits, N
+    )  # otherwise, it is replaced by another particle
+    idx = jnp.where(survived, if_survived, otherwise)
+    return idx
+
+
+@partial(conditional_resampling_decorator, name="Killing", desc=_DESCRIPTION)
+def conditional_resampling(
+    key: Array, logits: ArrayLike, n: int, pivot_in: int, pivot_out: int
+) -> Array:
+    # Unconditional resampling
+    key_resample, key_shuffle = random.split(key)
+    idx = resampling(key_resample, logits, n)
+
+    # Conditional rolling pivot
+    max_logit = jnp.max(logits)
+
+    pivot_logits = _log1mexp(logits - max_logit)
+    pivot_logits -= jnp.log(n)
+    pivot_logits = pivot_logits.at[pivot_out].set(-jnp.inf)
+    pivot_logits_i = _log1mexp(logsumexp(pivot_logits))
+    pivot_logits = pivot_logits.at[pivot_out].set(pivot_logits_i)
+
+    pivot_weights = jnp.exp(pivot_logits - logsumexp(pivot_logits))
+    pivot = random.choice(key_shuffle, n, p=pivot_weights)
+    idx = jnp.roll(idx, pivot_in - pivot)
+    idx = idx.at[pivot_in].set(pivot_out)
+    return idx
+
+
+def _log1mexp(x: ArrayLike) -> Array:
+    # There is probably a better way to do this
+    return jnp.log(1 - jnp.exp(x))

cuthbertlib/resampling/multinomial.py

@@ -0,0 +1,53 @@
+"""Implements multinomial resampling."""
+
+from functools import partial
+
+import jax
+from jax import numpy as jnp
+from jax import random
+
+from cuthbertlib.resampling.protocols import (
+    conditional_resampling_decorator,
+    resampling_decorator,
+)
+from cuthbertlib.resampling.utils import inverse_cdf
+from cuthbertlib.types import Array, ArrayLike
+
+_DESCRIPTION = """
+This has higher variance than other resampling schemes as it samples from
+the ancestors independently. It should only be used for illustration purposes,
+or if your algorithm *REALLY REALLY* needs independent samples.
+As a rule of thumb, you often don't."""
+
+
+@partial(resampling_decorator, name="Multinomial", desc=_DESCRIPTION)
+def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
+    # In practice we don't have to sort the generated uniforms, but searchsorted
+    # works faster and is more stable if both inputs are sorted, so we use the
+    # _sorted_uniforms from N. Chopin, but still use searchsorted instead of his
+    # O(N) loop as our code is meant to work on GPU where searchsorted is
+    # O(log(N)) anyway.
+    # We then permute the indices to enforce exchangeability.
+
+    key_uniforms, key_shuffle = random.split(key)
+    sorted_uniforms = _sorted_uniforms(key_uniforms, n)
+    idx = inverse_cdf(sorted_uniforms, logits)
+    return random.permutation(key_shuffle, idx)
+
+
+@partial(conditional_resampling_decorator, name="Multinomial", desc=_DESCRIPTION)
+def conditional_resampling(
+    key: Array, logits: ArrayLike, n: int, pivot_in: int, pivot_out: int
+) -> Array:
+    idx = resampling(key, logits, n)
+    idx = idx.at[pivot_in].set(pivot_out)
+    return idx
+
+
+@partial(jax.jit, static_argnames=("n",))
+def _sorted_uniforms(key: Array, n: int) -> Array:
+    # This is a small modification of the code from N. Chopin to output sorted
+    # log-uniforms *directly*. N. Chopin's code outputs sorted uniforms.
+    us = random.uniform(key, (n + 1,))
+    z = jnp.cumsum(-jnp.log(us))
+    return z[:-1] / z[-1]

cuthbertlib/resampling/protocols.py

@@ -0,0 +1,92 @@
+"""Shared protocols for resampling algorithms."""
+
+from typing import Protocol, runtime_checkable
+
+import jax
+
+from cuthbertlib.types import Array, ArrayLike, KeyArray
+
+
+@runtime_checkable
+class Resampling(Protocol):
+    """Protocol for resampling operations."""
+
+    def __call__(self, key: KeyArray, logits: ArrayLike, n: int) -> Array:
+        """Computes resampling indices according to given logits.
+
+        Args:
+            key: JAX PRNG key.
+            logits: Logits.
+            n: Number of indices to sample.
+
+        Returns:
+            Array of resampling indices.
+        """
+        ...
+
+
+@runtime_checkable
+class ConditionalResampling(Protocol):
+    """Protocol for conditional resampling operations."""
+
+    def __call__(
+        self,
+        key: KeyArray,
+        logits: ArrayLike,
+        n: int,
+        pivot_in: int,
+        pivot_out: int,
+    ) -> Array:
+        """Conditional resampling.
+
+        Args:
+            key: JAX PRNG key.
+            logits: Log-weights, possibly unnormalized.
+            n: Number of indices to sample.
+            pivot_in: Index of the particle to keep.
+            pivot_out: Value of the output at index `pivot_in`.
+
+        Returns:
+            Array of size n with indices to use for resampling.
+        """
+        ...
+
+
+def resampling_decorator(func: Resampling, name: str, desc: str = "") -> Resampling:
+    """Decorate Resampling function with unified docstring."""
+    doc = f"""
+    {name} resampling. {desc}
+
+    Args:
+        key: PRNGKey to use in resampling
+        logits: Log-weights, possibly unnormalized.
+        n: Number of indices to sample.
+    
+    Returns:
+        Array of size n with indices to use for resampling.
+    """
+
+    func.__doc__ = doc
+    return jax.jit(func, static_argnames=("n",))
+
+
+def conditional_resampling_decorator(
+    func: ConditionalResampling, name: str, desc: str = ""
+) -> ConditionalResampling:
+    """Decorate ConditionalResampling function with unified docstring."""
+    doc = f"""
+    {name} conditional resampling. {desc}
+
+    Args:
+        key: PRNGKey to use in resampling
+        logits: Log-weights, possibly unnormalized.
+        n: Number of indices to sample
+        pivot_in: Index of the particle to keep
+        pivot_out: Value of the output at index `pivot_in`
+    
+    Returns:
+        Array of size n with indices to use for resampling.
+    """
+
+    func.__doc__ = doc
+    return jax.jit(func, static_argnames=("n",))