---
title: "Competing Objectives in Hopf Network Fitting: FC vs. Frequency Gradient"
subtitle: "Parallel optimization from Pareto-selected initial states"
format:
html:
code-fold: false
toc: true
echo: false
embed-resources: true
fig-width: 8
out-width: "100%"
jupyter: python3
execute:
cache: true
---
Try this notebook interactively:
[Download .ipynb](https://github.com/virtual-twin/tvboptim/blob/main/docs/workflows/Hopf_Pareto_ParallelOpt.ipynb){.btn .btn-primary download="Hopf_Pareto_ParallelOpt.ipynb"}
[Download .qmd](Hopf_Pareto_ParallelOpt.qmd){.btn .btn-secondary download="Hopf_Pareto_ParallelOpt.qmd"}
[Open in Colab](https://colab.research.google.com/github/virtual-twin/tvboptim/blob/main/docs/workflows/Hopf_Pareto_ParallelOpt.ipynb){.btn .btn-warning target="_blank"}
## Introduction
We fit a delay-coupled Hopf network to two targets: an empirical **functional connectivity (FC)** matrix and a spatial **peak-frequency gradient**. This workflow examines their trade-off.
In this model and dataset, optimizing either target alone favors a limiting regime:
- **FC correlation alone** favors stronger coupling and greater synchronization, which reduces regional frequency differences.
- **The frequency-gradient term alone** favors weaker coupling, which preserves regional natural frequencies but reduces interaction-driven FC structure.
The combined loss penalizes both limits and favors an intermediate model regime. This is a model-based regularization effect, not evidence that the fitted regime uniquely characterizes cortical dynamics. Residual mismatch is expected because each brain region is represented by a single Hopf oscillator.
We first show the conflict directly on one starting point, then build the workflow around it in two stages:
0. **Genetic pre-search (NSGA-II)** over $(G, a, \sigma)$. With time measured in milliseconds, the isolated-node relation $f_{\mathrm{Hz}}=1000\omega/(2\pi)$ provides a direct initialization for $\omega$. The genetic search therefore estimates only the global parameters, while $\omega$ remains optimisable in the gradient stage.
1. **Parallel gradient optimization from each Pareto seed.** Every non-dominated solution initializes one gradient run. Axis `wrap=` callables construct and scale the differentiable parameters before the mapped optimization.
At the end we compare how the different starting points performed: their starting GA metrics, their post-optimization metrics, and where each of them ended up on the FC vs. frequency-gradient trade-off.
```{python}
#| output: false
#| echo: false
try:
import google.colab
!pip install -q tvboptim pymoo
except ImportError:
pass
```
```{python}
#| output: false
#| code-fold: true
#| code-summary: "Environment Setup and Imports"
#| echo: true
import os
from functools import partial
cpu = True
if cpu:
N = 8
os.environ['XLA_FLAGS'] = f'--xla_force_host_platform_device_count={N}'
import copy
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from scipy.signal import spectrogram
import jax
import jax.numpy as jnp
import optax
jax.config.update("jax_enable_x64", True)
from tvboptim.types import (
DataAxis,
SigmoidBoundedParameter,
)
from tvboptim.types.parameter import RescaledParameter
from tvboptim.types.spaces import Space
from tvboptim.execution import ParallelExecution
from tvboptim.utils import set_cache_path, cache
from tvboptim.optim.optax import OptaxOptimizer
from tvboptim.experimental.network_dynamics import Network, prepare
from tvboptim.experimental.network_dynamics.dynamics.tvb import SupHopf
from tvboptim.experimental.network_dynamics.coupling import LinearCoupling
from tvboptim.experimental.network_dynamics.graph import DenseDelayGraph
from tvboptim.experimental.network_dynamics.solvers import Heun
from tvboptim.experimental.network_dynamics.noise import AdditiveNoise
from tvboptim.data import load_structural_connectivity, load_functional_connectivity
from tvboptim.observations.tvb_monitors.bold import Bold
from tvboptim.observations.observation import compute_fc, fc_corr
import pandas as pd
from pymoo.core.problem import Problem
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.optimize import minimize as pymoo_minimize
from pymoo.indicators.hv import HV
set_cache_path("./hopf_pareto")
```
## Loading Data and Defining Target Frequency Gradient
```{python}
#| echo: true
#| output: false
weights, lengths, region_labels = load_structural_connectivity(name="dk_average")
fc_target = np.asarray(load_functional_connectivity(name="dk_average"))
weights = weights / np.max(weights)
n_nodes = weights.shape[0]
speed = 3.0
delays = lengths / speed
region_labels = np.array(region_labels)
idx_l = np.where(region_labels == "L.LOG")[0]
idx_r = np.where(region_labels == "R.LOG")[0]
dist_from_vc = np.array(np.squeeze(0.5 * (lengths[idx_l, :] + lengths[idx_r, :])))
f_min, f_max = 7.0, 11.0
min_dist, max_dist = dist_from_vc.min(), dist_from_vc.max()
peak_freqs_target = f_max - (f_max - f_min) / (max_dist - min_dist) * (dist_from_vc - min_dist)
# Per-region natural frequency in Hopf radians. With dt in ms and frequency in
# Hz, omega = 2*pi*f / 1000. Used as the initialization for omega in both the
# GA and gradient stages so that every region already starts from the right
# natural frequency.
omega_target = 2 * np.pi * peak_freqs_target / 1000.0
# Bounds for omega during optimization. The interval is wider than the target
# range of 7-11 Hz but excludes the zero-frequency limit. A sigmoid transform
# provides a smooth bounded coordinate.
OMEGA_HZ_LOW, OMEGA_HZ_HIGH = 1.0, 30.0
omega_low = 2 * np.pi * OMEGA_HZ_LOW / 1000.0
omega_high = 2 * np.pi * OMEGA_HZ_HIGH / 1000.0
# Adam updates the transformed parameter leaves. Under one learning rate, the
# corresponding physical step depends on each parameter transformation.
# Rescaling G and a makes their leaf coordinates comparable.
#
# G uses a characteristic magnitude of 0.15. The scale for a is 0.10, based on
# the interval it must traverse to cross the Hopf bifurcation at zero.
#
# Omega uses a logit coordinate. Increasing the common learning rate moves this
# coordinate farther; dividing the G and a scales by the same factor preserves
# their approximate physical step sizes.
OMEGA_STEP_BOOST = 3.0
G_SCALE = 0.15 / OMEGA_STEP_BOOST
A_SCALE = 0.10 / OMEGA_STEP_BOOST
```
## Building the Hopf Network
```{python}
#| echo: true
#| output: false
row_sums = weights.sum(axis=1)
weights_laplacian = weights - np.diag(row_sums)
graph = DenseDelayGraph(weights_laplacian, delays, region_labels=region_labels)
dynamics = SupHopf(a=0.01, omega=jnp.asarray(omega_target))
coupling = LinearCoupling(source=["x"], G=0.025)
noise = AdditiveNoise(sigma=1e-2, apply_to=["x", "y"])
network = Network(
dynamics=dynamics,
coupling={"instant": coupling},
graph=graph,
noise=noise,
)
```
## Preparing and Running the Simulation
```{python}
#| echo: true
#| output: false
#| code-fold: true
#| code-summary: "Show simulation prep"
t1 = 60_000 # ms
dt = 1.0
fs = 1000.0 / dt
solver = Heun()
model, state = prepare(network, solver, t1=t1, dt=dt)
result_init = model(state)
network.update_history(result_init)
model, state = prepare(network, solver, t1=t1, dt=dt)
bold_monitor = Bold(period=1000.0, downsample_period=dt, voi=0, history=result_init)
```
```{python}
result_init.plot(default_window = 300, dpi = 120)
```
## Defining Observations and Loss
The loss adds the two competing objectives, each written as a negative correlation so that *lower is better*:
$$\mathcal{L} = -\,\rho_{\mathrm{FC}} \;-\; \rho_{\text{freq-grad}}$$
The first term, `fc_loss`, rewards agreement with the empirical FC matrix. The second uses the Pearson correlation between soft peak-frequency estimates and the target gradient. In this example, the terms favor stronger and weaker coupling, respectively, so their sum penalizes either limiting solution.
```{python}
#| echo: true
fs_obs = 100.0
subsample = int(fs / fs_obs)
WELCH_NPERSEG = 512
WELCH_NFFT = 2048
def simulate(state):
result = model(state)
bold = bold_monitor(result)
fc = compute_fc(bold, skip_t=20)
x_sub = result.data[::subsample, 0, :]
_, psd = jax.scipy.signal.welch(x_sub.T, fs=fs_obs,
nperseg=WELCH_NPERSEG, nfft=WELCH_NFFT)
return result, bold, fc, psd
def soft_peak_freqs(psd, transient=3, beta=150.0):
f_eval = jnp.linspace(0, fs_obs / 2, psd.shape[1])
psd_t = psd[:, transient:]
f_t = f_eval[transient:]
psd_norm = psd_t / (jnp.max(psd_t, axis=1, keepdims=True) + 1e-12)
w = jax.nn.softmax(beta * psd_norm, axis=1)
return jnp.sum(w * f_t[None, :], axis=1)
peak_freqs_uniform = 9.0 * jnp.ones(n_nodes)
def loss(state):
_, _, fc, psd = simulate(state)
fc_loss = -fc_corr(fc, fc_target)
fitted_peaks = soft_peak_freqs(psd)
fgrad_corr = jnp.corrcoef(fitted_peaks, peak_freqs_target)[0, 1]
return fc_loss - fgrad_corr
```
## The Conflict, Demonstrated
We first run three 200-step optimizations from one moderate starting point: one for each individual term and one for their sum.
The starting point is a hand-picked moderate-coupling state, with $G$, $a$, and $\sigma$ close to a mid-front Pareto solution from Stage 0 but hardcoded here so the demonstration does not depend on the GA. All three runs free the same parameters as the main workflow: $G$ and per-region $a$ and $\omega$. Because $\omega$ is free, the gradient-only run can retune per-region frequencies without lowering $G$, so its pull may come out weaker than the FC-only run's.
```{python}
#| echo: true
#| output: false
#| code-fold: true
#| code-summary: "Show ablation setup"
LR = 0.005 * OMEGA_STEP_BOOST # leaf units; see the rescaling note
N_ABLATION_STEPS = 200
def loss_fc_only(state):
_, _, fc, _ = simulate(state)
return -fc_corr(fc, fc_target)
def loss_grad_only(state):
_, _, _, psd = simulate(state)
fitted_peaks = soft_peak_freqs(psd)
return -jnp.corrcoef(fitted_peaks, peak_freqs_target)[0, 1]
# Combined-loss optimizer; reused unchanged by the Stage 1 parallel sweep.
optimizer = OptaxOptimizer(loss, optax.adam(LR))
optimizer_fc = OptaxOptimizer(loss_fc_only, optax.adam(LR))
optimizer_grad = OptaxOptimizer(loss_grad_only, optax.adam(LR))
# A moderate-coupling starting point. These scalars are close to the
# median-`combined` Pareto seed that Stage 0 finds, but are hardcoded here so
# the demonstration runs before the genetic search and independently of it.
motivation_seed = {"G": 0.15, "a": 0.005, "sigma": 0.031}
```
Run experiment:
```{python}
#| echo: true
#| code-fold: true
#| code-summary: "Show ablation run"
def _ablation_metrics(eval_state):
_, _, fc_e, psd_e = simulate(eval_state)
peaks_e = soft_peak_freqs(psd_e)
return (float(fc_corr(fc_e, fc_target)),
float(jnp.corrcoef(peaks_e, peak_freqs_target)[0, 1]))
def run_one_ablation(opt):
"""Optimize the motivation seed under `opt` for N_ABLATION_STEPS, omega free."""
abl_state = copy.deepcopy(state)
abl_state.coupling.instant.G = RescaledParameter(
jnp.asarray(motivation_seed["G"]), G_SCALE
)
abl_state.dynamics.a = RescaledParameter(
motivation_seed["a"] * jnp.ones(n_nodes), A_SCALE
)
abl_state.dynamics.omega = SigmoidBoundedParameter(jnp.asarray(omega_target),
omega_low, omega_high)
abl_state.noise.sigma = jnp.asarray(motivation_seed["sigma"])
final_state, _ = opt.run(abl_state, max_steps=N_ABLATION_STEPS,
chunk_size=N_ABLATION_STEPS)
fc_c, grad_c = _ablation_metrics(final_state)
return {"fc_corr": fc_c, "freq_corr": grad_c,
"G": float(jnp.asarray(final_state.coupling.instant.G))}
@cache("ablation_motivation", redo=False)
def run_ablation():
# Start point: the motivation seed, omega at the target gradient.
start_state = copy.deepcopy(state)
start_state.coupling.instant.G = jnp.asarray(motivation_seed["G"])
start_state.dynamics.a = jnp.asarray(motivation_seed["a"])
start_state.noise.sigma = jnp.asarray(motivation_seed["sigma"])
start_fc, start_grad = _ablation_metrics(start_state)
return {
"start": {"fc_corr": start_fc, "freq_corr": start_grad},
"fc_only": run_one_ablation(optimizer_fc),
"grad_only": run_one_ablation(optimizer_grad),
"combined": run_one_ablation(optimizer),
}
ablation = run_ablation()
print(f"Motivation seed: G={motivation_seed['G']}, "
f"a={motivation_seed['a']}, sigma={motivation_seed['sigma']}")
for k in ["start", "fc_only", "grad_only", "combined"]:
v = ablation[k]
extra = "" if k == "start" else f", G={v['G']:.4f}"
print(f" {k:10s}: FC corr={v['fc_corr']:.3f}, "
f"freq corr={v['freq_corr']:.3f}{extra}")
```
```{python}
#| label: fig-ablation-fan
#| fig-cap: "**Each loss pulls the fit a different way.** Three 200-step Adam optimizations from one moderate-coupling seed (black circle), each minimizing a different loss; stars mark the endpoints. Lower-left is good on both objectives. All three free $G$ and per-region $a$ and $\\omega$."
#| echo: true
#| code-fold: true
#| code-summary: "Show plotting code"
start = ablation["start"]
sx, sy = 1.0 - start["freq_corr"], 1.0 - start["fc_corr"]
runs = [("FC only", "fc_only", "tab:red"),
("Gradient only", "grad_only", "tab:green"),
("Combined", "combined", "tab:blue")]
fig, ax = plt.subplots(figsize=(6, 5.5))
for label, key, color in runs:
r = ablation[key]
ex, ey = 1.0 - r["freq_corr"], 1.0 - r["fc_corr"]
ax.annotate("", xy=(ex, ey), xytext=(sx, sy),
arrowprops=dict(arrowstyle="->", color=color, lw=2.0,
shrinkA=6, shrinkB=6))
ax.scatter([ex], [ey], color=color, s=170, marker="*",
edgecolors="black", linewidths=0.6, zorder=3, label=label)
ax.scatter([sx], [sy], color="black", s=110, marker="o",
zorder=4, label="Starting seed")
ax.set_xlabel(r"$1 - \rho_{\text{freq-grad}}$")
ax.set_ylabel(r"$1 - \rho_{FC}$")
ax.set_title(f"Where each loss pulls the fit ({N_ABLATION_STEPS} Adam steps)")
ax.grid(alpha=0.3)
ax.legend(loc="best", fontsize=9)
plt.tight_layout()
plt.show()
```
Stage 0 estimates the Pareto front, and Stage 1 optimizes from every point on it.
## Stage 0: Genetic Pre-Search (NSGA-II)
A small NSGA-II search over $(G, a, \log_{10}\sigma)$ with two objectives: **PSD MAE to 9 Hz** and **$1 - \rho_{\mathrm{FC}}$**. $\omega$ is *not* searched; every region keeps its target natural frequency, set above as `omega_target`.
```{python}
#| echo: true
def ga_evaluate(state):
_, _, fc, psd = simulate(state)
fitted_peaks = soft_peak_freqs(psd)
# Stage 0's frequency objective is PSD MAE to a *uniform* 9 Hz, not to
# the spatial gradient: it only asks the global scalars to keep every
# region oscillating in the alpha band. Reproducing the per-region
# gradient is left to the Stage 1 gradient loss (`peak_freqs_target`).
psd_mae = jnp.mean(jnp.abs(fitted_peaks - peak_freqs_uniform))
fc_correlation = fc_corr(fc, fc_target)
return {"psd_mae": psd_mae, "fc_correlation": fc_correlation}
GA_PARAM_RANGES = {
"G": [0.001, 0.15],
"a": [-0.1, 0.01],
"log10_sigma": [-2.0, -0.5],
}
GA_POP_SIZE = N * 1
GA_N_GEN = 40
M1_SCALE = 9.0
GA_REF_POINT = np.array([1.5 * M1_SCALE, 1.5])
```
```{python}
#| echo: true
#| output: false
class HopfGAProblem(Problem):
"""NSGA-II problem: search (G, a, log10_sigma) on (PSD MAE, 1 - FC corr).
Omega is *not* searched. Every region keeps its target natural frequency
baked into `base_state` (via SupHopf), so the GA only varies the three
global scalars. Each population is evaluated in parallel across
`n_devices` pmap shards. Per-evaluation metrics are appended to
`eval_log` so the caller can inspect them after the run.
"""
def __init__(self, base_state, n_devices, eval_log):
# Decision-variable bounds; column order is [G, a, log10(sigma)].
xl = np.array([GA_PARAM_RANGES["G"][0],
GA_PARAM_RANGES["a"][0],
GA_PARAM_RANGES["log10_sigma"][0]])
xu = np.array([GA_PARAM_RANGES["G"][1],
GA_PARAM_RANGES["a"][1],
GA_PARAM_RANGES["log10_sigma"][1]])
super().__init__(n_var=3, n_obj=2, xl=xl, xu=xu)
self.base_state = base_state
self.n_devices = n_devices
self.eval_log = eval_log
def _evaluate(self, X, out, *args, **kwargs):
# Decode the third decision column out of log10-space.
sigma_phys = 10.0 ** X[:, 2]
# Build a batched state: each population row rides a DataAxis on
# G / a / sigma. omega is left untouched in base_state, where it
# already holds the per-region target natural frequency.
batch_state = copy.deepcopy(self.base_state)
batch_state.coupling.instant.G = DataAxis(jnp.asarray(X[:, 0]))
batch_state.dynamics.a = DataAxis(jnp.asarray(X[:, 1]))
batch_state.noise.sigma = DataAxis(jnp.asarray(sigma_phys))
# zip-mode walks the three DataAxes together (one config per row),
# which is what we want for a flat population sweep.
space = Space(batch_state, mode="zip")
execution = ParallelExecution(ga_evaluate, space, n_pmap=self.n_devices)
results = execution.run()
# Pack objectives + log every evaluation for the post-mortem plots.
F = np.empty((len(results), 2))
for i, r in enumerate(results):
psd_mae = float(r["psd_mae"])
fc_c = float(r["fc_correlation"])
F[i] = (psd_mae, 1.0 - fc_c)
self.eval_log.append({
"G": float(X[i, 0]),
"a": float(X[i, 1]),
"sigma": float(sigma_phys[i]),
"psd_mae": psd_mae,
"fc_correlation": fc_c,
})
# NaNs would poison non-dominated sorting; clip them out.
out["F"] = np.nan_to_num(F, nan=1e6, posinf=1e6, neginf=1e6)
def _make_hv_callback(hv_log):
"""Build an NSGA-II callback that logs hypervolume per generation."""
def cb(algorithm):
F = algorithm.pop.get("F")
try:
hv = HV(ref_point=GA_REF_POINT)(F)
except Exception:
hv = np.nan
hv_log.append(hv)
print(f"[Gen {algorithm.n_gen}] HV={hv:.4f} | ND={len(F)}")
return cb
@cache("ga_hopf_pareto", redo=False)
def run_ga(pop_size=GA_POP_SIZE, n_gen=GA_N_GEN):
eval_log = []
hv_log = []
res = pymoo_minimize(
problem=HopfGAProblem(base_state=state, n_devices=N, eval_log=eval_log),
algorithm=NSGA2(pop_size=pop_size),
termination=("n_gen", n_gen),
callback=_make_hv_callback(hv_log),
seed=42,
verbose=True,
)
return {
"pareto_X": np.asarray(res.X),
"pareto_F": np.asarray(res.F),
"all_evals": eval_log,
"hv_per_gen": np.asarray(hv_log),
}
ga_result = run_ga()
pareto_X = ga_result["pareto_X"].copy()
pareto_X[:, 2] = 10.0 ** pareto_X[:, 2] # decode log10(sigma) → physical sigma
pareto_F = ga_result["pareto_F"]
all_evals_df = pd.DataFrame(ga_result["all_evals"])
pareto_df = pd.DataFrame(
np.hstack([pareto_X, pareto_F]),
columns=["G", "a", "sigma", "psd_mae", "one_minus_fc"],
)
pareto_df["fc_correlation"] = 1.0 - pareto_df["one_minus_fc"]
pareto_df["combined"] = 0.5 * (pareto_df["psd_mae"] / M1_SCALE) + 0.5 * pareto_df["one_minus_fc"]
n_pareto = len(pareto_df)
print(f"Pareto front size: {n_pareto}")
```
### Pareto Front and Convergence
```{python}
#| label: fig-ga-pareto
#| fig-cap: "**Stage 0 NSGA-II output.** Pareto front in objective space (PSD MAE vs $1 - \\rho_{FC}$); grey dots are all evaluated solutions, colored dots are the non-dominated front. Right panel: hypervolume convergence."
#| echo: true
#| code-fold: true
#| code-summary: "Show plotting code"
fig, (ax_pf, ax_hv) = plt.subplots(1, 2, figsize=(11, 4))
ax_pf.scatter(all_evals_df["psd_mae"], 1.0 - all_evals_df["fc_correlation"],
c="lightgray", s=25, alpha=0.6, label="All evaluations")
sc = ax_pf.scatter(pareto_df["psd_mae"], pareto_df["one_minus_fc"],
c=pareto_df["combined"], cmap="viridis_r",
s=80, edgecolors="black", linewidths=0.8, label="Pareto front")
ax_pf.set_xlabel("PSD MAE to 9 Hz")
ax_pf.set_ylabel(r"$1 - \rho_{FC}$")
ax_pf.set_title("Pareto Front (Stage 0)")
ax_pf.grid(alpha=0.3)
ax_pf.legend(loc="best", fontsize=8)
plt.colorbar(sc, ax=ax_pf, label="Combined (0.5/0.5)")
ax_hv.plot(np.arange(1, len(ga_result["hv_per_gen"]) + 1),
ga_result["hv_per_gen"], marker="o")
ax_hv.set_xlabel("Generation")
ax_hv.set_ylabel("Hypervolume")
ax_hv.set_title("HV convergence")
ax_hv.grid(alpha=0.3)
plt.tight_layout()
plt.show()
```
## Stage 1: Parallel Gradient Optimization From Every Pareto Seed
Each Pareto-front entry initializes an independent multimodal optimization, executed in parallel with `ParallelExecution`.
Each target declares its parameter wrapper when the `Space` is constructed, leaving the mapped function to optimize and evaluate the state. Each run reinitializes its optimizer state. The result stores fitted parameters and scalar metrics for each seed.
```{python}
#| echo: true
N_OPT_STEPS = 200 # gradient steps per seed
# `optimizer` (combined loss, adam at LR) was built in the motivation section.
def optimize_from_seed(seed_state):
"""Run one full gradient optimization seeded by a single GA Pareto solution.
`seed_state` arrives ready to optimize: G, a and omega are already
Parameters because the axes carrying them declared a `wrap=`. Omega has no
axis and is wrapped once before constructing the Space. During traced
execution, `run` represents its step loop with `lax.scan`.
"""
final_state, _ = optimizer.run(seed_state, max_steps=N_OPT_STEPS)
_, _, fc_f, psd_f = simulate(final_state)
final_peaks = soft_peak_freqs(psd_f)
final_fc_corr = fc_corr(fc_f, fc_target)
final_freq_corr = jnp.corrcoef(final_peaks, peak_freqs_target)[0, 1]
# jnp.asarray on a SigmoidBoundedParameter triggers __jax_array__, so omega
# comes out already in its constrained physical-radians form rather than the
# unconstrained logit-space leaf the optimizer actually updates.
return {
"final_loss": -final_fc_corr - final_freq_corr,
"final_fc_corr": final_fc_corr,
"final_freq_corr": final_freq_corr,
"final_psd_mae": jnp.mean(jnp.abs(final_peaks - peak_freqs_uniform)),
"G_final": jnp.asarray(final_state.coupling.instant.G),
"a_final": jnp.asarray(final_state.dynamics.a),
"omega_final": jnp.asarray(final_state.dynamics.omega),
}
```
### Running the parallel sweep
Each Pareto seed is one row of a zip-mode `Space`. The `G` and `a` axes use `wrap=` because their values initialize gradient variables. The `sigma` axis remains a raw array because it is not optimized.
```{python}
#| echo: true
#| output: false
@cache("parallel_pareto_opt", redo=False)
def run_parallel_pareto_opt():
seed_state = copy.deepcopy(state)
seed_state.coupling.instant.G = DataAxis(
jnp.asarray(pareto_df["G"].to_numpy()),
wrap=partial(RescaledParameter, scale=G_SCALE),
)
seed_state.dynamics.a = DataAxis(
jnp.asarray(pareto_df["a"].to_numpy()),
wrap=lambda v: RescaledParameter(v * jnp.ones(n_nodes), A_SCALE),
)
seed_state.noise.sigma = DataAxis(jnp.asarray(pareto_df["sigma"].to_numpy()))
# Omega is not batched, so it is wrapped before constructing the Space.
seed_state.dynamics.omega = SigmoidBoundedParameter(
jnp.asarray(state.dynamics.omega), omega_low, omega_high,
)
space = Space(seed_state, mode="zip")
execution = ParallelExecution(optimize_from_seed, space, n_pmap=N)
results = execution.run()
rows = [{
"final_loss": float(results[i]["final_loss"]),
"final_fc_corr": float(results[i]["final_fc_corr"]),
"final_freq_corr": float(results[i]["final_freq_corr"]),
"final_psd_mae": float(results[i]["final_psd_mae"]),
"G_final": float(np.asarray(results[i]["G_final"])),
"a_final": np.asarray(results[i]["a_final"]),
"omega_final": np.asarray(results[i]["omega_final"]),
} for i in range(len(results))]
return pd.DataFrame(rows)
opt_df = run_parallel_pareto_opt()
```
```{python}
#| echo: true
#| output: false
#| code-fold: true
#| code-summary: "Show summary-table assembly"
# Merge GA seed table with optimization outcomes for a single overview frame.
summary_df = pd.concat(
[pareto_df.reset_index(drop=True),
opt_df.reset_index(drop=True).add_suffix("")],
axis=1,
)
summary_df["seed_idx"] = np.arange(len(summary_df))
# Parameter-drift columns: the gradient stage turns `a` and `omega` into
# per-region vectors, so summarize the spread `a` picked up and how far the
# fitted per-region omega moved from its target-gradient initialization (Hz).
omega_target_arr = np.asarray(omega_target)
summary_df["a_fit_mean"] = summary_df["a_final"].apply(lambda v: float(np.asarray(v).mean()))
summary_df["a_fit_std"] = summary_df["a_final"].apply(lambda v: float(np.asarray(v).std()))
summary_df["domega_mae_hz"] = summary_df["omega_final"].apply(
lambda v: float(np.mean(np.abs((np.asarray(v) - omega_target_arr) * 1000.0 / (2 * np.pi))))
)
```
## Overview: How Did the Different Starting Points Perform?
A single table per seed: the GA seed scalars, the gradient-optimized parameters (including the per-region spread the gradient stage gave $a$ and how far $\omega$ drifted from its target-gradient initialization), and the before/after fit metrics.
```{python}
#| echo: false
display_cols = [
"seed_idx",
"G", "a", "sigma",
"G_final", "a_fit_mean", "a_fit_std", "domega_mae_hz",
"psd_mae", "final_psd_mae",
"fc_correlation", "final_fc_corr", "final_freq_corr",
]
print(summary_df[display_cols].round(4).to_string(index=False))
```
```{python}
#| label: fig-pareto-trajectories
#| fig-cap: "**Where each seed landed after gradient optimization.** Circles show the GA Pareto front, stars show post-optimization objectives, and arrows connect each seed to its result. Color denotes coupling strength $G$. Values above the displayed range use the highest color."
#| echo: true
#| code-fold: true
#| code-summary: "Show plotting code"
final_one_minus_fc = 1.0 - summary_df["final_fc_corr"].to_numpy()
init_one_minus_fc = summary_df["one_minus_fc"].to_numpy()
init_psd_mae = summary_df["psd_mae"].to_numpy()
final_psd_mae = summary_df["final_psd_mae"].to_numpy()
init_G = summary_df["G"].to_numpy()
final_G = summary_df["G_final"].to_numpy()
# Shared color scale so the GA-seed G (circles) and optimized G (stars)
# are directly comparable. The max is capped at the GA-seed G range so the
# gradient across the initial Pareto seeds (circles) stays visible;
# optimized G values above this cap (stars) clip to the top color.
norm = plt.Normalize(vmin=init_G.min(),
vmax=0.165)
cmap = "plasma"
fig, ax = plt.subplots(figsize=(7.5, 5))
ax.scatter(all_evals_df["psd_mae"], 1.0 - all_evals_df["fc_correlation"],
c="lightgray", s=18, alpha=0.5, label="GA evaluations")
for i in range(len(summary_df)):
ax.annotate("",
xy=(final_psd_mae[i], final_one_minus_fc[i]),
xytext=(init_psd_mae[i], init_one_minus_fc[i]),
arrowprops=dict(arrowstyle="->", color="grey", alpha=0.55,
shrinkA=2, shrinkB=4))
ax.scatter(init_psd_mae, init_one_minus_fc, c=init_G, cmap=cmap, norm=norm,
s=70, marker="o", edgecolors="black", linewidths=1.0,
label="GA Pareto starts")
sc = ax.scatter(final_psd_mae, final_one_minus_fc, c=final_G, cmap=cmap, norm=norm,
s=170, marker="*", edgecolors="black", linewidths=0.6,
label="After parallel opt")
ax.set_xlabel("PSD MAE to 9 Hz")
ax.set_ylabel(r"$1 - \rho_{FC}$")
ax.set_xlim(0.0, 10)
ax.grid(alpha=0.3)
ax.legend(loc="best", fontsize=8)
plt.colorbar(sc, ax=ax, label=r"Coupling strength $G$", extend="max")
plt.tight_layout()
plt.show()
```
## Diagnostics for the Best Seed
We pick the seed with the lowest combined post-optimization loss and re-simulate its final state for a full diagnostic panel.
```{python}
#| echo: true
#| output: false
#| code-fold: true
#| code-summary: "Show best-seed reconstruction"
final_losses = summary_df["final_loss"].to_numpy()
best_i = int(np.nanargmin(np.where(np.isfinite(final_losses), final_losses, np.nan)))
best_row = summary_df.iloc[best_i]
print(f"Best seed index: {best_i}")
print(f" GA start : G={best_row['G']:.4f}, a={best_row['a']:.4f}, "
f"sigma={best_row['sigma']:.4f} (omega initialized from per-region target)")
print(f" GA metrics: PSD MAE={best_row['psd_mae']:.4f}, FC corr={best_row['fc_correlation']:.4f}")
print(f" Final : PSD MAE={best_row['final_psd_mae']:.4f}, FC corr={best_row['final_fc_corr']:.4f}, "
f"freq corr={best_row['final_freq_corr']:.4f}")
best_state = copy.deepcopy(state)
best_state.coupling.instant.G = jnp.asarray(best_row["G_final"])
best_state.dynamics.a = jnp.asarray(best_row["a_final"])
best_state.dynamics.omega = jnp.asarray(best_row["omega_final"])
result_best, bold_best, fc_best, psd_best = simulate(best_state)
fc_best = np.asarray(fc_best)
psd_best = np.asarray(psd_best)
ts_best = np.asarray(result_best.data)
bold_best = np.asarray(bold_best.data)
f_eval = np.linspace(0, fs_obs / 2, psd_best.shape[1])
transient = 3
peak_freqs_best = f_eval[transient:][np.argmax(psd_best[:, transient:], axis=1)]
mask_best = peak_freqs_best > 0
```
```{python}
#| label: fig-best-summary
#| fig-cap: "**Diagnostics for the best seed.** **a:** Spectrogram of one region, showing temporal variation in peak frequency. **b:** Fitted and target peak frequencies by region; the dashed line is identity and $\\rho_P$ / $\\rho_S$ are Pearson and Spearman correlations. The fitted values summarize the full recording and do not represent the temporal variation in panel a. **c:** BOLD time courses, with the selected region in blue. **d:** Target FC below the diagonal and estimated FC above it."
#| echo: true
#| code-fold: true
#| code-summary: "Show plotting code"
import scipy.stats
fig = plt.figure(figsize=(11, 7.5))
gs = gridspec.GridSpec(2, 5, height_ratios=[1, 1])
ax_ts = fig.add_subplot(gs[0, 0:3])
ax_fgrad = fig.add_subplot(gs[0, 3:])
ax_bold = fig.add_subplot(gs[1, 0:3])
ax_fc = fig.add_subplot(gs[1, 3:])
idx_node = 5 # region highlighted across panels a-c
def _panel_letter(ax, letter):
ax.set_title(letter, loc="left", fontweight="bold", fontsize=12)
# --- a: spectrogram of the highlighted region ---
ts_data = ts_best[-30000:, 0, idx_node]
nperseg = 128 * subsample
noverlap = nperseg - 8 * subsample # ~94% overlap → many, smoothly-spaced time slices
nfft = 2 * nperseg # zero-pad → finer, smoother frequency axis
f_sp, t_sp, Sxx = spectrogram(ts_data, fs=fs, nperseg=nperseg,
noverlap=noverlap, nfft=nfft)
Sxx = Sxx / np.max(Sxx)
im_spec = ax_ts.pcolormesh(t_sp, f_sp, Sxx, shading='gouraud', cmap='viridis')
ax_ts.set_ylim(0, 20)
ax_ts.set_xlabel('Time [s]')
ax_ts.set_ylabel('Frequency [Hz]')
ax_ts.set_title('Spectrogram (highlighted region)')
plt.colorbar(im_spec, ax=ax_ts, label='Power [a.u.]', shrink=0.85)
_panel_letter(ax_ts, "a")
# --- b: fitted vs target peak frequency, square axes with identity line ---
freq_lo, freq_hi = 5.5, 12.0
pr_f = np.corrcoef(peak_freqs_target[mask_best], peak_freqs_best[mask_best])[0, 1]
sr_f = scipy.stats.spearmanr(peak_freqs_target[mask_best], peak_freqs_best[mask_best])[0]
ax_fgrad.plot([freq_lo, freq_hi], [freq_lo, freq_hi], 'k--', lw=1, zorder=1)
ax_fgrad.scatter(peak_freqs_target[mask_best], peak_freqs_best[mask_best],
s=40, color="grey", edgecolors="white", linewidths=0.5,
zorder=2, label="region")
ax_fgrad.scatter(peak_freqs_target[mask_best][idx_node], peak_freqs_best[mask_best][idx_node],
color="royalblue", marker="s", s=120, edgecolors="k",
zorder=3, label="highlighted region")
ax_fgrad.set_xlim(freq_lo, freq_hi)
ax_fgrad.set_ylim(freq_lo, freq_hi)
ax_fgrad.set_aspect("equal")
ax_fgrad.set_xlabel("Target peak frequency [Hz]")
ax_fgrad.set_ylabel("Fitted peak frequency [Hz]")
ax_fgrad.set_title("Peak frequency")
ax_fgrad.legend(loc="upper left", fontsize=8, framealpha=0.9)
ax_fgrad.text(0.96, 0.04,
rf"$\rho_P = {pr_f:.2f}$" + "\n" + rf"$\rho_S = {sr_f:.2f}$",
transform=ax_fgrad.transAxes, ha="right", va="bottom", fontsize=9,
bbox=dict(boxstyle="round,pad=0.3", edgecolor="0.6", facecolor="white"))
ax_fgrad.grid(alpha=0.3)
_panel_letter(ax_fgrad, "b")
# --- c: BOLD time courses ---
ax_bold.plot(bold_best[-50:, 0, :10], color="grey", alpha=0.5)
ax_bold.plot(bold_best[-50:, 0, idx_node], color="royalblue", linewidth=1.5,
label='highlighted region')
ax_bold.set_xlabel("Time [TR]")
ax_bold.set_ylabel("BOLD [a.u.]")
ax_bold.set_title('BOLD time courses')
ax_bold.legend(loc="upper right", fontsize=8)
_panel_letter(ax_bold, "c")
# --- d: functional connectivity, target vs estimated ---
fc_dual = np.tril(fc_target) + np.triu(fc_best, k=1)
im_fc = ax_fc.imshow(fc_dual, cmap='cividis', vmin=0, vmax=1.0)
ax_fc.set_title(rf"FC ($\rho_P = {fc_corr(fc_best, fc_target):.3f}$)")
ax_fc.set_xlabel("Target FC (lower) / Estimated FC (upper)")
plt.colorbar(im_fc, ax=ax_fc, shrink=0.85)
_panel_letter(ax_fc, "d")
fig.tight_layout()
plt.show()
```
## Per-Region $\omega$ Drift for the Best Seed
The gradient stage turns $\omega$ into a per-region vector (the per-seed drift summary is in the overview table above). For the best seed, the plot below shows how far each region's fitted frequency moved from its target-gradient initialization.
```{python}
#| label: fig-omega-drift
#| fig-cap: "**Per-region $\\omega$ drift for the best seed.** Initial target frequency (open) vs gradient-optimized frequency (filled), in Hz. Regions are sorted by target frequency."
#| echo: true
#| code-fold: true
#| code-summary: "Show plotting code"
omega_init_hz = peak_freqs_target
omega_final_hz = np.asarray(best_row["omega_final"]) * 1000.0 / (2 * np.pi)
order_f = np.argsort(omega_init_hz)
xs = np.arange(n_nodes)
fig, ax = plt.subplots(figsize=(9, 3.5))
ax.scatter(xs, omega_init_hz[order_f], facecolors="none", edgecolors="black",
s=45, linewidths=1.0, label="target")
ax.scatter(xs, omega_final_hz[order_f], c="tab:blue", s=45, label="fitted")
ax.set_xlabel("Region (sorted by target frequency)")
ax.set_ylabel("Peak frequency [Hz]")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
plt.tight_layout()
plt.show()
```
## Summary
Seeds from across the Pareto front, high-FC and high-gradient extremes alike, converged into the same intermediate region after the parallel gradient stage. That convergence is the result: the combined loss has a single basin, and it sits where neither objective alone would land.
The fit is plausible, not exact. A single Hopf node carries one amplitude and one frequency, with no laminar structure and no E/I balance, so residual mismatch in both FC and the gradient is expected. The combined loss holds the model in that intermediate regime instead of letting either metric run to its degenerate optimum.
