moscot 0.3.3__tar.gz → 0.3.4__tar.gz

{moscot-0.3.3 → moscot-0.3.4}/.pre-commit-config.yaml

@@ -7,29 +7,29 @@ default_stages:
 minimum_pre_commit_version: 3.0.0
 repos:
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.5.1
+    rev: v1.8.0
     hooks:
       - id: mypy
         additional_dependencies: [numpy>=1.25.0]
         files: ^src
   - repo: https://github.com/psf/black
-    rev: 23.7.0
+    rev: 24.2.0
     hooks:
       - id: black
         additional_dependencies: [toml]
   - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v3.0.2
+    rev: v4.0.0-alpha.8
     hooks:
       - id: prettier
         language_version: system
   - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         additional_dependencies: [toml]
         args: [--order-by-type]
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: check-merge-conflict
       - id: check-ast
@@ -42,7 +42,7 @@ repos:
       - id: check-yaml
       - id: check-toml
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.10.1
+    rev: v3.15.1
     hooks:
       - id: pyupgrade
         args: [--py3-plus, --py38-plus, --keep-runtime-typing]
@@ -52,7 +52,7 @@ repos:
       - id: blacken-docs
         additional_dependencies: [black==23.1.0]
   - repo: https://github.com/rstcheck/rstcheck
-    rev: v6.1.2
+    rev: v6.2.0
     hooks:
       - id: rstcheck
         additional_dependencies: [tomli]
@@ -63,7 +63,7 @@ repos:
       - id: doc8
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.286
+    rev: v0.2.2
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{moscot-0.3.3 → moscot-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: moscot
-Version: 0.3.3
+Version: 0.3.4
 Summary: Multi-omic single-cell optimal transport tools
 Author: Dominik Klein, Giovanni Palla, Michal Klein, Zoe Piran, Marius Lange
 Maintainer-email: Dominik Klein <dominik.klein@helmholtz-muenchen.de>, Giovanni Palla <giovanni.palla@helmholtz-muenchen.de>, Michal Klein <michal.klein@helmholtz-muenchen.de>
@@ -66,7 +66,7 @@ Requires-Dist: anndata>=0.9.1
 Requires-Dist: scanpy>=1.9.3
 Requires-Dist: wrapt>=1.13.2
 Requires-Dist: docrep>=0.3.2
-Requires-Dist: ott-jax>=0.4.3
+Requires-Dist: ott-jax>=0.4.5
 Requires-Dist: cloudpickle>=2.2.0
 Requires-Dist: rich>=13.5
 Provides-Extra: spatial

{moscot-0.3.3 → moscot-0.3.4}/pyproject.toml

@@ -55,7 +55,7 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax>=0.4.3",
+    "ott-jax>=0.4.5",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
 ]
