# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# This notebook contains the code for the meta-analysis of healthy lung data for ACE2, TMPRSS2, and CTSL. It contains the hold-out analysis for the complex model with interaction terms that was run on the cell-level data. This script contains the code that was run on the full data and does not test for smoking associations.

import scanpy as sc
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import pandas as pd
from matplotlib import rcParams
from matplotlib import colors
from matplotlib import patches
import seaborn as sns
import batchglm
import diffxpy.api as de
import patsy as pat
from statsmodels.stats.multitest import multipletests
import logging, warnings
import statsmodels.api as sm

# +
plt.rcParams['figure.figsize']=(8,8) #rescale figures
sc.settings.verbosity = 3
#sc.set_figure_params(dpi=200, dpi_save=300)
sc.logging.print_versions()
de.__version__

logging.getLogger("tensorflow").setLevel(logging.ERROR)
logging.getLogger("batchglm").setLevel(logging.INFO)
logging.getLogger("diffxpy").setLevel(logging.INFO)

pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 35)
warnings.filterwarnings("ignore", category=DeprecationWarning, module="tensorflow")

# +
#User inputs
folder = '/storage/groups/ml01/workspace/malte.luecken/2020_cov19_study'

adata_diffxpy = '/storage/groups/ml01/workspace/malte.luecken/2020_cov19_study/COVID19_lung_atlas_revision_v3.h5ad'

output_folder = 'diffxpy_out/'

de_output_base = 'COVID19_lung_atlas_revision_v3_lung_cov19_poissonglm_holdouts_nUMIoffset_testInts'
# -

# # Read the data

adata = sc.read(adata_diffxpy)

adata

adata.obs.age = adata.obs.age.astype(float)

adata.obs.dtypes

adata.obs['dataset'] = adata.obs['last_author/PI']

adata.obs.dataset.value_counts()

# # Filter the data

# Keep only datsets with:
# - more than 1 donor
# - non-fetal
# - lung

# Remove fetal datasets
dats_to_remove = set(['Rawlins', 'Spence', 'Linnarsson'])

# +
dat = adata.obs.groupby(['donor']).agg({'sex':'first', 'age':'first', 'dataset':'first'})

# Single donor filter
don_tab = dat['dataset'].value_counts()
dats_to_remove.update(set(don_tab.index[don_tab == 1]))
# -

dats_to_remove = list(dats_to_remove)
dats_to_remove

adata = adata[~adata.obs.dataset.isin(dats_to_remove)].copy()

adata.obs.lung_vs_nasal.value_counts()

# Filter for only lung data
adata = adata[adata.obs.lung_vs_nasal.isin(['lung']),].copy()

adata

# Rename smoking status covariate
adata.obs['smoking_status'] = adata.obs.smoked_boolean

adata.obs.dataset.value_counts()
adata.obs['sample'].nunique()
adata.obs['donor'].nunique()

# # Check the data

np.mean(adata.X.astype(int) != adata.X)

# Check if any non-integer data in a particular dataset
for dat in adata.obs.dataset.unique():
    val = np.mean(adata[adata.obs.dataset.isin([dat]),:].X.astype(int) != adata[adata.obs.dataset.isin([dat]),:].X)
    if val != 0:
        print(f'dataset= {dat}; value= {val}')
        adata[adata.obs.dataset.isin([dat]),:].X[:20,:20].A

# All counts are integers

# # Fit models and perform DE

cluster_key = 'ann_level_2'
clust_tbl = adata.obs[cluster_key].value_counts()
clusters = clust_tbl.index[clust_tbl > 1000]
ct_to_rm = clusters[[ct.startswith('1') for ct in clusters]]
clusters = clusters.drop(ct_to_rm.tolist()).tolist()
clusters

# Calculate DE genes per cluster.

adata

adata.obs['total_counts_scaled'] = adata.obs['total_counts']/adata.obs['total_counts'].mean()

# Get interquartile range for ages to test
adata.obs.groupby(['donor']).agg({'age':'first'}).age.quantile([0.25,0.5,0.75])

# +
formula = "1 + sex + age + sex:age + dataset"
tested_coef = ["sex[T.male]", "age"]
dmat = de.utils.design_matrix(
    data=adata,
    formula="~" + formula,
    as_numeric=["age"],
    return_type="patsy"
)

