Optimize & enhance animation writing

Now we can specify a printf filename format and the save is made by PIL.
The image is build from array in main thread (avoid overhead of sending the
array through a pipe) and then saved by a Pool.

gte/__main__.py

@@ -7,7 +7,8 @@ import time
 import datetime
 import numpy as np
 from random import randint
-from multiprocessing import Pool, Process, Queue
+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")
@@ -91,6 +92,8 @@ 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',
@@ -104,7 +107,8 @@ parser_gen.add_argument('--gray', '-G', action='store_const',
 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")
+                        default=os.cpu_count()-1,
+                        help="Animation writting threads")
 
 args = parser.parse_args()
 
@@ -223,23 +227,16 @@ if 'pool_size' in args:
 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()
+        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)
@@ -256,21 +253,59 @@ else:
             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)
+        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)
-            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))
+            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:
-        for _ in wps:
-            data_q.put((None, None))
-        for wp in wps:
-            wp.join()
-        data_q.close()
+        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)