gwaslab 3.5.7__py3-none-any.whl → 3.6.0__py3-none-any.whl

gwaslab/util_ex_infer_ancestry.py

@@ -0,0 +1,65 @@
+
+import pandas as pd
+from gwaslab.g_Log import Log
+
+def _infer_ancestry(sumstats, 
+                    ancestry_af=None,
+                    build="19",
+                    log=Log(),
+                    verbose=True):
+    log.write("Start to infer ancestry based on Fst...", verbose=verbose)
+    ref_af = pd.read_csv(ancestry_af, sep="\t")
+    
+    data_af = pd.merge(sumstats[["CHR","POS","EA","NEA","EAF"]] ,ref_af,on=["CHR","POS"],how="inner") 
+
+    log.write(f"  -Estimating Fst using {len(data_af)} variants...", verbose=verbose)
+
+    is_filp = data_af["EA"] == data_af["ALT"]
+    data_af.loc[is_filp, ["EA","NEA"]] = data_af.loc[is_filp, ["NEA","EA"]]
+    data_af.loc[is_filp, "EAF"] = 1 - data_af.loc[is_filp, "EAF"]
+    
+    headers = []
+    for i in ['GBR', 'FIN', 'CHS', 'PUR', 'CDX',
+        'CLM', 'IBS', 'PEL', 'PJL', 'KHV', 'ACB', 'GWD', 'ESN', 'BEB', 'MSL',
+        'STU', 'ITU', 'CEU', 'YRI', 'CHB', 'JPT', 'LWK', 'ASW', 'MXL', 'TSI',
+        'GIH', 'EUR', 'EAS', 'AMR', 'SAS', 'AFR']:
+        headers.append(f"FST_{i}")
+        data_af[f"FST_{i}"] = data_af.apply(lambda x: calculate_fst(x["EAF"], x[i]), axis=1)
+    
+    for i,value in data_af[headers].mean().sort_values().items():
+        log.write( f"  -{i} : {value}", verbose=verbose)
+        
+        closest_ancestry = data_af[headers].mean().sort_values().idxmin()
+
+    log.write(f"  -Closest Ancestry: {closest_ancestry.split('_')[1]}", verbose=verbose)
+    log.write("Finished inferring ancestry.", verbose=verbose)
+    return closest_ancestry.split("_")[1]
+
+def calculate_fst(p_1, p_2):
+    # https://bios1140.github.io/understanding-fst-the-fixation-index.html
+    # calculate q1 and q2
+    q_1 = 1 - p_1
+    q_2 = 1 - p_2
+
+    # calculate total allele frequency
+    p_t = (p_1 + p_2)/2
+    q_t = 1 - p_t
+
+    # calculate expected heterozygosity
+    # first calculate expected heterozygosity for the two populations
+    # pop1
+    hs_1 = 2*p_1*q_1
+    # pop2
+    hs_2 = 2*p_2*q_2
+    # then take the mean of this
+    hs = (hs_1 + hs_2)/2
+
+    # next calculate expected heterozygosity for the metapopulations
+    ht = 2*p_t*q_t
+
+    # calculate fst
+    fst = (ht - hs)/ht
+
+    # return output
+    return fst
+    

gwaslab/util_ex_ldsc.py

@@ -11,9 +11,10 @@ from gwaslab.util_in_filter_value import filtervalues
 from gwaslab.util_in_filter_value import _filter_palindromic
 from gwaslab.util_in_filter_value import _exclude_hla
 from gwaslab.util_in_filter_value import _exclude_sexchr
+import copy
 
 class ARGS():
-    def __init__(self, **kwargs):
+    def __init__(self, kwargs=None):
         
         self.out = "ldsc"
 
@@ -257,11 +258,12 @@ class ARGS():
 ####################################################################################################################
 
 
-def _estimate_h2_by_ldsc(insumstats, log, verbose=True, munge=False, munge_args=None, **kwargs):
-    sumstats = insumstats.copy()
+def _estimate_h2_by_ldsc(insumstats,  log,  meta=None,verbose=True, munge=False, munge_args=None, **raw_kwargs):
+    sumstats = insumstats
+    kwargs = copy.deepcopy(raw_kwargs)
     
     if "N" in sumstats.columns:
