rpnifs/python_rpnifs/__main__.py

#!/usr/bin/env python3
import copy
import random
import multiprocessing

try:
    import pyrpn
except (ImportError, NameError) as e:
    print("Error importing pyrpn. Try to run make.",
          file=sys.stderr)
    raise e

from PIL import Image
from PIL.PngImagePlugin import PngImageFile, PngInfo
import numpy as np
from tqdm import tqdm

from .ifs1 import IFS1
from .ifs2 import IFS2
from .ifs3 import IFS3
from .ifs4 import IFS4
from .ifs5 import IFS5

from .shape_score import ShapeScore


#steps = 1024**2
steps=512**2
#steps=1024**2
#steps*=4
#steps*=10

#pool_sz = 30
#best_sz = 6

pool_sz = 32
best_sz = 8

n_mutation = 10
outfile='/tmp/rpnifs2_%02d.png'
outfile_grp='/tmp/rpnifs2_%02d_grp.png'
#world = IFS1.get_world()
#pool = [IFS1(init_sz=3, world=world) for _ in range(pool_sz)]
#world = IFS2.get_world()
#pool = [IFS2(init_sz=3, world=world) for _ in range(pool_sz)]
#world = IFS3.get_world()
#pool = [IFS3(init_sz=3, world=world) for _ in range(pool_sz)]
#world = IFS4.get_world()
#pool = [IFS4(init_sz=3, world=world) for _ in range(pool_sz)]
world = IFS5.get_world()
#pool = [IFS5(init_sz=3, world=world) for _ in range(pool_sz)]
pool = [IFS5(init_sz=6, world=world) for _ in range(pool_sz)]
print('POOL ready')

def stepper(ifs):
    for _ in range(steps):
        ifs.step()
    return ifs

def get_score(world):
    shapescore = ShapeScore(world, jobs=1)
    score = shapescore.score()
    return score, shapescore

generation=0
mp_pool = multiprocessing.Pool(2)
while True:
    generation+=1
    print('='*78)
    print(' '*10 + 'GENERATION #%d' % generation)
    print('='*78)

    scores = []
    i=0
    #for ifs in pool:
    #for ifs in tqdm(pool, unit='ifs'):
    #    ifs.raz_world()
    #    #for _ in tqdm(range(steps), total=steps, unit_scale=True):
    #    for _ in range(steps):
    #        ifs.step()
    pool = [ifs for ifs in tqdm(mp_pool.imap(stepper, pool),
                                total=len(pool),
                                unit='ifs')]
    
    for i, (score,shapescore) in tqdm(enumerate(mp_pool.imap(get_score,
                                                             [ifs._world for ifs in pool])),
                                      unit='ifs', total=len(pool)):
        ifs = pool[i]

        #score = ifs.score()
        scores.append((ifs, score))
        #print('%02d) %5.3f %s' % (i, score, ifs))
        print('%02d) %5.3f' % (i, score))
        im = Image.fromarray(ifs.get_image())
        im.save(outfile % i)
        shapescore.result_image(outfile_grp % i)

        im_png = PngImageFile(outfile % i)
        metadata = PngInfo()
        metadata.add_text('Score', '%5.3f' % score)
        metadata.add_text('Generation', '%04d' % generation)
        metadata.add_text('Pool_position', '%02d' % i)
        for key, text in ifs.expr_dict().items():
            metadata.add_text(key, text)
        im_png.save(outfile % i, pnginfo=metadata)

    # Selecting best
    scores = sorted(scores, key=lambda v: 0 if v[1] == np.nan else v[1],
                    reverse=True)
    best = scores[:best_sz]
    for i, (ifs, score) in enumerate(best):
        print('%02d) %5.3f %s' % (i, score, ifs))

    # Saving best
    im = Image.fromarray(best[0][0].get_image())
    outfile_gen = '/tmp/gen_%03d.png'
    im.save(outfile_gen % generation)

    im_png = PngImageFile(outfile_gen % generation)
    metadata = PngInfo()
    metadata.add_text('Score', '%5.3f' % score)
    metadata.add_text('Generation', '%04d' % generation)
    metadata.add_text('Pool_position', '%02d' % 1)
    for key, text in ifs.expr_dict().items():
        metadata.add_text(key, text)
    im_png.save(outfile_gen % generation, pnginfo=metadata)

    # Mutating best
    pool = [b for b, _ in best]
    for ifs, _ in best:
        nmut_max = (pool_sz//best_sz)-1
        for nmut in range(nmut_max):
            # needs >= IFS5
            again = 0
            while True:
                # half random mutation and half mutation by adding two IFS
                if again > 5 or nmut < nmut_max / 2:
                    new = copy.copy(ifs)
                    new.mutation(n_mutation)
                else:
                    new =  ifs + random.choice(best)[0]

                for cur in pool:
                    if cur - new == 0:
                        break
                else:
                    pool.append(new)
                    break 
                again += 1
            ###
            #new = copy.copy(ifs)
            #new.mutation(n_mutation)
            #pool.append(new)