smarter algorithm for positive and negative binomial rng

docs/src/root/distribution/beta.zig.html

@@ -178,7 +178,7 @@
 <span class="line" id="L64">    <span class="tok-kw">return</span> inverseIncompleteBeta(shape1, shape2, p);</span>
 <span class="line" id="L65">}</span>
 <span class="line" id="L66"></span>
-<span class="line" id="L67"><span class="tok-comment">/// Uses the relation to Gamma distribution or rejection sampling for lower α, β.</span></span>
+<span class="line" id="L67"><span class="tok-comment">/// Uses the relation to Gamma distribution or rejection sampling.</span></span>
 <span class="line" id="L68"><span class="tok-comment">/// http://luc.devroye.org/chapter_nine.pdf page 416.</span></span>
 <span class="line" id="L69"><span class="tok-kw">pub</span> <span class="tok-kw">const</span> random = <span class="tok-kw">struct</span> {</span>
 <span class="line" id="L70">    <span class="tok-kw">pub</span> <span class="tok-kw">fn</span> <span class="tok-fn">single</span>(generator: std.rand.Random, shape1: <span class="tok-type">f64</span>, shape2: <span class="tok-type">f64</span>) <span class="tok-type">f64</span> {</span>

src/distribution/beta.zig

@@ -64,7 +64,7 @@ pub fn quantile(p: f64, shape1: f64, shape2: f64) f64 {
     return inverseIncompleteBeta(shape1, shape2, p);
 }
 
-/// Uses the relation to Gamma distribution or rejection sampling for lower α, β.
+/// Uses the relation to Gamma distribution or rejection sampling.
 /// http://luc.devroye.org/chapter_nine.pdf page 416.
 pub const random = struct {
     pub fn single(generator: std.rand.Random, shape1: f64, shape2: f64) f64 {

src/distribution/binomial.zig

@@ -66,45 +66,69 @@ pub fn quantile(p: f64, size: u64, prob: f64) f64 {
         return n * prob;
     }
     const pq = prob / (1 - prob);
-    const initial_mass = std.math.pow(f64, 1 - prob, n);
-    return linearSearch(p, n, pq, initial_mass);
+    const mean = n * prob;
+    if (mean < 500) {
+        const initial_mass = std.math.pow(f64, 1 - prob, n);
+        return linearSearch(p, n, pq, 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);
+    }
 }
 
-/// Uses the quantile function or bit-counting when prob == 0.5.
+/// Uses the quantile function or bit-counting.
 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 mean = n * prob;
         if (prob == 0 or prob == 1 or size == 0) {
-            return n * prob;
+            return mean;
         } else if (prob == 0.5) {
             const mask = (@as(u64, 1) << @truncate(@mod(size, 64))) - 1;
             return bitCount(generator, mask, size);
-        } else {
-            const pq = prob / (1 - prob);
+        } 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 {
+            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);
         }
     }
 
     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 mean = n * prob;
         if (prob == 0 or prob == 1 or size == 0) {
-            @memset(buffer, n * prob);
+            @memset(buffer, mean);
         } else if (prob == 0.5) {
             const mask = (@as(u64, 1) << @truncate(@mod(size, 64))) - 1;
             for (buffer) |*x| {
                 x.* = bitCount(generator, mask, size);
             }
-        } else {
-            const pq = prob / (1 - prob);
+        } 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 {
+            const initial_bino = @ceil(mean);
+            const initial_mass = density(initial_bino, size, prob);
+            const initial_cumu = probability(initial_bino, size, prob);
+            for (buffer) |*x| {
+                const uni = generator.float(f64);
+                x.* = guidedSearch(uni, n, pq, initial_bino, initial_mass, initial_cumu);
+            }
         }
         return buffer;
     }
@@ -123,6 +147,31 @@ inline fn linearSearch(p: f64, n: f64, pq: f64, initial_mass: f64) f64 {
     return bin;
 }
 
