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, prepare_zscore_for_imputation, impg_like_imputation, realigned_zfiles_on_panel



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

    """
    def __init__(self, window_size=10000, imputation_style="batch", buf=2500,
                 lamb= 0.01, pinv_rcond = 0.01 ):
        """
        Initialise the imputation object. Fix the windows size, the buffer size
         and the king of imputation employed
        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 correlation matrices
        """
        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)
        zscore = prepare_zscore_for_imputation(ref_panel, zscore)
        zscore_results = zscore.copy(deep=True)
        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))
            ld_batch = imputer(ld_file)
            zscore_results = pd.concat([zscore_results, ld_batch])

        zscore_results.sort_values(by="pos")
        zscore_results = realigned_zfiles_on_panel(ref_panel, zscore_results)
        return zscore_results


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

        """

        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")