Merge commit 'b95dd5a46e8bdfdf984d11413dc86bc98306a3d8' into HEAD

stan/math/prim/prob.hpp

@@ -202,6 +202,9 @@
 #include <stan/math/prim/prob/multi_student_t_log.hpp>
 #include <stan/math/prim/prob/multi_student_t_lpdf.hpp>
 #include <stan/math/prim/prob/multi_student_t_rng.hpp>
+#include <stan/math/prim/prob/multinomial_logit_log.hpp>
+#include <stan/math/prim/prob/multinomial_logit_lpmf.hpp>
+#include <stan/math/prim/prob/multinomial_logit_rng.hpp>
 #include <stan/math/prim/prob/multinomial_log.hpp>
 #include <stan/math/prim/prob/multinomial_lpmf.hpp>
 #include <stan/math/prim/prob/multinomial_rng.hpp>

stan/math/prim/prob/multinomial_logit_log.hpp

@@ -0,0 +1,34 @@
+#ifndef STAN_MATH_PRIM_PROB_MULTINOMIAL_LOGIT_LOG_HPP
+#define STAN_MATH_PRIM_PROB_MULTINOMIAL_LOGIT_LOG_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/prob/multinomial_logit_lpmf.hpp>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/** \ingroup multivar_dists
+ * @deprecated use <code>multinomial_logit_lpmf</code>
+ */
+template <bool propto, typename T_beta, typename T_prob = scalar_type_t<T_beta>,
+          require_eigen_col_vector_t<T_beta>* = nullptr>
+return_type_t<T_prob> multinomial_logit_log(const std::vector<int>& ns,
+                                            const T_beta& beta) {
+  return multinomial_logit_lpmf<propto, T_beta>(ns, beta);
+}
+
+/** \ingroup multivar_dists
+ * @deprecated use <code>multinomial_logit_lpmf</code>
+ */
+template <typename T_beta, typename T_prob = scalar_type_t<T_beta>,
+          require_eigen_col_vector_t<T_beta>* = nullptr>
+return_type_t<T_prob> multinomial_logit_log(const std::vector<int>& ns,
+                                            const T_beta& beta) {
+  return multinomial_logit_lpmf<false>(ns, beta);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/prob/multinomial_logit_lpmf.hpp

@@ -0,0 +1,60 @@
+#ifndef STAN_MATH_PRIM_PROB_MULTINOMIAL_LOGIT_LPMF_HPP
+#define STAN_MATH_PRIM_PROB_MULTINOMIAL_LOGIT_LPMF_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/lgamma.hpp>
+#include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/** \ingroup multivar_dists
+ * Multinomial log PMF in log parametrization.
+ * Multinomial(ns| softmax(beta))
+ *
+ * @param ns Array of outcome counts
+ * @param beta Vector of unnormalized log probabilities
+ * @return log probability
+ */
+template <bool propto, typename T_beta, typename T_prob = scalar_type_t<T_beta>,
+          require_eigen_col_vector_t<T_beta>* = nullptr>
+return_type_t<T_prob> multinomial_logit_lpmf(const std::vector<int>& ns,
+                                             const T_beta& beta) {
+  static const char* function = "multinomial_logit_lpmf";
+  check_nonnegative(function, "Number of trials variable", ns);
+  check_finite(function, "log-probabilities parameter", beta);
+  check_size_match(function, "Size of number of trials variable", ns.size(),
+                   "rows of log-probabilities parameter", beta.rows());
+
+  return_type_t<T_prob> lp(0.0);
+
+  decltype(auto) ns_map = as_array_or_scalar(ns);
+
+  if (include_summand<propto>::value) {
+    lp += lgamma(1 + ns_map.sum()) - lgamma(1 + ns_map).sum();
+  }
+
+  if (include_summand<propto, T_prob>::value) {
+    decltype(auto) beta_ref = to_ref(beta);
+    T_prob alpha = log_sum_exp(beta_ref);
+    for (unsigned int i = 0; i < ns.size(); ++i) {
+      if (ns[i] != 0)
+        lp += ns[i] * (beta_ref[i] - alpha);
+    }
+  }
+
+  return lp;
+}
+
+template <typename T_beta, typename T_prob = scalar_type_t<T_beta>,
+          require_eigen_col_vector_t<T_beta>* = nullptr>
+return_type_t<T_prob> multinomial_logit_lpmf(const std::vector<int>& ns,
+                                             const T_beta& beta) {
+  return multinomial_logit_lpmf<false>(ns, beta);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/prob/multinomial_logit_rng.hpp

@@ -0,0 +1,52 @@
+#ifndef STAN_MATH_PRIM_PROB_MULTINOMIAL_LOGIT_RNG_HPP
+#define STAN_MATH_PRIM_PROB_MULTINOMIAL_LOGIT_RNG_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/prob/binomial_rng.hpp>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/** \ingroup multivar_dists
+ * Return a draw from a Multinomial distribution given a
+ * a vector of unnormalized log probabilities and a pseudo-random
+ * number generator.
+ *
+ * @tparam RNG Type of pseudo-random number generator.
+ * @param beta Vector of unnormalized log probabilities.
+ * @param N Total count
+ * @param rng Pseudo-random number generator.
+ * @return Multinomial random variate
+ */
+template <class RNG, typename T_beta,
+          require_eigen_col_vector_t<T_beta>* = nullptr>
+inline std::vector<int> multinomial_logit_rng(const T_beta& beta, int N,
+                                              RNG& rng) {
+  static const char* function = "multinomial_logit_rng";
+  check_finite(function, "Log-probabilities parameter", beta);
+  check_positive(function, "number of trials variables", N);
+
+  plain_type_t<T_beta> theta = softmax(beta);
+  std::vector<int> result(theta.size(), 0);
+  double mass_left = 1.0;
+  int n_left = N;
+
+  for (int k = 0; n_left > 0 && k < theta.size(); ++k) {
+    double p = theta[k] / mass_left;
+    if (p > 1.0) {
+      p = 1.0;
+    }
+    result[k] = binomial_rng(n_left, p, rng);
+    n_left -= result[k];
+    mass_left -= theta[k];
+  }
+
+  return result;
+}  // namespace math
+
+}  // namespace math
+}  // namespace stan
+#endif

