blackjax.mcmc.mclmc

blackjax.mcmc.mclmc#

Public API for the MCLMC Kernel

Classes#

MCLMCInfo

Additional information on the MCLMC transition.

Functions#

`init`(position, logdensity_fn, rng_key)
`build_kernel`(logdensity_fn, sqrt_diag_cov, integrator)	Build a HMC kernel.
`as_top_level_api`(→ blackjax.base.SamplingAlgorithm)	The general mclmc kernel builder (`blackjax.mcmc.mclmc.build_kernel()`, alias blackjax.mclmc.build_kernel) can be

Module Contents#

class MCLMCInfo[source]#

Additional information on the MCLMC transition.

logdensity: The log-density of the distribution at the current step of the MCLMC chain.
kinetic_change: The difference in kinetic energy between the current and previous step.
energy_change: The difference in energy between the current and previous step.

logdensity: float[source]#

kinetic_change: float[source]#

energy_change: float[source]#

init(position: blackjax.types.ArrayLike, logdensity_fn, rng_key)[source]#

build_kernel(logdensity_fn, sqrt_diag_cov, integrator)[source]#

Build a HMC kernel.

Parameters:

integrator – The symplectic integrator to use to integrate the Hamiltonian dynamics.
L – the momentum decoherence rate.
step_size – step size of the integrator.

Returns:

A kernel that takes a rng_key and a Pytree that contains the current state
of the chain and that returns a new state of the chain along with
information about the transition.

as_top_level_api(logdensity_fn: Callable, L, step_size, integrator=isokinetic_mclachlan, sqrt_diag_cov=1.0) → blackjax.base.SamplingAlgorithm[source]#

The general mclmc kernel builder (blackjax.mcmc.mclmc.build_kernel(), alias blackjax.mclmc.build_kernel) can be cumbersome to manipulate. Since most users only need to specify the kernel parameters at initialization time, we provide a helper function that specializes the general kernel.

We also add the general kernel and state generator as an attribute to this class so users only need to pass blackjax.mclmc to SMC, adaptation, etc. algorithms.

Examples

A new mclmc kernel can be initialized and used with the following code:

mclmc = blackjax.mcmc.mclmc.mclmc(
    logdensity_fn=logdensity_fn,
    L=L,
    step_size=step_size
)
state = mclmc.init(position)
new_state, info = mclmc.step(rng_key, state)

Kernels are not jit-compiled by default so you will need to do it manually:

step = jax.jit(mclmc.step)
new_state, info = step(rng_key, state)

Parameters:

logdensity_fn – The log-density function we wish to draw samples from.
L – the momentum decoherence rate
step_size – step size of the integrator
integrator – an integrator. We recommend using the default here.

Return type:

A SamplingAlgorithm.