def run_parametrised_hmc(model_config,
interceptor,
num_samples=2000,
burnin=1000,
num_leapfrog_steps=4,
num_adaptation_steps=500,
num_optimization_steps=2000):
"""Given a (centred) model, this function transforms it based on the provided
interceptor, and runs HMC on the reparameterised model.
"""
def model_ncp(*params):
with ed.interception(interceptor):
return model_config.model(*params)
log_joint_noncentered = ed.make_log_joint_fn(model_ncp)
with ed.tape() as model_tape:
_ = model_ncp(*model_config.model_args)
param_shapes = collections.OrderedDict()
target_ncp_kwargs = {}
for param in model_tape.keys():
if param not in model_config.observed_data.keys():
param_shapes[param] = model_tape[param].shape
else:
target_ncp_kwargs[param] = model_config.observed_data[param]
def target_ncp(*param_args):
i = 0
for param in model_tape.keys():
if param not in model_config.observed_data.keys():
target_ncp_kwargs[param] = param_args[i]
i = i + 1
return log_joint_noncentered(*model_config.model_args,
**target_ncp_kwargs)
stepsize_kwargs = {'num_leapfrog_steps': num_leapfrog_steps}
stepsize_kwargs = {'num_optimization_steps': num_optimization_steps}
for key in model_config.observed_data:
stepsize_kwargs[key] = model_config.observed_data[key]
(step_size_init_ncp, stepsize_elbo_ncp,
vi_time) = util.approximate_mcmc_step_size(model_ncp,
*model_config.model_args,
**stepsize_kwargs)
results = _run_hmc(target_ncp,
param_shapes,
step_size_init=step_size_init_ncp,
transform=model_config.to_centered,
num_samples=num_samples,
burnin=burnin,
num_adaptation_steps=num_adaptation_steps,
num_leapfrog_steps=num_leapfrog_steps)
results['elbo'] = stepsize_elbo_ncp
results['vi_time'] = vi_time
return results
def run_centered_hmc(model_config,
num_samples=2000,
burnin=1000,
num_leapfrog_steps=4,
num_adaptation_steps=500,
num_optimization_steps=2000):
"""Runs HMC on the provided (centred) model."""
tf.reset_default_graph()
log_joint_centered = ed.make_log_joint_fn(model_config.model)
with ed.tape() as model_tape:
_ = model_config.model(*model_config.model_args)
param_shapes = collections.OrderedDict()
target_cp_kwargs = {}
for param in model_tape.keys():
if param not in model_config.observed_data.keys():
param_shapes[param] = model_tape[param].shape
else:
target_cp_kwargs[param] = model_config.observed_data[param]
def target_cp(*param_args):
i = 0
for param in model_tape.keys():
if param not in model_config.observed_data.keys():
target_cp_kwargs[param] = param_args[i]
i = i + 1
return log_joint_centered(*model_config.model_args, **target_cp_kwargs)
stepsize_kwargs = {'num_leapfrog_steps': num_leapfrog_steps}
stepsize_kwargs = {'num_optimization_steps': num_optimization_steps}
for key in model_config.observed_data:
stepsize_kwargs[key] = model_config.observed_data[key]
(step_size_init_cp, stepsize_elbo_cp,
vi_time) = util.approximate_mcmc_step_size(model_config.model,
*model_config.model_args,
**stepsize_kwargs)
results = _run_hmc(target_cp,
param_shapes,
step_size_init=step_size_init_cp,
num_samples=num_samples,
burnin=burnin,
num_adaptation_steps=num_adaptation_steps,
num_leapfrog_steps=num_leapfrog_steps)
results['elbo'] = stepsize_elbo_cp
results['vi_time'] = vi_time
return results