feat: define updateFinset which updates a finite number components …

…of a vector (#7341)

* from the Sobolev project (formerly: marginal project)

---------

Co-authored-by: Mario Carneiro <di.gama@gmail.com>

Mathlib.lean

@@ -1469,6 +1469,7 @@ import Mathlib.Data.Finset.Sort
 import Mathlib.Data.Finset.Sum
 import Mathlib.Data.Finset.Sups
 import Mathlib.Data.Finset.Sym
+import Mathlib.Data.Finset.Update
 import Mathlib.Data.Finsupp.AList
 import Mathlib.Data.Finsupp.Antidiagonal
 import Mathlib.Data.Finsupp.Basic

Mathlib/Data/Finset/Basic.lean

@@ -837,6 +837,13 @@ instance [IsEmpty α] : Unique (Finset α) where
   default := ∅
   uniq _ := eq_empty_of_forall_not_mem isEmptyElim
 
+instance (i : α) : Unique ({i} : Finset α) where
+  default := ⟨i, mem_singleton_self i⟩
+  uniq j := Subtype.ext <| mem_singleton.mp j.2
+
+@[simp]
+lemma default_singleton (i : α) : ((default : ({i} : Finset α)) : α) = i := rfl
+
 end Singleton
 
 /-! ### cons -/
@@ -1040,6 +1047,11 @@ theorem mem_disjUnion {α s t h a} : a ∈ @disjUnion α s t h ↔ a ∈ s ∨ a
   rcases s with ⟨⟨s⟩⟩; rcases t with ⟨⟨t⟩⟩; apply List.mem_append
 #align finset.mem_disj_union Finset.mem_disjUnion
 
+@[simp, norm_cast]
+theorem coe_disjUnion {s t : Finset α} (h : Disjoint s t) :
+    (disjUnion s t h : Set α) = (s : Set α) ∪ t :=
+  Set.ext <| by simp
+
 theorem disjUnion_comm (s t : Finset α) (h : Disjoint s t) :
     disjUnion s t h = disjUnion t s h.symm :=
   eq_of_veq <| add_comm _ _
@@ -3853,6 +3865,8 @@ end Finset
 
 namespace Equiv
 
+open Finset
+
 /--
 Inhabited types are equivalent to `Option β` for some `β` by identifying `default α` with `none`.
 -/
@@ -3873,6 +3887,30 @@ def sigmaEquivOptionOfInhabited (α : Type u) [Inhabited α] [DecidableEq α] :
           · simp }⟩
 #align equiv.sigma_equiv_option_of_inhabited Equiv.sigmaEquivOptionOfInhabited
 
+variable [DecidableEq α] {s t : Finset α}
+
+/-- The disjoint union of finsets is a sum -/
+def Finset.union (s t : Finset α) (h : Disjoint s t) :
+    s ⊕ t ≃ (s ∪ t : Finset α) :=
+  Equiv.Set.ofEq (coe_union _ _) |>.trans (Equiv.Set.union (disjoint_coe.mpr h).le_bot) |>.symm
+
+@[simp]
+theorem Finset.union_symm_inl (h : Disjoint s t) (x : s) :
+    Equiv.Finset.union s t h (Sum.inl x) = ⟨x, Finset.mem_union.mpr <| Or.inl x.2⟩ :=
+  rfl
+
+@[simp]
+theorem Finset.union_symm_inr (h : Disjoint s t) (y : t) :
+    Equiv.Finset.union s t h (Sum.inr y) = ⟨y, Finset.mem_union.mpr <| Or.inr y.2⟩ :=
+  rfl
+
+/-- The type of dependent functions on the disjoint union of finsets `s ∪ t` is equivalent to the
+  type of pairs of functions on `s` and on `t`. This is similar to `Equiv.sumPiEquivProdPi`. -/
+def piFinsetUnion {ι} [DecidableEq ι] (α : ι → Type*) {s t : Finset ι} (h : Disjoint s t) :
+    ((∀ i : s, α i) × ∀ i : t, α i) ≃ ∀ i : (s ∪ t : Finset ι), α i :=
+  let e := Equiv.Finset.union s t h
+  sumPiEquivProdPi (fun b ↦ α (e b)) |>.symm.trans (.piCongrLeft (fun i : ↥(s ∪ t) ↦ α i) e)
+
 end Equiv
 
 namespace Multiset

Mathlib/Data/Finset/Update.lean

