PMF Utilities

NB: This module is temporary

Everything here is a general PMF bind/pure lemma with no dependence on any domain-specific structure. It should be upstreamed to Mathlib (likely Mathlib.Probability.ProbabilityMassFunction.Monad or a new Mathlib.Probability.ProbabilityMassFunction.Prod). Once accepted upstream, this file should be deleted and its consumers should import the Mathlib module instead.

Main results

• Cslib.Probability.PMF.bind_pair_apply: the "pairing" bind at (a, b) equals p a * f a b
• Cslib.Probability.PMF.bind_pair_tsum_fst: marginalizing over the first component
• Cslib.Probability.PMF.uniformOfFintype_map_equiv: a uniform distribution is invariant under equivalence
• Cslib.Probability.PMF.posteriorDist: the posterior as a PMF
• Cslib.Probability.PMF.posteriorDist_eq_prior_of_outputIndist: if the output distribution does not depend on the input, conditioning does not change the prior

theorem Cslib.Probability.PMF.bind_pair_apply {α : Type u} {β : Type v} (p : PMF α) (f : α → PMF β) (a : α) (b : β) : (p.bind fun (a' : α) => (f a').bind fun (b' : β) => PMF.pure (a', b')) (a, b) = p a * (f a) b

Evaluating the "pairing" bind (do let a ← p; return (a, ← f a)) at (a, b) gives the product p a * f a b.

theorem Cslib.Probability.PMF.bind_pair_tsum_fst {α : Type u} {β : Type v} (p : PMF α) (f : α → PMF β) (b : β) : ∑' (a : α), (p.bind fun (a' : α) => (f a').bind fun (b' : β) => PMF.pure (a', b')) (a, b) = (p.bind f) b

Summing the pairing bind over the first component gives the marginal.

theorem Cslib.Probability.PMF.uniformOfFintype_map_equiv {α : Type u} {γ : Type v} [Fintype α] [Fintype γ] [Nonempty α] [Nonempty γ] (e : α ≃ γ) : PMF.map (⇑e) (PMF.uniformOfFintype α) = PMF.uniformOfFintype γ

A uniform distribution on a finite type is invariant under any equivalence.

noncomputable def Cslib.Probability.PMF.posteriorDist {α : Type u} {β : Type v} (p : PMF α) (f : α → PMF β) (b : β) (hb : b ∈ (p.bind f).support) : PMF α

The posterior distribution Pr[A = a | B = b] as a PMF, given a ← p, b ← f a, and that b has positive marginal probability: the joint distribution's slice at b, normalized.

Equations

• Cslib.Probability.PMF.posteriorDist p f b hb = PMF.normalize (fun (a : α) => (p.bind fun (a' : α) => (f a').bind fun (b' : β) => PMF.pure (a', b')) (a, b)) ⋯ ⋯

@[simp]

theorem Cslib.Probability.PMF.posteriorDist_apply {α : Type u} {β : Type v} (p : PMF α) (f : α → PMF β) (b : β) (hb : b ∈ (p.bind f).support) (a : α) : (posteriorDist p f b hb) a = (p.bind fun (a' : α) => (f a').bind fun (b' : β) => PMF.pure (a', b')) (a, b) / (p.bind f) b

theorem Cslib.Probability.PMF.posteriorDist_eq_prior_of_outputIndist {α : Type u} {β : Type v} (p : PMF α) (f : α → PMF β) (h : ∀ (a₀ a₁ : α), f a₀ = f a₁) (b : β) (hb : b ∈ (p.bind f).support) : posteriorDist p f b hb = p

If the output distribution of a channel does not depend on the input, then conditioning on any output with positive probability leaves the prior unchanged.