AutoSamplingTheory.ExampleCases.ProximalBPS.GaussianAugmentation

- The law of independent `X ~ μ`, `Z ~ stdGaussian E` and `Y = X + sqrt η • Z` has the displayed joint density relative to `μ.prod volume`. The input law may be singular with respect to volume.

theorem augmentation_eq_withDensity (μ : Measure E) [IsProbabilityMeasure μ]
    (η : ℝ) (hη : 0 < η) :
    Measure.map (fun p : E × E => (p.1, p.1 + Real.sqrt η • p.2))
        (μ.prod (stdGaussian E)) =
      (μ.prod (volume : Measure E)).withDensity (fun p =>
        ENNReal.ofReal
          (((Real.sqrt (2 * Real.pi * η))⁻¹) ^ Module.finrank ℝ E *
            Real.exp (-‖p.2 - p.1‖ ^ 2 / (2 * η)))) := by
  let σ : E → E := fun z => Real.sqrt η • z
  let q : E → ℝ≥0∞ := fun z => ENNReal.ofReal
    (((Real.sqrt (2 * Real.pi * η))⁻¹) ^ Module.finrank ℝ E *
      Real.exp (-‖z‖ ^ 2 / (2 * η)))
  let S : E × E ≃ᵐ E × E := MeasurableEquiv.shearAddRight E
  have hq : Measurable q := by fun_prop
  have hnoise : (stdGaussian E).map σ = (volume : Measure E).withDensity q :=
    TechnicalLemmas.Measure.IsotropicGaussianDensity.map_sqrt_smul_stdGaussian_eq_withDensity
      η hη
  have hprod : Measure.map (Prod.map id σ) (μ.prod (stdGaussian E)) =
      μ.prod ((volume : Measure E).withDensity q) := by
    rw [← Measure.map_prod_map μ (stdGaussian E) measurable_id (by fun_prop),
      hnoise, Measure.map_id]
  have hΦ : (fun p : E × E => (p.1, p.1 + Real.sqrt η • p.2)) =
      S ∘ Prod.map id σ := rfl
  have hvol : (μ.prod (volume : Measure E)).map S = μ.prod volume :=
    (measurePreserving_prod_add μ volume).map_eq
  rw [hΦ, ← Measure.map_map S.measurable (by fun_prop), hprod,
    prod_withDensity_right hq,
    TechnicalLemmas.Measure.RadonNikodym.measurableEquiv_map_withDensity
      S _ (f := fun p : E × E => q p.2) (by fun_prop), hvol]
  congr 1
  funext p
  change q (-p.1 + p.2) = q (p.2 - p.1)
  rw [neg_add_eq_sub]

end AutoSamplingTheory.ExampleCases.ProximalBPS.GaussianAugmentation

AutoSamplingTheory/ExampleCases/ProximalBPS/GaussianAugmentation.lean:25published source at 0e31a3cda412