Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

LR Sinkhorn improvements #111

Merged
merged 40 commits into from
Sep 1, 2022
Merged

Conversation

meyerscetbon
Copy link
Collaborator

Here are the implementations of the update for the lr-sinkhorn algorithm:

  • New init using kmeans
  • Rescaling of the gamma
  • New stopping criterion
  • New implementation for the update of the cost matrices

@review-notebook-app
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@michalk8 michalk8 self-requested a review July 21, 2022 15:29
@michalk8 michalk8 added the enhancement New feature or request label Jul 21, 2022
@michalk8 michalk8 changed the title Branch meyer recovered LR Sinkhorn improvements Jul 21, 2022
Copy link
Collaborator

@michalk8 michalk8 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had a look the the notebook that causes the conflict, there are no new changes.
Could you please update the branch as git merge master -X theirs?

ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
Copy link
Collaborator

@michalk8 michalk8 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I got circular import error because of the kmeans, because it's placed in tools. Making it a relative import in LRSinkhorn.__call__ solved this.

ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
@michalk8 michalk8 mentioned this pull request Aug 4, 2022
@michalk8 michalk8 requested a review from marcocuturi August 31, 2022 07:14
@michalk8
Copy link
Collaborator

@meyerscetbon @marcocuturi you can both review it now, there some TODOs left for me, but not critical

@codecov-commenter
Copy link

codecov-commenter commented Aug 31, 2022

Codecov Report

Merging #111 (fb1a1df) into main (3d1c7d2) will decrease coverage by 6.27%.
The diff coverage is 87.60%.

Impacted file tree graph

@@            Coverage Diff             @@
##             main     #111      +/-   ##
==========================================
- Coverage   89.45%   83.17%   -6.28%     
==========================================
  Files          47       48       +1     
  Lines        4580     4713     +133     
  Branches      503      512       +9     
==========================================
- Hits         4097     3920     -177     
- Misses        364      650     +286     
- Partials      119      143      +24     
Impacted Files Coverage Δ
ott/core/quad_problems.py 86.97% <ø> (-3.73%) ⬇️
ott/core/initializers_lr.py 79.31% <79.31%> (ø)
ott/core/sinkhorn_lr.py 93.44% <94.61%> (+0.86%) ⬆️
ott/core/initializers.py 100.00% <100.00%> (ø)
ott/core/sinkhorn.py 97.23% <100.00%> (-0.80%) ⬇️
ott/tools/k_means.py 92.54% <100.00%> (-3.08%) ⬇️
ott/core/discrete_barycenter.py 15.58% <0.00%> (-70.13%) ⬇️
ott/tools/segment_sinkhorn.py 43.75% <0.00%> (-56.25%) ⬇️
ott/tools/sinkhorn_divergence.py 71.42% <0.00%> (-28.58%) ⬇️
ott/tools/soft_sort.py 69.90% <0.00%> (-25.25%) ⬇️
... and 23 more

Copy link
Contributor

@marcocuturi marcocuturi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot for this PR, it looks like we have a really great LR implementation now.

ott/core/initializers_lr.py Show resolved Hide resolved
ott/core/sinkhorn_lr.py Show resolved Hide resolved
ott/core/sinkhorn_lr.py Outdated Show resolved Hide resolved
ott/core/sinkhorn_lr.py Show resolved Hide resolved
ott/core/sinkhorn_lr.py Show resolved Hide resolved
tests/core/sinkhorn_lr_test.py Outdated Show resolved Hide resolved
tests/core/sinkhorn_test.py Show resolved Hide resolved
@michalk8 michalk8 force-pushed the branch_meyer_recovered branch from d7a2a16 to fb1a1df Compare September 1, 2022 13:10
@michalk8 michalk8 merged commit 1419513 into ott-jax:main Sep 1, 2022
@meyerscetbon meyerscetbon deleted the branch_meyer_recovered branch February 14, 2023 09:21
michalk8 pushed a commit that referenced this pull request Jun 27, 2024
* test

* update lr-sinkhorn

* restored_branch

* check

* review

* circular fixed

* update review

* Fix bugs in `LRSinkhorn`

* Use new `k-means` implementation

* Fix linter

* Refactor `LRSinkhorn` initializers

* Use `if` for `is_entropic`, remove dead variables

* Slightly improve types

* Do not use stateful `gamma`

* Fix typo in tests

* Fix using `state.gamma` instead of `self.gamma`

* Fix point cloud size in notebook

* Add assertion to k-means

* Use `jax.lax.cond` instead of `jax.numpy.where`

* Change convergence criterion

* Use safe log

* Fix more tests

* Fix tests

* Fix `tree_flatten` in `KMeansInitializer`

* Fix defaults, change `rank_2` -> `rank2`

* Simplify `apply`

* Update TODOs

* Update docs, make `lr_costs` private

* Increate tolerance in failing test

* Update LR notebook

* Address comments

* Remove LR Sinkhorn notebook from testing, to slow

Co-authored-by: Michal Klein
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants