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

cuthbertlib/resampling/systematic.py

@@ -0,0 +1,78 @@
+"""Implements systematic resampling."""
+
+from functools import partial
+
+from jax import numpy as jnp
+from jax import random
+from jax.lax import cond, select
+from jax.scipy.special import logsumexp
+
+from cuthbertlib.resampling.protocols import (
+    conditional_resampling_decorator,
+    resampling_decorator,
+)
+from cuthbertlib.resampling.utils import inverse_cdf
+from cuthbertlib.types import Array, ArrayLike
+
+_DESCRIPTION = """
+The Systematic resampling is a variance reduction which places marginally
+uniform samples into the [0, 1] interval but only requires one uniform random.
+"""
+
+
+@partial(resampling_decorator, name="Systematic", desc=_DESCRIPTION)
+def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
+    us = (random.uniform(key, ()) + jnp.arange(n)) / n
+    return inverse_cdf(us, logits)
+
+
+@partial(conditional_resampling_decorator, name="Systematic", desc=_DESCRIPTION)
+def conditional_resampling(
+    key: Array, logits: ArrayLike, n: int, pivot_in: int, pivot_out: int
+) -> Array:
+    logits = jnp.asarray(logits)
+    # FIXME: no need for normalizing in theory
+    N = logits.shape[0]
+    logits -= logsumexp(logits)
+
+    # FIXME: this rolling should be done in a single function, but this is killing me.
+    arange = jnp.arange(N)
+    logits = jnp.roll(logits, -pivot_out)
+    arange = jnp.roll(arange, -pivot_out)
+
+    idx = conditional_resampling_0_to_0(key, logits, n)
+    idx = arange[idx]
+    idx = jnp.roll(idx, pivot_in)
+    return idx
+
+
+def conditional_resampling_0_to_0(
+    key: Array,
+    logits: ArrayLike,
+    n: int,
+) -> Array:
+    logits = jnp.asarray(logits)
+
+    N = logits.shape[0]
+    weights = jnp.exp(logits - logsumexp(logits))
+    tmp = n * weights[0]
+    tmp_floor = jnp.floor(tmp)
+
+    U, V, W = random.uniform(key, (3,))
+
+    def _otherwise():
+        rem = tmp - tmp_floor
+        p_cond = rem * (tmp_floor + 1) / tmp
+        return select(V < p_cond, rem * U, rem + (1.0 - rem) * U)
+
+    uniform = cond(tmp <= 1, lambda: tmp * U, _otherwise)
+
+    linspace = (jnp.arange(n) + uniform) / n
+    idx = inverse_cdf(linspace, logits)
+
+    n_zero = jnp.sum(idx == 0)
+    zero_loc = jnp.flatnonzero(idx == 0, size=n, fill_value=-1)
+    roll_idx = jnp.floor(n_zero * W).astype(int)
+
+    idx = select(n_zero == 1, idx, jnp.roll(idx, -zero_loc[roll_idx]))
+    return jnp.clip(idx, 0, N - 1)

cuthbertlib/resampling/utils.py

