[MRG] Optimization enhancements: CMA-ES algorithm and correlation loss function#1221
[MRG] Optimization enhancements: CMA-ES algorithm and correlation loss function#1221asoplata merged 40 commits intojonescompneurolab:masterfrom
Conversation
|
Here is some testing code to try it out, you will need to 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)
|
| # "The current Network instance has external " | ||
| # + "drives, provide a Network object with no " | ||
| # + "external drives." | ||
| # ) |
There was a problem hiding this comment.
but aren't you then just adding on new drives with each iteration?
There was a problem hiding this comment.
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.
| f"Joblib will run {n_trials} trial(s) in parallel by " | ||
| f"distributing trials over {self.n_jobs} jobs." | ||
| ) | ||
| if net._verbose: |
There was a problem hiding this comment.
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...
| # I think that's also why you had to add the max(0.01, ...) for sigma in set_params_opt_drives | ||
| # _b_obj_func = cma.BoundDomainTransform(_obj_func, constraints) # evaluates fun only in the bounded domain | ||
|
|
||
| sigma = 1 / (np.array(constraints[1]) - np.array(constraints[0])) |
There was a problem hiding this comment.
@katduecker I actually forgot I already implemented a sigma that is scaled by the size of the bounds here!
This also removes ALL tests of the gen_opt.py file in order to continue investigating the failing test in test_batch_simulate.py
|
@asoplata the it was indeed the high |
|
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! |
|
|
||
| # Handle float and list inputs for relative_bandpower | ||
| if isinstance(relative_bandpower, float): | ||
| relative_bandpower = [relative_bandpower] * len(obj_fun_kwargs["f_bands"]) |
There was a problem hiding this comment.
TODO: Austin self-review of the math
| if initial_net.external_drives: | ||
| raise ValueError( | ||
| "The current Network instance has external " | ||
| + "drives, provide a Network object with no " | ||
| + "external drives." | ||
| ) |
There was a problem hiding this comment.
Has this warning been intended to remove for a while? I hope this is just a leftover old debug message that hasn't been true for a long time
There was a problem hiding this comment.
Maybe I'll ping @katduecker and @carolinafernandezp to see if they agree?
I think this should either be removed or demoted to a warning. I found it very useful to right set_params() function to modify external drives that have already been added as it allowed the function to exclusively pertain to the parameter updates
In the current form, you force set_params() to serve two purposed, update the parameters and add drives. While it works I don't think it's necessary to enforce.
| If None, the parameters will be set to the midpoints of parameter ranges. | ||
| solver : str | ||
| The optimizer, 'bayesian' or 'cobyla'. | ||
| The optimizer, 'bayesian', 'cobyla', or 'cma'. |
There was a problem hiding this comment.
Future refactor: Might want to add a small note recommending cma usage
asoplata
left a comment
asoplata
left a comment
There was a problem hiding this comment.
Looked at everything except the most mathematical of operations, and aside from the small comments I've made (which can be done after this PR is merged due to deadlines), this looks good to me.
| initial_params, | ||
| obj_fun, | ||
| max_iter, | ||
| obj_fun_kwargs, |
There was a problem hiding this comment.
Discussion on 1-on-1 on 2026-03-31: We discussed how obj_fun_kwargs is inconsistent with data passthrough, since it is used to based some simulation config parameters (like dt), but not others, in addition to objective-function-specific kwargs.
| obj_values.append(obj) | ||
| res = batch_simulation.run( | ||
| params_batch, | ||
| n_jobs=obj_fun_kwargs.get("n_jobs", 1), |
There was a problem hiding this comment.
@asoplata this should be the only parameter that controls if BatchSimulate runs simulations over multiple jobs (if I understood your above question correctly)
| relative_bandpower : list of float | float (Required if obj_fun='maximize_psd') | ||
| Weight for each frequency band in f_bands. If a single float is provided, | ||
| the same weight is applied to all frequency bands. | ||
| sigma0 : float| array-like (Only used if solver='cma') |
There was a problem hiding this comment.
@asoplata I've added documentation of the cma parameters here, please let me know if you agree with the current formatting!
|
@asoplata I made some of the updates we discussed! Quick note: I was wrong about the usage of |
|
I was having a quick look but don't have a ton of time unfortunately! Did you remove the set_params_opt_drives function because I had a comment on that! |
Apologies yes it was removed! @asoplata and I decided that it didn't really make sense to keep those functions in the source code, but would like to add them back later in an example on the textbook (or perhaps some "utility" functions file?) |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1221 +/- ##
==========================================
- Coverage 93.25% 93.01% -0.25%
==========================================
Files 28 28
Lines 5949 6616 +667
==========================================
+ Hits 5548 6154 +606
- Misses 401 462 +61 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
* debug: add first CMA try to GUI opt * maint: remove issue-metrics Github Action (jonescompneurolab#1279) As of the most recent U24 Steering Committee Meeting, we have effectively replaced measuring the metrics of our issues and PRs with Dylan's solution available here jonescompneurolab/hnn-tracking#1 This also allows us to query the information as-needed. We therefore decided in the HNN Code Review meeting on 2026-04-02 to remove usage of the `issue-metrics` Action. * DEBUG: attempt no backend with GUI CMA [no ci] @ntolley take a look at this. This is not a working solution yet - this takes all the code under the "with backend" in the main opt run function, and makes a separate copy that does NOT use a backend, exclusively for the CMA solver. Then, the only different for these blocks of code is solely that in the CMA case, "n_jobs" is being passed to obj_fun_kwargs. However, I'm getting some very confusing errors about serialization comin from Loky, so there could be some major problems. (For example, even though I do not believe we're passing anything non-serializable, something could be interfering with CMA's use of BatchSimulate's use of Joblib and maybe trying to serialize part of the GUI). * DEBUG: remove widgets from closure and hardcode popsize @ntolley Claude Opus found what I think was the existing issue, in that in the `set_params` closure, there were a few ipywidgets being used, but those were interfering with the serialization/pickling by Loky. That definitely makes sense that that could have been the problem. I've included Claude's changes here to strip the widgets from the closure for testing by us tomorrow, though I haven't looked at the changes very closely. Also, I have hardcoded the popsize to be lower since the testing was taking a while to run. Of course, popsize should maybe be a user-configurable value from a widget, but we can do that after we get everything working. This appears to pass `test_gui.py::test_gui_run_optimization` on my machine. We should discuss more tomorrow. * ref/debug: functionalize the main execution for QOL * Fix _hammfilt input validation and improve test coverage (jonescompneurolab#1266) * Fix _hammfilt input validation and improve test coverage * Address review: remove redundant validation and update tests * align and update tests * ref: minor changes * ref: replace util asserts with error-checks [no ci] * test: update test for better util error [no ci] * doc: improve docs and errors in utils [no ci] * test: 1 of 2 improving smooth_waveform tests [no ci] * test: more and better smoothing tests * doc: fix typos and clean comments --------- Co-authored-by: Austin E. Soplata <me@asoplata.com> * [MRG]: Re-enable Codecov (jonescompneurolab#1272) * Re-enable Codecov uploads via codecov-action v5 * CI: Remove Codecov token usage After following the instructions for https://docs.codecov.com/docs/codecov-tokens#uploading-without-a-token I used admin privileges to "install the Codecov Github app" for the HNN-Core repository. Then, I went to codecov.io, logged into the organization's account (after logging into my personal github account), went to Settings, went to Global Upload Token, and changed it to "Not required". In theory, everything should work now. This commit will also act as a test to see if it does. * debug: bump to investigate failing check * doc: add contribution [no ci] --------- Co-authored-by: Austin E. Soplata <me@asoplata.com> * [MRG] Optimization enhancements: CMA-ES algorithm and correlation loss function (jonescompneurolab#1221) * start opt drive function * Move CMA and corr functions to optimization files * read dt, allow n_trials, update external_drives * clean up optimization * add cma dependency and start tests * ruff formatting * fix tests * fix tests * add batch sim detection to all objective functions * reformatting * fix sigma * undo verbose checks * change default njobs * cleanup rebase * cleanup rebase * refactor is_batch check * clean up verbose and refactor preprocessing * reformat psd calculation * ruff cleanup * add warning * ruff * docstring * remove unecessary test * add docstrings about initial params for CMA * refactor dipole sampling * fix args * ruff * fix init imports * return weights * ruff * tests ofr sigma0 * typos * get rid of initial params error * remove obsolete test * ruff * remove commented out test * rename to anticorr * improve coverage * more tests * more tests --------- Co-authored-by: katduecker <katharina.duecker@gmail.com> * [MRG] Expose option to control random seed in CMA optimization (jonescompneurolab#1286) * update cma code with seed parameter * fix tests * fix tests * set_dt * numpy all in tests * add documentation * clean up tests * fix ruff complaints by proper opt exec passing * ref: doc and improve usage of "snapshot" drives "Snapshot drives" refers to the use of the new `_snapshot_drive_widgets`, which is used to extract all the values that we need for `_generate_constraints_and_func`, but without extracting any Widget objects, which are not picklable. * ref: simplify gui opt drive-snapshot code * test: decr test run time * feat: 1 of 2 for adding CMA solver widgets * feat 2/2 using CMA solver options in gui opt --------- Co-authored-by: Satvik Saluja <satviksaluja2507@gmail.com> Co-authored-by: vaishnavi baghel <vshnvibaghel45@gmail.com> Co-authored-by: Nicholas Tolley <55253912+ntolley@users.noreply.github.com> Co-authored-by: katduecker <katharina.duecker@gmail.com>
3 participants
This PR makes a lot of changes to implement new optimization functions, namely:
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.