Mercurial > repos > nick > dna_visualizer

comparison literal.py @ 0:5257ce9d9184

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Initial literal.py tool
author Nick Stoler <nstoler@psu.edu>
date Sun, 02 Mar 2014 13:51:03 -0500
parents
children 58160195728e
comparison
equal deleted inserted replaced
-1:000000000000 0:5257ce9d9184
1 #!/usr/bin/env python
2 from __future__ import division
3 import os
4 import sys
5 import Image
6 import argparse
7 import fastareader
8
9 OPT_DEFAULTS = {'size':'512x512', 'verbose':True,
10 'A':'0,255,0', 'T':'255,0,0', 'G':'255,255,255', 'C':'0,0,255'}
11 USAGE = "%(prog)s [options] genome.fasta"
12 DESCRIPTION = """Convert DNA sequence into a PNG image by representing each base
13 with one colored pixel."""
14 EPILOG = """"""
15
16 def main():
17
18 parser = argparse.ArgumentParser(
19 description=DESCRIPTION, usage=USAGE, epilog=EPILOG)
20 parser.set_defaults(**OPT_DEFAULTS)
21
22 parser.add_argument('fasta', metavar='genome.fasta',
23 help="""Input sequence. Can be in FASTA format or a plain text file
24 containing only the sequence. Any non-ATGC characters (case-insensitive)
25 will be skipped.""")
26 parser.add_argument('-s', '--size',
27 help="""The output image size, in pixels, in the format "widthxheight", e.g.
28 "640x480". If the sequence is larger than the number of pixels in the
29 image, it will be cut off. Default size: %(default)s""")
30 parser.add_argument('-o', '--outfile', metavar='image.png',
31 help="""Output filename. Overrides the default, which is
32 to use the input filename base plus .png.""")
33 parser.add_argument('-d', '--display', action='store_true',
34 help="""Display the image instead of saving it.""")
35 parser.add_argument('-c', '--clobber', action='store_true',
36 help="""If the output filename already exists, overwrite it instead of
37 throwing an error (the default).""")
38 parser.add_argument('-v', '--verbose', action='store_true',
39 help="""Verbose mode. On by default.""")
40 parser.add_argument('-q', '--quiet', action='store_false', dest='verbose',
41 help="""Quiet mode.""")
42 group = parser.add_argument_group('Color customization', """Use these options
43 to use custom colors for bases. Specify with a comma-delimited RGB value
44 like "100,150,10".""")
45 group.add_argument('-A', metavar='R,G,B',
46 help="""default: %(default)s""")
47 group.add_argument('-T', metavar='R,G,B',
48 help="""default: %(default)s""")
49 group.add_argument('-G', metavar='R,G,B',
50 help="""default: %(default)s""")
51 group.add_argument('-C', metavar='R,G,B',
52 help="""default: %(default)s""")
53
54 args = parser.parse_args()
55
56 try:
57 size = parse_size(args.size)
58 except ValueError:
59 parser.print_help()
60 fail('\nError: Invalid size string "%s".' % args.size)
61
62 fasta = fastareader.FastaLineGenerator(args.fasta)
63 bases = fasta.bases()
64
65 if not args.display:
66 outfile = args.outfile if args.outfile else outfile_name(args.fasta)
67 if os.path.exists(outfile) and not args.clobber:
68 fail('Error: Output filename already taken: "%s"' % outfile)
69
70 colors = {}
71 colors['A'] = parse_rgb(args.A)
72 colors['T'] = parse_rgb(args.T)
73 colors['G'] = parse_rgb(args.G)
74 colors['C'] = parse_rgb(args.C)
75
76 image = Image.new('RGB', size, 'white')
77 pixels = image.load()
78
79 done = False
80 for i in range(image.size[1]):
81 for j in range(image.size[0]):
82 try:
83 base = next(bases).upper()
84 except StopIteration:
85 done = True
86 break
87 if base in colors:
88 pixels[j,i] = colors[base]
89 if done:
90 break
91
92 if args.display:
93 image.show()
94 else:
95 image.save(outfile, 'PNG')
96
97
98
99 def parse_size(size_str):
100 """Parse size string, return a tuple of (width, height).
101 Accepts size strings in the format "640x480".
102 If not valid, raises ValueError."""
103 size = map(int, size_str.split('x'))
104 if len(size) != 2:
105 raise ValueError
106 else:
107 return tuple(size)
108
109
110 def parse_rgb(rgb_str):
111 """Parse RGB string, return a tuple of (R, G, B).
112 If not valid, raises ValueError."""
113 rgb = map(int, rgb_str.split(','))
114 if len(rgb) != 3:
115 raise ValueError
116 else:
117 return tuple(rgb)
118
119
120 def outfile_name(infilename):
121 base = infilename.split('.')[0]
122 if not base:
123 base = infilename
124 return base+'.png'
125
126
127 def fail(message):
128 sys.stderr.write(message+"\n")
129 sys.exit(1)
130
131 if __name__ == '__main__':
132 main()