optimize merging of the reference panel with the GWAS

jass_preprocessing/jass_preprocessing/map_reference/map_reference.py

@@ -20,19 +20,21 @@ def map_on_ref_panel(gw_df , ref_panel):
     Make sure that the same SNPs are in the reference panel and the gwas
 
     """
-    def key2(x):
-        return ('chr' + str(x["chr"]) + ":" + str(x["pos"]))
-
-    ref_panel['key2'] = ref_panel.apply(key2,1)
-
     inter_index = ref_panel.index.intersection(gw_df.index)
-    print("SNps {}".format(len(inter_index)))
-    merge_GWAS = pd.merge(ref_panel.loc[inter_index], gw_df.loc[inter_index], how='inner', indicator=True, left_index=True, right_index=True)
 
-    inter_index = gw_df.index.intersection(ref_panel.index)
-    other_snp = pd.merge(ref_panel.loc[inter_index], gw_df.loc[inter_index], how='inner', indicator=True, left_on ='key2', right_index=True)
+    merge_GWAS = pd.merge(ref_panel, gw_df,
+                          how='inner', indicator=True, left_index=True, right_index=True)
 
+    print("SNPs {}".format(merge_GWAS.shape[0]))
+
+    # If the GWAS file contains indexes of the kind chr*:position, perform a
+    # second join
+    if gw_df.index.map(str).str.contains("^chr*", case=False).any():
+        ref_panel['key2'] = "chr"+ref_panel.chr.map(str) + ":" +ref_panel.pos.map(str)
+        other_snp = pd.merge(ref_panel, gw_df, how='inner', indicator=True,
+                             left_on ='key2', right_index=True)
         merge_GWAS.loc[other_snp.index] = other_snp
+
     return(merge_GWAS)