<reponame>opentargets/genetics-finemapping
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# <NAME>
#

import utils as fm_utils
import gcta as fm_gcta
import os
import numpy as np
import pandas as pd
from scipy.stats import norm
import sys

def run_credible_set_for_locus(
            index_info,
            sumstats,
            top_loci,
            in_plink,
            temp_dir,
            fm_wind,
            cojo_window,
            cojo_collinear,
            pp_threshold,
            method='conditional',
            logger=None):
    ''' Run credible set analysis at a given locus (speficied by index_info)
    Args:
        index_info (dict): index locus information
        sumstats (pd.df): full summary stats
        top_loci (pd.df): table of top loci (for conditional analysis)
        method ([conditional|distance|none]): whether to perform conditional analysis
            or extract a set distance
    '''

    if logger:
        logger.info(
        '\n- Running credible set analysis for {0}'.format(index_info['variant_id']))

    temp_dir = os.path.join(temp_dir, 'credible_set')

    # Perform conditional analysis
    sumstat_cond = None
    if method == 'conditional':
        # Extract `cojo_window` region surrounding the index variant
        sumstat_wind = fm_utils.extract_window(
            sumstats, index_info['chrom'], index_info['pos'], cojo_window)
        if logger:
            logger.info('  {0} variants in {1}kb cojo window around index'
                        ' variant'.format(sumstat_wind.shape[0], cojo_window))

        # Get list of variants to condition on
        cond_list = make_list_to_condition_on(
            index_info['variant_id'],
            sumstat_wind.variant_id,
            top_loci.variant_id
        )

        if logger:
            logger.info('  conditioning on {0} variants'.format(
            len(cond_list)))

        # Only do conditional if there are variants to condition on
        if len(cond_list) > 0:
            sumstat_cond = fm_gcta.perform_conditional_adjustment(
                sumstat_wind,
                in_plink,
                temp_dir,
                index_info['variant_id'],
                index_info['chrom'],
                cond_list,
                cojo_window,
                cojo_collinear,
                logger=logger
            )
    else:
        sumstat_wind = sumstats
        if logger:
            logger.info('  not conditioning as method != conditional')

    # If conditional analysis was not performed, we need to make the sumstat df
    # look the same, add (beta_cond, se_cond, pval_cond) columns
    if sumstat_cond is None:
        sumstat_cond = sumstat_wind
        sumstat_cond['beta_cond'] = sumstat_cond['beta']
        sumstat_cond['se_cond'] = sumstat_cond['se']
        sumstat_cond['pval_cond'] = sumstat_cond['pval']

    # Extract `fm_wind` region surrounding the index variant
    # TODO: test results compared to Ed's code before, which didn't use sumstat_cond here
    sumstat_wind = fm_utils.extract_window(
        sumstat_cond, index_info['chrom'], index_info['pos'], fm_wind)
    if logger:
        logger.info('  {0} variants in {1}kb fine-mapping window around index'
                    ' variant'.format(sumstat_wind.shape[0], fm_wind))

    # Do credible set analysis
    if sumstat_wind.shape[0] > 0:

        if logger:
            logger.info('  calculating credible sets...')
        cred_sets = calc_credible_sets(sumstat_wind, pp_threshold=pp_threshold)

        if logger:
            logger.info('  found {0} in 95% and {1} in 99% cred sets'.format(
            cred_sets.is95_credset.sum(), cred_sets.is99_credset.sum()
            ))
            logger.info('  kept {0} vars with PP > {1}'.format(
            cred_sets.shape[0], pp_threshold
            ))

        # Script will fail if cred_sets is empty
        if cred_sets.shape[0] > 0:

            # Add index variant columns
            cred_sets.loc[:, 'lead_variant_id'] = index_info['variant_id']
            cred_sets[['lead_chrom', 'lead_pos', 'lead_ref', 'lead_alt']] = \
                cred_sets.lead_variant_id.str.split(':', expand=True)

            # Add column specifying method used
            cred_sets.loc[:, 'multisignal_method'] = method

            # Format output table
            cred_sets = format_credset_output(cred_sets)
        
        # Else None if credible set results is empty
        else:
            cred_sets = None

    # If df is empty skip analysis
    else:
        if logger:
            logger.warning('  skipping credible set analysis')
        cred_sets = None

    return cred_sets

