windows.py

"""
implement the imputation window is sliding along the genome:

- ImpG like: Non overlapping windows, the imputation is apply in batch to unknown snp in the window
- centered_window:  A sliding window centered on the Snp to impute
"""

from .stat_models import impg_model
from .ld_matrix import generate_sparse_matrix

import pandas as pd
import numpy as np

def parse_region_position(ld_file):
    """
    Retrieve the region definition from a ld-file generated by impute_jass
    Argument :
        ld_file : A ld file generated by jass_impute
    """
    (chrom, startpos, endpos ) = ld_file.split("/")[-1].split(".")[0].split('_')
    return (chrom, startpos, endpos)

def realigned_zfiles_on_panel(ref_panel, zscore):
    """
    Check if the counted allele is the same in the reference panel and
    the Zscore files.

    If not, the coded and other allele are inverted and the zscore sign
    is inverted also.
    """
    sub_ref_panel = ref_panel.loc[zscore.index]
    allele_inverted = (sub_ref_panel['Ref_all'] != zscore.A0)

    zscore.loc[allele_inverted, "A0"] = sub_ref_panel.loc[allele_inverted].Ref_all
    zscore.loc[allele_inverted, "A1"] = sub_ref_panel.loc[allele_inverted].alt_all
    zscore.loc[allele_inverted, "Z"] = - zscore.loc[allele_inverted, "Z"]

    return zscore

def prepare_zscore_for_imputation(ref_panel, zscore):
    """
    Prepare the known Z score by realigning them on the reference ref_panel
    the snps that are not present in the ref panel are filtered
    
    """
    zscore = realigned_zfiles_on_panel(ref_panel, zscore)
    zscore['Var'] = -1
    zscore['Nsnp_to_impute'] = -1
    zscore['ld_score'] = -1
    zscore = zscore.loc[zscore.index.intersection(ref_panel.index)]
    return zscore

def in_region(pos_vector, start, end):
    return ((start < pos_vector) & (pos_vector < end))


def compute_window_and_size(start_ld_block, end_ld_block, window_size):
    """
    Compute the number of window to pave the ld_block tightly
    """

    Nwindows = ((int(end_ld_block)) - (int(start_ld_block)))//window_size
    # adapt window size to cover the LD block
    if Nwindows > 0:
        window_resize = np.ceil((int(end_ld_block) - (int(start_ld_block)))/Nwindows)
    else:
        Nwindows = 1
        window_resize = np.ceil(int(end_ld_block) - (int(start_ld_block)))

    return Nwindows, window_resize

def format_result_df(imp, ref_panel, known, unknowns):

    result_dict = {
                'pos': ref_panel.loc[unknowns, 'pos'],
                'A0': ref_panel.loc[unknowns, "Ref_all"],
                "A1": ref_panel.loc[unknowns,"alt_all"],
                "Z" : imp['mu'],
                "Var": imp["var"],
                "ld_score" : imp["ld_score"],
                "condition_number": imp['condition_number'],
                "correct_inversion":imp["correct_inversion"],
                "Nsnp_to_impute" : len(known)
            }
    column_order = ['pos','A0',"A1","Z","Var", "ld_score", "condition_number", "correct_inversion", "Nsnp_to_impute"]
    batch_df = pd.DataFrame(result_dict, columns = column_order)
    return batch_df


def print_progression(i, Nwindows):
    """
    print the percentage of window treated in the ld block
    """
    if i%(np.ceil(Nwindows/10)) == 0:
        print("{0}\%".format(np.round(i/Nwindows,3)))

