You are reading an old version of the documentation (v1.4.2). For the latest version see

Previous topic

pylab_examples example code:

Next topic

pylab_examples example code:

This Page

pylab_examples example code: arrow_demo.pyΒΆ

(Source code, png, hires.png, pdf)

#!/usr/bin/env python
"""Arrow drawing example for the new fancy_arrow facilities.

Code contributed by: Rob Knight <[email protected]>


  python 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()
    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,\
    """Makes an arrow plot.


    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.

    max_text_size = size*12
    min_text_size = size
    label_text_size = size*2.5
    text_params={'ha':'center', 'va':'center', 'family':'sans-serif',\
    r2 = sqrt(2)

    deltas = {\
        'CA':(-1/r2, 1/r2),
        'AC':(1/r2, -1/r2),
        'GT':(1/r2, 1/r2),

    colors = {\

    label_positions = {\

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

    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, \
    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)*\
            length = max_arrow_length
        #set the transparency of the arrow
        if display == 'alph':
            alpha = min(data[pair]/sf, 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
            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,\
        elif where == 'center':
            orig_position = array([[length/2.0, 3*max_arrow_width]])
            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():

    #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 = {

extreme_data = {

sample_data = {

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
    if len(argv) > 2:
        display = argv[2]
        display = 'length'

    size = 4

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



Keywords: python, matplotlib, pylab, example, codex (see Search examples)