JointRefinement.lean

1 import KomorebiLSM.FinitePartitions
2 
3 /-! P9's actual iterative partition construction, termination and least refinement. -/
4 set_option autoImplicit false
5 namespace Komorebi.LSM.FinitePartitions
6 open scoped BigOperators
7 universe u v w
8 variable {I : Type u} [Fintype I] {U : Type v} {Y : Type w}
9 
10 def Lumpable (K : U → I → I → ℝ) (r : Setoid I) : Prop :=
11   ∀ u i j, r i j → blockAggregate r (K u i) = blockAggregate r (K u j)
12 
13 def LikelihoodConstant (l : Y → I → ℝ) (r : Setoid I) : Prop :=
14   ∀ y i j, r i j → l y i = l y j
15 
16 noncomputable def transitionSplit (K : U → I → I → ℝ) (r : Setoid I) : Setoid I where
17   r i j := r i j ∧ ∀ u k, blockMass r (K u i) k = blockMass r (K u j) k
18   iseqv := {
19     refl := fun _ => ⟨r.refl _, fun _ _ => rfl⟩
20     symm := fun h => ⟨r.symm h.1, fun u k => (h.2 u k).symm⟩
21     trans := fun h h' => ⟨r.trans h.1 h'.1, fun u k => (h.2 u k).trans (h'.2 u k)⟩ }
22 
23 @[simp] theorem transitionSplit_rel (K : U → I → I → ℝ) (r : Setoid I) (i j : I) :
24     transitionSplit K r i j ↔ r i j ∧
25       ∀ u k, blockMass r (K u i) k = blockMass r (K u j) k := Iff.rfl
26 
27 theorem transitionSplit_le (K : U → I → I → ℝ) (r : Setoid I) :
28     transitionSplit K r ≤ r := by
29   intro i j h
30   exact h.1
31 
32 theorem split_row_equality (K : U → I → I → ℝ) (r : Setoid I) (i j : I)
33     (h : transitionSplit K r i j) (u : U) :
34     blockAggregate r (K u i) = blockAggregate r (K u j) := by
35   funext q
36   induction q using Quotient.inductionOn with
37   | h k => simpa only [blockAggregate_mk] using h.2 u k
38 
39 theorem transitionSplit_deterministic [DecidableEq I] (f : U → I → I) (r : Setoid I)
40     (i j : I) :
41     transitionSplit (fun u x => FiniteBayes.pointMass (f u x)) r i j ↔
42       r i j ∧ ∀ u, r (f u i) (f u j) := by
43   classical
44   constructor
45   · intro h
46     refine ⟨h.1, ?_⟩
47     intro u
48     have he := split_row_equality (fun u x => FiniteBayes.pointMass (f u x)) r i j h u
49     change FiniteBayes.aggregate (Quotient.mk r) (FiniteBayes.pointMass (f u i)) =
50       FiniteBayes.aggregate (Quotient.mk r) (FiniteBayes.pointMass (f u j)) at he
51     rw [FiniteBayes.aggregate_pointMass, FiniteBayes.aggregate_pointMass] at he
52     exact Quotient.exact (FiniteBayes.pointMass_injective he)
53   · rintro ⟨hij, hf⟩
54     refine ⟨hij, ?_⟩
55     intro u k
56     have he : blockAggregate r (FiniteBayes.pointMass (f u i)) =
57         blockAggregate r (FiniteBayes.pointMass (f u j)) := by
58       simp only [blockAggregate, FiniteBayes.aggregate_pointMass, Quotient.sound (hf u)]
59     simpa only [blockAggregate_mk] using congrFun he (Quotient.mk r k)
60 
61 theorem transitionSplit_fixed_iff (K : U → I → I → ℝ) (r : Setoid I) :
62     transitionSplit K r = r ↔ Lumpable K r := by
63   constructor
64   · intro h u i j hij
65     exact split_row_equality K r i j (by rwa [h]) u
66   · intro h
67     apply le_antisymm (transitionSplit_le K r)
68     intro i j hij
69     refine ⟨hij, ?_⟩
70     intro u k
71     simpa only [blockAggregate_mk] using congrFun (h u i j hij) (Quotient.mk r k)
72 
73 theorem transitionSplit_mono (K : U → I → I → ℝ) {r s : Setoid I} (h : r ≤ s) :
74     transitionSplit K r ≤ transitionSplit K s := by
75   intro i j hij
76   refine ⟨h hij.1, ?_⟩
77   intro u k
78   have he := same_fine_implies_same_coarse h (K u i) (K u j)
79     (split_row_equality K r i j hij u)
80   simpa only [blockAggregate_mk] using congrFun he (Quotient.mk s k)
81 
82 noncomputable def refinementIter (K : U → I → I → ℝ) (r : Setoid I) : ℕ → Setoid I
83   | 0 => r
84   | n+1 => transitionSplit K (refinementIter K r n)
85 
86 @[simp] theorem refinementIter_zero (K : U → I → I → ℝ) (r : Setoid I) :
87     refinementIter K r 0 = r := rfl
88 
89 @[simp] theorem refinementIter_succ (K : U → I → I → ℝ) (r : Setoid I) (n : ℕ) :
90     refinementIter K r (n+1) = transitionSplit K (refinementIter K r n) := rfl
91 
92 theorem refinementIter_step_le (K : U → I → I → ℝ) (r : Setoid I) (n : ℕ) :
93     refinementIter K r (n+1) ≤ refinementIter K r n := transitionSplit_le _ _
94 
95 theorem refinementIter_le_initial (K : U → I → I → ℝ) (r : Setoid I) (n : ℕ) :
96     refinementIter K r n ≤ r := by
97   induction n with
98   | zero => exact le_rfl
99   | succ n ih => exact le_trans (refinementIter_step_le K r n) ih
100 
101 theorem refinementIter_stays_fixed (K : U → I → I → ℝ) (r : Setoid I) (n : ℕ)
102     (h : refinementIter K r (n+1) = refinementIter K r n) (k : ℕ) :
103     refinementIter K r (n+k) = refinementIter K r n := by
104   induction k with
105   | zero => simp
106   | succ k ih =>
107     rw [Nat.add_succ, refinementIter_succ, ih]
108     exact h
109 
110 theorem refinement_stabilizes (K : U → I → I → ℝ) (r : Setoid I) :
111     ∃ n ≤ Fintype.card I - blockCount r,
112       refinementIter K r (n+1) = refinementIter K r n := by
113   let N := Fintype.card I - blockCount r
114   have hcard := blockCount_le r
115   by_contra h
116   have hne : ∀ n ≤ N, refinementIter K r (n+1) ≠ refinementIter K r n := by
117     intro n hn he
118     exact h ⟨n, hn, he⟩
119   have hgrow : ∀ n ≤ N+1, blockCount r + n ≤ blockCount (refinementIter K r n) := by
120     intro n
121     induction n with
122     | zero => intro _; simp
123     | succ n ih =>
124       intro hn
125       have hprev := ih (by omega)
126       have hlt := strict_refinement_count (refinementIter_step_le K r n) (hne n (by omega))
127       omega
128   have hn := hgrow (N+1) le_rfl
129   have hbound := blockCount_le (refinementIter K r (N+1))
130   dsimp [N] at *
131   omega
132 
133 noncomputable def stableRefinement (K : U → I → I → ℝ) (r : Setoid I) : Setoid I :=
134   refinementIter K r (Fintype.card I - blockCount r)
135 
136 theorem stableRefinement_fixed (K : U → I → I → ℝ) (r : Setoid I) :
137     transitionSplit K (stableRefinement K r) = stableRefinement K r := by
138   obtain ⟨n, hn, he⟩ := refinement_stabilizes K r
139   have hh := refinementIter_stays_fixed K r n he (Fintype.card I - blockCount r - n)
140   have hnadd : n + (Fintype.card I - blockCount r - n) = Fintype.card I - blockCount r := by omega
141   rw [hnadd] at hh
142   unfold stableRefinement
143   rw [hh]
144   exact he
145 
146 theorem stableRefinement_le (K : U → I → I → ℝ) (r : Setoid I) :
147     stableRefinement K r ≤ r := refinementIter_le_initial _ _ _
148 
149 theorem stableRefinement_lumpable (K : U → I → I → ℝ) (r : Setoid I) :
150     Lumpable K (stableRefinement K r) :=
151   (transitionSplit_fixed_iff _ _).mp (stableRefinement_fixed K r)
152 
153 theorem stableRefinement_coarsest (K : U → I → I → ℝ) (r t : Setoid I)
154     (htr : t ≤ r) (ht : Lumpable K t) : t ≤ stableRefinement K r := by
155   have hfix := (transitionSplit_fixed_iff K t).mpr ht
156   have hi : ∀ n, t ≤ refinementIter K r n := by
157     intro n
158     induction n with
159     | zero => exact htr
160     | succ n ih =>
161       change t ≤ transitionSplit K (refinementIter K r n)
162       calc t = transitionSplit K t := hfix.symm
163            _ ≤ transitionSplit K (refinementIter K r n) := transitionSplit_mono K ih
164   exact hi _
165 
166 omit [Fintype I] in
167 def likelihoodSplit (r : Setoid I) (l : Y → I → ℝ) : Setoid I where
168   r i j := r i j ∧ ∀ y, l y i = l y j
169   iseqv := {
170     refl := fun _ => ⟨r.refl _, fun _ => rfl⟩
171     symm := fun h => ⟨r.symm h.1, fun y => (h.2 y).symm⟩
172     trans := fun h h' => ⟨r.trans h.1 h'.1, fun y => (h.2 y).trans (h'.2 y)⟩ }
173 
174 omit [Fintype I] in
175 @[simp] theorem likelihoodSplit_rel (r : Setoid I) (l : Y → I → ℝ) (i j : I) :
176     likelihoodSplit r l i j ↔ r i j ∧ ∀ y, l y i = l y j := Iff.rfl
177 
178 noncomputable def jointRefinement (K : U → I → I → ℝ) (l : Y → I → ℝ)
179     (r : Setoid I) : Setoid I := stableRefinement K (likelihoodSplit r l)
180 
181 theorem jointRefinement_refines (K : U → I → I → ℝ) (l : Y → I → ℝ) (r : Setoid I) :
182     jointRefinement K l r ≤ r := by
183   intro i j hij
184   exact (stableRefinement_le K (likelihoodSplit r l) hij).1
185 
186 theorem jointRefinement_likelihood (K : U → I → I → ℝ) (l : Y → I → ℝ) (r : Setoid I) :
187     LikelihoodConstant l (jointRefinement K l r) := by
188   intro y i j hij
189   exact (stableRefinement_le K (likelihoodSplit r l) hij).2 y
190 
191 theorem jointRefinement_lumpable (K : U → I → I → ℝ) (l : Y → I → ℝ) (r : Setoid I) :
192     Lumpable K (jointRefinement K l r) := stableRefinement_lumpable _ _
193 
194 theorem jointRefinement_coarsest (K : U → I → I → ℝ) (l : Y → I → ℝ)
195     (r t : Setoid I) (htr : t ≤ r) (hk : Lumpable K t) (hl : LikelihoodConstant l t) :
196     t ≤ jointRefinement K l r := by
197   apply stableRefinement_coarsest K (likelihoodSplit r l) t
198   · intro i j hij; exact ⟨htr hij, fun y => hl y i j hij⟩
199   · exact hk
200 
201 theorem jointRefinement_termination (K : U → I → I → ℝ) (l : Y → I → ℝ) (r : Setoid I) :
202     ∃ n ≤ Fintype.card I - blockCount (likelihoodSplit r l),
203       refinementIter K (likelihoodSplit r l) (n+1) = refinementIter K (likelihoodSplit r l) n :=
204   refinement_stabilizes _ _
205 
206 end Komorebi.LSM.FinitePartitions

Back to source controls