gensim_v2.py

import sys, os, pickle, gc, random

import numpy as np
import scanpy as sc
import pandas as pd

sys.path.append('/gstore/home/shih43/PROJECT/SCEPS/scripts/software') 

from src.utils import *
from src.scdata import *
from src.regression import *
from sim_utils_v2 import *

from tqdm import tqdm

import warnings
warnings.filterwarnings("ignore")

def main():
   
    args = get_command_line_sim()

    """
    Set up simulation directory
    """
    out_str = get_sim_str(args)
    out_dir = '{}/{}'.format(args.out_dir, out_str)
    if os.path.exists(out_dir) == False:
        os.mkdir(out_dir)
    out_prefix = '{}/sim'.format(out_dir)

    # set the seed
    logging.info('Using seed {}'.format(args.seed))
    np.random.seed(args.seed)
    random.seed(args.seed)

    """
    Load all the data
    """

    # load the single-cell data
    adata = sc.read_h5ad(args.adata)
    logging.info('Loaded {} cells for {} samples'.format(adata.shape[0],
        pd.unique(adata.obs[args.donor_id_col]).shape[0]))

    # intersect single-cell data genes with gwas genes
    gene_list = pd.read_table(args.gene_list, delim_whitespace=True)
    gene_list.columns = ['ID', 'SCORE']
    adata = adata[:,adata.var[args.gene_id_col].isin(gene_list['ID'])]
    gene_list = gene_list.sort_values(by=['SCORE'], ascending=False).reset_index(drop=True)
    logging.info('{} genes after intersecting gene list'.format(adata.shape[1]))

    # collect garbbage
    gc.collect()

    """
    Find the optimal number of random walk steps
    """
    # get the adjacency matrix
    adj_mat = get_connectivity(adata)
    obs = adata.obs.copy()

    # calculate the transition matrix
    trans_mat = get_transition_matrix(adj_mat)

    # calculate the optimal number of steps
    opt_nsteps, nam_matrix = choose_optimal_nsteps(args, trans_mat, obs)
    logging.info('Determined optimal step size is {}'.format(opt_nsteps))

    """
    Simulate the phenotype
    """
    all_raw_sim_params = []
    for i in tqdm(range(args.num_sim)):

        # generate simulations
        while True:
            sim_out = gen_sim(args, gene_list, trans_mat, nam_matrix, opt_nsteps, adata)
            if sim_out is not None:
                break
        
        # save simulated params
        sim_param_i = dict()
        sim_param_i['ngwas_genes'] = [sim_out['ngwas_genes']]
        sim_param_i['nctrl_genes'] = [sim_out['nctrl_genes']]
        sim_param_i['nrest_genes'] = [sim_out['nrest_genes']]
        sim_param_i['vareff_gwas'] = [sim_out['vareff_gwas']]
        sim_param_i['vareff_ctrl'] = [sim_out['vareff_ctrl']]
        sim_param_i['vareff_rest'] = [sim_out['vareff_rest']]
        sim_param_i['vareff_overall'] = [sim_out['vareff_overall']]
        sim_param_i['mean_var_gwas'] = [sim_out['mean_var_gwas']]
        sim_param_i['mean_var_ctrl'] = [sim_out['mean_var_ctrl']]
        sim_param_i['mean_var_rest'] = [sim_out['mean_var_rest']]
        sim_param_i['mean_var_overall'] = [sim_out['mean_var_overall']]
        sim_param_i['scale_factor_sq'] = [sim_out['scale_factor_sq']]
        all_raw_sim_params.append(pd.DataFrame(sim_param_i))

        # save simulated phenotype
        out_fnm = '{}.{}.txt.gz'.format(out_prefix, i+args.start_idx+1)
        df_outcome = sim_out['df_outcome']
        df_outcome.to_csv(out_fnm, sep='\t', index=False)
        
        # save causal neighborhood
        out_fnm = '{}.{}.causal_neighborhood.txt.gz'.format(out_prefix, i+args.start_idx+1)
        df_causal_nb = sim_out['df_causal_nb']
        df_causal_nb.to_csv(out_fnm, sep='\t', index=False)

        # save the causal cell
        out_fnm = '{}.{}.causal_cell.txt.gz'.format(out_prefix, i+args.start_idx+1)
        df_causal_cell = sim_out['df_causal_cell']
        df_causal_cell.to_csv(out_fnm, sep='\t', index=False)

        # save the gwas genes
        out_fnm = '{}.{}.gwas_genes.txt.gz'.format(out_prefix, i+args.start_idx+1)
        df_gwas_genes = sim_out['gwas_genes']
        df_gwas_genes.to_csv(out_fnm, sep='\t', index=False, header=False)

        # save the control genes
        out_fnm = '{}.{}.ctrl_genes.txt.gz'.format(out_prefix, i+args.start_idx+1)
        df_ctrl_genes = sim_out['ctrl_genes']
        df_ctrl_genes.to_csv(out_fnm, sep='\t', index=False, header=False)

    # save raw simulated parameters
    all_raw_sim_params = pd.concat(all_raw_sim_params, ignore_index=True)
    out_fnm = '{}.{}.raw_simulated_params.txt.gz'.format(out_prefix, args.start_idx)
    all_raw_sim_params.to_csv(out_fnm, sep='\t', index=False)

    # calculate simulated parameters
    mean_scale_factor_sq = np.mean(all_raw_sim_params['scale_factor_sq'])
    sigma_gwas = np.mean(all_raw_sim_params['vareff_gwas']) / mean_scale_factor_sq
    sigma_ctrl = np.mean(all_raw_sim_params['vareff_ctrl']) / mean_scale_factor_sq
    sigma_rest = np.mean(all_raw_sim_params['vareff_rest']) / mean_scale_factor_sq
    sigma_overall = np.mean(all_raw_sim_params['vareff_overall']) / mean_scale_factor_sq
    diff = sigma_gwas - sigma_ctrl

    weighted_sigma_gwas = np.mean(all_raw_sim_params['vareff_gwas'] * all_raw_sim_params['mean_var_gwas']) / mean_scale_factor_sq
    weighted_sigma_ctrl = np.mean(all_raw_sim_params['vareff_ctrl'] * all_raw_sim_params['mean_var_ctrl']) / mean_scale_factor_sq
    weighted_sigma_rest = np.mean(all_raw_sim_params['vareff_rest'] * all_raw_sim_params['mean_var_rest']) / mean_scale_factor_sq
    ngwas_genes = np.mean(all_raw_sim_params['ngwas_genes'])
    nctrl_genes = np.mean(all_raw_sim_params['nctrl_genes'])
    nrest_genes = np.mean(all_raw_sim_params['nrest_genes'])
    ntot_genes = ngwas_genes + nctrl_genes + nrest_genes
    weighted_sigma_overall = (weighted_sigma_gwas*ngwas_genes + weighted_sigma_ctrl*nctrl_genes + weighted_sigma_rest*nrest_genes) / ntot_genes

    weighted_diff = weighted_sigma_gwas - weighted_sigma_ctrl
    
    sim_params = {'SIGMA_GWAS': [sigma_gwas], \
                  'SIGMA_CONTROL': [sigma_ctrl], \
                  'SIGMA_REST': [sigma_rest], \
                  'SIGMA_OVERALL': [sigma_overall], \
                  'DIFF': [diff], \
                  'WEIGHTED_SIGMA_GWAS': [weighted_sigma_gwas], \
                  'WEIGHTED_SIGMA_CONTROL': [weighted_sigma_ctrl], \
                  'WEIGHTED_SIGMA_REST': [weighted_sigma_rest], \
                  'WEIGHTED_SIGMA_OVERALL': [weighted_sigma_overall], \
                  'WEIGHTED_DIFF': [weighted_diff],
                  'NSTEPS': [opt_nsteps]}
    
    sim_params = pd.DataFrame(sim_params)
    sim_params = sim_params[['SIGMA_GWAS', 'SIGMA_CONTROL', 'SIGMA_REST', 'SIGMA_OVERALL', 'DIFF', \
                              'WEIGHTED_SIGMA_GWAS', 'WEIGHTED_SIGMA_CONTROL', 'WEIGHTED_SIGMA_REST', 
                              'WEIGHTED_SIGMA_OVERALL', 'WEIGHTED_DIFF', 'NSTEPS']]
    sim_params = sim_params.transpose()
    out_fnm = '{}.simulated_params.txt'.format(out_prefix)
    sim_params.to_csv(out_fnm, sep='\t', header=False)


