improved data.py table generator, using vtable

1 files changed, 16 insertions, 8 deletions

diff --git a/data.py b/data.py
index d3ec25f..9aead79 100644
--- a/data.py
+++ b/data.py
@@ -116,7 +116,8 @@ arg = argv[1]
 if arg == 'img':
   for trial in trials:
     trial.plot(save=True)
-elif arg == 'nei':
+
+if arg == 'nei':
   x = []
   y = []
   for trial in trials:
@@ -136,7 +137,8 @@ elif arg == 'nei':
   plt.plot(x, y, 'bo')
   plt.plot(x, poly1d(polyfit(x, y, 1))(x))
   plt.savefig('nearest_neighbor.png', bbox_inches='tight')
-elif arg == 'depwid':
+
+if arg == 'depwid':
   depths, widths, sizes = [], [], []
   for trial in trials:
     size = sqrt(trial.size[0]*trial.size[1])
@@ -153,19 +155,25 @@ elif arg == 'depwid':
   plt.legend(loc='upper right')
   plt.text(10, 7, 'R^2 = %.3f\np=%.3f' % (pearsonr(sizes,depths)[0], pearsonr(sizes,depths)[1]), ha='center', va='center')
   plt.savefig('depth_width.png', bbox_inches='tight')
+
 if arg == 'table':
-  print('\\table{')
-  print('Dimensions (in)& Pit Depth (cm)& Pit Width (cm)& Nearest Neighbor (cm)')
+  print('\\vtable{')
+  table = ['Dimensions (in)', 'Pit Depth (cm)', 'Pit Width (cm)', 'Nearest Neighbor (cm)',''];
+  lastsize = 0;
   for trial in trials:
     size = trial.size
+    table[0] += '&\\multispan{' + str(len(trial.pits)) + '}\\vfil\\line{\\hfil' + '$\\times$'.join([str(el) for el in size]) + '\\hfil}\\vfil\\hrule'
     nei = trial.nearest_neighbor()
     for pitind in range(len(trial.pits)):
       pit = trial.pits[pitind]
-      print('& '.join(['x'.join([str(el) for el in size]), '%.1f' % pit.depth, '%.1f' % pit.diam, '%.2f' % nei[pitind]])+'\\cr\\noalign{\\hrule}')
-  print('}')
+      table[1] += '&' + '%.1f' % pit.depth
+      table[2] += '&' + '%.1f' % pit.diam
+      table[3] += '&' + '%.2f' % nei[pitind]
+  print('\\cr\\noalign{\\vrule}\n'.join(table) + '}', end='');
+
 if arg == 'deathtable':
   print('\\table{')
   print('Trial Size& Date& Introduced& Deaths& Pits formed\\cr\\noalign{\\hrule}')
   for trial in trials:
-    print('& '.join(['x'.join([str(el) for el in trial.size]), str(trial.date), str(trial.intro), str(trial.dead), str(len(trial.pits))])+'\\cr\\noalign{\\hrule}')
-  print('}')
+    print('& '.join(['$\\times$'.join([str(el) for el in trial.size]), str(trial.date), str(trial.intro), str(trial.dead), str(len(trial.pits))])+'\\cr\\noalign{\\hrule}')
+  print('}', end='')