GLUE and MOFA¶

from itertools import chain

import anndata as ad
import itertools
import networkx as nx
import pandas as pd
import scanpy as sc
import scglue
import seaborn as sns
from matplotlib import rcParams

import omicverse as ov
ov.utils.ov_plot_set()

guidance = scglue.genomics.rna_anchored_guidance_graph(rna, atac)
guidance

glue = scglue.models.fit_SCGLUE(
    {"rna": rna, "atac": atac}, guidance_hvf,
    fit_kws={"directory": "glue"}
)

dx = scglue.models.integration_consistency(
    glue, {"rna": rna, "atac": atac}, guidance_hvf
)
dx

_ = sns.lineplot(x="n_meta", y="consistency", data=dx).axhline(y=0.05, c="darkred", ls="--")

import scanpy as sc
rna=sc.read("rna-emb.h5ad")
atac=sc.read("atac-emb.h5ad")
rna,atac

import numpy as np

rna_com1=sc.AnnData(np.random.randint(0, 100, size=(61472, 1)),
                   obs=rna.obs)
rna_com1.obsm=rna.obsm.copy()

atac_com1=sc.AnnData(np.random.randint(0, 100, size=(130418, 1)),
                   obs=atac.obs)
atac_com1.obsm=atac.obsm.copy()

import anndata
combined1 = anndata.concat([rna_com1, atac_com1],merge='same')

from scvi.model.utils import mde
combined1.obsm["X_mde"] = mde(combined1.obsm["X_glue"])

import matplotlib.pyplot as plt
import omicverse as ov
fig, ax = plt.subplots(figsize=(3,3))
ov.pl.embedding(
    combined1,
    basis="X_mde",
    color=["domain"],
    title='Domain',
    frameon='small',
    ncols=2,
    palette=ov.utils.pyomic_palette()[11:],
    show=False,
    ax=ax,
    
)
plt.savefig("figures/mofa_umap_domain.png",dpi=300,bbox_inches = 'tight')
plt.savefig("pdf/mofa_umap_domain.pdf",dpi=300,bbox_inches = 'tight')

import matplotlib.pyplot as plt
import omicverse as ov
fig, ax = plt.subplots(figsize=(3,3))
ov.pl.embedding(
    rna,
    basis="X_mde",
    color=["Cell.Type"],
    title='Cell.Type',
    frameon='small',
    ncols=2,
    palette=ov.utils.pyomic_palette()[:],
    show=False,
    ax=ax,
    
)
plt.savefig("figures/mofa_rna_umap_ct.png",dpi=300,bbox_inches = 'tight')
plt.savefig("pdf/mofa_rna_umap_ct.pdf",dpi=300,bbox_inches = 'tight')

feature_embeddings = glue.encode_graph(guidance_hvf)
feature_embeddings = pd.DataFrame(feature_embeddings, index=glue.vertices)
feature_embeddings.iloc[:5, :5]

atac

rna

sc.pl.embedding(
    rna,
    basis="X_umap",
    color=["domain"],
    title='Domain',
    frameon=False,
    ncols=2,
    palette=ov.utils.pyomic_palette()[11:],
    show=False,
    #ax=ax,
    
)

