smcx 1.0.0__py3-none-any.whl

smcx/__init__.py

@@ -0,0 +1,94 @@
+# Copyright 2026 Michael Ellis
+# SPDX-License-Identifier: Apache-2.0
+
+"""Sequential Monte Carlo and particle filtering in JAX."""
+
+from importlib.metadata import PackageNotFoundError as _PackageNotFoundError
+from importlib.metadata import version as _version
+
+from smcx.auxiliary import auxiliary_filter
+from smcx.bootstrap import bootstrap_filter
+from smcx.containers import (
+    LiuWestPosterior,
+    ParticleFilterPosterior,
+    ParticleFilterResult,
+    ParticleState,
+    SMC2Posterior,
+    TemperedPosterior,
+)
+from smcx.diagnostics import (
+    crps,
+    cumulative_log_score,
+    diagnose,
+    log_bayes_factor,
+    log_ml_increments,
+    param_weighted_mean,
+    param_weighted_quantile,
+    pareto_k_diagnostic,
+    particle_diversity,
+    posterior_predictive_sample,
+    replicated_log_ml,
+    tail_ess,
+    weighted_mean,
+    weighted_quantile,
+    weighted_variance,
+)
+from smcx.exceptions import DegenerateWeightsError
+from smcx.guided import guided_filter
+from smcx.liu_west import liu_west_filter
+from smcx.resampling import (
+    multinomial,
+    residual,
+    stratified,
+    systematic,
+)
+from smcx.simulate import simulate
+from smcx.smc2 import smc2
+from smcx.tempering import temper
+from smcx.weights import ess, log_ess, log_normalize, normalize
+
+try:
+    __version__ = _version("smcx")
+except _PackageNotFoundError:
+    __version__ = "1.0.0"
+
+__all__ = [
+    "DegenerateWeightsError",
+    "LiuWestPosterior",
+    "ParticleFilterPosterior",
+    "ParticleFilterResult",
+    "ParticleState",
+    "SMC2Posterior",
+    "TemperedPosterior",
+    "__version__",
+    "auxiliary_filter",
+    "bootstrap_filter",
+    "crps",
+    "cumulative_log_score",
+    "diagnose",
+    "ess",
+    "guided_filter",
+    "liu_west_filter",
+    "log_bayes_factor",
+    "log_ess",
+    "log_ml_increments",
+    "log_normalize",
+    "multinomial",
+    "normalize",
+    "param_weighted_mean",
+    "param_weighted_quantile",
+    "pareto_k_diagnostic",
+    "particle_diversity",
+    "posterior_predictive_sample",
+    "replicated_log_ml",
+    "residual",
+    "simulate",
+    "smc2",
+    "stratified",
+    "systematic",
+    "tail_ess",
+    "temper",
+    "weighted_mean",
+    "weighted_quantile",
+    "weighted_variance",
+]

smcx/__main__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Michael Ellis
+# SPDX-License-Identifier: Apache-2.0
+
+"""Allow running the package with ``python -m smcx``."""
+
+from smcx import __version__
+
+
+def main() -> None:
+    """Print package version."""
+    print(f"smcx {__version__}")
+
+
+if __name__ == "__main__":
+    main()

smcx/_utils.py

