Drop dependency on OptimalTransport and implement LP problem (#15)

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

.github/workflows/CI.yml

@@ -18,7 +18,7 @@ jobs:
     strategy:
       matrix:
         version:
-          - '1.4'
+          - '1.3'
           - '1'
           - 'nightly'
         os:

Project.toml

@@ -1,15 +1,15 @@
 name = "CalibrationErrorsDistributions"
 uuid = "20087e1a-bb94-462b-b900-33d17a750383"
 authors = ["David Widmann <david.widmann@it.uu.se>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 CalibrationErrors = "33913031-fe46-5864-950f-100836f47845"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-OptimalTransport = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
+MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Tulip = "6dd1b50a-3aae-11e9-10b5-ef983d2400fa"
@@ -19,8 +19,8 @@ CalibrationErrors = "0.5"
 Distances = "0.10.1"
 Distributions = "0.23, 0.24"
 KernelFunctions = "0.8, 0.9"
-OptimalTransport = "0.2"
+MathOptInterface = "0.9"
 PDMats = "0.10, 0.11"
 Reexport = "0.2, 1.0"
 Tulip = "0.7"
-julia = "1.4"
+julia = "1.3"

src/CalibrationErrorsDistributions.jl

@@ -7,14 +7,18 @@ using Reexport
 @reexport using KernelFunctions
 
 using Distances: Distances
-using OptimalTransport: OptimalTransport
+using MathOptInterface: MathOptInterface
 using PDMats: PDMats
 using Tulip: Tulip
 
 using LinearAlgebra: LinearAlgebra
 
 export WassersteinExponentialKernel, MixtureWassersteinExponentialKernel
 
+const MOI = MathOptInterface
+
+include("optimaltransport.jl")
+
 include("distances/types.jl")
 include("distances/bures.jl")
 include("distances/wasserstein.jl")

src/distances/wasserstein.jl

@@ -73,7 +73,7 @@ end
 
 function (s::SqMixtureWasserstein)(a::AbstractMixtureModel, b::AbstractMixtureModel)
     C = Distances.pairwise(SqWasserstein(), components(a), components(b))
-    return OptimalTransport.emd2(probs(a), probs(b), C, s.lpsolver)
+    return sqwasserstein(probs(a), probs(b), C, deepcopy(s.lpsolver))
 end
 
 function (m::MixtureWasserstein)(a::AbstractMixtureModel, b::AbstractMixtureModel)

src/optimaltransport.jl

@@ -0,0 +1,84 @@
+function sqwasserstein(μ, ν, C, optimizer)
+    P = optimal_transport_map(μ, ν, C, optimizer)
+    return LinearAlgebra.dot(P, C)
+end
+
+"""
+    optimal_transport_map(μ, ν, C, optimizer)
+
+Solve the discrete optimal transport problem with source `μ`, target `ν`, and
+cost matrix `C` as a linear programming (LP) problem with the given `optimizer`.
+
+More concretely, this function returns a solution `P` of the LP problem
+```math
+\\begin{aligned}
+\\min_{p} c^T p & \\\\
+\\text{subject to } A_1p &= μ \\\\
+A_2p &= ν \\\\
+0 &\\leq p
+\\end{aligned}
+```
+where
+```math
+\\begin{aligned}
+p &= [P_{1,1},P_{2,1},\\ldots,P_{n,1},P_{2,1},\\ldots,P_{n,m}]^T, \\\\
+c &= [C_{1,1},C_{2,1},\\ldots,C_{n,1},C_{2,1},\\ldots,C_{n,m}]^T, \\\\
+A_1 &= \\begin{bmatrix}
+1_n^T \\otimes I_m 
+\\end{bmatrix}, \\\\
+A_2 &= \\begin{bmatrix}
+I_n \\otimes 1_m^T
+\\end{bmatrix}.
+\\end{aligned}
+```
+
+A possible choice of `optimizer` is `Tulip.Optimizer()` in the `Tulip` package.
+"""
+function optimal_transport_map(μ, ν, C, model::MOI.ModelLike)
+    nμ = length(μ)
+    nν = length(ν)
+    size(C) == (nμ, nν) || error("size of `C` does not match size of `μ` and `ν`")
+    nC = length(C)
+
+    # define variables
+    x = MOI.add_variables(model, nC)
+    xmat = reshape(x, nμ, nν)
+
+    # define objective function
+    T = eltype(C)
+    zero_T = zero(T)
+    MOI.set(
+        model,
+        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
+        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(vec(C), x), zero_T),
+    )
+    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
+
+    # add constraints
+    for xi in x
+        MOI.add_constraint(model, MOI.SingleVariable(xi), MOI.GreaterThan(zero_T))
+    end
+
+    # add constraints for source
+    for (xs, μi) in zip(eachrow(xmat), μ)
+        f = MOI.ScalarAffineFunction(
+            [MOI.ScalarAffineTerm(one(μi), xi) for xi in xs], zero(μi)
+        )
+        MOI.add_constraint(model, f, MOI.EqualTo(μi))
+    end
+
+    # add constraints for target
+    for (xs, νi) in zip(eachcol(xmat), ν)
+        f = MOI.ScalarAffineFunction(
+            [MOI.ScalarAffineTerm(one(νi), xi) for xi in xs], zero(νi)
+        )
+        MOI.add_constraint(model, f, MOI.EqualTo(νi))
+    end
+
+    # compute optimal solution
+    MOI.optimize!(model)
+    p = MOI.get(model, MOI.VariablePrimal(), x)
+    P = reshape(p, nμ, nν)
+
+    return P
+end

test/Project.toml

@@ -2,14 +2,17 @@
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tulip = "6dd1b50a-3aae-11e9-10b5-ef983d2400fa"
 
 [compat]
 Distances = "0.10.1"
 JuliaFormatter = "0.13"
+MathOptInterface = "0.9"
 PDMats = "0.10, 0.11"
 Tulip = "0.7"
 julia = "1.3"

test/optimaltransport.jl

@@ -0,0 +1,31 @@
+@testset "optimaltransport.jl" begin
+    M = 200
+    N = 250
+    μ = rand(M)
+    ν = rand(N)
+    μ ./= sum(μ)
+    ν ./= sum(ν)
+
+    # create random cost matrix
+    C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)
+
+    # compute optimal transport map and squared Wasserstein distance
+    lp = Tulip.Optimizer()
+    P = optimal_transport_map(μ, ν, C, lp)
+    @test size(C) == size(P)
+    @test MOI.get(lp, MOI.TerminationStatus()) == MOI.OPTIMAL
+
+    lp = Tulip.Optimizer()
+    cost = sqwasserstein(μ, ν, C, lp)
+    @test dot(C, P) ≈ cost atol = 1e-5
+    @test MOI.get(lp, MOI.TerminationStatus()) == MOI.OPTIMAL
+
+    # compute optimal transport map and cost with OptimalTransport.jl
+    @static if VERSION >= v"1.4"
+        P_ot = emd(μ, ν, C, Tulip.Optimizer())
+        @test maximum(abs, P .- P_ot) < 1e-2
+
+        cost_ot = emd2(μ, ν, C, Tulip.Optimizer())
+        @test cost ≈ cost_ot atol = 1e-5
+    end
+end

