imputation_launcher.py

"""
Function set to launch imputation on a complete chromosome or
on the genome
"""
import glob
import pandas as pd
from .windows import ld_region_centered_window_imputation, impg_like_imputation, realigned_zfiles_on_panel


def chromosome_imputation_entry_point(chrom, zscore, ref_panel, ld_folder, window_size=10000, imputation_style="online", buf=2500, lamb= 0.01, pinv_rcond = 0.01):
    """
    Function whose only purpose is to allow the calling of the ImputationLauncher.chromosome_imputation method
    from an entry point
    """
    imputer = ImputationLauncher( window_size=10000, imputation_style="online", buf=2500, lamb= 0.01, pinv_rcond = 0.01)
    imputer.chromosome_imputation(chrom, zscore, ref_panel, ld_folder)




class ImputationLauncher(object):
    """
    Class perform imputation of snp from summary statistic

    """
    def __init__(self, window_size=10000, imputation_style="online", buf=2500, lamb= 0.01, pinv_rcond = 0.01 ):
        """

        Args:
            window_size (int): size of the imputation window in bp
            imputation_style (str): define if the windows while span the genome in a non overlapping fashion ("batch") or
                                    by being centered on each snp to impute ('online')
            buffer (int): the size of the padding around the windows of imputation (relevant only for batch imputation)
            lamb (float): size of the increment added to snp correlation matrices to make it less singular
            pinv_rcond (float): the rcond scipy.linalg.pinv function argument. The scipy.linalg.pinv is used to invert
             the correlationmatrices

        """
        self.imputation_style = imputation_style
        self.window_size = window_size
        self.buffer = buf
        self.lamb = lamb
        self.rcond = pinv_rcond

    def chromosome_imputation(self, chrom, zscore, ref_panel, ld_folder):
        """
        Impute the panel zscore score for one chromosome and with the specified parameters

        Args:
            chrom : str specifying chromosome
            zscore : known zscore
            ref_panel : location of the folder of reference chromosome
            ld_folder: location of linkage desiquilibrium matrices
        Returns
            Imputed zscore dataframe
        """
        pattern = "{0}/{1}*.ld".format(ld_folder, chrom)
        if self.imputation_style == "online":
            def imputer(ld_file):
                return ld_region_centered_window_imputation(ld_file, ref_panel, zscore, self.window_size)
        elif self.imputation_style == "batch":
            def imputer(ld_file):
                return impg_like_imputation(ld_file, ref_panel, zscore, self.window_size, self.buffer, self.lamb, self.rcond)

        for ld_file in glob.glob(pattern):
            print("processing Region: {0}".format(ld_file))
            zscore = imputer(ld_file)

        zscore = realigned_zfiles_on_panel(ref_panel, zscore)
        return zscore


    def genome_imputation(self, gwas_tag, ref_panel_folder, ld_folder, zscore_folder, folder_output):
        """
        Launch imputation on all chromosome

        """

        for i in range(1, 23):

            ref_panel_file = ref_panel_folder + "/chr" + str(i) + ".eur.1pct.bim"
            ref_panel = pd.read_csv(ref_panel_file, sep="\t", names=['chr', "nothing", 'pos', 'Ref_all', 'alt_all'], index_col = 1)


            known_zscore_file = zscore_folder + "/z_" + gwas_tag + "_chr" + str(i) + ".txt"
            known_zscore = pd.read_csv(known_zscore_file, index_col=0, sep="\t")

            chrom = "chr"+str(i)
            z_imp = self.chromosome_imputation(chrom, known_zscore, ref_panel, ld_folder)

            imputed_zscore = folder_output + "/z_" + gwas_tag + "_chr" + str(i) + ".txt"
            z_imp.to_csv(imputed_zscore, sep="\t")