add the missing distributions for models (#45)

conjugate/distributions.py

@@ -150,6 +150,42 @@ def dist(self):
         return VectorizedDist(self.alpha, dist=stats.dirichlet)
 
 
+@dataclass
+class Multinomial(SliceMixin):
+    """Multinomial distribution.
+
+    Args:
+        n: number of trials
+        p: probability of success
+
+    """
+
+    n: NUMERIC
+    p: NUMERIC
+
+    @property
+    def dist(self):
+        return stats.multinomial(n=self.n, p=self.p)
+
+
+@dataclass
+class DirichletMultinomial(SliceMixin):
+    """Dirichlet multinomial distribution.
+
+    Args:
+        alpha: shape parameter
+        n: number of trials
+
+    """
+
+    alpha: NUMERIC
+    n: NUMERIC
+
+    @property
+    def dist(self):
+        return stats.dirichlet_multinomial(self.alpha, self.n)
+
+
 @dataclass
 class Exponential(ContinuousPlotDistMixin, SliceMixin):
     """Exponential distribution.

conjugate/models.py

@@ -10,6 +10,7 @@
 from conjugate.distributions import (
     Beta,
     Dirichlet,
+    DirichletMultinomial,
     Gamma,
     NegativeBinomial,
     BetaNegativeBinomial,
@@ -208,6 +209,20 @@ def categorical_dirichlet(x: NUMERIC, dirichlet_prior: Dirichlet) -> Dirichlet:
     return Dirichlet(alpha=alpha_post)
 
 
+def categorical_dirichlet_posterior_predictive(
+    dirichlet: Dirichlet, n: NUMERIC = 1
+) -> DirichletMultinomial:
+    """Posterior predictive distribution of Categorical model with Dirichlet prior.
+
+    Args:
+        dirichlet: Dirichlet distribution
+        n: Number of trials for each sample, defaults to 1.
+
+    """
+
+    return DirichletMultinomial(n=n, alpha=dirichlet.alpha)
+
+
 def get_multi_categorical_dirichlet_posterior_params(
     alpha_prior: NUMERIC, x: NUMERIC
 ) -> NUMERIC:
@@ -230,6 +245,20 @@ def multinomial_dirichlet(x: NUMERIC, dirichlet_prior: Dirichlet) -> Dirichlet:
     return Dirichlet(alpha=alpha_post)
 
 
+def multinomial_dirichlet_posterior_predictive(
+    dirichlet: Dirichlet, n: NUMERIC = 1
+) -> DirichletMultinomial:
+    """Posterior predictive distribution of Multinomial model with Dirichlet prior.
+
+    Args:
+        dirichlet: Dirichlet distribution
+        n: Number of trials for each sample, defaults to 1.
+
+    """
+
+    return DirichletMultinomial(n=n, alpha=dirichlet.alpha)
+
+
 def get_poisson_gamma_posterior_params(
     alpha: NUMERIC, beta: NUMERIC, x_total: NUMERIC, n: NUMERIC
 ) -> Tuple[NUMERIC, NUMERIC]: