#2946 and #3444

docs/conf.py

@@ -227,6 +227,7 @@ def setup(app: Sphinx):
     ("py:class", "scanpy._utils.Empty"),
     ("py:class", "numpy.random.mtrand.RandomState"),
     ("py:class", "scanpy.neighbors._types.KnnTransformerLike"),
+    ("py:class", "scanpy.tools._types.DensmapMethodKwds"),
 ]
 
 # Options for plot examples

docs/references.bib

@@ -533,6 +533,20 @@ @article{Muraro2016
   pages     = {385--394.e3},
 }
 
+@article{Narayan2021,
+  author    = {Narayan, Ashwin and Berger, Bonnie and Cho, Hyunghoon},
+  title     = {Assessing single-cell transcriptomic variability through density-preserving data visualization},
+  volume    = {39},
+  url       = {https://doi.org/10.1038/s41587-020-00801-7},
+  doi       = {10.1038/s41587-020-00801-7},
+  number    = {6},
+  journal   = {Nature Biotechnology},
+  publisher = {Springer Science and Business Media LLC},
+  year      = {2021},
+  month     = {jan},
+  pages     = {765--774},
+}
+
 @article{Nowotschin2019,
   author    = {Nowotschin, Sonja and Setty, Manu and Kuo, Ying-Yi and Liu, Vincent and Garg, Vidur and Sharma, Roshan and Simon, Claire S. and Saiz, Nestor and Gardner, Rui and Boutet, Stéphane C. and Church, Deanna M. and Hoodless, Pamela A. and Hadjantonakis, Anna-Katerina and Pe’er, Dana},
   title     = {The emergent landscape of the mouse gut endoderm at single-cell resolution},

docs/release-notes/2946.feature.md

@@ -0,0 +1 @@
+Add DensMAP support via `method="densmap"` in {func}`~scanpy.tl.umap` {smaller}`M Keller`

src/scanpy/neighbors/__init__.py

@@ -18,6 +18,7 @@
 from .._settings import settings
 from .._utils import NeighborsView, _doc_params, get_literal_vals
 from . import _connectivity
+from ._backends.pairwise import PairwiseDistancesTransformer
 from ._common import (
     _get_indices_distances_from_sparse_matrix,
     _get_sparse_matrix_from_indices_distances,
@@ -92,6 +93,12 @@ def neighbors(
     connectivities are computed according to :cite:t:`Coifman2005`, in the adaption of
     :cite:t:`Haghverdi2016`.
 
+    .. note::
+
+       Since scanpy 1.10, the results changed slightly.
+       We recommend to ensure reproducibility by pinning all package versions,
+       but you can get the old results by specifying `transformer='sklearn-pairwise'`.
+
     Parameters
     ----------
     adata
@@ -122,10 +129,13 @@ def neighbors(
         See :doc:`/how-to/knn-transformers` for more details.
         Also accepts the following known options:
 
-        `None` (the default)
+        `None` | `'sklearn'` (the default)
             Behavior depends on data size.
             For small data, we will calculate exact kNN, otherwise we use
             :class:`~pynndescent.pynndescent_.PyNNDescentTransformer`
+        `'sklearn-pairwise'`
+            For compatibility with scanpy <1.10, this allows to use
+            :class:`~sklearn.metrics.pairwise_distances`.
         `'pynndescent'`
             :class:`~pynndescent.pynndescent_.PyNNDescentTransformer`
         `'rapids'`
@@ -632,6 +642,7 @@ def _handle_transformer(
 
         If `transformer` is `None` and there are few data points,
         `transformer` will be set to a brute force
+        :class:`~from sklearn.metrics.pairwise_distances`.
         :class:`~sklearn.neighbors.KNeighborsTransformer`.
 
         If `transformer` is `None` and there are many data points,
@@ -642,7 +653,7 @@ def _handle_transformer(
         use_dense_distances = (
             kwds["metric"] == "euclidean" and self._adata.n_obs < 8192
         ) or not knn
-        shortcut = transformer == "sklearn" or (
+        shortcut = transformer in {"sklearn", "sklearn-pairwise"} or (
             transformer is None and (use_dense_distances or self._adata.n_obs < 4096)
         )
 
@@ -670,11 +681,19 @@ def _handle_transformer(
         if shortcut:
             from sklearn.neighbors import KNeighborsTransformer
 
-            assert transformer in {None, "sklearn"}
+            assert transformer in {None, "sklearn", "sklearn-pairwise"}
             n_neighbors = self._adata.n_obs - 1
             if knn:  # only obey n_neighbors arg if knn set
                 n_neighbors = min(n_neighbors, kwds["n_neighbors"])
-            transformer = KNeighborsTransformer(
+
+            # sklearn-pairwise is opt-in, because it takes more memory
+            transformer_cls = (
+                PairwiseDistancesTransformer
+                if transformer == "sklearn-pairwise"
+                else KNeighborsTransformer
+            )
+
+            transformer = transformer_cls(
                 algorithm="brute",
                 n_jobs=settings.n_jobs,
                 n_neighbors=n_neighbors,

src/scanpy/neighbors/_backends/pairwise.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from dataclasses import KW_ONLY, dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+from sklearn.base import TransformerMixin
+
+from .._common import (
+    _get_indices_distances_from_dense_matrix,
+    _get_indices_distances_from_sparse_matrix,
+    _get_sparse_matrix_from_indices_distances,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from typing import Literal, Self
+
+    from numpy.typing import NDArray
+
+    from ..._utils import _CSMatrix
+
+    _Metric = Literal["cityblock", "cosine", "euclidean", "l1", "l2", "manhattan"]
+    _MatrixLike = NDArray | _CSMatrix
+
+
+_DEBUG = False
+
+
+@dataclass
+class PairwiseDistancesTransformer(TransformerMixin):
+    _: KW_ONLY
+    algorithm: Literal["brute"]
+    n_jobs: int
+    n_neighbors: int
+    metric: _Metric
+    metric_params: Mapping[str, object]
+
+    def fit(self, x: _MatrixLike) -> Self:
+        self.x_ = x
+        return self
+
+    def transform(self, y: _MatrixLike | None) -> _CSMatrix:
+        from sklearn.metrics import pairwise_distances
+
+        d_arr = pairwise_distances(self.x_, y, metric=self.metric, **self.metric_params)
+        ind, dist = _get_indices_distances_from_dense_matrix(d_arr, self.n_neighbors)
+        rv = _get_sparse_matrix_from_indices_distances(ind, dist, keep_self=True)
+        if _DEBUG:
+            ind2, dist2 = _get_indices_distances_from_sparse_matrix(
+                rv, self.n_neighbors
+            )
+            np.testing.assert_equal(ind, ind2)
+            np.testing.assert_equal(dist, dist2)
+        return rv

src/scanpy/neighbors/_types.py

@@ -13,7 +13,7 @@
 
 # These two are used with get_literal_vals elsewhere
 _Method = Literal["umap", "gauss"]
-_KnownTransformer = Literal["pynndescent", "sklearn", "rapids"]
+_KnownTransformer = Literal["pynndescent", "sklearn", "sklearn-pairwise", "rapids"]
 
 # sphinx-autodoc-typehints can’t transitively import types from if TYPE_CHECKING blocks,
 # so these four needs to be importable

src/scanpy/tools/__init__.py

@@ -29,6 +29,8 @@
 if TYPE_CHECKING:
     from typing import Any
 
+    from ._types import DensmapMethodKwds  # noqa: F401
+
 
 def __getattr__(name: str) -> Any:
     if name == "pca":

src/scanpy/tools/_types.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from typing import TypedDict
+
+
+class DensmapMethodKwds(TypedDict, total=False):
+    dens_lambda: float
+    dens_frac: float
+    dens_var_shift: float

src/scanpy/tools/_umap.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 from sklearn.utils import check_array, check_random_state
@@ -18,6 +18,7 @@
     from anndata import AnnData
 
     from .._compat import _LegacyRandom
+    from ._types import DensmapMethodKwds
 
     _InitPos = Literal["paga", "spectral", "random"]
 
@@ -52,7 +53,8 @@ def umap(
     random_state: _LegacyRandom = 0,
     a: float | None = None,
     b: float | None = None,
-    method: Literal["umap", "rapids"] = "umap",
+    method: Literal["umap", "rapids", "densmap"] = "umap",
+    method_kwds: DensmapMethodKwds | None = None,
     key_added: str | None = None,
     neighbors_key: str = "neighbors",
     copy: bool = False,
@@ -128,6 +130,8 @@ def umap(
 
         ``'umap'``
             Umap’s simplical set embedding.
+        ``'densmap'``
+            Umap’s simplical set embedding with densmap=True :cite:p:`Narayan2021`.
         ``'rapids'``
             GPU accelerated implementation.
 
@@ -147,19 +151,51 @@ def umap(
     copy
         Return a copy instead of writing to adata.
 
+    method_kwds
+        Additional method parameters.
+
+        If method is ``'densmap'``, the following parameters are available:
+
+            ``dens_lambda`` : `float`, optional (default: 2.0)
+                Controls the regularization weight of the density correlation term
+                in densMAP. Higher values prioritize density preservation over the
+                UMAP objective, and vice versa for values closer to zero. Setting this
+                parameter to zero is equivalent to running the original UMAP algorithm.
+            ``dens_frac`` : `float`, optional (default: 0.3)
+                Controls the fraction of epochs (between 0 and 1) where the
+                density-augmented objective is used in densMAP. The first
+                (1 - dens_frac) fraction of epochs optimize the original UMAP objective
+                before introducing the density correlation term.
+            ``dens_var_shift`` : `float`, optional (default: 0.1)
+                A small constant added to the variance of local radii in the
+                embedding when calculating the density correlation objective to
+                prevent numerical instability from dividing by a small number
+
     Returns
     -------
-    Returns `None` if `copy=False`, else returns an `AnnData` object. Sets the following fields:
+    Returns `None` if `copy=False`, else returns an `AnnData` object. Sets the following fields unless method is 'densmap':
 
     `adata.obsm['X_umap' | key_added]` : :class:`numpy.ndarray` (dtype `float`)
         UMAP coordinates of data.
     `adata.uns['umap' | key_added]` : :class:`dict`
         UMAP parameters.
 
+    When method is 'densmap', sets the following fields:
+
+    `adata.obsm['X_densmap']` : :class:`numpy.ndarray` (dtype `float`)
+        densMAP coordinates of data.
+    `adata.uns['densmap']` : :class:`dict`
+        densMAP parameters.
+
     """
     adata = adata.copy() if copy else adata
 
-    key_obsm, key_uns = ("X_umap", "umap") if key_added is None else [key_added] * 2
+    key_obsm, key_uns = (
+        (("X_densmap", "densmap") if method == "densmap" else ("X_umap", "umap"))
+        if key_added is None
+        else [key_added] * 2
+    )
+    method_name = "DensMAP" if method == "densmap" else "UMAP"
 
     if neighbors_key is None:  # backwards compat
         neighbors_key = "neighbors"
@@ -185,6 +221,7 @@ def umap(
 
     if a is None or b is None:
         a, b = find_ab_params(spread, min_dist)
+
     adata.uns[key_uns] = dict(params=dict(a=a, b=b))
     if isinstance(init_pos, str) and init_pos in adata.obsm:
         init_coords = adata.obsm[init_pos]
@@ -208,11 +245,40 @@ def umap(
         n_pcs=neigh_params.get("n_pcs", None),
         silent=True,
     )
-    if method == "umap":
+
+    if method_kwds is None:
+        method_kwds = {}
+
+    densmap_kwds = (
+        {
+            "graph_dists": neighbors["distances"],
+            "n_neighbors": neigh_params.get("n_neighbors", 15),
+            # Default params from umap package
+            # Reference: https://github.com/lmcinnes/umap/blob/868e55cb614f361a0d31540c1f4a4b175136025c/umap/umap_.py#L1692
+            # If user provided method_kwds, the user-provided values should
+            # overwrite the default values specified above.
+            "lambda": method_kwds.get("dens_lambda", 2.0),
+            "frac": method_kwds.get("dens_frac", 0.3),
+            "var_shift": method_kwds.get("dens_var_shift", 0.1),
+        }
+        if method == "densmap"
+        else {}
+    )
+    if method == "densmap":
+        adata.uns[key_uns]["params"].update(
+            {
+                "dens_lambda": densmap_kwds["lambda"],
+                "dens_frac": densmap_kwds["frac"],
+                "dens_var_shift": densmap_kwds["var_shift"],
+            }
+        )
+
+    if method == "umap" or method == "densmap":
         # the data matrix X is really only used for determining the number of connected components
         # for the init condition in the UMAP embedding
         default_epochs = 500 if neighbors["connectivities"].shape[0] <= 10000 else 200
         n_epochs = default_epochs if maxiter is None else maxiter
+
         X_umap, _ = simplicial_set_embedding(
             data=X,
             graph=neighbors["connectivities"].tocoo(),
@@ -227,8 +293,8 @@ def umap(
             random_state=random_state,
             metric=neigh_params.get("metric", "euclidean"),
             metric_kwds=neigh_params.get("metric_kwds", {}),
-            densmap=False,
-            densmap_kwds={},
+            densmap=(method == "densmap"),
+            densmap_kwds=densmap_kwds,
             output_dens=False,
             verbose=settings.verbosity > 3,
         )
@@ -267,14 +333,15 @@ def umap(
             random_state=random_state,
         )
         X_umap = umap.fit_transform(X_contiguous)
+
     adata.obsm[key_obsm] = X_umap  # annotate samples with UMAP coordinates
     logg.info(
         "    finished",
         time=start,
         deep=(
             "added\n"
-            f"    {key_obsm!r}, UMAP coordinates (adata.obsm)\n"
-            f"    {key_uns!r}, UMAP parameters (adata.uns)"
+            f"    {key_obsm!r}, {method_name} coordinates (adata.obsm)\n"
+            f"    {key_uns!r}, {method_name} parameters (adata.uns)"
         ),
     )
     return adata if copy else None

tests/test_embedding.py

@@ -88,3 +88,24 @@ def test_diffmap():
     sc.tl.diffmap(pbmc, random_state=1234)
     d3 = pbmc.obsm["X_diffmap"].copy()
     assert_raises(AssertionError, assert_array_equal, d1, d3)
+
+
+def test_densmap():
+    pbmc = pbmc68k_reduced()
+
+    # Checking that the results are reproducible
+    sc.tl.umap(pbmc, method="densmap")
+    d1 = pbmc.obsm["X_densmap"].copy()
+    sc.tl.umap(pbmc, method="densmap")
+    d2 = pbmc.obsm["X_densmap"].copy()
+    assert_array_equal(d1, d2)
+
+    # Checking if specifying random_state  works, arrays shouldn't be equal
+    sc.tl.umap(pbmc, method="densmap", random_state=1234)
+    d3 = pbmc.obsm["X_densmap"].copy()
+    assert_raises(AssertionError, assert_array_equal, d1, d3)
+
+    # Checking if specifying dens_lambda  works, arrays shouldn't be equal
+    sc.tl.umap(pbmc, method="densmap", method_kwds=dict(dens_lambda=2.3456))
+    d4 = pbmc.obsm["X_densmap"].copy()
+    assert_raises(AssertionError, assert_array_equal, d1, d4)