add option to the centered windows function

impute_jass/impute_jass/windows.py

@@ -54,7 +54,7 @@ def prepare_Zscore_for_imputation(ref_panel, Zscores):
 def in_region(pos_vector, start, end):
     return ((start < pos_vector) & (pos_vector < end))
 
-def centered_window_imputation(LD_file, ref_panel_folder, Zfile, window_size):
+def centered_window_imputation(LD_file, ref_panel_folder, Zscores, window_size, unknowns=pd.Series([])):
     """
         Each missing Snp is imputed by known snp found in a window centered on the SNP to impute
         Argument
@@ -66,11 +66,13 @@ def centered_window_imputation(LD_file, ref_panel_folder, Zfile, window_size):
 
     LD_mat = generate_sparse_matrix(LD_file, ref_panel)
 
-    Zscores = pd.read_csv(Zfile, index_col=0, sep="\t")
+    #Zscores = pd.read_csv(Zfile, index_col=0, sep="\t")
     Zscores = prepare_Zscore_for_imputation(ref_panel, Zscores)
 
     # Find Snp to impute
-    unknowns = LD_mat.index.difference(Zscores.index)
+    if len(unknowns) == 0:
+        unknowns = LD_mat.index.difference(Zscores.index)
+
     N_snp = len(unknowns)
     print("### Imputation of {0} snps ###".format(len(unknowns)))
     i = 0
@@ -85,20 +87,20 @@ def centered_window_imputation(LD_file, ref_panel_folder, Zfile, window_size):
         in_LD_reg_n_window = in_region(Zscores.pos, int(start_ld_block), int(end_ld_block)) & in_region(Zscores.pos, start_pos, end_pos)
 
         known = Zscores.loc[in_LD_reg_n_window].index
-
         Sig_t = LD_mat.loc[known, known]
         Sig_i_t = LD_mat.loc[snp_unknown, known]
         Zt = Zscores.loc[known,'Z']
 
         if(len(known) > 0):
             imp = ImpG_model_snp(Zt, Sig_t, Sig_i_t)
+            Zscores.loc[snp_unknown, ["pos", "A0", "A1"]] = ref_panel.loc[snp_unknown, ['pos', "Ref_all", "alt_all"]]
             Zscores.loc[snp_unknown, "Z"] = imp['mu']
             Zscores.loc[snp_unknown, "Var"] = imp['Var']
-            Zscores[snp_unknown, 'Nsnp_to_impute'] = len(known)
+            Zscores.loc[snp_unknown, 'Nsnp_to_impute'] = len(known)
         i = i+1
 
-        if i%10 == 0:
+        if i%100 == 0:
             print("{0}\%".format(np.round(i/N_snp,4)))
 
 
-    return Zt.sort_values(by="pos")
+    return Zscores.sort_values(by="pos")