moscot 0.2.0__tar.gz → 0.3.1__tar.gz

{moscot-0.2.0 → moscot-0.3.1}/.pre-commit-config.yaml

@@ -7,22 +7,22 @@ default_stages:
 minimum_pre_commit_version: 3.0.0
 repos:
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.1.1
+    rev: v1.4.1
     hooks:
       - id: mypy
-        additional_dependencies: [numpy>=1.21.0, jax]
+        additional_dependencies: [numpy>=1.25.0]
         files: ^src
   - repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 23.7.0
     hooks:
       - id: black
         additional_dependencies: [toml]
   - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v3.0.0-alpha.6
+    rev: v3.0.0
     hooks:
       - id: prettier
         language_version: system
-  - repo: https://github.com/timothycrosley/isort
+  - repo: https://github.com/PyCQA/isort
     rev: 5.12.0
     hooks:
       - id: isort
@@ -42,12 +42,12 @@ repos:
       - id: check-yaml
       - id: check-toml
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.3.1
+    rev: v3.9.0
     hooks:
       - id: pyupgrade
         args: [--py3-plus, --py38-plus, --keep-runtime-typing]
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.13.0
+    rev: 1.15.0
     hooks:
       - id: blacken-docs
         additional_dependencies: [black==23.1.0]
@@ -61,9 +61,9 @@ repos:
     rev: v1.1.1
     hooks:
       - id: doc8
