Harmonize PSO and NM optimizers (#115)

README.md

@@ -64,6 +64,8 @@ Does other simulation software exist in Python? Sure, but most of them hold your
         > Validated the use of Python 3.11. Efficiency gains in simulation of jump processes. Ommitted dependency on Numba. All changes related to publishing our software manuscript in Journal of Computational Science. Improved nomenclature in model defenition.
     - Version 0.2.5 (2024-10-08, PR #106)
         > Validated the use of Python 3.12. Validated pySODM on macOS Sonoma 14.5. 'draw functions' only have 'parameters' as mandatory input, followed by an arbitrary number of additional parameters (PR #75). Tutorial environment can now be found in `tutorial_env.yml` and was renamed `PYSODM-TUTORIALS` (PR #76). Users can choose when the simulation starts when calibrating a model (PR #92). Initial model states can now be a function returning a dictionary of states. This initial condition function can have arguments, which become part of the model's parameters, and can therefore be optimised (PR #99).  Deprecation of legacy 'warmup' parameter (PR #100). Change 'states' --> 'initial_states' as input needed to initialize a model (PR #102).
+    - Version 0.2.6 (in preparation)
+        > Harmonize NM and PSO optimizer output (PR #115).
 - Version 0.1 (2022-12-23, PR #14)
     > Application pySODM to three use cases. Documentation website. Unit tests for ODE, JumpProcess and calibration. 
     - Version 0.1.1 (2023-01-09, PR #20)

docs/enzyme_kinetics.md

@@ -176,11 +176,11 @@ from pySODM.optimization import pso, nelder_mead
 if __name__ == '__main__':
 
     # PSO
-    theta = pso.optimize(objective_function, swarmsize=50, max_iter=30, processes=2)[0]    
+    theta, _ = pso.optimize(objective_function, swarmsize=50, max_iter=30, processes=2)    
 
     # Nelder-mead
     step = len(theta)*[0.05,]
-    theta = nelder_mead.optimize(objective_function, theta, step, processes=2, max_iter=30)[0]
+    theta, _ = nelder_mead.optimize(objective_function, theta, step, processes=2, max_iter=30)
 ```
 
 We find the following estimates for our parameters,

docs/optimization.md

@@ -163,7 +163,7 @@
 >    Perform a [Nelder-Mead minimization](https://en.wikipedia.org/wiki/Nelder%E2%80%93Mead_method).
 
 >    **Parameters:**
->    * **func** (function) - Callable function or class representing the objective function to be minimized. Recommended `pySODM.optimization.log_posterior_probability`.
+>    * **func** (function) - Callable function or class representing the objective function to be minimized. Recommended using `pySODM.optimization.log_posterior_probability`. 
 >    * **x_start** (list or 1D np.ndarray) - Starting estimate for the search algorithm. Length must equal the number of provided bounds. 
 >    * **step** (list or 1D np.ndarray) - (Relative) size of the initial search simplex. 
 >    * **bounds** (list) - optional - The bounds of the design variable(s). In form `[(lb_1, ub_1), ..., (lb_n, ub_n)]`. If class `log_posterior_probability` is used as `func`, it already contains bounds. If bounds are provided these will overwrite the bounds available in the 'log_posterior_probability' object.
@@ -181,7 +181,8 @@
 >    * **sigma** (float) - optional - Shrink coefficient
 
 >    **Returns:**
->    * **theta** (list) - Position 0: Estimated parameters. Position 1: corresponding score of `func`.
+>    * **theta** (list) - Optimised parameter values.
+>    * **score** (float) - Corresponding corresponding objective function value.
 
 ## pso.py
 
@@ -190,12 +191,12 @@
 >    Perform a [Particle Swarm Optimization](https://en.wikipedia.org/wiki/Particle_swarm_optimization).
 
 >    **Parameters:**
->    * **func** (function) - Callable function or class representing the objective function to be minimized. Recommended `log_posterior_probability`.
+>    * **func** (function) - Callable function or class representing the objective function to be minimized. Recommended using `log_posterior_probability`.
 >    * **bounds** (list) - optional - The bounds of the design variable(s). In form `[(lb_1, ub_1), ..., (lb_n, ub_n)]`. If class `log_posterior_probability` is used as `func`, it already contains bounds. If bounds are provided these will overwrite the bounds available in the 'log_posterior_probability' object.
->    * **ieqcons** (list) - A list of functions of length n such that ```ieqcons[j](x,*args) >= 0.0``` in a successfully optimized problem
->    * **f_ieqcons** (function) - Returns a 1-D array in which each element must be greater or equal to 0.0 in a successfully optimized problem. If f_ieqcons is specified, ieqcons is ignored
 >    * **args** (tuple) - optional - Additional arguments passed to objective function.
 >    * **kwargs** (dict) - optional - Additional keyworded arguments passed to objective function. Example use: To compute our log posterior probability (class `log_posterior_probability`) with the 'RK45' method, we must change the `method` argument of the `sim` function, which is called in `log_posterior_probability`. To achieve this, we can supply the keyworded argument `simulation_kwargs` of `log_posterior_probability`, which passes its arguments on to the `sim` function. To this end, use `kwargs={'simulation_kwargs':{'method': 'RK45'}}`.
+>    * **ieqcons** (list) - A list of functions of length n such that ```ieqcons[j](x,*args) >= 0.0``` in a successfully optimized problem
+>    * **f_ieqcons** (function) - Returns a 1-D array in which each element must be greater or equal to 0.0 in a successfully optimized problem. If f_ieqcons is specified, ieqcons is ignored
 >    * **processes** (int) - optional - Number of cores to use.
 
 >   **Hyperparameters:**
@@ -207,11 +208,11 @@
 >    * **phip** (float) - optional - Tendency to search away from the particles best known position.  A higher value means each particle has less confidence in it's own best value.
 >    * **phig** (float) - optional - Tendency to search away from the swarm's best known position. A higher value means each particle has less confidence in the swarm's best value.
 >    * **debug** (bool) - optional - If True, progress statements will be displayed every iteration
->    * **particle_output** (bool) - optional - If True, function additionally returns the best particles position and objective function score
 >    * **transform_pars** (func) - optional - Transform the parameter values. E.g. to integer values or to map to a list of possibilities
 
 >    **Returns:**
->    * **theta** (list) - Position 0: Estimated parameters. Position 1: Corresponding score of `func`. If `particle_output==True` then: Position 3: The best known position per particle. Position 4: Vorresponding score of `func`.
+>    * **theta** (list) - Optimised parameter values.
+>    * **score** (float) - Corresponding corresponding objective function value.
 
 ## mcmc.py
 

docs/workflow.md

@@ -246,7 +246,7 @@ if __name__ == '__main__':
     # Initial guess
     theta = [0.35,]
     # Run Nelder-Mead optimisation
-    theta = nelder_mead.optimize(objective_function, theta, step=[0.10,], processes=1, max_iter=10)[0]
+    theta, _ = nelder_mead.optimize(objective_function, theta, step=[0.10,], processes=1, max_iter=10)
 ```
 We find an optimal value of {math}`\beta \pm 0.27`. We can then asses the goodness-of-fit by updating the dictionary of model parameters with the newly found value for `beta`, simulating the model and visualizing the model prediction and dataset.
 

src/pySODM/optimization/nelder_mead.py

@@ -3,44 +3,55 @@
 import multiprocessing as mp
 from functools import partial
 
-'''
-    Pure Python/Numpy implementation of the Nelder-Mead algorithm with multiprocessing support.
+"""
+    A pure Python/Numpy implementation of the Nelder-Mead simplex algorithm with multiprocessing support.
     Reference: https://en.wikipedia.org/wiki/Nelder%E2%80%93Mead_method
-'''
+"""
 
 def _obj_wrapper(func, args, kwargs, x):
+    """
+    A posterior probability must be maximised; so we switch signs
+    """
     return -func(x, *args, **kwargs)
 
 def optimize(func, x_start, step,
-                bounds=None, args=(), kwargs={}, processes=1, no_improve_thr=1e-6,
+                bounds=None, args=(), kwargs={},
+                processes=1, no_improve_thr=1e-6,
                 no_improv_break=100, max_iter=1000,
                 alpha=1., gamma=2., rho=-0.5, sigma=0.5):
     """
-    Perform a Nelder-Mead minimization
+    Perform a Nelder-Mead simplex optimization -- minimization of an objective function
 
     Parameters
     ==========
+
     func : callable function or class 'log_posterior_probability' (~/src/optimization/objective_functions.py)
         The objective function to be minimized
     x_start: list or 1D np.ndarray
         Starting estimate for the search algorithm. Length must equal the number of provided bounds.
-     step: list or 1D np.ndarray
-        Size of the initial search simplex       
-    bounds: tuple array
-        The bounds of the design variable(s). In form [(lower, upper), ..., (lower, upper)]
-        Class 'log_posterior_probability' automatically contains bounds. If bounds are provided these overwrite the bounds available in the 'log_posterior_probability' object.
+    step: list or 1D np.ndarray
+        Size of the initial search simplex   
+
+    Optional
+    ========    
+
+    bounds: list containing tuples
+        The bounds of the parameter(s). In the form: [(lower, upper), ..., (lower, upper)].
+        If `func` is pySODM class 'log_posterior_probability', bounds are automatically inferred.
+        If bounds are provided anyway these overwrite the bounds available in the 'log_posterior_probability' object.
     args : tuple
         Additional arguments passed to objective function
-        (Default: empty tuple)
     kwargs : dict
         Additional keyword arguments passed to objective function
 
     Returns
     =======
-    theta: list
-        Position 0: estimated parameters
-        Position 1: corresponding score 
 
+    theta: list
+        optimised parameter values
+    
+    score: float
+        corresponding objective function value
     """
 
     ##################
@@ -52,7 +63,7 @@ def optimize(func, x_start, step,
             bounds = func.expanded_bounds
         except:
             raise Exception(
-                "'func' does not appear to be a pySODM model: 'expanded_bounds' not found. Provide bounds directly to `nelder_mead.optimize()`"
+                "please provide 'bounds'."
             )
 
     # Input check bounds
@@ -112,7 +123,7 @@ def optimize(func, x_start, step,
             score.append(obj(x))
     # Check bounds
     for i,x in enumerate(mp_args):
-        for j, x_j in enumerate(x):
+        for j, _ in enumerate(x):
             if ((x[j] < lb[j]) | (x[j] > ub[j])):
                 score[i] = np.inf 
     # Construct vector of inputs and scores
@@ -139,7 +150,7 @@ def optimize(func, x_start, step,
         # Break after max_iter
         if max_iter and iters >= max_iter:
             print('Maximum number of iteration reached. Quitting.\n')
-            return res[0]
+            return res[0][0], res[0][1]
         iters += 1
         # Break if no improvements for too long
         if best < prev_best - no_improve_thr:
@@ -149,7 +160,7 @@ def optimize(func, x_start, step,
             no_improv += 1
         if no_improv >= no_improv_break:
             print('Maximum number of iterations without improvement reached. Quitting.\n')
-            return res[0]
+            return res[0][0], res[0][1]
 
         ################
         ## Reflection ##
@@ -164,7 +175,7 @@ def optimize(func, x_start, step,
         xr = x0 + alpha*(x0 - res[-1][0])
         rscore = obj(xr)
         # Check bounds
-        for j, xr_j in enumerate(xr):
+        for j, _ in enumerate(xr):
             if ((xr[j] < lb[j]) | (xr[j] > ub[j])):
                 rscore = np.inf 
         if res[0][1] <= rscore < res[-2][1]:
@@ -180,7 +191,7 @@ def optimize(func, x_start, step,
             xe = x0 + gamma*(x0 - res[-1][0])
             escore = obj(xe)
             # Check bounds
-            for j, xe_j in enumerate(xe):
+            for j, _ in enumerate(xe):
                 if ((xe[j] < lb[j]) | (xe[j] > ub[j])):
                     escore = np.inf 
             if escore < rscore:
@@ -199,7 +210,7 @@ def optimize(func, x_start, step,
         xc = x0 + rho*(x0 - res[-1][0])
         cscore = obj(xc)
         # Check bounds
-        for j, xc_j in enumerate(xc):
+        for j, _ in enumerate(xc):
             if ((xc[j] < lb[j]) | (xc[j] > ub[j])):
                 escore = np.inf 
         if cscore < res[-1][1]:
@@ -229,7 +240,7 @@ def optimize(func, x_start, step,
                 score.append(obj(x))
         # Check bounds
         for i,x in enumerate(mp_args):
-            for j, x_j in enumerate(x):
+            for j, _ in enumerate(x):
                 if ((x[j] < lb[j]) | (x[j] > ub[j])):
                     score[i] = np.inf 
         # Construct nres

src/pySODM/optimization/pso.py

@@ -3,6 +3,9 @@
 from functools import partial
 
 def _obj_wrapper(func, args, kwargs, x):
+    """
+    A posterior probability must be maximised; so we switch signs
+    """
     return -func(x, *args, **kwargs)
 
 
@@ -22,34 +25,39 @@ def _cons_f_ieqcons_wrapper(f_ieqcons, args, kwargs, x):
     return np.array(f_ieqcons(x, *args, **kwargs))
 
 
-def optimize(func, bounds=None, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
-        processes=1, swarmsize=100, max_iter=100, minstep=1e-12, minfunc=1e-12, omega=0.8, phip=0.8, phig=0.8, 
-        debug=False, particle_output=False, transform_pars=None):
+def optimize(func, bounds=None, args=(), kwargs={},
+             ieqcons=[], f_ieqcons=None, processes=1,
+             swarmsize=100, max_iter=100, minstep=1e-12,
+             minfunc=1e-12, omega=0.8, phip=0.8, phig=0.8,
+             debug=False, transform_pars=None):
     """
-    Perform a particle swarm optimization (PSO)
+    Perform a particle swarm optimization (PSO) -- minimization of an objective function
 
     Parameters
     ==========
+
     func : callable function or class 'log_posterior_probability' (~/src/optimization/objective_functions.py)
         The objective function to be minimized
-    bounds: tuple array
-        The bounds of the design variable(s). In form [(lower, upper), ..., (lower, upper)]
-        Class 'log_posterior_probability' automatically contains bounds. If bounds are provided these overwrite the bounds available in the 'log_posterior_probability' object.
+    bounds: list containing tuples
+        The bounds of the parameter(s). In the form: [(lower, upper), ..., (lower, upper)].
+        If `func` is pySODM class 'log_posterior_probability', bounds are automatically inferred.
+        If bounds are provided anyway these overwrite the bounds available in the 'log_posterior_probability' object.
 
     Optional
     ========
+
+    args : tuple
+        Additional arguments passed to objective function
+        (Default: empty tuple)
+    kwargs : dict
+        Additional keyword arguments passed to objective function
     ieqcons : list
         A list of functions of length n such that ieqcons[j](x,*args) >= 0.0 in 
         a successfully optimized problem (Default: [])
     f_ieqcons : function
         Returns a 1-D array in which each element must be greater or equal 
         to 0.0 in a successfully optimized problem. If f_ieqcons is specified, 
         ieqcons is ignored (Default: None)
-    args : tuple
-        Additional arguments passed to objective function
-        (Default: empty tuple)
-    kwargs : dict
-        Additional keyword arguments passed to objective function
     phig : scalar
         Scaling factor to search away from the swarm's best known position
         (Default: 0.5)
@@ -67,32 +75,26 @@ def optimize(func, bounds=None, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
     processes : int
         The number of processes to use to evaluate objective function and 
         constraints (default: 1)
-    particle_output : boolean
-        Whether to include the best per-particle position and the objective
-        values at those.
     transform_pars : None / function
         Transform the parameter values. E.g. to integer values or to map to
         a list of possibilities.
 
     Returns
     =======
-    g : array
-        The swarm's best known position (optimal design)
-    f : scalar
-        The objective value at ``g``
-    p : array
-        The best known position per particle
-    pf: arrray
-        The objective values at each position in p
 
+    theta: list
+        optimised parameter values
+    
+    score: float
+        corresponding objective function value
     """
 
     if not bounds:
         try:
             bounds = func.expanded_bounds
         except:
             raise Exception(
-                "'func' does not appear to be a pySODM model: 'expanded_bounds' not found. Provide bounds directly to `pso.optimize()`"
+                "please provide 'bounds'."
             )
 
     lb, ub = [], []
@@ -146,7 +148,6 @@ def optimize(func, bounds=None, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
         import multiprocessing
         mp_pool = multiprocessing.Pool(processes)
 
-
     # Initialize the particle swarm ############################################
     S = swarmsize
     D = len(lb)  # the number of dimensions each particle has
@@ -241,18 +242,12 @@ def optimize(func, bounds=None, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
                 print('Stopping search: Swarm best objective change less than {:}'\
                     .format(minfunc))
                 sys.stdout.flush()
-                if particle_output:
-                    return p_min, fp[i_min], p, fp
-                else:
-                    return p_min, fp[i_min]
+                return p_min, fp[i_min]
             elif stepsize <= minstep:
                 print('Stopping search: Swarm best position change less than {:}'\
                     .format(minstep))
                 sys.stdout.flush()
-                if particle_output:
-                    return p_min, fp[i_min], p, fp
-                else:
-                    return p_min, fp[i_min]
+                return p_min, fp[i_min]
             else:
                 g = p_min.copy()
                 fg = fp[i_min]
@@ -271,7 +266,5 @@ def optimize(func, bounds=None, ieqcons=[], f_ieqcons=None, args=(), kwargs={},
     if not is_feasible(g):
         print("However, the optimization couldn't find a feasible design. Sorry")
         sys.stdout.flush()
-    if particle_output:
-        return g, fg, p, fp
-    else:
-        return g, fg
+
+    return g, fg