AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep

No declaration docstring.

private theorem measurable_fiber_kl {A B : Type u} [MeasurableSpace A] [MeasurableSpace B] [MeasurableSpace.CountableOrCountablyGenerated A B] (L K : Kernel A B) [IsMarkovKernel L] [IsMarkovKernel K] :
    Measurable (fun x => klDiv (L x) (K x)) := by
  classical
  have hf : Measurable (fun p : A × B =>
      ENNReal.ofReal (klFun ((Kernel.rnDeriv L K p.1 p.2).toReal))) := by
    exact ENNReal.measurable_ofReal.comp
      (continuous_klFun.measurable.comp (Kernel.measurable_rnDeriv L K).ennreal_toReal)
  have heq (x : A) : klDiv (L x) (K x) =
      if L x ≪ K x then ∫⁻ y, ENNReal.ofReal (klFun ((Kernel.rnDeriv L K x y).toReal)) ∂K x
      else ∞ := by
    rw [klDiv_eq_lintegral_klFun]
    split_ifs
    · apply lintegral_congr_ae
      filter_upwards [Kernel.rnDeriv_eq_rnDeriv_measure (κ := L) (η := K) (a := x)] with y hy
      rw [hy]
    · rfl
  simp_rw [heq]
  exact Measurable.ite (Kernel.measurableSet_absolutelyContinuous L K)
    hf.lintegral_kernel_prod_right' measurable_const

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:52published source at 0e31a3cda412

No declaration docstring.

private theorem conditional_kl_integral {A B : Type u} [MeasurableSpace A] [MeasurableSpace B] [MeasurableSpace.CountableOrCountablyGenerated A B] (μ : Measure A) [IsProbabilityMeasure μ]
    (L K : Kernel A B) [IsMarkovKernel L] [IsMarkovKernel K] :
    klDiv (μ ⊗ₘ L) (μ ⊗ₘ K) = ∫⁻ x, klDiv (L x) (K x) ∂μ := by
  classical
  by_cases hac : μ ⊗ₘ L ≪ μ ⊗ₘ K
  · have hfiber := hac.kernel_of_compProd
    have hwd : (μ ⊗ₘ K).withDensity (fun p => Kernel.rnDeriv L K p.1 p.2) = μ ⊗ₘ L := by
      rw [← Measure.compProd_withDensity (Kernel.measurable_rnDeriv L K)]
      apply Measure.compProd_congr
      filter_upwards [hfiber] with x hx
      exact Kernel.withDensity_rnDeriv_eq hx
    have hrn : (μ ⊗ₘ L).rnDeriv (μ ⊗ₘ K) =ᵐ[μ ⊗ₘ K]
        (fun p => Kernel.rnDeriv L K p.1 p.2) := by
      rw [← hwd]
      exact Measure.rnDeriv_withDensity _ (Kernel.measurable_rnDeriv L K)
    rw [klDiv_eq_lintegral_klFun_of_ac hac]
    calc
      _ = ∫⁻ p, ENNReal.ofReal (klFun ((Kernel.rnDeriv L K p.1 p.2).toReal)) ∂(μ ⊗ₘ K) := by
        apply lintegral_congr_ae
        filter_upwards [hrn] with p hp
        rw [hp]
      _ = ∫⁻ x, ∫⁻ y, ENNReal.ofReal (klFun ((Kernel.rnDeriv L K x y).toReal)) ∂K x ∂μ := by
        apply Measure.lintegral_compProd
        exact ENNReal.measurable_ofReal.comp
          (continuous_klFun.measurable.comp (Kernel.measurable_rnDeriv L K).ennreal_toReal)
      _ = _ := by
        apply lintegral_congr_ae
        filter_upwards [hfiber] with x hx
        rw [klDiv_eq_lintegral_klFun_of_ac hx]
        apply lintegral_congr_ae
        filter_upwards [Kernel.rnDeriv_eq_rnDeriv_measure (κ := L) (η := K) (a := x)] with y hy
        rw [hy]
  · rw [klDiv_of_not_ac hac]
    symm
    by_contra hfinite
    have hf := ae_lt_top (measurable_fiber_kl L K) hfinite
    apply hac
    apply Measure.AbsolutelyContinuous.compProd_right
    filter_upwards [hf] with x hx
    exact (klDiv_ne_top_iff.mp hx.ne).1

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:71published source at 0e31a3cda412

No declaration docstring.

