def check_X(self, mask):
"""
Check update equations of parent X.
Use a simple model that has parent nodes and a child node around
MatrixDot.
"""
# Parent nodes
A = Gaussian(self.a.reshape((self.D3,self.D2,1,self.M*self.N)),
np.identity(self.M*self.N),
name='A')
X = Gaussian(np.zeros(self.N),
np.identity(self.N),
plates=(self.D3,1,self.D1),
name='X')
# Node itself
AX = MatrixDot(A, X, name='AX')
# Child node
Y = Gaussian(AX,
self.Lambda,
plates=(self.D4,self.D3,self.D2,self.D1),
name='Y')
# Put in data
Y.observe(self.y, mask=mask)
# VB model
Q = VB(Y, X)
Q.update(X, repeat=1)
u = X.get_moments()
# Compute true solution
mask = np.ones((self.D4,self.D3,self.D2,self.D1)) * mask
vv = mask[...,np.newaxis,np.newaxis] * self.vv
aa = np.einsum('...ki,...kl,...lj->...ij', self.a, vv, self.a)
aa = aa + np.einsum('...ij,...kl,...kl->...ij',
np.identity(self.N),
np.identity(self.M),
vv)
aa = np.sum(aa, axis=2, keepdims=True)
aa = np.sum(aa, axis=0)
Cov_X = linalg.chol_inv(linalg.chol(aa + np.identity(self.N)))
mu_X = np.einsum('...kj,...kl,...l->...j',
self.a,
vv,
self.y)
mu_X = np.sum(mu_X, axis=2, keepdims=True)
mu_X = np.sum(mu_X, axis=0)
mu_X = np.einsum('...ij,...j->...i', Cov_X, mu_X)
# Compare VB results to the analytic solution:
Cov_vb = u[1] - u[0][...,np.newaxis,:]*u[0][...,:,np.newaxis]
testing.assert_allclose(Cov_vb, Cov_X,
err_msg="Incorrect second moment.")
testing.assert_allclose(u[0], mu_X,
err_msg="Incorrect first moment.")
def check_A(self, mask):
"""
Check update equations of parent A.
Use a simple model that has parent nodes and a child node around
MatrixDot.
"""
# Parent nodes
A = Gaussian(np.zeros(self.M*self.N),
np.identity(self.M*self.N),
plates=(self.D3,self.D2,1),
name='A')
X = Gaussian(self.x, np.identity(self.N), name='X')
# Node itself
AX = MatrixDot(A, X, name='AX')
# Child node
Y = Gaussian(AX,
self.Lambda,
plates=(self.D4,self.D3,self.D2,self.D1),
name='Y')
# Put in data
Y.observe(self.y, mask=mask)
# VB model
Q = VB(Y, A)
Q.update(A, repeat=1)
u = A.get_moments()
# Compute true solution
mask = np.ones((self.D4,self.D3,self.D2,self.D1)) * mask
vv = mask[...,np.newaxis,np.newaxis] * self.vv
xx = (self.x[...,:,np.newaxis] * self.x[...,np.newaxis,:]
+ np.identity(self.N))
Cov_A = (vv[...,:,np.newaxis,:,np.newaxis] *
xx[...,np.newaxis,:,np.newaxis,:]) # from data
Cov_A = np.sum(Cov_A, axis=(0,3), keepdims=True)
Cov_A = np.reshape(Cov_A,
(self.D3,self.D2,1,self.M*self.N,self.M*self.N))
Cov_A = Cov_A + np.identity(self.M*self.N) # add prior
Cov_A = linalg.chol_inv(linalg.chol(Cov_A))
mu_A = np.einsum('...ij,...j,...k->...ik', vv, self.y, self.x)
mu_A = np.sum(mu_A, axis=(0,3))
mu_A = np.reshape(mu_A, (self.D3,self.D2,1,self.M*self.N))
mu_A = np.einsum('...ij,...j->...i', Cov_A, mu_A)
# Compare VB results to the analytic solution:
Cov_vb = u[1] - u[0][...,np.newaxis,:]*u[0][...,:,np.newaxis]
testing.assert_allclose(Cov_vb, Cov_A,
err_msg="Incorrect second moment.")
testing.assert_allclose(u[0], mu_A,
err_msg="Incorrect first moment.")
def check_A(self, mask):
"""
Check update equations of parent A.
Use a simple model that has parent nodes and a child node around
MatrixDot.
"""
# Parent nodes
A = Gaussian(np.zeros(self.M * self.N),
np.identity(self.M * self.N),
plates=(self.D3, self.D2, 1),
name='A')
X = Gaussian(self.x, np.identity(self.N), name='X')
# Node itself
AX = MatrixDot(A, X, name='AX')
# Child node
Y = Gaussian(AX,
self.Lambda,
plates=(self.D4, self.D3, self.D2, self.D1),
name='Y')
# Put in data
Y.observe(self.y, mask=mask)
# VB model
Q = VB(Y, A)
Q.update(A, repeat=1)
u = A.get_moments()
# Compute true solution
mask = np.ones((self.D4, self.D3, self.D2, self.D1)) * mask
vv = mask[..., np.newaxis, np.newaxis] * self.vv
xx = (self.x[..., :, np.newaxis] * self.x[..., np.newaxis, :] +
np.identity(self.N))
Cov_A = (vv[..., :, np.newaxis, :, np.newaxis] *
xx[..., np.newaxis, :, np.newaxis, :]) # from data
Cov_A = np.sum(Cov_A, axis=(0, 3), keepdims=True)
Cov_A = np.reshape(
Cov_A, (self.D3, self.D2, 1, self.M * self.N, self.M * self.N))
Cov_A = Cov_A + np.identity(self.M * self.N) # add prior
Cov_A = linalg.chol_inv(linalg.chol(Cov_A))
mu_A = np.einsum('...ij,...j,...k->...ik', vv, self.y, self.x)
mu_A = np.sum(mu_A, axis=(0, 3))
mu_A = np.reshape(mu_A, (self.D3, self.D2, 1, self.M * self.N))
mu_A = np.einsum('...ij,...j->...i', Cov_A, mu_A)
# Compare VB results to the analytic solution:
Cov_vb = u[1] - u[0][..., np.newaxis, :] * u[0][..., :, np.newaxis]
testing.assert_allclose(Cov_vb,
Cov_A,
err_msg="Incorrect second moment.")
testing.assert_allclose(u[0], mu_A, err_msg="Incorrect first moment.")
def check_X(self, mask):
"""
Check update equations of parent X.
Use a simple model that has parent nodes and a child node around
MatrixDot.
"""
# Parent nodes
A = Gaussian(self.a.reshape((self.D3, self.D2, 1, self.M * self.N)),
np.identity(self.M * self.N),
name='A')
X = Gaussian(np.zeros(self.N),
np.identity(self.N),
plates=(self.D3, 1, self.D1),
name='X')
# Node itself
AX = MatrixDot(A, X, name='AX')
# Child node
Y = Gaussian(AX,
self.Lambda,
plates=(self.D4, self.D3, self.D2, self.D1),
name='Y')
# Put in data
Y.observe(self.y, mask=mask)
# VB model
Q = VB(Y, X)
Q.update(X, repeat=1)
u = X.get_moments()
# Compute true solution
mask = np.ones((self.D4, self.D3, self.D2, self.D1)) * mask
vv = mask[..., np.newaxis, np.newaxis] * self.vv
aa = np.einsum('...ki,...kl,...lj->...ij', self.a, vv, self.a)
aa = aa + np.einsum('...ij,...kl,...kl->...ij', np.identity(self.N),
np.identity(self.M), vv)
aa = np.sum(aa, axis=2, keepdims=True)
aa = np.sum(aa, axis=0)
Cov_X = linalg.chol_inv(linalg.chol(aa + np.identity(self.N)))
mu_X = np.einsum('...kj,...kl,...l->...j', self.a, vv, self.y)
mu_X = np.sum(mu_X, axis=2, keepdims=True)
mu_X = np.sum(mu_X, axis=0)
mu_X = np.einsum('...ij,...j->...i', Cov_X, mu_X)
# Compare VB results to the analytic solution:
Cov_vb = u[1] - u[0][..., np.newaxis, :] * u[0][..., :, np.newaxis]
testing.assert_allclose(Cov_vb,
Cov_X,
err_msg="Incorrect second moment.")
testing.assert_allclose(u[0], mu_X, err_msg="Incorrect first moment.")
def compute_fixed_moments(self, Lambda, gradient=None):
""" Compute moments for fixed x. """
L = linalg.chol(Lambda, ndim=self.ndim)
ldet = linalg.chol_logdet(L, ndim=self.ndim)
u = [Lambda, ldet]
if gradient is None:
return u
du0 = gradient[0]
du1 = (misc.add_trailing_axes(gradient[1], 2 * self.ndim) *
linalg.chol_inv(L, ndim=self.ndim))
du = du0 + du1
return (u, du)
def compute_fixed_moments(self, Lambda, gradient=None):
""" Compute moments for fixed x. """
L = linalg.chol(Lambda, ndim=self.ndim)
ldet = linalg.chol_logdet(L, ndim=self.ndim)
u = [Lambda,
ldet]
if gradient is None:
return u
du0 = gradient[0]
du1 = (
misc.add_trailing_axes(gradient[1], 2*self.ndim)
* linalg.chol_inv(L, ndim=self.ndim)
)
du = du0 + du1
return (u, du)
def compute_moments_and_cgf(self, phi, mask=True):
r"""
Return moments and cgf for given natural parameters
.. math::
\langle u \rangle =
\begin{bmatrix}
\phi_2 (-\phi_1)^{-1}
\\
-\log|-\phi_1| + \psi_k(\phi_2)
\end{bmatrix}
\\
g(\phi) = \phi_2 \log|-\phi_1| - \log \Gamma_k(\phi_2)
"""
U = linalg.chol(-phi[0])
k = np.shape(phi[0])[-1]
#k = self.dims[0][0]
logdet_phi0 = linalg.chol_logdet(U)
u0 = phi[1][...,np.newaxis,np.newaxis] * linalg.chol_inv(U)
u1 = -logdet_phi0 + misc.multidigamma(phi[1], k)
u = [u0, u1]
g = phi[1] * logdet_phi0 - special.multigammaln(phi[1], k)
return (u, g)
def compute_moments_and_cgf(self, phi, mask=True):
r"""
Return moments and cgf for given natural parameters
.. math::
\langle u \rangle =
\begin{bmatrix}
\phi_2 (-\phi_1)^{-1}
\\
-\log|-\phi_1| + \psi_k(\phi_2)
\end{bmatrix}
\\
g(\phi) = \phi_2 \log|-\phi_1| - \log \Gamma_k(\phi_2)
"""
U = linalg.chol(-phi[0])
k = np.shape(phi[0])[-1]
#k = self.dims[0][0]
logdet_phi0 = linalg.chol_logdet(U)
u0 = phi[1][..., np.newaxis, np.newaxis] * linalg.chol_inv(U)
u1 = -logdet_phi0 + misc.multidigamma(phi[1], k)
u = [u0, u1]
g = phi[1] * logdet_phi0 - special.multigammaln(phi[1], k)
return (u, g)
