mirror of
https://github.com/SheffieldML/GPy.git
synced 2026-05-04 09:12:38 +02:00
Refactor EP and EPDTC
This commit is contained in:
parent
7c4e2f21f8
commit
255d97a917
2 changed files with 233 additions and 209 deletions
|
|
@ -9,6 +9,107 @@ from .posterior import PosteriorEP as Posterior
|
|||
|
||||
log_2_pi = np.log(2*np.pi)
|
||||
|
||||
|
||||
#Four wrapper classes to help modularisation of different EP versions
|
||||
class marginalMoments(object):
|
||||
def __init__(self, num_data):
|
||||
self.Z_hat = np.empty(num_data,dtype=np.float64)
|
||||
self.mu_hat = np.empty(num_data,dtype=np.float64)
|
||||
self.sigma2_hat = np.empty(num_data,dtype=np.float64)
|
||||
|
||||
|
||||
class cavityParams(object):
|
||||
def __init__(self, num_data):
|
||||
self.tau = np.empty(num_data,dtype=np.float64)
|
||||
self.v = np.empty(num_data,dtype=np.float64)
|
||||
def _update_i(self, eta, ga_approx, post_params, i):
|
||||
self.tau[i] = 1./post_params.Sigma_diag[i] - eta*ga_approx.tau[i]
|
||||
self.v[i] = post_params.mu[i]/post_params.Sigma_diag[i] - eta*ga_approx.v[i]
|
||||
|
||||
|
||||
class gaussianApproximation(object):
|
||||
def __init__(self, v, tau):
|
||||
self.tau = tau
|
||||
self.v = v
|
||||
def _update_i(self, eta, delta, post_params, marg_moments, i):
|
||||
#Site parameters update
|
||||
delta_tau = delta/eta*(1./marg_moments.sigma2_hat[i] - 1./post_params.Sigma_diag[i])
|
||||
delta_v = delta/eta*(marg_moments.mu_hat[i]/marg_moments.sigma2_hat[i] - post_params.mu[i]/post_params.Sigma_diag[i])
|
||||
tau_tilde_prev = self.tau[i]
|
||||
self.tau[i] += delta_tau
|
||||
|
||||
# Enforce positivity of tau_tilde. Even though this is guaranteed for logconcave sites, it is still possible
|
||||
# to get negative values due to numerical errors. Moreover, the value of tau_tilde should be positive in order to
|
||||
# update the marginal likelihood without runnint into instabilities issues.
|
||||
if self.tau[i] < np.finfo(float).eps:
|
||||
self.tau[i] = np.finfo(float).eps
|
||||
delta_tau = self.tau[i] - tau_tilde_prev
|
||||
self.v[i] += delta_v
|
||||
|
||||
return (delta_tau, delta_v)
|
||||
|
||||
|
||||
class posteriorParamsBase(object):
|
||||
def __init__(self, mu, Sigma_diag):
|
||||
self.mu = mu
|
||||
self.Sigma_diag = Sigma_diag
|
||||
def _update_rank1(self, *arg):
|
||||
pass
|
||||
|
||||
def _recompute(self, *arg):
|
||||
pass
|
||||
|
||||
class posteriorParams(posteriorParamsBase):
|
||||
def __init__(self, mu, Sigma, L=None):
|
||||
self.Sigma = Sigma
|
||||
self.L = L
|
||||
Sigma_diag = np.diag(self.Sigma)
|
||||
super(posteriorParams, self).__init__(mu, Sigma_diag)
|
||||
|
||||
def _update_rank1(self, delta_tau, ga_approx, i):
|
||||
ci = delta_tau/(1.+ delta_tau*self.Sigma_diag[i])
|
||||
DSYR(self.Sigma, self.Sigma[:,i].copy(), -ci)
|
||||
self.mu = np.dot(self.Sigma, ga_approx.v)
|
||||
|
||||
@staticmethod
|
||||
def _recompute(K, ga_approx):
|
||||
num_data = len(ga_approx.tau)
|
||||
tau_tilde_root = np.sqrt(ga_approx.tau)
|
||||
Sroot_tilde_K = tau_tilde_root[:,None] * K
|
||||
B = np.eye(num_data) + Sroot_tilde_K * tau_tilde_root[None,:]
|
||||
L = jitchol(B)
|
||||
V, _ = dtrtrs(L, Sroot_tilde_K, lower=1)
|
||||
Sigma = K - np.dot(V.T,V) #K - KS^(1/2)BS^(1/2)K = (K^(-1) + \Sigma^(-1))^(-1)
|
||||
mu = np.dot(Sigma,ga_approx.v)
|
||||
return posteriorParams(mu=mu, Sigma=Sigma, L=L)
|
||||
|
||||
class posteriorParamsDTC(posteriorParamsBase):
|
||||
def __init__(self, mu, Sigma_diag):
|
||||
super(posteriorParamsDTC, self).__init__(mu, Sigma_diag)
|
||||
|
||||
def _update_rank1(self, LLT, Kmn, delta_v, delta_tau, i):
|
||||
#DSYR(Sigma, Sigma[:,i].copy(), -delta_tau/(1.+ delta_tau*Sigma[i,i]))
|
||||
DSYR(LLT,Kmn[:,i].copy(),delta_tau)
|
||||
L = jitchol(LLT)
|
||||
V,info = dtrtrs(L,Kmn,lower=1)
|
||||
self.Sigma_diag = np.maximum(np.sum(V*V,-2), np.finfo(float).eps) #diag(K_nm (L L^\top)^(-1)) K_mn
|
||||
si = np.sum(V.T*V[:,i],-1) #(V V^\top)[:,i]
|
||||
self.mu += (delta_v-delta_tau*self.mu[i])*si
|
||||
#mu = np.dot(Sigma, v_tilde)
|
||||
|
||||
@staticmethod
|
||||
def _recompute(LLT0, Kmn, ga_approx):
|
||||
LLT = LLT0 + np.dot(Kmn*ga_approx.tau[None,:],Kmn.T)
|
||||
L = jitchol(LLT)
|
||||
V, _ = dtrtrs(L,Kmn,lower=1)
|
||||
#Sigma_diag = np.sum(V*V,-2)
|
||||
#Knmv_tilde = np.dot(Kmn,v_tilde)
|
||||
#mu = np.dot(V2.T,Knmv_tilde)
|
||||
Sigma = np.dot(V.T,V)
|
||||
mu = np.dot(Sigma, ga_approx.v)
|
||||
Sigma_diag = np.diag(Sigma).copy()
|
||||
return posteriorParamsDTC(mu, Sigma_diag), LLT
|
||||
|
||||
class EPBase(object):
|
||||
def __init__(self, epsilon=1e-6, eta=1., delta=1., always_reset=False, max_iters=np.inf, ep_mode="alternated", parallel_updates=False):
|
||||
"""
|
||||
|
|
@ -28,6 +129,7 @@ class EPBase(object):
|
|||
:parallel_updates: boolean. If true, updates of the parameters of the sites in parallel
|
||||
"""
|
||||
super(EPBase, self).__init__()
|
||||
|
||||
self.always_reset = always_reset
|
||||
self.epsilon, self.eta, self.delta, self.max_iters = epsilon, eta, delta, max_iters
|
||||
self.ep_mode = ep_mode
|
||||
|
|
@ -35,7 +137,6 @@ class EPBase(object):
|
|||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.old_mutilde, self.old_vtilde = None, None
|
||||
self.ga_approx_old = None
|
||||
self._ep_approximation = None
|
||||
|
||||
|
|
@ -46,6 +147,11 @@ class EPBase(object):
|
|||
# TODO: update approximation in the end as well? Maybe even with a switch?
|
||||
pass
|
||||
|
||||
def _stop_criteria(self, ga_approx):
|
||||
tau_diff = np.mean(np.square(ga_approx.tau-self.ga_approx_old.tau))
|
||||
v_diff = np.mean(np.square(ga_approx.v-self.ga_approx_old.v))
|
||||
return ((tau_diff > self.epsilon) or (v_diff > self.epsilon))
|
||||
|
||||
def __setstate__(self, state):
|
||||
super(EPBase, self).__setstate__(state[0])
|
||||
self.epsilon, self.eta, self.delta = state[1]
|
||||
|
|
@ -68,19 +174,18 @@ class EP(EPBase, ExactGaussianInference):
|
|||
if self.ep_mode=="nested":
|
||||
#Force EP at each step of the optimization
|
||||
self._ep_approximation = None
|
||||
mu, Sigma, mu_tilde, tau_tilde, log_Z_tilde = self._ep_approximation = self.expectation_propagation(K, Y, likelihood, Y_metadata)
|
||||
post_params, ga_approx, log_Z_tilde = self._ep_approximation = self.expectation_propagation(K, Y, likelihood, Y_metadata)
|
||||
elif self.ep_mode=="alternated":
|
||||
if getattr(self, '_ep_approximation', None) is None:
|
||||
#if we don't yet have the results of runnign EP, run EP and store the computed factors in self._ep_approximation
|
||||
mu, Sigma, mu_tilde, tau_tilde, log_Z_tilde = self._ep_approximation = self.expectation_propagation(K, Y, likelihood, Y_metadata)
|
||||
post_params, ga_approx, log_Z_tilde = self._ep_approximation = self.expectation_propagation(K, Y, likelihood, Y_metadata)
|
||||
else:
|
||||
#if we've already run EP, just use the existing approximation stored in self._ep_approximation
|
||||
mu, Sigma, mu_tilde, tau_tilde, log_Z_tilde = self._ep_approximation
|
||||
post_params, ga_approx, log_Z_tilde = self._ep_approximation
|
||||
else:
|
||||
raise ValueError("ep_mode value not valid")
|
||||
|
||||
v_tilde = mu_tilde * tau_tilde
|
||||
return self._inference(K, tau_tilde, v_tilde, likelihood, Y_metadata=Y_metadata, Z_tilde=log_Z_tilde.sum())
|
||||
return self._inference(K, ga_approx, likelihood, Y_metadata=Y_metadata, Z_tilde=log_Z_tilde)
|
||||
|
||||
def expectation_propagation(self, K, Y, likelihood, Y_metadata):
|
||||
|
||||
|
|
@ -97,43 +202,51 @@ class EP(EPBase, ExactGaussianInference):
|
|||
ga_approx, post_params = self._init_approximations(K, num_data)
|
||||
|
||||
#Approximation
|
||||
tau_diff = self.epsilon + 1.
|
||||
v_diff = self.epsilon + 1.
|
||||
|
||||
stop = False
|
||||
iterations = 0
|
||||
while ((tau_diff > self.epsilon) or (v_diff > self.epsilon)) and (iterations < self.max_iters):
|
||||
self._update_cavity_params(num_data, cav_params, post_params, marg_moments, ga_approx, likelihood, Y, Y_metadata)
|
||||
while not stop and (iterations < self.max_iters):
|
||||
self._local_updates(num_data, cav_params, post_params, marg_moments, ga_approx, likelihood, Y, Y_metadata)
|
||||
|
||||
#(re) compute Sigma and mu using full Cholesky decompy
|
||||
post_params = self._ep_compute_posterior(K, ga_approx.tau, ga_approx.v)
|
||||
post_params = posteriorParams._recompute(K, ga_approx)
|
||||
|
||||
#monitor convergence
|
||||
if iterations > 0:
|
||||
tau_diff = np.mean(np.square(ga_approx.tau-self.ga_approx_old.tau))
|
||||
v_diff = np.mean(np.square(ga_approx.v-self.ga_approx_old.v))
|
||||
self.ga_approx_old = gaussianApproximation(ga_approx.mu.copy(), ga_approx.v.copy(), ga_approx.tau.copy())
|
||||
stop = self._stop_criteria(ga_approx)
|
||||
self.ga_approx_old = gaussianApproximation(ga_approx.v.copy(), ga_approx.tau.copy())
|
||||
iterations += 1
|
||||
|
||||
ga_approx.mu = ga_approx.v/ga_approx.tau
|
||||
|
||||
# Z_tilde after removing the terms that can lead to infinite terms due to tau_tilde close to zero.
|
||||
# This terms cancel with the coreresponding terms in the marginal loglikelihood
|
||||
log_Z_tilde = self._log_Z_tilde(marg_moments, ga_approx, cav_params)
|
||||
# - 0.5*np.log(tau_tilde) + 0.5*(v_tilde*v_tilde*1./tau_tilde)
|
||||
return (post_params, ga_approx, log_Z_tilde)
|
||||
|
||||
return post_params.mu, post_params.Sigma, ga_approx.mu, ga_approx.tau, log_Z_tilde
|
||||
|
||||
def _log_Z_tilde(self, marg_moments, ga_approx, cav_params):
|
||||
return (np.log(marg_moments.Z_hat) + 0.5*np.log(2*np.pi) + 0.5*np.log(1+ga_approx.tau/cav_params.tau) - 0.5 * ((ga_approx.v)**2 * 1./(cav_params.tau + ga_approx.tau))
|
||||
+ 0.5*(cav_params.v * ( ( (ga_approx.tau/cav_params.tau) * cav_params.v - 2.0 * ga_approx.v ) * 1./(cav_params.tau + ga_approx.tau))))
|
||||
|
||||
def _update_cavity_params(self, num_data, cav_params, post_params, marg_moments, ga_approx, likelihood, Y, Y_metadata, update_order=None):
|
||||
def _init_approximations(self, K, num_data):
|
||||
#initial values - Gaussian factors
|
||||
#Initial values - Posterior distribution parameters: q(f|X,Y) = N(f|mu,Sigma)
|
||||
if self.ga_approx_old is None:
|
||||
v_tilde, tau_tilde = np.zeros((2, num_data))
|
||||
ga_approx = gaussianApproximation(v_tilde, tau_tilde)
|
||||
Sigma = K.copy()
|
||||
diag.add(Sigma, 1e-7)
|
||||
mu = np.zeros(num_data)
|
||||
post_params = posteriorParams(mu, Sigma)
|
||||
else:
|
||||
assert self.ga_approx_old.v.size == num_data, "data size mis-match: did you change the data? try resetting!"
|
||||
ga_approx = gaussianApproximation(self.ga_approx_old.v, self.ga_approx_old.tau)
|
||||
post_params = posteriorParams._recompute(K, ga_approx)
|
||||
diag.add(post_params.Sigma, 1e-7)
|
||||
# TODO: Check the log-marginal under both conditions and choose the best one
|
||||
return (ga_approx, post_params)
|
||||
|
||||
def _local_updates(self, num_data, cav_params, post_params, marg_moments, ga_approx, likelihood, Y, Y_metadata, update_order=None):
|
||||
if update_order is None:
|
||||
update_order = np.random.permutation(num_data)
|
||||
for i in update_order:
|
||||
#Cavity distribution parameters
|
||||
cav_params.tau[i] = 1./post_params.Sigma[i,i] - self.eta*ga_approx.tau[i]
|
||||
cav_params.v[i] = post_params.mu[i]/post_params.Sigma[i,i] - self.eta*ga_approx.v[i]
|
||||
cav_params._update_i(self.eta, ga_approx, post_params, i)
|
||||
|
||||
if Y_metadata is not None:
|
||||
# Pick out the relavent metadata for Yi
|
||||
Y_metadata_i = {}
|
||||
|
|
@ -143,76 +256,39 @@ class EP(EPBase, ExactGaussianInference):
|
|||
Y_metadata_i = None
|
||||
#Marginal moments
|
||||
marg_moments.Z_hat[i], marg_moments.mu_hat[i], marg_moments.sigma2_hat[i] = likelihood.moments_match_ep(Y[i], cav_params.tau[i], cav_params.v[i], Y_metadata_i=Y_metadata_i)
|
||||
|
||||
#Site parameters update
|
||||
delta_tau = self.delta/self.eta*(1./marg_moments.sigma2_hat[i] - 1./post_params.Sigma[i,i])
|
||||
delta_v = self.delta/self.eta*(marg_moments.mu_hat[i]/marg_moments.sigma2_hat[i] - post_params.mu[i]/post_params.Sigma[i,i])
|
||||
tau_tilde_prev = ga_approx.tau[i]
|
||||
ga_approx.tau[i] += delta_tau
|
||||
|
||||
# Enforce positivity of tau_tilde. Even though this is guaranteed for logconcave sites, it is still possible
|
||||
# to get negative values due to numerical errors. Moreover, the value of tau_tilde should be positive in order to
|
||||
# update the marginal likelihood without inestability issues.
|
||||
if ga_approx.tau[i] < np.finfo(float).eps:
|
||||
ga_approx.tau[i] = np.finfo(float).eps
|
||||
delta_tau = ga_approx.tau[i] - tau_tilde_prev
|
||||
ga_approx.v[i] += delta_v
|
||||
#Site parameters update
|
||||
delta_tau, delta_v = ga_approx._update_i(self.eta, self.delta, post_params, marg_moments, i)
|
||||
|
||||
if self.parallel_updates == False:
|
||||
#Posterior distribution parameters update
|
||||
ci = delta_tau/(1.+ delta_tau*post_params.Sigma[i,i])
|
||||
DSYR(post_params.Sigma, post_params.Sigma[:,i].copy(), -ci)
|
||||
post_params.mu = np.dot(post_params.Sigma, ga_approx.v)
|
||||
|
||||
def _init_approximations(self, K, num_data):
|
||||
#initial values - Gaussian factors
|
||||
#Initial values - Posterior distribution parameters: q(f|X,Y) = N(f|mu,Sigma)
|
||||
if self.ga_approx_old is None:
|
||||
mu_tilde, v_tilde, tau_tilde = np.zeros((3, num_data))
|
||||
ga_approx = gaussianApproximation(mu_tilde, v_tilde, tau_tilde)
|
||||
Sigma = K.copy()
|
||||
diag.add(Sigma, 1e-7)
|
||||
mu = np.zeros(num_data)
|
||||
post_params = posteriorParams(mu, Sigma)
|
||||
else:
|
||||
assert self.ga_approx_old.mu.size == num_data, "data size mis-match: did you change the data? try resetting!"
|
||||
ga_approx = gaussianApproximation(self.ga_approx_old.mu, self.ga_approx_old.v)
|
||||
post_params = self._ep_compute_posterior(K, ga_approx.tau, ga_approx.v)
|
||||
diag.add(post_params.Sigma, 1e-7)
|
||||
# TODO: Check the log-marginal under both conditions and choose the best one
|
||||
return (ga_approx, post_params)
|
||||
|
||||
def _ep_compute_posterior(self, K, tau_tilde, v_tilde):
|
||||
num_data = len(tau_tilde)
|
||||
tau_tilde_root = np.sqrt(tau_tilde)
|
||||
Sroot_tilde_K = tau_tilde_root[:,None] * K
|
||||
B = np.eye(num_data) + Sroot_tilde_K * tau_tilde_root[None,:]
|
||||
L = jitchol(B)
|
||||
V, _ = dtrtrs(L, Sroot_tilde_K, lower=1)
|
||||
Sigma = K - np.dot(V.T,V) #K - KS^(1/2)BS^(1/2)K = (K^(-1) + \Sigma^(-1))^(-1)
|
||||
mu = np.dot(Sigma,v_tilde)
|
||||
return posteriorParams(mu, Sigma, L)
|
||||
post_params._update_rank1(delta_tau, ga_approx, i)
|
||||
|
||||
def _ep_marginal(self, K, tau_tilde, v_tilde, Z_tilde):
|
||||
post_params = self._ep_compute_posterior(K, tau_tilde, v_tilde)
|
||||
def _log_Z_tilde(self, marg_moments, ga_approx, cav_params):
|
||||
return np.sum((np.log(marg_moments.Z_hat) + 0.5*np.log(2*np.pi) + 0.5*np.log(1+ga_approx.tau/cav_params.tau) - 0.5 * ((ga_approx.v)**2 * 1./(cav_params.tau + ga_approx.tau))
|
||||
+ 0.5*(cav_params.v * ( ( (ga_approx.tau/cav_params.tau) * cav_params.v - 2.0 * ga_approx.v ) * 1./(cav_params.tau + ga_approx.tau)))))
|
||||
|
||||
|
||||
|
||||
def _ep_marginal(self, K, ga_approx, Z_tilde):
|
||||
post_params = posteriorParams._recompute(K, ga_approx)
|
||||
|
||||
# Gaussian log marginal excluding terms that can go to infinity due to arbitrarily small tau_tilde.
|
||||
# These terms cancel out with the terms excluded from Z_tilde
|
||||
B_logdet = np.sum(2.0*np.log(np.diag(post_params.L)))
|
||||
log_marginal = 0.5*(-len(tau_tilde) * log_2_pi - B_logdet + np.sum(v_tilde * np.dot(post_params.Sigma,v_tilde)))
|
||||
log_marginal = 0.5*(-len(ga_approx.tau) * log_2_pi - B_logdet + np.sum(ga_approx.v * np.dot(post_params.Sigma,ga_approx.v)))
|
||||
log_marginal += Z_tilde
|
||||
|
||||
return log_marginal, post_params.mu, post_params.Sigma, post_params.L
|
||||
return log_marginal, post_params
|
||||
|
||||
def _inference(self, K, tau_tilde, v_tilde, likelihood, Z_tilde, Y_metadata=None):
|
||||
log_marginal, mu, Sigma, L = self._ep_marginal(K, tau_tilde, v_tilde, Z_tilde)
|
||||
def _inference(self, K, ga_approx, likelihood, Z_tilde, Y_metadata=None):
|
||||
log_marginal, post_params = self._ep_marginal(K, ga_approx, Z_tilde)
|
||||
|
||||
tau_tilde_root = np.sqrt(tau_tilde)
|
||||
tau_tilde_root = np.sqrt(ga_approx.tau)
|
||||
Sroot_tilde_K = tau_tilde_root[:,None] * K
|
||||
|
||||
aux_alpha , _ = dpotrs(L, np.dot(Sroot_tilde_K, v_tilde), lower=1)
|
||||
alpha = (v_tilde - tau_tilde_root * aux_alpha)[:,None] #(K + Sigma^(\tilde))^(-1) /mu^(/tilde)
|
||||
LWi, _ = dtrtrs(L, np.diag(tau_tilde_root), lower=1)
|
||||
aux_alpha , _ = dpotrs(post_params.L, np.dot(Sroot_tilde_K, ga_approx.v), lower=1)
|
||||
alpha = (ga_approx.v - tau_tilde_root * aux_alpha)[:,None] #(K + Sigma^(\tilde))^(-1) /mu^(/tilde)
|
||||
LWi, _ = dtrtrs(post_params.L, np.diag(tau_tilde_root), lower=1)
|
||||
Wi = np.dot(LWi.T,LWi)
|
||||
symmetrify(Wi) #(K + Sigma^(\tilde))^(-1)
|
||||
|
||||
|
|
@ -245,24 +321,25 @@ class EPDTC(EPBase, VarDTC):
|
|||
if self.ep_mode=="nested":
|
||||
#Force EP at each step of the optimization
|
||||
self._ep_approximation = None
|
||||
mu, Sigma_diag, mu_tilde, tau_tilde, log_Z_tilde = self._ep_approximation = self.expectation_propagation(Kmm, Kmn, Y, likelihood, Y_metadata)
|
||||
post_params, ga_approx, log_Z_tilde = self._ep_approximation = self.expectation_propagation(Kmm, Kmn, Y, likelihood, Y_metadata)
|
||||
elif self.ep_mode=="alternated":
|
||||
if getattr(self, '_ep_approximation', None) is None:
|
||||
#if we don't yet have the results of runnign EP, run EP and store the computed factors in self._ep_approximation
|
||||
mu, Sigma_diag, mu_tilde, tau_tilde, log_Z_tilde = self._ep_approximation = self.expectation_propagation(Kmm, Kmn, Y, likelihood, Y_metadata)
|
||||
post_params, ga_approx, log_Z_tilde = self._ep_approximation = self.expectation_propagation(Kmm, Kmn, Y, likelihood, Y_metadata)
|
||||
else:
|
||||
#if we've already run EP, just use the existing approximation stored in self._ep_approximation
|
||||
mu, Sigma_diag, mu_tilde, tau_tilde, log_Z_tilde = self._ep_approximation
|
||||
post_params, ga_approx, log_Z_tilde = self._ep_approximation
|
||||
else:
|
||||
raise ValueError("ep_mode value not valid")
|
||||
|
||||
return super(EPDTC, self).inference(kern, X, Z, likelihood, mu_tilde,
|
||||
mu_tilde = ga_approx.v / ga_approx.tau.astype(float)
|
||||
|
||||
return super(EPDTC, self).inference(kern, X, Z, likelihood, ObsAr(mu_tilde[:,None]),
|
||||
mean_function=mean_function,
|
||||
Y_metadata=Y_metadata,
|
||||
precision=tau_tilde,
|
||||
precision=ga_approx.tau,
|
||||
Lm=Lm, dL_dKmm=dL_dKmm,
|
||||
psi0=psi0, psi1=psi1, psi2=psi2, Z_tilde=log_Z_tilde.sum())
|
||||
|
||||
psi0=psi0, psi1=psi1, psi2=psi2, Z_tilde=log_Z_tilde)
|
||||
|
||||
def expectation_propagation(self, Kmm, Kmn, Y, likelihood, Y_metadata):
|
||||
|
||||
|
|
@ -273,144 +350,91 @@ class EPDTC(EPBase, VarDTC):
|
|||
# than ObsArrays
|
||||
Y = Y.values.copy()
|
||||
|
||||
#Initial values - Marginal moments
|
||||
Z_hat = np.zeros(num_data,dtype=np.float64)
|
||||
mu_hat = np.zeros(num_data,dtype=np.float64)
|
||||
sigma2_hat = np.zeros(num_data,dtype=np.float64)
|
||||
#Initial values - Marginal moments, cavity params, gaussian approximation params and posterior params
|
||||
marg_moments = marginalMoments(num_data)
|
||||
cav_params = cavityParams(num_data)
|
||||
ga_approx, post_params, LLT0, LLT = self._init_approximations(Kmm, Kmn, num_data)
|
||||
|
||||
tau = np.empty(num_data,dtype=np.float64)
|
||||
v = np.empty(num_data,dtype=np.float64)
|
||||
#Approximation
|
||||
stop = False
|
||||
iterations = 0
|
||||
while not stop and (iterations < self.max_iters):
|
||||
self._local_updates(num_data, LLT0, LLT, Kmn, cav_params, post_params, marg_moments, ga_approx, likelihood, Y, Y_metadata)
|
||||
#(re) compute Sigma, Sigma_diag and mu using full Cholesky decompy
|
||||
post_params, LLT = posteriorParamsDTC._recompute(LLT0, Kmn, ga_approx)
|
||||
post_params.Sigma_diag = np.maximum(post_params.Sigma_diag, np.finfo(float).eps)
|
||||
|
||||
#monitor convergence
|
||||
if iterations > 0:
|
||||
stop = self._stop_criteria(ga_approx)
|
||||
self.ga_approx_old = gaussianApproximation(ga_approx.v.copy(), ga_approx.tau.copy())
|
||||
iterations += 1
|
||||
|
||||
log_Z_tilde = self._log_Z_tilde(marg_moments, ga_approx, cav_params)
|
||||
|
||||
return post_params, ga_approx, log_Z_tilde
|
||||
|
||||
def _log_Z_tilde(self, marg_moments, ga_approx, cav_params):
|
||||
mu_tilde = ga_approx.v/ga_approx.tau
|
||||
mu_cav = cav_params.v/cav_params.tau
|
||||
sigma2_sigma2tilde = 1./cav_params.tau + 1./ga_approx.tau
|
||||
|
||||
return np.sum((np.log(marg_moments.Z_hat) + 0.5*np.log(2*np.pi) + 0.5*np.log(sigma2_sigma2tilde)
|
||||
+ 0.5*((mu_cav - mu_tilde)**2) / (sigma2_sigma2tilde)))
|
||||
|
||||
def _init_approximations(self, Kmm, Kmn, num_data):
|
||||
#initial values - Gaussian factors
|
||||
#Initial values - Posterior distribution parameters: q(f|X,Y) = N(f|mu,Sigma)
|
||||
LLT0 = Kmm.copy()
|
||||
Lm = jitchol(LLT0) #K_m = L_m L_m^\top
|
||||
Vm,info = dtrtrs(Lm,Kmn,lower=1)
|
||||
Vm,info = dtrtrs(Lm, Kmn,lower=1)
|
||||
# Lmi = dtrtri(Lm)
|
||||
# Kmmi = np.dot(Lmi.T,Lmi)
|
||||
# KmmiKmn = np.dot(Kmmi,Kmn)
|
||||
# Qnn_diag = np.sum(Kmn*KmmiKmn,-2)
|
||||
Qnn_diag = np.sum(Vm*Vm,-2) #diag(Knm Kmm^(-1) Kmn)
|
||||
#diag.add(LLT0, 1e-8)
|
||||
if self.old_mutilde is None:
|
||||
if self.ga_approx_old is None:
|
||||
#Initial values - Posterior distribution parameters: q(f|X,Y) = N(f|mu,Sigma)
|
||||
LLT = LLT0.copy() #Sigma = K.copy()
|
||||
mu = np.zeros(num_data)
|
||||
Sigma_diag = Qnn_diag.copy() + 1e-8
|
||||
tau_tilde, mu_tilde, v_tilde = np.zeros((3, num_data))
|
||||
v_tilde, tau_tilde = np.zeros((2, num_data))
|
||||
ga_approx = gaussianApproximation(v_tilde, tau_tilde)
|
||||
post_params = posteriorParamsDTC(mu, Sigma_diag)
|
||||
|
||||
else:
|
||||
assert self.old_mutilde.size == num_data, "data size mis-match: did you change the data? try resetting!"
|
||||
mu_tilde, v_tilde = self.old_mutilde, self.old_vtilde
|
||||
tau_tilde = v_tilde/mu_tilde
|
||||
mu, Sigma_diag, LLT = self._ep_compute_posterior(LLT0, Kmn, tau_tilde, v_tilde)
|
||||
Sigma_diag += 1e-8
|
||||
assert self.ga_approx_old.v.size == num_data, "data size mis-match: did you change the data? try resetting!"
|
||||
ga_approx = gaussianApproximation(self.ga_approx_old.v, self.ga_approx_old.tau)
|
||||
post_params, LLT = posteriorParamsDTC._recompute(LLT0, Kmn, ga_approx)
|
||||
post_params.Sigma_diag += 1e-8
|
||||
|
||||
# TODO: Check the log-marginal under both conditions and choose the best one
|
||||
|
||||
#Approximation
|
||||
tau_diff = self.epsilon + 1.
|
||||
v_diff = self.epsilon + 1.
|
||||
tau_tilde_old = np.nan
|
||||
v_tilde_old = np.nan
|
||||
iterations = 0
|
||||
while ((tau_diff > self.epsilon) or (v_diff > self.epsilon)) and (iterations < self.max_iters):
|
||||
return (ga_approx, post_params, LLT0, LLT)
|
||||
|
||||
def _local_updates(self, num_data, LLT0, LLT, Kmn, cav_params, post_params, marg_moments, ga_approx, likelihood, Y, Y_metadata, update_order=None):
|
||||
if update_order is None:
|
||||
update_order = np.random.permutation(num_data)
|
||||
for i in update_order:
|
||||
#Cavity distribution parameters
|
||||
tau[i] = 1./Sigma_diag[i] - self.eta*tau_tilde[i]
|
||||
v[i] = mu[i]/Sigma_diag[i] - self.eta*v_tilde[i]
|
||||
if Y_metadata is not None:
|
||||
# Pick out the relavent metadata for Yi
|
||||
Y_metadata_i = {}
|
||||
for key in Y_metadata.keys():
|
||||
Y_metadata_i[key] = Y_metadata[key][i, :]
|
||||
else:
|
||||
Y_metadata_i = None
|
||||
for i in update_order:
|
||||
|
||||
#Marginal moments
|
||||
Z_hat[i], mu_hat[i], sigma2_hat[i] = likelihood.moments_match_ep(Y[i], tau[i], v[i], Y_metadata_i=Y_metadata_i)
|
||||
#Site parameters update
|
||||
delta_tau = self.delta/self.eta*(1./sigma2_hat[i] - 1./Sigma_diag[i])
|
||||
delta_v = self.delta/self.eta*(mu_hat[i]/sigma2_hat[i] - mu[i]/Sigma_diag[i])
|
||||
tau_tilde_prev = tau_tilde[i]
|
||||
tau_tilde[i] += delta_tau
|
||||
#Cavity distribution parameters
|
||||
cav_params._update_i(self.eta, ga_approx, post_params, i)
|
||||
|
||||
# Enforce positivity of tau_tilde. Even though this is guaranteed for logconcave sites, it is still possible
|
||||
# to get negative values due to numerical errors. Moreover, the value of tau_tilde should be positive in order to
|
||||
# update the marginal likelihood without inestability issues.
|
||||
if tau_tilde[i] < np.finfo(float).eps:
|
||||
tau_tilde[i] = np.finfo(float).eps
|
||||
delta_tau = tau_tilde[i] - tau_tilde_prev
|
||||
v_tilde[i] += delta_v
|
||||
|
||||
#Posterior distribution parameters update
|
||||
if self.parallel_updates == False:
|
||||
#DSYR(Sigma, Sigma[:,i].copy(), -delta_tau/(1.+ delta_tau*Sigma[i,i]))
|
||||
DSYR(LLT,Kmn[:,i].copy(),delta_tau)
|
||||
L = jitchol(LLT)
|
||||
V,info = dtrtrs(L,Kmn,lower=1)
|
||||
Sigma_diag = np.maximum(np.sum(V*V,-2), np.finfo(float).eps) #diag(K_nm (L L^\top)^(-1)) K_mn
|
||||
si = np.sum(V.T*V[:,i],-1) #(V V^\top)[:,i]
|
||||
mu += (delta_v-delta_tau*mu[i])*si
|
||||
#mu = np.dot(Sigma, v_tilde)
|
||||
if Y_metadata is not None:
|
||||
# Pick out the relavent metadata for Yi
|
||||
Y_metadata_i = {}
|
||||
for key in Y_metadata.keys():
|
||||
Y_metadata_i[key] = Y_metadata[key][i, :]
|
||||
else:
|
||||
Y_metadata_i = None
|
||||
|
||||
#(re) compute Sigma, Sigma_diag and mu using full Cholesky decompy
|
||||
mu, Sigma_diag, LLT = self._ep_compute_posterior(LLT0, Kmn, tau_tilde, v_tilde)
|
||||
Sigma_diag = np.maximum(Sigma_diag, np.finfo(float).eps)
|
||||
#Marginal moments
|
||||
marg_moments.Z_hat[i], marg_moments.mu_hat[i], marg_moments.sigma2_hat[i] = likelihood.moments_match_ep(Y[i], cav_params.tau[i], cav_params.v[i], Y_metadata_i=Y_metadata_i)
|
||||
#Site parameters update
|
||||
delta_tau, delta_v = ga_approx._update_i(self.eta, self.delta, post_params, marg_moments, i)
|
||||
|
||||
#monitor convergence
|
||||
if iterations>0:
|
||||
tau_diff = np.mean(np.square(tau_tilde-tau_tilde_old))
|
||||
v_diff = np.mean(np.square(v_tilde-v_tilde_old))
|
||||
tau_tilde_old = tau_tilde.copy()
|
||||
v_tilde_old = v_tilde.copy()
|
||||
iterations += 1
|
||||
|
||||
mu_tilde = v_tilde/tau_tilde
|
||||
mu_cav = v/tau
|
||||
sigma2_sigma2tilde = 1./tau + 1./tau_tilde
|
||||
|
||||
log_Z_tilde = (np.log(Z_hat) + 0.5*np.log(2*np.pi) + 0.5*np.log(sigma2_sigma2tilde)
|
||||
+ 0.5*((mu_cav - mu_tilde)**2) / (sigma2_sigma2tilde))
|
||||
|
||||
self.old_mutilde = mu_tilde
|
||||
self.old_vtilde = v_tilde
|
||||
|
||||
return mu, Sigma_diag, ObsAr(mu_tilde[:,None]), tau_tilde, log_Z_tilde
|
||||
|
||||
def _ep_compute_posterior(self, LLT0, Kmn, tau_tilde, v_tilde):
|
||||
LLT = LLT0 + np.dot(Kmn*tau_tilde[None,:],Kmn.T)
|
||||
L = jitchol(LLT)
|
||||
V, _ = dtrtrs(L,Kmn,lower=1)
|
||||
#Sigma_diag = np.sum(V*V,-2)
|
||||
#Knmv_tilde = np.dot(Kmn,v_tilde)
|
||||
#mu = np.dot(V2.T,Knmv_tilde)
|
||||
Sigma = np.dot(V.T,V)
|
||||
mu = np.dot(Sigma,v_tilde)
|
||||
Sigma_diag = np.diag(Sigma).copy()
|
||||
|
||||
return (mu, Sigma_diag, LLT)
|
||||
|
||||
#Four wrapper classes to help modularisation of different EP versions
|
||||
class marginalMoments(object):
|
||||
def __init__(self, num_data):
|
||||
#Initial values - Marginal moments
|
||||
self.Z_hat = np.empty(num_data,dtype=np.float64)
|
||||
self.mu_hat = np.empty(num_data,dtype=np.float64)
|
||||
self.sigma2_hat = np.empty(num_data,dtype=np.float64)
|
||||
|
||||
class cavityParams(object):
|
||||
def __init__(self, num_data):
|
||||
self.tau = np.empty(num_data,dtype=np.float64)
|
||||
self.v = np.empty(num_data,dtype=np.float64)
|
||||
|
||||
class gaussianApproximation(object):
|
||||
def __init__(self, mu, v, tau=None):
|
||||
self.mu = mu
|
||||
self.v = v
|
||||
self.tau = mu / v if tau is None else tau
|
||||
|
||||
class posteriorParams(object):
|
||||
def __init__(self, mu=None, Sigma=None, L=None):
|
||||
self.mu = mu
|
||||
self.Sigma = Sigma
|
||||
self.L = L
|
||||
#Posterior distribution parameters update
|
||||
if self.parallel_updates == False:
|
||||
post_params._update_rank1(LLT, Kmn, delta_v, delta_tau, i)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue