JuliaOptimalTransport / OptimalTransport.jl

Optimal transport algorithms for Julia
https://juliaoptimaltransport.github.io/OptimalTransport.jl/dev
MIT License
93 stars 8 forks source link

Implementing the Greenkhorn #159

Open davibarreira opened 2 years ago

davibarreira commented 2 years ago

This PR is related to #151 . I've implemented the Greenkhorn algorithm which is a greedy version of the Sinkhorn algorithm. The method I've implemented is actually the one in POT, which is a bit different from the one in the original paper (https://arxiv.org/pdf/1705.09634.pdf).

The implementation needs improvement. I was not able to get it to work with AD and with the batch tests. I was not very involved in the coding of the original Sinkhorn algorithm, and had some difficulty getting around the whole step!, solve! and cache structure.

Another point. The iteration in the Greenkhorn algorithm only updates one row of u and v each time, thus, it needs many more iterations in order to converge compared to the original Sinkhorn. Some preliminar benchmarks showed that this implementation of the Greenkhorn is slower then the original Sinkhorn_Gibbs, which seems to contradict the claims in the paper. I believe the reason for this might be that the Sinkhorn implementation is very optimized in the package compared to my version of Greenkhorn. Another possibility is that the Sinkhorn version of the paper was not very efficient ( if you read it here, they present the Sinkhron algorithm which computer diagm(u) K diagm(v) in each iteration).

I've compared the results from my algorithm against POT, and indeed it seems to be returning the exact same result each iteration, i.e. it seems that the Greenkhorn implementation is correct but not optimal.

coveralls commented 2 years ago

Pull Request Test Coverage Report for Build 1704769432

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

Details


Changes Missing Coverage Covered Lines Changed/Added Lines %
src/entropic/greenkhorn.jl 45 46 97.83%
<!-- Total: 45 46 97.83% -->
Totals Coverage Status
Change from base Build 1620585433: 0.2%
Covered Lines: 650
Relevant Lines: 680

💛 - Coveralls
codecov-commenter commented 2 years ago

Codecov Report

Merging #159 (82b1026) into master (de56119) will increase coverage by 0.16%. The diff coverage is 97.82%.

Impacted file tree graph

@@            Coverage Diff             @@
##           master     #159      +/-   ##
==========================================
+ Coverage   95.42%   95.58%   +0.16%     
==========================================
  Files          14       15       +1     
  Lines         634      680      +46     
==========================================
+ Hits          605      650      +45     
- Misses         29       30       +1     
Impacted Files Coverage Δ
src/entropic/greenkhorn.jl 97.82% <97.82%> (ø)

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update de56119...82b1026. Read the comment docs.