Precompute nn ranking for DR datasets (#841)

* precompute nn ranking on full data

* bugfix

* bugfix2

* fix sample_dataset

* just subsample it

* Remove nan check

Author: scottgigante-immunai

openproblems/tasks/dimensionality_reduction/README.md

@@ -48,7 +48,9 @@ WARNING: other than most tasks, `adata.X` should contain log CP10k-normalized da
    highly on these metrics.
 
 **Datasets** should provide *log CP10k normalized counts* in `adata.X` and store the
-original number of genes (i.e., `adata.shape[1]`) in `adata.uns["n_genes"]`.
+original number of genes (i.e., `adata.shape[1]`) in `adata.uns["n_genes"]`. Datasets
+should also contain the nearest-neighbor ranking matrix, required for the `nn_ranking`
+metrics, as computed by `_utils.ranking_matrix(adata.X)` on normalized counts.
 
 **Methods** should assign dimensionally-reduced 2D embedding coordinates to
 `adata.obsm['X_emb']`. They *should not* modify the dimensionality of `adata.X` (e.g.

openproblems/tasks/dimensionality_reduction/_utils.py

@@ -0,0 +1,27 @@
+from numba import njit
+
+import numpy as np
+
+
+@njit(cache=True, fastmath=True)
+def _ranking_matrix(D: np.ndarray) -> np.ndarray:  # pragma: no cover
+    assert D.shape[0] == D.shape[1]
+    R = np.zeros(D.shape)
+    m = len(R)
+    ks = np.arange(m)
+
+    for i in range(m):
+        for j in range(m):
+            R[i, j] = np.sum(
+                (D[i, :] < D[i, j]) | ((ks < j) & (np.abs(D[i, :] - D[i, j]) <= 1e-12))
+            )
+
+    return R
+
+
+def ranking_matrix(X):
+    from sklearn.metrics import pairwise_distances
+
+    D = pairwise_distances(X)
+    R = _ranking_matrix(D)
+    return R

openproblems/tasks/dimensionality_reduction/api.py

@@ -1,6 +1,7 @@
 from ...data.sample import load_sample_data
 from ...tools.decorators import dataset
 from ...tools.normalize import log_cp10k
+from . import _utils
 
 import numpy as np
 
@@ -31,6 +32,7 @@ def sample_dataset():
     adata = load_sample_data()
     adata = log_cp10k(adata)
     adata.uns["n_genes"] = adata.shape[1]
+    adata.obsm["X_ranking"] = _utils.ranking_matrix(adata.X)
     return adata
 
 

openproblems/tasks/dimensionality_reduction/datasets/mouse_blood_olsson_labelled.py

@@ -1,6 +1,7 @@
 from ....data.mouse_blood_olsson_labelled import load_olsson_2016_mouse_blood
 from ....tools.decorators import dataset
 from ....tools.normalize import log_cp10k
+from .._utils import ranking_matrix
 
 
 @dataset(
@@ -15,4 +16,6 @@
 def olsson_2016_mouse_blood(test=False):
     adata = load_olsson_2016_mouse_blood(test=test)
     adata.uns["n_genes"] = adata.shape[1]
-    return log_cp10k(adata)
+    adata = log_cp10k(adata)
+    adata.obsm["X_ranking"] = ranking_matrix(adata.X)
+    return adata

openproblems/tasks/dimensionality_reduction/datasets/mouse_hspc_nestorowa2016.py

@@ -1,6 +1,7 @@
 from ....data.mouse_hspc_nestorowa2016 import load_mouse_hspc_nestorowa2016
 from ....tools.decorators import dataset
 from ....tools.normalize import log_cp10k
+from .._utils import ranking_matrix
 
 
 @dataset(
@@ -15,4 +16,6 @@
 def mouse_hspc_nestorowa2016(test=False):
     adata = load_mouse_hspc_nestorowa2016(test=test)
     adata.uns["n_genes"] = adata.shape[1]
-    return log_cp10k(adata)
+    adata = log_cp10k(adata)
+    adata.obsm["X_ranking"] = ranking_matrix(adata.X)
+    return adata

openproblems/tasks/dimensionality_reduction/datasets/tenx_5k_pbmc.py

@@ -1,6 +1,7 @@
 from ....data.tenx import load_tenx_5k_pbmc
 from ....tools.decorators import dataset
 from ....tools.normalize import log_cp10k
+from .._utils import ranking_matrix
 
 
 @dataset(
@@ -16,4 +17,6 @@
 def tenx_5k_pbmc(test=False):
     adata = load_tenx_5k_pbmc(test=test)
     adata.uns["n_genes"] = adata.shape[1]
-    return log_cp10k(adata)
+    adata = log_cp10k(adata)
+    adata.obsm["X_ranking"] = ranking_matrix(adata.X)
+    return adata

openproblems/tasks/dimensionality_reduction/datasets/zebrafish.py

@@ -1,6 +1,7 @@
 from ....data.zebrafish import load_zebrafish
 from ....tools.decorators import dataset
 from ....tools.normalize import log_cp10k
+from .._utils import ranking_matrix
 
 
 @dataset(
@@ -15,6 +16,13 @@
     ),
 )
 def zebrafish_labs(test=False):
+    import scanpy as sc
+
     adata = load_zebrafish(test=test)
+    if not test:
+        # this dataset is too big
+        sc.pp.subsample(adata, n_obs=25000)
     adata.uns["n_genes"] = adata.shape[1]
-    return log_cp10k(adata)
+    adata = log_cp10k(adata)
+    adata.obsm["X_ranking"] = ranking_matrix(adata.X)
+    return adata

openproblems/tasks/dimensionality_reduction/metrics/nn_ranking.py

@@ -15,6 +15,7 @@
 """
 
 from ....tools.decorators import metric
+from .._utils import ranking_matrix
 from anndata import AnnData
 from numba import njit
 from typing import Tuple
@@ -33,22 +34,6 @@
 _K = 30
 
 
-@njit(cache=True, fastmath=True)
-def _ranking_matrix(D: np.ndarray) -> np.ndarray:  # pragma: no cover
-    assert D.shape[0] == D.shape[1]
-    R = np.zeros(D.shape)
-    m = len(R)
-    ks = np.arange(m)
-
-    for i in range(m):
-        for j in range(m):
-            R[i, j] = np.sum(
-                (D[i, :] < D[i, j]) | ((ks < j) & (np.abs(D[i, :] - D[i, j]) <= 1e-12))
-            )
-
-    return R
-
-
 @njit(cache=True, fastmath=True)
 def _coranking_matrix(R1: np.ndarray, R2: np.ndarray) -> np.ndarray:  # pragma: no cover
     assert R1.shape == R2.shape
@@ -63,22 +48,6 @@ def _coranking_matrix(R1: np.ndarray, R2: np.ndarray) -> np.ndarray:  # pragma:
     return Q
 
 
-@njit(cache=True, fastmath=True)
-def _trustworthiness(Q: np.ndarray, m: int) -> np.ndarray:  # pragma: no cover
-
-    T = np.zeros(m - 1)  # trustworthiness
-
-    for k in range(m - 1):
-        Qs = Q[k:, :k]
-        # a column vector of weights. weight = rank error = actual_rank - k
-        W = np.arange(Qs.shape[0]).reshape(-1, 1)
-        # 1 - normalized hard-k-intrusions. lower-left region.
-        # weighted by rank error (rank - k)
-        T[k] = 1 - np.sum(Qs * W) / ((k + 1) * m * (m - 1 - k))
-
-    return T
-
-
 @njit(cache=True, fastmath=True)
 def _continuity(Q: np.ndarray, m: int) -> np.ndarray:  # pragma: no cover
 
@@ -133,65 +102,38 @@ def _qnn_auc(QNN: np.ndarray) -> float:
     return AUC  # type: ignore
 
 
-def _metrics(
-    Q: np.ndarray,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float, np.ndarray, int, float, float]:
-    Q = Q[1:, 1:]
-    m = len(Q)
-
-    T = _trustworthiness(Q, m)
-    C = _continuity(Q, m)
-    QNN = _qnn(Q, m)
-    LCMC = _lcmc(QNN, m)
-    kmax = _kmax(LCMC)
-    Qlocal = _q_local(QNN, kmax)
-    Qglobal = _q_global(QNN, kmax, m)
-    AUC = _qnn_auc(QNN)
-
-    return T, C, QNN, AUC, LCMC, kmax, Qlocal, Qglobal
-
-
-def _high_dim(adata: AnnData) -> np.ndarray:
-    from scipy.sparse import issparse
-
-    high_dim = adata.X
-    return high_dim.A if issparse(high_dim) else high_dim
-
+def _fit(adata: AnnData) -> Tuple[float, float, float, float, float, float, float]:
+    Rx = adata.obsm["X_ranking"]
+    E = adata.obsm["X_emb"]
 
-def _fit(
-    X: np.ndarray, E: np.ndarray
-) -> Tuple[float, float, float, float, float, float, float]:
-    from sklearn.metrics import pairwise_distances
-
-    if np.any(np.isnan(E)):
-        return 0.0, 0.0, 0.0, 0.5, -np.inf, -np.inf, -np.inf
-
-    Dx = pairwise_distances(X)
-    De = pairwise_distances(E)
-    Rx, Re = _ranking_matrix(Dx), _ranking_matrix(De)
+    Re = ranking_matrix(E)
     Q = _coranking_matrix(Rx, Re)
+    Q = Q[1:, 1:]
+    m = len(Q)
 
-    T, C, QNN, AUC, LCMC, _kmax, Qlocal, Qglobal = _metrics(Q)
-
-    return T[_K], C[_K], QNN[_K], AUC, LCMC[_K], Qlocal, Qglobal
+    return Q, m
 
 
 @metric("continuity", paper_reference="zhang2021pydrmetrics", maximize=True)
 def continuity(adata: AnnData) -> float:
-    _, C, _, *_ = _fit(_high_dim(adata), adata.obsm["X_emb"])
+    Q, m = _fit(adata)
+    C = _continuity(Q, m)[_K]
     return float(np.clip(C, 0.0, 1.0))  # in [0, 1]
 
 
 @metric("co-KNN size", paper_reference="zhang2021pydrmetrics", maximize=True)
 def qnn(adata: AnnData) -> float:
-    _, _, QNN, *_ = _fit(_high_dim(adata), adata.obsm["X_emb"])
+    Q, m = _fit(adata)
+    QNN = _qnn(Q, m)[_K]
     # normalized in the code to [0, 1]
     return float(np.clip(QNN, 0.0, 1.0))
 
 
 @metric("co-KNN AUC", paper_reference="zhang2021pydrmetrics", maximize=True)
 def qnn_auc(adata: AnnData) -> float:
-    _, _, _, AUC, *_ = _fit(_high_dim(adata), adata.obsm["X_emb"])
+    Q, m = _fit(adata)
+    QNN = _qnn(Q, m)
+    AUC = _qnn_auc(QNN)
     return float(np.clip(AUC, 0.5, 1.0))  # in [0.5, 1]
 
 
@@ -201,19 +143,29 @@ def qnn_auc(adata: AnnData) -> float:
     maximize=True,
 )
 def lcmc(adata: AnnData) -> float:
-    *_, LCMC, _, _ = _fit(_high_dim(adata), adata.obsm["X_emb"])
+    Q, m = _fit(adata)
+    QNN = _qnn(Q, m)
+    LCMC = _lcmc(QNN, m)[_K]
     return LCMC
 
 
 @metric("local property", paper_reference="zhang2021pydrmetrics", maximize=True)
 def qlocal(adata: AnnData) -> float:
     # according to authors, this is usually preferred to
     # qglobal, because human are more sensitive to nearer neighbors
-    *_, Qlocal, _ = _fit(_high_dim(adata), adata.obsm["X_emb"])
+    Q, m = _fit(adata)
+    QNN = _qnn(Q, m)
+    LCMC = _lcmc(QNN, m)
+    kmax = _kmax(LCMC)
+    Qlocal = _q_local(QNN, kmax)
     return Qlocal
 
 
 @metric("global property", paper_reference="zhang2021pydrmetrics", maximize=True)
 def qglobal(adata: AnnData) -> float:
-    *_, Qglobal = _fit(_high_dim(adata), adata.obsm["X_emb"])
+    Q, m = _fit(adata)
+    QNN = _qnn(Q, m)
+    LCMC = _lcmc(QNN, m)
+    kmax = _kmax(LCMC)
+    Qglobal = _q_global(QNN, kmax, m)
     return Qglobal