Misc/project structure

ott/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """OTT library."""
-from . import core, geometry, tools
+from . import geometry, initializers, math, problems, solvers, tools, utils
 from ._version import __version__

ott/geometry/__init__.py

@@ -12,9 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """OTT ground geometries: Classes and cost functions to instantiate them."""
-from . import costs, low_rank, ops
-from .epsilon_scheduler import Epsilon
-from .geometry import Geometry
-from .graph import Graph
-from .grid import Grid
-from .pointcloud import PointCloud
+from . import costs, epsilon_scheduler, geometry, graph, grid, pointcloud, segment

ott/geometry/costs.py

@@ -17,13 +17,17 @@
 import abc
 import functools
 import math
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 
-from ott.core import fixed_point_loop
-from ott.geometry import matrix_square_root
+from ott.math import fixed_point_loop, matrix_square_root
+
+__all__ = [
+    "PNorm", "SqPNorm", "Euclidean", "SqEuclidean", "Cosine", "Bures",
+    "UnbalancedBures"
+]
 
 
 @jax.tree_util.register_pytree_node_class
@@ -32,9 +36,11 @@ class CostFn(abc.ABC):
 
   Cost functions evaluate a function on a pair of inputs. For convenience,
   that function is split into two norms -- evaluated on each input separately --
-  followed by a pairwise cost that involves both inputs, as in
+  followed by a pairwise cost that involves both inputs, as in:
+
+  .. math::
 
-  c(x,y) = norm(x) + norm(y) + pairwise(x,y)
+    c(x,y) = norm(x) + norm(y) + pairwise(x,y)
 
   If the norm function is not implemented, that value is handled as a 0.
   """
@@ -108,14 +114,14 @@ class TICost(CostFn):
 
   @abc.abstractmethod
   def h(self, z: jnp.ndarray) -> float:
-    """RBF function acting on difference of `x-y` to ouput cost."""
+    """TI function acting on difference of :math:`x-y` to output cost."""
 
   def h_legendre(self, z: jnp.ndarray) -> float:
-    """Legendre transform of RBF function `h` (when latter is convex)."""
+    """Legendre transform of TI function :func:`h` (when latter is convex)."""
     raise NotImplementedError("`h_legendre` not implemented.")
 
   def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
-    """Compute cost as evaluation of :func:`h` on `x-y`."""
+    """Compute cost as evaluation of :func:`h` on :math:`x-y`."""
     return self.h(x - y)
 
 
@@ -128,9 +134,10 @@ class SqPNorm(TICost):
   """
 
   def __init__(self, p: float):
+    super().__init__()
     assert p >= 1.0, "p parameter in sq. p-norm should be >= 1.0"
     self.p = p
-    self.q = 1. / (1 - 1 / self.p) if p > 1.0 else 'inf'
+    self.q = 1. / (1. - 1. / self.p) if p > 1.0 else "inf"
 
   def h(self, z: jnp.ndarray) -> float:
     return 0.5 * jnp.linalg.norm(z, self.p) ** 2
@@ -152,9 +159,10 @@ class PNorm(TICost):
   """p-norm (to the power p) of the difference of two vectors."""
 
   def __init__(self, p: float):
+    super().__init__()
     assert p >= 1.0, "p parameter in p-norm should be >= 1.0"
     self.p = p
-    self.q = 1. / (1 - 1 / self.p)
+    self.q = 1. / (1. - 1. / self.p) if p > 1. else "inf"
 
   def h(self, z: jnp.ndarray) -> float:
     return jnp.linalg.norm(z, self.p) ** self.p / self.p
@@ -175,8 +183,9 @@ def tree_unflatten(cls, aux_data, children):
 class Euclidean(CostFn):
   """Euclidean distance.
 
-  Note that the Euclidean distance is not cast as a `TICost`, because this
-  would correspond to `h = jnp.linalg.norm`, whose gradient is not invertible,
+  Note that the Euclidean distance is not cast as a
+  :class:`~ott.geometry.costs.TICost`, since this would correspond to :math:`h`
+  being :func:`jax.numpy.linalg.norm`, whose gradient is not invertible,
   because the function is not strictly convex (it is linear on rays).
   """
 
@@ -429,6 +438,7 @@ def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
     log_m_pi += -0.5 * ldet_c_ab
 
     # If all logdet signs are 1, output value, nan otherwise.
+    # TODO(michalk8): use lax.cond
     return jnp.where(
         sldet_c == 1 and sldet_c_ab == 1 and sldet_ab == 1 and sldet_t_ab == 1,
         2 * sig2 * mass_x * mass_y - 2 * (sig2 + gam) * jnp.exp(log_m_pi),
@@ -444,7 +454,8 @@ def tree_unflatten(cls, aux_data, children):
     return cls(aux_data[0], aux_data[1], aux_data[2], **aux_data[3])
 
 
-def x_to_means_and_covs(x: jnp.ndarray, dimension: jnp.ndarray) -> jnp.ndarray:
+def x_to_means_and_covs(x: jnp.ndarray,
+                        dimension: int) -> Tuple[jnp.ndarray, jnp.ndarray]:
   """Extract means and covariance matrices of Gaussians from raveled vector.
 
   Args:
@@ -456,7 +467,7 @@ def x_to_means_and_covs(x: jnp.ndarray, dimension: jnp.ndarray) -> jnp.ndarray:
     covariances: [num_gaussians, dimension] array that holds the covariances.
   """
   x = jnp.atleast_2d(x)
-  means = x[:, 0:dimension]
+  means = x[:, :dimension]
   covariances = jnp.reshape(
       x[:, dimension:dimension + dimension ** 2], (-1, dimension, dimension)
   )

ott/geometry/epsilon_scheduler.py

@@ -19,6 +19,8 @@
 import jax
 import jax.numpy as jnp
 
+__all__ = ["Epsilon"]
+
 
 @jax.tree_util.register_pytree_node_class
 class Epsilon: