Support for bernoulli

Author: gaetanserre

RandomDo/Tactic/Computable/Polymorphic.lean

@@ -6,11 +6,19 @@ Authors: Gaëtan Serré
 module
 
 public import Mathlib.Probability.Distributions.Gaussian.Real
+public import Mathlib.Probability.Distributions.Bernoulli
+public import Mathlib.MeasureTheory.Function.SpecialFunctions.Basic
 public import RandomDo.Monad.Instances
 public import RandomDo.NumLean.Distributions
 
 /-!
+# Polymorphic `rdo` programs
 
+A program written over an arbitrary `MeasurableSpaceMonad` `m`, drawing through the classes of this
+file, is read at `m := Measure` to prove things about it and run at `m := RandM` to sample from it.
+Each class has an instance of each kind: the distribution of Mathlib on `ℝ`, and the sampler of
+`NumLean` on `Float`. The scalar classes `HasExp`, `HasLog` and `HasSqrt` do the same for the
+functions a program computes with.
 -/
 
 @[expose] public section
@@ -35,3 +43,56 @@ instance instMeasurableSpaceFloat : MeasurableSpace Float := ⊤
 
 instance : HasGaussian RandM Float Float Float where
   gaussian μ v := normal' μ v
+
+/-- A typeclass for monads that can draw from a Bernoulli distribution. -/
+class HasBernoulli (m : (α : Type) → [MeasurableSpace α] → Type v) (P : Type) where
+  /-- Draw `true` with probability `p`, and `false` otherwise. -/
+  bernoulli : P → m Bool
+
+/-- A probability outside `[0, 1]` is clamped to it. -/
+noncomputable instance : HasBernoulli Measure ℝ where
+  bernoulli p := bernoulliMeasure true false (Set.projIcc 0 1 zero_le_one p)
+
+/-- 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
+
+/-- A typeclass for scalars with an exponential. -/
+class HasExp (R : Type) where
+  /-- The exponential. -/
+  exp : R → R
+
+noncomputable instance : HasExp ℝ := ⟨Real.exp⟩
+
+@[fun_prop]
+lemma HasExp.measurable_exp_real : Measurable (HasExp.exp : ℝ → ℝ) := Real.measurable_exp
+
+instance : HasExp Float := ⟨Float.exp⟩
+
+/-- A typeclass for scalars with a logarithm. -/
+class HasLog (R : Type) where
+  /-- The logarithm. -/
+  log : R → R
+
+noncomputable instance : HasLog ℝ := ⟨Real.log⟩
+
+@[fun_prop]
+lemma HasLog.measurable_log_real : Measurable (HasLog.log : ℝ → ℝ) := Real.measurable_log
+
+instance : HasLog Float := ⟨Float.log⟩
+
+/-- A typeclass for scalars with a square root. -/
+class HasSqrt (R : Type) where
+  /-- The square root. -/
+  sqrt : R → R
+
+noncomputable instance : HasSqrt ℝ := ⟨Real.sqrt⟩
+
+@[fun_prop]
+lemma HasSqrt.measurable_sqrt_real : Measurable (HasSqrt.sqrt : ℝ → ℝ) :=
+  Real.continuous_sqrt.measurable
+
+instance : HasSqrt Float := ⟨Float.sqrt⟩
+
+/-- A natural number as a `Float`, so that programs polymorphic in the scalars can cast counts. -/
+instance instNatCastFloat : NatCast Float := ⟨Nat.toFloat⟩

RandomDo/Tactic/IsMarkov/Elab.lean

@@ -136,6 +136,9 @@ def closeLeaf (g : MVarId) : MetaM (List MVarId) := do
   if let some gs ← observing? (g.applyConst ``IsMarkov.gaussianReal) then
     trace[is_markov] "`gaussianReal` leaf: handing back the measurability of its parameters"
     return gs
+  if let some gs ← observing? (g.applyConst ``IsMarkov.bernoulliMeasure) then
+    trace[is_markov] "`bernoulliMeasure` leaf: handing back the measurability of its parameters"
+    return gs
   return [g]
 
 /-- The constant heading the body of `κ`, when it is a definition the tactic could look through. -/
