optimization_bacteriarods: use bayesian optimization from sckit-optimize package and differential evolution method for comparison

Author: polinagaindrik

cr_bayesian_optim/__init__.py

@@ -29,6 +29,6 @@
 import cr_bayesian_optim.sim_branching as sim_branching
 import cr_bayesian_optim.plotting as plotting
 import cr_bayesian_optim.optimization as optimization
-from cr_bayesian_optim.optimize_bacterialrods import *
+import cr_bayesian_optim.optimize_bacterialrods as optimize_bacterialrods
 
 from .fractal_dim import fractal_dim_main

cr_bayesian_optim/optimization.py

@@ -11,6 +11,9 @@
 )
 
 import numpy as np
+from scipy.optimize import differential_evolution
+from skopt import gp_minimize, callbacks
+import pickle
 
 
 def rhs_fractal_dim(options: Options) -> tuple[float, float]:
@@ -35,3 +38,54 @@ def rhs_fractal_dim(options: Options) -> tuple[float, float]:
 
     _, _, popt, pcov = calculate_fractal_dim_for_pos(pos, options)
     return popt[0], pcov[0, 0] ** 0.5
+
+
+def optimization_diff_evolution(cost, bnds, args=(), workers=-1):
+    return differential_evolution(cost,
+                                  bounds=bnds,
+                                  tol=1e-3,
+                                  atol=1e-3,
+                                  maxiter=10,
+                                  #mutation=(0.3, 1.9),
+                                  #recombination=0.7,
+                                  popsize=5,
+                                  init='latinhypercube',
+                                  disp=True,
+                                  polish=False,
+                                  updating='deferred',
+                                  workers=workers,    
+                                  strategy='randtobest1bin',
+                                  callback=callback_diffevol)#, callback=callback_ll) #init='sobol'
+
+
+def callback_diffevol(intermediate_result):
+    with open("out/Optimization_result_diffevol.pkl", 'wb') as outp:  # Overwrites any existing file.
+        pickle.dump(intermediate_result, outp, pickle.HIGHEST_PROTOCOL)
+
+
+def optimization_bayes(cost, bnds, args=(), workers=-1):
+    return gp_minimize(cost,
+                      bnds,
+                      acq_func="EI",            # the acquisition function: EI, LCB, MES, gp_hedge, PVRS, PI, EIps, PIps
+                      n_calls=20,               # the number of evaluations of f
+                      n_random_starts=5,        # the number of random initialization points
+                      noise=0.,                 # the noise level (optional)
+                      random_state=1234,
+                      kappa=1.96,
+                      xi=0.01,
+                      acq_optimizer='lbfgs',    # is needed for parallelization
+                      n_restarts_optimizer=5,   # the number of restarts of the optimizer
+                      n_jobs=workers,           # the number of parallel evaluations of f
+                      callback=[callbacks.CheckpointSaver("out/Optimization_result.pkl")], # a callback function to be called after each iteration
+                      )
+
+
+def save_optimization_result(res, path='', add_filename=''):
+    with open(path+'Final_optimization_result'+add_filename+'.pkl', 'wb') as outp:  # Overwrites any existing file.
+        pickle.dump(res, outp, pickle.HIGHEST_PROTOCOL)
+
+
+def load_optimization_result(path='', add_filename=''):
+    with open(path+'Final_optimization_result'+add_filename+'.pkl', 'rb') as inp:
+        res = pickle.load(inp)
+    return res

cr_bayesian_optim/optimize_bacterialrods.py

@@ -1,10 +1,7 @@
 import cr_mech_coli as crm
 import cr_mech_coli.crm_fit as crm_fit
-
+from functools import partial
 import numpy as np
-import matplotlib.pyplot as plt
-
-from bayes_opt import BayesianOptimization, acquisition
 
 
 def extract_data(image_timesteps, n_vertices):
@@ -16,54 +13,36 @@ def extract_data(image_timesteps, n_vertices):
     return data
 
 
