Skip to content

Quickstart

This example simulates a one-dimensional linear-Gaussian model, runs a bootstrap filter, and then compares it with a guided proposal. Run the code blocks in order in one Python session.

The model

We use an AR(1) latent state with a standard-normal initial state and Gaussian observations,

\[ z_t = \rho\,z_{t-1} + \sigma_q\,\varepsilon_t, \qquad y_t = z_t + \sigma_r\,\eta_t, \]

with \(\rho = 0.95\), \(\sigma_q = 0.3\), and \(\sigma_r = 0.7\), so the observation noise dominates the process noise and filtering has real work to do. A model in smcx is a handful of closures, written per particle: each takes one state array or PyTree and smcx vmaps it over the particle cloud internally. Samplers take an explicit PRNG key; the observation model returns a log-density with at least float32 precision. The weight, ESS, and resampling APIs likewise reject float16 and bfloat16 arrays because those formats cannot reliably preserve particle-weight invariants.

import math

import jax.numpy as jnp
import jax.random as jr
import numpy as np
import smcx

rho, q_sd, r_sd = 0.95, 0.3, 0.7
key_sim, key_filt = jr.split(jr.key(0))


def initial_sampler(key, n):
    return jr.normal(key, (n, 1))


def transition_sampler(key, state):
    return rho * state + q_sd * jr.normal(key, state.shape)


def emission_sampler(key, state):
    return state + r_sd * jr.normal(key, state.shape)


def log_observation_fn(y, state):
    z = (y[0] - state[0]) / r_sd
    return -0.5 * z * z - math.log(r_sd * math.sqrt(2 * math.pi))

The key threading is JAX's standard functional PRNG: no global seed, no hidden state, and any run is reproducible from its key. smcx splits the key it receives across time steps and particles for you.

Simulate

simulate runs the same transition/emission closures forward to produce a latent path and the observations we will filter. It draws a single trajectory, so its initial sampler takes just a key — a lambda adapts the filter-style sampler:

states, observations = smcx.simulate(
    key_sim,
    lambda key: initial_sampler(key, 1)[0],
    transition_sampler,
    emission_sampler,
    num_timesteps=100,
)

Establish the exact baseline

This model is linear and Gaussian with known parameters, so its filtering distribution is available exactly. kalman_filter returns the prior and filtered moments at every time step plus the exact marginal likelihood. The separately callable smoother consumes that result rather than rerunning the model:

transition = jnp.array([[rho]])
exact = smcx.kalman_filter(
    initial_mean=jnp.array([0.0]),
    initial_covariance=jnp.array([[1.0]]),
    transition_matrix=transition,
    transition_covariance=jnp.array([[q_sd**2]]),
    observation_matrix=jnp.array([[1.0]]),
    observation_covariance=jnp.array([[r_sd**2]]),
    emissions=observations,
)
smoothed = smcx.rts_smoother(exact, transition)

exact.filtered_means.shape
smoothed.smoothed_means.shape

The filter and smoother are independent pieces: a compatible GaussianFilterPosterior produced by research code can be passed to rts_smoother directly.

The linear filter permits finite symmetric positive-semidefinite initial and transition covariances, while its observation covariance must be positive definite. The smoother also requires each positive-time predicted covariance to be positive definite because the backward recursion factors it. Concrete covariance entries must be zero or normal finite values. For any factored covariance, the diagonal-equilibrated spectrum must be non-indefinite within dtype-scaled roundoff and the active backend must produce a finite Cholesky factor with a strictly positive diagonal. These checks run at eager entry; outer JAX transformations can lower endpoint arithmetic differently.

Run a particle filter

The bootstrap filter proposes from the transition and weights by the observation density.

posterior = smcx.bootstrap_filter(
    key_filt,
    initial_sampler,
    transition_sampler,
    log_observation_fn,
    observations,
    num_particles=10_000,
)

means = smcx.weighted_mean(posterior)
rmse = float(np.sqrt(np.mean((np.array(means) - np.array(states)) ** 2)))
print("marginal loglik:", round(posterior.marginal_loglik.item(), 1))
print("filter RMSE:", round(rmse, 3), "observation sd:", round(r_sd, 3))

Particle filters accumulate marginal_loglik with float32/float64 compensation while retaining the original per-step log_evidence_increments. On long float32 series, a plain jnp.sum(posterior.log_evidence_increments) can lose information that the compensated scalar preserves. This numerical correction can change fixed-key totals produced by older smcx releases.

For the fixed seed above, the filtered RMSE is about 0.369, compared with an observation-noise scale of 0.7.

Draw one-step predictions

Posterior prediction resamples each filtered particle cloud, advances it once, and samples the corresponding observation:

predictions = smcx.posterior_predictive_sample(
    jr.key(2),
    posterior,
    transition_sampler,
    emission_sampler,
    num_samples=1_000,
)

predictions.shape

The result has shape (100, 1_000, 1): one predictive distribution after each retained filtering row. Models with exogenous inputs or learned parameters use the explicit callback contracts described in Custom models.

Diagnose

diagnose returns a dictionary of health summaries and a list of plain-language warnings. It runs host-side and is not meant for the hot loop — call it after filtering.

report = smcx.diagnose(posterior)
print("min ESS:", round(report["min_ess"], 1))
print("max Pareto-k:", round(report["max_pareto_k"], 2))
for w in report["warnings"]:
    print("warning:", w)

