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([integrator, ...])

Build an MCLMC 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]#
nonans: bool[source]#
init(position: blackjax.types.ArrayLike, logdensity_fn, rng_key)[source]#
build_kernel(integrator: Callable = isokinetic_mclachlan, desired_energy_var_max_ratio: float = jnp.inf, desired_energy_var: float = 0.0005)[source]#

Build an MCLMC kernel.

The returned kernel accepts inverse_mass_matrix as either a scalar / 1-D array (diagonal preconditioning) or a LowRankInverseMassMatrix NamedTuple (Low-Rank + Diagonal preconditioning, O(dk) per step).

Parameters:

  • integrator – The isokinetic integrator to use. The default isokinetic_mclachlan() automatically dispatches to the O(dk) LRD path when inverse_mass_matrix is a LowRankInverseMassMatrix.

  • desired_energy_var_max_ratio – Maximum ratio of energy variance to desired energy variance before rejecting a transition.

  • desired_energy_var – The target energy variance per dimension.

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, inverse_mass_matrix: blackjax.types.ArrayLike | blackjax.mcmc.metrics.LowRankInverseMassMatrix = 1.0, desired_energy_var_max_ratio=jnp.inf) → 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.