Entropy-regularised Gromov-Wasserstein

Project.toml

@@ -1,7 +1,7 @@
 name = "OptimalTransport"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
 authors = ["zsteve <stephenz@student.unimelb.edu.au>"]
-version = "0.3.19"
+version = "0.3.20"
 
 [deps]
 ExactOptimalTransport = "24df6009-d856-477c-ac5c-91f668376b31"

src/#gromov.jl#

@@ -0,0 +1,57 @@
+# Gromov-Wasserstein solver
+
+abstract type EntropicGromovWasserstein end
+
+struct EntropicGromovWassersteinGibbs <: EntropicGromovWasserstein 
+    alg_step::Sinkhorn
+end
+
+function entropic_gromov_wasserstein(μ::AbstractVector, ν::AbstractVector, Cμ::AbstractMatrix, Cν::AbstractMatrix, ε::Real,
+                                    alg::EntropicGromovWasserstein = EntropicGromovWassersteinGibbs(SinkhornGibbs()); atol = nothing, rtol = nothing, check_convergence = 10, maxiter::Int=1_000, kwargs...)
+    T = float(Base.promote_eltype(μ, one(eltype(Cμ)) / ε, eltype(Cν)))
+    C = similar(Cμ, T, size(μ, 1), size(ν, 1))
+    tmp = similar(C)
+    plan = similar(C)
+    @. plan = μ * ν'
+    plan_prev = similar(C)
+    plan_prev .= plan
+    norm_plan = sum(plan)
+
+    _atol = atol === nothing ? 0 : atol
+    _rtol = rtol === nothing ? (_atol > zero(_atol) ? zero(T) : sqrt(eps(T))) : rtol
+
+    function get_new_cost!(C, plan, tmp, Cμ, Cν)
+        A_batched_mul_B!(tmp, Cμ, plan)
+        A_batched_mul_B!(C, tmp, -4Cν)
+        # seems to be a missing factor of 4 (or something like that...) compared to the POT implementation?
+        # added the factor of 4 here to ensure reproducibility for the same value of ε.
+        # https://github.com/PythonOT/POT/blob/9412f0ad1c0003e659b7d779bf8b6728e0e5e60f/ot/gromov.py#L247
+    end
+
+    get_new_cost!(C, plan, tmp, Cμ, Cν)
+    to_check_step = check_convergence
+
+    isconverged = false
+    for iter in 1:maxiter
+        # perform Sinkhorn algorithm
+        solver = build_solver(μ, ν, C, ε, alg.alg_step; kwargs...)
+        solve!(solver)
+        # compute optimal transport plan
+        plan = sinkhorn_plan(solver)
+
+        to_check_step -= 1
+        if to_check_step == 0 || iter == maxiter
+            # reset counter
+            to_check_step = check_convergence
+            isconverged = sum(abs, plan - plan_prev) < max(_atol, _rtol * norm_plan)
+            if isconverged
+                @debug "$Gromov Wasserstein with $(solver.alg) ($iter/$maxiter): converged"
+                break
+            end
+            plan_prev .= plan
+        end
+        get_new_cost!(C, plan, tmp, Cμ, Cν)
+    end
+
+    return plan
+end

src/OptimalTransport.jl

@@ -13,16 +13,19 @@ using LinearAlgebra
 using IterativeSolvers
 using LogExpFunctions: LogExpFunctions
 using NNlib: NNlib
+using Logging
 
 export SinkhornGibbs, SinkhornStabilized, SinkhornEpsilonScaling
 export SinkhornBarycenterGibbs
 export QuadraticOTNewton
+export EntropicGromovWassersteinSinkhorn
 
 export sinkhorn, sinkhorn2
 export sinkhorn_stabilized, sinkhorn_stabilized_epsscaling, sinkhorn_barycenter
 export sinkhorn_unbalanced, sinkhorn_unbalanced2
 export sinkhorn_divergence
 export quadreg
+export entropic_gromov_wasserstein
 
 include("utils.jl")
 
@@ -42,4 +45,6 @@ include("quadratic_newton.jl")
 
 include("dual/entropic_dual.jl")
 
+include("gromov.jl")
+
 end

src/gromov.jl