test/runtests.jl

@@ -1,17 +1,39 @@
 using CalibrationErrorsDistributions
 using Distances
 using LinearAlgebra
+using MathOptInterface
 using PDMats
 using Random
 using Test
 
 using CalibrationErrorsDistributions:
-    Wasserstein, SqWasserstein, MixtureWasserstein, SqMixtureWasserstein
+    Wasserstein,
+    SqWasserstein,
+    MixtureWasserstein,
+    SqMixtureWasserstein,
+    sqwasserstein,
+    optimal_transport_map
 using Tulip: Tulip
 
+# only add OptimalTransport on >= Julia 1.4
+# Julia 1.3 is only supported by OptimalTransport 0.1 which requires Python
+@static if VERSION >= v"1.4"
+    using Pkg: Pkg
+    Pkg.add(;
+        name="OptimalTransport", uuid="7e02d93a-ae51-4f58-b602-d97af76e3b33", version="0.2"
+    )
+    using OptimalTransport
+end
+
 Random.seed!(1234)
 
+const MOI = MathOptInterface
+
 @testset "CalibrationErrorsDistributions" begin
+    @testset "optimal transport" begin
+        include("optimaltransport.jl")
+    end
+
     @testset "distances" begin
         @testset "Bures metric" begin
             include("distances/bures.jl")
@@ -21,7 +43,7 @@ Random.seed!(1234)
         end
     end
 
-    @testset "Kernels" begin
+    @testset "kernels" begin
         include("kernels.jl")
     end