-
Notifications
You must be signed in to change notification settings - Fork 0
/
write_col_aa.py
59 lines (49 loc) · 1.62 KB
/
write_col_aa.py
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
#!/usr/bin/env python
from Bio import SeqIO
import sys
import os
# Takes in two argument from command line
# fasta file name, and index of interests
def main():
GET_input = sys.argv[1]
index = int(sys.argv[2])
filepath = "fasta/" + GET_input + ".fasta"
# parse in all the file
i = 0
mapping = dict()
for seq_record in SeqIO.parse(filepath, "fasta"):
# Skipping the first one, the consensus
if i == 0:
i += 1
continue
curr_id = seq_record.id
color_coding = "" # To be filled
# Switch from index 1 to index 0
curr_aa = seq_record.seq[index - 1] + color_coding
# Assuming ID has form 1V##X### (# are numbers)
curr_version = curr_id[2: 4]
if curr_version not in mapping:
mapping[curr_version] = []
mapping[curr_version].append(curr_aa)
write_to_json(mapping, GET_input, index)
def write_to_json(mapping, GET_input, index):
if not os.path.exists("output/" + GET_input):
os.makedirs("output/" + GET_input)
f = open("output/" + GET_input + "/" + GET_input + "-" + str(index) + ".json", "wb")
i = 0
f.write("[\n")
for key, value in sorted(mapping.iteritems()):
if i != len(mapping) - 1:
f.write('\t{ "V": "' + str(key) + '", "aa": ' + to_string(value) + ' },\n')
else:
f.write('\t{ "V": "' + str(key) + '", "aa": ' + to_string(value) + ' }\n')
i += 1
f.write("]")
f.close()
def to_string(l):
res = '["' + l[0] + '"'
for i in range(1, len(l)):
res += ', "' + l[i] + '"'
res += "]"
return res
main()