@@ -0,0 +1,82 @@
+"""Utility functions (inverse CDF sampling) for resampling algorithms."""
+
+import jax
+import jax.numpy as jnp
+import numba as nb
+import numpy as np
+from jax.lax import platform_dependent
+from jax.scipy.special import logsumexp
+
+from cuthbertlib.types import Array, ArrayLike
+
+
+@jax.jit
+def inverse_cdf(sorted_uniforms: ArrayLike, logits: ArrayLike) -> Array:
+    """Inverse CDF sampling for resampling algorithms.
+
+    The implementation branches depending on the platform being CPU or GPU (and other parallel envs)
+    1. The CPU implementation is a numba-compiled specialized searchsorted(arr, vals) for *sorted* uniforms
+        which is guaranteed to run in O(N + M) where N is the size of logits and M that of sorted_uniforms.
+        This could be replaced mutatis mutandis by np.searchsorted (not jnp.searchsorted!) but the latter
+        does not guarantee this execution time.
+    2. On GPU, we use searchsorted with a sorting strategy: this proceeds by merge (arg) sorting,
+        which works in log(N+M) and is efficient for large arrays of sorted uniforms, typically our setting.
+
+    Args:
+        sorted_uniforms: Sorted uniforms.
+        logits: Log-weights, possibly un-normalized.
+
+    Returns:
+        Indices of the particles to be resampled.
+    """
+    weights = jnp.exp(logits - logsumexp(logits))
+    return platform_dependent(
+        sorted_uniforms, weights, cpu=_inverse_cdf_cpu, default=_inverse_cdf_default
+    )
+
+
+@jax.jit
+def _inverse_cdf_default(sorted_uniforms: ArrayLike, weights: ArrayLike) -> Array:
+    weights = jnp.asarray(weights)
+    M = weights.shape[0]
+    cs = jnp.cumsum(weights)
+    idx = jnp.searchsorted(cs, sorted_uniforms, method="sort")
+    return jnp.clip(idx, 0, M - 1).astype(
+        int
+    )  # Ensure indices are integers from the same dtype as basic jax ints.
+
+
+@jax.jit
+def _inverse_cdf_cpu(sorted_uniforms: ArrayLike, weights: ArrayLike) -> Array:
+    sorted_uniforms = jnp.asarray(sorted_uniforms)
+    weights = jnp.asarray(weights)
+    M = weights.shape[0]
+    N = sorted_uniforms.shape[0]
+    idx = jnp.zeros(N, dtype=int)
+
+    def callback(args):
+        su, w, idx_ = args
+        idx_ = np.array(idx_, dtype=idx.dtype)
+        su = np.asarray(su)
+        w = np.asarray(w)
+        _inverse_cdf_numba(su, w, idx_)
+        return idx_
+
+    idx = jax.pure_callback(
+        callback, idx, (sorted_uniforms, weights, idx), vmap_method="sequential"
+    )
+    return jnp.clip(idx, 0, M - 1)
+
+
+@nb.njit
+def _inverse_cdf_numba(su, ws, idx):
+    j = 0
+    s = ws[0]
+    M = su.shape[0]
+    N = ws.shape[0]
+
+    for n in range(M):
+        while su[n] > s and j < N - 1:
+            j += 1
+            s += ws[j]
+        idx[n] = j

cuthbertlib/smc/ess.py

@@ -0,0 +1,24 @@
+"""Importance sampling effective sample size (ESS) computation."""
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+from jax.typing import ArrayLike
+
+
+def log_ess(log_weights: ArrayLike) -> Array:
+    """Compute the logarithm of the effective sample size (ESS).
+
+    Args:
+        log_weights: Array of log weights for the particles.
+    """
+    return 2 * jax.nn.logsumexp(log_weights) - jax.nn.logsumexp(2 * log_weights)
+
+
+def ess(log_weights: ArrayLike) -> Array:
+    """Compute the effective sample size (ESS).
+
+    Args:
+        log_weights: Array of log weights for the particles.
+    """
+    return jnp.exp(log_ess(log_weights))

cuthbertlib/smc/smoothing/exact_sampling.py

