feat: use IndexedPartition.piecewise to create simple/measurable/strongly measurable functions (leanprover-community#32127)

The lemmas proved in this PR are needed in RemyDegenne/brownian-motion#304.

Author: CoolRmal

Mathlib.lean

@@ -5009,6 +5009,7 @@ public import Mathlib.MeasureTheory.Function.LpSpace.ContinuousFunctions
 public import Mathlib.MeasureTheory.Function.LpSpace.DomAct.Basic
 public import Mathlib.MeasureTheory.Function.LpSpace.DomAct.Continuous
 public import Mathlib.MeasureTheory.Function.LpSpace.Indicator
+public import Mathlib.MeasureTheory.Function.Piecewise
 public import Mathlib.MeasureTheory.Function.SimpleFunc
 public import Mathlib.MeasureTheory.Function.SimpleFuncDense
 public import Mathlib.MeasureTheory.Function.SimpleFuncDenseLp

Mathlib/Data/Setoid/Partition.lean

@@ -443,11 +443,12 @@ lemma piecewise_apply {β : Type*} {f : ι → α → β} (x : α) : hs.piecewis
 
 open Function
 
+variable {β : Type*} {f : ι → α → β}
+
 /-- A family of injective functions with pairwise disjoint
 domains and pairwise disjoint ranges can be glued together
 to form an injective function. -/
-theorem piecewise_inj {β : Type*} {f : ι → α → β}
-    (h_injOn : ∀ i, InjOn (f i) (s i))
+theorem piecewise_inj (h_injOn : ∀ i, InjOn (f i) (s i))
     (h_disjoint : PairwiseDisjoint (univ : Set ι) fun i => (f i) '' (s i)) :
     Injective (piecewise hs f) := by
   intro x y h
@@ -461,23 +462,64 @@ theorem piecewise_inj {β : Type*} {f : ι → α → β}
 /-- A family of bijective functions with pairwise disjoint
 domains and pairwise disjoint ranges can be glued together
 to form a bijective function. -/
-theorem piecewise_bij {β : Type*} {f : ι → α → β}
-    {t : ι → Set β} (ht : IndexedPartition t)
+theorem piecewise_bij {t : ι → Set β} (ht : IndexedPartition t)
     (hf : ∀ i, BijOn (f i) (s i) (t i)) :
     Bijective (piecewise hs f) := by
   set g := piecewise hs f with hg
-  have hg_bij : ∀ i, BijOn g (s i) (t i) := by
-    intro i
-    refine BijOn.congr (hf i) ?_
-    intro x hx
+  have hg_bij (i) : BijOn g (s i) (t i) := by
+    refine (hf i).congr fun x hx => ?_
     rw [hg, piecewise_apply, hs.mem_iff_index_eq.mp hx]
   have hg_inj : InjOn g (⋃ i, s i) := by
-    refine injOn_of_injective ?_
-    refine piecewise_inj hs (fun i ↦ BijOn.injOn (hf i)) ?h_disjoint
+    refine injOn_of_injective (piecewise_inj hs (fun i ↦ BijOn.injOn (hf i)) ?_)
     simp only [fun i ↦ BijOn.image_eq (hf i)]
     rintro i - j - hij
     exact ht.disjoint hij
   rw [← bijOn_univ, ← hs.iUnion, ← ht.iUnion]
   exact bijOn_iUnion hg_bij hg_inj
 
+theorem piecewise_preimage (f : ι → α → β) (t : Set β) :
+    hs.piecewise f ⁻¹' t = ⋃ i, s i ∩ (f i ⁻¹' t) := by
+  refine ext fun x => ⟨fun hx => ?_, fun ⟨a, ⟨i, hi⟩, ha⟩ => ?_⟩
+  · rw [mem_preimage, piecewise_apply, ← mem_preimage] at hx
+    exact mem_iUnion_of_mem (hs.index x) (mem_inter (hs.mem_index x) hx)
+  · rw [← hi, ← (mem_iff_index_eq hs).mp ha.1] at ha
+    simp_all [piecewise_apply]
+
+theorem range_piecewise (f : ι → α → β) : range (hs.piecewise f) = ⋃ i, f i '' s i := by
+  refine ext fun x => ⟨?_, fun ⟨t, ⟨i, hi⟩, ht⟩ ↦ ?_⟩
+  · rintro ⟨x, rfl⟩
+    exact mem_iUnion_of_mem (hs.index x) ⟨x, hs.mem_index x, rfl⟩
+  · simp only [← hi, mem_image] at ht
+    obtain ⟨a, ha1, ha2⟩ := ht
+    refine ⟨a, ?_⟩
+    simp only [hs.mem_iff_index_eq] at ha1
+    simpa [hs.mem_iff_index_eq, ← ha1] using ha2
+
+theorem range_piecewise_subset (f : ι → α → β) : range (hs.piecewise f) ⊆ ⋃ i, range (f i) :=
+  fun x ⟨y, hy⟩ => by simpa [IndexedPartition.piecewise_apply] using ⟨hs.index y, y, hy⟩
+
+/-- Given a collections of sets `s : ι → Set α` that forms an indexed partition, we can group
+some of the sets to obtain a coarser partition. -/
+noncomputable def coarserPartition (hs : IndexedPartition s) {κ : Type*} (g : ι → κ)
+    (hg : g.Surjective) :
+    IndexedPartition (fun k : κ => ⋃ i ∈ g ⁻¹' {k}, s i) where
+  eq_of_mem {_x _i _j} hxi hxj := by
+    obtain ⟨a, ⟨c, hc⟩, ha⟩ := hxi
+    obtain ⟨b, ⟨d, hd⟩, hb⟩ := hxj
+    simp only [← hc, mem_iUnion] at ha
+    simp only [← hd, mem_iUnion] at hb
+    have : c = d := hs.eq_of_mem ha.2 hb.2
+    by_contra!
+    have : c ≠ d := disjoint_iff_forall_ne.mp ((disjoint_singleton.mpr this).preimage g) ha.1 hb.1
+    grind
+  some k := hs.some ((singleton_nonempty k).preimage hg).some
+  some_mem k := by
+    refine mem_iUnion_of_mem ((singleton_nonempty k).preimage hg).some ?_
+    simp only [mem_preimage, mem_singleton_iff, mem_iUnion, exists_prop]
+    constructor
+    · simpa using ((singleton_nonempty k).preimage hg).some_mem
+    · exact hs.some_mem ((singleton_nonempty k).preimage hg).some
+  index x := g (hs.index x)
+  mem_index x := mem_iUnion_of_mem (hs.index x) (by simp [hs.mem_index])
+
 end IndexedPartition

Mathlib/MeasureTheory/Function/Piecewise.lean

@@ -0,0 +1,96 @@
+/-
+Copyright (c) 2026 Yongxi Lin. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Yongxi Lin
+-/
+module
+
+public import Mathlib.Data.Setoid.Partition
+public import Mathlib.MeasureTheory.Function.StronglyMeasurable.AEStronglyMeasurable
+
+/-!
+# Measurability of piecewise functions
+
+In this file, we prove some results about measurability of functions defined by using
+`IndexedPartition.piecewise`.
+
+-/
+
+@[expose] public section
+
+open MeasureTheory Set Filter
+
+namespace IndexedPartition
+
+variable {ι α β : Type*} [MeasurableSpace α] {s : ι → Set α} {f : ι → α → β}
+
+@[fun_prop]
+theorem measurable_piecewise [MeasurableSpace β] [Countable ι]
+    (hs : IndexedPartition s) (hm : ∀ i, MeasurableSet (s i)) (hf : ∀ i, Measurable (f i)) :
+    Measurable (hs.piecewise f) :=
+  fun t ht => by simpa [piecewise_preimage] using .iUnion (fun i => (hm i).inter ((hf i) ht))
+
+@[fun_prop]
+theorem aemeasurable_piecewise {μ : Measure α} [MeasurableSpace β] [Countable ι]
+    (hs : IndexedPartition s) (hm : ∀ i, MeasurableSet (s i)) (hf : ∀ i, AEMeasurable (f i) μ) :
+    AEMeasurable (hs.piecewise f) μ := by
+  choose p hp hq using hf
+  refine ⟨hs.piecewise p, hs.measurable_piecewise hm hp, ?_⟩
+  filter_upwards [ae_all_iff.2 hq] with x hx using hx (hs.index x)
+
+/-- This is the analogue of `SimpleFunc.piecewise` for `IndexedPartition`. -/
+def simpleFunc_piecewise [Finite ι] (hs : IndexedPartition s)
+    (hm : ∀ i, MeasurableSet (s i)) (f : ι → SimpleFunc α β) : SimpleFunc α β where
+  toFun := hs.piecewise (fun i => f i)
+  measurableSet_fiber' := fun _ =>
+    letI : MeasurableSpace β := ⊤
+    hs.measurable_piecewise hm (fun i => (f i).measurable) trivial
+  finite_range' := (finite_iUnion (fun i => (f i).finite_range)).subset
+    (hs.range_piecewise_subset _)
+
+@[fun_prop]
+theorem stronglyMeasurable_piecewise [Countable ι] (hs : IndexedPartition s)
+    (hm : ∀ i, MeasurableSet (s i)) [TopologicalSpace β] (hf : ∀ i, StronglyMeasurable (f i)) :
+    StronglyMeasurable (hs.piecewise f) := by
+  by_cases Fi : Finite ι
+  · refine ⟨fun n => simpleFunc_piecewise hs hm (fun i => (hf i).approx n), fun x => ?_⟩
+    simp [simpleFunc_piecewise, piecewise_apply, StronglyMeasurable.tendsto_approx]
+  simp only [not_finite_iff_infinite] at Fi
+  obtain ⟨e, -⟩ := exists_true_iff_nonempty.mpr (nonempty_equiv_of_countable (α := ℕ) (β := ι))
+  classical
+  let g (n : ℕ) (i : ι) : Fin (n + 1) :=
+    if hi : ∃ m < n, i = e m then ⟨hi.choose, by grind⟩ else Fin.last n
+  have sg (n : ℕ) : (g n).Surjective := by
+    intro b
+    unfold g
+    refine ⟨e b, ?_⟩
+    by_cases hb : b < n
+    · have : ∃ m < n, e b = e m := ⟨b, ⟨hb, rfl⟩⟩
+      simpa only [this, Fin.ext_iff] using e.injective this.choose_spec.2.symm
+    · simp [hb]
+      grind
+  have G (n : ℕ) := hs.coarserPartition (g n) (sg n)
+  refine ⟨fun n => (G n).simpleFunc_piecewise (fun i => ?_) (fun i => (hf (e i)).approx n),
+    fun x => ?_⟩
+  · exact .biUnion (to_countable _) fun _ _ ↦ hm _
+  simp only [simpleFunc_piecewise, SimpleFunc.coe_mk, piecewise_apply]
+  have : ∀ᶠ n in atTop, e ((G n).index x) = hs.index x := by
+    obtain ⟨y, hy⟩ := e.bijective.2 (hs.index x)
+    refine eventually_atTop.mpr ⟨y + 1, fun b hb => ?_⟩
+    have : y = (⟨y, by lia⟩ : Fin (b + 1)).1 := rfl
+    rw [← hy, EmbeddingLike.apply_eq_iff_eq, this, ← Fin.ext_iff, ← (G b).mem_iff_index_eq]
+    have : ∃ m < b, hs.index x = e m := ⟨y, ⟨by lia, hy.symm⟩⟩
+    simpa [g, hs.mem_iff_index_eq, this] using e.injective (hy.trans this.choose_spec.2).symm
+  have : ∀ᶠ n in atTop, (hf (hs.index x)).approx n x = (hf (e ((G n).index x))).approx n x := by
+    filter_upwards [this] with n hn using by rw [hn]
+  exact (Filter.tendsto_congr' this).mp (by simp [StronglyMeasurable.tendsto_approx])
+
+@[fun_prop]
+theorem aestronglyMeasurable_piecewise {μ : Measure α} [Countable ι] (hs : IndexedPartition s)
+    (hm : ∀ i, MeasurableSet (s i)) [TopologicalSpace β] (hf : ∀ i, AEStronglyMeasurable (f i) μ) :
+    AEStronglyMeasurable (hs.piecewise f) μ := by
+  choose p hp hq using hf
+  refine ⟨hs.piecewise p, hs.stronglyMeasurable_piecewise hm hp, ?_⟩
+  filter_upwards [ae_all_iff.2 hq] with x hx using hx (hs.index x)
+
+end IndexedPartition

Mathlib/MeasureTheory/MeasurableSpace/Basic.lean

@@ -293,10 +293,8 @@ protected theorem MeasurableSet.preimage {t : Set β} (ht : MeasurableSet t) (hf
 
 @[fun_prop]
 protected theorem Measurable.piecewise {_ : DecidablePred (· ∈ s)} (hs : MeasurableSet s)
-    (hf : Measurable f) (hg : Measurable g) : Measurable (piecewise s f g) := by
-  intro t ht
-  rw [piecewise_preimage]
-  exact hs.ite (hf ht) (hg ht)
+    (hf : Measurable f) (hg : Measurable g) : Measurable (piecewise s f g) :=
+  fun t ht => by simpa [piecewise_preimage] using hs.ite (hf ht) (hg ht)
 
 /-- This is slightly different from `Measurable.piecewise`. It can be used to show
 `Measurable (ite (x=0) 0 1)` by