def format_credset_output(cred_sets):
    ''' Formats the cred_sets table for output
    Args:
        cred_sets (pd.df)
    Returns:
        pd.df
    '''
    cols = fm_utils.get_credset_out_columns()
    meta = fm_utils.get_meta_info(type='cred_set')
    df = (
        cred_sets.loc[:, cols.keys()]
                 .rename(columns=cols)
                 .astype(dtype=meta)
    )
    return df

def calc_credible_sets(data, pp_threshold):
    ''' Calculates credible sets from provided sumstats
    Args:
        data (pd.df): sumstats to perform analysis on
        pp_threshold (float): returns any variant in ( (95% OR 99% threshold) OR pp > pp_threshold )
    Returns
        pd.df of results
    '''
    # Need to fix pC == 0.0 as this will produce inf logABF. Set to sys.float_info.min
    if (data.pval_cond == 0.0).any():
        print("Warning: some pval_cond == 0.0 in {0}\n - setting to sys.float_info.min")
        data.pval_cond[data.pval_cond == 0.0] = sys.float_info.min

    # Calculate case proportions
    data['case_prop'] = data['n_cases'] / data['n_total']

    # Calc ABFs
    # print(calc_abf(0.808621, 0.17690, 290365, 0.6203537)) # Should return -3.311501
    data["logABF"] = data.apply(
        lambda row: calc_abf(pval=row['pval_cond'],
                             maf=freq_to_maf(row['eaf']),
                             n=row['n_total'],
                             prop_cases=row['case_prop'] if row['is_cc'] else None
        ), axis=1)
    data = data.sort_values("logABF", ascending=False)

    # Calculate posterior probability for each SNP
    sum_lABF = log_sum(data["logABF"])
    data["postprob"] = (data["logABF"] - sum_lABF).apply(np.exp)

    # Calc cumulative sum of the posterior probabilities
    data["postprob_cumsum"] = data["postprob"].cumsum()

    # Find 99% and 95% credible sets - this is horrible
    set_idx = data["postprob_cumsum"].gt(0.95).tolist().index(True)
    data["is95_credset"] = [1] * (set_idx + 1) + [0] * (data.shape[0] - (set_idx + 1))
    data["is95_credset"] = data["is95_credset"].map({1:True, 0:False})
    set_idx = data["postprob_cumsum"].gt(0.99).tolist().index(True)
    data["is99_credset"] = [1] * (set_idx + 1) + [0] * (data.shape[0] - (set_idx + 1))
    data["is99_credset"] = data["is99_credset"].map({1:True, 0:False})

    # Only keep rows that are in the 95 or 99% credible sets
    to_keep = ((data["is95_credset"] | data["is99_credset"])
               & (data["postprob"] > pp_threshold) )
    cred_set_res = data.loc[to_keep, :]

    return cred_set_res

def make_list_to_condition_on(index_var, all_vars, top_vars):
    ''' Makes a list of variants on which to condition on
    Args:
        index_var (str): index variant at this locus
        all_vars (pd.Series): A series of all variant IDs in the locus window
        top_vars (pd.Series): A series of top loci variant IDs
    Returns:
        list of variants to condition on
    '''
    window_top_vars = set(all_vars).intersection(set(top_vars))
    cond_list = list(window_top_vars - set([index_var]))
    return cond_list

def calc_abf(pval, maf, n, prop_cases=None):
    """ Caluclate Approximate Bayes Factor (Wakefield, 2009, Genet Epidemiol.).
        Based on code from coloc: https://github.com/chr1swallace/coloc
    Args:
        pval (float): GWAS p-value
        maf (float): Minor allele freq
        n (int): Sample size
        prop_cases (float or None): number of cases, if left blank will assume
            quantitative trait
    Returns:
        natural log(ABF)
    """
    # Assert/set types
    pval = float(pval)
    maf = float(maf)
    n = int(n)
    prop_cases = float(prop_cases) if prop_cases else None

    # Estimate variance for quant trait
    if prop_cases is None:
        sd_prior = 0.15
        v = var_data(maf, n)
    # Estimate var for cc study
    else:
        sd_prior = 0.2
        v = var_data_cc(maf, n, prop_cases)

    # Calculate Z-score
    z = np.absolute(norm.ppf(pval / 2))

    # Calc shrinkage factor: ratio of the prior variance to the total variance
    r = sd_prior**2 / (sd_prior**2 + v)

    # Approximate BF - ln scale to compare in log natural scale with LR diff
    lABF = 0.5 * (np.log(1 - r) + (r * z**2))

    return lABF

