gmm.py 12.3 KB
Newer Older
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
1
import numpy as np
2
3
4
from model import *
from functions import multi_variate_normal
from scipy.linalg import block_diag
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
5
import scipy.sparse as ss
6
from termcolor import colored
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
7
from mvn import MVN, SparseMVN
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
8
9
10


class GMM(Model):
11
12
13
14
	def __init__(self, nb_states=1, nb_dim=None, init_zeros=False):
		Model.__init__(self, nb_states, nb_dim)
		# flag to indicate that publishing was not init
		self.publish_init = False
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
15

16
17
		if init_zeros:
			self.init_zeros()
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
18

19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
	def get_matching_mvn(self, max=False, mass=None):
		if max:
			priors = (self.priors == np.max(self.priors)).astype(np.float32)
			priors /= np.sum(priors)
		elif mass is not None:
			prior_lim = np.sort(self.priors)[::-1][np.max(
				[0, np.argmin(np.cumsum(np.sort(self.priors)[::-1]) < mass)])]

			priors = (self.priors >= prior_lim) * self.priors
			priors /= np.sum(priors)
		else:
			priors = self.priors
		# print priors, self.priors

		mus, sigmas = self.moment_matching(priors)
		mvn = MVN(nb_dim=self.nb_dim, mu=mus[0], sigma=sigmas[0])

		return mvn
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
37

38
39
40
41
42
43
44
45
46
	def moment_matching(self, h):
		"""
		Perform moment matching to approximate a mixture of Gaussian as a Gaussian
		:param h: 		np.array([nb_timesteps, nb_states])
			Activations of each states for different timesteps
		:return:
		"""
		if h.ndim == 1:
			h = h[None]
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
47

48
49
50
51
		mus = np.einsum('ak,ki->ai', h, self.mu)
		dmus = self.mu[None] - mus[:, None] # nb_timesteps, nb_states, nb_dim
		sigmas = np.einsum('ak,kij->aij', h, self.sigma) + \
				 np.einsum('ak,akij->aij',h , np.einsum('aki,akj->akij', dmus, dmus))
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
52

53
		return mus, sigmas
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
54

55
56
57
58
59
60
61
62
63
64
65
66
	def __add__(self, other):
		if isinstance(other, MVN):
			gmm = GMM(nb_dim=self.nb_dim, nb_states=self.nb_states)

			gmm.priors = self.priors
			gmm.mu = self.mu + other.mu[None]
			gmm.sigma = self.sigma + other.sigma[None]

			return gmm

		else:
			raise NotImplementedError
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
67

68
69
70
71
72
73
74
	def __mul__(self, other):
		"""
		Renormalized product of Gaussians, component by component

		:param other:
		:return:
		"""
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
		if isinstance(other, MVN):
			gmm = GMM(nb_dim=self.nb_dim, nb_states=self.nb_states)
			gmm.mu = np.einsum('aij,aj->ai', self.lmbda, self.mu) + \
					 np.einsum('ij,j->i', other.lmbda, other.mu)[None]

			gmm.lmbda = self.lmbda + other.lmbda[None]
			gmm.mu = np.einsum('aij,aj->ai', gmm.sigma, gmm.mu)

			Z = np.linalg.slogdet(self.lmbda)[1]\
				+ np.linalg.slogdet(other.lmbda)[1] \
				- 0.5 * np.linalg.slogdet(gmm.lmbda)[1] \
				- self.nb_dim / 2. * np.log(2 * np.pi) \
				+ 0.5 * (np.einsum('ai,aj->a',
								   np.einsum('ai,aij->aj', gmm.mu, gmm.lmbda), gmm.mu)
						-np.einsum('ai,aj->a',
								   np.einsum('ai,aij->aj', self.mu, self.lmbda), self.mu)
						-np.sum(np.einsum('i,ij->j', other.mu, other.lmbda) * other.mu)
						 )
			gmm.priors = np.exp(Z) * self.priors
			gmm.priors /= np.sum(gmm.priors)

		else:
			gmm = GMM(nb_dim=self.nb_dim, nb_states=self.nb_states)
			gmm.priors = self.priors
			gmm.mu = np.einsum('aij,aj->ai', self.lmbda, self.mu) + \
					 np.einsum('aij,aj->ai', other.lmbda, other.mu)

			gmm.lmbda = self.lmbda + other.lmbda

			gmm.mu = np.einsum('aij,aj->ai', gmm.sigma, gmm.mu)
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

		return gmm

	def marginal_model(self, dims):
		"""
		Get a GMM of a slice of this GMM
		:param dims:
		:type dims: slice
		:return:
		"""
		gmm = GMM(nb_dim=dims.stop-dims.start, nb_states=self.nb_states)
		gmm.priors = self.priors
		gmm.mu = self.mu[:, dims]
		gmm.sigma = self.sigma[:, dims, dims]

		return gmm

	def lintrans(self, A, b):
		"""
		Linear transformation of a GMM

		:param A:		[np.array(nb_dim, nb_dim)]
		:param b: 		[np.array(nb_dim)]
		:return:
		"""

		gmm = GMM(nb_dim=self.nb_dim, nb_states=self.nb_states)
		gmm.priors = self.priors
		gmm.mu = np.einsum('ij,aj->ai', A, self.mu) + b
		gmm.lmbda = np.einsum('aij,kj->aik',
								 np.einsum('ij,ajk->aik', A, self.lmbda), A)

		return gmm

Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
139
	def concatenate_gaussian(self, q, get_mvn=True, sparse=False):
