From b596def8038023f92aed32f92bd48b2dd8b4adc1 Mon Sep 17 00:00:00 2001 From: Holden Rohrer Date: Sat, 11 Jan 2020 21:05:28 -0500 Subject: added multi-function to data.py (incl. nearest-neighbor) --- data.py | 40 ++++++++++++++++++++++++++++++++++++---- 1 file 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() -- cgit