summarize.py

#!/usr/bin/env python3
"""Collect and summarize evaluation results"""

import argparse
import os
import sys
from pathlib import Path

import numpy as np

VALUE_CONVERSIONS = {'duration_time:u': lambda x: x / 1000 / 1000 / 1000}


def collect(result_dir, verbose=False):
    """Collect the data in result_dir and return calculated averages"""
    if not result_dir or not os.path.exists(result_dir) or not os.path.isdir(
            result_dir):
        print(f'{result_dir} is not a directory')
        return None

    result_dir = Path(result_dir)

    data = {}

    for result_file_path in result_dir.iterdir():
        if result_file_path.suffix != '.stats':
            continue
        target = result_file_path.name.split('.')[0]
        results = {}
        with open(result_file_path, 'r') as result_file:
            for line in result_file.readlines()[2:]:
                fields = line.split(';')
                key, value = fields[2], fields[0]

                try:
                    value = float(value)
                except ValueError as val_err:
                    if verbose:
                        print(
                            f'{val_err} occured during value conversion of {key}'
                        )
                    results[key] = np.nan
                    continue

                if key in VALUE_CONVERSIONS:
                    value = VALUE_CONVERSIONS[key](value)
                results[key] = value

        if target not in data:
            data[target] = {k: [v] for k, v in results.items()}
        else:
            target_data = data[target]
            for key in target_data:
                target_data[key].append(results[key])

    avgs = {
        t: {k: np.nanmean(vs)
            for k, vs in ms.items()}
        for t, ms in data.items()
    }

    return avgs


def summarize(avgs, keys=None):
    """Print a summary for each selected key of the collected averages"""
    if not keys:
        keys = ['duration_time:u']

    if not avgs:
        print('no averages to summarize')
        return False

    for key in keys:
        print(f'{key}:')
        for target in avgs:
            print(f'\t{target}: {avgs[target][key]}')

    return True


def collect_and_summarize(result_dir=None, keys=None, verbose=False):
    """Collect data and print a summary of the collected data"""
    avgs = collect(result_dir, verbose=verbose)
    if not avgs:
        return 1

    if not summarize(avgs, keys=keys):
        return 1

    return 0


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v',
                        '--verbose',
                        help='show build output',
                        action='store_true')
    parser.add_argument('-k', '--keys', help='keys to summarize', nargs='*')
    parser.add_argument('result_dir',
                        help='directory containing the results to summarize')

    _args = parser.parse_args()

    print('### Summary ###')
    sys.exit(collect_and_summarize(**vars(_args)))