def gen_sim(args, gene_list, trans_mat, nam_matrix, nstep, adata):

    # all donors
    all_donor = pd.unique(adata.obs[args.donor_id_col])

    # choose the causal cell neighborhood
    ncells = trans_mat.shape[0]
    causal_cell_idx = np.random.choice(np.array(range(ncells)))
    df_causal_cell = pd.DataFrame({args.cell_id_col: [adata.obs.iloc[causal_cell_idx][args.cell_id_col]]})
    donor_id_vec = adata.obs[args.donor_id_col].values

    # get cells in the neighborhood
    logging.info('Using step size: {}'.format(nstep))
    prob_i = prob_i_reach_all_j(trans_mat, causal_cell_idx, nstep)
    if args.neighborhood_definition_method == 'nam':
        nam_matrix_row = nam_matrix.iloc[causal_cell_idx][:]
        nam_neighborhood = get_neighborhood_cells_nam(args, prob_i, donor_id_vec,
            None, nam_matrix_row, check_outcome_in_neighborhood=False)
        if nam_neighborhood is None:
            logging.info('Neighborhood definition failed')
            return None
        neighborhood_cell_idx = nam_neighborhood[0]
    elif args.neighborhood_definition_method == 'soft threshold':
        prob_thres = nstep / ncells
        neighborhood_cell_idx = get_neighborhood_cells(prob_i, prob_thres=prob_thres)
    elif args.neighborhood_definition_method == 'hard threshold':
        neighborhood_cell_idx = get_neighborhood_cells(prob_i, prob_thres=args.prob_thres)
    else:
        logging.info('Neighborhood definition method does not exit')
        return None

    neighborhood_cells = adata.obs.iloc[neighborhood_cell_idx][args.cell_id_col]
    adata_neighbor = adata[adata.obs[args.cell_id_col].isin(neighborhood_cells), :].copy()
    logging.info('Neighborhood size {}'.format(neighborhood_cell_idx.shape[0]))

    # save causal neighborhood
    df_causal_nb = pd.DataFrame(neighborhood_cells)

    # get pseudobulk
    adata_pseudobulk = get_pseudobulk(args, adata_neighbor)
    logging.info('Number of donors {} for {} genes'.format(adata_pseudobulk.shape[0], \
                adata_pseudobulk.shape[1]))

    # prepare expression
    adata_pseudobulk = prep_pseudobulk(args, adata_pseudobulk)

    # simulate the outcome for donors in the neighborhood
    sim_out = sim_outcome(args, adata_pseudobulk, gene_list, all_donor)
    sim_out['df_causal_nb'] = df_causal_nb
    sim_out['df_causal_cell'] = df_causal_cell
    sim_out['causal_cell_idx'] = causal_cell_idx

    return sim_out


