gamma_distribution.h 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  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_GAMMA_DISTRIBUTION_H
  9. #define _LIBCPP___RANDOM_GAMMA_DISTRIBUTION_H
  10. #include <__config>
  11. #include <__random/exponential_distribution.h>
  12. #include <__random/uniform_real_distribution.h>
  13. #include <cmath>
  14. #include <iosfwd>
  15. #include <limits>
  16. #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
  17. # pragma GCC system_header
  18. #endif
  19. _LIBCPP_PUSH_MACROS
  20. #include <__undef_macros>
  21. _LIBCPP_BEGIN_NAMESPACE_STD
  22. template<class _RealType = double>
  23. class _LIBCPP_TEMPLATE_VIS gamma_distribution
  24. {
  25. public:
  26. // types
  27. typedef _RealType result_type;
  28. class _LIBCPP_TEMPLATE_VIS param_type
  29. {
  30. result_type __alpha_;
  31. result_type __beta_;
  32. public:
  33. typedef gamma_distribution distribution_type;
  34. _LIBCPP_INLINE_VISIBILITY
  35. explicit param_type(result_type __alpha = 1, result_type __beta = 1)
  36. : __alpha_(__alpha), __beta_(__beta) {}
  37. _LIBCPP_INLINE_VISIBILITY
  38. result_type alpha() const {return __alpha_;}
  39. _LIBCPP_INLINE_VISIBILITY
  40. result_type beta() const {return __beta_;}
  41. friend _LIBCPP_INLINE_VISIBILITY
  42. bool operator==(const param_type& __x, const param_type& __y)
  43. {return __x.__alpha_ == __y.__alpha_ && __x.__beta_ == __y.__beta_;}
  44. friend _LIBCPP_INLINE_VISIBILITY
  45. bool operator!=(const param_type& __x, const param_type& __y)
  46. {return !(__x == __y);}
  47. };
  48. private:
  49. param_type __p_;
  50. public:
  51. // constructors and reset functions
  52. #ifndef _LIBCPP_CXX03_LANG
  53. _LIBCPP_INLINE_VISIBILITY
  54. gamma_distribution() : gamma_distribution(1) {}
  55. _LIBCPP_INLINE_VISIBILITY
  56. explicit gamma_distribution(result_type __alpha, result_type __beta = 1)
  57. : __p_(param_type(__alpha, __beta)) {}
  58. #else
  59. _LIBCPP_INLINE_VISIBILITY
  60. explicit gamma_distribution(result_type __alpha = 1,
  61. result_type __beta = 1)
  62. : __p_(param_type(__alpha, __beta)) {}
  63. #endif
  64. _LIBCPP_INLINE_VISIBILITY
  65. explicit gamma_distribution(const param_type& __p)
  66. : __p_(__p) {}
  67. _LIBCPP_INLINE_VISIBILITY
  68. void reset() {}
  69. // generating functions
  70. template<class _URNG>
  71. _LIBCPP_INLINE_VISIBILITY
  72. result_type operator()(_URNG& __g)
  73. {return (*this)(__g, __p_);}
  74. template<class _URNG> result_type operator()(_URNG& __g, const param_type& __p);
  75. // property functions
  76. _LIBCPP_INLINE_VISIBILITY
  77. result_type alpha() const {return __p_.alpha();}
  78. _LIBCPP_INLINE_VISIBILITY
  79. result_type beta() const {return __p_.beta();}
  80. _LIBCPP_INLINE_VISIBILITY
  81. param_type param() const {return __p_;}
  82. _LIBCPP_INLINE_VISIBILITY
  83. void param(const param_type& __p) {__p_ = __p;}
  84. _LIBCPP_INLINE_VISIBILITY
  85. result_type min() const {return 0;}
  86. _LIBCPP_INLINE_VISIBILITY
  87. result_type max() const {return numeric_limits<result_type>::infinity();}
  88. friend _LIBCPP_INLINE_VISIBILITY
  89. bool operator==(const gamma_distribution& __x,
  90. const gamma_distribution& __y)
  91. {return __x.__p_ == __y.__p_;}
  92. friend _LIBCPP_INLINE_VISIBILITY
  93. bool operator!=(const gamma_distribution& __x,
  94. const gamma_distribution& __y)
  95. {return !(__x == __y);}
  96. };
  97. template <class _RealType>
  98. template<class _URNG>
  99. _RealType
  100. gamma_distribution<_RealType>::operator()(_URNG& __g, const param_type& __p)
  101. {
  102. result_type __a = __p.alpha();
  103. uniform_real_distribution<result_type> __gen(0, 1);
  104. exponential_distribution<result_type> __egen;
  105. result_type __x;
  106. if (__a == 1)
  107. __x = __egen(__g);
  108. else if (__a > 1)
  109. {
  110. const result_type __b = __a - 1;
  111. const result_type __c = 3 * __a - result_type(0.75);
  112. while (true)
  113. {
  114. const result_type __u = __gen(__g);
  115. const result_type __v = __gen(__g);
  116. const result_type __w = __u * (1 - __u);
  117. if (__w != 0)
  118. {
  119. const result_type __y = _VSTD::sqrt(__c / __w) *
  120. (__u - result_type(0.5));
  121. __x = __b + __y;
  122. if (__x >= 0)
  123. {
  124. const result_type __z = 64 * __w * __w * __w * __v * __v;
  125. if (__z <= 1 - 2 * __y * __y / __x)
  126. break;
  127. if (_VSTD::log(__z) <= 2 * (__b * _VSTD::log(__x / __b) - __y))
  128. break;
  129. }
  130. }
  131. }
  132. }
  133. else // __a < 1
  134. {
  135. while (true)
  136. {
  137. const result_type __u = __gen(__g);
  138. const result_type __es = __egen(__g);
  139. if (__u <= 1 - __a)
  140. {
  141. __x = _VSTD::pow(__u, 1 / __a);
  142. if (__x <= __es)
  143. break;
  144. }
  145. else
  146. {
  147. const result_type __e = -_VSTD::log((1-__u)/__a);
  148. __x = _VSTD::pow(1 - __a + __a * __e, 1 / __a);
  149. if (__x <= __e + __es)
  150. break;
  151. }
  152. }
  153. }
  154. return __x * __p.beta();
  155. }
  156. template <class _CharT, class _Traits, class _RT>
  157. basic_ostream<_CharT, _Traits>&
  158. operator<<(basic_ostream<_CharT, _Traits>& __os,
  159. const gamma_distribution<_RT>& __x)
  160. {
  161. __save_flags<_CharT, _Traits> __lx(__os);
  162. typedef basic_ostream<_CharT, _Traits> _OStream;
  163. __os.flags(_OStream::dec | _OStream::left | _OStream::fixed |
  164. _OStream::scientific);
  165. _CharT __sp = __os.widen(' ');
  166. __os.fill(__sp);
  167. __os << __x.alpha() << __sp << __x.beta();
  168. return __os;
  169. }
  170. template <class _CharT, class _Traits, class _RT>
  171. basic_istream<_CharT, _Traits>&
  172. operator>>(basic_istream<_CharT, _Traits>& __is,
  173. gamma_distribution<_RT>& __x)
  174. {
  175. typedef gamma_distribution<_RT> _Eng;
  176. typedef typename _Eng::result_type result_type;
  177. typedef typename _Eng::param_type param_type;
  178. __save_flags<_CharT, _Traits> __lx(__is);
  179. typedef basic_istream<_CharT, _Traits> _Istream;
  180. __is.flags(_Istream::dec | _Istream::skipws);
  181. result_type __alpha;
  182. result_type __beta;
  183. __is >> __alpha >> __beta;
  184. if (!__is.fail())
  185. __x.param(param_type(__alpha, __beta));
  186. return __is;
  187. }
  188. _LIBCPP_END_NAMESPACE_STD
  189. _LIBCPP_POP_MACROS
  190. #endif // _LIBCPP___RANDOM_GAMMA_DISTRIBUTION_H