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, Process, Queue

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('--max', '-m', action='store_const',
                           default=False, const=True,
                           help="Keep maximum score instead of average")

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('--exp-mutate', '-E', action='store_const',
                           default=False, const=True,
                           help="If True make 2**copy_num mutation")

parser_evolve.add_argument('--mod-steps', '-M', action='store_const',
                           default=False, const=True,
                           help="If True make steps change given try number")

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

parser_evolve.add_argument('--no-diagonals', '-d', action='store_const',
                           default=False, const=True)

parser_evolve.add_argument('--target', '-t', type=float, metavar="FLOAT",
                           default=None,
                           help="0 < Float <= 2.0 for targeted score")

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('--output-dir', '-A', type=str, default=None,
                        help="Save images for animation in this directory")
parser_gen.add_argument('--anim-div', '-D', type=int, default=2,
                        help='Save all X images (default 2) when -O given')
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)
parser_gen.add_argument('--gray', '-G', action='store_const',
                        default=False, const=True)
parser_gen.add_argument('--no-diagonals', '-d', action='store_const',
                        default=False, const=True)
parser_gen.add_argument('--threads', '-t', type=int, metavar='N',
                        default=2, help="Animation writting threads")

args = parser.parse_args()

if 'pool_size' in args:
    # Evolver
    if args.prog is None:
        prog = rpnlib.RpnExpr.random(args.prog_size)
        progs = [mutate(prog, force=True) for _ in range(args.pool_size-1)]
        #progs = [rpnlib.RpnExpr.random(args.prog_size)
        #        for _ in range(args.pool_size)]
    else:
        prog = rpnlib.RpnExpr.from_string(args.prog)
        if len(prog) == 0:
            prog = rpnlib.RpnExpr([rpnlib.RpnSymbol.from_string('0')])
        progs = [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
        if args.max:
            scores = [ [] for _ in range(len(progs))]
            for pid, score in res:
                scores[pid].append(score)
            scores = [(max(scores[i]), progs[i]) for i in range(len(progs))]
        else:
            scores = [0 for _ in range(len(progs))]
            sinfos = [dict() for _ in range(len(progs))]
            for pid, score, sinfo in res:
                scores[pid] += score
                # sum all score infos
                for k, s in sinfo.items():
                    if k not in sinfos[pid]:
                        sinfos[pid][k] = 0
                    sinfos[pid][k] += s
            # avg all infos
            for i in range(len(sinfos)):
                sinfo = sinfos[i]
                for k in sinfo:
                    sinfo[k] /= args.repeat_eval
                scores[i] /= args.repeat_eval
            scores = [(scores[i], progs[i], sinfos[i])
                      for i in range(len(progs))]

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

        if args.target is None:
            scores = sorted(scores, key=lambda x: x[0], reverse=True)
        else:
            scores = sorted(scores, key=lambda x: abs(x[0] - args.target),)
        for i, (score, prog, sinfo) in enumerate(scores):
            logger.info('P%d %.3f:(D:%.3f,F:%.3f) : "%s"' % (i, score,
                                                             sinfo['D'],
                                                             sinfo['F'],
                                                             str(prog)))
        if args.log_progs is not None:
            with open(args.log_progs, 'a') as logprogs:
                for i, (score, prog, sinfo) in enumerate(scores):
                    msg = 'Gen #%d P%d %.3f(D:%.3f,F:%.3f): "%s"\n'
                    logprogs.write(msg % (generation, i, score, sinfo['D'],
                                          sinfo['F'], 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]]:
            for p2 in progs:
                if prog == p2:
                    prog = mutate(prog, force=True)
                    break
            progs.append(prog)
            for cur_mut in range(keep_div - 1):
                if args.exp_mutate:
                    progs.append(mutate(prog, force=True, mutcount=2**cur_mut))
                else:
                    progs.append(mutate(prog, force=True, mutcount=1))
                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
    if args.output_dir is not None:
        #from PIL import Image
        import scipy
        def save_anim(data_q):
            while True:
                wval, fname = data_q.get()
                if wval is None:
                    return
                #im = Image.fromarray(wval, 'RGB')
                #im.save(fname)
                scipy.misc.imsave(fname, wval)
                logger.info('Written %s / %d' % (fname, args.steps//args.anim_div))

        data_q = Queue()
        wps = [Process(target=save_anim, args=(data_q,))
               for _ in range(args.threads)]
        for wp in wps:
            wp.start()
    w = World(args.world_height, args.world_width, gray=args.gray)
    prog = rpnlib.RpnExpr.from_string(args.prog)
    turmits = [LivingTurmit(world=w, prog=prog, diag=not args.no_diagonals)
               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()
        if args.output_dir is not None and step % args.anim_div == 0:
            fname = 'GTE_ANIM_%d.png' % (step/args.anim_div)
            fname = os.path.join(args.output_dir, fname)
            data_q.put((w._val, fname))
        #if args.output_dir is not None and step % args.anim_div == 0:
        #    #w.save('/tmp/anim/GTE_ANIM_%d.png' % (step/10))
        #    fname = 'GTE_ANIM_%d?png' % (step/args.anim_div)
        #    w.save(os.path.join(args.output_dir, fname))
    
    if args.output_dir is not None:
        for _ in wps:
            data_q.put((None, None))
        for wp in wps:
            wp.join()
        data_q.close()
    stop = time.time()

    score_dir = sum([t.dirvar() for t in turmits]) / len(turmits)
    msg = 'DirVar %.3f Fractdim %.3f after %d steps in %.2fs (%dus per step)'
    msg %= (score_dir, 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)