added multi-function to data.py (incl. nearest-neighbor)

1 files changed, 36 insertions, 4 deletions

diff --git a/data.py b/data.py
index 55a99f2..1933065 100644
--- a/data.py
+++ b/data.py
@@ -1,9 +1,8 @@
 # Data
 
-from scipy.spatial import Voronoi, voronoi_plot_2d
+from scipy.spatial import Voronoi, voronoi_plot_2d, KDTree, distance
 import matplotlib.pyplot as plt
 
-
 class Pit:
   def __init__(self, loc, depth, diam):
     self.loc = loc
@@ -35,6 +34,7 @@ class Trial:
     self.dead = dead
     self.size = size
     self.pits = pits
+    self.pitlocs = [pit.loc for pit in self.pits]
 
   def __repr__(self):
     return str(self.date) + ' ' + str(self.pits)
@@ -50,6 +50,11 @@ class Trial:
     else:
       plt.show()
 
+  def nearest_neighbor(self):
+    tree = KDTree(self.pitlocs)
+    sumnn = 0
+    return [dists[1] for dists in tree.query(self.pitlocs,2)[0]]
+
 trials = [
   Trial(Date(2019, 10, 16), 31, 6, [31,32], [
     Pit([4,25],1.3,4.2),
@@ -98,5 +103,32 @@ trials = [
   ]),
 ]
 
-for trial in trials:
-  trial.plot(save=True)
+import sys
+import math
+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`')
+  quit()
+
+arg = sys.argv[1]
+
+if arg == 'image':
+  for trial in trials:
+    trial.plot(save=True)
+elif arg == 'neighbor':
+  x = []
+  y = []
+  for trial in trials:
+    nei = trial.nearest_neighbor()
+    for n in range(len(nei)):
+      nei[n] += ((-1)**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))
+  plt.plot(x, y, 'bo')
+  plt.plot(x, np.poly1d(np.polyfit(x, y, 1))(x))
+  plt.show()