randInt and uniform

Author: gaetanserre

RandomDo/NumLean/Distributions.lean

@@ -7,15 +7,19 @@ module
 
 public import RandomDo.NumLean.PCG64
 public meta import RandomDo.NumLean.PCG64
+public import Batteries.Data.Float.Basic
 
 /-!
 # Sample from specific distributions using the PCG-64 generator.
 
-This files provides samplers for specific distributions using the PCG-64 generator.
+This file provides samplers for specific distributions using the PCG-64 generator. Each one draws
+exactly what its numpy counterpart draws, so that a `RandPCG` program and a numpy `Generator`
+seeded alike produce the same values.
 
 ## Main definitions
-* `randUInt64`: sample a `UInt64`.
+* `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`.
 -/
 
 @[expose] public section
@@ -28,9 +32,84 @@ def randUInt64 : RandPCG IO UInt64 := do
   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
+  let rngExcl := rng + 1
+  let mut m := (← randUInt32).toUInt64 * rngExcl.toUInt64
+  if m.toUInt32 < rngExcl then
+    let threshold := (0xFFFFFFFF - rng) % rngExcl
+    while m.toUInt32 < threshold do
+      m := (← randUInt32).toUInt64 * rngExcl.toUInt64
+  return (m >>> 32).toUInt32
+
+/-- 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
+  let rngExcl := rng + 1
+  let mut x ← randUInt64
+  let mut leftover := x * rngExcl
+  if leftover < rngExcl then
+    let threshold := (0xFFFFFFFFFFFFFFFF - rng) % rngExcl
+    while leftover < threshold do
+      x ← randUInt64
+      leftover := x * rngExcl
+  return PCG64.mulHi x rngExcl
+
+/-- 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
+  if rng == 0 then return 0
+  else if rng == 0xFFFFFFFF then return (← randUInt32).toUInt64
+  else if rng < 0xFFFFFFFF then return (← randLemireUInt32 rng.toUInt32).toUInt64
+  else if rng == 0xFFFFFFFFFFFFFFFF then randUInt64
+  else randLemireUInt64 rng
+
+/-- 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
+  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"
+  if low > high then throw <| IO.userError (if endpoint then "low > high" else "low >= high")
+  let x ← randBoundedUInt64 (high - low).toNat.toUInt64
+  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
+
+/-- 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 :=
+    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))
 
+/-- `x * y + z`, rounded once, as the C `fma`: the product and the sum are computed exactly, as
+integers scaled by a power of two, and only the final value is rounded to a `Float`. -/
+def fma (x y z : Float) : Float :=
+  match x.toRatParts, y.toRatParts, z.toRatParts with
+  | some (vx, ex), some (vy, ey), some (vz, ez) =>
+    -- the product is `vx * vy * 2 ^ (ex + ey)`, so the sum is exact over the common exponent `e`
+    let ep := ex + ey
+    let e := min ep ez
+    let n := vx * vy * 2 ^ (ep - e).toNat + vz * 2 ^ (ez - e).toNat
+    if e ≥ 0 then Int.divFloat (n * 2 ^ e.toNat) 1 else Int.divFloat n (2 ^ (-e).toNat)
+  | _, _, _ => x * y + z
+
+/-- Sample a `Float` uniformly in `[low, high)`. -/
+def uniform (low high : Float) : RandPCG IO Float := do
+  if low > high then throw <| IO.userError "low > high"
+  let x ← random
+  return fma x (high - low) low
+
 end NumLean

RandomDo/NumLean/PCG64.lean

@@ -7,6 +7,7 @@ module
 
 public import Mathlib.Control.Random
 public import RandomDo.NumLean.SeedSequence
+public import Batteries.Lean.LawfulMonad
 
 /-!
 # The PCG-64 pseudo-random number generator
@@ -24,6 +25,8 @@ native machine arithmetic.
 ## Main definitions
 
 * `PCG64`: the generator state, and its `RandomGen` instance
