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

univariate dual vectors for arbitrary sizes #505

Merged
merged 47 commits into from
Jun 4, 2024
Merged

univariate dual vectors for arbitrary sizes #505

merged 47 commits into from
Jun 4, 2024

Conversation

marcocuturi
Copy link
Contributor

this PR implements an extension of the univariate solver that outputs dual vectors, regardless of input sizes / weights.

the implementation uses a lineax sparse solver, instead of going through a manually coded Gauss-seidel loop as in e.g. Alg. 1 of https://arxiv.org/pdf/2201.00730.pdf

the linear system is mostly build by looking at transported mass (returned by the primal solver) and inverting the corresponding linear system.

at the moment, i still observe a bug for the (arguably) least interesting case when comparing two univariate distributions of the same size (in the test, when weight vectors associated to both x and z are unif_n for instance). note this was the case that the https://arxiv.org/abs/2206.07630 paper was dwelling a bit deeper into, but with a quadratic iteration.

note: it might be interesting to actually code the Alg. 1 of https://arxiv.org/pdf/2201.00730.pdf loop with single entry updates and see if it does run faster than this iterative solver.

src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
tests/solvers/linear/univariate_test.py Outdated Show resolved Hide resolved
tests/solvers/linear/univariate_test.py Outdated Show resolved Hide resolved
tests/solvers/linear/univariate_test.py Outdated Show resolved Hide resolved
tests/solvers/linear/univariate_test.py Outdated Show resolved Hide resolved
tests/solvers/linear/univariate_test.py Outdated Show resolved Hide resolved
@michalk8 michalk8 added the enhancement New feature or request label Mar 27, 2024
Copy link
Contributor Author

@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 Michal for the implementation of NW

src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Show resolved Hide resolved
tests/solvers/linear/univariate_test.py Outdated Show resolved Hide resolved
Copy link

codecov bot commented May 30, 2024

Codecov Report

Attention: Patch coverage is 94.53125% with 7 lines in your changes missing coverage. Please review.

Project coverage is 91.34%. Comparing base (c0a61e8) to head (2bd0cdb).
Report is 39 commits behind head on main.

Files with missing lines Patch % Lines
src/ott/solvers/linear/_solve.py 40.00% 6 Missing ⚠️
src/ott/solvers/linear/univariate.py 99.09% 1 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main     #505      +/-   ##
==========================================
+ Coverage   91.20%   91.34%   +0.13%     
==========================================
  Files          68       68              
  Lines        7075     7083       +8     
  Branches     1000      991       -9     
==========================================
+ Hits         6453     6470      +17     
+ Misses        468      463       -5     
+ Partials      154      150       -4     
Files with missing lines Coverage Δ
src/ott/geometry/distrib_costs.py 100.00% <100.00%> (ø)
src/ott/solvers/quadratic/lower_bound.py 84.61% <ø> (+9.61%) ⬆️
src/ott/solvers/linear/univariate.py 97.63% <99.09%> (+11.19%) ⬆️
src/ott/solvers/linear/_solve.py 57.89% <40.00%> (-25.44%) ⬇️

... and 3 files with indirect coverage changes

@michalk8 michalk8 self-requested a review June 4, 2024 13:32
Copy link
Contributor Author

@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 Michal!! a few questions here and there, about naming but maybe this hides some API choices :)

src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
src/ott/solvers/linear/univariate.py Outdated Show resolved Hide resolved
@michalk8 michalk8 merged commit 1b7248e into main Jun 4, 2024
11 of 12 checks passed
@michalk8 michalk8 deleted the 1d_dual branch June 4, 2024 17:11
michalk8 pushed a commit that referenced this pull request Jun 27, 2024
* univariate dual vectors for arbitrary sizes

* Remove old impl.

* Simplify outputs

* Better type

* Add initial north-west solver

* Update NW solver

* Clean the solver implementation

* Remove unnecessary sorting

* Add TODO

* Remove call to `jnp.minimum`

* Use functional interface

* Remove `UnivariateSolver` class

* Update tests

* Use `BCOO`

* Add `dual_costs` property

* Fix NW not undoing the sort

* Fix `test_dual_vectors`

* Fix last test

* Fix links in docs

* Update citation

* Update `UnivariateWasserstein`'s docs

* More grad tests

* Polish another test

* Parametrize the last test

* Revert back to `eps=1.5e-3`

* Add more LB tests

* Make Sinkhorn test more stable

* Start improving docs in `UnivariateOutput`

* Update title in the docs

* Continue polishing the docs

* Fix remaining TODOs

* Nicer impl.

* Better docs

* Last pass over the docs

* Test different costs in `test_dual_vectors`

* [ci-skip] Fix typo in docs

* Rename functions

* Add `solve_univariate` high-level function

* Polish docs

* Incorporate the last comments

* Fix spelling

* Update NW solver docs

* Unify phrasing in the docs

* Update `UnivariateWasserstein` docs

* Remove extra `.`

---------

Co-authored-by: James Thornton <JTT94@users.noreply.github.com>
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.

3 participants