miloDE ideas for python

[MaximilianNuber]

Dear all,

Thank you for this great tool. I am very interested in exploring the combination of milo and miloDE in a variety of datasets.
However, analyzing datasets of 200k cells and above, I just haven´t been able to make miloDE fast enough, or at times prevent it from crashing.

Therefore, I have been attempting to bring at least some functionality to python, using the pertpy-implementation of milo.

I am not computationally trained by background learning everything on the go, so I have questions regarding algorithms and implementation and would appreciate any help you can give me.

My first questions regards assign_neighbourhoods:
To my understanding, the function builds the milo-graph at k_init for sampling vertices, rebuilds the graph at k, and then refines the vertices with the vertices previously sampled at k_init. Are these steps not the same, but similar, to building a graph with miloR and running makeNhoods with graph-refinement?
I had significant issues calling RcppGreedySetCover with rpy2, so I use the greedy set cover approach from https://www.researchgate.net/publication/362015207_ML2R_Coding_Nuggets_Greedy_Set_Cover_with_Native_Python_Data_Types

import pertpy as pt
def assign_neighbourhoods(
    adata,
    k = 25,
    reducedDim_name = "X_pca",
    prop = 0.1,
):
    milo = pt.tl.Milo()
    mdata = milo.load(adata)
    # from sklearn_ann.kneighbors.annoy import AnnoyTransformer
    # transformer = AnnoyTransformer(n_neighbors=k)
    sc.pp.neighbors(mdata["rna"], use_rep=reducedDim_name,# , transformer = transformer
                    n_neighbors = k
                   )
    
    milo.make_nhoods(mdata["rna"], prop=prop)
    
    # nhood_size = np.array(mdata["rna"].obsm["nhoods"].sum(0)).ravel()
    # plt.hist(nhood_size, bins=100)
    # plt.xlabel("# cells in nhood")
    # plt.ylabel("# nhoods");
    # plt.show()
    return mdata

def greedySetCoverV2(U, S):
    # build inverted index

    I = {x : [i for i in S if x in S[i]] for x in U}

    # pre-select necessary sets Sj into C
    C = {j : S[j] for j in [I[x][0] for x in I if len(I[x]) == 1]}

    U = U - set().union(*[C[j] for j in C])

    while U:
        j = max(S.keys(), key = lambda k: len(S[k] & U))
        C[j] = S[j]
        U = U - S[j]

    return C


def _filter_neighbourhoods(mdata):
    nhood_df = pd.DataFrame()

    for nhood_idx in range(mdata["rna"].obsm["nhoods"].shape[1]):
        cur_nhoods = mdata["rna"].obsm["nhoods"][:, nhood_idx].toarray().ravel()
    
        cur_cells = adata.obs_names[cur_nhoods.astype(bool)]
        
        cur_df = pd.DataFrame({
            "set": nhood_idx,
            "elements": cur_cells
        })
        nhood_df = pd.concat([nhood_df, cur_df], axis = 0)

    S = []
    for s in nhood_df.set.unique():
        S.append(set(nhood_df.query('set == @s').elements.unique()))
    S = dict(enumerate(S))
    U = nhood_df.elements.unique()
    U = set(U)

    gsc_result = greedySetCoverV2(U, S)
    nhs_to_use = list(gsc_result.keys())
    nhs_to_use = np.asarray(nhs_to_use)

    index_for_refined = adata.obs[adata.obs.nhood_ixs_refined == 1].index[nhs_to_use.astype(int)]
    
    mdata["rna"].obs["nhood_ixs_refined"] = 0
    mdata["rna"].obs.loc[index_for_refined, "nhood_ixs_refined"] = 1
    
    mdata["rna"].obsm["nhoods"] = mdata["rna"].obsm["nhoods"][:, nhs_to_use.astype(int)]

