Ctrl+K

blackjax.adaptation.mclmc_adaptation

blackjax.adaptation.mclmc_adaptation#

Algorithms to adapt the MCLMC kernel parameters, namely step size and L.

Module Contents#

Classes#

MCLMCAdaptationState

Represents the tunable parameters for MCLMC adaptation.

Functions#

mclmc_find_L_and_step_size(mclmc_kernel, num_steps, ...)

Finds the optimal value of the parameters for the MCLMC algorithm.

make_L_step_size_adaptation(kernel, dim, frac_tune1, ...)

Adapts the stepsize and L of the MCLMC kernel. Designed for the unadjusted MCLMC

make_adaptation_L(kernel, frac, Lfactor)

determine L by the autocorrelations (around 10 effective samples are needed for this to be accurate)

handle_nans(previous_state, next_state, step_size, ...)

if there are nans, let's reduce the stepsize, and not update the state. The
class MCLMCAdaptationState[source]#

Represents the tunable parameters for MCLMC adaptation.

L

The momentum decoherent rate for the MCLMC algorithm.

step_size

The step size used for the MCLMC algorithm.

L: float[source]#
step_size: float[source]#
mclmc_find_L_and_step_size(mclmc_kernel, num_steps, state, rng_key, frac_tune1=0.1, frac_tune2=0.1, frac_tune3=0.1, desired_energy_var=0.0005, trust_in_estimate=1.5, num_effective_samples=150)[source]#

Finds the optimal value of the parameters for the MCLMC algorithm.

Parameters:

  • mclmc_kernel – The kernel function used for the MCMC algorithm.

  • num_steps – The number of MCMC steps that will subsequently be run, after tuning.

  • state – The initial state of the MCMC algorithm.

  • rng_key – The random number generator key.

  • frac_tune1 – The fraction of tuning for the first step of the adaptation.

  • frac_tune2 – The fraction of tuning for the second step of the adaptation.

  • frac_tune3 – The fraction of tuning for the third step of the adaptation.

  • desired_energy_va – The desired energy variance for the MCMC algorithm.

  • trust_in_estimate – The trust in the estimate of optimal stepsize.

  • num_effective_samples – The number of effective samples for the MCMC algorithm.

Return type:

A tuple containing the final state of the MCMC algorithm and the final hyperparameters.

Examples

# Define the kernel function
def kernel(x):
    return x ** 2

# Define the initial state
initial_state = MCMCState(position=0, momentum=1)

# Generate a random number generator key
rng_key = jax.random.key(0)

# Find the optimal parameters for the MCLMC algorithm
final_state, final_params = mclmc_find_L_and_step_size(
    mclmc_kernel=kernel,
    num_steps=1000,
    state=initial_state,
    rng_key=rng_key,
    frac_tune1=0.2,
    frac_tune2=0.3,
    frac_tune3=0.1,
    desired_energy_var=1e-4,
    trust_in_estimate=2.0,
    num_effective_samples=200,
)
make_L_step_size_adaptation(kernel, dim, frac_tune1, frac_tune2, desired_energy_var=0.001, trust_in_estimate=1.5, num_effective_samples=150)[source]#

Adapts the stepsize and L of the MCLMC kernel. Designed for the unadjusted MCLMC

make_adaptation_L(kernel, frac, Lfactor)[source]#

determine L by the autocorrelations (around 10 effective samples are needed for this to be accurate)

handle_nans(previous_state, next_state, step_size, step_size_max, kinetic_change)[source]#

if there are nans, let’s reduce the stepsize, and not update the state. The function returns the old state in this case.

Contents