-def cost(data, settings, init_pos, *param):
+def cost_bacterialrods(data, settings, init_pos, param):
     (days, x_target) = data
     container = crm_fit.predict(param, init_pos, settings)
     if container is None:
         print("Simulation Failed")
-        exit()
-    iterations = container.get_all_iterations()
-    x_prediction = np.zeros(np.shape(x_target))
-    delta_iter = np.mean(np.array(iterations)[1:]-np.array(iterations)[:-1])
-    ## TODO why is  saved iterations step changes from 952 to 953 ??
-    iter_data = delta_iter*(np.array(days)-days[0]+1)
-    ind_last = np.argmin(np.abs(iter_data[-1]-iterations))
-    i = 0
-    for iter in iterations[:ind_last+1]:
-        if np.any(np.abs(iter_data-iter) <= 1.):
-            cells = container.get_cells_at_iteration(iter)
-            keys = sorted(cells.keys())
-            # what is the last dimension: why 3 and not 2 ?
-            pos = np.array([cells[key][0].pos for key in keys])[:, :, :-1]
-            x_prediction[i] = pos
-            i += 1
-    return np.mean(squared_difference(x_target, x_prediction))
+        return 1e10
+    else:
+        iterations = container.get_all_iterations()
+        x_prediction = np.zeros(np.shape(x_target))
+        delta_iter = np.mean(np.array(iterations)[1:]-np.array(iterations)[:-1])
+        ## TODO why is  saved iterations step changes from 952 to 953 ??
+        iter_data = delta_iter*(np.array(days)-days[0]+1)
+        ind_last = np.argmin(np.abs(iter_data[-1]-iterations))
+        i = 0
+        for iter in iterations[:ind_last+1]:
+            if np.any(np.abs(iter_data-iter) <= 1.):
+                cells = container.get_cells_at_iteration(iter)
+                keys = sorted(cells.keys())
+                # what is the last dimension: why 3 and not 2 ?
+                pos = np.array([cells[key][0].pos for key in keys])[:, :, :-1]
+                x_prediction[i] = pos
+                i += 1
+        return np.mean(squared_difference(x_target, x_prediction))
 
 
 def squared_difference(x_target, x_prediction):
     return (x_target-x_prediction)**2
 
 
-def posterior(optimizer, grid):
-    mu, sigma = optimizer._gp.predict(grid, return_std=True)
-    return mu, sigma
-
-
-def plot_objective_GP(optimizer, bnds, name=''):
-    for k in bnds.keys():
-        fig, ax = plt.subplots()
-        x_gp = np.linspace(*bnds[k], 100)
-        mean_gp, sigma_gp = posterior(optimizer, x_gp.reshape(-1, 1))
-        ax.plot(x_gp, mean_gp, label=k)
-        ax.fill_between(x_gp, mean_gp + sigma_gp, mean_gp - sigma_gp, alpha=0.1)
-        ax.scatter(optimizer.space.params.flatten(), optimizer.space.target, c="red", s=50, zorder=10)
-        ax.legend(fontsize=12)
-        plt.savefig(f'{k}_{name}'+'.png', bbox_inches='tight')
-        plt.close(fig)
-
-
-
-def optimize_bacterialrods_main():
+def create_test_ABM_framework():
     n_vertices = 8
     # Extract data from masks which have been previously generated
     image_timesteps = ['42', '43', '44', '45', '46', '47', '48', '49', '52']
@@ -72,41 +51,20 @@ def optimize_bacterialrods_main():
     # Target/model/simulation 
     # Define settings required to run simulation 
     settings = crm_fit.Settings.from_toml("data/crm_fit/0001/settings.toml")
-    settings.constants.n_vertices = n_vertices
     settings.constants.n_saves = 15
     settings.others = crm_fit.Others(True)
+    return data, settings
+
 