@@ -0,0 +1,111 @@
+"""Implements exact backward sampling for smoothing in SMC."""
+
+import jax
+from jax import numpy as jnp
+from jax import random, vmap
+from jax.scipy.special import logsumexp
+
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    KeyArray,
+    LogConditionalDensity,
+)
+
+
+def log_weights_single(
+    x0: ArrayTreeLike,
+    x1: ArrayTreeLike,
+    log_weight_x0: ArrayLike,
+    log_density: LogConditionalDensity,
+) -> Array:
+    """Compute smoothing weight for a single sample x0 given a single sample x1.
+
+    Args:
+        x0: The previous state.
+        x1: The current state.
+        log_weight_x0: The log weights of the previous state.
+        log_density: The log density function of x1 given x0.
+
+    Returns:
+        The smoothing weight for sample x0 given a single sample x1.
+    """
+    return jnp.asarray(log_weight_x0) + log_density(x0, x1)
+
+
+def log_weights(x0_all, x1, log_weight_x0_all, log_density) -> Array:
+    """Compute log smoothing weights over a collection of x0 given a single x1.
+
+    Args:
+        x0_all: Collection of previous states.
+        x1: The current state.
+        log_weight_x0_all: Collection of log weights of the previous state.
+        log_density: The log density function of x1 given x0.
+
+    Returns:
+        Log normalized smoothing weights for each sample x0 given single sample x1.
+    """
+    backward_log_weights_all = vmap(
+        lambda x0, log_weight_x0: log_weights_single(x0, x1, log_weight_x0, log_density)
+    )(x0_all, log_weight_x0_all)
+
+    # Log normalize
+    backward_log_weights_all = backward_log_weights_all - logsumexp(
+        backward_log_weights_all, axis=0
+    )
+    return backward_log_weights_all
+
+
+def simulate_single(
+    key, x0_all, x1, log_weight_x0_all, log_density
+) -> tuple[ArrayTree, Array]:
+    """Sample x0 from a collection given a single x1.
+
+    Args:
+        key: A JAX random key.
+        x0_all: Collection of previous states.
+        x1: The current state.
+        log_weight_x0_all: Collection of log weights of the previous state.
+        log_density: The log density function of x1 given x0.
+
+    Returns:
+        A single sample x0 from the smoothing trajectory along with its index.
+    """
+    backward_log_weights_all = log_weights(x0_all, x1, log_weight_x0_all, log_density)
+    sampled_index = random.categorical(key, backward_log_weights_all)
+    return jax.tree.map(lambda z: z[sampled_index], x0_all), sampled_index
+
+
+def simulate(
+    key: KeyArray,
+    x0_all: ArrayTreeLike,
+    x1_all: ArrayTreeLike,
+    log_weight_x0_all: ArrayLike,
+    log_density: LogConditionalDensity,
+    x1_ancestor_indices: ArrayLike,
+) -> tuple[ArrayTree, Array]:
+    """Implements the exact backward sampling algorithm for smoothing in SMC.
+
+    Some arguments are only included for protocol compatibility and not used in this
+    implementation.
+
+    Args:
+        key: JAX PRNG key.
+        x0_all: A collection of previous states $x_0$.
+        x1_all: A collection of current states $x_1$.
+        log_weight_x0_all: The log weights of $x_0$.
+        log_density: The log density function of $x_1$ given $x_0$.
+        x1_ancestor_indices: The ancestor indices of $x_1$. Not used.
+
+    Returns:
+        A collection of samples $x_0$ and their sampled indices.
+    """
+    log_weight_x0_all = jnp.asarray(log_weight_x0_all)
+    n_smoother_particles = jax.tree.leaves(x1_all)[0].shape[0]
+    keys = random.split(key, n_smoother_particles)
+
+    return vmap(
+        lambda k, x1: simulate_single(k, x0_all, x1, log_weight_x0_all, log_density)
+    )(keys, x1_all)

cuthbertlib/smc/smoothing/mcmc.py

