Binomial distribution

Author: gaetanserre

RandomDo.lean

@@ -6,6 +6,7 @@ public import RandomDo.Monad.ForInInstances
 public import RandomDo.Monad.Instances
 public import RandomDo.Monad.MeasurableSpace
 public import RandomDo.Monad.Notation
+public import RandomDo.NumLean.Binomial
 public import RandomDo.NumLean.Distributions
 public import RandomDo.NumLean.PCG64
 public import RandomDo.NumLean.SeedSequence

RandomDo/NumLean/Binomial.lean

@@ -0,0 +1,194 @@
+/-
+Copyright (c) 2026 Gaëtan Serré. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Gaëtan Serré
+-/
+module
+
+public import RandomDo.NumLean.PCG64
+
+/-!
+# The binomial distribution
+
+`binomial n p` draws exactly what numpy's `Generator.binomial` draws: the inversion of the
+cumulative distribution where the mean `n * p` is at most `30`, the BTPE algorithm of
+Kachitvichyanukul and Schmeiser beyond, and the mirror image of either when `p > 1 / 2`.
+
+`binomial`, which chooses between the two and mirrors them, is in
+`RandomDo.NumLean.Distributions`, with the other distributions.
+
+## Main definitions
+
+* `Inversion`, `inversionSetup`, `inversionDraw`: numpy's `random_binomial_inversion`.
+* `Btpe`, `btpeSetup`, `Btpe.accept`, `btpeDraw`: numpy's `random_binomial_btpe`.
+
+## References
+
+* V. Kachitvichyanukul and B. W. Schmeiser, *Binomial random variate generation*, Communications
+  of the ACM 31 (1988), 216-222.
+* numpy's `numpy/random/src/distributions/distributions.c`.
+-/
+
+@[expose] public section
+
+namespace NumLean
+
+/-! ## Inversion -/
+
+/-- The constants the inversion reads a draw against. -/
+structure Inversion where
+  /-- The number of trials. -/
+  n : Float
+  /-- The probability of a success, at most one half. -/
+  p : Float
+  /-- The probability of a failure, `1 - p`. -/
+  q : Float
+  /-- `q ^ n`, the probability that no trial succeeds. -/
+  qn : Float
+  /-- The number of successes past which the walk gives up and starts over. -/
+  bound : Float
+
+/-- The constants of `random_binomial_inversion`, for `p ≤ 0.5` and `n * p ≤ 30`. -/
+@[inline] def inversionSetup (n p : Float) : Inversion :=
+  let q := 1.0 - p
+  let np := n * p
+  let b := np + 10.0 * Float.sqrt (np * q + 1)
+  { n, p, q, qn := Float.exp (n * Float.log q), bound := if n < b then n else b }
+
+/-- Walk up the cumulative distribution from zero until it passes `u`, as the loop of
+`random_binomial_inversion`. Answers `-1` where the walk runs past `bound`, which the tail beyond
+it is too thin to reach and where numpy starts the draw over. -/
+partial def inversionWalk (s : Inversion) (x px u : Float) : Float :=
+  if u > px then
+    let x := x + 1
+    if x > s.bound then -1
+    else inversionWalk s x (((s.n - x + 1) * s.p * px) / (x * s.q)) (u - px)
+  else x
+
+/-- Sample by inverting the cumulative distribution, as numpy's `random_binomial_inversion`. -/
+partial def inversionDraw (s : Inversion) : RandPCG IO Float := do
+  let x := inversionWalk s 0 s.qn (← random)
+  if x < 0 then inversionDraw s else return x
+
+/-! ## BTPE -/
+
+/-- The constants BTPE reads a draw against. -/
+structure Btpe where
+  /-- The number of trials. -/
+  n : Float
+  /-- The probability of a success, at most one half. -/
+  r : Float
+  /-- The probability of a failure, `1 - r`. -/
+  q : Float
+  /-- The mode of the distribution. -/
+  m : Float
+  /-- The middle of the triangle, `m + 1 / 2`. -/
+  xm : Float
+  /-- The half-width of the triangle. -/
+  p1 : Float
+  /-- The left end of the parallelogram. -/
+  xl : Float
+  /-- The right end of the parallelogram. -/
+  xr : Float
+  /-- The height of the parallelogram, relative to the triangle. -/
+  c : Float
+  /-- The rate of the left exponential tail. -/
+  laml : Float
+  /-- The rate of the right exponential tail. -/
+  lamr : Float
+  /-- The area of the triangle and the parallelogram. -/
+  p2 : Float
+  /-- The area of the triangle, the parallelogram and the left tail. -/
+  p3 : Float
+  /-- The area of all four regions, which a draw is scaled by. -/
+  p4 : Float
+  /-- The variance `n * r * q`. -/
+  nrq : Float
+
+/-- The constants of `random_binomial_btpe`, for `p ≤ 0.5` and `n * p > 30`. -/
+@[inline] def btpeSetup (n p : Float) : Btpe :=
+  let r := if p < 1.0 - p then p else 1.0 - p
+  let q := 1.0 - r
+  let fm := n * r + r
+  let m := Float.floor fm
+  let p1 := Float.floor (2.195 * Float.sqrt (n * r * q) - 4.6 * q) + 0.5
+  let xm := m + 0.5
+  let xl := xm - p1
+  let xr := xm + p1
+  let c := 0.134 + 20.5 / (15.3 + m)
+  let al := (fm - xl) / (fm - xl * r)
+  let ar := (xr - fm) / (xr * q)
+  let laml := al * (1.0 + al / 2.0)
+  let lamr := ar * (1.0 + ar / 2.0)
+  let p2 := p1 * (1.0 + 2.0 * c)
+  let p3 := p2 + c / laml
+  { n, r, q, m, xm, p1, xl, xr, c, laml, lamr, p2, p3, p4 := p3 + c / lamr, nrq := n * r * q }
+
+/-- One term of the Stirling series bounding `log` of a factorial, as the last test of BTPE spells
+it out. `u2` is `u * u`. -/
+@[inline] def btpeStirling (u u2 : Float) : Float :=
+  (13680.0 - (462.0 - (132.0 - (99.0 - 140.0 / u2) / u2) / u2) / u2) / u / 166320.0
+
+/-- The ratios of the probabilities from the mode up to `y`, multiplied into `f` one at a time as
+the step 50 of `random_binomial_btpe` takes them. -/
+partial def btpeUp (a s f i y : Float) : Float :=
+  if i ≤ y then btpeUp a s (f * (a / i - s)) (i + 1) y else f
+
+/-- The ratios from `y` up to the mode, divided out of `f` one at a time. Dividing the running
+value and dividing by the product do not round alike, and BTPE reads the first. -/
+partial def btpeDown (a s f i m : Float) : Float :=
+  if i ≤ m then btpeDown a s (f / (a / i - s)) (i + 1) m else f
+
+/-- Whether BTPE accepts the candidate `y` drawn with `v`, as the steps 50 and 52 of
+`random_binomial_btpe`: by the explicit product of the ratios of the probabilities between the mode
+and `y` when the two are close, and by a squeeze then the Stirling bound otherwise. -/
+def Btpe.accept (b : Btpe) (y v : Float) : Bool := Id.run do
+  let k := Float.abs (y - b.m)
+  unless k > 20 && k < b.nrq / 2.0 - 1 do
+    let s := b.r / b.q
+    let a := s * (b.n + 1)
+    if b.m < y then return !(v > btpeUp a s 1.0 (b.m + 1) y)
+    if b.m > y then return !(v > btpeDown a s 1.0 (y + 1) b.m)
+    return !(v > 1.0)
+  let rho := (k / b.nrq) * ((k * (k / 3.0 + 0.625) + 0.16666666666666666) / b.nrq + 0.5)
+  let t := -k * k / (2 * b.nrq)
+  let a := Float.log v
+  if a < t - rho then return true
+  if a > t + rho then return false
+  let x1 := y + 1
+  let f1 := b.m + 1
+  let z := b.n + 1 - b.m
+  let w := b.n - y + 1
+  return !(a > b.xm * Float.log (f1 / x1) + (b.n - b.m + 0.5) * Float.log (z / w)
+    + (y - b.m) * Float.log (w * b.r / (x1 * b.q))
+    + btpeStirling f1 (f1 * f1) + btpeStirling z (z * z)
+    + btpeStirling x1 (x1 * x1) + btpeStirling w (w * w))
+
+/-- Draw a candidate from the triangle, the parallelogram or one of the two exponential tails, and
+start over until one is accepted, as the steps 10 to 60 of `random_binomial_btpe`. -/
+partial def btpeDraw (b : Btpe) : RandPCG IO Float := do
+  let u := (← random) * b.p4
+  let v ← random
+  if u ≤ b.p1 then
+    return Float.floor (b.xm - b.p1 * v + u)
+  else if u ≤ b.p2 then
+    let x := b.xl + (u - b.p1) / b.c
+    let v := v * b.c + 1.0 - Float.abs (b.m - x + 0.5) / b.p1
+    if v > 1.0 then btpeDraw b else
+      let y := Float.floor x
+      if b.accept y v then return y else btpeDraw b
+  else if u ≤ b.p3 then
+    let y := Float.floor (b.xl + Float.log v / b.laml)
+    -- `v` can be zero, and the floor of the resulting infinity is no candidate.
+    if y < 0 || v == 0.0 then btpeDraw b else
+      let v := v * (u - b.p2) * b.laml
+      if b.accept y v then return y else btpeDraw b
+  else
+    let y := Float.floor (b.xr - Float.log v / b.lamr)
+    if y > b.n || v == 0.0 then btpeDraw b else
+      let v := v * (u - b.p3) * b.lamr
+      if b.accept y v then return y else btpeDraw b
+
+end NumLean
+
+end

