1 // Copyright 2017 The Abseil Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 #ifndef ABSL_RANDOM_POISSON_DISTRIBUTION_H_
16 #define ABSL_RANDOM_POISSON_DISTRIBUTION_H_
17
18 #include <cassert>
19 #include <cmath>
20 #include <cstdint>
21 #include <istream>
22 #include <limits>
23 #include <ostream>
24
25 #include "absl/base/config.h"
26 #include "absl/random/internal/fast_uniform_bits.h"
27 #include "absl/random/internal/fastmath.h"
28 #include "absl/random/internal/generate_real.h"
29 #include "absl/random/internal/iostream_state_saver.h"
30 #include "absl/random/internal/traits.h"
31
32 namespace absl {
33 ABSL_NAMESPACE_BEGIN
34
35 // absl::poisson_distribution:
36 // Generates discrete variates conforming to a Poisson distribution.
37 // p(n) = (mean^n / n!) exp(-mean)
38 //
39 // Depending on the parameter, the distribution selects one of the following
40 // algorithms:
41 // * The standard algorithm, attributed to Knuth, extended using a split method
42 // for larger values
43 // * The "Ratio of Uniforms as a convenient method for sampling from classical
44 // discrete distributions", Stadlober, 1989.
45 // http://www.sciencedirect.com/science/article/pii/0377042790903495
46 //
47 // NOTE: param_type.mean() is a double, which permits values larger than
48 // poisson_distribution<IntType>::max(), however this should be avoided and
49 // the distribution results are limited to the max() value.
50 //
51 // The goals of this implementation are to provide good performance while still
52 // being thread-safe: This limits the implementation to not using lgamma
53 // provided by <math.h>.
54 //
55 template <typename IntType = int>
56 class poisson_distribution {
57 public:
58 using result_type = IntType;
59
60 class param_type {
61 public:
62 using distribution_type = poisson_distribution;
63 explicit param_type(double mean = 1.0);
64
mean()65 double mean() const { return mean_; }
66
67 friend bool operator==(const param_type& a, const param_type& b) {
68 return a.mean_ == b.mean_;
69 }
70
71 friend bool operator!=(const param_type& a, const param_type& b) {
72 return !(a == b);
73 }
74
75 private:
76 friend class poisson_distribution;
77
78 double mean_;
79 double emu_; // e ^ -mean_
80 double lmu_; // ln(mean_)
81 double s_;
82 double log_k_;
83 int split_;
84
85 static_assert(random_internal::IsIntegral<IntType>::value,
86 "Class-template absl::poisson_distribution<> must be "
87 "parameterized using an integral type.");
88 };
89
poisson_distribution()90 poisson_distribution() : poisson_distribution(1.0) {}
91
poisson_distribution(double mean)92 explicit poisson_distribution(double mean) : param_(mean) {}
93
poisson_distribution(const param_type & p)94 explicit poisson_distribution(const param_type& p) : param_(p) {}
95
reset()96 void reset() {}
97
98 // generating functions
99 template <typename URBG>
operator()100 result_type operator()(URBG& g) { // NOLINT(runtime/references)
101 return (*this)(g, param_);
102 }
103
104 template <typename URBG>
105 result_type operator()(URBG& g, // NOLINT(runtime/references)
106 const param_type& p);
107
param()108 param_type param() const { return param_; }
param(const param_type & p)109 void param(const param_type& p) { param_ = p; }
110
result_type(min)111 result_type(min)() const { return 0; }
result_type(max)112 result_type(max)() const { return (std::numeric_limits<result_type>::max)(); }
113
mean()114 double mean() const { return param_.mean(); }
115
116 friend bool operator==(const poisson_distribution& a,
117 const poisson_distribution& b) {
118 return a.param_ == b.param_;
119 }
120 friend bool operator!=(const poisson_distribution& a,
121 const poisson_distribution& b) {
122 return a.param_ != b.param_;
123 }
124
125 private:
126 param_type param_;
127 random_internal::FastUniformBits<uint64_t> fast_u64_;
128 };
129
130 // -----------------------------------------------------------------------------
131 // Implementation details follow
132 // -----------------------------------------------------------------------------
133
134 template <typename IntType>
param_type(double mean)135 poisson_distribution<IntType>::param_type::param_type(double mean)
136 : mean_(mean), split_(0) {
137 assert(mean >= 0);
138 assert(mean <=
139 static_cast<double>((std::numeric_limits<result_type>::max)()));
140 // As a defensive measure, avoid large values of the mean. The rejection
141 // algorithm used does not support very large values well. It my be worth
142 // changing algorithms to better deal with these cases.
143 assert(mean <= 1e10);
144 if (mean_ < 10) {
145 // For small lambda, use the knuth method.
146 split_ = 1;
147 emu_ = std::exp(-mean_);
148 } else if (mean_ <= 50) {
149 // Use split-knuth method.
150 split_ = 1 + static_cast<int>(mean_ / 10.0);
151 emu_ = std::exp(-mean_ / static_cast<double>(split_));
152 } else {
153 // Use ratio of uniforms method.
154 constexpr double k2E = 0.7357588823428846;
155 constexpr double kSA = 0.4494580810294493;
156
157 lmu_ = std::log(mean_);
158 double a = mean_ + 0.5;
159 s_ = kSA + std::sqrt(k2E * a);
160 const double mode = std::ceil(mean_) - 1;
161 log_k_ = lmu_ * mode - absl::random_internal::StirlingLogFactorial(mode);
162 }
163 }
164
165 template <typename IntType>
166 template <typename URBG>
167 typename poisson_distribution<IntType>::result_type
operator()168 poisson_distribution<IntType>::operator()(
169 URBG& g, // NOLINT(runtime/references)
170 const param_type& p) {
171 using random_internal::GeneratePositiveTag;
172 using random_internal::GenerateRealFromBits;
173 using random_internal::GenerateSignedTag;
174
175 if (p.split_ != 0) {
176 // Use Knuth's algorithm with range splitting to avoid floating-point
177 // errors. Knuth's algorithm is: Ui is a sequence of uniform variates on
178 // (0,1); return the number of variates required for product(Ui) <
179 // exp(-lambda).
180 //
181 // The expected number of variates required for Knuth's method can be
182 // computed as follows:
183 // The expected value of U is 0.5, so solving for 0.5^n < exp(-lambda) gives
184 // the expected number of uniform variates
185 // required for a given lambda, which is:
186 // lambda = [2, 5, 9, 10, 11, 12, 13, 14, 15, 16, 17]
187 // n = [3, 8, 13, 15, 16, 18, 19, 21, 22, 24, 25]
188 //
189 result_type n = 0;
190 for (int split = p.split_; split > 0; --split) {
191 double r = 1.0;
192 do {
193 r *= GenerateRealFromBits<double, GeneratePositiveTag, true>(
194 fast_u64_(g)); // U(-1, 0)
195 ++n;
196 } while (r > p.emu_);
197 --n;
198 }
199 return n;
200 }
201
202 // Use ratio of uniforms method.
203 //
204 // Let u ~ Uniform(0, 1), v ~ Uniform(-1, 1),
205 // a = lambda + 1/2,
206 // s = 1.5 - sqrt(3/e) + sqrt(2(lambda + 1/2)/e),
207 // x = s * v/u + a.
208 // P(floor(x) = k | u^2 < f(floor(x))/k), where
209 // f(m) = lambda^m exp(-lambda)/ m!, for 0 <= m, and f(m) = 0 otherwise,
210 // and k = max(f).
211 const double a = p.mean_ + 0.5;
212 for (;;) {
213 const double u = GenerateRealFromBits<double, GeneratePositiveTag, false>(
214 fast_u64_(g)); // U(0, 1)
215 const double v = GenerateRealFromBits<double, GenerateSignedTag, false>(
216 fast_u64_(g)); // U(-1, 1)
217
218 const double x = std::floor(p.s_ * v / u + a);
219 if (x < 0) continue; // f(negative) = 0
220 const double rhs = x * p.lmu_;
221 // clang-format off
222 double s = (x <= 1.0) ? 0.0
223 : (x == 2.0) ? 0.693147180559945
224 : absl::random_internal::StirlingLogFactorial(x);
225 // clang-format on
226 const double lhs = 2.0 * std::log(u) + p.log_k_ + s;
227 if (lhs < rhs) {
228 return x > static_cast<double>((max)())
229 ? (max)()
230 : static_cast<result_type>(x); // f(x)/k >= u^2
231 }
232 }
233 }
234
235 template <typename CharT, typename Traits, typename IntType>
236 std::basic_ostream<CharT, Traits>& operator<<(
237 std::basic_ostream<CharT, Traits>& os, // NOLINT(runtime/references)
238 const poisson_distribution<IntType>& x) {
239 auto saver = random_internal::make_ostream_state_saver(os);
240 os.precision(random_internal::stream_precision_helper<double>::kPrecision);
241 os << x.mean();
242 return os;
243 }
244
245 template <typename CharT, typename Traits, typename IntType>
246 std::basic_istream<CharT, Traits>& operator>>(
247 std::basic_istream<CharT, Traits>& is, // NOLINT(runtime/references)
248 poisson_distribution<IntType>& x) { // NOLINT(runtime/references)
249 using param_type = typename poisson_distribution<IntType>::param_type;
250
251 auto saver = random_internal::make_istream_state_saver(is);
252 double mean = random_internal::read_floating_point<double>(is);
253 if (!is.fail()) {
254 x.param(param_type(mean));
255 }
256 return is;
257 }
258
259 ABSL_NAMESPACE_END
260 } // namespace absl
261
262 #endif // ABSL_RANDOM_POISSON_DISTRIBUTION_H_
263