FiniteBayes.lean

1 import Mathlib.Algebra.BigOperators.Field
2 import Mathlib.Basic.Real.Basic
3 import Mathlib.Tactic.FieldSimp
4 import Mathlib.Tactic.Ring
5 import Mathlib.Tactic.Linarith
6 import Mathlib.Tactic.Positivity
7 
8 /-! Finite Bayesian updates and aggregation. Priors and denominators are explicit. -/
9 set_option autoImplicit false
10 namespace Komorebi.LSM.FiniteBayes
11 open scoped BigOperators
12 universe u v w
13 variable {I : Type u} {A : Type v} {B : Type w}
14 
15 noncomputable def aggregate [Fintype I] [DecidableEq A]
16     (c : I → A) (p : I → ℝ) (a : A) : ℝ :=
17   ∑ i, if c i = a then p i else 0
18 
19 noncomputable def evidence [Fintype I] (p l : I → ℝ) : ℝ := ∑ i, p i * l i
20 
21 noncomputable def posterior [Fintype I] (p l : I → ℝ) (i : I) : ℝ :=
22   p i * l i / evidence p l
23 
24 def IsPrior [Fintype I] (p : I → ℝ) : Prop :=
25   (∀ i, 0 ≤ p i) ∧ ∑ i, p i = 1
26 
27 section Aggregation
28 variable [Fintype I] [Fintype A] [DecidableEq A]
29 
30 theorem aggregate_sum (c : I → A) (p : I → ℝ) : ∑ a, aggregate c p a = ∑ i, p i := by
31   simp only [aggregate]
32   rw [Finset.sum_comm]
33   simp
34 
35 omit [Fintype A] in
36 theorem aggregate_nonneg (c : I → A) (p : I → ℝ) (hp : ∀ i, 0 ≤ p i) (a : A) :
37     0 ≤ aggregate c p a := by
38   apply Finset.sum_nonneg
39   intro i _
40   split_ifs
41   · exact hp i
42   · exact le_rfl
43 
44 theorem aggregate_prior (c : I → A) (p : I → ℝ) (hp : IsPrior p) :
45     IsPrior (aggregate c p) := by
46   exact ⟨aggregate_nonneg c p hp.1, (aggregate_sum c p).trans hp.2⟩
47 
48 omit [Fintype A] in
49 theorem aggregate_add (c : I → A) (p q : I → ℝ) :
50     aggregate c (fun i => p i + q i) = fun a => aggregate c p a + aggregate c q a := by
51   funext a
52   simp only [aggregate, ← Finset.sum_add_distrib]
53   apply Finset.sum_congr rfl
54   intro i _
55   split_ifs <;> simp
56 
57 omit [Fintype A] in
58 theorem aggregate_scale (c : I → A) (p : I → ℝ) (s : ℝ) :
59     aggregate c (fun i => s * p i) = fun a => s * aggregate c p a := by
60   funext a
61   simp only [aggregate, Finset.mul_sum]
62   apply Finset.sum_congr rfl
63   intro i _
64   split_ifs <;> simp
65 
66 theorem aggregate_weighted_sum (c : I → A) (p : I → ℝ) (f : A → ℝ) :
67     (∑ a, aggregate c p a * f a) = ∑ i, p i * f (c i) := by
68   simp only [aggregate, Finset.sum_mul]
69   rw [Finset.sum_comm]
70   apply Finset.sum_congr rfl
71   intro i _
72   simp [ite_mul]
73 
74 theorem aggregate_comp [Fintype B] [DecidableEq B]
75     (c : I → A) (d : A → B) (p : I → ℝ) :
76     aggregate d (aggregate c p) = aggregate (d ∘ c) p := by
77   funext b
78   have h := aggregate_weighted_sum c p (fun a => if d a = b then 1 else 0)
79   simpa [aggregate, mul_ite] using h
80 
81 omit [Fintype A] [DecidableEq A] in
82 theorem aggregate_id [DecidableEq I] (p : I → ℝ) : aggregate id p = p := by
83   funext i
84   simp [aggregate]
85 
86 end Aggregation
87 
88 section Updates
89 variable [Fintype I]
90 
91 theorem evidence_pos (p l : I → ℝ) (hp : IsPrior p) (hl : ∀ i, 0 < l i) :
92     0 < evidence p l := by
93   obtain ⟨i, hi⟩ : ∃ i, 0 < p i := by
94     by_contra h
95     push Not at h
96     have hz : ∑ i, p i = 0 := Finset.sum_eq_zero (fun i _ => le_antisymm (h i) (hp.1 i))
97     linarith [hp.2]
98   apply Finset.sum_pos'
99   · intro j _; exact mul_nonneg (hp.1 j) (le_of_lt (hl j))
100   · exact ⟨i, Finset.mem_univ i, mul_pos hi (hl i)⟩
101 
102 theorem posterior_sum (p l : I → ℝ) (hz : evidence p l ≠ 0) :
103     ∑ i, posterior p l i = 1 := by
104   simp only [posterior, ← Finset.sum_div]
105   exact div_self hz
106 
107 theorem posterior_prior (p l : I → ℝ) (hp : IsPrior p)
108     (hl : ∀ i, 0 ≤ l i) (hz : 0 < evidence p l) : IsPrior (posterior p l) := by
109   exact ⟨fun i => div_nonneg (mul_nonneg (hp.1 i) (hl i)) hz.le, posterior_sum p l (ne_of_gt hz)⟩
110 
111 theorem posterior_likelihood_scale (p l : I → ℝ) (s : ℝ) (hs : s ≠ 0) :
112     posterior p (fun i => s * l i) = posterior p l := by
113   have hz : evidence p (fun i => s * l i) = s * evidence p l := by
114     simp only [evidence, Finset.mul_sum]; apply Finset.sum_congr rfl; intro i _; ring
115   funext i
116   simp only [posterior, hz]
117   rw [show p i * (s * l i) = s * (p i * l i) by ring, mul_div_mul_left _ _ hs]
118 
119 end Updates
120 section Coarsening
121 variable [Fintype I] [Fintype A] [DecidableEq A]
122 
123 omit [Fintype A] in
124 theorem aggregate_div (c : I → A) (p : I → ℝ) (s : ℝ) :
125     aggregate c (fun i => p i / s) = fun a => aggregate c p a / s := by
126   simpa only [div_eq_mul_inv, mul_comm] using aggregate_scale c p s⁻¹
127 
128 omit [Fintype A] in
129 theorem aggregate_mul_pullback (c : I → A) (p : I → ℝ) (b : A → ℝ) :
130     aggregate c (fun i => p i * b (c i)) = fun a => aggregate c p a * b a := by
131   funext a
132   simp only [aggregate, Finset.sum_mul]
133   apply Finset.sum_congr rfl
134   intro i _
135   by_cases h : c i = a <;> simp [h]
136 
137 theorem evidence_pullback (c : I → A) (p : I → ℝ) (b : A → ℝ) :
138     evidence p (b ∘ c) = evidence (aggregate c p) b := by
139   exact (aggregate_weighted_sum c p b).symm
140 
141 theorem posterior_coarsening_pullback (c : I → A) (p : I → ℝ) (b : A → ℝ) :
142     aggregate c (posterior p (b ∘ c)) = posterior (aggregate c p) b := by
143   unfold posterior
144   rw [aggregate_div, evidence_pullback]
145   simp only [Function.comp_apply, aggregate_mul_pullback]
146 
147 theorem posterior_coarsening_scaled (c : I → A) (p l : I → ℝ) (b : A → ℝ)
148     (s : ℝ) (hs : s ≠ 0) (hl : ∀ i, l i = s * b (c i)) :
149     aggregate c (posterior p l) = posterior (aggregate c p) b := by
150   have heq : l = fun i => s * (b ∘ c) i := funext hl
151   rw [heq, posterior_likelihood_scale p (b ∘ c) s hs]
152   exact posterior_coarsening_pullback c p b
153 
154 end Coarsening
155 
156 section PointPriors
157 variable [Fintype I] [DecidableEq I]
158 
159 def pointMass (i : I) (j : I) : ℝ := if j = i then 1 else 0
160 
161 omit [Fintype I] in
162 theorem pointMass_injective : Function.Injective (pointMass : I → I → ℝ) := by
163   intro i j hij
164   by_contra h
165   have he := congrFun hij i
166   norm_num [pointMass, h] at he
167 
168 noncomputable def pairPrior (i j : I) (k : I) : ℝ := (1/2) * pointMass i k + (1/2) * pointMass j k
169 
170 theorem pointMass_test (i : I) (f : I → ℝ) : ∑ k, pointMass i k * f k = f i := by
171   simp [pointMass, ite_mul]
172 
173 theorem pairPrior_test (i j : I) (f : I → ℝ) :
174     ∑ k, pairPrior i j k * f k = (f i + f j) / 2 := by
175   simp only [pairPrior, add_mul, mul_assoc, Finset.sum_add_distrib, ← Finset.mul_sum,
176     pointMass_test]
177   ring
178 
179 theorem pairPrior_isPrior (i j : I) : IsPrior (pairPrior i j) := by
180   constructor
181   · intro k; simp only [pairPrior, pointMass]; positivity
182   · have h := pairPrior_test i j (fun _ => 1)
183     norm_num at h ⊢
184     exact h
185 
186 theorem pairPrior_evidence (i j : I) (l : I → ℝ) :
187     evidence (pairPrior i j) l = (l i + l j)/2 := pairPrior_test i j l
188 
189 variable [Fintype A] [DecidableEq A]
190 
191 omit [Fintype A] in
192 theorem aggregate_pointMass (c : I → A) (i : I) :
193     aggregate c (pointMass i) = pointMass (c i) := by
194   funext a
195   simp only [aggregate]
196   rw [Finset.sum_eq_single i]
197   · simp [pointMass, eq_comm]
198   · intro j _ h; simp [pointMass, h]
199   · intro h; exact False.elim (h (Finset.mem_univ i))
200 
201 omit [Fintype A] in
202 theorem aggregate_pairPrior (c : I → A) (i j : I) :
203     aggregate c (pairPrior i j) = pairPrior (c i) (c j) := by
204   unfold pairPrior
205   rw [aggregate_add, aggregate_scale, aggregate_scale, aggregate_pointMass, aggregate_pointMass]
206 
207 omit [Fintype A] in
208 theorem aggregate_pair_weight (c : I → A) (i j : I) (l : I → ℝ) (a : A) :
209     aggregate c (fun k => pairPrior i j k * l k) a =
210       ((if c i = a then l i else 0) + (if c j = a then l j else 0))/2 := by
211   have h := pairPrior_test i j (fun k => if c k = a then l k else 0)
212   simpa only [aggregate, mul_ite, mul_zero] using h
213 
214 omit [Fintype A] in
215 theorem aggregate_pair_posterior (c : I → A) (i j : I) (hij : c i ≠ c j)
216     (l : I → ℝ) (hl : ∀ k, 0 < l k) :
217     aggregate c (posterior (pairPrior i j) l) (c i) = l i/(l i+l j) := by
218   have hz : l i + l j ≠ 0 := ne_of_gt (add_pos (hl i) (hl j))
219   unfold posterior
220   simp only [aggregate_div, aggregate_pair_weight, pairPrior_evidence]
221   simp [Ne.symm hij]
222   field_simp
223 
224 end PointPriors
225 end Komorebi.LSM.FiniteBayes

Back to source controls