-  - repo: https://github.com/charliermarsh/ruff-pre-commit
+  - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.257
+    rev: v0.0.280
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{moscot-0.2.0 → moscot-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: moscot
-Version: 0.2.0
+Version: 0.3.1
 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>
@@ -62,16 +62,18 @@ Provides-Extra: test
 Provides-Extra: docs
 License-File: LICENSE
 
-|Codecov|
+|PyPI| |Downloads| |CI| |Pre-commit| |Codecov| |Docs|
 
 moscot - multi-omic single-cell optimal transport tools
 =======================================================
 
 **moscot** is a scalable framework for Optimal Transport (OT) applications in
 single-cell genomics. It can be used for
-- temporal and spatio-temporal trajectory inference
-- spatial mapping
-- spatial alignment
+
+- 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
@@ -85,27 +87,44 @@ You can install **moscot** via::
 
     pip install moscot
 
-In order to install **moscot** from source, run::
+In order to install **moscot** from in editable mode, run::
 
     git clone https://github.com/theislab/moscot
     cd moscot
-    pip install -e .'[dev]'
+    pip install -e .
+
+For further instructions how to install jax, please refer to https://github.com/google/jax.
 
-If used with GPU, additionally run::
+Resources
+---------
 
-    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+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>`_.
 
 .. |Codecov| image:: https://codecov.io/gh/theislab/moscot/branch/master/graph/badge.svg?token=Rgtm5Tsblo
     :target: https://codecov.io/gh/theislab/moscot
     :alt: Coverage
 
-Resources
----------
+.. |PyPI| image:: https://img.shields.io/pypi/v/moscot.svg
+    :target: https://pypi.org/project/moscot/
+    :alt: PyPI
 
-Please have a look at our `documentation <https://moscot.readthedocs.io>`_
+.. |CI| image:: https://img.shields.io/github/actions/workflow/status/theislab/moscot/test.yml?branch=main
+    :target: https://github.com/theislab/moscot/actions
+    :alt: CI
 
-Reference
----------
+.. |Pre-commit| image:: https://results.pre-commit.ci/badge/github/theislab/moscot/main.svg
+   :target: https://results.pre-commit.ci/latest/github/theislab/moscot/main
+   :alt: pre-commit.ci status
+
+.. |Docs| image:: https://img.shields.io/readthedocs/moscot
+    :target: https://moscot.readthedocs.io/en/stable/
+    :alt: Documentation
 
-Our manuscript will be available soon.
+.. |Downloads| image:: https://pepy.tech/badge/moscot
+    :target: https://pepy.tech/project/moscot
+    :alt: Downloads

moscot-0.3.1/README.rst

@@ -0,0 +1,66 @@
+|PyPI| |Downloads| |CI| |Pre-commit| |Codecov| |Docs|
+
+moscot - multi-omic single-cell optimal transport tools
+=======================================================
+
+**moscot** is a scalable framework for Optimal Transport (OT) applications in
+single-cell genomics. It can be used for
+
+- 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
+Transport toolkit that supports just-in-time compilation, GPU acceleration, automatic
+differentiation and linear memory complexity for OT problems.
+
+Installation
+------------
+You can install **moscot** via::
+
+    pip install moscot
+
+In order to install **moscot** from in editable mode, run::
+
+    git clone https://github.com/theislab/moscot
+    cd moscot
+    pip install -e .
+
+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>`_.
+
+.. |Codecov| image:: https://codecov.io/gh/theislab/moscot/branch/master/graph/badge.svg?token=Rgtm5Tsblo
+    :target: https://codecov.io/gh/theislab/moscot
+    :alt: Coverage
+
+.. |PyPI| image:: https://img.shields.io/pypi/v/moscot.svg
+    :target: https://pypi.org/project/moscot/
+    :alt: PyPI
+
+.. |CI| image:: https://img.shields.io/github/actions/workflow/status/theislab/moscot/test.yml?branch=main
+    :target: https://github.com/theislab/moscot/actions
+    :alt: CI
+
+.. |Pre-commit| image:: https://results.pre-commit.ci/badge/github/theislab/moscot/main.svg
+   :target: https://results.pre-commit.ci/latest/github/theislab/moscot/main
+   :alt: pre-commit.ci status
+
+.. |Docs| image:: https://img.shields.io/readthedocs/moscot
+    :target: https://moscot.readthedocs.io/en/stable/
+    :alt: Documentation
+
+.. |Downloads| image:: https://pepy.tech/badge/moscot
+    :target: https://pepy.tech/project/moscot
+    :alt: Downloads

{moscot-0.2.0 → moscot-0.3.1}/pyproject.toml

@@ -47,15 +47,14 @@ maintainers = [
 dependencies = [
     "numpy>=1.20.0",
     "scipy>=1.7.0",
-    "pandas>=1.4.0",
+    "pandas>=2.0.1",
     "networkx>=2.6.3",
     # https://github.com/scverse/scanpy/issues/2411
     "matplotlib>=3.5.0",
-    "anndata>=0.8.0",
+    "anndata>=0.9.1",
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
-    "docrep>=0.3.2",
-    "ott-jax>=0.4.0",
+    "ott-jax==0.4.0",
     "cloudpickle>=2.2.0",
 ]
 
@@ -79,9 +78,10 @@ docs = [
     "sphinx_copybutton>=0.5.0",
     "sphinxcontrib-bibtex>=2.3.0",
     "sphinxcontrib-spelling>=7.6.2",
+    "sphinx-autodoc-typehints",
     "furo>=2022.09.29",
+    "sphinx-tippy>=0.4.1",
     "myst-nb>=0.17.1",
-    "nbsphinx>=0.8.1",
     "ipython>=7.20.0",
     "sphinx_design>=0.3.0",
 ]
@@ -151,9 +151,6 @@ target-version = "py38"
 "src/moscot/utils/subset_policy.py" = ["D101", "D102"]
 [tool.ruff.pydocstyle]
 convention = "numpy"
-[tool.ruff.pyupgrade]
-# Preserve types, even if a file imports `from __future__ import annotations`.
-keep-runtime-typing = true
 [tool.ruff.flake8-tidy-imports]
 ban-relative-imports = "all"
 [tool.ruff.flake8-quotes]
@@ -168,9 +165,10 @@ include = '\.pyi?$'
 profile = "black"
 include_trailing_comma = true
 multi_line_output = 3
-sections = ["FUTURE", "STDLIB", "THIRDPARTY", "NUMERIC", "PLOTTING", "BIO", "FIRSTPARTY", "LOCALFOLDER"]
+sections = ["FUTURE", "STDLIB", "THIRDPARTY", "GENERIC", "NUMERIC", "PLOTTING", "BIO", "FIRSTPARTY", "LOCALFOLDER"]
 # also contains what we import in notebooks
-known_numeric = ["numpy", "scipy", "jax", "ott", "pandas", "sklearn", "networkx"]
+known_generic = ["wrapt", "joblib"]
+known_numeric = ["numpy", "scipy", "jax", "ott", "pandas", "sklearn", "networkx", "statsmodels"]
 known_bio = ["anndata", "scanpy", "squidpy"]
 known_plotting = ["IPython", "matplotlib", "mpl_toolkits", "seaborn"]
 
@@ -178,9 +176,8 @@ known_plotting = ["IPython", "matplotlib", "mpl_toolkits", "seaborn"]
 markers = ["fast: marks tests as fask"]
 xfail_strict = true
 filterwarnings = [
-    "ignore:X.dtype being converted:FutureWarning",
     "ignore:No data for colormapping:UserWarning",
-    "ignore:jax\\.experimental\\.pjit\\.PartitionSpec:DeprecationWarning",
+    "ignore:The dtype argument will be deprecated in anndata:PendingDeprecationWarning"
 ]
 
 [tool.coverage.run]
@@ -214,6 +211,7 @@ ignore_directives = [
     "automodule",
     "autoclass",
     "bibliography",
+    "glossary",
     "card",
     "grid",
 ]
@@ -251,7 +249,7 @@ show_column_numbers = true
 error_summary = true
 ignore_missing_imports = true
 
-disable_error_code = ["assignment", "comparison-overlap", "no-untyped-def"]
+disable_error_code = ["assignment", "comparison-overlap", "no-untyped-def", "override"]
 
 [tool.doc8]
 max_line_length = 120
@@ -294,6 +292,7 @@ commands =
 
 [testenv:clean-docs]
 description = Remove the documentation.
+deps =
 skip_install = true
 changedir = {tox_root}{/}docs
 allowlist_externals = make
@@ -302,7 +301,6 @@ commands =
 
 [testenv:build-docs]
 description = Build the documentation.
-use_develop = true
 deps =
 extras = docs
 allowlist_externals = make
@@ -324,4 +322,14 @@ commands =
     python -m twine check {tox_root}{/}dist{/}*
 commands_post =
     python -c 'import pathlib; print(f"Package is under:", pathlib.Path("{tox_root}") / "dist")'
+
+[testenv:format-references]
+description = Format references.bib.
+deps =
+skip_install = true
+allowlist_externals = biber
+commands = biber --tool --output_file={tox_root}{/}docs{/}references.bib --nolog \
+    --output_align --output_indent=2 --output_fieldcase=lower \
+    --output_legacy_dates --output-field-replace=journaltitle:journal,thesis:phdthesis,institution:school \
+    {tox_root}{/}docs{/}references.bib
 """

{moscot-0.2.0 → moscot-0.3.1}/src/moscot/_constants.py

@@ -4,7 +4,6 @@ STAR = "star"
 EXTERNAL_STAR = "external_star"
 TRIU = "triu"
 TRIL = "tril"
-PAIRWISE = "pairwise"
 EXPLICIT = "explicit"
 DUMMY = "dummy"
 # plotting

{moscot-0.2.0 → moscot-0.3.1}/src/moscot/_types.py

@@ -16,7 +16,7 @@ except (ImportError, TypeError):
 ProblemKind_t = Literal["linear", "quadratic", "unknown"]
 Numeric_t = Union[int, float]  # type of `time_key` arguments
 Filter_t = Optional[Union[str, Mapping[str, Sequence[Any]]]]  # type how to filter adata
-Str_Dict_t = Union[str, Mapping[str, Sequence[Any]]]  # type for `cell_transition`
+Str_Dict_t = Optional[Union[str, Mapping[str, Sequence[Any]]]]  # type for `cell_transition`
 SinkFullRankInit = Literal["default", "gaussian", "sorting"]
 LRInitializer_t = Literal["random", "rank2", "k-means", "generalized-k-means"]
 
@@ -24,20 +24,35 @@ SinkhornInitializer_t = Optional[Union[SinkFullRankInit, LRInitializer_t]]
 QuadInitializer_t = Optional[LRInitializer_t]
 
 Initializer_t = Union[SinkhornInitializer_t, LRInitializer_t]
-ProblemStage_t = Literal["initialized", "prepared", "solved"]
-Device_t = Literal["cpu", "gpu", "tpu"]
+ProblemStage_t = Literal["prepared", "solved"]
+Device_t = Union[Literal["cpu", "gpu", "tpu"], str]
 
 # TODO(michalk8): autogenerate from the enums
-ScaleCost_t = Optional[Union[float, Literal["mean", "max_cost", "max_bound", "max_norm", "median"]]]
-OttCostFn_t = Literal["euclidean", "sq_euclidean", "cosine", "bures", "unbalanced_bures"]
+ScaleCost_t = Union[float, Literal["mean", "max_cost", "max_bound", "max_norm", "median"]]
+OttCostFn_t = Literal[
+    "euclidean",
+    "sq_euclidean",
+    "cosine",
+    "PNormP",
+    "SqPNorm",
+    "Euclidean",
+    "SqEuclidean",
+    "Cosine",
+    "ElasticL1",
+    "ElasticSTVS",
+    "ElasticSqKOverlap",
+]
+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]
 Policy_t = Literal[
     "sequential",
     "star",
     "external_star",
+    "explicit",
     "triu",
     "tril",
-    "explicit",
 ]
+CostKwargs_t = Union[Mapping[str, Any], Mapping[Literal["x", "y", "xy"], Mapping[str, Any]]]

{moscot-0.2.0 → moscot-0.3.1}/src/moscot/backends/ott/__init__.py

@@ -1,10 +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.solver import GWSolver, SinkhornSolver
 from moscot.costs import register_cost
 
-__all__ = ["OTTOutput", "GWSolver", "SinkhornSolver"]
+__all__ = ["OTTOutput", "GWSolver", "SinkhornSolver", "sinkhorn_divergence"]
 
 register_cost("euclidean", backend="ott")(costs.Euclidean)
 register_cost("sq_euclidean", backend="ott")(costs.SqEuclidean)

moscot-0.3.1/src/moscot/backends/ott/_utils.py

@@ -0,0 +1,88 @@
+from typing import Any, Optional
+
+import jax
+import jax.numpy as jnp
+import scipy.sparse as sp
+from ott.geometry import geometry, pointcloud
+from ott.tools import sinkhorn_divergence as sdiv
+
+from moscot._logging import logger
+from moscot._types import ArrayLike, ScaleCost_t
+
+__all__ = ["sinkhorn_divergence"]
+
+
+def sinkhorn_divergence(
+    point_cloud_1: ArrayLike,
+    point_cloud_2: ArrayLike,
+    a: Optional[ArrayLike] = None,
+    b: Optional[ArrayLike] = None,
+    epsilon: Optional[float] = 1e-1,
+    scale_cost: ScaleCost_t = 1.0,
+    **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(
+        pointcloud.PointCloud,
+        x=point_cloud_1,
+        y=point_cloud_2,
+        a=a,
+        b=b,
+        epsilon=epsilon,
+        scale_cost=scale_cost,
+        **kwargs,
+    )
+    xy_conv, xx_conv, *yy_conv = output.converged
+
+    if not xy_conv:
+        logger.warning("Solver did not converge in the `x/y` term.")
+    if not xx_conv:
+        logger.warning("Solver did not converge in the `x/x` term.")
+    if len(yy_conv) and not yy_conv[0]:
+        logger.warning("Solver did not converge in the `y/y` term.")
+
+    return float(output.divergence)
+
+
+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]
+    if n != n_:
+        raise ValueError(f"Expected the first geometry to have `{n}` points, found `{n_}`.")
+    if m != m_:
+        raise ValueError(f"Expected the second geometry to have `{m}` points, found `{m_}`.")
+
+
+def alpha_to_fused_penalty(alpha: float) -> float:
+    """Convert."""
+    if not (0 < alpha <= 1):
+        raise ValueError(f"Expected `alpha` to be in interval `(0, 1]`, found `{alpha}`.")
+    return (1 - alpha) / alpha
+
+
+def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
+    """Ensure that an array is 2-dimensional.
+
+    Parameters
+    ----------
+    arr
+        Array to check.
+    reshape
+        Allow reshaping 1-dimensional array to ``[n, 1]``.
+
+    Returns
+    -------
+    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

{moscot-0.2.0 → moscot-0.3.1}/src/moscot/backends/ott/output.py

@@ -5,14 +5,12 @@ import jaxlib.xla_extension as xla_ext
 import jax
 import jax.numpy as jnp
 import numpy as np
-from ott.solvers.linear.sinkhorn import SinkhornOutput as OTTSinkhornOutput
-from ott.solvers.linear.sinkhorn_lr import LRSinkhornOutput as OTTLRSinkhornOutput
-from ott.solvers.quadratic.gromov_wasserstein import GWOutput as OTTGWOutput
+from ott.solvers.linear import sinkhorn, sinkhorn_lr
+from ott.solvers.quadratic import gromov_wasserstein
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 
-from moscot._docs._docs import d
 from moscot._types import ArrayLike, Device_t
 from moscot.base.output import BaseSolverOutput
 
@@ -20,58 +18,59 @@ __all__ = ["OTTOutput"]
 
 
 class OTTOutput(BaseSolverOutput):
-    """Output of various optimal transport problems.
+    """Output of various :term:`OT` problems.
 
     Parameters
     ----------
     output
-        Output of :mod:`ott` backend.
+        Output of the :mod:`ott` backend.
     """
 
-    _NOT_COMPUTED = -1.0
+    _NOT_COMPUTED = -1.0  # sentinel value used in `ott`
 
-    def __init__(self, output: Union[OTTSinkhornOutput, OTTLRSinkhornOutput, OTTGWOutput]):
+    def __init__(
+        self, output: Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput, gromov_wasserstein.GWOutput]
+    ):
         super().__init__()
         self._output = output
-        self._costs = None if isinstance(output, OTTSinkhornOutput) else output.costs
+        self._costs = None if isinstance(output, sinkhorn.SinkhornOutput) else output.costs
         self._errors = output.errors
 
-    @d.get_sections(base="plot_costs", sections=["Parameters", "Returns"])
     def plot_costs(
         self,
         last: Optional[int] = None,
         title: Optional[str] = None,
         return_fig: bool = False,
+        ax: Optional[mpl.axes.Axes] = None,
         figsize: Optional[Tuple[float, float]] = None,
         dpi: Optional[int] = None,
         save: Optional[str] = None,
-        ax: Optional[mpl.axes.Axes] = None,
         **kwargs: Any,
     ) -> Optional[mpl.figure.Figure]:
-        """Plot regularized OT costs during the iterations.
+        """Plot regularized :term:`OT` costs during the iterations.
 
         Parameters
         ----------
         last
-            How many of the last steps of the algorithm to plot. If `None`, plot the full curve.
+            How many of the last steps of the algorithm to plot. If :obj:`None`, plot the full curve.
         title
-            Title of the plot. If `None`, it is determined automatically.
+            Title of the plot. If :obj:`None`, it is determined automatically.
+        return_fig
+            Whether to return the figure.
+        ax
+            Axes on which to plot.
         figsize
             Size of the figure.
         dpi
             Dots per inch.
         save
             Path where to save the figure.
-        return_fig
-            Whether to return the figure.
-        ax
-            Axes on which to plot.
         kwargs
-            Keyword arguments for :meth:`~matplotlib.axes.Axes.plot`.
+            Keyword arguments for :meth:`matplotlib.axes.Axes.plot`.
 
         Returns
         -------
-        The figure if ``return_fig = True``.
+        If ``return_fig = True``, return the figure.
         """
         if self._costs is None:
             raise RuntimeError("No costs to plot.")
@@ -83,7 +82,6 @@ class OTTOutput(BaseSolverOutput):
             fig.savefig(save)
         return fig if return_fig else None
 
-    @d.dedent
     def plot_errors(
         self,
         last: Optional[int] = None,
@@ -96,15 +94,34 @@ class OTTOutput(BaseSolverOutput):
         ax: Optional[mpl.axes.Axes] = None,
         **kwargs: Any,
     ) -> Optional[mpl.figure.Figure]:
-        """Plot errors during the iterations.
+        """Plot errors along iterations.
 
         Parameters
         ----------
-        %(plot_costs.parameters)s
+        last
+            Number of errors corresponding at the ``last`` steps of the algorithm to plot. If :obj:`None`,
+            plot the full curve.
+        title
+            Title of the plot. If :obj:`None`, it is determined automatically.
+        outer_iteration
+            Which outermost iteration's errors to plot.
+            Only used when this is the solution to the :term:`quadratic problem`.
+        return_fig
+            Whether to return the figure.
+        ax
+            Axes on which to plot.
+        figsize
+            Size of the figure.
+        dpi
+            Dots per inch.
+        save
+            Path where to save the figure.
+        kwargs
+            Keyword arguments for :meth:`matplotlib.axes.Axes.plot`.
 
         Returns
         -------
-        %(plot_costs.returns)s
+        If ``return_fig = True``, return the figure.
         """
         if self._errors is None:
             raise RuntimeError("No errors to plot.")
@@ -155,7 +172,7 @@ class OTTOutput(BaseSolverOutput):
 
     @property
     def shape(self) -> Tuple[int, int]:  # noqa: D102
-        if isinstance(self._output, OTTSinkhornOutput):
+        if isinstance(self._output, sinkhorn.SinkhornOutput):
             return self._output.f.shape[0], self._output.g.shape[0]
         return self._output.geom.shape
 
@@ -165,9 +182,12 @@ class OTTOutput(BaseSolverOutput):
 
     @property
     def is_linear(self) -> bool:  # noqa: D102
-        return isinstance(self._output, (OTTSinkhornOutput, OTTLRSinkhornOutput))
+        return isinstance(self._output, (sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput))
 
     def to(self, device: Optional[Device_t] = None) -> "OTTOutput":  # noqa: D102
+        if device is None:
+            return OTTOutput(jax.device_put(self._output, device=device))
+
         if isinstance(device, str) and ":" in device:
             device, ix = device.split(":")
             idx = int(ix)
@@ -184,9 +204,7 @@ class OTTOutput(BaseSolverOutput):
 
     @property
     def cost(self) -> float:  # noqa: D102
-        if isinstance(self._output, (OTTSinkhornOutput, OTTLRSinkhornOutput)):
-            return float(self._output.reg_ot_cost)
-        return float(self._output.reg_gw_cost)
+        return float(self._output.reg_ot_cost if self.is_linear else self._output.reg_gw_cost)
 
     @property
     def converged(self) -> bool:  # noqa: D102
@@ -194,14 +212,14 @@ class OTTOutput(BaseSolverOutput):
 
     @property
     def potentials(self) -> Optional[Tuple[ArrayLike, ArrayLike]]:  # noqa: D102
-        if isinstance(self._output, OTTSinkhornOutput):
+        if isinstance(self._output, sinkhorn.SinkhornOutput):
             return self._output.f, self._output.g
         return None
 
     @property
     def rank(self) -> int:  # noqa: D102
         lin_output = self._output if self.is_linear else self._output.linear_state
-        return len(lin_output.g) if isinstance(lin_output, OTTLRSinkhornOutput) else -1
+        return len(lin_output.g) if isinstance(lin_output, sinkhorn_lr.LRSinkhornOutput) else -1
 
     def _ones(self, n: int) -> ArrayLike:  # noqa: D102
-        return jnp.ones((n,))  # type: ignore[return-value]
+        return jnp.ones((n,))