moscot 0.3.4__tar.gz → 0.3.5__tar.gz

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

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

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

@@ -7,13 +7,13 @@ default_stages:
 minimum_pre_commit_version: 3.0.0
 repos:
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.8.0
+    rev: v1.10.1
     hooks:
       - id: mypy
         additional_dependencies: [numpy>=1.25.0]
         files: ^src
   - repo: https://github.com/psf/black
-    rev: 24.2.0
+    rev: 24.4.2
     hooks:
       - id: black
         additional_dependencies: [toml]
@@ -29,7 +29,7 @@ repos:
         additional_dependencies: [toml]
         args: [--order-by-type]
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
       - id: check-merge-conflict
       - id: check-ast
@@ -42,12 +42,12 @@ repos:
       - id: check-yaml
       - id: check-toml
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.15.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]
@@ -63,7 +63,7 @@ repos:
       - id: doc8
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.2.2
+    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.4 → moscot-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: moscot
-Version: 0.3.4
+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.5
+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.4 → 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.5",
+    "ott-jax>=0.4.6",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
+    "docstring_inheritance>=2.0.0"
 ]
 
 [project.optional-dependencies]
@@ -282,7 +283,6 @@ commands =
 
 [testenv:lint-docs]
 description = Lint the documentation.
-deps =
 extras = docs
 ignore_errors = true
 allowlist_externals = make
@@ -294,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.4 → moscot-0.3.5}/src/moscot/_types.py

@@ -36,10 +36,6 @@ OttCostFn_t = Literal[
     "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]]

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

@@ -12,7 +12,3 @@ 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_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.4 → moscot-0.3.5}/src/moscot/backends/ott/_utils.py

@@ -1,6 +1,7 @@
 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, geodesic, geometry, pointcloud
@@ -67,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.
 
@@ -81,9 +101,6 @@ 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:
@@ -91,6 +108,13 @@ def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
     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],

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

@@ -16,6 +16,8 @@ from moscot.backends.ott._utils import (
     _instantiate_geodesic_cost,
     alpha_to_fused_penalty,
     check_shapes,
+    convert_scipy_sparse,
+    densify,
     ensure_2d,
 )
 from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
@@ -76,8 +78,8 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
             if not isinstance(cost_fn, costs.CostFn):
                 raise TypeError(f"Expected `cost_fn` to be `ott.geometry.costs.CostFn`, found `{type(cost_fn)}`.")
 
-            y = None if x.data_tgt is None else ensure_2d(x.data_tgt, reshape=True)
-            x = ensure_2d(x.data_src, reshape=True)
+            y = None if x.data_tgt is None else densify(ensure_2d(x.data_tgt, reshape=True))
+            x = densify(ensure_2d(x.data_src, reshape=True))
             if y is not None and x.shape[1] != y.shape[1]:
                 raise ValueError(
                     f"Expected `x/y` to have the same number of dimensions, found `{x.shape[1]}/{y.shape[1]}`."
@@ -94,6 +96,8 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
             )
 
         arr = ensure_2d(x.data_src, reshape=False)