to_test = dict()
to_test['age'] = [32,62]
to_test['sex[T.male]'] = [0,1]

dmat[1]


# -

# ### Function definition to test effect sizes at particular covariate values

def calc_effects(dmat, cov_mat, params, effect, coefs):
    from patsy.design_info import DesignMatrix
    from diffxpy.api.stats import wald_test_chisq
    
    dmat_cond = isinstance(dmat, tuple) and isinstance(dmat[0], DesignMatrix)

    if not dmat_cond:
        raise ValueError("`dmat` should be a patsy output Design Matrix.")

    effect_list = ['sex[T.male]', 'age', 'smoking_status[T.True]']
    if not effect in effect_list:
        raise ValueError(f'{effect} is not one of: '
                         f'{effect_list}')

    if not isinstance(coefs, dict):
        raise TypeError('`coefs` should contain a dictionary of coefficients '
                        'where the effects should be evaluated.')
    
    ## Note: this is only correct when 3 covariates are tested in combinations
    #if np.sum([coef in coefs for coef in effect_list]) < 2:
    #    raise ValueError('The `coefs` dict must contain values for the two '
    #                     'coefficient not tested in:'
    #                     f'{effect_list}')
        
    if 'smoking_status[T.True]' in coefs and coefs['smoking_status[T.True]'] not in [0,1]:
        raise ValueError('Smoking status should be encoded as 0 or 1.')

    if 'sex[T.male]' in coefs and coefs['sex[T.male]'] not in [0,1]:
        raise ValueError('Sex should be encoded as 0 or 1.')

    if 'age' in coefs and not (isinstance(coefs['age'], float) or isinstance(coefs['age'], int)):
        raise ValueError('Age should be a numerical value.')

    coef_list = []

    for term in dmat[1]:
        if effect not in term:
            coef_list.append(0)
        elif term == effect:
            coef_list.append(1)
        else:
            t_list = term.split(':')
            t_list.remove(effect)
            coef_list.append(coefs[t_list[0]])

    C = np.array(coef_list)
    val = np.matmul(C,np.array(params))
    stderr = np.sqrt(np.matmul(np.matmul(C.T,cov_mat),C))
    pval = wald_test_chisq(np.array([val]).reshape(1,1), np.array([stderr**2]).reshape(1,1,1))[0]
    return (val, stderr, pval)


# ## Poisson GLM

# +
# Poisson GLM loop
de_results_lvl2_glm = dict()

