use Semiring + unitInterval

Author: EtienneC30

Mathlib/MeasureTheory/Measure/Dirac.lean

@@ -141,6 +141,15 @@ lemma sum_smul_dirac_singleton [MeasurableSingletonClass α] {f : α → ℝ≥0
     sum (fun b : α ↦ f b • dirac b) {a} = f a := by
   simp +contextual [tsum_eq_single a]
 
+/-- The sum of scaled Dirac measures applied to a singleton is the coefficient of that singleton. -/
+lemma sum_smul_dirac_singleton' [MeasurableSingletonClass α] {f : δ → ℝ≥0∞}
+    {x : δ → α} (hx : Injective x) {d : δ} :
+    sum (fun d ↦ f d • dirac (x d)) {x d} = f d := by
+  simp only [MeasurableSet.singleton, sum_apply, smul_apply, dirac_apply', smul_eq_mul]
+  rw [tsum_eq_single d]
+  · simp
+  · simp; grind
+
 /-- A measure on a countable type is a sum of Dirac measures.
 If `α` has measurable singletons, `sum_smul_dirac` gives a simpler sum. -/
 lemma exists_sum_smul_dirac [Countable α] (μ : Measure α) :

Mathlib/Probability/Distributions/Geometric.lean

@@ -10,20 +10,37 @@ public import Mathlib.Probability.ProbabilityMassFunction.Basic
 
 import Mathlib.MeasureTheory.Integral.Bochner.SumMeasure
 
-/-! # Geometric distributions over ℕ
+/-! # Geometric distributions
 
-Define the geometric measure over the natural numbers. For `0 < p ≤ 1`, `geometricMeasure p` is the
-measure which to `{n}` associates `(1 - p) ^ n * n`.
+We define the geometric measure over an arbitrary semiring. For `0 < p ≤ 1`, `geometricMeasure ℕ p`
+is the measure over natural numbers which to `{n}` associates `(1 - p) ^ n * n`. To be allowed
+to apply this definition for other domains such as `ℝ`, `ℕ∞` or `ℝ≥0∞`, we define
+`geometricMeasure R p` for any semiring `R` as the measure over `R` which
+gives mass `(1 - p) ^ n * n` to `Nat.cast n`.
+
+As the parameter `p` needs to lie between `0` and `1`, we define `geometricMeasure R p` with
+`p : unitInterval`.
+
+Imagine a certain experience which has success probability `p`. If you repeat this experience
+infintely many times and independently, the number of failures before the first success
+follows a geometric distribution with parameter `p`.
+
+While the geometric measure with parameter `p := 0` is not really meaningful over `ℕ` or `ℝ`,
+it makes sense to define it as `Measure.dirac ⊤` over `ℕ∞` or `ℝ≥0∞`, since if the experience
+always failq we need infinitely many failures before succeeding.
+Therefore we require `SupSet R` in the definition
+so as to define `geometric R p := Measure.dirac (sSup univ)`
+
+TODO: define a `SupSet` instance for `ℕ∞`.
 
 ## Main definition
 
-* `geometricMeasure p`: a geometric measure on `ℕ`, parametrized by its success probability `p`.
+* `geometricMeasure R p`: a geometric measure on a semiring `R`,
+  parametrized by its success probability `p`.
 
-## Implementation note
+## Tags
 
-In order to automatically infer `IsProbabilityMeasure (geometricMeasure p)`, we define
-`geometricMeasure p` for a general `p : ℝ` and set it to `Measure.dirac 0` if `p ≤ 0` or `1 < p`.
-Use `geometricMeasure_eq` to rewrite the measure as a sum of Dirac masses.
+geometric distribution
 -/
 
 @[expose] public section
@@ -34,88 +51,95 @@ open MeasureTheory Real Set Filter Topology
 
 namespace ProbabilityTheory
 
-variable {p : ℝ}
+variable {R : Type*} [Semiring R] [SupSet R] {mR : MeasurableSpace R} {p : unitInterval}
 
 /-- The geometric measure with success probability `p` as a measure over `ℕ`. -/
-noncomputable def geometricMeasure (p : ℝ) : Measure ℕ := if 0 < p ∧ p ≤ 1
+noncomputable def geometricMeasure (R : Type*) [Semiring R] [SupSet R] [MeasurableSpace R]
+    (p : unitInterval) : Measure R := if p ≠ 0
   then
-    Measure.sum (fun n ↦ ENNReal.ofReal ((1 - p) ^ n * p) • .dirac n)
+    Measure.sum (fun (n : ℕ) ↦ ENNReal.ofReal ((1 - p) ^ n * p) • .dirac n)
   else
-    .dirac 0
+    .dirac (sSup univ)
 
-lemma geometricMeasure_eq (h1 : 0 < p) (h2 : p ≤ 1) :
-    geometricMeasure p = Measure.sum (fun n ↦ ENNReal.ofReal ((1 - p) ^ n * p) • .dirac n) :=
-  if_pos ⟨h1, h2⟩
+lemma geometricMeasure_eq (hp : p ≠ 0) :
+    geometricMeasure R p =
+      Measure.sum (fun (n : ℕ) ↦ ENNReal.ofReal ((1 - p) ^ n * p) • .dirac (n : R)) :=
+  if_pos hp
 
 /-- The `positivty` tactic does not work for this goal. Use this lemma to rewrite
-`(ENNReal.ofReal ((1 - p) ^ n * p)).toReal = (1 - p) ^ n * p)`. -/
-lemma geometricMeasure_nonneg (h1 : 0 < p) (h2 : p ≤ 1) (n : ℕ) :
-    0 ≤ (1 - p) ^ n * p := by positivity [pow_nonneg (a := 1 - p) (by linarith) n]
-
-lemma geometricMeasure_pos (h1 : 0 < p) (h2 : p < 1) (n : ℕ) :
-    0 < (1 - p) ^ n * p := by positivity [pow_pos (a := 1 - p) (by linarith) n]
-
-lemma geometricMeasure_real_singleton (h1 : 0 < p) (h2 : p ≤ 1) (n : ℕ) :
-    (geometricMeasure p).real {n} = (1 - p) ^ n * p := by
-  rw [geometricMeasure, if_pos ⟨h1, h2⟩, measureReal_def, Measure.sum_smul_dirac_singleton,
-    ENNReal.toReal_ofReal (geometricMeasure_nonneg h1 h2 n)]
-
-lemma geometricMeasure_real_singleton_pos (n : ℕ) (h1 : 0 < p) (h2 : p < 1) :
-    0 < (geometricMeasure p).real {n} := by
-  rw [geometricMeasure_real_singleton h1 h2.le]
+`(ENNReal.ofReal ((1 - p) ^ n * p)).toReal = (1 - p) ^ n * p`. -/
+lemma geometricMeasure_nonneg (p : unitInterval) (n : ℕ) :
+    0 ≤ (1 - p : ℝ) ^ n * p := mul_nonneg (pow_nonneg (by grind) n) p.2.1
+
+lemma geometricMeasure_pos (h1 : p ≠ 0) (h2 : p ≠ 1) (n : ℕ) :
+    0 < (1 - p : ℝ) ^ n * p := mul_pos (pow_pos (by grind) n) (by grind)
+
+lemma geometricMeasure_real_singleton [MeasurableSingletonClass R] [CharZero R]
+    (hp : p ≠ 0) (n : ℕ) :
+    (geometricMeasure R p).real {(n : R)} = (1 - p) ^ n * p := by
+  rw [geometricMeasure_eq hp, measureReal_def,
+    Measure.sum_smul_dirac_singleton' (Nat.cast_injective),
+    ENNReal.toReal_ofReal (geometricMeasure_nonneg p n)]
+
+lemma geometricMeasure_real_singleton_pos [MeasurableSingletonClass R] [CharZero R]
+    (h1 : p ≠ 0) (h2 : p ≠ 1) (n : ℕ) :
+    0 < (geometricMeasure R p).real {(n : R)} := by
+  rw [geometricMeasure_real_singleton h1]
   exact geometricMeasure_pos h1 h2 n
 
-lemma hasSum_one_geometricMeasure (h1 : 0 < p) (h2 : p ≤ 1) :
-    HasSum (fun n ↦ (1 - p) ^ n * p) 1 := by
-  convert (hasSum_geometric_of_lt_one (r := 1 - p) (by linarith) (by linarith)).mul_right p
-  simp [field]
+lemma hasSum_one_geometricMeasure (hp : p ≠ 0) :
+    HasSum (fun n ↦ (1 - p : ℝ) ^ n * p) 1 := by
+  convert (hasSum_geometric_of_lt_one (r := 1 - p) (by grind) (by grind)).mul_right (p : ℝ)
+  grind
 
 instance isProbabilityMeasure_geometricMeasure :
-    IsProbabilityMeasure (geometricMeasure p) := by
+    IsProbabilityMeasure (geometricMeasure R p) := by
   rw [geometricMeasure]
   split_ifs with h
-  · exact (hasSum_one_geometricMeasure h.1 h.2).isProbabilityMeasure_sum_dirac
-      (geometricMeasure_nonneg h.1 h.2)
+  · exact (hasSum_one_geometricMeasure h).isProbabilityMeasure_sum_dirac
+      (geometricMeasure_nonneg p)
   · infer_instance
 
 section Integral
 
-variable {E : Type*} [NormedAddCommGroup E]
+variable {E : Type*} [NormedAddCommGroup E] [MeasurableSingletonClass R]
 
-lemma integrable_geometricMeasure_iff (h1 : 0 < p) (h2 : p ≤ 1) {f : ℕ → E} :
-    Integrable f (geometricMeasure p) ↔ Summable (fun n ↦ (1 - p) ^ n * p * ‖f n‖) := by
-  rw [geometricMeasure_eq h1 h2, integrable_sum_dirac_iff (by simp)]
+lemma integrable_geometricMeasure_iff (hp : p ≠ 0) {f : R → E} :
+    Integrable f (geometricMeasure R p) ↔ Summable (fun (n : ℕ) ↦ (1 - p : ℝ) ^ n * p * ‖f n‖) := by
+  rw [geometricMeasure_eq hp, integrable_sum_dirac_iff (by simp)]
   congrm Summable (fun n ↦ ?_ * _)
-  rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg h1 h2 n)]
+  rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg p n)]
 
 variable [NormedSpace ℝ E]
 