@@ -0,0 +1,172 @@
+# Copyright 2026 Michael Ellis
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared private helpers for particle filters.
+
+These utilities are extracted from the individual filter modules to
+eliminate duplication.  They are not part of the public API.
+"""
+
+from collections.abc import Callable
+
+import jax.numpy as jnp
+from jax import lax, vmap
+from jaxtyping import Array, Float, Int
+
+from smcx.containers import ParticleState
+from smcx.types import PRNGKeyT
+from smcx.weights import ess as compute_ess
+from smcx.weights import log_normalize, normalize
+
+
+def _prepend(first: Array, rest: Array) -> Array:
+    """Prepend a single leading slice to an array along axis 0.
+
+    Args:
+        first: Array of shape ``(...)``.
+        rest: Array of shape ``(T, ...)``.
+
+    Returns:
+        Concatenated array of shape ``(T+1, ...)``.
+    """
+    return jnp.concatenate([jnp.expand_dims(first, 0), rest], axis=0)
+
+
+def _weighted_quantile_1d(
+    particles: Float[Array, " num_particles"],
+    weights: Float[Array, " num_particles"],
+    q: Float[Array, " num_quantiles"],
+) -> Float[Array, " num_quantiles"]:
+    """Compute weighted quantiles for a single 1-D vector.
+
+    Sorts particles, builds a midpoint CDF from the normalised
+    weights, and interpolates at the requested quantile levels.
+
+    Args:
+        particles: Particle values for one dimension.
+        weights: Normalised weights (sum to one).
+        q: Quantile levels in [0, 1].
+
+    Returns:
+        Interpolated quantile values.
+    """
+    sort_idx = jnp.argsort(particles)
+    p_sorted = particles[sort_idx]
+    w_sorted = weights[sort_idx]
+    cum_w = jnp.cumsum(w_sorted)
+    # Midpoint CDF: centre each particle's mass in its interval
+    # so that zero-weight particles don't create flat regions.
+    mid_cdf = (jnp.concatenate([jnp.zeros(1), cum_w[:-1]]) + cum_w) / 2
+    # Tiny tiebreaker ensures strict monotonicity for jnp.interp.
+    n = p_sorted.shape[0]
+    eps = jnp.arange(n, dtype=p_sorted.dtype) * 1e-12
+    return jnp.interp(q, mid_cdf + eps, p_sorted)
+
+
+def _init_standard(
+    init_key: PRNGKeyT,
+    initial_sampler: Callable,
+    log_observation_fn: Callable,
+    first_emission: Array,
+    num_particles: int,
+    log_n: Array,
+) -> tuple[Array, Array, Array, Array, Array, ParticleState]:
+    """Initialise a standard (bootstrap/auxiliary) filter at t=0.
+
+    Samples from the prior, weights by the first observation, and
+    builds the initial :class:`ParticleState`.
+
+    Args:
+        init_key: PRNG key for initialisation.
+        initial_sampler: State prior sampler ``(key, N) -> particles``.
+        log_observation_fn: Observation log-density
+            ``(emission, state) -> log_prob``.
+        first_emission: First observation y_0.
+        num_particles: Number of particles N.
+        log_n: Precomputed ``log(N)`` as a scalar array in the
+            default float dtype.
+
+    Returns:
+        Tuple of ``(particles_0, log_w_0, log_ev_0, ess_0,
+        identity_ancestors, init_state)``.
+    """
+    particles_0 = initial_sampler(init_key, num_particles)
+    log_obs_0 = vmap(lambda z: log_observation_fn(first_emission, z))(
+        particles_0
+    )
+    log_w_0, log_sum_0 = log_normalize(log_obs_0)
+    log_ev_0 = log_sum_0 - log_n
+    ess_0: Array = jnp.asarray(compute_ess(log_w_0))
+    identity_ancestors = jnp.arange(num_particles, dtype=jnp.int32)
+
+    init_state = ParticleState(
+        particles=particles_0,
+        log_weights=log_w_0,
+        log_marginal_likelihood=log_ev_0,
+    )
+    return (
+        particles_0,
+        log_w_0,
+        log_ev_0,
+        ess_0,
+        identity_ancestors,
+        init_state,
+    )
+
+
+def _conditional_resample(
+    key: PRNGKeyT,
+    log_weights: Float[Array, " num_particles"],
+    resampling_fn: Callable,
+    threshold: float,
+    num_particles: int,
+    identity: Int[Array, " num_particles"],
+) -> tuple[Array, Int[Array, " num_particles"]]:
+    """Conditionally resample particles based on ESS.
+
+    Computes the effective sample size of the given log weights
+    and resamples only when ESS falls below the threshold.
+
+    Args:
+        key: PRNG key for resampling.
+        log_weights: Normalised log weights (logsumexp = 0).
+        resampling_fn: Blackjax-compatible resampling function.
+        threshold: Absolute ESS threshold (e.g. ``0.5 * N``).
+        num_particles: Number of particles N.
+        identity: Identity ancestor indices ``arange(N)``.
+
+    Returns:
+        Tuple of ``(do_resample, ancestors)`` where *do_resample*
+        is a boolean scalar and *ancestors* are the resampled (or
+        identity) indices.
+    """
+    cur_ess = compute_ess(log_weights)
+    w_norm = normalize(log_weights)
+    do_resample: Array = jnp.asarray(cur_ess < threshold)
+    ancestors = lax.cond(
+        do_resample,
+        lambda: resampling_fn(key, w_norm, num_particles),
+        lambda: identity,
+    )
+    return do_resample, ancestors
+
+
+def _raise_if_degenerate(marginal_loglik) -> None:
+    """Raise :class:`DegenerateWeightsError` on a collapsed filter.
+
+    Host-side check: fires only in eager execution. Under a user
+    ``jax.jit`` the value is a tracer and the check is skipped — the
+    ``-inf``/NaN marginal propagates instead (see the exception's
+    docstring).
+    """
+    from jax.core import Tracer
+
+    from smcx.exceptions import DegenerateWeightsError
+
+    if isinstance(marginal_loglik, Tracer):
+        return
+    value = float(marginal_loglik)
+    if value != value or value == float("-inf"):
+        raise DegenerateWeightsError(
+            f"all particle weights collapsed (marginal log-likelihood {value})"
+        )

smcx/auxiliary.py

@@ -0,0 +1,245 @@
+# Copyright 2026 Michael Ellis
+# SPDX-License-Identifier: Apache-2.0
+
+r"""Auxiliary particle filter (Pitt & Shephard, 1999).
+
+The auxiliary particle filter (APF) improves on the bootstrap filter by
+using a *look-ahead* step that biases resampling towards particles
+likely to match the next observation **before** propagation.
+
+At each time step the APF:
+
+1. **First-stage weights** — combines the current normalised weights
+   with the look-ahead log-density
+   :math:`\log g(y_{t+1} \mid x_t^i)` to form first-stage weights.
+2. **Resamples** (conditionally on ESS) using the first-stage weights.
+3. **Propagates** resampled particles through the transition prior.
+4. **Second-stage weights** — corrects for the look-ahead bias:
+   :math:`w_t^{(2)} = p(y_{t+1} \mid x_{t+1}^i) / g(y_{t+1} \mid x_t^{a_i})`.
+
+When ``log_auxiliary_fn`` returns zero for all inputs, the APF
+reduces to the bootstrap filter.
+
+The implementation uses :func:`jax.lax.scan` so the full time-loop is
+compiled into a single XLA program.
+"""
+
+import math
+from collections.abc import Callable
+
+import jax.numpy as jnp
+import jax.random as jr
+from jax import lax, vmap
+from jaxtyping import Array, Float
+
+from smcx._utils import (
+    _conditional_resample,
+    _init_standard,
+    _prepend,
+    _raise_if_degenerate,
+)
+from smcx.containers import ParticleFilterPosterior, ParticleState
+from smcx.resampling import systematic
+from smcx.types import PRNGKeyT
+from smcx.weights import ess as compute_ess
+from smcx.weights import log_normalize
+
+
+def auxiliary_filter(
+    key: PRNGKeyT,
+    initial_sampler: Callable,
+    transition_sampler: Callable,
+    log_observation_fn: Callable,
+    log_auxiliary_fn: Callable,
+    emissions: Float[Array, "ntime emission_dim"],
+    num_particles: int,
+    resampling_fn: Callable = systematic,
+    resampling_threshold: float = 0.5,
+    *,
+    store_history: bool = True,
+) -> ParticleFilterPosterior:
+    r"""Run an auxiliary particle filter (Pitt & Shephard, 1999).
+
+    Args:
+        key: JAX PRNG key.
+        initial_sampler: Function ``(key, num_particles) -> particles``
+            that draws from the initial state distribution
+            :math:`p(z_1)`.
+        transition_sampler: Function ``(key, state) -> state`` that
+            draws from the transition distribution
+            :math:`p(z_t \mid z_{t-1})`.  Will be ``vmap``-ped over
+            the particle dimension internally.
+        log_observation_fn: Function
+            ``(emission, state) -> log_prob`` that evaluates the
+            observation log-density :math:`\log p(y_t \mid z_t)`.
+            Will be ``vmap``-ped over the particle dimension (second
+            argument) internally.
+        log_auxiliary_fn: Function
+            ``(emission, state) -> log_prob`` that evaluates the
+            look-ahead log-density
+            :math:`\log g(y_{t+1} \mid x_t)`.
+            Will be ``vmap``-ped over the particle dimension (second
+            argument) internally.  When this returns zero for all
+            inputs the APF reduces to the bootstrap filter.
+        emissions: Observed emissions, shape ``(T, D)``.
+        num_particles: Number of particles :math:`N`.
+        resampling_fn: Resampling algorithm matching the Blackjax
+            signature ``(key, weights, num_samples) -> indices``.
+            Defaults to :func:`~smcx.resampling.systematic`.
+        resampling_threshold: Fraction of ``num_particles`` below
+            which resampling is triggered (e.g. 0.5 means resample
+            when ``ESS < 0.5 * N``).
+        store_history: When False (ADR-0011), the scan stacks no
+            per-step particle/weight/ancestor histories — the returned
+            arrays cover only the final step (time axis length 1)
+            while ``ess``/``log_evidence_increments`` stay full.
+
+    Returns:
+        :class:`~smcx.containers.ParticleFilterPosterior` containing
+        filtered particles, log weights, ancestor indices, the
+        marginal log-likelihood estimate, and ESS trace.
+    """
+    key, init_key = jr.split(key)
+    log_n = jnp.asarray(math.log(num_particles))
+
+    # --- Initialise at t=0 -------------------------------------------------
+    (
+        particles_0,
+        log_w_0,
+        log_ev_0,
+        ess_0,
+        identity_ancestors,
+        init_state,
+    ) = _init_standard(
+        init_key,
+        initial_sampler,
+        log_observation_fn,
+        emissions[0],
+        num_particles,
+        log_n,
+    )
+
+    # --- Scan body for t = 1, ..., T-1 -------------------------------------
+    def _step(
+        carry: tuple[ParticleState, Array],
+        args: tuple[PRNGKeyT, Float[Array, " emission_dim"]],
+    ):
+        (state, _prev_ancestors), (step_key, y_t) = carry, args
+        k1, k2 = jr.split(step_key)
+        # Invariant: state.log_weights are normalized (logsumexp = 0).
+
+        # 1. First-stage weights: combine current weights with
+        #    look-ahead g(y_{t+1} | x_t)
+        log_aux = vmap(lambda z: log_auxiliary_fn(y_t, z))(state.particles)
+        log_first_stage = state.log_weights + log_aux
+
+        # Normalise first-stage weights for resampling
+        log_first_norm, log_first_sum = log_normalize(log_first_stage)
+
+        # 2. Conditionally resample using first-stage weights
+        threshold = resampling_threshold * num_particles
+        do_resample, ancestors = _conditional_resample(
+            k1,
+            log_first_norm,
+            resampling_fn,
+            threshold,
+            num_particles,
+            identity_ancestors,
+        )
+        resampled_particles = state.particles[ancestors]
+
+        # Store the look-ahead values for ancestors (needed for
+        # second-stage correction)
+        log_aux_ancestors = log_aux[ancestors]
+
+        # 3. Propagate through transition
+        keys = jr.split(k2, num_particles)
+        propagated = vmap(transition_sampler)(keys, resampled_particles)
+
+        # 4. Second-stage weights: observation / look-ahead adjustment
+        log_obs = vmap(lambda z: log_observation_fn(y_t, z))(propagated)
+        log_second_stage = log_obs - log_aux_ancestors
+
+        # Compute evidence increment and normalize.
+        # If resampled: first-stage weights were used for resampling,
+        #   the evidence increment is the product of two factors:
+        #   (a) E_W[g] = sum_i W_i * g_i  (first-stage normaliser)
+        #   (b) (1/N) sum_j w_j^(2)       (mean second-stage weight)
+        # If not resampled: standard importance weighting,
+        #   increment = logsumexp(log_w_old + log_obs)
+        log_w_unnorm = jnp.where(
+            do_resample,
+            log_second_stage,
+            state.log_weights + log_obs,
+        )
+        log_w_norm, log_sum = log_normalize(log_w_unnorm)
+
+        # log_first_sum = logsumexp(log_w_norm_old + log_aux)
+        #   = log(sum W_i g_i) = log E_W[g]  (no 1/N needed)
+        # log_sum for second stage = logsumexp(log_second_stage)
+        #   so mean = log_sum - log_n
+        log_ev_inc_resample = log_first_sum + log_sum - log_n
+        log_ev_inc_no_resample = log_sum
+        log_ev_inc = jnp.where(
+            do_resample, log_ev_inc_resample, log_ev_inc_no_resample
+        )
+
+        new_state = ParticleState(
+            particles=propagated,
+            log_weights=log_w_norm,
+            log_marginal_likelihood=(
+                state.log_marginal_likelihood + log_ev_inc
+            ),
+        )
+        ess_t: Array = jnp.asarray(compute_ess(log_w_norm))
+        if store_history:
+            return (new_state, ancestors), (
+                propagated,
+                log_w_norm,
+                ancestors,
+                ess_t,
+                log_ev_inc,
+            )
+        # Final-only mode (ADR-0011): ancestors ride the carry (O(N));
+        # the scan stacks just the scalar traces.
+        return (new_state, ancestors), (ess_t, log_ev_inc)
+
+    # Run the scan over t = 1 ... T-1
+    step_keys = jr.split(key, emissions.shape[0] - 1)
+    init_carry = (init_state, identity_ancestors)
+    if store_history:
+        (
+            (final_state, _),
+            (
+                particles_rest,
+                log_w_rest,
+                ancestors_rest,
+                ess_rest,
+                log_ev_inc_rest,
+            ),
+        ) = lax.scan(_step, init_carry, (step_keys, emissions[1:]))
+        all_particles = _prepend(particles_0, particles_rest)
+        all_log_w = _prepend(log_w_0, log_w_rest)
+        all_ancestors = _prepend(identity_ancestors, ancestors_rest)
+    else:
+        (
+            (final_state, final_ancestors),
+            (ess_rest, log_ev_inc_rest),
+        ) = lax.scan(_step, init_carry, (step_keys, emissions[1:]))
+        all_particles = final_state.particles[None]
+        all_log_w = final_state.log_weights[None]
+        all_ancestors = final_ancestors[None]
+    ess_0_arr: Array = jnp.asarray(ess_0)
+    all_ess = _prepend(ess_0_arr, ess_rest)
+    all_log_ev_inc = _prepend(jnp.asarray(log_ev_0), log_ev_inc_rest)
+
+    _raise_if_degenerate(final_state.log_marginal_likelihood)
+
+    return ParticleFilterPosterior(
+        marginal_loglik=final_state.log_marginal_likelihood,
+        filtered_particles=all_particles,
+        filtered_log_weights=all_log_w,
+        ancestors=all_ancestors,
+        ess=all_ess,
+        log_evidence_increments=all_log_ev_inc,
+    )