added legend to depwid

1 files changed, 5 insertions, 4 deletions

diff --git a/data.py b/data.py
index e2d197e..3884910 100644
--- a/data.py
+++ b/data.py
@@ -134,16 +134,17 @@ elif arg == 'nei':
   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()
+  plt.savefig('nearest_neighbor.png', bbox='tight')
 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.plot(size, pit.depth, 'bo', label='depth')
+      plt.plot(size, pit.diam, 'ro', label='width')
   plt.xlabel('Square root of Trial Area (cm)')
   plt.ylabel('Depth/Width of Antlion Pits (cm)')
-  plt.show()
+  plt.legend(loc='upper right')
+  plt.savefig('depth_width.png', bbox='tight')
 if arg == 'table':
   print('Dimensions (in)\tPit Depth (cm)\tPit Width (cm)\tNearest Neighbor (cm)')
   for trial in trials: