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