Note
Click here to download the full example code
Multilevel Descending-Drive Optimisation on a Trapezoidal Ramp¶
Optimizes DD network parameters to match Fuglevand [fuglevand1993] recruitment targets across multiple %MVC levels simultaneously. Uses a single trapezoidal ramp per trial instead of separate constant-drive simulations.
Usage: cd examples/02_finetune python 06_optimize_dd_multilevel.py --n-mu 200 --n-trials 100 python 06_optimize_dd_multilevel.py --n-mu 150 --n-trials 100 python 06_optimize_dd_multilevel.py --sweep # runs all 4 MU counts
import argparse
import json
import math
import os
import subprocess
import sys
import time
os.environ["MPLBACKEND"] = "Agg"
from pathlib import Path
from types import SimpleNamespace
import numpy as np
import optuna
import quantities as pq
from neuron import h
SCRIPT_DIR = Path(__file__).resolve().parent
from myogen import set_random_seed
from myogen.simulator import RecruitmentThresholds
from myogen.simulator.core.force.force_model import ForceModel
from myogen.simulator.core.force.force_utils import FUGLEVAND_FORCE_FREQUENCY, HUMAN_UNIT_FORCE_FREQUENCY
from myogen.utils.fr_optimization_targets import compute_fr_targets
from myogen.utils.nmodl import load_nmodl_mechanisms
from myogen.utils.ramp_analysis import extract_ramp_recruitment
from myogen.simulator.neuron.descending_drive_network import ( # noqa: E402 (one authority, shared with the profiles)
NETWORK, build_dd_mn_network, excitation_to_drive_hz, run_descending_drive,
insert_scaled_gfluctdv as _insert_gfluctdv,
)
# ─── Constants ────────────────────────────────────────────────
GAMMA_SHAPE = 3.0
SYNAPTIC_WEIGHT = 0.05 # µS, DD input weight on the smallest (best-coupled) motor neuron
WEIGHT_EXPONENT_RANGE = (0.0, 1.0) # fitted range of the DD weight exponent (--fit-weight-exponent)
FIT_WEIGHT_EXPONENT = False
CEILING_WEIGHT = 3.0 # Penalty weight for the recruitment ceiling (force at the last recruitment vs LAST_RECRUIT_PCT_MVC)
CEILING_RAMP_PER_S = 0.1 # Excitation per second in the MVC ramp (also used by 07_fdi_spikes_force.py)
MVC_REST_S, MVC_HOLD_S = 0.5, 3.0 # MVC ramp: rest, ramp 0 -> 1 at CEILING_RAMP_PER_S, hold at full excitation
FORCE_FS_HZ = 2048.0 # ForceModel sampling rate
FORCE_SMOOTH_S = 0.25 # Moving-average window on the ramp force
GFLUCTDV_RANGE = (0.05, 0.5) # Optuna search range for Gfluctdv scale (FDI best was 0.243 at the old 0.25 cap)
ISI_CV_TARGET = (0.15, 0.22) # Literature target: 15-22% at 30 % MVC [enokaduchateau2017] [dideriksen2012] (range not checked against the papers)
ISI_CV_WEIGHT = 2.0 # Penalty weight for ISI CV outside target
TIMESTEP_MS = 0.1
RAMP_MVC_PCT = 30 # Ramp target: 0→30% MVC
REST_S = 3.0
RAMP_S = 4.0
PLATEAU_S = 6.0
CHECKPOINTS = [5, 10, 20, 25, 30]
CHECKPOINT_WEIGHTS = {5: 1.0, 10: 3.0, 20: 2.0, 25: 2.0, 30: 1.0}
# Threshold and target options; set from the command line in main().
THRESHOLD_MODE = "combined" # RecruitmentThresholds mode for the pool
RECRUITMENT_RANGE = 30.0 # largest / smallest threshold
MAX_RECRUITMENT = 1.0 # excitation at which the last unit is recruited (rate-coding targets only)
LAST_RECRUIT_PCT_MVC: float | None = None # force (% MVC) at which the last unit is recruited; None: no ceiling term
CEILING_RANK = 0 # score the unit with CEILING_RANK units recruited above it (0: the last unit)
FORCE_FREQUENCY_SHAPE = FUGLEVAND_FORCE_FREQUENCY # (a, c) of the ForceModel force-frequency saturation
FORCE_FREQUENCY_SHAPES = {"fuglevand": FUGLEVAND_FORCE_FREQUENCY, "human": HUMAN_UNIT_FORCE_FREQUENCY}
# Mean FDI peak firing rates (Hz) in 40 and 80 % MVC contractions [deluca1982 p. 118 Table 1].
HIGH_FORCE_RATE_TARGETS = {40: 25.3, 80: 41.4}
HIGH_FORCE_RATE_WEIGHT = 1.0 # per level, relative rate error; scored on the MVC ramp
SCORE_HIGH_FORCE_RATES = False
DE_LUCA_PARAMS: dict = {} # overrides for de_luca_targets
RESULTS_TAG = "polygon"
def _thresholds(profile, n_mu):
"""Pool thresholds (normalised to 1) and the excitation thresholds used for the targets.
``profile`` is the muscle key; on this branch thresholds come straight from ``RecruitmentThresholds``
(the muscle-profile library is not part of main).
"""
thr, _ = RecruitmentThresholds(N=n_mu, recruitment_range__ratio=RECRUITMENT_RANGE, mode=THRESHOLD_MODE)
thr = np.asarray(thr, dtype=float)
thr = thr / thr.max()
return thr, thr * MAX_RECRUITMENT
# Fuglevand model parameters [fuglevand1993] (back-compat)
MIN_FR_HZ = 8.0
MAX_FR_HZ = 17.0
RESULTS_BASE = SCRIPT_DIR.parent.parent / "results" # same root as the other 02_finetune examples
def build_ramp_drive(threshold_hz, max_drive_hz, timestep_ms=TIMESTEP_MS):
"""Build a 0→30% MVC trapezoidal ramp drive signal."""
peak_hz = threshold_hz + (RAMP_MVC_PCT / 100.0) * (max_drive_hz - threshold_hz)
rest_hz = max(threshold_hz * 0.8, 0)
ramp_ms = RAMP_S * 1000
plat_ms = PLATEAU_S * 1000
total_ms = (2 * REST_S + 2 * RAMP_S + PLATEAU_S) * 1000
n_steps = int(total_ms / timestep_ms)
signal = np.full(n_steps, rest_hz)
rest1_end = int(REST_S * 1000 / timestep_ms)
ramp_up_end = rest1_end + int(ramp_ms / timestep_ms)
plat_end = ramp_up_end + int(plat_ms / timestep_ms)
ramp_down_end = plat_end + int(ramp_ms / timestep_ms)
signal[rest1_end:ramp_up_end] = np.linspace(rest_hz, peak_hz, ramp_up_end - rest1_end)
signal[ramp_up_end:plat_end] = peak_hz
signal[plat_end:ramp_down_end] = np.linspace(peak_hz, rest_hz, ramp_down_end - plat_end)
rng = np.random.default_rng(42)
signal += np.abs(rng.normal(0, peak_hz * 0.01, size=n_steps))
return signal, total_ms
def detect_threshold(thresholds, dd_params, min_fr=5.0):
"""Binary search for the drive Hz that produces sustained firing."""
max_drive = dd_params["dd_drive__Hz"]
lo, hi = 0.05, 0.60
threshold_hz = hi * max_drive
for _ in range(8):
mid = (lo + hi) / 2.0
drive_hz = mid * max_drive
block = run_dd_constant(thresholds, drive_hz, dd_params, duration_ms=3000.0)
has_sustained = any(
len(st) >= int(min_fr * 3.0)
for st in block.segments[0].spiketrains
)
if has_sustained:
threshold_hz = drive_hz
hi = mid
else:
lo = mid
return float(threshold_hz)
def network_block(weight_exponent=None):
"""The ``network`` block of a parameter set: NETWORK, plus the DD weight exponent when it was fitted."""
return dict(NETWORK) if weight_exponent is None else {**NETWORK, "weight_exponent": float(weight_exponent)}
def recruitment_times(block, window_ms=500.0, min_spikes=3):
"""Recruitment time (s) of each unit: the first spike of the first ``min_spikes`` spikes within ``window_ms``
(sustained firing, as in ``extract_ramp_recruitment``). NaN for units never recruited."""
out = []
for st in block.segments[0].spiketrains:
t = np.asarray(st.rescale("s").magnitude)
hits = np.nonzero(t[min_spikes - 1:] - t[:len(t) - min_spikes + 1] < window_ms / 1000.0)[0]
out.append(float(t[hits[0]]) if len(hits) else np.nan)
return np.array(out)
def recruitment_excitation(block, excitation, timestep_ms=TIMESTEP_MS, window_ms=500.0, min_spikes=3):
"""Excitation at each unit's recruitment (see ``recruitment_times``). NaN for units never recruited."""
t_rec = recruitment_times(block, window_ms, min_spikes)
idx = np.clip(np.round(np.nan_to_num(t_rec) * 1000.0 / timestep_ms).astype(int), 0, len(excitation) - 1)
return np.where(np.isnan(t_rec), np.nan, np.asarray(excitation, dtype=float)[idx])
def smoothed_force(force, fs):
w = max(int(FORCE_SMOOTH_S * fs), 1)
return np.convolve(force, np.ones(w) / w, mode="same")
def last_recruitment_pct_mvc(t_rec, force, fs, mvc, rank=0):
"""Force (% MVC, smoothed) at the recruitment of the unit with ``rank`` units recruited after it (0: the last
unit), and the number of units never recruited. Unrecruited units count as recruited after it; more of them than
``rank`` push the value past 100 % (100 * (1 + (n_unrecruited - rank) / n))."""
t_rec = np.asarray(t_rec, dtype=float)
n_unrec = int(np.isnan(t_rec).sum())
if n_unrec > rank:
return 100.0 * (1.0 + (n_unrec - rank) / len(t_rec)), n_unrec
t = np.sort(t_rec[~np.isnan(t_rec)])[::-1][rank - n_unrec]
i = min(int(round(t * fs)), len(force) - 1)
return float(100.0 * smoothed_force(force, fs)[i] / mvc), n_unrec
def rates_at_pct_mvc(trains, force, fs, mvc, levels=tuple(HIGH_FORCE_RATE_TARGETS), window_s=1.0, min_spikes=3):
"""Mean firing rate (Hz) of the active units (>= min_spikes in window_s) centred where the smoothed force first
reaches each level (% MVC). ``trains`` are spike times in s. NaN for a level never reached."""
sm = smoothed_force(force, fs)
out = {}
for lev in levels:
hit = np.nonzero(sm >= lev / 100.0 * mvc)[0]
if not len(hit):
out[lev] = float("nan")
continue
t0 = hit[0] / fs
counts = [np.sum((np.asarray(st) > t0 - window_s / 2) & (np.asarray(st) <= t0 + window_s / 2)) for st in trains]
rates = [c / window_s for c in counts if c >= min_spikes]
out[lev] = float(np.mean(rates)) if rates else 0.0
return out
def high_force_rate_penalty(rates):
"""HIGH_FORCE_RATE_WEIGHT * relative rate error at each HIGH_FORCE_RATE_TARGETS level (1 per missing level)."""
err = 0.0
for lev, tgt in HIGH_FORCE_RATE_TARGETS.items():
r = rates.get(lev, float("nan"))
err += 1.0 if not np.isfinite(r) else abs(r - tgt) / tgt
return HIGH_FORCE_RATE_WEIGHT * err
def ceiling_penalty(pct_last, target_pct=None):
"""CEILING_WEIGHT * relative miss of the last recruitment force; 0 when no target is set."""
if target_pct is None:
return 0.0
return CEILING_WEIGHT * abs(pct_last - target_pct) / target_pct
def make_force_model(thresholds):
"""ForceModel (Fuglevand 1993 [fuglevand1993] twitches) at FORCE_FS_HZ with the FORCE_FREQUENCY_SHAPE saturation."""
return ForceModel(recruitment_thresholds=thresholds, recording_frequency__Hz=FORCE_FS_HZ * pq.Hz,
force_frequency_shape=FORCE_FREQUENCY_SHAPE)
def simulate_force(thresholds, dd_params, threshold_hz, segments):
"""Run the drive of ``mvc_drive``; return (excitation, block, force at FORCE_FS_HZ) (see ``make_force_model``)."""
e, hz = mvc_drive(segments, threshold_hz, dd_params["dd_drive__Hz"])
block = run_dd_ramp(thresholds, hz, dd_params, len(e) * TIMESTEP_MS)
return e, block, np.asarray(make_force_model(thresholds).generate_force(block).magnitude)[:, 0]
def run_mvc_ramp(thresholds, dd_params, threshold_hz):
"""MVC ramp: MVC_REST_S rest, excitation 0 -> 1 at CEILING_RAMP_PER_S, MVC_HOLD_S hold at 1.
MVC is the mean force over the last 2 s of the hold. Returns a namespace with e, block, force, mvc,
t_rec (recruitment times, s), pct_last / n_unrecruited (see ``last_recruitment_pct_mvc``), pct_ranked
(the CEILING_RANK unit) and rates (``rates_at_pct_mvc`` at the HIGH_FORCE_RATE_TARGETS levels)."""
ramp_s = 1.0 / CEILING_RAMP_PER_S
e, block, f = simulate_force(thresholds, dd_params, threshold_hz,
[(MVC_REST_S, 0.0, 0.0), (ramp_s, 0.0, 1.0), (MVC_HOLD_S, 1.0, 1.0)])
t = np.arange(f.size) / FORCE_FS_HZ
t_end = MVC_REST_S + ramp_s + MVC_HOLD_S
mvc = float(np.mean(f[(t > t_end - 2.0) & (t < t_end)]))
t_rec = recruitment_times(block)
pct_last, n_unrec = last_recruitment_pct_mvc(t_rec, f, FORCE_FS_HZ, mvc)
pct_ranked, _ = last_recruitment_pct_mvc(t_rec, f, FORCE_FS_HZ, mvc, rank=CEILING_RANK)
trains = [np.asarray(st.rescale("s").magnitude) for st in block.segments[0].spiketrains]
rates = rates_at_pct_mvc(trains, f, FORCE_FS_HZ, mvc)
return SimpleNamespace(e=e, block=block, force=f, mvc=mvc, t_rec=t_rec, pct_last=pct_last,
n_unrecruited=n_unrec, pct_ranked=pct_ranked, rates=rates)
def run_dd_constant(thresholds, drive_hz, dd_params, duration_ms=3000.0):
"""Run a constant-drive DD simulation (for threshold detection)."""
n_steps = int(duration_ms / TIMESTEP_MS)
drive = np.full(n_steps, drive_hz) + np.clip(
np.random.default_rng(42).normal(0, drive_hz * 0.02, size=n_steps), 0, None,
)
return run_descending_drive(thresholds, drive, dd_params, duration_ms, timestep_ms=TIMESTEP_MS)
def run_dd_ramp(thresholds, drive_signal, dd_params, total_ms):
"""Run a DD simulation with time-varying drive (``run_descending_drive`` at TIMESTEP_MS)."""
return run_descending_drive(thresholds, drive_signal, dd_params, total_ms, timestep_ms=TIMESTEP_MS)
def mvc_drive(segments, threshold_hz, max_hz):
"""``excitation_to_drive_hz`` at TIMESTEP_MS (see there)."""
return excitation_to_drive_hz(segments, threshold_hz, max_hz, timestep_ms=TIMESTEP_MS)
def make_objective(thresholds, fug_targets):
"""Create the Optuna objective function for a given MU count."""
def objective(trial):
try:
dd_neurons = trial.suggest_int("dd_neurons", 50, 500)
conn_probability = trial.suggest_float("conn_prob", 0.05, 0.5)
dd_drive__Hz = trial.suggest_float("dd_drive", 100.0, 1000.0)
gfluctdv_scale = trial.suggest_float(
"gfluctdv_scale", *GFLUCTDV_RANGE,
)
weight_exponent = (trial.suggest_float("weight_exponent", *WEIGHT_EXPONENT_RANGE)
if FIT_WEIGHT_EXPONENT else None)
dd_params = {
"dd_neurons": dd_neurons,
"conn_probability": conn_probability,
"dd_drive__Hz": dd_drive__Hz,
"synaptic_weight": SYNAPTIC_WEIGHT,
"gamma_shape": GAMMA_SHAPE,
"gfluctdv_scale": gfluctdv_scale,
"network": network_block(weight_exponent),
}
# Detect threshold for this param set
threshold_hz = detect_threshold(thresholds, dd_params)
# Build and run trapezoidal ramp
drive_signal, total_ms = build_ramp_drive(threshold_hz, dd_drive__Hz)
block = run_dd_ramp(thresholds, drive_signal, dd_params, total_ms)
# Extract recruitment curve at checkpoints
ramp_data = extract_ramp_recruitment(
block, drive_signal, threshold_hz, dd_drive__Hz,
checkpoints_pct=CHECKPOINTS,
)
# Score against Fuglevand targets
total_error = 0.0
for mvc_pct in CHECKPOINTS:
sim = ramp_data[mvc_pct]
tgt = fug_targets[mvc_pct]
w = CHECKPOINT_WEIGHTS.get(mvc_pct, 1.0)
if tgt["n_active"] == 0:
total_error += w * sim["n_active"] * 0.1
continue
n_err = abs(sim["n_active"] - tgt["n_active"]) / max(tgt["n_active"], 1)
fr_err = (
abs(sim["mean_fr_hz"] - tgt["mean_fr_hz"]) / max(tgt["mean_fr_hz"], 1.0)
if sim["n_active"] > 0 else 1.0
)
total_error += w * (n_err + fr_err)
# ISI CV measurement at plateau (ARCH-23)
plat_start_s = (REST_S + RAMP_S)
plat_end_s = plat_start_s + PLATEAU_S
isi_cvs = []
for st in block.segments[0].spiketrains:
times_s = np.array(st.rescale("s").magnitude)
plat = times_s[(times_s >= plat_start_s) & (times_s < plat_end_s)]
if len(plat) < 5:
continue
isis = np.diff(plat) * 1000
m = float(np.mean(isis))
if m > 0:
isi_cvs.append(float(np.std(isis)) / m)
mean_cv = float(np.mean(isi_cvs)) if isi_cvs else 0.0
# Penalize ISI CV outside the 15-22% target range
cv_lo, cv_hi = ISI_CV_TARGET
if mean_cv < cv_lo:
total_error += ISI_CV_WEIGHT * (cv_lo - mean_cv) / cv_lo
elif mean_cv > cv_hi:
total_error += ISI_CV_WEIGHT * (mean_cv - cv_hi) / cv_hi
# Recruitment ceiling: the ramp above stops at RAMP_MVC_PCT, so the last unit is found on a separate
# MVC ramp, scored by the force (% MVC) at its recruitment. Units never recruited push it past 100 %.
# The same ramp gives the firing rates at high force (De Luca 1982) [deluca1982], which the ramp above never reaches.
pct_last, pct_ranked, n_unrecruited, hf_rates = float("nan"), float("nan"), None, {}
if LAST_RECRUIT_PCT_MVC is not None or SCORE_HIGH_FORCE_RATES:
mvc_ramp = run_mvc_ramp(thresholds, dd_params, threshold_hz)
pct_last, n_unrecruited = mvc_ramp.pct_last, mvc_ramp.n_unrecruited
pct_ranked, hf_rates = mvc_ramp.pct_ranked, mvc_ramp.rates
total_error += ceiling_penalty(pct_ranked, LAST_RECRUIT_PCT_MVC)
if SCORE_HIGH_FORCE_RATES:
total_error += high_force_rate_penalty(hf_rates)
# Store metadata
trial.set_user_attr("pct_mvc_last_recruited", pct_last)
trial.set_user_attr("pct_mvc_ceiling_ranked", pct_ranked)
for lev, r in hf_rates.items():
trial.set_user_attr(f"mean_fr_{lev}_mvc_ramp", r)
trial.set_user_attr("n_unrecruited", n_unrecruited)
trial.set_user_attr("threshold_hz", threshold_hz)
trial.set_user_attr("isi_cv_mean", mean_cv)
trial.set_user_attr("gfluctdv_scale", gfluctdv_scale)
for mvc_pct in CHECKPOINTS:
trial.set_user_attr(f"n_active_{mvc_pct}", ramp_data[mvc_pct]["n_active"])
trial.set_user_attr(f"mean_fr_{mvc_pct}", ramp_data[mvc_pct]["mean_fr_hz"])
if trial.number % 5 == 0:
active_str = " | ".join(
f"{p}%:{ramp_data[p]['n_active']}/{fug_targets[p]['n_active']}"
for p in CHECKPOINTS
)
print(f" Trial {trial.number}: err={total_error:.3f} | {active_str} | "
f"drive={dd_drive__Hz:.0f} thr={threshold_hz:.0f} "
f"gf={gfluctdv_scale:.3f} cv={mean_cv*100:.1f}% last={pct_last:.0f}%MVC"
+ (f" p={weight_exponent:.2f}" if weight_exponent is not None else ""))
return total_error
except Exception as e:
print(f" Trial {trial.number} failed: {e}")
return 1000.0
return objective
def optimize_for_mu_count(n_mu, n_trials, profile, target_model="de_luca",
fresh=False,
target_muscle="vl",
sampler_seed=42,
save_json=True):
"""Run the full optimization for one MU count.
Parameters
----------
n_mu : int
Motor unit count.
n_trials : int
Optuna trial count.
profile : MuscleProfile
Profile to generate thresholds from.
target_model : {"fuglevand", "de_luca", "experimental"}
Rate coding target model (see fr_optimization_targets).
target_muscle : str
Muscle name for experimental lookup (ignored by other models).
sampler_seed : int
TPE sampler seed. Use different values per parallel worker.
save_json : bool
Whether to save the best-trial JSON. False for parallel workers
(the coordinator process saves after all workers finish).
"""
print(f"\n{'='*60}")
print(f"Optimizing DD for {n_mu} MUs (density=312/mm2, Lexell anchor)")
print(f"Target model: {target_model}")
print(f"{'='*60}")
set_random_seed(42)
thresholds, target_thresholds = _thresholds(profile, n_mu)
# Rate coding targets from selected model
if target_model == "fuglevand":
fug = compute_fr_targets(
target_thresholds, mvc_levels=CHECKPOINTS, model="fuglevand",
min_fr_hz=MIN_FR_HZ, max_fr_hz=MAX_FR_HZ,
)
elif target_model == "experimental":
fug = compute_fr_targets(
target_thresholds, mvc_levels=CHECKPOINTS, model="experimental",
muscle_name=target_muscle,
)
else: # de_luca (default)
fug = compute_fr_targets(
target_thresholds, mvc_levels=CHECKPOINTS, model="de_luca", **DE_LUCA_PARAMS,
)
print(f"Rate coding targets ({target_model}):")
for mvc_pct in CHECKPOINTS:
t = fug[mvc_pct]
print(f" {mvc_pct}% MVC: {t['n_active']} active, {t['mean_fr_hz']:.1f} Hz")
# Optuna study
results_dir = RESULTS_BASE / f"dd_{target_muscle}_{n_mu}mu_{RESULTS_TAG}"
results_dir.mkdir(exist_ok=True, parents=True)
db_path = results_dir / "optuna_multilevel.db"
if fresh and db_path.exists():
db_path.unlink()
print(f" Removed existing {db_path} (--fresh)")
storage = f"sqlite:///{db_path}"
study = optuna.create_study(
direction="minimize",
sampler=optuna.samplers.TPESampler(
seed=sampler_seed, multivariate=True,
warn_independent_sampling=False,
),
study_name=f"dd_multilevel_{n_mu}mu",
storage=storage,
load_if_exists=True,
)
objective_fn = make_objective(thresholds, fug)
t0 = time.time()
study.optimize(objective_fn, n_trials=n_trials, show_progress_bar=False)
elapsed = time.time() - t0
# Best trial
best = study.best_trial
print(f"\nBest trial {best.number} (err={best.value:.3f}, {elapsed:.0f}s):")
print(f" dd_neurons={best.params['dd_neurons']}, "
f"conn_prob={best.params['conn_prob']:.3f}, "
f"dd_drive={best.params['dd_drive']:.0f} Hz, "
f"gfluctdv={best.params.get('gfluctdv_scale', 0):.3f}, "
f"weight exponent={best.params.get('weight_exponent', 1.0):.2f}")
print(f" ISI CV: {best.user_attrs.get('isi_cv_mean', 0)*100:.1f}% "
f"(target: {ISI_CV_TARGET[0]*100:.0f}-{ISI_CV_TARGET[1]*100:.0f}%)")
print(f" Last unit recruited at {best.user_attrs.get('pct_mvc_last_recruited', float('nan')):.1f} % MVC; "
f"unit with {CEILING_RANK} above it at {best.user_attrs.get('pct_mvc_ceiling_ranked', float('nan')):.1f} % "
f"(target {LAST_RECRUIT_PCT_MVC}), unrecruited {best.user_attrs.get('n_unrecruited', '?')}")
print(" MVC-ramp rates: " + ", ".join(
f"{lev} % MVC {best.user_attrs.get(f'mean_fr_{lev}_mvc_ramp', float('nan')):.1f}/{tgt:.1f} Hz"
for lev, tgt in HIGH_FORCE_RATE_TARGETS.items()))
for p in CHECKPOINTS:
sim_n = best.user_attrs.get(f"n_active_{p}", "?")
sim_fr = best.user_attrs.get(f"mean_fr_{p}", 0)
tgt = fug[p]
print(f" {p}% MVC: sim={sim_n}/{tgt['n_active']} active, "
f"FR={float(sim_fr):.1f}/{tgt['mean_fr_hz']:.1f} Hz")
# Save results (skip for parallel workers — coordinator saves)
result = {
"target": {
"model": target_model,
"threshold_mode": THRESHOLD_MODE,
"recruitment_range": RECRUITMENT_RANGE,
"max_recruitment": MAX_RECRUITMENT,
"last_recruit_pct_mvc": LAST_RECRUIT_PCT_MVC,
"ceiling_rank": CEILING_RANK,
"force_frequency_shape": list(FORCE_FREQUENCY_SHAPE),
"high_force_rate_targets": HIGH_FORCE_RATE_TARGETS if SCORE_HIGH_FORCE_RATES else None,
"de_luca_params": DE_LUCA_PARAMS,
"targets": {str(k): v for k, v in fug.items()},
"min_fr_hz": MIN_FR_HZ,
"max_fr_hz": MAX_FR_HZ,
},
"best_trial": {
"trial_number": best.number,
"objective_value": best.value,
"dd_neurons": best.params["dd_neurons"],
"conn_probability": best.params["conn_prob"],
"dd_drive__Hz": best.params["dd_drive"],
"synaptic_weight": SYNAPTIC_WEIGHT,
"gamma_shape": GAMMA_SHAPE,
"gfluctdv_scale": best.params.get("gfluctdv_scale", 0.0),
"isi_cv_mean": best.user_attrs.get("isi_cv_mean", 0.0),
"network": network_block(best.params.get("weight_exponent")),
"pct_mvc_last_recruited": best.user_attrs.get("pct_mvc_last_recruited"),
"pct_mvc_ceiling_ranked": best.user_attrs.get("pct_mvc_ceiling_ranked"),
**{f"mean_fr_{lev}_mvc_ramp": best.user_attrs.get(f"mean_fr_{lev}_mvc_ramp")
for lev in HIGH_FORCE_RATE_TARGETS},
"n_unrecruited": best.user_attrs.get("n_unrecruited"),
"threshold_hz": best.user_attrs.get("threshold_hz"),
**{f"n_active_{p}": best.user_attrs.get(f"n_active_{p}") for p in CHECKPOINTS},
**{f"mean_fr_{p}": best.user_attrs.get(f"mean_fr_{p}") for p in CHECKPOINTS},
},
"optimization": {
"n_motor_units": n_mu,
"n_trials": len(study.trials),
"ramp_mvc_pct": RAMP_MVC_PCT,
"checkpoints": CHECKPOINTS,
"checkpoint_weights": CHECKPOINT_WEIGHTS,
"ceiling_weight": CEILING_WEIGHT,
"ceiling_ramp_per_s": CEILING_RAMP_PER_S,
"mvc_rest_s": MVC_REST_S,
"high_force_rate_weight": HIGH_FORCE_RATE_WEIGHT,
"mvc_hold_s": MVC_HOLD_S,
"elapsed_s": elapsed,
},
}
if save_json:
out_path = results_dir / "dd_optimized_params.json"
with open(out_path, "w") as f:
json.dump(result, f, indent=2, default=float)
print(f"Saved: {out_path}")
return result
def _save_best_from_study(study, n_mu, target_muscle, target_model,
fug_targets, elapsed):
"""Extract best trial from an Optuna study and save to JSON."""
best = study.best_trial
results_dir = RESULTS_BASE / f"dd_{target_muscle}_{n_mu}mu_{RESULTS_TAG}"
print(f"\nBest trial {best.number} (err={best.value:.3f}, {elapsed:.0f}s):")
print(f" dd_neurons={best.params['dd_neurons']}, "
f"conn_prob={best.params['conn_prob']:.3f}, "
f"dd_drive={best.params['dd_drive']:.0f} Hz, "
f"gfluctdv={best.params.get('gfluctdv_scale', 0):.3f}, "
f"weight exponent={best.params.get('weight_exponent', 1.0):.2f}")
print(f" ISI CV: {best.user_attrs.get('isi_cv_mean', 0)*100:.1f}% "
f"(target: {ISI_CV_TARGET[0]*100:.0f}-{ISI_CV_TARGET[1]*100:.0f}%)")
print(f" Last unit recruited at {best.user_attrs.get('pct_mvc_last_recruited', float('nan')):.1f} % MVC; "
f"unit with {CEILING_RANK} above it at {best.user_attrs.get('pct_mvc_ceiling_ranked', float('nan')):.1f} % "
f"(target {LAST_RECRUIT_PCT_MVC}), unrecruited {best.user_attrs.get('n_unrecruited', '?')}")
print(" MVC-ramp rates: " + ", ".join(
f"{lev} % MVC {best.user_attrs.get(f'mean_fr_{lev}_mvc_ramp', float('nan')):.1f}/{tgt:.1f} Hz"
for lev, tgt in HIGH_FORCE_RATE_TARGETS.items()))
for p in CHECKPOINTS:
sim_n = best.user_attrs.get(f"n_active_{p}", "?")
sim_fr = best.user_attrs.get(f"mean_fr_{p}", 0)
tgt = fug_targets[p]
print(f" {p}% MVC: sim={sim_n}/{tgt['n_active']} active, "
f"FR={float(sim_fr):.1f}/{tgt['mean_fr_hz']:.1f} Hz")
result = {
"target": {
"model": target_model,
"threshold_mode": THRESHOLD_MODE,
"recruitment_range": RECRUITMENT_RANGE,
"max_recruitment": MAX_RECRUITMENT,
"last_recruit_pct_mvc": LAST_RECRUIT_PCT_MVC,
"ceiling_rank": CEILING_RANK,
"force_frequency_shape": list(FORCE_FREQUENCY_SHAPE),
"high_force_rate_targets": HIGH_FORCE_RATE_TARGETS if SCORE_HIGH_FORCE_RATES else None,
"de_luca_params": DE_LUCA_PARAMS,
"targets": {str(k): v for k, v in fug_targets.items()},
"min_fr_hz": MIN_FR_HZ,
"max_fr_hz": MAX_FR_HZ,
},
"best_trial": {
"trial_number": best.number,
"objective_value": best.value,
"dd_neurons": best.params["dd_neurons"],
"conn_probability": best.params["conn_prob"],
"dd_drive__Hz": best.params["dd_drive"],
"synaptic_weight": SYNAPTIC_WEIGHT,
"gamma_shape": GAMMA_SHAPE,
"gfluctdv_scale": best.params.get("gfluctdv_scale", 0.0),
"isi_cv_mean": best.user_attrs.get("isi_cv_mean", 0.0),
"network": network_block(best.params.get("weight_exponent")),
"pct_mvc_last_recruited": best.user_attrs.get("pct_mvc_last_recruited"),
"pct_mvc_ceiling_ranked": best.user_attrs.get("pct_mvc_ceiling_ranked"),
**{f"mean_fr_{lev}_mvc_ramp": best.user_attrs.get(f"mean_fr_{lev}_mvc_ramp")
for lev in HIGH_FORCE_RATE_TARGETS},
"n_unrecruited": best.user_attrs.get("n_unrecruited"),
"threshold_hz": best.user_attrs.get("threshold_hz"),
**{f"n_active_{p}": best.user_attrs.get(f"n_active_{p}") for p in CHECKPOINTS},
**{f"mean_fr_{p}": best.user_attrs.get(f"mean_fr_{p}") for p in CHECKPOINTS},
},
"optimization": {
"n_motor_units": n_mu,
"n_trials": len([t for t in study.trials
if t.state == optuna.trial.TrialState.COMPLETE]),
"ramp_mvc_pct": RAMP_MVC_PCT,
"checkpoints": CHECKPOINTS,
"checkpoint_weights": CHECKPOINT_WEIGHTS,
"ceiling_weight": CEILING_WEIGHT,
"ceiling_ramp_per_s": CEILING_RAMP_PER_S,
"mvc_rest_s": MVC_REST_S,
"high_force_rate_weight": HIGH_FORCE_RATE_WEIGHT,
"mvc_hold_s": MVC_HOLD_S,
"elapsed_s": elapsed,
},
}
out_path = results_dir / "dd_optimized_params.json"
with open(out_path, "w") as f:
json.dump(result, f, indent=2, default=float)
print(f"Saved: {out_path}")
return result
def run_parallel(args):
"""Launch parallel worker processes against a shared Optuna DB.
Each worker is a separate Python/NEURON process. Optuna coordinates
trial assignment through the shared SQLite database. Workers use
different TPE sampler seeds for diverse exploration.
"""
from myogen.utils.parallel_budget import estimate_safe_njobs
n_mu = args.n_mu
n_trials = args.n_trials
target_muscle = args.target_muscle
target_model = args.target_model
profile = target_muscle
# Determine worker count
if args.n_workers:
n_workers = args.n_workers
else:
n_workers = estimate_safe_njobs(
n_mu=n_mu, dd_neurons=500, # worst-case from search range
max_jobs=8,
)
trials_per_worker = math.ceil(n_trials / n_workers)
# Set up results directory and DB
results_dir = RESULTS_BASE / f"dd_{target_muscle}_{n_mu}mu_{RESULTS_TAG}"
results_dir.mkdir(exist_ok=True, parents=True)
db_path = results_dir / "optuna_multilevel.db"
if args.fresh and db_path.exists():
db_path.unlink()
print(f" Removed existing {db_path} (--fresh)")
storage = f"sqlite:///{db_path}"
study_name = f"dd_multilevel_{n_mu}mu"
# Pre-create study so workers don't race on creation
optuna.logging.set_verbosity(optuna.logging.WARNING)
optuna.create_study(
direction="minimize",
sampler=optuna.samplers.TPESampler(seed=42, multivariate=True),
study_name=study_name,
storage=storage,
load_if_exists=True,
)
# Compute targets for final reporting
set_random_seed(42)
thresholds, target_thresholds = _thresholds(profile, n_mu)
if target_model == "fuglevand":
fug = compute_fr_targets(target_thresholds, mvc_levels=CHECKPOINTS,
model="fuglevand",
min_fr_hz=MIN_FR_HZ, max_fr_hz=MAX_FR_HZ)
elif target_model == "experimental":
fug = compute_fr_targets(target_thresholds, mvc_levels=CHECKPOINTS,
model="experimental",
muscle_name=target_muscle)
else:
fug = compute_fr_targets(target_thresholds, mvc_levels=CHECKPOINTS,
model="de_luca", **DE_LUCA_PARAMS)
print(f"\n{'='*60}")
print(f"Parallel DD optimization: {target_muscle} {n_mu} MUs")
print(f"Workers: {n_workers}, trials/worker: {trials_per_worker}")
print(f"Total trials: ~{n_workers * trials_per_worker} "
f"(requested {n_trials})")
print(f"Target model: {target_model}")
print(f"{'='*60}\n")
# Spawn workers
script = str(Path(__file__).resolve())
python = sys.executable
log_dir = results_dir / "worker_logs"
log_dir.mkdir(exist_ok=True)
procs = []
log_files = []
t0 = time.time()
for w in range(n_workers):
log_path = log_dir / f"worker_{w}.log"
cmd = [
python, script,
"--n-mu", str(n_mu),
"--n-trials", str(trials_per_worker),
"--target-model", target_model,
"--target-muscle", target_muscle,
"--threshold-mode", THRESHOLD_MODE,
"--recruitment-range", str(RECRUITMENT_RANGE),
"--max-recruitment", str(MAX_RECRUITMENT),
*(["--last-recruit-mvc", str(LAST_RECRUIT_PCT_MVC)] if LAST_RECRUIT_PCT_MVC is not None else []),
"--ceiling-rank", str(CEILING_RANK),
"--force-frequency", next(k for k, v in FORCE_FREQUENCY_SHAPES.items() if v == FORCE_FREQUENCY_SHAPE),
*(["--high-force-rates"] if SCORE_HIGH_FORCE_RATES else []),
*(["--fit-weight-exponent"] if FIT_WEIGHT_EXPONENT else []),
"--de-luca-params", json.dumps(DE_LUCA_PARAMS),
"--results-tag", RESULTS_TAG,
"--_worker-seed", str(42 + w),
]
lf = open(log_path, "w")
p = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT)
procs.append((w, p))
log_files.append(lf)
print(f" Worker {w} started (PID {p.pid}, seed {42 + w})")
# Wait for all workers with periodic progress reporting
target_total = n_workers * trials_per_worker
while True:
# Check every 30s, but break immediately when all are done
for _ in range(30):
time.sleep(1)
if all(p.poll() is not None for _, p in procs):
break
try:
study = optuna.load_study(study_name=study_name, storage=storage)
completed = len([t for t in study.trials
if t.state == optuna.trial.TrialState.COMPLETE])
best_val = study.best_value if completed > 0 else float("inf")
elapsed = time.time() - t0
rate = completed / elapsed * 3600 if elapsed > 0 else 0
print(f" [{elapsed / 60:.0f}m] {completed}/{target_total} trials "
f"complete, best err={best_val:.3f}, "
f"rate={rate:.0f} trials/hr", flush=True)
except Exception:
pass
if all(p.poll() is not None for _, p in procs):
break
# Close log files
for lf in log_files:
lf.close()
elapsed = time.time() - t0
failures = [(w, p.returncode) for w, p in procs if p.returncode != 0]
if failures:
print(f"\n WARNING: {len(failures)} worker(s) failed:")
for w, rc in failures:
print(f" Worker {w}: exit code {rc} "
f"(see {log_dir / f'worker_{w}.log'})")
# Load final study and save best trial
study = optuna.load_study(study_name=study_name, storage=storage)
completed = len([t for t in study.trials
if t.state == optuna.trial.TrialState.COMPLETE])
print(f"\nOptimization complete: {completed} trials in {elapsed:.0f}s")
_save_best_from_study(study, n_mu, target_muscle, target_model,
fug, elapsed)
def main():
load_nmodl_mechanisms()
h.secondorder = 2
ap = argparse.ArgumentParser()
ap.add_argument("--n-mu", type=int, default=200,
help="MU count to optimize for")
ap.add_argument("--n-trials", type=int, default=100,
help="Optuna trials")
ap.add_argument("--sweep", action="store_true",
help="Sweep all MU counts [150, 200, 250, 300]")
ap.add_argument("--fresh", action="store_true",
help="Delete existing Optuna database and start fresh")
ap.add_argument("--parallel", action="store_true",
help="Run optimization with multiple parallel worker "
"processes. Each worker gets its own NEURON instance.")
ap.add_argument("--n-workers", type=int, default=None,
help="Number of parallel workers (default: auto from "
"parallel_budget, capped at 8)")
ap.add_argument("--target-model", type=str, default="de_luca",
choices=["fuglevand", "de_luca", "experimental"],
help="Rate coding target model. de_luca (default) is "
"more realistic than fuglevand linear. experimental "
"uses a muscle-specific lookup table.")
ap.add_argument("--target-muscle", type=str, default="vl",
help="Muscle profile key. Used for results directory, "
"experimental FR lookup, and profile selection.")
ap.add_argument("--_worker-seed", type=int, default=None,
help=argparse.SUPPRESS) # internal: per-worker sampler seed
ap.add_argument("--threshold-mode", default="combined",
choices=["fuglevand", "deluca", "konstantin", "combined"])
ap.add_argument("--recruitment-range", type=float, default=30.0)
ap.add_argument("--max-recruitment", type=float, default=1.0,
help="Excitation fraction at which the last unit is recruited (rate-coding targets).")
ap.add_argument("--last-recruit-mvc", type=float, default=None,
help="Force (%% MVC) at which the last unit should be recruited, scored on a separate MVC ramp. "
"Default: no ceiling term.")
ap.add_argument("--ceiling-rank", type=int, default=0,
help="Score the unit with K units recruited above it instead of the last unit (default 0).")
ap.add_argument("--force-frequency", choices=sorted(FORCE_FREQUENCY_SHAPES), default="fuglevand",
help="ForceModel force-frequency saturation: fuglevand (1993, default) [fuglevand1993] or human (fit to human "
"single motor units, force_utils.HUMAN_UNIT_FORCE_FREQUENCY).")
ap.add_argument("--fit-weight-exponent", action="store_true",
help="Also fit the DD weight exponent p in [0, 1]: weights scale as (k_max / k_i) ** p "
"(default p = 1, equal current per input at rest; see build_dd_mn_network).")
ap.add_argument("--high-force-rates", action="store_true",
help="Also score MVC-ramp firing rates at 40 and 80 %% MVC against De Luca et al. 1982 [deluca1982].")
ap.add_argument("--de-luca-params", type=str, default="{}",
help="JSON overrides for de_luca_targets, e.g. fr_onset_hz.")
ap.add_argument("--results-tag", type=str, default="polygon")
args = ap.parse_args()
global THRESHOLD_MODE, RECRUITMENT_RANGE, MAX_RECRUITMENT, LAST_RECRUIT_PCT_MVC, DE_LUCA_PARAMS, RESULTS_TAG
global CEILING_RANK, FORCE_FREQUENCY_SHAPE, SCORE_HIGH_FORCE_RATES, FIT_WEIGHT_EXPONENT
THRESHOLD_MODE, RECRUITMENT_RANGE = args.threshold_mode, args.recruitment_range
MAX_RECRUITMENT, RESULTS_TAG = args.max_recruitment, args.results_tag
DE_LUCA_PARAMS = json.loads(args.de_luca_params)
LAST_RECRUIT_PCT_MVC = args.last_recruit_mvc
CEILING_RANK, SCORE_HIGH_FORCE_RATES = args.ceiling_rank, args.high_force_rates
FIT_WEIGHT_EXPONENT = args.fit_weight_exponent
FORCE_FREQUENCY_SHAPE = FORCE_FREQUENCY_SHAPES[args.force_frequency]
if LAST_RECRUIT_PCT_MVC is None:
print("Recruitment ceiling term off (pass --last-recruit-mvc PCT to score the last unit's recruitment force).")
if args.parallel:
run_parallel(args)
return
profile = args.target_muscle
is_worker = args._worker_seed is not None
if args.sweep:
for n_mu in [150, 200, 250, 300]:
optimize_for_mu_count(n_mu, args.n_trials, profile,
target_model=args.target_model,
target_muscle=args.target_muscle,
fresh=args.fresh)
else:
optimize_for_mu_count(
args.n_mu, args.n_trials, profile,
target_model=args.target_model,
target_muscle=args.target_muscle,
fresh=args.fresh if not is_worker else False,
sampler_seed=args._worker_seed if is_worker else 42,
save_json=not is_worker,
)
if __name__ == "__main__":
main()
Total running time of the script: ( 0 minutes 0.000 seconds)