Probability mass functions with finite support and non-negative rational weights

This is a special case of PMF that suffices for denotational semantics of OracleComp with finite oracle specifications.

Design

We use a two-layer approach:

• FinRatPMF.Raw α is an array-based representation storing pairs (a, p) of outcomes and probabilities. It has a computable Monad and LawfulMonad instance but non-canonical equality: two values representing the same distribution may differ in array order or duplicate entries.

• FinRatPMF α is the quotient of FinRatPMF.Raw α by distributional equality (SameDist). It has canonical equality (two values are equal iff they define the same distribution) and a LawfulMonad instance (necessarily noncomputable since bind must extract from the quotient).

For computation, use FinRatPMF.Raw. For reasoning about distributional equality, use FinRatPMF. Both connect to Mathlib's PMF via toPMF.

Raw representation

structure FinRatPMF.Raw (α : Type u) : Type u

Raw probability mass function with finite support and non-negative rational weights.

Stores an array of (outcome, weight) pairs whose weights sum to 1. This representation is computable but not canonical: many different FinRatPMF.Raw values can represent the same probability distribution.

• data : Array (α × ℚ≥0)
• sum_eq_one : (List.map Prod.snd self.data.toList).sum = 1

def FinRatPMF.Raw.toList {α : Type u} (p : Raw α) : List (α × ℚ≥0)

View the raw array representation as a list, which is convenient for proofs.

def FinRatPMF.Raw.probOfList {α : Type u} [DecidableEq α] (l : List (α × ℚ≥0)) (x : α) : ℚ≥0

def FinRatPMF.Raw.supportOfList {α : Type u} [DecidableEq α] (l : List (α × ℚ≥0)) : Finset α

theorem FinRatPMF.Raw.ext {α : Type u} {p q : Raw α} (h : p.data = q.data) : p = q

theorem FinRatPMF.Raw.ext_iff {α : Type u} {p q : Raw α} : p = q ↔ p.data = q.data

theorem FinRatPMF.Raw.ext_toList {α : Type u} {p q : Raw α} (h : p.toList = q.toList) : p = q

Bind sum auxiliary lemma

Monad operations

def FinRatPMF.Raw.pure {α : Type u} (a : α) : Raw α

@[implicit_reducible]

instance FinRatPMF.Raw.instInhabited {α : Type u} [Inhabited α] : Inhabited (Raw α)

@[implicit_reducible]

instance FinRatPMF.Raw.instRepr {α : Type u} [Repr α] : Repr (Raw α)

@[implicit_reducible]

instance FinRatPMF.Raw.instToStringOfRepr {α : Type u} [Repr α] : ToString (Raw α)

def FinRatPMF.Raw.bind {α β : Type u} (f : Raw α) (g : α → Raw β) : Raw β

@[implicit_reducible]

instance FinRatPMF.Raw.instMonad : Monad Raw

@[simp]

theorem FinRatPMF.Raw.pure_toList {α : Type u} (a : α) : (pure a).toList = [(a, 1)]

@[simp]

theorem FinRatPMF.Raw.bind_toList {α β : Type u} (f : Raw α) (g : α → Raw β) : (f >>= g).toList = List.flatMap (fun (x : α × ℚ≥0) => match x with | (a, p) => List.map (fun (x : β × ℚ≥0) => match x with | (b, q) => (b, p * q)) (g a).toList) f.toList

instance FinRatPMF.Raw.instLawfulMonad : LawfulMonad Raw

def FinRatPMF.Raw.uniformList {α : Type u} (l : List α) (hl : l ≠ []) : Raw α

Uniform distribution over a nonempty list, treating repeated entries as repeated tickets.

def FinRatPMF.Raw.uniform {α : Type u} [FinEnum α] [Inhabited α] : Raw α

Uniform distribution over a finite inhabited type.

def FinRatPMF.Raw.coin : Raw Bool

Fair coin.

def FinRatPMF.Raw.bernoulli (p : ℚ≥0) (hp : p ≤ 1) : Raw Bool

Bernoulli distribution with probability p of true.

Probability and support

