Add uncertainty score in KNN label_transfer in PerturbationSpace

[stefanpeidli]

PR Checklist

• Referenced issue is linked (NA; addresses reviewer comment)
• If you've fixed a bug or added code that should be tested, add tests!
• Documentation in docs is updated (NA)

Description of changes

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. This is saved in a new specified column in adata.obs as a score between 0 and 1. Cells that had a label to begin with will get highest certainty assigned.

Additional context

[stefanpeidli]

@Zethson Norman et al. analysis has to be updated with this. I can do that if needed, just LMK.

[Zethson]

Great, thank you!

I was wondering whether it's more useful to model a certainty score or an uncertainty score. Don't people usually look at the latter? They are of course the same but maybe a 1- is the way to go?

Also, I think that this is Occam's razor. But I think we can make a better case if we implement something slightly more sophisticated that actually isn't much more work. Issues are:

1. Dependence on n_neighbors: The certainty measure is highly dependent on the choice of n_neighbors. A small number might lead to overfitting, while a large number might smooth out important local structures.
2. Lack of distance weighting: All neighbors are treated equally, regardless of their distance in the embedding space. This might not accurately represent the data structure, especially in areas of varying density.
3. Binary treatment of labels: The method doesn't consider similarities between different labels. For example, if there are labels A, B, and C, and B is more similar to A than C is, this information is not used.

There's more issues but these we can't easily tackle. I'm copy and pasting a solution by Claude here. Not because we should copy and paste this but so that you see how we can potentially do that without a lot of work.

import numpy as np
import pandas as pd
from pynndescent import NNDescent
from sklearn.neighbors import KernelDensity
from scipy.spatial.distance import pdist, squareform

def improved_label_transfer(
    adata,
    column="perturbation",
    column_certainty_score="perturbation_transfer_certainty",
    target_val="unknown",
    max_neighbors=50,
    use_rep="X_umap",
    label_similarity_matrix=None
):
    """Improved imputation of missing values using adaptive KNN and weighted voting.
    
    Args:
        adata: AnnData object containing single-cell data.
        column: Column name in AnnData object to perform imputation on.
        column_certainty_score: Column name to store the certainty score.
        target_val: The target value to impute.
        max_neighbors: Maximum number of neighbors to consider.
        use_rep: Key in adata.obsm where the embedding is stored.
        label_similarity_matrix: Dictionary of label similarities. If None, treats all labels as equally dissimilar.
    """
    if use_rep not in adata.obsm:
        raise ValueError(f"Representation {use_rep} not found in the AnnData object.")
    
    embedding = adata.obsm[use_rep]
    
    # Estimate local density
    kde = KernelDensity(kernel='gaussian', bandwidth=0.5).fit(embedding)
    local_density = np.exp(kde.score_samples(embedding))
    
    # Adaptive neighbor selection
    adaptive_n_neighbors = np.clip(np.int32(max_neighbors * (1 - local_density)), 5, max_neighbors)
    
    # Find neighbors
    nnd = NNDescent(embedding, n_neighbors=max_neighbors)
    indices, distances = nnd.query(embedding, k=max_neighbors)
    
    perturbations = np.array(adata.obs[column])
    certainty = np.ones(adata.n_obs)
    missing_mask = perturbations == target_val
    
    # Create label similarity matrix if not provided
    if label_similarity_matrix is None:
        unique_labels = np.unique(perturbations[~missing_mask])
        label_similarity_matrix = {l1: {l2: 1.0 if l1 == l2 else 0.0 for l2 in unique_labels} for l1 in unique_labels}
    
    for idx in np.where(missing_mask)[0]:
        n_neighbors = adaptive_n_neighbors[idx]
        neighbor_indices = indices[idx][:n_neighbors]
        neighbor_distances = distances[idx][:n_neighbors]
        neighbor_categories = perturbations[neighbor_indices]
        
        # Distance-weighted voting
        weights = 1 / (neighbor_distances + 1e-8)  # Add small epsilon to avoid division by zero
        
        # Consider label similarities
        vote_dict = {}
        for cat, weight in zip(neighbor_categories, weights):
            for possible_cat in label_similarity_matrix[cat]:
                similarity = label_similarity_matrix[cat][possible_cat]
                vote_dict[possible_cat] = vote_dict.get(possible_cat, 0) + weight * similarity
        
        most_common = max(vote_dict, key=vote_dict.get)
        perturbations[idx] = most_common
        
        # Calculate certainty based on weighted votes
        total_votes = sum(vote_dict.values())
        certainty[idx] = vote_dict[most_common] / total_votes if total_votes > 0 else 0

    adata.obs[column] = perturbations
    adata.obs[column_certainty_score] = certainty

