const std = @import("../std.zig");
const builtin = @import("builtin");
const math = std.math;
const Random = std.rand.Random;
pub fn next_f64(random: Random, comptime tables: ZigTable) f64 {
while (true) {
const bits = random.int(u64);
const i = @as(usize, @truncate(u8, bits));
const u = blk: {
if (tables.is_symmetric) {
const repr = ((0x3ff + 1) << 52) | (bits >> 12);
break :blk @bitCast(f64, repr) - 3.0;
} else {
const repr = (0x3ff << 52) | (bits >> 12);
break :blk @bitCast(f64, repr) - (1.0 - math.floatEps(f64) / 2.0);
}
};
const x = u * tables.x[i];
const test_x = if (tables.is_symmetric) @fabs(x) else x;
if (test_x < tables.x[i + 1]) {
return x;
}
if (i == 0) {
return tables.zero_case(random, u);
}
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,
pdf: fn (f64) f64,
is_symmetric: bool,
zero_case: fn (Random, f64) f64,
};
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;
}
pub const NormDist = blk: {
@setEvalBranchQuota(30000);
break :blk ZigTableGen(true, norm_r, norm_v, norm_f, norm_f_inv, norm_zero_case);
};
const norm_r = 3.6541528853610088;
const norm_v = 0.00492867323399;
fn norm_f(x: f64) f64 {
return @exp(-x * x / 2.0);
}
fn norm_f_inv(y: f64) f64 {
return @sqrt(-2.0 * @log(y));
}
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 = std.rand.DefaultPrng.init(0);
const random = prng.random();
var i: usize = 0;
while (i < 1000) : (i += 1) {
_ = random.floatNorm(f64);
}
}
pub const ExpDist = blk: {
@setEvalBranchQuota(30000);
break :blk ZigTableGen(false, exp_r, exp_v, exp_f, exp_f_inv, exp_zero_case);
};
const exp_r = 7.69711747013104972;
const exp_v = 0.0039496598225815571993;
fn exp_f(x: f64) f64 {
return @exp(-x);
}
fn exp_f_inv(y: f64) f64 {
return -@log(y);
}
fn exp_zero_case(random: Random, _: f64) f64 {
return exp_r - @log(random.float(f64));
}
test "exp dist smoke test" {
var prng = std.rand.DefaultPrng.init(0);
const random = prng.random();
var i: usize = 0;
while (i < 1000) : (i += 1) {
_ = random.floatExp(f64);
}
}
test {
_ = NormDist;
}