@@ -0,0 +1,76 @@
+"""Implements MCMC backward smoothing in SMC."""
+
+import jax
+from jax import numpy as jnp
+from jax import random
+
+from cuthbertlib.resampling import multinomial
+from cuthbertlib.smc.smoothing.tracing import simulate as ancestor_tracing_simulate
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    KeyArray,
+    LogConditionalDensity,
+)
+
+
+def simulate(
+    key: KeyArray,
+    x0_all: ArrayTreeLike,
+    x1_all: ArrayTreeLike,
+    log_weight_x0_all: ArrayLike,
+    log_density: LogConditionalDensity,
+    x1_ancestor_indices: ArrayLike,
+    n_steps: int,
+) -> tuple[ArrayTree, Array]:
+    """Implements the IMH algorithm for smoothing in SMC.
+
+    Args:
+        key: JAX PRNG key.
+        x0_all: A collection of previous states $x_0$.
+        x1_all: A collection of current states $x_1$.
+        log_weight_x0_all: The log weights of $x_0$.
+        log_density: The log density function of $x_1$ given $x_0$.
+        x1_ancestor_indices: The ancestor indices of $x_1$.
+        n_steps: Number of MCMC steps to perform.
+
+    Returns:
+        A collection of samples $x_0$ and their sampled indices.
+
+    References:
+        https://arxiv.org/abs/2207.00976
+    """
+    key, subkey = random.split(key)
+    x0_init, x1_ancestor_indices = ancestor_tracing_simulate(
+        subkey, x0_all, x1_all, log_weight_x0_all, log_density, x1_ancestor_indices
+    )
+    n_samples = x1_ancestor_indices.shape[0]
+
+    keys = random.split(key, (n_steps * 2)).reshape((n_steps, 2))
+
+    def body(carry, keys_t):
+        # IMH proposal
+        idx, x0_res, idx_log_p = carry
+        key_prop, key_acc = keys_t
+
+        prop_idx = multinomial.resampling(key_prop, log_weight_x0_all, n_samples)
+        x0_prop = jax.tree.map(lambda z: z[prop_idx], x0_all)
+        prop_log_p = jax.vmap(log_density)(x0_prop, x1_all)
+
+        log_alpha = prop_log_p - idx_log_p
+
+        lu = jnp.log(random.uniform(key_acc, (n_samples,)))
+        acc = lu < log_alpha
+
+        idx: Array = jnp.where(acc, prop_idx, idx)
+        x0_res: ArrayTreeLike = jax.tree.map(lambda z: z[idx], x0_all)
+        idx_log_p: Array = jnp.where(acc, prop_log_p, idx_log_p)
+        return (idx, x0_res, idx_log_p), None
+
+    x0_init = jax.tree.map(lambda z: z[x1_ancestor_indices], x0_all)
+    init_log_p = jax.vmap(log_density)(x1_all, x0_init)
+    init = (x1_ancestor_indices, x0_init, init_log_p)
+    (out_index, out_samples, _), _ = jax.lax.scan(body, init, keys)
+    return out_samples, out_index

cuthbertlib/smc/smoothing/protocols.py

@@ -0,0 +1,44 @@
+"""Shared protocols for backward smoothing functions in SMC."""
+
+from typing import Protocol, runtime_checkable
+
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    KeyArray,
+    LogConditionalDensity,
+)
+
+
+@runtime_checkable
+class BackwardSampling(Protocol):
+    """Protocol for backward sampling functions."""
+
+    def __call__(
+        self,
+        key: KeyArray,
+        x0_all: ArrayTreeLike,
+        x1_all: ArrayTreeLike,
+        log_weight_x0_all: ArrayLike,
+        log_density: LogConditionalDensity,
+        x1_ancestor_indices: ArrayLike,
+    ) -> tuple[ArrayTree, Array]:
+        """Performs a backward sampling step.
+
+        Samples a collection of $x_0$ that combine with the provided $x_1$ to
+        give a collection of pairs $(x_0, x_1)$ from the smoothing distribution.
+
+        Args:
+            key: JAX PRNG key.
+            x0_all: A collection of previous states $x_0$.
+            x1_all: A collection of current states $x_1$.
+            log_weight_x0_all: The log weights of $x_0$.
+            log_density: The log density function of $x_1$ given $x_0$.
+            x1_ancestor_indices: The ancestor indices of $x_1$.
+
+        Returns:
+            A collection of samples $x_0$ and their sampled indices.
+        """
+        ...

cuthbertlib/smc/smoothing/tracing.py

