8266511767

Committed 13 Mar 2024 02:36PM UTC coverage: 83.028% (+0.003%) from 83.025%

Build # 8266511767

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1071 from griffithlab/vector_png

Fix blurry pVACvector visualization images

Run Details

1 of 9 new or added lines in 1 file covered. (11.11%)

23 existing lines in 2 files now uncovered.

6027 of 7259 relevant lines covered (83.03%)

4.14 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

11.86

/pvactools/lib/vector_visualization.py

import turtle
import os
import sys
import math
from PIL import Image

class VectorVisualization:
    def __init__(self, input_fasta, output_directory, spacers):

        self.input_fasta = input_fasta
        self.output_directory = output_directory
        self.spacers = spacers

        self.max_pep_length = 100

        self.pep_seqs, self.pep_ids, self.junct_scores = self.parse_input()

        #Error if not a peptide sequence for every peptide ID
        if len(self.pep_ids) != len(self.pep_seqs):
            sys.exit("Error: Not an equal number of peptide sequences and peptide IDs")

        self.conversion_factor = self.get_conversion_factor()
        self.num_peptides = self.get_peptide_num()

        self.turtle = turtle.Turtle()

        self.circle_radius = -200
        self.circle_pos = (-200,0)
        self.header_pos = (0,0)
        self.wht_space_angle = 15
        self.pen_thick = 5
        self.pen_thin = 2
        self.pep_id_space = 45 + self.num_peptides
        self.junct_seq_space = 25

    def parse_input(self):
        pep_seqs = []
        with open(self.input_fasta, 'r') as input_f:
            header = input_f.readline().strip()
            for line in input_f:
                pep_seqs.append(line.strip())
        #remove >, get peptide names and junction scores from FASTA input
        edited_header = header[1:]
        fields = edited_header.split("|")
        pep_ids_joined = fields[0].split(",")
        pep_ids = []
        for pep_id in pep_ids_joined:
            if "." in pep_id:
                desc = pep_id.split(".")
                if len(desc) == 3: #if in expected "MT.GENE.AACHANGE" format, simplify
                    mt, gene, var = pep_id.split(".")
                    pep_id = "-".join((gene,var))
            pep_ids.append(pep_id)
        junct_scores = fields[3].split(":")
        junct_scores = junct_scores[1].split(",")
        return(pep_seqs, pep_ids, junct_scores)

    #determine what proportion of circle each peptide should take up
    def get_conversion_factor(self):
        total_len = 0
        for pep in self.pep_seqs:
            if len(pep) <= self.max_pep_length:
                total_len += len(pep)
            else:
                total_len += 100
        #30 degrees reserved for white space
        conversion_factor = 330 / total_len
        return conversion_factor

    #determine number of peptides (not including junction additions)
    def get_peptide_num(self):
        num_peptides = 0
        for pep in self.pep_seqs:
            length = len(pep)
            if pep not in self.spacers:
                num_peptides += 1
        return(num_peptides)

    def draw(self):
        self.setup_turtle()
        self.draw_header()
        self.reset_turtle()

        angle_parsed = 0
        #add white space in circle before genes
        self.draw_wht_space()
        angle_parsed += self.wht_space_angle
        self.draw_junction()

        #draw main part of circle
        junctions_parsed = 0
        peptides_parsed = 0
        for pep in self.pep_seqs:
            junction_drawn, angle_parsed = self.draw_peptide(pep, peptides_parsed, junctions_parsed, angle_parsed)
            if junction_drawn:
                junctions_parsed += 1
            peptides_parsed += 1

        #add white space in circle after genes
        self.draw_junction()
        self.draw_wht_space()
        self.output_screen()

        #keeps turtle screen open until closed by user
        #turtle.mainloop()

    def setup_turtle(self):
        turtle.setup(1200,600)
        #myWin = turtle.Screen()
        turtle.colormode(255)
        turtle.mode("logo")
        self.turtle.speed(0)
        self.turtle.hideturtle()

    def draw_header(self):
        self.turtle.pu()
        self.turtle.setpos(self.header_pos)
        self.turtle.write("Vector Design", align="center", font=("Arial", 18, "bold"))
        self.turtle.pd()

    def reset_turtle(self):
        self.turtle.pu()
        self.turtle.setpos(self.circle_pos)
        self.turtle.pd()
        self.turtle.pensize(self.pen_thick)

    def draw_wht_space(self):
        self.turtle.pencolor("white")
        self.turtle.circle(self.circle_radius, self.wht_space_angle)

    #draw second line of junctions with amino acid additions
    def draw_junction(self):
        reset = self.turtle.heading()
        self.turtle.rt(90)
        self.turtle.pencolor("black")
        self.turtle.pensize(self.pen_thin)
        self.turtle.forward(10)
        self.turtle.back(20)
        self.turtle.forward(10)
        self.turtle.setheading(reset)

    def draw_peptide(self, peptide, peptide_index, junction_index, angle_parsed):
        junction_drawn = False
        peptide_length = len(peptide)
        peptide_id = self.pep_ids[peptide_index]
        self.turtle.pensize(self.pen_thick)
        angle_parsed += self.conversion_factor * peptide_length
        #if pep is in the list of spacers, draw and label arc for junction
        if peptide in self.spacers:
            self.draw_arc_junct(peptide_id, peptide_length)
            self.draw_junction()
        else:
            #if below reasonable range, draw and label arc for peptide with actual length
            if peptide_length <= self.max_pep_length:
                self.draw_arc_peptide(peptide_id, peptide_length, junction_index, angle_parsed)
            #otherwise, draw and label arc for peptide with 100 length
            else:
                self.draw_arc_peptide(peptide_id, 100, junction_index, angle_parsed)
            if junction_index < len(self.junct_scores):
                self.draw_junction_w_label(self.junct_scores[junction_index], angle_parsed)
                junction_drawn = True
        return (junction_drawn, angle_parsed)

    #draw arc for peptide
    def draw_arc_peptide(self, peptide, length, count, angle):
        self.turtle.pencolor(self.get_color(count))
        self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)
        self.turtle.pu()
        reset = self.turtle.heading()
        self.turtle.left(90)
        self.turtle.forward(self.pep_id_space)
        if (angle > 80 and angle < 100) or (angle > 260 and angle < 280):
            self.turtle.write(peptide, align="center", font=("Arial", 10, "bold"))
        elif (angle > 0 and angle < 90) or (angle > 270 and angle < 360):
            self.turtle.write(peptide, align="right", font=("Arial", 10, "bold"))
        else:
            self.turtle.write(peptide, align="left", font=("Arial", 10, "bold"))
        self.turtle.back(self.pep_id_space)
        self.turtle.setheading(reset)
        self.turtle.pd()
        self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)

    #draw perpindicular line to arc to mark junction
    def draw_junction_w_label(self, junct_score, angle):
        reset = self.turtle.heading()
        self.turtle.rt(90)
        self.turtle.pencolor("black")
        self.turtle.pensize(self.pen_thin)
        self.turtle.forward(10)
        self.turtle.back(20)
        self.turtle.pu()
        if (angle >= 0 and angle < 70):
            self.write_junct_score(junct_score, 15)
        elif (angle >= 65 and angle < 115):
            self.write_junct_score(junct_score, 10)
        elif (angle >= 115 and angle < 165):
            self.write_junct_score(junct_score, 20)
        elif (angle >= 165 and angle < 195):
            self.write_junct_score(junct_score, 25)
        elif (angle >= 195 and angle < 245):
            self.write_junct_score(junct_score, 35)
        elif (angle >= 245 and angle < 295):
            self.write_junct_score(junct_score, 15)
        else:
            self.write_junct_score(junct_score, 20)
        self.turtle.pd()
        self.turtle.forward(10)
        self.turtle.setheading(reset)

    def write_junct_score(self, junct_score, size):
        self.turtle.back(size)
        self.turtle.write("{}nM".format(round(float(junct_score), 5)) , align="center")
        self.turtle.forward(size)

    #draw arc for amino acid inserts to junctions
    def draw_arc_junct(self, peptide, length):
        #t.pencolor("black")
        self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)
        self.turtle.pu()
        reset = self.turtle.heading()
        self.turtle.left(90)
        self.turtle.back(self.junct_seq_space)
        self.turtle.write(peptide, align="center")
        self.turtle.forward(self.junct_seq_space)
        self.turtle.setheading(reset)
        self.turtle.pd()
        self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)

    #print turtle screen to a postscript file, convert to pdf
    def output_screen(self):
        ps_file = os.path.join(self.output_directory, "vector.ps")
        out_file = os.path.join(self.output_directory, "vector.png")
        ts = self.turtle.getscreen()
        ts.getcanvas().postscript(file=ps_file)
        with Image.open(ps_file, formats=["EPS"]) as img:
            original = [float(d) for d in img.size]
            dpi = 300
            scale = dpi / 72.0
            img.load(scale = math.ceil(scale))
            img.thumbnail([round(scale * d) for d in original], Image.Resampling.LANCZOS)
            img.save(out_file, dpi=(300.0, 300.0))
        os.remove(ps_file)

    #select color from scheme
    def get_color(self, count):
        #option 1: 3 blue/green color scheme
        #scheme = [(161,218,180),(65,182,196),(44,127,184),(37,52,148)]
        #option 2: 6 color scheme
        scheme = [(31,120,180),(51,160,44),(227,26,28),(255,127,0),(106,61,154),(177,89,40)]
        count =  count % len(scheme)
        return scheme[count]

