Inline and separate Ziggurat

Author: gaetanserre

RandomDo.lean

@@ -10,6 +10,7 @@ public import RandomDo.NumLean.Distributions
 public import RandomDo.NumLean.PCG64
 public import RandomDo.NumLean.SeedSequence
 public import RandomDo.NumLean.Ziggurat
+public import RandomDo.NumLean.ZigguratSampler
 public import RandomDo.Tactic.Deriving
 public import RandomDo.Tactic.Elab
 public import RandomDo.Tactic.ForInStep

RandomDo/NumLean/Distributions.lean

@@ -9,6 +9,7 @@ public import RandomDo.NumLean.PCG64
 public meta import RandomDo.NumLean.PCG64
 public import FFI.Float
 public import RandomDo.NumLean.Ziggurat
+public import RandomDo.NumLean.ZigguratSampler
 
 /-!
 # Sample from specific distributions using the PCG-64 generator.
@@ -17,37 +18,25 @@ This file provides samplers for specific distributions using the PCG-64 generato
 exactly what its numpy counterpart draws, so that a `RandPCG` program and a numpy `Generator`
 seeded alike produce the same values.
 
+The draws taken straight from the generator, `randUInt64`, `randUInt32` and `random`, are in
+`RandomDo.NumLean.PCG64`; the ziggurat the normal and the exponential share is in
+`RandomDo.NumLean.ZigguratSampler` and its tables in `RandomDo.NumLean.Ziggurat`.
+
 ## Main definitions
-* `randUInt64` / `randUInt32`: sample a `UInt64`, or a `UInt32` as numpy's `next_uint32` does.
-* `random`: sample a `Float` in `[0, 1)`.
 * `randInt`: sample an integer in `[low, high)` or `[low, high]`, as numpy's `Generator.integers`.
 * `uniform`: sample a `Float` in `[low, high)`, as numpy's `Generator.uniform`.
 * `standardNormal` / `normal`: sample a normal deviate, as numpy's `Generator.normal`.
 * `standardExponential` / `exponential`: sample an exponential deviate, as numpy's
   `Generator.exponential`.
-* `ziggurat`: the sampler shape both of those share.
 -/
 
 @[expose] public section
 
 namespace NumLean
 
-/-- Sample a `UInt64` from a PCG-64 generator. -/
-def randUInt64 : RandPCG IO UInt64 := do
-  let (x, g) := (← get).down.nextUInt64
-  set (ULift.up g)
-  return x
-
-/-- Sample a `UInt32` from a PCG-64 generator, as numpy's `next_uint32` does for `PCG64`: the two
-halves of each 64-bit output are handed out in turn, see `PCG64.nextUInt32`. -/
-def randUInt32 : RandPCG IO UInt32 := do
-  let (x, g) := (← get).down.nextUInt32
-  set (ULift.up g)
-  return x
-
 /-- Sample a `UInt32` uniformly in `[0, rng]` by Lemire's nearly divisionless rejection method, as
 numpy's `buffered_bounded_lemire_uint32`. -/
-def randLemireUInt32 (rng : UInt32) : RandPCG IO UInt32 := do
+@[inline] def randLemireUInt32 (rng : UInt32) : RandPCG IO UInt32 := do
   let rngExcl := rng + 1
   let mut m := (← randUInt32).toUInt64 * rngExcl.toUInt64
   if m.toUInt32 < rngExcl then
@@ -58,7 +47,7 @@ def randLemireUInt32 (rng : UInt32) : RandPCG IO UInt32 := do
 
 /-- Sample a `UInt64` uniformly in `[0, rng]` by Lemire's method, as numpy's
 `bounded_lemire_uint64`. -/
-def randLemireUInt64 (rng : UInt64) : RandPCG IO UInt64 := do
+@[inline] def randLemireUInt64 (rng : UInt64) : RandPCG IO UInt64 := do
   let rngExcl := rng + 1
   let mut x ← randUInt64
   let mut leftover := x * rngExcl
@@ -71,7 +60,7 @@ def randLemireUInt64 (rng : UInt64) : RandPCG IO UInt64 := do
 
 /-- Sample a `UInt64` uniformly in `[0, rng]`, as numpy's `random_bounded_uint64` with Lemire's
 rejection. -/
