feat(LinearAlgebra/ConvexSpace): StdSimplex bind operation

Mathlib/LinearAlgebra/ConvexSpace.lean

@@ -134,6 +134,14 @@ lemma map_map (f : StdSimplex R M) (g₁ : M → N) (g₂ : N → P) :
     (f.map g₁).map g₂ = f.map (fun x ↦ g₂ (g₁ x)) :=
   (map_comp ..).symm
 
+lemma sum_map (s : StdSimplex R M) (f : M → N) {g : N → R → R}
+    (hadd : ∀ (a : N) (b₁ b₂ : R), g a (b₁ + b₂) = g a b₁ + g a b₂) :
+    (map f s).sum g = s.sum (fun m r ↦ g (f m) r) := by
+  have hzero (n : N) : g n 0 = 0 := by simpa using hadd n 0 0
+  simp only [map, Finsupp.mapDomain, Finsupp.sum_sum_index hzero hadd]
+  congr with m r
+  rw [Finsupp.sum_single_index (hzero (f m))]
+
 /--
 Join operation for standard simplices (monadic join).
 Given a distribution over distributions, flattens it to a single distribution.
@@ -146,6 +154,38 @@ def join (f : StdSimplex R (StdSimplex R M)) : StdSimplex R M where
     convert f.total
     rw [Finsupp.sum_smul_index (fun _ ↦ rfl), ← Finsupp.mul_sum, StdSimplex.total, mul_one]
 
+/--
+Monadic bind operation for standard simplices. Given weights `f : StdSimplex R M` and
+a family of weights `g : M → StdSimplex R N'`, `StdSimplex.bind f g` is the convex
+combination of the `g m`, weighted by `f`.
+-/
+def bind (f : StdSimplex R M) (g : M → StdSimplex R N) : StdSimplex R N := (f.map g).join
+
+@[simp]
+lemma bind_single (m : M) (g : M → StdSimplex R N) : bind (single m) g = g m := by
+  simp [bind, join]
+
+@[simp]
+lemma bind_const (f : StdSimplex R M) (g : StdSimplex R N) : bind f (fun _ ↦ g) = g := by
+  simp [bind, join]
+
+lemma weights_bind (f : StdSimplex R M) (g : M → StdSimplex R N) (n : N) :
+    (bind f g).weights n = ∑ k ∈ f.support, f.weights k * (g k).weights n := by
+  simp only [bind, join, map, Finsupp.sum_apply]
+  rw [Finsupp.sum_mapDomain_index (by simp) (by simp [add_mul])]
+  simp [Finsupp.sum]
+
+lemma support_subset_support_bind {f : StdSimplex R M} (g : M → StdSimplex R N) {m : M}
+    (hm : m ∈ f.support) : (g m).support ⊆ (bind f g).support := by
+  intro n hn
+  rw [Finsupp.mem_support_iff, weights_bind]
+  refine ne_of_gt ?_
+  have hpos : 0 < f.weights m * (g m).weights n :=
+    mul_pos ((f.nonneg m).lt_of_ne' (by grind)) (((g m).nonneg n).lt_of_ne' (by grind))
+  have hnonneg (k : M) (hk : k ∈ f.support) : 0 ≤ f.weights k * (g k).weights n := by
+    exact mul_nonneg (f.nonneg k) ((g k).nonneg n)
+  exact lt_of_lt_of_le hpos (Finset.single_le_sum hnonneg hm)
+
 end StdSimplex
 
 /--