LSHTM_analysis/scripts/combining_dfs.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Created on Tue Aug  6 12:56:03 2019

@author: tanu
'''
#=======================================================================
# Task: combining all dfs to a single one

# Input:  8 dfs
#1) <gene>.lower()'_complex_mcsm_norm.csv'
#2) <gene>.lower()_foldx.csv'
#3) <gene>.lower()_dssp.csv'
#4) <gene>.lower()_kd.csv'
#5) <gene>.lower()_rd.csv'
#6) 'ns' + <gene>.lower()_snp_info.csv'
#7) <gene>.lower()_af_or.csv'
#8) <gene>.lower() _af_or_kinship.csv

# combining order
#Merge1 = 1 + 2

#Merge2 = 3 + 4
#Merge3 = Merge2 + 5

#Merge4 = Merge1 + Merge3

#Merge5 = 6 + 7
#Merge6 = Merge5 + 8

#Merge7 = Merge4 + Merge6

# Output: single csv of all 8 dfs combined
# useful link
# https://stackoverflow.com/questions/23668427/pandas-three-way-joining-multiple-dataframes-on-columns
#=======================================================================
#%% load packages
import sys, os
import pandas as pd
from pandas import DataFrame
import numpy as np
import argparse
from functools import reduce
#=======================================================================
#%% specify input and curr dir
homedir = os.path.expanduser('~')

# set working dir
os.getcwd()
os.chdir(homedir + '/git/LSHTM_analysis/scripts')
os.getcwd()

# FIXME: local imports
#from combining import combine_dfs_with_checks
from combining_FIXME import detect_common_cols
from reference_dict import oneletter_aa_dict 
from reference_dict import low_3letter_dict 

from aa_code import get_aa_3lower
from aa_code import get_aa_1upper

# REGEX: as required
# mcsm_regex = r'^([A-Za-z]{1})([0-9]+)([A-Za-z]{1})$'
# mcsm_wt = mcsm_df['mutationinformation'].str.extract(mcsm_regex)[0]
# mcsm_mut = mcsm_df['mutationinformation'].str.extract(mcsm_regex)[2]

# gwas_regex = r'^([A-Za-z]{3})([0-9]+)([A-Za-z]{3})$'
# gwas_wt = mcsm_df['mutation'].str.extract(gwas_regex)[0]
# gwas_pos = mcsm_df['mutation'].str.extract(gwas_regex)[1]
# gwas_mut = mcsm_df['mutation'].str.extract(gwas_regex)[2]
#=======================================================================
#%% command line args: case sensitive
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-d', '--drug', help = 'drug name', default = '')
arg_parser.add_argument('-g', '--gene', help = 'gene name', default = '') 

arg_parser.add_argument('--datadir', help = 'Data Directory. By default, it assmumes homedir + git/Data')
arg_parser.add_argument('-i', '--input_dir', help = 'Input dir containing pdb files. By default, it assmumes homedir + <drug> + input')
arg_parser.add_argument('-o', '--output_dir', help = 'Output dir for results. By default, it assmes homedir + <drug> + output')

arg_parser.add_argument('--debug', action ='store_true', help = 'Debug Mode')

args = arg_parser.parse_args()
#=======================================================================
#%% variable assignment: input and output 
drug    = args.drug
gene    = args.gene
datadir = args.datadir
indir   = args.input_dir
outdir  = args.output_dir

gene_match = gene + '_p.'
print('mut pattern for gene', gene, ':',  gene_match)

#%%=======================================================================
#==============
# directories
#==============
if not datadir:
    datadir = homedir + '/git/Data/'
    
if not indir:
    indir = datadir + drug + '/input/'
    
if not outdir:
    outdir = datadir + drug + '/output/'

#=======
# input
#=======
#in_filename_mcsm = gene.lower() + '_complex_mcsm_norm.csv' 
in_filename_mcsm     = gene.lower() + '_complex_mcsm_norm_SAM.csv' # gidb
in_filename_foldx    = gene.lower() + '_foldx.csv'
in_filename_deepddg  = gene.lower() + '_ni_deepddg.csv' # change to decent filename and put it in the correct dir
in_filename_dssp     = gene.lower() + '_dssp.csv'
in_filename_kd       = gene.lower() + '_kd.csv'
in_filename_rd       = gene.lower() + '_rd.csv'
#in_filename_snpinfo = 'ns' + gene.lower() + '_snp_info_f.csv' # gwas f info
in_filename_afor     = gene.lower() + '_af_or.csv'
#in_filename_afor_kin = gene.lower() + '_af_or_kinship.csv'
infilename_dynamut   = gene.lower() + '_complex_dynamut_norm.csv'
infilename_dynamut2  = gene.lower() + '_complex_dynamut2_norm.csv'
infilename_mcsm_na   = gene.lower() + '_complex_mcsm_na_norm.csv'
infilename_mcsm_f_snps = gene.lower() + '_mcsm_formatted_snps.csv'

infile_mcsm    = outdir + in_filename_mcsm
infile_foldx   = outdir + in_filename_foldx
infile_deepddg = outdir + in_filename_deepddg
infile_dssp    = outdir + in_filename_dssp
infile_kd      = outdir + in_filename_kd
infile_rd      = outdir + in_filename_rd
#infile_snpinfo = outdir + in_filename_snpinfo 
infile_afor    = outdir + in_filename_afor
#infile_afor_kin = outdir + in_filename_afor_kin
infile_dynamut = outdir + 'dynamut_results/' + infilename_dynamut
infile_dynamut2 = outdir + 'dynamut_results/dynamut2/' + infilename_dynamut2
infile_mcsm_na = outdir + 'mcsm_na_results/' + infilename_mcsm_na
infile_mcsm_f_snps = outdir + infilename_mcsm_f_snps

# read csv
mcsm_df      = pd.read_csv(infile_mcsm, sep = ',')
foldx_df     = pd.read_csv(infile_foldx , sep = ',')
deepddg_df   = pd.read_csv(infile_deepddg, sep = ',')
dssp_df      = pd.read_csv(infile_dssp, sep = ',')
kd_df        = pd.read_csv(infile_kd, sep = ',')
rd_df        = pd.read_csv(infile_rd, sep = ',')
afor_df      = pd.read_csv(infile_afor, sep = ',') 
dynamut_df   = pd.read_csv(infile_dynamut, sep = ',')
dynamut2_df  = pd.read_csv(infile_dynamut2, sep = ',')
mcsm_na_df   = pd.read_csv(infile_mcsm_na, sep = ',')
mcsm_f_snps  = pd.read_csv(infile_mcsm_f_snps, sep = ',', names = ['mutationinformation'], header = None)

#=======
# output 
#=======
out_filename_comb = gene.lower() + '_all_params.csv'
outfile_comb =  outdir + out_filename_comb
print('Output filename:', outfile_comb
      , '\n===================================================================')

o_join = 'outer'
l_join = 'left'
r_join = 'right'
i_join = 'inner'

# end of variable assignment for input and output files
#%%############################################################################  
#=====================
# some preprocessing
#=====================
#-------------
# FoldX
#-------------
foldx_df.shape
#=======================
# scale foldx values
#=======================
  
# Rescale values in Foldx_change col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
foldx_min = foldx_df['ddg'].min() 
foldx_max = foldx_df['ddg'].max() 
foldx_min
foldx_max

foldx_scale = lambda x : x/abs(foldx_min) if x < 0 else (x/foldx_max if x >= 0 else 'failed')

foldx_df['foldx_scaled'] = foldx_df['ddg'].apply(foldx_scale)
print('Raw foldx scores:\n', foldx_df['ddg']
    , '\n---------------------------------------------------------------'
    , '\nScaled foldx scores:\n', foldx_df['foldx_scaled'])

# additional check added
fsmi = foldx_df['foldx_scaled'].min()
fsma = foldx_df['foldx_scaled'].max()

c = foldx_df[foldx_df['ddg']>=0].count()
foldx_pos = c.get(key = 'ddg')

c2 = foldx_df[foldx_df['foldx_scaled']>=0].count()
foldx_pos2 = c2.get(key = 'foldx_scaled')

if foldx_pos == foldx_pos2 and fsmi == -1 and fsma == 1:
    print('\nPASS: Foldx values scaled correctly b/w -1 and 1')
else:
    print('\nFAIL: Foldx values scaled numbers MISmatch'
          , '\nExpected number:', foldx_pos
          , '\nGot:', foldx_pos2
          , '\n======================================================')

# rename ddg column to ddg_foldx
foldx_df['ddg']
foldx_df = foldx_df.rename(columns = {'ddg':'ddg_foldx'})   
foldx_df['ddg_foldx']

#-------------
# Deepddg
#-------------
deepddg_df.shape

#=======================
# scale Deepddg values
#=======================

# Rescale values in deepddg_change col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
deepddg_min = deepddg_df['deepddg'].min() 
deepddg_max = deepddg_df['deepddg'].max() 

deepddg_scale = lambda x : x/abs(deepddg_min) if x < 0 else (x/deepddg_max if x >= 0 else 'failed')

deepddg_df['deepddg_scaled'] = deepddg_df['deepddg'].apply(deepddg_scale)
print('Raw deepddg scores:\n', deepddg_df['deepddg']
    , '\n---------------------------------------------------------------'
    , '\nScaled deepddg scores:\n', deepddg_df['deepddg_scaled'])
    
# additional check added
dsmi = deepddg_df['deepddg_scaled'].min()
dsma = deepddg_df['deepddg_scaled'].max()

c = deepddg_df[deepddg_df['deepddg']>=0].count()
deepddg_pos = c.get(key = 'deepddg')

c2 = deepddg_df[deepddg_df['deepddg_scaled']>=0].count()
deepddg_pos2 = c2.get(key = 'deepddg_scaled')
    
if deepddg_pos == deepddg_pos2 and dsmi == -1 and dsma == 1:
    print('\nPASS: deepddg values scaled correctly b/w -1 and 1')
else:
    print('\nFAIL: deepddg values scaled numbers MISmatch'
          , '\nExpected number:', deepddg_pos
          , '\nGot:', deepddg_pos2
          , '\n======================================================')
#%%=============================================================================
# Now merges begin
#%%=============================================================================
print('==================================='
      , '\nFirst merge: mcsm + foldx'
      , '\n===================================')

mcsm_df.shape

# add 3 lowercase aa code for wt and mutant
get_aa_3lower(df = mcsm_df
              , wt_colname = 'wild_type'
              , mut_colname = 'mutant_type'
              , col_wt = 'wt_aa_3lower'
              , col_mut = 'mut_aa_3lower')

#mcsm_df.columns = mcsm_df.columns.str.lower()
# foldx_df.shape

#mcsm_foldx_dfs = combine_dfs_with_checks(mcsm_df, foldx_df, my_join = o_join)
merging_cols_m1  = detect_common_cols(mcsm_df, foldx_df)
mcsm_foldx_dfs = pd.merge(mcsm_df, foldx_df, on = merging_cols_m1,  how =  o_join)
ncols_m1 = len(mcsm_foldx_dfs.columns)

print('\n\nResult of first merge:', mcsm_foldx_dfs.shape
      , '\n===================================================================')
mcsm_foldx_dfs[merging_cols_m1].apply(len)
mcsm_foldx_dfs[merging_cols_m1].apply(len) == len(mcsm_foldx_dfs)
#%%
print('==================================='
      , '\nSecond merge: mcsm_foldx_dfs + deepddg'
      , '\n===================================')

#deepddg_df =  pd.read_csv(infile_deepddg, sep = ',')
#deepddg_df.columns

# merge with mcsm_foldx_dfs and deepddg_df
mcsm_foldx_deepddg_dfs = pd.merge(mcsm_foldx_dfs, deepddg_df, on = 'mutationinformation',  how = l_join)
mcsm_foldx_deepddg_dfs['deepddg_outcome'].value_counts()

ncols_deepddg_merge = len(mcsm_foldx_deepddg_dfs.columns)
#%%============================================================================
print('==================================='
      , '\Third merge: dssp + kd'
      , '\n===================================')

dssp_df.shape
kd_df.shape
rd_df.shape

#dssp_kd_dfs = combine_dfs_with_checks(dssp_df, kd_df, my_join = o_join)
merging_cols_m2 = detect_common_cols(dssp_df, kd_df)
dssp_kd_dfs = pd.merge(dssp_df, kd_df, on = merging_cols_m2,  how = o_join)

print('\n\nResult of third merge:', dssp_kd_dfs.shape
      , '\n===================================================================')
#%%============================================================================
print('==================================='
      , '\nFourth merge: third merge + rd_df' 
      , '\ndssp_kd_dfs + rd_df'
      , '\n===================================')
#dssp_kd_rd_dfs = combine_dfs_with_checks(dssp_kd_dfs, rd_df, my_join = o_join)
merging_cols_m3 = detect_common_cols(dssp_kd_dfs,  rd_df)
dssp_kd_rd_dfs = pd.merge(dssp_kd_dfs, rd_df, on = merging_cols_m3
                          , how = o_join)

ncols_m3 = len(dssp_kd_rd_dfs.columns)

print('\n\nResult of Third merge:', dssp_kd_rd_dfs.shape
      , '\n===================================================================')
dssp_kd_rd_dfs[merging_cols_m3].apply(len)
dssp_kd_rd_dfs[merging_cols_m3].apply(len) == len(dssp_kd_rd_dfs)
#%%============================================================================
print('======================================='
      , '\nFifth merge: Second merge + fourth merge'
      , '\nmcsm_foldx_dfs + dssp_kd_rd_dfs'
      , '\n=======================================')
#combined_df = combine_dfs_with_checks(mcsm_foldx_dfs, dssp_kd_rd_dfs, my_join = i_join)
#merging_cols_m4 = detect_common_cols(mcsm_foldx_dfs, dssp_kd_rd_dfs)
#combined_df = pd.merge(mcsm_foldx_dfs, dssp_kd_rd_dfs, on = merging_cols_m4, how  = i_join)
#combined_df_expected_cols = ncols_m1 + ncols_m3 - len(merging_cols_m4)

# with deepddg values
merging_cols_m4 = detect_common_cols(mcsm_foldx_deepddg_dfs, dssp_kd_rd_dfs)
combined_df = pd.merge(mcsm_foldx_deepddg_dfs, dssp_kd_rd_dfs, on = merging_cols_m4, how  = i_join)

combined_df_expected_cols = ncols_deepddg_merge + ncols_m3 - len(merging_cols_m4)

if len(combined_df) == len(mcsm_df) and len(combined_df.columns) == combined_df_expected_cols:
    print('PASS: successfully combined 5 dfs'
          , '\nNo. of rows combined_df:', len(combined_df)
          , '\nNo. of cols combined_df:', len(combined_df.columns))   
else:
    sys.exit('FAIL: check individual df merges')

print('\nResult of Fourth merge:', combined_df.shape
      , '\n===================================================================')

combined_df[merging_cols_m4].apply(len)
combined_df[merging_cols_m4].apply(len) == len(combined_df)
#%%============================================================================
# Format the combined df columns
combined_df_colnames = combined_df.columns

# check redundant columns
combined_df['chain'].equals(combined_df['chain_id'])
combined_df['wild_type'].equals(combined_df['wild_type_kd']) # has nan
combined_df['wild_type'].equals(combined_df['wild_type_dssp'])

#sanity check
foo = combined_df[['wild_type', 'wild_type_kd', 'wt_3letter_caps', 'wt_aa_3lower', 'mut_aa_3lower']] 

# Drop cols
cols_to_drop = ['chain_id', 'wild_type_kd', 'wild_type_dssp', 'wt_3letter_caps' ]
combined_df_clean = combined_df.drop(cols_to_drop, axis = 1)

del(foo)
#%%============================================================================
# Output columns
out_filename_stab_struc = gene.lower() + '_comb_stab_struc_params.csv'
outfile_stab_struc =  outdir + '/' + out_filename_stab_struc
print('Output filename:', outfile_stab_struc
      , '\n===================================================================')

# write csv
print('\nWriting file: combined stability and structural parameters')
combined_df_clean.to_csv(outfile_stab_struc, index = False)
print('\nFinished writing file:'
      , '\nNo. of rows:', combined_df_clean.shape[0]
      , '\nNo. of cols:', combined_df_clean.shape[1])
#%%=====================================================================
print('\n======================================='
      , '\nFifth merge:'
      , '\ncombined_df_clean + afor_df '
      , '\n=======================================')

afor_cols = afor_df.columns
afor_df.shape

# create a mapping from the gwas mutation column i.e <gene_match>_abcXXXrst
#----------------------
# call get_aa_upper(): 
# adds 3 more cols with one letter aa code
#----------------------
get_aa_1upper(df = afor_df
               , gwas_mut_colname = 'mutation'
               , wt_colname = 'wild_type'
               , pos_colname = 'position'
               , mut_colname = 'mutant_type')

afor_df['mutationinformation'] = afor_df['wild_type'] + afor_df['position'].map(str) + afor_df['mutant_type']
afor_cols = afor_df.columns

merging_cols_m5 = detect_common_cols(combined_df_clean, afor_df) 

# remove position so that merging can take place without dtype conflicts
merging_cols_m5.remove('position')

# drop position column from afor_df
afor_df = afor_df.drop(['position'], axis = 1)
afor_cols = afor_df.columns

# merge 
combined_stab_afor = pd.merge(combined_df_clean, afor_df, on = merging_cols_m5, how  = l_join)
comb_afor_df_cols = combined_stab_afor.columns

comb_afor_expected_cols = len(combined_df_clean.columns) + len(afor_df.columns) - len(merging_cols_m5)

if len(combined_stab_afor) == len(combined_df_clean) and len(combined_stab_afor.columns) == comb_afor_expected_cols:
    print('\nPASS: successfully combined 6 dfs'
          , '\nNo. of rows combined_stab_afor:', len(combined_stab_afor)
          , '\nNo. of cols combined_stab_afor:', len(combined_stab_afor.columns))   
else:
    sys.exit('\nFAIL: check individual df merges')

print('\n\nResult of Fourth merge:', combined_stab_afor.shape
      , '\n===================================================================')

combined_stab_afor[merging_cols_m5].apply(len)
combined_stab_afor[merging_cols_m5].apply(len) == len(combined_stab_afor)

if len(combined_stab_afor) -  combined_stab_afor['mutation'].isna().sum() == len(afor_df):
     print('\nPASS: Merge successful for af and or'
          , '\nNo. of nsSNPs with valid ORs: ', len(afor_df))
else:
    sys.exit('\nFAIL: merge unsuccessful for af and or')

#%%============================================================================
# Output columns: when dynamut, dynamut2 and others weren't being combined
out_filename_comb_afor = gene.lower() + '_comb_afor.csv'
outfile_comb_afor =  outdir + '/' + out_filename_comb_afor
print('Output filename:', outfile_comb_afor
      , '\n===================================================================')

# # write csv
print('Writing file: combined stability and afor')
combined_stab_afor.to_csv(outfile_comb_afor, index = False)
print('\nFinished writing file:'
      , '\nNo. of rows:', combined_stab_afor.shape[0]
      , '\nNo. of cols:', combined_stab_afor.shape[1])
#%%============================================================================
# combine dynamut, dynamut2, and mcsm_na
dfs_list = [dynamut_df, dynamut2_df, mcsm_na_df]

dfs_merged = reduce(lambda  left,right: pd.merge(left
                                                , right
                                                , on = ['mutationinformation']
                                                , how = 'inner')
                   , dfs_list)
# drop excess columns
drop_cols = detect_common_cols(dfs_merged, combined_stab_afor)
drop_cols.remove('mutationinformation')

dfs_merged_clean = dfs_merged.drop(drop_cols, axis = 1)
merging_cols_m6 = detect_common_cols(dfs_merged_clean, combined_stab_afor)

len(dfs_merged_clean.columns)
len(combined_stab_afor.columns)

combined_all_params = pd.merge(combined_stab_afor
                               , dfs_merged_clean
                               , on = merging_cols_m6
                               , how  = i_join)

expected_ncols = len(dfs_merged_clean.columns) + len(combined_stab_afor.columns) - len(merging_cols_m6)
expected_nrows = len(combined_stab_afor)

if len(combined_all_params.columns) == expected_ncols and len(combined_all_params) == expected_nrows:
    print('\nPASS: All dfs combined')
else:
    print('\nFAIL:lengths mismatch'
          , '\nExpected ncols:', expected_ncols
          , '\nGot:', len(dfs_merged_clean.columns) 
          , '\nExpected nrows:', expected_nrows
          , '\nGot:', len(dfs_merged_clean) )
    
#%% Done for gid on 10/09/2021
# write csv
print('Writing file: all params')
combined_all_params.to_csv(outfile_comb, index = False)

print('\nFinished writing file:'
      , '\nNo. of rows:', combined_all_params.shape[0]
      , '\nNo. of cols:', combined_all_params.shape[1])
#%% end of script