make_latex_table.py

#!/usr/bin/env python

import sys

from collections import OrderedDict
from argparse import ArgumentParser


def argparser():
    ap = ArgumentParser(description='make latex table from results in TSV')
    ap.add_argument(
        '--caption',
        default='TODO: Caption',
        help='table caption'
    )
    ap.add_argument(
        '--label',
        default='tab:todo-label',
        help='table label'
    )
    ap.add_argument(
        '--wide',
        action='store_true',
        help='use table*'
    )
    ap.add_argument(
        '--decimals',
        type=int,
        default=1,
        help='number of decimal places'
    )
    ap.add_argument(
        '--percentage',
        action='store_true',
        help='format values as percentages'
    )
    ap.add_argument(
        '--name-map',
        metavar='NAME:LABEL,...',
        help='map names to labels'
    )
    ap.add_argument(
        '--order',
        metavar='NAME,NAME,...',
        help='name order'
    )
    ap.add_argument(
        'data',
        nargs='+',
        metavar='NAME:FILE',
    )
    return ap


def load_tsv(fn):
    data = OrderedDict()
    with open(fn) as f:
        for ln, l in enumerate(f, start=1):
            label, value = l.split('\t')
            assert label not in data
            data[label] = float(value)
    return data


def format_value(value, args):
    if args.percentage:
        return f'{value:.{args.decimals}%}'.replace('%', r'\%')
    else:
        return f'{value:.{args.decimals}}'


def main(argv):
    args = argparser().parse_args(argv[1:])

    if args.name_map:
        name_map = dict(kv.split(':') for kv in args.name_map.split(','))
    else:
        name_map = {}

    names, data = [], {}
    for d in args.data:
        name, fn = d.split(':')
        name = name_map.get(name, name)

        names.append(name)
        assert name not in data, f'duplicate name {name}'
        data[name] = load_tsv(fn)

    # check that labels match
    for n in names[1:]:
        assert data[n].keys() == data[names[0]].keys()
    labels = list(data[names[0]].keys())

    # reorder if --order provided
    if args.order:
        assert set(names) == set(args.order.split(','))
        names = args.order.split(',')

    env_name = 'table' if not args.wide else 'table*'

    TABLE_START = '''
% NOTE: this latex table is generated by the script
% https://github.com/LumiOpen/evals/blob/master/make_latex_table.py.
% To change the table, you should probably edit and rerun the script.
\\begin{'''+env_name+'''}[t!]
\\centering
\\begin{tabular}{l|'''+'l'*len(names)+'''}
'''.strip('\n')
    print(TABLE_START)

    print(' & '.join(['']+names) + r' \\ \hline')
    for i in labels:
        values = [format_value(data[n][i], args) for n in names]
        label = i.replace('_', r'\_')
        print(' & '.join([label]+values) + r' \\')

    TABLE_END = '''
\\end{tabular}
\\caption{'''+args.caption+'''}
\\label{'''+args.label+'''}
\\end{'''+env_name+'''}
'''.strip('\n')
    print(TABLE_END)


if __name__ == '__main__':
    sys.exit(main(sys.argv))