# Copyright 2020- The Blackjax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# """Public API for Late Adjusted Parallel Sampler (LAPS)"""
from typing import Any, NamedTuple
import jax
import jax.numpy as jnp
from blackjax.adaptation import laps_burn_in
from blackjax.adaptation.laps_burn_in import (
equipartition_diagonal,
equipartition_diagonal_loss,
)
from blackjax.adaptation.step_size import bisection_monotonic_fn
from blackjax.eca import run_eca
from blackjax.mcmc.adjusted_mclmc import build_kernel as build_kernel_mclmc
from blackjax.mcmc.hmc import HMCState
from blackjax.mcmc.integrators import (
generate_isokinetic_integrator,
mclachlan_coefficients,
omelyan_coefficients,
)
[docs]
class AdaptationState(NamedTuple):
[docs]
steps_per_sample: float
[docs]
stepsize_adaptation_state: Any # the state of the bisection algorithm to find a stepsize
[docs]
class Adaptation:
def __init__(
self,
adaptation_state,
num_adaptation_samples,
steps_per_sample=15,
acc_prob_target=0.8,
observables=lambda x: 0.0,
observables_for_bias=lambda x: 0.0,
contract=lambda x: 0.0,
):
[docs]
self.num_adaptation_samples = num_adaptation_samples
[docs]
self.observables = observables
[docs]
self.observables_for_bias = observables_for_bias
[docs]
self.contract = contract
"""
num_adaptation_samples: amount of tuning in the adjusted phase before fixing params
steps_per_sample: number of steps per sample
acc_prob_target: target acceptance probability
observables: function to compute observables, for diagnostics
observables_for_bias: function to compute observables for bias, for diagnostics
contract: function to contract observables, for diagnostics
"""
# Determine the initial hyperparameters #
step_size = adaptation_state.step_size
# Initialize the bisection for finding the step size
[docs]
self.epsadap_update = bisection_monotonic_fn(acc_prob_target)
stepsize_adaptation_state = (jnp.array([-jnp.inf, jnp.inf]), False)
[docs]
self.initial_state = AdaptationState(
steps_per_sample, step_size, stepsize_adaptation_state, 0
)
[docs]
def summary_statistics_fn(self, state, info, rng_key):
return {
"acceptance_probability": info.acceptance_rate,
"equipartition_diagonal": equipartition_diagonal(state),
"observables": self.observables(state.position),
"observables_for_bias": self.observables_for_bias(state.position),
}
[docs]
def update(self, adaptation_state, Etheta):
acc_prob = Etheta["acceptance_probability"]
equi_diag = equipartition_diagonal_loss(Etheta["equipartition_diagonal"])
true_bias = self.contract(Etheta["observables_for_bias"])
info_to_be_stored = {
"L": adaptation_state.step_size * adaptation_state.steps_per_sample,
"steps_per_sample": adaptation_state.steps_per_sample,
"step_size": adaptation_state.step_size,
"acc_prob": acc_prob,
"equi_diag": equi_diag,
"bias": true_bias,
"observables": Etheta["observables"],
}
# Bisection to find step size
stepsize_adaptation_state, step_size = self.epsadap_update(
adaptation_state.stepsize_adaptation_state,
adaptation_state.step_size,
acc_prob,
)
return (
AdaptationState(
adaptation_state.steps_per_sample,
step_size,
stepsize_adaptation_state,
adaptation_state.iteration + 1,
),
info_to_be_stored,
)
[docs]
def bias(model):
"""should be transfered to benchmarks/"""
def observables(position):
return jnp.square(model.transform(position))
def contract(sampler_E_x2):
bsq = jnp.square(sampler_E_x2 - model.E_x2) / model.Var_x2
return jnp.array([jnp.max(bsq), jnp.average(bsq)])
return observables, contract
[docs]
def while_steps_num(cond):
if jnp.all(cond):
return len(cond)
else:
return jnp.argmin(cond) + 1
[docs]
def laps(
logdensity_fn,
sample_init,
ndims,
num_steps1,
num_steps2,
num_chains,
mesh,
rng_key,
microcanonical=True,
alpha=1.9,
save_frac=0.2,
C=0.1,
early_stop=True,
r_end=0.01,
bias_type=3,
diagonal_preconditioning=True,
integrator_coefficients=None,
steps_per_sample=15,
acc_prob=None,
observables_for_bias=lambda x: x,
all_chains_info=None,
diagnostics=True,
contract=lambda x: 0.0,
superchain_size=1,
):
"""
model: the target density object
num_steps1: number of steps in the first phase
num_steps2: number of steps in the second phase
num_chains: number of chains
mesh: the mesh object, used for distributing the computation across cpus and nodes
rng_key: the random key
alpha: L = sqrt{d} * alpha * variances
save_frac: the fraction of samples used to estimate the fluctuation in the first phase
C: constant in stage 1 that determines step size (eq (9) of EMAUS paper)
early_stop: whether to stop the first phase early
r_end
diagonal_preconditioning: whether to use diagonal preconditioning
integrator_coefficients: the coefficients of the integrator
steps_per_sample: the number of steps per sample
acc_prob: the acceptance probability
observables: the observables (for diagnostic use)
all_chains_info: summary statistics calculated and stored for all chain at each iteration so it can be memory intensive
diagnostics: whether to return diagnostics
"""
key_init, key1, key2 = jax.random.split(rng_key, 3)
# initialize the chains
initial_state = laps_burn_in.initialize(
key_init,
logdensity_fn,
microcanonical,
sample_init,
num_chains,
mesh,
superchain_size,
)
# burn-in with the unadjusted method #
kernel = laps_burn_in.build_kernel(logdensity_fn, ndims, microcanonical)
save_num = (jnp.rint(save_frac * num_steps1)).astype(int)
adap = laps_burn_in.Adaptation(
ndims,
microcanonical=microcanonical,
alpha=alpha,
bias_type=bias_type,
save_num=save_num,
C=C,
r_end=r_end,
observables_for_bias=observables_for_bias,
contract=contract,
)
final_state, final_adaptation_state, info1 = run_eca(
key1,
initial_state,
kernel,
adap,
num_steps1,
num_chains,
mesh,
superchain_size,
all_chains_info,
early_stop=early_stop,
)
# refine the results with the adjusted method
_acc_prob = acc_prob
if integrator_coefficients is None:
high_dims = ndims > 200
_integrator_coefficients = (
omelyan_coefficients if high_dims else mclachlan_coefficients
)
if acc_prob is None:
_acc_prob = 0.9 if high_dims else 0.7
else:
_integrator_coefficients = integrator_coefficients
if acc_prob is None:
_acc_prob = 0.9
gradient_calls_per_step = (
len(_integrator_coefficients) // 2
) # The number of B updates in scheme = BABAB..AB is len(scheme)//2 + 1. The number of gradient calls is then len(scheme)//2, because the last B's gradient can be reused in the next step.
if diagonal_preconditioning:
inverse_mass_matrix = final_adaptation_state.inverse_mass_matrix
# scale the stepsize so that it reflects averag scale change of the preconditioning
average_scale_change = jnp.sqrt(jnp.average(inverse_mass_matrix))
final_adaptation_state = final_adaptation_state._replace(
step_size=final_adaptation_state.step_size / average_scale_change
)
else:
inverse_mass_matrix = 1.0
if microcanonical:
integrator = generate_isokinetic_integrator(_integrator_coefficients)
built_kernel = build_kernel_mclmc(
integrator=integrator,
)
kernel = lambda key, state, adap: built_kernel(
rng_key=key,
state=state,
logdensity_fn=logdensity_fn,
step_size=adap.step_size,
integration_steps_params=(adap.steps_per_sample,),
inverse_mass_matrix=inverse_mass_matrix,
L_proposal_factor=1.25,
)
else:
raise ValueError("Only microcanonical mode is supported for LAPS.")
initial_state = HMCState(
final_state.position, final_state.logdensity, final_state.logdensity_grad
)
num_samples = num_steps2 // (gradient_calls_per_step * steps_per_sample)
num_adaptation_samples = (
num_samples // 2
) # number of samples after which the stepsize is fixed.
final_adaptation_state = final_adaptation_state._replace(
step_size=final_adaptation_state.step_size.item()
)
adap = Adaptation(
final_adaptation_state,
num_adaptation_samples,
steps_per_sample,
_acc_prob,
contract=contract,
observables_for_bias=observables_for_bias,
)
final_state, final_adaptation_state, info2 = run_eca(
key2,
initial_state,
kernel,
adap,
num_samples,
num_chains,
mesh,
superchain_size,
all_chains_info,
)
if diagnostics:
info = {"phase_1": info1, "phase_2": info2}
else:
info = None
return info, gradient_calls_per_step, _acc_prob, final_state