-lemma hasSum_integral_geometricMeasure [CompleteSpace E] {f : ℕ → E}
-    (h1 : 0 < p) (h2 : p ≤ 1) (hf : Integrable f (geometricMeasure p)) :
-    HasSum (fun n ↦ ((1 - p) ^ n * p) • f n) (∫ n, f n ∂geometricMeasure p) := by
-  have : (fun n ↦ ((1 - p) ^ n * p) • f n) =
-      fun n ↦ (ENNReal.ofReal ((1 - p) ^ n * p)).toReal • f n := by
-    ext n; rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg h1 h2 n)]
-  rw [this, geometricMeasure_eq h1 h2]
+lemma hasSum_integral_geometricMeasure [CompleteSpace E] {f : R → E}
+    (hp : p ≠ 0) (hf : Integrable f (geometricMeasure R p)) :
+    HasSum (fun (n : ℕ) ↦ ((1 - p : ℝ) ^ n * p) • f n) (∫ n, f n ∂geometricMeasure R p) := by
+  have : (fun (n : ℕ) ↦ ((1 - p : ℝ) ^ n * p) • f n) =
+      fun (n : ℕ) ↦ (ENNReal.ofReal ((1 - p) ^ n * p)).toReal • f n := by
+    ext n; rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg p n)]
+  rw [this, geometricMeasure_eq hp]
   apply hasSum_integral_sum_dirac (by simp)
