model.py

#!/usr/bin/env python
"""Model fit example program."""
import argparse
import collections
import itertools
import os
import time


RESULTS_DIR = 'results'
PARAMETERS = collections.OrderedDict((
    ('x', [0.1 * i for i in range(1, 11)]),
    ('y', [-0.1, 0.1]),
    ('z', ['control', 'positive', 'negative'])))
COMBINATIONS = list(itertools.product(*PARAMETERS.values()))


def n_remaining():
    """Return number of completed combinations."""
    total = len(COMBINATIONS)
    if not os.path.exists(RESULTS_DIR):
        completed = 0
    else:
        completed = len([name
                         for name in os.listdir(RESULTS_DIR)
                         if os.path.isfile(os.path.join(RESULTS_DIR, name))])
    return total - completed


def run(index):
    """Run example model fit."""
    prefix = '{}: '.format(index)
    record = collections.OrderedDict(
        zip(PARAMETERS.keys(),
            COMBINATIONS[index - 1]))
    print(prefix + 'Fitting model to parameters: {} ...'.format(
        ", ".join(["{} = {}".format(param, value)
                   for param, value in record.items()])))
    value = {'positive': 0.8, 'negative': 0.1, 'control': 0.3}
    result = record['x'] ** record['y'] / value[record['z']]
    time.sleep(5)
    path_result = os.path.join(RESULTS_DIR, '{:02d}.dat'.format(index))
    if not os.path.exists(RESULTS_DIR):
        os.mkdir(RESULTS_DIR)
    with open(path_result, 'w') as handle:
        handle.writelines([str(result)])
    print(prefix + '... done!  Saved result {:6.3f} to {}'.format(result, path_result))


def parse_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description='Model fit example program.')
    parser.add_argument('--sim', action='store_true',
                        help='Number of simulations to run')
    parser.add_argument('--remaining', action='store_true',
                        help='Remaining simulations to run')
    parser.add_argument('indices', metavar='I', nargs='*', type=int,
                        help='Model indices to run')
    return parser


def main(argv=None):
    """Entry point for model.py"""
    parser = parse_args()
    args = parser.parse_args(argv)
    if args.sim:
        print(len(COMBINATIONS))
        return
    if args.remaining:
        print(n_remaining())
        return
    if not args.indices:
        parser.error('Need at least 1 index to run!')
    for index in args.indices:
        run(index)


if __name__ == '__main__':
    main()