[aggregator] add mpi backend

aggregator.py

@@ -2,14 +2,21 @@
 
 import numpy as np
 from functools import partial
-from mypool import MyPool
+import postprocessor
 
-from postprocessor import PostProcessor
+try:
+    from mpi4py.futures import MPIPoolExecutor
+
+    mpi = True
+except ModuleNotFoundError:
+    from mypool import MyPool
+
+    mpi = False
 
 
 def _map_aux(fun, path, path_out, pp_params, run_num, **kwargs):
     try:
-        pp = PostProcessor(
+        pp = postprocessor.PostProcessor(
             path + "/" + run_num[0], run_num[1], path_out + "/" + run_num[0], pp_params
         )
     except Exception as e:
@@ -46,11 +53,19 @@ class Aggregator:
             map_fn = partial(
                 _map_aux, func, self.path, self.path_out, self.pp_params, **kwargs
             )
-
-            pool = MyPool(processes=num_process, maxtasksperchild=1)
-            result = pool.map(map_fn, run_num)
-            pool.close()
-            pool.join()
+            if mpi:
+                executor = MPIPoolExecutor(max_workers=num_process)
+                try:
+                    result = list(executor.map(map_fn, run_num, unordered=True))
+                finally:
+                    executor.shutdown()
+            else:
+                pool = MyPool(processes=num_process)
+                try:
+                    result = pool.map(map_fn, run_num)
+                finally:
+                    pool.close()
+                    pool.join()
 
         return result