-def randBoundedUInt64 (rng : UInt64) : RandPCG IO UInt64 := do
+@[inline] def randBoundedUInt64 (rng : UInt64) : RandPCG IO UInt64 := do
   if rng == 0 then return 0
   else if rng == 0xFFFFFFFF then return (← randUInt32).toUInt64
   else if rng < 0xFFFFFFFF then return (← randLemireUInt32 rng.toUInt32).toUInt64
@@ -80,7 +69,7 @@ def randBoundedUInt64 (rng : UInt64) : RandPCG IO UInt64 := do
 
 /-- Sample an integer uniformly in `[low, high)`, or in `[low, high]` when `endpoint` is set, as
 numpy's `Generator.integers`. -/
-def randInt₀ (low high : Int) (endpoint : Bool := false) : RandPCG IO Int := do
+@[inline] def randInt₀ (low high : Int) (endpoint : Bool := false) : RandPCG IO Int := do
   let high := if endpoint then high else high - 1
   if low < Int64.minValue.toInt then throw <| IO.userError "low is out of bounds for int64"
   if high > Int64.maxValue.toInt then throw <| IO.userError "high is out of bounds for int64"
@@ -89,49 +78,19 @@ def randInt₀ (low high : Int) (endpoint : Bool := false) : RandPCG IO Int := d
   return low + x.toNat
 
 /-- Sample an integer uniformly in `[0, high)`, or in `[0, high]` when `endpoint` is set. -/
-def randInt (high : Int) (endpoint : Bool := false) : RandPCG IO Int := randInt₀ 0 high endpoint
+@[inline] def randInt (high : Int) (endpoint : Bool := false) : RandPCG IO Int :=
+    randInt₀ 0 high endpoint
 
 /-- Sample an integer uniformly in `[low, high)`, or in `[low, high]` when `endpoint` is set. -/
-def randBoundedInt (low high : Int) (endpoint : Bool := false) : RandPCG IO Int :=
+@[inline] def randBoundedInt (low high : Int) (endpoint : Bool := false) : RandPCG IO Int :=
     randInt₀ low high endpoint
 
-/-- Sample a `Float` in `[0, 1)` from a PCG-64 generator. -/
-def random : RandPCG IO Float := do
-  let x ← randUInt64
-  return (x >>> 11).toFloat * (Float.ofBits <| 0x3CA <<< (52 : UInt64))
-
 /-- Sample a `Float` uniformly in `[low, high)`. -/
-def uniform (low high : Float) : RandPCG IO Float := do
+@[inline] def uniform (low high : Float) : RandPCG IO Float := do
   if low > high then throw <| IO.userError "low > high"
   let x ← random
   return Float.fma x (high - low) low
 
-/-- The rejection test on a strip that sticks out of the curve: the point drawn at height `u`
-between the density at the strip's two edges lies under the curve. -/
-@[inline] def zigguratWedge (f : Array Float) (idx : Nat) (u density : Float) : Bool :=
-  Float.fma (f[idx - 1]! - f[idx]!) u f[idx]! < density
-
-/-- The sampler shape shared by numpy's `random_standard_normal` and
-`random_standard_exponential`, the ziggurat of Marsaglia and Tsang: the density is covered by 256
-strips of equal area, and one 64-bit output supplies at once the strip `idx` and the integer `ri`
-that `w` scales to an abscissa, `split` saying how those bits are laid out and whether the deviate
-comes out negated.
-
-The draw is returned as it stands when `ri` falls below `k[idx]`, which is where about 99% of the
-draws end. Otherwise the strip either is the base one, whose unbounded part `tail` samples from the
-integer drawn, or sticks out of the curve, and then the point is tested against `density` and the
-whole draw is started over on rejection. -/
-@[specialize] partial def ziggurat (k : Array UInt64) (w f : Array Float)
-    (split : UInt64 → Nat × UInt64 × Bool) (density : Float → Float)
-    (tail : UInt64 → RandPCG IO Float) : RandPCG IO Float := do
-  let (idx, ri, negate) := split (← randUInt64)
-  let x := ri.toFloat * w[idx]!
-  let x := if negate then -x else x
-  if ri < k[idx]! then return x
-  if idx == 0 then tail ri
-  else if zigguratWedge f idx (← random) (density x) then return x
-  else ziggurat k w f split density tail
-
 /-- Sample the tail of the standard normal beyond `Ziggurat.norR`, as the `idx == 0` branch of
 numpy's `random_standard_normal`: draw from an exponential tail until the pair of draws falls under
 the normal's, which is Marsaglia's method for the tail. `negate` carries the sign numpy reads off