@@ -0,0 +1,61 @@
+/-
+Copyright (c) 2023 Floris van Doorn. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Floris van Doorn
+-/
+import Mathlib.Data.Finset.Basic
+
+/-!
+# Update a function on a set of values
+
+This file defines `Function.updateFinset`, the operation that updates a function on a
+(finite) set of values.
+
+This is a very specific function used for `MeasureTheory.marginal`, and possibly not that useful
+for other purposes.
+-/
+variable {ι : Sort _} {π : ι → Sort _} {x : ∀ i, π i} [DecidableEq ι]
+
+namespace Function
+
+/-- `updateFinset x s y` is the vector `x` with the coordinates in `s` changed to the values of `y`.
+-/
+def updateFinset (x : ∀ i, π i) (s : Finset ι) (y : ∀ i : ↥s, π i) (i : ι) : π i :=
+  if hi : i ∈ s then y ⟨i, hi⟩ else x i
+
+open Finset Equiv
+
+@[simp] theorem updateFinset_empty {y} : updateFinset x ∅ y = x :=
+  rfl
+
+theorem updateFinset_singleton {i y} :
+    updateFinset x {i} y = Function.update x i (y ⟨i, mem_singleton_self i⟩) := by
+  congr with j
+  by_cases hj : j = i
+  · cases hj
+    simp only [dif_pos, Finset.mem_singleton, update_same, updateFinset]
+  · simp [hj, updateFinset]
+
+theorem update_eq_updateFinset {i y} :
+    Function.update x i y = updateFinset x {i} (uniqueElim y) := by
+  congr with j
+  by_cases hj : j = i
+  · cases hj
+    simp only [dif_pos, Finset.mem_singleton, update_same, updateFinset]
+    exact uniqueElim_default (α := fun j : ({i} : Finset ι) => π j) y
+  · simp [hj, updateFinset]
+
+theorem updateFinset_updateFinset {s t : Finset ι} (hst : Disjoint s t)
+    {y : ∀ i : ↥s, π i} {z : ∀ i : ↥t, π i} :
+    updateFinset (updateFinset x s y) t z =
+    updateFinset x (s ∪ t) (Equiv.piFinsetUnion π hst ⟨y, z⟩) := by
+  set e := Equiv.Finset.union s t hst
+  congr with i
+  by_cases his : i ∈ s <;> by_cases hit : i ∈ t <;>
+    simp only [updateFinset, his, hit, dif_pos, dif_neg, Finset.mem_union, true_or_iff,
+      false_or_iff, not_false_iff]
+  · exfalso; exact Finset.disjoint_left.mp hst his hit
+  · exact piCongrLeft_sum_inl (fun b : ↥(s ∪ t) => π b) e y z ⟨i, his⟩ |>.symm
+  · exact piCongrLeft_sum_inr (fun b : ↥(s ∪ t) => π b) e y z ⟨i, hit⟩ |>.symm
+
+end Function

Mathlib/Data/Set/Prod.lean

@@ -833,6 +833,9 @@ theorem Disjoint.set_pi (hi : i ∈ s) (ht : Disjoint (t₁ i) (t₂ i)) : Disjo
   disjoint_left.2 fun _ h₁ h₂ => disjoint_left.1 ht (h₁ _ hi) (h₂ _ hi)
 #align set.disjoint.set_pi Set.Disjoint.set_pi
 
+theorem uniqueElim_preimage [Unique ι] (t : ∀ i, Set (α i)) :
+    uniqueElim ⁻¹' pi univ t = t (default : ι) := by ext; simp [Unique.forall_iff]
+
 section Nonempty
 
 variable [∀ i, Nonempty (α i)]

Mathlib/Logic/Function/Basic.lean

@@ -674,6 +674,9 @@ theorem apply_update₂ {ι : Sort*} [DecidableEq ι] {α β γ : ι → Sort*}
   · simp [h]
 #align function.apply_update₂ Function.apply_update₂
 
+theorem pred_update (P : ∀ ⦃a⦄, β a → Prop) (f : ∀ a, β a) (a' : α) (v : β a') (a : α) :
+    P (update f a' v a) ↔ a = a' ∧ P v ∨ a ≠ a' ∧ P (f a) := by
+  rw [apply_update P, update_apply, ite_prop_iff_or]
 
 theorem comp_update {α' : Sort*} {β : Sort*} (f : α' → β) (g : α → α') (i : α) (v : α') :
     f ∘ update g i v = update (f ∘ g) i (f v) :=

Mathlib/Logic/Unique.lean

@@ -252,6 +252,20 @@ def Surjective.uniqueOfSurjectiveConst (α : Type*) {β : Type*} (b : β)
 
 end Function
 
+section Pi
+
+variable {ι : Type*} {α : ι → Type*}
+/-- Given one value over a unique, we get a dependent function. -/
+def uniqueElim [Unique ι] (x : α (default : ι)) (i : ι) : α i := by
+  rw [Unique.eq_default i]
+  exact x
+
+@[simp]
+theorem uniqueElim_default {_ : Unique ι} (x : α (default : ι)) : uniqueElim x (default : ι) = x :=
+  rfl
+
+end Pi
+
 -- TODO: Mario turned this off as a simp lemma in Std, wanting to profile it.
 attribute [simp] eq_iff_true_of_subsingleton in
 theorem Unique.bijective {A B} [Unique A] [Unique B] {f : A → B} : Function.Bijective f := by