Variation.lean

1 import MachineTheory.Bayes
2 import Mathlib.Algebra.Order.BigOperators.Group.Finset
3 import Mathlib.Tactic.Ring
4 import Mathlib.Data.Fintype.BigOperators
5 
6 /-!
7 Finite total variation on the exact real-valued distributions of Blocks 01--03.
8 The event and conditioning bounds refer to the same two laws and the same event.
9 No posterior is asserted at zero event mass.
10 -/
11 namespace Komorebi.MachineTheory.Variation
12 open scoped BigOperators
13 open Instrument
14 noncomputable section
15 variable {X Y Z : Type*}
16 
17 def l1 [Fintype X] (p : X → ℝ) : ℝ := ∑ x, |p x|
18 
19 def tv [Fintype X] (p q : X → ℝ) : ℝ := l1 (fun x => p x - q x) / 2
20 
21 theorem l1_nonneg [Fintype X] (p : X → ℝ) : 0 ≤ l1 p :=
22   Finset.sum_nonneg (fun x _ => abs_nonneg (p x))
23 
24 theorem tv_nonneg [Fintype X] (p q : X → ℝ) : 0 ≤ tv p q :=
25   div_nonneg (l1_nonneg _) (by norm_num)
26 
27 theorem tv_self [Fintype X] (p : X → ℝ) : tv p p = 0 := by simp [tv, l1]
28 
29 theorem tv_symm [Fintype X] (p q : X → ℝ) : tv p q = tv q p := by
30   simp only [tv, l1, abs_sub_comm]
31 
32 theorem tv_triangle [Fintype X] (p q r : X → ℝ) : tv p r ≤ tv p q + tv q r := by
33   have h : l1 (fun x => p x-r x) ≤ l1 (fun x => p x-q x) + l1 (fun x => q x-r x) := by
34     unfold l1
35     rw [← Finset.sum_add_distrib]
36     exact Finset.sum_le_sum (fun x _ => abs_sub_le (p x) (q x) (r x))
37   unfold tv
38   linarith
39 
40 theorem l1_mass [Fintype X] (p : X → ℝ) (hp : ∀ x, 0 ≤ p x) : l1 p = mass p := by
41   simp only [l1, mass, abs_of_nonneg (hp _)]
42 
43 theorem abs_mass_le_l1 [Fintype X] (p : X → ℝ) : |mass p| ≤ l1 p := by
44   exact Finset.abs_sum_le_sum_abs _ _
45 
46 theorem l1_scale [Fintype X] (c : ℝ) (p : X → ℝ) : l1 (scale c p) = |c| * l1 p := by
47   simp [l1, scale, abs_mul, Finset.mul_sum]
48 
49 theorem tv_scale [Fintype X] (c : ℝ) (p q : X → ℝ) :
50     tv (scale c p) (scale c q) = |c| * tv p q := by
51   unfold tv
52   have h : (fun x => scale c p x - scale c q x) = scale c (fun x => p x-q x) := by
53     funext x; simp [scale, mul_sub]
54   rw [h, l1_scale]; ring
55 
56 theorem tv_le_one [Fintype X] (p q : X → ℝ) (hp : IsDist p) (hq : IsDist q) : tv p q ≤ 1 := by
57   have h : l1 (fun x => p x-q x) ≤ mass p + mass q := by
58     unfold l1 mass
59     rw [← Finset.sum_add_distrib]
60     apply Finset.sum_le_sum
61     intro x _
62     exact (abs_sub (p x) (q x)).trans_eq (by rw [abs_of_nonneg (hp.1 x), abs_of_nonneg (hq.1 x)])
63   rw [hp.2, hq.2] at h
64   unfold tv; linarith
65 
66 /-- A signed finite sum is contracted by exact deterministic aggregation. -/
67 theorem l1_push_le [Fintype X] [Fintype Y] [DecidableEq Y]
68     (f : X → Y) (p : X → ℝ) : l1 (push f p) ≤ l1 p := by
69   calc
70     l1 (push f p) ≤ ∑ y, ∑ x, if f x = y then |p x| else 0 := by
71       apply Finset.sum_le_sum
72       intro y _
73       simpa only [push, abs_ite, abs_zero] using
74         (Finset.abs_sum_le_sum_abs (fun x => if f x = y then p x else 0) Finset.univ)
75     _ = l1 p := by
76       rw [Finset.sum_comm]
77       simp [l1]
78 
79 theorem push_sub [Fintype X] [DecidableEq Y] (f : X → Y) (p q : X → ℝ) :
80     push f (fun x => p x-q x) = fun y => push f p y-push f q y := by
81   funext y
82   simp only [push, ← Finset.sum_sub_distrib]
83   apply Finset.sum_congr rfl
84   intro x _
85   split_ifs <;> simp
86 
87 /-- Forgetting coordinates cannot increase total variation. -/
88 theorem tv_push_le [Fintype X] [Fintype Y] [DecidableEq Y]
89     (f : X → Y) (p q : X → ℝ) : tv (push f p) (push f q) ≤ tv p q := by
90   unfold tv
91   rw [← push_sub]
92   exact div_le_div_of_nonneg_right (l1_push_le f _) (by norm_num)
93 
94 /-- Column-stochastic kernels, oriented like the inherited joint instrument. -/
95 structure Kernel (X Y : Type*) [Fintype Y] where
96   entry : Y → X → ℝ
97   nonnegative : ∀ y x, 0 ≤ entry y x
98   normalized : ∀ x, ∑ y, entry y x = 1
99 
100 theorem kernel_mass [Fintype X] [Fintype Y] (K : Kernel X Y) (p : X → ℝ) :
101     mass (applyKernel K.entry p) = mass p := by
102   unfold mass applyKernel
103   rw [Finset.sum_comm]
104   simp only [← Finset.sum_mul, K.normalized, one_mul]
105 
106 theorem kernel_dist [Fintype X] [Fintype Y] (K : Kernel X Y)
107     (p : X → ℝ) (hp : IsDist p) : IsDist (applyKernel K.entry p) :=
108   ⟨apply_nonnegative _ K.nonnegative _ hp.1, (kernel_mass K p).trans hp.2⟩
109 
110 theorem kernel_contraction [Fintype X] [Fintype Y] (K : Kernel X Y) (p q : X → ℝ) :
111     tv (applyKernel K.entry p) (applyKernel K.entry q) ≤ tv p q := by
112   have h : l1 (fun y => applyKernel K.entry p y-applyKernel K.entry q y) ≤ l1 (fun x => p x-q x) := by
113     calc
114       _ = ∑ y, |∑ x, K.entry y x * (p x-q x)| := by
115         simp [l1, applyKernel, Finset.sum_sub_distrib, mul_sub]
116       _ ≤ ∑ y, ∑ x, K.entry y x * |p x-q x| := by
117         apply Finset.sum_le_sum; intro y _
118         simpa only [abs_mul, abs_of_nonneg (K.nonnegative _ _)] using
119           (Finset.abs_sum_le_sum_abs (fun x => K.entry y x*(p x-q x)) Finset.univ)
120       _ = _ := by
121         rw [Finset.sum_comm]
122         simp only [← Finset.sum_mul, K.normalized, one_mul, l1]
123   unfold tv; exact div_le_div_of_nonneg_right h (by norm_num)
124 
125 /-- Convex mixing charges the full per-column law discrepancy. -/
126 theorem kernel_perturbation [Fintype X] [Fintype Y]
127     (K L : Y → X → ℝ) (p : X → ℝ) (hp : IsDist p) (ε : ℝ)
128     (hε : ∀ x, tv (fun y => K y x) (fun y => L y x) ≤ ε) :
129     tv (applyKernel K p) (applyKernel L p) ≤ ε := by
130   have h : l1 (fun y => applyKernel K p y-applyKernel L p y) ≤ 2*ε := by
131     calc
132       _ = ∑ y, |∑ x, (K y x-L y x)*p x| := by
133         simp [l1, applyKernel, Finset.sum_sub_distrib, sub_mul]
134       _ ≤ ∑ y, ∑ x, |K y x-L y x| *p x := by
135         apply Finset.sum_le_sum; intro y _
136         simpa only [abs_mul, abs_of_nonneg (hp.1 _)] using
137           (Finset.abs_sum_le_sum_abs (fun x => (K y x-L y x)*p x) Finset.univ)
138       _ = ∑ x, l1 (fun y => K y x-L y x)*p x := by
139         rw [Finset.sum_comm]; simp [l1, Finset.sum_mul]
140       _ ≤ ∑ x, (2*ε)*p x := by
141         apply Finset.sum_le_sum; intro x _
142         have he := hε x
143         unfold tv at he
144         exact mul_le_mul_of_nonneg_right (by linarith) (hp.1 x)
145       _ = 2*ε := by rw [← Finset.mul_sum]; change (2*ε)*mass p=2*ε; rw [hp.2, mul_one]
146   unfold tv; linarith
147 
148 /-- Approximate simulation across different finite state types. -/
149 theorem push_kernel_step [Fintype X] [Fintype Y] [Fintype Z] [DecidableEq Z]
150     {W : Type*} [Fintype W] [DecidableEq W]
151     (π : X → Z) (φ : Y → W) (K : Kernel X Y) (L : Kernel Z W)
152     (p : X → ℝ) (q : Z → ℝ) (hp : IsDist p) (ε : ℝ)
153     (hε : ∀ x, tv (push φ (fun y => K.entry y x)) (fun w => L.entry w (π x)) ≤ ε) :
154     tv (push φ (applyKernel K.entry p)) (applyKernel L.entry q) ≤ tv (push π p) q + ε := by
155   let P : W → X → ℝ := fun w x => push φ (fun y => K.entry y x) w
156   let Q : W → X → ℝ := fun w x => L.entry w (π x)
157   have eqP : push φ (applyKernel K.entry p) = applyKernel P p := by
158     funext w
159     unfold push applyKernel
160     have hi : ∀ y, (if φ y = w then ∑ x, K.entry y x * p x else 0) =
161         ∑ x, (if φ y = w then K.entry y x else 0) * p x := by
162       intro y
163       split_ifs <;> simp_all
164     simp_rw [hi]
165     rw [Finset.sum_comm]
166     apply Finset.sum_congr rfl
167     intro x _
168     simp [P, push, Finset.sum_mul, ite_mul]
169   have eqQ : applyKernel Q p = applyKernel L.entry (push π p) := by
170     funext w
171     unfold applyKernel push
172     simp_rw [Finset.mul_sum, mul_ite, mul_zero]
173     rw [Finset.sum_comm]
174     simp [Q]
175   rw [eqP]
176   have h1 := kernel_perturbation P Q p hp ε hε
177   have h2 := kernel_contraction L (push π p) q
178   have h3 := tv_triangle (applyKernel P p) (applyKernel Q p) (applyKernel L.entry q)
179   rw [eqQ] at h3 h1
180   linarith
181 
182 /-- Exact restriction, with zero outside the named event. -/
183 def restrict [DecidableEq X] (E : Finset X) (p : X → ℝ) : X → ℝ :=
184   fun x => if x ∈ E then p x else 0
185 
186 theorem restrict_nonneg [DecidableEq X] (E : Finset X) (p : X → ℝ)
187     (hp : ∀ x, 0 ≤ p x) : ∀ x, 0 ≤ restrict E p x := by
188   intro x; unfold restrict; split_ifs <;> simp_all
189 
190 theorem mass_sub [Fintype X] (p q : X → ℝ) :
191     mass (fun x => p x-q x) = mass p-mass q := by simp [mass, Finset.sum_sub_distrib]
192 
193 theorem mass_split [Fintype X] [DecidableEq X] (E : Finset X) (p : X → ℝ) :
194     mass (restrict E p) + mass (restrict Eᶜ p) = mass p := by
195   unfold mass
196   rw [← Finset.sum_add_distrib]
197   apply Finset.sum_congr rfl
198   intro x _
199   by_cases h : x ∈ E <;> simp [restrict, h]
200 
201 theorem l1_split [Fintype X] [DecidableEq X] (E : Finset X) (p q : X → ℝ) :
202     l1 (fun x => restrict E p x-restrict E q x) +
203       l1 (fun x => restrict Eᶜ p x-restrict Eᶜ q x) = l1 (fun x => p x-q x) := by
204   unfold l1
205   rw [← Finset.sum_add_distrib]
206   apply Finset.sum_congr rfl
207   intro x _
208   by_cases h : x ∈ E <;> simp [restrict, h]
209 
210 /-- The complement pays for the difference of event masses; this yields the sharp constant. -/
211 theorem restricted_discrepancy [Fintype X] [DecidableEq X]
212     (E : Finset X) (p q : X → ℝ) (hm : mass p = mass q) :
213     l1 (fun x => restrict E p x-restrict E q x) +
214       |mass (restrict E p)-mass (restrict E q)| ≤ 2 * tv p q := by
215   have hc := abs_mass_le_l1 (fun x => restrict Eᶜ p x-restrict Eᶜ q x)
216   rw [mass_sub] at hc
217   have he : mass (restrict Eᶜ p)-mass (restrict Eᶜ q) =
218       -(mass (restrict E p)-mass (restrict E q)) := by
219     have h1 := mass_split E p; have h2 := mass_split E q; linarith
220   rw [he, abs_neg] at hc
221   have hs := l1_split E p q
222   unfold tv
223   linarith
224 
225 /-- Event probabilities differ by at most total variation, not twice that value. -/
226 theorem event_bound [Fintype X] [DecidableEq X]
227     (E : Finset X) (p q : X → ℝ) (hm : mass p = mass q) :
228     |mass (restrict E p)-mass (restrict E q)| ≤ tv p q := by
229   have hi := abs_mass_le_l1 (fun x => restrict E p x-restrict E q x)
230   rw [mass_sub] at hi
231   have hc := restricted_discrepancy E p q hm
232   linarith
233 
234 theorem normalize_eq_scale [Fintype X] (p : X → ℝ) :
235     normalize p = scale (1 / mass p) p := by
236   funext x; simp [normalize, scale, div_eq_mul_inv, mul_comm]
237 
238 theorem tv_scaled_same [Fintype X] (p : X → ℝ) (hp : ∀ x, 0 ≤ p x) (c d : ℝ) :
239     tv (scale c p) (scale d p) = |c-d| * mass p / 2 := by
240   unfold tv
241   have he : (fun x => scale c p x-scale d p x) = scale (c-d) p := by
242     funext x; simp [scale, sub_mul]
243   rw [he, l1_scale, l1_mass p hp]
244 
245 /-- Normalization is controlled by unnormalized error and the discrepancy in mass. -/
246 theorem normalize_bound [Fintype X] (u v : X → ℝ) (hv : ∀ x, 0 ≤ v x)
247     (huM : 0 < mass u) (hvM : 0 < mass v) :
248     tv (normalize u) (normalize v) ≤
249       (l1 (fun x => u x-v x) + |mass u-mass v|) / (2*mass u) := by
250   have h := tv_triangle (scale (1/mass u) u) (scale (1/mass u) v) (scale (1/mass v) v)
251   rw [tv_scale, tv_scaled_same v hv, abs_of_pos (one_div_pos.mpr huM)] at h
252   have hid : (1 / mass u - 1 / mass v) = (mass v-mass u)/(mass u*mass v) := by
253     field_simp
254   rw [hid, abs_div, abs_mul, abs_of_pos huM, abs_of_pos hvM, abs_sub_comm (mass v)] at h
255   have he : 1 / mass u * tv u v +
256       |mass u-mass v| / (mass u*mass v) * mass v / 2 =
257       (l1 (fun x => u x-v x) + |mass u-mass v|) / (2*mass u) := by
258     unfold tv
259     field_simp
260   rw [he] at h
261   simpa only [← normalize_eq_scale] using h
262 
263 /-- Either denominator gives a valid conditioning bound. -/
264 theorem condition_left [Fintype X] [DecidableEq X]
265     (E : Finset X) (p q : X → ℝ) (hp : IsDist p) (hq : IsDist q)
266     (hEp : 0 < mass (restrict E p)) (hEq : 0 < mass (restrict E q)) :
267     tv (normalize (restrict E p)) (normalize (restrict E q)) ≤
268       tv p q / mass (restrict E p) := by
269   have hn := normalize_bound (restrict E p) (restrict E q) (restrict_nonneg E q hq.1) hEp hEq
270   have hr := restricted_discrepancy E p q (hp.2.trans hq.2.symm)
271   have hd := div_le_div_of_nonneg_right hr (show 0 ≤ 2*mass (restrict E p) by positivity)
272   have he : (2*tv p q)/(2*mass (restrict E p)) = tv p q/mass (restrict E p) := by ring
273   rw [he] at hd
274   exact hn.trans hd
275 
276 /-- Core 3.8's sharp symmetric conditioning factor. -/
277 theorem condition_bound [Fintype X] [DecidableEq X]
278     (E : Finset X) (p q : X → ℝ) (hp : IsDist p) (hq : IsDist q)
279     (hEp : 0 < mass (restrict E p)) (hEq : 0 < mass (restrict E q)) :
280     tv (normalize (restrict E p)) (normalize (restrict E q)) ≤
281       min 1 (tv p q / max (mass (restrict E p)) (mass (restrict E q))) := by
282   apply le_min
283   · exact tv_le_one _ _
284       (normalize_dist _ (restrict_nonneg E p hp.1) hEp)
285       (normalize_dist _ (restrict_nonneg E q hq.1) hEq)
286   · by_cases he : mass (restrict E p) ≤ mass (restrict E q)
287     · rw [max_eq_right he, tv_symm p q, tv_symm (normalize (restrict E p))]
288       exact condition_left E q p hq hp hEq hEp
289     · rw [max_eq_left (le_of_not_ge he)]
290       exact condition_left E p q hp hq hEp hEq
291 
292 /-- Positive probability in either law above the whole-law defect excludes an impossible counterpart. -/
293 theorem event_positive [Fintype X] [DecidableEq X]
294     (E : Finset X) (p q : X → ℝ) (hp : IsDist p) (hq : IsDist q)
295     (d : ℝ) (hd : tv p q ≤ d) (he : d < mass (restrict E p)) :
296     0 < mass (restrict E q) := by
297   have h := event_bound E p q (hp.2.trans hq.2.symm)
298   have h1 := le_abs_self (mass (restrict E p)-mass (restrict E q))
299   linarith
300 
301 /-- Condition on the same event, then retain only the requested final output. -/
302 theorem conditional_output_bound [Fintype X] [Fintype Y] [DecidableEq X] [DecidableEq Y]
303     (f : X → Y) (E : Finset X) (p q : X → ℝ) (hp : IsDist p) (hq : IsDist q)
304     (hEp : 0 < mass (restrict E p)) (hEq : 0 < mass (restrict E q))
305     (d : ℝ) (hd : tv p q ≤ d) :
306     tv (push f (normalize (restrict E p))) (push f (normalize (restrict E q))) ≤
307       min 1 (d / max (mass (restrict E p)) (mass (restrict E q))) := by
308   apply (tv_push_le f _ _).trans
309   apply (condition_bound E p q hp hq hEp hEq).trans
310   exact min_le_min_left 1
311     (div_le_div_of_nonneg_right hd (le_trans (le_of_lt hEp) (le_max_left _ _)))
312 
313 /-- A bounded common loss is Lipschitz in total variation with its full range as constant. -/
314 theorem expectation_bound [Fintype X] (p q f : X → ℝ)
315     (hm : mass p = mass q) (B : ℝ) (_hB : 0 ≤ B) (hf : ∀ x, 0 ≤ f x ∧ f x ≤ B) :
316     |(∑ x, f x*p x) - (∑ x, f x*q x)| ≤ B * tv p q := by
317   have hz : ∑ x, (p x-q x) = 0 := by
318     rw [Finset.sum_sub_distrib]; change mass p-mass q=0; rw [hm, sub_self]
319   have he : (∑ x, f x*p x) - (∑ x, f x*q x) =
320       ∑ x, (f x-B/2)*(p x-q x) := by
321     calc
322       _ = ∑ x, f x*(p x-q x) := by simp [mul_sub, Finset.sum_sub_distrib]
323       _ = (∑ x, f x*(p x-q x)) - (B/2)*(∑ x, (p x-q x)) := by rw [hz, mul_zero, sub_zero]
324       _ = _ := by
325         rw [Finset.mul_sum, ← Finset.sum_sub_distrib]
326         apply Finset.sum_congr rfl
327         intro x _
328         ring
329   rw [he]
330   calc
331     _ ≤ ∑ x, |(f x-B/2)*(p x-q x)| := Finset.abs_sum_le_sum_abs _ _
332     _ ≤ ∑ x, (B/2)*|p x-q x| := by
333       apply Finset.sum_le_sum; intro x _
334       rw [abs_mul]
335       have hb : |f x-B/2| ≤ B/2 := by rw [abs_le]; constructor <;> linarith [hf x]
336       exact mul_le_mul_of_nonneg_right hb (abs_nonneg _)
337     _ = B*tv p q := by simp only [← Finset.mul_sum, tv, l1]; ring
338 
339 /-- Composition of exact aggregation maps, including maps that are not surjective. -/
340 theorem push_comp [Fintype X] [Fintype Y] [DecidableEq Y] [DecidableEq Z]
341     (f : X → Y) (g : Y → Z) (p : X → ℝ) : push g (push f p) = push (g ∘ f) p := by
342   funext z
343   unfold push
344   have hi : ∀ y, (if g y=z then ∑ x, if f x=y then p x else 0 else 0) =
345       ∑ x, if f x=y then (if g y=z then p x else 0) else 0 := by
346     intro y
347     by_cases h : g y=z <;> simp [h]
348   simp_rw [hi]
349   rw [Finset.sum_comm]
350   simp
351 
352 /-- Pushing the second coordinate retains its actual association with the first. -/
353 theorem push_pair_right [Fintype X] [Fintype Y] [DecidableEq Y] [DecidableEq Z]
354     (f : X → Z) (p : Y → X → ℝ) (y : Y) (z : Z) :
355     push (fun t : Y × X => (t.1, f t.2)) (fun t => p t.1 t.2) (y,z) = push f (p y) z := by
356   simp [push, Fintype.sum_prod_type, Prod.mk.injEq, ite_and]
357 
358 /-- Joint-law error with a common mixing distribution is the weighted row error. -/
359 theorem tv_weighted_rows [Fintype X] [Fintype Y]
360     (w : Y → ℝ) (hw : ∀ y, 0 ≤ w y) (p q : Y → X → ℝ) :
361     tv (fun t : Y × X => w t.1*p t.1 t.2) (fun t => w t.1*q t.1 t.2) =
362       ∑ y, w y * tv (p y) (q y) := by
363   simp only [tv, l1, Fintype.sum_prod_type, ← mul_sub, abs_mul, abs_of_nonneg (hw _),
364     ← Finset.mul_sum, Finset.sum_div, mul_div_assoc]
365 
366 /-- Mixing over the same policy cannot enlarge a uniform instrument-row error. -/
367 theorem tv_weighted_rows_le [Fintype X] [Fintype Y]
368     (w : Y → ℝ) (hw : IsDist w) (p q : Y → X → ℝ) (ε : ℝ)
369     (hε : ∀ y, tv (p y) (q y) ≤ ε) :
370     tv (fun t : Y × X => w t.1*p t.1 t.2) (fun t => w t.1*q t.1 t.2) ≤ ε := by
371   rw [tv_weighted_rows w hw.1]
372   calc
373     _ ≤ ∑ y, w y*ε := Finset.sum_le_sum (fun y _ => mul_le_mul_of_nonneg_left (hε y) (hw.1 y))
374     _ = ε := by rw [← Finset.sum_mul]; change mass w*ε=ε; rw [hw.2, one_mul]
375 
376 end
377 end Komorebi.MachineTheory.Variation

Back to source controls