Complete normalized probability rewriting laws

src/Probability/Kernel.lean

@@ -17,9 +17,11 @@ This represents a "conditional probability distribution" `p_{Y|X}` [1, §3.6.1,
 `bindK_bindK` composes successive draws; `bindK_single_one` removes a point kernel.
 `fTransform_bindK` and `bindK_fTransform` move deterministic maps through those draws.
 
-`ProbDist.map` and `ProbDist.bind` retain normalization. `ProbDist.map_map` composes
-initializations; `ProbDist.map_pi_apply` selects one coordinate of an independent table.
-`prodProbDist_map` combines independent initialized draws:
+`ProbDist.map` and `ProbDist.bind` retain normalization. `ProbDist.map_id` removes the
+identity map; `ProbDist.map_map` composes deterministic transformations, and `ProbDist.map_pi_apply`
+selects one coordinate of an independent table. `prodProbDist_uniform` identifies independent
+uniform coordinates with a uniform pair. `prodProbDist_map` and its one-coordinate forms
+`prodProbDist_map_left`, `prodProbDist_map_right` transport independent draws:
 
 ```text
 (f_* P) × (g_* Q) = (f × g)_* (P × Q).
@@ -365,6 +367,10 @@ noncomputable def map
 /-- The underlying distribution of `f_* P` is `fTransform f P.1`. -/
 @[simp] lemma map_val (f : A → B) (P : ProbDist A) : (map f P).1 = fTransform f P.1 := rfl
 
+/-- The identity map leaves a law unchanged: `id_* P = P`. -/
+@[simp] lemma map_id (P : ProbDist A) : P.map id = P :=
+  Subtype.ext (fTransform_id P.1)
+
 /-- Successive deterministic maps compose: `f_* (g_* P) = (f ∘ g)_* P`. -/
 lemma map_map
     -- `g` transforms the draw from `P` first, then `f` transforms that result.
@@ -497,4 +503,25 @@ lemma prodProbDist_map
       ProbDist.map (Prod.map f g) (prodProbDist P Q) :=
   Subtype.ext (fTransform_prod f g P.1 Q.1).symm
 
+/-- Independent uniform coordinates have the uniform law on their product type. -/
+lemma prodProbDist_uniform
+    -- `A`, `B` are finite nonempty alphabets, so both uniform laws are normalized.
+    (A B : Type*) [Fintype A] [Nonempty A] [Fintype B] [Nonempty B] :
+    prodProbDist 𝒰[A] 𝒰[B] = 𝒰[A × B] :=
+  Subtype.ext prod_uniform
+
+/-- Transform only the first independent draw: `(f_* P) × Q = (f × id)_*(P × Q)`. -/
+lemma prodProbDist_map_left
+    -- `f` transforms outcomes in `A` to `C`; `P` and `Q` supply the independent draws.
+    (f : A → C) (P : ProbDist A) (Q : ProbDist B) :
+    prodProbDist (P.map f) Q = (prodProbDist P Q).map (Prod.map f id) := by
+  simpa only [ProbDist.map_id] using prodProbDist_map f id P Q
+
+/-- Transform only the second independent draw: `P × (g_* Q) = (id × g)_*(P × Q)`. -/
+lemma prodProbDist_map_right
+    -- `g` transforms outcomes in `B` to `C`; `P` and `Q` supply the independent draws.
+    (g : B → C) (P : ProbDist A) (Q : ProbDist B) :
+    prodProbDist P (Q.map g) = (prodProbDist P Q).map (Prod.map id g) := by
+  simpa only [ProbDist.map_id] using prodProbDist_map id g P Q
+
 end Probability.Distribution

src/Probability/README.md

@@ -34,6 +34,10 @@ it gives $`\sum_a \mu(a)\, f(a)`$:
 [constant probability kernel][ProbDist.bind_const] returns its law.
 [Successive maps][ProbDist.map_map] compose, and [independent maps][prodProbDist_map] satisfy
 $`(f_*P) \times (g_*Q) = (f \times g)_*(P \times Q)`$.
+[The identity map][ProbDist.map_id] leaves a law unchanged, so $`g=\mathrm{id}`$ or
+$`f=\mathrm{id}`$ gives the [left][prodProbDist_map_left] and
+[right][prodProbDist_map_right] one-coordinate forms.
+[Independent uniform coordinates][prodProbDist_uniform] form a uniform pair.
 A [joint law][ProbDist.jointLaw] retains both successive draws:
 $`J_{\mu,K}(h,z)=\mu(h)K(h,z)`$, with [first marginal][ProbDist.jointLaw_fst] $`\mu`$. For
 arbitrary maps $`d_H, d_Z`$ with $`(d_H)_*\mu'=\mu`$ and $`(d_Z)_*K'(h)=K(d_H h)`$,
@@ -152,10 +156,10 @@ injectivity of all replay and public sites; neither asserts independence on bad
 [piDep]: DependentProduct.lean#L53
 [uniform]: Distribution.lean#L656
 [ProbDist.uniform]: Distribution.lean#L689
-[bindK]: Kernel.lean#L49
-[bindK_bindK]: Kernel.lean#L95
-[bindK_single_comp]: Kernel.lean#L117
-[pi_eq_bindK_update]: Kernel.lean#L344
+[bindK]: Kernel.lean#L51
+[bindK_bindK]: Kernel.lean#L97
+[bindK_single_comp]: Kernel.lean#L119
+[pi_eq_bindK_update]: Kernel.lean#L346
 [ProbDist.map]: Kernel.lean#L360
 [ProbDist.map_map]: Kernel.lean#L369
 [prodProbDist_map]: Kernel.lean#L491
@@ -198,3 +202,8 @@ injectivity of all replay and public sites; neither asserts independence on bad
 [uniform_mass_fresh_outputs_eq]: Counting.lean#L729
 
 [uniform_mass_transcript_collision_le_card]: Counting/Collision.lean
+
+[ProbDist.map_id]: Kernel.lean#L371
+[prodProbDist_uniform]: Kernel.lean
+[prodProbDist_map_left]: Kernel.lean
+[prodProbDist_map_right]: Kernel.lean

src/Tests.lean

@@ -46,6 +46,7 @@ public import Tests.DSL.Inlining
 public import Tests.DSL.ParallelOutside
 public import Tests.DSL.ProbabilityNotation
 public import Tests.DSL.Semantics
+public import Tests.NormalizedProbability
 public import Tests.ProbabilityNotation
 public import Tests.ProbabilityTablePatterns
 public import Tests.QueryRouting

src/Tests/NormalizedProbability.lean

@@ -0,0 +1,32 @@
+module
+
+public import Probability.Kernel
+
+/-!
+# Normalized probability calculations
+
+`map_uniform_coordinates` starts from one uniform pair and transforms its coordinates in
+successive steps. The result is the product of the two transformed uniform laws. Every step
+uses normalized operations and equations, so a client keeps normalization throughout.
+-/
+
+@[expose] public section
+
+namespace Tests.NormalizedProbability
+
+open Probability.Distribution
+
+/-- Successive coordinate transformations preserve the independence of uniform inputs. -/
+lemma map_uniform_coordinates
+    -- `A`, `B` are finite nonempty input alphabets; `X`, `Y` are the output alphabets.
+    {A B X Y : Type*} [Fintype A] [Nonempty A] [Fintype B] [Nonempty B]
+    -- `f` transforms the first uniform coordinate and `g` the second.
+    (f : A → X) (g : B → Y) :
+    ((𝒰[A × B]).map (Prod.map f id)).map (Prod.map id g) =
+      prodProbDist ((𝒰[A]).map f) ((𝒰[B]).map g) := by
+  -- Separate the uniform coordinates, then move each transformation to its own factor.
+  rw [← prodProbDist_uniform, ← prodProbDist_map_left, ← prodProbDist_map_right]
+
+#print axioms map_uniform_coordinates
+
+end Tests.NormalizedProbability