Batch running of multiple optimizations in parallel

In the previous tutorials we fit homogeneous and heterogeneous models to the pooled empirical functional data of a single group.

However, we often want to perform multiple optimizations, especially if we are fitting subject-specific (individualized) BNMs. Furthermore, even when fitting the model to data of a single individual or group, we should perform multiple optimization runs using different random seeds to reduce the chance of convergence to local optima.

At the same time, the CMA-ES runs we performed earlier used a relatively number of simulations in parallel (128 particles). Depending on the GPU hardware, this may underutilize the available GPU resources.

To address this, the cuBNM toolbox includes functionality for batch-running multiple optimizations in parallel: across subjects, across runs for the same subject, or both. This helps fully leverage the GPU and improves efficiency compared to running the optimizations independently and serially.

In this tutorial, we’ll use this batch optimization functionality to run CMA-ES optimizations of node-based heterogeneous models (based on canonical resting state networks) in parallel for two subjects from the HCP-YA dataset, performing two optimizations (with different sampling seeds) per subject (i.e., a total of four paralllel optimization runs).

Loading the data

We will use the data of the first two subjects from the HCP-YA dataset. The data of these two subjects (and only these two) is included in the toolbox and can be loaded as follows:

from cubnm import datasets, utils

subjects = ['100206', '100307']

scs = {}
emp_bolds = {}

for sub in subjects:
    scs[sub] = datasets.load_sc('strength', 'schaefer-100', sub)
    emp_bolds[sub] = datasets.load_bold('schaefer-100', sub, ses='REST1_LR')

Let’s fist visualize the SC matrices and BOLD signals:

from cubnm import datasets
import seaborn as sns
import matplotlib.pyplot as plt
import scipy.stats


fig, axes = plt.subplots(1, 2, figsize=(8, 4))
for i_sub, sub in enumerate(subjects):
    sns.heatmap(scs[sub], cmap='Blues', cbar=False, ax=axes[i_sub])
    axes[i_sub].set_title(f"SC sub-{sub}")

fig, axes = plt.subplots(1, 2, figsize=(15, 2))
for i_sub, sub in enumerate(subjects):
    ax = axes[i_sub]
    # plot z-scored BOLD signal
    sns.heatmap(scipy.stats.zscore(emp_bolds[sub], axis=1), cbar=False, ax=ax)
    ax.set_ylabel("Nodes")
    ax.set_xlabel("Time")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(f"BOLD sub-{sub}")

Running batch optimization

Now we’ll create separate BNMProblem objects for each subject. These will define heterogeneous rWW models with FIC, similar to the model used in the node-based heterogeneous CMA-ES tutorial. Note that each subject has a different sc, emp_bold, and out_dir, so we pass sc directly to the BNMProblem instead of including it in sim_options dictionary.

import os
from cubnm import optimize

# define simulation options
sim_options = dict(
    duration=900,
    bold_remove_s=30,
    TR=0.72,
    sc_dist=None,
    dt='0.1',
    bw_dt='1.0',
    ext_out=True,
    states_ts=False,
    states_sampling=None,
    noise_out=False,
    sim_seed=0,
    noise_segment_length=30,
    gof_terms=['+fc_corr', '-fcd_ks'],
    do_fc=True,
    do_fcd=True,
    window_size=30,
    window_step=5,
    fcd_drop_edges=True,
    exc_interhemispheric=True,
    bw_params='heinzle2016-3T',
    sim_verbose=False,
    do_fic=True,
    max_fic_trials=0,
    fic_penalty_scale=0.5,
)

# load the Yeo networks map
yeo_map = datasets.load_maps('yeo7', parc='schaefer-100').astype(int)[0]

# define a map-based heterogeneous problem for each subject
problems = {}
for sub in subjects:
    problems[sub] = optimize.BNMProblem(
        model = 'rWW',
        params = {
            'G': (0.001, 10.0),
            'wEE': (0.001, 5.0),
            'wEI': (0.001, 5.0)
        },
        emp_bold = emp_bolds[sub], # subject-specific bold signal
        sc = scs[sub], # subject-specific SC
        het_params = ['wEE', 'wEI'],
        node_grouping = yeo_map,
        out_dir = os.path.join('./batch_cmaes', sub), # subject-specific output directories
        **sim_options
    )

problems