from tqdm import tqdm,trange
import anndata
def glue_pair(rna:anndata.AnnData,
              atac:anndata.AnnData,depth:int=20)->pd.DataFrame:
    r"""
    Pair the cells between RNA and ATAC using result of GLUE.

    Arguments:
        rna: the AnnData of RNA-seq.
        atac: the AnnData of ATAC-seq.
        depth: the depth of the search for the nearest neighbor.
    
    """


    #提取GLUE层结果
    print('......Extract GLUE layer from obs')
    rna_loc=pd.DataFrame(rna.obsm['X_glue'], index=rna.obs.index)
    atac_loc=pd.DataFrame(atac.obsm['X_glue'], index=atac.obs.index)

    #对GLUE层进行Pearson系数分析
    print('......Prepare for pair')
    import gc
    len1=(len(rna_loc)//5000)+1
    len2=(len(atac_loc)//5000)+1
    if len1>len2:
        len1=len2
    p_pd=pd.DataFrame(columns=['rank_'+str(i) for i in range(50)])
    n_pd=pd.DataFrame(columns=['rank_'+str(i) for i in range(50)])
    print('......Start to calculate the Pearson coef')
    for j in range(len1):
        c=pd.DataFrame()
        with trange(len1) as tt:
            for i in tt:
                t1=rna_loc.iloc[5000*(i):5000*(i+1)]
                t2=atac_loc.iloc[5000*(j):5000*(j+1)]
                a=np.corrcoef(t1,t2)[len(t1):,0:len(t1)]
                b=pd.DataFrame(a,index=t2.index,columns=t1.index)  
                c=pd.concat([c,b],axis=1)
                del t1
                del t2
                del a
                del b
                gc.collect()
                tt.set_description('Now Pearson block is {}/{}'.format(i,len1))
        with trange(len(c)) as t:
            for i in t:
                t_c=c.iloc[i]
                p_pd.loc[t_c.name]=c.iloc[i].sort_values(ascending=False)[:50].values
                n_pd.loc[t_c.name]=c.iloc[i].sort_values(ascending=False)[:50].index.tolist()
                t.set_description('Now rna_index is {}/{}, all is {}'.format(i+j*5000,i+j*5000+len(c),len(atac_loc)))
        print('Now epoch is {}, {}/{}'.format(j,j*5000+len(c),len(atac_loc))) 
        del c
        gc.collect()
        
    #寻找最近的细胞，其中depth的灵活调整可以使得配对成功的细胞数变大，同时精度有所下降
    def find_neighbor_cell(p_pd,n_pd,depth=10):
        if depth>50:
            print('......depth limited to 50')
            depth=50
        rubish_c=[]
        finish_c=[]
        with trange(depth) as dt:
            for d in dt:
                p_pd=p_pd.loc[p_pd['rank_{}'.format(d)]>0.9]
                p_pd=p_pd.sort_values('rank_{}'.format(d),ascending=False)
                for i in p_pd.index:
                    name=n_pd.loc[i,'rank_{}'.format(d)]
                    if name not in rubish_c:
                        finish_c.append(i)
                        rubish_c.append(name)
                    else:
                        continue
                p_pd=p_pd.loc[~p_pd.index.isin(finish_c)]
                n_pd=n_pd.loc[~n_pd.index.isin(finish_c)]
                dt.set_description('Now depth is {}/{}'.format(d,depth))
        result=pd.DataFrame()
        result['omic_1']=rubish_c
        result['omic_2']=finish_c
        result.index=['cell_{}'.format(i) for i in range(len(result))]
        return result
    print('......Start to find neighbor')
    res_pair=find_neighbor_cell(p_pd,n_pd,depth=depth)
    return res_pair
res_pair=glue_pair(rna,atac,depth=10)

res_pair.shape

rna1=rna[res_pair['omic_1']]
rna1

atac1=atac[res_pair['omic_2']]
atac1

sc.pl.embedding(
    rna1,
    basis="X_mde",
    color=["Cell.Type", "domain"],
    frameon=False,
    ncols=2,
)

sc.pl.embedding(
    atac1,
    basis="X_mde",
    color=["Cell.Type", "domain"],
    frameon=False,
    ncols=2,
)

from mudata import MuData

mdata = MuData({'rna': rna1, 'atac': atac1})
mdata

from sklearn.metrics import adjusted_rand_score as ari
ari(mdata.obs['rna:Cell.Type'], mdata.obs['atac:Cell.Type'])

from sklearn.metrics import adjusted_rand_score as ari
ari(rna1[random_obs_index].obs['Cell.Type'], atac1[random_obs_index].obs['Cell.Type'])

test_mofa.mofa_preprocess()
test_mofa.mofa_run(outfile='models/ad_rna_atac.hdf5')

from mudata import MuData

mdata = MuData({'rna': rna1, 'atac': atac1})
mdata

mdata.obsm['X_mofa'].shape

import muon as mu
mu.pl.mofa(mdata, color="rna:Cell.Type", components=["1,2", "3,4"])

a1.head()

import pandas as pd
plot_data=pd.DataFrame()
for i in m.get_r2()['View'].unique():
    plot_data[i]=m.get_r2().loc[m.get_r2()['View']==i,'R2'].values
plot_data

import seaborn as sns
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(2,3))
sns.heatmap(plot_data,cmap='Greens',ax=ax,xticklabels=True,yticklabels=True,
            cbar_kws={'shrink':0.5})
plt.xticks(fontsize=10)
plt.yticks(fontsize=10)
plt.ylabel('Factor',fontsize=12)
plt.xlabel('View',fontsize=12)
plt.title('Varience',fontsize=12)
plt.savefig("figures/mofa_varience.png",dpi=300,bbox_inches = 'tight')
plt.savefig("pdf/mofa_varience.pdf",dpi=300,bbox_inches = 'tight')

f1_w=ov.single.get_weights(hdf5_path='models/ad_rna_atac.hdf5',view='RNA',factor=1)
f1_w.sort_values('weights',ascending=False)

f1_w.index=[str(i,"utf8").replace('RNA_','') for i in f1_w['feature']]
f1_w

ov.single.get_weights(hdf5_path='models/ad_rna_atac.hdf5',view='RNA',factor=1)

factor_w.head()

rna1=ov.single.factor_exact(rna1,hdf5_path='models/ad_rna_atac.hdf5')
rna1

import seaborn as sns
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(6,3))
sns.heatmap(plot_data1,cmap='Purples',ax=ax,square=True,cbar_kws={'shrink':0.5})
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.xlabel('Factor',fontsize=12)
plt.ylabel('Cell.Type',fontsize=12)
plt.title('Correlation',fontsize=12)
plt.savefig("figures/mofa_correlation.png",dpi=300,bbox_inches = 'tight')
plt.savefig("pdf/mofa_correlation.pdf",dpi=300,bbox_inches = 'tight')