# Test over clusters
for clust in clusters:
    res_list = []
    
    adata_tmp = adata[adata.obs[cluster_key] == clust,:]

    hold_outs = np.unique(adata_tmp.obs["dataset"].values)
    for ho in hold_outs:
        
        adata_tmp_ho = adata_tmp[~adata_tmp.obs.dataset.isin([ho]),:].copy()    

        print(f'Holdout {ho} in cluster {clust}:')

        print(pd.crosstab(adata_tmp_ho.obs['smoking_status'], adata_tmp_ho.obs['sex']))

        # Filter out genes to reduce multiple testing burden
        sc.pp.filter_genes(adata_tmp_ho, min_cells=10)
        if adata_tmp_ho.n_vars == 0:
            print('No genes expressed in more than 10 cells!')
            continue
        if len(adata_tmp_ho.obs.sex.value_counts())==1:
            print(f'{clust} only has 1 type of male/female sample.')
            continue
        
        print(f'Testing {adata_tmp_ho.n_vars} genes...')
        print("")
        
        # List to store results
        de_results_list = []        
        
        # Set up design matrix
        dmat = de.utils.design_matrix(
            data=adata_tmp_ho, #[idx_train],
            formula="~" + formula,
            as_numeric=["age"],
            return_type="patsy"
        )
        
        # Test if model is full rank
        if np.linalg.matrix_rank(np.asarray(dmat[0])) < np.min(dmat[0].shape):
            print(f'Cannot test {clust} as design matrix is not full rank.')
            continue
        
        for i, gene in enumerate(adata_tmp_ho.var_names):
            # Specify model
            pois_model = sm.GLM(
                endog=adata_tmp_ho.X[:, i].todense(), #[idx_train, :], 
                exog=dmat[0], 
                offset=np.log(adata_tmp_ho.obs['total_counts_scaled'].values),
                family=sm.families.Poisson()
            )
            
            # Fit the model
            pois_results = pois_model.fit()

            # Get the covariance matrix
            cov_mat = pois_results.cov_params()

            # Test over coefs
            for coef in tested_coef:
                iter_coefs = tested_coef.copy()
                iter_coefs.remove(coef)

                for c1 in to_test[iter_coefs[0]]:
                    coef_vals = {iter_coefs[0]:c1}
                    val, stderr, pval = calc_effects(
                        dmat = dmat,
                        cov_mat = cov_mat,
                        params = pois_results.params,
                        effect = coef,
                        coefs = coef_vals)

                    case = iter_coefs[0]+':'+str(c1)
                    case = case.replace('sex[T.male]:0','F').replace('sex[T.male]:1','M')
                    case = case.replace('age:32','32yr').replace('age:62','62yr')
                    case = case.replace('_',' ')                

                    # Output the results nicely
                    de_results_temp = pd.DataFrame({
                        "gene": gene,
                        "cell_identity": clust,
                        "covariate": coef,
                        "eval_at": case,
                        "holdout": ho,
                        "coef": val,
                        "coef_sd": stderr,                 
                        "pval": pval
                    }, index= [clust+"_"+gene+"_"+coef])
                    
                    de_results_list.append(de_results_temp)
            
        de_results = pd.concat(de_results_list)
        de_results['adj_pvals'] = multipletests(de_results['pval'].tolist(), method='fdr_bh')[1]
        res_list.append(de_results)

    # Store the results
    if len(res_list) > 0:
        de_results_lvl2_glm[clust] = pd.concat(res_list, ignore_index=True)
    
# Join the dataframes:
full_res_lvl2_glm = pd.concat([de_results_lvl2_glm[i] for i in de_results_lvl2_glm.keys()], ignore_index=True)
# -

# ## Inspect some results

de_results_lvl2_glm.keys()

full_res_lvl2_glm = full_res_lvl2_glm.sort_values(by=['gene', 'cell_identity', 'covariate'])
full_res_lvl2_glm

full_res_lvl2_glm.loc[(full_res_lvl2_glm['gene'] == 'ACE2') & (full_res_lvl2_glm['adj_pvals'] < 0.05),]


# ### Aggregate hold-out results statistics

# +
def prop_signif(series):
    return (series < 0.05).mean()

def prop_pos(series):
    return (series > 0).mean()

def prop_pos_zero(series):
    return (series >= 0).mean()

def prop_neg_zero(series):
    return (series <= 0).mean()



# +
res_summary_lvl2 = full_res_lvl2_glm.groupby(['gene', 'cell_identity', 'covariate', 'eval_at']).agg({
    'adj_pvals':prop_signif,
    'coef':['mean', 'std', prop_pos],
    'holdout':'count'
}).reset_index()

res_summary_lvl2
# -



# # Level 3 annotation

cluster_key = 'ann_level_3'
clust_tbl = adata.obs[cluster_key].value_counts()
clusters = clust_tbl.index[clust_tbl > 1000]
ct_to_rm = clusters[[ct.startswith('1') or ct.startswith('2') for ct in clusters]]
clusters = clusters.drop(ct_to_rm.tolist()).tolist()
clusters

# +
adata_sub = adata[adata.obs.ann_level_3.isin(clusters),:]

adata_sub
adata_sub.obs.donor.nunique()
adata_sub.obs['sample'].nunique()
# -

# ## Poisson GLM

# +
# Poisson GLM loop
de_results_lvl3_glm = dict()

