from data import trials
from scipy.stats import pearsonr
from numpy import poly1d, polyfit
from math import sqrt
import matplotlib.pyplot as plt

x = []
y = []
for trial in trials:
  size = sqrt(trial.size[0]*trial.size[1])
  nei = trial.nearest_neighbor()
  x.append(size)
  y.append(sum(nei)/len(nei))
fig = plt.figure()
ax = fig.add_subplot(111)
plt.text(0.1, 0.9, 'R^2 = %.3f\np=%.3f' % (pearsonr(x,y)[0]**2, pearsonr(x,y)[1]), ha='center', va='center', transform=ax.transAxes)
plt.xlabel('Square root of Trial Area (cm)')
plt.ylabel('Nearest Neighbor for Individual Pits (cm)')
plt.plot(x, y, 'bo')
plt.plot(x, poly1d(polyfit(x, y, 1))(x))
plt.savefig('nearest_neighbor.png', bbox_inches='tight')