m.get_factors().shape

fig, ax = plt.subplots(figsize=(3,3))
factor1,factor2=4,6
sc.pl.embedding(
    adata=rna1,
    basis='X_mofa',
    color='Cell.Type',
    title='INH',
    components="{},{}".format(factor1,factor2),
    palette=ov.utils.pyomic_palette(),
    #frameon=False,
    ncols=2,
    ax=ax
)
fig.savefig("figures/mofa_factor_{}_{}.png".format(factor1,factor2),dpi=300,bbox_inches = 'tight')

fig, ax = plt.subplots(figsize=(3,3))
sc.pl.embedding(
    rna1,
    basis="X_mde",
    color=["factor6"],
    frameon=False,
    ncols=2,
    #palette=ov.utils.pyomic_palette(),
    show=False,
    cmap='Greens',
    vmin=0,
    ax=ax
)
plt.savefig("figures/umap_factor6.png",dpi=300,bbox_inches = 'tight')

fig, ax = plt.subplots(figsize=(3,3))
sc.pl.embedding(
    rna1,
    basis="X_mde",
    color=["Cell.Type"],
    frameon=False,
    ncols=2,
    #palette=ov.utils.pyomic_palette(),
    show=False,
    cmap='Purples_r',
    ax=ax
)
plt.savefig("figures/umap_celltype_rna.png",dpi=300,bbox_inches = 'tight')

combined

sc.pp.neighbors(combined, use_rep="X_glue", metric="cosine")
#sc.tl.umap(combined)
#sc.pl.umap(combined, color=["Cell.Type", "domain"], wspace=0.65)

sc.pl.embedding(
    combined,
    basis="X_mde",
    color=["Cell.Type"],
    frameon=False,
    ncols=2,
    palette=ov.utils.pyomic_palette(),
    show=False,
)
plt.savefig("figures/umap_celltype.png",dpi=300,bbox_inches = 'tight')

sc.pl.embedding(
    combined,
    basis="X_mde",
    color=["domain"],
    title='Omics',
    frameon=False,
    ncols=2,
    palette=ov.utils.pyomic_palette()[11:],
    show=False,
)
plt.savefig("figures/umap_domain.png",dpi=300,bbox_inches = 'tight')

factor_w[['factor_4','factor_6']]

factor_w['factor_1'].sort_values(ascending=False)

fig, ax = plt.subplots(figsize=(3,4))
plot_data4=pd.DataFrame()
plot_data4['weight']=factor_w['factor_1'].sort_values(ascending=False)
plot_data4['rank']=range(len(plot_data4['weight']))
plt.plot(plot_data4['rank'],plot_data4['weight'],color='#a51616')

hub_gene=plot_data4.index[:5]
plt.scatter(plot_data4.loc[hub_gene,'rank'],
           plot_data4.loc[hub_gene,'weight'],color='#a51616',
            alpha=0.5)

from adjustText import adjust_text
texts=[ax.text(plot_data4.loc[i,'rank'],
                   plot_data4.loc[i,'weight'],
                   i,
                   fontdict={'size':10,'weight':'normal','color':'black'}
                   ) for i in hub_gene]

adjust_text(texts,only_move={'text': 'xy'},arrowprops=dict(arrowstyle='->', color='grey'),)

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(True)
ax.spines['left'].set_visible(True)

plt.title('factor1',fontsize=12,)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.xlabel('Feature rank',fontsize=12)
plt.ylabel('Weight',fontsize=12)


plt.grid(False)

