Refactor sinkhorn and sinkhorn2 (#100)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
JuliaOptimalTransport · Jun 27, 2021 · f03b05c · f03b05c · devmotion · Jun 27, 2021
1 parent a2224ee
commit f03b05c
Show file tree

Hide file tree

Showing 21 changed files with 1,643 additions and 1,094 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "OptimalTransport"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
 authors = ["zsteve <stephenz@student.unimelb.edu.au>"]
-version = "0.3.11"
+version = "0.3.12"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
@@ -10,6 +10,7 @@ IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -21,6 +22,7 @@ Distributions = "0.24, 0.25"
 IterativeSolvers = "0.8.4, 0.9"
 LogExpFunctions = "0.2"
 MathOptInterface = "0.9"
+NNlib = "0.6, 0.7"
 PDMats = "0.11"
 QuadGK = "2"
 StatsBase = "0.33.8"

diff --git a/docs/src/index.md b/docs/src/index.md
@@ -21,18 +21,33 @@ squared2wasserstein
 ```@docs
 sinkhorn
 sinkhorn2
-sinkhorn_stabilized_epsscaling
-sinkhorn_stabilized
 sinkhorn_barycenter
 ```
 
+Currently the following variants of the Sinkhorn algorithm are supported:
+
+```@docs
+SinkhornGibbs
+SinkhornStabilized
+SinkhornEpsilonScaling
+```
+
+The following methods are deprecated and will be removed:
+
+```@docs
+sinkhorn_stabilized
+sinkhorn_stabilized_epsscaling
+```
+
 ## Unbalanced optimal transport
+
 ```@docs
 sinkhorn_unbalanced
 sinkhorn_unbalanced2
 ```
 
 ## Quadratically regularised optimal transport
+
 ```@docs
 quadreg
 ```
diff --git a/src/OptimalTransport.jl b/src/OptimalTransport.jl
@@ -12,8 +12,11 @@ using MathOptInterface
 using Distributions
 using PDMats
 using QuadGK
+using NNlib: NNlib
 using StatsBase: StatsBase
 
+export SinkhornGibbs, SinkhornStabilized, SinkhornEpsilonScaling
+
 export sinkhorn, sinkhorn2
 export emd, emd2
 export sinkhorn_stabilized, sinkhorn_stabilized_epsscaling, sinkhorn_barycenter
@@ -27,8 +30,14 @@ include("distances/bures.jl")
 include("utils.jl")
 include("exact.jl")
 include("wasserstein.jl")
+
 include("entropic/sinkhorn.jl")
+include("entropic/sinkhorn_gibbs.jl")
 include("entropic/sinkhorn_stabilized.jl")
+include("entropic/sinkhorn_epsscaling.jl")
+include("entropic/sinkhorn_unbalanced.jl")
+include("entropic/sinkhorn_barycenter.jl")
+
 include("quadratic.jl")
 
 end