def sim_y_expr(args, adata, prop_cau, genes, varexp):

    ngenes = genes.shape[0]
    nsample = adata.shape[0]
    ncau_genes = int(np.floor(prop_cau * ngenes))
    adata_genes = adata[:,adata.var[args.gene_id_col].isin(genes)].copy()
    sum_var_genes = np.sum(adata_genes.X.var(axis=0))
    vareff = varexp/sum_var_genes

    if varexp == 0:
        return np.zeros(nsample), np.zeros(ngenes)

    while True:

        causal_genes = genes.sample(ncau_genes)

        causal_gene_idx = np.where(adata_genes.var[args.gene_id_col].isin(causal_genes))[0]
        X_causal = adata_genes.X[:, causal_gene_idx]
        vareff_causal = varexp / np.sum(X_causal.var(axis=0))

        eff_cau = np.random.normal(scale=np.sqrt(vareff_causal), size=ncau_genes)
        eff = np.zeros(ngenes)
        eff[causal_gene_idx] = eff_cau
        
        y_expr = np.dot(X_causal, eff_cau)

        err_varexp = abs(np.var(y_expr) - varexp)
        err_vareff = abs(np.var(eff) - vareff)
        if (err_varexp <= ERROR_TOL*varexp) and (err_vareff <= ERROR_TOL*vareff):
            break

    gc.collect()

    return y_expr, eff