1	import turtle	5✔
2	import os	5✔
3	import sys	5✔
4	import math	5✔
5	from PIL import Image	5✔
6
7	class VectorVisualization:	5✔
8	def __init__(self, input_fasta, output_directory, spacers):	5✔
9
10	self.input_fasta = input_fasta	×
11	self.output_directory = output_directory	×
12	self.spacers = spacers	×
13
14	self.max_pep_length = 100	×
15
16	self.pep_seqs, self.pep_ids, self.junct_scores = self.parse_input()	×
17
18	#Error if not a peptide sequence for every peptide ID
19	if len(self.pep_ids) != len(self.pep_seqs):	×
20	sys.exit("Error: Not an equal number of peptide sequences and peptide IDs")	×
21
22	self.conversion_factor = self.get_conversion_factor()	×
23	self.num_peptides = self.get_peptide_num()	×
24
25	self.turtle = turtle.Turtle()	×
26
27	self.circle_radius = -200	×
28	self.circle_pos = (-200,0)	×
29	self.header_pos = (0,0)	×
30	self.wht_space_angle = 15	×
31	self.pen_thick = 5	×
32	self.pen_thin = 2	×
33	self.pep_id_space = 45 + self.num_peptides	×
34	self.junct_seq_space = 25	×
35
36	def parse_input(self):	5✔
37	pep_seqs = []	×
38	with open(self.input_fasta, 'r') as input_f:	×
39	header = input_f.readline().strip()	×
40	for line in input_f:	×
41	pep_seqs.append(line.strip())	×
42	#remove >, get peptide names and junction scores from FASTA input
43	edited_header = header[1:]	×
44	fields = edited_header.split("\|")	×
45	pep_ids_joined = fields[0].split(",")	×
46	pep_ids = []	×
47	for pep_id in pep_ids_joined:	×
48	if "." in pep_id:	×
49	desc = pep_id.split(".")	×
50	if len(desc) == 3: #if in expected "MT.GENE.AACHANGE" format, simplify	×
51	mt, gene, var = pep_id.split(".")	×
52	pep_id = "-".join((gene,var))	×
53	pep_ids.append(pep_id)	×
54	junct_scores = fields[3].split(":")	×
55	junct_scores = junct_scores[1].split(",")	×
56	return(pep_seqs, pep_ids, junct_scores)	×
57
58	#determine what proportion of circle each peptide should take up
59	def get_conversion_factor(self):	5✔
60	total_len = 0	×
61	for pep in self.pep_seqs:	×
62	if len(pep) <= self.max_pep_length:	×
63	total_len += len(pep)	×
64	else:
65	total_len += 100	×
66	#30 degrees reserved for white space
67	conversion_factor = 330 / total_len	×
68	return conversion_factor	×
69
70	#determine number of peptides (not including junction additions)
71	def get_peptide_num(self):	5✔
72	num_peptides = 0	×
73	for pep in self.pep_seqs:	×
74	length = len(pep)	×
75	if pep not in self.spacers:	×
76	num_peptides += 1	×
77	return(num_peptides)	×
78
79	def draw(self):	5✔
80	self.setup_turtle()	×
81	self.draw_header()	×
82	self.reset_turtle()	×
83
84	angle_parsed = 0	×
85	#add white space in circle before genes
86	self.draw_wht_space()	×
87	angle_parsed += self.wht_space_angle	×
88	self.draw_junction()	×
89
90	#draw main part of circle
91	junctions_parsed = 0	×
92	peptides_parsed = 0	×
93	for pep in self.pep_seqs:	×
94	junction_drawn, angle_parsed = self.draw_peptide(pep, peptides_parsed, junctions_parsed, angle_parsed)	×
95	if junction_drawn:	×
96	junctions_parsed += 1	×
97	peptides_parsed += 1	×
98
99	#add white space in circle after genes
100	self.draw_junction()	×
101	self.draw_wht_space()	×
102	self.output_screen()	×
103
104	#keeps turtle screen open until closed by user
105	#turtle.mainloop()
106
107	def setup_turtle(self):	5✔
108	turtle.setup(1200,600)	×
109	#myWin = turtle.Screen()
110	turtle.colormode(255)	×
111	turtle.mode("logo")	×
112	self.turtle.speed(0)	×
113	self.turtle.hideturtle()	×
114
115	def draw_header(self):	5✔
116	self.turtle.pu()	×
117	self.turtle.setpos(self.header_pos)	×
118	self.turtle.write("Vector Design", align="center", font=("Arial", 18, "bold"))	×
119	self.turtle.pd()	×
120
121	def reset_turtle(self):	5✔
122	self.turtle.pu()	×
123	self.turtle.setpos(self.circle_pos)	×
124	self.turtle.pd()	×
125	self.turtle.pensize(self.pen_thick)	×
126
127	def draw_wht_space(self):	5✔
128	self.turtle.pencolor("white")	×
129	self.turtle.circle(self.circle_radius, self.wht_space_angle)	×
130
131	#draw second line of junctions with amino acid additions
132	def draw_junction(self):	5✔
133	reset = self.turtle.heading()	×
134	self.turtle.rt(90)	×
135	self.turtle.pencolor("black")	×
136	self.turtle.pensize(self.pen_thin)	×
137	self.turtle.forward(10)	×
138	self.turtle.back(20)	×
139	self.turtle.forward(10)	×
140	self.turtle.setheading(reset)	×
141
142	def draw_peptide(self, peptide, peptide_index, junction_index, angle_parsed):	5✔
143	junction_drawn = False	×
144	peptide_length = len(peptide)	×
145	peptide_id = self.pep_ids[peptide_index]	×
146	self.turtle.pensize(self.pen_thick)	×
147	angle_parsed += self.conversion_factor * peptide_length	×
148	#if pep is in the list of spacers, draw and label arc for junction
149	if peptide in self.spacers:	×
150	self.draw_arc_junct(peptide_id, peptide_length)	×
151	self.draw_junction()	×
152	else:
153	#if below reasonable range, draw and label arc for peptide with actual length
154	if peptide_length <= self.max_pep_length:	×
155	self.draw_arc_peptide(peptide_id, peptide_length, junction_index, angle_parsed)	×
156	#otherwise, draw and label arc for peptide with 100 length
157	else:
158	self.draw_arc_peptide(peptide_id, 100, junction_index, angle_parsed)	×
159	if junction_index < len(self.junct_scores):	×
160	self.draw_junction_w_label(self.junct_scores[junction_index], angle_parsed)	×
161	junction_drawn = True	×
162	return (junction_drawn, angle_parsed)	×
163
164	#draw arc for peptide
165	def draw_arc_peptide(self, peptide, length, count, angle):	5✔
166	self.turtle.pencolor(self.get_color(count))	×
167	self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)	×
168	self.turtle.pu()	×
169	reset = self.turtle.heading()	×
170	self.turtle.left(90)	×
171	self.turtle.forward(self.pep_id_space)	×
172	if (angle > 80 and angle < 100) or (angle > 260 and angle < 280):	×
173	self.turtle.write(peptide, align="center", font=("Arial", 10, "bold"))	×
174	elif (angle > 0 and angle < 90) or (angle > 270 and angle < 360):	×
175	self.turtle.write(peptide, align="right", font=("Arial", 10, "bold"))	×
176	else:
177	self.turtle.write(peptide, align="left", font=("Arial", 10, "bold"))	×
178	self.turtle.back(self.pep_id_space)	×
179	self.turtle.setheading(reset)	×
180	self.turtle.pd()	×
181	self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)	×
182
183	#draw perpindicular line to arc to mark junction
184	def draw_junction_w_label(self, junct_score, angle):	5✔
185	reset = self.turtle.heading()	×
186	self.turtle.rt(90)	×
187	self.turtle.pencolor("black")	×
188	self.turtle.pensize(self.pen_thin)	×
189	self.turtle.forward(10)	×
190	self.turtle.back(20)	×
191	self.turtle.pu()	×
192	if (angle >= 0 and angle < 70):	×
193	self.write_junct_score(junct_score, 15)	×
194	elif (angle >= 65 and angle < 115):	×
195	self.write_junct_score(junct_score, 10)	×
196	elif (angle >= 115 and angle < 165):	×
197	self.write_junct_score(junct_score, 20)	×
198	elif (angle >= 165 and angle < 195):	×
199	self.write_junct_score(junct_score, 25)	×
200	elif (angle >= 195 and angle < 245):	×
201	self.write_junct_score(junct_score, 35)	×
202	elif (angle >= 245 and angle < 295):	×
203	self.write_junct_score(junct_score, 15)	×
204	else:
205	self.write_junct_score(junct_score, 20)	×
206	self.turtle.pd()	×
207	self.turtle.forward(10)	×
208	self.turtle.setheading(reset)	×
209
210	def write_junct_score(self, junct_score, size):	5✔
211	self.turtle.back(size)	×
212	self.turtle.write("{}nM".format(round(float(junct_score), 5)) , align="center")	×
213	self.turtle.forward(size)	×
214
215	#draw arc for amino acid inserts to junctions
216	def draw_arc_junct(self, peptide, length):	5✔
217	#t.pencolor("black")
218	self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)	×
219	self.turtle.pu()	×
220	reset = self.turtle.heading()	×
221	self.turtle.left(90)	×
222	self.turtle.back(self.junct_seq_space)	×
223	self.turtle.write(peptide, align="center")	×
224	self.turtle.forward(self.junct_seq_space)	×
225	self.turtle.setheading(reset)	×
226	self.turtle.pd()	×
227	self.turtle.circle(self.circle_radius, (self.conversion_factor * length) / 2)	×
228
229	#print turtle screen to a postscript file, convert to pdf
230	def output_screen(self):	5✔
231	ps_file = os.path.join(self.output_directory, "vector.ps")	×
NEW 232	out_file = os.path.join(self.output_directory, "vector.png")	×
233	ts = self.turtle.getscreen()	×
234	ts.getcanvas().postscript(file=ps_file)	×
NEW 235	with Image.open(ps_file, formats=["EPS"]) as img:	×
NEW 236	original = [float(d) for d in img.size]	×
NEW 237	dpi = 300	×
NEW 238	scale = dpi / 72.0	×
NEW 239	img.load(scale = math.ceil(scale))	×
NEW 240	img.thumbnail([round(scale * d) for d in original], Image.Resampling.LANCZOS)	×
NEW 241	img.save(out_file, dpi=(300.0, 300.0))	×
UNCOV 242	os.remove(ps_file)	×
243
244	#select color from scheme
245	def get_color(self, count):	5✔
246	#option 1: 3 blue/green color scheme
247	#scheme = [(161,218,180),(65,182,196),(44,127,184),(37,52,148)]
248	#option 2: 6 color scheme
249	scheme = [(31,120,180),(51,160,44),(227,26,28),(255,127,0),(106,61,154),(177,89,40)]	×
250	count = count % len(scheme)	×
251	return scheme[count]	×

griffithlab / pVACtools / 8266511767

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous