Feature/geodesic sinkhorn #457
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This module contains the implementation of the Geodesic Sinkhorn algorithm [1]. [1] Huguet, G., Tong, A., Zapatero, M. R., Wolf, G., & Krishnaswamy, S. (2022). Geodesic Sinkhorn: optimal transport for high-dimensional datasets. arXiv preprint arXiv:2211.00805.
Create `geodesic` module
…eodesic-sinkhorn
Closed
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #457 +/- ##
==========================================
+ Coverage 81.97% 90.60% +8.63%
==========================================
Files 59 60 +1
Lines 6345 6464 +119
Branches 615 913 +298
==========================================
+ Hits 5201 5857 +656
+ Misses 1024 465 -559
- Partials 120 142 +22
|
|
Hey @michalk8 let me know if you have any suggestions on the PR. I am not sure what is the best way to deal with the dtype in |
michalk8
requested changes
Nov 27, 2023
src/ott/geometry/geodesic.py
Outdated
| self.eigval = eigval | ||
| self.chebyshev_coeffs = chebyshev_coeffs | ||
| self.t = t | ||
| self.order = order |
There was a problem hiding this comment.
I think we should remove this, as it's not being used anywhere in the class.
The docs will need to be adapted a bit as well.
There was a problem hiding this comment.
I only removed self.order, because self.t is used for the attribute cost_matrix, and self.eigval, self.chebyshev_coeffs are used inexpm_multiply. I renamed phi to eigval in expm_multiply to make it clearer.
michalk8
approved these changes
Nov 28, 2023
michalk8
pushed a commit
that referenced
this pull request
Jun 27, 2024
* Create `geodesic` module This module contains the implementation of the Geodesic Sinkhorn algorithm [1]. [1] Huguet, G., Tong, A., Zapatero, M. R., Wolf, G., & Krishnaswamy, S. (2022). Geodesic Sinkhorn: optimal transport for high-dimensional datasets. arXiv preprint arXiv:2211.00805. * Lint code * Lint code * Add Geodesic kernel citation to `docs/references.bib` * Remove unused functions; update docstrings; remove `n_steps`; remove forced symm * Remove hardcoded random key at `compute_largest_eigenvalue` * Fix docstrings of functions inside `apply_kernel` * Add chebyshev coeff computation to `from_graph` * Change `jax.exp.sparse` import * Change `tree_flatten` outputs * Change input of `lobpcg_std` to be a sparsified product * Change definition of cost from kernel * Add `Geodesic` to `docs/geometry` * Remove `np.max` from max eigenval computation * Add `default_prng_key` to eigenval computation * Add `safe_log` to cost matrix computation * Change to `jesp.BCOO` at chebyshev approx * Restructure `from_graph`; coeffs are computed earlier now * fn outside of the class * mv fn & process L once * wip tests geo * jax pure_callback & new fn cheb * symmetric kernel & wip tests * fix formatting ruff * wrap eigenval fn * rm num_scheme & _scale * type dense or sparse * expm with scan & fix hardcode dty & lobpcg iter * test compare with BE and CN * rm lap_mn_id & test sink spd & fix tree_flatten * default and type of Cheb. co. & docstrings * dtype in purecallback depending on Lap * rm laplacian & update test * simpler wrapper for `ive` * lint-docs spell check * fix mistake rm setup * typo & rm fn scale lap & * rm t order init & typing * use rng util * t for cost & test with ground truth * rm order & change `phi` to `eigval` & type hints * differentiability test * condition to symmetrize kernel * fix indentation --------- Co-authored-by: Diego Baptista Theuerkauf <[email protected]> Co-authored-by: diegoabt <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Implementation in: https://arxiv.org/abs/2211.00805
previous PR by @diegoabt can be found in #425