poisson_distribution.h 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241
  1. //===----------------------------------------------------------------------===//
  2. //
  3. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  4. // See https://llvm.org/LICENSE.txt for license information.
  5. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  6. //
  7. //===----------------------------------------------------------------------===//
  8. #ifndef _LIBCPP___RANDOM_POISSON_DISTRIBUTION_H
  9. #define _LIBCPP___RANDOM_POISSON_DISTRIBUTION_H
  10. #include <__config>
  11. #include <__random/clamp_to_integral.h>
  12. #include <__random/exponential_distribution.h>
  13. #include <__random/is_valid.h>
  14. #include <__random/normal_distribution.h>
  15. #include <__random/uniform_real_distribution.h>
  16. #include <cmath>
  17. #include <iosfwd>
  18. #include <limits>
  19. #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
  20. # pragma GCC system_header
  21. #endif
  22. _LIBCPP_PUSH_MACROS
  23. #include <__undef_macros>
  24. _LIBCPP_BEGIN_NAMESPACE_STD
  25. template <class _IntType = int>
  26. class _LIBCPP_TEMPLATE_VIS poisson_distribution {
  27. static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");
  28. public:
  29. // types
  30. typedef _IntType result_type;
  31. class _LIBCPP_TEMPLATE_VIS param_type {
  32. double __mean_;
  33. double __s_;
  34. double __d_;
  35. double __l_;
  36. double __omega_;
  37. double __c0_;
  38. double __c1_;
  39. double __c2_;
  40. double __c3_;
  41. double __c_;
  42. public:
  43. typedef poisson_distribution distribution_type;
  44. _LIBCPP_HIDE_FROM_ABI explicit param_type(double __mean = 1.0);
  45. _LIBCPP_HIDE_FROM_ABI double mean() const { return __mean_; }
  46. friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {
  47. return __x.__mean_ == __y.__mean_;
  48. }
  49. friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }
  50. friend class poisson_distribution;
  51. };
  52. private:
  53. param_type __p_;
  54. public:
  55. // constructors and reset functions
  56. #ifndef _LIBCPP_CXX03_LANG
  57. _LIBCPP_HIDE_FROM_ABI poisson_distribution() : poisson_distribution(1.0) {}
  58. _LIBCPP_HIDE_FROM_ABI explicit poisson_distribution(double __mean) : __p_(__mean) {}
  59. #else
  60. _LIBCPP_HIDE_FROM_ABI explicit poisson_distribution(double __mean = 1.0) : __p_(__mean) {}
  61. #endif
  62. _LIBCPP_HIDE_FROM_ABI explicit poisson_distribution(const param_type& __p) : __p_(__p) {}
  63. _LIBCPP_HIDE_FROM_ABI void reset() {}
  64. // generating functions
  65. template <class _URNG>
  66. _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {
  67. return (*this)(__g, __p_);
  68. }
  69. template <class _URNG>
  70. _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);
  71. // property functions
  72. _LIBCPP_HIDE_FROM_ABI double mean() const { return __p_.mean(); }
  73. _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }
  74. _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }
  75. _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }
  76. _LIBCPP_HIDE_FROM_ABI result_type max() const { return numeric_limits<result_type>::max(); }
  77. friend _LIBCPP_HIDE_FROM_ABI bool operator==(const poisson_distribution& __x, const poisson_distribution& __y) {
  78. return __x.__p_ == __y.__p_;
  79. }
  80. friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const poisson_distribution& __x, const poisson_distribution& __y) {
  81. return !(__x == __y);
  82. }
  83. };
  84. template <class _IntType>
  85. poisson_distribution<_IntType>::param_type::param_type(double __mean)
  86. // According to the standard `inf` is a valid input, but it causes the
  87. // distribution to hang, so we replace it with the maximum representable
  88. // mean.
  89. : __mean_(isinf(__mean) ? numeric_limits<double>::max() : __mean) {
  90. if (__mean_ < 10) {
  91. __s_ = 0;
  92. __d_ = 0;
  93. __l_ = std::exp(-__mean_);
  94. __omega_ = 0;
  95. __c3_ = 0;
  96. __c2_ = 0;
  97. __c1_ = 0;
  98. __c0_ = 0;
  99. __c_ = 0;
  100. } else {
  101. __s_ = std::sqrt(__mean_);
  102. __d_ = 6 * __mean_ * __mean_;
  103. __l_ = std::trunc(__mean_ - 1.1484);
  104. __omega_ = .3989423 / __s_;
  105. double __b1 = .4166667E-1 / __mean_;
  106. double __b2 = .3 * __b1 * __b1;
  107. __c3_ = .1428571 * __b1 * __b2;
  108. __c2_ = __b2 - 15. * __c3_;
  109. __c1_ = __b1 - 6. * __b2 + 45. * __c3_;
  110. __c0_ = 1. - __b1 + 3. * __b2 - 15. * __c3_;
  111. __c_ = .1069 / __mean_;
  112. }
  113. }
  114. template <class _IntType>
  115. template <class _URNG>
  116. _IntType poisson_distribution<_IntType>::operator()(_URNG& __urng, const param_type& __pr) {
  117. static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");
  118. double __tx;
  119. uniform_real_distribution<double> __urd;
  120. if (__pr.__mean_ < 10) {
  121. __tx = 0;
  122. for (double __p = __urd(__urng); __p > __pr.__l_; ++__tx)
  123. __p *= __urd(__urng);
  124. } else {
  125. double __difmuk;
  126. double __g = __pr.__mean_ + __pr.__s_ * normal_distribution<double>()(__urng);
  127. double __u;
  128. if (__g > 0) {
  129. __tx = std::trunc(__g);
  130. if (__tx >= __pr.__l_)
  131. return std::__clamp_to_integral<result_type>(__tx);
  132. __difmuk = __pr.__mean_ - __tx;
  133. __u = __urd(__urng);
  134. if (__pr.__d_ * __u >= __difmuk * __difmuk * __difmuk)
  135. return std::__clamp_to_integral<result_type>(__tx);
  136. }
  137. exponential_distribution<double> __edist;
  138. for (bool __using_exp_dist = false; true; __using_exp_dist = true) {
  139. double __e;
  140. if (__using_exp_dist || __g <= 0) {
  141. double __t;
  142. do {
  143. __e = __edist(__urng);
  144. __u = __urd(__urng);
  145. __u += __u - 1;
  146. __t = 1.8 + (__u < 0 ? -__e : __e);
  147. } while (__t <= -.6744);
  148. __tx = std::trunc(__pr.__mean_ + __pr.__s_ * __t);
  149. __difmuk = __pr.__mean_ - __tx;
  150. __using_exp_dist = true;
  151. }
  152. double __px;
  153. double __py;
  154. if (__tx < 10 && __tx >= 0) {
  155. const double __fac[] = {1, 1, 2, 6, 24, 120, 720, 5040, 40320, 362880};
  156. __px = -__pr.__mean_;
  157. __py = std::pow(__pr.__mean_, (double)__tx) / __fac[static_cast<int>(__tx)];
  158. } else {
  159. double __del = .8333333E-1 / __tx;
  160. __del -= 4.8 * __del * __del * __del;
  161. double __v = __difmuk / __tx;
  162. if (std::abs(__v) > 0.25)
  163. __px = __tx * std::log(1 + __v) - __difmuk - __del;
  164. else
  165. __px = __tx * __v * __v *
  166. (((((((.1250060 * __v + -.1384794) * __v + .1421878) * __v + -.1661269) * __v + .2000118) * __v +
  167. -.2500068) *
  168. __v +
  169. .3333333) *
  170. __v +
  171. -.5) -
  172. __del;
  173. __py = .3989423 / std::sqrt(__tx);
  174. }
  175. double __r = (0.5 - __difmuk) / __pr.__s_;
  176. double __r2 = __r * __r;
  177. double __fx = -0.5 * __r2;
  178. double __fy = __pr.__omega_ * (((__pr.__c3_ * __r2 + __pr.__c2_) * __r2 + __pr.__c1_) * __r2 + __pr.__c0_);
  179. if (__using_exp_dist) {
  180. if (__pr.__c_ * std::abs(__u) <= __py * std::exp(__px + __e) - __fy * std::exp(__fx + __e))
  181. break;
  182. } else {
  183. if (__fy - __u * __fy <= __py * std::exp(__px - __fx))
  184. break;
  185. }
  186. }
  187. }
  188. return std::__clamp_to_integral<result_type>(__tx);
  189. }
  190. template <class _CharT, class _Traits, class _IntType>
  191. _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
  192. operator<<(basic_ostream<_CharT, _Traits>& __os, const poisson_distribution<_IntType>& __x) {
  193. __save_flags<_CharT, _Traits> __lx(__os);
  194. typedef basic_ostream<_CharT, _Traits> _OStream;
  195. __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);
  196. return __os << __x.mean();
  197. }
  198. template <class _CharT, class _Traits, class _IntType>
  199. _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
  200. operator>>(basic_istream<_CharT, _Traits>& __is, poisson_distribution<_IntType>& __x) {
  201. typedef poisson_distribution<_IntType> _Eng;
  202. typedef typename _Eng::param_type param_type;
  203. __save_flags<_CharT, _Traits> __lx(__is);
  204. typedef basic_istream<_CharT, _Traits> _Istream;
  205. __is.flags(_Istream::dec | _Istream::skipws);
  206. double __mean;
  207. __is >> __mean;
  208. if (!__is.fail())
  209. __x.param(param_type(__mean));
  210. return __is;
  211. }
  212. _LIBCPP_END_NAMESPACE_STD
  213. _LIBCPP_POP_MACROS
  214. #endif // _LIBCPP___RANDOM_POISSON_DISTRIBUTION_H