add Hungarian solver from optax #598
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
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #598 +/- ##
==========================================
+ Coverage 87.90% 88.23% +0.32%
==========================================
Files 73 74 +1
Lines 7770 7808 +38
Branches 556 556
==========================================
+ Hits 6830 6889 +59
+ Misses 799 778 -21
Partials 141 141
|
michalk8
requested changes
Nov 19, 2024
| @@ -347,6 +347,34 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 | |||
| return cls(*aux_data) | |||
|
|
|||
|
|
|||
| @jtu.register_pytree_node_class | |||
| class EuclideanP(TICost): | |||
There was a problem hiding this comment.
ok. Don't want to surprise the user because there's no Legendre transform, but can add.
| the first geometry sends mass to point :math:`j` in the second. | ||
| """ | ||
| geom: geometry.Geometry | ||
| paired_indices: Optional[jnp.ndarray] = None |
There was a problem hiding this comment.
Would remove the optional and default none.
src/ott/tools/unreg.py
Outdated
| geom: geometry object | ||
| paired_indices: Array of shape ``[2, n]``, of :math:`n` pairs | ||
| of indices, for which the optimal transport assigns mass. Namely, for each | ||
| index :math:`0 <= k < n`, if one has |
src/ott/tools/unreg.py
Outdated
|
|
||
|
|
||
| def wassdis_p(x: jnp.ndarray, y: jnp.ndarray, p: float = 2.0) -> float: | ||
| """Convenience wrapper on `hungarian` to get :term:`Wasserstein distance`. |
There was a problem hiding this comment.
Would remove the convenience wrapper part and rephrase as
Compute the :term:`Wasserstein distance` using the Hungarian algoritm.
or similar.
src/ott/tools/unreg.py
Outdated
|
|
||
|
|
||
| def hungarian(geom: geometry.Geometry) -> Tuple[jnp.ndarray, HungarianOutput]: | ||
| """Solve assignment problem using Hungarian as implemented in optax. |
There was a problem hiding this comment.
Also please add this to intersphinx_mapping in conf.py
There was a problem hiding this comment.
isn't it there already?
intersphinx_mapping = {
"python": ("https://docs.python.org/3", None),
"numpy": ("https://numpy.org/doc/stable/", None),
"jax": ("https://jax.readthedocs.io/en/latest/", None),
"jaxopt": ("https://jaxopt.github.io/stable", None),
"lineax": ("https://docs.kidger.site/lineax/", None),
"flax": ("https://flax.readthedocs.io/en/latest/", None),
"optax": ("https://optax.readthedocs.io/en/latest/", None),
"diffrax": ("https://docs.kidger.site/diffrax/", None),
"scipy": ("https://docs.scipy.org/doc/scipy/", None),
"pot": ("https://pythonot.github.io/", None),
"matplotlib": ("https://matplotlib.org/stable/", None),
}
tests/tools/unreg_test.py
Outdated
| rng1, rng2 = jax.random.split(rng, 2) | ||
| x, y = gen_data(rng1, n, m, dim) | ||
| geom = pointcloud.PointCloud(x, y, cost_fn=cost_fn, epsilon=.0005) | ||
| cost_hung, out_hung = unreg.hungarian(geom) |
tests/tools/unreg_test.py
Outdated
| x, y = gen_data(rng1, n, m, dim) | ||
| geom = pointcloud.PointCloud(x, y, cost_fn=costs.EuclideanP(p=p)) | ||
| cost_hung, _ = unreg.hungarian(geom) | ||
| w_p = unreg.wassdis_p(x, y, p) |
src/ott/tools/unreg.py
Outdated
| def wassdis_p(x: jnp.ndarray, y: jnp.ndarray, p: float = 2.0) -> float: | ||
| """Convenience wrapper on `hungarian` to get :term:`Wasserstein distance`. | ||
|
|
||
| Uses :func:`~ott.tools.unreg.hungarian` to solve the |
src/ott/tools/unreg.py
Outdated
| """Solve assignment problem using Hungarian as implemented in optax. | ||
|
|
||
| Args: | ||
| geom: (square) geometry object. |
There was a problem hiding this comment.
Would say
Geometry of shape ``[n, n]``.
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.
Take advantage of google-deepmind/optax#1083 to include unregularized OT solver for point clouds of equal size.