plot_naacl_bar.py

from __future__ import division
from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
import os
import string
import operator


def sorted_mean(l):
    return sorted(l, key=lambda key: np.mean(np.array(key[1])[:,1]))

# change data here
data_fname = "vary_len_scores.tsv"
data_lines = map(lambda x: x.strip().split('\t'), open(data_fname, 'r').readlines())
dataset_intmap = {}
data_map = {}
for line in data_lines:
    # so that we can look up the name later
    if line[0] not in dataset_intmap.values():
        data_map[line[0]] = []
        dataset_intmap[len(dataset_intmap)] = line[0]
    data_map[line[0]].append([int(line[1]), float(line[2])])

num_datum = len(data_map)
print(num_datum)

bar_width = 1/(num_datum+1)

fig, ax = plt.subplots()
start_loc = 0
rects = []
for d, data in sorted_mean(data_map.iteritems()):
    data = np.array(data)
    labels = data[:,0]
    values = data[:,1]
    rect = ax.bar(np.arange(len(values))+start_loc, height=values, width=bar_width)
    rects.append(rect)
    start_loc += bar_width

lgd = ax.legend(map(lambda r: r[0], rects), data_map.keys(), bbox_to_anchor=(1.0, 1.0))

ax.set_ylabel('Score')
ax.set_xlabel('Distance')
example_datum = np.array(data_map[dataset_intmap[0]])
print(example_datum[:,0])
ax.set_xticks(np.arange(len(example_datum)) + bar_width*1.5)
ax.set_xticklabels(example_datum[:,0])

# ax.set_title("Totals")
fig.tight_layout()

# change filename here
fig_name = 'fine.pdf'
plt.savefig(fig_name, bbox_extra_artists=(lgd,), bbox_inches='tight')

plt.show()