test/unit/math/mix/prob/multinomial_logit_test.cpp

@@ -0,0 +1,13 @@
+#include <test/unit/math/test_ad.hpp>
+
+TEST(mathMixScalFun, multinomialLogit) {
+  std::vector<int> ns{0, 1, 2, 3};
+  Eigen::VectorXd beta(4);
+  beta << 0.1, 0.1, 0.5, 0.3;
+
+  auto f = [&ns](const auto& b) {
+    return stan::math::multinomial_logit_lpmf(ns, b);
+  };
+
+  stan::test::expect_ad(f, beta);
+}

test/unit/math/prim/prob/multinomial_logit_log_test.cpp

@@ -0,0 +1,21 @@
+#include <stan/math/prim.hpp>
+#include <gtest/gtest.h>
+#include <vector>
+
+TEST(ProbMultinomialLogit, log_matches_lpmf) {
+  using stan::math::multinomial_logit_log;
+  using stan::math::multinomial_logit_lpmf;
+
+  std::vector<int> ns;
+  ns.push_back(1);
+  ns.push_back(2);
+  ns.push_back(3);
+  Eigen::Matrix<double, Eigen::Dynamic, 1> theta(3, 1);
+  theta << log(0.2), log(0.3), log(0.5);
+  EXPECT_FLOAT_EQ((multinomial_logit_lpmf(ns, theta)),
+                  (multinomial_logit_log(ns, theta)));
+  EXPECT_FLOAT_EQ((multinomial_logit_lpmf<true>(ns, theta)),
+                  (multinomial_logit_log<true>(ns, theta)));
+  EXPECT_FLOAT_EQ((multinomial_logit_lpmf<false>(ns, theta)),
+                  (multinomial_logit_log<false>(ns, theta)));
+}

test/unit/math/prim/prob/multinomial_logit_test.cpp

