1 files changed, 67 insertions, 0 deletions
diff --git a/data.py b/data.py
new file mode 100644
index 0000000..cff030b
--- /dev/null
+++ b/data.py
@@ -0,0 +1,67 @@
+# Data
+
+from scipy.spatial import Voronoi, voronoi_plot_2d
+import matplotlib.pyplot as plt
+
+
+class Pit:
+  def __init__(self, loc, depth, diam):
+    self.loc = loc
+    self.depth = depth
+    self.diam = diam
+
+  def __repr__(self):
+    return '%s @ %.1fcm deep, %.1fcm wide' % (str(self.loc), self.depth, self.diam)
+
+  def disp(self):
+    return '%.1fcm dp\n %.1fcm wd' % (self.depth, self.diam)
+
+  def __getitem__(self,ind):
+    return self.loc[ind]
+
+class Date:
+  def __init__(self, year, month, day):
+    self.year = year
+    self.month = month
+    self.day = day
+    
+  def __repr__(self):
+    return '%d-%d-%d' % (self.year, self.month, self.day)
+
+class Trial:
+  def __init__(self, date, intro, dead, size, pits):
+    self.date = date
+    self.intro = intro
+    self.dead = dead
+    self.size = size
+    self.pits = pits
+
+  def __repr__(self):
+    return str(self.date) + ' ' + str(self.pits)
+
+  def plot(self, save=False):
+    vor = Voronoi([pit.loc for pit in self.pits])
+    voronoi_plot_2d(vor)
+    for pit in self.pits:
+      plt.text(pit[0], pit[1], pit.disp(), ha='center', va='bottom')
+    plt.xlabel('%s (dimension %dx%dcm)' % (str(self.date), self.size[0], self.size[1]))
+    if save:
+      plt.savefig(str(self.date)+'.svg')
+    else:
+      plt.show()
+
+trials = [
+  Trial(Date(2019, 10, 16), 31, 6, [31,32], [
+    Pit([4,25],1.3,4.2),
+    Pit([3,13],1.4,3.7),
+    Pit([10,25],1.1,3.0),
+    Pit([18,18],1.8,2.3),
+    Pit([30,14],2.2,3.1),
+    Pit([29,11],1.4,2.5),
+    Pit([27,10],1.2,2.1),
+    Pit([29,6],2.4,3.9),
+    Pit([26,5],1.8,3.6),
+  ]),
+]
+
+trials[0].plot()