exposed dirichlet_multinomial distribution

stan-dev · Dec 17, 2023 · 073e29f · 073e29f
1 parent 6de5732
commit 073e29f
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 0 deletions.
diff --git a/src/middle/Stan_math_signatures.ml b/src/middle/Stan_math_signatures.ml
@@ -235,6 +235,7 @@ let distributions =
   ; (full_lpdf_depr, "cauchy", [DVReal; DVReal; DVReal], SoA)
   ; (full_lpdf_depr, "chi_square", [DVReal; DVReal], SoA)
   ; ([Lpdf; Log], "dirichlet", [DVectors; DVectors], SoA)
+  ; ([Lpmf], "dirichlet_multinomial", [DIntArray; DVector], AoS)
   ; (full_lpmf_depr, "discrete_range", [DVInt; DVInt; DVInt], SoA)
   ; (full_lpdf_depr, "double_exponential", [DVReal; DVReal; DVReal], SoA)
   ; (full_lpdf_depr, "exp_mod_normal", [DVReal; DVReal; DVReal; DVReal], SoA)
@@ -1150,6 +1151,8 @@ let () =
         bare_types)
     (List.range 0 8);
   add_unqualified ("dirichlet_rng", ReturnType UVector, [UVector], AoS);
+  add_unqualified
+    ("dirichlet_multinomial_rng", ReturnType (UArray UInt), [UVector; UInt], AoS);
   add_unqualified ("distance", ReturnType UReal, [UVector; UVector], SoA);
   add_unqualified ("distance", ReturnType UReal, [URowVector; URowVector], SoA);
   add_unqualified ("distance", ReturnType UReal, [UVector; URowVector], SoA);

diff --git a/...distributions/multivariate/discrete/dirichlet_multinomial/dirichlet_multinomial_lpmf.stan b/...distributions/multivariate/discrete/dirichlet_multinomial/dirichlet_multinomial_lpmf.stan
@@ -0,0 +1,21 @@
+data {
+  int d_int;
+  array[d_int] int d_int_array;
+  vector[d_int] d_vector;
+}
+transformed data {
+  real transformed_data_real;
+  transformed_data_real = dirichlet_multinomial_lpmf(d_int_array | d_vector);
+}
+parameters {
+  vector[d_int] p_vector;
+  real y_p;
+}
+transformed parameters {
+  real transformed_param_real;
+  transformed_param_real = dirichlet_multinomial_lpmf(d_int_array | d_vector);
+  transformed_param_real = dirichlet_multinomial_lpmf(d_int_array | p_vector);
+}
+model {
+  y_p ~ normal(0, 1);
+}
diff --git a/test/integration/good/function-signatures/distributions/rngs.stan b/test/integration/good/function-signatures/distributions/rngs.stan
@@ -324,6 +324,7 @@ generated quantities {
   n = categorical_logit_rng(theta);
   ns = multinomial_rng(theta, 20);
   ns = multinomial_logit_rng(theta, 20);
+  ns = dirichlet_multinomial_rng(alpha, 20);
   x = chi_square_rng(alpha);
   x = chi_square_rng(nu);
   x = chi_square_rng(ns);

diff --git a/test/integration/signatures/stan_math_distributions.t b/test/integration/signatures/stan_math_distributions.t
@@ -15,6 +15,7 @@ Display all Stan math distributions exposed in the language
   cauchy: lpdf, rng, ccdf, cdf, log (deprecated)
   chi_square: lpdf, rng, ccdf, cdf, log (deprecated)
   dirichlet: lpdf, log (deprecated)
+  dirichlet_multinomial: lpmf
   discrete_range: lpmf, rng, ccdf, cdf, log (deprecated)
   double_exponential: lpdf, rng, ccdf, cdf, log (deprecated)
   exp_mod_normal: lpdf, rng, ccdf, cdf, log (deprecated)

diff --git a/test/integration/signatures/stan_math_signatures.t b/test/integration/signatures/stan_math_signatures.t
@@ -4254,6 +4254,8 @@ Display all Stan math signatures exposed in the language
   dirichlet_lpdf(array[] row_vector, row_vector) => real
   dirichlet_lpdf(array[] row_vector, array[] vector) => real
   dirichlet_lpdf(array[] row_vector, array[] row_vector) => real
+  dirichlet_multinomial_lpmf(array[] int, vector) => real
+  dirichlet_multinomial_rng(vector, int) => array[] int
   dirichlet_rng(vector) => vector
   discrete_range_ccdf_log(int, int, int) => real
   discrete_range_ccdf_log(int, int, array[] int) => real