Merge pull request #315 from mnneveau/vaccine_visualization

Vaccine Visualization

.travis.yml

@@ -1,8 +1,19 @@
 language: python
 python:
   - "3.5"
+addons:
+  apt:
+   packages:
+     - xvfb
+before_install:
+  - sudo apt-get -qq update
+  - sudo apt-get install -y ghostscript
 install:
   - pip install -e .
+before_script:
+  - "export DISPLAY=:99.0"
+  - "sh -e /etc/init.d/xvfb start"
+  - sleep 3 
 script:
   - TEST_FLAG=1 python3 -m unittest discover -v
   - prove --recurse --verbose

pvacseq/lib/vaccine_design.py

@@ -21,6 +21,7 @@
 
 from lib.prediction_class import *
 
+import turtle
 
 def define_parser():
 
@@ -158,6 +159,234 @@ def anneal(self):
         # Return best state and energy
         return self.best_state, self.best_energy
 
+def parse_input(input_file):
+    pep_seqs = []
+    with open(input_file, '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:
+            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 number of peptides (not including junction additions)
+def get_peptide_num(pep_seqs, min_pep_length, max_pep_length):    
+    num_peptides = 0
+    for pep in pep_seqs:
+        length = len(pep)
+        if length > min_pep_length and length < max_pep_length:
+            num_peptides += 1
+    return(num_peptides)
+
+#determine what proportion of circle each peptide should take up
+def get_conversion_factor(pep_seqs):
+    total_len = 0
+    for pep in pep_seqs:
+        total_len += len(pep)   
+    #30 degrees reserved for white space
+    conversion_factor = 330 / total_len
+    return(conversion_factor)
+
+def draw_header(t, header_pos):
+    t.pu()
+    t.setpos(header_pos)
+    t.write("Vaccine Design", align="center", font=("Arial", 18, "bold"))
+    t.pd()
+    return()
+
+def draw_wht_space(t, circle_radius, wht_space_angle):
+    t.pencolor("white")
+    t.circle(circle_radius,wht_space_angle)
+    return()
+
+#select color from scheme
+def get_color(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]\
+
+def write_junct_score(t, junct_score, size):
+    t.back(size)
+    t.write(junct_score + 'nM', align="center")
+    t.forward(size)
+
+#draw perpindicular line to arc to mark junction
+def draw_junction_w_label(junct_score, t, pen_thin, angle):
+    reset = t.heading()
+    t.rt(90)
+    t.pencolor("black")
+    t.pensize(pen_thin)
+    t.forward(10)
+    t.back(20)
+    t.pu()
+    if (angle >= 0 and angle < 70):
+        write_junct_score(t, junct_score, 15)
+    elif (angle >= 65 and angle < 115):
+        write_junct_score(t, junct_score, 10)
+    elif (angle >= 115 and angle < 165):
+        write_junct_score(t, junct_score, 20)
+    elif (angle >= 165 and angle < 195):
+        write_junct_score(t, junct_score, 25)
+    elif (angle >= 195 and angle < 245):
+        write_junct_score(t, junct_score, 35)
+    elif (angle >= 245 and angle < 295):
+        write_junct_score(t, junct_score, 15)
+    else:
+        write_junct_score(t, junct_score, 20)
+    t.pd()
+    t.forward(10)
+    t.setheading(reset)
+    return()
+
+#draw second line of junctions with amino acid additions
+def draw_junction(t, pen_thin):
+    reset = t.heading()
+    t.rt(90)
+    t.pencolor("black")
+    t.pensize(pen_thin)
+    t.forward(10)
+    t.back(20)
+    t.forward(10)
+    t.setheading(reset)
+    return()
+
+#draw arc for peptide
+def draw_arc_peptide(peptide, length, count, angle, t, circle_radius, conversion_factor, pep_id_space):
+    t.pencolor(get_color(count))
+    t.circle(circle_radius, (conversion_factor * length) / 2)
+    t.pu()
+    reset = t.heading()
+    t.left(90)
+    t.forward(pep_id_space)
+    if (angle > 80 and angle < 100) or (angle > 260 and angle < 280):
+        t.write(peptide, align="center", font=("Arial", 10, "bold"))
+    elif (angle > 0 and angle < 90) or (angle > 270 and angle < 360):
+        t.write(peptide, align="right", font=("Arial", 10, "bold"))
+    else:
+        t.write(peptide, align="left", font=("Arial", 10, "bold"))
+    t.back(pep_id_space)
+    t.setheading(reset)
+    t.pd()
+    t.circle(circle_radius, (conversion_factor * length) / 2)
+
+#draw arc for amino acid inserts to junctions
+def draw_arc_junct(peptide, length, t, conversion_factor, junct_seq_space, circle_radius):
+    #t.pencolor("black")
+    t.circle(circle_radius, (conversion_factor * length) / 2)
+    t.pu()
+    reset = t.heading()
+    t.left(90)
+    t.back(junct_seq_space)
+    t.write(peptide, align="center")
+    t.forward(junct_seq_space)
+    t.setheading(reset)
+    t.pd()
+    t.circle(circle_radius, (conversion_factor * length) / 2)
+
+def draw_peptide(t, pep, pep_ids, peptides_parsed, pen_thick, angle_parsed, conversion_factor, min_pep_length, max_pep_length, junctions_parsed, junct_scores, circle_radius, pep_id_space, junct_seq_space, pen_thin):
+    junction_parsed = 0
+    pep_length = len(pep)
+    peptide = pep_ids[peptides_parsed]
+    t.pensize(pen_thick)
+    angle_parsed += conversion_factor * pep_length
+    #if length within reasonable range, draw and label arc for peptide
+    if pep_length > min_pep_length and pep_length < max_pep_length:
+        draw_arc_peptide(peptide, pep_length, junctions_parsed, angle_parsed, t, circle_radius, conversion_factor, pep_id_space)
+        if junctions_parsed < len(junct_scores):
+            draw_junction_w_label(junct_scores[junctions_parsed], t, pen_thin, angle_parsed)
+            junction_parsed += 1
+    #if length is less than minimum peptide length, assume amino acid addition to junction
+    elif pep_length < min_pep_length:
+        draw_arc_junct(peptide, pep_length, t, conversion_factor, junct_seq_space, circle_radius)
+        draw_junction(t, pen_thin)
+    else:
+        print("Error: Peptide sequence over 100 amino acids inputted")
+        sys.exit()
+    return(junction_parsed, angle_parsed)
+
+#print turtle screen to a postscript file, convert to pdf
+def output_screen(t, out_f):
+    ps_file = "/".join((out_f, "vaccine.ps"))
+    out_file = "/".join((out_f, "vaccine.jpg"))
+    ts = t.getscreen()
+    ts.getcanvas().postscript(file=ps_file)
+    os.system('convert -density 300 -quality 100 ' + ps_file + " " + out_file)
+    os.system('rm ' + ps_file)
+    return()
+
+def output_vaccine_png(input_file, out_f):
+    min_pep_length = 8
+    max_pep_length = 100
+
+    pep_seqs, pep_ids, junct_scores = parse_input(input_file)
+
+    #Error if not a peptide sequence for every peptide ID
+    if len(pep_ids) != len(pep_seqs):
+        print("Error: Not an equal number of peptide sequences and peptide IDs")
+        sys.exit()
+
+    conversion_factor = get_conversion_factor(pep_seqs)
+    num_peptides = get_peptide_num(pep_seqs, min_pep_length, max_pep_length)
+
+    #draw vaccine
+    #800,600
+    turtle.setup(800,600)
+    t = turtle.Turtle()
+    myWin = turtle.Screen()
+    turtle.colormode(255)
+    turtle.mode("logo")
+    t.speed(0)
+    t.hideturtle()
+
+    #negative radius draws circle clockwise
+    circle_radius = -200
+    circle_pos = (-200,0)
+    header_pos = (0,0)
+    wht_space_angle = 15
+    pen_thick = 5
+    pen_thin = 2
+    pep_id_space = 45 + num_peptides
+    junct_seq_space = 25
+
+    draw_header(t, header_pos)
+    t.pu()
+    t.setpos(circle_pos)
+    t.pd()
+    t.pensize(pen_thick)
+    angle_parsed = 0
+    #add white space in circle before genes
+    draw_wht_space(t, circle_radius, wht_space_angle)
+    angle_parsed += wht_space_angle
+    draw_junction(t, pen_thin)
+
+    #draw main part of circle
+    junctions_parsed = 0
+    peptides_parsed = 0
+    for pep in pep_seqs:
+        junction_parsed, angle_parsed = draw_peptide(t, pep, pep_ids, peptides_parsed, pen_thick, angle_parsed, conversion_factor, min_pep_length, max_pep_length, junctions_parsed, junct_scores, circle_radius, pep_id_space, junct_seq_space, pen_thin)
+        junctions_parsed += junction_parsed 
+        peptides_parsed += 1
+
+    #add white space in circle after genes    
+    draw_junction(t, pen_thin)
+    draw_wht_space(t, circle_radius, wht_space_angle)
+    output_screen(t, out_f)
+
+    #keeps turtle screen open until closed by user
+    #turtle.mainloop()
 
 def main(args_input=sys.argv[1:]):
 
@@ -383,6 +612,10 @@ def main(args_input=sys.argv[1:]):
     if not args.keep_tmp:
         shutil.rmtree(tmp_dir)
 
+    out_f = os.path.join(base_output_dir, runname) 
+
+    output_vaccine_png(results_file, out_f)
+
 if __name__ == "__main__":
     main()
 

tests/test_vaccine_design.py

@@ -47,5 +47,10 @@ def test_vaccine_design_runs_and_produces_expected_output(self):
             os.path.join(self.test_data_dir, "Test.vaccine.results.output.fa")
         ))
 
+        image_out = os.path.join(output_dir.name,self.test_run_name, 'vaccine.jpg')
+
+        self.assertTrue(os.path.exists(image_out))
+        self.assertTrue(os.stat(image_out).st_size > 0)
+
         output_dir.cleanup()