def FinRatPMF.Raw.prob {α : Type u} [DecidableEq α] (p : Raw α) (x : α) : ℚ≥0

The probability assigned to x by the raw PMF: sum of weights for all entries with key x.

theorem FinRatPMF.Raw.prob_eq_prob {α : Type u} (inst1 inst2 : DecidableEq α) (p : Raw α) (x : α) : p.prob x = p.prob x

prob is independent of the DecidableEq instance, since decide (a = b) gives the same Bool for any Decidable instance on the same Prop.

def FinRatPMF.Raw.support {α : Type u} [DecidableEq α] (p : Raw α) : Finset α

The finite support: the set of outcomes appearing in the array.

theorem FinRatPMF.Raw.prob_eq_zero_of_not_mem_support {α : Type u} [DecidableEq α] (p : Raw α) {x : α} (hx : x ∉ p.support) : p.prob x = 0

theorem FinRatPMF.Raw.mem_support_iff {α : Type u} [DecidableEq α] (p : Raw α) (x : α) : x ∈ p.support ↔ p.prob x ≠ 0

theorem FinRatPMF.Raw.sum_prob_eq_sum {α : Type u} [DecidableEq α] (p : Raw α) : p.support.sum p.prob = (List.map Prod.snd p.toList).sum

Sum of prob over the support recovers the total weight.

theorem FinRatPMF.Raw.sum_prob_eq_one {α : Type u} [DecidableEq α] (p : Raw α) : p.support.sum p.prob = 1

@[simp]

theorem FinRatPMF.Raw.prob_pure {α : Type u} [DecidableEq α] (a x : α) : (Raw.pure a).prob x = if x = a then 1 else 0

@[simp]

theorem FinRatPMF.Raw.support_pure {α : Type u} [DecidableEq α] (a : α) : (Raw.pure a).support = {a}

theorem FinRatPMF.Raw.prob_uniformList_of_nodup {α : Type u} [DecidableEq α] {l : List α} (hl : l ≠ []) (hnd : l.Nodup) (x : α) : (Raw.uniformList l hl).prob x = if x ∈ l then (↑l.length)⁻¹ else 0

Point probabilities for uniformList on a duplicate-free list.

theorem FinRatPMF.Raw.support_uniformList_of_nodup {α : Type u} [DecidableEq α] {l : List α} (hl : l ≠ []) (hnd : l.Nodup) : (Raw.uniformList l hl).support = l.toFinset

Uniform distributions over duplicate-free lists have support exactly the listed elements.

@[simp]

theorem FinRatPMF.Raw.prob_uniform {α : Type u} [FinEnum α] [Inhabited α] (x : α) : Raw.uniform.prob x = (↑(Fintype.card α))⁻¹

Point probabilities of Raw.uniform.

@[simp]

theorem FinRatPMF.Raw.support_uniform {α : Type u} [DecidableEq α] [FinEnum α] [Inhabited α] : Raw.uniform.support = Finset.univ

Raw.uniform has full finite support.

@[simp]

theorem FinRatPMF.Raw.prob_coin (b : Bool) : Raw.coin.prob b = 2⁻¹

Point probabilities of Raw.coin.

@[simp]

theorem FinRatPMF.Raw.support_coin : Raw.coin.support = {true, false}

@[simp]

theorem FinRatPMF.Raw.prob_bernoulli_true (p : ℚ≥0) (hp : p ≤ 1) : (Raw.bernoulli p hp).prob true = p

@[simp]

theorem FinRatPMF.Raw.prob_bernoulli_false (p : ℚ≥0) (hp : p ≤ 1) : (Raw.bernoulli p hp).prob false = 1 - p

theorem FinRatPMF.Raw.prob_le_one {α : Type u} [DecidableEq α] (p : Raw α) (x : α) : p.prob x ≤ 1

@[simp]

theorem FinRatPMF.Raw.prob_bind {α β : Type u} [DecidableEq α] [DecidableEq β] (m : Raw α) (f : α → Raw β) (y : β) : (m >>= f).prob y = ∑ a ∈ m.support, m.prob a * (f a).prob y

