arrow_demo.py :  » Chart-Report » Matplotlib » matplotlib-0.99.1.1 » examples » pylab_examples » Python Open Source

Home
Python Open Source
1.3.1.2 Python
2.Ajax
3.Aspect Oriented
4.Blog
5.Build
6.Business Application
7.Chart Report
8.Content Management Systems
9.Cryptographic
10.Database
11.Development
12.Editor
13.Email
14.ERP
15.Game 2D 3D
16.GIS
17.GUI
18.IDE
19.Installer
20.IRC
21.Issue Tracker
22.Language Interface
23.Log
24.Math
25.Media Sound Audio
26.Mobile
27.Network
28.Parser
29.PDF
30.Project Management
31.RSS
32.Search
33.Security
34.Template Engines
35.Test
36.UML
37.USB Serial
38.Web Frameworks
39.Web Server
40.Web Services
41.Web Unit
42.Wiki
43.Windows
44.XML
Python Open Source » Chart Report » Matplotlib 
Matplotlib » matplotlib 0.99.1.1 » examples » pylab_examples » arrow_demo.py
#!/usr/bin/env python
"""Arrow drawing example for the new fancy_arrow facilities.

Code contributed by: Rob Knight <rob@spot.colorado.edu>

usage:

  python arrow_demo.py realistic|full|sample|extreme


"""
from pylab import *

rates_to_bases={'r1':'AT', 'r2':'TA', 'r3':'GA','r4':'AG','r5':'CA','r6':'AC', \
            'r7':'GT', 'r8':'TG', 'r9':'CT','r10':'TC','r11':'GC','r12':'CG'}
numbered_bases_to_rates = dict([(v,k) for k, v in rates_to_bases.items()])
lettered_bases_to_rates = dict([(v, 'r'+v) for k, v in rates_to_bases.items()])
def add_dicts(d1, d2):
    """Adds two dicts and returns the result."""
    result = d1.copy()
    result.update(d2)
    return result

def make_arrow_plot(data, size=4, display='length', shape='right', \
        max_arrow_width=0.03, arrow_sep = 0.02, alpha=0.5, \
        normalize_data=False, ec=None, labelcolor=None, \
        head_starts_at_zero=True, rate_labels=lettered_bases_to_rates,\
        **kwargs):
    """Makes an arrow plot.

    Parameters:

    data: dict with probabilities for the bases and pair transitions.
    size: size of the graph in inches.
    display: 'length', 'width', or 'alpha' for arrow property to change.
    shape: 'full', 'left', or 'right' for full or half arrows.
    max_arrow_width: maximum width of an arrow, data coordinates.
    arrow_sep: separation between arrows in a pair, data coordinates.
    alpha: maximum opacity of arrows, default 0.8.

    **kwargs can be anything allowed by a Arrow object, e.g.
    linewidth and edgecolor.
    """

    xlim(-0.5,1.5)
    ylim(-0.5,1.5)
    gcf().set_size_inches(size,size)
    xticks([])
    yticks([])
    max_text_size = size*12
    min_text_size = size
    label_text_size = size*2.5
    text_params={'ha':'center', 'va':'center', 'family':'sans-serif',\
        'fontweight':'bold'}
    r2 = sqrt(2)

    deltas = {\
        'AT':(1,0),
        'TA':(-1,0),
        'GA':(0,1),
        'AG':(0,-1),
        'CA':(-1/r2, 1/r2),
        'AC':(1/r2, -1/r2),
        'GT':(1/r2, 1/r2),
        'TG':(-1/r2,-1/r2),
        'CT':(0,1),
        'TC':(0,-1),
        'GC':(1,0),
        'CG':(-1,0)
        }

    colors = {\
        'AT':'r',
        'TA':'k',
        'GA':'g',
        'AG':'r',
        'CA':'b',
        'AC':'r',
        'GT':'g',
        'TG':'k',
        'CT':'b',
        'TC':'k',
        'GC':'g',
        'CG':'b'
        }

    label_positions = {\
        'AT':'center',
        'TA':'center',
        'GA':'center',
        'AG':'center',
        'CA':'left',
        'AC':'left',
        'GT':'left',
        'TG':'left',
        'CT':'center',
        'TC':'center',
        'GC':'center',
        'CG':'center'
        }


    def do_fontsize(k):
        return float(clip(max_text_size*sqrt(data[k]),\
            min_text_size,max_text_size))

    A = text(0,1, '$A_3$', color='r', size=do_fontsize('A'), **text_params)
    T = text(1,1, '$T_3$', color='k', size=do_fontsize('T'), **text_params)
    G = text(0,0, '$G_3$', color='g', size=do_fontsize('G'), **text_params)
    C = text(1,0, '$C_3$', color='b', size=do_fontsize('C'), **text_params)

    arrow_h_offset = 0.25  #data coordinates, empirically determined
    max_arrow_length = 1 - 2*arrow_h_offset

    max_arrow_width = max_arrow_width
    max_head_width = 2.5*max_arrow_width
    max_head_length = 2*max_arrow_width
    arrow_params={'length_includes_head':True, 'shape':shape, \
        'head_starts_at_zero':head_starts_at_zero}
    ax = gca()
    sf = 0.6   #max arrow size represents this in data coords

    d = (r2/2 + arrow_h_offset - 0.5)/r2    #distance for diags
    r2v = arrow_sep/r2                      #offset for diags

    #tuple of x, y for start position
    positions = {\
        'AT': (arrow_h_offset, 1+arrow_sep),
        'TA': (1-arrow_h_offset, 1-arrow_sep),
        'GA': (-arrow_sep, arrow_h_offset),
        'AG': (arrow_sep, 1-arrow_h_offset),
        'CA': (1-d-r2v, d-r2v),
        'AC': (d+r2v, 1-d+r2v),
        'GT': (d-r2v, d+r2v),
        'TG': (1-d+r2v, 1-d-r2v),
        'CT': (1-arrow_sep, arrow_h_offset),
        'TC': (1+arrow_sep, 1-arrow_h_offset),
        'GC': (arrow_h_offset, arrow_sep),
        'CG': (1-arrow_h_offset, -arrow_sep),
        }

    if normalize_data:
        #find maximum value for rates, i.e. where keys are 2 chars long
        max_val = 0
        for k, v in data.items():
            if len(k) == 2:
                max_val = max(max_val, v)
        #divide rates by max val, multiply by arrow scale factor
        for k, v in data.items():
            data[k] = v/max_val*sf

    def draw_arrow(pair, alpha=alpha, ec=ec, labelcolor=labelcolor):
        #set the length of the arrow
        if display == 'length':
            length = max_head_length+(max_arrow_length-max_head_length)*\
                data[pair]/sf
        else:
            length = max_arrow_length
        #set the transparency of the arrow
        if display == 'alph':
            alpha = min(data[pair]/sf, alpha)
        else:
            alpha=alpha
        #set the width of the arrow
        if display == 'width':
            scale = data[pair]/sf
            width = max_arrow_width*scale
            head_width = max_head_width*scale
            head_length = max_head_length*scale
        else:
            width = max_arrow_width
            head_width = max_head_width
            head_length = max_head_length

        fc = colors[pair]
        ec = ec or fc

        x_scale, y_scale = deltas[pair]
        x_pos, y_pos = positions[pair]
        arrow(x_pos, y_pos, x_scale*length, y_scale*length, \
            fc=fc, ec=ec, alpha=alpha, width=width, head_width=head_width, \
            head_length=head_length, **arrow_params)

        #figure out coordinates for text
        #if drawing relative to base: x and y are same as for arrow
        #dx and dy are one arrow width left and up
        #need to rotate based on direction of arrow, use x_scale and y_scale
        #as sin x and cos x?
        sx, cx = y_scale, x_scale

        where = label_positions[pair]
        if where == 'left':
            orig_position = 3*array([[max_arrow_width, max_arrow_width]])
        elif where == 'absolute':
            orig_position = array([[max_arrow_length/2.0, 3*max_arrow_width]])
        elif where == 'right':
            orig_position = array([[length-3*max_arrow_width,\
                3*max_arrow_width]])
        elif where == 'center':
            orig_position = array([[length/2.0, 3*max_arrow_width]])
        else:
            raise ValueError, "Got unknown position parameter %s" % where



        M = array([[cx, sx],[-sx,cx]])
        coords = dot(orig_position, M) + [[x_pos, y_pos]]
        x, y = ravel(coords)
        orig_label = rate_labels[pair]
        label = '$%s_{_{\mathrm{%s}}}$' % (orig_label[0], orig_label[1:])

        text(x, y, label, size=label_text_size, ha='center', va='center', \
            color=labelcolor or fc)

    for p in positions.keys():
        draw_arrow(p)

    #test data