As the python implementation does not include graph refinement, given that my assumptions above are true, my python implementation first builds a kNN-graph with scanpy/annoy. Then we run make_nhoods, which uses (non-graph) sampling for refined neighbourhoods as in the default miloR/milopy workflow. Refined sampling is (I think) performed on the adjacency matrix.

I am aware this is quite a deviation from the approach in R, but for the time being I would like to build upon that. Or do you think the non-graph refinement is ill-suited for this task? One thing I noticed when running the python implementation is, that neighbourhood sizes are a lot smaller when compared to your original package for the same values of k on the embryo dataset from the tutorial (about 10-fold, see example below). This also leads to a higher number neighbourhoods after filtering. But is this plausible?
I tested graph building in R and making the graph in python, so I found out that the higher neighbourhood sizes stem from the R implementation of buildGraph.
Is it sufficient to simply set k higher to achieve the desired neighbourhood sizes?

For calculating the AUC, pertpy provides an implementation of Augur. However, in your implementation, calculate_auc is called per neighbourhood with the parameter n_subsamples = 0. This is not possible in the python version. I do not fully understand what exactly this option does, except being a special case within Augur. Could you explain why n_subsamples = 0, or if other values might be possible?
Subsequently, I would post a feature request on the repository.
However, some tests I did suggest, that AUC values are generally high for most, if not all neighbourhoods when using python.

Then the DE testing per neighbourhood. For this, I call edgeR using an adapted version of the EdgeR-class in pertpy. Even parallelization worked well in my tests.
And finally, the spatial adjustment of p-values across neighbourhoods. Below I am posting my python-implementation of your code. This is my own translation, so I would appreciate any help to correct it.

import numpy as np

def get_weights(nhoods_x):
  

  intersect_mat = np.dot(nhoods_x.T, nhoods_x)
  t_connect = np.sum(intersect_mat, axis=1)
  weights = 1 / t_connect
  return weights

def _return_null_df():
    null_df = pd.DataFrame(index = adata.var_names)
    null_df["log_fc"] = np.nan
    null_df["logCPM"] = np.nan
    null_df["F"] = np.nan
    null_df["p_value"] = np.nan
    null_df["adj_p_value"] = np.nan
    null_df = null_df.reset_index(names = ["variable"])
    return null_df

def _run_edger(ad, design, contrast, model):
    try:
        mod = model(ad, design)
        mod.fit()
        res_df = mod._test_single_contrast(contrast)
        return res_df
    except:
        return _return_null_df()

