Improve KNN label_transfer in PerturbationSpace (#658)

* Add uncertainty score in KNN label_transfer in PerturbationSpace
Certainty is quantified as the fraction of nearest neighbors belonging to the classified (i.e. the most abundant) label compared to the total number of nearest neighbors.

* Update pre-commit-config.yaml
Replaces yanked dependency of mypy "types-pkg-resources" with "types-setuptools" as recommended: https://pypi.org/project/types-pkg-resources/

* Improve label imputation in PerturbationSpace class
Key changes:
- Now uses KNN graph in adata: saves cost and increases consistency
- Vectorized operations instead of expensive for loop
- Distance weighting for KNN imputation
- Quantifies uncertainty as local KNN label entropy

.pre-commit-config.yaml

@@ -32,7 +32,7 @@ repos:
  - id: mypy
  args: [--no-strict-optional, --ignore-missing-imports]
  additional_dependencies:
- ["types-pkg-resources", "types-requests", "types-attrs"]
+ ["types-setuptools", "types-requests", "types-attrs"]
  - repo: local
  hooks:
  - id: forbid-to-commit

pertpy/tools/_perturbation_space/_perturbation_space.py

@@ -7,6 +7,7 @@
 from anndata import AnnData
 from lamin_utils import logger
 from rich import print
+from scipy.stats import entropy
 
 if TYPE_CHECKING:
  from collections.abc import Iterable
@@ -364,50 +365,58 @@ def label_transfer(
  self,
  adata: AnnData,
  column: str = "perturbation",
+ column_uncertainty_score_key: str = "perturbation_transfer_uncertainty",
  target_val: str = "unknown",
- n_neighbors: int = 5,
- use_rep: str = "X_umap",
+ neighbors_key: str = "neighbors",
  ) -> None:
  """Impute missing values in the specified column using KNN imputation in the space defined by `use_rep`.
 
+ Uncertainty is calculated as the entropy of the label distribution in the neighborhood of the target cell.
+ In other words, a cell where all neighbors have the same set of labels will have an uncertainty of 0, whereas a cell
+ where all neighbors have many different labels will have high uncertainty.
+
  Args:
  adata: The AnnData object containing single-cell data.
- column: The column name in AnnData object to perform imputation on.
+ column: The column name in adata.obs to perform imputation on.
+ column_uncertainty_score_key: The column name in adata.obs to store the uncertainty score of the label transfer.
  target_val: The target value to impute.
- n_neighbors: Number of neighbors to use for imputation.
- use_rep: The key in `adata.obsm` where the embedding (UMAP, PCA, etc.) is stored.
+ neighbors_key: The key in adata.uns where the neighbors are stored.
 
  Examples:
  >>> import pertpy as pt
  >>> import scanpy as sc
  >>> import numpy as np
  >>> adata = sc.datasets.pbmc68k_reduced()
- >>> rng = np.random.default_rng()
- >>> adata.obs["perturbation"] = rng.choice(
- ... ["A", "B", "C", "unknown"], size=adata.n_obs, p=[0.33, 0.33, 0.33, 0.01]
- ... )
+ >>> # randomly dropout 10% of the data annotations
+ >>> adata.obs["perturbation"] = adata.obs["louvain"].astype(str).copy()
+ >>> random_cells = np.random.choice(adata.obs.index, int(adata.obs.shape[0] * 0.1), replace=False)
+ >>> adata.obs.loc[random_cells, "perturbation"] = "unknown"
  >>> sc.pp.neighbors(adata)
  >>> sc.tl.umap(adata)
  >>> ps = pt.tl.PseudobulkSpace()
- >>> ps.label_transfer(adata, n_neighbors=5, use_rep="X_umap")
+ >>> ps.label_transfer(adata)
  """
- if use_rep not in adata.obsm:
- raise ValueError(f"Representation {use_rep} not found in the AnnData object.")
-
- embedding = adata.obsm[use_rep]
-
- from pynndescent import NNDescent
-
- nnd = NNDescent(embedding, n_neighbors=n_neighbors)
- indices, _ = nnd.query(embedding, k=n_neighbors)
-
- perturbations = np.array(adata.obs[column])
- missing_mask = perturbations == target_val
-
- for idx in np.where(missing_mask)[0]:
- neighbor_indices = indices[idx]
- neighbor_categories = perturbations[neighbor_indices]
- most_common = pd.Series(neighbor_categories).mode()[0]
- perturbations[idx] = most_common
-
- adata.obs[column] = perturbations
+ if neighbors_key not in adata.uns:
+ raise ValueError(f"Key {neighbors_key} not found in adata.uns. Please run `sc.pp.neighbors` first.")
+
+ labels = adata.obs[column].astype(str)
+ target_cells = labels == target_val
+
+ connectivities = adata.obsp[adata.uns[neighbors_key]["connectivities_key"]]
+ # convert labels to an incidence matrix
+ one_hot_encoded_labels = adata.obs[column].astype(str).str.get_dummies()
+ # convert to distance-weighted neighborhood incidence matrix
+ weighted_label_occurence = pd.DataFrame(
+ (one_hot_encoded_labels.values.T * connectivities).T,
+ index=adata.obs_names,
+ columns=one_hot_encoded_labels.columns,
+ )
+ # choose best label for each target cell
+ best_labels = weighted_label_occurence.drop(target_val, axis=1)[target_cells].idxmax(axis=1)
+ adata.obs[column] = labels
+ adata.obs.loc[target_cells, column] = best_labels
+
+ # calculate uncertainty
+ uncertainty = np.zeros(adata.n_obs)
+ uncertainty[target_cells] = entropy(weighted_label_occurence.drop(target_val, axis=1)[target_cells], axis=1)
+ adata.obs[column_uncertainty_score_key] = uncertainty

tests/tools/_perturbation_space/test_simple_perturbation_space.py

@@ -240,9 +240,17 @@ def test_label_transfer():
  adata = AnnData(X)
  perturbations = np.array(["A", "B", "C"] * 22 + ["unknown"] * 3)
  adata.obs["perturbation"] = perturbations
+
+ with pytest.raises(ValueError):
+ ps = pt.tl.PseudobulkSpace()
+ ps.label_transfer(adata)
+
  sc.pp.neighbors(adata, use_rep="X")
  sc.tl.umap(adata)
-
  ps = pt.tl.PseudobulkSpace()
- ps.label_transfer(adata, n_neighbors=5, use_rep="X_umap")
+ ps.label_transfer(adata)
  assert "unknown" not in adata.obs["perturbation"]
+ assert all(adata.obs["perturbation_transfer_uncertainty"] >= 0)
+ assert not all(adata.obs["perturbation_transfer_uncertainty"] == 0)
+ is_known = perturbations != "unknown"
+ assert all(adata.obs.loc[is_known, "perturbation_transfer_uncertainty"] == 0)