Repository: theislab/cellrank (language: python)

Issue:
Encountered NaN values in pseudotime
Hello,

I am very interested in using your package in some single-cell datasets. So far your documentation looks great and I look forward to fully test this tool in my data. My datasets do not contain the mappings for unspliced and spliced (scVelo), that is why I am trying to follow the steps described in the 'CellRank beyond RNA velocity' section.

Here is my code:

```
import scvelo as scv
import scanpy as sc
import cellrank as cr
import numpy as np

#adata was already processed for normalization and highly variable genes 

adata.layers["spliced"] = adata.X
adata.layers["unspliced"] = adata.X
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)
```

The problem comes when I try to run these lines:
```
from cellrank.tl.kernels import CytoTRACEKernel
ctk = CytoTRACEKernel(adata)
```

I get the error:
ValueError: Encountered NaN values in pseudotime.

I have tried to change the Nan to zero and did not work. I was wondering if I am missing any important details, or if you have any comments.


Thank you so much for your work and help!

Best,
Gaby B.


Discussion / hints:

Fix patch:
diff --git a/cellrank/tl/_lineage.py b/cellrank/tl/_lineage.py
index d027c35de..655ab6d7d 100644
--- a/cellrank/tl/_lineage.py
+++ b/cellrank/tl/_lineage.py
@@ -198,7 +198,7 @@ def _register_handled_functions():
 
 class LineageMeta(type):
     """
-    Metaclass for Lineaage.
+    Metaclass for Lineage.
 
     Registers functions which are handled by us and overloads common attributes, such as `.sum` with these functions.
     """
diff --git a/cellrank/tl/kernels/_cytotrace_kernel.py b/cellrank/tl/kernels/_cytotrace_kernel.py
index fffb6629b..20c444db5 100644
--- a/cellrank/tl/kernels/_cytotrace_kernel.py
+++ b/cellrank/tl/kernels/_cytotrace_kernel.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Optional
 from typing_extensions import Literal
 
 from enum import auto
@@ -8,7 +8,7 @@
 from cellrank._key import Key
 from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import _correlation_test_helper
+from cellrank.tl._utils import _correlation_test, _correlation_test_helper
 from cellrank.tl.kernels._pseudotime_kernel import PseudotimeKernel
 
 import numpy as np
@@ -31,7 +31,8 @@ class CytoTRACEKernel(PseudotimeKernel):
 
     The KNN graph contains information about the (undirected) connectivities among cells, reflecting their similarity.
     CytoTRACE can be used to estimate cellular plasticity and in turn, a pseudotemporal ordering of cells from more
-    plastic to less plastic states.
+    plastic to less plastic states. It relies on the assumption that differentiated cells express, on average,
+    less genes than naive cells.
     This kernel internally uses the :class:`cellrank.tl.kernels.PseudotimeKernel` to direct the KNN graph
     on the basis of the CytoTRACE-derived pseudotime.
 
@@ -42,7 +43,6 @@ class CytoTRACEKernel(PseudotimeKernel):
     %(adata)s
     %(backward)s
     %(cytotrace.parameters)s
-
     %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
