Add implementation of discrete 1D optimal transport (#88)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: devmotion

Project.toml

@@ -1,7 +1,7 @@
 name = "OptimalTransport"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
 authors = ["zsteve <[email protected]>"]
-version = "0.3.5"
+version = "0.3.6"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
@@ -12,6 +12,7 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 Distances = "0.9.0, 0.10"
@@ -20,6 +21,7 @@ IterativeSolvers = "0.8.4, 0.9"
 LogExpFunctions = "0.2"
 MathOptInterface = "0.9"
 QuadGK = "2"
+StatsBase = "0.33.8"
 julia = "1"
 
 [extras]

src/OptimalTransport.jl

@@ -11,6 +11,7 @@ using LogExpFunctions: LogExpFunctions
 using MathOptInterface
 using Distributions
 using QuadGK
+using StatsBase: StatsBase
 
 export sinkhorn, sinkhorn2
 export emd, emd2
@@ -21,110 +22,14 @@ export ot_cost, ot_plan
 
 const MOI = MathOptInterface
 
+include("exact.jl")
+
 dot_matwise(x::AbstractMatrix, y::AbstractMatrix) = dot(x, y)
 function dot_matwise(x::AbstractArray, y::AbstractMatrix)
     xmat = reshape(x, size(x, 1) * size(x, 2), :)
     return reshape(reshape(y, 1, :) * xmat, size(x)[3:end])
 end
 
-"""
-    emd(μ, ν, C, optimizer)
-
-Compute the optimal transport plan `γ` for the Monge-Kantorovich problem with source
-histogram `μ`, target histogram `ν`, and cost matrix `C` of size `(length(μ), length(ν))`
-which solves
-```math
-\\inf_{γ ∈ Π(μ, ν)} \\langle γ, C \\rangle.
-```
-
-The corresponding linear programming problem is solved with the user-provided `optimizer`.
-Possible choices are `Tulip.Optimizer()` and `Clp.Optimizer()` in the `Tulip` and `Clp`
-packages, respectively.
-"""
-function emd(μ, ν, C, model::MOI.ModelLike)
-    # check size of cost matrix
-    nμ = length(μ)
-    nν = length(ν)
-    size(C) == (nμ, nν) || error("cost matrix `C` must be of size `(length(μ), length(ν))`")
-    nC = length(C)
-
-    # define variables
-    x = MOI.add_variables(model, nC)
-    xmat = reshape(x, nμ, nν)
-
-    # define objective function
-    T = float(eltype(C))
-    zero_T = zero(T)
-    MOI.set(
-        model,
-        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
-        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(float.(vec(C)), x), zero_T),
-    )
-    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
-
-    # add non-negativity constraints
-    for xi in x
-        MOI.add_constraint(model, MOI.SingleVariable(xi), MOI.GreaterThan(zero_T))
-    end
-
-    # add constraints for source
-    for (i, μi) in zip(axes(xmat, 1), μ) # eachrow(xmat) is not available on Julia 1.0
-        f = MOI.ScalarAffineFunction(
-            [MOI.ScalarAffineTerm(one(μi), xi) for xi in view(xmat, i, :)], zero(μi)
-        )
-        MOI.add_constraint(model, f, MOI.EqualTo(μi))
-    end
-
-    # add constraints for target
-    for (i, νi) in zip(axes(xmat, 2), ν) # eachcol(xmat) is not available on Julia 1.0
-        f = MOI.ScalarAffineFunction(
-            [MOI.ScalarAffineTerm(one(νi), xi) for xi in view(xmat, :, i)], zero(νi)
-        )
-        MOI.add_constraint(model, f, MOI.EqualTo(νi))
-    end
-
-    # compute optimal solution
-    MOI.optimize!(model)
-    status = MOI.get(model, MOI.TerminationStatus())
-    status === MOI.OPTIMAL || error("failed to compute optimal transport plan: ", status)
-    p = MOI.get(model, MOI.VariablePrimal(), x)
-    γ = reshape(p, nμ, nν)
-
-    return γ
-end
-
-"""
-    emd2(μ, ν, C, optimizer; plan=nothing)
-
-Compute the optimal transport cost (a scalar) for the Monge-Kantorovich problem with source
-histogram `μ`, target histogram `ν`, and cost matrix `C` of size `(length(μ), length(ν))`
-which is given by
-```math
-\\inf_{γ ∈ Π(μ, ν)} \\langle γ, C \\rangle.
-```
-
-The corresponding linear programming problem is solved with the user-provided `optimizer`.
-Possible choices are `Tulip.Optimizer()` and `Clp.Optimizer()` in the `Tulip` and `Clp`
-packages, respectively.
-
-A pre-computed optimal transport `plan` may be provided.
-"""
-function emd2(μ, ν, C, optimizer; plan=nothing)
-    γ = if plan === nothing
-        # compute optimal transport plan
-        emd(μ, ν, C, optimizer)
-    else
-        # check dimensions
-        size(C) == (length(μ), length(ν)) ||
-            error("cost matrix `C` must be of size `(length(μ), length(ν))`")
-        size(plan) == size(C) || error(
-            "optimal transport plan `plan` and cost matrix `C` must be of the same size",
-        )
-        plan
-    end
-    return dot(γ, C)
-end
-
 """
     sinkhorn_gibbs(
         μ, ν, K; atol=0, rtol=atol > 0 ? 0 : √eps, check_convergence=10, maxiter=1_000
@@ -879,62 +784,4 @@ 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