def add_columns(self, other, columns):
"""Given another document `other' find the matching tables and add
the columns in `columns' to the table"""
# TODO: This is similar to join, refactoring should be possible
d = LatexDoc()
for it1 in self.items:
labels1 = set(it1.row_labels)
for it2 in other.items:
labels2 = set(it2.row_labels)
common = labels1.intersection(labels2)
if common:
t = Table()
common = list(common)
projected_it2 = it2.select_columns(columns)
ind1 = utilities.indices(it1.row_labels, common)
ind2 = utilities.indices(projected_it2.row_labels, common)
rows1 = [it1[row].values for row in ind1]
rows2 = [projected_it2[row].values for row in ind2]
new_rows = [r1 + r2 for (r1, r2) in zip(rows1, rows2)]
# print "new rows: ", new_rows
t.column_labels = it1.column_labels + columns
for (prob_idx, prob) in enumerate(common):
t.add_row_in_order(new_rows[prob_idx], prob)
d.add_item(t)
return d
def join(self, other, op=operator.concat):
"""Join all the items that has at least one common row label. Pivot is the first column"""
# TODO: Pivot as other columns?
d = LatexDoc()
for it1 in self.items:
if type(it1) != Table:
raise TypeError
labels1 = set(it1.row_labels)
for it2 in other.items:
if type(it2) != Table:
raise TypeError
labels2 = set(it2.row_labels)
common = labels1.intersection(labels2)
if common:
t = Table()
common = list(common)
ind1 = utilities.indices(it1.row_labels, common)
ind2 = utilities.indices(it2.row_labels, common)
rows1 = [it1[row].values for row in ind1]
rows2 = [it2[row].values for row in ind2]
new_rows = [op(r1, r2) for (r1, r2) in zip(rows1, rows2)]
t.column_labels = op(it1.column_labels, it2.column_labels[1:])
for (prob_idx, prob) in enumerate(common):
t.add_row_in_order(new_rows[prob_idx], prob)
d.add_item(t)
return d
def names_to_indices(self, names):
"""Return the indices of the column names in the rows"""
return [ind - 1 for ind in utilities.indices(self.column_labels, names)]