notation and law of total expectation

LeanMDPs/FinPr.lean

@@ -7,47 +7,49 @@ import Mathlib.Logic.Function.Defs -- Function.Injective
 
 import Mathlib.Data.Finsupp.Indicator
 
-universe u
+--universe u
 
-variable {τ τ₁ τ₂: Type u} 
+variable {τ τ₁ τ₂: Type } 
 variable {T₁ : Finset τ₁} {T₂ : Finset τ₂}
 
 open NNReal
 
-/-- Finite probability space -/
-structure Findist (Ω : Finset τ) : Type u where
+/-- Finite probability distribution -/
+structure Findist (Ω : Finset τ) : Type  where
   p : τ → ℝ≥0 -- TODO: {p : ℝ // 0 ≤ p ∧ p ≤ 1}
   sumsto : (∑ ω ∈ Ω, p ω ) = 1
 
-abbrev Δ : Finset τ → Type u := Findist
+abbrev Δ : Finset τ → Type  := Findist
 
-structure Finprob (τ : Type u) : Type u where
+/-- Finite probability space -/
+structure Finprob (τ : Type ) : Type  where
   Ω : Finset τ
   prob : Findist Ω
 
-
-
+/-- Random variable defined on a finite probability space -/
+structure Finrvar (P : Finprob τ) (ρ : Type ) : Type  where
+  x : τ → ρ   -- actual value of the random variable
+
 /- --------------------------------------------------------------- -/
 namespace Finprob
 
 -- This is the random variable output type
 variable {ρ : Type}
 variable [HMul ℝ≥0 ρ ρ] [HMul ℕ ρ ρ] [AddCommMonoid ρ] 
 
-/-- Handles the necessary product in the expectation -/
+/-- Handles the products in the expectation -/
 instance HMul_NN_R : HMul ℝ≥0 ℝ ℝ where
   hMul := fun a b => ↑a * b
 
 /-- Probability of a sample -/
 def pr (pr : Finprob τ) (t : pr.Ω) := pr.prob.p t.1
 
-/-- Expected value of random variable x : Ω → ρ -/
-def expect (pr : Finprob τ) (x : τ → ρ) : ρ := ∑ ω ∈ pr.Ω, pr.prob.p ω * x ω
+
+/- ---------------------- Index -----------------/
 
 /-- Boolean indicator function -/
 def 𝕀 (cond : τ → Bool) (ω : τ) : ℕ := (cond ω).rec 0 1
 
-
 /-- Indicator is 0 or 1 -/
 theorem ind_zero_one (cond : τ → Bool) (ω : τ) : (𝕀 cond ω = 1) ∨ (𝕀 cond ω = 0) := 
   if h : (cond ω) then 
@@ -70,20 +72,30 @@ theorem indicator_in_zero_one (cond : τ → Bool) :
         (by simp [Finset.mem_insert_self, Finset.pair_comm]) (cond ω) 
 -/
 
-/-- Probability -/
-abbrev ℙ (pr : Finprob τ) (c : τ → Bool) : ℝ≥0 := expect pr (fun ω ↦ ↑(𝕀 c ω))
+
+/- ---------------------- Expectation -----------------/
+
+/-- Expected value of random variable x : Ω → ρ -/
+def expect (pr : Finprob τ) (x : τ → ρ) : ρ := ∑ ω ∈ pr.Ω, pr.prob.p ω * x ω
+
+abbrev expectP {P : Finprob τ} (X : Finrvar P ρ) : ρ := expect P X.x
+
+--scoped[Finprob] 
+notation "𝔼[" X "]" => expectP X 
 
 /-- 
 Conditional expected value E[x | c ] where x is an indicator function
 IMPORTANT: conditional expectation for zero probability event is zero 
 -/
 noncomputable
 def expect_cnd (pr : Finprob τ) (x : τ → ρ) (c : τ → Bool) : ρ :=
-    let f := (fun ω ↦ (𝕀 c ω) * x ω) 
-    (1:ℝ≥0)/(ℙ pr c) * (expect pr f)    
+    let F : Finrvar pr ρ := ⟨fun ω ↦ 𝕀 c ω * x ω⟩
+    let I : Finrvar pr ℝ≥0 := ⟨fun ω ↦ ↑(𝕀 c ω)⟩
+    ((1:ℝ≥0) / 𝔼[ I ]) * 𝔼[ F ]
 
-noncomputable
-abbrev 𝔼c : Finprob τ → (τ → ρ) → (τ → Bool) → ρ := expect_cnd
+
+--noncomputable
+--abbrev 𝔼c : Finprob τ → (τ → ρ) → (τ → Bool) → ρ := expect_cnd
 
 /-- Conditional expectation on a random variable --/
 noncomputable
@@ -92,6 +104,22 @@ def expect_cnd_rv {V : Finset τ₁} [DecidableEq τ₁]
     expect_cnd pr x (fun ω' ↦ if y ω = y ω' then Bool.true else Bool.false)
 
 
+noncomputable 
+def expect_cnd_rv_P  {V : Finset τ₁} [DecidableEq τ₁] {P : Finprob τ} 
+                        (X : Finrvar P ρ) (Y : Finrvar P V) : Finrvar P ρ := 
+    ⟨expect_cnd_rv P X.x Y.x⟩ 
+
+notation "𝔼[" X "|" Y "]" => expect_cnd_rv_P X Y
+
+
+/- ------------ Law of total expectation ----------/
+
+theorem total_expectation {V : Finset τ₁} [DecidableEq τ₁] {P : Finprob τ} (X : Finrvar P ρ) (Y : Finrvar P V) :
+        𝔼[ 𝔼[ X | Y] ] = 𝔼[ X ] := sorry
+
+
+/- ---------------------- Supporting Results -----------------/
+
 /--
 Product of a probability distribution with a dependent probability 
 distributions is a probability distribution. 

lake-manifest.json

@@ -5,7 +5,7 @@
    "type": "git",
    "subDir": null,
    "scope": "",
-   "rev": "e47321100938852afffd41c3c882ba637662d354",
+   "rev": "9c324ade60adf9b37f4f7b49776bd2bb3a82d5e9",
    "name": "«doc-gen4»",
    "manifestFile": "lake-manifest.json",
    "inputRev": "main",
@@ -25,7 +25,7 @@
    "type": "git",
    "subDir": null,
    "scope": "leanprover-community",
-   "rev": "fb40f7550d02605e9a7bf195b85184714c5d06dc",
+   "rev": "69ec65dc81216c48154474008c059388fac15246",
    "name": "mathlib",
    "manifestFile": "lake-manifest.json",
    "inputRev": "master",
@@ -115,7 +115,7 @@
    "type": "git",
    "subDir": null,
    "scope": "leanprover-community",
-   "rev": "43bcb1964528411e47bfa4edd0c87d1face1fce4",
+   "rev": "a4a08d92be3de00def5298059bf707c72dfd3c66",
    "name": "aesop",
    "manifestFile": "lake-manifest.json",
    "inputRev": "master",