feat: integrals and conditional expectations against a conditional distribution

LMLExplore/Probability/HasCondDistrib.lean

@@ -8,16 +8,21 @@ module
 public import Mathlib.Probability.HasCondDistrib
 public import LeanMachineLearning.ForMathlib.Probability.HasCondDistrib
 public import TauCeti.Probability.Kernel.Composition.MeasureCompProd
+public import LMLExplore.Probability.Kernel.Composition.IntegralCompProd
 
 /-!
-# Conditional distributions: joint masses and composition-products
+# Conditional distributions: joint masses, integrals and composition-products
 
 If `κ` is the conditional distribution of `Y` given `X` under `P`, then the joint mass of a
 rectangle is obtained by integrating the conditional mass: for measurable sets `s` and `t`,
 `P(X ∈ s, Y ∈ t) = ∫_{X ∈ s} κ X t dP`, and more generally
 `∫_{X ∈ s} g(Y) dP = ∫_{X ∈ s} (∫ g dκ(X)) dP` and `∫ g(X, Y) dP = ∫ (∫ g(X, y) dκ(X)(y)) dP`
-for measurable `g ≥ 0`. The second coordinate of a composition-product `μ ⊗ₘ η.comap f` has
-conditional distribution `η` given `f` of the first coordinate.
+for measurable `g ≥ 0`; the last identity characterizes the conditional distribution, and its
+Bochner form holds for every integrable `g`. The conditional expectation of `g(X, Y)` given `X` is
+then `∫ g(X, y) dκ(X)(y)`, with no assumption on the space of `Y`: the kernel `κ` is given, rather
+than obtained from the disintegration theorem as Mathlib's `condDistrib`. The second coordinate of
+a composition-product `μ ⊗ₘ η.comap f` has conditional distribution `η` given `f` of the first
+coordinate, and conditional distributions under each `κ a` glue into one under `μ ⊗ₘ κ`.
 
 ## Main statements
 
@@ -27,16 +32,31 @@ conditional distribution `η` given `f` of the first coordinate.
   `∫⁻ ω in X ⁻¹' s, g (Y ω) ∂P = ∫⁻ ω in X ⁻¹' s, ∫⁻ y, g y ∂(κ (X ω)) ∂P`.
 * `ProbabilityTheory.HasCondDistrib.lintegral_comp_prodMk_eq_lintegral_lintegral`:
   `∫⁻ ω, g (X ω, Y ω) ∂P = ∫⁻ ω, ∫⁻ y, g (X ω, y) ∂(κ (X ω)) ∂P`.
+* `ProbabilityTheory.hasCondDistrib_iff_forall_lintegral_eq`: conversely, this identity for every
+  measurable `g ≥ 0` characterizes the conditional distribution.
+* `ProbabilityTheory.HasCondDistrib.integral_comp_prodMk_eq_integral_integral`: the Bochner form
+  `∫ ω, g (X ω, Y ω) ∂P = ∫ ω, ∫ y, g (X ω, y) ∂(κ (X ω)) ∂P` for integrable `g`.
+* `ProbabilityTheory.HasCondDistrib.condExp_prod_ae_eq_integral`:
+  `P[g (X, Y) | σ(X)] = ∫ y, g (X, y) ∂(κ X)` almost everywhere, for integrable `g`.
 * `ProbabilityTheory.HasCondDistrib.restrict_preimage`: `κ` is still the conditional
   distribution of `Y` given `X` under the restriction of `P` to an event `X ⁻¹' s`.
 * `ProbabilityTheory.HasCondDistrib.ae_mem`: if every conditional law is concentrated on a set,
   then `Y` lies in that set almost surely.
 * `ProbabilityTheory.hasCondDistrib_snd_compProd_comap`: under `μ ⊗ₘ η.comap f`, the second
   coordinate has conditional distribution `η` given `f` of the first coordinate.
+* `ProbabilityTheory.hasCondDistrib_compProd_of_forall_sectR`: if `Y` has conditional
+  distribution `η (a, ·)` given `X` under every `κ a`, then under `μ ⊗ₘ κ` the variable `Y` of the
+  second coordinate has conditional distribution `η` given the first coordinate and `X`.
+* `ProbabilityTheory.HasCondDistrib.of_comp_hasLaw`: a conditional distribution of `Y ∘ f` given
+  `X ∘ f` under `P` is one of `Y` given `X` under the law of `f`, the converse of LML's
+  `HasCondDistrib.comp_hasLaw`.
 
