data.py should run faster

1 files changed, 9 insertions, 9 deletions

diff --git a/data.py b/data.py
index f345b1e..3363dbd 100644
--- a/data.py
+++ b/data.py
@@ -103,16 +103,16 @@ trials = [
   ]),
 ]
 
-import sys
-import math
-import numpy as np
+from sys import argv
+from math import sqrt
+from numpy import poly1d, polyfit
 from scipy.stats import pearsonr
 
-if len(sys.argv) < 2:
+if len(argv) < 2:
   print('You must provide at least one argument to choose the function of this program: `img` or `nei` or `depwid')
   quit()
 
-arg = sys.argv[1]
+arg = argv[1]
 
 if arg == 'img':
   for trial in trials:
@@ -121,11 +121,11 @@ elif arg == 'nei':
   x = []
   y = []
   for trial in trials:
-    size = math.sqrt(trial.size[0]*trial.size[1])
+    size = sqrt(trial.size[0]*trial.size[1])
     nei = trial.nearest_neighbor()
     #for n in range(len(nei)):
     #  nei[n] += n*.05
-    #x += [math.sqrt(trial.size[0]*trial.size[1])]*len(nei)
+    #x += [sqrt(trial.size[0]*trial.size[1])]*len(nei)
     #y += nei
     x.append(size)
     y.append(sum(nei)/len(nei))
@@ -135,12 +135,12 @@ elif arg == 'nei':
   plt.xlabel('Square root of Trial Area (cm)')
   plt.ylabel('Nearest Neighbor for Individual Pits (cm)')
   plt.plot(x, y, 'bo')
-  plt.plot(x, np.poly1d(np.polyfit(x, y, 1))(x))
+  plt.plot(x, poly1d(polyfit(x, y, 1))(x))
   plt.savefig('nearest_neighbor.png', bbox_inches='tight')
 elif arg == 'depwid':
   depths, widths, sizes = [], [], []
   for trial in trials:
-    size = math.sqrt(trial.size[0]*trial.size[1])
+    size = sqrt(trial.size[0]*trial.size[1])
     for pit in trial.pits:
       sizes.append(size)
       depths.append(pit.depth)