@@ -122,6 +122,8 @@ ignore = [
     "D107",
     # Missing docstring in magic method
     "D105",
+    # Use `X | Y` for type annotations
+    "UP007",
 ]
 line-length = 120
 select = [

{moscot-0.3.3 → moscot-0.3.4}/src/moscot/_types.py

@@ -8,7 +8,7 @@ import numpy as np
 try:
     from numpy.typing import DTypeLike, NDArray
 
-    ArrayLike = NDArray[np.float_]
+    ArrayLike = NDArray[np.float64]
 except (ImportError, TypeError):
     ArrayLike = np.ndarray  # type: ignore[misc]
     DTypeLike = np.dtype  # type: ignore[misc]
@@ -33,20 +33,20 @@ OttCostFn_t = Literal[
     "euclidean",
     "sq_euclidean",
     "cosine",
-    "PNormP",
-    "SqPNorm",
-    "Euclidean",
-    "SqEuclidean",
-    "Cosine",
-    "ElasticL1",
-    "ElasticSTVS",
-    "ElasticSqKOverlap",
+    "pnorm_p",
+    "sq_pnorm",
+    "cosine",
+    "elastic_l1",
+    "elastic_l2",
+    "elastic_stvs",
+    "elastic_sqk_overlap",
+    "geodesic",
 ]
 OttCostFnMap_t = Union[OttCostFn_t, Mapping[Literal["xy", "x", "y"], OttCostFn_t]]
 GenericCostFn_t = Literal["barcode_distance", "leaf_distance", "custom"]
 CostFn_t = Union[str, GenericCostFn_t, OttCostFn_t]
 CostFnMap_t = Union[Union[OttCostFn_t, GenericCostFn_t], Mapping[str, Union[OttCostFn_t, GenericCostFn_t]]]
-PathLike = Union[os.PathLike, str]
+PathLike = Union[os.PathLike, str]  # type: ignore[type-arg]
 Policy_t = Literal[
     "sequential",
     "star",

{moscot-0.3.3 → moscot-0.3.4}/src/moscot/backends/ott/__init__.py

@@ -1,11 +1,11 @@
 from ott.geometry import costs
 
 from moscot.backends.ott._utils import sinkhorn_divergence
-from moscot.backends.ott.output import OTTOutput
+from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
 from moscot.backends.ott.solver import GWSolver, SinkhornSolver
 from moscot.costs import register_cost
 
-__all__ = ["OTTOutput", "GWSolver", "SinkhornSolver", "sinkhorn_divergence"]
+__all__ = ["OTTOutput", "GraphOTTOutput", "GWSolver", "SinkhornSolver", "sinkhorn_divergence"]
 
 register_cost("euclidean", backend="ott")(costs.Euclidean)
 register_cost("sq_euclidean", backend="ott")(costs.SqEuclidean)
@@ -13,4 +13,6 @@ register_cost("cosine", backend="ott")(costs.Cosine)
 register_cost("pnorm_p", backend="ott")(costs.PNormP)
 register_cost("sq_pnorm", backend="ott")(costs.SqPNorm)
 register_cost("elastic_l1", backend="ott")(costs.ElasticL1)
+register_cost("elastic_l2", backend="ott")(costs.ElasticL2)
 register_cost("elastic_stvs", backend="ott")(costs.ElasticSTVS)
+register_cost("elastic_sqk_overlap", backend="ott")(costs.ElasticSqKOverlap)

{moscot-0.3.3 → moscot-0.3.4}/src/moscot/backends/ott/_utils.py

@@ -1,14 +1,17 @@
-from typing import Any, Optional, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 import scipy.sparse as sp
-from ott.geometry import epsilon_scheduler, geometry, pointcloud
+from ott.geometry import epsilon_scheduler, geodesic, geometry, pointcloud
 from ott.tools import sinkhorn_divergence as sdiv
 
 from moscot._logging import logger
 from moscot._types import ArrayLike, ScaleCost_t
 
+Scale_t = Union[float, Literal["mean", "median", "max_cost", "max_norm", "max_bound"]]
+
+
 __all__ = ["sinkhorn_divergence"]
 
 
@@ -86,3 +89,23 @@ def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
     if arr.ndim != 2:
         raise ValueError(f"Expected array to have 2 dimensions, found `{arr.ndim}`.")
     return arr
+
+
+def _instantiate_geodesic_cost(
+    arr: jax.Array,
+    problem_shape: Tuple[int, int],
+    t: Optional[float],
+    is_linear_term: bool,
+    epsilon: Union[float, epsilon_scheduler.Epsilon] = None,
+    relative_epsilon: Optional[bool] = None,
+    scale_cost: Scale_t = 1.0,
+    directed: bool = True,
+    **kwargs: Any,
+) -> geometry.Geometry:
+    n_src, n_tgt = problem_shape
+    if is_linear_term and n_src + n_tgt != arr.shape[0]:
+        raise ValueError(f"Expected `x` to have `{n_src + n_tgt}` points, found `{arr.shape[0]}`.")
+    t = epsilon / 4.0 if t is None else t
+    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)

{moscot-0.3.3 → moscot-0.3.4}/src/moscot/backends/ott/output.py

@@ -14,7 +14,7 @@ import matplotlib.pyplot as plt
 from moscot._types import ArrayLike, Device_t
 from moscot.base.output import BaseSolverOutput
 
-__all__ = ["OTTOutput"]
+__all__ = ["OTTOutput", "GraphOTTOutput"]
 
 
 class OTTOutput(BaseSolverOutput):
@@ -174,7 +174,10 @@ class OTTOutput(BaseSolverOutput):
     def _apply(self, x: ArrayLike, *, forward: bool) -> ArrayLike:
         if x.ndim == 1:
             return self._output.apply(x, axis=1 - forward)
-        return self._output.apply(x.T, axis=1 - forward).T  # convert to batch first
+        return self._output.apply(
+            x.T,
+            axis=1 - forward,
+        ).T  # convert to batch first
 
     @property
     def shape(self) -> Tuple[int, int]:  # noqa: D102
@@ -233,3 +236,63 @@ class OTTOutput(BaseSolverOutput):
 
     def _ones(self, n: int) -> ArrayLike:  # noqa: D102
         return jnp.ones((n,))
+
+
+class GraphOTTOutput(OTTOutput):
+    """Output of :term:`OT` problems with a graph geometry in the linear term.
+
+    Parameters
+    ----------
+    output
+        Output of the :mod:`ott` backend.
+    shape
+        Shape of the problem.
+    """
+
+    def __init__(
+        self,
+        output: Union[
+            sinkhorn.SinkhornOutput,
+            sinkhorn_lr.LRSinkhornOutput,
+            gromov_wasserstein.GWOutput,
+            gromov_wasserstein_lr.LRGWOutput,
+        ],
+        shape: Tuple[int, int],
+    ):
+        super().__init__(output)
+        self._shape = shape
+
+    @property
+    def shape(self) -> Tuple[int, int]:  # noqa: D102
+        return self._shape
+
+    def _expand_data(self, x: jnp.ndarray, forward: bool) -> jnp.ndarray:
+        if forward:
+            shape = (self.shape[1],) if x.ndim == 1 else (self.shape[1], x.shape[1])
+            return jnp.concatenate((x, jnp.zeros(shape)))
+        shape = (self.shape[0],) if x.ndim == 1 else (self.shape[0], x.shape[1])
+        return jnp.concatenate((jnp.zeros(shape), x))
+
+    def _apply(self, x: ArrayLike, *, forward: bool) -> ArrayLike:
+        x_expanded = self._expand_data(x, forward=forward)
+        # ott-jax only supports lse_mode=False with graph geometry
+        res = self._output.apply(x_expanded.T, axis=1 - forward, lse_mode=False).T
+        return res[len(x) :] if forward else res[: -len(x)]
+
+    def to(self, device: Optional[Device_t] = None) -> "GraphOTTOutput":  # noqa: D102
+        if device is None:
+            return GraphOTTOutput(jax.device_put(self._output, device=device), shape=self.shape)
+
+        if isinstance(device, str) and ":" in device:
+            device, ix = device.split(":")
+            idx = int(ix)
+        else:
+            idx = 0
+
+        if not isinstance(device, xla_ext.Device):
+            try:
+                device = jax.devices(device)[idx]
+            except IndexError:
+                raise IndexError(f"Unable to fetch the device with `id={idx}`.") from None
+
+        return GraphOTTOutput(jax.device_put(self._output, device), shape=self.shape)

{moscot-0.3.3 → moscot-0.3.4}/src/moscot/backends/ott/solver.py

@@ -4,15 +4,22 @@ import types
 from typing import Any, Literal, Mapping, Optional, Set, Tuple, Union
 
 import jax
-from ott.geometry import costs, epsilon_scheduler, geometry, pointcloud
+import jax.numpy as jnp
+from ott.geometry import costs, epsilon_scheduler, geodesic, geometry, pointcloud
 from ott.problems.linear import linear_problem
 from ott.problems.quadratic import quadratic_problem
 from ott.solvers.linear import sinkhorn, sinkhorn_lr
 from ott.solvers.quadratic import gromov_wasserstein, gromov_wasserstein_lr
 
 from moscot._types import ProblemKind_t, QuadInitializer_t, SinkhornInitializer_t
-from moscot.backends.ott._utils import alpha_to_fused_penalty, check_shapes, ensure_2d
-from moscot.backends.ott.output import OTTOutput
+from moscot.backends.ott._utils import (
+    _instantiate_geodesic_cost,
+    alpha_to_fused_penalty,
+    check_shapes,
+    ensure_2d,
+)
+from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
+from moscot.base.problems._utils import TimeScalesHeatKernel
 from moscot.base.solver import OTSolver
 from moscot.costs import get_cost
 from moscot.utils.tagged_array import TaggedArray
@@ -43,14 +50,21 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
         self._solver: Optional[OTTSolver_t] = None
         self._problem: Optional[OTTProblem_t] = None
         self._jit = jit
+        self._a: Optional[jnp.ndarray] = None
+        self._b: Optional[jnp.ndarray] = None
 
     def _create_geometry(
         self,
         x: TaggedArray,
+        *,
+        is_linear_term: bool,
         epsilon: Union[float, epsilon_scheduler.Epsilon] = None,
         relative_epsilon: Optional[bool] = None,
         scale_cost: Scale_t = 1.0,
         batch_size: Optional[int] = None,
+        problem_shape: Optional[Tuple[int, int]] = None,
+        t: Optional[float] = None,
+        directed: bool = True,
         **kwargs: Any,
     ) -> geometry.Geometry:
         if x.is_point_cloud:
@@ -88,17 +102,80 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
             return geometry.Geometry(
                 kernel_matrix=arr, epsilon=epsilon, relative_epsilon=relative_epsilon, scale_cost=scale_cost
             )
+        if x.is_graph:  # we currently only support this for the linear term.
+            return self._create_graph_geometry(
+                is_linear_term=is_linear_term,
+                x=x,
+                arr=arr,
+                problem_shape=problem_shape,
+                t=t,
+                epsilon=epsilon,
+                relative_epsilon=relative_epsilon,
+                scale_cost=scale_cost,
+                directed=directed,
+                **kwargs,
+            )
         raise NotImplementedError(f"Creating geometry from `tag={x.tag!r}` is not yet implemented.")
 
     def _solve(  # type: ignore[override]
         self,
         prob: OTTProblem_t,
         **kwargs: Any,
-    ) -> OTTOutput:
+    ) -> Union[OTTOutput, GraphOTTOutput]:
         solver = jax.jit(self.solver) if self._jit else self.solver
         out = solver(prob, **kwargs)
