gte/gte/__main__.py

import sys
import os
import os.path
import argparse
import logging
import time
import datetime
import numpy as np
from random import randint
from multiprocessing import Pool

logging.basicConfig(level=logging.INFO,
                    format="%(created)d %(asctime)s:%(module)s(%(lineno)d):%(levelname)s:%(message)s")
logger = logging.getLogger()
logger.debug('Logger started')

from . import turmit, rpnlib
from .turmit import Turmit
from .world import World, LivingTurmit, eval_prog
from .mutator import mutate

parser = argparse.ArgumentParser(description='Genetic Turmit Evolver')

subparsers = parser.add_subparsers(help='sub-commands help')

parser_evolve = subparsers.add_parser('evolve', help='evolving help')

parser_evolve.add_argument('--steps', '-s', type=int, metavar='N',
                           default=30000, help="Number of steps")

parser_evolve.add_argument('--repeat-eval', '-R', type=int, metavar='N',
                           default=5)

parser_evolve.add_argument('--world-height', '-y', type=int, metavar='HEIGHT',
                           default=512)
parser_evolve.add_argument('--world-width', '-x', type=int, metavar='WIDTH',
                           default=512)

parser_evolve.add_argument('--turmit-count', '-T', type=int, metavar='N',
                           default=3)

parser_evolve.add_argument('--pool-size', '-p', type=int, metavar='N',
                           default=10)

parser_evolve.add_argument('--pool-div', '-D', type=int, metavar='N',
                           default=2,
                           help='Each generation keep 1/N of the pool')

parser_evolve.add_argument('--prog-size', type=int, metavar='SIZE',
                           default=5)

parser_evolve.add_argument('--prog', '-P', type=str, metavar='EXPR',
                           default=None)

parser_evolve.add_argument('--threads', type=int, metavar='N',
                           default=os.cpu_count())

parser_evolve.add_argument('--max-generation', '-G', type=int, metavar='N',
                           default=0)

parser_evolve.add_argument('--log-progs', '-L', type=str, metavar='FILENAME',
                           default=None)

parser_evolve.add_argument('--quiet', '-q', action='store_const',
                           default=False, const=True)



parser_gen = subparsers.add_parser('generate', help='evolving help')

parser_gen.add_argument('--prog', '-P', type=str, default=None)
parser_gen.add_argument('--output', '-o', type=str, default='gte.png')
parser_gen.add_argument('--world-height', '-y', type=int, metavar='HEIGHT',
                        default=512)
parser_gen.add_argument('--world-width', '-x', type=int, metavar='WIDTH',
                        default=512)
parser_gen.add_argument('--turmit-count', '-T', type=int, metavar='N',
                        default=3)
parser_gen.add_argument('--steps', '-s', type=int, metavar='N',
                        default=30000)


args = parser.parse_args()

if 'pool_size' in args:
    # Evolver
    if args.prog is None:
        progs = [rpnlib.RpnExpr.random(args.prog_size)
                for _ in range(args.pool_size)]
    else:
        prog = rpnlib.RpnExpr.from_string(args.prog)
        progs = [mutate(prog, force=True) for _ in range(args.pool_size-1)]

    generation = 0

    pool = Pool(args.threads)

    if args.log_progs is not None:
        logprogs = open(args.log_progs, 'a')
        logprogs.write('%s Run args %s\n' % (datetime.datetime.now(),
                                             sys.argv))
        logprogs.write('%s\n' % args)
        logprogs.close()
    else:
        logger.warning('No log specified (--log-progs or -L)')

    while True:
        if args.max_generation > 0 and generation >= args.max_generation:
            exit(0)
        scores = []
        msg = 'Gen#%d Running %d eval with %d turmits and %d steps for each %d progs'
        logger.info(msg % (generation, args.repeat_eval, args.turmit_count,
                           args.steps, args.pool_size))
        genstart = time.time()
        
        # Preparing work
        works = []
        for pid, prog in enumerate(progs):
            for i in range(args.repeat_eval):
                works.append((generation, pid, prog, i, args))
        # running jobs
        res = pool.imap(eval_prog, works)
        # Processing results
        scores = [0 for _ in range(len(progs))]
        for pid, score in res:
            scores[pid] += score
        scores = [(scores[i] / args.repeat_eval, progs[i])
                  for i in range(len(progs))]

        genstop = time.time()
        # Displaying eval results
        logger.info('Generation evaluating ended in %.2fs' % (genstop - genstart))

        scores = sorted(scores, key=lambda x: x[0], reverse=True)
        for i, (score, prog) in enumerate(scores):
            logger.info('P%d %.3f : "%s"' % (i, score, str(prog)))
        if args.log_progs is not None:
            with open(args.log_progs, 'a') as logprogs:
                for i, (score, prog) in enumerate(scores):
                    msg = 'Gen #%d P%d %.3f : "%s"\n'
                    logprogs.write(msg % (generation, i, score, str(prog)))
                    logprogs.flush()
                logprogs.write('\n')

        # Split pool & mutate the one we kept
        keep_div = args.pool_div
        progs = []
        for prog in [ prog for _, prog in scores[:len(scores)//keep_div]]:
            progs.append(prog)
            for cur_mut in range(keep_div - 1):
                progs.append(mutate(prog, force=True, mutcount= 2**cur_mut))
                if len(progs) >= args.pool_size:
                    break
            if len(progs) >= args.pool_size:
                break
        for pid, prog in enumerate(progs):
            logger.debug('P%d : %s' % (pid, str(prog)))
        generation += 1
    exit(0)
else:
    # Generate
    w = World(args.world_height, args.world_width)
    prog = rpnlib.RpnExpr.from_string(args.prog)
    turmits = [LivingTurmit(world=w, prog=prog)
               for _ in range(args.turmit_count)]
    msg = 'Generating image for program %s with %d steps and %d turmits'
    logger.info(msg % (str(prog), args.steps, args.turmit_count))
    start = time.time()
    for step in range(args.steps):
        if step % 512 == 1:
            msg = 'Step %d/%d %dus/step fractdim=%.3f'
            msg %= (step, args.steps,
                    ((time.time() - start)*1000000)//step//len(turmits),
                    w.fractdim())
            logger.info(msg)
        for turmit in turmits:
            turmit()
    stop = time.time()

    msg = 'Fractdim %.3f after %d steps in %.2fs (%dus per step)'
    msg %= (w.fractdim(), args.steps,
            stop - start,
            ((stop - start)*1000000)//args.steps//args.turmit_count)
    logger.info(msg)

    w.save(args.output)
    logger.info('Image saved in %s' % args.output)

exit(0)