@@ -83,11 +83,12 @@ def __init__(
         self,
         adata: AnnData,
         backward: bool = False,
-        layer: str = "Ms",
+        layer: Optional[str] = "Ms",
         aggregation: Literal[
             "mean", "median", "hmean", "gmean"
         ] = CytoTRACEAggregation.MEAN,
         use_raw: bool = False,
+        n_top_genes: int = 200,
         compute_cond_num: bool = False,
         check_connectivity: bool = False,
         **kwargs: Any,
@@ -100,15 +101,17 @@ def __init__(
             check_connectivity=check_connectivity,
             layer=layer,
             aggregation=aggregation,
+            n_top_genes=n_top_genes,
             use_raw=use_raw,
             **kwargs,
         )
-        self._time_key = Key.cytotrace("pseudotime")  # quirk or PT kernel
+        # TODO(michlak8): needed by the PT kernel?
+        self._time_key = Key.cytotrace("pseudotime")
 
     def _read_from_adata(
         self,
         time_key: str,
-        layer: str = "Ms",
+        layer: Optional[str] = "Ms",
         aggregation: Literal[
             "mean", "median", "hmean", "gmean"
         ] = CytoTRACEAggregation.MEAN,
@@ -123,11 +126,12 @@ def _read_from_adata(
     @inject_docs(ct=CytoTRACEAggregation)
     def compute_cytotrace(
         self,
-        layer: str = "Ms",
+        layer: Optional[str] = "Ms",
         aggregation: Literal[
             "mean", "median", "hmean", "gmean"
         ] = CytoTRACEAggregation.MEAN,
         use_raw: bool = False,
+        n_top_genes: int = 200,
     ) -> None:
         """
         Re-implementation of the CytoTRACE algorithm :cite:`gulati:20` to estimate cellular plasticity.
@@ -149,10 +153,11 @@ def compute_cytotrace(
                 - `{ct.MEDIAN!r}` - median.
                 - `{ct.HMEAN!r}` - harmonic mean.
                 - `{ct.GMEAN!r}` - geometric mean.
-
         use_raw
             Whether to use the :attr:`anndata.AnnData.raw` to compute the number of genes expressed per cell
             (#genes/cell) and the correlation of gene expression across cells with #genes/cell.
+        n_top_genes
+            Number of genes used to compute the CytoTRACE score.
 
         Returns
         -------
@@ -173,57 +178,61 @@ def compute_cytotrace(
         This will not exactly reproduce the results of the original CytoTRACE algorithm :cite:`gulati:20` because we
         allow for any normalization and imputation techniques whereas CytoTRACE has built-in specific methods for that.
         """
-        # check use_raw
+        from cellrank.tl import Lineage
+
         aggregation = CytoTRACEAggregation(aggregation)
+
+        if n_top_genes <= 0:
+            raise ValueError(
+                f"Expected `n_top_genes` to be `> 0`, found `{n_top_genes}`."
+            )
         if use_raw and self.adata.raw is None:
             logg.warning("`adata.raw` is `None`. Setting `use_raw=False`")
             use_raw = False
-        if use_raw and self.adata.raw.n_vars != self.adata.n_vars:
-            logg.warning(
-                f"`adata.raw` has different number of genes ({self.adata.raw.n_vars}) "
-                f"than `adata` ({self.adata.n_vars}). Setting `use_raw=False`"
-            )
-            use_raw = False
-
-        adata_mraw = self.adata.raw if use_raw else self.adata
-        if layer != "X" and layer not in self.adata.layers:
+        if layer not in (None, "X") and layer not in self.adata.layers:
             raise KeyError(
                 f"Unable to find `{layer!r}` in `adata.layers`. "
                 f"Valid option are: `{sorted({'X'} | set(self.adata.layers.keys()))}`."
             )
 
-        msg = f"Computing CytoTRACE score with `{self.adata.n_vars}` genes"
-        if self.adata.n_vars < 10000:
+        adata_mraw = self.adata.raw if use_raw else self.adata
+        msg = f"Computing CytoTRACE score with `{adata_mraw.n_vars}` genes"
+        if adata_mraw.n_vars < 10000:
             msg += ". Consider using more than `10000` genes"
         start = logg.info(msg)
 
         # compute number of expressed genes per cell
-        logg.debug(
-            f"Computing number of genes expressed per cell with `use_raw={use_raw}`"
-        )
-        num_exp_genes = np.array((adata_mraw.X > 0).sum(axis=1)).reshape(-1)
-        self.adata.obs[Key.cytotrace("num_exp_genes")] = num_exp_genes
+        num_exp_genes = np.asarray((adata_mraw.X > 0).sum(axis=1)).squeeze()
 
-        # fmt: off
-        # compute correlation with all genes
         logg.debug("Correlating all genes with number of genes expressed per cell")
-        gene_corr, _, _, _ = _correlation_test_helper(adata_mraw.X.T, num_exp_genes[:, None])
+        gene_corr = _correlation_test(
+            adata_mraw.X,
+            Lineage(num_exp_genes[:, None], names=["gene"]),
+            gene_names=adata_mraw.var_names,
+        )["gene_corr"]
 
-        # annotate the top 200 genes in terms of correlation
-        logg.debug("Finding the top `200` most correlated genes")
-        self.adata.var[Key.cytotrace("gene_corr")] = gene_corr
-        top_200 = self.adata.var.sort_values(by=Key.cytotrace("gene_corr"), ascending=False).index[:200]
-        self.adata.var[Key.cytotrace("correlates")] = False
-        self.adata.var.loc[top_200, Key.cytotrace("correlates")] = True
+        # fmt: off
+        top_genes = [g for g in gene_corr.sort_values(ascending=False).index if g in self.adata.var_names]
+        top_genes = top_genes[:n_top_genes]
+        null_genes = int(gene_corr.loc[top_genes].isnull().sum())
+        if null_genes:
+            raise ValueError(f"Top `{len(top_genes)}` correlated genes contain `{null_genes}` NaN values.")
+        if len(top_genes) != n_top_genes:
+            logg.warning(f"Unable to get requested top correlated `{n_top_genes}`. "
+                         f"Using top `{len(top_genes)}` genes")
 
-        # compute mean/median over top 200 genes, aggregate over genes and shift to [0, 1] range
-        logg.debug(f"Aggregating imputed gene expression using aggregation `{aggregation}` in layer `{layer}`")
-        corr_mask = self.adata.var[Key.cytotrace("correlates")]
-        imputed_exp = self.adata[:, corr_mask].X if layer == "X" else self.adata[:, corr_mask].layers[layer]
+        logg.debug(
+            f"Aggregating imputed gene expression using "
+            f"`{aggregation}` top `{len(top_genes)}` in layer `{layer}`"
+        )
+        # compute mean/median over top n genes, aggregate over genes
+        if layer == "X":
+            imputed_exp = self.adata[:, top_genes].X
+        else:
+            imputed_exp = self.adata[:, top_genes].layers[layer]
         if issparse(imputed_exp):
             imputed_exp = imputed_exp.A
 
-        # aggregate across the top 200 genes
         if aggregation == CytoTRACEAggregation.MEAN:
             cytotrace_score = np.mean(imputed_exp, axis=1)
         elif aggregation == CytoTRACEAggregation.MEDIAN:
@@ -236,16 +245,19 @@ def compute_cytotrace(
             raise NotImplementedError(f"Aggregation method `{aggregation}` is not yet implemented.")
         # fmt: on
 
-        # scale to 0-1 range
+        self.adata.obs[Key.cytotrace("num_exp_genes")] = num_exp_genes
+        self.adata.var[Key.cytotrace("gene_corr")] = gene_corr
+        self.adata.var[Key.cytotrace("correlates")] = False
+        self.adata.var.loc[top_genes, Key.cytotrace("correlates")] = True
         cytotrace_score -= np.min(cytotrace_score)
         cytotrace_score /= np.max(cytotrace_score)
         self.adata.obs[Key.cytotrace("score")] = cytotrace_score
         self.adata.obs[Key.cytotrace("pseudotime")] = 1 - cytotrace_score
-
         self.adata.uns[Key.cytotrace("params")] = {
             "aggregation": aggregation,
             "layer": layer,
             "use_raw": use_raw,
+            "n_top_genes": len(top_genes),
         }
 
         logg.info(
diff --git a/cellrank/tl/kernels/_pseudotime_kernel.py b/cellrank/tl/kernels/_pseudotime_kernel.py
index 06b7765d2..7b6e1e70e 100644
--- a/cellrank/tl/kernels/_pseudotime_kernel.py
+++ b/cellrank/tl/kernels/_pseudotime_kernel.py
@@ -68,7 +68,7 @@ def __init__(
     def _read_from_adata(self, time_key: str, **kwargs: Any) -> None:
         super()._read_from_adata(**kwargs)
 
-        if time_key not in self.adata.obs.keys():
+        if time_key not in self.adata.obs:
             raise KeyError(f"Could not find time key in `adata.obs[{time_key!r}]`.")
 
         self._pseudotime = np.array(self.adata.obs[time_key]).astype(_dtype)


Test patch:
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index 148537bc0..bfbdf80a9 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -1484,9 +1484,7 @@ def test_aggregation(self, adata: AnnData, agg: CytoTRACEAggregation):
     @pytest.mark.parametrize("use_raw", [False, True])
     def test_raw(self, adata: AnnData, use_raw: bool):
         _ = CytoTRACEKernel(adata, use_raw=use_raw)
-        assert adata.uns[Key.cytotrace("params")]["use_raw"] == (
-            adata.raw.n_vars == adata.n_vars if use_raw else False
-        )
+        assert adata.uns[Key.cytotrace("params")]["use_raw"] == use_raw
 
     def test_correct_class(self, adata: AnnData):
         k = CytoTRACEKernel(adata)
@@ -1499,6 +1497,7 @@ def test_writes_params(self, adata: AnnData):
         assert adata.uns[Key.cytotrace("params")] == {
             "layer": "X",
             "aggregation": "mean",
+            "n_top_genes": 143,
             "use_raw": False,
         }
 
@@ -1517,6 +1516,28 @@ def test_writes_params(self, adata: AnnData):
         np.testing.assert_array_equal(k.pseudotime.min(), 0.0)
         np.testing.assert_array_equal(k.pseudotime.max(), 1.0)
 
+    def test_raw_less_genes(self, adata: AnnData):
+        adata.raw = adata.raw.to_adata()[:, :20]
+        _ = CytoTRACEKernel(adata, use_raw=True, n_top_genes=31)
+        assert adata.uns[Key.cytotrace("params")] == {
+            "layer": "Ms",
+            "aggregation": "mean",
+            "n_top_genes": 20,
+            "use_raw": True,
+        }
+
+    @pytest.mark.parametrize("n_top_genes", [0, 10, 300])
+    @pytest.mark.parametrize("use_raw", [False, True])
+    def test_n_top_genes(self, adata: AnnData, use_raw: bool, n_top_genes: int):
+        n_genes = 50 if use_raw else 143
+        if n_top_genes < 0:
+            with pytest.raises(ValueError, match=r"Expected `n_top_genes`"):
+                _ = CytoTRACEKernel(adata, use_raw=use_raw, n_top_genes=n_top_genes)
+        else:
+            _ = CytoTRACEKernel(adata, use_raw=use_raw, n_top_genes=n_top_genes)
+            assert adata.var[Key.cytotrace("correlates")].sum() == n_genes
+            assert adata.uns[Key.cytotrace("params")]["n_top_genes"] == n_genes
+
     def test_compute_transition_matrix(self, adata: AnnData):
         k = CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
         k.compute_transition_matrix()