@@ -0,0 +1,57 @@
+# Gromov-Wasserstein solver
+
+abstract type EntropicGromovWasserstein end
+
+struct EntropicGromovWassersteinSinkhorn <: EntropicGromovWasserstein 
+    alg_step::Sinkhorn
+end
+
+function entropic_gromov_wasserstein(μ::AbstractVector, ν::AbstractVector, Cμ::AbstractMatrix, Cν::AbstractMatrix, ε::Real,
+                                    alg::EntropicGromovWasserstein = EntropicGromovWassersteinSinkhorn(SinkhornGibbs()); atol = nothing, rtol = nothing, check_convergence = 10, maxiter::Int=1_000, kwargs...)
+    T = float(Base.promote_eltype(μ, one(eltype(Cμ)) / ε, eltype(Cν)))
+    C = similar(Cμ, T, size(μ, 1), size(ν, 1))
+    tmp = similar(C)
+    plan = similar(C)
+    @. plan = μ * ν'
+    plan_prev = similar(C)
+    plan_prev .= plan
+    norm_plan = sum(plan)
+
+    _atol = atol === nothing ? 0 : atol
+    _rtol = rtol === nothing ? (_atol > zero(_atol) ? zero(T) : sqrt(eps(T))) : rtol
+
+    function get_new_cost!(C, plan, tmp, Cμ, Cν)
+        A_batched_mul_B!(tmp, Cμ, plan)
+        A_batched_mul_B!(C, tmp, -4Cν)
+        # seems to be a missing factor of 4 (or something like that...) compared to the POT implementation?
+        # added the factor of 4 here to ensure reproducibility for the same value of ε.
+        # https://github.com/PythonOT/POT/blob/9412f0ad1c0003e659b7d779bf8b6728e0e5e60f/ot/gromov.py#L247
+    end
+
+    get_new_cost!(C, plan, tmp, Cμ, Cν)
+    to_check_step = check_convergence
+
+    isconverged = false
+    for iter in 1:maxiter
+        # perform Sinkhorn algorithm
+        solver = build_solver(μ, ν, C, ε, alg.alg_step; kwargs...)
+        solve!(solver)
+        # compute optimal transport plan
+        plan = sinkhorn_plan(solver)
+
+        to_check_step -= 1
+        if to_check_step == 0 || iter == maxiter
+            # reset counter
+            to_check_step = check_convergence
+            isconverged = sum(abs, plan - plan_prev) < max(_atol, _rtol * norm_plan)
+            if isconverged
+                @debug "Gromov Wasserstein with $(solver.alg) ($iter/$maxiter): converged"
+                break
+            end
+            plan_prev .= plan
+        end
+        get_new_cost!(C, plan, tmp, Cμ, Cν)
+    end
+
+    return plan
+end

test/Project.toml

@@ -0,0 +1,2 @@
+[deps]
+OptimalTransport = "7e02d93a-ae51-4f58-b602-d97af76e3b33"

test/gromov.jl

@@ -0,0 +1,31 @@
+using OptimalTransport
+
+using Distances
+using PythonOT: PythonOT
+
+using Random
+using Test
+using LinearAlgebra
+
+const POT = PythonOT
+
+Random.seed!(100)
+
+@testset "gromov.jl" begin
+    @testset "entropic_gromov_wasserstein" begin
+        M, N = 250, 200
+
+        μ = fill(1/M, M)
+        μ_spt = rand(M)
+        ν = fill(1/N, N)
+        ν_spt = rand(N)
+
+        Cμ = pairwise(SqEuclidean(), μ_spt)
+        Cν = pairwise(SqEuclidean(), ν_spt)
+
+        γ = entropic_gromov_wasserstein(μ, ν, Cμ, Cν, 0.01; check_convergence = 10)
+        γ_pot = PythonOT.entropic_gromov_wasserstein(μ, ν, Cμ, Cν, 0.01)
+
+        @test γ ≈ γ_pot rtol = 1e-6
+    end
+end

test/runtests.jl

@@ -36,6 +36,10 @@ const GROUP = get(ENV, "GROUP", "All")
         @safetestset "Quadratically regularized OT" begin
             include("quadratic.jl")
         end
+
+        @safetestset "Gromov-Wasserstein OT" begin
+            include("gromov.jl")
+        end
     end
 
     # CUDA requires Julia >= 1.6