@@ -172,7 +175,9 @@ def abstractLoopVars (vars : Array FVarId) (g : MVarId) : MetaM MVarId := do
 
 /-- Turn a goal `IsMarkov κ` into the list of goals the user is left with. -/
 partial def isMarkovCore (g : MVarId) (fuel : Nat) : MetaM (List MVarId) := g.withContext do
-  let target ← instantiateMVars (← g.getType)
+  /- The annotations a goal may carry, e.g. the one a tactic `have` leaves on the goal after it,
+  would hide the head of the statement. -/
+  let target := (← instantiateMVars (← g.getType)).cleanupAnnotations
   -- `IsMarkov` takes five arguments: `γ`, `α`, their `MeasurableSpace` instances, and `κ`.
   unless target.isAppOfArity ``IsMarkov 5 do
     trace[is_markov] "not an `IsMarkov` goal, handed back: {target}"
@@ -251,7 +256,8 @@ partial def isMarkovCore (g : MVarId) (fuel : Nat) : MetaM (List MVarId) := g.wi
         let mut goals := []
         let mut side := []
         for g' in gs do
-          if (← instantiateMVars (← g'.getType)).isAppOfArity ``IsMarkov 5 then
+          let t := (← instantiateMVars (← g'.getType)).cleanupAnnotations
+          if t.isAppOfArity ``IsMarkov 5 then
             goals := goals ++ (← isMarkovCore g' fuel)
           else
             side := side ++ [g']
@@ -280,7 +286,7 @@ partial def isMarkovCore (g : MVarId) (fuel : Nat) : MetaM (List MVarId) := g.wi
       if ← g.isAssigned then return leftover
       /- The goal was not closed, so we try to unfold names in the head of the program until we
       reach a known shape. If that fails, we leave the goal to the user. -/
-      match ← unfoldToKnownShape (← instantiateMVars (← g.getType)) fuel with
+      match ← unfoldToKnownShape (← instantiateMVars (← g.getType)).cleanupAnnotations fuel with
       | some target =>
         trace[is_markov] "unfolded the head definition to: {target.appArg!}"
         isMarkovCore (← g.change target) (fuel - 1)
@@ -316,7 +322,7 @@ lemma _root_.isProbabilityMeasure_of_isMarkov {α : Type*} [MeasurableSpace α]
 /-- Bring a goal of the form `IsProbabilityMeasure μ` into the form `IsMarkov fun _ : Unit ↦ μ`, so
 that `isMarkovCore` can be applied. -/
 def toIsMarkovGoal (g : MVarId) : MetaM MVarId := do
-  let target ← instantiateMVars (← g.getType)
+  let target := (← instantiateMVars (← g.getType)).cleanupAnnotations
   unless target.isAppOfArity ``IsProbabilityMeasure 3 do
     return g
   match ← g.applyConst ``isProbabilityMeasure_of_isMarkov with

RandomDo/Tactic/IsMarkov/Lemmas.lean

@@ -12,6 +12,7 @@ public import RandomDo.Tactic.IsMarkov.ForInStep
 public import Mathlib.MeasureTheory.Measure.ProbabilityMeasure
 public import Mathlib.Data.List.OfFn
 public import Mathlib.Probability.Distributions.Gaussian.Real
+public import Mathlib.Probability.Distributions.Bernoulli
 
 /-!
 # Markov property of `rdo` programs
@@ -31,6 +32,8 @@ complex program to the Markov property/measurability of its underlying mathemati
   the bound variable is Markovian in the parameter.
 * `gaussianReal`: A Gaussian distribution whose mean and variance depend measurably on the parameter
   is Markovian in the parameter.
+* `bernoulliMeasure`: A Bernoulli distribution whose two outcomes and probability depend measurably
+  on the parameter is Markovian in the parameter.
 * `comp`: Composing a Markov kernel `κ` with a measurable function `g` is Markovian in the
   parameter.
 * `ite`: A conditional `rdo` program that chooses between two Markov kernels `κ` and `η` based on a
@@ -43,6 +46,7 @@ complex program to the Markov property/measurability of its underlying mathemati
 * `forInList_comp`, `forInArray_comp`, `forInVector_comp`: The same three, for a loop over a
   collection the program takes as an argument. The body is then asked to be Markovian jointly in the
   parameter and in the element, which the fixed collections do not need.
+* `forIn_nil`, `forIn_cons`: A `for` loop over a list, unrolled one element at a time.
 * `breakRunK`: The case analysis a program performs after a loop that returns early, on the `Option`
   slot holding the returned value, is Markovian as soon as both of its branches are.
 -/
@@ -51,6 +55,7 @@ complex program to the Markov property/measurability of its underlying mathemati
 
 open MeasureTheory ProbabilityTheory Function
 open MeasurableSpacePure
+open scoped ENNReal
 
 namespace IsMarkov
 
@@ -91,6 +96,18 @@ lemma gaussianReal {m : γ → ℝ} {v : γ → NNReal} (hm : Measurable m) (hv
     IsMarkov fun c ↦ ProbabilityTheory.gaussianReal (m c) (v c) :=
   ⟨ProbabilityTheory.measurable_gaussianReal.comp (hm.prodMk hv), fun _ ↦ inferInstance⟩
 
+lemma bernoulliMeasure {x y : γ → α} {p : γ → unitInterval} (hx : Measurable x)
+    (hy : Measurable y) (hp : Measurable p) :
+    IsMarkov fun c ↦ ProbabilityTheory.bernoulliMeasure (x c) (y c) (p c) := by
+  refine ⟨Measure.measurable_of_measurable_coe _ fun s hs ↦ ?_, fun _ ↦ inferInstance⟩
+  simp only [bernoulliMeasure_def, Measure.add_apply, Measure.smul_apply,
+    Measure.dirac_apply' _ hs, ENNReal.smul_def, smul_eq_mul]
+  have hp' : Measurable fun c ↦ ((unitInterval.toNNReal (p c) : ℝ≥0∞)) := by fun_prop
+  have hq' : Measurable fun c ↦ ((unitInterval.toNNReal (unitInterval.symm (p c)) : ℝ≥0∞)) := by
+    fun_prop
+  exact (hp'.mul ((measurable_one.indicator hs).comp hx)).add
+    (hq'.mul ((measurable_one.indicator hs).comp hy))
+
 lemma comp {κ : γ → Measure α} (hκ : IsMarkov κ) {g : σ → γ} (hg : Measurable g) :
     IsMarkov fun c ↦ κ (g c) := ⟨hκ.measurable.comp hg, fun _ ↦ hκ.isProbabilityMeasure _⟩
 
@@ -228,11 +245,14 @@ private lemma forIn_eq_listLoop (l : List ι) (b : σ) (g : ι → σ → Measur
     MeasurableSpaceForIn.forIn (m := Measure) l b g = listLoop g l b :=
   loop_eq_listLoop g l b l _ (fun _ _ _ ↦ rfl) ⟨[], rfl⟩
 
-private lemma forIn_nil (b : σ) (g : ι → σ → Measure (ForInStep σ)) :
+/-- A `for` loop over the empty list returns its initial state. -/
+lemma forIn_nil (b : σ) (g : ι → σ → Measure (ForInStep σ)) :
     MeasurableSpaceForIn.forIn (m := Measure) ([] : List ι) b g = mPure b :=
   forIn_eq_listLoop _ _ _
 
-private lemma forIn_cons (a : ι) (l : List ι) (b : σ) (g : ι → σ → Measure (ForInStep σ)) :
+/-- A `for` loop over `a :: l` runs its body on `a`, then stops or carries on with the loop over
+`l`. -/
+lemma forIn_cons (a : ι) (l : List ι) (b : σ) (g : ι → σ → Measure (ForInStep σ)) :
     MeasurableSpaceForIn.forIn (m := Measure) (a :: l) b g
       = g a b >>=ₘ fun step ↦ ForInStep.casesOn (motive := fun _ ↦ Measure σ) step mPure
           fun b' ↦ MeasurableSpaceForIn.forIn (m := Measure) l b' g := by