Fix emd with integer-valued cost matrices (#72)

Author: devmotion

Project.toml

@@ -1,7 +1,7 @@
 name = "OptimalTransport"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
 authors = ["zsteve <[email protected]>"]
-version = "0.3.0"
+version = "0.3.1"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"

src/OptimalTransport.jl

@@ -43,12 +43,12 @@ function emd(μ, ν, C, model::MOI.ModelLike)
     xmat = reshape(x, nμ, nν)
 
     # define objective function
-    T = eltype(C)
+    T = float(eltype(C))
     zero_T = zero(T)
     MOI.set(
         model,
         MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
-        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(vec(C), x), zero_T),
+        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(float.(vec(C)), x), zero_T),
     )
     MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
 

test/runtests.jl

@@ -23,29 +23,45 @@ Random.seed!(100)
     μ ./= sum(μ)
     ν ./= sum(ν)
 
-    # create random cost matrix
-    C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)
-
-    # compute optimal transport map and cost with POT
-    pot_P = POT.emd(μ, ν, C)
-    pot_cost = POT.emd2(μ, ν, C)
-
-    # compute optimal transport map and cost with Tulip
-    lp = Tulip.Optimizer()
-    P = emd(μ, ν, C, lp)
-    @test size(C) == size(P)
-    @test MOI.get(lp, MOI.TerminationStatus()) == MOI.OPTIMAL
-    @test maximum(abs, P .- pot_P) < 1e-2
-
-    lp = Tulip.Optimizer()
-    cost = emd2(μ, ν, C, lp)
-    @test dot(C, P) ≈ cost atol = 1e-5
-    @test MOI.get(lp, MOI.TerminationStatus()) == MOI.OPTIMAL
-    @test cost ≈ pot_cost atol = 1e-5
-
-    # ensure that provided map is used
-    cost2 = emd2(similar(μ), similar(ν), C, lp; plan=P)
-    @test cost2 ≈ cost
+    @testset "example" begin
+        # create random cost matrix
+        C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)
+
+        # compute optimal transport map and cost with POT
+        pot_P = POT.emd(μ, ν, C)
+        pot_cost = POT.emd2(μ, ν, C)
+
+        # compute optimal transport map and cost with Tulip
+        lp = Tulip.Optimizer()
+        P = emd(μ, ν, C, lp)
+        @test size(C) == size(P)
+        @test MOI.get(lp, MOI.TerminationStatus()) == MOI.OPTIMAL
+        @test maximum(abs, P .- pot_P) < 1e-2
+
+        lp = Tulip.Optimizer()
+        cost = emd2(μ, ν, C, lp)
+        @test dot(C, P) ≈ cost atol = 1e-5
+        @test MOI.get(lp, MOI.TerminationStatus()) == MOI.OPTIMAL
+        @test cost ≈ pot_cost atol = 1e-5
+    end
+
+    @testset "pre-computed plan" begin
+        # create random cost matrix
+        C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)
+
+        # compute optimal transport map
+        P = emd(μ, ν, C, Tulip.Optimizer())
+
+        # do not use μ and ν to ensure that provided map is used
+        cost = emd2(similar(μ), similar(ν), C, Tulip.Optimizer(); plan=P)
+        @test cost ≈ emd2(μ, ν, C, Tulip.Optimizer())
+    end
+
+    # https://github.com/JuliaOptimalTransport/OptimalTransport.jl/issues/71
+    @testset "cost matrix with integers" begin
+        C = pairwise(SqEuclidean(), rand(1:10, 1, M), rand(1:10, 1, N); dims=2)
+        emd2(μ, ν, C, Tulip.Optimizer())
+    end
 end
 
 @testset "entropically regularized transport" begin