moscot 0.4.0__tar.gz → 0.4.2__tar.gz

{moscot-0.4.0 → moscot-0.4.2}/.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.13.0
+    rev: v1.15.0
     hooks:
       - id: mypy
         additional_dependencies: [numpy>=1.25.0]
         files: ^src
   - repo: https://github.com/psf/black
-    rev: 24.10.0
+    rev: 25.1.0
     hooks:
       - id: black
         additional_dependencies: [toml]
@@ -23,7 +23,7 @@ repos:
       - id: prettier
         language_version: system
   - repo: https://github.com/PyCQA/isort
-    rev: 5.13.2
+    rev: 6.0.1
     hooks:
       - id: isort
         additional_dependencies: [toml]
@@ -42,7 +42,7 @@ repos:
       - id: check-yaml
       - id: check-toml
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.19.0
+    rev: v3.19.1
     hooks:
       - id: pyupgrade
         args: [--py3-plus, --py38-plus, --keep-runtime-typing]
@@ -55,7 +55,7 @@ repos:
     rev: v6.2.4
     hooks:
       - id: rstcheck
-        additional_dependencies: [tomli]
+        additional_dependencies: [toml, sphinx]
         args: [--config=pyproject.toml]
   - repo: https://github.com/PyCQA/doc8
     rev: v1.1.2
@@ -63,7 +63,7 @@ repos:
       - id: doc8
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.7.2
+    rev: v0.9.10
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{moscot-0.4.0 → moscot-0.4.2}/.readthedocs.yml

@@ -14,7 +14,7 @@ python:
   install:
     - method: pip
       path: .
-      extra_requirements: [docs]
+      extra_requirements: [docs, neural]
 
 submodules:
   include: [docs/notebooks]