140
141
142
143
144
145
146
147
148
149
150
151
		"""
		Get a concatenated-block-diagonal replication of the GMM with sequence of state
		given by q.

		:param q: 			[list of int]
		:param get_mvn: 	[bool]


		:return:
		"""
		if not get_mvn:
			return np.concatenate([self.mu[i] for i in q]), block_diag(*[self.sigma[i] for i in q])
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
152
		else:
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
153
154
155
156
157
158
159
160
161
162
163
			if not sparse:
				mvn = MVN()
				mvn.mu = np.concatenate([self.mu[i] for i in q])
				mvn._sigma = block_diag(*[self.sigma[i] for i in q])
				mvn._lmbda = block_diag(*[self.lmbda[i] for i in q])
			else:
				mvn = SparseMVN()

				mvn.mu = np.concatenate([self.mu[i] for i in q])
				mvn._sigma = ss.block_diag([self.sigma[i] for i in q])
				mvn._lmbda = ss.block_diag([self.lmbda[i] for i in q])
164
			return mvn
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
165

166
167
168
169
	def compute_resp(self, demo=None, dep=None, table=None, marginal=None, ):
		sample_size = demo.shape[0]

		B = np.ones((self.nb_states, sample_size))
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
170

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
		if marginal != []:
			for i in range(self.nb_states):
				mu, sigma = (self.mu, self.sigma)

				if marginal is not None:
					mu, sigma = self.get_marginal(marginal)

				if dep is None:
					B[i, :] = multi_variate_normal(demo,
												   mu[i],
												   sigma[i], log=False)
				else:  # block diagonal computation
					B[i, :] = 1.0
					for d in dep:
						dGrid = np.ix_([i], d, d)
						B[[i], :] *= multi_variate_normal(demo, mu[i, d],
														  sigma[dGrid][:, :, 0], log=False)
		B *= self.priors[:, None]
		return B/np.sum(B, axis=0)

	def init_params_scikit(self, data):
		from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
		gmm_init = GaussianMixture(self.nb_states, 'full', n_init=10, init_params='random')
		gmm_init.fit(data)

		self.mu = gmm_init.means_
		self.sigma = gmm_init.covariances_
		self.priors = gmm_init.weights_

		self.Trans = np.ones((self.nb_states, self.nb_states)) * 0.01

		self.init_priors = np.ones(self.nb_states) * 1. / self.nb_states

	def init_params_kmeans(self, data):
		from sklearn.cluster import KMeans
		km_init = KMeans(n_clusters=self.nb_states)
		km_init.fit(data)
		self.mu = km_init.cluster_centers_
		self.priors = np.ones(self.nb_states)/ self.nb_states
		self.sigma = np.array([np.eye(self.nb_dim) for i in range(self.nb_states)])
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
211

212
		self.Trans = np.ones((self.nb_states, self.nb_states)) * 0.01
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
213

214
		self.init_priors = np.ones(self.nb_states) * 1. / self.nb_states
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
215
216