For this run, no warnings are returned. The diagnostics describe the particle weights; they do not measure the Monte Carlo variance of the evidence estimate. If a per-step Pareto-k estimate is undefined, diagnose warns and computes max_pareto_k from the remaining finite estimates, or returns NaN when none are finite.

For model comparison, smcx.cumulative_log_score(posterior) reconstructs every predictive-score prefix with the same float32 Neumaier compensation used by compensated filter totals. This is a fixed-input numerical correction: prefixes can differ from earlier releases that used an ordinary cumulative sum, particularly for long or cancellation-heavy traces; the returned shape, dtype, and JIT contract are unchanged. Following NEP 23, smcx treats the change as a clear bug fix.

smcx.crps scores its prediction vector as an equally weighted empirical distribution. Its power-of-two-scaled empirical-CDF spacing calculation is translation-stable and nonnegative for finite float32 inputs whose represented spacings remain in the normal range, and it remains safe when the sample count squared would exceed the backend integer range. Scores affected by the earlier raw-value calculation can change at fixed inputs; following NEP 23, smcx treats this as the correction of a clear numerical bug. The signature and \(O(N \log N)\) complexity are unchanged. Backend-addressable forecasts must contain fewer than \(2^{31}\) samples; larger static shapes are rejected before materialization. A fixed-width uint32 accumulator assembles a correctly rounded result for subnormal or top-bin float32 inputs and whenever the conservative largest-error-over-\(N^2\) bound can enter the subnormal range. This includes mixed narrower observation dtypes, needs no backend uint64, and prevents represented translations from exposing flush-to-zero behavior.

Weighted mean and variance summaries likewise reduce in coordinates centered on a maximum-weight particle. Particles whose normalized linear weights are represented as zero are masked before centered products and powers, so finite outliers cannot poison summaries through overflow. Material coordinates whose centered difference would overflow are reduced after an exact one-bit binary downshift. The mean carries each centered subtraction's same-dtype compensation term; backend flush-to-zero can still erase subnormal terms. Variance retains square-first arithmetic unless a finite deviation square overflows. It then applies the weight before the second deviation factor to recover a representable variance. Affected fixed-input summaries can change as a clear numerical correction under NEP 23; their shapes, dtypes, and JIT contracts are unchanged.

Cut the variance with a guided proposal

The bootstrap proposal ignores the current observation. The guided filter proposes from the locally optimal density \(p(z_t \mid z_{t-1}, y_t)\), which for this linear-Gaussian model is available in closed form. The proposal precision is the sum of the process and observation precisions:

prop_var = 1.0 / (1.0 / q_sd**2 + 1.0 / r_sd**2)
prop_sd = math.sqrt(prop_var)


def proposal_sampler(key, state, y):
    mean = prop_var * (rho * state / q_sd**2 + y[0] / r_sd**2)
    return mean + prop_sd * jr.normal(key, state.shape)


def log_proposal_fn(y, new_state, old_state):
    mean = prop_var * (rho * old_state[0] / q_sd**2 + y[0] / r_sd**2)
    z = (new_state[0] - mean) / prop_sd
    return -0.5 * z * z - math.log(prop_sd * math.sqrt(2 * math.pi))


def log_transition_fn(new_state, old_state):
    z = (new_state[0] - rho * old_state[0]) / q_sd
    return -0.5 * z * z - math.log(q_sd * math.sqrt(2 * math.pi))


guided = smcx.guided_filter(
    key_filt,
    initial_sampler,
    proposal_sampler,
    log_proposal_fn,
    log_transition_fn,
    log_observation_fn,
    observations,
    num_particles=10_000,
)

g = smcx.diagnose(guided)
print("guided marginal loglik:", round(guided.marginal_loglik.item(), 1))
print("guided min ESS:", round(g["min_ess"], 1))

replicated_log_ml runs independent filters without retaining particle histories:

def boot_lml(key):
    return smcx.bootstrap_filter(
        key,
        initial_sampler,
        transition_sampler,
        log_observation_fn,
        observations,
        num_particles=10_000,
        store_history=False,
    ).marginal_loglik


def guided_lml(key):
    return smcx.guided_filter(
        key,
        initial_sampler,
        proposal_sampler,
        log_proposal_fn,
        log_transition_fn,
        log_observation_fn,
        observations,
        num_particles=10_000,
        store_history=False,
    ).marginal_loglik


lml_b = smcx.replicated_log_ml(jr.key(7), boot_lml, 20)
lml_g = smcx.replicated_log_ml(jr.key(7), guided_lml, 20)
print("bootstrap log-ML sd:", round(float(np.std(np.array(lml_b))), 3))
print("guided    log-ML sd:", round(float(np.std(np.array(lml_g))), 3))

With these seeds, the minimum ESS is about 1,122 for the bootstrap filter and 1,480 for the guided filter. Across the 20 replications above, the standard deviation of the log-evidence estimate is 0.074 and 0.053 respectively. These numbers describe this example, not a general performance guarantee.

What next

  • The stochastic volatility guide adds an unknown static parameter and learns it online.
  • The custom-model guide covers structured latent states and time-varying inputs.
  • bootstrap_init, bootstrap_step, and bootstrap_update support checkpointed or chunked filtering. Resumed checkpoints retain normalized log weights, their matching ESS, and finite evidence state. Concrete calls validate those invariants; a transformed bootstrap_step skips the data-dependent checks, while bootstrap_update is host-only.
  • Every function used here — including kalman_filter, rts_smoother, bootstrap_filter, guided_filter, simulate, and diagnose — has a full contract in the API reference.