def impg_like_imputation(ld_file, ref_panel, zscore, window_size, buffer, lamb, rcond, unknowns=pd.Series([])):
    """
        Each missing Snp is imputed by known snp found in a window centered on the SNP to impute
        Argument
    """
    (chrom, start_ld_block, end_ld_block) = parse_region_position(ld_file)

    LD_mat = generate_sparse_matrix(ld_file, ref_panel)

    Nwindows, window_resize = compute_window_and_size(start_ld_block, end_ld_block, window_size)

    all_unknowns = ref_panel.loc[ref_panel.index.difference(zscore.index)]

    zscore = prepare_zscore_for_imputation(ref_panel, zscore)
    zscore_results = zscore.copy(deep=True)

    print("### Imputation of {0} snps ###".format(all_unknowns.shape[0]))

    for i in range(Nwindows):
        # Boundary of the sliding_window
        start_windows = int(start_ld_block) + i*window_resize - buffer
        end_windows = int(start_ld_block) + (i+1)*window_resize + buffer

        start_pos = max(start_windows, float(start_ld_block))
        end_pos = min(end_windows, float(end_ld_block))

        in_LD_reg_n_window =  in_region(zscore.pos, start_pos, end_pos)
        unknown_in_LD_reg_n_window =  in_region(all_unknowns.pos, start_pos, end_pos)

        known = zscore.loc[in_LD_reg_n_window].index
        unknowns = all_unknowns.loc[unknown_in_LD_reg_n_window].index

        sig_t = LD_mat.loc[known, known]
        sig_i_t = LD_mat.loc[unknowns, known]
        zt = zscore.loc[known,'Z']

        if(len(known) > 0):
            imp = impg_model(zt, sig_t, sig_i_t, lamb=lamb, rcond=rcond, batch=True)
            batch_df = format_result_df(imp, ref_panel, known, unknowns)

            # keep only snp in the core window (and not in the buffer)
            start_core_window = int(start_ld_block) + i*window_resize
            end_core_window = int(start_ld_block) + (i+1)*window_resize
            in_core_window = in_region(batch_df.pos, start_core_window, end_core_window)

            # keep only SNP with non negligible explained variance
            snp_well_predicted = (batch_df.Var < 0.9)
            batch_df_filt = batch_df.loc[in_core_window & snp_well_predicted, zscore_results.columns]
            zscore_results = pd.concat([zscore_results, batch_df_filt])
        i = i+1
        print_progression(i, Nwindows)

    return zscore_results.sort_values(by="pos")


def ld_region_centered_window_imputation(ld_file, ref_panel, zscore, window_size, unknowns=pd.Series([])):
    """
        Each missing Snp is imputed by known snp found in a window centered on the SNP to impute
        Argument
    """
    (chrom, start_ld_block, end_ld_block) = parse_region_position(ld_file)

    LD_mat = generate_sparse_matrix(ld_file, ref_panel)
    zscore = prepare_zscore_for_imputation(ref_panel, zscore)

    # Find Snp to impute
    if len(unknowns) == 0:
        unknowns = LD_mat.index.difference(zscore.index)

    N_snp = len(unknowns)
    print("### Imputation of {0} snps ###".format(len(unknowns)))

    for i,snp_unknown in enumerate(unknowns):
        # Boundary of the centered_window
        start_pos = max((ref_panel.loc[snp_unknown,'pos'] - window_size), float(start_ld_block))
        end_pos = min(ref_panel.loc[snp_unknown,'pos'] + window_size, float(end_ld_block))

        in_LD_reg_n_window =  in_region(zscore.pos, start_pos, end_pos)

        known = zscore.loc[in_LD_reg_n_window].index
        sig_t = LD_mat.loc[known, known]
        sig_i_t = LD_mat.loc[snp_unknown, known]
        zt = zscore.loc[known,'Z']

        if(len(known) > 0):
            imp = impg_model(zt, sig_t, sig_i_t, batch=False)
            zscore.loc[snp_unknown] = [ref_panel.loc[snp_unknown, 'pos'], ref_panel.loc[snp_unknown, "Ref_all"],  ref_panel.loc[snp_unknown, "alt_all"], imp['mu'], imp['var'], len(known)]

        if i%300 == 0:
            print("{0}\%".format(np.round(i/N_snp,4)))

    return zscore.sort_values(by="pos")