217
218
219
220
	def init_params_random(self, data):
		mu = np.mean(data, axis=0)
		sigma = np.dot((data - mu).T, (data - mu)) / \
				(data.shape[0] - 1)
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
221

222
223
		self.mu = np.array([np.random.multivariate_normal(mu, sigma)
			 for i in range(self.nb_states)])
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
224

225
		self.sigma = np.array([sigma + self.reg for i in range(self.nb_states)])
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
226

227
		self.priors = np.ones(self.nb_states) / self.nb_states
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
228

229
	def em(self, data, reg=1e-8, maxiter=100, minstepsize=1e-5, diag=False, reg_finish=False,
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
230
		   kmeans_init=False, random_init=True, dep_mask=None, verbose=False, only_scikit=False):
231
		"""
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
232

233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
		:param data:	 		[np.array([nb_timesteps, nb_dim])]
		:param reg:				[list([nb_dim]) or float]
			Regulariazation for EM
		:param maxiter:
		:param minstepsize:
		:param diag:			[bool]
			Use diagonal covariance matrices
		:param reg_finish:		[np.array([nb_dim]) or float]
			Regulariazation for finish step
		:param kmeans_init:		[bool]
			Init components with k-means.
		:param random_init:		[bool]
			Init components randomely.
		:param dep_mask: 		[np.array([nb_dim, nb_dim])]
			Composed of 0 and 1. Mask given the dependencies in the covariance matrices
		:return:
		"""
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
250

251
		self.reg = reg
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
252
253
254
255
256
257
258
259

		nb_min_steps = 5  # min num iterations
		nb_max_steps = maxiter  # max iterations
		max_diff_ll = minstepsize # max log-likelihood increase

		nb_samples = data.shape[0]


260
261
262
263
264
265
266

		if random_init:
			self.init_params_random(data)
		elif kmeans_init:
			self.init_params_kmeans(data)
		else:
			self.init_params_scikit(data)
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
267

Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
268
		if only_scikit: return
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
269
270
271
272
273
274
275
276
277
278
		data = data.T

		LL = np.zeros(nb_max_steps)
		for it in range(nb_max_steps):

			# E - step
			L = np.zeros((self.nb_states, nb_samples))
			L_log = np.zeros((self.nb_states, nb_samples))

			for i in range(self.nb_states):
279
280
				L_log[i, :] = np.log(self.priors[i]) + multi_variate_normal(data.T, self.mu[i],
												   self.sigma[i], log=True)
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
281
282
283
284
285
286

			L = np.exp(L_log)
			GAMMA = L / np.sum(L, axis=0)
			GAMMA2 = GAMMA / np.sum(GAMMA, axis=1)[:, np.newaxis]

			# M-step
287
			self.mu = np.einsum('ac,ic->ai', GAMMA2,
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
288
289
									data)  # a states, c sample, i dim

290
			dx = data[None, :] - self.mu[:, :, None]  # nb_dim, nb_states, nb_samples
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
291

292
			self.sigma = np.einsum('acj,aic->aij', np.einsum('aic,ac->aci', dx, GAMMA2),
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
293
294
									   dx)  # a states, c sample, i-j dim

295
296
297
298
299
300
301
			self.sigma += self.reg

			if diag:
				self.sigma *= np.eye(self.nb_dim)

			if dep_mask is not None:
				self.sigma *= dep_mask
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
302

303
			# print self.Sigma[:,u :, i]
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
304
305

			# Update initial state probablility vector
306
			self.priors = np.mean(GAMMA, axis=1)
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
307
308


309
			LL[it] = np.mean(np.log(np.sum(L, axis=0)))
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
310
311
312
			# Check for convergence
			if it > nb_min_steps:
				if LL[it] - LL[it - 1] < max_diff_ll:
313
314
315
316
					if reg_finish is not False:
						self.sigma = np.einsum(
							'acj,aic->aij', np.einsum('aic,ac->aci', dx, GAMMA2), dx) + reg_finish

317
318
					if verbose:
						print colored('Converged after %d iterations: %.3e' % (it, LL[it]), 'red', 'on_white')
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
319
					return GAMMA
320
321
		if verbose:
			print "GMM did not converge before reaching max iteration. Consider augmenting the number of max iterations."
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
322
323
324
		return GAMMA


