1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
|
//! Implements [ZIGNOR][1] (Jurgen A. Doornik, 2005, Nuffield College, Oxford).
//!
//! [1]: https://www.doornik.com/research/ziggurat.pdf
//!
//! rust/rand used as a reference;
//!
//! NOTE: This seems interesting but reference code is a bit hard to grok:
//! https://sbarral.github.io/etf.
const std = @import("../std.zig");
const builtin = @import("builtin");
const math = std.math;
const Random = std.Random;
pub fn next_f64(random: Random, comptime tables: ZigTable) f64 {
while (true) {
// We manually construct a float from parts as we can avoid an extra random lookup here by
// using the unused exponent for the lookup table entry.
const bits = random.int(u64);
const i = @as(usize, @as(u8, @truncate(bits)));
const u = blk: {
if (tables.is_symmetric) {
// Generate a value in the range [2, 4) and scale into [-1, 1)
const repr = ((0x3ff + 1) << 52) | (bits >> 12);
break :blk @as(f64, @bitCast(repr)) - 3.0;
} else {
// Generate a value in the range [1, 2) and scale into (0, 1)
const repr = (0x3ff << 52) | (bits >> 12);
break :blk @as(f64, @bitCast(repr)) - (1.0 - math.floatEps(f64) / 2.0);
}
};
const x = u * tables.x[i];
const test_x = if (tables.is_symmetric) @abs(x) else x;
// equivalent to |u| < tables.x[i+1] / tables.x[i] (or u < tables.x[i+1] / tables.x[i])
if (test_x < tables.x[i + 1]) {
return x;
}
if (i == 0) {
return tables.zero_case(random, u);
}
// equivalent to f1 + DRanU() * (f0 - f1) < 1
if (tables.f[i + 1] + (tables.f[i] - tables.f[i + 1]) * random.float(f64) < tables.pdf(x)) {
return x;
}
}
}
pub const ZigTable = struct {
r: f64,
x: [257]f64,
f: [257]f64,
// probability density function used as a fallback
pdf: fn (f64) f64,
// whether the distribution is symmetric
is_symmetric: bool,
// fallback calculation in the case we are in the 0 block
zero_case: fn (Random, f64) f64,
};
// zigNorInit
pub fn ZigTableGen(
comptime is_symmetric: bool,
comptime r: f64,
comptime v: f64,
comptime f: fn (f64) f64,
comptime f_inv: fn (f64) f64,
comptime zero_case: fn (Random, f64) f64,
) ZigTable {
var tables: ZigTable = undefined;
tables.is_symmetric = is_symmetric;
tables.r = r;
tables.pdf = f;
tables.zero_case = zero_case;
tables.x[0] = v / f(r);
tables.x[1] = r;
for (tables.x[2..256], 0..) |*entry, i| {
const last = tables.x[2 + i - 1];
entry.* = f_inv(v / last + f(last));
}
tables.x[256] = 0;
for (tables.f[0..], 0..) |*entry, i| {
entry.* = f(tables.x[i]);
}
return tables;
}
// N(0, 1)
pub const NormDist = blk: {
@setEvalBranchQuota(30000);
break :blk ZigTableGen(true, norm_r, norm_v, norm_f, norm_f_inv, norm_zero_case);
};
pub const norm_r = 3.6541528853610088;
pub const norm_v = 0.00492867323399;
pub fn norm_f(x: f64) f64 {
return @exp(-x * x / 2.0);
}
pub fn norm_f_inv(y: f64) f64 {
return @sqrt(-2.0 * @log(y));
}
pub fn norm_zero_case(random: Random, u: f64) f64 {
var x: f64 = 1;
var y: f64 = 0;
while (-2.0 * y < x * x) {
x = @log(random.float(f64)) / norm_r;
y = @log(random.float(f64));
}
if (u < 0) {
return x - norm_r;
} else {
return norm_r - x;
}
}
test "normal dist sanity" {
var prng = Random.DefaultPrng.init(0);
const random = prng.random();
var i: usize = 0;
while (i < 1000) : (i += 1) {
_ = random.floatNorm(f64);
}
}
// Exp(1)
pub const ExpDist = blk: {
@setEvalBranchQuota(30000);
break :blk ZigTableGen(false, exp_r, exp_v, exp_f, exp_f_inv, exp_zero_case);
};
pub const exp_r = 7.69711747013104972;
pub const exp_v = 0.0039496598225815571993;
pub fn exp_f(x: f64) f64 {
return @exp(-x);
}
pub fn exp_f_inv(y: f64) f64 {
return -@log(y);
}
pub fn exp_zero_case(random: Random, _: f64) f64 {
return exp_r - @log(random.float(f64));
}
test "exp dist smoke test" {
var prng = Random.DefaultPrng.init(0);
const random = prng.random();
var i: usize = 0;
while (i < 1000) : (i += 1) {
_ = random.floatExp(f64);
}
}
test {
_ = NormDist;
}
|