+inline fn guidedSearch(p: f64, n: f64, pq: f64, initial_bino: f64, initial_mass: f64, initial_cumu: f64) f64 {
+    var bino = initial_bino;
+    var mass = initial_mass;
+    var cumu = initial_cumu;
+    if (initial_cumu <= p) {
+        while (cumu <= p) {
+            const num = n - bino;
+            bino += 1;
+            mass *= pq * num / bino;
+            cumu += mass;
+        }
+    } else {
+        while (true) {
+            cumu -= mass;
+            if (cumu <= p) {
+                break;
+            }
+            const num = bino;
+            bino -= 1;
+            mass *= num / (pq * (n - bino));
+        }
+    }
+    return bino;
+}
+
 inline fn bitCount(generator: std.rand.Random, mask: u64, size: u64) f64 {
     var bino: usize = 0;
     var i: usize = 64;

src/distribution/negativeBinomial.zig

@@ -1,10 +1,12 @@
 //! Support: X ∈ {0,1,2,⋯}
 //!
 //! Parameters:
-//! - n: `size` ∈ {0,1,2,⋯}
+//! - n: `size` ∈ {1,2,⋯}
 //! - p: `prob` ∈ (0,1]
 
 const std = @import("std");
+const gamma = @import("gamma.zig");
+const poisson = @import("poisson.zig");
 const math = @import("../math.zig");
 const incompleteBeta = @import("../thirdyparty/prob.zig").incompleteBeta;
 const assert = std.debug.assert;
@@ -58,45 +60,65 @@ pub fn quantile(p: f64, size: u64, prob: f64) f64 {
         return inf;
     }
     const n = @as(f64, @floatFromInt(size));
-    const q = 1 - prob;
-    var mass = std.math.pow(f64, prob, n);
-    var cumu = mass;
-    var nbi: f64 = 0;
-    while (p >= cumu) : (nbi += 1) {
-        mass *= q * (n + nbi) / (nbi + 1);
-        cumu += mass;
+    const mean = n * (1 - prob) / prob;
+    if (mean < 250) {
+        const initial_mass = std.math.pow(f64, prob, n);
+        return linearSearch(p, n, 1 - prob, initial_mass);
+    } else {
+        const initial_nbin = @ceil(mean);
+        const initial_mass = density(initial_nbin, size, prob);
+        const initial_cumu = probability(initial_nbin, size, prob);
+        return guidedSearch(p, n, 1 - prob, initial_nbin, initial_mass, initial_cumu);
     }
-    return nbi;
 }
 
-/// Uses the quantile function.
+/// Uses the quantile function or the relation to Poisson distribution.
 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 initial_mass = std.math.pow(f64, prob, n);
+            const uni = generator.float(f64);
+            return linearSearch(uni, n, 1 - prob, initial_mass);
+        } else {
+            const lambda = gamma.random.single(generator, n, z);
+            return poisson.random.single(generator, lambda);
         }
-        const n = @as(f64, @floatFromInt(size));
-        const q = 1 - prob;
-        const initial_mass = std.math.pow(f64, prob, n);
-        const uni = generator.float(f64);
-        return linearSearch(uni, n, q, initial_mass);
     }
 
     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);
-            return buffer;
-        }
-        const n = @as(f64, @floatFromInt(size));
-        const q = 1 - prob;
-        const initial_mass = std.math.pow(f64, prob, n);
-        for (buffer) |*x| {
-            const uni = generator.float(f64);
-            x.* = linearSearch(uni, n, q, initial_mass);
+        } else 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);
+            }
+        } 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);
+            }
+        } else {
+            for (buffer) |*x| {
+                const lambda = gamma.random.single(generator, n, z);
+                x.* = poisson.random.single(generator, lambda);
+            }
         }
         return buffer;
     }
@@ -115,6 +137,31 @@ inline fn linearSearch(p: f64, n: f64, q: f64, initial_mass: f64) f64 {
     return nbin;
 }
 
+inline fn guidedSearch(p: f64, n: f64, q: f64, initial_nbin: f64, initial_mass: f64, initial_cumu: f64) f64 {
+    var nbin = initial_nbin;
+    var mass = initial_mass;
+    var cumu = initial_cumu;
+    if (initial_cumu <= p) {
+        while (cumu <= p) {
+            const num = q * (n + nbin);
+            nbin += 1;
+            mass *= num / nbin;
+            cumu += mass;
+        }
+    } else {
+        while (true) {
+            cumu -= mass;
+            if (cumu <= p) {
+                break;
+            }
+            const num = nbin;
+            nbin -= 1;
+            mass *= num / (q * (n + nbin));
+        }
+    }
+    return nbin;
+}
+
 const expectEqual = std.testing.expectEqual;
 const expectApproxEqRel = std.testing.expectApproxEqRel;
 const eps = 10 * std.math.floatEps(f64); // 2.22 × 10^-15

src/distribution/poisson.zig

@@ -57,7 +57,7 @@ pub fn quantile(p: f64, lambda: f64) f64 {
     }
 }
 
-/// Uses the quantile function or rejection sampling for higher λ.
+/// Uses the quantile function or rejection sampling.
 /// https://www.jstor.org/stable/2346807
 pub const random = struct {
     pub fn single(generator: std.rand.Random, lambda: f64) f64 {
@@ -106,8 +106,9 @@ inline fn linearSearch(p: f64, lambda: f64) f64 {
     var pois: f64 = 0;
     var mass = @exp(-lambda);
     var cumu = mass;
-    while (cumu <= p) : (pois += 1) {
-        mass *= lambda / (pois + 1);
+    while (cumu <= p) {
+        pois += 1;
+        mass *= lambda / pois;
         cumu += mass;
     }
     return pois;
@@ -126,10 +127,10 @@ inline fn guidedSearch(p: f64, lambda: f64, initial_pois: f64, initial_mass: f64
     } else {
         while (true) {
             cumu -= mass;
-            mass *= pois / lambda;
             if (cumu <= p) {
                 break;
             }
+            mass *= pois / lambda;
             pois -= 1;
         }
     }