# Example usage:
# improved_label_transfer(adata, label_similarity_matrix={'A': {'A': 1.0, 'B': 0.5, 'C': 0.1}, 'B': {'A': 0.5, 'B': 1.0, 'C': 0.3}, 'C': {'A': 0.1, 'B': 0.3, 'C': 1.0}})

What do you think? We can also keep your simple implementation as a flavor but I think we should do this in a better way to argue better in the manuscript

[Zethson]

FAILED tests/tools/_perturbation_space/test_simple_perturbation_space.py::test_label_transfer - AttributeError: 'AnnData' object has no attribute 'loc'

The test failed by the way

[stefanpeidli]

Great suggestions @Zethson. I tried to keep the (un)certainty measure as simple as the label transfer method (which is too simplistic). Let's improve the label transfer itself then first?

My senf:

• Yes let's make this an uncertainty score instead.
• If we use distance weighting (which I agree we should), we can increase the number of neighbors more safely so the choices for this parameter have less impact in general. I think this alone should suffice to improve the label transfer and uncertainty score.
• Instead of label similarity (which uses global information that might not necessarily hold locally, I do not like this) I propose we use the entropy of the local class distribution for the uncertainty. If two classes are very similar (locally!), this case should still get a higher certainty score than a case where every class gets has the same number of neighbors to the query node. Is that understandable?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Zethson]

Instead of label similarity (which uses global information that might not necessarily hold locally, I do not like this) I propose we use the entropy of the local class distribution for the uncertainty. If two classes are very similar (locally!), this case should still get a higher certainty score than a case where every class gets has the same number of neighbors to the query node. Is that understandable?

Yes! +1

[stefanpeidli]

I reworked the function, it's now much better I think. Key changes:

• Now uses KNN graph in adata: saves cost and increases consistency e.g. with UMAPs generated from same KNN.
• Now vectorized operations instead of expensive for loop, MUCH faster
• Using distance weighting for KNN imputation --> more robust
• Quantifies uncertainty as local KNN label entropy. Basically measures how sure the classifier is with its prediction. We do not account for global similarity between classes.
• (tests passed at least locally :))))

Also, I changed the example from random labels to louvain with randomly dropped labels.

Here is some example result:

And somewhat hidden in the commits: I removed a faulty dependency in pre-commit config that was preventing me from using it. See commit message: "Replaces yanked dependency of mypy "types-pkg-resources" with "types-setuptools" as recommended: https://pypi.org/project/types-pkg-resources/".

[Zethson]

Thanks! I'll have a look soon. But always use Leiden instead of Louvain, please :)

[stefanpeidli]

I'm using the louvain clusters because they came with the preprocessed dataset. Otherwise I'd 100% agree with always using leiden!

[Zethson]

Thank you very much! This is much better.

1. The test still fails. The only source of truth is the CI and "it works on my machine" is sadly not good enough ^_^

FAILED tests/tools/_perturbation_space/test_simple_perturbation_space.py::test_label_transfer - assert False
 +  where False = all(0     0.000000\n1     0.000000\n2     0.000000\n3     0.000000\n4     0.000000\n        ...   \n64    0.000000\n65    0.000000\n66    0.812785\n67    0.933854\n68    0.981642\nName: perturbation_transfer_uncertainty, Length: 69, dtype: float64 == 0)

3. Could you please always resolve my comments? Then I know what you've already done and what not.
4. I stopped commenting on those at some point but you have many lines of code where you added a single comment above outlining what the following line of code does. This is not soo useful and noisy. Either the line of code is well written with clear variable names and it's obvious what it does or it needs to be improved. Always document WHY you're doing certain things if it matters but don't duplicate the "what"

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 65.56%. Comparing base (99dcd18) to head (8660301).
Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #658      +/-   ##
==========================================
+ Coverage   65.55%   65.56%   +0.01%     
==========================================
  Files          47       47              
  Lines        6105     6105              
==========================================
+ Hits         4002     4003       +1     
+ Misses       2103     2102       -1     

Files with missing lines	Coverage Δ
...y/tools/_perturbation_space/_perturbation_space.py	87.11% <100.00%> (+0.51%)	⬆️

... and 1 file with indirect coverage changes

[Zethson]

Great! I think that the function description might benefit from a 2 sentence description on how the uncertainty is calculated.

But this is ready to go now!