def test_leapfrog_reversible():
n = 3
start, model, _ = models.non_normal(n)
step = BaseHMC(vars=model.vars, model=model)
bij = DictToArrayBijection(step.ordering, start)
q0 = bij.map(start)
p0 = floatX(np.ones(n) * .05)
precision = select_by_precision(float64=1E-8, float32=1E-4)
for epsilon in [.01, .1, 1.2]:
for n_steps in [1, 2, 3, 4, 20]:
q, p = q0, p0
q, p, _ = step.leapfrog(q, p, floatX(np.array(epsilon)), np.array(n_steps, dtype='int32'))
q, p, _ = step.leapfrog(q, -p, floatX(np.array(epsilon)), np.array(n_steps, dtype='int32'))
close_to(q, q0, precision, str((n_steps, epsilon)))
close_to(-p, p0, precision, str((n_steps, epsilon)))
def test_leapfrog_reversible_single():
n = 3
start, model, _ = models.non_normal(n)
integrators = ['leapfrog', 'two-stage', 'three-stage']
steps = [BaseHMC(vars=model.vars, model=model, integrator=method, use_single_leapfrog=True)
for method in integrators]
for method, step in zip(integrators, steps):
bij = DictToArrayBijection(step.ordering, start)
q0 = bij.map(start)
p0 = floatX(np.ones(n) * .05)
precision = select_by_precision(float64=1E-8, float32=1E-5)
for epsilon in [0.01, 0.1, 1.2]:
for n_steps in [1, 2, 3, 4, 20]:
dlogp0 = step.dlogp(q0)
q, p = q0, p0
dlogp = dlogp0
energy = step.compute_energy(q, p)
for _ in range(n_steps):
q, p, v, dlogp, _ = step.leapfrog(q, p, dlogp, floatX(np.array(epsilon)))
p = -p
for _ in range(n_steps):
q, p, v, dlogp, _ = step.leapfrog(q, p, dlogp, floatX(np.array(epsilon)))
close_to(q, q0, precision, str(('q', method, n_steps, epsilon)))
close_to(-p, p0, precision, str(('p', method, n_steps, epsilon)))
def test_leapfrog_reversible():
n = 3
start, model, _ = models.non_normal(n)
step = BaseHMC(vars=model.vars, model=model)
bij = DictToArrayBijection(step.ordering, start)
q0 = bij.map(start)
p0 = np.ones(n) * .05
for epsilon in [.01, .1, 1.2]:
for n_steps in [1, 2, 3, 4, 20]:
q, p = q0, p0
q, p, _ = step.leapfrog(q, p, np.array(epsilon), np.array(n_steps, dtype='int32'))
q, p, _ = step.leapfrog(q, -p, np.array(epsilon), np.array(n_steps, dtype='int32'))
close_to(q, q0, 1e-8, str((n_steps, epsilon)))
close_to(-p, p0, 1e-8, str((n_steps, epsilon)))
def test_leapfrog_reversible():
n = 3
np.random.seed(42)
start, model, _ = models.non_normal(n)
size = model.ndim
scaling = floatX(np.random.rand(size))
step = BaseHMC(vars=model.vars, model=model, scaling=scaling)
step.integrator._logp_dlogp_func.set_extra_values({})
p = floatX(step.potential.random())
q = floatX(np.random.randn(size))
start = step.integrator.compute_state(p, q)
for epsilon in [0.01, 0.1]:
for n_steps in [1, 2, 3, 4, 20]:
state = start
for _ in range(n_steps):
state = step.integrator.step(epsilon, state)
for _ in range(n_steps):
state = step.integrator.step(-epsilon, state)
npt.assert_allclose(state.q, start.q, rtol=1e-5)
npt.assert_allclose(state.p, start.p, rtol=1e-5)