# Test over clusters
for clust in clusters:
    res_list = []
    
    adata_tmp = adata_sub[adata_sub.obs[cluster_key] == clust,:]

    hold_outs = np.unique(adata_tmp.obs["dataset"].values)
    for ho in hold_outs:
        
        adata_tmp_ho = adata_tmp[~adata_tmp.obs.dataset.isin([ho]),:].copy()    

        print(f'Holdout {ho} in cluster {clust}:')

        print(pd.crosstab(adata_tmp_ho.obs['smoking_status'], adata_tmp_ho.obs['sex']))

        # Filter out genes to reduce multiple testing burden
        sc.pp.filter_genes(adata_tmp_ho, min_cells=10)
        if adata_tmp_ho.n_vars == 0:
            print('No genes expressed in more than 10 cells!')
            continue
        if len(adata_tmp_ho.obs.sex.value_counts())==1:
            print(f'{clust} only has 1 type of male/female sample.')
            continue
        
        print(f'Testing {adata_tmp_ho.n_vars} genes...')
        print("")
        
        # List to store results
        de_results_list = []
        
        # Set up design matrix
        dmat = de.utils.design_matrix(
            data=adata_tmp_ho,
            formula="~" + formula,
            as_numeric=["age"],
            return_type="patsy"
        )
        
        # Test if model is full rank
        if np.linalg.matrix_rank(np.asarray(dmat[0])) < np.min(dmat[0].shape):
            print(f'Cannot test {clust} as design matrix is not full rank.')
            continue
        
        for i, gene in enumerate(adata_tmp_ho.var_names):
            # Specify model
            pois_model = sm.GLM(
                endog=adata_tmp_ho.X[:, i].todense(),
                exog=dmat[0],
                offset=np.log(adata_tmp_ho.obs['total_counts_scaled'].values),
                family=sm.families.Poisson()
            )
            
            # Fit the model
            pois_results = pois_model.fit()
            
            # Get the covariance matrix
            cov_mat = pois_results.cov_params()

            # Test over coefs
            for coef in tested_coef:
                iter_coefs = tested_coef.copy()
                iter_coefs.remove(coef)

                for c1 in to_test[iter_coefs[0]]:
                    coef_vals = {iter_coefs[0]:c1}
                    val, stderr, pval = calc_effects(
                        dmat = dmat,
                        cov_mat = cov_mat,
                        params = pois_results.params,
                        effect = coef,
                        coefs = coef_vals)

                    case = iter_coefs[0]+':'+str(c1)
                    case = case.replace('sex[T.male]:0','F').replace('sex[T.male]:1','M')
                    case = case.replace('age:32','32yr').replace('age:62','62yr')
                    case = case.replace('_',' ')
                    
                    # Output the results nicely
                    de_results_temp = pd.DataFrame({
                        "gene": gene,
                        "cell_identity": clust,
                        "covariate": coef,
                        "eval_at": case,
                        "holdout": ho,                            
                        "coef": val,
                        "coef_sd": stderr,                 
                        "pval": pval
                    }, index= [clust+"_"+gene+"_"+coef])
                    
                    de_results_list.append(de_results_temp)
            
        de_results = pd.concat(de_results_list)
        de_results['adj_pvals'] = multipletests(de_results['pval'].tolist(), method='fdr_bh')[1]
        res_list.append(de_results)

    # Store the results
    if len(res_list) > 0:
        de_results_lvl3_glm[clust] = pd.concat(res_list, ignore_index=True)
    
# Join the dataframes:
full_res_lvl3_glm = pd.concat([de_results_lvl3_glm[i] for i in de_results_lvl3_glm.keys()], ignore_index=True)
# -

# ## Inspect some results

de_results_lvl3_glm.keys()

full_res_lvl3_glm = full_res_lvl3_glm.sort_values(by=['gene', 'cell_identity', 'covariate'])


full_res_lvl3_glm.loc[full_res_lvl3_glm['gene'] == 'ACE2']

full_res_lvl3_glm.loc[full_res_lvl3_glm['gene'] == 'TMPRSS2']

full_res_lvl3_glm.loc[full_res_lvl3_glm['gene'] == 'CTSL']

full_res_lvl3_glm.loc[(full_res_lvl3_glm['gene'] == 'ACE2') & (full_res_lvl3_glm['adj_pvals'] < 0.05),]

# ### Aggregate hold-out results statistics

# +
res_summary_lvl3 = full_res_lvl3_glm.groupby(['gene', 'cell_identity', 'covariate', 'eval_at']).agg({
    'adj_pvals':prop_signif,
    'coef':['mean', 'std', prop_pos],
    'holdout':'count'
}).reset_index()

