updated data

1 files changed, 22 insertions, 10 deletions

diff --git a/data.py b/data.py
index 1144e9e..2018cd4 100644
--- a/data.py
+++ b/data.py
@@ -109,26 +109,38 @@ import numpy as np
 from scipy.stats import pearsonr
 
 if len(sys.argv) < 2:
-  print('You must provide at least one argument to choose the function of this program: `image` or `neighbor`')
+  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]
 
-if arg == 'image':
+if arg == 'img':
   for trial in trials:
     trial.plot(save=True)
-elif arg == 'neighbor':
+elif arg == 'nei':
   x = []
   y = []
   for trial in trials:
+    size = math.sqrt(trial.size[0]*trial.size[1])
     nei = trial.nearest_neighbor()
-    for n in range(len(nei)):
-      nei[n] += n*.05
-    plt.xlabel('Square root of Trial Area (cm)')
-    plt.ylabel('Nearest Neighbor for Individual Pits (cm)')
-    x += [math.sqrt(trial.size[0]*trial.size[1])]*len(nei)
-    y += nei
-  #print('p-value', pearsonr(x,y))
+    #for n in range(len(nei)):
+    #  nei[n] += n*.05
+    #x += [math.sqrt(trial.size[0]*trial.size[1])]*len(nei)
+    #y += nei
+    x.append(size)
+    y.append(sum(nei)/len(nei))
+  print('p-value', pearsonr(x,y))
+  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.show()
+elif arg == 'depwid':
+  for trial in trials:
+    size = math.sqrt(trial.size[0]*trial.size[1])
+    for pit in trial.pits:
+      plt.plot(size, pit.depth, 'bo')
+      plt.plot(size, pit.diam, 'ro')
+  plt.xlabel('Square root of Trial Area (cm)')
+  plt.ylabel('Depth/Width of Antlion Pits (cm)')
+  plt.show()