+* `PCG64.nextUInt64` / `PCG64.nextUInt32`: the 64-bit output, and numpy's 32-bit output, which
+  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`
 
@@ -40,7 +43,8 @@ native machine arithmetic.
 namespace NumLean
 
 /-- The state of a PCG-64 generator: a 128-bit LCG state and a 128-bit increment (which selects the
-stream and must be odd), each stored as a pair of 64-bit words. -/
+stream and must be odd), each stored as a pair of 64-bit words, together with the one-word buffer
+numpy's bit generator keeps for its 32-bit outputs, see `nextUInt32`. -/
 structure PCG64 where
   /-- High 64 bits of the LCG state. -/
   stateHi : UInt64
@@ -50,6 +54,12 @@ structure PCG64 where
   incHi : UInt64
   /-- Low 64 bits of the increment; it is odd for any generator built through the API below. -/
   incLo : UInt64
+  /-- Whether `uinteger` holds the unused upper half of the last output drawn through `nextUInt32`.
+  Mirrors the `has_uint32` field of numpy's `pcg64_state`. -/
+  hasUInt32 : Bool := false
+  /-- The upper 32 bits of the last output drawn through `nextUInt32`, meaningful only while
+  `hasUInt32` is set. Mirrors the `uinteger` field of numpy's `pcg64_state`. -/
+  uinteger : UInt32 := 0
   deriving Repr, DecidableEq
 
 namespace PCG64
@@ -97,6 +107,17 @@ implementation, the state is stepped before the output permutation is applied. -
   let g := g.step
   (g.output, g)
 
+/-- The next 32-bit output, together with the advanced generator, as numpy's `pcg64_next32`: a
+64-bit output is drawn and its low half returned, its high half being kept in `uinteger` to be
+returned by the next call, which then does not step the generator. The buffer survives
+intermediate `nextUInt64` calls, as it does in numpy. -/
+@[inline] def nextUInt32 (g : PCG64) : UInt32 × PCG64 :=
+  if g.hasUInt32 then
+    (g.uinteger, { g with hasUInt32 := false })
+  else
+    let (x, g) := g.nextUInt64
+    (x.toUInt32, { g with hasUInt32 := true, uinteger := (x >>> 32).toUInt32 })
+
 /-- Seed a generator from four 64-bit words, following the reference
 `pcg_setseq_128_srandom_r`: `stateHi:stateLo` is the initial state and `seqHi:seqLo` selects the
 stream, whose increment is `2 * seq + 1`. -/

scripts/check_numpy.py scripts/check_float.pyscripts/check_numpy.py renamed to scripts/check_float.py

@@ -16,6 +16,9 @@ def main (args : List String) : IO Unit := do
 shutil.rmtree(DIR, ignore_errors=True)
 os.makedirs(DIR)
 open(f"{DIR}/dump.lean", "w").write(LEAN.replace("@DIR@", DIR).replace("@N@", str(N)))
+subprocess.run(
+    ["lake", "build"], check=True
+)
 subprocess.run(
     ["lake", "env", "lean", "--run", f"{DIR}/dump.lean", *map(str, SEEDS)], check=True
 )

scripts/check_int.py

@@ -0,0 +1,45 @@
+import os, shutil, subprocess, sys
+from decimal import Decimal
+import numpy as np
+
+N, SEEDS, DIR = 1_000_000, np.random.randint(1_000_000_000, size=5), "test_data"
+
+LEAN = """import RandomDo
+import Batteries.Data.Float.Basic
+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@ do h.putStrLn <| toString (← NumLean.randInt 1000000)
+"""
+
+shutil.rmtree(DIR, ignore_errors=True)
+os.makedirs(DIR)
+open(f"{DIR}/dump.lean", "w").write(LEAN.replace("@DIR@", DIR).replace("@N@", str(N)))
+subprocess.run(
+    ["lake", "build"], check=True
+)
+subprocess.run(
+    ["lake", "env", "lean", "--run", f"{DIR}/dump.lean", *map(str, SEEDS)], check=True
+)
+
+
+def check_distrib(path, xs):
+    with open(path) as f:
+        for i, (line, x) in enumerate(zip(f, xs)):
+            if line.rstrip("\n") != format(Decimal(float(x)), "f"):
+                return i
+    return None
+
+
+ok = True
+for seed in SEEDS:
+    xs = np.random.default_rng(seed).integers(1000000, size=N)
+    bad_line = check_distrib(f"{DIR}/pcg64-{seed}.txt", xs)
+    if bad_line is None:
+        print(f"seed {seed} {N} identical draws")
+    else:
+        print(f"seed {seed} DIVERGENCE at line {bad_line}")
+    ok = ok and bad_line is None
+
+sys.exit(0 if ok else 1)

scripts/check_uniform.py

@@ -0,0 +1,45 @@
+import os, shutil, subprocess, sys
+from decimal import Decimal
+import numpy as np
+
+N, SEEDS, DIR = 1_000_000, np.random.randint(1_000_000_000, size=5), "test_data"
+
+LEAN = """import RandomDo
+import Batteries.Data.Float.Basic
+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@ do h.putStrLn (← NumLean.uniform (-1000000) 1000000).toStringFull
+"""
+
+shutil.rmtree(DIR, ignore_errors=True)
+os.makedirs(DIR)
+open(f"{DIR}/dump.lean", "w").write(LEAN.replace("@DIR@", DIR).replace("@N@", str(N)))
+subprocess.run(
+    ["lake", "build"], check=True
+)
+subprocess.run(
+    ["lake", "env", "lean", "--run", f"{DIR}/dump.lean", *map(str, SEEDS)], check=True
+)
+
+
+def check_distrib(path, xs):
+    with open(path) as f:
+        for i, (line, x) in enumerate(zip(f, xs)):
+            if line.rstrip("\n") != format(Decimal(float(x)), "f"):
+                return i
+    return None
+
+
+ok = True
+for seed in SEEDS:
+    xs = np.random.default_rng(seed).uniform(-1000000, 1000000, size=N)
+    bad_line = check_distrib(f"{DIR}/pcg64-{seed}.txt", xs)
+    if bad_line is None:
+        print(f"seed {seed} {N} identical draws")
+    else:
+        print(f"seed {seed} DIVERGENCE at line {bad_line}")
+    ok = ok and bad_line is None
+
+sys.exit(0 if ok else 1)