add "starmap" parallel execution method, with tests

Author: JoshKarpel

pymoo/model/problem.py

@@ -360,6 +360,16 @@ def func(_x):
         if _type is None:
             [ret.append(func(x)) for x in X]
 
+        elif _type is "starmap":
+
+            if len(_params) != 1:
+                raise Exception("The starmap parallelization method must be accompanied by a starmapping callable")
+
+            params = [[X[k], calc_gradient, self._evaluate, args, kwargs] for k in range(len(X))]
+
+            starmapper = _params[0]
+            ret = starmapper(evaluate_in_parallel, params)
+
         elif _type == "threads":
 
             if len(_params) == 0:
@@ -368,11 +378,9 @@ def func(_x):
                 n_threads = _params[0]
 
             with ThreadPool(n_threads) as pool:
-                params = []
-                for k in range(len(X)):
-                    params.append([X[k], calc_gradient, self._evaluate, args, kwargs])
+                params = [[X[k], calc_gradient, self._evaluate, args, kwargs] for k in range(len(X))]
 
-                ret = np.array(pool.starmap(evaluate_in_parallel, params))
+                ret = pool.starmap(evaluate_in_parallel, params)
 
         elif _type == "dask":
 
@@ -389,7 +397,7 @@ def func(_x):
             ret = [job.result() for job in jobs]
 
         else:
-            raise Exception("Unknown parallelization method: %s (None, threads, dask)" % self.parallelization)
+            raise Exception("Unknown parallelization method: %s (should be one of: None, starmap, threads, dask)" % _type)
 
         # stack all the single outputs together
         for key in ret[0].keys():
@@ -449,6 +457,10 @@ def calc_constraint_violation(G):
         else:
             return np.sum(G * (G > 0).astype(np.float), axis=1)[:, None]
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["parallelization"] = None
+        return state
 
 # makes all the output at least 2-d dimensional
 def at_least2d(d):

test.py

@@ -0,0 +1,30 @@
+import autograd.numpy as anp
+import numpy as np
+from pymoo.algorithms.nsga2 import NSGA2
+from pymoo.factory import get_sampling, get_crossover, get_mutation, get_termination
+from pymoo.optimize import minimize
+from pymoo.util.misc import stack
+from pymoo.model.problem import Problem
+
+import multiprocessing
+
+class MyProblem(Problem):
+    def __init__(self, **kwargs):
+        super().__init__(n_var = 10, n_obj = 1, n_constr = 0, xl = -5, xu = 5,
+                         **kwargs)
+
+    def _evaluate(self, x, out, *args, **kwargs):
+        out["F"] = (x ** 2).sum(axis = -1)
+
+with multiprocessing.Pool() as pool:
+    problem = MyProblem(elementwise_evaluation = True, parallelization = ('starmap', pool.starmap))
+    r = np.random.RandomState(seed = 1)
+    X = r.random((5, problem.n_var))
+    result = problem.evaluate(X).squeeze()
+expected = (X ** 2).sum(axis = -1)
+
+print(result)
+print(expected)
+
+assert np.all(result == expected)
+print('yay')

tests/problems/test_parallel.py

@@ -1,5 +1,8 @@
 import unittest
 
+import multiprocessing
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+
 import numpy as np
 
 from pymoo.model.problem import Problem
@@ -23,6 +26,18 @@ def test_evaluation_in_threads_number(self):
         _F = MyProblemElementwise(parallelization=("threads", 2)).evaluate(X)
         self.assertTrue(np.all(np.abs(_F - F) < 0.00001))
 
+    def test_evaluation_with_multiprocessing_process_pool_starmap(self):
+        X, F = self.get_data()
+        with multiprocessing.Pool() as pool:
+            _F = MyProblemElementwise(parallelization = ("starmap", pool.starmap)).evaluate(X)
+        self.assertTrue(np.all(np.abs(_F - F) < 0.00001))
+
+    def test_evaluation_with_multiprocessing_thread_pool_starmap(self):
+        X, F = self.get_data()
+        with multiprocessing.pool.ThreadPool() as pool:
+            _F = MyProblemElementwise(parallelization = ("starmap", pool.starmap)).evaluate(X)
+        self.assertTrue(np.all(np.abs(_F - F) < 0.00001))
+
 
 class MyProblemElementwise(Problem):