-  convert (integrable_geometricMeasure_iff h1 h2).1 hf with n
-  rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg h1 h2 n)]
+  convert (integrable_geometricMeasure_iff hp).1 hf with n
+  rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg p n)]
 
-lemma integral_geometricMeasure' [CompleteSpace E] {f : ℕ → E} (h1 : 0 < p) (h2 : p ≤ 1)
-    (hf : Integrable f (geometricMeasure p)) :
-    ∫ n, f n ∂geometricMeasure p = ∑' n, ((1 - p) ^ n * p) • f n :=
-  (hasSum_integral_geometricMeasure h1 h2 hf).tsum_eq.symm
+lemma integral_geometricMeasure' [CompleteSpace E] {f : R → E} (hp : p ≠ 0)
+    (hf : Integrable f (geometricMeasure R p)) :
+    ∫ n, f n ∂geometricMeasure R p = ∑' n : ℕ, ((1 - p : ℝ) ^ n * p) • f n :=
+  (hasSum_integral_geometricMeasure hp hf).tsum_eq.symm
 
-lemma integral_geometricMeasure [FiniteDimensional ℝ E] (h1 : 0 < p) (h2 : p ≤ 1) (f : ℕ → E) :
-    ∫ n, f n ∂geometricMeasure p = ∑' n, ((1 - p) ^ n * p) • f n := by
-  rw [geometricMeasure_eq h1 h2, integral_sum_dirac (by simp)]
+lemma integral_geometricMeasure [FiniteDimensional ℝ E] (hp : p ≠ 0) (f : R → E) :
+    ∫ n, f n ∂geometricMeasure R p = ∑' n : ℕ, ((1 - p : ℝ) ^ n * p) • f n := by
+  rw [geometricMeasure_eq hp, integral_sum_dirac (by simp)]
   congr with n
-  rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg h1 h2 n)]
+  rw [ENNReal.toReal_ofReal (geometricMeasure_nonneg p n)]
 
 end Integral
 
 section GeometricPMF
 
+variable {p : ℝ}
+
 /-- The pmf of the geometric distribution depending on its success probability. -/
 @[deprecated geometricMeasure (since := "2026-03-08")]
 noncomputable