blackjax.mcmc.adjusted_mclmc_dynamic

blackjax.mcmc.adjusted_mclmc_dynamic#

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”.

Functions#

init(→ blackjax.mcmc.dynamic_hmc.DynamicHMCState)

Create an initial state for the dynamic MHMCHMC kernel.

build_kernel([integration_steps_fn, integrator, ...])

Build a Dynamic MHMCHMC kernel where the number of integration steps is chosen randomly.

as_top_level_api([1], integration_steps_fn, 1, 10), ...)

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

Module Contents#

init(position: blackjax.types.ArrayLikeTree, logdensity_fn: Callable, random_generator_arg: blackjax.types.Array) → blackjax.mcmc.dynamic_hmc.DynamicHMCState[source]#

Create an initial state for the dynamic MHMCHMC kernel.

Parameters:

  • position – Initial position of the chain.

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

  • random_generator_arg – Argument passed to integration_steps_fn and next_random_arg_fn to generate the number of integration steps.

Return type:

The initial DynamicHMCState.

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

Build a Dynamic MHMCHMC kernel where the number of integration steps is chosen randomly.

Parameters:

  • integration_steps_fn – Callable with signature (random_generator_arg, *integration_steps_params) -> int that draws the number of integration steps for a single transition. Extra positional arguments beyond random_generator_arg are supplied at call time via integration_steps_params on the inner kernel, so tunable parameters (e.g. average number of steps, distribution bounds) can be adapted without rebuilding the kernel.

  • integrator – The 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.

  • next_random_arg_fn – Function that generates the next random_generator_arg from its previous value.

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, next_random_arg_fn: Callable = lambda key: ..., integration_steps_fn: Callable = lambda key: ..., integration_steps_params: tuple = ()) → blackjax.base.SamplingAlgorithm[source]#

Implements the (basic) user interface for the dynamic 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.

  • divergence_threshold – The absolute value of the difference in energy between two states above which we say that the transition is divergent. The default value is commonly found in other libraries, and yet is arbitrary.

  • integrator – (algorithm parameter) The symplectic integrator to use to integrate the trajectory.

  • next_random_arg_fn – Function that generates the next random_generator_arg from its previous value.

  • integration_steps_fn – Callable with signature (random_generator_arg, *integration_steps_params) -> int that draws the number of integration steps for a single transition.

  • integration_steps_params – Extra positional arguments unpacked into integration_steps_fn after random_generator_arg on every step. Use this to pass tunable parameters (e.g. (avg_num_integration_steps,) or (lower_bound, upper_bound)) without rebuilding the kernel. Defaults to () so that a plain 1-arg integration_steps_fn works unchanged.

Return type:

A SamplingAlgorithm.