{'100206': <cubnm.optimize.BNMProblem at 0x148f9cb18940>,
 '100307': <cubnm.optimize.BNMProblem at 0x149001dfa470>}

Next, we define two CMA-ES optimizers with different seeds. These will be reused across subjects, and subject-specific copies of the optimizers will be created internally.

optimizers = {}
for seed in [1, 2]:
    optimizers[seed] = optimize.CMAESOptimizer(
        popsize=128,
        n_iter=120,
        seed=seed,
        algorithm_kws=dict(tolfun=5e-3),
        print_history=False
    )

optimizers

{1: <cubnm.optimize.CMAESOptimizer at 0x148f9bb9beb0>,
 2: <cubnm.optimize.CMAESOptimizer at 0x148f9bb99f00>}

Note that unlike single-run optimizations, we do not need to call setup_problem() as it is handled internally during batch optimization.

We can now run all four optimization runs (2 subjects × 2 seeds) in parallel using the optimize.batch_optimize() function. This function takes lists of problems and optimizers corresponding to each run. After running the optimizations, it returns fitted copies of the provided Optimizer (e.g. CMAESOptimizer) instances.

Technical Note: Internally, batch_optimize() creates a MultiSimGroup by merging the SimGroup objects of all provided BNMProblems. This allows the simulations to be run in parallel on GPU. After the optimization is complete, it re-runs the optimal simulations (by default, when save=True) and saves all results to disk, just like in individual runs.

%%time
# define the corresponding lists of problems and optimizers for all 4 runs
batch_problems   = [problems['100206'], problems['100206'], problems['100307'], problems['100307']]
batch_optimizers = [optimizers[1],      optimizers[2],      optimizers[1],      optimizers[2]]

# run batch optimization
fitted_optimizers = optimize.batch_optimize(batch_optimizers, batch_problems)

Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Initializing GPU session...
took 3.537090 s
Running 512 simulations...
took 125.100445 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.592713 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.663813 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.704441 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.969926 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.477227 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.549147 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.839053 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.667680 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.631570 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.873895 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.704827 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.354186 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 122.929431 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.484529 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.717410 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.733929 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.807886 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.278482 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.902069 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.515891 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.348063 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.386154 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.830351 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.922564 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.636614 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.993901 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 126.006616 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 126.005143 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.322111 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.433722 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.592587 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.334307 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.818956 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.375597 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.808496 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.216640 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.919699 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.728686 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.857746 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.693779 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.236669 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.450333 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.642894 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.079206 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.510397 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.632290 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.414723 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.095354 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.590547 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.021558 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.959324 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.840188 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.279497 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.522079 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.801905 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.766108 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.770043 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.477744 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.675950 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.983362 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.021716 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.893851 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.979849 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 126.075110 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.688075 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.611358 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.885648 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.345514 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.397241 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.468664 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.603721 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.392503 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.322818 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.751567 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.523709 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.518871 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.938100 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 126.113676 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.780034 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.824127 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.511608 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.612495 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.936340 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.350764 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.863299 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.384618 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.841633 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.420809 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.621160 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.528535 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.955580 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.155829 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.378009 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.353493 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.520158 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.346317 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.464978 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.026636 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.337337 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.343647 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.653681 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 126.146944 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.527061 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.726167 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.710777 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.062110 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 123.440224 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.805390 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 126.005931 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.651837 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.680127 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.375171 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.706333 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 122.783462 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.954332 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.764869 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.009186 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 125.439525 s
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Current session is already initialized
Running 512 simulations...
took 124.507975 s
Rerunning and saving the optimal simulation(s)
Warning: Progress cannot be shown in real time when using Jupyter, but simulations are running in background. Please wait...
Initializing GPU session...
took 3.277639 s
Running 4 simulations...
took 29.241990 s
CPU times: user 4h 23min 47s, sys: 1min 40s, total: 4h 25min 27s
Wall time: 4h 35min 4s

We can observe that, in each iteration (before any optimizer converges), a total of 128 × 4 = 512 simulations are run in parallel. On an NVIDIA A100 (40 GB) GPU, this takes approximately 125 seconds per generation. For comparison, in the single-run CMA-ES tutorial with 128 particles, running 128 simulations per generation took around 40 seconds on average. So, by running all 4 optimizations as a batch (taking around 125 seconds per generation total), instead of running them individually (which would take 40 × 4 = 160 seconds per generation), we have made the process approximately ~20% more efficient.

