Specific Constructions of Probability Mass Functions

This file gives a number of different PMF constructions for common probability distributions.

map and seq allow pushing a PMF α along a function f : α → β (or distribution of functions f : PMF (α → β)) to get a PMF β.

ofFinset and ofFintype simplify the construction of a PMF α from a function f : α → ℝ≥0∞, by allowing the "sum equals 1" constraint to be in terms of Finset.sum instead of tsum.

normalize constructs a PMF α by normalizing a function f : α → ℝ≥0∞ by its sum, and filter uses this to filter the support of a PMF and re-normalize the new distribution.

bernoulli represents the bernoulli distribution on Bool.

def PMF.map {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) : PMF β

The functorial action of a function on a PMF.

theorem PMF.monad_map_eq_map {α β : Type u} (f : α → β) (p : PMF α) : f <$> p = map f p

@[simp]

theorem PMF.map_apply {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) (b : β) : (map f p) b = ∑' (a : α), if b = f a then p a else 0

@[simp]

theorem PMF.support_map {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) : (map f p).support = f '' p.support

theorem PMF.mem_support_map_iff {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) (b : β) : b ∈ (map f p).support ↔ ∃ a ∈ p.support, f a = b

theorem PMF.bind_pure_comp {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) : p.bind (pure ∘ f) = map f p

theorem PMF.map_id {α : Type u_1} (p : PMF α) : map id p = p

theorem PMF.map_comp {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : α → β) (p : PMF α) (g : β → γ) : map g (map f p) = map (g ∘ f) p

theorem PMF.pure_map {α : Type u_1} {β : Type u_2} (f : α → β) (a : α) : map f (pure a) = pure (f a)

theorem PMF.map_bind {α : Type u_1} {β : Type u_2} {γ : Type u_3} (p : PMF α) (q : α → PMF β) (f : β → γ) : map f (p.bind q) = p.bind fun (a : α) => map f (q a)

@[simp]

theorem PMF.bind_map {α : Type u_1} {β : Type u_2} {γ : Type u_3} (p : PMF α) (f : α → β) (q : β → PMF γ) : (map f p).bind q = p.bind (q ∘ f)

@[simp]

theorem PMF.map_const {α : Type u_1} {β : Type u_2} (p : PMF α) (b : β) : map (Function.const α b) p = pure b

@[simp]