theorem FinRatPMF.Raw.mem_support_bind_iff {α β : Type u} [DecidableEq α] [DecidableEq β] (m : Raw α) (f : α → Raw β) (y : β) : y ∈ (m >>= f).support ↔ ∃ x ∈ m.support, y ∈ (f x).support

@[simp]

theorem FinRatPMF.Raw.support_bind {α β : Type u} [DecidableEq α] [DecidableEq β] (m : Raw α) (f : α → Raw β) : (m >>= f).support = m.support.biUnion fun (x : α) => (f x).support

def FinRatPMF.Raw.addWeight {α : Type u} [BEq α] [Hashable α] (m : Std.HashMap α ℚ≥0) (a : α × ℚ≥0) : Std.HashMap α ℚ≥0

def FinRatPMF.Raw.accumulateWeights {α : Type u} [BEq α] [Hashable α] (xs : Array (α × ℚ≥0)) : Std.HashMap α ℚ≥0

def FinRatPMF.Raw.normalizeMap {α : Type u} [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) : Std.HashMap α ℚ≥0

def FinRatPMF.Raw.normalize {α : Type u} [DecidableEq α] [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) : Raw α

Merge duplicate outcomes and discard zero-weight entries. This implementation is linear in the number of raw entries up to hash-map operations, so it avoids the repeated support/prob rescans that would make normalization quadratic.

@[simp]

theorem FinRatPMF.Raw.normalize_toList {α : Type u} [DecidableEq α] [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) : p.normalize.toList = p.normalizeMap.toList

@[simp]

theorem FinRatPMF.Raw.prob_normalize {α : Type u} [DecidableEq α] [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) (x : α) : p.normalize.prob x = p.prob x

@[simp]

theorem FinRatPMF.Raw.support_normalize {α : Type u} [DecidableEq α] [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) : p.normalize.support = p.support

noncomputable def FinRatPMF.Raw.toPMF {α : Type u} [DecidableEq α] (p : Raw α) : PMF α

Convert to a PMF by mapping weights through ℚ≥0 → ℝ≥0 → ℝ≥0∞.

@[simp]

theorem FinRatPMF.Raw.toPMF_apply {α : Type u} [DecidableEq α] (p : Raw α) (x : α) : p.toPMF x = ↑↑(p.prob x)

@[simp]

theorem FinRatPMF.Raw.toPMF_pure {α : Type u} [DecidableEq α] (a : α) : (Raw.pure a).toPMF = PMF.pure a

@[simp]

theorem FinRatPMF.Raw.toPMF_uniform {α : Type u} [FinEnum α] [Inhabited α] : Raw.uniform.toPMF = PMF.uniformOfFintype α

@[simp]

theorem FinRatPMF.Raw.toPMF_coin : Raw.coin.toPMF = PMF.uniformOfFintype Bool

@[simp]

theorem FinRatPMF.Raw.toPMF_bind {α β : Type u} [DecidableEq α] [DecidableEq β] (m : Raw α) (f : α → Raw β) : (m >>= f).toPMF = do let a ← m.toPMF (f a).toPMF

noncomputable def FinRatPMF.Raw.toPMFHom : Raw →ᵐ PMF

Distributional equality and quotient

def FinRatPMF.SameDist {α : Type u} (p q : Raw α) : Prop

Two raw PMFs represent the same distribution if they assign the same probability to every outcome. Uses classical decidable equality since this is a Prop.

theorem FinRatPMF.SameDist.refl {α : Type u} (p : Raw α) : SameDist p p

theorem FinRatPMF.SameDist.symm {α : Type u} {p q : Raw α} (h : SameDist p q) : SameDist q p

theorem FinRatPMF.SameDist.trans {α : Type u} {p q r : Raw α} (hpq : SameDist p q) (hqr : SameDist q r) : SameDist p r

theorem FinRatPMF.SameDist.prob_eq {α : Type u} [DecidableEq α] {p q : Raw α} (h : SameDist p q) (x : α) : p.prob x = q.prob x

theorem FinRatPMF.SameDist.support_eq {α : Type u} [DecidableEq α] {p q : Raw α} (h : SameDist p q) : p.support = q.support