-The last statement is adapted from the LMLPapers project (same author), and
-`HasCondDistrib.restrict_preimage` from the Essakine2026Tight project
-(`RemyDegenne/Essakine2026Tight`, Apache 2.0, `ForLMLExplore/Probability/CondDistrib.lean`).
+`hasCondDistrib_snd_compProd_comap` is adapted from the LMLPapers project (same author). The
+characterization, the Bochner integral, the conditional expectation,
+`HasCondDistrib.restrict_preimage`, the gluing over `μ ⊗ₘ κ` and the transport by a law are
+adapted from the Essakine2026Tight project
+(`RemyDegenne/Essakine2026Tight`, Apache 2.0, files `HasCondDistribIntegral.lean`,
+`HasCondDistribCondExp.lean` and `CondDistrib.lean` of `ForLMLExplore/Probability/`).
 -/
 
 public section
@@ -99,6 +119,28 @@ lemma HasCondDistrib.lintegral_comp_prodMk_eq_lintegral_lintegral (h : HasCondDi
     _ = ∫⁻ ω, ∫⁻ y, g (X ω, y) ∂(κ (X ω)) ∂P :=
         lintegral_map' hg.lintegral_kernel_prod_right'.aemeasurable h.aemeasurable_fst
 
+/-- A kernel `κ` is the conditional distribution of `Y` given `X` if and only if, for every
+measurable function `g ≥ 0` of both, `∫ g(X, Y) dP = ∫ (∫ g(X, y) dκ(X)(y)) dP`. -/
+lemma hasCondDistrib_iff_forall_lintegral_eq (hX : AEMeasurable X P) (hY : AEMeasurable Y P) :
+    HasCondDistrib Y X κ P ↔ ∀ g : 𝓧 × 𝓨 → ℝ≥0∞, Measurable g →
+      ∫⁻ ω, g (X ω, Y ω) ∂P = ∫⁻ ω, ∫⁻ y, g (X ω, y) ∂(κ (X ω)) ∂P := by
+  refine ⟨fun h g hg ↦ h.lintegral_comp_prodMk_eq_lintegral_lintegral hg,
+    fun h ↦ ⟨hX.prodMk hY, Measure.ext_of_lintegral _ fun g hg ↦ ?_⟩⟩
+  rw [lintegral_map' hg.aemeasurable (hX.prodMk hY), h g hg, Measure.lintegral_compProd hg,
+    lintegral_map' hg.lintegral_kernel_prod_right'.aemeasurable hX]
+
+/-- If `κ` is the conditional distribution of `Y` given `X`, then for a function `g` of both,
+integrable for the joint law, `∫ g(X, Y) dP = ∫ (∫ g(X, y) dκ(X)(y)) dP`. -/
+lemma HasCondDistrib.integral_comp_prodMk_eq_integral_integral {E : Type*}
+    [NormedAddCommGroup E] [NormedSpace ℝ E] (h : HasCondDistrib Y X κ P) {g : 𝓧 × 𝓨 → E}
+    (hg : Integrable g (P.map X ⊗ₘ κ)) :
+    ∫ ω, g (X ω, Y ω) ∂P = ∫ ω, ∫ y, g (X ω, y) ∂(κ (X ω)) ∂P := by
+  calc ∫ ω, g (X ω, Y ω) ∂P
+      = ∫ p, g p ∂(P.map X ⊗ₘ κ) := h.integral_comp hg.aestronglyMeasurable
+    _ = ∫ x, ∫ y, g (x, y) ∂(κ x) ∂(P.map X) := Measure.integral_compProd hg
+    _ = ∫ ω, ∫ y, g (X ω, y) ∂(κ (X ω)) ∂P :=
+        integral_map h.aemeasurable_fst hg.integral_measure_compProd.aestronglyMeasurable
+
 /-- If `κ` is the conditional distribution of `Y` given `X` under `P`, it is also the conditional
 distribution of `Y` given `X` under the restriction of `P` to an event `X ⁻¹' s` determined by
 `X`. -/
@@ -147,6 +189,72 @@ lemma hasCondDistrib_snd_compProd_comap (μ : Measure α) [SFinite μ]
   rw [Measure.fst, Measure.map_map hf measurable_fst]
   rfl
 
+variable {δ : Type*} {mδ : MeasurableSpace δ}
+
+/-- If `Y` has conditional distribution `η.sectR a = η (a, ·)` given `X` under every `κ a`, then
+under `μ ⊗ₘ κ` the variable `Y` of the second coordinate has conditional distribution `η` given
+the first coordinate and `X` of the second. -/
+lemma hasCondDistrib_compProd_of_forall_sectR {μ : Measure α} [SFinite μ] {κ : Kernel α β}
+    [IsSFiniteKernel κ] {X : β → γ} {Y : β → δ} {η : Kernel (α × γ) δ} [IsSFiniteKernel η]
+    (hX : Measurable X) (hY : Measurable Y) (h : ∀ a, HasCondDistrib Y X (η.sectR a) (κ a)) :
+    HasCondDistrib (fun p ↦ Y p.2) (fun p ↦ (p.1, X p.2)) η (μ ⊗ₘ κ) := by
+  refine (hasCondDistrib_iff_forall_lintegral_eq (by fun_prop) (by fun_prop)).mpr fun g hg ↦ ?_
+  rw [Measure.lintegral_compProd (f := fun p ↦ g ((p.1, X p.2), Y p.2)) (by fun_prop),
+    Measure.lintegral_compProd (f := fun p ↦ ∫⁻ y, g ((p.1, X p.2), y) ∂η (p.1, X p.2))
+      (hg.lintegral_kernel_prod_right'.comp (by fun_prop))]
+  refine lintegral_congr fun a ↦ ?_
+  simpa using (h a).lintegral_comp_prodMk_eq_lintegral_lintegral
+    (g := fun q ↦ g ((a, q.1), q.2)) (by fun_prop)
+
 end CompProd
 
+section Transport
+
+variable {Ω Ω' 𝓧 𝓨 : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'}
+  {m𝓧 : MeasurableSpace 𝓧} {m𝓨 : MeasurableSpace 𝓨} {P : Measure Ω} {P' : Measure Ω'}
+  {f : Ω → Ω'} {X : Ω' → 𝓧} {Y : Ω' → 𝓨} {κ : Kernel 𝓧 𝓨}
+
+/-- If `Y ∘ f` has conditional distribution `κ` given `X ∘ f` under `P` and `f` has law `P'`
+under `P`, then `Y` has conditional distribution `κ` given `X` under `P'`. This is the converse
+of LML's `HasCondDistrib.comp_hasLaw`. -/
+lemma HasCondDistrib.of_comp_hasLaw (h : HasCondDistrib (Y ∘ f) (X ∘ f) κ P) (hf : HasLaw f P' P)
+    (hX : AEMeasurable X P') (hY : AEMeasurable Y P') :
+    HasCondDistrib Y X κ P' := by
+  refine ⟨hX.prodMk hY, ?_⟩
+  rw [← hf.map_eq, AEMeasurable.map_map_of_aemeasurable (hf.map_eq ▸ hX.prodMk hY) hf.aemeasurable,
+    AEMeasurable.map_map_of_aemeasurable (hf.map_eq ▸ hX) hf.aemeasurable]
+  exact h.map_eq
+
+end Transport
+
+section CondExp
+
+variable {Ω 𝓧 𝓨 E : Type*} {mΩ : MeasurableSpace Ω} {m𝓧 : MeasurableSpace 𝓧}
+  {m𝓨 : MeasurableSpace 𝓨} [NormedAddCommGroup E] [NormedSpace ℝ E] [CompleteSpace E]
+  {P : Measure Ω} [IsFiniteMeasure P] {X : Ω → 𝓧} {Y : Ω → 𝓨} {κ : Kernel 𝓧 𝓨}
+  [IsSFiniteKernel κ]
+
+/-- If `κ` is the conditional distribution of `Y` given `X`, then for a function `g` of both,
+integrable for the joint law, `𝔼[g(X, Y) | σ(X)] = ∫ g(X, y) dκ(X)(y)` almost everywhere.
+Unlike Mathlib's `condExp_prod_ae_eq_integral_condDistrib'`, this needs no standard Borel
+assumption on the space of `Y`; it generalizes LML's `HasCondDistrib.condExp_comp_ae_eq` to
+functions of `(X, Y)`, on any measurable space. -/
+lemma HasCondDistrib.condExp_prod_ae_eq_integral (h : HasCondDistrib Y X κ P) (hX : Measurable X)
+    {g : 𝓧 × 𝓨 → E} (hg : Integrable g (P.map X ⊗ₘ κ)) :
+    P[fun ω ↦ g (X ω, Y ω) | m𝓧.comap X] =ᵐ[P] fun ω ↦ ∫ y, g (X ω, y) ∂(κ (X ω)) := by
+  have hg' := hg.integral_measure_compProd
+  refine (ae_eq_condExp_of_forall_setIntegral_eq hX.comap_le (h.integrable_comp hg)
+    (fun _ _ _ ↦ (hg'.comp_aemeasurable hX.aemeasurable).integrableOn) ?_
+    (hg'.aestronglyMeasurable.comp_ae_measurable' hX.aemeasurable)).symm
+  rintro _ ⟨t, ht, rfl⟩ _
+  -- Both set integrals over `X ⁻¹' t` are integrals of `g` restricted to `t ×ˢ univ`.
+  rw [← integral_indicator (hX ht), ← integral_indicator (hX ht)]
+  convert (h.integral_comp_prodMk_eq_integral_integral (hg.indicator (ht.prod .univ))).symm
+    using 1
+  all_goals
+    refine integral_congr_ae (.of_forall fun ω ↦ ?_)
+    by_cases hω : X ω ∈ t <;> simp [hω]
+
+end CondExp
+
 end ProbabilityTheory

LMLExplore/Probability/Kernel/Composition/IntegralCompProd.lean

@@ -10,14 +10,17 @@ public import Mathlib.Probability.Kernel.Composition.IntegralCompProd
 import LMLExplore.MeasureTheory.Integral.Bochner.Basic
 
 /-!
-# Integrals of differences against a composition-product
+# Integrals against a composition-product
 
-Complements to Mathlib's `Mathlib.Probability.Kernel.Composition.IntegralCompProd`. For a Markov
-kernel `κ`, the integral against `μ` of `f a - ∫ g (a, b) dκ_a(b)` is the integral of
-`f p.1 - g p` against the composition-product `μ ⊗ₘ κ`.
+Complements to Mathlib's `Mathlib.Probability.Kernel.Composition.IntegralCompProd`. If `g` is
+integrable for the composition-product `μ ⊗ₘ κ`, then `a ↦ ∫ g (a, b) dκ_a(b)` is integrable for
+`μ`. For a Markov kernel `κ`, the integral against `μ` of `f a - ∫ g (a, b) dκ_a(b)` is the
+integral of `f p.1 - g p` against `μ ⊗ₘ κ`.
 
 ## Main statements
 
+* `MeasureTheory.Integrable.integral_measure_compProd`: if `g` is integrable for `μ ⊗ₘ κ`, then
+  `a ↦ ∫ g (a, b) dκ_a(b)` is integrable for `μ`.
 * `MeasureTheory.Measure.integral_sub_integral_eq_integral_compProd`:
   `∫ (f a - ∫ g (a, b) dκ_a(b)) dμ(a) = ∫ (f p.1 - g p) d(μ ⊗ₘ κ)(p)`.
 -/
@@ -26,11 +29,20 @@ public section
 
 open ProbabilityTheory
 
-namespace MeasureTheory.Measure
+namespace MeasureTheory
 
 variable {α β E : Type*} {mα : MeasurableSpace α} {mβ : MeasurableSpace β}
   [NormedAddCommGroup E] [NormedSpace ℝ E] {μ : Measure α} {κ : Kernel α β}
 
+/-- If `g` is integrable for the composition-product `μ ⊗ₘ κ`, then `a ↦ ∫ g (a, b) dκ_a(b)` is
+integrable for `μ`. This is the measure form of Mathlib's `Integrable.integral_compProd`. -/
+lemma Integrable.integral_measure_compProd [SFinite μ] [IsSFiniteKernel κ] {g : α × β → E}
+    (hg : Integrable g (μ ⊗ₘ κ)) : Integrable (fun a ↦ ∫ b, g (a, b) ∂(κ a)) μ := by
+  rw [Measure.compProd] at hg
+  simpa only [Kernel.prodMkLeft_apply, Kernel.const_apply] using hg.integral_compProd
+
+namespace Measure
+
 /-- For a Markov kernel `κ`, the integral against `μ` of `f a - ∫ g (a, b) dκ_a(b)` is the
 integral of `f p.1 - g p` against the composition-product `μ ⊗ₘ κ`. -/
 lemma integral_sub_integral_eq_integral_compProd [CompleteSpace E] [SFinite μ] [IsMarkovKernel κ]
@@ -44,4 +56,6 @@ lemma integral_sub_integral_eq_integral_compProd [CompleteSpace E] [SFinite μ]
   filter_upwards [((integrable_compProd_iff hg.1).1 hg).1] with a ha
   rw [integral_const_sub _ ha]
 
-end MeasureTheory.Measure
+end Measure
+
+end MeasureTheory