junkcode-for-jplace-analysis/placement_analysis.py at master · JonEilers/junkcode-for-jplace-analysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#!/usr/bin/python3.5

import json
import argparse
import re
import os
import pandas as pd

internal_count = 0
external_count = 0
total_placement_count = 0
totalEdgeCount = 0
leafCount = 0
internalCount = 0


def get_files(root_directory, extension='.jplace'):
    filtered_files = []
    extension_length = len(extension)
    for root, subdirList, files in os.walk(os.path.abspath(root_directory)):
        for name in files:
            if name[-extension_length:] == extension:
                filtered_files.append(os.path.join(root, name))
    return filtered_files


def tree_splitter(jplace):
    tree = jplace["tree"]
    branches = re.split('[, ( )]',tree)  ##splitting at ')' '(' and ',' eliminates all delimiters that are not curly braces--
    # leaves either empty elements ex: u'' or elements withedge numbers/labels/lengths in each element of the list
    return branches


def edge_counter(split_tree):
    global totalEdgeCount, leafCount, internalCount
    totalEdgeCount -= 1
    leafEdges = []
    internalEdges = []
    branches = tree_splitter(split_tree)
    # i is integer index of each element in list
    # v is value (string) at each index
    for i, v in enumerate(branches):
        if '{' in v:
            totalEdgeCount += 1
        if "|" in v:
            leafCount += 1
            leaf_edgeNum = int(v.split('{')[1].split('}')[0])
            leafEdges.append(leaf_edgeNum)
        elif v != '':
            internal_edgeNum = int(v.split('{')[1].split('}')[0])
            internalCount += 1
            internalEdges.append(internal_edgeNum)
    return {"internalEdges": internalEdges, "internalCount": internalCount, "leafCount": leafCount,
            "leafEdges": leafEdges}


def number_of_placements(file):
    global total_placement_count
    placements = file['placements']
    total_placement_count += len(placements)
    return total_placement_count

def edge_indice(file):
    fields = file["fields"]
    for index, items in enumerate(fields):
        if items == "edge_num":
            return index

def placement_location(file):
    internal_edge_list = edge_counter(file)["internalEdges"]
    leaf_edge_list = edge_counter(file)["leafEdges"]
    placements = file['placements']
    edge_index = edge_indice(file)
    for i in placements:
        placement_edge = i["p"][0][edge_index]
        if placement_edge in internal_edge_list:
            global internal_count
            internal_count += 1
        elif placement_edge in leaf_edge_list:
            global external_count
            external_count += 1
    return internal_count, external_count


def internal_vs_leaf(dir, out_file):
    jplace_files = get_files(dir)
    for file in jplace_files:
        with open(file) as json_data:
            jplace = json.load(json_data)
        placement_count_total = number_of_placements(jplace)
        placement_int_vs_ext = placement_location(jplace)
    ivl_dict = {'Leaf Count': external_count, 'Internal Count':internal_count, 'Total Number of Placements':total_placement_count}
    ivl_series = pd.Series(ivl_dict)
    output = ivl_series.to_csv(out_file)
    return output

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Count the number of internal placements vs leaf placements on a phylogenetic tree. Takes .jplace files")
    parser.add_argument('-directory', help = 'directory with .jplace files', required = True)
    parser.add_argument('-out_file', help = 'output file (txt formart)', required = True)
    args = parser.parse_args()

    internal_vs_leaf(dir = args.directory, out_file=args.out_file)