Finite transport plans as matrices

On finite spaces, a probability mass function on a product is the same data as a nonnegative matrix of total mass one. Prescribing its two marginals says exactly that the row and column sums of this matrix are the prescribed probability vectors.

This file packages that correspondence. TauCeti.TransportMatrix μ ν is the type of ℝ≥0∞-valued matrices with row sums μ and column sums ν, and TauCeti.transportMatrixEquiv identifies it with the subtype of product PMFs whose two pushforwards are μ and ν. The codomain ℝ≥0∞ makes nonnegativity intrinsic, rather than a separate side condition.

The entries are also read as a real-valued function TauCeti.TransportMatrix.toRealFun on the product, whose row and column sums are the real marginal masses, and the total TauCeti.TransportMatrix.cost of a matrix against a real cost function is recorded here. Both are basic transportation-matrix API rather than duality results, and the cost is allowed to be negative because the finite transport problem is a linear program.

This is the finite transportation-matrix acceptance case of Layer 0 of the optimal-transport roadmap. It is also the representation used by later finite primal and dual problems.

structure TauCeti.TransportMatrix {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) : Type (max u v)

A finite transportation matrix with prescribed row distribution μ and column distribution ν. Nonnegativity is built into the ℝ≥0∞-valued matrix.

• matrix : Matrix ι κ ENNReal

  The mass assigned to each source-target pair.

• row_sum (i : ι) : ∑ j : κ, self.matrix i j = μ i

  Every row has the mass prescribed by the source distribution.

• col_sum (j : κ) : ∑ i : ι, self.matrix i j = ν j

  Every column has the mass prescribed by the target distribution.

@[reducible, inline]

abbrev TauCeti.RealPlans {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) : Set (ι × κ → ℝ)

Nonnegative real arrays with the row and column sums prescribed by μ and ν.

Equations

• TauCeti.RealPlans μ ν = {f : ι × κ → ℝ | (∀ (q : ι × κ), 0 ≤ f q) ∧ (∀ (i : ι), ∑ j : κ, f (i, j) = (μ i).toReal) ∧ ∀ (j : κ), ∑ i : ι, f (i, j) = (ν j).toReal}

@[instance_reducible]

instance TauCeti.TransportMatrix.instCoeFunForallForallENNReal {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} : CoeFun (TransportMatrix μ ν) fun (x : TransportMatrix μ ν) => ι → κ → ENNReal

Coerce a transportation matrix to its entries, allowing the notation A i j.

Equations

• TauCeti.TransportMatrix.instCoeFunForallForallENNReal = { coe := TauCeti.TransportMatrix.matrix }

theorem TauCeti.TransportMatrix.ext {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} {A B : TransportMatrix μ ν} (h : ∀ (i : ι) (j : κ), A.matrix i j = B.matrix i j) : A = B

Transportation matrices are equal when all their entries are equal.

theorem TauCeti.TransportMatrix.ext_iff {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} {A B : TransportMatrix μ ν} : A = B ↔ ∀ (i : ι) (j : κ), A.matrix i j = B.matrix i j

noncomputable def TauCeti.TransportMatrix.independent {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) : TransportMatrix μ ν

The independent transportation matrix, whose entries are products of marginal masses.

Equations

• TauCeti.TransportMatrix.independent μ ν = { matrix := fun (i : ι) (j : κ) => μ i * ν j, row_sum := ⋯, col_sum := ⋯ }

@[simp]

theorem TauCeti.TransportMatrix.independent_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) (i : ι) (j : κ) : (independent μ ν).matrix i j = μ i * ν j