-    #settings.parameters.damping = crm_fit.SampledFloat(min=0, max=2.5, initial=1.5)
-    settings.parameters.damping = 2.0
-    settings.parameters.potential_type.Mie.en = 10.
-    settings.parameters.potential_type.Mie.em = 1.5
+def main_bacterialrods_optimization(optimizer, update_ABM=None):
+    # Define ABM framework
+    data, settings = create_test_ABM_framework()
+    if update_ABM is not None:
+        settings = update_ABM(settings)
     lower, upper, x0, param_infos, constants, constant_infos = settings.generate_optimization_infos(len(data[1][0]))
+    bnds_dict = {p_inf[0]: (u_b, l_b) for u_b, l_b, p_inf in zip(lower, upper, param_infos)}
     print(param_infos)
-
-    # Define the cost function with arguments as optimizes parameters:
-    #cost_for_optimization = lambda Damping, Strength: cost(data, settings, data[1][0], Damping, Strength)
-    #cost_for_optimization = lambda Damping: cost(data, settings, data[1][0], Damping)
-    cost_for_optimization = lambda Strength: cost(data, settings, data[1][0], Strength)
-
-    N_iter = 20
-    acq = acquisition.ExpectedImprovement(1.) #ProbabilityOfImprovement(1.) #UpperConfidenceBound(kappa=1.)#
-    bnds = {p_inf[0]: (u_b, l_b) for u_b, l_b, p_inf in zip(lower, upper, param_infos)}
-    optimizer = BayesianOptimization(
-        f=None,
-        acquisition_function=acq,
-        pbounds=bnds,
-        verbose=2,
-        random_state=17695,
-    )
-    for j in range(N_iter):
-        next_params = optimizer.suggest()
-        target = cost_for_optimization(**next_params)
-        optimizer.register(
-            params=next_params,
-            target=target,
-        )
-        plot_objective_GP(optimizer, bnds, name=f'EI_{j}')
-
-   
-if __name__ == "__main__":
-    optimize_bacterialrods_main()
+    cost_for_optimization = partial(cost_bacterialrods, data, settings, data[1][0])
+    res = optimizer(cost_for_optimization, [bnds_dict[k] for k in bnds_dict.keys()])
+    print(res.x, res.fun)
+    return res

cr_bayesian_optim/plotting.py

@@ -12,6 +12,7 @@
 import tqdm
 import multiprocessing as mp
 import itertools
+from skopt.plots import plot_gaussian_process, plot_convergence, plot_objective
 
 # Define colors
 COLOR1 = "#0c457d"
@@ -185,3 +186,24 @@ def generate_movie(opath: Path, play_movie: bool = True):
         print("Playing Movie")
         bashcmd2 = f"firefox ./{opath}/movie.mp4"
         os.system(bashcmd2)
+
+
+# Plotting for bacteria rods optimization:
+def plot_optimization_convergence_bayes(res, path='', add_name=''):
+    fig, ax = plt.subplots()
+    plot_convergence(res)
+    plt.savefig(path+'skopt_convergence'+add_name+'.png', bbox_inches='tight')
+    plt.close(fig)
+
+
+def plot_1D_cost_approximation_bayes(res, path='', add_name=''):
+    fig, ax = plt.subplots()
+    _ = plot_gaussian_process(res)
+    plt.savefig(path+'cost_approximation'+add_name+'.png', bbox_inches='tight')
+    plt.close(fig)
+
+def plot_objective_projection_bayes(res, path='', add_name=''):
+    fig, ax = plt.subplots()
+    _ = plot_objective(res, n_points=10)
+    plt.savefig(path+'objective_projection'+add_name+'.png', bbox_inches='tight')
+    plt.close(fig)

docs/requirements.txt

@@ -5,3 +5,9 @@ setuptools==80.9.0
 myst_parser==4.0.1
 sphinx_subfigure==0.2.4
 bayesian-optimization==3.0.0
+scikit-optimize==0.10.2
+numpy==2.3.0
+scipy==1.15.3
+joblib==1.5.1
+scikit-learn==1.7.0
+matplotlib==3.10.3