(feat): Linear GENOT

docs/conf.py

@@ -63,12 +63,24 @@
 nitpicky = True
 nitpick_ignore = [
     ("py:class", "numpy.float64"),
+    # see: https://github.com/numpy/numpydoc/issues/275
+    ("py:class", "None.  Remove all items from D."),
+    ("py:class", "a set-like object providing a view on D's items"),
+    ("py:class", "a set-like object providing a view on D's keys"),
+    ("py:class", "v, remove specified key and return the corresponding value."),  # noqa: E501
+    ("py:class", "None.  Update D from dict/iterable E and F."),
+    ("py:class", "an object providing a view on D's values"),
+    ("py:class", "a shallow copy of D"),
 ]
 # TODO(michalk8): remove once typing has been cleaned-up
 nitpick_ignore_regex = [
     (r"py:class", r"moscot\..*(K|B|O)"),
     (r"py:class", r"numpy\._typing.*"),
     (r"py:class", r"moscot\..*Protocol.*"),
+    (
+        r"py:class",
+        r"moscot.base.output.BaseSolverOutput",
+    ),  # https://github.com/sphinx-doc/sphinx/issues/10974 means there is simply no way around this with generics
 ]
 
 

docs/developer.rst

@@ -12,6 +12,8 @@ Backends
     backends.ott.GWSolver
     backends.ott.OTTOutput
     backends.ott.GraphOTTOutput
+    backends.ott.GENOTLinSolver
+    backends.ott.output.OTTNeuralOutput
     backends.utils.get_solver
     backends.utils.get_available_backends
 
@@ -44,6 +46,7 @@ Problems
     problems.BaseCompoundProblem
     problems.CompoundProblem
     cost.BaseCost
+    problems.CondOTProblem
 
 Mixins
 ^^^^^^
@@ -62,14 +65,13 @@ Solvers
 
     solver.BaseSolver
     solver.OTSolver
-    output.BaseSolverOutput
 
 Output
 ^^^^^^
 .. autosummary::
     :toctree: genapi
 
-    output.BaseSolverOutput
+    output.BaseDiscreteSolverOutput
     output.MatrixSolverOutput
 
 Utils
@@ -100,6 +102,8 @@ Miscellaneous
     data.apoptosis_markers
     tagged_array.TaggedArray
     tagged_array.Tag
+    tagged_array.DistributionCollection
+    tagged_array.DistributionContainer
 
 .. currentmodule:: moscot.base.problems
 .. autosummary::

docs/references.bib

@@ -477,3 +477,12 @@ @article{srivatsan:20
   year={2020},
   publisher={American Association for the Advancement of Science}
 }
+
+@misc{klein2023generative,
+      title={Generative Entropic Neural Optimal Transport To Map Within and Across Spaces},
+      author={Dominik Klein and Théo Uscidda and Fabian Theis and Marco Cuturi},
+      year={2023},
+      eprint={2310.09254},
+      archivePrefix={arXiv},
+      primaryClass={stat.ML}
+}

docs/user.rst

@@ -27,6 +27,7 @@ Generic Problems
     generic.SinkhornProblem
     generic.GWProblem
     generic.FGWProblem
+    generic.GENOTLinProblem
 
 Plotting
 ~~~~~~~~

pyproject.toml

@@ -55,7 +55,7 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax>=0.4.6",
+    "ott-jax[neural]>=0.4.6",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
 ]
@@ -267,11 +267,11 @@ skip_missing_interpreters = true
 
 [testenv]
 extras = test
-pass_env = PYTEST_*,CI
 commands =
     python -m pytest {tty:--color=yes} {posargs: \
         --cov={env_site_packages_dir}{/}moscot --cov-config={tox_root}{/}pyproject.toml \
         --no-cov-on-fail --cov-report=xml --cov-report=term-missing:skip-covered}
+passenv = PYTEST_*,CI
 
 [testenv:lint-code]
 description = Lint the code.

src/moscot/backends/ott/__init__.py

@@ -1,11 +1,12 @@
 from ott.geometry import costs
 
 from moscot.backends.ott._utils import sinkhorn_divergence
-from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
-from moscot.backends.ott.solver import GWSolver, SinkhornSolver
+from moscot.backends.ott.output import GraphOTTOutput, OTTNeuralOutput, OTTOutput
+from moscot.backends.ott.solver import GENOTLinSolver, GWSolver, SinkhornSolver
 from moscot.costs import register_cost
 
-__all__ = ["OTTOutput", "GraphOTTOutput", "GWSolver", "SinkhornSolver", "sinkhorn_divergence"]
+__all__ = ["OTTOutput", "GWSolver", "SinkhornSolver", "OTTNeuralOutput", "sinkhorn_divergence", "GENOTLinSolver"]
+
 
 register_cost("euclidean", backend="ott")(costs.Euclidean)
 register_cost("sq_euclidean", backend="ott")(costs.SqEuclidean)

src/moscot/backends/ott/_utils.py

@@ -1,11 +1,16 @@
-from typing import Any, Literal, Optional, Tuple, Union
+from collections import defaultdict
+from functools import partial
+from typing import Any, Dict, Iterable, Literal, Optional, Tuple, Union
 
 import jax
 import jax.experimental.sparse as jesp
 import jax.numpy as jnp
+import numpy as np
 import scipy.sparse as sp
 from ott.geometry import epsilon_scheduler, geodesic, geometry, pointcloud