def log_sum(l):
    """ Calculates the log of the sum of the exponentiated logs taking out the
        max, i.e. insuring that the sum is not Inf
    Args:
        l (pandas Series)
    Returns:
        Sum of exponentiated logs
    """
    l_max = l.max()
    l_logsum = l_max + np.log(np.sum(np.exp(l - l_max)))
    return l_logsum

def freq_to_maf(freq):
    """ Convert allele frequency to minor allele freq
    """
    return min(freq, 1-freq)

def var_data(maf, n):
    """ Calc variance of MLE of beta for quantitative trait, assuming var(y)=0
    """
    var = 1 / (2 * n * maf * (1 - maf))
    return var

def var_data_cc(maf, n, prop_cases):
    """ Calc variance of MLE of beta for case-control
    """
    var = 1 / (2 * n * maf * (1 - maf) * prop_cases * (1 - prop_cases))
    return var

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# <NAME>
#

import utils as fm_utils
import gcta as fm_gcta
import os
import numpy as np
import pandas as pd
from scipy.stats import norm
import sys

def run_credible_set_for_locus(
            index_info,
            sumstats,
            top_loci,
            in_plink,
            temp_dir,
            fm_wind,
            cojo_window,
            cojo_collinear,
            pp_threshold,
            method='conditional',
            logger=None):
    ''' Run credible set analysis at a given locus (speficied by index_info)
    Args:
        index_info (dict): index locus information
        sumstats (pd.df): full summary stats
        top_loci (pd.df): table of top loci (for conditional analysis)
        method ([conditional|distance|none]): whether to perform conditional analysis
            or extract a set distance
    '''

    if logger:
        logger.info(
        '\n- Running credible set analysis for {0}'.format(index_info['variant_id']))

    temp_dir = os.path.join(temp_dir, 'credible_set')

    # Perform conditional analysis
    sumstat_cond = None
    if method == 'conditional':
        # Extract `cojo_window` region surrounding the index variant
        sumstat_wind = fm_utils.extract_window(
            sumstats, index_info['chrom'], index_info['pos'], cojo_window)
        if logger:
            logger.info('  {0} variants in {1}kb cojo window around index'
                        ' variant'.format(sumstat_wind.shape[0], cojo_window))

        # Get list of variants to condition on
        cond_list = make_list_to_condition_on(
            index_info['variant_id'],
            sumstat_wind.variant_id,
            top_loci.variant_id
        )

        if logger:
            logger.info('  conditioning on {0} variants'.format(
            len(cond_list)))

        # Only do conditional if there are variants to condition on
        if len(cond_list) > 0:
            sumstat_cond = fm_gcta.perform_conditional_adjustment(
                sumstat_wind,
                in_plink,
                temp_dir,
                index_info['variant_id'],
                index_info['chrom'],
                cond_list,
                cojo_window,
                cojo_collinear,
                logger=logger
            )
    else:
        sumstat_wind = sumstats
        if logger:
            logger.info('  not conditioning as method != conditional')

    # If conditional analysis was not performed, we need to make the sumstat df
    # look the same, add (beta_cond, se_cond, pval_cond) columns
    if sumstat_cond is None:
        sumstat_cond = sumstat_wind
        sumstat_cond['beta_cond'] = sumstat_cond['beta']
        sumstat_cond['se_cond'] = sumstat_cond['se']
        sumstat_cond['pval_cond'] = sumstat_cond['pval']

    # Extract `fm_wind` region surrounding the index variant
    # TODO: test results compared to Ed's code before, which didn't use sumstat_cond here
    sumstat_wind = fm_utils.extract_window(
        sumstat_cond, index_info['chrom'], index_info['pos'], fm_wind)
    if logger:
        logger.info('  {0} variants in {1}kb fine-mapping window around index'
                    ' variant'.format(sumstat_wind.shape[0], fm_wind))

    # Do credible set analysis
    if sumstat_wind.shape[0] > 0:

        if logger:
            logger.info('  calculating credible sets...')
        cred_sets = calc_credible_sets(sumstat_wind, pp_threshold=pp_threshold)

        if logger:
            logger.info('  found {0} in 95% and {1} in 99% cred sets'.format(
            cred_sets.is95_credset.sum(), cred_sets.is99_credset.sum()
            ))
            logger.info('  kept {0} vars with PP > {1}'.format(
            cred_sets.shape[0], pp_threshold
            ))

        # Script will fail if cred_sets is empty
        if cred_sets.shape[0] > 0:

            # Add index variant columns
            cred_sets.loc[:, 'lead_variant_id'] = index_info['variant_id']
            cred_sets[['lead_chrom', 'lead_pos', 'lead_ref', 'lead_alt']] = \
                cred_sets.lead_variant_id.str.split(':', expand=True)

            # Add column specifying method used
            cred_sets.loc[:, 'multisignal_method'] = method

            # Format output table
            cred_sets = format_credset_output(cred_sets)
        
        # Else None if credible set results is empty
        else:
            cred_sets = None

    # If df is empty skip analysis
    else:
        if logger:
            logger.warning('  skipping credible set analysis')
        cred_sets = None

    return cred_sets