325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
	def init_hmm_kbins(self, demos, dep=None, reg=1e-8, dep_mask=None):
		"""
		Init HMM by splitting each demos in K bins along time. Each K states of the HMM will
		be initialized with one of the bin. It corresponds to a left-to-right HMM.

		:param demos:	[list of np.array([nb_timestep, nb_dim])]
		:param dep:
		:param reg:		[float]
		:return:
		"""

		# delimit the cluster bins for first demonstration
		self.nb_dim = demos[0].shape[1]

		self.init_zeros()
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
340

341
		t_sep = []
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
342

343
344
345
346
		for demo in demos:
			t_sep += [map(int, np.round(np.linspace(0, demo.shape[0], self.nb_states + 1)))]

		# print t_sep
Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
347
		for i in range(self.nb_states):
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
			data_tmp = np.empty((0, self.nb_dim))
			inds = []
			states_nb_data = 0   # number of datapoints assigned to state i

			# Get bins indices for each demonstration
			for n, demo in enumerate(demos):
				inds = range(t_sep[n][i], t_sep[n][i+1])

				data_tmp = np.concatenate([data_tmp, demo[inds]], axis=0)
				states_nb_data += t_sep[n][i+1]-t_sep[n][i]

			self.priors[i] = states_nb_data
			self.mu[i] = np.mean(data_tmp, axis=0)

			if dep_mask is not None:
				self.sigma *= dep_mask

			if dep is None:
				self.sigma[i] = np.cov(data_tmp.T) + np.eye(self.nb_dim) * reg
			else:
				for d in dep:
					dGrid = np.ix_([i], d, d)
					self.sigma[dGrid] = (np.cov(data_tmp[:, d].T) + np.eye(
						len(d)) * reg)[:, :, np.newaxis]
				# print self.Sigma[:,:,i]

		# normalize priors
		self.priors = self.priors / np.sum(self.priors)

		# Hmm specific init
		self.Trans = np.ones((self.nb_states, self.nb_states)) * 0.01

		nb_data = np.mean([d.shape[0] for d in demos])

		for i in range(self.nb_states - 1):
			self.Trans[i, i] = 1.0 - float(self.nb_states) / nb_data
			self.Trans[i, i + 1] = float(self.nb_states) / nb_data

		self.Trans[-1, -1] = 1.0
		self.init_priors = np.ones(self.nb_states) * 1. / self.nb_states

Emmanuel PIGNAT's avatar
Emmanuel PIGNAT committed
389
390
391
392
393
394
395
396
397
398
399
400
401
402
	def add_trash_component(self, data, scale=2.):
		if isinstance(data, list):
			data = np.concatenate(data, axis=0)

		mu_new = np.mean(data, axis=0)
		sigma_new = scale ** 2 * np.cov(data.T)

		self.priors = np.concatenate([self.priors, 0.01 * np.ones(1)])
		self.priors /= np.sum(self.priors)
		self.mu = np.concatenate([self.mu, mu_new[None]], axis=0)
		self.sigma = np.concatenate([self.sigma, sigma_new[None]], axis=0)



403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
	def mvn_pdf(self, x, reg=None):
		"""

		:param x: 			np.array([nb_samples, nb_dim])
			samples
		:param mu: 			np.array([nb_states, nb_dim])
			mean vector
		:param sigma_chol: 	np.array([nb_states, nb_dim, nb_dim])
			cholesky decomposition of covariance matrices
		:param lmbda: 		np.array([nb_states, nb_dim, nb_dim])
			precision matrices
		:return: 			np.array([nb_states, nb_samples])
			log mvn
		"""
		# if len(x.shape) > 1:  # TODO implement mvn for multiple xs
		# 	raise NotImplementedError
		mu, lmbda_, sigma_chol_ = self.mu, self.lmbda, self.sigma_chol

		if x.ndim > 1:
			dx = mu[None] - x[:, None]  # nb_timesteps, nb_states, nb_dim
		else:
			dx = mu - x

		eins_idx = ('baj,baj->ba', 'ajk,baj->bak') if x.ndim > 1 else (
		'aj,aj->a', 'ajk,aj->ak')

		return -0.5 * np.einsum(eins_idx[0], dx, np.einsum(eins_idx[1], lmbda_, dx)) \
			   - mu.shape[1] / 2. * np.log(2 * np.pi) - np.sum(
			np.log(sigma_chol_.diagonal(axis1=1, axis2=2)), axis=1)