AutoSamplingTheory.TechnicalLemmas.Measure.QuadraticOptimalMap

- Joint law induced by a measurable transport map.

noncomputable def graphCoupling
    {E F : Type*} [MeasurableSpace E] [MeasurableSpace F]
    (T : E → F) (mu : Measure E) : Measure (E × F) :=
  Measure.map (fun x => (x, T x)) mu

/-- A graph pushforward has first marginal `mu` and second marginal `nu` as soon
as `T` is measurable and pushes `mu` to `nu`. -/

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:29published source at 0e31a3cda412

- A graph pushforward has first marginal `mu` and second marginal `nu` as soon as `T` is measurable and pushes `mu` to `nu`.

theorem isCoupling_graphCoupling
    {E F : Type*} [MeasurableSpace E] [MeasurableSpace F]
    {T : E → F} {mu : Measure E} {nu : Measure F}
    (hT : Measurable T) (hmap : Measure.map T mu = nu) :
    IsCoupling (graphCoupling T mu) mu nu := by
  have hGraph : Measurable (fun x : E => (x, T x)) :=
    measurable_id.prodMk hT
  constructor
  · rw [Measure.fst, graphCoupling, Measure.map_map measurable_fst hGraph]
    simpa [Function.comp_def]
  · rw [Measure.snd, graphCoupling, Measure.map_map measurable_snd hGraph]
    simpa [Function.comp_def] using hmap

/-- Equality of two measurable graph couplings determines the underlying maps
almost everywhere with respect to their common first marginal. -/

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:36published source at 0e31a3cda412

- Equality of two measurable graph couplings determines the underlying maps almost everywhere with respect to their common first marginal.

theorem ae_eq_of_graphCoupling_eq
    {E F : Type*} [MeasurableSpace E] [MeasurableSpace F] [MeasurableEq F]
    {T S : E → F} {mu : Measure E}
    (hT : Measurable T) (hS : Measurable S)
    (hEq : graphCoupling T mu = graphCoupling S mu) :
    T =ᵐ[mu] S := by
  have hGraphT : Measurable (fun x : E => (x, T x)) :=
    measurable_id.prodMk hT
  have hGraphS : Measurable (fun x : E => (x, S x)) :=
    measurable_id.prodMk hS
  have hGraphSet : MeasurableSet {z : E × F | z.2 = T z.1} :=
    measurableSet_eq_fun measurable_snd (hT.comp measurable_fst)
  have hOnT : ∀ᵐ z ∂graphCoupling T mu, z.2 = T z.1 := by
    change ∀ᵐ z ∂Measure.map (fun x : E => (x, T x)) mu, z.2 = T z.1
    exact (ae_map_iff hGraphT.aemeasurable hGraphSet).2 <| by
      filter_upwards with x
      rfl
  have hOnS : ∀ᵐ z ∂graphCoupling S mu, z.2 = T z.1 := by
    rw [← hEq]
    exact hOnT
  have hPull : ∀ᵐ x ∂mu, S x = T x := by
    have := ae_of_ae_map hGraphS.aemeasurable (by
      simpa [graphCoupling] using hOnS)
    simpa using this
  exact hPull.mono fun _ hx => hx.symm

/-- A quadratic-optimal transport map is a measurable map whose pushforward is
the prescribed target and whose induced graph coupling attains the quadratic
Kantorovich optimum. -/

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:51published source at 0e31a3cda412

- A quadratic-optimal transport map is a measurable map whose pushforward is the prescribed target and whose induced graph coupling attains the quadratic Kantorovich optimum.

def IsQuadraticOptimalMap
    {E : Type*} [NormedAddCommGroup E] [MeasurableSpace E]
    (T : E → E) (mu nu : Measure E) : Prop :=
  Measurable T ∧
    Measure.map T mu = nu ∧
    IsQuadraticOptimalCoupling (graphCoupling T mu) mu nu

/-- Expansion of the optimal-map interface. -/

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:80published source at 0e31a3cda412

- Expansion of the optimal-map interface.

theorem isQuadraticOptimalMap_iff
    {E : Type*} [NormedAddCommGroup E] [MeasurableSpace E]
    (T : E → E) (mu nu : Measure E) :
    IsQuadraticOptimalMap T mu nu ↔
      Measurable T ∧
        Measure.map T mu = nu ∧
        IsQuadraticOptimalCoupling (graphCoupling T mu) mu nu :=
  Iff.rfl

/-- Package a measurable pushforward map as optimal once its graph coupling is
identified with an already optimal coupling. -/

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:88published source at 0e31a3cda412

- Package a measurable pushforward map as optimal once its graph coupling is identified with an already optimal coupling.

theorem isQuadraticOptimalMap_of_eq_graphCoupling
    {E : Type*} [NormedAddCommGroup E] [MeasurableSpace E]
    {T : E → E} {mu nu : Measure E} {gamma : Measure (E × E)}
    (hT : Measurable T)
    (hmap : Measure.map T mu = nu)
    (hopt : IsQuadraticOptimalCoupling gamma mu nu)
    (hgamma : gamma = graphCoupling T mu) :
    IsQuadraticOptimalMap T mu nu := by
  refine ⟨hT, hmap, ?_⟩
  rw [← hgamma]
  exact hopt

/-- Forgetting the map packaging recovers the optimality of its graph coupling. -/

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:99published source at 0e31a3cda412

- Forgetting the map packaging recovers the optimality of its graph coupling.

theorem isQuadraticOptimalCoupling_graphCoupling
    {E : Type*} [NormedAddCommGroup E] [MeasurableSpace E]
    {T : E → E} {mu nu : Measure E}
    (hT : IsQuadraticOptimalMap T mu nu) :
    IsQuadraticOptimalCoupling (graphCoupling T mu) mu nu :=
  hT.2.2

end

end QuadraticOptimalMap
end Measure
end TechnicalLemmas
end AutoSamplingTheory

AutoSamplingTheory/TechnicalLemmas/Measure/QuadraticOptimalMap.lean:112published source at 0e31a3cda412