RandomDo/NumLean/Distributions.lean

@@ -8,6 +8,7 @@ module
 public import RandomDo.NumLean.PCG64
 public meta import RandomDo.NumLean.PCG64
 public import FFI.Float
+public import RandomDo.NumLean.Binomial
 public import RandomDo.NumLean.Ziggurat
 public import RandomDo.NumLean.ZigguratSampler
 
@@ -141,4 +142,22 @@ deviation. -/
   if scale < 0 then throw <| IO.userError "scale < 0"
   return scale * (← standardExponential)
 
+/-- Draw samples from a binomial distribution. -/
+def binomial (n : Nat) (p : Float) : RandPCG IO Nat := do
+  -- The comparisons are the `Bool` ones: through `Decidable`, each costs more than a draw.
+  if p.lt 0.0 || Float.lt 1.0 p || p.isNaN then
+    throw <| IO.userError "p < 0, p > 1 or p is NaN"
+  let n := n.toUInt64.toFloat
+  if n == 0 || p == 0.0 then return 0
+  if Float.le p 0.5 then
+    if Float.le (p * n) 30.0 then return (← inversionDraw (inversionSetup n p)).toUInt64.toNat
+    else return (← btpeDraw (btpeSetup n p)).toUInt64.toNat
+  else
+    let q := 1.0 - p
+    if Float.le (q * n) 30.0 then return (n - (← inversionDraw (inversionSetup n q))).toUInt64.toNat
+    else return (n - (← btpeDraw (btpeSetup n q))).toUInt64.toNat
+
+/-- Draw samples from a Bernoulli distribution. -/
+@[inline] def bernoulli (p : Float) := binomial 1 p
+
 end NumLean

