fixed dirichlet and added example

Author: jsalvatier

examples/dirichlet.py

@@ -0,0 +1,25 @@
+import numpy as np
+from pymc import  *
+from pymc.distributions.multivariate import Dirichlet
+
+
+model = Model()
+Var = model.Var
+Data = model.Data 
+
+k = 5
+a = constant(np.array([2,3.,4, 2,2]))
+
+p, p_m1 = model.TransformedVar(
+                 'p', Dirichlet(k,a), 
+                 transforms.simplex, shape = k - 1)
+
+
+H = model.d2logpc()
+
+s = find_MAP(model)
+
+step = hmc_step(model, model.vars, H(s))
+trace, _,t = sample(1000, step, s) 
+
+

pymc/distributions/multivariate.py

@@ -44,8 +44,12 @@ def Dirichlet(k, a):
         \cdot\left(1-\sum_{i=1}^{k-1}x_i\right)^\theta_k
 
     :Parameters:
-        theta : array
-            An (n,k) or (1,k) array > 0.
+        k : scalar int
+            k > 1
+        a : float tensor 
+            a > 0
+            concentration parameters
+            last index is the k index
 
     :Support: 
         x : vector 
@@ -55,16 +59,17 @@ def Dirichlet(k, a):
         Only the first `k-1` elements of `x` are expected. Can be used
         as a parent of Multinomial and Categorical nevertheless.
     """
+
     support = 'continuous'
 
-    a = ones(k) * a
+    a = ones([k]) * a
     def logp(value):
 
         #only defined for sum(value) == 1 
         return bound(
-                sum((a -1)*log(value)) + gammaln(sum(a)) - sum(gammaln(a)),
+                sum(logpow(value, a -1) - gammaln(a), axis = 0) + gammaln(sum(a)),
 
-                k > 2,
+                k > 1,
                 a > 0)
 
     mean = a/sum(a)