-        sumstats["N"] = sumstats["N"].astype("int64")
+        sumstats["N"] = sumstats["N"].fillna(sumstats["N"].median()).apply("int64")
 
     if munge:
         if munge_args is None:
@@ -291,21 +293,25 @@ def _estimate_h2_by_ldsc(insumstats, log, verbose=True, munge=False, munge_args=
     log.write("  -Adopted from LDSC source code: https://github.com/bulik/ldsc", verbose=verbose)
     log.write("  -Please cite LDSC: Bulik-Sullivan, et al. LD Score Regression Distinguishes Confounding from Polygenicity in Genome-Wide Association Studies. Nature Genetics, 2015.", verbose=verbose)
     
-
-
+    if meta["gwaslab"]["sample_prevalence"]!="Unknown" and meta["gwaslab"]["population_prevalence"]!="Unknown" :
+        if "samp_prev" not in kwargs.keys():
+            kwargs["samp_prev"] = "{}".format(meta["gwaslab"]["sample_prevalence"])
+        if "pop_prev" not in kwargs.keys():
+            kwargs["pop_prev"] =  "{}".format(meta["gwaslab"]["population_prevalence"])
 
     log.write(" -Arguments:", verbose=verbose)
     for key, value in kwargs.items():
         log.write("  -{}:{}".format(key, value), verbose=verbose)
-    default_args = ARGS(**kwargs)
+    
+    default_args = ARGS(kwargs = kwargs)
 
     if "Z" not in sumstats.columns:
         sumstats["Z"] = sumstats["BETA"]/sumstats["SE"]
 
     sumstats = sumstats.rename(columns={"EA":"A1","NEA":"A2","rsID":"SNP"})
-    
+
     log.write(" -LDSC log:", verbose=verbose)
-    summary = estimate_h2(sumstats, default_args, log)
+    summary = estimate_h2(sumstats, args = default_args, log = log)
     
     results_table = None
     if type(summary) is tuple:
@@ -321,10 +327,11 @@ def _estimate_h2_by_ldsc(insumstats, log, verbose=True, munge=False, munge_args=
 
 ####################################################################################################################
 
-def _estimate_partitioned_h2_by_ldsc(insumstats, log, verbose=True, **kwargs):
+def _estimate_partitioned_h2_by_ldsc(insumstats,  log,  meta=None,verbose=True, **raw_kwargs):
     sumstats = insumstats.copy()
+    kwargs = copy.deepcopy(raw_kwargs)
     if "N" in sumstats.columns:
-        sumstats["N"] = sumstats["N"].astype("int64")
+        sumstats["N"] = sumstats["N"].fillna(sumstats["N"].median()).apply("int64")
     ##start function with col checking##########################################################
     _start_line = "run LD score regression"
     _end_line = "running LD score regression"
@@ -347,10 +354,16 @@ def _estimate_partitioned_h2_by_ldsc(insumstats, log, verbose=True, **kwargs):
     log.write("  -Please cite LDSC: Bulik-Sullivan, et al. LD Score Regression Distinguishes Confounding from Polygenicity in Genome-Wide Association Studies. Nature Genetics, 2015.", verbose=verbose)
     log.write(" -Arguments:", verbose=verbose)
     
+    if meta["gwaslab"]["sample_prevalence"]!="Unknown" and meta["gwaslab"]["population_prevalence"]!="Unknown" :
+        if "samp_prev" not in kwargs.keys():
+            kwargs["samp_prev"] = "{}".format(meta["gwaslab"]["sample_prevalence"])
+        if "pop_prev" not in kwargs.keys():
+            kwargs["pop_prev"] =  "{}".format(meta["gwaslab"]["population_prevalence"])
+
     for key, value in kwargs.items():
         log.write("  -{}:{}".format(key, value), verbose=verbose)
     
-    default_args = ARGS(**kwargs)
+    default_args = ARGS(kwargs = kwargs)
 
     if "Z" not in sumstats.columns:
         sumstats["Z"] = sumstats["BETA"]/sumstats["SE"]
@@ -369,10 +382,11 @@ def _estimate_partitioned_h2_by_ldsc(insumstats, log, verbose=True, **kwargs):
 
 
 
-def _estimate_rg_by_ldsc(insumstats, other_traits ,log, verbose=True, **kwargs):
+def _estimate_rg_by_ldsc(insumstats,  other_traits ,log, meta=None, verbose=True, **raw_kwargs):
     sumstats = insumstats.copy()
+    kwargs = copy.deepcopy(raw_kwargs)
     if "N" in sumstats.columns:
-        sumstats["N"] = sumstats["N"].astype("int64")
+        sumstats["N"] = sumstats["N"].fillna(sumstats["N"].median()).apply("int64")
     ##start function with col checking##########################################################
     _start_line = "run LD score regression for genetic correlation"
     _end_line = "running LD score regression for genetic correlation"
@@ -395,18 +409,37 @@ def _estimate_rg_by_ldsc(insumstats, other_traits ,log, verbose=True, **kwargs):
     log.write("  -Please cite LDSC: Bulik-Sullivan, B., et al. An Atlas of Genetic Correlations across Human Diseases and Traits. Nature Genetics, 2015.", verbose=verbose)
     log.write(" -Arguments:", verbose=verbose)
     
+
+    
+    samp_prev_string=""
+    pop_prev_string=""
+
+    if meta["gwaslab"]["sample_prevalence"]!="Unknown" and meta["gwaslab"]["population_prevalence"]!="Unknown" :
+
+        if "samp_prev" not in kwargs.keys():
+            samp_prev_string =  "{}".format(meta["gwaslab"]["sample_prevalence"])
+        if "pop_prev" not in kwargs.keys():
+            pop_prev_string =  "{}".format(meta["gwaslab"]["population_prevalence"])
+    
+    if "rg" in kwargs.keys():
+        alias = kwargs["rg"].split(",")[1:]
+    else:
+        alias=[]
+        for index, each_other_sumstats in enumerate(other_traits):
+            alias.append(each_other_sumstats.meta["gwaslab"]["study_name"])
+        kwargs["rg"]=",".join([meta["gwaslab"]["study_name"]]+alias)
+    
     for key, value in kwargs.items():
         log.write("  -{}:{}".format(key, value), verbose=verbose)
-    
-    default_args = ARGS(**kwargs)
+
+    default_args = ARGS(kwargs = kwargs)
 
     if "Z" not in sumstats.columns:
         sumstats["Z"] = sumstats["BETA"]/sumstats["SE"]
 
     sumstats = sumstats.rename(columns={"EA":"A1","NEA":"A2","rsID":"SNP"})
-    
+
     other_traits_to_use = []
-    alias = default_args.rg.split(",")[1:]
 
     for index, each_other_sumstats in enumerate(other_traits):
         log.write(" -Processing sumstats with alias {} ({})".format(alias[index], each_other_sumstats.meta["gwaslab"]["study_name"]))
@@ -419,6 +452,18 @@ def _estimate_rg_by_ldsc(insumstats, other_traits ,log, verbose=True, **kwargs):
             to_append["Z"] = to_append["BETA"]/to_append["SE"]
 
         other_traits_to_use.append(to_append[["SNP","A1","A2","Z","N"]])    
+        
+        if each_other_sumstats.meta["gwaslab"]["sample_prevalence"]!="Unknown" and each_other_sumstats.meta["gwaslab"]["population_prevalence"]!="Unknown" :
+                samp_prev_string +=  ",{}".format(meta["gwaslab"]["sample_prevalence"])
+                pop_prev_string += ",{}".format(meta["gwaslab"]["population_prevalence"])
+
+    if len(pop_prev_string.split(",")) == len(other_traits)+1 and len(samp_prev_string.split(",")) == len(other_traits)+1:
+        if "samp_prev" not in kwargs.keys():
+            log.write(" -{}:{}".format("samp_prev", samp_prev_string), verbose=verbose)
+            default_args.samp_prev = samp_prev_string
+        if "pop_prev" not in kwargs.keys():
+            log.write(" -{}:{}".format("pop_prev", pop_prev_string), verbose=verbose)
+            default_args.pop_prev =  pop_prev_string    
 
     log.write(" -LDSC log:", verbose=verbose)
     summary = estimate_rg(sumstats[["SNP","A1","A2","Z","N"]], other_traits_to_use, default_args, log)[1]
@@ -431,10 +476,11 @@ def _estimate_rg_by_ldsc(insumstats, other_traits ,log, verbose=True, **kwargs):
 ####################################################################################################################
 
 
-def _estimate_h2_cts_by_ldsc(insumstats, log, verbose=True, **kwargs):
+def _estimate_h2_cts_by_ldsc(insumstats, log, verbose=True, **raw_kwargs):
     sumstats = insumstats.copy()
+    kwargs = copy.deepcopy(raw_kwargs)
     if "N" in sumstats.columns:
-        sumstats["N"] = sumstats["N"].astype("int64")
+        sumstats["N"] = sumstats["N"].fillna(sumstats["N"].median()).apply("int64")
     ##start function with col checking##########################################################
     _start_line = "run LD score regression"
     _end_line = "running LD score regression"
@@ -460,7 +506,7 @@ def _estimate_h2_cts_by_ldsc(insumstats, log, verbose=True, **kwargs):
     for key, value in kwargs.items():
         log.write("  -{}:{}".format(key, value), verbose=verbose)
     
-    default_args = ARGS(**kwargs)
+    default_args = ARGS(kwargs = kwargs)
 
     if "Z" not in sumstats.columns:
         sumstats["Z"] = sumstats["BETA"]/sumstats["SE"]

gwaslab/util_ex_match_ldmatrix.py

@@ -0,0 +1,396 @@
+import scipy.sparse as sparse
+import numpy as np
+import pandas as pd
+from gwaslab.hm_casting import _merge_mold_with_sumstats_by_chrpos
+import subprocess
+import os
+import re
+import gc
+import pandas as pd
+import numpy as np
+from gwaslab.g_Log import Log
+from gwaslab.qc_fix_sumstats import start_to
+from gwaslab.qc_fix_sumstats import finished
+from gwaslab.util_in_get_sig import getsig
+from gwaslab.util_ex_process_ref import _process_plink_input_files
+from gwaslab.g_version import _checking_plink_version
+from gwaslab.util_in_filter_value import _exclude_hla
+from gwaslab.util_ex_calculate_ldmatrix import _extract_variants_in_locus
+
+
+
+
+def tofinemapping_m(sumstats, 
+                    studies=None, 
+                    group=None,
+                    ld_paths = None,
+                    ld_types = None, 
+                    ld_maps = None,
+                    ld_map_dics = None,
+                    bfile=None, 
+                    vcf=None, 
+                    locus=None,
+                    loci=None,
+                    loci_chrpos=None,
+                    out="./",
+                    plink="plink",
+                    plink2="plink2",
+                    windowsizekb=1000,
+                    n_cores=1, 
+                    mode="r",
+                    exclude_hla=False, 
+                    getlead_args=None, 
+                    memory=None, 
+                    overwrite=False,
+                    log=Log(),
+                    suffixes=None,
+                    ld_map_kwargs=None,
+                    extra_plink_option="",
+                    verbose=True,
+                    **kwargs):
+    
+    ##start function with col checking##########################################################
+    _start_line = "calculate LD matrix"
+    _end_line =   "calculating LD matrix"
+    _start_cols =["SNPID","CHR","POS","EA","NEA"]
+    _start_function = ".calculate_ld_matrix()"
+    _must_args ={}
+
+    is_enough_info = start_to(sumstats=sumstats,
+                            log=log,
+                            verbose=verbose,
+                            start_line=_start_line,
+                            end_line=_end_line,
+                            start_cols=_start_cols,
+                            start_function=_start_function,
+                            **_must_args)
+    
+    if is_enough_info == False: raise ValueError("Not enough columns for calculating LD matrix")
+
+    ############################################################################################
+    if suffixes is None:
+        suffixes=[""]
+    if getlead_args is None:
+        getlead_args={"windowsizekb":1000}
+    if ld_map_kwargs is None:
+        ld_map_kwargs={}
+    
+    if loci is None:
+        log.write(" -Loci were not provided. All significant loci will be automatically extracted...",verbose=verbose)
+        sig_df = getsig(sumstats,id="SNPID",chrom="CHR",pos="POS",p="P"+suffixes[0],**getlead_args)
+    else:
+        sig_df = sumstats.loc[sumstats["SNPID"].isin(loci),:]
+
+    # Drop duplicate!!!!
+    log.write(" -Dropping duplicated SNPIDs...",verbose=verbose)
+    sumstats = sumstats.drop_duplicates(subset=["SNPID"]).copy()
+
+    # init Filelist DataFrame
+    output_file_list = pd.DataFrame(columns=["SNPID","SNPID_LIST","LD_R_MATRIX","LOCUS_SUMSTATS"])
+    
+    plink_log=""
+
+    if exclude_hla==True:
+        sig_df = _exclude_hla(sig_df, log=log, verbose=verbose)
+    
+    sig_df = sig_df.reset_index()
+    row = sig_df.iloc[0,:]
+
+    matched_sumstats = _extract_variants_in_locus(sumstats, windowsizekb, locus = (row["CHR"],row["POS"]))
+
+    for i in range(2):
+    # for each study
+        gc.collect()
+        # get ld path and dic
+        ld_map_path = ld_maps[i]
+        ld_map_rename_dic =  ld_map_dics[i]
+        
+        
+        log.write(" -Processing locus with lead variant {} at CHR {} POS {} ...".format(row["SNPID"],row["CHR"],row["POS"]))
+        ld_map = _load_ld_map(ld_map_path, 
+                              ld_map_rename_dic = ld_map_rename_dic, 
+                              **ld_map_kwargs )
+
+        ## check available snps with reference file
+        matched_sumstats = _merge_ld_map_with_sumstats_m(row=row, 
+                                                        locus_sumstats=matched_sumstats, 
+                                                        ld_map=ld_map,
+                                                        log=log,
+                                                        index=i)
+        
+        if len(matched_sumstats)==0:
+            log.write(" -No matching LD information... Skipping...")
+            continue
+    
+    # drop na
+    matched_sumstats = matched_sumstats.dropna()
+
+    # export common variants list
+    matched_snp_list_path, matched_sumstats_paths=_export_snplist_and_locus_sumstats_m(matched_sumstats=matched_sumstats, 
+                                                                                    out=out, 
+                                                                                    group=group, 
+                                                                                    row=row, 
+                                                                                    windowsizekb=windowsizekb,
+                                                                                    log=log)  
+    
+    for i in range(2):
+        ld_path = ld_paths[i]
+
+        r_matrix = _load_ld_matrix(ld_path, 
+                                   fmt="txt", 
+                                   if_square=False, 
+                                   if_add_T=False, 
+                                   log=log, 
+                                   verbose=verbose)
+        
+        matched_ld_matrix_path = _extract_variants_from_ld_matrix_m(merged_sumstats = matched_sumstats, 
+                                                   r_matrix = r_matrix, 
+                                                   out = out, 
+                                                   group = group, 
+                                                   row = row, 
+                                                   windowsizekb = windowsizekb, 
+                                                   index=i,
+                                                   log=log, verbose=verbose)
+#     #########################################################################################################
+
+        row_dict={}
+        row_dict["SNPID"]=row["SNPID"]
+        row_dict["SNPID_LIST"] = matched_snp_list_path
+        row_dict["LD_R_MATRIX"] = matched_ld_matrix_path
+        row_dict["LOCUS_SUMSTATS"] = matched_sumstats_paths[i] + ".gz"
+        row_dict["LOCUS"] = row["SNPID"]
+        row_dict["SUBSTUDY"]= i+1
+        row_dict["STUDY"] = studies[i]
+        file_row = pd.Series(row_dict).to_frame().T
+        output_file_list = pd.concat([output_file_list, file_row],ignore_index=True)
+
+    if len(output_file_list)>0:
+        
+        output_file_list["GROUP"] = group
+        nloci = len(output_file_list)
+        output_file_list_path =  "{}/{}_{}study_{}_{}kb.filelist".format(out.rstrip("/"), group,nloci, row["SNPID"], windowsizekb)
+        output_file_list.to_csv(output_file_list_path,index=None,sep="\t")
+        log.write(" -File list is saved to: {}".format(output_file_list_path),verbose=verbose)
+        log.write(" -Finished LD matrix calculation.",verbose=verbose)
+    else:
+        output_file_list_path=None
+        log.write(" -No avaialable lead variants.",verbose=verbose)
+        log.write(" -Stopped LD matrix calculation.",verbose=verbose)
+
+    finished(log=log, verbose=verbose, end_line=_end_line)
+
+    return output_file_list_path, output_file_list,  plink_log
+
+
+def _export_snplist_and_locus_sumstats_m(matched_sumstats, out, group, row, windowsizekb,log):
+        # study suffixes starting from 1
+        suffixes=["_{}".format(i+1) for i in range(2)]
+        
+        matched_snp_list_path = "{}/{}_{}_{}.snplist.raw".format(out.rstrip("/"), group, row["SNPID"] ,windowsizekb)
+        
+        matched_sumstats["SNPID"].to_csv(matched_snp_list_path, index=None, header=None)
+        log.write(" -Exporting SNP list of {}  to: {}...".format(len(matched_sumstats) ,matched_snp_list_path))
+
+        # create locus-sumstats EA, NEA, (BETA, SE), Z 
+        matched_sumstats_paths =[]
+
+        
+        for i in range(2):
+            # export sumstats for each study
+            suffix = suffixes[i]
+            
+            matched_sumstats_path =  "{}/{}_{}_{}_{}.sumstats".format(out.rstrip("/"), group, row["SNPID"] ,windowsizekb, i + 1)
+            matched_sumstats_paths.append(matched_sumstats_path)
+            to_export_columns=["CHR","POS","EA","NEA"]
+
+            if "Z"+suffix in matched_sumstats.columns :
+                to_export_columns.append("Z"+suffix)
+            if ("BETA"+suffix in matched_sumstats.columns) and ("SE"+suffix in matched_sumstats.columns):
+                to_export_columns.append("BETA"+suffix)
+                to_export_columns.append("SE"+suffix)
+            if "EAF"+suffix in matched_sumstats.columns :
+                to_export_columns.append("EAF"+suffix)
+            if "N"+suffix in matched_sumstats.columns:
+                to_export_columns.append("N"+suffix)
+
+            log.write(" -Exporting locus sumstats to: {}...".format(matched_sumstats_path))
+            log.write(" -Exported columns: {}...".format(["SNPID"]+to_export_columns))
+            #matched_sumstats[ ["SNPID"]+to_export_columns].to_csv(matched_sumstats_path, sep="\t",index=None)
+            rename_dic={
+                "BETA"+suffix:"Beta",
+                "SE"+suffix:"Se",
+                "SNPID":"SNP"
+            }
+            matched_sumstats[ ["SNPID"]+to_export_columns].rename(columns=rename_dic).to_csv(matched_sumstats_path, sep="\t",index=None)
+            matched_sumstats[ ["SNPID"]+to_export_columns].rename(columns=rename_dic).to_csv(matched_sumstats_path+".gz", sep="\t",index=None)
+        
+        return matched_snp_list_path, matched_sumstats_paths
+
+###################################################################################################################################################################
+####################################################################################################
+def _load_ld_matrix(path, 
+                    fmt="npz", 
+                    if_square=False, 
+                    if_add_T=False,
+                    log=Log(),
+                    verbose=True):
+    
+    if fmt == "npz":
+        log.write("   -Loading LD matrix from npz file...",verbose=verbose)
+        r_matrix = sparse.load_npz(path).toarray()
+    if fmt == "txt":
+        log.write("   -Loading LD matrix from text file...",verbose=verbose)
+        r_matrix = np.loadtxt(path,delimiter="\t")
+        log.write("   -LD matrix shape : {}".format(r_matrix.shape) ,verbose=verbose)
+
+    if if_add_T==True:
+        log.write("   -Transforming LD matrix by adding its transpose...",verbose=verbose)
+        r_matrix += r_matrix.T
+    if if_square==True:
+        log.write("   -Transforming LD matrix by squaring all elements...",verbose=verbose)
+        r_matrix = np.power(r_matrix,2)
+    return r_matrix
+    
+def _load_ld_map(path, 
+                 snpid="rsid", 
+                 chrom="chromosome", 
+                 pos="position", 
+                 ref="allele1", 
+                 alt="allele2",
+                 ld_map_rename_dic = None,
+                 **ld_map_kwargs):
+    
+    if ld_map_rename_dic is not None:
+        # ld map format
+        # SNPID_bim,CHR，POS, NEA_bim, EA_bim 
+        if type(ld_map_rename_dic) is dict:
+            ld_map_rename_dic_to_use={ld_map_rename_dic["EA"]:'EA_bim', 
+                                        ld_map_rename_dic["NEA"]:'NEA_bim', 
+                                        ld_map_rename_dic["POS"]:'POS', 
+                                        ld_map_rename_dic["CHR"]:'CHR',
+                                        ld_map_rename_dic["SNPID"]:'SNPID_bim'
+                                        }
+            ld_map_kwargs["usecols"]=list(ld_map_rename_dic.values())
+        else:
+            ld_map_rename_dic_to_use={ld_map_rename_dic[4]:'EA_bim', 
+                                        ld_map_rename_dic[3]:'NEA_bim', 
+                                        ld_map_rename_dic[2]:'POS', 
+                                        ld_map_rename_dic[1]:'CHR',
+                                        ld_map_rename_dic[0]:'SNPID_bim'
+                                        }
+            ld_map_kwargs["usecols"]=ld_map_rename_dic
+    else:
+        ld_map_rename_dic_to_use={alt:'EA_bim', 
+                                  ref:'NEA_bim', 
+                                  pos:'POS', 
+                                  chrom:'CHR',
+                                  snpid:"SNPID_bim"
+                                    }
+        ld_map_kwargs["usecols"]=[chrom, pos, ref, alt, snpid]
+    #rsid    chromosome      position        allele1 allele2
+    if "sep" not in ld_map_kwargs:
+        ld_map_kwargs["sep"] = "\s+"
+    
+    ld_map = pd.read_csv(path,**ld_map_kwargs)
+    ld_map = ld_map.rename(columns=ld_map_rename_dic_to_use, errors='ignore')
+    # "SNPID",0:"CHR_bim",3:"POS_bim",4:"EA_bim",5:"NEA_bim"
+    return ld_map
+
+def _extract_variants_from_ld_matrix_m(merged_sumstats, r_matrix, out, group, row, windowsizekb, log, verbose, index):
+    # study suffixes starting from 1
+    index_bim_header = "_INDEX_BIM_{}".format(index + 1) 
+    flipped_header = "_FLIPPED_{}".format(index + 1) 
+    
+    # a series of int to indicate if the variant index in raw ld matrix
+    avaiable_index = merged_sumstats[index_bim_header].values 
+
+    # a series of boolean values to indicate if the variants is flipped
+    flipped = merged_sumstats[flipped_header].values 
+
+    # extract the sub-matrix
+    reduced_r_matrix = r_matrix[np.ix_(avaiable_index, avaiable_index)]
+    
+    log.write(" -Flipping LD matrix for {} variants...".format(sum(flipped)),verbose=verbose)
+    reduced_r_matrix[flipped,:] = -1 * reduced_r_matrix[flipped,:]
+    reduced_r_matrix[:,flipped] = -1 * reduced_r_matrix[:,flipped]
+
+    output_prefix =  "{}/{}_{}_{}_{}".format(out.rstrip("/"),group,row["SNPID"],windowsizekb, index + 1)
+    output_path = "{}.ld.gz".format(output_prefix)
+    
+    pd.DataFrame(reduced_r_matrix).to_csv(output_path,sep="\t",index=None,header=None)
+    #reduced_r_matrix.to_csv("{}.ld.gz".format(output_prefix),se="\t")
+    return output_path
+
+def _merge_ld_map_with_sumstats_m(row, 
+                             locus_sumstats, 
+                             ld_map, 
+                             log=Log(),
+                             index=None):
+    '''
+    align sumstats with ld map
+    '''
+    # study suffixes starting from 1
+    index_suffix = "_{}".format(index+1)
+    
+    index1= "_INDEX_SUMSTATS" 
+    index2= "_INDEX_BIM" +index_suffix
+    
+    # Sumstats index
+    locus_sumstats[index1] = locus_sumstats.index
+    
+    # ld map index
+    ld_map[index2] =  ld_map.index
+    
+    # init a column to show if the variants in LD map are flipped or not
+    locus_sumstats["_FLIPPED"+index_suffix] = False
+    
+    
+    log.write("   -Variants in locus ({}): {}".format(row["SNPID"],len(locus_sumstats)))
+    # convert category to string
+    locus_sumstats["EA"] = locus_sumstats["EA"].astype("string")
+    locus_sumstats["NEA"] = locus_sumstats["NEA"].astype("string")
+    
+    # matching by SNPID
+    # preserve bim keys (use intersection of keys from both frames, similar to a SQL inner join; preserve the order of the left keys.)
+    # vairants without a match were removed
+    combined_df = pd.merge(ld_map, locus_sumstats, on=["CHR","POS"],how="inner")
+
+    # match allele
+    perfect_match =  ((combined_df["EA"] == combined_df["EA_bim"]) & (combined_df["NEA"] == combined_df["NEA_bim"]) ) 
+    log.write("   -Variants with perfect matched alleles:{}".format(sum(perfect_match)))
+
+    # fliipped allele
+    #ea_mis_match = combined_df["EA"] != combined_df["EA_bim"]
+    flipped_match = ((combined_df["EA"] == combined_df["NEA_bim"])& (combined_df["NEA"] == combined_df["EA_bim"]))
+    log.write("   -Variants with flipped alleles:{}".format(sum(flipped_match)))
+    
+    allele_match = perfect_match | flipped_match
+    log.write("   -Total Variants matched:{}".format(sum(allele_match)))
+    
+    combined_df.loc[flipped_match,"_FLIPPED"+index_suffix] = True
+
+    if row["SNPID"] not in combined_df.loc[allele_match,"SNPID"].values:
+        log.warning("Lead variant was not available in reference!")
+    
+    # adjust output columns
+    output_columns=["SNPID","CHR","POS","EA","NEA"]
+    for i in combined_df.columns:
+        if "_INDEX_BIM" in i:
+            output_columns.append(i)
+        if "_FLIPPED" in i:
+            output_columns.append(i)
+    
+    for i in range(2):
+        # study suffixes starting from 1
+        index_suffix = "_{}".format(i+1)
+        if ("BETA"+index_suffix in combined_df.columns) and ("SE"+index_suffix in combined_df.columns):
+            output_columns.append("BETA"+index_suffix)
+            output_columns.append("SE"+index_suffix)
+        if "Z"+index_suffix in combined_df.columns:
+            output_columns.append("Z"+index_suffix)
+        if "EAF"+index_suffix in combined_df.columns:
+            output_columns.append("EAF"+index_suffix)
+        if "N"+index_suffix in combined_df.columns:
+            output_columns.append("N"+index_suffix)
+    
+    return combined_df.loc[allele_match,output_columns]

gwaslab/util_ex_run_2samplemr.py

@@ -99,8 +99,6 @@ def _run_two_sample_mr(sumstatspair_object,
     write.csv(results_directionality, "{prefix}.directionality", row.names = FALSE)
     '''.format(prefix = prefix)
     
-
-    
     # Two Sample MR
     # Tests
     ## Pleiotropy