Probability: kernels, pushforward and sequential composition of laws

README.md

@@ -14,7 +14,7 @@ summarizes the theory layer by layer, with links to the Lean declarations.
 
 | Library | Content |
 |---|---|
-| `Probability` | finitely supported distributions, statistical distance, couplings, expectation, counting |
+| `Probability` | finitely supported distributions, kernels, statistical distance, couplings, expectation, counting |
 | `RandomSystems` | systems and automata; deterministic converters and programs; distributions over deterministic systems; random systems given by their cumulative probabilities; probabilistic converters, attachment and parallel composition; transcript distance and distinguishers; games |
 | `ConstructiveCryptography.Specification`, `.Construction` | specifications, constructions and relaxations over any carrier |
 | `ConstructiveCryptography.CryptographicAlgebra` | the classes of cryptographic algebras: resource theories, parallel composition, compatible pseudo-metrics and distinguisher classes, relaxations |

THEORY.md

@@ -842,7 +842,13 @@ R \;\simeq_{\mathrm{ptxt}}\; P R \;\simeq_{\mathrm{cca}}\; P C R \;\simeq_{\math
 A [**distribution**][Distribution] is a finitely supported function to $`\mathbb R`$, and a
 [**probability distribution**][ProbDist] one that is [nonnegative with weight one][isProbDist].
 Distributions have [masses][mass] of events, [pushforwards][fTransform], [products][prod] of two
-laws and of [families][pi], and [uniform][uniform] laws. The [**statistical distance**][statDist]
+laws and of [families][pi], and [uniform][uniform] laws. A [**kernel**][bindK] $`f`$ is a conditional
+distribution [3, §3.6.1]; mixed over a law $`\mu`$ it gives $`\sum_a \mu(a)\, f(a)`$:
+[mixtures compose][bindK_bindK], a pushforward
+[is the point kernel of its function][bindK_single_comp], and an independent product
+[disintegrates][pi_eq_bindK_update] at a coordinate. Probability laws have
+[pushforwards][ProbDist.map] and [sequential compositions][ProbDist.bind]. The
+[**statistical distance**][statDist]
 satisfies the [triangle inequality][statDist_triangle] and [data processing][statDist_fTransform_le],
 $`\delta(fX, fY) \le \delta(X, Y)`$; with [positive parts][statDist_eq_weight_posPart] it is
 $`\delta(X, Y) = \lvert (X - Y)^{+} \rvert`$, and with [overlaps][statDist_eq_weight_sub_weight_inf]
@@ -1287,10 +1293,16 @@ agreeing with a map on a finite set $`S`$: there are $`\lvert Y \rvert^{\lvert X
 [ProbDist]: src/Probability/Distribution.lean#L149
 [isProbDist]: src/Probability/Distribution.lean#L128
 [mass]: src/Probability/Distribution.lean#L217
-[fTransform]: src/Probability/Distribution.lean#L587
-[prod]: src/Probability/Distribution.lean#L942
-[pi]: src/Probability/Distribution.lean#L1583
-[uniform]: src/Probability/Distribution.lean#L521
+[fTransform]: src/Probability/Distribution.lean#L632
+[prod]: src/Probability/Distribution.lean#L1005
+[pi]: src/Probability/Distribution.lean#L1646
+[bindK]: src/Probability/Kernel.lean#L49
+[bindK_bindK]: src/Probability/Kernel.lean#L95
+[bindK_single_comp]: src/Probability/Kernel.lean#L117
+[pi_eq_bindK_update]: src/Probability/Kernel.lean#L283
+[ProbDist.map]: src/Probability/Kernel.lean#L329
+[ProbDist.bind]: src/Probability/Kernel.lean#L341
+[uniform]: src/Probability/Distribution.lean#L566
 [statDist]: src/Probability/StatisticalDistance.lean#L88
 [statDist_triangle]: src/Probability/StatisticalDistance.lean#L243
 [statDist_fTransform_le]: src/Probability/StatisticalDistance.lean#L615

src/Probability.lean

@@ -2,8 +2,10 @@ import Probability.Counting
 import Probability.Coupling
 import Probability.Distribution
 import Probability.Expectation
+import Probability.Kernel
 import Probability.Lift
 import Probability.MultiCoupling
 import Probability.StatisticalDistance
 
-/-! Finitely supported distributions, statistical distance, couplings, expectation and counting. -/
+/-! Finitely supported distributions, kernels, statistical distance, couplings, expectation and
+counting. -/

src/Probability/Distribution.lean

@@ -280,6 +280,51 @@ theorem mass_restrict {A : Type*} (X : Distribution A) (P Q : A → Prop) :
   intro a _ha
   by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
 
+/-- Restricting keeps only mass that was there. -/
+lemma support_restrict_subset {A : Type*} (μ : Distribution A) (P : A → Prop) :
+    (μ.restrict P).support ⊆ μ.support := fun a ha => by
+  -- The restriction at `a` is `μ a` or zero; it is nonzero, so it is `μ a`.
+  rw [Finsupp.mem_support_iff, restrict_apply] at ha
+  rw [Finsupp.mem_support_iff]
+  split_ifs at ha
+  · exact ha
+  · exact absurd rfl ha
+
+
+/-- Restricting twice is restricting to both events. -/
+lemma restrict_restrict {A : Type*} (μ : Distribution A) (P Q : A → Prop) :
+    (μ.restrict P).restrict Q = μ.restrict fun a => P a ∧ Q a := by
+  -- Pointwise: `μ a` exactly when both events hold.
+  ext a
+  simp only [restrict_apply]
+  by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
+
+
+/-- Restriction depends only on the event up to equivalence. -/
+lemma restrict_congr {A : Type*} (μ : Distribution A) {P Q : A → Prop} (h : ∀ a, P a ↔ Q a) :
+    μ.restrict P = μ.restrict Q := by
+  -- Pointwise, the two events are the same condition.
+  ext a
+  simp only [restrict_apply, h]
+
+
+/-- Restricting to the impossible event leaves nothing. -/
+lemma restrict_false {A : Type*} (μ : Distribution A) : μ.restrict (fun _ => False) = 0 := by
+  -- Pointwise, the condition never holds.
+  ext a
+  simp
+
+
+/-- A point mass restricted to an event: itself or nothing. -/
+lemma single_restrict {A : Type*} (a : A) (w : ℝ) (P : A → Prop) :
+    Distribution.restrict (Finsupp.single a w : Distribution A) P =
+      if P a then Finsupp.single a w else 0 := by
+  -- The only point with mass is `a`: kept if `P a`, dropped otherwise.
+  by_cases h : P a
+  · simp [restrict, Finsupp.filter_single_of_pos _ h, h]
+  · simp [restrict, Finsupp.filter_single_of_neg _ h, h]
+
+
 theorem mass_true {A : Type*} (X : Distribution A) :
     X.mass (fun _ => True) = X.weight := by
   unfold mass weight
@@ -651,6 +696,24 @@ theorem mass_fTransform {A B : Type*}
     (fun b => by by_cases hb : P b <;> simp [hb])
     (fun b m₁ m₂ => by by_cases hb : P b <;> simp [hb])]
 
+/-- The pushforward of a point mass. -/
+lemma fTransform_single {A B : Type*} (f : A → B) (a : A) (w : ℝ) :
+    fTransform f (Finsupp.single a w) = Finsupp.single (f a) w :=
+  Finsupp.mapDomain_single
+
+
+/-- Restricting a pushforward is pushing forward the restriction to the preimage. -/
+lemma restrict_fTransform {A B : Type*} (f : A → B) (μ : Distribution A) (P : B → Prop) :
+    (fTransform f μ).restrict P = fTransform f (μ.restrict fun a => P (f a)) := by
+  -- At `b`, both sides are masses of fibres of `f`.
+  ext b
+  rw [restrict_apply, fTransform_apply_eq_mass, fTransform_apply_eq_mass, mass_restrict]
+  split_ifs with h
+  · -- `P b`: on the fibre of `b`, the condition `P (f a)` holds
+    exact mass_congr _ fun a => ⟨fun ha => ⟨ha, ha ▸ h⟩, fun ha => ha.1⟩
+  · -- not `P b`: the fibre of `b` meets no point with `P (f a)`
+    exact (mass_eq_zero_of_forall_not _ fun a (ha : f a = b ∧ P (f a)) => h (ha.1 ▸ ha.2)).symm
+
 /-- Pushforward preserves total mass. -/
 theorem weight_fTransform {A B : Type*}
     (f : A → B) (X : Distribution A) :