RandomDo/Tactic/Computable/Counterparts.lean

@@ -8,6 +8,7 @@ module
 public import RandomDo.Tactic.Computable.Defs
 public import RandomDo.NumLean.Distributions
 public import Mathlib.Probability.Distributions.Gaussian.Real
+public import Mathlib.Probability.Distributions.Bernoulli
 
 /-!
 # Computable counterparts of the pieces an `rdo` program is made of
@@ -18,7 +19,7 @@ counterpart recorded here through `@[computable_as]`. There is one entry per pie
 one distribution they draw from.
 -/
 
-public meta section
+@[expose] public section
 
 /-! ## Types -/
 
@@ -29,12 +30,16 @@ attribute [computable_as Float] NNReal
 
 attribute [computable_as NumLean.normal'] ProbabilityTheory.gaussianReal
 
+def bernoulliChoice (α : Type) [MeasurableSpace α] (x y : α) (p : Float) :
+    NumLean.RandPCG IO α := do
+  return if (← NumLean.bernoulli p) == 1 then x else y
+
+attribute [computable_as bernoulliChoice] ProbabilityTheory.bernoulliMeasure
+
 /-! ## Classical functions -/
 
 attribute [computable_as Float.sqrt] Real.sqrt
 
 attribute [computable_as Float.log] Real.log
 
 attribute [computable_as Float.exp] Real.exp
-
-end

RandomDo/Tactic/Computable/Deriving.lean

@@ -62,6 +62,10 @@ partial def translate (σ : FVarSubst) (e : Expr) : MetaM Expr :=
       mkAppOptM ``Bind.bind #[← computableMonad, none, none, none,
         ← translate σ p, ← translate σ k]
     | MeasureTheory.Measure α _ => return mkApp (← computableMonad) (← translate σ α)
