Merge pull request #138 from ihincks/feature-dirichlet-distribution

DirichletDistribution

Author: cgranade

doc/source/apiref/distributions.rst

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

src/qinfer/distributions.py

@@ -76,6 +76,7 @@
     'SlantedNormalDistribution',
     'LogNormalDistribution',
     'BetaDistribution',
+    'DirichletDistribution',
     'BetaBinomialDistribution',
     'GammaDistribution',
     'GinibreUniform',
@@ -1006,6 +1007,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_concrete_models.py

@@ -64,7 +64,7 @@
     NormalDistribution,
     BetaDistribution, UniformDistribution,
     PostselectedDistribution,
-    ConstrainedSumDistribution,
+    ConstrainedSumDistribution, DirichletDistribution,
     DirectViewParallelizedModel,
     GaussianHyperparameterizedModel
 )
@@ -291,8 +291,7 @@ class TestMultinomialModel(ConcreteModelTest, DerandomizedTestCase):
     def instantiate_model(self):
         return MultinomialModel(NDieModel(n=6))
     def instantiate_prior(self):
-        unif = UniformDistribution(np.array([[0,1],[0,1],[0,1],[0,1],[0,1],[0,1]]))
-        return ConstrainedSumDistribution(unif, desired_total=1)
+        return DirichletDistribution([1,2,3,10,1,3])
     def instantiate_expparams(self):
         return np.arange(10).astype(self.model.expparams_dtype)
 

src/qinfer/tests/test_distributions.py

@@ -347,6 +347,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 dirichlet 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):
     """