blackjax.smc.base#

Module Contents#

Classes#

SMCState

State of the SMC sampler

SMCInfo

Additional information on the tempered SMC step.

Functions#

init(particles)

step(→ Tuple[SMCState, SMCInfo])

General SMC sampling step.
class SMCState[source]#

State of the SMC sampler

particles: blackjax.types.ArrayTree[source]#
weights: blackjax.types.Array[source]#
class SMCInfo[source]#

Additional information on the tempered SMC step.

proposals: PyTree

The particles that were proposed by the MCMC pass.

ancestors: Array

The index of the particles proposed by the MCMC pass that were selected by the resampling step.

log_likelihood_increment: float

The log-likelihood increment due to the current step of the SMC algorithm.

ancestors: blackjax.types.Array[source]#
log_likelihood_increment: float[source]#
update_info: NamedTuple[source]#
init(particles: blackjax.types.ArrayLikeTree)[source]#
step(rng_key: blackjax.types.PRNGKey, state: SMCState, update_fn: Callable, weigh_fn: Callable, resample_fn: Callable, num_resampled: Optional[int] = None) → Tuple[SMCState, SMCInfo][source]#

General SMC sampling step.

update_fn here corresponds to the Markov kernel \(M_{t+1}\), and weigh_fn corresponds to the potential function \(G_t\). We first use update_fn to generate new particles from the current ones, weigh these particles using weigh_fn and resample them with resample_fn.

The update_fn and weigh_fn functions must be batched by the called either using jax.vmap or jax.pmap.

In Feynman-Kac terms, the algorithm goes roughly as follows:

M_t: update_fn
G_t: weigh_fn
R_t: resample_fn
idx = R_t(weights)
x_t = x_tm1[idx]
x_{t+1} = M_t(x_t)
weights = G_t(x_{t+1})
Parameters:

  • rng_key – Key used to generate pseudo-random numbers.

  • state – Current state of the SMC sampler: particles and their respective log-weights

  • update_fn – Function that takes an array of keys and particles and returns new particles.

  • weigh_fn – Function that assigns a weight to the particles.

  • resample_fn – Function that resamples the particles.

  • num_resampled – The number of particles to resample. This can be used to implement Waste-Free SMC [DC20], in which case we resample a number \(M<N\) of particles, and the update function is in charge of returning \(N\) samples.

Returns:

  • new_particles – An array that contains the new particles generated by this SMC step.

  • info – An SMCInfo object that contains extra information about the SMC transition.