[MRG] Optimization enhancements: CMA-ES algorithm and correlation loss function

[ntolley]

This PR makes a lot of changes to implement new optimization functions, namely:

• Covariance matrix adaption evolutionary strategy (ENH: Add CMA-ES algorithm as an optimization solver #1196)
• Correlation coefficient as an objective function (already started here: WIP: Add correlation optimization objective #1206)

Rather than trying to merge this PR in one go, I think it should be broken into smaller tasks as this required a decent amount of changes all over the code base. However there is value to have a fully functioning branch to refer back to and play around with.

[ntolley]

Here is some testing code to try it out, you will need to pip install cma for the code to run. I'm still working on improving the default optimization parameters and fixing random seeds to make everything reproducible.

With more intelligently chosen constraints you can do better, but honestly it's nice to see how well this approach does under such extreme conditions.

NOTE: This is a simulation hungry algorithm, the reason it works well is it runs many simulations in a batch on every epoch (100 by default). I'd recommend running this on an HPC with access to many cores (64 cores in my case).

import numpy as np
import matplotlib.pyplot as plt
from hnn_core import jones_2009_model, simulate_dipole, JoblibBackend
from hnn_core.network_models import add_erp_drives_to_jones_model
from hnn_core.optimization import Optimizer, add_opt_drives, set_params_opt_drives

tstop = 200
dt = 0.5  # Just used for testing, 0.025 is the default but will be slower
scaling_factor = 3000
smooth_win = 30

# Create target dipole
net_target = jones_2009_model()
add_erp_drives_to_jones_model(net_target)
target_dpl = simulate_dipole(net_target, tstop=tstop, dt=dt, verbose=False)
target_dpl = target_dpl[0].copy().smooth(smooth_win).scale(scaling_factor)

# Create base network with drives to be optimized
net_base = jones_2009_model()
net_base._verbose = False
constraints, initial_params = add_opt_drives(net_base, n_prox=2, n_dist=1)

# Run optimization
max_iter = 100
optim = Optimizer(net_base, tstop=tstop, constraints=constraints, solver='cma',
                set_params=set_params_opt_drives, initial_params=initial_params, max_iter=max_iter, obj_fun="dipole_corr")
                
popsize = 500  # number of simulations per epoch, bigger is often better but very expensive!
optim.fit(target=target_dpl, n_trials=1, scale_factor=scaling_factor,
        smooth_window_len=smooth_win, dt=dt, popsize=popsize)

# Simulate best parameters
with JoblibBackend(n_jobs=10):
    dpl_opt = simulate_dipole(optim.net_, tstop=tstop, dt=dt, n_trials=10)

opt_data = np.stack([dpl.copy().scale(scaling_factor).smooth(smooth_win).data['agg'] for dpl in dpl_opt])

target_std = np.std(target_dpl.data['agg'])
opt_std = np.std(opt_data)

opt_scaling =  (target_std / opt_std)
opt_data *= opt_scaling

# Get best results
fontsize = 14
ticksize = 10
plt.figure(figsize=(8, 3))
plt.subplot(1,2,1)
plt.plot(target_dpl.data['agg'] , color='k', label='Target')
plt.plot(opt_data.T, color='C2', linewidth=1, alpha=0.4)
plt.plot(np.mean(opt_data, axis=0), color='C2', label='Optimized')

plt.xlabel('Time (ms)', fontsize=fontsize)
plt.ylabel('Dipole (nAm)', fontsize=fontsize)
plt.xticks(fontsize=ticksize)
plt.yticks(fontsize=ticksize)
plt.legend(fontsize=10)

plt.subplot(1,2,2)
plt.plot(1 - np.array(optim.obj_))
plt.xlabel('Epochs', fontsize=fontsize)
plt.ylabel('Correlation', fontsize=fontsize)
plt.xticks(fontsize=ticksize)
plt.yticks(fontsize=ticksize)
plt.tight_layout()
plt.ylim(None, 1.01)

[katduecker]

but aren't you then just adding on new drives with each iteration?

[ntolley]

In this instance, I found that it is easier to write a set_params function that directly modifies the properties of existing drives. Technically I write a workaround, but in any case I don't necessarily see why we should force users to make a function to either add brand new drives every time or modify existing ones.

[ntolley]

A batch simulate test is failing that I'm at a loss to explain, it indicates that simulating serially is faster in parallel

The only real changes to batch simulate or parallel backends have to do with the verbose setting

Is it possible that these if blocks are slowing down parallel execution in a non-trivial way? But it doesn't make sense to me why serial would be faster...

[ntolley]

maybe the default value should be lower

[ntolley]

@katduecker I actually forgot I already implemented a sigma that is scaled by the size of the bounds here!

[ntolley]

@katduecker this is the right way to set it! the one earlier was wayyyy to small

the performance gain is quite impressive, you can use substantially smaller popsizes

[ntolley]

@asoplata the it was indeed the high njobs that was causing the issues!

[asoplata]

Awesome, congrats! Yes I did my own secondary run of the same code here https://github.com/asoplata/hnn-core/actions/runs/22916219851 to test stochasticity and looks like you solved it! 😌 congrats!

[ntolley]

@asoplata @katduecker what do you think about removing these functions from the PR and including them in a textbook example in a different PR?

They're mostly convenience functions but as @katduecker brought up down below they need to be documented better

[ntolley]

@katduecker @asoplata @carolinafernandezp this PR is ready to review!

I'm thinking that we can add examples of use in a separate PR (perhaps on the textbook website?). The main goal here is just to make sure that the newly added solver="cma" and obj_fun="dipole_corr" work with the appropriate tests.

Important notes on testing:

• initial_params serves no role in the "cma" solver so I've added a warning in case it is passed
• I'm not sure how to approach testing the user defined objective function. While def test_user_obj_fun(solver): exists, it's mainly verifying it's possible to pass a function. I think what is actually necessary is a function like validate_obj_fun() which tells a user if their custom objective function works for all solvers (CMA requires the ability to batch simulate).

Related to making a future textbook example, it's unclear to me ifadd_opt_drives() and set_params_opt_drives() should be in the source code, as they are primarily convenience functions. Maybe just for a textbook example?