|
| 1 | +#ifndef STAN_MATH_PRIM_FUN_HYPERGEOMETRIC_3F2_HPP |
| 2 | +#define STAN_MATH_PRIM_FUN_HYPERGEOMETRIC_3F2_HPP |
| 3 | + |
| 4 | +#include <stan/math/prim/meta.hpp> |
| 5 | +#include <stan/math/prim/err.hpp> |
| 6 | +#include <stan/math/prim/fun/append_row.hpp> |
| 7 | +#include <stan/math/prim/fun/as_array_or_scalar.hpp> |
| 8 | +#include <stan/math/prim/fun/to_vector.hpp> |
| 9 | +#include <stan/math/prim/fun/constants.hpp> |
| 10 | +#include <stan/math/prim/fun/fabs.hpp> |
| 11 | +#include <stan/math/prim/fun/hypergeometric_pFq.hpp> |
| 12 | +#include <stan/math/prim/fun/sum.hpp> |
| 13 | +#include <stan/math/prim/fun/sign.hpp> |
| 14 | +#include <stan/math/prim/fun/value_of_rec.hpp> |
| 15 | + |
| 16 | +namespace stan { |
| 17 | +namespace math { |
| 18 | +namespace internal { |
| 19 | +template <typename Ta, typename Tb, typename Tz, |
| 20 | + typename T_return = return_type_t<Ta, Tb, Tz>, |
| 21 | + typename ArrayAT = Eigen::Array<scalar_type_t<Ta>, 3, 1>, |
| 22 | + typename ArrayBT = Eigen::Array<scalar_type_t<Ta>, 3, 1>, |
| 23 | + require_all_vector_t<Ta, Tb>* = nullptr, |
| 24 | + require_stan_scalar_t<Tz>* = nullptr> |
| 25 | +T_return hypergeometric_3F2_infsum(const Ta& a, const Tb& b, const Tz& z, |
| 26 | + double precision = 1e-6, |
| 27 | + int max_steps = 1e5) { |
| 28 | + ArrayAT a_array = as_array_or_scalar(a); |
| 29 | + ArrayBT b_array = append_row(as_array_or_scalar(b), 1.0); |
| 30 | + check_3F2_converges("hypergeometric_3F2", a_array[0], a_array[1], a_array[2], |
| 31 | + b_array[0], b_array[1], z); |
| 32 | + |
| 33 | + T_return t_acc = 1.0; |
| 34 | + T_return log_t = 0.0; |
| 35 | + T_return log_z = log(fabs(z)); |
| 36 | + Eigen::ArrayXi a_signs = sign(value_of_rec(a_array)); |
| 37 | + Eigen::ArrayXi b_signs = sign(value_of_rec(b_array)); |
| 38 | + plain_type_t<decltype(a_array)> apk = a_array; |
| 39 | + plain_type_t<decltype(b_array)> bpk = b_array; |
| 40 | + int z_sign = sign(value_of_rec(z)); |
| 41 | + int t_sign = z_sign * a_signs.prod() * b_signs.prod(); |
| 42 | + |
| 43 | + int k = 0; |
| 44 | + while (k <= max_steps && log_t >= log(precision)) { |
| 45 | + // Replace zero values with 1 prior to taking the log so that we accumulate |
| 46 | + // 0.0 rather than -inf |
| 47 | + const auto& abs_apk = math::fabs((apk == 0).select(1.0, apk)); |
| 48 | + const auto& abs_bpk = math::fabs((bpk == 0).select(1.0, bpk)); |
| 49 | + T_return p = sum(log(abs_apk)) - sum(log(abs_bpk)); |
| 50 | + if (p == NEGATIVE_INFTY) { |
| 51 | + return t_acc; |
| 52 | + } |
| 53 | + |
| 54 | + log_t += p + log_z; |
| 55 | + t_acc += t_sign * exp(log_t); |
| 56 | + |
| 57 | + if (is_inf(t_acc)) { |
| 58 | + throw_domain_error("hypergeometric_3F2", "sum (output)", t_acc, |
| 59 | + "overflow hypergeometric function did not converge."); |
| 60 | + } |
| 61 | + k++; |
| 62 | + apk.array() += 1.0; |
| 63 | + bpk.array() += 1.0; |
| 64 | + a_signs = sign(value_of_rec(apk)); |
| 65 | + b_signs = sign(value_of_rec(bpk)); |
| 66 | + t_sign = a_signs.prod() * b_signs.prod() * t_sign; |
| 67 | + } |
| 68 | + if (k == max_steps) { |
| 69 | + throw_domain_error("hypergeometric_3F2", "k (internal counter)", max_steps, |
| 70 | + "exceeded iterations, hypergeometric function did not ", |
| 71 | + "converge."); |
| 72 | + } |
| 73 | + return t_acc; |
| 74 | +} |
| 75 | +} // namespace internal |
| 76 | + |
| 77 | +/** |
| 78 | + * Hypergeometric function (3F2). |
| 79 | + * |
| 80 | + * Function reference: http://dlmf.nist.gov/16.2 |
| 81 | + * |
| 82 | + * \f[ |
| 83 | + * _3F_2 \left( |
| 84 | + * \begin{matrix}a_1 a_2 a[2] \\ b_1 b_2\end{matrix}; z |
| 85 | + * \right) = \sum_k=0^\infty |
| 86 | + * \frac{(a_1)_k(a_2)_k(a_3)_k}{(b_1)_k(b_2)_k}\frac{z^k}{k!} \f] |
| 87 | + * |
| 88 | + * Where $(a_1)_k$ is an upper shifted factorial. |
| 89 | + * |
| 90 | + * Calculate the hypergeometric function (3F2) as the power series |
| 91 | + * directly to within <code>precision</code> or until |
| 92 | + * <code>max_steps</code> terms. |
| 93 | + * |
| 94 | + * This function does not have a closed form but will converge if: |
| 95 | + * - <code>|z|</code> is less than 1 |
| 96 | + * - <code>|z|</code> is equal to one and <code>b[0] + b[1] < a[0] + a[1] + |
| 97 | + * a[2]</code> This function is a rational polynomial if |
| 98 | + * - <code>a[0]</code>, <code>a[1]</code>, or <code>a[2]</code> is a |
| 99 | + * non-positive integer |
| 100 | + * This function can be treated as a rational polynomial if |
| 101 | + * - <code>b[0]</code> or <code>b[1]</code> is a non-positive integer |
| 102 | + * and the series is terminated prior to the final term. |
| 103 | + * |
| 104 | + * @tparam Ta type of Eigen/Std vector 'a' arguments |
| 105 | + * @tparam Tb type of Eigen/Std vector 'b' arguments |
| 106 | + * @tparam Tz type of z argument |
| 107 | + * @param[in] a Always called with a[1] > 1, a[2] <= 0 |
| 108 | + * @param[in] b Always called with int b[0] < |a[2]|, <= 1) |
| 109 | + * @param[in] z z (is always called with 1 from beta binomial cdfs) |
| 110 | + * @param[in] precision precision of the infinite sum. defaults to 1e-6 |
| 111 | + * @param[in] max_steps number of steps to take. defaults to 1e5 |
| 112 | + * @return The 3F2 generalized hypergeometric function applied to the |
| 113 | + * arguments {a1, a2, a3}, {b1, b2} |
| 114 | + */ |
| 115 | +template <typename Ta, typename Tb, typename Tz, |
| 116 | + require_all_vector_t<Ta, Tb>* = nullptr, |
| 117 | + require_stan_scalar_t<Tz>* = nullptr> |
| 118 | +auto hypergeometric_3F2(const Ta& a, const Tb& b, const Tz& z) { |
| 119 | + check_3F2_converges("hypergeometric_3F2", a[0], a[1], a[2], b[0], b[1], z); |
| 120 | + // Boost's pFq throws convergence errors in some cases, fallback to naive |
| 121 | + // infinite-sum approach (tests pass for these) |
| 122 | + if (z == 1.0 && (sum(b) - sum(a)) < 0.0) { |
| 123 | + return internal::hypergeometric_3F2_infsum(a, b, z); |
| 124 | + } |
| 125 | + return hypergeometric_pFq(to_vector(a), to_vector(b), z); |
| 126 | +} |
| 127 | + |
| 128 | +/** |
| 129 | + * Hypergeometric function (3F2). |
| 130 | + * |
| 131 | + * Overload for initializer_list inputs |
| 132 | + * |
| 133 | + * @tparam Ta type of scalar 'a' arguments |
| 134 | + * @tparam Tb type of scalar 'b' arguments |
| 135 | + * @tparam Tz type of z argument |
| 136 | + * @param[in] a Always called with a[1] > 1, a[2] <= 0 |
| 137 | + * @param[in] b Always called with int b[0] < |a[2]|, <= 1) |
| 138 | + * @param[in] z z (is always called with 1 from beta binomial cdfs) |
| 139 | + * @param[in] precision precision of the infinite sum. defaults to 1e-6 |
| 140 | + * @param[in] max_steps number of steps to take. defaults to 1e5 |
| 141 | + * @return The 3F2 generalized hypergeometric function applied to the |
| 142 | + * arguments {a1, a2, a3}, {b1, b2} |
| 143 | + */ |
| 144 | +template <typename Ta, typename Tb, typename Tz, |
| 145 | + require_all_stan_scalar_t<Ta, Tb, Tz>* = nullptr> |
| 146 | +auto hypergeometric_3F2(const std::initializer_list<Ta>& a, |
| 147 | + const std::initializer_list<Tb>& b, const Tz& z) { |
| 148 | + return hypergeometric_3F2(std::vector<Ta>(a), std::vector<Tb>(b), z); |
| 149 | +} |
| 150 | + |
| 151 | +} // namespace math |
| 152 | +} // namespace stan |
| 153 | +#endif |
0 commit comments