PMF Utilities

NB: This module is temporary

The bind/pure and posterior lemmas here have no dependence on any domain-specific structure. They should be upstreamed to Mathlib (likely Mathlib.Probability.ProbabilityMassFunction.Monad or a new Mathlib.Probability.ProbabilityMassFunction.Prod). Once accepted upstream, these lemmas should be removed and their consumers should import the Mathlib module instead.

The uniform samplers retain the PMF interface removed in mathlib#42909, pending a decision on CSLib's probability API. Their definitions and proofs are adapted from Mathlib.Probability.Distributions.Uniform before that PR, by Josha Dekker, Devon Tuma, and Kexing Ying, under the Apache 2.0 license. Original sampler copyright (c) 2024 Josha Dekker. All rights reserved.

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

noncomputable def Cslib.Probability.PMF.uniformOfFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) : PMF α

Uniform probability mass function on a nonempty finite set.

Equations

• Cslib.Probability.PMF.uniformOfFinset s hs = PMF.ofFinset (fun (a : α) => if a ∈ s then (↑s.card)⁻¹ else 0) s ⋯ ⋯

@[simp]

theorem Cslib.Probability.PMF.uniformOfFinset_apply {α : Type u} (s : Finset α) (hs : s.Nonempty) (a : α) : (uniformOfFinset s hs) a = if a ∈ s then (↑s.card)⁻¹ else 0

theorem Cslib.Probability.PMF.mem_support_uniformOfFinset_iff {α : Type u} {s : Finset α} (hs : s.Nonempty) (a : α) : a ∈ (uniformOfFinset s hs).support ↔ a ∈ s

noncomputable def Cslib.Probability.PMF.uniformOfFintype (α : Type u_1) [Fintype α] [Nonempty α] : PMF α

Uniform probability mass function on a nonempty finite type.

Equations

• Cslib.Probability.PMF.uniformOfFintype α = Cslib.Probability.PMF.uniformOfFinset Finset.univ ⋯

@[simp]

theorem Cslib.Probability.PMF.uniformOfFintype_apply {α : Type u} [Fintype α] [Nonempty α] (a : α) : (uniformOfFintype α) a = (↑(Fintype.card α))⁻¹

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) (uniformOfFintype α) = uniformOfFintype γ

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

theorem Cslib.Probability.PMF.uniformOfFintype_prod (α : Type u_1) (β : Type u_2) [Fintype α] [Nonempty α] [Fintype β] [Nonempty β] : ((uniformOfFintype α).bind fun (a : α) => PMF.map (fun (b : β) => (a, b)) (uniformOfFintype β)) = uniformOfFintype (α × β)

Independent uniform sampling on α and β equals uniform sampling on α × β.

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.