-
Notifications
You must be signed in to change notification settings - Fork 79
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
Conversation
There was a problem hiding this 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
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ 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
|
There was a problem hiding this 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 :)
* 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>
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
andz
areunif_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.