res_summary_lvl3

# +
prop_agreement = (res_summary_lvl3[('coef','prop_pos')] >= 0.8) | (res_summary_lvl3[('coef','prop_pos')] <= 0.2)
gene_mask = (res_summary_lvl3['gene'] == 'ACE2')
signif = (res_summary_lvl3[('adj_pvals', 'prop_signif')] >= 0.5)

res_summary_lvl3.loc[(prop_agreement & gene_mask)]
res_summary_lvl3.loc[(prop_agreement & gene_mask & signif)]

# +
prop_agreement = (res_summary_lvl3[('coef','prop_pos')] >= 0.8) | (res_summary_lvl3[('coef','prop_pos')] <= 0.2)
gene_mask = (res_summary_lvl3['gene'] == 'TMPRSS2')
signif = (res_summary_lvl3[('adj_pvals', 'prop_signif')] >= 0.5)

res_summary_lvl3.loc[(prop_agreement & gene_mask)]
res_summary_lvl3.loc[(prop_agreement & gene_mask & signif)]

# +
prop_agreement = (res_summary_lvl3[('coef','prop_pos')] >= 0.8) | (res_summary_lvl3[('coef','prop_pos')] <= 0.2)
gene_mask = (res_summary_lvl3['gene'] == 'CTSL')
signif = (res_summary_lvl3[('adj_pvals', 'prop_signif')] >= 0.5)

res_summary_lvl3.loc[(prop_agreement & gene_mask)]
res_summary_lvl3.loc[(prop_agreement & gene_mask & signif)]

# +
# Find number of disagreeing holdout datasets
h_count = res_summary_lvl3[('holdout','count')]
prop_pos = res_summary_lvl3[('coef','prop_pos')]

dat_diff_pos = h_count - prop_pos*h_count
dat_diff_neg = prop_pos*h_count
dat_diff = pd.concat([dat_diff_pos, dat_diff_neg], axis=1).min(axis=1).astype(int)
res_summary_lvl3['holdout_dataset_dis'] = dat_diff
#dat_diff = np.min(dat_diff_pos, dat_diff_neg)
# -

res_summary_lvl3.columns

# +
gene_mask = (res_summary_lvl3['gene'] == 'ACE2')
ct_mask = (res_summary_lvl3['cell_identity'].isin(['AT2', 'Basal', 'Multiciliated lineage', 'Secretory']))

ace2_holdout_res = res_summary_lvl3.loc[(gene_mask & ct_mask)][[('holdout_dataset_dis',''), ('covariate',''), ('cell_identity',''), ('eval_at','')]]
ace2_holdout_res['cov_eval'] = [' '.join([i1, i2]) for i1,i2 in zip(ace2_holdout_res[('covariate','')], ace2_holdout_res[('eval_at','')])]
ace2_holdout_res = ace2_holdout_res.pivot(index='cov_eval', columns='cell_identity')
ace2_holdout_res = ace2_holdout_res.drop(columns=['eval_at', 'covariate'])

ace2_holdout_res.columns = ace2_holdout_res.columns.get_level_values(2)
ace2_holdout_res.index = [item.replace('sex[T.male]', 'Sex').replace('smoking_status[T.True]', 'Smoking status') for item in ace2_holdout_res.index.tolist()]

rcParams['figure.figsize'] = (6,6)
p1 = sns.heatmap(ace2_holdout_res, cbar=False, cmap='Blues', annot=True, linewidths=.5)
plt.tick_params(axis='both', which='major', labelsize=10, labelbottom = False, bottom=False, top = False, labeltop=True)
p1.set_yticklabels(ace2_holdout_res.index, rotation=0)
plt.ylabel('')
plt.xlabel('')
plt.savefig(folder+'/'+output_folder+de_output_base+'_annlvl3_ace2_dataset_disagreements.pdf', dpi=300, bbox_inches='tight')
plt.show()
rcParams['figure.figsize'] = (8,8)

# +
gene_mask = (res_summary_lvl3['gene'] == 'TMPRSS2')
ct_mask = (res_summary_lvl3['cell_identity'].isin(['AT2', 'Multiciliated lineage']))

