1d Optimal Transport (#45)

Co-authored-by: David Widmann <[email protected]>
Co-authored-by: David Widmann <[email protected]>

Author: davibarreira

Project.toml

@@ -5,17 +5,21 @@ version = "0.3.2"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Distances = "0.9.0, 0.10"
+Distributions = "0.25"
 IterativeSolvers = "0.8.4, 0.9"
 LogExpFunctions = "0.2"
 MathOptInterface = "0.9"
+QuadGK = "2"
 julia = "1"
 
 [extras]

docs/src/index.md

@@ -2,9 +2,12 @@
 
 
 ## Exact optimal transport (Kantorovich) problem
+
 ```@docs
 emd
 emd2
+ot_plan
+ot_cost
 ```
 
 ## Entropically regularised optimal transport

src/OptimalTransport.jl

@@ -9,12 +9,15 @@ using LinearAlgebra
 using IterativeSolvers, SparseArrays
 using LogExpFunctions: LogExpFunctions
 using MathOptInterface
+using Distributions
+using QuadGK
 
 export sinkhorn, sinkhorn2
 export emd, emd2
 export sinkhorn_stabilized, sinkhorn_stabilized_epsscaling, sinkhorn_barycenter
 export sinkhorn_unbalanced, sinkhorn_unbalanced2
 export quadreg
+export ot_cost, ot_plan
 
 const MOI = MathOptInterface
 
@@ -649,4 +652,62 @@ function quadreg(mu, nu, C, ϵ; θ=0.1, tol=1e-5, maxiter=50, κ=0.5, δ=1e-5)
     return sparse(γ')
 end
 
+"""
+    ot_cost(
+        c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution; plan=nothing
+    )
+
+Compute the optimal transport cost for the Monge-Kantorovich problem with univariate
+distributions `μ` and `ν` as source and target marginals and cost function `c` of
+the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function.
+
+In this setting, the optimal transport cost can be computed as
+```math
+\\int_0^1 c(F_\\mu^{-1}(x), F_\\nu^{-1}(x)) \\mathrm{d}x
+```
+where ``F_\\mu^{-1}`` and ``F_\\nu^{-1}`` are the quantile functions of `μ` and `ν`,
+respectively.
+
+A pre-computed optimal transport `plan` may be provided.
+
+See also: [`ot_plan`](@ref), [`emd2`](@ref)
+"""
+function ot_cost(
+    c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution; plan=nothing
+)
+    cost, _ = if plan === nothing
+        quadgk(0, 1) do q
+            return c(quantile(μ, q), quantile(ν, q))
+        end
+    else
+        quadgk(0, 1) do q
+            x = quantile(μ, q)
+            return c(x, plan(x))
+        end
+    end
+    return cost
+end
+
+"""
+    ot_plan(c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution)
+
+Compute the optimal transport plan for the Monge-Kantorovich problem with univariate
+distributions `μ` and `ν` as source and target marginals and cost function `c` of
+the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function.
+
+In this setting, the optimal transport plan is the Monge map
+```math
+T = F_\\nu^{-1} \\circ F_\\mu
+```
+where ``F_\\mu`` is the cumulative distribution function of `μ` and ``F_\\nu^{-1}`` is the
+quantile function of `ν`.
+
+See also: [`ot_cost`](@ref), [`emd`](@ref)
+"""
+function ot_plan(c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution)
+    # Use T instead of γ to indicate that this is a Monge map.
+    T(x) = quantile(ν, cdf(μ, x))
+    return T
+end
+
 end

test/runtests.jl

@@ -4,6 +4,7 @@ using Distances
 using PythonOT: PythonOT
 using Tulip
 using MathOptInterface
+using Distributions
 using SparseArrays
 
 using LinearAlgebra
@@ -64,6 +65,51 @@ Random.seed!(100)
     end
 end
 
+@testset "1D Optimal Transport for Convex Cost" begin
+    @testset "continuous distributions" begin
+        # two normal distributions (has analytical solution)
+        μ = Normal(randn(), rand())
+        ν = Normal(randn(), rand())
+
+        # compute OT plan
+        γ = ot_plan(sqeuclidean, μ, ν)
+        x = randn()
+        @test γ(x) ≈ quantile(ν, cdf(μ, x))
+
+        # compute OT cost
+        c = ot_cost(sqeuclidean, μ, ν)
+        @test c ≈ (mean(μ) - mean(ν))^2 + (std(μ) - std(ν))^2
+
+        # do not use ν to ensure that the provided plan is used
+        @test ot_cost(sqeuclidean, μ, Normal(randn(), rand()); plan=γ) ≈ c
+    end
+
+    @testset "semidiscrete case" begin
+        μ = Normal(randn(), rand())
+        νprobs = rand(30)
+        νprobs ./= sum(νprobs)
+        ν = Categorical(νprobs)
+
+        # compute OT plan
+        γ = ot_plan(euclidean, μ, ν)
+        x = randn()
+        @test γ(x) ≈ quantile(ν, cdf(μ, x))
+
+        # compute OT cost, without and with provided plan
+        # do not use ν in the second case to ensure that the provided plan is used
+        c = ot_cost(euclidean, μ, ν)
+        @test ot_cost(euclidean, μ, Categorical(reverse(νprobs)); plan=γ) ≈ c
+
+        # check that OT cost is consistent with OT cost of a discretization
+        m = 500
+        xs = rand(μ, m)
+        μdiscrete = fill(1 / m, m)
+        C = pairwise(Euclidean(), xs', (1:length(νprobs))'; dims=2)
+        c2 = emd2(μdiscrete, νprobs, C, Tulip.Optimizer())
+        @test c2 ≈ c rtol = 1e-1
+    end
+end
+
 @testset "entropically regularized transport" begin
     M = 250
     N = 200