Merge pull request #91 from martinghunt/phandango_colours

Phandango colours
sanger-pathogens · May 27, 2016 · 24c1521 · 24c1521
2 parents 4b2c976 + b907d4a
commit 24c1521
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 1 deletion.
diff --git a/ariba/summary.py b/ariba/summary.py
@@ -1,4 +1,5 @@
 import os
+import copy
 import re
 import sys
 import openpyxl
@@ -257,6 +258,31 @@ def _filter_matrix_columns(cls, matrix, phandango_header, csv_header):
         return phandango_header, csv_header, matrix
 
 
+    @classmethod
+    def _add_phandango_colour_columns(cls, header, matrix):
+        header = copy.deepcopy(header)
+        matrix = copy.deepcopy(matrix)
+        cols_to_add_colour_col = [i for i in range(len(header)) if header[i].endswith(':o1')]
+        field_to_col = {
+            'yes': '#1f78b4',
+            'yes_nonunique': '#a6cee3',
+            'no': '#33a02c',
+            'NA': '#b2df8a',
+        }
+
+        cols_to_add_colour_col.reverse()
+
+        for col_index in cols_to_add_colour_col:
+            header[col_index] = header[col_index][:-3]
+            header.insert(col_index + 1, header[col_index] + ':colour')
+
+            for row_index in range(len(matrix)):
+                colour = field_to_col[matrix[row_index][col_index]]
+                matrix[row_index].insert(col_index + 1, colour)
+
+        return header, matrix
+
+
     @classmethod
     def _matrix_to_csv(cls, matrix, header, outfile):
         f = pyfastaq.utils.open_file_write(outfile)
@@ -356,7 +382,8 @@ def run(self):
             if self.verbose:
                 print('Making Phandango csv file', csv_file, flush=True)
             csv_file = self.outprefix + '.phandango.csv'
-            Summary._matrix_to_csv(matrix, phandango_header, csv_file)
+            phandango_header, phandango_matrix = Summary._add_phandango_colour_columns(phandango_header, matrix)
+            Summary._matrix_to_csv(phandango_matrix, phandango_header, csv_file)
             dist_matrix_file = self.outprefix + '.phandango.distance_matrix'
             tree_file = self.outprefix + '.phandango.tre'
 

diff --git a/ariba/tests/summary_test.py b/ariba/tests/summary_test.py
@@ -343,6 +343,28 @@ def test_filter_matrix_columns(self):
         self.assertEqual(expected_matrix, got_matrix)
 
 
+    def test_add_phandango_colour_columns(self):
+        '''Test _add_phandango_colour_columns'''
+        header = ['head1', 'head2:o1', 'head3:o1', 'head4', 'head5:o1']
+        matrix = [
+            ['yes', 'yes', 'yes_nonunique', 'yes', 'no'],
+            ['yes', 'yes_nonunique', 'no', 'yes', 'NA'],
+            ['yes', 'no', 'NA', 'yes', 'yes'],
+            ['yes', 'NA', 'yes', 'yes', 'yes_nonunique'],
+        ]
+
+        expected_header = ['head1', 'head2', 'head2:colour', 'head3', 'head3:colour', 'head4', 'head5', 'head5:colour']
+        expected_matrix = [
+            ['yes', 'yes', '#1f78b4', 'yes_nonunique', '#a6cee3', 'yes', 'no', '#33a02c'],
+            ['yes', 'yes_nonunique', '#a6cee3', 'no', '#33a02c', 'yes', 'NA', '#b2df8a'],
+            ['yes', 'no', '#33a02c', 'NA', '#b2df8a', 'yes', 'yes', '#1f78b4'],
+            ['yes', 'NA', '#b2df8a', 'yes', '#1f78b4', 'yes', 'yes_nonunique', '#a6cee3'],
+        ]
+        got_header, got_matrix = summary.Summary._add_phandango_colour_columns(header, matrix)
+        self.assertEqual(expected_header, got_header)
+        self.assertEqual(expected_matrix, got_matrix)
+
+
     def test_matrix_to_csv(self):
         '''Test _matrix_to_csv'''
         matrix = [