Add WishartBartlett example.

Author: twiecki

pymc3/distributions/multivariate.py

@@ -4,12 +4,15 @@
 import warnings
 
 from .dist_math import *
+from . import ChiSquared, Normal
+from .. import Deterministic
 
 import numpy as np
+import scipy
 from . import transforms
 
 from theano.tensor.nlinalg import det, matrix_inverse, trace, eigh
-from theano.tensor import dot, cast, eye, diag, eq, le, ge, gt, all
+from theano.tensor import dot, cast, eye, diag, eq, le, ge, gt, all, zeros, sqrt, set_subtensor
 from theano.printing import Print
 from pymc3.distributions.distribution import draw_values, generate_samples
 import scipy.stats as st
@@ -270,18 +273,18 @@ def WishartBartlett(name, S, nu, is_cholesky=False, return_cholesky=False):
     tril_idx = np.tril_indices_from(S, k=-1)
     n_diag = len(diag_idx[0])
     n_tril = len(tril_idx[0])
-    c = T.sqrt(pm.ChiSquared('c', nu - np.arange(2, 2+n_diag), shape=n_diag))
-    z = pm.Normal('z', 0, 1, shape=n_tril)
+    c = sqrt(ChiSquared('c', nu - np.arange(2, 2+n_diag), shape=n_diag))
+    z = Normal('z', 0, 1, shape=n_tril)
     # Construct A matrix
-    A = T.zeros(S.shape, dtype=np.float32)
-    A = T.set_subtensor(A[diag_idx], c)
-    A = T.set_subtensor(A[tril_idx], z)
+    A = zeros(S.shape, dtype=np.float32)
+    A = set_subtensor(A[diag_idx], c)
+    A = set_subtensor(A[tril_idx], z)
 
     # L * A * A.T * L.T ~ Wishart(L*L.T, nu)
     if return_cholesky:
-        return pm.Deterministic(name, T.dot(L, A))
+        return Deterministic(name, dot(L, A))
     else:
-        return pm.Deterministic(name, T.dot(T.dot(T.dot(L, A), A.T), L.T))
+        return Deterministic(name, dot(dot(dot(L, A), A.T), L.T))
 
 
 class LKJCorr(Continuous):

pymc3/examples/wishart.py

@@ -0,0 +1,39 @@
+import pymc3 as pm
+import numpy as np
+import theano
+import theano.tensor as T
+import scipy.stats
+import matplotlib.pyplot as plt
+
+# Covariance matrix we want to recover
+covariance = np.matrix([[2, .5, -.5],
+                        [.5, 1.,  0.],
+                        [-.5, 0., 0.5]])
+
+prec = np.linalg.inv(covariance)
+
+mean = [.5, 1, .2]
+data = scipy.stats.multivariate_normal(mean, covariance).rvs(5000)
+
+plt.scatter(data[:, 0], data[:, 1])
+
+with pm.Model() as model:
+    S = np.eye(3)
+    nu = 5
+    mu = pm.Normal('mu', mu=0, sd=1, shape=3)
+    
+    # Use the transformed Wishart distribution
+    # Under the hood this will do a Cholesky decomposition
+    # of S and add two RVs to the sampler: c and z
+    prec = pm.WishartBartlett('prec', S, nu)
+
+    # To be able to compare it to truth, convert precision to covariance
+    cov = pm.Deterministic('cov', T.nlinalg.matrix_inverse(prec))
+
+    lp = pm.MvNormal('likelihood', mu=mu, tau=prec, observed=data)
+
+    start = pm.find_MAP()
+    step = pm.NUTS(scaling=start)
+    trace = pm.sample(500, step)
+
+pm.traceplot(trace[100:]);