#!/usr/bin/env python3 # -*- coding: utf-8 -*- ''' @author: tanu ''' #%% Description # FIXME: include error checking to enure you only # concentrate on positions that have structural info? # FIXME: import dirs.py to get the basic dir paths available #======================================================================= # TASK: extract ALL matched mutations from GWAS data # Input data file has the following format: each row = unique sample id # id,country,lineage,sublineage,drtype,drug,dr_muts_col,other_muts_col... # 0,sampleID,USA,lineage2,lineage2.2.1,Drug-resistant,0.0,WT,gene_matchPOS; pncA_c.POS... # where multiple mutations and multiple mutation types are separated by ';'. # We are interested in the protein coding region i.e mutation with the_'p.' format. # This script splits the mutations on the ';' and extracts protein coding muts only # where each row is a separate mutation # sample ids AND mutations are NOT unique, but the COMBINATION (sample id + mutation) = unique # output files: all lower case # 1) _gwas.csv # 2) _common_ids.csv # 3) _ambiguous_muts.csv # 4) _mcsm_formatted_snps.csv # 5) _metadata_poscounts.csv # 6) _metadata.csv # 7) _all_muts_msa.csv # 8) _mutational_positons.csv #------------ # NOTE #----------- # drtype is renamed to 'resistance' in the 35k dataset # all colnames in the ouput files lowercase #------------- # requires #------------- #reference_dict.py #tidy_split.py # bash counting for cross checks: example #grep -oP 'pncA_p.[A-Za-z]{3}[0-9]+[A-Za-z]{3}' mtb_gwas_meta_v6.csv | sort | uniq -c | wc -l #======================================================================= #%% load libraries import os, sys import re import pandas as pd import numpy as np import argparse from statistics import mean, median, mode from statistics import multimode #======================================================================= #%% dir and local imports homedir = os.path.expanduser('~') # set working dir os.getcwd() os.chdir(homedir + '/git/LSHTM_analysis/scripts') os.getcwd() #======================================================================= #%% command line args arg_parser = argparse.ArgumentParser() arg_parser.add_argument('-d', '--drug', help='drug name (case sensitive)', default = None) arg_parser.add_argument('-g', '--gene', help='gene name (case sensitive)', default = None) 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 + + input') arg_parser.add_argument('-o', '--output_dir', help = 'Output dir for results. By default, it assmes homedir + + output') arg_parser.add_argument('-m', '--make_dirs', help = 'Make dir for input and output', action='store_true') arg_parser.add_argument('--debug', action ='store_true', help = 'Debug Mode') args = arg_parser.parse_args() ############################################################################### #%% variable assignment: input and output paths & filenames drug = args.drug gene = args.gene datadir = args.datadir indir = args.input_dir outdir = args.output_dir make_dirs = args.make_dirs ############################################################################### #%% variable assignment: input and output dirs and files #======= # dirs #======= if not datadir: datadir = homedir + '/' + 'git/Data' if not indir: indir = datadir + '/' + drug + '/input_v2' if not outdir: outdir = datadir + '/' + drug + '/output_v2' if make_dirs: print('make_dirs is turned on, creating data dir:', datadir) try: os.makedirs(datadir, exist_ok = True) print("Directory '%s' created successfully" %datadir) except OSError as error: print("Directory '%s' can not be created") print('make_dirs is turned on, creating indir:', indir) try: os.makedirs(indir, exist_ok = True) print("Directory '%s' created successfully" %indir) except OSError as error: print("Directory '%s' can not be created") print('make_dirs is turned on, creating outdir:', outdir) try: os.makedirs(outdir, exist_ok = True) print("Directory '%s' created successfully" %outdir) except OSError as error: print("Directory '%s' can not be created") # handle missing dirs here if not os.path.isdir(datadir): print('ERROR: Data directory does not exist:', datadir , '\nPlease create and ensure gwas data is present and then rerun\nelse specify cmd option --make_dirs') sys.exit() if not os.path.isdir(indir): print('ERROR: Input directory does not exist:', indir , '\nPlease either create or specify indir and rerun\nelse specify cmd option --make_dirs') sys.exit() if not os.path.isdir(outdir): print('ERROR: Output directory does not exist:', outdir , '\nPlease create or specify outdir and rerun\nelse specify cmd option --make_dirs') sys.exit() ############################################################################### #%% required local imports from reference_dict import my_aa_dict # CHECK DIR STRUC THERE! from tidy_split import tidy_split ############################################################################### #%% creating required variables: gene and drug dependent, and input filename gene_match = gene + '_p.' print('mut pattern for gene', gene, ':', gene_match) nssnp_match = gene_match +'[A-Za-z]{3}[0-9]+[A-Za-z]{3}' print('nsSNP for gene', gene, ':', nssnp_match) nssnp_match_captureG = "(" + nssnp_match + ")" wt_regex = gene_match.lower()+'([A-Za-z]{3})' print('wt regex:', wt_regex) mut_regex = r'[0-9]+(\w{3})$' print('mt regex:', mut_regex) pos_regex = r'([0-9]+)' print('position regex:', pos_regex) # building cols to extract dr_muts_col = 'dr_mutations_' + drug other_muts_col = 'other_mutations_' + drug resistance_col = 'drtype' print('Extracting columns based on variables:\n' , drug , '\n' , dr_muts_col , '\n' , other_muts_col , '\n' , resistance_col , '\n===============================================================') #======= # input #======= #in_filename_master_master = 'original_tanushree_data_v2.csv' #19k in_filename_master = 'mtb_gwas_meta_v6.csv' #35k infile_master = datadir + '/' + in_filename_master print('Input file: ', infile_master , '\n============================================================') #======= # output #======= # several output files: in respective sections at the time of outputting files print('Output filename: in the respective sections' , '\nOutput path: ', outdir , '\n=============================================================') #end of variable assignment for input and output files ############################################################################### #%% Read input file master_data = pd.read_csv(infile_master, sep = ',') # column names #list(master_data.columns) # extract elevant columns to extract from meta data related to the drug if in_filename_master == 'original_tanushree_data_v2.csv': meta_data = master_data[['id' , 'country' , 'lineage' , 'sublineage' , 'drtype' , drug , dr_muts_col , other_muts_col]] else: core_cols = ['id' , 'sample' #, 'patient_id' #, 'strain' , 'lineage' , 'sublineage' , 'country_code' #, 'geographic_source' , resistance_col] variable_based_cols = [drug , dr_muts_col , other_muts_col] cols_to_extract = core_cols + variable_based_cols print('Extracting', len(cols_to_extract), 'columns from master data') meta_data = master_data[cols_to_extract] del(master_data, variable_based_cols, cols_to_extract) print('Extracted meta data from filename:', in_filename_master , '\nDim:', meta_data.shape) # checks and results total_samples = meta_data['id'].nunique() print('RESULT: Total samples:', total_samples , '\n===========================================================') # counts NA per column meta_data.isna().sum() print('No. of NAs/column:' + '\n', meta_data.isna().sum() , '\n===========================================================\n') ############################################################################### #%% Quick checks: meta_data['lineage'].value_counts() meta_data['lineage'].value_counts().sum() meta_data['lineage'].nunique() # replace lineage with 'L' in lineage_labels #meta_data['lineage_labels'] = meta_data['lineage'] #meta_data['lineage_labels'].equals(meta_data['lineage']) #all(meta_data['lineage_labels'].value_counts() == meta_data['lineage'].value_counts()) #meta_data['lineage_labels'] = meta_data['lineage_labels'].str.replace('lineage', 'L') #meta_data['lineage'].value_counts() #meta_data['lineage_labels'].value_counts() meta_data['lineage'] = meta_data['lineage'].str.replace('lineage', 'L') meta_data['lineage'].value_counts() print("\n================================" , "\nLineage numbers" , "\nComplete lineage samples:", meta_data['lineage'].value_counts().sum() , "\nMissing lineage samples:", meta_data['id'].nunique() - meta_data['lineage'].value_counts().sum() , "\n================================") meta_data['id'].nunique() meta_data['sample'].nunique() meta_data['id'].equals(meta_data['sample']) foo = meta_data.copy() foo['Diff'] = np.where( foo['id'] == foo['sample'] , '1', '0') foo['Diff'].value_counts() meta_data['drtype'].value_counts() meta_data['drtype'].value_counts().sum() print("\n================================" , "\ndrtype numbers" , "\nComplete drtype samples:", meta_data['drtype'].value_counts().sum() , "\nMissing drtype samples:", meta_data['id'].nunique() - meta_data['drtype'].value_counts().sum() , "\n================================") meta_data['drug_name'] = meta_data[drug].map({1:'R' , 0:'S'}) meta_data['drug_name'].value_counts() meta_data[drug].value_counts() meta_data[drug].value_counts().sum() print("\n================================" , "\ndrug", drug, "numbers" , "\nComplete drug samples:", meta_data[drug].value_counts().sum() , "\nMissing drug samples:", meta_data['id'].nunique() - meta_data[drug].value_counts().sum() , "\n================================") print("\n================================" , "\ndrug", drug, "numbers" , "\nComplete drug samples:", meta_data['drug_name'].value_counts().sum() , "\nMissing drug samples:", meta_data['id'].nunique() - meta_data['drug_name'].value_counts().sum() , "\n================================") #%% Quick counts: All samples, drug: print('===========================================================\n' , 'RESULT: Total no. of samples tested for', drug, ':', meta_data[drug].value_counts().sum() , '\n===========================================================\n' , 'RESULT: No. of Sus and Res', drug, 'samples:\n', meta_data[drug].value_counts() , '\n===========================================================\n' , 'RESULT: Percentage of Sus and Res', drug, 'samples:\n', meta_data[drug].value_counts(normalize = True)*100 , '\n===========================================================') print('===========================================================\n' , 'RESULT: Total no. of samples tested for', drug, ':', meta_data['drug_name'].value_counts().sum() , '\n===========================================================\n' , 'RESULT: No. of Sus and Res', drug, 'samples:\n', meta_data['drug_name'].value_counts() , '\n===========================================================\n' , 'RESULT: Percentage of Sus and Res', drug, 'samples:\n', meta_data['drug_name'].value_counts(normalize = True)*100 , '\n===========================================================') ############################################################################## #%% Extracting nsSNP for gene (meta_gene_all): from dr_muts col and other_muts_col #meta_data_gene = meta_data.loc[ (meta_data[dr_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False) ) | (meta_data[other_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False) ) ] # so downstream code doesn't change meta_gene_all = meta_data.loc[ (meta_data[dr_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False) ) | (meta_data[other_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False) ) ] #%% DF: with dr_muts_col meta_gene_dr = meta_gene_all.loc[meta_gene_all[dr_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False)] meta_gene_dr1 = meta_data.loc[meta_data[dr_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False)] if meta_gene_dr.equals(meta_gene_dr1): print('\nPASS: DF with column:', dr_muts_col, 'extracted successfully' , '\ngene_snp_match in column:',dr_muts_col, meta_gene_dr.shape) else: sys.exit('\nFAIL: DF with column:', dr_muts_col,'could not be extracted' , '\nshape of df1:', meta_gene_dr.shape , '\nshape of df2:', meta_gene_dr1.shape , '\nCheck again!') ############################################################################## #%% DF: with other_muts_col meta_gene_other = meta_gene_all.loc[meta_gene_all[other_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False)] meta_gene_other1 = meta_data.loc[meta_data[other_muts_col].str.contains(nssnp_match, na = False, regex = True, case = False)] print('gene_snp_match in dr:', len(meta_gene_other)) meta_gene_other.equals(meta_gene_other1) if meta_gene_other.equals(meta_gene_other1): print('\nPASS: DF with column:', other_muts_col,'extracted successfully' , '\ngene_snp_match in column:',other_muts_col, meta_gene_other.shape) else: sys.exit('\nFAIL: DF with column:', other_muts_col,'could not be extracted' , '\nshape of df1:', meta_gene_other.shape , '\nshape of df2:', meta_gene_other1.shape , '\nCheck again!') ############################################################################## #%% Quick counts: nsSNP samples, drug meta_gene_all[drug].value_counts() print('===========================================================\n' , 'RESULT: Total no. of samples for', drug, 'with nsSNP mutations:', meta_gene_all['id'].nunique() , '\n===========================================================\n' , '===========================================================\n' , 'RESULT: Total no. of samples tested for', drug, 'with nsSNP mutations:', meta_gene_all[drug].value_counts().sum() , '\n===========================================================\n' , '===========================================================\n' , 'RESULT: Total no. of samples tested for', drug, 'with nsSNP mutations:', meta_gene_all['drug_name'].value_counts().sum() , '\n===========================================================\n' , 'RESULT: No. of Sus and Res', drug, 'samples with nsSNP:\n', meta_gene_all['drug_name'].value_counts() , '\n===========================================================\n' , 'RESULT: Percentage of Sus and Res', drug, 'samples with nsSNP mutations:\n', meta_gene_all['drug_name'].value_counts(normalize = True)*100 , '\n===========================================================') ############################################################################### #%% Create a copy of indices for downstream mergeing meta_gene_all['index_orig'] = meta_gene_all.index meta_gene_all['index_orig_copy'] = meta_gene_all.index all(meta_gene_all.index.values == meta_gene_all['index_orig'].values) all(meta_gene_all.index.values == meta_gene_all['index_orig_copy'].values) ############################################################################## #%% Important sanity checks: Dr muts column for tidy split(), nsSNPs, etc. # Split based on semi colon dr_muts_col and other_muts_col search = ";" # count of occurrence of ";" in dr_muts_col: No.of semicolons + 1 is no. of rows created * occurence count_df_dr = meta_gene_dr[['id', dr_muts_col]] count_df_dr['dr_semicolon_count'] = meta_gene_dr.loc[:, dr_muts_col].str.count(search, re.I) dr_sc_C = count_df_dr['dr_semicolon_count'].value_counts().reset_index() dr_sc_C dr_sc_C['index_semicolon'] = (dr_sc_C['index'] + 1) *dr_sc_C['dr_semicolon_count'] dr_sc_C expected_drC = dr_sc_C['index_semicolon'].sum() expected_drC # count no. of nsSNPs and extract those nsSNPs count_df_dr['dr_geneSNP_count'] = meta_gene_dr.loc[:, dr_muts_col].str.count(nssnp_match, re.I) dr_gene_count1 = count_df_dr['dr_geneSNP_count'].sum() # DEL: need to count by find_all # dr_snps_df = meta_gene_dr.loc[:, dr_muts_col].str.extract(nssnp_match_captureG, re.I) # dr_snps_df.columns = [dr_muts_col] # dr_snps_countU = dr_snps_df[dr_muts_col].nunique() # dr_snps_countU # dr_snps_count_dict = dr_snps_df[dr_muts_col].value_counts().to_dict() # dr_snps_count_dict ############################################################################### # This is to find out how many samples have 1 and more than 1 mutation,so you # can use it to check if your data extraction process for dr_muts # and other_muts has worked correctly AND also to check the dim of the # final formatted data. # This will have: unique COMBINATION of sample id and mutations #======== # First: counting mutations in dr_muts_col column #======== print('Now counting WT &', nssnp_match, 'muts within the column:', dr_muts_col) # drop na and extract a clean df clean_df = meta_data.dropna(subset=[dr_muts_col]) # sanity check: count na na_count = meta_data[dr_muts_col].isna().sum() if len(clean_df) == (total_samples - na_count): print('PASS: clean_df extracted: length is', len(clean_df) , '\nNo.of NAs in', dr_muts_col, '=', na_count, '/', total_samples , '\n==========================================================') else: sys.exit('FAIL: Could not drop NAs') dr_gene_count = 0 wt = 0 id_dr = [] id2_dr = [] dr_gene_mutsL = [] #nssnp_match_regex = re.compile(nssnp_match) for i, id in enumerate(clean_df.id): #print (i, id) #id_dr.append(id) #count_gene_dr = clean_df[dr_muts_col].iloc[i].count(gene_match) # can include stop muts count_gene_dr = len(re.findall(nssnp_match, clean_df[dr_muts_col].iloc[i], re.IGNORECASE)) gene_drL = re.findall(nssnp_match, clean_df[dr_muts_col].iloc[i], re.IGNORECASE) #print(count_gene_dr) if count_gene_dr > 0: id_dr.append(id) if count_gene_dr > 1: id2_dr.append(id) #print(id, count_gene_dr) dr_gene_count = dr_gene_count + count_gene_dr dr_gene_mutsL = dr_gene_mutsL + gene_drL count_wt = clean_df[dr_muts_col].iloc[i].count('WT') wt = wt + count_wt print('RESULTS:') print('Total WT in dr_muts_col:', wt) print('Total matches of', gene, 'SNP matches in', dr_muts_col, ':', dr_gene_count) print('Total samples with > 1', gene, 'nsSNPs in dr_muts_col:', len(id2_dr) ) print('Total matches of UNIQUE', gene, 'SNP matches in', dr_muts_col, ':', len(set(dr_gene_mutsL))) print('=================================================================') if dr_gene_count == dr_gene_count1: print('\nPass: dr gene SNP count match') else: sys.exit('\nFAIL: dr gene SNP count MISmatch') del(clean_df, na_count, i, id, wt, id2_dr, count_gene_dr, count_wt) #%% tidy_split(): dr_muts_col #========= # DF1: dr_muts_col # and remove leading white spaces #========= col_to_split1 = dr_muts_col print ('Firstly, applying tidy split on dr muts df', meta_gene_dr.shape , '\ncolumn name to apply tidy_split():' , col_to_split1 , '\n============================================================') # apply tidy_split() dr_WF0 = tidy_split(meta_gene_dr, col_to_split1, sep = ';') # remove leading white space else these are counted as distinct mutations as well dr_WF0[dr_muts_col] = dr_WF0[dr_muts_col].str.strip() if len(dr_WF0) == expected_drC: print('\nPass: tidy split on', dr_muts_col, 'worked' , '\nExpected nrows:', expected_drC , '\nGot nrows:', len(dr_WF0) ) else: print('\nFAIL: tidy split on', dr_muts_col, 'did not work' , '\nExpected nrows:', expected_drC , '\nGot nrows:', len(dr_WF0) ) # Extract only the samples/rows with nssnp_match #dr_gene_WF0 = dr_WF0.loc[dr_WF0[dr_muts_col].str.contains(gene_match)] dr_gene_WF0 = dr_WF0.loc[dr_WF0[dr_muts_col].str.contains(nssnp_match, regex = True, case = False)] dr_gene_WF0_v2 = dr_WF0.loc[dr_WF0[dr_muts_col].str.contains(nssnp_match_captureG, regex = True, case = False)] print('Lengths after tidy split and extracting', nssnp_match, 'muts:' , '\nOld length:' , len(meta_gene_dr) , '\nLength after split:', len(dr_WF0) , '\nLength of nssnp df:', len(dr_gene_WF0) , '\nExpected len:', dr_gene_count , '\n=============================================================') if len(dr_gene_WF0) == dr_gene_count: print('PASS: length matches expected length' , '\n===============================================================') else: sys.exit('FAIL: lengths mismatch') # count the freq of 'dr_muts' samples dr_muts_df = dr_gene_WF0 [['id', dr_muts_col]] print('Dim of dr_muts_df:', dr_muts_df.shape) # add freq column dr_muts_df['dr_sample_freq'] = dr_muts_df.groupby('id')['id'].transform('count') #dr_muts_df['dr_sample_freq'] = dr_muts_df.loc[dr_muts_df.groupby('id')].transform('count') print('Revised dim of dr_muts_df:', dr_muts_df.shape) c1 = dr_muts_df.dr_sample_freq.value_counts() print('Counting no. of sample frequency:\n', c1 , '\n===================================================================') # sanity check: length of gene samples if len(dr_gene_WF0) == c1.sum(): print('PASS: WF data has expected length' , '\nLength of dr_gene WFO:', c1.sum() , '\n===============================================================') else: sys.exit('FAIL: length mismatch!') # Important: Assign 'column name' on which split was performed as an extra column # This is so you can identify if mutations are dr_type or other in the final df dr_df = dr_gene_WF0.assign(mutation_info = dr_muts_col) print('Dim of dr_df:', dr_df.shape , '\n==============================================================' , '\nEnd of tidy split() on dr_muts, and added an extra column relecting mut_category' , '\n===============================================================') ######################################################################## #%% Important sanity checks: other muts column for tidy split(), nsSNPs, etc. # count of occurrence of ";" in other_muts_col: No.of semicolons + 1 is no. of rows created * occurence count_df_other = meta_gene_other[['id', other_muts_col]] count_df_other['other_semicolon_count'] = meta_gene_other.loc[:, other_muts_col].str.count(search, re.I) other_sc_C = count_df_other['other_semicolon_count'].value_counts().reset_index() other_sc_C other_sc_C['index_semicolon'] = (other_sc_C['index']+1)*other_sc_C['other_semicolon_count'] other_sc_C expected_otherC = other_sc_C['index_semicolon'].sum() expected_otherC # count no. of nsSNPs and extract those nsSNPs count_df_other['other_geneSNP_count'] = meta_gene_other.loc[:, other_muts_col].str.count(nssnp_match, re.I) # DEL: need to count by find_all # other_snps_df = meta_gene_other.loc[:, other_muts_col].str.extract(nssnp_match_captureG, re.I) # other_snps_df # other_snps_df.columns = [other_muts_col] # other_snps_countU = other_snps_df[other_muts_col].nunique() # other_snps_countU # other_snps_count_dict = other_snps_df[other_muts_col].value_counts().to_dict() # other_snps_count_dict ###############################################################################