Efficiency gains of batch optimization will depend on the GPU hardware, but it typically becomes more efficient as more optimizations are included, up to the hardware’s parallelization limits.

Visualization of the results

Just like the individually-run optimizers from previous tutorials, the batch-fitted optimizers provide:

• optimizer.history: The optimization history (free parameters and scores per particle).

• optimizer.opt: The set of free parameters with the lowest cost in addition to its associated scores.

• optimizer.problem.sim_group: A SimGroup (N = 1) object which contains the simulation data (BOLD, FC/D, states (time series), regional parameters) of the optimal simulation.

We will now look at the results of these optimization runs. To make this easier, we first convert the list of returned optimizers into a nested dictionary indexed by subject and seed:

fitted_optimizers_dict = {}
i_run = 0
for sub in subjects:
    fitted_optimizers_dict[sub] = {}
    for seed in [1, 2]:
        fitted_optimizers_dict[sub][seed] = fitted_optimizers[i_run]
        i_run += 1

Negative cost trajectories

We can visualize how negative cost evolved over the course of the 4 CMA-ES optimization runs. Here, to indicate better simulations with higher values, rather than the cost function we will visualize negative cost function.

fig, axes = plt.subplots(2, 2, figsize=(12, 8))

for i_sub, sub in enumerate(subjects):
    for i_seed, seed in enumerate([1, 2]):
        ax = axes[i_sub, i_seed]
        fitted_optimizers_dict[sub][seed].plot_history('-cost', ax=ax)
        ax.set_title(f'sub-{sub} (seed = {seed})')

fig.tight_layout(pad=0.5)

Each subplot shows one optimization run. We can use these plots to evaluate convergence behavior across seeds and subjects. As we can see in all runs the optimization is close to convergence, but not yet converged.

Goodness-of-fit in the best runs across subjects

For each subject, we select the best optimization run, i.e., the run with the lowest cost across the two CMA-ES runs:

import numpy as np

best_optimizers = {}

for sub in subjects:
    best_cost = np.inf
    for seed in [1, 2]:
        if fitted_optimizers_dict[sub][seed].opt['cost'] < best_cost:
            best_cost = fitted_optimizers_dict[sub][seed].opt['cost']
            best_seed = seed
            best_optimizers[sub] = fitted_optimizers_dict[sub][seed]
    print(f'{sub} best run has seed = {best_seed}')
    print(fitted_optimizers_dict[sub][best_seed].opt)

100206 best run has seed = 1
index            69.000000
G                 3.531223
wEE0              0.598576
wEE1              0.016630
wEE2              0.007659
wEE3              0.154302
wEE4              2.820660
wEE5              0.005342
wEE6              0.007981
wEI0              1.888817
wEI1              0.121716
wEI2              0.175818
wEI3              0.461653
wEI4              2.428353
wEI5              0.628592
wEI6              1.807846
cost             -0.270489
+fc_corr          0.430831
-fcd_ks          -0.130346
+gof              0.300485
-fic_penalty     -0.029996
gen             120.000000
Name: 15301, dtype: float64
100307 best run has seed = 1
index            10.000000
G                 4.788896
wEE0              0.001915
wEE1              0.028818
wEE2              0.004678
wEE3              0.002975
wEE4              0.128203
wEE5              0.044878
wEE6              0.005709
wEI0              0.475662
wEI1              0.276742
wEI2              0.162929
wEI3              3.317446
wEI4              2.532542
wEI5              2.343312
wEI6              2.781872
cost             -0.126429
+fc_corr          0.283956
-fcd_ks          -0.097847
+gof              0.186109
-fic_penalty     -0.059680
gen             119.000000
Name: 15114, dtype: float64

The best run happened to be the one with seed = 1 for both subjects.

For the best run of each subject, we can then visualize:

• The regression plot comparing simulated and empirical functional connectivity (FC)

• The distribution comparison between simulated and empirical functional connectivity dynamics (FCD)

