箭头符号演示

新的花式箭头工具的箭头绘制示例。

代码由此人贡献： Rob Knight < rob@spot.colorado.edu >

用法:

箭头符号图示例

  1. import matplotlib.pyplot as plt
  2. import numpy as np
  4. rates_to_bases = {'r1': 'AT', 'r2': 'TA', 'r3': 'GA', 'r4': 'AG', 'r5': 'CA',
  5.                   'r6': 'AC', 'r7': 'GT', 'r8': 'TG', 'r9': 'CT', 'r10': 'TC',
  6.                   'r11': 'GC', 'r12': 'CG'}
  7. numbered_bases_to_rates = {v: k for k, v in rates_to_bases.items()}
  8. lettered_bases_to_rates = {v: 'r' + v for k, v in rates_to_bases.items()}
  11. def add_dicts(d1, d2):
  12.     """Adds two dicts and returns the result."""
  13.     result = d1.copy()
  14.     result.update(d2)
  15.     return result
  18. def make_arrow_plot(data, size=4, display='length', shape='right',
  19.                     max_arrow_width=0.03, arrow_sep=0.02, alpha=0.5,
  20.                     normalize_data=False, ec=None, labelcolor=None,
  21.                     head_starts_at_zero=True,
  22.                     rate_labels=lettered_bases_to_rates,
  23.                     **kwargs):
  24.     """Makes an arrow plot.
  26.     Parameters:
  28.     data: dict with probabilities for the bases and pair transitions.
  29.     size: size of the graph in inches.
  30.     display: 'length', 'width', or 'alpha' for arrow property to change.
  31.     shape: 'full', 'left', or 'right' for full or half arrows.
  32.     max_arrow_width: maximum width of an arrow, data coordinates.
  33.     arrow_sep: separation between arrows in a pair, data coordinates.
  34.     alpha: maximum opacity of arrows, default 0.8.
  36.     **kwargs can be anything allowed by a Arrow object, e.g.
  37.     linewidth and edgecolor.
  38.     """
  40.     plt.xlim(-0.5, 1.5)
  41.     plt.ylim(-0.5, 1.5)
  42.     plt.gcf().set_size_inches(size, size)
  43.     plt.xticks([])
  44.     plt.yticks([])
  45.     max_text_size = size * 12
  46.     min_text_size = size
  47.     label_text_size = size * 2.5
  48.     text_params = {'ha': 'center', 'va': 'center', 'family': 'sans-serif',
  49.                    'fontweight': 'bold'}
  50.     r2 = np.sqrt(2)
  52.     deltas = {
  53.         'AT': (1, 0),
  54.         'TA': (-1, 0),
  55.         'GA': (0, 1),
  56.         'AG': (0, -1),
  57.         'CA': (-1 / r2, 1 / r2),
  58.         'AC': (1 / r2, -1 / r2),
  59.         'GT': (1 / r2, 1 / r2),
  60.         'TG': (-1 / r2, -1 / r2),
  61.         'CT': (0, 1),
  62.         'TC': (0, -1),
  63.         'GC': (1, 0),
  64.         'CG': (-1, 0)}
  66.     colors = {
  67.         'AT': 'r',
  68.         'TA': 'k',
  69.         'GA': 'g',
  70.         'AG': 'r',
  71.         'CA': 'b',
  72.         'AC': 'r',
  73.         'GT': 'g',
  74.         'TG': 'k',
  75.         'CT': 'b',
  76.         'TC': 'k',
  77.         'GC': 'g',
  78.         'CG': 'b'}
  80.     label_positions = {
  81.         'AT': 'center',
  82.         'TA': 'center',
  83.         'GA': 'center',
  84.         'AG': 'center',
  85.         'CA': 'left',
  86.         'AC': 'left',
  87.         'GT': 'left',
  88.         'TG': 'left',
  89.         'CT': 'center',
  90.         'TC': 'center',
  91.         'GC': 'center',
  92.         'CG': 'center'}
  94.     def do_fontsize(k):
  95.         return float(np.clip(max_text_size * np.sqrt(data[k]),
  96.                              min_text_size, max_text_size))
  98.     A = plt.text(0, 1, '$A_3$', color='r', size=do_fontsize('A'),
  99.                  **text_params)
  100.     T = plt.text(1, 1, '$T_3$', color='k', size=do_fontsize('T'),
  101.                  **text_params)
  102.     G = plt.text(0, 0, '$G_3$', color='g', size=do_fontsize('G'),
  103.                  **text_params)
  104.     C = plt.text(1, 0, '$C_3$', color='b', size=do_fontsize('C'),
  105.                  **text_params)
  107.     arrow_h_offset = 0.25  # data coordinates, empirically determined
  108.     max_arrow_length = 1 - 2 * arrow_h_offset
  109.     max_head_width = 2.5 * max_arrow_width
  110.     max_head_length = 2 * max_arrow_width
  111.     arrow_params = {'length_includes_head': True, 'shape': shape,
  112.                     'head_starts_at_zero': head_starts_at_zero}
  113.     ax = plt.gca()
  114.     sf = 0.6  # max arrow size represents this in data coords
  116.     d = (r2 / 2 + arrow_h_offset - 0.5) / r2  # distance for diags
  117.     r2v = arrow_sep / r2  # offset for diags
  119.     # tuple of x, y for start position
  120.     positions = {
  121.         'AT': (arrow_h_offset, 1 + arrow_sep),
  122.         'TA': (1 - arrow_h_offset, 1 - arrow_sep),
  123.         'GA': (-arrow_sep, arrow_h_offset),
  124.         'AG': (arrow_sep, 1 - arrow_h_offset),
  125.         'CA': (1 - d - r2v, d - r2v),
  126.         'AC': (d + r2v, 1 - d + r2v),
  127.         'GT': (d - r2v, d + r2v),
  128.         'TG': (1 - d + r2v, 1 - d - r2v),
  129.         'CT': (1 - arrow_sep, arrow_h_offset),
  130.         'TC': (1 + arrow_sep, 1 - arrow_h_offset),
  131.         'GC': (arrow_h_offset, arrow_sep),
  132.         'CG': (1 - arrow_h_offset, -arrow_sep)}
  134.     if normalize_data:
  135.         # find maximum value for rates, i.e. where keys are 2 chars long
  136.         max_val = 0
  137.         for k, v in data.items():
  138.             if len(k) == 2:
  139.                 max_val = max(max_val, v)
  140.         # divide rates by max val, multiply by arrow scale factor
  141.         for k, v in data.items():
  142.             data[k] = v / max_val * sf
  144.     def draw_arrow(pair, alpha=alpha, ec=ec, labelcolor=labelcolor):
  145.         # set the length of the arrow
  146.         if display == 'length':
  147.             length = max_head_length + data[pair] / sf * (max_arrow_length -
  148.                                                           max_head_length)
  149.         else:
  150.             length = max_arrow_length
  151.         # set the transparency of the arrow
  152.         if display == 'alpha':
  153.             alpha = min(data[pair] / sf, alpha)
  155.         # set the width of the arrow
  156.         if display == 'width':
  157.             scale = data[pair] / sf
  158.             width = max_arrow_width * scale
  159.             head_width = max_head_width * scale
  160.             head_length = max_head_length * scale
  161.         else:
  162.             width = max_arrow_width
  163.             head_width = max_head_width
  164.             head_length = max_head_length
  166.         fc = colors[pair]
  167.         ec = ec or fc
  169.         x_scale, y_scale = deltas[pair]
  170.         x_pos, y_pos = positions[pair]
  171.         plt.arrow(x_pos, y_pos, x_scale * length, y_scale * length,
  172.                   fc=fc, ec=ec, alpha=alpha, width=width,
  173.                   head_width=head_width, head_length=head_length,
  174.                   **arrow_params)
  176.         # figure out coordinates for text
  177.         # if drawing relative to base: x and y are same as for arrow
  178.         # dx and dy are one arrow width left and up
  179.         # need to rotate based on direction of arrow, use x_scale and y_scale
  180.         # as sin x and cos x?
  181.         sx, cx = y_scale, x_scale
  183.         where = label_positions[pair]
  184.         if where == 'left':
  185.             orig_position = 3 * np.array([[max_arrow_width, max_arrow_width]])
  186.         elif where == 'absolute':
  187.             orig_position = np.array([[max_arrow_length / 2.0,
  188.                                        3 * max_arrow_width]])
  189.         elif where == 'right':
  190.             orig_position = np.array([[length - 3 * max_arrow_width,
  191.                                        3 * max_arrow_width]])
  192.         elif where == 'center':
  193.             orig_position = np.array([[length / 2.0, 3 * max_arrow_width]])
  194.         else:
  195.             raise ValueError("Got unknown position parameter %s" % where)
  197.         M = np.array([[cx, sx], [-sx, cx]])
  198.         coords = np.dot(orig_position, M) + [[x_pos, y_pos]]
  199.         x, y = np.ravel(coords)
  200.         orig_label = rate_labels[pair]
  201.         label = r'$%s_{_{\mathrm{%s}}}$' % (orig_label[0], orig_label[1:])
  203.         plt.text(x, y, label, size=label_text_size, ha='center', va='center',
  204.                  color=labelcolor or fc)
  206.     for p in sorted(positions):
  207.         draw_arrow(p)
  210. # test data
  211. all_on_max = dict([(i, 1) for i in 'TCAG'] +
  212.                   [(i + j, 0.6) for i in 'TCAG' for j in 'TCAG'])
  214. realistic_data = {
  215.     'A': 0.4,
  216.     'T': 0.3,
  217.     'G': 0.5,
  218.     'C': 0.2,
  219.     'AT': 0.4,
  220.     'AC': 0.3,
  221.     'AG': 0.2,
  222.     'TA': 0.2,
  223.     'TC': 0.3,
  224.     'TG': 0.4,
  225.     'CT': 0.2,
  226.     'CG': 0.3,
  227.     'CA': 0.2,
  228.     'GA': 0.1,
  229.     'GT': 0.4,
  230.     'GC': 0.1}
  232. extreme_data = {
  233.     'A': 0.75,
  234.     'T': 0.10,
  235.     'G': 0.10,
  236.     'C': 0.05,
  237.     'AT': 0.6,
  238.     'AC': 0.3,
  239.     'AG': 0.1,
  240.     'TA': 0.02,
  241.     'TC': 0.3,
  242.     'TG': 0.01,
  243.     'CT': 0.2,
  244.     'CG': 0.5,
  245.     'CA': 0.2,
  246.     'GA': 0.1,
  247.     'GT': 0.4,
  248.     'GC': 0.2}
  250. sample_data = {
  251.     'A': 0.2137,
  252.     'T': 0.3541,
  253.     'G': 0.1946,
  254.     'C': 0.2376,
  255.     'AT': 0.0228,
  256.     'AC': 0.0684,
  257.     'AG': 0.2056,
  258.     'TA': 0.0315,
  259.     'TC': 0.0629,
  260.     'TG': 0.0315,
  261.     'CT': 0.1355,
  262.     'CG': 0.0401,
  263.     'CA': 0.0703,
  264.     'GA': 0.1824,
  265.     'GT': 0.0387,
  266.     'GC': 0.1106}
  269. if __name__ == '__main__':
  270.     from sys import argv
  271.     d = None
  272.     if len(argv) > 1:
  273.         if argv[1] == 'full':
  274.             d = all_on_max
  275.             scaled = False
  276.         elif argv[1] == 'extreme':
  277.             d = extreme_data
  278.             scaled = False
  279.         elif argv[1] == 'realistic':
  280.             d = realistic_data
  281.             scaled = False
  282.         elif argv[1] == 'sample':
  283.             d = sample_data
  284.             scaled = True
  285.     if d is None:
  286.         d = all_on_max
  287.         scaled = False
  288.     if len(argv) > 2:
  289.         display = argv[2]
  290.     else:
  291.         display = 'length'
  293.     size = 4
  294.     plt.figure(figsize=(size, size))
  296.     make_arrow_plot(d, display=display, linewidth=0.001, edgecolor=None,
  297.                     normalize_data=scaled, head_starts_at_zero=True, size=size)
  299.     plt.show()

下载这个示例