--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/visualization.py	Thu Aug 02 11:57:26 2018 -0400
@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+"""Visualise the output of Heinz.
+
+This script is used to visualize the output of Heinz, which is in the form of
+DOT language:
+    1. Clear the output of Heinz, extract the DOT source code.
+    2. Visualize the DOT source code and save it into file.
+
+The function of this script is rather simple, for more advanced visualization,
+please adopt other solutions mentioned in the paper
+doi: 10.1093/bioinformatics/btv526
+
+This tool is only designed for visualizing the output of Heinz tool.
+"""
+
+# Author: Cico Zhang
+# Date: 2 Aug 2017
+# Version: 0.2
+
+import argparse
+import sys
+
+from graphviz import Source
+
+
+def get_args():
+    """Collect the inputs."""
+    parser = argparse.ArgumentParser(
+        description='Visualise the output of Heinz')
+    parser.add_argument('-i', '--input', required=True, dest='heinz',
+                        metavar='Heinz_output.txt', type=str,
+                        help='Output file of Heinz as the input')
+    parser.add_argument('-o', '--output', required=True, dest='output',
+                        metavar='graph.pdf', type=str,
+                        help='The output file that saves the visualisation')
+    args = parser.parse_args()
+
+    if args.heinz is None:
+        sys.exit('Input file must be designated.')
+
+    return args
+
+
+def main():
+    """Main function."""
+    args = get_args()
+    # Read the whole output file
+    with open(args.heinz) as r:
+        graph_dot = r.readlines()
+
+    # Remove the redundant lines
+    while not graph_dot[0].startswith('graph G {'):
+        graph_dot.pop(0)
+
+    src = Source(''.join(graph_dot))
+    data_pdf = src.pipe('pdf')
+    # Redirect the output (very important)
+    with open(args.output, 'wb') as w:
+        w.write(data_pdf)
+    print('The visualization is saved as PDF!')
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()