all_on_max = dict([(i, 1) for i in 'TCAG'] + \
        [(i+j, 0.6) for i in 'TCAG' for j in 'TCAG'])

realistic_data = {
        'A':0.4,
        'T':0.3,
        'G':0.5,
        'C':0.2,
        'AT':0.4,
        'AC':0.3,
        'AG':0.2,
        'TA':0.2,
        'TC':0.3,
        'TG':0.4,
        'CT':0.2,
        'CG':0.3,
        'CA':0.2,
        'GA':0.1,
        'GT':0.4,
        'GC':0.1,
    }

extreme_data = {
        'A':0.75,
        'T':0.10,
        'G':0.10,
        'C':0.05,
        'AT':0.6,
        'AC':0.3,
        'AG':0.1,
        'TA':0.02,
        'TC':0.3,
        'TG':0.01,
        'CT':0.2,
        'CG':0.5,
        'CA':0.2,
        'GA':0.1,
        'GT':0.4,
        'GC':0.2,
    }

sample_data = {
        'A':0.2137,
        'T':0.3541,
        'G':0.1946,
        'C':0.2376,
        'AT':0.0228,
        'AC':0.0684,
        'AG':0.2056,
        'TA':0.0315,
        'TC':0.0629,
        'TG':0.0315,
        'CT':0.1355,
        'CG':0.0401,
        'CA':0.0703,
        'GA':0.1824,
        'GT':0.0387,
        'GC':0.1106,
    }


if __name__ == '__main__':
    from sys import argv
    d = None
    if len(argv) > 1:
        if argv[1] == 'full':
            d = all_on_max
            scaled = False
        elif argv[1] == 'extreme':
            d = extreme_data
            scaled = False
        elif argv[1] == 'realistic':
            d = realistic_data
            scaled = False
        elif argv[1] == 'sample':
            d = sample_data
            scaled = True
    if d is None:
        d = all_on_max
        scaled=False
    if len(argv) > 2:
        display = argv[2]
    else:
        display = 'length'

    size = 4
    figure(figsize=(size,size))

    make_arrow_plot(d, display=display, linewidth=0.001, edgecolor=None,
        normalize_data=scaled, head_starts_at_zero=True, size=size)

    draw()

    show()
www.java2java.com | Contact Us
Copyright 2009 - 12 Demo Source and Support. All rights reserved.
All other trademarks are property of their respective owners.