def format_credset_output(cred_sets):
    ''' Formats the cred_sets table for output
    Args:
        cred_sets (pd.df)
    Returns:
        pd.df
    '''
    cols = fm_utils.get_credset_out_columns()
    meta = fm_utils.get_meta_info(type='cred_set')
    df = (
        cred_sets.loc[:, cols.keys()]
                 .rename(columns=cols)
                 .astype(dtype=meta)
    )
    return df

def calc_credible_sets(data, pp_threshold):
    ''' Calculates credible sets from provided sumstats
    Args:
        data (pd.df): sumstats to perform analysis on
        pp_threshold (float): returns any variant in ( (95% OR 99% threshold) OR pp > pp_threshold )
    Returns
        pd.df of results
    '''
    # Need to fix pC == 0.0 as this will produce inf logABF. Set to sys.float_info.min
    if (data.pval_cond == 0.0).any():
        print("Warning: some pval_cond == 0.0 in {0}\n - setting to sys.float_info.min")
        data.pval_cond[data.pval_cond == 0.0] = sys.float_info.min

    # Calculate case proportions
    data['case_prop'] = data['n_cases'] / data['n_total']

    # Calc ABFs
    # print(calc_abf(0.808621, 0.17690, 290365, 0.6203537)) # Should return -3.311501
    data["logABF"] = data.apply(
        lambda row: calc_abf(pval=row['pval_cond'],
                             maf=freq_to_maf(row['eaf']),
                             n=row['n_total'],
                             prop_cases=row['case_prop'] if row['is_cc'] else None
        ), axis=1)
    data = data.sort_values("logABF", ascending=False)

    # Calculate posterior probability for each SNP
    sum_lABF = log_sum(data["logABF"])
    data["postprob"] = (data["logABF"] - sum_lABF).apply(np.exp)

    # Calc cumulative sum of the posterior probabilities
    data["postprob_cumsum"] = data["postprob"].cumsum()

    # Find 99% and 95% credible sets - this is horrible
    set_idx = data["postprob_cumsum"].gt(0.95).tolist().index(True)
    data["is95_credset"] = [1] * (set_idx + 1) + [0] * (data.shape[0] - (set_idx + 1))
    data["is95_credset"] = data["is95_credset"].map({1:True, 0:False})
    set_idx = data["postprob_cumsum"].gt(0.99).tolist().index(True)
    data["is99_credset"] = [1] * (set_idx + 1) + [0] * (data.shape[0] - (set_idx + 1))
    data["is99_credset"] = data["is99_credset"].map({1:True, 0:False})

    # Only keep rows that are in the 95 or 99% credible sets
    to_keep = ((data["is95_credset"] | data["is99_credset"])
               & (data["postprob"] > pp_threshold) )
    cred_set_res = data.loc[to_keep, :]

    return cred_set_res

def make_list_to_condition_on(index_var, all_vars, top_vars):
    ''' Makes a list of variants on which to condition on
    Args:
        index_var (str): index variant at this locus
        all_vars (pd.Series): A series of all variant IDs in the locus window
        top_vars (pd.Series): A series of top loci variant IDs
    Returns:
        list of variants to condition on
    '''
    window_top_vars = set(all_vars).intersection(set(top_vars))
    cond_list = list(window_top_vars - set([index_var]))
    return cond_list