@@ -0,0 +1,45 @@
+"""Implements the ancestor/genealogy tracing algorithm for smoothing in SMC."""
+
+import jax
+from jax import numpy as jnp
+
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    KeyArray,
+    LogConditionalDensity,
+)
+
+
+def simulate(
+    key: KeyArray,
+    x0_all: ArrayTreeLike,
+    x1_all: ArrayTreeLike,
+    log_weight_x0_all: ArrayLike,
+    log_density: LogConditionalDensity,
+    x1_ancestor_indices: ArrayLike,
+) -> tuple[ArrayTree, Array]:
+    """Implements the ancestor/genealogy tracing algorithm for smoothing in SMC.
+
+    Some arguments are only included for protocol compatibility and not used in this
+    implementation.
+
+    Args:
+        key: JAX PRNG key. Not used
+        x0_all: A collection of previous states $x_0$.
+        x1_all: A collection of current states $x_1$. Not used.
+        log_weight_x0_all: The log weights of $x_0$. Not used.
+        log_density: The log density function of $x_1$ given $x_0$. Not used.
+        x1_ancestor_indices: The ancestor indices of $x_1$.
+
+    Returns:
+        A collection of samples $x_0$ and their sampled indices.
+
+    References:
+        https://arxiv.org/abs/2207.00976
+    """
+    x1_ancestor_indices = jnp.asarray(x1_ancestor_indices)
+    x0 = jax.tree.map(lambda z: z[x1_ancestor_indices], x0_all)
+    return x0, x1_ancestor_indices

cuthbertlib/stats/multivariate_normal.py

