FEAT: helper functions for parallel computing
@@ -6,13 +6,18 @@ Module with utilities.
from __future__ import division, print_function
-__author__ = 'Carlos Alberto Gomez Gonzalez'
-__all__ = ['Progressbar']
+__author__ = 'Carlos Alberto Gomez Gonzalez, Ralf Farkas'
+__all__ = ['Progressbar',
+ 'pool_map', 'pool_imap', 'fixed',
+ 'make_chunks']
import os
import sys
import numpy as np
+import itertools as itt
+from multiprocessing import Pool
+
sep = '-' * 80
vip_figsize = (10, 5)
vip_print_np = {'float_kind': '{0:.3f}'.format}
@@ -152,6 +157,84 @@ def eval_func_tuple(f_args):
return f_args[0](*f_args[1:])
+class FixedObj():
+ def __init__(self, v):
+ self.v = v
+
+def fixed(v):
+ return FixedObj(v)
+
+
+def pool_map(nproc, fkt, *args, msg=None, verbose=True,
+ progressbar_single=True, generator=False):
+ """
+ pool.map which uses `fixed()` for ...
+
+ """
+
+ args_r = [a.v if isinstance(a, FixedObj) else itt.repeat(a) for a in args]
+ z = zip(itt.repeat(fkt), *args_r)
+
+ if nproc == 1:
+ if progressbar_single:
+ total = len([a.v for a in args if isinstance(a, FixedObj)][0])
+ z = Progressbar(z, desc=msg, verbose=verbose, total=total)
+
+ res = map(eval_func_tuple, z)
+ if not generator:
+ res = list(res)
+
+
+ else:
+ if verbose and msg is not None:
+ print(msg+" with {} processes".format(nproc))
+ pool = Pool(processes=nproc)
+ if generator:
+ res = pool.imap(eval_func_tuple, z)
+ else:
+ res = pool.map(eval_func_tuple, z)
+ pool.close()
+ pool.join()
+
+ return res
+
+
+def pool_imap(nproc, fkt, *args, msg=None, verbose=True,
+ progressbar_single=True):
+ return pool_map(nproc, fkt, *args, msg=msg, verbose=verbose,
+ progressbar_single=progressbar_single, generator=True)
+
+
+
+
+def repeat(*args):
+ """
+
+ Examples
+ --------
+
+ # instead of using
+
+ import itertools as itt
+
+ my_fkt(itt.repeat(a), itt.repeat(b), itt.repeat(c), d, itt.repeat(e))
+
+ # you could use `repeat`:
+
+ my_fkt(*repeat(a, b, c), d, *repeat(e))
+
+ """
+ return [itt.repeat(a) for a in args]
+
+
+def make_chunks(l, n):
+ """
+ make_chunks(range(13), 3)
+ -> [[0, 3, 6, 9, 12], [1, 4, 7, 10], [2, 5, 8, 11]]
+ """
+ return [l[i::n] for i in range(n)]
+
+
class redirect_output(object):
""" Context manager for redirecting stdout/err to files"""
def __init__(self, stdout='', stderr=''):
