libcxx/include/__random/gamma_distribution.h

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