@@ -0,0 +1,126 @@
+#include <stan/math/prim.hpp>
+#include <gtest/gtest.h>
+#include <boost/random/mersenne_twister.hpp>
+#include <boost/math/distributions.hpp>
+#include <limits>
+#include <vector>
+
+using Eigen::Dynamic;
+using Eigen::Matrix;
+
+TEST(ProbDistributionsMultinomialLogit, RNGSize) {
+  boost::random::mt19937 rng;
+  Matrix<double, Dynamic, 1> beta(5);
+  beta << log(0.3), log(0.1), log(0.2), log(0.2), log(0.2);
+  std::vector<int> sample = stan::math::multinomial_logit_rng(beta, 10, rng);
+  EXPECT_EQ(5U, sample.size());
+}
+
+TEST(ProbDistributionsMultinomialLogit, MultinomialLogit) {
+  std::vector<int> ns;
+  ns.push_back(1);
+  ns.push_back(2);
+  ns.push_back(3);
+  Matrix<double, Dynamic, 1> beta(3, 1);
+  beta << log(0.2), log(0.3), log(0.5);
+  EXPECT_FLOAT_EQ(-2.002481, stan::math::multinomial_logit_log(ns, beta));
+}
+TEST(ProbDistributionsMultinomialLogit, Propto) {
+  std::vector<int> ns;
+  ns.push_back(1);
+  ns.push_back(2);
+  ns.push_back(3);
+  Matrix<double, Dynamic, 1> beta(3, 1);
+  beta << log(0.2), log(0.3), log(0.5);
+  EXPECT_FLOAT_EQ(0.0, stan::math::multinomial_logit_log<true>(ns, beta));
+}
+
+using stan::math::multinomial_logit_log;
+
+TEST(ProbDistributionsMultinomialLogit, error) {
+  double nan = std::numeric_limits<double>::quiet_NaN();
+  double inf = std::numeric_limits<double>::infinity();
+
+  std::vector<int> ns;
+  ns.push_back(1);
+  ns.push_back(2);
+  ns.push_back(3);
+  Matrix<double, Dynamic, 1> beta(3, 1);
+  beta << log(0.2), log(0.3), log(0.5);
+
+  EXPECT_NO_THROW(multinomial_logit_log(ns, beta));
+
+  ns[1] = 0;
+  EXPECT_NO_THROW(multinomial_logit_log(ns, beta));
+  ns[1] = -1;
+  EXPECT_THROW(multinomial_logit_log(ns, beta), std::domain_error);
+  ns[1] = 1;
+
+  beta(0) = nan;
+  EXPECT_THROW(multinomial_logit_log(ns, beta), std::domain_error);
+  beta(0) = inf;
+  EXPECT_THROW(multinomial_logit_log(ns, beta), std::domain_error);
+  beta(0) = -inf;
+  EXPECT_THROW(multinomial_logit_log(ns, beta), std::domain_error);
+
+  beta(0) = 0.2;
+  beta(1) = 0.3;
+  beta(2) = 0.5;
+
+  ns.resize(2);
+  EXPECT_THROW(multinomial_logit_log(ns, beta), std::invalid_argument);
+}
+
+TEST(ProbDistributionsMultinomialLogit, zeros) {
+  double result;
+  std::vector<int> ns;
+  ns.push_back(0);
+  ns.push_back(1);
+  ns.push_back(2);
+  Matrix<double, Dynamic, 1> beta(3, 1);
+  beta << log(0.2), log(0.3), log(0.5);
+
+  result = multinomial_logit_log(ns, beta);
+  EXPECT_FALSE(std::isnan(result));
+
+  std::vector<int> ns2;
+  ns2.push_back(0);
+  ns2.push_back(0);
+  ns2.push_back(0);
+
+  double result2 = multinomial_logit_log(ns2, beta);
+  EXPECT_FLOAT_EQ(0.0, result2);
+}
+
+TEST(ProbDistributionsMultinomialLogit, chiSquareGoodnessFitTest) {
+  boost::random::mt19937 rng;
+  int M = 10;
+  int trials = 1000;
+  int N = M * trials;
+
+  int K = 3;
+  Matrix<double, Dynamic, 1> beta(K);
+  beta << -1, 1, -10;
+  Eigen::VectorXd theta = stan::math::softmax(beta);
+  boost::math::chi_squared mydist(K - 1);
+
+  double expect[K];
+  for (int i = 0; i < K; ++i)
+    expect[i] = N * theta(i);
+
+  int bin[K];
+  for (int i = 0; i < K; ++i)
+    bin[i] = 0;
+
+  for (int count = 0; count < M; ++count) {
+    std::vector<int> a = stan::math::multinomial_logit_rng(beta, trials, rng);
+    for (int i = 0; i < K; ++i)
+      bin[i] += a[i];
+  }
+
+  double chi = 0;
+  for (int j = 0; j < K; j++)
+    chi += ((bin[j] - expect[j]) * (bin[j] - expect[j])) / expect[j];
+
+  EXPECT_TRUE(chi < quantile(complement(mydist, 1e-6)));
+}

test/unit/math/rev/prob/multinomial_logit_test.cpp

@@ -0,0 +1,46 @@
+#include <stan/math/rev.hpp>
+#include <test/unit/math/rev/util.hpp>
+#include <test/unit/math/rev/prob/expect_eq_diffs.hpp>
+#include <gtest/gtest.h>
+#include <string>
+#include <vector>
+
+using stan::math::multinomial_logit_lpmf;
+
+template <typename T_prob>
+void expect_propto(std::vector<int>& ns1, T_prob beta1, std::vector<int>& ns2,
+                   T_prob beta2, std::string message) {
+  expect_eq_diffs(multinomial_logit_lpmf<false>(ns1, beta1),
+                  multinomial_logit_lpmf<false>(ns2, beta2),
+                  multinomial_logit_lpmf<true>(ns1, beta1),
+                  multinomial_logit_lpmf<true>(ns2, beta2), message);
+}
+
+using Eigen::Dynamic;
+using Eigen::Matrix;
+using stan::math::var;
+
+TEST(AgradDistributionsMultinomialLogit, Propto) {
+  std::vector<int> ns;
+  ns.push_back(1);
+  ns.push_back(2);
+  ns.push_back(3);
+  Matrix<var, Dynamic, 1> beta1(3, 1);
+  beta1 << log(0.3), log(0.5), log(0.2);
+  Matrix<var, Dynamic, 1> beta2(3, 1);
+  beta2 << log(0.1), log(0.2), log(0.7);
+
+  expect_propto(ns, beta1, ns, beta2, "var: beta");
+}
+
+TEST(AgradDistributionsMultinomialLogit, check_varis_on_stack) {
+  std::vector<int> ns;
+  ns.push_back(1);
+  ns.push_back(2);
+  ns.push_back(3);
+  Matrix<var, Dynamic, 1> beta(3, 1);
+  beta << log(0.3), log(0.5), log(0.2);
+
+  test::check_varis_on_stack(multinomial_logit_lpmf<false>(ns, beta));
+  test::check_varis_on_stack(multinomial_logit_lpmf<true>(ns, beta));
+}