import dendropy
from dendropy.model.parsimony import fitch_down_pass
from dendropy.model.parsimony import fitch_up_pass
import pandas
import csv
import numpy
import sys
import re
from Bio import Phylo
from Bio import SeqIO
from statistics import mean
from statistics import stdev
from scipy.stats import wilcoxon
from scipy.stats import mannwhitneyu
import bisect
import operator
import argparse
import time
import os
###############################################################################
## Load up the trees ##########################################################
###############################################################################
def get_options():
purpose = '''This is a python script to output each branch's mutation rate.
Usage: node_balance_and_mutation_get.py <gubbins_dir> <hit_locs_csv> <out_csv> '''
parser = argparse.ArgumentParser(description=purpose, prog='node_mutation.py')
parser.add_argument('--gubbins_res', required=True, help='Directory where all cluster dirs of gubbins res are stored"', type=str)
parser.add_argument('--hit_locs_csv', required=True, help='Hit csv file from hit allocator', type=str)
parser.add_argument('--out_name', required=True, help='Prefix to append to out out_files', type=str)
args = parser.parse_args()
return args
def search_function(value, df, col_name):
current_col = df[col_name]
indexio = current_col.loc[current_col == value].index
if len(indexio) > 1:
indexio = indexio[0]
if indexio.size > 0:
row_indy = indexio
return row_indy
else:
return "Not_here"
def find_nth(haystack, needle, n):
start = haystack.find(needle)
while start >= 0 and n > 1:
start = haystack.find(needle, start + len(needle))
n -= 1
return start
def node_reconstruct(tree_loc, hit_csv):
## Function to reconstruct the insertion node of the individual cluster for a specific gps cluster
tree = dendropy.Tree.get(path=tree_loc,
schema="newick", preserve_underscores=True)
## hit df ##
cluster_csv = hit_csv.copy()
cluster_csv = cluster_csv.reset_index(drop=True)
# print(tree.taxon_namespace)
## Now we'll create the mapping of the tree isolates to the cluster nums, with 0 being no hit.
cluster_num_col = [1]
tree_names = ["start"]
for taxon in tree.taxon_namespace:
current_name = taxon.label
current_name_string = str(current_name)
indy = search_function(current_name_string, cluster_csv, 'id')
if isinstance(indy, str):
multi_hit_name = current_name_string + "_1"
indy_2 = search_function(multi_hit_name, cluster_csv, 'id')
if isinstance(indy_2, str):
cluster_num_col.append(0)
else:
cluster_val = cluster_csv.iloc[indy_2, cluster_csv.columns.get_loc("insert_name")]
cluster_num_col.append(int(cluster_val))
else:
cluster_val = cluster_csv.iloc[indy.values[0], cluster_csv.columns.get_loc("insert_name")]
cluster_num_col.append(int(cluster_val))
tree_names.append(current_name_string)
node_changer = tree.find_node_with_taxon_label(label=current_name)
node_changer.taxon.label = dendropy.Taxon(label=current_name_string)
cluster_num_col = cluster_num_col[1:]
tree_names = tree_names[1:]
## Heres the dictionary with our insertion numbers and the tree tip ids
mega_characters = dict(zip(tree_names, cluster_num_col))
with open("./fasta_data.tsv", 'w') as outfile:
csv_writer = csv.writer(outfile, delimiter='\t', quotechar='|', quoting=csv.QUOTE_MINIMAL)
for k, v in mega_characters.items():
csv_writer.writerow([k] + [v])
SeqIO.convert("./fasta_data.tsv",
"tab", "./fasta_data.fasta",
"fasta")
taxa = dendropy.TaxonNamespace()
data_mega = dendropy.StandardCharacterMatrix.get_from_path("./fasta_data.fasta",
"fasta", taxon_namespace=taxa)
taxon_state_sets_map = data_mega.taxon_state_sets_map(gaps_as_missing=True)
tree = dendropy.Tree.get_from_path(tree_loc,
schema="newick", preserve_underscores=True, taxon_namespace=taxa)
os.remove("./fasta_data.tsv")
os.remove("./fasta_data.fasta")
score = fitch_down_pass(postorder_nodes=tree.postorder_node_iter(),
taxon_state_sets_map=taxon_state_sets_map)
fitch_up_pass(tree.postorder_node_iter())
###########################################################################
## Ok so now the tree has been reconstructed with the fitch up pass down ##
## pass algorithm, we'll return the node labelled tree #################
###########################################################################
return tree
def branch_mutations(tree, embl_csv, embl_reccy):
## Function to work out all the mutations and their positions along a branch.
## The embl csv gives you total mutations, both in and outside recombinations, so we first
## calculate the total, then to get just the base substitutions we check if any recombinations
## along the branches and remove those snps inside these events from the total to get the clonal or (outside)
## mutations.
total_isolates = []
for taxon in tree.leaf_node_iter():
leaf_label = taxon.taxon
# leaf_label = re.sub("\"","", leaf_label)
total_isolates.append(leaf_label)
start_id = []
isolates_ids = []
start_node = []
finish_node = []
tag = []
AT = []
AC = []
AG = []
TA = []
TC = []
TG = []
CA = []
CT = []
CG = []
GA = []
GT = []
GC = []
AT_outside = []
AC_outside = []
AG_outside = []
TA_outside = []
TC_outside = []
TG_outside = []
CA_outside = []
CT_outside = []
CG_outside = []
GA_outside = []
GT_outside = []
GC_outside = []
cluster_num = []
node_chain_len = []
end_nodes = []
particular_end_node = []
node_already_tested = []
edge_lengths = []
for isolate in range(len(total_isolates)):
## This loops through all the isolates to get all the branches
##
current_row = total_isolates[isolate]
isolate_test = str(current_row)
isolate_test = re.sub("\'", "", isolate_test)
nodes_in_insertion = []
edge_length = []
## First part of the loop establishes the chain of nodes leading to the insertion
## node for the element.
isolate_id = str(current_row)
current_row = re.sub("\'", "", isolate_id)
nodes_in_insertion.append(isolate_test)
tree_node = tree.find_node_with_taxon_label(label=current_row)
# print(tree_node)
edge_length.append(tree_node.edge.length)
parent_node = tree_node.parent_node
while parent_node != None:
new_parent = parent_node.parent_node
old_parent = parent_node
parent_node = new_parent
nodes_in_insertion.append(old_parent.label)
edge_length.append(old_parent.edge.length)
if parent_node == None:
end_node = old_parent.label
if end_node not in end_nodes:
end_nodes.append(end_node)
for nodeys in range(len(nodes_in_insertion) - 1):
current_end_id = nodes_in_insertion[nodeys]
if current_end_id not in node_already_tested:
source_node = nodes_in_insertion[nodeys + 1]
target_node = current_end_id
start_node.append(source_node)
finish_node.append(target_node)
tag.append("No")
subset_embl_csv = embl_csv[embl_csv['end_node'] == current_end_id]
A_T = 0
A_C = 0
A_G = 0
T_A = 0
T_C = 0
T_G = 0
C_A = 0
C_T = 0
C_G = 0
G_A = 0
G_T = 0
G_C = 0
A_T_outside = 0
A_C_outside = 0
A_G_outside = 0
T_A_outside = 0
T_C_outside = 0
T_G_outside = 0
C_A_outside = 0
C_T_outside = 0
C_G_outside = 0
G_A_outside = 0
G_T_outside = 0
G_C_outside = 0
if subset_embl_csv.empty:
A_T += 0
A_C += 0
A_G += 0
T_A += 0
T_C += 0
T_G += 0
C_A += 0
C_T += 0
C_G += 0
G_A += 0
G_T += 0
G_C += 0
A_T_outside += 0
A_C_outside += 0
A_G_outside += 0
T_A_outside += 0
T_C_outside += 0
T_G_outside += 0
C_A_outside += 0
C_T_outside += 0
C_G_outside += 0
G_A_outside += 0
G_T_outside += 0
G_C_outside += 0
else:
A_to_T = len(subset_embl_csv[(subset_embl_csv['start_base'] == "A") & (
subset_embl_csv['end_base'] == "T")].index)
A_to_G = len(subset_embl_csv[(subset_embl_csv['start_base'] == "A") & (
subset_embl_csv['end_base'] == "G")].index)
A_to_C = len(subset_embl_csv[(subset_embl_csv['start_base'] == "A") & (
subset_embl_csv['end_base'] == "C")].index)
T_to_A = len(subset_embl_csv[(subset_embl_csv['start_base'] == "T") & (
subset_embl_csv['end_base'] == "A")].index)
T_to_G = len(subset_embl_csv[(subset_embl_csv['start_base'] == "T") & (
subset_embl_csv['end_base'] == "G")].index)
T_to_C = len(subset_embl_csv[(subset_embl_csv['start_base'] == "T") & (
subset_embl_csv['end_base'] == "C")].index)
G_to_A = len(subset_embl_csv[(subset_embl_csv['start_base'] == "G") & (
subset_embl_csv['end_base'] == "A")].index)
G_to_T = len(subset_embl_csv[(subset_embl_csv['start_base'] == "G") & (
subset_embl_csv['end_base'] == "T")].index)
G_to_C = len(subset_embl_csv[(subset_embl_csv['start_base'] == "G") & (
subset_embl_csv['end_base'] == "C")].index)
C_to_A = len(subset_embl_csv[(subset_embl_csv['start_base'] == "C") & (
subset_embl_csv['end_base'] == "A")].index)
C_to_T = len(subset_embl_csv[(subset_embl_csv['start_base'] == "C") & (
subset_embl_csv['end_base'] == "T")].index)
C_to_G = len(subset_embl_csv[(subset_embl_csv['start_base'] == "C") & (
subset_embl_csv['end_base'] == "G")].index)
A_T += (A_to_T)
A_C += (A_to_C)
A_G += (A_to_G)
T_A += (T_to_A)
T_C += (T_to_C)
T_G += (T_to_G)
C_A += (C_to_A)
C_T += (C_to_T)
C_G += (C_to_G)
G_A += (G_to_A)
G_T += (G_to_T)
G_C += (G_to_C)
subset_reccy_csv = embl_reccy[embl_reccy['end_node'] == current_end_id]
if subset_reccy_csv.empty:
A_T_outside += A_to_T
A_C_outside += A_to_C
A_G_outside += A_to_G
T_A_outside += T_to_A
T_C_outside += T_to_C
T_G_outside += T_to_G
C_A_outside += C_to_A
C_T_outside += C_to_T
C_G_outside += C_to_G
G_A_outside += G_to_A
G_T_outside += G_to_T
G_C_outside += G_to_C
else:
for reccy_row in range(len(subset_reccy_csv.index)):
current_start_reccy = subset_reccy_csv.iloc[reccy_row, 2]
current_end_reccy = subset_reccy_csv.iloc[reccy_row, 3]
for mut_row in range(len(subset_embl_csv.index)):
base_pos = subset_embl_csv.iloc[mut_row, 4]
start_base = subset_embl_csv.iloc[mut_row, 2]
end_base = subset_embl_csv.iloc[mut_row, 3]
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "A" and end_base == "T":
A_to_T -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "A" and end_base == "C":
A_to_C -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "A" and end_base == "G":
A_to_G -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "T" and end_base == "A":
T_to_A -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "T" and end_base == "C":
T_to_C -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "T" and end_base == "G":
T_to_G -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "C" and end_base == "A":
C_to_A -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "C" and end_base == "T":
C_to_T -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "C" and end_base == "G":
C_to_G -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "G" and end_base == "A":
G_to_A -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "G" and end_base == "T":
G_to_T -= 1
if base_pos >= current_start_reccy and base_pos <= current_end_reccy and start_base == "G" and end_base == "C":
G_to_C -= 1
A_T_outside += A_to_T
A_C_outside += A_to_C
A_G_outside += A_to_G
T_A_outside += T_to_A
T_C_outside += T_to_C
T_G_outside += T_to_G
C_A_outside += C_to_A
C_T_outside += C_to_T
C_G_outside += C_to_G
G_A_outside += G_to_A
G_T_outside += G_to_T
G_C_outside += G_to_C
AT.append(A_T)
AC.append(A_C)
AG.append(A_G)
TA.append(T_A)
TC.append(T_C)
TG.append(T_G)
CA.append(C_A)
CT.append(C_T)
CG.append(C_G)
GA.append(G_A)
GT.append(G_T)
GC.append(G_C)
AT_outside.append(A_T_outside)
AC_outside.append(A_C_outside)
AG_outside.append(A_G_outside)
TA_outside.append(T_A_outside)
TC_outside.append(T_C_outside)
TG_outside.append(T_G_outside)
CA_outside.append(C_A_outside)
CT_outside.append(C_T_outside)
CG_outside.append(C_G_outside)
GA_outside.append(G_A_outside)
GT_outside.append(G_T_outside)
GC_outside.append(G_C_outside)
edge_lengths.append(edge_length[nodeys])
node_already_tested.append(current_end_id)
start_id.append(current_row)
if isolate % 10 == 0:
print(isolate / len(total_isolates) * 100)
non_mge_mutations_out = pandas.DataFrame({'start_node': start_node,
'end_node': finish_node,
'tag': tag,
'A-T': AT,
'A-C': AC,
'A-G': AG,
'T-A': TA,
'T-C': TC,
'T-G': TG,
'C-A': CA,
'C-T': CT,
'C-G': CG,
'G-A': GA,
'G-T': GT,
'G-C': GC,
'A-T_clonal': AT_outside,
'A-C_clonal': AC_outside,
'A-G_clonal': AG_outside,
'T-A_clonal': TA_outside,
'T-C_clonal': TC_outside,
'T-G_clonal': TG_outside,
'C-A_clonal': CA_outside,
'C-T_clonal': CT_outside,
'C-G_clonal': CG_outside,
'G-A_clonal': GA_outside,
'G-T_clonal': GT_outside,
'G-C_clonal': GC_outside,
'branch_lengths': edge_lengths,
'starting_isolate': start_id})
return non_mge_mutations_out
if __name__ == '__main__':
start_overall = time.perf_counter()
input_args = get_options()
###############################################################################
## Load up the csv and then run through each of the clusters to check through #
## the res ####################################################################
###############################################################################
cluster_csv = pandas.read_csv(input_args.hit_locs_csv)
base_loc = input_args.gubbins_res
unique_clusters = cluster_csv['cluster_name'].unique()
tot_reccy_csv = pandas.DataFrame()
tot_non_reccy = pandas.DataFrame()
seq_clus = 1
for cluster in unique_clusters:
print("On cluster: %s, %s of %s" % (cluster, seq_clus, len(unique_clusters)))
tic_cluster = time.perf_counter()
current_dat = cluster_csv[cluster_csv['cluster_name'] == cluster]
current_dir = base_loc + cluster
try:
cluster_files = os.listdir(current_dir)
except:
current_dir = current_dir + "_run_data"
cluster_files = os.listdir(current_dir)
tree_indexio = [k for k, s in enumerate(cluster_files) if "node_labelled.final_tree.tre" in s]
embl_branch = [k for k, s in enumerate(cluster_files) if "_branch_base.csv" in s]
embl_reccy = [k for k, s in enumerate(cluster_files) if "_recombinations.csv" in s]
per_branch_file = [k for k, s in enumerate(cluster_files) if "_per_branch_mutations.csv" in s]
if len(per_branch_file) > 0:
print("Using already formed per_branch_mutations.csv")
continue
tree_loc = current_dir + "/" + cluster_files[tree_indexio[0]]
embl_branch_loc = current_dir + "/" + cluster_files[embl_branch[0]]
embl_rec_loc = current_dir + "/" + cluster_files[embl_reccy[0]]
embl_csv = pandas.read_csv(embl_branch_loc)
embl_reccy_csv = pandas.read_csv(embl_rec_loc)
## So now we've got all the files we need for this particular cluster, we'll run through
## the tree to get the nodes labelled with the inserts. Then we'll run through the
## branches and get the summary of the mutations present
#tree = node_reconstruct(tree_loc, current_dat)
tree = dendropy.Tree.get(path=tree_loc,
schema="newick", preserve_underscores=True)
branches_csv = branch_mutations(tree,embl_csv, embl_reccy_csv)
branches_csv['cluster_name'] = cluster
branches_out_loc = current_dir + "/" + cluster + "_per_branch_mutations.csv"
branches_csv.to_csv(path_or_buf=branches_out_loc, index=False)
seq_clus += 1
toc_tot = time.perf_counter()
print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
print("Mutation finder took: %s (seconds)" % (toc_tot - start_overall))
print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Branch mutations found ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