theorem PMF.toOuterMeasure_map_apply {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) (s : Set β) : (map f p).toOuterMeasure s = p.toOuterMeasure (f ⁻¹' s)

@[simp]

theorem PMF.toMeasure_map_apply {α : Type u_1} {β : Type u_2} (f : α → β) (p : PMF α) (s : Set β) {mα : MeasurableSpace α} {mβ : MeasurableSpace β} (hf : Measurable f) (hs : MeasurableSet s) : (map f p).toMeasure s = p.toMeasure (f ⁻¹' s)

@[simp]

theorem PMF.toMeasure_map {α : Type u_1} {β : Type u_2} (f : α → β) {mα : MeasurableSpace α} {mβ : MeasurableSpace β} (p : PMF α) (hf : Measurable f) : MeasureTheory.Measure.map f p.toMeasure = (map f p).toMeasure

def PMF.seq {α : Type u_1} {β : Type u_2} (q : PMF (α → β)) (p : PMF α) : PMF β

The monadic sequencing operation for PMF.

theorem PMF.monad_seq_eq_seq {α β : Type u} (q : PMF (α → β)) (p : PMF α) : q <*> p = q.seq p

@[simp]

theorem PMF.seq_apply {α : Type u_1} {β : Type u_2} (q : PMF (α → β)) (p : PMF α) (b : β) : (q.seq p) b = ∑' (f : α → β) (a : α), if b = f a then q f * p a else 0

@[simp]

theorem PMF.support_seq {α : Type u_1} {β : Type u_2} (q : PMF (α → β)) (p : PMF α) : (q.seq p).support = ⋃ f ∈ q.support, f '' p.support

theorem PMF.mem_support_seq_iff {α : Type u_1} {β : Type u_2} (q : PMF (α → β)) (p : PMF α) (b : β) : b ∈ (q.seq p).support ↔ ∃ f ∈ q.support, b ∈ f '' p.support

instance PMF.instLawfulFunctor : LawfulFunctor PMF

instance PMF.instLawfulMonad : LawfulMonad PMF

instance PMF.instULiftable : ULiftable PMF PMF

This instance allows do notation for PMF to be used across universes, for instance as

example {R : Type u} [Ring R] (x : PMF ℕ) : PMF R := do
  let ⟨n⟩ ← ULiftable.up x
  pure n

where x is in universe 0, but the return value is in universe u.

def PMF.ofFinset {α : Type u_1} (f : α → ENNReal) (s : Finset α) (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) : PMF α

Given a finset s and a function f : α → ℝ≥0∞ with sum 1 on s, such that f a = 0 for a ∉ s, we get a PMF.

@[simp]

theorem PMF.ofFinset_apply {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) (a : α) : (ofFinset f s h h') a = f a

@[simp]

theorem PMF.support_ofFinset {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) : (ofFinset f s h h').support = ↑s ∩ Function.support f

theorem PMF.mem_support_ofFinset_iff {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) (a : α) : a ∈ (ofFinset f s h h').support ↔ a ∈ s ∧ f a ≠ 0

theorem PMF.ofFinset_apply_of_notMem {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) {a : α} (ha : a ∉ s) : (ofFinset f s h h') a = 0

@[deprecated PMF.ofFinset_apply_of_notMem (since := "2025-05-23")]

theorem PMF.ofFinset_apply_of_not_mem {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) {a : α} (ha : a ∉ s) : (ofFinset f s h h') a = 0

Alias of PMF.ofFinset_apply_of_notMem.

@[simp]

theorem PMF.toOuterMeasure_ofFinset_apply {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) (t : Set α) : (ofFinset f s h h').toOuterMeasure t = ∑' (x : α), t.indicator f x

@[simp]

theorem PMF.toMeasure_ofFinset_apply {α : Type u_1} {f : α → ENNReal} {s : Finset α} (h : ∑ a ∈ s, f a = 1) (h' : ∀ a ∉ s, f a = 0) (t : Set α) [MeasurableSpace α] (ht : MeasurableSet t) : (ofFinset f s h h').toMeasure t = ∑' (x : α), t.indicator f x

def PMF.ofFintype {α : Type u_1} [Fintype α] (f : α → ENNReal) (h : ∑ a : α, f a = 1) : PMF α

Given a finite type α and a function f : α → ℝ≥0∞ with sum 1, we get a PMF.

@[simp]

theorem PMF.ofFintype_apply {α : Type u_1} [Fintype α] {f : α → ENNReal} (h : ∑ a : α, f a = 1) (a : α) : (ofFintype f h) a = f a

@[simp]

theorem PMF.support_ofFintype {α : Type u_1} [Fintype α] {f : α → ENNReal} (h : ∑ a : α, f a = 1) : (ofFintype f h).support = Function.support f

theorem PMF.mem_support_ofFintype_iff {α : Type u_1} [Fintype α] {f : α → ENNReal} (h : ∑ a : α, f a = 1) (a : α) : a ∈ (ofFintype f h).support ↔ f a ≠ 0

@[simp]

theorem PMF.map_ofFintype {α : Type u_1} {β : Type u_2} [Fintype α] [Fintype β] (f : α → ENNReal) (h : ∑ a : α, f a = 1) (g : α → β) : map g (ofFintype f h) = ofFintype (fun (b : β) => ∑ a : α with g a = b, f a) ⋯

@[simp]

theorem PMF.toOuterMeasure_ofFintype_apply {α : Type u_1} [Fintype α] {f : α → ENNReal} (h : ∑ a : α, f a = 1) (s : Set α) : (ofFintype f h).toOuterMeasure s = ∑' (x : α), s.indicator f x

@[simp]

theorem PMF.toMeasure_ofFintype_apply {α : Type u_1} [Fintype α] {f : α → ENNReal} (h : ∑ a : α, f a = 1) (s : Set α) [MeasurableSpace α] (hs : MeasurableSet s) : (ofFintype f h).toMeasure s = ∑' (x : α), s.indicator f x

def PMF.normalize {α : Type u_1} (f : α → ENNReal) (hf0 : tsum f ≠ 0) (hf : tsum f ≠ ⊤) : PMF α

Given an f with non-zero and non-infinite sum, get a PMF by normalizing f by its tsum.

@[simp]

theorem PMF.normalize_apply {α : Type u_1} {f : α → ENNReal} (hf0 : tsum f ≠ 0) (hf : tsum f ≠ ⊤) (a : α) : (normalize f hf0 hf) a = f a * (∑' (x : α), f x)⁻¹

@[simp]

theorem PMF.support_normalize {α : Type u_1} {f : α → ENNReal} (hf0 : tsum f ≠ 0) (hf : tsum f ≠ ⊤) : (normalize f hf0 hf).support = Function.support f

theorem PMF.mem_support_normalize_iff {α : Type u_1} {f : α → ENNReal} (hf0 : tsum f ≠ 0) (hf : tsum f ≠ ⊤) (a : α) : a ∈ (normalize f hf0 hf).support ↔ f a ≠ 0

def PMF.filter {α : Type u_1} (p : PMF α) (s : Set α) (h : ∃ a ∈ s, a ∈ p.support) : PMF α

Create new PMF by filtering on a set with non-zero measure and normalizing.

@[simp]

theorem PMF.filter_apply {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) (a : α) : (p.filter s h) a = s.indicator (⇑p) a * (∑' (a' : α), s.indicator (⇑p) a')⁻¹

theorem PMF.filter_apply_eq_zero_of_notMem {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) {a : α} (ha : a ∉ s) : (p.filter s h) a = 0

@[deprecated PMF.filter_apply_eq_zero_of_notMem (since := "2025-05-23")]

theorem PMF.filter_apply_eq_zero_of_not_mem {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) {a : α} (ha : a ∉ s) : (p.filter s h) a = 0

Alias of PMF.filter_apply_eq_zero_of_notMem.

theorem PMF.mem_support_filter_iff {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) {a : α} : a ∈ (p.filter s h).support ↔ a ∈ s ∧ a ∈ p.support

@[simp]

theorem PMF.support_filter {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) : (p.filter s h).support = s ∩ p.support

theorem PMF.filter_apply_eq_zero_iff {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) (a : α) : (p.filter s h) a = 0 ↔ a ∉ s ∨ a ∉ p.support

theorem PMF.filter_apply_ne_zero_iff {α : Type u_1} {p : PMF α} {s : Set α} (h : ∃ a ∈ s, a ∈ p.support) (a : α) : (p.filter s h) a ≠ 0 ↔ a ∈ s ∧ a ∈ p.support

def PMF.bernoulli (p : NNReal) (h : p ≤ 1) : PMF Bool

A PMF which assigns probability p to true and 1 - p to false.

@[simp]

theorem PMF.bernoulli_apply {p : NNReal} (h : p ≤ 1) (b : Bool) : (bernoulli p h) b = ↑(bif b then p else 1 - p)

@[simp]

theorem PMF.support_bernoulli {p : NNReal} (h : p ≤ 1) : (bernoulli p h).support = {b : Bool | bif b then p ≠ 0 else p ≠ 1}

theorem PMF.mem_support_bernoulli_iff {p : NNReal} (h : p ≤ 1) (b : Bool) : b ∈ (bernoulli p h).support ↔ bif b then p ≠ 0 else p ≠ 1