+        if isinstance(prob, linear_problem.LinearProblem) and isinstance(prob.geom, geodesic.Geodesic):
+            return GraphOTTOutput(out, shape=(len(self._a), len(self._b)))  # type: ignore[arg-type]
         return OTTOutput(out)
 
+    def _create_graph_geometry(
+        self,
+        is_linear_term: bool,
+        x: TaggedArray,
+        arr: jax.Array,
+        problem_shape: Optional[Tuple[int, int]],
+        t: Optional[float],
+        epsilon: Union[float, epsilon_scheduler.Epsilon] = None,
+        relative_epsilon: Optional[bool] = None,
+        scale_cost: Scale_t = 1.0,
+        directed: bool = True,
+        **kwargs: Any,
+    ) -> geometry.Geometry:
+        if x.cost == "geodesic":
+            if self.problem_kind == "linear":
+                if t is None:
+                    if epsilon is None:
+                        raise ValueError("`epsilon` cannot be `None`.")
+                    return geodesic.Geodesic.from_graph(arr, t=epsilon / 4.0, directed=directed, **kwargs)
+
+                return _instantiate_geodesic_cost(
+                    arr=arr,
+                    problem_shape=problem_shape,  # type: ignore[arg-type]
+                    t=t,
+                    is_linear_term=True,
+                    epsilon=epsilon,
+                    relative_epsilon=relative_epsilon,
+                    scale_cost=scale_cost,
+                    directed=directed,
+                    **kwargs,
+                )
+            if self.problem_kind == "quadratic":
+                problem_shape = x.shape if problem_shape is None else problem_shape
+                return _instantiate_geodesic_cost(
+                    arr=arr,
+                    problem_shape=problem_shape,
+                    t=t,
+                    is_linear_term=is_linear_term,
+                    epsilon=epsilon,
+                    relative_epsilon=relative_epsilon,
+                    scale_cost=scale_cost,
+                    directed=directed,
+                    **kwargs,
+                )
+
+            raise NotImplementedError(f"Invalid problem kind `{self.problem_kind}`.")
+        raise NotImplementedError(f"If the geometry is a graph, `cost` must be `geodesic`, found `{x.cost}`.")
+
     @property
     def solver(self) -> OTTSolver_t:
         """:mod:`ott` solver."""
