[Merged by Bors] - feat(measure_theory/probability_mass_function): Add definitions for filtering pmfs on a predicate

src/algebra/big_operators/order.lean

@@ -136,6 +136,9 @@ variables [canonically_ordered_add_monoid β]
 @[simp] lemma sum_eq_zero_iff : ∑ x in s, f x = 0 ↔ ∀ x ∈ s, f x = 0 :=
 sum_eq_zero_iff_of_nonneg $ λ x hx, zero_le (f x)
 
+lemma sum_ne_zero_iff : ∑ x in s, f x ≠ 0 ↔ ∃ x ∈ s, f x ≠ 0 :=
+by simp
+
 lemma sum_le_sum_of_subset (h : s₁ ⊆ s₂) : (∑ x in s₁, f x) ≤ (∑ x in s₂, f x) :=
 sum_le_sum_of_subset_of_nonneg h $ assume x h₁ h₂, zero_le _
 

src/measure_theory/probability_mass_function.lean

@@ -47,6 +47,9 @@ lemma summable_coe (p : pmf α) : summable p := (p.has_sum_coe_one).summable
 /-- The support of a `pmf` is the set where it is nonzero. -/
 def support (p : pmf α) : set α := {a | p.1 a ≠ 0}
 
+@[simp] lemma mem_support_iff {p : pmf α} {a : α} :
+  a ∈ p.support ↔ p a ≠ 0 := iff.rfl
+
 /-- 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. -/
 def pure (a : α) : pmf α := ⟨λ a', if a' = a then 1 else 0, has_sum_ite_eq _ _⟩
@@ -148,6 +151,50 @@ def of_multiset (s : multiset α) (hs : s ≠ 0) : pmf α :=
 def of_fintype [fintype α] (f : α → ℝ≥0) (h : ∑ x, f x = 1) : pmf α :=
 ⟨f, h ▸ has_sum_sum_of_ne_finset_zero (by simp)⟩
 
+/-- Given a fintype type `α` and a function `f : α → ℝ≥0` not identically 0,
+  we get a `pmf` by normalizing `f` -/
+def of_fintype' [fintype α] (f : α → ℝ≥0) (h : ∃ a, f a ≠ 0) : pmf α :=
+⟨λ a, f a * (∑ x, f x)⁻¹,
+  have 1 = ∑ a, f a * (∑ x, f x)⁻¹, by simp only [← finset.sum_mul,
+    mul_inv_cancel (finset.sum_ne_zero_iff.2 (let ⟨a, ha⟩ := h in ⟨a, finset.mem_univ a, ha⟩))],
+  this.symm ▸ has_sum_fintype (λ (a : α), f a / ∑ x, f x)⟩
+
+/-- Given a `f` with non-zero sum, we get a `pmf` by normalizing `f` by its `sum` -/
+def normalize (f : α → ℝ≥0) (hf0 : tsum f ≠ 0) : pmf α :=
+⟨λ a, f a * (∑' x, f x)⁻¹,
+  (mul_inv_cancel hf0) ▸ has_sum.mul_right (∑' x, f x)⁻¹
+    (not_not.mp (mt tsum_eq_zero_of_not_summable hf0 : ¬¬summable f)).has_sum⟩
+
+lemma normalize_apply (f : α → ℝ≥0) {hf0 : tsum f ≠ 0} {a : α} :
+  (normalize f hf0) a = f a * (∑' x, f x)⁻¹ := rfl
+
+/-- Given a `pmf` on a fintype `α`, create new `pmf` by filtering on a set and normalizing -/
+def filter (p : pmf α) (s : set α) (h : ∃ a ∈ s, p a ≠ 0) : pmf α :=
+pmf.normalize (s.indicator p) $ nnreal.tsum_indicator_ne_zero p.2.summable h
+
+lemma filter_apply (p : pmf α) {s : set α} {h : ∃ a ∈ s, p a ≠ 0} {a : α} :
+  (p.filter s h) a = (s.indicator p a) * (∑' x, (s.indicator p) x)⁻¹ :=
+by rw [filter, normalize_apply]
+
+lemma filter_apply_eq_zero_of_not_mem (p : pmf α) {s : set α} {h : ∃ a ∈ s, p a ≠ 0}
+  {a : α} (ha : a ∉ s) : (p.filter s h) a = 0 :=
+by rw [filter_apply, set.indicator_apply_eq_zero.mpr (λ ha', absurd ha' ha), zero_mul]
+
+lemma filter_apply_eq_zero_iff (p : pmf α) {s : set α} {h : ∃ a ∈ s, p a ≠ 0} {a : α} :
+  (p.filter s h) a = 0 ↔ p a = 0 ∨ a ∉ s :=
+begin
+  split; intro ha,
+  { rw [filter_apply, mul_eq_zero] at ha,
+    refine ha.by_cases
+      (λ ha, (em (a ∈ s)).by_cases (λ h, or.inl ((set.indicator_apply_eq_zero.mp ha) h)) or.inr)
+      (λ ha, absurd (inv_eq_zero.1 ha) (nnreal.tsum_indicator_ne_zero p.2.summable h)) },
+  { rw [filter_apply, set.indicator_apply_eq_zero.2 (λ h, ha.by_cases id (absurd h)), zero_mul] }
+end
+
+lemma filter_apply_ne_zero_iff (p : pmf α) {s : set α} {h : ∃ a ∈ s, p a ≠ 0} {a : α} :
+  (p.filter s h) a ≠ 0 ↔ a ∈ p.support ∧ a ∈ s :=
+by rw [← not_iff, filter_apply_eq_zero_iff, not_iff, not_or_distrib, not_not, mem_support_iff]
+
 /-- A `pmf` which assigns probability `p` to `tt` and `1 - p` to `ff`. -/
 def bernoulli (p : ℝ≥0) (h : p ≤ 1) : pmf bool :=
 of_fintype (λ b, cond b p (1 - p)) (nnreal.eq $ by simp [h])

src/topology/instances/ennreal.lean

@@ -672,6 +672,19 @@ begin
       using h₂ }
 end
 
+lemma tsum_indicator_ne_zero {f : α → ℝ≥0} (hf : summable f) {s : set α} (h : ∃ a ∈ s, f a ≠ 0) :
+  ∑' x, (s.indicator f) x ≠ 0 :=
+begin
+  have : summable (s.indicator f),
+  { refine nnreal.summable_of_le (λ a, le_trans (le_of_eq (s.indicator_apply f a)) _) hf,
+    split_ifs,
+    exact le_refl (f a),
+    exact nnreal.zero_le_coe },
+  exact λ h', let ⟨a, ha, hap⟩ := h in
+    hap (trans (set.indicator_apply_eq_self.mpr (absurd ha)).symm
+      (((tsum_eq_zero_iff this).1 h') a)),
+end
+
 open finset
 
 /-- For `f : ℕ → ℝ≥0`, then `∑' k, f (k + i)` tends to zero. This does not require a summability