Merge pull request #3285 from lingium/feature/issue-3239-yule-simon-rng

add yule_simon_rng

Author: lingium

stan/math/prim/prob.hpp

@@ -320,5 +320,9 @@
 #include <stan/math/prim/prob/wishart_cholesky_rng.hpp>
 #include <stan/math/prim/prob/wishart_lpdf.hpp>
 #include <stan/math/prim/prob/wishart_rng.hpp>
+#include <stan/math/prim/prob/yule_simon_cdf.hpp>
+#include <stan/math/prim/prob/yule_simon_lccdf.hpp>
+#include <stan/math/prim/prob/yule_simon_lcdf.hpp>
 #include <stan/math/prim/prob/yule_simon_lpmf.hpp>
+#include <stan/math/prim/prob/yule_simon_rng.hpp>
 #endif

stan/math/prim/prob/yule_simon_rng.hpp

@@ -0,0 +1,52 @@
+#ifndef STAN_MATH_PRIM_PROB_YULE_SIMON_RNG_HPP
+#define STAN_MATH_PRIM_PROB_YULE_SIMON_RNG_HPP
+
+#include <utility>
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/exp.hpp>
+#include <stan/math/prim/fun/log.hpp>
+#include <stan/math/prim/fun/log1m.hpp>
+#include <stan/math/prim/prob/exponential_rng.hpp>
+#include <stan/math/prim/prob/neg_binomial_rng.hpp>
+
+namespace stan {
+namespace math {
+
+/** \ingroup prob_dists
+ * Return a yule-simon random variate with the given shape parameter,
+ * using the given random number generator.
+ *
+ * alpha can be a scalar or a one-dimensional container.
+ *
+ * @tparam T_alpha type of shape parameter
+ * @tparam RNG type of random number generator
+ *
+ * @param alpha (Sequence of) shape parameter(s)
+ * @param rng random number generator
+ * @return (Sequence of) yule-simon random variate(s)
+ * @throw std::domain_error if alpha is nonpositive
+ */
+template <typename T_alpha, typename RNG>
+inline auto yule_simon_rng(T_alpha&& alpha, RNG& rng) {
+  static constexpr const char* function = "yule_simon_rng";
+  decltype(auto) alpha_ref = to_ref(std::forward<T_alpha>(alpha));
+  check_positive_finite(function, "Shape parameter", alpha_ref);
+
+  auto w = exponential_rng(std::forward<decltype(alpha_ref)>(alpha_ref), rng);
+  auto w_arr = as_array_or_scalar(w);
+  const auto p = stan::math::exp(-w_arr);
+  const auto odds_ratio_p
+      = stan::math::exp(stan::math::log(p) - stan::math::log1m(p));
+
+  if constexpr (is_stan_scalar_v<T_alpha>) {
+    return neg_binomial_rng(1.0, odds_ratio_p, rng) + 1;
+  } else {
+    return to_array_1d(
+        as_array_or_scalar(neg_binomial_rng(1.0, std::move(odds_ratio_p), rng))
+        + 1);
+  }
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

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

@@ -0,0 +1,54 @@
+#include <stan/math/prim.hpp>
+#include <stan/math/prim/prob/yule_simon_rng.hpp>
+#include <test/unit/math/prim/prob/vector_rng_test_helper.hpp>
+#include <test/unit/math/prim/prob/VectorIntRNGTestRig.hpp>
+#include <gtest/gtest.h>
+#include <boost/random/mersenne_twister.hpp>
+#include <boost/math/distributions.hpp>
+#include <limits>
+#include <vector>
+
+class YuleSimonTestRig : public VectorIntRNGTestRig {
+ public:
+  YuleSimonTestRig()
+      : VectorIntRNGTestRig(10000, 10, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+                            {2.1, 4.1, 10.1}, {1, 2, 3}, {-3.0, -2.0, 0.0},
+                            {-3, -1, 0}) {}
+
+  template <typename T1, typename T2, typename T3, typename T_rng>
+  auto generate_samples(const T1& alpha, const T2&, const T3&,
+                        T_rng& rng) const {
+    return stan::math::yule_simon_rng(alpha, rng);
+  }
+
+  template <typename T1>
+  double pmf(int y, T1 alpha, double, double) const {
+    return std::exp(stan::math::yule_simon_lpmf(y, alpha));
+  }
+};
+
+TEST(ProbDistributionsYuleSimon, errorCheck) {
+  check_dist_throws_all_types(YuleSimonTestRig());
+}
+
+TEST(ProbDistributionsYuleSimon, distributionCheck) {
+  check_counts_real(YuleSimonTestRig());
+}
+
+TEST(ProbDistributionsYuleSimon, error_check) {
+  boost::random::mt19937 rng;
+
+  EXPECT_NO_THROW(stan::math::yule_simon_rng(1.0, rng));
+  EXPECT_NO_THROW(stan::math::yule_simon_rng(2.0, rng));
+  EXPECT_NO_THROW(stan::math::yule_simon_rng(10.0, rng));
+
+  EXPECT_THROW(stan::math::yule_simon_rng(-0.5, rng), std::domain_error);
+  EXPECT_THROW(stan::math::yule_simon_rng(0.0, rng), std::domain_error);
+  EXPECT_THROW(stan::math::yule_simon_rng(-10.0, rng), std::domain_error);
+
+  EXPECT_THROW(stan::math::yule_simon_rng(stan::math::positive_infinity(), rng),
+               std::domain_error);
+
+  EXPECT_THROW(stan::math::yule_simon_rng(stan::math::not_a_number(), rng),
+               std::domain_error);
+}