xref: /llvm-project/libcxx/include/__random/binomial_distribution.h (revision 84ce230e4298672bb5247170d6183b31aa06fc4b)
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_BINOMIAL_DISTRIBUTION_H
10 #define _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H
11 
12 #include <__config>
13 #include <__random/is_valid.h>
14 #include <__random/uniform_real_distribution.h>
15 #include <cmath>
16 #include <iosfwd>
17 
18 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
19 #  pragma GCC system_header
20 #endif
21 
22 _LIBCPP_PUSH_MACROS
23 #include <__undef_macros>
24 
25 _LIBCPP_BEGIN_NAMESPACE_STD
26 
27 template <class _IntType = int>
28 class _LIBCPP_TEMPLATE_VIS binomial_distribution {
29   static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");
30 
31 public:
32   // types
33   typedef _IntType result_type;
34 
35   class _LIBCPP_TEMPLATE_VIS param_type {
36     result_type __t_;
37     double __p_;
38     double __pr_;
39     double __odds_ratio_;
40     result_type __r0_;
41 
42   public:
43     typedef binomial_distribution distribution_type;
44 
45     _LIBCPP_HIDE_FROM_ABI explicit param_type(result_type __t = 1, double __p = 0.5);
46 
47     _LIBCPP_HIDE_FROM_ABI result_type t() const { return __t_; }
48     _LIBCPP_HIDE_FROM_ABI double p() const { return __p_; }
49 
50     friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {
51       return __x.__t_ == __y.__t_ && __x.__p_ == __y.__p_;
52     }
53     friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }
54 
55     friend class binomial_distribution;
56   };
57 
58 private:
59   param_type __p_;
60 
61 public:
62   // constructors and reset functions
63 #ifndef _LIBCPP_CXX03_LANG
64   _LIBCPP_HIDE_FROM_ABI binomial_distribution() : binomial_distribution(1) {}
65   _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(result_type __t, double __p = 0.5)
66       : __p_(param_type(__t, __p)) {}
67 #else
68   _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(result_type __t = 1, double __p = 0.5)
69       : __p_(param_type(__t, __p)) {}
70 #endif
71   _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(const param_type& __p) : __p_(__p) {}
72   _LIBCPP_HIDE_FROM_ABI void reset() {}
73 
74   // generating functions
75   template <class _URNG>
76   _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {
77     return (*this)(__g, __p_);
78   }
79   template <class _URNG>
80   _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);
81 
82   // property functions
83   _LIBCPP_HIDE_FROM_ABI result_type t() const { return __p_.t(); }
84   _LIBCPP_HIDE_FROM_ABI double p() const { return __p_.p(); }
85 
86   _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }
87   _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }
88 
89   _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }
90   _LIBCPP_HIDE_FROM_ABI result_type max() const { return t(); }
91 
92   friend _LIBCPP_HIDE_FROM_ABI bool operator==(const binomial_distribution& __x, const binomial_distribution& __y) {
93     return __x.__p_ == __y.__p_;
94   }
95   friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const binomial_distribution& __x, const binomial_distribution& __y) {
96     return !(__x == __y);
97   }
98 };
99 
100 // The LLVM C library provides this with conflicting `noexcept` attributes.
101 #if !defined(_LIBCPP_MSVCRT_LIKE) && !defined(__LLVM_LIBC__)
102 extern "C" double lgamma_r(double, int*);
103 #endif
104 
105 inline _LIBCPP_HIDE_FROM_ABI double __libcpp_lgamma(double __d) {
106 #if defined(_LIBCPP_MSVCRT_LIKE) || defined(__LLVM_LIBC__)
107   return lgamma(__d);
108 #else
109   int __sign;
110   return lgamma_r(__d, &__sign);
111 #endif
112 }
113 
114 template <class _IntType>
115 binomial_distribution<_IntType>::param_type::param_type(result_type __t, double __p) : __t_(__t), __p_(__p) {
116   if (0 < __p_ && __p_ < 1) {
117     __r0_ = static_cast<result_type>((__t_ + 1) * __p_);
118     __pr_ = std::exp(
119         std::__libcpp_lgamma(__t_ + 1.) - std::__libcpp_lgamma(__r0_ + 1.) - std::__libcpp_lgamma(__t_ - __r0_ + 1.) +
120         __r0_ * std::log(__p_) + (__t_ - __r0_) * std::log(1 - __p_));
121     __odds_ratio_ = __p_ / (1 - __p_);
122   }
123 }
124 
125 // Reference: Kemp, C.D. (1986). `A modal method for generating binomial
126 //           variables', Commun. Statist. - Theor. Meth. 15(3), 805-813.
127 template <class _IntType>
128 template <class _URNG>
129 _IntType binomial_distribution<_IntType>::operator()(_URNG& __g, const param_type& __pr) {
130   static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");
131   if (__pr.__t_ == 0 || __pr.__p_ == 0)
132     return 0;
133   if (__pr.__p_ == 1)
134     return __pr.__t_;
135   uniform_real_distribution<double> __gen;
136   double __u = __gen(__g) - __pr.__pr_;
137   if (__u < 0)
138     return __pr.__r0_;
139   double __pu      = __pr.__pr_;
140   double __pd      = __pu;
141   result_type __ru = __pr.__r0_;
142   result_type __rd = __ru;
143   while (true) {
144     bool __break = true;
145     if (__rd >= 1) {
146       __pd *= __rd / (__pr.__odds_ratio_ * (__pr.__t_ - __rd + 1));
147       __u -= __pd;
148       __break = false;
149       if (__u < 0)
150         return __rd - 1;
151     }
152     if (__rd != 0)
153       --__rd;
154     ++__ru;
155     if (__ru <= __pr.__t_) {
156       __pu *= (__pr.__t_ - __ru + 1) * __pr.__odds_ratio_ / __ru;
157       __u -= __pu;
158       __break = false;
159       if (__u < 0)
160         return __ru;
161     }
162     if (__break)
163       return 0;
164   }
165 }
166 
167 template <class _CharT, class _Traits, class _IntType>
168 _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
169 operator<<(basic_ostream<_CharT, _Traits>& __os, const binomial_distribution<_IntType>& __x) {
170   __save_flags<_CharT, _Traits> __lx(__os);
171   typedef basic_ostream<_CharT, _Traits> _OStream;
172   __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);
173   _CharT __sp = __os.widen(' ');
174   __os.fill(__sp);
175   return __os << __x.t() << __sp << __x.p();
176 }
177 
178 template <class _CharT, class _Traits, class _IntType>
179 _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
180 operator>>(basic_istream<_CharT, _Traits>& __is, binomial_distribution<_IntType>& __x) {
181   typedef binomial_distribution<_IntType> _Eng;
182   typedef typename _Eng::result_type result_type;
183   typedef typename _Eng::param_type param_type;
184   __save_flags<_CharT, _Traits> __lx(__is);
185   typedef basic_istream<_CharT, _Traits> _Istream;
186   __is.flags(_Istream::dec | _Istream::skipws);
187   result_type __t;
188   double __p;
189   __is >> __t >> __p;
190   if (!__is.fail())
191     __x.param(param_type(__t, __p));
192   return __is;
193 }
194 
195 _LIBCPP_END_NAMESPACE_STD
196 
197 _LIBCPP_POP_MACROS
198 
199 #endif // _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H
200