faster binomial and logarithmic rng

PauloCampana · Feb 3, 2024 · fcf9383 · fcf9383
1 parent 061bd6a
commit fcf9383
Show file tree

Hide file tree

Showing 13 changed files with 557 additions and 521 deletions.
diff --git a/docs/data-astNodes.js b/docs/data-astNodes.js
diff --git a/docs/data-calls.js b/docs/data-calls.js
diff --git a/docs/data-comptimeExprs.js b/docs/data-comptimeExprs.js
diff --git a/docs/data-decls.js b/docs/data-decls.js
diff --git a/docs/data-exprs.js b/docs/data-exprs.js
diff --git a/docs/data-types.js b/docs/data-types.js
diff --git a/docs/src/builtin/builtin.zig.html b/docs/src/builtin/builtin.zig.html
@@ -116,7 +116,7 @@
 <span class="line" id="L2"><span class="tok-comment">/// Zig version. When writing code that supports multiple versions of Zig, prefer</span></span>
 <span class="line" id="L3"><span class="tok-comment">/// feature detection (i.e. with `@hasDecl` or `@hasField`) over version checks.</span></span>
 <span class="line" id="L4"><span class="tok-kw">pub</span> <span class="tok-kw">const</span> zig_version = std.SemanticVersion.parse(zig_version_string) <span class="tok-kw">catch</span> <span class="tok-kw">unreachable</span>;</span>
-<span class="line" id="L5"><span class="tok-kw">pub</span> <span class="tok-kw">const</span> zig_version_string = <span class="tok-str">&quot;0.12.0-dev.2341+92211135f&quot;</span>;</span>
+<span class="line" id="L5"><span class="tok-kw">pub</span> <span class="tok-kw">const</span> zig_version_string = <span class="tok-str">&quot;0.12.0-dev.2543+9eda6ccef&quot;</span>;</span>
 <span class="line" id="L6"><span class="tok-kw">pub</span> <span class="tok-kw">const</span> zig_backend = std.builtin.CompilerBackend.stage2_llvm;</span>
 <span class="line" id="L7"></span>
 <span class="line" id="L8"><span class="tok-kw">pub</span> <span class="tok-kw">const</span> output_mode = std.builtin.OutputMode.Exe;</span>

