Merge pull request #43 from zsteve/dw/simplex

Simplify implementation of LP formulation

Author: zsteve

Project.toml

@@ -1,23 +1,19 @@
 name = "OptimalTransport"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
 authors = ["zsteve <[email protected]>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
-FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
-LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Distances = "0.9.0, 0.10"
-FillArrays = "0.9, 0.10, 0.11"
 IterativeSolvers = "0.8.4, 0.9"
-LazyArrays = "0.18, 0.19, 0.20, 0.21"
 MathOptInterface = "0.9"
 Requires = "1.1"
 julia = "1"

src/OptimalTransport.jl

@@ -8,8 +8,6 @@ using Distances
 using LinearAlgebra
 using IterativeSolvers, SparseArrays
 using Requires
-using FillArrays
-using LazyArrays
 using MathOptInterface
 
 export sinkhorn, sinkhorn2
@@ -27,8 +25,6 @@ function __init__()
     end
 end
 
-include("simplex.jl")
-
 """
     emd(μ, ν, C, optimizer)
 
@@ -43,16 +39,54 @@ The corresponding linear programming problem is solved with the user-provided `o
 Possible choices are `Tulip.Optimizer()` and `Clp.Optimizer()` in the `Tulip` and `Clp`
 packages, respectively.
 """
-function emd(μ, ν, C, optimizer)
+function emd(μ, ν, C, model::MOI.ModelLike)
     # check size of cost matrix
-    m = length(μ)
-    n = length(ν)
-    size(C) == (m, n) || error("cost matrix `C` must be of size `(length(μ), length(ν))`")
-
-    # solve linear programming problem
-    c, A, b = toSimplexFormat(μ, ν, C)
-    p = solveLP(c, A, b, optimizer)
-    γ = reshape(p, m, n)
+    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 = 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 non-negativity 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)
+    status = MOI.get(model, MOI.TerminationStatus())
+    status === MOI.OPTIMAL || error("failed to compute optimal transport map: ", status)
+    p = MOI.get(model, MOI.VariablePrimal(), x)
+    γ = reshape(p, nμ, nν)
 
     return γ
 end