-
Notifications
You must be signed in to change notification settings - Fork 6
/
split_hmc.py
80 lines (51 loc) · 2.13 KB
/
split_hmc.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
'''
Created on Mar 7, 2011
@author: johnsalvatier
'''
from numpy import floor, concatenate, exp, zeros, diag, real
from numpy.linalg import inv, cholesky, eig
from utils import *
from ..core import *
def split_hmc_step(model, vars, C, approx_loc, approx_C, step_size_scaling = .25, trajectory_length = 2. ):
mapp = DASpaceMap( vars)
approx_loc = mapp.project(approx_loc)
n = C.shape[0]
logp_d_dict = model_logp_dlogp(model, vars)
step_size = step_size_scaling * n**(1/4.)
invC = inv(C)
cholInvC = cholesky(invC)
A = zeros((2*n,2*n))
A[:n,n:] = C
A[n:,:n] = -approx_C
D, gamma = eig(A)
e = step_size
R = real(gamma.dot(diag(exp(D * e))).dot(gamma.conj().T))
def step(logp_d, state, q0):
if state is None:
state = SamplerHist()
nstep = int(floor(trajectory_length / step_size))
q = q0
logp0, dlogp = logp_d(q)
dlogp = dlogp + dot(approx_C, q - approx_loc)
logp = logp0
# momentum scale proportional to inverse of parameter scale (basically sqrt(covariance))
p = p0 = cholesky_normal( q.shape, cholInvC)
#use the leapfrog method
for i in range(nstep):
#alternate full variable and momentum updates
p = p - (e/2) * -dlogp # half momentum update
x = concatenate((q - approx_loc, p))
x = dot(R, x)
q = x[:n] + approx_loc
p = x[n:]
logp, dlogp = logp_d(q)
dlogp = dlogp + dot(approx_C, q - approx_loc)
p = p - (e/2) * -dlogp # do a half step momentum update to finish off
p = -p
# - H(q*, p*) + H(q, p) = -H(q, p) + H(q0, p0) = -(- logp(q) + K(p)) + (-logp(q0) + K(p0))
mr = (-logp0) + K(C, p0) - ((-logp) + K(C, p))
state.metrops.append(mr)
return state, metrop_select(mr, q, q0)
return array_step(step, logp_d_dict, vars)
def K (cov, x):
return .5 * dot(x,dot(cov, x))