blackjax.mcmc.adjusted_mclmc

blackjax.mcmc.adjusted_mclmc#

Public API for the Metropolis Hastings Microcanonical Hamiltonian Monte Carlo (MHMCHMC) Kernel. This is closely related to the Microcanonical Langevin Monte Carlo (MCLMC) Kernel, which is an unadjusted method. This kernel adds a Metropolis-Hastings correction to the MCLMC kernel. It also only refreshes the momentum variable after each MH step, rather than during the integration of the trajectory. Hence “Hamiltonian” and not “Langevin”.

NOTE: For best performance, we recommend using adjusted_mclmc_dynamic instead of this module, which is primarily intended for use in parallelized versions of the algorithm.

Functions#

init(→ blackjax.mcmc.hmc.HMCState)

Create an initial state for the MHMCHMC kernel.

build_kernel([integrator, divergence_threshold])

Build an MHMCHMC kernel.

as_top_level_api(→ blackjax.base.SamplingAlgorithm)

Implements the (basic) user interface for the MHMCHMC kernel.

Module Contents#

init(position: blackjax.types.ArrayLikeTree, logdensity_fn: Callable) → blackjax.mcmc.hmc.HMCState[source]#

Create an initial state for the MHMCHMC kernel.

Parameters:

  • position – Initial position of the chain.

  • logdensity_fn – Log-density function of the target distribution.

Return type:

The initial HMCState.

build_kernel(integrator: Callable = integrators.isokinetic_mclachlan, divergence_threshold: float = 1000)[source]#

Build an MHMCHMC kernel.

Parameters:

  • integrator – The symplectic integrator to use to integrate the Hamiltonian dynamics.

  • divergence_threshold – Value of the difference in energy above which we consider that the transition is divergent.

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, step_size: float, L_proposal_factor: float = jnp.inf, inverse_mass_matrix=1.0, *, divergence_threshold: int = 1000, integrator: Callable = integrators.isokinetic_mclachlan, num_integration_steps=None, integration_steps_params: tuple | None = None) → blackjax.base.SamplingAlgorithm[source]#

Implements the (basic) user interface for the MHMCHMC kernel.

Parameters:

  • logdensity_fn – The log-density function we wish to draw samples from.

  • step_size – The value to use for the step size in the symplectic integrator.

  • L_proposal_factor – Factor controlling partial momentum refreshment. jnp.inf disables refreshment (standard HMC-like behavior).

  • inverse_mass_matrix – Inverse mass matrix for the isokinetic integrator. Scalar or array.

  • divergence_threshold – The absolute value of the difference in energy between two states above which we say that the transition is divergent.

  • integrator – The symplectic integrator to use to integrate the trajectory.

  • num_integration_steps – Number of integration steps per transition. Deprecated in favour of integration_steps_params=(num_integration_steps,). Providing both raises a DeprecationWarning and integration_steps_params takes precedence.

  • integration_steps_params – Tuple of parameters unpacked into the kernel’s integration_steps_params argument. For the static kernel this must be a 1-tuple (num_steps,). Defaults to (num_integration_steps,) when only num_integration_steps is provided.

Return type:

A SamplingAlgorithm.