+    /- A subtype is its carrier, and one of its values is the value it carries: the constraint and
+    the proof of it are what a computable counterpart does not have. -/
+    | Subtype α _ => translate σ α
+    | Subtype.mk _ _ v _ => translate σ v
     | _ => match e with
       | .fvar x => return σ.get x
       | .sort .. | .lit .. => return e

Test/Computable.lean

@@ -1,11 +1,8 @@
 module
 
 public import Test.IsMarkov
+public import Test.Bind
 import Batteries.Data.Float.Basic
-/- A `run_cmd` runs at elaboration time, so what it calls has to be imported as `meta` too: the
-sampler it draws with, and `Float.toStringFull` it prints with. -/
-meta import RandomDo.NumLean.Distributions
-meta import Batteries.Data.Float.Basic
 
 set_option linter.style.header false
 
@@ -15,19 +12,19 @@ namespace Test.Computable
 
 open Test.IsMarkov NumLean Lean.Elab.Command
 
-def logComputable (prog : RandPCG IO Float) : CommandElabM Unit := do
-  let x ← (IO.runRandPCG prog : IO Float)
-  let y ← (IO.runRandPCGWith 42 prog : IO Float)
-  Lean.logInfo m!"x = {x.toStringFull}"
-  Lean.logInfo m!"y (seed 42) = {y.toStringFull}"
+def logComputable {α : Type} [Lean.ToMessageData α] (prog : RandPCG IO α) : CommandElabM Unit := do
+  let x ← (IO.runRandPCG prog : IO α)
+  let y ← (IO.runRandPCGWith 42 prog : IO α)
+  Lean.logInfo m!"x = {x}"
+  Lean.logInfo m!"y (seed 42) = {y}"
 
---attribute [computable] sumTwo
+attribute [computable] sumTwo
 
---run_cmd do logComputable sumTwoComputable
+run_cmd do logComputable sumTwoComputable
 
 @[computable]
 noncomputable
-def test : MeasureTheory.Measure ℝ := rdo
+def unfoldSumTwo : MeasureTheory.Measure ℝ := rdo
   let y ← sumTwo
   let x ← ProbabilityTheory.gaussianReal 0 1
   return x + y
@@ -42,4 +39,12 @@ run_cmd do logComputable (branchOnComputable 20)
 
 run_cmd do logComputable (branchOnComputable (-1))
 
+attribute [computable] fairCoin
+
+run_cmd do logComputable (fairCoinComputable)
+
+attribute [computable] Bind.twoCoins
+
+run_cmd do logComputable (Bind.twoCoinsComputable)
+
 end Test.Computable

scripts/check_binomial.py

@@ -0,0 +1,27 @@
+import numpy as np
+from common import compare
+
+print("Checking binomial distribution...")
+
+# One pair per branch of numpy's `random_binomial`: inversion and BTPE, each on both sides of
+# `p = 1/2`, on both sides of the `n * p = 30` threshold, and the three degenerate cases.
+PARAMS = [(10, 0.3), (60, 0.5), (100, 0.31), (1000, 0.5), (1000, 0.9), (5, 0.99),
+          (0, 0.5), (10, 0.0), (10, 1.0)]
+
+LEAN = """import RandomDo
+def params : List (Nat × Float) :=
+  [%s]
+def main (args : List String) : IO Unit := do
+  for s in args do
+    IO.FS.withFile (System.FilePath.mk s!"@DIR@/pcg64-{s}.txt") .write fun h ↦
+      IO.runRandPCGWith s.toNat! do
+        for _ in List.range (@N@ / %d) do
+          for (n, p) in params do
+            h.putStrLn (toString (← NumLean.binomial n p))
+""" % (",\n   ".join(f"({n}, {p})" for n, p in PARAMS), len(PARAMS))
+
+def distrib(seed, N):
+    g = np.random.default_rng(seed)
+    return [g.binomial(n, p) for _ in range(N // len(PARAMS)) for n, p in PARAMS]
+
+compare(LEAN, distrib)