@@ -146,7 +105,7 @@ partial def normalTail (negate : Bool) : RandPCG IO Float := do
 /-- Sample from the standard normal, as numpy's `random_standard_normal`. The 64-bit output gives
 the strip in its low byte, then the sign, then a 52-bit abscissa; the tail takes its sign from a
 further bit of that same abscissa, as numpy does. -/
-def standardNormal : RandPCG IO Float :=
+@[inline] def standardNormal : RandPCG IO Float :=
   ziggurat Ziggurat.ki Ziggurat.wi Ziggurat.fi
     (fun r =>
       let idx := (r &&& 0xFF).toNat
@@ -158,7 +117,7 @@ def standardNormal : RandPCG IO Float :=
 /-- Sample from the standard exponential, as numpy's `random_standard_exponential`. The 64-bit
 output is first shifted by three, then gives the strip in its low byte and a 53-bit abscissa; the
 tail is the exponential's own, memoryless, so one draw beyond `Ziggurat.expR` suffices. -/
-def standardExponential : RandPCG IO Float :=
+@[inline] def standardExponential : RandPCG IO Float :=
   ziggurat Ziggurat.ke Ziggurat.we Ziggurat.fe
     (fun r =>
       let r := r >>> 3
@@ -167,12 +126,12 @@ def standardExponential : RandPCG IO Float :=
     (fun _ => do return Ziggurat.expR - Float.log1p (-(← random)))
 
 /-- Draw random samples from a normal (Gaussian) distribution. -/
-def normal (loc : Float := 0) (scale : Float := 1) : RandPCG IO Float := do
+@[inline] def normal (loc : Float := 0) (scale : Float := 1) : RandPCG IO Float := do
   if scale < 0 then throw <| IO.userError "scale < 0"
   return Float.fma scale (← standardNormal) loc
 
 /-- Draw samples from an exponential distribution. -/
-def exponential (scale : Float := 1) : RandPCG IO Float := do
+@[inline] def exponential (scale : Float := 1) : RandPCG IO Float := do
   if scale < 0 then throw <| IO.userError "scale < 0"
   return scale * (← standardExponential)
 

RandomDo/NumLean/PCG64.lean

@@ -29,6 +29,8 @@ native machine arithmetic.
   hands out the two halves of a 64-bit output in turn
 * `PCG64.seedWords` / `PCG64.seed` / `mkPCG64`: seeding, following the reference
   `pcg_setseq_128_srandom_r`
+* `randUInt64` / `randUInt32` / `random`: the three outputs a `RandPCG` computation draws straight
+  from the generator.
 
 ## References
 
@@ -200,6 +202,25 @@ instance : RandomGen PCG64 where
 `RandPCG m α` should be thought of a random value in `m α`. -/
 abbrev RandPCG := RandGT PCG64
 
+/-- Sample a `UInt64` from a PCG-64 generator. -/
+@[inline] def randUInt64 : RandPCG IO UInt64 := do
+  let (x, g) := (← get).down.nextUInt64
+  set (ULift.up g)
+  return x
+
+/-- Sample a `UInt32` from a PCG-64 generator, as numpy's `next_uint32` does for `PCG64`: the two
+halves of each 64-bit output are handed out in turn, see `PCG64.nextUInt32`. -/
+@[inline] def randUInt32 : RandPCG IO UInt32 := do
+  let (x, g) := (← get).down.nextUInt32
+  set (ULift.up g)
+  return x
+
+/-- Sample a `Float` in `[0, 1)` from a PCG-64 generator, as numpy's `next_double`: the top 53 bits
+of a 64-bit output, scaled by `2 ^ (-53)`. -/
+@[inline] def random : RandPCG IO Float := do
+  let x ← randUInt64
+  return (x >>> 11).toFloat * (Float.ofBits <| 0x3CA <<< (52 : UInt64))
+
 end NumLean
 
 namespace IO

RandomDo/NumLean/ZigguratSampler.lean

@@ -0,0 +1,87 @@
+/-
+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
+public import FFI.Float
+
+/-!
+# The ziggurat sampler
+
+The sampler shape numpy's `random_standard_normal` and `random_standard_exponential` share, the
+ziggurat of Marsaglia and Tsang. The density is covered by 256 strips of equal area; a draw picks a
+strip and a point in it, and is accepted at once when the point falls in the strip's rectangular
+part, which is where about 99% of the draws end.
+
+The tables the two samplers read live in `RandomDo.NumLean.Ziggurat`, and the samplers themselves
+in `RandomDo.NumLean.Distributions`; this file holds only the shape they share, which is generic in
+the tables `k`, `w` and `f` and in the three functions that say how a 64-bit output is cut up, what
+the density is, and how the base strip's tail is sampled.
+
+## Main definitions
+
+* `ziggurat`: the sampler, whose fast path is inlined into its caller
+* `zigguratSlow`: the branches the remaining ~1% of the draws take
+* `zigguratWedge`: the rejection test on a strip that sticks out of the curve
+
+## References
+
+* G. Marsaglia and W. W. Tsang, *The Ziggurat Method for Generating Random Variables*, Journal of
+  Statistical Software, 2000.
+* numpy's samplers: `numpy/random/src/distributions/distributions.c`
+-/
+
+@[expose] public section
+
+namespace NumLean
+
+/-- The rejection test on a strip that sticks out of the curve: the point drawn at height `u`
+between the density at the strip's two edges lies under the curve. -/
+@[inline] def zigguratWedge (f : Array Float) (idx : Nat) (u density : Float) : Bool :=
+  Float.fma (f[idx - 1]! - f[idx]!) u f[idx]! < density
+
+/-- The rare branches of `ziggurat`, reached by about 1% of the draws: the strip either is the base
+one, whose unbounded part `tail` samples from the integer drawn, or sticks out of the curve, and
+then the point is tested against `density` and the whole draw is started over on rejection.
+
+`idx`, `ri` and `x` are the strip, the integer and the abscissa `ziggurat` has already drawn and
+found not to land in the rectangular part of its strip. Each redraw retries the fast path here
+rather than returning to `ziggurat`, so the two together run exactly the loop of numpy's samplers.
+
+This is kept apart from `ziggurat` so that the fast path can be inlined into its caller: a
+recursive function cannot be, and behind a call boundary every draw would have to box the generator
+state and its result, which costs several times the draw itself. -/
+@[specialize] partial def zigguratSlow (k : Array UInt64) (w f : Array Float)
+    (split : UInt64 → Nat × UInt64 × Bool) (density : Float → Float)
+    (tail : UInt64 → RandPCG IO Float) (idx : Nat) (ri : UInt64) (x : Float) :
+    RandPCG IO Float := do
+  if idx == 0 then tail ri
+  else if zigguratWedge f idx (← random) (density x) then return x
+  else
+    let (idx, ri, negate) := split (← randUInt64)
+    let x := ri.toFloat * w[idx]!
+    let x := if negate then -x else x
+    if ri < k[idx]! then return x
+    else zigguratSlow k w f split density tail idx ri x
+
+/-- The sampler shape shared by numpy's `random_standard_normal` and
+`random_standard_exponential`, the ziggurat of Marsaglia and Tsang: the density is covered by 256
+strips of equal area, and one 64-bit output supplies at once the strip `idx` and the integer `ri`
+that `w` scales to an abscissa, `split` saying how those bits are laid out and whether the deviate
+comes out negated.
+
+The draw is returned as it stands when `ri` falls below `k[idx]`, which is where about 99% of the
+draws end; `zigguratSlow` takes over the remaining ones. -/
+@[inline] def ziggurat (k : Array UInt64) (w f : Array Float)
+    (split : UInt64 → Nat × UInt64 × Bool) (density : Float → Float)
+    (tail : UInt64 → RandPCG IO Float) : RandPCG IO Float := do
+  let (idx, ri, negate) := split (← randUInt64)
+  let x := ri.toFloat * w[idx]!
+  let x := if negate then -x else x
+  if ri < k[idx]! then return x
+  else zigguratSlow k w f split density tail idx ri x
+
+end NumLean

lakefile.toml

@@ -24,8 +24,7 @@ name = "RandomDo"
 [[lean_lib]]
 name = "Test"
 
-# The check scripts generate `test_data/Dump.lean` and run it through this target, so that the
-# sampler runs as compiled code: `lake env lean --run` would interpret it, which is far slower.
+# Used to run the tests in `scripts`
 [[lean_exe]]
 name = "dump"
 root = "Dump"