fig, axes = plt.subplots(2, 2, figsize=(9, 6))
axes = axes.flatten()
for i_sub, sub in enumerate(subjects):
    # regression plot
    ax = axes[i_sub*2]
    sns.regplot(
        x=best_optimizers[sub].problem.sim_group.sim_fc_trils[0],
        y=best_optimizers[sub].problem.emp_fc_tril,
        scatter_kws=dict(s=5, alpha=0.2, color="grey"),
        line_kws=dict(color="red"),
        ax=ax
    )
    ax.set_xlabel("Simulated FC")
    ax.set_ylabel("Empirical FC")
    # add fc_corr statistic
    ax.set_title(f'{sub} (r = {best_optimizers[sub].opt["+fc_corr"]:.3f})')
    # histograms
    ax = axes[i_sub*2+1]
    sns.histplot(best_optimizers[sub].problem.sim_group.sim_fcd_trils[0], element='step', alpha=0.5, label='Simulated', stat='density', ax=ax)
    sns.histplot(best_optimizers[sub].problem.emp_fcd_tril, element='step', alpha=0.5, label='Empirical', stat='density', ax=ax)
    ax.legend()
    # add fcd_ks statistic
    # -fcd_ks in .opt needs to be negated to have positive values
    ax.set_title(f'{sub} (KS = {-best_optimizers[sub].opt["-fcd_ks"]:.3f})')
fig.tight_layout()

Parameters and state variables of optimal simulations

We now visualize the global coupling parameter \(G\) and the regional values of \(w^{EE}\) and \(w^{EI}\) for the optimal simulations:

# create the axes
fig, axd = plt.subplot_mosaic(
    [["G"] + ["space0"] + ["wEE"] * 100 + ["space1"] + ["wEI"] * 100],
    gridspec_kw=dict(width_ratios=[10] + [1] * 202),
    figsize=(12, 1.5)
)
# plot heatmaps of parameters
sns.heatmap(
    np.array([o.problem.sim_group.param_lists['G'][0] for o in best_optimizers.values()])[:, None],
    cbar=False, cmap='viridis', ax=axd['G']
)
axd['G'].set_title(r'$G$')
sns.heatmap(
    np.array([o.problem.sim_group.param_lists['wEE'][0] for o in best_optimizers.values()]),
    cbar=False, cmap='viridis', ax=axd['wEE']
)
axd['wEE'].set_title(r'$w^{EE}$')
sns.heatmap(
    np.array([o.problem.sim_group.param_lists['wEI'][0] for o in best_optimizers.values()]),
    cbar=False, cmap='viridis', ax=axd['wEI']
)
axd['wEI'].set_title(r'$w^{EI}$')
# aesthetics and labeling
for ax in list(axd.values())[1:]:
    ax.axis('off')
axd['G'].set_xticks([])
labels = []
for sub in subjects:
    labels.append(f'sub-{sub}')
axd['G'].set_yticklabels(labels, rotation=0)

[Text(0, 0.5, 'sub-100206'), Text(0, 1.5, 'sub-100307')]

Similarly, we can also visualize node-wise time-averaged values of a model state, such as the input to excitatory neuronal ensembles \(I^E\), for each optimal simulation.

# take time-averaged node-wise I_E in optimal simulations
I_E_avgs = np.array([
    o.problem.sim_group.sim_states['I_E'][0, :]
    for o in best_optimizers.values()
])
# plot
fig, ax = plt.subplots(figsize=(7, 1.5))
sns.heatmap(I_E_avgs, cmap='viridis', ax=ax)
ax.set_title(r'$\langle I^E \rangle$')
labels = []
for sub in subjects:
    labels.append(f'sub-{sub}')
ax.set_yticklabels(labels, rotation=0)
ax.set_xticks([]);

Limitations of batch optimization

Currently, batch optimization supports varying only the following attributes across the batched BNMProblem instances:

• sc: Structural connectivity strength matrix

• emp_bold: Empirical BOLD signal

• out_dir: Output directory

Note: When multiple optimizers are run on the same problem, their outputs are saved in the same directory and automatically numbered, e.g., out_dir/cmaes_run-0, out_dir/cmaes_run-1, etc.

All other model and simulation parameters must be identical across batched problems. This includes, for example:

• sc_dist: Structural connectivity lengths (used in simulations with conduction delays)

• sim_seed: Random seed for generating simulation noise

In addition, all optimizers that are batched together must be:

• Subclasses of PymooOptimizer (e.g., CMAESOptimizer). Therefore, GridOptimizer for example is not supported.

• Of the same optimizer type across all runs, therefore it is not possible to mix e.g., CMA-ES with NSGA-II.

However, optimizer-specific properties, including seed, popsize and n_iter can differ between runs.

The toolbox may support varying additional features across batched problems and optimizers in the future versions.