Monad Operations for Probability Mass Functions

This file constructs two operations on PMF that give it a monad structure. pure a is the distribution where a single value a has probability 1. bind pa pb : PMF β is the distribution given by sampling a : α from pa : PMF α, and then sampling from pb a : PMF β to get a final result b : β.

bindOnSupport generalizes bind to allow binding to a partial function, so that the second argument only needs to be defined on the support of the first argument.

def PMF.pure {α : Type u_1} (a : α) : PMF α

The pure PMF is the PMF where all the mass lies in one point. The value of pure a is 1 at a and 0 elsewhere.

Equations

• PMF.pure a = { val := fun (a' : α) => if a' = a then 1 else 0, property := ⋯ }

@[simp]

theorem PMF.pure_apply {α : Type u_1} (a : α) (a' : α) : (PMF.pure a) a' = if a' = a then 1 else 0

@[simp]

theorem PMF.support_pure {α : Type u_1} (a : α) : PMF.support (PMF.pure a) = {a}

theorem PMF.mem_support_pure_iff {α : Type u_1} (a : α) (a' : α) : a' ∈ PMF.support (PMF.pure a) ↔ a' = a

theorem PMF.pure_apply_self {α : Type u_1} (a : α) : (PMF.pure a) a = 1

theorem PMF.pure_apply_of_ne {α : Type u_1} (a : α) (a' : α) (h : a' ≠ a) : (PMF.pure a) a' = 0

instance PMF.instInhabitedPMF {α : Type u_1} [Inhabited α] : Inhabited (PMF α)

Equations

• PMF.instInhabitedPMF = { default := PMF.pure default }

@[simp]

theorem PMF.toOuterMeasure_pure_apply {α : Type u_1} (a : α) (s : Set α) : ↑(PMF.toOuterMeasure (PMF.pure a)) s = if a ∈ s then 1 else 0

@[simp]

theorem PMF.toMeasure_pure_apply {α : Type u_1} (a : α) (s : Set α) [MeasurableSpace α] (hs : MeasurableSet s) : ↑↑(PMF.toMeasure (PMF.pure a)) s = if a ∈ s then 1 else 0

The measure of a set under pure a is 1 for sets containing a and 0 otherwise.

theorem PMF.toMeasure_pure {α : Type u_1} (a : α) [MeasurableSpace α] : PMF.toMeasure (PMF.pure a) = MeasureTheory.Measure.dirac a

@[simp]

theorem PMF.toPMF_dirac {α : Type u_1} (a : α) [MeasurableSpace α] [Countable α] [h : MeasurableSingletonClass α] : MeasureTheory.Measure.toPMF (MeasureTheory.Measure.dirac a) = PMF.pure a

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

The monadic bind operation for PMF.

Equations

• PMF.bind p f = { val := fun (b : β) => ∑' (a : α), p a * (f a) b, property := ⋯ }

@[simp]

theorem PMF.bind_apply {α : Type u_1} {β : Type u_2} (p : PMF α) (f : α → PMF β) (b : β) : (PMF.bind p f) b = ∑' (a : α), p a * (f a) b

@[simp]

theorem PMF.support_bind {α : Type u_1} {β : Type u_2} (p : PMF α) (f : α → PMF β) : PMF.support (PMF.bind p f) = ⋃ a ∈ PMF.support p, PMF.support (f a)

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

@[simp]

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

@[simp]

theorem PMF.bind_pure {α : Type u_1} (p : PMF α) : PMF.bind p PMF.pure = p

@[simp]

theorem PMF.bind_const {α : Type u_1} {β : Type u_2} (p : PMF α) (q : PMF β) : (PMF.bind p fun (x : α) => q) = q

@[simp]

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

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

@[simp]

theorem PMF.toOuterMeasure_bind_apply {α : Type u_1} {β : Type u_2} (p : PMF α) (f : α → PMF β) (s : Set β) : ↑(PMF.toOuterMeasure (PMF.bind p f)) s = ∑' (a : α), p a * ↑(PMF.toOuterMeasure (f a)) s

@[simp]

theorem PMF.toMeasure_bind_apply {α : Type u_1} {β : Type u_2} (p : PMF α) (f : α → PMF β) (s : Set β) [MeasurableSpace β] (hs : MeasurableSet s) : ↑↑(PMF.toMeasure (PMF.bind p f)) s = ∑' (a : α), p a * ↑↑(PMF.toMeasure (f a)) s

The measure of a set under p.bind f is the sum over a : α of the probability of a under p times the measure of the set under f a.

instance PMF.instMonadPMF : Monad PMF

Equations

• PMF.instMonadPMF = Monad.mk

def PMF.bindOnSupport {α : Type u_1} {β : Type u_2} (p : PMF α) (f : (a : α) → a ∈ PMF.support p → PMF β) : PMF β

Generalized version of bind allowing f to only be defined on the support of p. p.bind f is equivalent to p.bindOnSupport (fun a _ ↦ f a), see bindOnSupport_eq_bind.

Equations

• PMF.bindOnSupport p f = { val := fun (b : β) => ∑' (a : α), p a * if h : p a = 0 then 0 else (f a h) b, property := ⋯ }

@[simp]

theorem PMF.bindOnSupport_apply {α : Type u_1} {β : Type u_2} {p : PMF α} (f : (a : α) → a ∈ PMF.support p → PMF β) (b : β) : (PMF.bindOnSupport p f) b = ∑' (a : α), p a * if h : p a = 0 then 0 else (f a h) b

@[simp]

theorem PMF.support_bindOnSupport {α : Type u_1} {β : Type u_2} {p : PMF α} (f : (a : α) → a ∈ PMF.support p → PMF β) : PMF.support (PMF.bindOnSupport p f) = ⋃ (a : α), ⋃ (h : a ∈ PMF.support p), PMF.support (f a h)

theorem PMF.mem_support_bindOnSupport_iff {α : Type u_1} {β : Type u_2} {p : PMF α} (f : (a : α) → a ∈ PMF.support p → PMF β) (b : β) : b ∈ PMF.support (PMF.bindOnSupport p f) ↔ ∃ (a : α) (h : a ∈ PMF.support p), b ∈ PMF.support (f a h)

@[simp]

theorem PMF.bindOnSupport_eq_bind {α : Type u_1} {β : Type u_2} (p : PMF α) (f : α → PMF β) : (PMF.bindOnSupport p fun (a : α) (x : a ∈ PMF.support p) => f a) = PMF.bind p f

bindOnSupport reduces to bind if f doesn't depend on the additional hypothesis.

theorem PMF.bindOnSupport_eq_zero_iff {α : Type u_1} {β : Type u_2} {p : PMF α} (f : (a : α) → a ∈ PMF.support p → PMF β) (b : β) : (PMF.bindOnSupport p f) b = 0 ↔ ∀ (a : α) (ha : p a ≠ 0), (f a ha) b = 0

@[simp]

theorem PMF.pure_bindOnSupport {α : Type u_1} {β : Type u_2} (a : α) (f : (a' : α) → a' ∈ PMF.support (PMF.pure a) → PMF β) : PMF.bindOnSupport (PMF.pure a) f = f a ⋯

theorem PMF.bindOnSupport_pure {α : Type u_1} (p : PMF α) : (PMF.bindOnSupport p fun (a : α) (x : a ∈ PMF.support p) => PMF.pure a) = p

@[simp]

theorem PMF.bindOnSupport_bindOnSupport {α : Type u_1} {β : Type u_2} {γ : Type u_3} (p : PMF α) (f : (a : α) → a ∈ PMF.support p → PMF β) (g : (b : β) → b ∈ PMF.support (PMF.bindOnSupport p f) → PMF γ) : PMF.bindOnSupport (PMF.bindOnSupport p f) g = PMF.bindOnSupport p fun (a : α) (ha : a ∈ PMF.support p) => PMF.bindOnSupport (f a ha) fun (b : β) (hb : b ∈ PMF.support (f a ha)) => g b ⋯

theorem PMF.bindOnSupport_comm {α : Type u_1} {β : Type u_2} {γ : Type u_3} (p : PMF α) (q : PMF β) (f : (a : α) → a ∈ PMF.support p → (b : β) → b ∈ PMF.support q → PMF γ) : (PMF.bindOnSupport p fun (a : α) (ha : a ∈ PMF.support p) => PMF.bindOnSupport q (f a ha)) = PMF.bindOnSupport q fun (b : β) (hb : b ∈ PMF.support q) => PMF.bindOnSupport p fun (a : α) (ha : a ∈ PMF.support p) => f a ha b hb

@[simp]

theorem PMF.toOuterMeasure_bindOnSupport_apply {α : Type u_1} {β : Type u_2} {p : PMF α} (f : (a : α) → a ∈ PMF.support p → PMF β) (s : Set β) : ↑(PMF.toOuterMeasure (PMF.bindOnSupport p f)) s = ∑' (a : α), p a * if h : p a = 0 then 0 else ↑(PMF.toOuterMeasure (f a h)) s

@[simp]

theorem PMF.toMeasure_bindOnSupport_apply {α : Type u_1} {β : Type u_2} {p : PMF α} (f : (a : α) → a ∈ PMF.support p → PMF β) (s : Set β) [MeasurableSpace β] (hs : MeasurableSet s) : ↑↑(PMF.toMeasure (PMF.bindOnSupport p f)) s = ∑' (a : α), p a * if h : p a = 0 then 0 else ↑↑(PMF.toMeasure (f a h)) s

The measure of a set under p.bindOnSupport f is the sum over a : α of the probability of a under p times the measure of the set under f a _. The additional if statement is needed since f is only a partial function.