test(Probability): add concrete instances for SideInformation and Embeddability (#949)

ComplexityTests/Embeddability.lean

@@ -5,11 +5,15 @@ Authors: Samuel Schlesinger
 -/
 
 import Complexity.Probability.CorrelatedSampling.Embeddability
+import Complexity.Probability.PMF.Uniform
+import Complexity.Probability.StatisticalDistance
 
 /-! # Local embeddings
 
 Independent private samples embed exactly, and a common value known to both parties is embedded
-by Lemma 8.
+by Lemma 8. A perfectly correlated pair embeds with error `0`, whereas a target requiring one
+party to reproduce the other party's uniform input without communication is not locally embeddable
+below error `1 / 2`.
 -/
 
 namespace ComplexityTests.Embeddability
@@ -23,11 +27,61 @@ example (P₀ : PMF (Bool × Bool)) (KS : Bool → PMF Bool) (KT : Bool → PMF
       (P₀.bind fun xy => (KS xy.1).bind fun s => (KT xy.2).map fun t => (xy, (s, t))) 0 :=
   isEmbeddable_bind_prod P₀ KS KT
 
--- Lemma 8 applies to every joint law.
-example (J : PMF ((Bool × Bool) × Bool)) :
-    IsEmbeddable (J.map Prod.fst) (J.map fun z => (z.1, (z.2, z.2)))
-      (2 * dist J (withKernel (J.map Prod.fst) fun xy => condX J xy.1) +
-        2 * dist J (withKernel (J.map Prod.fst) fun xy => condY J xy.2)) :=
-  isEmbeddable_diag J
+/-- A perfectly correlated triple of bits `((b, b), b)`. -/
+noncomputable def correlatedTriple : PMF ((Bool × Bool) × Bool) :=
+  (Cslib.Probability.PMF.uniformOfFintype Bool).map fun b => ((b, b), b)
+
+-- Lemma 8 instantiated on `correlatedTriple`.
+example :
+    IsEmbeddable (correlatedTriple.map Prod.fst)
+      (correlatedTriple.map fun z => (z.1, (z.2, z.2)))
+      (2 * dist correlatedTriple
+          (withKernel (correlatedTriple.map Prod.fst) fun xy => condX correlatedTriple xy.1) +
+        2 * dist correlatedTriple
+          (withKernel (correlatedTriple.map Prod.fst) fun xy => condY correlatedTriple xy.2)) :=
+  isEmbeddable_diag correlatedTriple
+
+-- A perfectly correlated pair embeds its marginal in `((b, b), (b, b))` with error `0`.
+example :
+    IsEmbeddable (correlatedTriple.map Prod.fst)
+      ((correlatedTriple.map Prod.fst).bind fun xy =>
+        ((PMF.pure xy.1).map fun s => (s, s)).map (xy, ·)) 0 := by
+  have h := isEmbeddable_withKernel (correlatedTriple.map Prod.fst) PMF.pure PMF.pure
+    (fun xy => PMF.pure xy.1)
+  have hYX : withKernel (correlatedTriple.map Prod.fst) (fun xy => PMF.pure xy.2) =
+      withKernel (correlatedTriple.map Prod.fst) (fun xy => PMF.pure xy.1) := by
+    simp only [correlatedTriple, withKernel, PMF.map_comp, PMF.bind_map]
+    rfl
+  simpa [hYX] using h
+
+/-- A uniform input bit for Alice and a trivial input for Bob. -/
+noncomputable def aliceInput : PMF (Bool × Unit) :=
+  Cslib.Probability.PMF.uniformOfFintype (Bool × Unit)
+
+/-- Target law where Bob's output bit equals Alice's input bit. -/
+noncomputable def bobGuessesAlice : PMF ((Bool × Unit) × (Unit × Bool)) :=
+  aliceInput.bind fun xy => (PMF.pure ((), xy.1)).map (xy, ·)
+
+-- Without communication, Bob cannot match a uniform bit held only by Alice below error `1 / 2`.
+example {ε : ℝ} (hε : ε < 1 / 2) : ¬ IsEmbeddable aliceInput bobGuessesAlice ε := by
+  rintro ⟨R, ρ, fA, fB, h⟩
+  set q : PMF (Unit × Bool) := ρ.map fun r => ((), fB () r)
+  have hmap : (aliceInput.bind fun xy => ρ.map fun r => (xy, (fA xy.1 r, fB xy.2 r))) =
+      aliceInput.bind fun xy => q.map (xy, ·) := by
+    congr 1
+    funext xy
+    exact (PMF.map_comp _ _ _).symm
+  rw [hmap, bobGuessesAlice, Cslib.Probability.PMF.dist_bind_pair] at h
+  have hsum : (∑ a : Bool × Unit, (aliceInput a).toReal * dist q (PMF.pure ((), a.1))) =
+      2⁻¹ * (dist q (PMF.pure ((), true)) + dist q (PMF.pure ((), false))) := by
+    rw [Fintype.sum_prod_type, Fintype.sum_bool]
+    simp [aliceInput, Cslib.Probability.PMF.uniformOfFintype_apply]
+    ring
+  have htri : dist (PMF.pure ((), true)) (PMF.pure ((), false)) ≤
+      dist q (PMF.pure ((), true)) + dist q (PMF.pure ((), false)) :=
+    dist_triangle_left _ _ _
+  have hpure : dist (PMF.pure ((), true)) (PMF.pure ((), false)) = 1 := by
+    simp [Cslib.Probability.PMF.dist_eq, Fintype.sum_prod_type]
+  linarith
 
 end ComplexityTests.Embeddability

ComplexityTests/SideInformation.lean

@@ -10,7 +10,7 @@ import Complexity.Probability.ProductConditioning.SideInformation
 
 Without a conditioning variable or side information, Corollary 6 specializes to Raz's
 conditioning lemma; and the side-information form applies to any joint distribution with a product
-first marginal.
+first marginal, including two fair bits conditioned on the first bit being `true`.
 -/
 
 namespace ComplexityTests.SideInformation
@@ -31,4 +31,34 @@ example {ι : Type} [Fintype ι] {β : ι → Type} [∀ i, Fintype (β i)] (p :
   sum_dist_filter_sideJoint_le (PMF.pure ()) (fun i _ => p i) (fun _ => (pi p).map (·, ()))
     (fun _ => by rw [PMF.map_comp]; exact PMF.map_id _) W hW
 
+/-- Two fair bits. -/
+noncomputable def bits : Fin 2 → PMF Bool := fun _ => uniformOfFintype Bool
+
+/-- Two fair bits with trivial conditioning variable and side information. -/
+noncomputable def bitsJoint : PMF ((Unit × Unit) × (Fin 2 → Bool)) :=
+  sideJoint (PMF.pure ()) fun _ => (pi bits).map (·, ())
+
+theorem first_true_mem_bitsJoint :
+    ∃ x ∈ {x : (Unit × Unit) × (Fin 2 → Bool) | x.2 0 = true}, x ∈ bitsJoint.support := by
+  refine ⟨(((), ()), fun _ => true), rfl, ?_⟩
+  rw [PMF.mem_support_iff, bitsJoint, sideJoint_apply]
+  refine mul_ne_zero (by simp) ?_
+  rw [← PMF.mem_support_iff, PMF.support_map]
+  refine ⟨fun _ => true, ?_, rfl⟩
+  rw [mem_support_pi_iff]
+  intro i
+  simp [bits]
+
+-- Conditioning two fair bits on the first bit being `true` under `sideJoint`.
+example :
+    ∑ j, dist ((bitsJoint.filter {x | x.2 0 = true} first_true_mem_bitsJoint).map
+          fun z => (z.1, z.2 j))
+        (((bitsJoint.filter {x | x.2 0 = true} first_true_mem_bitsJoint).map Prod.fst).bind
+          fun s => (bits j).map (s, ·)) ≤
+      Real.sqrt (2 * Real.log (1 / (bitsJoint.toOuterMeasure {x | x.2 0 = true}).toReal) / 2) := by
+  simpa [bitsJoint] using
+    sum_dist_filter_sideJoint_le (PMF.pure ()) (fun i _ => bits i)
+      (fun _ => (pi bits).map (·, ())) (fun _ => by rw [PMF.map_comp]; exact PMF.map_id _)
+      {x | x.2 0 = true} first_true_mem_bitsJoint
+
 end ComplexityTests.SideInformation