def calc_abf(pval, maf, n, prop_cases=None):
    """ Caluclate Approximate Bayes Factor (Wakefield, 2009, Genet Epidemiol.).
        Based on code from coloc: https://github.com/chr1swallace/coloc
    Args:
        pval (float): GWAS p-value
        maf (float): Minor allele freq
        n (int): Sample size
        prop_cases (float or None): number of cases, if left blank will assume
            quantitative trait
    Returns:
        natural log(ABF)
    """
    # Assert/set types
    pval = float(pval)
    maf = float(maf)
    n = int(n)
    prop_cases = float(prop_cases) if prop_cases else None

    # Estimate variance for quant trait
    if prop_cases is None:
        sd_prior = 0.15
        v = var_data(maf, n)
    # Estimate var for cc study
    else:
        sd_prior = 0.2
        v = var_data_cc(maf, n, prop_cases)

    # Calculate Z-score
    z = np.absolute(norm.ppf(pval / 2))

    # Calc shrinkage factor: ratio of the prior variance to the total variance
    r = sd_prior**2 / (sd_prior**2 + v)

    # Approximate BF - ln scale to compare in log natural scale with LR diff
    lABF = 0.5 * (np.log(1 - r) + (r * z**2))

    return lABF

def log_sum(l):
    """ Calculates the log of the sum of the exponentiated logs taking out the
        max, i.e. insuring that the sum is not Inf
    Args:
        l (pandas Series)
    Returns:
        Sum of exponentiated logs
    """
    l_max = l.max()
    l_logsum = l_max + np.log(np.sum(np.exp(l - l_max)))
    return l_logsum

def freq_to_maf(freq):
    """ Convert allele frequency to minor allele freq
    """
    return min(freq, 1-freq)

def var_data(maf, n):
    """ Calc variance of MLE of beta for quantitative trait, assuming var(y)=0
    """
    var = 1 / (2 * n * maf * (1 - maf))
    return var

def var_data_cc(maf, n, prop_cases):
    """ Calc variance of MLE of beta for case-control
    """
    var = 1 / (2 * n * maf * (1 - maf) * prop_cases * (1 - prop_cases))
    return var
#!/usr/bin/env python # -*- coding: utf-8 -*- # # <NAME> # Run credible set analysis at a given locus (speficied by index_info)     Args:         index_info (dict): index locus information         sumstats (pd.df): full summary stats         top_loci (pd.df): table of top loci (for conditional analysis)         method ([conditional|distance|none]): whether to perform conditional analysis             or extract a set distance # Perform conditional analysis # Extract `cojo_window` region surrounding the index variant # Get list of variants to condition on # Only do conditional if there are variants to condition on # If conditional analysis was not performed, we need to make the sumstat df # look the same, add (beta_cond, se_cond, pval_cond) columns # Extract `fm_wind` region surrounding the index variant # TODO: test results compared to Ed's code before, which didn't use sumstat_cond here # Do credible set analysis # Script will fail if cred_sets is empty # Add index variant columns # Add column specifying method used # Format output table # Else None if credible set results is empty # If df is empty skip analysis Formats the cred_sets table for output     Args:         cred_sets (pd.df)     Returns:         pd.df Calculates credible sets from provided sumstats     Args:         data (pd.df): sumstats to perform analysis on         pp_threshold (float): returns any variant in ( (95% OR 99% threshold) OR pp > pp_threshold )     Returns         pd.df of results # Need to fix pC == 0.0 as this will produce inf logABF. Set to sys.float_info.min # Calculate case proportions # Calc ABFs # print(calc_abf(0.808621, 0.17690, 290365, 0.6203537)) # Should return -3.311501 # Calculate posterior probability for each SNP # Calc cumulative sum of the posterior probabilities # Find 99% and 95% credible sets - this is horrible # Only keep rows that are in the 95 or 99% credible sets Makes a list of variants on which to condition on     Args:         index_var (str): index variant at this locus         all_vars (pd.Series): A series of all variant IDs in the locus window         top_vars (pd.Series): A series of top loci variant IDs     Returns:         list of variants to condition on Caluclate Approximate Bayes Factor (Wakefield, 2009, Genet Epidemiol.).         Based on code from coloc: https://github.com/chr1swallace/coloc     Args:         pval (float): GWAS p-value         maf (float): Minor allele freq         n (int): Sample size         prop_cases (float or None): number of cases, if left blank will assume             quantitative trait     Returns:         natural log(ABF) # Assert/set types # Estimate variance for quant trait # Estimate var for cc study # Calculate Z-score # Calc shrinkage factor: ratio of the prior variance to the total variance # Approximate BF - ln scale to compare in log natural scale with LR diff Calculates the log of the sum of the exponentiated logs taking out the         max, i.e. insuring that the sum is not Inf     Args:         l (pandas Series)     Returns:         Sum of exponentiated logs Convert allele frequency to minor allele freq Calc variance of MLE of beta for quantitative trait, assuming var(y)=0 Calc variance of MLE of beta for case-control