-from ott.tools import sinkhorn_divergence as sdiv
+from ott.neural import datasets
+from ott.solvers import utils as solver_utils
+from ott.tools.sinkhorn_divergence import sinkhorn_divergence as sinkhorn_div
 
 from moscot._logging import logger
 from moscot._types import ArrayLike, ScaleCost_t
@@ -22,22 +27,27 @@ def sinkhorn_divergence(
     a: Optional[ArrayLike] = None,
     b: Optional[ArrayLike] = None,
     epsilon: Union[float, epsilon_scheduler.Epsilon] = 1e-1,
+    tau_a: float = 1.0,
+    tau_b: float = 1.0,
     scale_cost: ScaleCost_t = 1.0,
+    batch_size: Optional[int] = None,
     **kwargs: Any,
 ) -> float:
     point_cloud_1 = jnp.asarray(point_cloud_1)
     point_cloud_2 = jnp.asarray(point_cloud_2)
     a = None if a is None else jnp.asarray(a)
     b = None if b is None else jnp.asarray(b)
 
-    output = sdiv.sinkhorn_divergence(
+    output = sinkhorn_div(
         pointcloud.PointCloud,
         x=point_cloud_1,
         y=point_cloud_2,
+        batch_size=batch_size,
         a=a,
         b=b,
-        epsilon=epsilon,
+        sinkhorn_kwargs={"tau_a": tau_a, "tau_b": tau_b},
         scale_cost=scale_cost,
+        epsilon=epsilon,
         **kwargs,
     )
     xy_conv, xx_conv, *yy_conv = output.converged
@@ -52,6 +62,17 @@ def sinkhorn_divergence(
     return float(output.divergence)
 
 
+@partial(jax.jit, static_argnames=["k"])
+def get_nearest_neighbors(
+    input_batch: jnp.ndarray, target: jnp.ndarray, k: int = 30
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Get the k nearest neighbors of the input batch in the target."""
+    if target.shape[0] < k:
+        raise ValueError(f"k is {k}, but must be smaller or equal than {target.shape[0]}.")
+    pairwise_euclidean_distances = pointcloud.PointCloud(input_batch, target).cost_matrix
+    return jax.lax.approx_min_k(pairwise_euclidean_distances, k=k, recall_target=0.95, aggregate_to_topk=True)
+
+
 def check_shapes(geom_x: geometry.Geometry, geom_y: geometry.Geometry, geom_xy: geometry.Geometry) -> None:
     n, m = geom_xy.shape
     n_, m_ = geom_x.shape[0], geom_y.shape[0]
@@ -133,3 +154,83 @@ def _instantiate_geodesic_cost(
     cm_full = geodesic.Geodesic.from_graph(arr, t=t, directed=directed, **kwargs).cost_matrix
     cm = cm_full[:n_src, n_src:] if is_linear_term else cm_full
     return geometry.Geometry(cm, epsilon=epsilon, relative_epsilon=relative_epsilon, scale_cost=scale_cost)
+
+
+def data_match_fn(
+    src_lin: Optional[jnp.ndarray] = None,
+    tgt_lin: Optional[jnp.ndarray] = None,
+    src_quad: Optional[jnp.ndarray] = None,
+    tgt_quad: Optional[jnp.ndarray] = None,
+    *,
+    typ: Literal["lin", "quad", "fused"],
+    **data_match_fn_kwargs,
+) -> jnp.ndarray:
+    if typ == "lin":
+        return solver_utils.match_linear(x=src_lin, y=tgt_lin, **data_match_fn_kwargs)
+    if typ == "quad":
+        return solver_utils.match_quadratic(xx=src_quad, yy=tgt_quad, **data_match_fn_kwargs)
+    if typ == "fused":
+        return solver_utils.match_quadratic(xx=src_quad, yy=tgt_quad, x=src_lin, y=tgt_lin, **data_match_fn_kwargs)
+    raise NotImplementedError(f"Unknown type: {typ}.")
+
+
+class Loader:
+
+    def __init__(self, dataset: datasets.OTDataset, batch_size: int, seed: Optional[int] = None):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self._rng = np.random.default_rng(seed)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> Dict[str, jnp.ndarray]:
+        data = defaultdict(list)
+        for _ in range(self.batch_size):
+            ix = self._rng.integers(0, len(self.dataset))
+            for k, v in self.dataset[ix].items():
+                data[k].append(v)
+        return {k: jnp.vstack(v) for k, v in data.items()}
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class MultiLoader:
+    """Dataset for OT problems with conditions.
+
+    This data loader wraps several data loaders and samples from them.
+
+    Args:
+      datasets: Datasets to sample from.
+      seed: Random seed.
+    """
+
+    def __init__(
+        self,
+        datasets: Iterable[Loader],
+        seed: Optional[int] = None,
+    ):
+        self.datasets = tuple(datasets)
+        self._rng = np.random.default_rng(seed)
+        self._iterators: list[MultiLoader] = []
+        self._it = 0
+
+    def __next__(self) -> Dict[str, jnp.ndarray]:
+        self._it += 1
+
+        ix = self._rng.choice(len(self._iterators))
+        iterator = self._iterators[ix]
+        if self._it < len(self):
+            return next(iterator)
+        # reset the consumed iterator and return it's first element
+        self._iterators[ix] = iterator = iter(self.datasets[ix])
+        return next(iterator)
+
+    def __iter__(self) -> "MultiLoader":
+        self._it = 0
+        self._iterators = [iter(ds) for ds in self.datasets]
+        return self
+
+    def __len__(self) -> int:
+        return max((len(ds) for ds in self.datasets), default=0)