From 55fae0cf01ad6f73f1f48beca8de7ccc8fd29898 Mon Sep 17 00:00:00 2001
From: Jameson Graef Rollins <jrollins@finestructure.net>
Date: Tue, 14 Aug 2018 08:12:30 -0700
Subject: [PATCH] Accept pool argument, instead of instantiating
 multiprocessing

This lets users provide a pool with a map() method, rather than
handling this via multiprocessing internally.  This allows for more
flexibility for users to provide alternate pools (such as from
ipyparallel).
---
 pyswarm/pso.py | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/pyswarm/pso.py b/pyswarm/pso.py
index 1811cd9..9ec1915 100644
--- a/pyswarm/pso.py
+++ b/pyswarm/pso.py
@@ -18,7 +18,7 @@ def _cons_f_ieqcons_wrapper(f_ieqcons, args, kwargs, x):
     
 def pso(func, lb, ub, ieqcons=[], f_ieqcons=None, args=(), kwargs={}, 
         swarmsize=100, omega=0.5, phip=0.5, phig=0.5, maxiter=100, 
-        minstep=1e-8, minfunc=1e-8, debug=False, processes=1,
+        minstep=1e-8, minfunc=1e-8, debug=False, pool=None,
         particle_output=False):
     """
     Perform a particle swarm optimization (PSO)
@@ -68,9 +68,10 @@ def pso(func, lb, ub, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
     debug : boolean
         If True, progress statements will be displayed every iteration
         (Default: False)
-    processes : int
-        The number of processes to use to evaluate objective function and 
-        constraints (default: 1)
+    pool : object
+        An instantiated multiprocessing.Pool or similar with a map
+        method for evaluating objective function and constraints
+        (default: None)
     particle_output : boolean
         Whether to include the best per-particle position and the objective
         values at those.
@@ -116,11 +117,6 @@ def pso(func, lb, ub, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
         cons = partial(_cons_f_ieqcons_wrapper, f_ieqcons, args, kwargs)
     is_feasible = partial(_is_feasible_wrapper, cons)
 
-    # Initialize the multiprocessing module if necessary
-    if processes > 1:
-        import multiprocessing
-        mp_pool = multiprocessing.Pool(processes)
-        
     # Initialize the particle swarm ############################################
     S = swarmsize
     D = len(lb)  # the number of dimensions each particle has
@@ -137,9 +133,9 @@ def pso(func, lb, ub, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
     x = lb + x*(ub - lb)
 
     # Calculate objective and constraints for each particle
-    if processes > 1:
-        fx = np.array(mp_pool.map(obj, x))
-        fs = np.array(mp_pool.map(is_feasible, x))
+    if pool:
+        fx = np.array(pool.map(obj, x))
+        fs = np.array(pool.map(is_feasible, x))
     else:
         for i in range(S):
             fx[i] = obj(x[i, :])
@@ -179,9 +175,9 @@ def pso(func, lb, ub, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
         x = x*(~np.logical_or(maskl, masku)) + lb*maskl + ub*masku
 
         # Update objectives and constraints
-        if processes > 1:
-            fx = np.array(mp_pool.map(obj, x))
-            fs = np.array(mp_pool.map(is_feasible, x))
+        if pool:
+            fx = np.array(pool.map(obj, x))
+            fs = np.array(pool.map(is_feasible, x))
         else:
             for i in range(S):
                 fx[i] = obj(x[i, :])