tmprss2_holdout_res = res_summary_lvl3.loc[(gene_mask & ct_mask)][[('holdout_dataset_dis',''), ('covariate',''), ('cell_identity',''), ('eval_at','')]]
tmprss2_holdout_res['cov_eval'] = [' '.join([i1, i2]) for i1,i2 in zip(tmprss2_holdout_res[('covariate','')], tmprss2_holdout_res[('eval_at','')])]
tmprss2_holdout_res = tmprss2_holdout_res.pivot(index='cov_eval', columns='cell_identity')
tmprss2_holdout_res = tmprss2_holdout_res.drop(columns=['eval_at', 'covariate'])

tmprss2_holdout_res.columns = tmprss2_holdout_res.columns.get_level_values(2)
tmprss2_holdout_res.index = [item.replace('sex[T.male]', 'Sex').replace('smoking_status[T.True]', 'Smoking status') for item in tmprss2_holdout_res.index.tolist()]

rcParams['figure.figsize'] = (3,6)
p1 = sns.heatmap(tmprss2_holdout_res, cbar=False, cmap='Blues', annot=True, linewidths=.5)
plt.tick_params(axis='both', which='major', labelsize=10, labelbottom = False, bottom=False, top = False, labeltop=True)
p1.set_yticklabels(tmprss2_holdout_res.index, rotation=0)
plt.ylabel('')
plt.xlabel('')
plt.savefig(folder+'/'+output_folder+de_output_base+'_annlvl3_tmprss2_dataset_disagreements.pdf', dpi=300, bbox_inches='tight')
plt.show()
rcParams['figure.figsize'] = (8,8)

# +
gene_mask = (res_summary_lvl3['gene'] == 'CTSL')
ct_mask = (res_summary_lvl3['cell_identity'].isin(['AT2', 'Multiciliated lineage']))

ctsl_holdout_res = res_summary_lvl3.loc[(gene_mask & ct_mask)][[('holdout_dataset_dis',''), ('covariate',''), ('cell_identity',''), ('eval_at','')]]
ctsl_holdout_res['cov_eval'] = [' '.join([i1, i2]) for i1,i2 in zip(ctsl_holdout_res[('covariate','')], ctsl_holdout_res[('eval_at','')])]
ctsl_holdout_res = ctsl_holdout_res.pivot(index='cov_eval', columns='cell_identity')
ctsl_holdout_res = ctsl_holdout_res.drop(columns=['eval_at', 'covariate'])

ctsl_holdout_res.columns = ctsl_holdout_res.columns.get_level_values(2)
ctsl_holdout_res.index = [item.replace('sex[T.male]', 'Sex').replace('smoking_status[T.True]', 'Smoking status') for item in ctsl_holdout_res.index.tolist()]

rcParams['figure.figsize'] = (3,6)
p1 = sns.heatmap(ctsl_holdout_res, cbar=False, cmap='Blues', annot=True, linewidths=.5)
plt.tick_params(axis='both', which='major', labelsize=10, labelbottom = False, bottom=False, top = False, labeltop=True)
p1.set_yticklabels(ctsl_holdout_res.index, rotation=0)
plt.ylabel('')
plt.xlabel('')
plt.savefig(folder+'/'+output_folder+de_output_base+'_annlvl3_ctsl_dataset_disagreements.pdf', dpi=300, bbox_inches='tight')
plt.show()
rcParams['figure.figsize'] = (8,8)
# -

# # Store results

res_summary_lvl2.columns = ['_'.join(col).strip('_') for col in res_summary_lvl2.columns.values]
res_summary_lvl3.columns = ['_'.join(col).strip('_') for col in res_summary_lvl3.columns.values]

res_summary_lvl2.to_csv(folder+'/'+output_folder+de_output_base+'_lvl2_summary.csv')
full_res_lvl2_glm.to_csv(folder+'/'+output_folder+de_output_base+'_lvl2_full.csv')

res_summary_lvl3.to_csv(folder+'/'+output_folder+de_output_base+'_lvl3_summary.csv')
full_res_lvl3_glm.to_csv(folder+'/'+output_folder+de_output_base+'_lvl3_full.csv')

StarcoderdataJupyterScriptsDedupFiltered