def sim_outcome(args, adata, gene_list, all_donor):
    """
    Simulate the outcome from pseudo-bulk expression using 
    """
    
    # intersect genes
    gwas_genes = gene_list['ID'].iloc[0:args.num_gwas_genes]
    gwas_genes = gwas_genes[gwas_genes.isin(adata.var[args.gene_id_col])]
    ngwas_genes = gwas_genes.shape[0]

    ctrl_genes = sample_control_genes(args, adata.var.copy(), gwas_genes, nbins=10)
    nctrl_genes = ctrl_genes.shape[0]

    rest_genes = gene_list['ID'][gene_list['ID'].isin(gwas_genes)==False]
    rest_genes = rest_genes[rest_genes.isin(ctrl_genes)==False]
    rest_genes = rest_genes[rest_genes.isin(adata.var[args.gene_id_col])]
    nrest_genes = rest_genes.shape[0]

    nnongwas_genes = nctrl_genes + nrest_genes
    ntot_genes = ngwas_genes + nctrl_genes + nrest_genes

    # get expression of gwas and non gwas genes
    X_gwas = adata.X[:, np.where(adata.var[args.gene_id_col].isin(gwas_genes))[0]]
    X_ctrl = adata.X[:, np.where(adata.var[args.gene_id_col].isin(ctrl_genes))[0]]
    X_rest = adata.X[:, np.where(adata.var[args.gene_id_col].isin(rest_genes))[0]]

    # calculate summary information for the 
    mean_var_gwas = np.mean(X_gwas.var(axis=0))
    mean_var_ctrl = np.mean(X_ctrl.var(axis=0))
    mean_var_rest = np.mean(X_rest.var(axis=0))
    mean_var_gwas_ctrl = (mean_var_gwas + mean_var_ctrl) / 2.0

    sum_var_gwas = np.sum(X_gwas.var(axis=0))
    sum_var_ctrl = np.sum(X_ctrl.var(axis=0))
    sum_var_rest = np.sum(X_rest.var(axis=0))

    sum_var_all = sum_var_gwas + sum_var_ctrl + sum_var_rest
    mean_var_overall = sum_var_all / ntot_genes

    # derive simulation parameters
    prop_cau_gwas = args.prop_causal_gwas
    prop_cau_nongwas = args.prop_causal_non_gwas

    varexp_tot = args.varexp_total
    varexp_enrich_per_gwas = args.varexp_enrich_per_gwas_gene
    varexp_enrich_per_nongwas = (ntot_genes - varexp_enrich_per_gwas*ngwas_genes) / nnongwas_genes
    varexp_enrich_per_nongwas = max(varexp_enrich_per_nongwas, 0.0)

    vareff_overall = varexp_tot / sum_var_all
    vareff_gwas = varexp_enrich_per_gwas * vareff_overall / (mean_var_gwas_ctrl / mean_var_overall)
    vareff_gwas = min(vareff_gwas, varexp_tot/sum_var_gwas)
    vareff_ctrl = varexp_enrich_per_nongwas * vareff_overall / (mean_var_gwas_ctrl / mean_var_overall)
    vareff_ctrl = max(vareff_ctrl, 0.0)
    vareff_rest = varexp_enrich_per_nongwas * vareff_overall / (mean_var_rest / mean_var_overall)
    vareff_rest = max(vareff_rest, 0.0)

    varexp_gwas = vareff_gwas * sum_var_gwas
    varexp_ctrl = vareff_ctrl * sum_var_ctrl
    varexp_rest = vareff_rest * sum_var_rest

    # simulate component of disease due to expression of gwas genes
    y_gwas, eff_gwas = sim_y_expr(args, adata, prop_cau_gwas, gwas_genes, varexp_gwas)
    y_ctrl, eff_ctrl = sim_y_expr(args, adata, prop_cau_nongwas, ctrl_genes, varexp_ctrl)
    y_rest, eff_rest = sim_y_expr(args, adata, prop_cau_nongwas, rest_genes, varexp_rest)

    # scale to total variance explained
    y_expr = y_gwas + y_ctrl + y_rest
    std_y_expr = np.std(y_expr)
    scale_factor_sq = 1.0
    if varexp_tot > 0:
        y_expr = y_expr * np.sqrt(varexp_tot) / std_y_expr
        scale_factor =  std_y_expr / np.sqrt(varexp_tot)
        scale_factor_sq = scale_factor * scale_factor

    # simulate the random noise
    nbhood_donor = adata.obs[args.donor_id_col]
    ndonor = nbhood_donor.shape[0]
    var_env = 1.0 - varexp_tot
    while True:
        y_env = np.random.normal(size=ndonor)
        y_env = y_env / (np.std(y_env) + EPS) * np.sqrt(var_env)
        y = y_expr + y_env
        if abs(np.var(y)-1) < ERROR_TOL:
            break

    # construct the dataframe
    df_y = pd.DataFrame({args.donor_id_col: nbhood_donor, \
                         'IS_IN_CAUSAL_NEIGHBORHOOD': ['True']*y.shape[0], \
                         'Y': y, 'Y_gwas': y_gwas, \
                         'Y_ctrl': y_ctrl, 'Y_rest': y_rest})

    # simulate the outcome for donors not in the neighborhood
    other_donor = all_donor[~all_donor.isin(nbhood_donor)]
    num_other_donor = other_donor.shape[0]
    y_gwas, y_ctrl, y_rest = np.zeros(num_other_donor), np.zeros(num_other_donor), np.zeros(num_other_donor)
    y = np.random.normal(size=num_other_donor, scale=np.sqrt(var_env))
    df_other_donor = pd.DataFrame({args.donor_id_col: other_donor, \
                                  'IS_IN_CAUSAL_NEIGHBORHOOD': ['False'] * other_donor.shape[0], \
                                  'Y': y, 'Y_gwas': y_gwas, 'Y_ctrl': y_ctrl, \
                                  'Y_rest': y_rest})

    # concat the results
    df_outcome = pd.concat([df_y, df_other_donor], ignore_index=True)

    # create output dictionary
    out_dict = dict()
    out_dict['df_outcome'] = df_outcome
    out_dict['scale_factor_sq'] = scale_factor_sq
    out_dict['ngwas_genes'] = ngwas_genes
    out_dict['nctrl_genes'] = nctrl_genes
    out_dict['nrest_genes'] = nrest_genes
    out_dict['gwas_genes'] = gwas_genes
    out_dict['ctrl_genes'] = ctrl_genes
    out_dict['vareff_gwas'] = np.var(eff_gwas)
    out_dict['vareff_ctrl'] = np.var(eff_ctrl)
    out_dict['vareff_rest'] = np.var(eff_rest)
    out_dict['vareff_overall'] = np.var(np.concatenate((eff_gwas, eff_ctrl, eff_rest)))
    out_dict['mean_var_gwas'] = mean_var_gwas
    out_dict['mean_var_ctrl'] = mean_var_ctrl
    out_dict['mean_var_rest'] = mean_var_rest
    out_dict['mean_var_overall'] = mean_var_overall

    gc.collect()

    return out_dict