@@ -165,6 +242,8 @@ class SinkhornSolver(OTTJaxSolver):
 
     def _prepare(
         self,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
         xy: Optional[TaggedArray] = None,
         x: Optional[TaggedArray] = None,
         y: Optional[TaggedArray] = None,
@@ -175,24 +254,35 @@ class SinkhornSolver(OTTJaxSolver):
         scale_cost: Scale_t = 1.0,
         cost_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         cost_matrix_rank: Optional[int] = None,
+        time_scales_heat_kernel: Optional[TimeScalesHeatKernel] = None,
         # problem
         **kwargs: Any,
     ) -> linear_problem.LinearProblem:
         del x, y
+        time_scales_heat_kernel = (
+            TimeScalesHeatKernel(None, None, None) if time_scales_heat_kernel is None else time_scales_heat_kernel
+        )
         if xy is None:
             raise ValueError(f"Unable to create geometry from `xy={xy}`.")
-
+        self._a = a
+        self._b = b
         geom = self._create_geometry(
             xy,
+            is_linear_term=True,
             epsilon=epsilon,
             relative_epsilon=relative_epsilon,
             batch_size=batch_size,
+            problem_shape=(len(self._a), len(self._b)),
             scale_cost=scale_cost,
+            t=time_scales_heat_kernel.xy,
             **cost_kwargs,
         )
         if cost_matrix_rank is not None:
             geom = geom.to_LRCGeometry(rank=cost_matrix_rank)