+        arr = densify(arr) if x.is_graph else convert_scipy_sparse(arr)
+
         if x.is_cost_matrix:
             return geometry.Geometry(
                 cost_matrix=arr, epsilon=epsilon, relative_epsilon=relative_epsilon, scale_cost=scale_cost

{moscot-0.3.4 → moscot-0.3.5}/src/moscot/base/problems/_utils.py

@@ -158,6 +158,17 @@ def _validate_args_cell_transition(
     raise TypeError(f"Expected argument to be either `str` or `dict`, found `{type(arg)}`.")
 
 
+def _assert_series_match(a: pd.Series, b: pd.Series) -> None:
+    """Assert that two series are equal ignoring the names."""
+    pd.testing.assert_series_equal(a, b, check_names=False)
+
+
+def _assert_columns_and_index_match(a: pd.Series, b: pd.DataFrame) -> None:
+    """Assert that a series and a dataframe's index and columns are matching."""
+    _assert_series_match(a, b.index.to_series())
+    _assert_series_match(a, b.columns.to_series())
+
+
 def _get_cell_indices(
     adata: AnnData,
     key: Optional[str] = None,

{moscot-0.3.4 → moscot-0.3.5}/src/moscot/base/problems/birth_death.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -166,7 +167,14 @@ class BirthDeathProblem(BirthDeathMixin, OTProblem):
         proliferation_key: Optional[str] = None,
         apoptosis_key: Optional[str] = None,
         scaling: Optional[float] = None,
-        **kwargs: Any,
+        beta_max: float = 1.7,
+        beta_min: float = 0.3,
+        beta_center: float = 0.25,
+        beta_width: float = 0.5,
+        delta_max: float = 1.7,
+        delta_min: float = 0.3,
+        delta_center: float = 0.1,
+        delta_width: float = 0.2,
     ) -> ArrayLike:
         """Estimate the source or target :term:`marginals` based on marker genes, either with the
         `birth-death process <https://en.wikipedia.org/wiki/Birth%E2%80%93death_process>`_,
@@ -189,9 +197,22 @@ class BirthDeathProblem(BirthDeathMixin, OTProblem):
             If :obj:`float` is passed, it will be used as a scaling parameter in an exponential kernel
             with proliferation and apoptosis scores.
             If :obj:`None`, parameters corresponding to the birth and death processes will be used.
-        kwargs
-            Keyword arguments for :func:`~moscot.base.problems.birth_death.beta` and
-            :func:`~moscot.base.problems.birth_death.delta`.
+        beta_max
+            Argument for :func:`~moscot.base.problems.birth_death.beta`
+        beta_min
+            Argument for :func:`~moscot.base.problems.birth_death.beta`
+        beta_center
+            Argument for :func:`~moscot.base.problems.birth_death.beta`
+        beta_width
+            Argument for :func:`~moscot.base.problems.birth_death.beta`
+        delta_max
+            Argument for :func:`~moscot.base.problems.birth_death.delta`
+        delta_min
+            Argument for :func:`~moscot.base.problems.birth_death.delta`
+        delta_center
+            Argument for :func:`~moscot.base.problems.birth_death.delta`
+        delta_width
+            Argument for :func:`~moscot.base.problems.birth_death.delta`
 
         Returns
         -------
@@ -223,12 +244,18 @@ class BirthDeathProblem(BirthDeathMixin, OTProblem):
         self.apoptosis_key = apoptosis_key
 
         if scaling:
-            beta_fn = delta_fn = lambda x, *_, **__: x
+            beta_fn = delta_fn = lambda x: x
         else:
-            beta_fn, delta_fn = beta, delta
+            beta_fn = partial(
+                beta, beta_max=beta_max, beta_min=beta_min, beta_center=beta_center, beta_width=beta_width
+            )
+            delta_fn = partial(
+                delta, delta_max=delta_max, delta_min=delta_min, delta_center=delta_center, delta_width=delta_width
+            )
+
             scaling = 1.0
-        birth = estimate(proliferation_key, fn=beta_fn, **kwargs)
-        death = estimate(apoptosis_key, fn=delta_fn, **kwargs)
+        birth = estimate(proliferation_key, fn=beta_fn)
+        death = estimate(apoptosis_key, fn=delta_fn)
 
         prior_growth = np.exp((birth - death) * self.delta / scaling)
 
@@ -287,7 +314,6 @@ def beta(
     beta_min: float = 0.3,
     beta_center: float = 0.25,
     beta_width: float = 0.5,
-    **_: Any,
 ) -> ArrayLike:
     """Birth process."""
     return _gen_logistic(p, beta_max, beta_min, beta_center, beta_width)
@@ -299,7 +325,6 @@ def delta(
     delta_min: float = 0.3,
     delta_center: float = 0.1,
     delta_width: float = 0.2,
-    **_: Any,
 ) -> ArrayLike:
     """Death process."""
     return _gen_logistic(a, delta_max, delta_min, delta_center, delta_width)

{moscot-0.3.4 → moscot-0.3.5}/src/moscot/base/problems/compound_problem.py

@@ -149,9 +149,10 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
         xy_callback_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         x_callback_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         y_callback_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
-        **kwargs: Any,
+        a: Optional[Union[bool, str, ArrayLike]] = None,
+        b: Optional[Union[bool, str, ArrayLike]] = None,
+        marginal_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
     ) -> Dict[Tuple[K, K], B]:
-        from moscot.base.problems.birth_death import BirthDeathProblem
 
         if TYPE_CHECKING:
             assert isinstance(self._policy, SubsetPolicy)
@@ -187,10 +188,7 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
             if y_data:
                 y = dict(y)
                 y["tagged_array"] = y_data
-            if isinstance(problem, BirthDeathProblem):
-                kwargs["proliferation_key"] = self.proliferation_key  # type: ignore[attr-defined]
-                kwargs["apoptosis_key"] = self.apoptosis_key  # type: ignore[attr-defined]
-            problems[src_name, tgt_name] = problem.prepare(xy=xy, x=x, y=y, **kwargs)
+            problems[src_name, tgt_name] = problem.prepare(xy=xy, x=x, y=y, a=a, b=b, marginal_kwargs=marginal_kwargs)
 
         return problems
 
@@ -200,13 +198,18 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
         key: Optional[str],
         subset: Optional[Sequence[Tuple[K, K]]] = None,
         reference: Optional[Any] = None,
+        xy: Mapping[str, Any] = types.MappingProxyType({}),
+        x: Mapping[str, Any] = types.MappingProxyType({}),
+        y: Mapping[str, Any] = types.MappingProxyType({}),
         xy_callback: Optional[Union[Literal["local-pca"], Callback_t]] = None,
         x_callback: Optional[Union[Literal["local-pca"], Callback_t]] = None,
         y_callback: Optional[Union[Literal["local-pca"], Callback_t]] = None,
         xy_callback_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         x_callback_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         y_callback_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
-        **kwargs: Any,
+        a: Optional[Union[bool, str, ArrayLike]] = None,
+        b: Optional[Union[bool, str, ArrayLike]] = None,
+        marginal_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
     ) -> "BaseCompoundProblem[K, B]":
         """Prepare the individual :term:`OT` subproblems.
 
@@ -224,6 +227,12 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
             for the :class:`~moscot.utils.subset_policy.ExplicitPolicy`. Only used when ``policy = 'explicit'``.
         reference
             Reference for the :class:`~moscot.utils.subset_policy.SubsetPolicy`. Only used when ``policy = 'star'``.
+        xy
+            Data for the :term:`linear term`.
+        x
+            Data for the source :term:`quadratic term`.
+        y
+            Data for the target :term:`quadratic term`.
         xy_callback
             Callback function used to prepare the data in the :term:`linear term`.
         x_callback
@@ -236,8 +245,24 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
             Keyword arguments for the ``x_callback``.
         y_callback_kwargs
             Keyword arguments for the ``y_callback``.
-        kwargs
-            Keyword arguments for the subproblems' :meth:`~moscot.base.problems.OTProblem.prepare` method.
+        a
+            Source :term:`marginals`. Valid options are:
+
+            - :class:`str` - key in :attr:`~anndata.AnnData.obs` where the source marginals are stored.
+            - :class:`bool` - if :obj:`True`,
+              :meth:`estimate the marginals <moscot.base.problems.OTProblem.estimate_marginals>`,
+              otherwise use uniform marginals.
+            - :obj:`None` - uniform marginals.
+        b
+            Target :term:`marginals`. Valid options are:
+
+            - :class:`str` - key in :attr:`~anndata.AnnData.obs` where the target marginals are stored.
+            - :class:`bool` - if :obj:`True`,
+              :meth:`estimate the marginals <moscot.base.problems.OTProblem.estimate_marginals>`,
+              otherwise use uniform marginals.
+            - :obj:`None` - uniform marginals.
+        marginal_kwargs
+            Keyword arguments for the :meth:`~moscot.base.problems.OTProblem.estimate_marginals` method.
 
         Returns
         -------
@@ -264,13 +289,18 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
         # when refactoring the callback, consider changing this
         self._problem_manager = ProblemManager(self, policy=policy)
         problems = self._create_problems(
-            xy_callback=xy_callback,
+            x=x,
+            y=y,
+            xy=xy,
+            a=a,
+            b=b,
             x_callback=x_callback,
             y_callback=y_callback,
-            xy_callback_kwargs=xy_callback_kwargs,
+            xy_callback=xy_callback,
             x_callback_kwargs=x_callback_kwargs,
             y_callback_kwargs=y_callback_kwargs,
-            **kwargs,
+            xy_callback_kwargs=xy_callback_kwargs,
+            marginal_kwargs=marginal_kwargs,
         )
         self._problem_manager.add_problems(problems)
 
@@ -313,6 +343,11 @@ class BaseCompoundProblem(BaseProblem, abc.ABC, Generic[K, B]):
         problems = self._problem_manager.get_problems(stage=stage)
 
         logger.info(f"Solving `{len(problems)}` problems")
+        # expose min/max iterations to the user but remove them if they are None
+        if "min_iterations" in kwargs and kwargs["min_iterations"] is None:
+            kwargs.pop("min_iterations")
+        if "max_iterations" in kwargs and kwargs["max_iterations"] is None:
+            kwargs.pop("max_iterations")
         for problem in problems.values():
             logger.info(f"Solving problem {problem}.")
             _ = problem.solve(**kwargs)
@@ -587,14 +622,14 @@ class CompoundProblem(BaseCompoundProblem[K, B], abc.ABC):
             linear_cost_matrix = data[mask, :][:, mask_2]
             if sp.issparse(linear_cost_matrix):
                 logger.warning("Linear cost matrix being densified.")
-                linear_cost_matrix = linear_cost_matrix.A
+                linear_cost_matrix = linear_cost_matrix.toarray()
             return TaggedArray(linear_cost_matrix, tag=Tag.COST_MATRIX)
 
         if term in ("x", "y"):
             quad_cost_matrix = data[mask, :][:, mask]
             if sp.issparse(quad_cost_matrix):
                 logger.warning("Quadratic cost matrix being densified.")
-                quad_cost_matrix = quad_cost_matrix.A
+                quad_cost_matrix = quad_cost_matrix.toarray()
             return TaggedArray(quad_cost_matrix, tag=Tag.COST_MATRIX)
 
         raise ValueError(f"Expected `term` to be one of `x`, `y`, or `xy`, found `{term!r}`.")