def sample_control_genes(args, adata_var_orig, gwas_genes, nbins=10):
    """
    Sample causal control genes matched for mean and variance
    """

    # create bins
    adata_var = adata_var_orig.copy()
    out_mean = pd.qcut(adata_var['mean'].values, nbins, duplicates='drop')
    adata_var.loc[:,'bin_mean'] = out_mean.astype(str)

    out_std = pd.qcut(adata_var['std'].values, nbins, duplicates='drop')
    adata_var.loc[:,'bin_std'] = out_std.astype(str)

    adata_var.loc[:,'bin'] = adata_var['bin_mean']

    # count occurance of top genes in each bin
    adata_var_top = adata_var[adata_var[args.gene_id_col].isin(gwas_genes)].copy()
    adata_var.loc[:,'weight'] = 0.0
    top_cnt = adata_var_top['bin'].value_counts()
    for idx, cnt in enumerate(top_cnt):
        adata_var.loc[adata_var['bin']==top_cnt.index[idx],'weight'] = cnt

    # remove gwas genes gene from adata_var
    adata_var = adata_var[~adata_var[args.gene_id_col].isin(gwas_genes)]
    adata_var = adata_var.reset_index(drop=True)

    # sample causal control genes
    nctrl_genes = gwas_genes.shape[0]
    controls = adata_var.sample(n=nctrl_genes, weights=adata_var['weight'])[args.gene_id_col]

    return controls


def prep_pseudobulk(args, adata):
    """
    Prepare the pseudo-bulk data
    """

    mean_X = np.mean(adata.X, axis=0)
    std_X = np.std(adata.X, axis=0)
    adata.X = (adata.X - mean_X) #/ (std_X + EPS)
    adata.var['GENE_INDEX'] = np.array(range(adata.var.shape[0])).astype(int)
    adata.var = adata.var.set_index(args.gene_id_col, drop=False)

    return adata



if __name__ == "__main__":
    main()