private theorem composed_kl_bound {A B : Type u} [MeasurableSpace A] [MeasurableSpace B] [MeasurableSpace.CountableOrCountablyGenerated A B] (P Q : Measure A) [IsProbabilityMeasure P] [IsProbabilityMeasure Q]
    (L K : Kernel A B) [IsMarkovKernel L] [IsMarkovKernel K] :
    klDiv (L ∘ₘ P) (K ∘ₘ Q) ≤ klDiv P Q + ∫⁻ x, klDiv (L x) (K x) ∂P := by
  calc
    klDiv (L ∘ₘ P) (K ∘ₘ Q) = klDiv (P ⊗ₘ L).snd (Q ⊗ₘ K).snd := by
      rw [Measure.snd_compProd, Measure.snd_compProd]
    _ ≤ klDiv (P ⊗ₘ L) (Q ⊗ₘ K) := klDiv_map_le _ _ measurable_snd
    _ = _ := by rw [klDiv_compProd_eq_add, conditional_kl_integral]

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:111published source at 0e31a3cda412

No declaration docstring.

private theorem mixture_kl_bound {A B : Type u} [MeasurableSpace A] [MeasurableSpace B]
    [MeasurableSpace.CountableOrCountablyGenerated A B]
    (μ : Measure A) [IsProbabilityMeasure μ] (ν : Measure B) [IsProbabilityMeasure ν]
    (R T : Kernel A B) [IsMarkovKernel R] [IsMarkovKernel T]
    (hT : ∀ᵐ x ∂μ, T x=ν) :
    klDiv (R ∘ₘ μ) ν ≤ ∫⁻ x, klDiv (R x) (T x) ∂μ := by
  let K : Kernel A B := Kernel.const A ν
  have hc := composed_kl_bound μ μ R K
  have hconst : K ∘ₘ μ=ν := by
    change μ.bind (fun _ => ν)=ν
    rw [Measure.bind_const,measure_univ,one_smul]
  rw [hconst,klDiv_self,zero_add] at hc
  refine hc.trans_eq ?_
  apply lintegral_congr_ae
  filter_upwards [hT] with x hx
  simp only [K,Kernel.const_apply,hx]

end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep


namespace AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep
open MeasureTheory ProbabilityTheory InformationTheory
open scoped ENNReal NNReal
variable {E S : Type u} [MeasurableSpace E] [MeasurableSpace S]
  [StandardBorelSpace S] [Nonempty S]
  [MeasurableSpace.CountableOrCountablyGenerated S E]
  [MeasurableSpace.CountableOrCountablyGenerated E E]

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:120published source at 0e31a3cda412

No declaration docstring.

