ENH Add Wishart with Bartlett decomposition.

Author: twiecki

pymc3/distributions/multivariate.py

@@ -8,7 +8,7 @@
 from theano.tensor import dot, cast, eye, diag, eq, le, ge, gt, all
 from theano.printing import Print
 
-__all__ = ['MvNormal', 'Dirichlet', 'Multinomial', 'Wishart', 'LKJCorr']
+__all__ = ['MvNormal', 'Dirichlet', 'Multinomial', 'Wishart', 'WishartBartlett', 'LKJCorr']
 
 class MvNormal(Continuous):
     """
@@ -187,6 +187,62 @@ def logp(self, X):
             eq(X, X.T)
         )
 
+def WishartBartlett(name, S, nu, is_cholesky=False, return_cholesky=False):
+    """
+    Bartlett decomposition of the Wishart distribution. As the Wishart
+    distribution requires the matrix to be symmetric positive semi-definite
+    it is impossible for MCMC to ever propose acceptable matrices.
+
+    Instead, we can use the Barlett decomposition which samples a lower
+    diagonal matrix. Specifically:
+
+    If L ~ [[sqrt(c_1), 0, ...],
+             [z_21, sqrt(c_1), 0, ...],
+             [z_31, z32, sqrt(c3), ...]]
+    with c_i ~ Chi²(n-i+1) and n_ij ~ N(0, 1), then
+    L * A * A.T * L.T ~ Wishart(L * L.T, nu)
+
+    See http://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition
+    for more information.
+
+    :Parameters:
+      S : ndarray
+        p x p positive definite matrix
+        Or:
+        p x p lower-triangular matrix that is the Cholesky factor
+        of the covariance matrix.
+      nu : int
+        Degrees of freedom, > dim(S).
+      is_cholesky : bool (default=False)
+        Input matrix S is already Cholesky decomposed as S.T * S
+      return_cholesky : bool (default=False)
+        Only return the Cholesky decomposed matrix.
+
+    :Note:
+      This is not a standard Distribution class but follows a similar
+      interface. Besides the Wishart distribution, it will add RVs
+      c and z to your model which make up the matrix.
+    """
+
+    L = S if is_cholesky else scipy.linalg.cholesky(S)
+
+    diag_idx = np.diag_indices_from(S)
+    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)
+    # 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)
+
+    # L * A * A.T * L.T ~ Wishart(L*L.T, nu)
+    if return_cholesky:
+        return pm.Deterministic(name, T.dot(L, A))
+    else:
+        return pm.Deterministic(name, T.dot(T.dot(T.dot(L, A), A.T), L.T))
+
 
 class LKJCorr(Continuous):
     """