Merge branch 'master' into ot_reg

Author: zsteve

.github/workflows/CI.yml

@@ -14,6 +14,7 @@ jobs:
     strategy:
       matrix:
         version:
+          - '1.0'
           - '1'
           - 'nightly'
         os:

Project.toml

@@ -19,11 +19,10 @@ Requires = "1.1"
 julia = "1"
 
 [extras]
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tulip = "6dd1b50a-3aae-11e9-10b5-ef983d2400fa"
 
 [targets]
-test = ["Tulip", "CUDA", "PyCall", "Random", "Test"]
+test = ["Tulip", "PyCall", "Random", "Test"]

src/OptimalTransport.jl

@@ -66,17 +66,17 @@ function emd(μ, ν, C, model::MOI.ModelLike)
     end
 
     # add constraints for source
-    for (xs, μi) in zip(eachrow(xmat), μ)
+    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 xs], zero(μi)
+            [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 (xs, νi) in zip(eachcol(xmat), ν)
+    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 xs], zero(νi)
+            [MOI.ScalarAffineTerm(one(νi), xi) for xi in view(xmat, :, i)], zero(νi)
         )
         MOI.add_constraint(model, f, MOI.EqualTo(νi))
     end
@@ -415,16 +415,11 @@ function sinkhorn_stabilized(
     absorb_tol=1e3,
     max_iter=1000,
     tol=1e-9,
-    alpha=nothing,
-    beta=nothing,
+    alpha=zeros(size(mu)),
+    beta=zeros(size(nu)),
     return_duals=false,
     verbose=false,
 )
-    if isnothing(alpha) || isnothing(beta)
-        alpha = zeros(size(mu))
-        beta = zeros(size(nu))
-    end
-
     u = ones(size(mu))
     v = ones(size(nu))
     K = getK(C, alpha, beta, eps, mu, nu)

test/runtests.jl

@@ -1,6 +1,5 @@
 using OptimalTransport
 
-using CUDA
 using Distances
 using PyCall
 using Tulip
@@ -108,27 +107,6 @@ end
         @test c_pot isa Float64 # POT does not respect input types
         @test c ≈ c_pot atol = Base.eps(Float32)
     end
-
-    # computation on the GPU
-    if CUDA.functional()
-        @testset "CUDA" begin
-            # create two uniform histograms
-            μ = CUDA.fill(Float32(1 / M), M)
-            ν = CUDA.fill(Float32(1 / N), N)
-
-            # create random cost matrix
-            C = abs2.(CUDA.rand(M) .- CUDA.rand(1, N))
-
-            # compute optimal transport map
-            eps = 0.01f0
-            γ = sinkhorn(μ, ν, C, eps)
-            @test γ isa CuArray{Float32}
-
-            # compute optimal transport cost
-            c = sinkhorn2(μ, ν, C, eps)
-            @test c isa Float32
-        end
-    end
 end
 
 @testset "unbalanced transport" begin
@@ -210,7 +188,7 @@ end
 @testset "sinkhorn barycenter" begin
     @testset "example" begin
         # set up support
-        support = range(-1, 1; length=250)
+        support = range(-1; stop=1, length=250)
         μ1 = exp.(-(support .+ 0.5) .^ 2 ./ 0.1^2)
         μ1 ./= sum(μ1)
         μ2 = exp.(-(support .- 0.5) .^ 2 ./ 0.1^2)