poisson_distribution.h 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280
  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. {
  28. static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be an integer type larger than char");
  29. public:
  30. // types
  31. typedef _IntType result_type;
  32. class _LIBCPP_TEMPLATE_VIS param_type
  33. {
  34. double __mean_;
  35. double __s_;
  36. double __d_;
  37. double __l_;
  38. double __omega_;
  39. double __c0_;
  40. double __c1_;
  41. double __c2_;
  42. double __c3_;
  43. double __c_;
  44. public:
  45. typedef poisson_distribution distribution_type;
  46. explicit param_type(double __mean = 1.0);
  47. _LIBCPP_INLINE_VISIBILITY
  48. double mean() const {return __mean_;}
  49. friend _LIBCPP_INLINE_VISIBILITY
  50. bool operator==(const param_type& __x, const param_type& __y)
  51. {return __x.__mean_ == __y.__mean_;}
  52. friend _LIBCPP_INLINE_VISIBILITY
  53. bool operator!=(const param_type& __x, const param_type& __y)
  54. {return !(__x == __y);}
  55. friend class poisson_distribution;
  56. };
  57. private:
  58. param_type __p_;
  59. public:
  60. // constructors and reset functions
  61. #ifndef _LIBCPP_CXX03_LANG
  62. _LIBCPP_INLINE_VISIBILITY
  63. poisson_distribution() : poisson_distribution(1.0) {}
  64. _LIBCPP_INLINE_VISIBILITY
  65. explicit poisson_distribution(double __mean)
  66. : __p_(__mean) {}
  67. #else
  68. _LIBCPP_INLINE_VISIBILITY
  69. explicit poisson_distribution(double __mean = 1.0)
  70. : __p_(__mean) {}
  71. #endif
  72. _LIBCPP_INLINE_VISIBILITY
  73. explicit poisson_distribution(const param_type& __p) : __p_(__p) {}
  74. _LIBCPP_INLINE_VISIBILITY
  75. void reset() {}
  76. // generating functions
  77. template<class _URNG>
  78. _LIBCPP_INLINE_VISIBILITY
  79. result_type operator()(_URNG& __g)
  80. {return (*this)(__g, __p_);}
  81. template<class _URNG> result_type operator()(_URNG& __g, const param_type& __p);
  82. // property functions
  83. _LIBCPP_INLINE_VISIBILITY
  84. double mean() const {return __p_.mean();}
  85. _LIBCPP_INLINE_VISIBILITY
  86. param_type param() const {return __p_;}
  87. _LIBCPP_INLINE_VISIBILITY
  88. void param(const param_type& __p) {__p_ = __p;}
  89. _LIBCPP_INLINE_VISIBILITY
  90. result_type min() const {return 0;}
  91. _LIBCPP_INLINE_VISIBILITY
  92. result_type max() const {return numeric_limits<result_type>::max();}
  93. friend _LIBCPP_INLINE_VISIBILITY
  94. bool operator==(const poisson_distribution& __x,
  95. const poisson_distribution& __y)
  96. {return __x.__p_ == __y.__p_;}
  97. friend _LIBCPP_INLINE_VISIBILITY
  98. bool operator!=(const poisson_distribution& __x,
  99. const poisson_distribution& __y)
  100. {return !(__x == __y);}
  101. };
  102. template<class _IntType>
  103. poisson_distribution<_IntType>::param_type::param_type(double __mean)
  104. // According to the standard `inf` is a valid input, but it causes the
  105. // distribution to hang, so we replace it with the maximum representable
  106. // mean.
  107. : __mean_(isinf(__mean) ? numeric_limits<double>::max() : __mean)
  108. {
  109. if (__mean_ < 10)
  110. {
  111. __s_ = 0;
  112. __d_ = 0;
  113. __l_ = _VSTD::exp(-__mean_);
  114. __omega_ = 0;
  115. __c3_ = 0;
  116. __c2_ = 0;
  117. __c1_ = 0;
  118. __c0_ = 0;
  119. __c_ = 0;
  120. }
  121. else
  122. {
  123. __s_ = _VSTD::sqrt(__mean_);
  124. __d_ = 6 * __mean_ * __mean_;
  125. __l_ = _VSTD::trunc(__mean_ - 1.1484);
  126. __omega_ = .3989423 / __s_;
  127. double __b1_ = .4166667E-1 / __mean_;
  128. double __b2_ = .3 * __b1_ * __b1_;
  129. __c3_ = .1428571 * __b1_ * __b2_;
  130. __c2_ = __b2_ - 15. * __c3_;
  131. __c1_ = __b1_ - 6. * __b2_ + 45. * __c3_;
  132. __c0_ = 1. - __b1_ + 3. * __b2_ - 15. * __c3_;
  133. __c_ = .1069 / __mean_;
  134. }
  135. }
  136. template <class _IntType>
  137. template<class _URNG>
  138. _IntType
  139. poisson_distribution<_IntType>::operator()(_URNG& __urng, const param_type& __pr)
  140. {
  141. static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");
  142. double __tx;
  143. uniform_real_distribution<double> __urd;
  144. if (__pr.__mean_ < 10)
  145. {
  146. __tx = 0;
  147. for (double __p = __urd(__urng); __p > __pr.__l_; ++__tx)
  148. __p *= __urd(__urng);
  149. }
  150. else
  151. {
  152. double __difmuk;
  153. double __g = __pr.__mean_ + __pr.__s_ * normal_distribution<double>()(__urng);
  154. double __u;
  155. if (__g > 0)
  156. {
  157. __tx = _VSTD::trunc(__g);
  158. if (__tx >= __pr.__l_)
  159. return _VSTD::__clamp_to_integral<result_type>(__tx);
  160. __difmuk = __pr.__mean_ - __tx;
  161. __u = __urd(__urng);
  162. if (__pr.__d_ * __u >= __difmuk * __difmuk * __difmuk)
  163. return _VSTD::__clamp_to_integral<result_type>(__tx);
  164. }
  165. exponential_distribution<double> __edist;
  166. for (bool __using_exp_dist = false; true; __using_exp_dist = true)
  167. {
  168. double __e;
  169. if (__using_exp_dist || __g <= 0)
  170. {
  171. double __t;
  172. do
  173. {
  174. __e = __edist(__urng);
  175. __u = __urd(__urng);
  176. __u += __u - 1;
  177. __t = 1.8 + (__u < 0 ? -__e : __e);
  178. } while (__t <= -.6744);
  179. __tx = _VSTD::trunc(__pr.__mean_ + __pr.__s_ * __t);
  180. __difmuk = __pr.__mean_ - __tx;
  181. __using_exp_dist = true;
  182. }
  183. double __px;
  184. double __py;
  185. if (__tx < 10 && __tx >= 0)
  186. {
  187. const double __fac[] = {1, 1, 2, 6, 24, 120, 720, 5040,
  188. 40320, 362880};
  189. __px = -__pr.__mean_;
  190. __py = _VSTD::pow(__pr.__mean_, (double)__tx) / __fac[static_cast<int>(__tx)];
  191. }
  192. else
  193. {
  194. double __del = .8333333E-1 / __tx;
  195. __del -= 4.8 * __del * __del * __del;
  196. double __v = __difmuk / __tx;
  197. if (_VSTD::abs(__v) > 0.25)
  198. __px = __tx * _VSTD::log(1 + __v) - __difmuk - __del;
  199. else
  200. __px = __tx * __v * __v * (((((((.1250060 * __v + -.1384794) *
  201. __v + .1421878) * __v + -.1661269) * __v + .2000118) *
  202. __v + -.2500068) * __v + .3333333) * __v + -.5) - __del;
  203. __py = .3989423 / _VSTD::sqrt(__tx);
  204. }
  205. double __r = (0.5 - __difmuk) / __pr.__s_;
  206. double __r2 = __r * __r;
  207. double __fx = -0.5 * __r2;
  208. double __fy = __pr.__omega_ * (((__pr.__c3_ * __r2 + __pr.__c2_) *
  209. __r2 + __pr.__c1_) * __r2 + __pr.__c0_);
  210. if (__using_exp_dist)
  211. {
  212. if (__pr.__c_ * _VSTD::abs(__u) <= __py * _VSTD::exp(__px + __e) -
  213. __fy * _VSTD::exp(__fx + __e))
  214. break;
  215. }
  216. else
  217. {
  218. if (__fy - __u * __fy <= __py * _VSTD::exp(__px - __fx))
  219. break;
  220. }
  221. }
  222. }
  223. return _VSTD::__clamp_to_integral<result_type>(__tx);
  224. }
  225. template <class _CharT, class _Traits, class _IntType>
  226. basic_ostream<_CharT, _Traits>&
  227. operator<<(basic_ostream<_CharT, _Traits>& __os,
  228. const poisson_distribution<_IntType>& __x)
  229. {
  230. __save_flags<_CharT, _Traits> __lx(__os);
  231. typedef basic_ostream<_CharT, _Traits> _OStream;
  232. __os.flags(_OStream::dec | _OStream::left | _OStream::fixed |
  233. _OStream::scientific);
  234. return __os << __x.mean();
  235. }
  236. template <class _CharT, class _Traits, class _IntType>
  237. basic_istream<_CharT, _Traits>&
  238. operator>>(basic_istream<_CharT, _Traits>& __is,
  239. poisson_distribution<_IntType>& __x)
  240. {
  241. typedef poisson_distribution<_IntType> _Eng;
  242. typedef typename _Eng::param_type param_type;
  243. __save_flags<_CharT, _Traits> __lx(__is);
  244. typedef basic_istream<_CharT, _Traits> _Istream;
  245. __is.flags(_Istream::dec | _Istream::skipws);
  246. double __mean;
  247. __is >> __mean;
  248. if (!__is.fail())
  249. __x.param(param_type(__mean));
  250. return __is;
  251. }
  252. _LIBCPP_END_NAMESPACE_STD
  253. _LIBCPP_POP_MACROS
  254. #endif // _LIBCPP___RANDOM_POISSON_DISTRIBUTION_H