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
import multiprocessing
from multiprocessing import Pool, Process, Pipe

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('--anim-fmt', '-F', type=str, default='GTE_ANIM_%d.png',
                        help='Save all X images (default "GTE_ANIM_%%d.png") 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=os.cpu_count()-1,
                        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
        def save_pool(im, fname):
            st1 = time.time()
            im.save(fname)
            logger.debug('Image %s saved in %.3fs' % (fname, 
                                                      time.time() - st1))

        write_pool = Pool(args.threads)
        async_res = []
        img_count = 0
    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 (args.anim_div <= 1 or
                                            step % (args.anim_div-1) == 0):
            imgid = step // (1 if args.anim_div <= 1 else args.anim_div)
            img_count += 1
            fname = args.anim_fmt % imgid
            fname = os.path.join(args.output_dir, fname)
            im = Image.fromarray(np.uint8(w._val))
            async_res.append(write_pool.apply_async(save_pool, (im, fname)))
            while async_res[0].ready():
                async_res[0].get()
                async_res.pop(0)
    
    if args.output_dir is not None:
        while async_res[0].ready():
            async_res[0].get()
            sync_res.pop(0)
        write_pool.close()
        logger.info('Waiting for animation to be written on disk')
        img_st1 = None
        while len(async_res) > 0:
            async_res[0].get()
            async_res.pop(0)
            stat_upd = 100
            if len(async_res) % stat_upd == 0:
                msg = '%d images left' % (len(async_res))
                stupd = stat_upd
                if img_st1 is None:
                    img_st1 = start
                    stupd = img_count - len(async_res)
                img_st2 = time.time()
                img_st = (img_st2 - img_st1) / stupd
                img_ts = stupd / (img_st2 - img_st1)
                etcs = ''
                if img_st1 != start:
                    etc = len(async_res) * img_st
                    etcs = ' ETC:'
                    if etc >= 60:
                        etcs += '%d min ' % (etc // 60)
                        etc %= 60
                    etcs += '%ds' % etc
                msg += ' %.1fimg/s%s' % (img_ts, etcs)
                img_st1 = img_st2
                logger.info(msg)
                    


        write_pool.join()
        #for _ in wps:
        #    data_q.put((None, None))
        #for wp in wps:
        #    wp.join()
        #data_q.close()
        logger.info('Animation written in "%s"' % args.output_dir)
    stop = time.time()

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