theorem FinRatPMF.SameDist.bind_congr {α β : Type u} {p q : Raw α} {f g : α → Raw β} (hpq : SameDist p q) (hfg : ∀ (x : α), SameDist (f x) (g x)) : SameDist (p >>= f) (q >>= g)

theorem FinRatPMF.SameDist.map_congr {α β : Type u} {p q : Raw α} (hpq : SameDist p q) (f : α → β) : SameDist (f <$> p) (f <$> q)

instance FinRatPMF.instIsEquivRawSameDist {α : Type u} : IsEquiv (Raw α) SameDist

@[implicit_reducible]

instance FinRatPMF.Raw.instSetoid {α : Type u} : Setoid (Raw α)

theorem FinRatPMF.Raw.sameDist_normalize {α : Type u} [DecidableEq α] [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) : SameDist p p.normalize

theorem FinRatPMF.Raw.sameDist_normalize_normalize {α : Type u} [DecidableEq α] [BEq α] [Hashable α] [LawfulBEq α] [LawfulHashable α] (p : Raw α) : SameDist p.normalize.normalize p.normalize

def FinRatPMF (α : Type u) : Type u

Probability mass function with finite support and rational weights, with canonical equality.

Defined as the quotient of FinRatPMF.Raw by distributional equality. Two values are equal iff they assign the same probability to every outcome.

def FinRatPMF.mk {α : Type u} (p : Raw α) : FinRatPMF α

Construct from a raw PMF.

theorem FinRatPMF.eq_iff {α : Type u} {p q : Raw α} : mk p = mk q ↔ SameDist p q

FinRatPMF equality is distributional equality.

Monad operations

@[implicit_reducible]

noncomputable instance FinRatPMF.instMonad : Monad FinRatPMF

bind lifts Raw.bind across the quotient in the first argument. The codomain remains noncomputable because the Kleisli function returns quotient values, so we still choose a raw representative for each branch.

theorem FinRatPMF.bind_eq_out {α β : Type u} (ma : FinRatPMF α) (f : α → FinRatPMF β) : ma >>= f = mk ((Quotient.out ma).bind fun (a : α) => Quotient.out (f a))

@[simp]

theorem FinRatPMF.bind_mk {α β : Type u} (p : Raw α) (f : α → FinRatPMF β) : mk p >>= f = mk (p.bind fun (a : α) => Quotient.out (f a))

Connection to PMF

noncomputable def FinRatPMF.toPMF {α : Type u} (p : FinRatPMF α) : PMF α

Convert a FinRatPMF to a PMF. Well-defined since SameDist implies equal toPMF.

@[simp]

theorem FinRatPMF.toPMF_mk {α : Type u} (p : Raw α) : (mk p).toPMF = p.toPMF

@[simp]

theorem FinRatPMF.toPMF_pure {α : Type u} (a : α) : (pure a).toPMF = PMF.pure a

theorem FinRatPMF.toPMF_out {α : Type u} (p : FinRatPMF α) : p.toPMF = (Quotient.out p).toPMF

@[simp]

theorem FinRatPMF.toPMF_bind {α β : Type u} (ma : FinRatPMF α) (f : α → FinRatPMF β) : (ma >>= f).toPMF = do let a ← ma.toPMF (f a).toPMF

noncomputable def FinRatPMF.toPMFHom : FinRatPMF →ᵐ PMF

theorem FinRatPMF.toPMF_injective {α : Type u} : Function.Injective toPMF

instance FinRatPMF.instLawfulMonad : LawfulMonad FinRatPMF

@[implicit_reducible]

instance FinRatPMF.instDecidableSameDist {α : Type u} [DecidableEq α] (p q : Raw α) : Decidable (SameDist p q)

@[implicit_reducible]

instance FinRatPMF.instBEqRawOfDecidableEq {α : Type u} [DecidableEq α] : BEq (Raw α)

@[simp]

theorem FinRatPMF.beq_iff_sameDist {α : Type u} [DecidableEq α] (p q : Raw α) : (p == q) = true ↔ SameDist p q