.. note:: :class: sphx-glr-download-link-note Click :ref:here  to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_gallery_text_labels_and_annotations_arrow_demo.py: ========== Arrow Demo ========== Arrow drawing example for the new fancy_arrow facilities. Code contributed by: Rob Knight usage: python arrow_demo.py realistic|full|sample|extreme .. image:: /gallery/text_labels_and_annotations/images/sphx_glr_arrow_demo_001.png :class: sphx-glr-single-img .. code-block:: default import matplotlib.pyplot as plt import numpy as np 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 = {v: k for k, v in rates_to_bases.items()} lettered_bases_to_rates = {v: 'r' + v for k, v in rates_to_bases.items()} 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', 'alpha'} The arrow property to change. shape : {'full', 'left', 'right'} For full or half arrows. max_arrow_width : float Maximum width of an arrow, data coordinates. arrow_sep : float Separation between arrows in a pair, data coordinates. alpha : float Maximum opacity of arrows. **kwargs Can be anything allowed by a Arrow object, e.g. *linewidth* or *edgecolor*. """ plt.xlim(-0.5, 1.5) plt.ylim(-0.5, 1.5) plt.gcf().set_size_inches(size, size) plt.xticks([]) plt.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 = np.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(np.clip(max_text_size * np.sqrt(data[k]), min_text_size, max_text_size)) A = plt.text(0, 1, '$A_3$', color='r', size=do_fontsize('A'), **text_params) T = plt.text(1, 1, '$T_3$', color='k', size=do_fontsize('T'), **text_params) G = plt.text(0, 0, '$G_3$', color='g', size=do_fontsize('G'), **text_params) C = plt.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_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} 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 = max((v for k, v in data.items() if len(k) == 2), default=0) # 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 + data[pair] / sf * (max_arrow_length - max_head_length)) else: length = max_arrow_length # set the transparency of the arrow if display == 'alpha': 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 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] plt.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 * np.array([[max_arrow_width, max_arrow_width]]) elif where == 'absolute': orig_position = np.array([[max_arrow_length / 2.0, 3 * max_arrow_width]]) elif where == 'right': orig_position = np.array([[length - 3 * max_arrow_width, 3 * max_arrow_width]]) elif where == 'center': orig_position = np.array([[length / 2.0, 3 * max_arrow_width]]) else: raise ValueError("Got unknown position parameter %s" % where) M = np.array([[cx, sx], [-sx, cx]]) coords = np.dot(orig_position, M) + [[x_pos, y_pos]] x, y = np.ravel(coords) orig_label = rate_labels[pair] label = r'$%s_{_{\mathrm{%s}}}$' % (orig_label[0], orig_label[1:]) plt.text(x, y, label, size=label_text_size, ha='center', va='center', color=labelcolor or fc) for p in sorted(positions): 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 plt.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) plt.show() .. _sphx_glr_download_gallery_text_labels_and_annotations_arrow_demo.py: .. only :: html .. container:: sphx-glr-footer :class: sphx-glr-footer-example .. container:: sphx-glr-download :download:Download Python source code: arrow_demo.py  .. container:: sphx-glr-download :download:Download Jupyter notebook: arrow_demo.ipynb  .. only:: html .. rst-class:: sphx-glr-signature Keywords: matplotlib code example, codex, python plot, pyplot Gallery generated by Sphinx-Gallery _