-        self._problem = linear_problem.LinearProblem(geom, **kwargs)
+        if isinstance(geom, geodesic.Geodesic):
+            a = jnp.concatenate((a, jnp.zeros_like(self._b)), axis=0)
+            b = jnp.concatenate((jnp.zeros_like(self._a), b), axis=0)
+        self._problem = linear_problem.LinearProblem(geom, a=a, b=b, **kwargs)
         return self._problem
 
     @property
@@ -206,7 +296,15 @@ class SinkhornSolver(OTTJaxSolver):
 
     @classmethod
     def _call_kwargs(cls) -> Tuple[Set[str], Set[str]]:
-        geom_kwargs = {"epsilon", "relative_epsilon", "batch_size", "scale_cost", "cost_kwargs", "cost_matrix_rank"}
+        geom_kwargs = {
+            "epsilon",
+            "relative_epsilon",
+            "batch_size",
+            "scale_cost",
+            "cost_kwargs",
+            "cost_matrix_rank",
+            "t",
+        }
         problem_kwargs = set(inspect.signature(linear_problem.LinearProblem).parameters.keys())
         problem_kwargs -= {"geom"}
         return geom_kwargs | problem_kwargs, {"epsilon"}
@@ -270,6 +368,8 @@ class GWSolver(OTTJaxSolver):
 
     def _prepare(
         self,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
         xy: Optional[TaggedArray] = None,
         x: Optional[TaggedArray] = None,
         y: Optional[TaggedArray] = None,
@@ -280,32 +380,45 @@ class GWSolver(OTTJaxSolver):
         scale_cost: Scale_t = 1.0,
         cost_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         cost_matrix_rank: Optional[int] = None,
+        time_scales_heat_kernel: Optional[TimeScalesHeatKernel] = None,
         # problem
         alpha: float = 0.5,
         **kwargs: Any,
     ) -> quadratic_problem.QuadraticProblem:
+        self._a = a
+        self._b = b
+        time_scales_heat_kernel = (
+            TimeScalesHeatKernel(None, None, None) if time_scales_heat_kernel is None else time_scales_heat_kernel
+        )
         if x is None or y is None:
             raise ValueError(f"Unable to create geometry from `x={x}`, `y={y}`.")
-        geom_kwargs: Any = {
+        geom_kwargs: dict[str, Any] = {
             "epsilon": epsilon,
             "relative_epsilon": relative_epsilon,
             "batch_size": batch_size,
             "scale_cost": scale_cost,
-            "cost_matrix_rank": cost_matrix_rank,
             **cost_kwargs,
         }
-        geom_xx = self._create_geometry(x, **geom_kwargs)
-        geom_yy = self._create_geometry(y, **geom_kwargs)
+        if cost_matrix_rank is not None:
+            geom_kwargs["cost_matrix_rank"] = cost_matrix_rank
+        geom_xx = self._create_geometry(x, t=time_scales_heat_kernel.x, is_linear_term=False, **geom_kwargs)
+        geom_yy = self._create_geometry(y, t=time_scales_heat_kernel.y, is_linear_term=False, **geom_kwargs)
         if alpha == 1.0 or xy is None:  # GW
             # arbitrary fused penalty; must be positive
             geom_xy, fused_penalty = None, 1.0
         else:  # FGW
             fused_penalty = alpha_to_fused_penalty(alpha)
-            geom_xy = self._create_geometry(xy, **geom_kwargs)
+            geom_xy = self._create_geometry(
+                xy,
+                t=time_scales_heat_kernel.xy,
+                problem_shape=(x.shape[0], y.shape[0]),
+                is_linear_term=True,
+                **geom_kwargs,
+            )
             check_shapes(geom_xx, geom_yy, geom_xy)
 
         self._problem = quadratic_problem.QuadraticProblem(
-            geom_xx, geom_yy, geom_xy, fused_penalty=fused_penalty, **kwargs
+            geom_xx, geom_yy, geom_xy, fused_penalty=fused_penalty, a=self._a, b=self._b, **kwargs
         )
         return self._problem