diff --git a/docs/src/root/distribution/binomial.zig.html b/docs/src/root/distribution/binomial.zig.html
diff --git a/docs/src/root/distribution/logarithmic.zig.html b/docs/src/root/distribution/logarithmic.zig.html
diff --git a/docs/src/root/distribution/negativeBinomial.zig.html b/docs/src/root/distribution/negativeBinomial.zig.html
diff --git a/src/distribution/binomial.zig b/src/distribution/binomial.zig
@@ -59,22 +59,30 @@ pub fn quantile(p: f64, size: u64, prob: f64) f64 {
     assert(0 <= prob and prob <= 1);
     assert(0 <= p and p <= 1);
     const n = @as(f64, @floatFromInt(size));
-    if (p == 0 or p == 1) {
-        return n * p;
-    }
-    if (prob == 0 or prob == 1 or size == 0) {
-        return n * prob;
+    if (p == 0) {
+        return 0;
+    } else if (p == 1) {
+        return n;
     }
-    const pq = prob / (1 - prob);
+    const qrob = 1 - prob;
     const mean = n * prob;
-    if (mean < 500) {
-        const initial_mass = std.math.pow(f64, 1 - prob, n);
-        return linearSearch(p, n, pq, initial_mass);
+    if (prob == 0 or size == 0) {
+        return 0;
+    } else if (prob == 1) {
+        return n;
+    } else if (mean < 500) {
+        if (prob < 0.5) {
+            const initial_mass = std.math.pow(f64, qrob, n);
+            return linearSearch(p, n, prob / qrob, initial_mass);
+        } else {
+            const initial_mass = std.math.pow(f64, prob, n);
+            return n - linearSearch(p, n, qrob / prob, initial_mass);
+        }
     } else {
         const initial_bino = @ceil(mean);
         const initial_mass = density(initial_bino, size, prob);
         const initial_cumu = probability(initial_bino, size, prob);
-        return guidedSearch(p, n, pq, initial_bino, initial_mass, initial_cumu);
+        return guidedSearch(p, n, prob / qrob, initial_bino, initial_mass, initial_cumu);
     }
 }
 
@@ -83,45 +91,66 @@ pub const random = struct {
     pub fn single(generator: std.rand.Random, size: u64, prob: f64) f64 {
         assert(0 <= prob and prob <= 1);
         const n = @as(f64, @floatFromInt(size));
-        const pq = prob / (1 - prob);
+        const qrob = 1 - prob;
         const mean = n * prob;
-        if (prob == 0 or prob == 1 or size == 0) {
-            return mean;
+        if (prob == 0 or size == 0) {
+            return 0;
+        } else if (prob == 1) {
+            return n;
         } else if (prob == 0.5) {
             const mask = (@as(u64, 1) << @truncate(@mod(size, 64))) - 1;
             return bitCount(generator, mask, size);
-        } else if (mean < 500) {
-            const initial_mass = std.math.pow(f64, 1 - prob, n);
-            const uni = generator.float(f64);
-            return linearSearch(uni, n, pq, initial_mass);
+        } else if (mean < 1000) {
+            if (prob < 0.5) {
+                const initial_mass = std.math.pow(f64, qrob, n);
+                const uni = generator.float(f64);
+                return linearSearch(uni, n, prob / qrob, initial_mass);
+            } else {
+                const initial_mass = std.math.pow(f64, prob, n);
+                const uni = generator.float(f64);
+                return n - linearSearch(uni, n, qrob / prob, initial_mass);
+            }
         } else {
             const initial_bino = @ceil(mean);
             const initial_mass = density(initial_bino, size, prob);
             const initial_cumu = probability(initial_bino, size, prob);
             const uni = generator.float(f64);
-            return guidedSearch(uni, n, pq, initial_bino, initial_mass, initial_cumu);
+            return guidedSearch(uni, n, prob / qrob, initial_bino, initial_mass, initial_cumu);
         }
     }
 
     pub fn fill(buffer: []f64, generator: std.rand.Random, size: u64, prob: f64) []f64 {
         assert(0 <= prob and prob <= 1);
         const n = @as(f64, @floatFromInt(size));
-        const pq = prob / (1 - prob);
+        const qrob = 1 - prob;
         const mean = n * prob;
-        if (prob == 0 or prob == 1 or size == 0) {
-            @memset(buffer, mean);
+        if (prob == 0 or size == 0) {
+            @memset(buffer, 0);
+        } else if (prob == 1) {
+            @memset(buffer, n);
         } else if (prob == 0.5) {
             const mask = (@as(u64, 1) << @truncate(@mod(size, 64))) - 1;
             for (buffer) |*x| {
                 x.* = bitCount(generator, mask, size);
             }
-        } else if (buffer.len < 10) {
-            const initial_mass = std.math.pow(f64, 1 - prob, n);
-            for (buffer) |*x| {
-                const uni = generator.float(f64);
-                x.* = linearSearch(uni, n, pq, initial_mass);
+        } else if (buffer.len < 20) {
+            if (prob < 0.5) {
+                const pq = prob / qrob;
+                const initial_mass = std.math.pow(f64, qrob, n);
+                for (buffer) |*x| {
+                    const uni = generator.float(f64);
+                    x.* = linearSearch(uni, n, pq, initial_mass);
+                }
+            } else {
+                const qp = qrob / prob;
+                const initial_mass = std.math.pow(f64, prob, n);
+                for (buffer) |*x| {
+                    const uni = generator.float(f64);
+                    x.* = n - linearSearch(uni, n, qp, initial_mass);
+                }
             }
         } else {
+            const pq = prob / qrob;
             const initial_bino = @ceil(mean);
             const initial_mass = density(initial_bino, size, prob);
             const initial_cumu = probability(initial_bino, size, prob);

diff --git a/src/distribution/logarithmic.zig b/src/distribution/logarithmic.zig
@@ -34,9 +34,9 @@ pub fn probability(q: f64, prob: f64) f64 {
     if (q == inf) {
         return 1;
     }
+    var loga: f64 = 1;
     var mass = prob / -std.math.log1p(-prob);
     var cumu = mass;
-    var loga: f64 = 1;
     for (1..@intFromFloat(q)) |_| {
         const num = prob * loga;
         loga += 1;
@@ -60,29 +60,30 @@ pub fn quantile(p: f64, prob: f64) f64 {
     return linearSearch(p, prob, initial_mass);
 }
 
-/// Uses Kemp's algorithm.
-/// http://luc.devroye.org/chapter_ten.pdf page 547.
+/// Uses the quantile function.
 pub const random = struct {
     pub fn single(generator: std.rand.Random, prob: f64) f64 {
         assert(0 < prob and prob < 1);
-        const r = std.math.log1p(-prob);
-        return kemp(generator, prob, r);
+        const initial_mass = prob / -std.math.log1p(-prob);
+        const uni = generator.float(f64);
+        return linearSearch(uni, prob, initial_mass);
     }
 
     pub fn fill(buffer: []f64, generator: std.rand.Random, prob: f64) []f64 {
         assert(0 < prob and prob < 1);
-        const r = std.math.log1p(-prob);
+        const initial_mass = prob / -std.math.log1p(-prob);
         for (buffer) |*x| {
-            x.* = kemp(generator, prob, r);
+            const uni = generator.float(f64);
+            x.* = linearSearch(uni, prob, initial_mass);
         }
         return buffer;
     }
 };
 
 inline fn linearSearch(p: f64, prob: f64, initial_mass: f64) f64 {
+    var loga: f64 = 1;
     var mass = initial_mass;
     var cumu = mass;
-    var loga: f64 = 1;
     while (cumu <= p) {
         const num = prob * loga;
         loga += 1;
@@ -92,22 +93,6 @@ inline fn linearSearch(p: f64, prob: f64, initial_mass: f64) f64 {
     return loga;
 }
 
-inline fn kemp(generator: std.rand.Random, prob: f64, r: f64) f64 {
-    const uni1 = generator.float(f64);
-    if (uni1 >= prob) {
-        return 1;
-    }
-    const uni2 = generator.float(f64);
-    const q = -std.math.expm1(r * uni2);
-    if (uni1 > q) {
-        return 2;
-    }
-    if (q * q < uni1 and uni1 <= q) {
-        return 1;
-    }
-    return @floor(1 + @log(uni1) / @log(q));
-}
-
 const expectEqual = std.testing.expectEqual;
 const expectApproxEqRel = std.testing.expectApproxEqRel;
 const eps = 10 * std.math.floatEps(f64); // 2.22 × 10^-15

diff --git a/src/distribution/negativeBinomial.zig b/src/distribution/negativeBinomial.zig
@@ -77,46 +77,50 @@ pub const random = struct {
     pub fn single(generator: std.rand.Random, size: u64, prob: f64) f64 {
         assert(0 < prob and prob <= 1);
         assert(size != 0);
-        const n = @as(f64, @floatFromInt(size));
-        const z = (1 - prob) / prob;
-        const mean = n * z;
         if (prob == 1) {
             return 0;
-        } else if (mean < 150) {
+        }
+        const n = @as(f64, @floatFromInt(size));
+        const qrob = 1 - prob;
+        const qp = qrob / prob;
+        const mean = n * qp;
+        if (mean < 150) {
             const initial_mass = std.math.pow(f64, prob, n);
             const uni = generator.float(f64);
-            return linearSearch(uni, n, 1 - prob, initial_mass);
+            return linearSearch(uni, n, qrob, initial_mass);
         } else {
-            const lambda = gamma.random.single(generator, n, z);
+            const lambda = gamma.random.single(generator, n, qp);
             return poisson.random.single(generator, lambda);
         }
     }
 
     pub fn fill(buffer: []f64, generator: std.rand.Random, size: u64, prob: f64) []f64 {
         assert(0 < prob and prob <= 1);
         assert(size != 0);
-        const n = @as(f64, @floatFromInt(size));
-        const z = (1 - prob) / prob;
-        const mean = n * z;
         if (prob == 1) {
             @memset(buffer, 0);
-        } else if (buffer.len < 15 and mean < 15) {
+        }
+        const n = @as(f64, @floatFromInt(size));
+        const qrob = 1 - prob;
+        const qp = qrob / prob;
+        const mean = n * qp;
+        if (buffer.len < 15 and mean < 15) {
             const initial_mass = std.math.pow(f64, prob, n);
             for (buffer) |*x| {
                 const uni = generator.float(f64);
-                x.* = linearSearch(uni, n, 1 - prob, initial_mass);
+                x.* = linearSearch(uni, n, qrob, initial_mass);
             }
         } else if (mean < 15000) {
             const initial_nbin = @ceil(mean);
             const initial_mass = density(initial_nbin, size, prob);
             const initial_cumu = probability(initial_nbin, size, prob);
             for (buffer) |*x| {
                 const uni = generator.float(f64);
-                x.* = guidedSearch(uni, n, 1 - prob, initial_nbin, initial_mass, initial_cumu);
+                x.* = guidedSearch(uni, n, qrob, initial_nbin, initial_mass, initial_cumu);
             }
         } else {
             for (buffer) |*x| {
-                const lambda = gamma.random.single(generator, n, z);
+                const lambda = gamma.random.single(generator, n, qp);
                 x.* = poisson.random.single(generator, lambda);
             }
         }