diff --git a/README.md b/README.md index ac648f8..28ee142 100644 --- a/README.md +++ b/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 | diff --git a/THEORY.md b/THEORY.md index b717af8..448c9b1 100644 --- a/THEORY.md +++ b/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 diff --git a/src/Probability.lean b/src/Probability.lean index 9e99c2c..9447640 100644 --- a/src/Probability.lean +++ b/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. -/ diff --git a/src/Probability/Distribution.lean b/src/Probability/Distribution.lean index 1c74683..ede7487 100644 --- a/src/Probability/Distribution.lean +++ b/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) : diff --git a/src/Probability/Kernel.lean b/src/Probability/Kernel.lean new file mode 100644 index 0000000..3177344 --- /dev/null +++ b/src/Probability/Kernel.lean @@ -0,0 +1,353 @@ +import Probability.Distribution +import Mathlib.LinearAlgebra.Finsupp.LinearCombination + +/-! +# Kernels + +A kernel `f : A → Distribution B` gives a law of a second outcome for every value of a first one: +a conditional probability distribution `p_{Y|X}` [1, §3.6.1, p. 65]. Its mixture over a law `μ` +of the first outcome is the law of the second (`bindK`), + +``` + (bindK μ f)(b) = ∑ₐ μ(a) · f(a)(b). +``` + +A pushforward is the point kernel of its function, an independent product samples its factors one +after the other, and kernels of probability laws mix probability laws into probability laws. + +## Main definitions + +* `Distribution.bindK μ f`: the mixture of the kernel `f` over `μ` +* `Distribution.ProbDist.map f P`: the law of `f a` for `a ← P` +* `Distribution.ProbDist.bind P f`: the law of `b ← f a` for `a ← P` + +## Main results + +* `Distribution.bindK_bindK`, `Distribution.bindK_single_one`: mixtures compose, and the point + kernel leaves a law unchanged +* `Distribution.bindK_single_comp`, `Distribution.fTransform_bindK`, + `Distribution.bindK_fTransform`: a pushforward is the point kernel of its function, and passes + through mixtures +* `Distribution.weight_bindK_of_weight`, `Distribution.bindK_nonNeg`: kernels of laws preserve the + weight and nonnegativity +* `Distribution.prod_eq_bindK`, `Distribution.pi_eq_bindK_update`: independent products as + mixtures +* `Distribution.ProbDist.map_map`: pushing forward twice is pushing forward along the composite + +## References + +1. U. Maurer. *Cryptography Foundations*. Lecture notes, ETH Zurich, Spring 2018. +-/ + +namespace Probability.Distribution + +variable {A B C : Type*} + +/-! ### Kernels -/ + +/-- **The mixture of a kernel** over a distribution: `∑ a, μ a • f a`. -/ +noncomputable def bindK (μ : Distribution A) (f : A → Distribution B) : Distribution B := + Finsupp.linearCombination ℝ f μ + +/-- Nothing mixes to nothing. -/ +@[simp] lemma bindK_zero (f : A → Distribution B) : bindK 0 f = 0 := + map_zero _ + +/-- A point mass of weight `w` mixes to `w` times its kernel law. -/ +@[simp] lemma bindK_single (a : A) (w : ℝ) (f : A → Distribution B) : + bindK (Finsupp.single a w) f = w • f a := + Finsupp.linearCombination_single _ _ _ + +/-- The mixture is additive in the first law. -/ +lemma bindK_add (μ ν : Distribution A) (f : A → Distribution B) : + bindK (μ + ν) f = bindK μ f + bindK ν f := + map_add _ _ _ + +/-- The mixture at `b`: `∑ₐ μ(a) · f(a)(b)`. -/ +lemma bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) : + bindK μ f b = μ.sum fun a w => w * f a b := by + -- A linear combination evaluated at `b` is the combination of the values at `b`. + unfold bindK + rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply] + rfl + +/-- A kernel that vanishes mixes to nothing. -/ +@[simp] lemma bindK_zero_right (μ : Distribution A) : + bindK μ (fun _ => (0 : Distribution B)) = 0 := by + -- Induction on `μ` as a sum of point masses. + induction μ using Finsupp.induction_linear with + | zero => simp + -- A sum of two laws: the mixture is additive, and `0 + 0 = 0`. + | add μ ν hμ hν => rw [bindK_add, hμ, hν, add_zero] + -- A point mass of weight `w`: it mixes to `w • 0 = 0`. + | single a w => simp + +/-- A mixture depends only on the kernel on the support. -/ +lemma bindK_congr {μ : Distribution A} {f g : A → Distribution B} + (h : ∀ a ∈ μ.support, f a = g a) : bindK μ f = bindK μ g := by + -- Both are sums over the support of `μ`, term by term equal. + unfold bindK + rw [Finsupp.linearCombination_apply, Finsupp.linearCombination_apply] + exact Finset.sum_congr rfl fun a ha => congrArg (μ a • ·) (h a ha) + +/-- **Mixtures compose**: mixing `g` over the mixture of `f` is mixing, over `μ`, the mixtures of +`g` over each `f a`. -/ +lemma bindK_bindK (μ : Distribution A) (f : A → Distribution B) (g : B → Distribution C) : + bindK (bindK μ f) g = bindK μ fun a => bindK (f a) g := by + -- Induction on `μ`: both sides are additive in `μ`, and agree on point masses. + induction μ using Finsupp.induction_linear with + | zero => simp + -- A sum of two laws: every mixture involved is additive in its first law. + | add μ ν hμ hν => rw [bindK_add, bindK_add, bindK_add, hμ, hν] + -- A point mass of weight `w` at `a`: both sides are `w • bindK (f a) g`, as the mixture is + -- linear in its first law. + | single a w => simp [bindK] + +/-- **The point kernel** leaves a distribution unchanged. -/ +lemma bindK_single_one (μ : Distribution A) : bindK μ (fun a => Finsupp.single a 1) = μ := by + -- Induction on `μ`: a point mass of weight `w` mixes to `w` times the point mass. + induction μ using Finsupp.induction_linear with + | zero => simp + -- A sum of two laws: the mixture is additive. + | add μ ν hμ hν => rw [bindK_add, hμ, hν] + -- A point mass of weight `w` at `a`: `w • single a 1 = single a w`. + | single a w => simp + +/-- **A point kernel along `φ` is the pushforward** along `φ`. -/ +lemma bindK_single_comp (μ : Distribution A) (φ : A → B) : + bindK μ (fun a => Finsupp.single (φ a) 1) = fTransform φ μ := by + -- Induction on `μ`: both sides are additive, and move a point mass from `a` to `φ a`. + induction μ using Finsupp.induction_linear with + | zero => simp [fTransform] + | add μ ν hμ hν => + -- the mixture is additive; then the induction hypotheses ... + rw [bindK_add, hμ, hν] + -- ... and the pushforward of a sum is the sum of the pushforwards + exact (Finsupp.mapDomain_add).symm + -- A point mass of weight `w` at `a`: both sides are the point mass of weight `w` at `φ a`. + | single a w => simp [fTransform] + +/-- **Pushing a mixture forward** pushes each law of the kernel forward. -/ +lemma fTransform_bindK (φ : B → C) (μ : Distribution A) (f : A → Distribution B) : + fTransform φ (bindK μ f) = bindK μ fun a => fTransform φ (f a) := by + -- Induction on `μ`: the pushforward is linear. + induction μ using Finsupp.induction_linear with + | zero => simp [fTransform] + | add μ ν hμ hν => + -- both mixtures split over `μ + ν`; the right one is folded back by the hypotheses ... + rw [bindK_add, bindK_add, ← hμ, ← hν] + -- ... leaving the additivity of the pushforward + exact Finsupp.mapDomain_add + | single a w => + -- a point mass of weight `w` at `a` mixes to `w • f a` and `w • fTransform φ (f a)` ... + rw [bindK_single, bindK_single] + -- ... and the pushforward commutes with scaling by `w` + exact Finsupp.mapDomain_smul _ _ + +/-- **Mixing over a pushforward** is mixing the composite kernel. -/ +lemma bindK_fTransform (φ : A → B) (μ : Distribution A) (f : B → Distribution C) : + bindK (fTransform φ μ) f = bindK μ (f ∘ φ) := by + -- Induction on `μ`: a point mass at `a` is moved to `φ a`, where the kernel is `f (φ a)`. + induction μ using Finsupp.induction_linear with + | zero => simp [fTransform] + | add μ ν hμ hν => + -- the pushforward and both mixtures are additive; then the induction hypotheses + rw [show fTransform φ (μ + ν) = fTransform φ μ + fTransform φ ν from Finsupp.mapDomain_add, + bindK_add, bindK_add, hμ, hν] + -- A point mass of weight `w` at `a`: both sides are `w • f (φ a)`. + | single a w => simp [fTransform] + +/-- A pushforward depends only on the function on the support. -/ +lemma fTransform_congr_support {f g : A → B} {μ : Distribution A} + (h : ∀ a ∈ μ.support, f a = g a) : fTransform f μ = fTransform g μ := by + -- Both are point kernels, which agree on the support. + rw [← bindK_single_comp, ← bindK_single_comp] + exact bindK_congr fun a ha => by rw [h a ha] + +/-- **Restricting a mixture** restricts each law of the kernel. -/ +lemma restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribution B) : + (bindK μ f).restrict P = bindK μ fun a => (f a).restrict P := by + classical + -- Induction on `μ`: restriction is linear. + induction μ using Finsupp.induction_linear with + | zero => simp [restrict, Finsupp.filter_zero] + -- A sum of two laws: the mixtures and the restriction (a filter) are all additive. + | add μ ν hμ hν => rw [bindK_add, bindK_add, restrict, Finsupp.filter_add, ← restrict, + ← restrict, hμ, hν] + -- A point mass of weight `w` at `a`: restricting `w • f a` is `w •` the restriction of `f a`. + | single a w => simp [restrict, Finsupp.filter_smul] + +open scoped Classical in +/-- A kernel that vanishes off an event is the kernel on the restriction to it. -/ +lemma bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution B) : + bindK μ (fun a => if P a then f a else 0) = bindK (μ.restrict P) f := by + -- Induction on `μ`: a point mass at `a` contributes `f a` exactly when `P a`. + induction μ using Finsupp.induction_linear with + | zero => simp [restrict, Finsupp.filter_zero] + -- A sum of two laws: the mixtures are additive, and the restriction of a sum is the sum of + -- the restrictions. + | add μ ν hμ hν => rw [bindK_add, hμ, hν, ← bindK_add, restrict, restrict, restrict, + Finsupp.filter_add] + | single a w => + -- a point mass of weight `w` at `a`: case on the event at `a` + by_cases h : P a + -- `P a`: the restriction keeps the point mass, and both sides are `w • f a` + · simp [h, restrict, Finsupp.filter_single_of_pos _ h] + -- not `P a`: the restriction drops the point mass, and both sides are zero + · simp [h, restrict, Finsupp.filter_single_of_neg _ h] + +/-- **The weight of a mixture**: each weight of the kernel, weighted by `μ`. -/ +lemma weight_bindK (μ : Distribution A) (f : A → Distribution B) : + (bindK μ f).weight = μ.sum fun a w => w * (f a).weight := by + -- Induction on `μ`: the weight is additive, and scales with a point mass. + induction μ using Finsupp.induction_linear with + | zero => simp [weight] + | add μ ν hμ hν => + -- the mixture and its weight are additive; the sum over `μ + ν` splits, as its summand + -- `w * (f a).weight` is additive in `w` + rw [bindK_add, weight_add, hμ, hν, + Finsupp.sum_add_index' (by simp) (fun _ _ _ => add_mul _ _ _)] + -- A point mass of weight `w` at `a`: the weight of `w • f a` is `w` times that of `f a`. + | single a w => simp [weight_smul] + +/-- **A kernel of weight one preserves the weight.** -/ +lemma weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B} + (hf : ∀ a ∈ μ.support, (f a).weight = 1) : (bindK μ f).weight = μ.weight := by + -- Each `μ a` is weighted by the weight of `f a`, which is one. + rw [weight_bindK, weight] + -- Term by term over the support of `μ`: `μ(a) · 1 = μ(a)`. + exact Finset.sum_congr rfl fun a ha => by simp only [hf a ha, mul_one] + +/-- **A nonnegative kernel mixed over a nonnegative distribution** is nonnegative. -/ +lemma bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.NonNeg) + (hf : ∀ a, (f a).NonNeg) : (bindK μ f).NonNeg := by + -- At each `b`, a sum of products of nonnegative numbers. + intro b + -- The mixture at `b` is `∑ₐ μ(a) · f(a)(b)`. + unfold bindK + rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply] + -- Each term is `μ(a) · f(a)(b)`, with `μ(a) ≥ 0` and `f(a)(b) ≥ 0`. + exact Finset.sum_nonneg fun a _ => by + simpa [Finsupp.smul_apply] using mul_nonneg (hμ a) (hf a b) + +/-- A possible outcome of a mixture is possible under the kernel at a possible first outcome. -/ +lemma mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b : B} + (hb : b ∈ (bindK μ f).support) : ∃ a ∈ μ.support, b ∈ (f a).support := by + -- Otherwise every term `μ a · f a b` of the mixture at `b` vanishes. + by_contra hn + push Not at hn + -- So `f(a)(b) = 0` at every possible `a`; we show the mixture at `b` is zero, against `hb`. + apply Finsupp.mem_support_iff.mp hb + -- The mixture at `b` is `∑ₐ μ(a) · f(a)(b)` over the support of `μ` ... + unfold bindK + rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply] + -- ... and every term vanishes, since `f(a)(b) = 0`. + exact Finset.sum_eq_zero fun a ha => by + simp [Finsupp.notMem_support_iff.mp (hn a ha)] + +/-! ### Independent products -/ + +/-- **An independent product as a mixture**: sample the first factor, then the second beside it. -/ +lemma prod_eq_bindK (μ : Distribution A) (ν : Distribution B) : + prod μ ν = bindK μ fun a => fTransform (fun b => (a, b)) ν := by + classical + -- Compare at a pair `(a, b)`: `μ a · ν b` on the left. + ext ⟨a, b⟩ + rw [prod_apply, bindK_apply, Finsupp.sum] + -- The law `ν` placed beside `a'` puts `ν b` at `(a, b)` if `a' = a`, and nothing otherwise. + have hpt : ∀ a', fTransform (fun b' => (a', b')) ν (a, b) = if a' = a then ν b else 0 := by + intro a' + -- The pushforward at `(a, b)` is the `ν`-mass of the event `(a', b') = (a, b)`. + rw [fTransform_apply_eq_mass] + split_ifs with h + -- `a' = a`: the event is `b' = b`, of mass `ν b` (zero when `b` is outside the support). + · subst h + rw [mass, Finsupp.sum] + simp only [Prod.mk.injEq, true_and] + rw [Finset.sum_ite_eq'] + split_ifs with hb + · rfl + · exact (Finsupp.notMem_support_iff.mp hb).symm + -- `a' ≠ a`: no pair `(a', b')` is `(a, b)`, so the event is empty. + · exact mass_eq_zero_of_forall_not _ fun b' h' => h (Prod.mk.inj h').1 + -- So the mixture at `(a, b)` keeps only the term `a' = a`: `μ a · ν b`. + simp only [hpt, mul_ite, mul_zero] + rw [Finset.sum_ite_eq'] + -- The sum keeps that term if `a` is in the support of `μ`; otherwise `μ a = 0`. + split_ifs with ha + · rfl + · rw [Finsupp.notMem_support_iff.mp ha, zero_mul] + +/-- **Disintegration of an independent product** at one coordinate: sample that coordinate first, +then the product with that coordinate fixed. -/ +lemma pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι → Distribution A) + (n : ι) : pi μ = bindK (μ n) fun z => pi (Function.update μ n (Finsupp.single z 1)) := by + classical + -- Compare at a sample `t`: `∏ᵢ μᵢ(tᵢ)` on the left. + ext t + -- With coordinate `n` fixed to `z`, the product is `[t n = z] · ∏_{i ≠ n} μᵢ(tᵢ)`. + have hrest : ∀ z : A, (∏ i, Function.update μ n (Finsupp.single z 1) i (t i)) = + (Finsupp.single z (1 : ℝ)) (t n) * ∏ i ∈ Finset.univ.erase n, μ i (t i) := by + intro z + -- Split off the factor at `n`, which the update sets to the point mass at `z` ... + rw [← Finset.mul_prod_erase Finset.univ _ (Finset.mem_univ n), Function.update_self] + congr 1 + -- ... and off `n` the update leaves the factors `μ i` unchanged. + exact Finset.prod_congr rfl fun i hi => by + rw [Function.update_of_ne (Finset.ne_of_mem_erase hi)] + -- Mixing over `z ← μ n` keeps the term `z = t n`: `μ n (t n) · ∏_{i ≠ n} μᵢ(tᵢ)`. + rw [pi_apply, bindK_apply] + simp only [pi_apply, hrest, Finsupp.single_apply] + -- Split the factor at `n` off the left side too: `μ n (t n) · ∏_{i ≠ n} μᵢ(tᵢ)`. + rw [← Finset.mul_prod_erase Finset.univ _ (Finset.mem_univ n), Finsupp.sum] + -- On the right, the term at `z` is `μ n z · ∏_{i ≠ n} μᵢ(tᵢ)` if `z = t n`, and zero otherwise. + simp only [← mul_assoc, mul_ite, mul_one, mul_zero, ite_mul, zero_mul] + rw [Finset.sum_ite_eq'] + -- The sum keeps that term if `t n` is in the support of `μ n`; otherwise `μ n (t n) = 0`. + split_ifs with h + · rfl + · rw [Finsupp.notMem_support_iff.mp h, zero_mul] + +/-- With one coordinate fixed to `z`, every possible sample has `z` there. -/ +lemma eq_of_mem_support_pi_update {ι : Type*} [Fintype ι] [DecidableEq ι] + (μ : ι → Distribution A) (n : ι) (z : A) {t : ι → A} + (ht : t ∈ (pi (Function.update μ n (Finsupp.single z 1))).support) : t n = z := by + classical + -- Otherwise the factor at `n`, the point mass at `z`, vanishes at `t n`. + rw [Finsupp.mem_support_iff, pi_apply] at ht + by_contra hne + apply ht + -- The product `∏ᵢ` vanishes: its factor at `n`, `single z 1 (t n)`, is zero as `t n ≠ z`. + refine Finset.prod_eq_zero (Finset.mem_univ n) ?_ + rw [Function.update_self, Finsupp.single_apply, if_neg (Ne.symm hne)] + +/-! ### Probability laws -/ + +namespace ProbDist + +/-- **The pushforward** of a law: the law of `f a` for `a ← P`. -/ +noncomputable def map (f : A → B) (P : ProbDist A) : ProbDist B := + ⟨fTransform f P.1, fTransform_isProbDist f P.2⟩ + +@[simp] lemma map_val (f : A → B) (P : ProbDist A) : (map f P).1 = fTransform f P.1 := rfl + +/-- **Pushing forward twice** is pushing forward along the composite. -/ +lemma map_map (f : B → C) (g : A → B) (P : ProbDist A) : + map f (map g P) = map (f ∘ g) P := + -- Both sides have the same underlying distribution: `fTransform f ∘ fTransform g`. + Subtype.ext (fTransform_fTransform _ _ _) + +/-- **Sequential composition**: the law of `b ← f a` for `a ← P`, the mixture of `f` over `P`. -/ +noncomputable def bind (P : ProbDist A) (f : A → ProbDist B) : ProbDist B := + ⟨bindK P.1 fun a => (f a).1, + -- nonnegative: a mixture of nonnegative laws over the nonnegative law `P` + bindK_nonNeg P.2.1 fun a => (f a).2.1, + -- weight one: each `f a` has weight one, so the mixture keeps the weight of `P`, which is one + (weight_bindK_of_weight fun a _ => (f a).2.2).trans P.2.2⟩ + +@[simp] lemma bind_val (P : ProbDist A) (f : A → ProbDist B) : + (bind P f).1 = bindK P.1 fun a => (f a).1 := rfl + +end ProbDist + +end Probability.Distribution