def check(plates_X, plates_mu=(), plates_Lambda=(), plates_B=(), plates_S=(), plates_v=()): D = 3 K = 2 N = 4 np.random.seed(42) mu = Gaussian(np.random.randn(*(plates_mu+(D,))), random.covariance(D)) Lambda = Wishart(D+np.ones(plates_Lambda), random.covariance(D)) B = GaussianARD(np.random.randn(*(plates_B+(D,D,K))), 1+np.random.rand(*(plates_B+(D,D,K))), shape=(D,K), plates=plates_B+(D,)) S = GaussianARD(np.random.randn(*(plates_S+(N,K))), 1+np.random.rand(*(plates_S+(N,K))), shape=(K,), plates=plates_S+(N,)) v = Gamma(1+np.random.rand(*(plates_v+(1,D))), 1+np.random.rand(*(plates_v+(1,D)))) X = VaryingGaussianMarkovChain(mu, Lambda, B, S, v, name="X") self.assertEqual(plates_X, X.plates, msg="Incorrect plates deduced") pass
def _run_checks(self, check): # Basic test check(1, 1, 1) check(2, 1, 1) check(1, 2, 1) check(1, 1, 2) check(3, 4, 2) # Test mu check(2, 3, 4, mu=GaussianARD(2, 4, shape=(2,), plates=())) # Test Lambda check(2, 3, 4, Lambda=Wishart(3, random.covariance(2))) # Test Lambda and mu check(2, 3, 4, mu=GaussianARD(2, 4, shape=(2,), plates=()), Lambda=Wishart(2, random.covariance(2))) # TODO: Test plates pass
def check(plates_X, plates_mu=(), plates_Lambda=(), plates_B=(), plates_S=(), plates_v=()): D = 3 K = 2 N = 4 np.random.seed(42) mu = Gaussian(np.random.randn(*(plates_mu + (D,))), random.covariance(D)) Lambda = Wishart(D + np.ones(plates_Lambda), random.covariance(D)) B = GaussianARD( np.random.randn(*(plates_B + (D, D, K))), 1 + np.random.rand(*(plates_B + (D, D, K))), shape=(D, K), plates=plates_B + (D,), ) S = GaussianARD( np.random.randn(*(plates_S + (N, K))), 1 + np.random.rand(*(plates_S + (N, K))), shape=(K,), plates=plates_S + (N,), ) v = Gamma(1 + np.random.rand(*(plates_v + (1, D))), 1 + np.random.rand(*(plates_v + (1, D)))) X = VaryingGaussianMarkovChain(mu, Lambda, B, S, v, name="X") self.assertEqual(plates_X, X.plates, msg="Incorrect plates deduced") pass
def test_message_to_parents(self): np.random.seed(42) N = 5 D1 = 3 D2 = 4 D3 = 2 X1 = Gaussian(np.random.randn(N, D1), random.covariance(D1)) X2 = Gaussian(np.random.randn(N, D2), random.covariance(D2)) X3 = np.random.randn(N, D3) Z = ConcatGaussian(X1, X2, X3) Y = Gaussian(Z, random.covariance(D1 + D2 + D3)) Y.observe(np.random.randn(*(Y.plates + Y.dims[0]))) self.assert_message_to_parent( Y, X1, eps=1e-7, rtol=1e-5, atol=1e-5 ) self.assert_message_to_parent( Y, X2, eps=1e-7, rtol=1e-5, atol=1e-5 ) pass
def _run_checks(self, check): # Basic test check(1, 1, 1) check(2, 1, 1) check(1, 2, 1) check(1, 1, 2) check(3, 4, 2) # Test mu check(2, 3, 4, mu=GaussianARD(2, 4, shape=(2, ), plates=())) # Test Lambda check(2, 3, 4, Lambda=Wishart(3, random.covariance(2))) # Test Lambda and mu check(2, 3, 4, mu=GaussianARD(2, 4, shape=(2, ), plates=()), Lambda=Wishart(2, random.covariance(2))) # TODO: Test plates pass
def test_riemannian_gradient(self): """Test Riemannian gradient of a Gaussian node.""" D = 3 # # Without observations # # Construct model mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda) # Random initialization mu0 = np.random.randn(D) Lambda0 = random.covariance(D) X.initialize_from_parameters(mu0, Lambda0) # Initial parameters phi0 = X.phi # Gradient g = X.get_riemannian_gradient() # Parameters after VB-EM update X.update() phi1 = X.phi # Check self.assertAllClose(g[0], phi1[0] - phi0[0]) self.assertAllClose(g[1], phi1[1] - phi0[1]) # TODO/FIXME: Actually, gradient should be zero because cost function # is zero without observations! Use the mask! # # With observations # # Construct model mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda) V = random.covariance(D) Y = Gaussian(X, V) Y.observe(np.random.randn(D)) # Random initialization mu0 = np.random.randn(D) Lambda0 = random.covariance(D) X.initialize_from_parameters(mu0, Lambda0) # Initial parameters phi0 = X.phi # Gradient g = X.get_riemannian_gradient() # Parameters after VB-EM update X.update() phi1 = X.phi # Check self.assertAllClose(g[0], phi1[0] - phi0[0]) self.assertAllClose(g[1], phi1[1] - phi0[1]) pass
def test_moments(self): """ Test the moments of Wishart node """ np.random.seed(42) # Test prior moments D = 3 n = (D-1) + np.random.uniform(0.1,2) V = random.covariance(D) Lambda = Wishart(n, V) Lambda.update() u = Lambda.get_moments() self.assertAllClose(u[0], n*np.linalg.inv(V), msg='Mean incorrect') self.assertAllClose(u[1], (np.sum(special.digamma((n - np.arange(D))/2)) + D*np.log(2) - np.linalg.slogdet(V)[1]), msg='Log determinant incorrect') # Test posterior moments D = 3 n = (D-1) + np.random.uniform(0.1,2) V = random.covariance(D) Lambda = Wishart(n, V) mu = np.random.randn(D) Y = Gaussian(mu, Lambda) y = np.random.randn(D) Y.observe(y) Lambda.update() u = Lambda.get_moments() n = n + 1 V = V + np.outer(y-mu, y-mu) self.assertAllClose(u[0], n*np.linalg.inv(V), msg='Mean incorrect') self.assertAllClose(u[1], (np.sum(special.digamma((n - np.arange(D))/2)) + D*np.log(2) - np.linalg.slogdet(V)[1]), msg='Log determinant incorrect') pass
def test_lower_bound(self): """ Test the Wishart VB lower bound """ # # By having the Wishart node as the only latent node, VB will give exact # results, thus the VB lower bound is the true marginal log likelihood. # Thus, check that they are equal. The true marginal likelihood is the # multivariate Student-t distribution. # np.random.seed(42) D = 3 n = (D-1) + np.random.uniform(0.1, 0.5) V = random.covariance(D) Lambda = Wishart(n, V) mu = np.random.randn(D) Y = Gaussian(mu, Lambda) y = np.random.randn(D) Y.observe(y) Lambda.update() L = Y.lower_bound_contribution() + Lambda.lower_bound_contribution() mu = mu nu = n + 1 - D Cov = V / nu self.assertAllClose(L, _student_logpdf(y, mu, Cov, nu)) pass
def check(N, D, plates=None, mu=None, Lambda=None, A=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if A is None: A = np.random.randn(D, D) if V is None: V = np.random.rand(D) X = GaussianMarkovChain(mu, Lambda, A, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = Gaussian._ensure_moments(mu, GaussianMoments, ndim=1).get_moments() (Lambda, _) = Wishart._ensure_moments(Lambda, WishartMoments, ndim=1).get_moments() (a, aa) = Gaussian._ensure_moments(A, GaussianMoments, ndim=1).get_moments() a = a * np.ones( (N - 1, D, D)) # explicit broadcasting for simplicity aa = aa * np.ones( (N - 1, D, D, D)) # explicit broadcasting for simplicity (v, _) = Gamma._ensure_moments(V, GammaMoments).get_moments() v = v * np.ones((N - 1, D)) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N, D, N, D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N, D)) m[..., 0, :] = np.einsum('...ij,...j->...i', Lambda, mu) C[..., 0, :, 0, :] = Lambda + np.einsum( '...dij,...d->...ij', aa[..., 0, :, :, :], v[..., 0, :]) for n in range(1, N - 1): C[..., n, :, n, :] = (np.einsum( '...dij,...d->...ij', aa[..., n, :, :, :], v[..., n, :]) + v[..., n, :, None] * np.identity(D)) for n in range(N - 1): C[..., n, :, n + 1, :] = -np.einsum( '...di,...d->...id', a[..., n, :, :], v[..., n, :]) C[..., n + 1, :, n, :] = -np.einsum( '...di,...d->...di', a[..., n, :, :], v[..., n, :]) C[..., -1, :, -1, :] = v[..., -1, :, None] * np.identity(D) C = np.reshape(C, plates_C + (N * D, N * D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C + (N, D, N, D)) m0 = np.einsum('...minj,...nj->...mi', Cov, m) m1 = np.zeros(plates_C + (N, D, D)) m2 = np.zeros(plates_C + (N - 1, D, D)) for n in range(N): m1[..., n, :, :] = Cov[..., n, :, n, :] + np.einsum( '...i,...j->...ij', m0[..., n, :], m0[..., n, :]) for n in range(N - 1): m2[..., n, :, :] = Cov[..., n, :, n + 1, :] + np.einsum( '...i,...j->...ij', m0[..., n, :], m0[..., n + 1, :]) self.assertAllClose(m0, u0 * np.ones(np.shape(m0))) self.assertAllClose(m1, u1 * np.ones(np.shape(m1))) self.assertAllClose(m2, u2 * np.ones(np.shape(m2))) pass
def test_moments(self): np.random.seed(42) N = 4 D1 = 2 D2 = 3 X1 = Gaussian(np.random.randn(N, D1), random.covariance(D1)) X2 = Gaussian(np.random.randn(N, D2), random.covariance(D2)) Z = ConcatGaussian(X1, X2) u = Z._message_to_child() # First moment self.assertAllClose( u[0][...,:D1], X1.u[0] ) self.assertAllClose( u[0][...,D1:], X2.u[0] ) # Second moment self.assertAllClose( u[1][...,:D1,:D1], X1.u[1] ) self.assertAllClose( u[1][...,D1:,D1:], X2.u[1] ) self.assertAllClose( u[1][...,:D1,D1:], X1.u[0][...,:,None] * X2.u[0][...,None,:] ) self.assertAllClose( u[1][...,D1:,:D1], X2.u[0][...,:,None] * X1.u[0][...,None,:] ) pass
def check(N, D, K, plates=None, mu=None, Lambda=None, B=None, S=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if B is None: B = np.random.randn(D, D, K) if S is None: S = np.random.randn(N - 1, K) if V is None: V = np.random.rand(D) X = VaryingGaussianMarkovChain(mu, Lambda, B, S, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = X.parents[0].get_moments() (Lambda, _) = X.parents[1].get_moments() (b, bb) = X.parents[2].get_moments() (s, ss) = X.parents[3].get_moments() (v, _) = X.parents[4].get_moments() v = v * np.ones((N - 1, D)) # V = np.atleast_3d(v)[...,-1,:,None]*np.identity(D) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N, D, N, D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N, D)) m[..., 0, :] = np.einsum("...ij,...j->...i", Lambda, mu) # m = np.reshape(m, plates_mu + (N*D,)) A = np.einsum("...dik,...nk->...ndi", b, s) AA = np.einsum("...dikjl,...nkl->...ndij", bb, ss) C[..., 0, :, 0, :] = Lambda + np.einsum("...dij,...d->...ij", AA[..., 0, :, :, :], v[..., 0, :]) for n in range(1, N - 1): C[..., n, :, n, :] = np.einsum("...dij,...d->...ij", AA[..., n, :, :, :], v[..., n, :]) + v[ ..., n, :, None ] * np.identity(D) for n in range(N - 1): C[..., n, :, n + 1, :] = -np.einsum("...di,...d->...id", A[..., n, :, :], v[..., n, :]) C[..., n + 1, :, n, :] = -np.einsum("...di,...d->...di", A[..., n, :, :], v[..., n, :]) C[..., -1, :, -1, :] = v[..., -1, :, None] * np.identity(D) C = np.reshape(C, plates_C + (N * D, N * D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C + (N, D, N, D)) m0 = np.einsum("...minj,...nj->...mi", Cov, m) m1 = np.zeros(plates_C + (N, D, D)) m2 = np.zeros(plates_C + (N - 1, D, D)) for n in range(N): m1[..., n, :, :] = Cov[..., n, :, n, :] + np.einsum("...i,...j->...ij", m0[..., n, :], m0[..., n, :]) for n in range(N - 1): m2[..., n, :, :] = Cov[..., n, :, n + 1, :] + np.einsum( "...i,...j->...ij", m0[..., n, :], m0[..., n + 1, :] ) self.assertAllClose(m0, u0 * np.ones(np.shape(m0))) self.assertAllClose(m1, u1 * np.ones(np.shape(m1))) self.assertAllClose(m2, u2 * np.ones(np.shape(m2))) pass
def create_model(self, N, D): # Construct the model Mu = Gaussian(np.random.randn(D), np.identity(D)) Lambda = Wishart(D, random.covariance(D)) A = Gaussian(np.random.randn(D, D), np.identity(D)) V = Gamma(D, np.random.rand(D)) X = GaussianMarkovChain(Mu, Lambda, A, V, n=N) Y = Gaussian(X, np.identity(D)) return (Y, X, Mu, Lambda, A, V)
def test_message_to_parents(self): """ Check gradient passed to inputs parent node """ D = 3 X = Gaussian(np.random.randn(D), random.covariance(D)) V = Wishart(D + np.random.rand(), random.covariance(D)) Y = Gaussian(X, V) self.assert_moments( Y, lambda u: [u[0], u[1] + u[1].T] ) Y.observe(np.random.randn(D)) self.assert_message_to_parent(Y, X) #self.assert_message_to_parent(Y, V) pass
def test_message_to_parents(self): """ Check gradient passed to inputs parent node """ N = 3 D = 2 Mu = Gaussian(np.random.randn(D), random.covariance(D)) Lambda = Wishart(D, random.covariance(D)) A = Gaussian(np.random.randn(D,D), random.covariance(D)) V = Gamma(D, np.random.rand(D)) X = GaussianMarkovChain(Mu, Lambda, A, V, n=N+1) Y = Gaussian(X, random.covariance(D)) self.assert_moments( X, postprocess=lambda u: [ u[0], u[1] + linalg.transpose(u[1], ndim=1), u[2] ] ) Y.observe(np.random.randn(N+1, D)) self.assert_message_to_parent(X, Mu, eps=1e-8) self.assert_message_to_parent( X, Lambda, eps=1e-8, postprocess=lambda u: [ u[0] + linalg.transpose(u[0], ndim=1), u[1], ] ) self.assert_message_to_parent(X, A) self.assert_message_to_parent(X, V, eps=1e-10, atol=1e-5) pass
def test_message_to_parents(self): """ Check gradient passed to inputs parent node """ D = 3 X = Gaussian(np.random.randn(D), random.covariance(D)) a = Gamma(np.random.rand(D), np.random.rand(D)) Y = GaussianARD(X, a) Y.observe(np.random.randn(D)) self.assert_message_to_parent(Y, X) self.assert_message_to_parent(Y, a) pass
def check(N, D, plates=None, mu=None, Lambda=None, A=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if A is None: A = np.random.randn(D, D) if V is None: V = np.random.rand(D) X = GaussianMarkovChain(mu, Lambda, A, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = Gaussian._ensure_moments(mu, GaussianMoments, ndim=1).get_moments() (Lambda, _) = Wishart._ensure_moments(Lambda, WishartMoments, ndim=1).get_moments() (a, aa) = Gaussian._ensure_moments(A, GaussianMoments, ndim=1).get_moments() a = a * np.ones((N - 1, D, D)) # explicit broadcasting for simplicity aa = aa * np.ones((N - 1, D, D, D)) # explicit broadcasting for simplicity (v, _) = Gamma._ensure_moments(V, GammaMoments).get_moments() v = v * np.ones((N - 1, D)) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N, D, N, D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N, D)) m[..., 0, :] = np.einsum("...ij,...j->...i", Lambda, mu) C[..., 0, :, 0, :] = Lambda + np.einsum("...dij,...d->...ij", aa[..., 0, :, :, :], v[..., 0, :]) for n in range(1, N - 1): C[..., n, :, n, :] = np.einsum("...dij,...d->...ij", aa[..., n, :, :, :], v[..., n, :]) + v[ ..., n, :, None ] * np.identity(D) for n in range(N - 1): C[..., n, :, n + 1, :] = -np.einsum("...di,...d->...id", a[..., n, :, :], v[..., n, :]) C[..., n + 1, :, n, :] = -np.einsum("...di,...d->...di", a[..., n, :, :], v[..., n, :]) C[..., -1, :, -1, :] = v[..., -1, :, None] * np.identity(D) C = np.reshape(C, plates_C + (N * D, N * D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C + (N, D, N, D)) m0 = np.einsum("...minj,...nj->...mi", Cov, m) m1 = np.zeros(plates_C + (N, D, D)) m2 = np.zeros(plates_C + (N - 1, D, D)) for n in range(N): m1[..., n, :, :] = Cov[..., n, :, n, :] + np.einsum("...i,...j->...ij", m0[..., n, :], m0[..., n, :]) for n in range(N - 1): m2[..., n, :, :] = Cov[..., n, :, n + 1, :] + np.einsum( "...i,...j->...ij", m0[..., n, :], m0[..., n + 1, :] ) self.assertAllClose(m0, u0 * np.ones(np.shape(m0))) self.assertAllClose(m1, u1 * np.ones(np.shape(m1))) self.assertAllClose(m2, u2 * np.ones(np.shape(m2))) pass
def generate_data(N, D, K, seed=1, spread=3): """ Generate data from a mixture of Gaussians model """ np.random.seed(seed) mu = spread*np.random.randn(K, D) # Lambda is actually precision matrix (inverse covariance) Lambda = random.covariance(D, size=K, nu=2*D) pi = random.dirichlet(5*np.ones(K)) y = np.zeros((N,D)) for n in range(N): ind = nodes.Categorical(pi).random() y[n] = nodes.Gaussian(mu[ind], Lambda[ind]).random() np.savetxt('mog-data-%02d.csv' % seed, y, delimiter=',', fmt='%f') return y
def check(Mu, Lambda, A, V, U): X = GaussianMarkovChain(Mu, Lambda, A, V, inputs=U) Y = Gaussian(X, random.covariance(D)) # Check moments self.assert_moments( X, postprocess=lambda u: [ u[0], u[1] + linalg.transpose(u[1], ndim=1), u[2] ] ) Y.observe(np.random.randn(N+1, D)) X.update() # Check gradient messages to parents self.assert_message_to_parent(X, Mu) self.assert_message_to_parent( X, Lambda, postprocess=lambda phi: [ phi[0] + linalg.transpose(phi[0], ndim=1), phi[1] ] ) self.assert_message_to_parent( X, A, postprocess=lambda phi: [ phi[0], phi[1] + linalg.transpose(phi[1], ndim=1), ] ) self.assert_message_to_parent(X, V) self.assert_message_to_parent(X, U)
def _run_checks(self, check): # Basic test check(2, 3) # Test mu check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=())) check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=(5,))) # Test Lambda check(2, 3, Lambda=Wishart(3, random.covariance(2))) check(2, 3, Lambda=Wishart(3, random.covariance(2), plates=(5,))) # Test A check(2, 3, A=GaussianARD(2, 4, shape=(2,), plates=(2,))) check(2, 3, A=GaussianARD(2, 4, shape=(2,), plates=(3,2))) check(2, 3, A=GaussianARD(2, 4, shape=(2,), plates=(5,3,2))) # Test Lambda and mu check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=()), Lambda=Wishart(2, random.covariance(2))) check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=(5,)), Lambda=Wishart(2, random.covariance(2), plates=(5,))) # Test mu and A check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=()), A=GaussianARD(2, 4, shape=(2,), plates=(2,))) check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=(5,)), A=GaussianARD(2, 4, shape=(2,), plates=(5,1,2,))) # Test Lambda and A check(2, 3, Lambda=Wishart(2, random.covariance(2)), A=GaussianARD(2, 4, shape=(2,), plates=(2,))) check(2, 3, Lambda=Wishart(2, random.covariance(2), plates=(5,)), A=GaussianARD(2, 4, shape=(2,), plates=(5,1,2,))) # Test mu, Lambda and A check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=()), Lambda=Wishart(2, random.covariance(2)), A=GaussianARD(2, 4, shape=(2,), plates=(2,))) check(2, 3, mu=GaussianARD(2, 4, shape=(2,), plates=(5,)), Lambda=Wishart(2, random.covariance(2), plates=(5,)), A=GaussianARD(2, 4, shape=(2,), plates=(5,1,2,))) pass
def test_message_to_child(self): # A very simple check before the more complex ones: # 1-D process, k=1, fixed constant parameters m = 1.0 l = 4.0 b = 2.0 s = [3.0, 8.0] v = 5.0 X = VaryingGaussianMarkovChain([m], [[l]], [[[b]]], [[s[0]],[s[1]]], [v]) (u0, u1, u2) = X._message_to_child() C = np.array([[l+b**2*s[0]**2*v, -b*s[0]*v, 0], [ -b*s[0]*v, v+b**2*s[1]**2*v, -b*s[1]*v], [ 0, -b*s[1]*v, v]]) Cov = np.linalg.inv(C) m0 = np.dot(Cov, [[l*m], [0], [0]]) m1 = np.diag(Cov)[:,None,None] + m0[:,:,None]**2 m2 = np.diag(Cov, k=1)[:,None,None] + m0[1:,:,None]*m0[:-1,:,None] self.assertAllClose(m0, u0) self.assertAllClose(m1, u1) self.assertAllClose(m2, u2) def check(N, D, K, plates=None, mu=None, Lambda=None, B=None, S=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if B is None: B = np.random.randn(D,D,K) if S is None: S = np.random.randn(N-1,K) if V is None: V = np.random.rand(D) X = VaryingGaussianMarkovChain(mu, Lambda, B, S, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = X.parents[0].get_moments() (Lambda, _) = X.parents[1].get_moments() (b, bb) = X.parents[2].get_moments() (s, ss) = X.parents[3].get_moments() (v, _) = X.parents[4].get_moments() v = v * np.ones((N-1,D)) #V = np.atleast_3d(v)[...,-1,:,None]*np.identity(D) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N,D,N,D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N,D)) m[...,0,:] = np.einsum('...ij,...j->...i', Lambda, mu) #m = np.reshape(m, plates_mu + (N*D,)) A = np.einsum('...dik,...nk->...ndi', b, s) AA = np.einsum('...dikjl,...nkl->...ndij', bb, ss) C[...,0,:,0,:] = Lambda + np.einsum('...dij,...d->...ij', AA[...,0,:,:,:], v[...,0,:]) for n in range(1,N-1): C[...,n,:,n,:] = (np.einsum('...dij,...d->...ij', AA[...,n,:,:,:], v[...,n,:]) + v[...,n,:,None] * np.identity(D)) for n in range(N-1): C[...,n,:,n+1,:] = -np.einsum('...di,...d->...id', A[...,n,:,:], v[...,n,:]) C[...,n+1,:,n,:] = -np.einsum('...di,...d->...di', A[...,n,:,:], v[...,n,:]) C[...,-1,:,-1,:] = v[...,-1,:,None]*np.identity(D) C = np.reshape(C, plates_C+(N*D,N*D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C+(N,D,N,D)) m0 = np.einsum('...minj,...nj->...mi', Cov, m) m1 = np.zeros(plates_C+(N,D,D)) m2 = np.zeros(plates_C+(N-1,D,D)) for n in range(N): m1[...,n,:,:] = Cov[...,n,:,n,:] + np.einsum('...i,...j->...ij', m0[...,n,:], m0[...,n,:]) for n in range(N-1): m2[...,n,:,:] = Cov[...,n,:,n+1,:] + np.einsum('...i,...j->...ij', m0[...,n,:], m0[...,n+1,:]) self.assertAllClose(m0, u0*np.ones(np.shape(m0))) self.assertAllClose(m1, u1*np.ones(np.shape(m1))) self.assertAllClose(m2, u2*np.ones(np.shape(m2))) pass check(2,1,1) check(2,3,1) check(2,1,3) check(4,3,2) # # Test mu # # Simple check(4,3,2, mu=Gaussian(np.random.randn(3), random.covariance(3))) # Plates check(4,3,2, mu=Gaussian(np.random.randn(5,6,3), random.covariance(3), plates=(5,6))) # Plates with moments broadcasted over plates check(4,3,2, mu=Gaussian(np.random.randn(3), random.covariance(3), plates=(5,))) check(4,3,2, mu=Gaussian(np.random.randn(1,3), random.covariance(3), plates=(5,))) # Plates broadcasting check(4,3,2, plates=(5,), mu=Gaussian(np.random.randn(3), random.covariance(3), plates=())) check(4,3,2, plates=(5,), mu=Gaussian(np.random.randn(1,3), random.covariance(3), plates=(1,))) # # Test Lambda # # Simple check(4,3,2, Lambda=Wishart(10+np.random.rand(), random.covariance(3))) # Plates check(4,3,2, Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=(5,6))) # Plates with moments broadcasted over plates check(4,3,2, Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=(5,))) check(4,3,2, Lambda=Wishart(10+np.random.rand(1), random.covariance(3), plates=(5,))) # Plates broadcasting check(4,3,2, plates=(5,), Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=())) check(4,3,2, plates=(5,), Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=(1,))) # # Test B # # Simple check(4,3,2, B=GaussianARD(np.random.randn(3,3,2), np.random.rand(3,3,2), shape=(3,2), plates=(3,))) # Plates check(4,3,2, B=GaussianARD(np.random.randn(5,6,3,3,2), np.random.rand(5,6,3,3,2), shape=(3,2), plates=(5,6,3))) # Plates with moments broadcasted over plates check(4,3,2, B=GaussianARD(np.random.randn(3,3,2), np.random.rand(3,3,2), shape=(3,2), plates=(5,3))) check(4,3,2, B=GaussianARD(np.random.randn(1,3,3,2), np.random.rand(1,3,3,2), shape=(3,2), plates=(5,3))) # Plates broadcasting check(4,3,2, plates=(5,), B=GaussianARD(np.random.randn(3,3,2), np.random.rand(3,3,2), shape=(3,2), plates=(3,))) check(4,3,2, plates=(5,), B=GaussianARD(np.random.randn(3,3,2), np.random.rand(3,3,2), shape=(3,2), plates=(1,3))) # # Test S # # Simple check(4,3,2, S=GaussianARD(np.random.randn(4-1,2), np.random.rand(4-1,2), shape=(2,), plates=(4-1,))) # Plates check(4,3,2, S=GaussianARD(np.random.randn(5,6,4-1,2), np.random.rand(5,6,4-1,2), shape=(2,), plates=(5,6,4-1,))) # Plates with moments broadcasted over plates check(4,3,2, S=GaussianARD(np.random.randn(4-1,2), np.random.rand(4-1,2), shape=(2,), plates=(5,4-1,))) check(4,3,2, S=GaussianARD(np.random.randn(1,4-1,2), np.random.rand(1,4-1,2), shape=(2,), plates=(5,4-1,))) # Plates broadcasting check(4,3,2, plates=(5,), S=GaussianARD(np.random.randn(4-1,2), np.random.rand(4-1,2), shape=(2,), plates=(4-1,))) check(4,3,2, plates=(5,), S=GaussianARD(np.random.randn(4-1,2), np.random.rand(4-1,2), shape=(2,), plates=(1,4-1,))) # # Test v # # Simple check(4,3,2, V=Gamma(np.random.rand(1,3), np.random.rand(1,3), plates=(1,3))) check(4,3,2, V=Gamma(np.random.rand(3), np.random.rand(3), plates=(3,))) # Plates check(4,3,2, V=Gamma(np.random.rand(5,6,1,3), np.random.rand(5,6,1,3), plates=(5,6,1,3))) # Plates with moments broadcasted over plates check(4,3,2, V=Gamma(np.random.rand(1,3), np.random.rand(1,3), plates=(5,1,3))) check(4,3,2, V=Gamma(np.random.rand(1,1,3), np.random.rand(1,1,3), plates=(5,1,3))) # Plates broadcasting check(4,3,2, plates=(5,), V=Gamma(np.random.rand(1,3), np.random.rand(1,3), plates=(1,3))) check(4,3,2, plates=(5,), V=Gamma(np.random.rand(1,1,3), np.random.rand(1,1,3), plates=(1,1,3))) # # Uncertainty in both B and S # check(4,3,2, B=GaussianARD(np.random.randn(3,3,2), np.random.rand(3,3,2), shape=(3,2), plates=(3,)), S=GaussianARD(np.random.randn(4-1,2), np.random.rand(4-1,2), shape=(2,), plates=(4-1,))) pass
def test_message_to_child(self): # A very simple check before the more complex ones: # 1-D process, k=1, fixed constant parameters m = 1.0 l = 4.0 b = 2.0 s = [3.0, 8.0] v = 5.0 X = VaryingGaussianMarkovChain([m], [[l]], [[[b]]], [[s[0]], [s[1]]], [v]) (u0, u1, u2) = X._message_to_child() C = np.array( [ [l + b ** 2 * s[0] ** 2 * v, -b * s[0] * v, 0], [-b * s[0] * v, v + b ** 2 * s[1] ** 2 * v, -b * s[1] * v], [0, -b * s[1] * v, v], ] ) Cov = np.linalg.inv(C) m0 = np.dot(Cov, [[l * m], [0], [0]]) m1 = np.diag(Cov)[:, None, None] + m0[:, :, None] ** 2 m2 = np.diag(Cov, k=1)[:, None, None] + m0[1:, :, None] * m0[:-1, :, None] self.assertAllClose(m0, u0) self.assertAllClose(m1, u1) self.assertAllClose(m2, u2) def check(N, D, K, plates=None, mu=None, Lambda=None, B=None, S=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if B is None: B = np.random.randn(D, D, K) if S is None: S = np.random.randn(N - 1, K) if V is None: V = np.random.rand(D) X = VaryingGaussianMarkovChain(mu, Lambda, B, S, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = X.parents[0].get_moments() (Lambda, _) = X.parents[1].get_moments() (b, bb) = X.parents[2].get_moments() (s, ss) = X.parents[3].get_moments() (v, _) = X.parents[4].get_moments() v = v * np.ones((N - 1, D)) # V = np.atleast_3d(v)[...,-1,:,None]*np.identity(D) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N, D, N, D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N, D)) m[..., 0, :] = np.einsum("...ij,...j->...i", Lambda, mu) # m = np.reshape(m, plates_mu + (N*D,)) A = np.einsum("...dik,...nk->...ndi", b, s) AA = np.einsum("...dikjl,...nkl->...ndij", bb, ss) C[..., 0, :, 0, :] = Lambda + np.einsum("...dij,...d->...ij", AA[..., 0, :, :, :], v[..., 0, :]) for n in range(1, N - 1): C[..., n, :, n, :] = np.einsum("...dij,...d->...ij", AA[..., n, :, :, :], v[..., n, :]) + v[ ..., n, :, None ] * np.identity(D) for n in range(N - 1): C[..., n, :, n + 1, :] = -np.einsum("...di,...d->...id", A[..., n, :, :], v[..., n, :]) C[..., n + 1, :, n, :] = -np.einsum("...di,...d->...di", A[..., n, :, :], v[..., n, :]) C[..., -1, :, -1, :] = v[..., -1, :, None] * np.identity(D) C = np.reshape(C, plates_C + (N * D, N * D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C + (N, D, N, D)) m0 = np.einsum("...minj,...nj->...mi", Cov, m) m1 = np.zeros(plates_C + (N, D, D)) m2 = np.zeros(plates_C + (N - 1, D, D)) for n in range(N): m1[..., n, :, :] = Cov[..., n, :, n, :] + np.einsum("...i,...j->...ij", m0[..., n, :], m0[..., n, :]) for n in range(N - 1): m2[..., n, :, :] = Cov[..., n, :, n + 1, :] + np.einsum( "...i,...j->...ij", m0[..., n, :], m0[..., n + 1, :] ) self.assertAllClose(m0, u0 * np.ones(np.shape(m0))) self.assertAllClose(m1, u1 * np.ones(np.shape(m1))) self.assertAllClose(m2, u2 * np.ones(np.shape(m2))) pass check(2, 1, 1) check(2, 3, 1) check(2, 1, 3) check(4, 3, 2) # # Test mu # # Simple check(4, 3, 2, mu=Gaussian(np.random.randn(3), random.covariance(3))) # Plates check(4, 3, 2, mu=Gaussian(np.random.randn(5, 6, 3), random.covariance(3), plates=(5, 6))) # Plates with moments broadcasted over plates check(4, 3, 2, mu=Gaussian(np.random.randn(3), random.covariance(3), plates=(5,))) check(4, 3, 2, mu=Gaussian(np.random.randn(1, 3), random.covariance(3), plates=(5,))) # Plates broadcasting check(4, 3, 2, plates=(5,), mu=Gaussian(np.random.randn(3), random.covariance(3), plates=())) check(4, 3, 2, plates=(5,), mu=Gaussian(np.random.randn(1, 3), random.covariance(3), plates=(1,))) # # Test Lambda # # Simple check(4, 3, 2, Lambda=Wishart(10 + np.random.rand(), random.covariance(3))) # Plates check(4, 3, 2, Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=(5, 6))) # Plates with moments broadcasted over plates check(4, 3, 2, Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=(5,))) check(4, 3, 2, Lambda=Wishart(10 + np.random.rand(1), random.covariance(3), plates=(5,))) # Plates broadcasting check(4, 3, 2, plates=(5,), Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=())) check(4, 3, 2, plates=(5,), Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=(1,))) # # Test B # # Simple check(4, 3, 2, B=GaussianARD(np.random.randn(3, 3, 2), np.random.rand(3, 3, 2), shape=(3, 2), plates=(3,))) # Plates check( 4, 3, 2, B=GaussianARD( np.random.randn(5, 6, 3, 3, 2), np.random.rand(5, 6, 3, 3, 2), shape=(3, 2), plates=(5, 6, 3) ), ) # Plates with moments broadcasted over plates check(4, 3, 2, B=GaussianARD(np.random.randn(3, 3, 2), np.random.rand(3, 3, 2), shape=(3, 2), plates=(5, 3))) check( 4, 3, 2, B=GaussianARD(np.random.randn(1, 3, 3, 2), np.random.rand(1, 3, 3, 2), shape=(3, 2), plates=(5, 3)) ) # Plates broadcasting check( 4, 3, 2, plates=(5,), B=GaussianARD(np.random.randn(3, 3, 2), np.random.rand(3, 3, 2), shape=(3, 2), plates=(3,)), ) check( 4, 3, 2, plates=(5,), B=GaussianARD(np.random.randn(3, 3, 2), np.random.rand(3, 3, 2), shape=(3, 2), plates=(1, 3)), ) # # Test S # # Simple check(4, 3, 2, S=GaussianARD(np.random.randn(4 - 1, 2), np.random.rand(4 - 1, 2), shape=(2,), plates=(4 - 1,))) # Plates check( 4, 3, 2, S=GaussianARD( np.random.randn(5, 6, 4 - 1, 2), np.random.rand(5, 6, 4 - 1, 2), shape=(2,), plates=(5, 6, 4 - 1) ), ) # Plates with moments broadcasted over plates check( 4, 3, 2, S=GaussianARD(np.random.randn(4 - 1, 2), np.random.rand(4 - 1, 2), shape=(2,), plates=(5, 4 - 1)) ) check( 4, 3, 2, S=GaussianARD(np.random.randn(1, 4 - 1, 2), np.random.rand(1, 4 - 1, 2), shape=(2,), plates=(5, 4 - 1)), ) # Plates broadcasting check( 4, 3, 2, plates=(5,), S=GaussianARD(np.random.randn(4 - 1, 2), np.random.rand(4 - 1, 2), shape=(2,), plates=(4 - 1,)), ) check( 4, 3, 2, plates=(5,), S=GaussianARD(np.random.randn(4 - 1, 2), np.random.rand(4 - 1, 2), shape=(2,), plates=(1, 4 - 1)), ) # # Test v # # Simple check(4, 3, 2, V=Gamma(np.random.rand(1, 3), np.random.rand(1, 3), plates=(1, 3))) check(4, 3, 2, V=Gamma(np.random.rand(3), np.random.rand(3), plates=(3,))) # Plates check(4, 3, 2, V=Gamma(np.random.rand(5, 6, 1, 3), np.random.rand(5, 6, 1, 3), plates=(5, 6, 1, 3))) # Plates with moments broadcasted over plates check(4, 3, 2, V=Gamma(np.random.rand(1, 3), np.random.rand(1, 3), plates=(5, 1, 3))) check(4, 3, 2, V=Gamma(np.random.rand(1, 1, 3), np.random.rand(1, 1, 3), plates=(5, 1, 3))) # Plates broadcasting check(4, 3, 2, plates=(5,), V=Gamma(np.random.rand(1, 3), np.random.rand(1, 3), plates=(1, 3))) check(4, 3, 2, plates=(5,), V=Gamma(np.random.rand(1, 1, 3), np.random.rand(1, 1, 3), plates=(1, 1, 3))) # # Uncertainty in both B and S # check( 4, 3, 2, B=GaussianARD(np.random.randn(3, 3, 2), np.random.rand(3, 3, 2), shape=(3, 2), plates=(3,)), S=GaussianARD(np.random.randn(4 - 1, 2), np.random.rand(4 - 1, 2), shape=(2,), plates=(4 - 1,)), ) pass
def test_message_to_child(self): """ Test the updating of GaussianMarkovChain. Check that the moments and the lower bound contribution are computed correctly. """ # TODO: Add plates and missing values! # Dimensionalities D = 3 N = 5 (Y, X, Mu, Lambda, A, V) = self.create_model(N, D) # Inference with arbitrary observations y = np.random.randn(N, D) Y.observe(y) X.update() (x_vb, xnxn_vb, xpxn_vb) = X.get_moments() # Get parameter moments (mu0, mumu0) = Mu.get_moments() (icov0, logdet0) = Lambda.get_moments() (a, aa) = A.get_moments() (icov_x, logdetx) = V.get_moments() icov_x = np.diag(icov_x) # Prior precision Z = np.einsum("...kij,...kk->...ij", aa, icov_x) U_diag = [icov0 + Z] + (N - 2) * [icov_x + Z] + [icov_x] U_super = (N - 1) * [-np.dot(a.T, icov_x)] U = misc.block_banded(U_diag, U_super) # Prior mean mu_prior = np.zeros(D * N) mu_prior[:D] = np.dot(icov0, mu0) # Data Cov = np.linalg.inv(U + np.identity(D * N)) mu = np.dot(Cov, mu_prior + y.flatten()) # Moments xx = mu[:, np.newaxis] * mu[np.newaxis, :] + Cov mu = np.reshape(mu, (N, D)) xx = np.reshape(xx, (N, D, N, D)) # Check results self.assertAllClose(x_vb, mu, msg="Incorrect mean") for n in range(N): self.assertAllClose(xnxn_vb[n, :, :], xx[n, :, n, :], msg="Incorrect second moment") for n in range(N - 1): self.assertAllClose(xpxn_vb[n, :, :], xx[n, :, n + 1, :], msg="Incorrect lagged second moment") # Compute the entropy H(X) ldet = linalg.logdet_cov(Cov) H = random.gaussian_entropy(-ldet, N * D) # Compute <log p(X|...)> xx = np.reshape(xx, (N * D, N * D)) mu = np.reshape(mu, (N * D,)) ldet = -logdet0 - np.sum(np.ones((N - 1, D)) * logdetx) P = random.gaussian_logpdf( np.einsum("...ij,...ij", xx, U), np.einsum("...i,...i", mu, mu_prior), np.einsum("...ij,...ij", mumu0, icov0), -ldet, N * D, ) # The VB bound from the net l = X.lower_bound_contribution() self.assertAllClose(l, H + P) # Compute the true bound <log p(X|...)> + H(X) # # Simple tests # def check(N, D, plates=None, mu=None, Lambda=None, A=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if A is None: A = np.random.randn(D, D) if V is None: V = np.random.rand(D) X = GaussianMarkovChain(mu, Lambda, A, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = Gaussian._ensure_moments(mu, GaussianMoments, ndim=1).get_moments() (Lambda, _) = Wishart._ensure_moments(Lambda, WishartMoments, ndim=1).get_moments() (a, aa) = Gaussian._ensure_moments(A, GaussianMoments, ndim=1).get_moments() a = a * np.ones((N - 1, D, D)) # explicit broadcasting for simplicity aa = aa * np.ones((N - 1, D, D, D)) # explicit broadcasting for simplicity (v, _) = Gamma._ensure_moments(V, GammaMoments).get_moments() v = v * np.ones((N - 1, D)) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N, D, N, D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N, D)) m[..., 0, :] = np.einsum("...ij,...j->...i", Lambda, mu) C[..., 0, :, 0, :] = Lambda + np.einsum("...dij,...d->...ij", aa[..., 0, :, :, :], v[..., 0, :]) for n in range(1, N - 1): C[..., n, :, n, :] = np.einsum("...dij,...d->...ij", aa[..., n, :, :, :], v[..., n, :]) + v[ ..., n, :, None ] * np.identity(D) for n in range(N - 1): C[..., n, :, n + 1, :] = -np.einsum("...di,...d->...id", a[..., n, :, :], v[..., n, :]) C[..., n + 1, :, n, :] = -np.einsum("...di,...d->...di", a[..., n, :, :], v[..., n, :]) C[..., -1, :, -1, :] = v[..., -1, :, None] * np.identity(D) C = np.reshape(C, plates_C + (N * D, N * D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C + (N, D, N, D)) m0 = np.einsum("...minj,...nj->...mi", Cov, m) m1 = np.zeros(plates_C + (N, D, D)) m2 = np.zeros(plates_C + (N - 1, D, D)) for n in range(N): m1[..., n, :, :] = Cov[..., n, :, n, :] + np.einsum("...i,...j->...ij", m0[..., n, :], m0[..., n, :]) for n in range(N - 1): m2[..., n, :, :] = Cov[..., n, :, n + 1, :] + np.einsum( "...i,...j->...ij", m0[..., n, :], m0[..., n + 1, :] ) self.assertAllClose(m0, u0 * np.ones(np.shape(m0))) self.assertAllClose(m1, u1 * np.ones(np.shape(m1))) self.assertAllClose(m2, u2 * np.ones(np.shape(m2))) pass check(4, 1) check(4, 3) # # Test mu # # Simple check(4, 3, mu=Gaussian(np.random.randn(3), random.covariance(3))) # Plates check(4, 3, mu=Gaussian(np.random.randn(5, 6, 3), random.covariance(3), plates=(5, 6))) # Plates with moments broadcasted over plates check(4, 3, mu=Gaussian(np.random.randn(3), random.covariance(3), plates=(5,))) check(4, 3, mu=Gaussian(np.random.randn(1, 3), random.covariance(3), plates=(5,))) # Plates broadcasting check(4, 3, plates=(5,), mu=Gaussian(np.random.randn(3), random.covariance(3), plates=())) check(4, 3, plates=(5,), mu=Gaussian(np.random.randn(1, 3), random.covariance(3), plates=(1,))) # # Test Lambda # # Simple check(4, 3, Lambda=Wishart(10 + np.random.rand(), random.covariance(3))) # Plates check(4, 3, Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=(5, 6))) # Plates with moments broadcasted over plates check(4, 3, Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=(5,))) check(4, 3, Lambda=Wishart(10 + np.random.rand(1), random.covariance(3), plates=(5,))) # Plates broadcasting check(4, 3, plates=(5,), Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=())) check(4, 3, plates=(5,), Lambda=Wishart(10 + np.random.rand(), random.covariance(3), plates=(1,))) # # Test A # # Simple check(4, 3, A=GaussianARD(np.random.randn(3, 3), np.random.rand(3, 3), shape=(3,), plates=(3,))) # Plates on time axis check(5, 3, A=GaussianARD(np.random.randn(4, 3, 3), np.random.rand(4, 3, 3), shape=(3,), plates=(4, 3))) # Plates on time axis with broadcasted moments check(5, 3, A=GaussianARD(np.random.randn(1, 3, 3), np.random.rand(1, 3, 3), shape=(3,), plates=(4, 3))) check(5, 3, A=GaussianARD(np.random.randn(3, 3), np.random.rand(3, 3), shape=(3,), plates=(4, 3))) # Plates check( 4, 3, A=GaussianARD( np.random.randn(5, 6, 1, 3, 3), np.random.rand(5, 6, 1, 3, 3), shape=(3,), plates=(5, 6, 1, 3) ), ) # Plates with moments broadcasted over plates check(4, 3, A=GaussianARD(np.random.randn(3, 3), np.random.rand(3, 3), shape=(3,), plates=(5, 1, 3))) check( 4, 3, A=GaussianARD(np.random.randn(1, 1, 3, 3), np.random.rand(1, 1, 3, 3), shape=(3,), plates=(5, 1, 3)) ) # Plates broadcasting check(4, 3, plates=(5,), A=GaussianARD(np.random.randn(3, 3), np.random.rand(3, 3), shape=(3,), plates=(3,))) check( 4, 3, plates=(5,), A=GaussianARD(np.random.randn(3, 3), np.random.rand(3, 3), shape=(3,), plates=(1, 1, 3)) ) # # Test v # # Simple check(4, 3, V=Gamma(np.random.rand(1, 3), np.random.rand(1, 3), plates=(1, 3))) check(4, 3, V=Gamma(np.random.rand(3), np.random.rand(3), plates=(3,))) # Plates check(4, 3, V=Gamma(np.random.rand(5, 6, 1, 3), np.random.rand(5, 6, 1, 3), plates=(5, 6, 1, 3))) # Plates with moments broadcasted over plates check(4, 3, V=Gamma(np.random.rand(1, 3), np.random.rand(1, 3), plates=(5, 1, 3))) check(4, 3, V=Gamma(np.random.rand(1, 1, 3), np.random.rand(1, 1, 3), plates=(5, 1, 3))) # Plates broadcasting check(4, 3, plates=(5,), V=Gamma(np.random.rand(1, 3), np.random.rand(1, 3), plates=(1, 3))) check(4, 3, plates=(5,), V=Gamma(np.random.rand(1, 1, 3), np.random.rand(1, 1, 3), plates=(1, 1, 3))) # # Check with input signals # mu = 2 Lambda = 3 A = 4 B = 5 v = 6 inputs = [[-2], [3]] X = GaussianMarkovChain([mu], [[Lambda]], [[A, B]], [v], inputs=inputs) V = np.array([[v * A ** 2, -v * A, 0], [-v * A, v * A ** 2, -v * A], [0, -v * A, 0]]) + np.array( [[Lambda, 0, 0], [0, v, 0], [0, 0, v]] ) m = ( np.array([Lambda * mu, 0, 0]) + np.array([0, v * B * inputs[0][0], v * B * inputs[1][0]]) - np.array([v * A * B * inputs[0][0], v * A * B * inputs[1][0], 0]) ) Cov = np.linalg.inv(V) mean = np.dot(Cov, m) X.update() u = X.get_moments() self.assertAllClose(u[0], mean[:, None]) self.assertAllClose(u[1] - u[0][..., None, :] * u[0][..., :, None], Cov[(0, 1, 2), (0, 1, 2), None, None]) self.assertAllClose(u[2] - u[0][..., :-1, :, None] * u[0][..., 1:, None, :], Cov[(0, 1), (1, 2), None, None]) pass
def test_message_to_child(self): """ Test the message from SumMultiply to its children. """ def compare_moments(u0, u1, *args): Y = SumMultiply(*args) u_Y = Y.get_moments() self.assertAllClose(u_Y[0], u0) self.assertAllClose(u_Y[1], u1) # Test constant parent y = np.random.randn(2, 3, 4) compare_moments(y, linalg.outer(y, y, ndim=2), 'ij->ij', y) # Do nothing for 2-D array Y = GaussianARD(np.random.randn(5, 2, 3), np.random.rand(5, 2, 3), plates=(5, ), shape=(2, 3)) y = Y.get_moments() compare_moments(y[0], y[1], 'ij->ij', Y) compare_moments(y[0], y[1], Y, [0, 1], [0, 1]) # Sum over the rows of a matrix Y = GaussianARD(np.random.randn(5, 2, 3), np.random.rand(5, 2, 3), plates=(5, ), shape=(2, 3)) y = Y.get_moments() mu = np.einsum('...ij->...j', y[0]) cov = np.einsum('...ijkl->...jl', y[1]) compare_moments(mu, cov, 'ij->j', Y) compare_moments(mu, cov, Y, [0, 1], [1]) # Inner product of three vectors X1 = GaussianARD(np.random.randn(2), np.random.rand(2), plates=(), shape=(2, )) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(6, 1, 2), np.random.rand(6, 1, 2), plates=(6, 1), shape=(2, )) x2 = X2.get_moments() X3 = GaussianARD(np.random.randn(7, 6, 5, 2), np.random.rand(7, 6, 5, 2), plates=(7, 6, 5), shape=(2, )) x3 = X3.get_moments() mu = np.einsum('...i,...i,...i->...', x1[0], x2[0], x3[0]) cov = np.einsum('...ij,...ij,...ij->...', x1[1], x2[1], x3[1]) compare_moments(mu, cov, 'i,i,i', X1, X2, X3) compare_moments(mu, cov, 'i,i,i->', X1, X2, X3) compare_moments(mu, cov, X1, [9], X2, [9], X3, [9]) compare_moments(mu, cov, X1, [9], X2, [9], X3, [9], []) # Outer product of two vectors X1 = GaussianARD(np.random.randn(2), np.random.rand(2), plates=(5, ), shape=(2, )) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(6, 1, 2), np.random.rand(6, 1, 2), plates=(6, 1), shape=(2, )) x2 = X2.get_moments() mu = np.einsum('...i,...j->...ij', x1[0], x2[0]) cov = np.einsum('...ik,...jl->...ijkl', x1[1], x2[1]) compare_moments(mu, cov, 'i,j->ij', X1, X2) compare_moments(mu, cov, X1, [9], X2, [7], [9, 7]) # Matrix product Y1 = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), plates=(), shape=(3, 2)) y1 = Y1.get_moments() Y2 = GaussianARD(np.random.randn(5, 2, 3), np.random.rand(5, 2, 3), plates=(5, ), shape=(2, 3)) y2 = Y2.get_moments() mu = np.einsum('...ik,...kj->...ij', y1[0], y2[0]) cov = np.einsum('...ikjl,...kmln->...imjn', y1[1], y2[1]) compare_moments(mu, cov, 'ik,kj->ij', Y1, Y2) compare_moments(mu, cov, Y1, ['i', 'k'], Y2, ['k', 'j'], ['i', 'j']) # Trace of a matrix product Y1 = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), plates=(), shape=(3, 2)) y1 = Y1.get_moments() Y2 = GaussianARD(np.random.randn(5, 2, 3), np.random.rand(5, 2, 3), plates=(5, ), shape=(2, 3)) y2 = Y2.get_moments() mu = np.einsum('...ij,...ji->...', y1[0], y2[0]) cov = np.einsum('...ikjl,...kilj->...', y1[1], y2[1]) compare_moments(mu, cov, 'ij,ji', Y1, Y2) compare_moments(mu, cov, 'ij,ji->', Y1, Y2) compare_moments(mu, cov, Y1, ['i', 'j'], Y2, ['j', 'i']) compare_moments(mu, cov, Y1, ['i', 'j'], Y2, ['j', 'i'], []) # Vector-matrix-vector product X1 = GaussianARD(np.random.randn(3), np.random.rand(3), plates=(), shape=(3, )) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(6, 1, 2), np.random.rand(6, 1, 2), plates=(6, 1), shape=(2, )) x2 = X2.get_moments() Y = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), plates=(), shape=(3, 2)) y = Y.get_moments() mu = np.einsum('...i,...ij,...j->...', x1[0], y[0], x2[0]) cov = np.einsum('...ia,...ijab,...jb->...', x1[1], y[1], x2[1]) compare_moments(mu, cov, 'i,ij,j', X1, Y, X2) compare_moments(mu, cov, X1, [1], Y, [1, 2], X2, [2]) # Complex sum-product of 0-D, 1-D, 2-D and 3-D arrays V = GaussianARD(np.random.randn(7, 6, 5), np.random.rand(7, 6, 5), plates=(7, 6, 5), shape=()) v = V.get_moments() X = GaussianARD(np.random.randn(6, 1, 2), np.random.rand(6, 1, 2), plates=(6, 1), shape=(2, )) x = X.get_moments() Y = GaussianARD(np.random.randn(3, 4), np.random.rand(3, 4), plates=(5, ), shape=(3, 4)) y = Y.get_moments() Z = GaussianARD(np.random.randn(4, 2, 3), np.random.rand(4, 2, 3), plates=(6, 5), shape=(4, 2, 3)) z = Z.get_moments() mu = np.einsum('...,...i,...kj,...jik->...k', v[0], x[0], y[0], z[0]) cov = np.einsum('...,...ia,...kjcb,...jikbac->...kc', v[1], x[1], y[1], z[1]) compare_moments(mu, cov, ',i,kj,jik->k', V, X, Y, Z) compare_moments(mu, cov, V, [], X, ['i'], Y, ['k', 'j'], Z, ['j', 'i', 'k'], ['k']) # Test with constant nodes N = 10 D = 5 a = np.random.randn(N, D) B = Gaussian( np.random.randn(D), random.covariance(D), ) X = SumMultiply('i,i->', B, a) np.testing.assert_allclose( X.get_moments()[0], np.einsum('ni,i->n', a, B.get_moments()[0]), ) np.testing.assert_allclose( X.get_moments()[1], np.einsum('ni,nj,ij->n', a, a, B.get_moments()[1]), ) # # Gaussian-gamma parents # # Outer product of vectors X1 = GaussianARD(np.random.randn(2), np.random.rand(2), shape=(2, )) x1 = X1.get_moments() X2 = GaussianGamma(np.random.randn(6, 1, 2), random.covariance(2), np.random.rand(6, 1), np.random.rand(6, 1), plates=(6, 1)) x2 = X2.get_moments() Y = SumMultiply('i,j->ij', X1, X2) u = Y._message_to_child() y = np.einsum('...i,...j->...ij', x1[0], x2[0]) yy = np.einsum('...ik,...jl->...ijkl', x1[1], x2[1]) self.assertAllClose(u[0], y) self.assertAllClose(u[1], yy) self.assertAllClose(u[2], x2[2]) self.assertAllClose(u[3], x2[3]) # Test with constant nodes N = 10 M = 8 D = 5 a = np.random.randn(N, 1, D) B = GaussianGamma( np.random.randn(M, D), random.covariance(D, size=(M, )), np.random.rand(M), np.random.rand(M), ndim=1, ) X = SumMultiply('i,i->', B, a) np.testing.assert_allclose( X.get_moments()[0], np.einsum('nmi,mi->nm', a, B.get_moments()[0]), ) np.testing.assert_allclose( X.get_moments()[1], np.einsum('nmi,nmj,mij->nm', a, a, B.get_moments()[1]), ) np.testing.assert_allclose( X.get_moments()[2], B.get_moments()[2], ) np.testing.assert_allclose( X.get_moments()[3], B.get_moments()[3], ) pass
def test_message_to_parent(self): """ Test the message from SumMultiply node to its parents. """ data = 2 tau = 3 def check_message(true_m0, true_m1, parent, *args, F=None): if F is None: A = SumMultiply(*args) B = GaussianARD(A, tau) B.observe(data*np.ones(A.plates + A.dims[0])) else: A = F (A_m0, A_m1) = A._message_to_parent(parent) self.assertAllClose(true_m0, A_m0) self.assertAllClose(true_m1, A_m1) pass # Check: different message to each of multiple parents X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x2 = X2.get_moments() m0 = tau * data * x2[0] m1 = -0.5 * tau * x2[1] * np.identity(2) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, [9], X2, [9], [9]) m0 = tau * data * x1[0] m1 = -0.5 * tau * x1[1] * np.identity(2) check_message(m0, m1, 1, 'i,i->i', X1, X2) check_message(m0, m1, 1, X1, [9], X2, [9], [9]) # Check: key not in output X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x1 = X1.get_moments() m0 = tau * data * np.ones(2) m1 = -0.5 * tau * np.ones((2,2)) check_message(m0, m1, 0, 'i', X1) check_message(m0, m1, 0, 'i->', X1) check_message(m0, m1, 0, X1, [9]) check_message(m0, m1, 0, X1, [9], []) # Check: key not in some input X1 = GaussianARD(np.random.randn(), np.random.rand()) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x2 = X2.get_moments() m0 = tau * data * np.sum(x2[0], axis=-1) m1 = -0.5 * tau * np.sum(x2[1] * np.identity(2), axis=(-1,-2)) check_message(m0, m1, 0, ',i->i', X1, X2) check_message(m0, m1, 0, X1, [], X2, [9], [9]) m0 = tau * data * x1[0] * np.ones(2) m1 = -0.5 * tau * x1[1] * np.identity(2) check_message(m0, m1, 1, ',i->i', X1, X2) check_message(m0, m1, 1, X1, [], X2, [9], [9]) # Check: keys in different order Y1 = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), ndim=2) y1 = Y1.get_moments() Y2 = GaussianARD(np.random.randn(2,3), np.random.rand(2,3), ndim=2) y2 = Y2.get_moments() m0 = tau * data * y2[0].T m1 = -0.5 * tau * np.einsum('ijlk->jikl', y2[1] * misc.identity(2,3)) check_message(m0, m1, 0, 'ij,ji->ij', Y1, Y2) check_message(m0, m1, 0, Y1, ['i','j'], Y2, ['j','i'], ['i','j']) m0 = tau * data * y1[0].T m1 = -0.5 * tau * np.einsum('ijlk->jikl', y1[1] * misc.identity(3,2)) check_message(m0, m1, 1, 'ij,ji->ij', Y1, Y2) check_message(m0, m1, 1, Y1, ['i','j'], Y2, ['j','i'], ['i','j']) # Check: plates when different dimensionality X1 = GaussianARD(np.random.randn(5), np.random.rand(5), shape=(), plates=(5,)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(5,3), np.random.rand(5,3), shape=(3,), plates=(5,)) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((5,3)) * x2[0], axis=-1) m1 = -0.5 * tau * np.sum(x2[1] * misc.identity(3), axis=(-1,-2)) check_message(m0, m1, 0, ',i->i', X1, X2) check_message(m0, m1, 0, X1, [], X2, ['i'], ['i']) m0 = tau * data * x1[0][:,np.newaxis] * np.ones((5,3)) m1 = -0.5 * tau * x1[1][:,np.newaxis,np.newaxis] * misc.identity(3) check_message(m0, m1, 1, ',i->i', X1, X2) check_message(m0, m1, 1, X1, [], X2, ['i'], ['i']) # Check: other parent's moments broadcasts over plates when node has the # same plates X1 = GaussianARD(np.random.randn(5,4,3), np.random.rand(5,4,3), shape=(3,), plates=(5,4)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3,), plates=(5,4)) x2 = X2.get_moments() m0 = tau * data * np.ones((5,4,3)) * x2[0] m1 = -0.5 * tau * x2[1] * misc.identity(3) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i']) # Check: other parent's moments broadcasts over plates when node does # not have that plate X1 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3,), plates=()) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3,), plates=(5,4)) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((5,4,3)) * x2[0], axis=(0,1)) m1 = -0.5 * tau * np.sum(np.ones((5,4,1,1)) * misc.identity(3) * x2[1], axis=(0,1)) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i']) # Check: other parent's moments broadcasts over plates when the node # only broadcasts that plate X1 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3,), plates=(1,1)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3,), plates=(5,4)) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((5,4,3)) * x2[0], axis=(0,1), keepdims=True) m1 = -0.5 * tau * np.sum(np.ones((5,4,1,1)) * misc.identity(3) * x2[1], axis=(0,1), keepdims=True) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i']) # Check: broadcasted dimensions X1 = GaussianARD(np.random.randn(1,1), np.random.rand(1,1), ndim=2) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), ndim=2) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((3,2)) * x2[0], keepdims=True) m1 = -0.5 * tau * np.sum(misc.identity(3,2) * x2[1], keepdims=True) check_message(m0, m1, 0, 'ij,ij->ij', X1, X2) check_message(m0, m1, 0, X1, [0,1], X2, [0,1], [0,1]) m0 = tau * data * np.ones((3,2)) * x1[0] m1 = -0.5 * tau * misc.identity(3,2) * x1[1] check_message(m0, m1, 1, 'ij,ij->ij', X1, X2) check_message(m0, m1, 1, X1, [0,1], X2, [0,1], [0,1]) # Check: non-ARD observations X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x1 = X1.get_moments() Lambda = np.array([[2, 1.5], [1.5, 2]]) F = SumMultiply('i->i', X1) Y = Gaussian(F, Lambda) y = np.random.randn(2) Y.observe(y) m0 = np.dot(Lambda, y) m1 = -0.5 * Lambda check_message(m0, m1, 0, 'i->i', X1, F=F) check_message(m0, m1, 0, X1, ['i'], ['i'], F=F) # Check: mask with same shape X1 = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), shape=(2,), plates=(3,)) x1 = X1.get_moments() mask = np.array([True, False, True]) F = SumMultiply('i->i', X1) Y = GaussianARD(F, tau, ndim=1) Y.observe(data*np.ones((3,2)), mask=mask) m0 = tau * data * mask[:,np.newaxis] * np.ones(2) m1 = -0.5 * tau * mask[:,np.newaxis,np.newaxis] * np.identity(2) check_message(m0, m1, 0, 'i->i', X1, F=F) check_message(m0, m1, 0, X1, ['i'], ['i'], F=F) # Check: mask larger X1 = GaussianARD(np.random.randn(2), np.random.rand(2), shape=(2,), plates=()) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), shape=(2,), plates=(3,)) x2 = X2.get_moments() mask = np.array([True, False, True]) F = SumMultiply('i,i->i', X1, X2) Y = GaussianARD(F, tau, plates=(3,), ndim=1) Y.observe(data*np.ones((3,2)), mask=mask) m0 = tau * data * np.sum(mask[:,np.newaxis] * x2[0], axis=0) m1 = -0.5 * tau * np.sum(mask[:,np.newaxis,np.newaxis] * x2[1] * np.identity(2), axis=0) check_message(m0, m1, 0, 'i,i->i', X1, X2, F=F) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i'], F=F) # Check: mask for broadcasted plate X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1, plates=(1,)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1, plates=(3,)) x2 = X2.get_moments() mask = np.array([True, False, True]) F = SumMultiply('i,i->i', X1, X2) Y = GaussianARD(F, tau, plates=(3,), ndim=1) Y.observe(data*np.ones((3,2)), mask=mask) m0 = tau * data * np.sum(mask[:,np.newaxis] * x2[0], axis=0, keepdims=True) m1 = -0.5 * tau * np.sum(mask[:,np.newaxis,np.newaxis] * x2[1] * np.identity(2), axis=0, keepdims=True) check_message(m0, m1, 0, 'i->i', X1, F=F) check_message(m0, m1, 0, X1, ['i'], ['i'], F=F) # Check: Gaussian-gamma parents X1 = GaussianGamma( np.random.randn(2), random.covariance(2), np.random.rand(), np.random.rand() ) x1 = X1.get_moments() X2 = GaussianGamma( np.random.randn(2), random.covariance(2), np.random.rand(), np.random.rand() ) x2 = X2.get_moments() F = SumMultiply('i,i->i', X1, X2) V = random.covariance(2) y = np.random.randn(2) Y = Gaussian(F, V) Y.observe(y) m0 = np.dot(V, y) * x2[0] m1 = -0.5 * V * x2[1] m2 = -0.5 * np.einsum('i,ij,j', y, V, y) * x2[2]#linalg.inner(V, x2[2], ndim=2) m3 = 0.5 * 2 #linalg.chol_logdet(linalg.chol(V)) + 2*x2[3] m = F._message_to_parent(0) self.assertAllClose(m[0], m0) self.assertAllClose(m[1], m1) self.assertAllClose(m[2], m2) self.assertAllClose(m[3], m3) pass
def test_message_to_parent(self): """ Test the message from SumMultiply node to its parents. """ data = 2 tau = 3 def check_message(true_m0, true_m1, parent, *args, F=None): if F is None: A = SumMultiply(*args) B = GaussianARD(A, tau) B.observe(data * np.ones(A.plates + A.dims[0])) else: A = F (A_m0, A_m1) = A._message_to_parent(parent) self.assertAllClose(true_m0, A_m0) self.assertAllClose(true_m1, A_m1) pass # Check: different message to each of multiple parents X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x2 = X2.get_moments() m0 = tau * data * x2[0] m1 = -0.5 * tau * x2[1] * np.identity(2) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, [9], X2, [9], [9]) m0 = tau * data * x1[0] m1 = -0.5 * tau * x1[1] * np.identity(2) check_message(m0, m1, 1, 'i,i->i', X1, X2) check_message(m0, m1, 1, X1, [9], X2, [9], [9]) # Check: key not in output X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x1 = X1.get_moments() m0 = tau * data * np.ones(2) m1 = -0.5 * tau * np.ones((2, 2)) check_message(m0, m1, 0, 'i', X1) check_message(m0, m1, 0, 'i->', X1) check_message(m0, m1, 0, X1, [9]) check_message(m0, m1, 0, X1, [9], []) # Check: key not in some input X1 = GaussianARD(np.random.randn(), np.random.rand()) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x2 = X2.get_moments() m0 = tau * data * np.sum(x2[0], axis=-1) m1 = -0.5 * tau * np.sum(x2[1] * np.identity(2), axis=(-1, -2)) check_message(m0, m1, 0, ',i->i', X1, X2) check_message(m0, m1, 0, X1, [], X2, [9], [9]) m0 = tau * data * x1[0] * np.ones(2) m1 = -0.5 * tau * x1[1] * np.identity(2) check_message(m0, m1, 1, ',i->i', X1, X2) check_message(m0, m1, 1, X1, [], X2, [9], [9]) # Check: keys in different order Y1 = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), ndim=2) y1 = Y1.get_moments() Y2 = GaussianARD(np.random.randn(2, 3), np.random.rand(2, 3), ndim=2) y2 = Y2.get_moments() m0 = tau * data * y2[0].T m1 = -0.5 * tau * np.einsum('ijlk->jikl', y2[1] * misc.identity(2, 3)) check_message(m0, m1, 0, 'ij,ji->ij', Y1, Y2) check_message(m0, m1, 0, Y1, ['i', 'j'], Y2, ['j', 'i'], ['i', 'j']) m0 = tau * data * y1[0].T m1 = -0.5 * tau * np.einsum('ijlk->jikl', y1[1] * misc.identity(3, 2)) check_message(m0, m1, 1, 'ij,ji->ij', Y1, Y2) check_message(m0, m1, 1, Y1, ['i', 'j'], Y2, ['j', 'i'], ['i', 'j']) # Check: plates when different dimensionality X1 = GaussianARD(np.random.randn(5), np.random.rand(5), shape=(), plates=(5, )) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(5, 3), np.random.rand(5, 3), shape=(3, ), plates=(5, )) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((5, 3)) * x2[0], axis=-1) m1 = -0.5 * tau * np.sum(x2[1] * misc.identity(3), axis=(-1, -2)) check_message(m0, m1, 0, ',i->i', X1, X2) check_message(m0, m1, 0, X1, [], X2, ['i'], ['i']) m0 = tau * data * x1[0][:, np.newaxis] * np.ones((5, 3)) m1 = -0.5 * tau * x1[1][:, np.newaxis, np.newaxis] * misc.identity(3) check_message(m0, m1, 1, ',i->i', X1, X2) check_message(m0, m1, 1, X1, [], X2, ['i'], ['i']) # Check: other parent's moments broadcasts over plates when node has the # same plates X1 = GaussianARD(np.random.randn(5, 4, 3), np.random.rand(5, 4, 3), shape=(3, ), plates=(5, 4)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3, ), plates=(5, 4)) x2 = X2.get_moments() m0 = tau * data * np.ones((5, 4, 3)) * x2[0] m1 = -0.5 * tau * x2[1] * misc.identity(3) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i']) # Check: other parent's moments broadcasts over plates when node does # not have that plate X1 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3, ), plates=()) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3, ), plates=(5, 4)) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((5, 4, 3)) * x2[0], axis=(0, 1)) m1 = -0.5 * tau * np.sum(np.ones( (5, 4, 1, 1)) * misc.identity(3) * x2[1], axis=(0, 1)) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i']) # Check: other parent's moments broadcasts over plates when the node # only broadcasts that plate X1 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3, ), plates=(1, 1)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3), np.random.rand(3), shape=(3, ), plates=(5, 4)) x2 = X2.get_moments() m0 = tau * data * np.sum( np.ones((5, 4, 3)) * x2[0], axis=(0, 1), keepdims=True) m1 = -0.5 * tau * np.sum(np.ones( (5, 4, 1, 1)) * misc.identity(3) * x2[1], axis=(0, 1), keepdims=True) check_message(m0, m1, 0, 'i,i->i', X1, X2) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i']) # Check: broadcasted dimensions X1 = GaussianARD(np.random.randn(1, 1), np.random.rand(1, 1), ndim=2) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), ndim=2) x2 = X2.get_moments() m0 = tau * data * np.sum(np.ones((3, 2)) * x2[0], keepdims=True) m1 = -0.5 * tau * np.sum(misc.identity(3, 2) * x2[1], keepdims=True) check_message(m0, m1, 0, 'ij,ij->ij', X1, X2) check_message(m0, m1, 0, X1, [0, 1], X2, [0, 1], [0, 1]) m0 = tau * data * np.ones((3, 2)) * x1[0] m1 = -0.5 * tau * misc.identity(3, 2) * x1[1] check_message(m0, m1, 1, 'ij,ij->ij', X1, X2) check_message(m0, m1, 1, X1, [0, 1], X2, [0, 1], [0, 1]) # Check: non-ARD observations X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1) x1 = X1.get_moments() Lambda = np.array([[2, 1.5], [1.5, 2]]) F = SumMultiply('i->i', X1) Y = Gaussian(F, Lambda) y = np.random.randn(2) Y.observe(y) m0 = np.dot(Lambda, y) m1 = -0.5 * Lambda check_message(m0, m1, 0, 'i->i', X1, F=F) check_message(m0, m1, 0, X1, ['i'], ['i'], F=F) # Check: mask with same shape X1 = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), shape=(2, ), plates=(3, )) x1 = X1.get_moments() mask = np.array([True, False, True]) F = SumMultiply('i->i', X1) Y = GaussianARD(F, tau, ndim=1) Y.observe(data * np.ones((3, 2)), mask=mask) m0 = tau * data * mask[:, np.newaxis] * np.ones(2) m1 = -0.5 * tau * mask[:, np.newaxis, np.newaxis] * np.identity(2) check_message(m0, m1, 0, 'i->i', X1, F=F) check_message(m0, m1, 0, X1, ['i'], ['i'], F=F) # Check: mask larger X1 = GaussianARD(np.random.randn(2), np.random.rand(2), shape=(2, ), plates=()) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(3, 2), np.random.rand(3, 2), shape=(2, ), plates=(3, )) x2 = X2.get_moments() mask = np.array([True, False, True]) F = SumMultiply('i,i->i', X1, X2) Y = GaussianARD(F, tau, plates=(3, ), ndim=1) Y.observe(data * np.ones((3, 2)), mask=mask) m0 = tau * data * np.sum(mask[:, np.newaxis] * x2[0], axis=0) m1 = -0.5 * tau * np.sum( mask[:, np.newaxis, np.newaxis] * x2[1] * np.identity(2), axis=0) check_message(m0, m1, 0, 'i,i->i', X1, X2, F=F) check_message(m0, m1, 0, X1, ['i'], X2, ['i'], ['i'], F=F) # Check: mask for broadcasted plate X1 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1, plates=(1, )) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(2), np.random.rand(2), ndim=1, plates=(3, )) x2 = X2.get_moments() mask = np.array([True, False, True]) F = SumMultiply('i,i->i', X1, X2) Y = GaussianARD(F, tau, plates=(3, ), ndim=1) Y.observe(data * np.ones((3, 2)), mask=mask) m0 = tau * data * np.sum( mask[:, np.newaxis] * x2[0], axis=0, keepdims=True) m1 = -0.5 * tau * np.sum( mask[:, np.newaxis, np.newaxis] * x2[1] * np.identity(2), axis=0, keepdims=True) check_message(m0, m1, 0, 'i->i', X1, F=F) check_message(m0, m1, 0, X1, ['i'], ['i'], F=F) # Test with constant nodes N = 10 M = 8 D = 5 K = 3 a = np.random.randn(N, D) B = Gaussian( np.random.randn(D), random.covariance(D), ) C = GaussianARD(np.random.randn(M, 1, D, K), np.random.rand(M, 1, D, K), ndim=2) F = SumMultiply('i,i,ij->', a, B, C) tau = np.random.rand(M, N) Y = GaussianARD(F, tau, plates=(M, N)) y = np.random.randn(M, N) Y.observe(y) (m0, m1) = F._message_to_parent(1) np.testing.assert_allclose( m0, np.einsum('mn,ni,mnik->i', tau * y, a, C.get_moments()[0]), ) np.testing.assert_allclose( m1, np.einsum('mn,ni,nj,mnikjl->ij', -0.5 * tau, a, a, C.get_moments()[1]), ) # Check: Gaussian-gamma parents X1 = GaussianGamma(np.random.randn(2), random.covariance(2), np.random.rand(), np.random.rand()) x1 = X1.get_moments() X2 = GaussianGamma(np.random.randn(2), random.covariance(2), np.random.rand(), np.random.rand()) x2 = X2.get_moments() F = SumMultiply('i,i->i', X1, X2) V = random.covariance(2) y = np.random.randn(2) Y = Gaussian(F, V) Y.observe(y) m0 = np.dot(V, y) * x2[0] m1 = -0.5 * V * x2[1] m2 = -0.5 * np.einsum('i,ij,j', y, V, y) * x2[2] #linalg.inner(V, x2[2], ndim=2) m3 = 0.5 * 2 #linalg.chol_logdet(linalg.chol(V)) + 2*x2[3] m = F._message_to_parent(0) self.assertAllClose(m[0], m0) self.assertAllClose(m[1], m1) self.assertAllClose(m[2], m2) self.assertAllClose(m[3], m3) # Delta moments N = 10 M = 8 D = 5 a = np.random.randn(N, D) B = GaussianGamma(np.random.randn(D), random.covariance(D), np.random.rand(), np.random.rand(), ndim=1) F = SumMultiply('i,i->', a, B) tau = np.random.rand(M, N) Y = GaussianARD(F, tau, plates=(M, N)) y = np.random.randn(M, N) Y.observe(y) (m0, m1, m2, m3) = F._message_to_parent(1) np.testing.assert_allclose( m0, np.einsum('mn,ni->i', tau * y, a), ) np.testing.assert_allclose( m1, np.einsum('mn,ni,nj->ij', -0.5 * tau, a, a), ) np.testing.assert_allclose( m2, np.einsum('mn->', -0.5 * tau * y**2), ) np.testing.assert_allclose( m3, np.einsum('mn->', 0.5 * np.ones(np.shape(tau))), ) pass
def test_gradient(self): """Test standard gradient of a Gaussian node.""" D = 3 np.random.seed(42) # # Without observations # # Construct model mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda) # Random initialization mu0 = np.random.randn(D) Lambda0 = random.covariance(D) X.initialize_from_parameters(mu0, Lambda0) Q = VB(X) # Initial parameters phi0 = X.phi # Gradient rg = X.get_riemannian_gradient() g = X.get_gradient(rg) # Numerical gradient eps = 1e-6 p0 = X.get_parameters() l0 = Q.compute_lowerbound(ignore_masked=False) g_num = [np.zeros(D), np.zeros((D, D))] for i in range(D): e = np.zeros(D) e[i] = eps p1 = p0[0] + e X.set_parameters([p1, p0[1]]) l1 = Q.compute_lowerbound(ignore_masked=False) g_num[0][i] = (l1 - l0) / eps for i in range(D): for j in range(i + 1): e = np.zeros((D, D)) e[i, j] += eps e[j, i] += eps p1 = p0[1] + e X.set_parameters([p0[0], p1]) l1 = Q.compute_lowerbound(ignore_masked=False) g_num[1][i, j] = (l1 - l0) / (2 * eps) g_num[1][j, i] = (l1 - l0) / (2 * eps) # Check self.assertAllClose(g[0], g_num[0]) self.assertAllClose(g[1], g_num[1]) # # With observations # # Construct model mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda) # Random initialization mu0 = np.random.randn(D) Lambda0 = random.covariance(D) X.initialize_from_parameters(mu0, Lambda0) V = random.covariance(D) Y = Gaussian(X, V) Y.observe(np.random.randn(D)) Q = VB(Y, X) # Initial parameters phi0 = X.phi # Gradient rg = X.get_riemannian_gradient() g = X.get_gradient(rg) # Numerical gradient eps = 1e-6 p0 = X.get_parameters() l0 = Q.compute_lowerbound() g_num = [np.zeros(D), np.zeros((D, D))] for i in range(D): e = np.zeros(D) e[i] = eps p1 = p0[0] + e X.set_parameters([p1, p0[1]]) l1 = Q.compute_lowerbound() g_num[0][i] = (l1 - l0) / eps for i in range(D): for j in range(i + 1): e = np.zeros((D, D)) e[i, j] += eps e[j, i] += eps p1 = p0[1] + e X.set_parameters([p0[0], p1]) l1 = Q.compute_lowerbound() g_num[1][i, j] = (l1 - l0) / (2 * eps) g_num[1][j, i] = (l1 - l0) / (2 * eps) # Check self.assertAllClose(g[0], g_num[0]) self.assertAllClose(g[1], g_num[1]) # # With plates # # Construct model K = D + 1 mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda, plates=(K, )) V = random.covariance(D, size=(K, )) Y = Gaussian(X, V) Y.observe(np.random.randn(K, D)) Q = VB(Y, X) # Random initialization mu0 = np.random.randn(*(X.get_shape(0))) Lambda0 = random.covariance(D, size=X.plates) X.initialize_from_parameters(mu0, Lambda0) # Initial parameters phi0 = X.phi # Gradient rg = X.get_riemannian_gradient() g = X.get_gradient(rg) # Numerical gradient eps = 1e-6 p0 = X.get_parameters() l0 = Q.compute_lowerbound() g_num = [np.zeros(X.get_shape(0)), np.zeros(X.get_shape(1))] for k in range(K): for i in range(D): e = np.zeros(X.get_shape(0)) e[k, i] = eps p1 = p0[0] + e X.set_parameters([p1, p0[1]]) l1 = Q.compute_lowerbound() g_num[0][k, i] = (l1 - l0) / eps for i in range(D): for j in range(i + 1): e = np.zeros(X.get_shape(1)) e[k, i, j] += eps e[k, j, i] += eps p1 = p0[1] + e X.set_parameters([p0[0], p1]) l1 = Q.compute_lowerbound() g_num[1][k, i, j] = (l1 - l0) / (2 * eps) g_num[1][k, j, i] = (l1 - l0) / (2 * eps) # Check self.assertAllClose(g[0], g_num[0]) self.assertAllClose(g[1], g_num[1]) pass
def test_message_to_child(self): """ Test the message from SumMultiply to its children. """ def compare_moments(u0, u1, *args): Y = SumMultiply(*args) u_Y = Y.get_moments() self.assertAllClose(u_Y[0], u0) self.assertAllClose(u_Y[1], u1) # Test constant parent y = np.random.randn(2,3,4) compare_moments(y, linalg.outer(y, y, ndim=2), 'ij->ij', y) # Do nothing for 2-D array Y = GaussianARD(np.random.randn(5,2,3), np.random.rand(5,2,3), plates=(5,), shape=(2,3)) y = Y.get_moments() compare_moments(y[0], y[1], 'ij->ij', Y) compare_moments(y[0], y[1], Y, [0,1], [0,1]) # Sum over the rows of a matrix Y = GaussianARD(np.random.randn(5,2,3), np.random.rand(5,2,3), plates=(5,), shape=(2,3)) y = Y.get_moments() mu = np.einsum('...ij->...j', y[0]) cov = np.einsum('...ijkl->...jl', y[1]) compare_moments(mu, cov, 'ij->j', Y) compare_moments(mu, cov, Y, [0,1], [1]) # Inner product of three vectors X1 = GaussianARD(np.random.randn(2), np.random.rand(2), plates=(), shape=(2,)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(6,1,2), np.random.rand(6,1,2), plates=(6,1), shape=(2,)) x2 = X2.get_moments() X3 = GaussianARD(np.random.randn(7,6,5,2), np.random.rand(7,6,5,2), plates=(7,6,5), shape=(2,)) x3 = X3.get_moments() mu = np.einsum('...i,...i,...i->...', x1[0], x2[0], x3[0]) cov = np.einsum('...ij,...ij,...ij->...', x1[1], x2[1], x3[1]) compare_moments(mu, cov, 'i,i,i', X1, X2, X3) compare_moments(mu, cov, 'i,i,i->', X1, X2, X3) compare_moments(mu, cov, X1, [9], X2, [9], X3, [9]) compare_moments(mu, cov, X1, [9], X2, [9], X3, [9], []) # Outer product of two vectors X1 = GaussianARD(np.random.randn(2), np.random.rand(2), plates=(5,), shape=(2,)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(6,1,2), np.random.rand(6,1,2), plates=(6,1), shape=(2,)) x2 = X2.get_moments() mu = np.einsum('...i,...j->...ij', x1[0], x2[0]) cov = np.einsum('...ik,...jl->...ijkl', x1[1], x2[1]) compare_moments(mu, cov, 'i,j->ij', X1, X2) compare_moments(mu, cov, X1, [9], X2, [7], [9,7]) # Matrix product Y1 = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), plates=(), shape=(3,2)) y1 = Y1.get_moments() Y2 = GaussianARD(np.random.randn(5,2,3), np.random.rand(5,2,3), plates=(5,), shape=(2,3)) y2 = Y2.get_moments() mu = np.einsum('...ik,...kj->...ij', y1[0], y2[0]) cov = np.einsum('...ikjl,...kmln->...imjn', y1[1], y2[1]) compare_moments(mu, cov, 'ik,kj->ij', Y1, Y2) compare_moments(mu, cov, Y1, ['i','k'], Y2, ['k','j'], ['i','j']) # Trace of a matrix product Y1 = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), plates=(), shape=(3,2)) y1 = Y1.get_moments() Y2 = GaussianARD(np.random.randn(5,2,3), np.random.rand(5,2,3), plates=(5,), shape=(2,3)) y2 = Y2.get_moments() mu = np.einsum('...ij,...ji->...', y1[0], y2[0]) cov = np.einsum('...ikjl,...kilj->...', y1[1], y2[1]) compare_moments(mu, cov, 'ij,ji', Y1, Y2) compare_moments(mu, cov, 'ij,ji->', Y1, Y2) compare_moments(mu, cov, Y1, ['i','j'], Y2, ['j','i']) compare_moments(mu, cov, Y1, ['i','j'], Y2, ['j','i'], []) # Vector-matrix-vector product X1 = GaussianARD(np.random.randn(3), np.random.rand(3), plates=(), shape=(3,)) x1 = X1.get_moments() X2 = GaussianARD(np.random.randn(6,1,2), np.random.rand(6,1,2), plates=(6,1), shape=(2,)) x2 = X2.get_moments() Y = GaussianARD(np.random.randn(3,2), np.random.rand(3,2), plates=(), shape=(3,2)) y = Y.get_moments() mu = np.einsum('...i,...ij,...j->...', x1[0], y[0], x2[0]) cov = np.einsum('...ia,...ijab,...jb->...', x1[1], y[1], x2[1]) compare_moments(mu, cov, 'i,ij,j', X1, Y, X2) compare_moments(mu, cov, X1, [1], Y, [1,2], X2, [2]) # Complex sum-product of 0-D, 1-D, 2-D and 3-D arrays V = GaussianARD(np.random.randn(7,6,5), np.random.rand(7,6,5), plates=(7,6,5), shape=()) v = V.get_moments() X = GaussianARD(np.random.randn(6,1,2), np.random.rand(6,1,2), plates=(6,1), shape=(2,)) x = X.get_moments() Y = GaussianARD(np.random.randn(3,4), np.random.rand(3,4), plates=(5,), shape=(3,4)) y = Y.get_moments() Z = GaussianARD(np.random.randn(4,2,3), np.random.rand(4,2,3), plates=(6,5), shape=(4,2,3)) z = Z.get_moments() mu = np.einsum('...,...i,...kj,...jik->...k', v[0], x[0], y[0], z[0]) cov = np.einsum('...,...ia,...kjcb,...jikbac->...kc', v[1], x[1], y[1], z[1]) compare_moments(mu, cov, ',i,kj,jik->k', V, X, Y, Z) compare_moments(mu, cov, V, [], X, ['i'], Y, ['k','j'], Z, ['j','i','k'], ['k']) # # Gaussian-gamma parents # # Outer product of vectors X1 = GaussianARD(np.random.randn(2), np.random.rand(2), shape=(2,)) x1 = X1.get_moments() X2 = GaussianGamma( np.random.randn(6,1,2), random.covariance(2), np.random.rand(6,1), np.random.rand(6,1), plates=(6,1) ) x2 = X2.get_moments() Y = SumMultiply('i,j->ij', X1, X2) u = Y._message_to_child() y = np.einsum('...i,...j->...ij', x1[0], x2[0]) yy = np.einsum('...ik,...jl->...ijkl', x1[1], x2[1]) self.assertAllClose(u[0], y) self.assertAllClose(u[1], yy) self.assertAllClose(u[2], x2[2]) self.assertAllClose(u[3], x2[3]) pass
def _run_checks(self, check): # Basic test check(2, 3) # Test mu check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=())) check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=(5, ))) # Test Lambda check(2, 3, Lambda=Wishart(3, random.covariance(2))) check(2, 3, Lambda=Wishart(3, random.covariance(2), plates=(5, ))) # Test A check(2, 3, A=GaussianARD(2, 4, shape=(2, ), plates=(2, ))) check(2, 3, A=GaussianARD(2, 4, shape=(2, ), plates=(3, 2))) check(2, 3, A=GaussianARD(2, 4, shape=(2, ), plates=(5, 3, 2))) # Test Lambda and mu check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=()), Lambda=Wishart(2, random.covariance(2))) check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=(5, )), Lambda=Wishart(2, random.covariance(2), plates=(5, ))) # Test mu and A check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=()), A=GaussianARD(2, 4, shape=(2, ), plates=(2, ))) check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=(5, )), A=GaussianARD(2, 4, shape=(2, ), plates=( 5, 1, 2, ))) # Test Lambda and A check(2, 3, Lambda=Wishart(2, random.covariance(2)), A=GaussianARD(2, 4, shape=(2, ), plates=(2, ))) check(2, 3, Lambda=Wishart(2, random.covariance(2), plates=(5, )), A=GaussianARD(2, 4, shape=(2, ), plates=( 5, 1, 2, ))) # Test mu, Lambda and A check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=()), Lambda=Wishart(2, random.covariance(2)), A=GaussianARD(2, 4, shape=(2, ), plates=(2, ))) check(2, 3, mu=GaussianARD(2, 4, shape=(2, ), plates=(5, )), Lambda=Wishart(2, random.covariance(2), plates=(5, )), A=GaussianARD(2, 4, shape=(2, ), plates=( 5, 1, 2, ))) pass
def test_message_to_parents_with_inputs(self): """ Check gradient passed to inputs parent node """ def check(Mu, Lambda, A, V, U): X = GaussianMarkovChain(Mu, Lambda, A, V, inputs=U) Y = Gaussian(X, random.covariance(D)) # Check moments self.assert_moments( X, postprocess=lambda u: [ u[0], u[1] + linalg.transpose(u[1], ndim=1), u[2] ] ) Y.observe(np.random.randn(N+1, D)) X.update() # Check gradient messages to parents self.assert_message_to_parent(X, Mu) self.assert_message_to_parent( X, Lambda, postprocess=lambda phi: [ phi[0] + linalg.transpose(phi[0], ndim=1), phi[1] ] ) self.assert_message_to_parent( X, A, postprocess=lambda phi: [ phi[0], phi[1] + linalg.transpose(phi[1], ndim=1), ] ) self.assert_message_to_parent(X, V) self.assert_message_to_parent(X, U) N = 4 D = 2 K = 3 check( Gaussian( np.random.randn(D), random.covariance(D) ), Wishart( D, random.covariance(D) ), Gaussian( np.random.randn(D,D+K), random.covariance(D+K) ), Gamma( D, np.random.rand(D) ), Gaussian( np.random.randn(N,K), random.covariance(K) ) ) check( Gaussian( np.random.randn(D), random.covariance(D) ), Wishart( D, random.covariance(D) ), GaussianGamma( np.random.randn(D,D+K), random.covariance(D+K), D, np.random.rand(D), ndim=1 ), Gamma( D, np.random.rand(D) ), Gaussian( np.random.randn(N,K), random.covariance(K) ) ) pass
def check(N, D, K, plates=None, mu=None, Lambda=None, B=None, S=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if B is None: B = np.random.randn(D,D,K) if S is None: S = np.random.randn(N-1,K) if V is None: V = np.random.rand(D) X = VaryingGaussianMarkovChain(mu, Lambda, B, S, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = X.parents[0].get_moments() (Lambda, _) = X.parents[1].get_moments() (b, bb) = X.parents[2].get_moments() (s, ss) = X.parents[3].get_moments() (v, _) = X.parents[4].get_moments() v = v * np.ones((N-1,D)) #V = np.atleast_3d(v)[...,-1,:,None]*np.identity(D) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N,D,N,D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N,D)) m[...,0,:] = np.einsum('...ij,...j->...i', Lambda, mu) #m = np.reshape(m, plates_mu + (N*D,)) A = np.einsum('...dik,...nk->...ndi', b, s) AA = np.einsum('...dikjl,...nkl->...ndij', bb, ss) C[...,0,:,0,:] = Lambda + np.einsum('...dij,...d->...ij', AA[...,0,:,:,:], v[...,0,:]) for n in range(1,N-1): C[...,n,:,n,:] = (np.einsum('...dij,...d->...ij', AA[...,n,:,:,:], v[...,n,:]) + v[...,n,:,None] * np.identity(D)) for n in range(N-1): C[...,n,:,n+1,:] = -np.einsum('...di,...d->...id', A[...,n,:,:], v[...,n,:]) C[...,n+1,:,n,:] = -np.einsum('...di,...d->...di', A[...,n,:,:], v[...,n,:]) C[...,-1,:,-1,:] = v[...,-1,:,None]*np.identity(D) C = np.reshape(C, plates_C+(N*D,N*D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C+(N,D,N,D)) m0 = np.einsum('...minj,...nj->...mi', Cov, m) m1 = np.zeros(plates_C+(N,D,D)) m2 = np.zeros(plates_C+(N-1,D,D)) for n in range(N): m1[...,n,:,:] = Cov[...,n,:,n,:] + np.einsum('...i,...j->...ij', m0[...,n,:], m0[...,n,:]) for n in range(N-1): m2[...,n,:,:] = Cov[...,n,:,n+1,:] + np.einsum('...i,...j->...ij', m0[...,n,:], m0[...,n+1,:]) self.assertAllClose(m0, u0*np.ones(np.shape(m0))) self.assertAllClose(m1, u1*np.ones(np.shape(m1))) self.assertAllClose(m2, u2*np.ones(np.shape(m2))) pass
def test_message_to_child(self): """ Test the updating of GaussianMarkovChain. Check that the moments and the lower bound contribution are computed correctly. """ # TODO: Add plates and missing values! # Dimensionalities D = 3 N = 5 (Y, X, Mu, Lambda, A, V) = self.create_model(N, D) # Inference with arbitrary observations y = np.random.randn(N,D) Y.observe(y) X.update() (x_vb, xnxn_vb, xpxn_vb) = X.get_moments() # Get parameter moments (mu0, mumu0) = Mu.get_moments() (icov0, logdet0) = Lambda.get_moments() (a, aa) = A.get_moments() (icov_x, logdetx) = V.get_moments() icov_x = np.diag(icov_x) # Prior precision Z = np.einsum('...kij,...kk->...ij', aa, icov_x) U_diag = [icov0+Z] + (N-2)*[icov_x+Z] + [icov_x] U_super = (N-1) * [-np.dot(a.T, icov_x)] U = misc.block_banded(U_diag, U_super) # Prior mean mu_prior = np.zeros(D*N) mu_prior[:D] = np.dot(icov0,mu0) # Data Cov = np.linalg.inv(U + np.identity(D*N)) mu = np.dot(Cov, mu_prior + y.flatten()) # Moments xx = mu[:,np.newaxis]*mu[np.newaxis,:] + Cov mu = np.reshape(mu, (N,D)) xx = np.reshape(xx, (N,D,N,D)) # Check results self.assertAllClose(x_vb, mu, msg="Incorrect mean") for n in range(N): self.assertAllClose(xnxn_vb[n,:,:], xx[n,:,n,:], msg="Incorrect second moment") for n in range(N-1): self.assertAllClose(xpxn_vb[n,:,:], xx[n,:,n+1,:], msg="Incorrect lagged second moment") # Compute the entropy H(X) ldet = linalg.logdet_cov(Cov) H = random.gaussian_entropy(-ldet, N*D) # Compute <log p(X|...)> xx = np.reshape(xx, (N*D, N*D)) mu = np.reshape(mu, (N*D,)) ldet = -logdet0 - np.sum(np.ones((N-1,D))*logdetx) P = random.gaussian_logpdf(np.einsum('...ij,...ij', xx, U), np.einsum('...i,...i', mu, mu_prior), np.einsum('...ij,...ij', mumu0, icov0), -ldet, N*D) # The VB bound from the net l = X.lower_bound_contribution() self.assertAllClose(l, H+P) # Compute the true bound <log p(X|...)> + H(X) # # Simple tests # def check(N, D, plates=None, mu=None, Lambda=None, A=None, V=None): if mu is None: mu = np.random.randn(D) if Lambda is None: Lambda = random.covariance(D) if A is None: A = np.random.randn(D,D) if V is None: V = np.random.rand(D) X = GaussianMarkovChain(mu, Lambda, A, V, plates=plates, n=N) (u0, u1, u2) = X._message_to_child() (mu, mumu) = Gaussian._ensure_moments(mu, GaussianMoments, ndim=1).get_moments() (Lambda, _) = Wishart._ensure_moments(Lambda, WishartMoments, ndim=1).get_moments() (a, aa) = Gaussian._ensure_moments(A, GaussianMoments, ndim=1).get_moments() a = a * np.ones((N-1,D,D)) # explicit broadcasting for simplicity aa = aa * np.ones((N-1,D,D,D)) # explicit broadcasting for simplicity (v, _) = Gamma._ensure_moments(V, GammaMoments).get_moments() v = v * np.ones((N-1,D)) plates_C = X.plates plates_mu = X.plates C = np.zeros(plates_C + (N,D,N,D)) plates_mu = np.shape(mu)[:-1] m = np.zeros(plates_mu + (N,D)) m[...,0,:] = np.einsum('...ij,...j->...i', Lambda, mu) C[...,0,:,0,:] = Lambda + np.einsum('...dij,...d->...ij', aa[...,0,:,:,:], v[...,0,:]) for n in range(1,N-1): C[...,n,:,n,:] = (np.einsum('...dij,...d->...ij', aa[...,n,:,:,:], v[...,n,:]) + v[...,n,:,None] * np.identity(D)) for n in range(N-1): C[...,n,:,n+1,:] = -np.einsum('...di,...d->...id', a[...,n,:,:], v[...,n,:]) C[...,n+1,:,n,:] = -np.einsum('...di,...d->...di', a[...,n,:,:], v[...,n,:]) C[...,-1,:,-1,:] = v[...,-1,:,None]*np.identity(D) C = np.reshape(C, plates_C+(N*D,N*D)) Cov = np.linalg.inv(C) Cov = np.reshape(Cov, plates_C+(N,D,N,D)) m0 = np.einsum('...minj,...nj->...mi', Cov, m) m1 = np.zeros(plates_C+(N,D,D)) m2 = np.zeros(plates_C+(N-1,D,D)) for n in range(N): m1[...,n,:,:] = Cov[...,n,:,n,:] + np.einsum('...i,...j->...ij', m0[...,n,:], m0[...,n,:]) for n in range(N-1): m2[...,n,:,:] = Cov[...,n,:,n+1,:] + np.einsum('...i,...j->...ij', m0[...,n,:], m0[...,n+1,:]) self.assertAllClose(m0, u0*np.ones(np.shape(m0))) self.assertAllClose(m1, u1*np.ones(np.shape(m1))) self.assertAllClose(m2, u2*np.ones(np.shape(m2))) pass check(4,1) check(4,3) # # Test mu # # Simple check(4,3, mu=Gaussian(np.random.randn(3), random.covariance(3))) # Plates check(4,3, mu=Gaussian(np.random.randn(5,6,3), random.covariance(3), plates=(5,6))) # Plates with moments broadcasted over plates check(4,3, mu=Gaussian(np.random.randn(3), random.covariance(3), plates=(5,))) check(4,3, mu=Gaussian(np.random.randn(1,3), random.covariance(3), plates=(5,))) # Plates broadcasting check(4,3, plates=(5,), mu=Gaussian(np.random.randn(3), random.covariance(3), plates=())) check(4,3, plates=(5,), mu=Gaussian(np.random.randn(1,3), random.covariance(3), plates=(1,))) # # Test Lambda # # Simple check(4,3, Lambda=Wishart(10+np.random.rand(), random.covariance(3))) # Plates check(4,3, Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=(5,6))) # Plates with moments broadcasted over plates check(4,3, Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=(5,))) check(4,3, Lambda=Wishart(10+np.random.rand(1), random.covariance(3), plates=(5,))) # Plates broadcasting check(4,3, plates=(5,), Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=())) check(4,3, plates=(5,), Lambda=Wishart(10+np.random.rand(), random.covariance(3), plates=(1,))) # # Test A # # Simple check(4,3, A=GaussianARD(np.random.randn(3,3), np.random.rand(3,3), shape=(3,), plates=(3,))) # Plates on time axis check(5,3, A=GaussianARD(np.random.randn(4,3,3), np.random.rand(4,3,3), shape=(3,), plates=(4,3))) # Plates on time axis with broadcasted moments check(5,3, A=GaussianARD(np.random.randn(1,3,3), np.random.rand(1,3,3), shape=(3,), plates=(4,3))) check(5,3, A=GaussianARD(np.random.randn(3,3), np.random.rand(3,3), shape=(3,), plates=(4,3))) # Plates check(4,3, A=GaussianARD(np.random.randn(5,6,1,3,3), np.random.rand(5,6,1,3,3), shape=(3,), plates=(5,6,1,3))) # Plates with moments broadcasted over plates check(4,3, A=GaussianARD(np.random.randn(3,3), np.random.rand(3,3), shape=(3,), plates=(5,1,3))) check(4,3, A=GaussianARD(np.random.randn(1,1,3,3), np.random.rand(1,1,3,3), shape=(3,), plates=(5,1,3))) # Plates broadcasting check(4,3, plates=(5,), A=GaussianARD(np.random.randn(3,3), np.random.rand(3,3), shape=(3,), plates=(3,))) check(4,3, plates=(5,), A=GaussianARD(np.random.randn(3,3), np.random.rand(3,3), shape=(3,), plates=(1,1,3))) # # Test v # # Simple check(4,3, V=Gamma(np.random.rand(1,3), np.random.rand(1,3), plates=(1,3))) check(4,3, V=Gamma(np.random.rand(3), np.random.rand(3), plates=(3,))) # Plates check(4,3, V=Gamma(np.random.rand(5,6,1,3), np.random.rand(5,6,1,3), plates=(5,6,1,3))) # Plates with moments broadcasted over plates check(4,3, V=Gamma(np.random.rand(1,3), np.random.rand(1,3), plates=(5,1,3))) check(4,3, V=Gamma(np.random.rand(1,1,3), np.random.rand(1,1,3), plates=(5,1,3))) # Plates broadcasting check(4,3, plates=(5,), V=Gamma(np.random.rand(1,3), np.random.rand(1,3), plates=(1,3))) check(4,3, plates=(5,), V=Gamma(np.random.rand(1,1,3), np.random.rand(1,1,3), plates=(1,1,3))) # # Check with input signals # mu = 2 Lambda = 3 A = 4 B = 5 v = 6 inputs = [[-2], [3]] X = GaussianMarkovChain([mu], [[Lambda]], [[A,B]], [v], inputs=inputs) V = (np.array([[v*A**2, -v*A, 0], [-v*A, v*A**2, -v*A], [0, -v*A, 0]]) + np.array([[Lambda, 0, 0], [0, v, 0], [0, 0, v]])) m = (np.array([Lambda*mu, 0, 0]) + np.array([0, v*B*inputs[0][0], v*B*inputs[1][0]]) - np.array([v*A*B*inputs[0][0], v*A*B*inputs[1][0], 0])) Cov = np.linalg.inv(V) mean = np.dot(Cov, m) X.update() u = X.get_moments() self.assertAllClose(u[0], mean[:,None]) self.assertAllClose(u[1] - u[0][...,None,:]*u[0][...,:,None], Cov[(0,1,2),(0,1,2),None,None]) self.assertAllClose(u[2] - u[0][...,:-1,:,None]*u[0][...,1:,None,:], Cov[(0,1),(1,2),None,None]) pass
def test_gradient(self): """Test standard gradient of a Gaussian node.""" D = 3 np.random.seed(42) # # Without observations # # Construct model mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda) # Random initialization mu0 = np.random.randn(D) Lambda0 = random.covariance(D) X.initialize_from_parameters(mu0, Lambda0) Q = VB(X) # Initial parameters phi0 = X.phi # Gradient rg = X.get_riemannian_gradient() g = X.get_gradient(rg) # Numerical gradient eps = 1e-6 p0 = X.get_parameters() l0 = Q.compute_lowerbound(ignore_masked=False) g_num = [np.zeros(D), np.zeros((D,D))] for i in range(D): e = np.zeros(D) e[i] = eps p1 = p0[0] + e X.set_parameters([p1, p0[1]]) l1 = Q.compute_lowerbound(ignore_masked=False) g_num[0][i] = (l1 - l0) / eps for i in range(D): for j in range(i+1): e = np.zeros((D,D)) e[i,j] += eps e[j,i] += eps p1 = p0[1] + e X.set_parameters([p0[0], p1]) l1 = Q.compute_lowerbound(ignore_masked=False) g_num[1][i,j] = (l1 - l0) / (2*eps) g_num[1][j,i] = (l1 - l0) / (2*eps) # Check self.assertAllClose(g[0], g_num[0]) self.assertAllClose(g[1], g_num[1]) # # With observations # # Construct model mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda) # Random initialization mu0 = np.random.randn(D) Lambda0 = random.covariance(D) X.initialize_from_parameters(mu0, Lambda0) V = random.covariance(D) Y = Gaussian(X, V) Y.observe(np.random.randn(D)) Q = VB(Y, X) # Initial parameters phi0 = X.phi # Gradient rg = X.get_riemannian_gradient() g = X.get_gradient(rg) # Numerical gradient eps = 1e-6 p0 = X.get_parameters() l0 = Q.compute_lowerbound() g_num = [np.zeros(D), np.zeros((D,D))] for i in range(D): e = np.zeros(D) e[i] = eps p1 = p0[0] + e X.set_parameters([p1, p0[1]]) l1 = Q.compute_lowerbound() g_num[0][i] = (l1 - l0) / eps for i in range(D): for j in range(i+1): e = np.zeros((D,D)) e[i,j] += eps e[j,i] += eps p1 = p0[1] + e X.set_parameters([p0[0], p1]) l1 = Q.compute_lowerbound() g_num[1][i,j] = (l1 - l0) / (2*eps) g_num[1][j,i] = (l1 - l0) / (2*eps) # Check self.assertAllClose(g[0], g_num[0]) self.assertAllClose(g[1], g_num[1]) # # With plates # # Construct model K = D+1 mu = np.random.randn(D) Lambda = random.covariance(D) X = Gaussian(mu, Lambda, plates=(K,)) V = random.covariance(D, size=(K,)) Y = Gaussian(X, V) Y.observe(np.random.randn(K,D)) Q = VB(Y, X) # Random initialization mu0 = np.random.randn(*(X.get_shape(0))) Lambda0 = random.covariance(D, size=X.plates) X.initialize_from_parameters(mu0, Lambda0) # Initial parameters phi0 = X.phi # Gradient rg = X.get_riemannian_gradient() g = X.get_gradient(rg) # Numerical gradient eps = 1e-6 p0 = X.get_parameters() l0 = Q.compute_lowerbound() g_num = [np.zeros(X.get_shape(0)), np.zeros(X.get_shape(1))] for k in range(K): for i in range(D): e = np.zeros(X.get_shape(0)) e[k,i] = eps p1 = p0[0] + e X.set_parameters([p1, p0[1]]) l1 = Q.compute_lowerbound() g_num[0][k,i] = (l1 - l0) / eps for i in range(D): for j in range(i+1): e = np.zeros(X.get_shape(1)) e[k,i,j] += eps e[k,j,i] += eps p1 = p0[1] + e X.set_parameters([p0[0], p1]) l1 = Q.compute_lowerbound() g_num[1][k,i,j] = (l1 - l0) / (2*eps) g_num[1][k,j,i] = (l1 - l0) / (2*eps) # Check self.assertAllClose(g[0], g_num[0]) self.assertAllClose(g[1], g_num[1]) pass
def test_messages(self): D = 2 M = 3 np.random.seed(42) def check(mu, Lambda, alpha, beta, ndim): X = GaussianGamma( mu, ( Lambda if isinstance(Lambda._moments, WishartMoments) else Lambda.as_wishart(ndim=ndim) ), alpha, beta, ndim=ndim ) self.assert_moments( X, postprocess=lambda u: [ u[0], u[1] + linalg.transpose(u[1], ndim=ndim), u[2], u[3] ], rtol=1e-5, atol=1e-6, eps=1e-8 ) X.observe( ( np.random.randn(*(X.plates + X.dims[0])), np.random.rand(*X.plates) ) ) self.assert_message_to_parent(X, mu) self.assert_message_to_parent( X, Lambda, postprocess=lambda m: [ m[0] + linalg.transpose(m[0], ndim=ndim), m[1], ] ) self.assert_message_to_parent(X, beta) check( Gaussian(np.random.randn(M, D), random.covariance(D), plates=(M,)), Wishart(D + np.random.rand(M), random.covariance(D), plates=(M,)), np.random.rand(M), Gamma(np.random.rand(M), np.random.rand(M), plates=(M,)), ndim=1 ) check( GaussianARD(np.random.randn(M, D), np.random.rand(M, D), ndim=0), Gamma(np.random.rand(M, D), np.random.rand(M, D)), np.random.rand(M, D), Gamma(np.random.rand(M, D), np.random.rand(M, D)), ndim=0 ) pass