def de_stat_neighbourhoods(
    mdata,
    sample_col = "patient",
    design = "~condition",
    covariates = "condition",
    contrast = None,
    # subset_nhoods = stat_auc$Nhood[!is.na(stat_auc$auc)],
    layout = "X_umap", 
    n_jobs = 1,
):
    def get_cells_in_nhoods(mdata, nhood_ids):
        '''
        Get cells in neighbourhoods of interest '''
        in_nhoods = np.array(mdata["rna"].obsm['nhoods'][:,nhood_ids.astype('int')].sum(1))
        ad = mdata["rna"][in_nhoods.astype(bool), :].copy()
        return ad

    covs = covariates

    from itertools import repeat
    from joblib import Parallel, delayed

    nhoods = mdata["rna"].obsm["nhoods"]
    n_nhoods = nhoods.shape[1]
    # n_nhoods = 100
    # get generator with all nhoods

    func = "sum"
    layer = "counts"
    
    all_nhoods = (get_cells_in_nhoods(mdata, nhood_ids = np.asarray([i])) for i in range(n_nhoods))
    aggregated_nhoods = (
        sc.get.aggregate(ad, by = [sample_col, covs], func = func, layer=layer)
        for ad in all_nhoods
    )
    def _func_to_X(ad, func):
        ad.X = ad.layers[func].copy()
        return ad
    aggregated_nhoods = (
        _func_to_X(ad, func)
        for ad in aggregated_nhoods
    )

    import diff_exp

    if contrast is None:
        from formulaic import model_matrix
        mm = model_matrix("~condition", mdata["rna"].obs)
        n_coefs = mm.shape[1]
        contrast = np.zeros(n_coefs)
        contrast[n_coefs-1] = 1
        contrast = list(contrast)
        del mm
        
    args = zip(
        aggregated_nhoods, 
        repeat(design),
        repeat(contrast),
        repeat(diff_exp.EdgeR)
    )

    start = time.time()
    pval_df_list = Parallel(n_jobs=n_jobs)(delayed(_run_edger)(ad, design, contrast, model) for ad, design, contrast, model in args)
    end = time.time()
    print("edger done in")
    print(end-start)
    print("seconds")
    # return pval_df_list

    pval_by_nhood = pd.concat([df.set_index("variable")[["p_value"]] for df in pval_df_list], axis = 1)
    FDR_across_genes = pd.concat([df.set_index("variable")[["adj_p_value"]] for df in pval_df_list], axis = 1)
    logfc_nhoods = pd.concat([df.set_index("variable")[["log_fc"]] for df in pval_df_list], axis = 1)
    
    print("Calculating FDR across nhoods")

    n_genes = pval_by_nhood.shape[0]
    FDR_across_nhoods = pval_by_nhood.copy()
    for i in range(n_genes):
        
        pvalues = pval_by_nhood.iloc[i, :].values
        idx_not_nan = np.isnan(pvalues)
        
        _weights = get_weights(mdata["rna"].obsm["nhoods"][:, ~idx_not_nan])
        _weights = np.asarray(_weights)
        _weights = _weights.ravel()
        # o = order_matrix(pvalues)
        pvalues = pvalues[~idx_not_nan]
        
        o = np.argsort(pvalues)
        weights = _weights[o]
        pvalues = pvalues[o]
        _weights
        # weights = weights[o]
        adjp = np.zeros(len(o))
        adjp[o] = np.flip(np.minimum.accumulate(np.flip(np.sum(weights) * pvalues / np.cumsum(weights))))
        adjp = np.minimum(adjp, 1)
        FDR_across_nhoods.iloc[i, :] = np.nan
    
        FDR_across_nhoods.iloc[i, ~idx_not_nan] = adjp

    return pval_by_nhood, logfc_nhoods, FDR_across_genes, FDR_across_nhoods

I would appreciate any help to make this work, and thank you in advance.
Best,
max

Here an example using the miloDE tutorial data:

r('''
library(miloDE)

# library containing toy data
suppressMessages(library(MouseGastrulationData))

# analysis libraries
library(scuttle)
suppressMessages(library(miloR))
suppressMessages(library(uwot))
library(scran)
suppressMessages(library(dplyr))
library(reshape2)

# packages for gene cluster analysis
# library(scWGCNA)
#> 
# suppressMessages(library(Seurat))

# plotting libraries
library(ggplot2)
library(viridis)
#> Loading required package: viridisLite
library(ggpubr)

''')

r('''
library(BiocParallel)

ncores = 4
mcparam = MulticoreParam(workers = ncores)
register(mcparam)

# load chimera Tal1
sce = suppressMessages(MouseGastrulationData::Tal1ChimeraData())

# downsample to few selected cell types
cts = c("Spinal cord" , "Haematoendothelial progenitors", "Endothelium" , "Blood progenitors 1" , "Blood progenitors 2")
sce = sce[, sce$celltype.mapped %in% cts]
# let's rename Haematoendothelial progenitors
sce$celltype.mapped[sce$celltype.mapped == "Haematoendothelial progenitors"] = "Haem. prog-s."

# delete row for tomato
sce = sce[!rownames(sce) == "tomato-td" , ]

# add logcounts
sce = logNormCounts(sce)

# update tomato field to be more interpretable 
sce$tomato = sapply(sce$tomato , function(x) ifelse(isTRUE(x) , "Tal1_KO" , "WT"))

# for this exercise, we focus on 3000 highly variable genes (for computational efficiency)
dec.sce = modelGeneVar(sce)
hvg.genes = getTopHVGs(dec.sce, n = 3000)
sce = sce[hvg.genes , ]
# change rownames to symbol names
rowdata = as.data.frame(rowData(sce))
rownames(sce) = rowdata$SYMBOL

# add UMAPs
set.seed(32)
umaps = as.data.frame(uwot::umap(reducedDim(sce , "pca.corrected")))
# let's store UMAPs - we will use them for visualisation
reducedDim(sce , "UMAP") = umaps

# plot UMAPs, colored by cell types
umaps = cbind(as.data.frame(colData(sce)) , reducedDim(sce , "UMAP"))
names(EmbryoCelltypeColours)[names(EmbryoCelltypeColours) == "Haematoendothelial progenitors"] = "Haem. prog-s."
# cols_ct = EmbryoCelltypeColours[names(EmbryoCelltypeColours) %in% unique(umaps$celltype.mapped)]
''')

