added DirichletDistribution and test

Author: ihincks

doc/source/apiref/distributions.rst

@@ -56,6 +56,9 @@ Specific Distributions
 
 .. autoclass:: BetaBinomialDistribution
     :members:
+    
+.. autoclass:: DirichletDistribution
+    :members:
 
 .. autoclass:: GammaDistribution
     :members:

src/qinfer/distributions.py

@@ -66,6 +66,7 @@
     'SlantedNormalDistribution',
     'LogNormalDistribution',
     'BetaDistribution',
+    'DirichletDistribution',
     'BetaBinomialDistribution',
     'GammaDistribution',
     'GinibreUniform',
@@ -996,6 +997,31 @@ def n_rvs(self):
 
     def sample(self, n=1):
         return self.dist.rvs(size=n)[:, np.newaxis]
+        
+class DirichletDistribution(Distribution):
+    r"""
+    The dirichlet distribution, whose pdf at :math:`x` is proportional to
+    :math:`\prod_i x_i^{\alpha_i-1}`.
+    
+    :param alpha: The list of concentration parameters.
+    """
+    def __init__(self, alpha):
+        self._alpha = np.array(alpha)
+        if self.alpha.ndim != 1:
+            raise ValueError('The input alpha must be a 1D list of concentration parameters.')
+
+        self._dist = st.dirichlet(alpha=self.alpha)
+        
+    @property
+    def alpha(self):
+        return self._alpha
+
+    @property
+    def n_rvs(self):
+        return self._alpha.size
+
+    def sample(self, n=1):
+        return self._dist.rvs(size=n)
 
 class BetaBinomialDistribution(Distribution):
     r"""

src/qinfer/tests/test_distributions.py

@@ -337,6 +337,46 @@ def test_betabinomial_n_rvs(self):
         """
         dist = BetaBinomialDistribution(10, alpha=10,beta=42)
         assert(dist.n_rvs == 1)
+        
+class TestDirichletDistribution(DerandomizedTestCase):
+    """
+    Tests ``DirichletDistribution``
+    """
+
+    ## TEST METHODS ##
+
+    def test_dirichlet_moments(self):
+        """
+        Distributions: Checks that the beta distribution has the right
+        moments, with either of the two input formats
+        """
+        alpha = [1,2,3,4]
+        alpha_np = np.array(alpha)
+        alpha_0 = alpha_np.sum()
+        mean = alpha_np / alpha_0
+        var = alpha_np * (alpha_0 - alpha_np) / (alpha_0 **2 * (alpha_0+1))
+        
+        dist = DirichletDistribution(alpha)
+        samples = dist.sample(100000)
+
+        assert samples.shape == (100000, alpha_np.size)
+        assert_almost_equal(samples.mean(axis=0), mean, 2)
+        assert_almost_equal(samples.var(axis=0), var, 2)
+        
+        alpha = np.array([8,7,5,2,2])
+        alpha_np = np.array(alpha)
+        alpha_0 = alpha_np.sum()
+        mean = alpha_np / alpha_0
+        var = alpha_np * (alpha_0 - alpha_np) / (alpha_0 **2 * (alpha_0+1))
+        
+        dist = DirichletDistribution(alpha)
+        samples = dist.sample(100000)
+
+        assert samples.shape == (100000, alpha_np.size)
+        assert_almost_equal(samples.mean(axis=0), mean, 2)
+        assert_almost_equal(samples.var(axis=0), var, 2)
+
+
 
 class TestGammaDistribution(DerandomizedTestCase):
     """