diff --git a/stan/math/prim/prob.hpp b/stan/math/prim/prob.hpp index e04c7e87172..be55ca6a8ef 100644 --- a/stan/math/prim/prob.hpp +++ b/stan/math/prim/prob.hpp @@ -316,5 +316,5 @@ #include #include #include - +#include #endif diff --git a/stan/math/prim/prob/yule_simon_lpmf.hpp b/stan/math/prim/prob/yule_simon_lpmf.hpp new file mode 100644 index 00000000000..441d47dab13 --- /dev/null +++ b/stan/math/prim/prob/yule_simon_lpmf.hpp @@ -0,0 +1,95 @@ +#ifndef STAN_MATH_PRIM_PROB_YULE_SIMON_LPMF_HPP +#define STAN_MATH_PRIM_PROB_YULE_SIMON_LPMF_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** \ingroup prob_dists + * Returns the log PMF of the Yule-Simon distribution with shape parameter. + * Given containers of matching sizes, returns the log sum of probabilities. + * + * @tparam T_n type of outcome variable + * @tparam T_alpha type of shape parameter + * + * @param n outcome variable + * @param alpha shape parameter + * @return log probability or log sum of probabilities + * @throw std::domain_error if alpha fails to be positive + * @throw std::invalid_argument if container sizes mismatch + */ +template * = nullptr> +inline return_type_t yule_simon_lpmf(const T_n &n, + const T_alpha &alpha) { + using std::log; + using T_partials_return = partials_return_t; + using T_n_ref = ref_type_t; + using T_alpha_ref = ref_type_t; + static constexpr const char *function = "yule_simon_lpmf"; + check_consistent_sizes(function, "Failures variable", n, "Shape parameter", + alpha); + if (size_zero(n, alpha)) { + return 0.0; + } + + T_n_ref n_ref = n; + T_alpha_ref alpha_ref = alpha; + check_greater_or_equal(function, "Outcome variable", n_ref, 1); + check_positive_finite(function, "Shape parameter", alpha_ref); + + if constexpr (!include_summand::value) { + return 0.0; + } + + auto ops_partials = make_partials_propagator(alpha_ref); + + scalar_seq_view n_vec(n_ref); + scalar_seq_view alpha_vec(alpha_ref); + const size_t max_size_seq_view = max_size(n_ref, alpha_ref); + T_partials_return logp(0.0); + if constexpr (include_summand::value) { + if constexpr (is_stan_scalar_v) { + logp += lgamma(n_ref) * max_size_seq_view; + } + } + for (size_t i = 0; i < max_size_seq_view; i++) { + if constexpr (include_summand::value) { + if constexpr (!is_stan_scalar_v) { + logp += lgamma(n_vec.val(i)); + } + } + T_partials_return alpha_plus_one = alpha_vec.val(i) + 1.0; + logp += log(alpha_vec.val(i)) + lgamma(alpha_plus_one) + - lgamma(n_vec.val(i) + alpha_plus_one); + if constexpr (!is_constant::value) { + partials<0>(ops_partials)[i] += 1.0 / alpha_vec.val(i) + + digamma(alpha_plus_one) + - digamma(n_vec.val(i) + alpha_plus_one); + } + } + return ops_partials.build(logp); +} + +template +inline return_type_t yule_simon_lpmf(const T_n &n, + const T_alpha &alpha) { + return yule_simon_lpmf(n, alpha); +} + +} // namespace math +} // namespace stan +#endif diff --git a/test/prob/yule_simon/yule_simon_test.hpp b/test/prob/yule_simon/yule_simon_test.hpp new file mode 100644 index 00000000000..61784486944 --- /dev/null +++ b/test/prob/yule_simon/yule_simon_test.hpp @@ -0,0 +1,71 @@ +// Arguments: Ints, Doubles +#include +#include +#include + +using std::numeric_limits; +using std::vector; + +class AgradDistributionsYuleSimon : public AgradDistributionTest { + public: + void valid_values(vector >& parameters, + vector& log_prob) { + vector param(2); + + param[0] = 5; // n + param[1] = 20.0; // alpha + parameters.push_back(param); + log_prob.push_back(-9.494202658325099); // expected log_prob + + param[0] = 10; // n + param[1] = 5.5; // alpha + parameters.push_back(param); + log_prob.push_back(-9.108616882863778); // expected log_prob + } + + void invalid_values(vector& index, vector& value) { + // n + index.push_back(0U); + value.push_back(-1); + + index.push_back(0U); + value.push_back(0); + + // alpha + index.push_back(1U); + value.push_back(0.0); + + index.push_back(1U); + value.push_back(-1.0); + + index.push_back(1U); + value.push_back(std::numeric_limits::infinity()); + } + + template + stan::return_type_t log_prob(const T_n& n, const T_alpha& alpha, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::yule_simon_lpmf(n, alpha); + } + + template + stan::return_type_t log_prob(const T_n& n, const T_alpha& alpha, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::yule_simon_lpmf(n, alpha); + } + + template + stan::return_type_t log_prob_function(const T_n& n, + const T_alpha& alpha, + const T2&, const T3&, + const T4&, const T5&) { + using stan::math::lbeta; + using std::log; + return log(alpha) + lbeta(n, alpha + 1.0); + } +};