@@ -0,0 +1,102 @@
+"""Multivariate normal distribution functions with chol_cov input."""
+
+from functools import partial
+
+import numpy as np
+from jax import lax
+from jax import numpy as jnp
+from jax._src.numpy.util import promote_dtypes_inexact
+from jax._src.typing import Array, ArrayLike
+
+from cuthbertlib.linalg import collect_nans_chol
+
+
+def logpdf(
+    x: ArrayLike, mean: ArrayLike, chol_cov: ArrayLike, nan_support: bool = True
+) -> Array:
+    """Multivariate normal log probability distribution function with chol_cov input.
+
+    Here `chol_cov` is the (generalized) Cholesky factor of the covariance matrix.
+    Modified version of `jax.scipy.stats.multivariate_normal.logpdf` which takes
+    the full covariance matrix as input.
+
+    Args:
+        x: Value at which to evaluate the PDF.
+        mean: Mean of the distribution.
+        chol_cov: Generalized Cholesky factor of the covariance matrix of the distribution.
+        nan_support: If `True`, ignores NaNs in `x` by projecting the distribution onto the
+            lower-dimensional subspace spanned by the non-NaN entries of `x`. Note that
+            `nan_support=True` uses the [tria][cuthbertlib.linalg.tria] operation (QR
+            decomposition), and therefore increases the internal complexity of the function
+            from $O(n^2)$ to $O(n^3)$, where $n$ is the dimension of `x`.
+
+    Returns:
+        Array of logpdf values.
+    """
+    x, mean, chol_cov = promote_dtypes_inexact(x, mean, chol_cov)
+
+    # If nan_support is True, we need to collect the NaNs at the top of the covariance matrix
+    # this uses a QR decomposition so is more expensive
+    if nan_support:
+        flag = jnp.isnan(x)
+        flag, chol_cov, x, mean = collect_nans_chol(flag, chol_cov, x, mean)
+        mean = jnp.asarray(mean)
+        x = jnp.asarray(x)
+        chol_cov = jnp.asarray(chol_cov)
+
+    if not mean.shape and not np.shape(x):
+        # Both mean and x are scalars
+        return -1 / 2 * jnp.square(x - mean) / chol_cov**2 - 1 / 2 * (
+            jnp.log(2 * np.pi) + 2 * jnp.log(chol_cov)
+        )
+    else:
+        n = mean.shape[-1] if mean.shape else x.shape[-1]
+        if not np.shape(chol_cov):
+            y = x - mean
+            return -1 / 2 * jnp.einsum("...i,...i->...", y, y) / chol_cov**2 - n / 2 * (
+                jnp.log(2 * np.pi) + 2 * jnp.log(chol_cov)
+            )
+        elif chol_cov.ndim == 1:
+            y = (x - mean) / chol_cov
+            return (
+                -1 / 2 * jnp.einsum("...i,...i->...", y, y)
+                - n / 2 * jnp.log(2 * np.pi)
+                - jnp.log(jnp.abs(chol_cov)).sum(-1)
+            )
+        else:
+            if chol_cov.ndim < 2 or chol_cov.shape[-2:] != (n, n):
+                raise ValueError("multivariate_normal.logpdf got incompatible shapes")
+            y = jnp.vectorize(
+                partial(lax.linalg.triangular_solve, lower=True, transpose_a=True),
+                signature="(n,n),(n)->(n)",
+            )(chol_cov, x - mean)
+            return (
+                -1 / 2 * jnp.einsum("...i,...i->...", y, y)
+                - n / 2 * jnp.log(2 * np.pi)
+                - jnp.log(jnp.abs(chol_cov.diagonal(axis1=-1, axis2=-2))).sum(-1)
+            )
+
+
+def pdf(
+    x: ArrayLike, mean: ArrayLike, chol_cov: ArrayLike, nan_support: bool = True
+) -> Array:
+    """Multivariate normal probability distribution function with chol_cov input.
+
+    Here `chol_cov` is the (generalized) Cholesky factor of the covariance matrix.
+    Modified version of `jax.scipy.stats.multivariate_normal.pdf` which takes
+    the full covariance matrix as input.
+
+    Args:
+        x: Value at which to evaluate the PDF.
+        mean: Mean of the distribution.
+        chol_cov: Generalized Cholesky factor of the covariance matrix of the distribution.
+        nan_support: If `True`, ignores NaNs in `x` by projecting the distribution onto the
+            lower-dimensional subspace spanned by the non-NaN entries of `x`. Note that
+            `nan_support=True` uses the [tria][cuthbertlib.linalg.tria] operation (QR
+            decomposition), and therefore increases the internal complexity of the function
+            from $O(n^2)$ to $O(n^3)$, where $n$ is the dimension of `x`.
+
+    Returns:
+        Array of pdf values.
+    """
+    return lax.exp(logpdf(x, mean, chol_cov, nan_support))

cuthbert-0.0.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-cuthbert/__init__.py,sha256=60_FB1nfduZLPphbjEc7WRpebCnVYiwMyKuD-7yZBvw,281
-cuthbert/filtering.py,sha256=HaUPJWhBO8P5IWlRYhLwCeOEtlYenAvWjYsQaF_cPX0,2700
-cuthbert/inference.py,sha256=u02wVKGu7mIdqn2XhcSZ93xXyPIQkxzSvxJXMH7Zo5k,7600
-cuthbert/smoothing.py,sha256=qvNAWYTGGaEUEBp7HAtYLtF1UOQK9E7u3V_l4XgxeaI,4960
-cuthbert/utils.py,sha256=0JQgRiyVs4SXZ0ullh4OmAMtyoOtuCzvYuDmpu9jXOE,1023
-cuthbert-0.0.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-cuthbertlib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cuthbertlib/types.py,sha256=sgHC0J7W_0wOzMm5dlLZXwJB3W1xvzq2J6cxl6U831w,1020
-cuthbert-0.0.1.dist-info/METADATA,sha256=Q-MikVBB27Lne_4q7HvrBVNgln6zY23RdY5PdnIeE4M,7040
-cuthbert-0.0.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-cuthbert-0.0.1.dist-info/top_level.txt,sha256=R7-G6fUQZSMNMcM4-IcpHKtY9CZOQeejEpVW9q-Sarw,21
-cuthbert-0.0.1.dist-info/RECORD,,