-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathvariational_GMM.py
189 lines (127 loc) · 5.42 KB
/
variational_GMM.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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 4 14:16:07 2016
@author: hrs13
"""
import numpy as np
import scipy.stats as stats
from scipy.special import psi
from scipy.special import gamma, gammaln
from utils import log_B
W_0 = np.eye(2) # 2D, but would work just as well for higher dimensions
v_0 = 3. # this number being higher (min is 2) makes the clusters more irregular in shape
m_0 = np.zeros(2) # clusters to be around the origin
b_0 = 0.5 # larger makes the clusters closer to the origin
a_0 = 0.4 # larger makes the mixing coefficients more similar
def variational_E_step(X, a_k, b_k, m_k, W_k, v_k):
N, D = X.shape
K = len(a_k)
r = np.empty((N, K))
a_hat = np.sum(a_k)
E_log_det = np.zeros(K)
E_log_pis = np.zeros(K)
for k in range(K):
E_log_det[k] = psi(v_k[k]/2) + psi((v_k[k]+1)/2)
E_log_det[k] += 2 * np.log(2) + np.log(np.linalg.det(W_k[k, :, :]))
E_log_pis[k] = psi(a_k[k]) - psi(a_hat)
for n in range(N): # vectorise this later
x = np.reshape(X[n, :] - m_k[k, :], (2, 1))
E_exponent = 2/b_k[k] + v_k[k] * x.T.dot(W_k[k, :, :]).dot(x)
r[n, k] = np.exp(E_log_pis[k]) * np.exp(0.5*E_log_det[k]) * np.exp(-0.5*E_exponent)
return (r.T/np.sum(r, axis=1)).T
def variational_LB(X, r, a_k, b_k, m_k, W_k, v_k):
N, D = X.shape
K = len(a_k)
N_k = np.sum(r, axis=0)
x_k = ((r.T.dot(X)).T/N_k).T
X_c = np.reshape([X[n, :] - x_k[k, :] for n in range(N) for k in range(K)], (N, K, 2))
S_k = np.einsum('nk,nkd,nke, k->kde', r, X_c, X_c, 1/N_k)
a_hat = np.sum(a_k)
E_log_det = np.zeros(K)
E_log_pis = np.zeros(K)
L = 0.
for k in range(K):
E_log_det[k] = psi(v_k[k]/2) + psi((v_k[k]+1)/2)
E_log_det[k] += 2 * np.log(2) + np.log(np.linalg.det(W_k[k, :, :]))
E_log_pis[k] = psi(a_k[k]) - psi(a_hat)
x = np.reshape(x_k[k, :] - m_k[k, :], (2, 1))
L += 0.5*N_k[k]*(E_log_det[k] - 2/b_k[k] - v_k[k]*np.sum(S_k[k, :, :] * W_k[k, : :]) - v_k[k]*x.T.dot(W_k[k, :, :]).dot(x) - 2*np.log(2))
m = np.reshape(m_k[k, :] - m_0, (2, 1))
L += 0.5*(2*np.log(b_0/(2*np.pi)) + E_log_det[k] - 2*b_0/b_k[k] - b_0*v_k[k]*m.T.dot(W_k[k, :, :]).dot(m))
L -= 0.5*v_k[k]*np.sum(np.linalg.inv(W_0)*W_k[k, :, :])
L -= (a_k[k] - 1)*E_log_pis[k] - gammaln(a_k[k])
L -= 0.5*E_log_det[k] + np.log(b_k[k]) - np.log(2*np.pi) - 1. - stats.wishart.entropy(v_k[k], W_k[k])
L -= gammaln(a_hat)
L += K*log_B(W_0, v_0)
L += ((v_0-2.)/2)*np.sum(E_log_det)
L -= np.sum(r*np.log(r))
L += np.sum(r.dot(np.reshape(E_log_pis, (K, 1))))
L += gammaln(K*a_0) - K*gammaln(a_0) + (a_0 - 1.)*np.sum(E_log_pis)
return float(L)
def variational_M_step(X, r):
N, K = r.shape
N_k = np.sum(r, axis=0)
x_k = ((r.T.dot(X)).T/N_k).T
X_c = np.reshape([X[n, :] - x_k[k, :] for n in range(N) for k in range(K)], (N, K, 2))
S_k = np.einsum('nk,nkd,nke, k->kde', r, X_c, X_c, 1/N_k)
a_k = np.empty(K)
b_k = np.empty(K)
m_k = np.empty((K, 2))
W_k = np.empty((K, 2, 2))
v_k = np.empty(K)
for k in range(K):
a_k[k] = a_0 + N_k[k]
b_k[k] = b_0 + N_k[k]
m_k[k, :] = (b_0 * m_0 + N_k[k] * x_k[k, :])/b_k[k]
temp = np.linalg.inv(W_0) + N_k[k]*S_k[k, :, :]
x = np.reshape(x_k[k, :] - m_0, (2, 1))
temp += (b_0*N_k[k]/(b_0 + N_k[k]))*x.dot(x.T)
W_k[k, :, :] = np.linalg.inv(temp)
v_k[k] = v_0 + N_k[k] +1
return a_k, b_k, m_k, W_k, v_k
from utils import generate_parameters, generate_data
from plotting import double_panel_demo
from utils import generate_rand_samples
if __name__ == '__main__':
K = 3
N = 100
num_its = 50
X = generate_data(N, generate_parameters(K))[0]
plt = double_panel_demo(K)
while True:
X = generate_data(N, generate_parameters(K))[0]
plt.set_new_lims(X, num_its)
params = generate_parameters(K)
# these initial parameters are an independent draw from the prior
objective = []
plt.cla('ax1')
plt.cla('ax2')
plt.plot_points_black(X)
plt.draw()
plt.pause(2.)
rand = np.reshape(np.random.rand(N*K), (N, K))
r = (rand.T/np.sum(rand, axis=1)).T
for i in range(num_its):
a_k, b_k, m_k, W_k, v_k = variational_M_step(X, r)
r = variational_E_step(X, a_k, b_k, m_k, W_k, v_k)
objective.append(variational_LB(X, r, a_k, b_k, m_k, W_k, v_k))
plt.plot_GMM_objective(objective)
for i in range(10):
params = generate_rand_samples(a_k, b_k, m_k, W_k, v_k)
plt.cla('ax1')
plt.plot_parameters(params)
plt.plot_data_coloured(X, r)
plt.pause(1/60.)
#for i in range(10):
# r = r_variational_GMM(X, a_k, b_k, m_k, W_k, v_k)
# print LB_variational_GMM(X, r, a_k, b_k, m_k, W_k, v_k)
#
# for j in range(10):
# plt.cla()
# params = generate_samples(a_k, b_k, m_k, W_k, v_k)
# plot_parameters(params)
# plot_data_coloured(X, r)
# show_plot(X)
# plt.pause(0.5)
#
#