Merge pull request #4660 from domspad/MEP12-on-arrow_demo.py

jenshnielsen · jenshnielsen · commit b8d7e254484b · 2015-07-19T12:22:56.000+01:00
MEP12-on-arrow_demo.py
diff --git a/examples/pylab_examples/arrow_demo.py b/examples/pylab_examples/arrow_demo.py
@@ -9,10 +9,12 @@
 
 
 """
-from pylab import *
+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'}
+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()])
 
@@ -45,17 +47,17 @@ def make_arrow_plot(data, size=4, display='length', shape='right',
     linewidth and edgecolor.
     """
 
-    xlim(-0.5, 1.5)
-    ylim(-0.5, 1.5)
-    gcf().set_size_inches(size, size)
-    xticks([])
-    yticks([])
+    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 = sqrt(2)
+    r2 = np.sqrt(2)
 
     deltas = {
         'AT': (1, 0),
@@ -103,13 +105,13 @@ def make_arrow_plot(data, size=4, display='length', shape='right',
         }
 
     def do_fontsize(k):
-        return float(clip(max_text_size*sqrt(data[k]),
+        return float(np.clip(max_text_size*np.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)
+    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
@@ -119,7 +121,7 @@ def do_fontsize(k):
     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()
+    ax = plt.gca()
     sf = 0.6  # max arrow size represents this in data coords
 
     d = (r2/2 + arrow_h_offset - 0.5)/r2  # distance for diags
@@ -179,7 +181,7 @@ def draw_arrow(pair, alpha=alpha, ec=ec, labelcolor=labelcolor):
 
         x_scale, y_scale = deltas[pair]
         x_pos, y_pos = positions[pair]
-        arrow(x_pos, y_pos, x_scale*length, y_scale*length,
+        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)
 
@@ -192,24 +194,24 @@ def draw_arrow(pair, alpha=alpha, ec=ec, labelcolor=labelcolor):
 
         where = label_positions[pair]
         if where == 'left':
-            orig_position = 3*array([[max_arrow_width, max_arrow_width]])
+            orig_position = 3*np.array([[max_arrow_width, max_arrow_width]])
         elif where == 'absolute':
-            orig_position = array([[max_arrow_length/2.0, 3*max_arrow_width]])
+            orig_position = np.array([[max_arrow_length/2.0, 3*max_arrow_width]])
         elif where == 'right':
-            orig_position = array([[length - 3*max_arrow_width,
+            orig_position = np.array([[length - 3*max_arrow_width,
                                     3*max_arrow_width]])
         elif where == 'center':
-            orig_position = array([[length/2.0, 3*max_arrow_width]])
+            orig_position = np.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)
+        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 = '$%s_{_{\mathrm{%s}}}$' % (orig_label[0], orig_label[1:])
 
-        text(x, y, label, size=label_text_size, ha='center', va='center',
+        plt.text(x, y, label, size=label_text_size, ha='center', va='center',
              color=labelcolor or fc)
 
     for p in positions.keys():
@@ -302,11 +304,11 @@ def draw_arrow(pair, alpha=alpha, ec=ec, labelcolor=labelcolor):
         display = 'length'
 
     size = 4
-    figure(figsize=(size, size))
+    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)
 
-    draw()
+    plt.draw()
 
-    show()
+    plt.show()