fig, ax = plt.subplots(figsize=(3,3))
ax.scatter(plot_data3.loc[plot_data3['sig']=='normal','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='normal','factor_{}'.format(factor2)],
          color='#c2c2c2',alpha=0.5)

ax.scatter(plot_data3.loc[plot_data3['sig']=='up-up','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='up-up','factor_{}'.format(factor2)],
           color='#a51616',alpha=0.5)

ax.scatter(plot_data3.loc[plot_data3['sig']=='up-down','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='up-down','factor_{}'.format(factor2)],
           color='#e25d5d',alpha=0.5)

ax.scatter(plot_data3.loc[plot_data3['sig']=='down-up','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='down-up','factor_{}'.format(factor2)],
           color='#3b9868',alpha=0.5)

ax.scatter(plot_data3.loc[plot_data3['sig']=='down-down','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='down-down','factor_{}'.format(factor2)],
           color='#0d6a3b',alpha=0.5)

plt.vlines(x=0.5,ymin=-1,ymax=1,color='#a51616',linestyles='dashed')
plt.hlines(y=0.5,xmin=-1,xmax=1,color='#a51616',linestyles='dashed')

plt.vlines(x=-0.5,ymin=-1,ymax=1,color='#0d6a3b',linestyles='dashed')
plt.hlines(y=-0.5,xmin=-1,xmax=1,color='#0d6a3b',linestyles='dashed')

ax.spines['top'].set_visible(True)
ax.spines['right'].set_visible(True)
ax.spines['bottom'].set_visible(True)
ax.spines['left'].set_visible(True)

plt.grid(False)

plt.xlabel('factor_{}'.format(factor1),fontsize=12)
plt.ylabel('factor_{}'.format(factor2),fontsize=12)

plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

from adjustText import adjust_text

for sig,color in zip(['up-up','up-down','down-up','down-down'],
                     ['#a51616','#e25d5d','#3b9868','#0d6a3b']):
    hub_gene=plot_data3.loc[plot_data3['sig']==sig].index.tolist()
    if len(hub_gene)==0:
        continue
    texts=[ax.text(plot_data3.loc[i,'factor_{}'.format(factor1)],
                   plot_data3.loc[i,'factor_{}'.format(factor2)],
                   i,
                   fontdict={'size':10,'weight':'bold','color':'black'}
                   ) for i in hub_gene[:5]]

    adjust_text(texts,only_move={'text': 'xy'},arrowprops=dict(arrowstyle='->', color='grey'),)

plt.title('INH',fontsize=12)
plt.savefig("figures/factor_gene_inh.png",dpi=300,bbox_inches = 'tight')

fig, ax = plt.subplots(figsize=(3,3))
ax.scatter(plot_data3.loc[plot_data3['sig']=='normal','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='normal','factor_{}'.format(factor2)],
          color='#c2c2c2',alpha=0.5)

ax.scatter(plot_data3.loc[plot_data3['sig']=='up','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='up','factor_{}'.format(factor2)],
           color='#a51616',alpha=0.5)


ax.scatter(plot_data3.loc[plot_data3['sig']=='down','factor_{}'.format(factor1)],
           plot_data3.loc[plot_data3['sig']=='down','factor_{}'.format(factor2)],
           color='#0d6a3b',alpha=0.5)

plt.vlines(x=0.5,ymin=-1,ymax=1,color='#a51616',linestyles='dashed')
plt.hlines(y=0.5,xmin=-1,xmax=1,color='#a51616',linestyles='dashed')

plt.vlines(x=-0.5,ymin=-1,ymax=1,color='#0d6a3b',linestyles='dashed')
plt.hlines(y=-0.5,xmin=-1,xmax=1,color='#0d6a3b',linestyles='dashed')

ax.spines['top'].set_visible(True)
ax.spines['right'].set_visible(True)
ax.spines['bottom'].set_visible(True)
ax.spines['left'].set_visible(True)

plt.grid(False)

plt.xlabel('factor_{}'.format(factor1),fontsize=12)
plt.ylabel('factor_{}'.format(factor2),fontsize=12)

plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

from adjustText import adjust_text

for sig,color in zip(['up','down'],
                     ['#a51616','#0d6a3b']):
    if 'up' in sig:
        hub_gene=plot_data3.loc[plot_data3['sig']==sig].sort_values('factor_{}'.format(factor1),ascending=False).index.tolist()
    else:
        hub_gene=plot_data3.loc[plot_data3['sig']==sig].sort_values('factor_{}'.format(factor1),ascending=True).index.tolist()
    if len(hub_gene)==0:
        continue
    texts=[ax.text(plot_data3.loc[i,'factor_{}'.format(factor1)],
                   plot_data3.loc[i,'factor_{}'.format(factor2)],
                   i,
                   fontdict={'size':10,'weight':'bold','color':'black'}
                   ) for i in hub_gene[:5]]

    adjust_text(texts,only_move={'text': 'xy'},arrowprops=dict(arrowstyle='->', color='grey'),)

plt.title('INH',fontsize=12)
plt.savefig("figures/factor_gene_inh.png",dpi=300,bbox_inches = 'tight')