Cleanup and tests for polymorphic programs

Author: gaetanserre

Compute.lean

@@ -16,7 +16,8 @@ public import RandomDo.Tactic.Computable.Counterparts
 public import RandomDo.Tactic.Computable.Defs
 public import RandomDo.Tactic.Computable.Deriving
 public import RandomDo.Tactic.Computable.Example
-public import RandomDo.Tactic.Computable.Polymorphic
+public import RandomDo.Tactic.Computable.Polymorphic.Polymorphic
+public import RandomDo.Tactic.Computable.Polymorphic.Scalar
 public import RandomDo.Tactic.IsMarkov.Defs
 public import RandomDo.Tactic.IsMarkov.Deriving
 public import RandomDo.Tactic.IsMarkov.Elab

Polymorphic.lean

@@ -145,8 +145,8 @@ deviation. -/
 /-- 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"
+  if p < 0.0 || 1.0 < p || p.isNaN then
+    throw <| IO.userError s!"expected 0 ≤ p ≤ 1, got {p}"
   let n := n.toUInt64.toFloat
   if n == 0 || p == 0.0 then return 0
   if Float.le p 0.5 then

RandomDo.lean

@@ -55,7 +55,7 @@ noncomputable instance : HasBernoulli Measure ℝ where
 
 /-- A probability outside `[0, 1]` is clamped to it, as in the instance on `Measure`. -/
 instance : HasBernoulli RandM Float where
-  bernoulli p := show RandPCG IO Bool from return (← bernoulli (max 0 (min 1 p))) == 1
+  bernoulli p := return (← NumLean.bernoulli p) == 1
 
 /-- A typeclass for scalars with an exponential. -/
 class HasExp (R : Type) where

RandomDo/NumLean/Distributions.lean

@@ -9,4 +9,5 @@ public import Test.Instances
 public import Test.IsMarkov
 public import Test.Loops
 public import Test.MonadLaws
+public import Test.Polymorphic
 public import Test.Sample

…domDo/Tactic/Computable/Polymorphic.lean …/Computable/Polymorphic/Polymorphic.leanRandomDo/Tactic/Computable/Polymorphic.lean renamed to RandomDo/Tactic/Computable/Polymorphic/Polymorphic.lean RandomDo/Tactic/Computable/Polymorphic.lean renamed to RandomDo/Tactic/Computable/Polymorphic/Polymorphic.lean

@@ -0,0 +1,119 @@
+module
+
+public import Test.IsMarkov
+public meta import RandomDo
+import Batteries.Data.Float.Basic
+
+set_option linter.style.header false
+
+/-!
+# Polymorphic `rdo` programs
+
+The programs of `Test.Computable`, written once over an arbitrary `MeasurableSpaceMonad` `m` and
+drawing through `HasGaussian` and `HasBernoulli`. Read at `m := Measure`, each one is a probability
+measure, checked by `is_markov`, and is the program of `Test.IsMarkov` when there is one. Run at
+`m := RandM`, it samples.
+-/
+
+@[expose] public section
+
+namespace Test.Polymorphic
+
+open Test.IsMarkov NumLean Lean.Elab.Command MeasureTheory ProbabilityTheory
+
+universe v
+
+variable {m : (α : Type) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad m]
+  {R V : Type} [Scalar R] [Scalar V]
+
+def logPolymorphic {α : Type} [MeasurableSpace α] [Lean.ToMessageData α] (prog : RandM α) :
+    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}"
+
+def sumTwo [HasGaussian m R V R] : m R := rdo
+  let x ← HasGaussian.gaussian (m := m) (0 : R) (1 : V)
+  let y ← HasGaussian.gaussian (m := m) (0 : R) (1 : V)
+  return x + y
+
+example : sumTwo (m := Measure) (R := ℝ) (V := NNReal) = Test.IsMarkov.sumTwo := rfl
+
+run_cmd logPolymorphic (sumTwo (m := RandM) (R := Float) (V := Float))
+
+def unfoldSumTwo [HasGaussian m R V R] : m R := rdo
+  let y ← sumTwo (m := m) (V := V)
+  let x ← HasGaussian.gaussian (m := m) (0 : R) (1 : V)
+  return x + y
+
+example : IsProbabilityMeasure (unfoldSumTwo (m := Measure) (R := ℝ) (V := NNReal)) := by
+  is_markov
+
+run_cmd logPolymorphic (unfoldSumTwo (m := RandM) (R := Float) (V := Float))
+
+def centred [HasGaussian m R V R] (c : R) : m R := rdo
+  let x ← HasGaussian.gaussian (m := m) c (1 : V)
+  return x
+
+example : centred (m := Measure) (R := ℝ) (V := NNReal) = Test.IsMarkov.centred := rfl
+
+run_cmd logPolymorphic (centred (m := RandM) (R := Float) (V := Float) 20)
+
+def branchOn [LT R] [DecidableLT R] [HasGaussian m R V R] (c : R) : m R := rdo
+  if 0 < c then
+    let x ← HasGaussian.gaussian (m := m) c (1 : V)
+    return x
+  else
+    let x ← HasGaussian.gaussian (m := m) (0 : R) (1 : V)
+    return x
+
+example : branchOn (m := Measure) (R := ℝ) (V := NNReal) = Test.IsMarkov.branchOn := rfl
+
+run_cmd logPolymorphic (branchOn (m := RandM) (R := Float) (V := Float) 20)
+
+run_cmd logPolymorphic (branchOn (m := RandM) (R := Float) (V := Float) (-1))
+
+def coin [HasBernoulli m R] (p : R) : m Bool := rdo
+  let b ← HasBernoulli.bernoulli (m := m) p
+  return b
+
+example : IsProbabilityMeasure (coin (m := Measure) (1 / 2 : ℝ)) := by is_markov
+
+run_cmd logPolymorphic (coin (m := RandM) (0.5 : Float))
+
+/--
+error: expected 0 ≤ p ≤ 1, got 2.000000
+-/
+#guard_msgs in
+run_cmd logPolymorphic (coin (m := RandM) (2 : Float))
+
+/--
+error: expected 0 ≤ p ≤ 1, got -1.000000
+-/
+#guard_msgs in
+run_cmd logPolymorphic (coin (m := RandM) (-1 : Float))
+
+def twoCoins [HasBernoulli m R] (p : R) : m Bool := rdo
+  let x ← coin (m := m) p
+  let y ← coin (m := m) p
+  return x && y
+
+example : IsProbabilityMeasure (twoCoins (m := Measure) (1 / 2 : ℝ)) := by is_markov
+
+run_cmd logPolymorphic (twoCoins (m := RandM) (0.5 : Float))
+
+def ex1 [HasGaussian m R V R] : m R := rdo
+  let mut x : R := 0
+  for _ in List.range 1000 rdo
+    let y ← HasGaussian.gaussian (m := m) (0 : R) (1 : V)
+    x := x + y
+  return x
+
+example : IsProbabilityMeasure (ex1 (m := Measure) (R := ℝ) (V := NNReal)) := by is_markov
+
+run_cmd logPolymorphic (ex1 (m := RandM) (R := Float) (V := Float))
+
+end Test.Polymorphic
+
+end

Test.lean

@@ -29,11 +29,3 @@ name = "Test"
 name = "dump"
 root = "Dump"
 srcDir = "test_data"
-
-[[lean_exe]]
-name = "compute"
-root = "Compute"
-
-[[lean_exe]]
-name = "polymorphic"
-root = "Polymorphic"