DensMAP support

docs/api/tools.md

@@ -20,6 +20,7 @@ Any transformation of the data matrix that is not *preprocessing*. In contrast t
    pp.pca
    tl.tsne
    tl.umap
+   tl.densmap
    tl.draw_graph
    tl.diffmap
 ```

docs/conf.py

@@ -225,6 +225,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.UmapMethodKwds"),
     # Will work once scipy 1.8 is released
     ("py:class", "scipy.sparse.base.spmatrix"),
     ("py:class", "scipy.sparse.csr.csr_matrix"),

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 with {func}`~scanpy.tl.densmap` {smaller}`M Keller`

src/scanpy/tools/__init__.py

@@ -24,11 +24,13 @@
 from ._score_genes import score_genes, score_genes_cell_cycle
 from ._sim import sim
 from ._tsne import tsne
-from ._umap import umap
+from ._umap import densmap, umap
 
 if TYPE_CHECKING:
     from typing import Any
 
+    from ._types import UmapMethodKwds  # noqa: F401
+
 
 def __getattr__(name: str) -> Any:
     if name == "pca":
@@ -56,4 +58,5 @@ def __getattr__(name: str) -> Any:
     "sim",
     "tsne",
     "umap",
+    "densmap",
 ]

src/scanpy/tools/_types.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from typing import TypedDict
+
+
+class UmapMethodKwds(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 UmapMethodKwds
 
     _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: UmapMethodKwds | 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.
         ``'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"
@@ -186,6 +222,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]
@@ -209,11 +246,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(),
@@ -228,8 +294,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,154 @@ 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
+
+
+# Convenience function for densMAP
+
+
+def densmap(
+    adata: AnnData,
+    *,
+    min_dist: float = 0.5,
+    spread: float = 1.0,
+    n_components: int = 2,
+    maxiter: int | None = None,
+    alpha: float = 1.0,
+    gamma: float = 1.0,
+    negative_sample_rate: int = 5,
+    init_pos: _InitPos | np.ndarray | None = "spectral",
+    random_state: _LegacyRandom = 0,
+    a: float | None = None,
+    b: float | None = None,
+    copy: bool = False,
+    neighbors_key: str | None = "neighbors",
+    dens_lambda: float | None = 2.0,
+    dens_frac: float | None = 0.3,
+    dens_var_shift: float | None = 0.1,
+) -> AnnData | None:
+    """\
+    Embed the neighborhood graph using densMAP :cite:p:`Narayan2021`.
+
+    We use the implementation of densMAP defined in
+    umap-learn_ :cite:p:`McInnes2018`.
+
+    .. _umap-learn: https://github.com/lmcinnes/umap
+
+    Parameters
+    ----------
+    adata
+        Annotated data matrix.
+    min_dist
+        The effective minimum distance between embedded points. Smaller values
+        will result in a more clustered/clumped embedding where nearby points on
+        the manifold are drawn closer together, while larger values will result
+        on a more even dispersal of points. The value should be set relative to
+        the ``spread`` value, which determines the scale at which embedded
+        points will be spread out. The default of in the `umap-learn` package is
+        0.1.
+    spread
+        The effective scale of embedded points. In combination with `min_dist`
+        this determines how clustered/clumped the embedded points are.
+    n_components
+        The number of dimensions of the embedding.
+    maxiter
+        The number of iterations (epochs) of the optimization. Called `n_epochs`
+        in the original UMAP.
+    alpha
+        The initial learning rate for the embedding optimization.
+    gamma
+        Weighting applied to negative samples in low dimensional embedding
+        optimization. Values higher than one will result in greater weight
+        being given to negative samples.
+    negative_sample_rate
+        The number of negative edge/1-simplex samples to use per positive
+        edge/1-simplex sample in optimizing the low dimensional embedding.
+    init_pos
+        How to initialize the low dimensional embedding. Called `init` in the
+        original UMAP. Options are:
+
+        * Any key for `adata.obsm`.
+        * 'paga': positions from :func:`~scanpy.pl.paga`.
+        * 'spectral': use a spectral embedding of the graph.
+        * 'random': assign initial embedding positions at random.
+        * A numpy array of initial embedding positions.
+    random_state
+        If `int`, `random_state` is the seed used by the random number generator;
+        If `RandomState` or `Generator`, `random_state` is the random number generator;
+        If `None`, the random number generator is the `RandomState` instance used
+        by `np.random`.
+    a
+        More specific parameters controlling the embedding. If `None` these
+        values are set automatically as determined by `min_dist` and
+        `spread`.
+    b
+        More specific parameters controlling the embedding. If `None` these
+        values are set automatically as determined by `min_dist` and
+        `spread`.
+    copy
+        Return a copy instead of writing to adata.
+    neighbors_key
+        If not specified, umap looks .uns['neighbors'] for neighbors settings
+        and .obsp['connectivities'] for connectivities
+        (default storage places for pp.neighbors).
+        If specified, umap looks .uns[neighbors_key] for neighbors settings and
+        .obsp[.uns[neighbors_key]['connectivities_key']] for connectivities.
+    dens_lambda
+        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
+        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
+        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:
+
+    `adata.obsm['X_densmap']` : :class:`numpy.ndarray` (dtype `float`)
+        DensMAP coordinates of data.
+    `adata.uns['densmap']` : :class:`dict`
+        DensMAP parameters.
+
+    """
+    return umap(
+        adata,
+        min_dist=min_dist,
+        spread=spread,
+        n_components=n_components,
+        maxiter=maxiter,
+        alpha=alpha,
+        gamma=gamma,
+        negative_sample_rate=negative_sample_rate,
+        init_pos=init_pos,
+        random_state=random_state,
+        a=a,
+        b=b,
+        copy=copy,
+        method="densmap",
+        method_kwds={
+            "dens_lambda": dens_lambda,
+            "dens_frac": dens_frac,
+            "dens_var_shift": dens_var_shift,
+        },
+        neighbors_key=neighbors_key,
+    )