{moscot-0.4.0 → moscot-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: moscot
-Version: 0.4.0
+Version: 0.4.2
 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>
@@ -65,7 +65,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[neural]>=0.5.0
+Requires-Dist: ott-jax>=0.5.0
 Requires-Dist: cloudpickle>=2.2.0
 Requires-Dist: rich>=13.5
 Requires-Dist: docstring_inheritance>=2.0.0
@@ -76,6 +76,7 @@ Provides-Extra: neural
 Requires-Dist: optax; extra == "neural"
 Requires-Dist: flax; extra == "neural"
 Requires-Dist: diffrax; extra == "neural"
+Requires-Dist: ott-jax[neural]>=0.5.0; extra == "neural"
 Provides-Extra: dev
 Requires-Dist: pre-commit>=3.0.0; extra == "dev"
 Requires-Dist: tox>=4; extra == "dev"
@@ -85,6 +86,7 @@ Requires-Dist: pytest-xdist>=3; extra == "test"
 Requires-Dist: pytest-mock>=3.5.0; extra == "test"
 Requires-Dist: pytest-cov>=4; extra == "test"
 Requires-Dist: coverage[toml]>=7; extra == "test"
+Requires-Dist: moscot[neural]; extra == "test"
 Provides-Extra: docs
 Requires-Dist: sphinx>=5.1.1; extra == "docs"
 Requires-Dist: sphinx_copybutton>=0.5.0; extra == "docs"
@@ -96,20 +98,37 @@ Requires-Dist: sphinx-tippy>=0.4.1; extra == "docs"
 Requires-Dist: myst-nb>=0.17.1; extra == "docs"
 Requires-Dist: ipython>=7.20.0; extra == "docs"
 Requires-Dist: sphinx_design>=0.3.0; extra == "docs"
+Dynamic: license-file
 
 |PyPI| |Downloads| |CI| |Pre-commit| |Codecov| |Docs|
 
-moscot - multi-omic single-cell optimal transport tools
+Moscot - Multiomics Single-cell Optimal Transport
 =======================================================
 
-**moscot** is a scalable framework for Optimal Transport (OT) applications in
-single-cell genomics. It can be used for
+.. image:: docs/_static/img/light_mode_concept_revised.png
+    :width: 800px
+    :align: center
+    :class: only-light
+
+.. image:: docs/_static/img/dark_mode_concept_revised.png
+    :width: 800px
+    :align: center
+    :class: only-dark
+
+
+**moscot** is a framework for Optimal Transport (OT) applications in
+single-cell genomics. It scales to large datasets and can be used for a
+variety of applications across different modalities.
+
+moscot's key applications
+---------------------------
+- Trajectory inference (incorporating spatial and lineage information).
+- Mapping cells to their spatial organisation.
+- Aligning spatial transcriptomics slides.
+- Translating modalities.
+- prototyping of new OT models in single-cell genomics.
+- ... and more, check out the `documentation <https://moscot.readthedocs.io>`_ for more information.
 
-- trajectory inference (incorporating spatial and lineage information)
-- mapping cells to their spatial organisation
-- aligning spatial transcriptomics slides
-- translating modalities
-- prototyping of new OT models in single-cell genomics
 
 **moscot** is powered by
 `OTT <https://ott-jax.readthedocs.io>`_ which is a JAX-based Optimal
@@ -118,7 +137,7 @@ differentiation and linear memory complexity for OT problems.
 
 Installation
 ------------
-You can install **moscot** via::
+Install **moscot** by running::
 
     pip install moscot
 
@@ -130,15 +149,10 @@ In order to install **moscot** from in editable mode, run::
 
 For further instructions how to install jax, please refer to https://github.com/google/jax.
 
-Resources
----------
-
-Please have a look at our `documentation <https://moscot.readthedocs.io>`_
-
-Reference
----------
-
-Our preprint "Mapping cells through time and space with moscot" can be found `here <https://www.biorxiv.org/content/10.1101/2023.05.11.540374v1>`_.
+Citing moscot
+-------------
+If you find a model useful for your research, please consider citing the `Klein et al., 2025`_ manuscript as
+well as the publication introducing the model, which can be found in the corresponding documentation.
 
 .. |Codecov| image:: https://codecov.io/gh/theislab/moscot/branch/master/graph/badge.svg?token=Rgtm5Tsblo
     :target: https://codecov.io/gh/theislab/moscot
@@ -163,3 +177,5 @@ Our preprint "Mapping cells through time and space with moscot" can be found `he
 .. |Downloads| image:: https://static.pepy.tech/badge/moscot
     :target: https://pepy.tech/project/moscot
     :alt: Downloads
+
+.. _Klein et al., 2025: https://www.nature.com/articles/s41586-024-08453-2

{moscot-0.4.0 → moscot-0.4.2}/README.rst

@@ -1,16 +1,32 @@
 |PyPI| |Downloads| |CI| |Pre-commit| |Codecov| |Docs|
 
-moscot - multi-omic single-cell optimal transport tools
+Moscot - Multiomics Single-cell Optimal Transport
 =======================================================
 
-**moscot** is a scalable framework for Optimal Transport (OT) applications in
-single-cell genomics. It can be used for
+.. image:: docs/_static/img/light_mode_concept_revised.png
+    :width: 800px
+    :align: center
+    :class: only-light
+
+.. image:: docs/_static/img/dark_mode_concept_revised.png
+    :width: 800px
+    :align: center
+    :class: only-dark
+
+
+**moscot** is a framework for Optimal Transport (OT) applications in
+single-cell genomics. It scales to large datasets and can be used for a
+variety of applications across different modalities.
+
+moscot's key applications
+---------------------------
+- Trajectory inference (incorporating spatial and lineage information).
+- Mapping cells to their spatial organisation.
+- Aligning spatial transcriptomics slides.
+- Translating modalities.
+- prototyping of new OT models in single-cell genomics.
+- ... and more, check out the `documentation <https://moscot.readthedocs.io>`_ for more information.
 
-- trajectory inference (incorporating spatial and lineage information)
-- mapping cells to their spatial organisation
-- aligning spatial transcriptomics slides
-- translating modalities
-- prototyping of new OT models in single-cell genomics
 
 **moscot** is powered by
 `OTT <https://ott-jax.readthedocs.io>`_ which is a JAX-based Optimal
@@ -19,7 +35,7 @@ differentiation and linear memory complexity for OT problems.
 
 Installation
 ------------
-You can install **moscot** via::
+Install **moscot** by running::
 
     pip install moscot
 
@@ -31,15 +47,10 @@ In order to install **moscot** from in editable mode, run::
 
 For further instructions how to install jax, please refer to https://github.com/google/jax.
 
-Resources
----------
-
-Please have a look at our `documentation <https://moscot.readthedocs.io>`_
-
-Reference
----------
-
-Our preprint "Mapping cells through time and space with moscot" can be found `here <https://www.biorxiv.org/content/10.1101/2023.05.11.540374v1>`_.
+Citing moscot
+-------------
+If you find a model useful for your research, please consider citing the `Klein et al., 2025`_ manuscript as
+well as the publication introducing the model, which can be found in the corresponding documentation.
 
 .. |Codecov| image:: https://codecov.io/gh/theislab/moscot/branch/master/graph/badge.svg?token=Rgtm5Tsblo
     :target: https://codecov.io/gh/theislab/moscot
@@ -64,3 +75,5 @@ Our preprint "Mapping cells through time and space with moscot" can be found `he
 .. |Downloads| image:: https://static.pepy.tech/badge/moscot
     :target: https://pepy.tech/project/moscot
     :alt: Downloads
+
+.. _Klein et al., 2025: https://www.nature.com/articles/s41586-024-08453-2

{moscot-0.4.0 → moscot-0.4.2}/pyproject.toml

@@ -54,7 +54,7 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax[neural]>=0.5.0",
+    "ott-jax>=0.5.0",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
     "docstring_inheritance>=2.0.0",
@@ -70,6 +70,7 @@ neural = [
     "optax",
     "flax",
     "diffrax",
+    "ott-jax[neural]>=0.5.0",
 
 ]
 dev = [
@@ -82,6 +83,7 @@ test = [
     "pytest-mock>=3.5.0",
     "pytest-cov>=4",
     "coverage[toml]>=7",
+    "moscot[neural]"
 ]
 docs = [
     "sphinx>=5.1.1",
@@ -274,7 +276,7 @@ env_list = lint-code,py{3.10,3.11,3.12}
 skip_missing_interpreters = true
 
 [testenv]
-extras = test
+extras = test,neural
 commands =
     python -m pytest {tty:--color=yes} {posargs: \
         --cov={env_site_packages_dir}{/}moscot --cov-config={tox_root}{/}pyproject.toml \
@@ -290,7 +292,7 @@ commands =
 
 [testenv:lint-docs]
 description = Lint the documentation.
-extras = docs
+extras = docs,neural
 ignore_errors = true
 allowlist_externals = make
 pass_env = PYENCHANT_LIBRARY_PATH
@@ -310,7 +312,7 @@ deps =
     jupytext
     nbconvert
     leidenalg
-extras = docs
+extras = docs,neural
 changedir = {tox_root}{/}docs
 commands =
     python -m ipykernel install --user --name=moscot
@@ -328,7 +330,7 @@ commands =
 [testenv:build-docs]
 description = Build the documentation.
 deps =
-extras = docs
+extras = docs,neural
 allowlist_externals = make
 changedir = {tox_root}{/}docs
 commands =

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/_types.py

@@ -2,19 +2,18 @@ import os
 from typing import Any, Literal, Mapping, Optional, Sequence, Union
 
 import numpy as np
+from jax import Array as JaxArray
+from numpy.typing import DTypeLike as DTypeLikeNumpy
+from numpy.typing import NDArray
 from ott.initializers.linear.initializers import SinkhornInitializer
 from ott.initializers.linear.initializers_lr import LRInitializer
 from ott.initializers.quadratic.initializers import BaseQuadraticInitializer
 
 # TODO(michalk8): polish
 
-try:
-    from numpy.typing import DTypeLike, NDArray
 
-    ArrayLike = NDArray[np.float64]
-except (ImportError, TypeError):
-    ArrayLike = np.ndarray  # type: ignore[misc]
-    DTypeLike = np.dtype  # type: ignore[misc]
+ArrayLike = Union[NDArray[np.floating], JaxArray]
+DTypeLike = DTypeLikeNumpy
 
 ProblemKind_t = Literal["linear", "quadratic", "unknown"]
 Numeric_t = Union[int, float]  # type of `time_key` arguments

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/backends/ott/__init__.py

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

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/backends/ott/_utils.py

@@ -184,7 +184,7 @@ def alpha_to_fused_penalty(alpha: float) -> float:
     return (1 - alpha) / alpha
 
 
-def densify(arr: ArrayLike) -> jax.Array:
+def densify(arr: Union[ArrayLike, sp.sparray, sp.spmatrix]) -> jax.Array:
     """If the input is sparse, convert it to dense.
 
     Parameters
@@ -197,7 +197,8 @@ def densify(arr: ArrayLike) -> jax.Array:
     dense :mod:`jax` array.
     """
     if sp.issparse(arr):
-        arr = arr.toarray()  # type: ignore[attr-defined]
+        arr_sp: Union[sp.sparray, sp.spmatrix] = arr
+        arr = arr_sp.toarray()
     elif isinstance(arr, jesp.BCOO):
         arr = arr.todense()
     return jnp.asarray(arr)

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/backends/ott/output.py

@@ -17,7 +17,7 @@ from moscot._types import ArrayLike, Device_t
 from moscot.backends.ott._utils import get_nearest_neighbors
 from moscot.base.output import BaseDiscreteSolverOutput, BaseNeuralOutput
 
-__all__ = ["OTTOutput", "GraphOTTOutput", "OTTNeuralOutput"]
+__all__ = ["OTTOutput", "GraphOTTOutput", "NeuralOutput"]
 
 
 class OTTOutput(BaseDiscreteSolverOutput):
@@ -182,6 +182,9 @@ class OTTOutput(BaseDiscreteSolverOutput):
             axis=1 - forward,
         ).T  # convert to batch first
 
+    def _apply_forward(self, x: ArrayLike) -> ArrayLike:
+        return self._apply(x, forward=True)
+
     @property
     def shape(self) -> Tuple[int, int]:  # noqa: D102
         if isinstance(self._output, sinkhorn.SinkhornOutput):
@@ -241,11 +244,11 @@ class OTTOutput(BaseDiscreteSolverOutput):
         return jnp.ones((n,))
 
 
-class OTTNeuralOutput(BaseNeuralOutput):
+class NeuralOutput(BaseNeuralOutput):
     """Output wrapper for GENOT."""
 
     def __init__(self, model: GENOT, logs: dict[str, list[float]]):
-        """Initialize `OTTNeuralOutput`.
+        """Initialize `NeuralOutput`.
 
         Parameters
         ----------
@@ -269,8 +272,7 @@ class OTTNeuralOutput(BaseNeuralOutput):
         self,
         src_dist: ArrayLike,
         tgt_dist: ArrayLike,
-        forward: bool,
-        func: Callable[[jnp.ndarray], jnp.ndarray],
+        func: Callable[[ArrayLike], ArrayLike],
         save_transport_matrix: bool = False,  # TODO(@MUCDK) adapt order of arguments
         batch_size: int = 1024,
         k: int = 30,
@@ -279,9 +281,9 @@ class OTTNeuralOutput(BaseNeuralOutput):
         recall_target: float = 0.95,
         aggregate_to_topk: bool = True,
     ) -> sp.csr_matrix:
-        row_indices: Union[jnp.ndarray, List[jnp.ndarray]] = []
-        column_indices: Union[jnp.ndarray, List[jnp.ndarray]] = []
-        distances_list: Union[jnp.ndarray, List[jnp.ndarray]] = []
+        row_indices: List[ArrayLike] = []
+        column_indices: List[ArrayLike] = []
+        distances_list: List[ArrayLike] = []
         if length_scale is None:
             key = jax.random.PRNGKey(seed)
             src_batch = src_dist[jax.random.choice(key, src_dist.shape[0], shape=((batch_size,)))]
@@ -306,20 +308,14 @@ class OTTNeuralOutput(BaseNeuralOutput):
         row_indices = jnp.concatenate(row_indices)
         column_indices = jnp.concatenate(column_indices)
         tm = sp.csr_matrix((distances, (row_indices, column_indices)), shape=[len(src_dist), len(tgt_dist)])
-        if forward:
-            if save_transport_matrix:
-                self._transport_matrix = tm
-        else:
-            tm = tm.T
-            if save_transport_matrix:
-                self._inverse_transport_matrix = tm
+        if save_transport_matrix:
+            self._transport_matrix = tm
         return tm
 
     def project_to_transport_matrix(  # type:ignore[override]
         self,
         src_cells: ArrayLike,
         tgt_cells: ArrayLike,
-        forward: bool = True,
         condition: ArrayLike = None,
         save_transport_matrix: bool = False,  # TODO(@MUCDK) adapt order of arguments
         batch_size: int = 1024,
@@ -351,7 +347,7 @@ class OTTNeuralOutput(BaseNeuralOutput):
         save_transport_matrix
             Whether to save the transport matrix.
         batch_size
-            Number of data points in the source distribution the neighborhoodgraph is computed
+            Number of data points in the source distribution the neighborhood graph is computed
             for in parallel.
         k
             Number of neighbors to construct the k-nearest neighbor graph of a mapped cell.
@@ -375,13 +371,12 @@ class OTTNeuralOutput(BaseNeuralOutput):
         The projected transport matrix.
         """
         src_cells, tgt_cells = jnp.asarray(src_cells), jnp.asarray(tgt_cells)
-        push = self.push if condition is None else lambda x: self.push(x, condition)
-        pull = self.pull if condition is None else lambda x: self.pull(x, condition)
-        func, src_dist, tgt_dist = (push, src_cells, tgt_cells) if forward else (pull, tgt_cells, src_cells)
+        conditioned_fn: Callable[[ArrayLike], ArrayLike] = lambda x: self.push(x, condition)
+        push = self.push if condition is None else conditioned_fn
+        func, src_dist, tgt_dist = (push, src_cells, tgt_cells)
         return self._project_transport_matrix(
             src_dist=src_dist,
             tgt_dist=tgt_dist,
-            forward=forward,
             func=func,
             save_transport_matrix=save_transport_matrix,  # TODO(@MUCDK) adapt order of arguments
             batch_size=batch_size,
@@ -406,31 +401,13 @@ class OTTNeuralOutput(BaseNeuralOutput):
         -------
         Pushed distribution.
         """
+        if isinstance(x, (bool, int, float, complex)):
+            raise ValueError("Expected array, found scalar value.")
         if x.ndim not in (1, 2):
             raise ValueError(f"Expected 1D or 2D array, found `{x.ndim}`.")
-        return self._apply(x, cond=cond, forward=True)
-
-    def pull(self, x: ArrayLike, cond: Optional[ArrayLike] = None) -> ArrayLike:
-        """Pull distribution `x` conditioned on condition `cond`.
-
-        This does not make sense for some neural models and is therefore left unimplemented.
-
-        Parameters
-        ----------
-         x
-            Distribution to push.
-        cond
-            Condition of conditional neural OT.
-
-        Raises
-        ------
-        NotImplementedError
-        """
-        raise NotImplementedError("`pull` does not make sense for neural OT.")
+        return self._apply_forward(x, cond=cond)
 
-    def _apply(self, x: ArrayLike, forward: bool, cond: Optional[ArrayLike] = None) -> ArrayLike:
-        if not forward:
-            raise NotImplementedError("Backward i.e., pull on neural OT is not supported.")
+    def _apply_forward(self, x: ArrayLike, cond: Optional[ArrayLike] = None) -> ArrayLike:
         return self._model.transport(x, condition=cond)
 
     @property
@@ -445,7 +422,7 @@ class OTTNeuralOutput(BaseNeuralOutput):
     def to(
         self,
         device: Optional[Device_t] = None,
-    ) -> "OTTNeuralOutput":
+    ) -> "NeuralOutput":
         """Transfer the output to another device or change its data type.
 
         Parameters
@@ -471,7 +448,7 @@ class OTTNeuralOutput(BaseNeuralOutput):
         #         raise IndexError(f"Unable to fetch the device with `id={idx}`.") from err
 
         # out = jax.device_put(self._model, device)
-        # return OTTNeuralOutput(out)
+        # return NeuralOutput(out)
         return self  # TODO(ilan-gold) move model to device
 
     @property

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/backends/ott/solver.py

@@ -53,7 +53,7 @@ from moscot.backends.ott._utils import (
     densify,
     ensure_2d,
 )
-from moscot.backends.ott.output import GraphOTTOutput, OTTNeuralOutput, OTTOutput
+from moscot.backends.ott.output import GraphOTTOutput, NeuralOutput, OTTOutput
 from moscot.base.problems._utils import TimeScalesHeatKernel
 from moscot.base.solver import OTSolver
 from moscot.costs import get_cost
@@ -216,7 +216,7 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
                 problem_shape = x.shape if problem_shape is None else problem_shape
                 return _instantiate_geodesic_cost(
                     arr=arr,
-                    problem_shape=problem_shape,
+                    problem_shape=problem_shape,  # type: ignore[arg-type]
                     t=t,
                     is_linear_term=is_linear_term,
                     epsilon=epsilon,
@@ -699,10 +699,10 @@ class GENOTLinSolver(OTSolver[OTTOutput]):
     def _call_kwargs(cls) -> Tuple[Set[str], Set[str]]:
         return {"batch_size", "train_size", "trainloader", "validloader", "seed"}, {}  # type: ignore[return-value]
 
-    def _solve(self, data_samplers: Tuple[MultiLoader, MultiLoader]) -> OTTNeuralOutput:  # type: ignore[override]
+    def _solve(self, data_samplers: Tuple[MultiLoader, MultiLoader]) -> NeuralOutput:  # type: ignore[override]
         seed = self._neural_kwargs.get("seed", 0)  # TODO(ilan-gold): unify rng hadnling like OTT tests
         rng = jax.random.PRNGKey(seed)
         logs = self.solver(
             data_samplers[0], n_iters=self._neural_kwargs.get("n_iters", 100), rng=rng
         )  # TODO(ilan-gold): validation and figure out defualts
-        return OTTNeuralOutput(self.solver, logs)
+        return NeuralOutput(self.solver, logs)

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/backends/utils.py

@@ -42,8 +42,7 @@ def register_solver(
     return _REGISTRY.register(backend)  # type: ignore[return-value]
 
 
-# TODO(@MUCDK) fix mypy error
-@register_solver("ott")  # type: ignore[arg-type]
+@register_solver("ott")
 def _(
     problem_kind: Literal["linear", "quadratic"],
     solver_name: Optional[Literal["GENOTLinSolver"]] = None,

{moscot-0.4.0 → moscot-0.4.2}/src/moscot/base/cost.py

@@ -58,7 +58,7 @@ class BaseCost(abc.ABC):
                 f"Cost matrix contains `{np.sum(np.isnan(cost))}` NaN values, "
                 f"setting them to the maximum value `{maxx}`."
             )
-            cost = np.nan_to_num(cost, nan=maxx)  # type: ignore[call-overload]
+            cost = np.nan_to_num(cost, nan=maxx)  # type: ignore[arg-type, type-var]
         if np.any(cost < 0):
             raise ValueError(f"Cost matrix contains `{np.sum(cost < 0)}` negative values.")
         return cost