private theorem conditional_step (J : Kernel S (E × S)) [IsMarkovKernel J]
    (P : Kernel S S) [IsMarkovKernel P]
    (T H : Kernel S E) [IsMarkovKernel T] [IsMarkovKernel H]
    (K : Kernel (S × E) E) [IsMarkovKernel K]
    (R : Kernel S E) [IsMarkovKernel R] (s : S)
    (hP : P s=J.snd s)
    (hrecover : (K.comap (Prod.mk s) measurable_prodMk_left) ∘ₘ H s = T s)
    (halign : ∀ᵐ y ∂J.fst s, ∀ᵐ t ∂J.condKernel (s,y), T t=K (s,y)) :
    klDiv ((R ∘ₖ P) s) (T s) ≤
      klDiv (J.fst s) (H s) + ∫⁻ t, klDiv (R t) (T t) ∂P s := by
  let C := J.condKernel
  let LR := (R ∘ₖ C).comap (Prod.mk s) measurable_prodMk_left
  let KR := K.comap (Prod.mk s) measurable_prodMk_left
  have hdecomp (f : S → ℝ≥0∞) (hf : Measurable f) :
      (∫⁻ t, f t ∂P s) = ∫⁻ y, ∫⁻ t, f t ∂C (s,y) ∂J.fst s := by
    rw [hP,Kernel.snd_apply,lintegral_map hf measurable_snd]
    conv_lhs => rw [← Kernel.disintegrate J J.condKernel]
    exact Kernel.lintegral_compProd J.fst J.condKernel s (hf.comp measurable_snd)
  have hout : LR ∘ₘ J.fst s = (R ∘ₖ P) s := by
    ext A hA
    rw [Measure.bind_apply hA LR.aemeasurable,Kernel.comp_apply' R P s hA]
    rw [hdecomp (fun t => R t A) (R.measurable_coe hA)]
    apply lintegral_congr
    intro y
    change (R ∘ₖ C) (s,y) A = _
    exact Kernel.comp_apply' R C (s,y) hA
  have hc := composed_kl_bound (J.fst s) (H s) LR KR
  rw [hout,hrecover] at hc
  have hm : ∀ᵐ y ∂J.fst s, klDiv (LR y) (KR y) ≤
      ∫⁻ t, klDiv (R t) (T t) ∂C (s,y) := by
    filter_upwards [halign] with y hy
    have hh := mixture_kl_bound (C (s,y)) (K (s,y)) R T hy
    exact hh
  calc
    klDiv ((R ∘ₖ P) s) (T s) ≤ klDiv (J.fst s) (H s) +
        ∫⁻ y, klDiv (LR y) (KR y) ∂J.fst s := hc
    _ ≤ klDiv (J.fst s) (H s) +
        ∫⁻ y, ∫⁻ t, klDiv (R t) (T t) ∂C (s,y) ∂J.fst s :=
      add_le_add le_rfl (lintegral_mono_ae hm)
    _ = _ := congrArg (fun z => klDiv (J.fst s) (H s)+z)
      (hdecomp (fun t => klDiv (R t) (T t)) (measurable_fiber_kl R T)).symm

end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep
-- Source excerpt truncated; follow the exact source link.

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:148published source at 0e31a3cda412

Excerpt truncated; the exact source link is authoritative.

No declaration docstring.

private theorem actual_gibbs_posterior {V : E → ℝ} {α β : ℝ≥0}
    (hα : 0 < α) (hαβ : α ≤ β) (hV : ContDiff ℝ 2 V)
    (hHess : ∀ x w : E, (α:ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ V) x w w ∧
      fderiv ℝ (fderiv ℝ V) x w w ≤ (β:ℝ)*‖w‖^2)
    (b a : S → ℝ) (u : S → E) (hb : Measurable b) (ha : Measurable a)
    (hu : Measurable u) (hb0 : ∀ s, 0 ≤ b s) (ha0 : ∀ s, 0 < a s) :
    let μ := (volume : Measure E).tilted (fun x => -V x)
    IsProbabilityMeasure μ ∧
    ∃ (T H : Kernel S E) (K : Kernel (S × E) E),
      IsMarkovKernel T ∧ IsMarkovKernel H ∧ IsMarkovKernel K ∧
      (∀ s, T s=μ.tilted (fun x => -(b s/2)*‖x-u s‖^2)) ∧
      (∀ s, H s=AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing
        (T s) (Real.sqrt (a s))) ∧
      (∀ s y, K (s,y)=μ.tilted (fun x => -((b s+(a s)⁻¹)/2)*
        ‖x-(b s+(a s)⁻¹)⁻¹ • (b s • u s+(a s)⁻¹ • y)‖^2)) ∧
      ∀ s, (K.comap (Prod.mk s) measurable_prodMk_left) ∘ₘ H s = T s := by
  intro μ
  have hμ : IsProbabilityMeasure μ :=
    (AutoSamplingTheory.TechnicalLemmas.Analysis.GibbsGradientMoment.gibbs_gradient_moment
      hα hαβ hV hHess).1
  let := hμ
  obtain ⟨T,H,B,hT,hH,hB,hTf,hHf,hBpost,hBup,hrec⟩ :=
    AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.StateDependentRGO.state_dependent_recovery
      μ b a u hb ha hu hb0 ha0
  let := hT
  let := hH
  let := hB
  let K := B.snd
  have hK : IsMarkovKernel K := by dsimp [K]; infer_instance
  let := hK
  have hKf (s : S) (y : E) : K (s,y)=μ.tilted (fun x => -((b s+(a s)⁻¹)/2)*
      ‖x-(b s+(a s)⁻¹)⁻¹ • (b s • u s+(a s)⁻¹ • y)‖^2) := by
    dsimp only [K]
    rw [Kernel.snd_apply,hBup,Measure.map_map measurable_snd measurable_prodMk_left]
    change Measure.map id _ = _
    exact Measure.map_id
  refine ⟨hμ,T,H,K,hT,hH,hK,hTf,hHf,hKf,?_⟩
  intro s
  have he := hrec (Measure.dirac s) inferInstance
  ext A hA
  have hset : MeasurableSet (Prod.snd ⁻¹' A : Set (S × E)) := measurable_snd hA
  have hei := congrArg (fun m : Measure (S × E) => m (Prod.snd ⁻¹' A)) he
  have hm : Measurable (fun r : S => ∫⁻ y, B (r,y) (Prod.snd ⁻¹' A) ∂H r) :=
    (B.measurable_coe hset).lintegral_kernel_prod_right'
-- Source excerpt truncated; follow the exact source link.

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:200published source at 0e31a3cda412

Excerpt truncated; the exact source link is authoritative.

No declaration docstring.

private def firstIndex (q : ℕ → S → ℝ) (b : S → ℝ) (s : S) : ℕ := by
  classical
  exact if h : ∃ n, q n s ≤ b s then Nat.find h else 0

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:260published source at 0e31a3cda412

No declaration docstring.

theorem one_step_kl_error {V : E → ℝ} {α β : ℝ≥0}
    (hα : 0 < α) (hαβ : α ≤ β) (hV : ContDiff ℝ 2 V)
    (hHess : ∀ x w : E, (α:ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ V) x w w ∧
      fderiv ℝ (fderiv ℝ V) x w w ≤ (β:ℝ)*‖w‖^2)
    (hd : 0 < (Module.finrank ℝ E : ℝ))
    (η τ : RefState E → ℝ) (hη : Measurable η) (hτ : Measurable τ)
    (hη0 : ∀ s, 0 < η s) (hτ0 : ∀ s, 0 < τ s)
    (M : Kernel (RefState E) E) [IsMarkovKernel M] (threshold : ℝ≥0) :
    let d : ℝ := Module.finrank ℝ E
    let v := fun s : RefState E => (η s+τ s)/((β:ℝ)+s.1)
    let bp := fun s : RefState E => s.1+Real.toNNReal (v s)⁻¹
    let obs := fun p : RefState E × (E × E) => p.2.1+Real.sqrt (τ p.1/((β:ℝ)+p.1.1)) • p.2.2
    let center := fun p : RefState E × (E × E) => (bp p.1:ℝ)⁻¹ •
      ((p.1.1:ℝ) • p.1.2.1+(v p.1)⁻¹ • obs p)
    let F := fun p x => V x+(bp p.1:ℝ)/2*‖x-center p‖^2
    let T := fun p x => x-((β:ℝ)+bp p.1)⁻¹ • gradient (F p) x
    let q := fun n p => ‖gradient (F p) ((T p)^[n] p.2.1)‖^2
    let N := firstIndex q (fun p => ((α:ℝ)+bp p.1)*d)
    let out := fun p => (T p)^[N p] p.2.1
    let update := fun p : RefState E × (E × E) =>
      (bp p.1,center p,out p,p.1.2.2.2.1+1,fun n => Nat.casesOn n (obs p) p.1.2.2.2.2)
    let μ := (volume : Measure E).tilted (fun x => -V x)
    IsProbabilityMeasure μ ∧
    ∃ (T H : Kernel (RefState E) E) (P : Kernel (RefState E) (RefState E)),
      ∃ hT : IsMarkovKernel T, ∃ hH : IsMarkovKernel H, ∃ hP : IsMarkovKernel P,
      letI := hT
      letI := hH
      letI := hP
      (∀ s, T s=μ.tilted (fun x => -((s.1:ℝ)/2)*‖x-s.2.1‖^2)) ∧
      (∀ s, H s=AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing
        (T s) (Real.sqrt (v s))) ∧
      (∀ s, P s=if threshold ≤ s.1 then Measure.dirac s else
        ((M s).prod (stdGaussian E)).map (fun z => update (s,z))) ∧
      let Q := fun s => AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing
        (M s) (Real.sqrt (τ s/((β:ℝ)+s.1)))
      let e := fun s => if threshold ≤ s.1 then 0 else klDiv (Q s) (H s)
      Measurable e ∧ ∀ (R : Kernel (RefState E) E) [IsMarkovKernel R],
        Measurable (fun s => klDiv (R s) (T s)) ∧ ∀ s,
        klDiv ((R ∘ₖ P) s) (T s) ≤ e s+∫⁻ t, klDiv (R t) (T t) ∂P s := by
  classical
  intro d v bp obs center F Tstep q N out update μ
  have hβ : 0 < (β:ℝ) := lt_of_lt_of_le (show 0 < (α:ℝ) from hα)
    (show (α:ℝ) ≤ β from hαβ)
  have hv (s : RefState E) : 0 < v s :=
-- Source excerpt truncated; follow the exact source link.

AutoSamplingTheory/ExampleCases/SmoothedPicardHMC/EnhancedKLOneStep.lean:266published source at 0e31a3cda412

Excerpt truncated; the exact source link is authoritative.