# _r_to_py is a self-written helper function for rpy conversion
counts = _r_to_py(r('counts(sce)'))
counts = counts.T
counts
obs = _r_to_py(r('as.data.frame(colData(sce))'))
var = _r_to_py(r('as.data.frame(rowData(sce))'))
X = _r_to_py(r('logcounts(sce)'))
X = X.T
pca_corrected = _r_to_py(r('reducedDim(sce, "pca.corrected")'))
X_umap = _r_to_py(r('reducedDim(sce, "UMAP")'))

adata = anndata.AnnData(
    X = X,
    obs = obs,
    var = var
)
adata.layers["counts"] = counts
adata.obsm["pca_corrected"] = pca_corrected
adata.obsm["X_umap"] = X_umap
adata.obsm["X_umap"] = adata.obsm["X_umap"].values

sc.pl.umap(adata,
          color = "celltype.mapped")

%%time
plot = True
LATENT_KEY = "pca_corrected"
k_grid = list(range(10, 40, 5))
# k_grid.append(100)
prop = 0.5
filtering = True

size_df = pd.DataFrame()
for k in k_grid:
    print(k)
    mdata = assign_neighbourhoods(adata, k = k, prop = prop)
    if filtering:
        print("filtering")
        _filter_neighbourhoods(mdata)
    nhood_size = np.array(mdata["rna"].obsm["nhoods"].sum(0)).ravel()
    _size_df = pd.DataFrame(dict(nhood_size= nhood_size))
    _size_df = _size_df.assign(k = str(k))
    size_df = pd.concat([size_df, _size_df], axis = 0)

if plot:
    sns.boxplot(
        size_df,
        x = "k",
        y = "nhood_size",
        palette = sns.color_palette("Paired", len(k_grid))
    )
    sns.despine()

%%time
pval, logfc, FDR_genes, FDR_nhoods = de_stat_neighbourhoods(
    mdata,
    sample_col = "sample",
    design = "0+tomato",
    covariates = "tomato",
    contrast=[-1, 1],
    n_jobs = 5
)

nhood_size = np.array(mdata["rna"].obsm["nhoods"].sum(0)).ravel()
mdata["milo"].var["Nhood_size"] = nhood_size
mdata["milo"].var["nhood_adjusted"] = np.sum(FDR_nhoods <= 0.05, axis = 0).values
mdata["milo"].var["gene_adjusted"] = np.sum(FDR_genes <= 0.05, axis = 0).values

sc.pl.embedding(mdata["milo"].T,
                basis = "X_milo_graph",
                size = mdata["milo"].var["Nhood_size"],
                color = "nhood_adjusted"
               )
sc.pl.embedding(mdata["milo"].T,
                basis = "X_milo_graph",
                size = mdata["milo"].var["Nhood_size"],
                color = "gene_adjusted"
               )

[Zethson]

Dear @MaximilianNuber

would you be interested in getting back to working on this and contributing it to pertpy? If yes, let's create an issue on pertpy and continue the discussion there.

Thank you!