Add mixscape

Author: jaydu1

paper/methods/cinemaot.py

@@ -7,6 +7,7 @@
 import scanpy as sc
 from scipy.stats import wilcoxon
 from statsmodels.stats.multitest import multipletests
+from sklearn.neighbors import NearestNeighbors
 
 
 def run_cinemaot(Y, A, raw=False, weighted=False, thres=0.15, smoothness=1e-3, **kwargs):
@@ -69,3 +70,69 @@ def run_cinemaot(Y, A, raw=False, weighted=False, thres=0.15, smoothness=1e-3, *
 
     df = pd.DataFrame({'stat':stat, 'pvalue':pvalue, 'padj':padj})
     return {'cinemaot.df':df, 'cinemaot.res':{'W':cf, 'Y_hat_0':Y_hat_0, 'Y_hat_1':Y_hat_1, 'de':de}}
+
+
+
+def run_mixscape(Y, A, raw=False, nn=20, **kwargs):
+    '''
+    The Python wrapper for running Mixscape.
+
+    Parameters
+    ----------
+    Y : np.ndarray
+        The (preprocessed) gene expression matrix of shape [n,p].
+    A : np.ndarray
+        The treatment vector of shape [n,].
+    raw : bool
+        Whether the input data is raw or preprocessed.
+    nn : int
+        The number of nearest neighbors to use.
+
+    Returns
+    -------
+    A dictionary containing the results dataframe and other matrices.
+    '''
+    adata = sc.AnnData(Y.copy())
+    if A.ndim > 1:
+        A = A[:, 0]
+    adata.obs['A'] = A
+
+    if raw:
+        sc.pp.normalize_total(adata, target_sum=1e4)
+        sc.pp.log1p(adata)
+        sc.pp.scale(adata, max_value=10)
+    
+    sc.pp.pca(adata)
+
+    Y_hat_0 = np.zeros_like(Y)
+    Y_hat_1 = np.zeros_like(Y)
+    de = np.zeros_like(Y)
+
+    # Calculate counterfactual for treated group (A=1)
+    X_pca_ctrl = adata.obsm['X_pca'][adata.obs['A'] == 0, :]
+    X_pca_trt = adata.obsm['X_pca'][adata.obs['A'] == 1, :]
+    nbrs_ctrl = NearestNeighbors(n_neighbors=nn, algorithm='ball_tree').fit(X_pca_ctrl)
+    mixscape_matrix_trt = nbrs_ctrl.kneighbors_graph(X_pca_trt).toarray()
+    
+    Y_hat_0[A == 1] = (mixscape_matrix_trt / np.sum(mixscape_matrix_trt, axis=1, keepdims=True)) @ Y[A == 0]
+    Y_hat_0[A == 0] = Y[A == 0]
+
+    # Calculate counterfactual for control group (A=0)
+    X_pca_trt = adata.obsm['X_pca'][adata.obs['A'] == 1, :]
+    X_pca_ctrl = adata.obsm['X_pca'][adata.obs['A'] == 0, :]
+    nbrs_trt = NearestNeighbors(n_neighbors=nn, algorithm='ball_tree').fit(X_pca_trt)
+    mixscape_matrix_ctrl = nbrs_trt.kneighbors_graph(X_pca_ctrl).toarray()
+
+    Y_hat_1[A == 0] = (mixscape_matrix_ctrl / np.sum(mixscape_matrix_ctrl, axis=1, keepdims=True)) @ Y[A == 1]
+    Y_hat_1[A == 1] = Y[A == 1]
+
+    de[A == 1] = Y[A == 1] - Y_hat_0[A == 1]
+    de[A == 0] = Y_hat_1[A == 0] - Y[A == 0]
+
+    stat, pvalue = list(zip(*[wilcoxon(de[:, j], zero_method='zsplit') for j in range(de.shape[1])]))
+    padj = multipletests(pvalue, alpha=0.05, method='fdr_bh')[1]
+
+    df = pd.DataFrame({'stat': stat, 'pvalue': pvalue, 'padj': padj})
+    
+    # The concept of a coupling matrix 'W' is specific to OT, returning None.
+    return {'mixscape.df': df, 'mixscape.res': {'W': adata.obsm['X_pca'], 'Y_hat_0': Y_hat_0, 'Y_hat_1': Y_hat_1, 'de': de}}

paper/methods/metrics.py

@@ -25,7 +25,7 @@ def comp_stat(true, pred, c):
 
 
 def comp_score(Y, CF, celltype, Z, W_hat=None):
-    adata = sc.AnnData(CF)
+    adata = sc.AnnData(CF.astype('float'))
     # adata.obs['covariate'] = W[:,-1]
     # adata.obs['trt'] = A