univariate dual vectors for arbitrary sizes

docs/geometry.rst

@@ -56,6 +56,7 @@ Cost Functions
     :toctree: _autosummary
 
     costs.CostFn
+    costs.TICost
     costs.SqPNorm
     costs.PNormP
     costs.SqEuclidean

docs/references.bib

@@ -169,6 +169,7 @@ @article{vayer:20
 
 @article{demetci:22,
   author  = {Demetci, Pinar and Santorella, Rebecca and Sandstede, Björn and Noble, William Stafford and Singh, Ritambhara},
+  doi     = {10.1089/cmb.2021.0446},
   journal = {Journal of Computational Biology},
   note    = {PMID: 35050714},
   number  = {1},
@@ -435,18 +436,13 @@ @article{schmitz:18
   year    = {2018},
 }
 
-@article{alvarez-esteban:16,
-  author   = {Álvarez-Esteban, Pedro C. and {del Barrio}, E. and Cuesta-Albertos, J.A. and Matrán, C.},
-  url      = {https://www.sciencedirect.com/science/article/pii/S0022247X16300907},
-  doi      = {10.1016/j.jmaa.2016.04.045},
-  issn     = {0022-247X},
-  journal  = {Journal of Mathematical Analysis and Applications},
-  keywords = {Mass transportation problem,-Wasserstein distance,Wasserstein barycenter,Fréchet mean,Fixed-point iteration,Location-scatter families},
-  number   = {2},
-  pages    = {744--762},
-  title    = {A fixed-point approach to barycenters in Wasserstein space},
-  volume   = {441},
-  year     = {2016},
+@misc{alvarez-esteban:16,
+  author      = {Álvarez-Esteban, Pedro C. and del Barrio, E. and Cuesta-Albertos, J. A. and Matrán, C.},
+  eprint      = {1511.05355},
+  eprintclass = {stat.CO},
+  eprinttype  = {arXiv},
+  title       = {A fixed-point approach to barycenters in Wasserstein space},
+  year        = {2016},
 }
 
 @article{lehmann:21,

docs/solvers/linear.rst

@@ -33,12 +33,15 @@ Barycenter Solvers
     discrete_barycenter.FixedBarycenter
     discrete_barycenter.SinkhornBarycenterOutput
 
-Other Solvers
--------------
+Univariate Solvers
+------------------
 .. autosummary::
     :toctree: _autosummary
 
-    univariate.UnivariateSolver
+    univariate.uniform_distance
+    univariate.quantile_distance
+    univariate.north_west_distance
+    univariate.UnivariateOutput
 
 Sinkhorn Acceleration
 ---------------------

src/ott/geometry/distrib_costs.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
+from typing import Callable, Optional
 
 import jax.numpy as jnp
 import jax.tree_util as jtu
@@ -20,9 +20,7 @@
 from ott.problems.linear import linear_problem
 from ott.solvers.linear import univariate
 
-__all__ = [
-    "UnivariateWasserstein",
-]
+__all__ = ["UnivariateWasserstein"]
 
 
 @jtu.register_pytree_node_class
@@ -34,31 +32,24 @@ class UnivariateWasserstein(costs.CostFn):
   ground cost.
 
   Args:
+    solve_fn: 1D optimal transport solver, e.g.,
+      :func:`~ott.solvers.linear.univariate.uniform_distance`.
     ground_cost: Cost used to compute the 1D optimal transport between vector,
       should be a translation-invariant (TI) cost for correctness.
       If :obj:`None`, defaults to :class:`~ott.geometry.costs.SqEuclidean`.
-    solver: 1D optimal transport solver.
-    kwargs: Arguments passed on when calling the
-      :class:`~ott.solvers.linear.univariate.UnivariateSolver`. May include
-      random key, or specific instructions to subsample or compute using
-      quantiles.
   """
 
   def __init__(
       self,
+      solve_fn: Callable[[linear_problem.LinearProblem],
+                         univariate.UnivariateOutput],
       ground_cost: Optional[costs.TICost] = None,
-      solver: Optional[univariate.UnivariateSolver] = None,
-      **kwargs: Any
   ):
     super().__init__()
-
     self.ground_cost = (
         costs.SqEuclidean() if ground_cost is None else ground_cost
     )
-    self._solver = univariate.UnivariateSolver() if solver is None else solver
-    self._kwargs_solve = kwargs
-    # ensure transport solutions are neither computed nor stored
-    self._kwargs_solve["return_transport"] = False
+    self._solve_fn = solve_fn
 
   def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
     """Wasserstein distance between :math:`x` and :math:`y` seen as a 1D dist.
@@ -70,20 +61,16 @@ def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
     Returns:
       The transport cost.
     """
-    out = self._solver(
-        linear_problem.LinearProblem(
-            pointcloud.PointCloud(
-                x[:, None], y[:, None], cost_fn=self.ground_cost
-            )
-        ), **self._kwargs_solve
+    geom = pointcloud.PointCloud(
+        x[:, None], y[:, None], cost_fn=self.ground_cost
     )
+    prob = linear_problem.LinearProblem(geom)
+    out = self._solve_fn(prob)
     return jnp.squeeze(out.ot_costs)
 
   def tree_flatten(self):  # noqa: D102
-    return (self.ground_cost,), (self._solver, self._kwargs_solve)
+    return (self.ground_cost,), (self._solve_fn,)
 
   @classmethod
   def tree_unflatten(cls, aux_data, children):  # noqa: D102
-    ground_cost, = children
-    solver, solve_kwargs = aux_data
-    return cls(ground_cost, solver, **solve_kwargs)
+    return cls(solve_fn=aux_data[0], ground_cost=children[0])