File: venn_mpl.py

package info (click to toggle)
python-pybedtools 0.8.0-5
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 16,140 kB
  • sloc: python: 9,589; cpp: 899; makefile: 149; sh: 116
file content (147 lines) | stat: -rw-r--r-- 4,952 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
#!/usr/bin/env python
"""
Given 3 files, creates a 3-way Venn diagram of intersections using matplotlib; \
see :mod:`pybedtools.contrib.venn_maker` for more flexibility.

Numbers are placed on the diagram.  If you don't have matplotlib installed.
try venn_gchart.py to use the Google Chart API instead.

The values in the diagram assume:

    * unstranded intersections
    * no features that are nested inside larger features
"""

import argparse
import sys
import os
import pybedtools

def venn_mpl(a, b, c, colors=None, outfn='out.png', labels=None):
    """
    *a*, *b*, and *c* are filenames to BED-like files.

    *colors* is a list of matplotlib colors for the Venn diagram circles.

    *outfn* is the resulting output file.  This is passed directly to
    fig.savefig(), so you can supply extensions of .png, .pdf, or whatever your
    matplotlib installation supports.

    *labels* is a list of labels to use for each of the files; by default the
    labels are ['a','b','c']
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.patches import Circle
    except ImportError:
        sys.stderr.write('matplotlib is required to make a Venn diagram with %s\n' % os.path.basename(sys.argv[0]))
        sys.exit(1)

    a = pybedtools.BedTool(a)
    b = pybedtools.BedTool(b)
    c = pybedtools.BedTool(c)

    if colors is None:
        colors = ['r','b','g']

    radius = 6.0
    center = 0.0
    offset = radius / 2

    if labels is None:
        labels = ['a','b','c']

    circle_a = Circle(xy = (center-offset, center+offset), radius=radius, edgecolor=colors[0], label=labels[0])
    circle_b = Circle(xy = (center+offset, center+offset), radius=radius, edgecolor=colors[1], label=labels[1])
    circle_c = Circle(xy = (center,        center-offset), radius=radius, edgecolor=colors[2], label=labels[2])


    fig = plt.figure(facecolor='w')
    ax = fig.add_subplot(111)

    for circle in (circle_a, circle_b, circle_c):
        circle.set_facecolor('none')
        circle.set_linewidth(3)
        ax.add_patch(circle)

    ax.axis('tight')
    ax.axis('equal')
    ax.set_axis_off()


    kwargs = dict(horizontalalignment='center')

    # Unique to A
    ax.text( center-2*offset, center+offset, str((a - b - c).count()), **kwargs)

    # Unique to B
    ax.text( center+2*offset, center+offset, str((b - a - c).count()), **kwargs)

    # Unique to C
    ax.text( center, center-2*offset, str((c - a - b).count()), **kwargs)

    # A and B not C
    ax.text( center, center+2*offset-0.5*offset, str((a + b - c).count()), **kwargs)

    # A and C not B
    ax.text( center-1.2*offset, center-0.5*offset, str((a + c - b).count()), **kwargs)

    # B and C not A
    ax.text( center+1.2*offset, center-0.5*offset, str((b + c - a).count()), **kwargs)

    # all
    ax.text( center, center, str((a + b + c).count()), **kwargs)

    ax.legend(loc='best')

    fig.savefig(outfn)

    plt.close(fig)

def main():
    """Create a 3-way Venn diagram, using matplotlib"""
    op = argparse.ArgumentParser(description=__doc__, prog=sys.argv[0])
    op.add_argument('-a', help='File to use for the left-most circle')
    op.add_argument('-b', help='File to use for the right-most circle')
    op.add_argument('-c', help='File to use for the bottom circle')
    op.add_argument('--labels',
                  help='Optional comma-separated list of '
                       'labels for a, b, and c', default='a,b,c')
    op.add_argument('--colors', default='r,b,g',
                  help='Comma-separated list of matplotlib-valid colors '
                       'for circles a, b, and c.  E.g., --colors=r,b,k')
    op.add_argument('-o', default='out.png', 
                  help='Output file to save as.  Extension is '
                       'meaningful, e.g., out.pdf, out.png, out.svg.  Default is "%(default)s"')
    op.add_argument('--test', action='store_true', help='run test, overriding all other options.')
    options = op.parse_args()

    reqd_args = ['a','b','c']
    if not options.test:
        for ra in reqd_args:
            if not getattr(options,ra):
                op.print_help()
                sys.stderr.write('Missing required arg "%s"\n' % ra)
                sys.exit(1)

    if options.test:
        pybedtools.bedtool.random.seed(1)
        a = pybedtools.example_bedtool('rmsk.hg18.chr21.small.bed')
        b = pybedtools.example_bedtool('venn.b.bed')
        c = pybedtools.example_bedtool('venn.c.bed')
        options.a = a.fn
        options.b = b.fn
        options.c = c.fn
        options.colors='r,b,g'
        options.o = 'out.png'
        options.labels = 'a,b,c'

    venn_mpl(a=options.a, b=options.b, c=options.c, 
             colors=options.colors.split(','),
             labels=options.labels.split(','), 
             outfn=options.o)

if __name__ == "__main__":
    import doctest
    if doctest.testmod(optionflags=doctest.ELLIPSIS).failed == 0:
        main()