moscot 0.3.3__tar.gz → 0.3.5__tar.gz

{moscot-0.3.3 → moscot-0.3.5}/.gitignore

@@ -154,3 +154,4 @@ packages.dot
 
 # plotting tests
 tests/plotting/actual_figures/
+tests/plotting/figures/

{moscot-0.3.3 → moscot-0.3.5}/.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.10.1
     hooks:
       - id: mypy
         additional_dependencies: [numpy>=1.25.0]
         files: ^src
   - repo: https://github.com/psf/black
-    rev: 23.7.0
+    rev: 24.4.2
     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.6.0
     hooks:
       - id: check-merge-conflict
       - id: check-ast
@@ -42,17 +42,17 @@ repos:
       - id: check-yaml
       - id: check-toml
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.10.1
+    rev: v3.16.0
     hooks:
       - id: pyupgrade
         args: [--py3-plus, --py38-plus, --keep-runtime-typing]
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.16.0
+    rev: 1.18.0
     hooks:
       - 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.5.0
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

moscot-0.3.5/.run_notebooks.sh

@@ -0,0 +1,62 @@
+#!/bin/bash
+
+# Check if the base directory is provided as an argument
+if [ "$#" -ne 1 ]; then
+    echo "Usage: $0 <base_notebook_directory>"
+    exit 1
+fi
+
+# Base directory for notebooks
+base_dir=$1
+
+# Define notebook directories or patterns
+declare -a notebooks=(
+    "$base_dir/examples/plotting/*.ipynb"
+    "$base_dir/examples/problems/*.ipynb"
+    "$base_dir/examples/solvers/*.ipynb"
+)
+
+# Initialize an array to hold valid notebook paths
+declare -a valid_notebooks
+
+# Gather all valid notebook files from the patterns
+echo "Gathering notebooks..."
+for pattern in "${notebooks[@]}"; do
+    for nb in $pattern; do
+        if [[ -f "$nb" ]]; then  # Check if the file exists
+            valid_notebooks+=("$nb")  # Add to the list of valid notebooks
+        fi
+    done
+done
+
+# Check if we have any notebooks to run
+if [ ${#valid_notebooks[@]} -eq 0 ]; then
+    echo "No notebooks found to run."
+    exit 1
+fi
+
+# Echo the notebooks that will be run for clarity
+echo "Preparing to run the following notebooks:"
+for nb in "${valid_notebooks[@]}"; do
+    echo "$nb"
+done
+
+# Initialize a flag to track the success of all commands
+all_success=true
+
+# Execute all valid notebooks
+for nb in "${valid_notebooks[@]}"; do
+    echo "Running $nb"
+    jupytext -k moscot --execute "$nb" || {
+        echo "Failed to run $nb"
+        all_success=false
+    }
+done
+
+# Check if any executions failed
+if [ "$all_success" = false ]; then
+    echo "One or more notebooks failed to execute."
+    exit 1
+fi
+
+echo "All notebooks executed successfully."

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: moscot
-Version: 0.3.3
+Version: 0.3.5
 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,9 +66,10 @@ 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.6
 Requires-Dist: cloudpickle>=2.2.0
 Requires-Dist: rich>=13.5
+Requires-Dist: docstring_inheritance>=2.0.0
 Provides-Extra: spatial
 Requires-Dist: squidpy>=1.2.3; extra == "spatial"
 Provides-Extra: dev

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

@@ -55,9 +55,10 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax>=0.4.3",
+    "ott-jax>=0.4.6",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
+    "docstring_inheritance>=2.0.0"
 ]
 
 [project.optional-dependencies]
@@ -122,6 +123,8 @@ ignore = [
     "D107",
     # Missing docstring in magic method
     "D105",
+    # Use `X | Y` for type annotations
+    "UP007",
 ]
 line-length = 120
 select = [
@@ -280,7 +283,6 @@ commands =
 
 [testenv:lint-docs]
 description = Lint the documentation.
-deps =
 extras = docs
 ignore_errors = true
 allowlist_externals = make
@@ -292,6 +294,21 @@ commands =
     # TODO(michalk8): uncomment after https://github.com/theislab/moscot/issues/490
     # make spelling {posargs}
 
+[testenv:examples-docs]
+allowlist_externals = bash
+description = Run the notebooks.
+use_develop = true
+deps =
+    ipykernel
+    jupytext
+    nbconvert
+    leidenalg
+extras = docs
+changedir = {tox_root}{/}docs
+commands =
+    python -m ipykernel install --user --name=moscot
+    bash {tox_root}/.run_notebooks.sh {tox_root}{/}docs/notebooks
+
 [testenv:clean-docs]
 description = Remove the documentation.
 deps =

{moscot-0.3.3 → moscot-0.3.5}/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,16 @@ OttCostFn_t = Literal[
     "euclidean",
     "sq_euclidean",
     "cosine",
-    "PNormP",
-    "SqPNorm",
-    "Euclidean",
-    "SqEuclidean",
-    "Cosine",
-    "ElasticL1",
-    "ElasticSTVS",
-    "ElasticSqKOverlap",
+    "pnorm_p",
+    "sq_pnorm",
+    "cosine",
+    "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.5}/src/moscot/backends/ott/__init__.py

@@ -1,16 +1,14 @@
 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)
 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_stvs", backend="ott")(costs.ElasticSTVS)

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

@@ -1,14 +1,18 @@
-from typing import Any, Optional, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 import jax
+import jax.experimental.sparse as jesp
 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"]
 
 
@@ -64,6 +68,25 @@ def alpha_to_fused_penalty(alpha: float) -> float:
     return (1 - alpha) / alpha
 
 
+def densify(arr: ArrayLike) -> jax.Array:
+    """If the input is sparse, convert it to dense.
+
+    Parameters
+    ----------
+    arr
+        Array to check.
+
+    Returns
+    -------
+    dense :mod:`jax` array.
+    """
+    if sp.issparse(arr):
+        arr = arr.toarray()  # type: ignore[attr-defined]
+    elif isinstance(arr, jesp.BCOO):
+        arr = arr.todense()
+    return jnp.asarray(arr)
+
+
 def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
     """Ensure that an array is 2-dimensional.
 
@@ -78,11 +101,35 @@ def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
     -------
     2-dimensional :mod:`jax` array.
     """
-    if sp.issparse(arr):
-        arr = arr.A  # type: ignore[attr-defined]
-    arr = jnp.asarray(arr)
     if reshape and arr.ndim == 1:
         return jnp.reshape(arr, (-1, 1))
     if arr.ndim != 2:
         raise ValueError(f"Expected array to have 2 dimensions, found `{arr.ndim}`.")
     return arr
+
+
+def convert_scipy_sparse(arr: Union[sp.spmatrix, jesp.BCOO]) -> jesp.BCOO:
+    """If the input is a scipy sparse matrix, convert it to a jax BCOO."""
+    if sp.issparse(arr):
+        return jesp.BCOO.from_scipy_sparse(arr)
+    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.5}/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)