forked from wuami/draw_rna
-
Notifications
You must be signed in to change notification settings - Fork 1
/
draw_rna.py
114 lines (99 loc) · 3.77 KB
/
draw_rna.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
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
104
105
106
107
108
109
110
111
112
113
114
from __future__ import absolute_import, division, print_function
from . import render_rna
from . import svg
from . import inv_utils
import argparse
import re
from matplotlib import cm
NODE_R = 10
PRIMARY_SPACE = 20
PAIR_SPACE = 23
CELL_PADDING = 40
TEXT_SIZE = 50
RENDER_IN_LETTERS = True
COLORS = {#"r": [255, 0, 0],
"r": [255, 102, 102],
"g": [113, 188, 120],
"b": [51, 153, 255],
#"b": [0, 153, 255],
"k": [1, 0, 0],
"y": [255, 211, 0],
"c": [0, 255, 255],
"m": [255, 0, 255],
"w": [255, 255, 255],
"e": [100, 100, 100],
"o": [231, 115, 0],
"i": [51, 204, 204],
"h": [51, 153, 255]}
#"i": [0, 204, 153],
#"h": [46, 184, 46]}
def draw_rna(sequence, secstruct, colors, filename="secstruct", line=False, square=True, flipped=True, draw_pairs=True):
r = render_rna.RNARenderer()
pairmap = render_rna.get_pairmap_from_secstruct(secstruct)
pairs = []
for i in range(len(pairmap)):
if pairmap[i] > i:
pairs.append({"from":i, "to":pairmap[i], "p":1.0, "color":COLORS["e"]})
r.setup_tree(secstruct, NODE_R, PRIMARY_SPACE, PAIR_SPACE, False)
size = r.get_size()
cell_size = max(size) + CELL_PADDING * 2
# if colors are numeric, create color scale
try:
colors = [float(x) for x in colors.split()]
gist_earth = cm.get_cmap('gist_earth_r')
min_ = min(colors)
range_ = max(colors) - min_
colors = [gist_earth((x - min_)/range_)[:-1] for x in colors]
colors = [[c * 256 for c in color] for color in colors]
# otherwise, colors are indexes to COLORS dict
except:
colors = [COLORS[x] for x in list(colors)]
width = cell_size if square else size[0]
height = cell_size if square else size[1]
svgobj = svg.svg("%s.svg" % filename, width, height)
r.draw(svgobj, CELL_PADDING, CELL_PADDING, colors, pairs if draw_pairs else None, sequence, RENDER_IN_LETTERS, line)
def parse_colors(color_string):
colorings = color_string.strip().split(",")
colors = []
for coloring in colorings:
if "x" in coloring:
[n, color] = coloring.split("x")
colors += int(n) * [color]
else:
colors += [coloring]
return colors
def reorder_strands(order, seq, colors):
breaks = [-1] + [m.start() for m in re.finditer("&", seq)] + [len(seq)]
seq_segments = seq.split("&")
color_segments = [colors[breaks[i]+1:breaks[i+1]] for i in range(len(breaks)-1)]
seq = ""
colors = []
for strand in order:
seq += seq_segments[strand-1] + "&"
colors += color_segments[strand-1] + ["w"]
return [seq[:-1], colors[:-1]]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("filename", help="filename specifying sequences to fold", type=str)
parser.add_argument("-f", "--fold", help="automatic folding", type=str)
args = parser.parse_args()
with open(args.filename) as f:
n = int(f.readline())
for i in range(n):
seq = f.readline().strip()
if args.fold == "nupack":
result = inv_utils.nupack_fold(seq, 1e-7)
secstruct = result[0]
colors = parse_colors(f.readline())
seq, colors = reorder_strands(result[2], seq, colors)
elif args.fold == "vienna":
secstruct = inv_utils.vienna_fold(seq)[0]
colors = parse_colors(f.readline())
else:
secstruct = f.readline().strip()
colors = parse_colors(f.readline())
seq.replace("&", "")
secstruct.replace("&", "")
draw_rna(seq, secstruct, colors, "%s_%d" % (args.filename, i))
if __name__ == "__main__":
main()