def TauCeti.TransportMatrix.toPMF {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : PMF (ι × κ)

A transportation matrix gives a probability mass function on the product by reading its entries as point masses.

Equations

• A.toPMF = PMF.ofFintype (fun (p : ι × κ) => A.matrix p.1 p.2) ⋯

@[simp]

theorem TauCeti.TransportMatrix.toPMF_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (p : ι × κ) : A.toPMF p = A.matrix p.1 p.2

theorem TauCeti.TransportMatrix.apply_ne_top {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (i : ι) (j : κ) : A.matrix i j ≠ ⊤

Every entry of a transportation matrix is finite: it is bounded by a marginal mass.

theorem TauCeti.TransportMatrix.apply_le_row {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (i : ι) (j : κ) : A.matrix i j ≤ μ i

Every entry of a transportation matrix is at most its row mass.

theorem TauCeti.TransportMatrix.apply_le_col {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (i : ι) (j : κ) : A.matrix i j ≤ ν j

Every entry of a transportation matrix is at most its column mass.

def TauCeti.TransportMatrix.toRealFun {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (q : ι × κ) : ℝ

The entries of a transportation matrix, read as a real-valued function on the product.

Equations

• A.toRealFun q = (A.matrix q.1 q.2).toReal

@[simp]

theorem TauCeti.TransportMatrix.toRealFun_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (q : ι × κ) : A.toRealFun q = (A.matrix q.1 q.2).toReal

The defining formula for the real-valued entries.

theorem TauCeti.TransportMatrix.toRealFun_nonneg {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (q : ι × κ) : 0 ≤ A.toRealFun q

The real-valued entries of a transportation matrix are nonnegative.

theorem TauCeti.TransportMatrix.sum_toRealFun_row {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (i : ι) : ∑ j : κ, A.toRealFun (i, j) = (μ i).toReal

The real-valued row sums are the masses of the source distribution.

theorem TauCeti.TransportMatrix.sum_toRealFun_col {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) (j : κ) : ∑ i : ι, A.toRealFun (i, j) = (ν j).toReal

The real-valued column sums are the masses of the target distribution.

theorem TauCeti.TransportMatrix.toRealFun_mem_realPlans {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : A.toRealFun ∈ RealPlans μ ν

The real entries of a transportation matrix form a real plan.

def TauCeti.TransportMatrix.ofRealFun {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} {f : ι × κ → ℝ} (hf : f ∈ RealPlans μ ν) : TransportMatrix μ ν

Convert a nonnegative real plan with prescribed marginals to a transportation matrix.

Equations

• TauCeti.TransportMatrix.ofRealFun hf = { matrix := fun (i : ι) (j : κ) => ENNReal.ofReal (f (i, j)), row_sum := ⋯, col_sum := ⋯ }

@[simp]

theorem TauCeti.TransportMatrix.ofRealFun_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} {f : ι × κ → ℝ} (hf : f ∈ RealPlans μ ν) (i : ι) (j : κ) : (ofRealFun hf).matrix i j = ENNReal.ofReal (f (i, j))

The entries of a matrix converted from a real plan.

@[simp]

theorem TauCeti.TransportMatrix.toRealFun_ofRealFun {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} {f : ι × κ → ℝ} (hf : f ∈ RealPlans μ ν) : (ofRealFun hf).toRealFun = f

Converting a real plan to a matrix preserves every real entry.

@[simp]

theorem TauCeti.TransportMatrix.ofRealFun_toRealFun {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : ofRealFun ⋯ = A

Converting the real entries of a transportation matrix back recovers the matrix.

theorem TauCeti.TransportMatrix.sum_toRealFun {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : ∑ p : ι × κ, A.toRealFun p = 1

The real-valued entries of a transportation matrix have total mass one.

def TauCeti.TransportMatrix.cost {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (c : ι × κ → ℝ) (A : TransportMatrix μ ν) : ℝ

The total cost of a finite transportation matrix. The cost function is real-valued, so negative costs are allowed.

Equations

• TauCeti.TransportMatrix.cost c A = ∑ q : ι × κ, c q * A.toRealFun q

theorem TauCeti.TransportMatrix.cost_def {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (c : ι × κ → ℝ) (A : TransportMatrix μ ν) : cost c A = ∑ q : ι × κ, c q * A.toRealFun q

The defining formula for the cost. The body of the definition is not exposed, so this is the lemma downstream modules should rewrite with.

theorem TauCeti.TransportMatrix.cost_ofRealFun {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (c : ι × κ → ℝ) {f : ι × κ → ℝ} (hf : f ∈ RealPlans μ ν) : cost c (ofRealFun hf) = ∑ q : ι × κ, c q * f q

The cost of a matrix converted from a real plan is the cost of that plan.

@[simp]

theorem TauCeti.TransportMatrix.cost_add_const {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (c : ι × κ → ℝ) (A : TransportMatrix μ ν) (a : ℝ) : cost (fun (p : ι × κ) => c p + a) A = cost c A + a

Adding a constant to every entry of a cost adds that constant to the matrix cost.

@[simp]

theorem TauCeti.TransportMatrix.map_fst_toPMF {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : PMF.map Prod.fst A.toPMF = μ

The PMF associated to a transportation matrix has the prescribed first marginal.

@[simp]

theorem TauCeti.TransportMatrix.map_snd_toPMF {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : PMF.map Prod.snd A.toPMF = ν

The PMF associated to a transportation matrix has the prescribed second marginal.

def TauCeti.TransportMatrix.ofPMF {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (π : PMF (ι × κ)) (hμ : PMF.map Prod.fst π = μ) (hν : PMF.map Prod.snd π = ν) : TransportMatrix μ ν

The transportation matrix obtained by recording the point masses of a finite product PMF with prescribed marginals.

Equations

• TauCeti.TransportMatrix.ofPMF π hμ hν = { matrix := fun (i : ι) (j : κ) => π (i, j), row_sum := ⋯, col_sum := ⋯ }

@[simp]

theorem TauCeti.TransportMatrix.ofPMF_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (π : PMF (ι × κ)) (hμ : PMF.map Prod.fst π = μ) (hν : PMF.map Prod.snd π = ν) (i : ι) (j : κ) : (ofPMF π hμ hν).matrix i j = π (i, j)

@[simp]

theorem TauCeti.TransportMatrix.toPMF_ofPMF {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (π : PMF (ι × κ)) (hμ : PMF.map Prod.fst π = μ) (hν : PMF.map Prod.snd π = ν) : (ofPMF π hμ hν).toPMF = π

@[simp]

theorem TauCeti.TransportMatrix.ofPMF_toPMF {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] {μ : PMF ι} {ν : PMF κ} (A : TransportMatrix μ ν) : ofPMF A.toPMF ⋯ ⋯ = A

def TauCeti.transportMatrixEquiv {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) : { π : PMF (ι × κ) // PMF.map Prod.fst π = μ ∧ PMF.map Prod.snd π = ν } ≃ TransportMatrix μ ν

Finite product PMFs with prescribed marginals are equivalent to nonnegative transportation matrices with the corresponding row and column sums.

Equations

• One or more equations did not get rendered due to their size.

@[simp]

theorem TauCeti.transportMatrixEquiv_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) (π : { π : PMF (ι × κ) // PMF.map Prod.fst π = μ ∧ PMF.map Prod.snd π = ν }) (i : ι) (j : κ) : ((transportMatrixEquiv μ ν) π).matrix i j = ↑π (i, j)

@[simp]

theorem TauCeti.transportMatrixEquiv_symm_apply {ι : Type u} {κ : Type v} [Fintype ι] [Fintype κ] (μ : PMF ι) (ν : PMF κ) (A : TransportMatrix μ ν) (p : ι × κ) : ↑((transportMatrixEquiv μ ν).symm A) p = A.matrix p.1 p.2