Skip to content

Note

Click here to download the full example code

Optimize Oscillating DC Offset to Match Reference Force

This example optimizes the DC offset component of an oscillating descending drive to match the reference force produced by constant drive (from script 01). The drive pattern is: DC_offset + 20*sin(2π*20*t).

This finds the equivalent of the original 58 pps value (Phase 3) that produces the same force as constant drive (Phase 1) when combined with 20 Hz oscillation.

Note

Purpose: Find DC offset for Phase 3 that matches Phase 1 force

  • Reference: Constant drive force from script 01 (configurable, default 40 Hz)
  • Optimize: DC offset + 20 Hz oscillation to match that force
  • Result: DC offset < constant drive (oscillation adds activation)

Important

Prerequisites: Run 01_compute_baseline_force.py first

Output: results/watanabe_optimization/dc_offset_optimized.json

MyoGen Components Used:

  • AlphaMN__Pool: Same motor neuron pool as script 01. Created fresh for each optimization trial to ensure independent network realizations.

  • DescendingDrive__Pool: Poisson spike generators driven by time-varying rate: DC_offset + amplitude*sin(2πft). The integrate() method advances the internal Poisson process by one timestep.

  • Network: Rebuilds the network for each trial with same connectivity parameters (30%).

  • ForceModel: Computes steady-state force from spike trains to evaluate objective function.

  • External dependency: optuna for Bayesian optimization (TPE sampler).

Key Concept: The oscillation itself contributes to motor neuron activation, so the DC offset in Phase 3 must be lower than the constant drive in Phase 1 to produce the same mean force. The original paper found 58 Hz offset vs 65 Hz constant.

Workflow Position: Step 2 of 6

Next Step: Run 03_10pct_mvc_simulation.py to run the full 3-phase simulation.

Import Libraries

import json
import os
import subprocess
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import optuna

sys.path.insert(0, str(Path(__file__).parent))
from _oscillating_dc_helpers import (  # noqa: E402
    DD_CONNECTIVITY,
    MAX_FORCE_N,
    N_DD_NEURONS,
    N_MOTOR_UNITS,
    N_TRIALS,
    OSC_AMPLITUDE__HZ,
    OSC_FREQUENCY__HZ,
    REFERENCE_DRIVE__HZ,
    REFERENCE_FORCE__N,
    RESULTS_DIR,
    STUDY_NAME,
    SYNAPTIC_WEIGHT,
    TARGET_FORCE__N,
    make_storage,
    objective,
    recruitment_thresholds,
)

plt.style.use("fivethirtyeight")

Configuration Banner

print("\nLoading Reference Force (Constant Drive)")
print("=" * 50)
print(f"Reference force: {REFERENCE_FORCE__N:.2f} N ({REFERENCE_DRIVE__HZ:.1f} Hz constant)")
print(f"Target force: {TARGET_FORCE__N:.2f} N (match with oscillation)")
print(f"Force scaling: {MAX_FORCE_N:.0f} N maximum")
print(f"Oscillation: {OSC_FREQUENCY__HZ} Hz, amplitude {OSC_AMPLITUDE__HZ} Hz")
print(f"DD neurons: {N_DD_NEURONS}")
print(f"Connection probability: {DD_CONNECTIVITY:.1%}")
print("=" * 50 + "\n")

Run Optimization (Parallel Workers)

Each trial is a fully independent NEURON simulation (~80 s), so we fan out across worker processes. NEURON's h object is a process-global singleton and cannot be shared across threads. Workers are separate interpreter processes (_optimize_dc_worker.py) — they do not re-import this gallery file, avoiding the macOS spawn fork-bomb problem. All workers write to a single shared SQLite Optuna study.

Override the worker count via MYOGEN_OPTUNA_WORKERS (e.g. =1 for strictly sequential, or a small number on memory-constrained CI).

print(f"\nOptimizing DC Offset to Match Reference Force (Oscillating Drive)")
print("=" * 50)
print(f"Target force: {TARGET_FORCE__N:.2f} N ({REFERENCE_DRIVE__HZ:.1f} Hz reference)")
print(f"Oscillation: {OSC_FREQUENCY__HZ} Hz (amplitude {OSC_AMPLITUDE__HZ} Hz)")
print(f"Optimizing: DC offset component")
print(f"Trials: {N_TRIALS}\n")

colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]

# Start each pipeline run from a clean study so a leftover SQLite DB from a
# previous run doesn't accumulate stale trials into best_trial.
try:
    optuna.delete_study(study_name=STUDY_NAME, storage=make_storage())
except KeyError:
    pass

study = optuna.create_study(
    direction="minimize",
    sampler=optuna.samplers.TPESampler(seed=42),  # workers set their own seeded samplers
    study_name=STUDY_NAME,
    storage=make_storage(),
    load_if_exists=True,
)

# Conservative default: each worker builds a memory-heavy 800-MN + 400-DD NEURON
# network, so 8 concurrent workers can OOM a laptop/CI runner. Raise it via
# MYOGEN_OPTUNA_WORKERS on a machine with enough RAM for the full speedup.
n_workers = int(os.environ.get("MYOGEN_OPTUNA_WORKERS", min(4, N_TRIALS)))
n_workers = max(1, min(n_workers, N_TRIALS))

per_worker = [N_TRIALS // n_workers + (1 if i < N_TRIALS % n_workers else 0) for i in range(n_workers)]
worker_script = str(Path(__file__).parent / "_optimize_dc_worker.py")
print(f"Launching {n_workers} worker process(es): trial split {per_worker}")
procs = [
    subprocess.Popen([sys.executable, worker_script, "--n-trials", str(n), "--seed", str(42 + i)])
    for i, n in enumerate(per_worker) if n > 0
]
exit_codes = [p.wait() for p in procs]
if any(code != 0 for code in exit_codes):
    raise RuntimeError(f"Optuna worker(s) failed with exit codes {exit_codes}")

# Reload to see all trials written by the workers
study = optuna.load_study(study_name=STUDY_NAME, storage=make_storage())

Analyze Results

best_trial = study.best_trial

print("\nOptimization Complete")
print("=" * 50)
print(f"Best trial: {best_trial.number}")
print(f"Force error: {best_trial.value:.1%}")
print("\nForce Results:")
print(f"  Target:   {TARGET_FORCE__N:.2f} N ({REFERENCE_DRIVE__HZ:.1f} Hz reference)")
print(f"  Achieved: {best_trial.user_attrs['force_achieved']:.2f} N")
print(f"  Error:    {best_trial.user_attrs['force_error']:.1%}")
print("\nOptimized Parameters:")
print(f"  DC offset: {best_trial.user_attrs['dc_offset__Hz']:.2f} Hz")
print(f"  Oscillation: {OSC_FREQUENCY__HZ} Hz (amplitude {OSC_AMPLITUDE__HZ} Hz)")
print(f"  Active MUs: {best_trial.user_attrs['n_active']}/{N_MOTOR_UNITS}")
print(
    f"  Firing rate: {best_trial.user_attrs['FR_mean']:.1f}±{best_trial.user_attrs['FR_std']:.1f} Hz"
)

Save Results

dd_parameters = {
    "dd_neurons": N_DD_NEURONS,
    "dd_connectivity": DD_CONNECTIVITY,
    "synaptic_weight__uS": SYNAPTIC_WEIGHT,
    "dc_offset__Hz": best_trial.user_attrs["dc_offset__Hz"],
    "process_type": "poisson",
}

results = {
    "reference_drive__Hz": REFERENCE_DRIVE__HZ,
    "reference_force__N": REFERENCE_FORCE__N,
    "target_force__N": TARGET_FORCE__N,
    "achieved_force__N": best_trial.user_attrs["force_achieved"],
    "force_error": best_trial.user_attrs["force_error"],
    "force_scaling": {
        "max_force__N": MAX_FORCE_N,
    },
    "oscillation": {
        "frequency__Hz": OSC_FREQUENCY__HZ,
        "amplitude__Hz": OSC_AMPLITUDE__HZ,
    },
    "dd_parameters": dd_parameters,
}

json_path = RESULTS_DIR / "dc_offset_optimized.json"
with open(json_path, "w") as f:
    json.dump(results, f, indent=2)

import joblib
joblib.dump(study, RESULTS_DIR / "study_oscillating_dc.pkl")

print(f"\nSaved results: {json_path}")
print(f"\nNext step: Run 03_10pct_mvc_simulation.py to run the full 3-phase simulation")

Visualize Optimization History

fig = plt.figure(figsize=(14, 10))
gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)

# Extract trial data
trial_numbers = [t.number for t in study.trials]
force_achieved = [t.user_attrs.get("force_achieved", None) for t in study.trials]
force_errors = [t.value for t in study.trials]
dc_offset_vals = [t.params.get("dc_offset", None) for t in study.trials]

# 1. Optimization progress
ax_error = fig.add_subplot(gs[0, 0])
ax_error.plot(trial_numbers, force_errors, "o-", alpha=0.6, markersize=4, color=colors[0])
ax_error.axhline(
    best_trial.value, linestyle="--", color=colors[1], label=f"Best: {best_trial.value:.2%}"
)
ax_error.set_ylabel("Relative Error")
ax_error.set_title(f"Optimization Progress (DC Offset for {REFERENCE_DRIVE__HZ:.0f} Hz Match)")
ax_error.set_yscale("log")
ax_error.grid(True, alpha=0.3)
ax_error.legend(framealpha=1.0, edgecolor="none")
ax_error.set_xlabel("Trial")

# 2. Force convergence
ax_force = fig.add_subplot(gs[1, 0])
ax_force.plot(trial_numbers, force_achieved, "o-", alpha=0.6, markersize=4, color=colors[0])
ax_force.axhline(
    TARGET_FORCE__N, linestyle="--", color=colors[1], label=f"Target: {TARGET_FORCE__N:.2f} N"
)
ax_force.set_xlabel("Trial")
ax_force.set_ylabel("Force (N)")
ax_force.set_title("Force Convergence")
ax_force.legend(framealpha=1.0, edgecolor="none")

# 3. DC offset distribution
ax_hist = fig.add_subplot(gs[0, 1])
ax_hist.hist(dc_offset_vals, bins=20, alpha=0.7, color=colors[2])
ax_hist.axvline(
    best_trial.params["dc_offset"],
    linestyle="--",
    color=colors[1],
    label=f"Best: {best_trial.params['dc_offset']:.1f}Hz",
)
ax_hist.set_xlabel("DC Offset (Hz)")
ax_hist.set_ylabel("Number of Trials")
ax_hist.set_title("DC Offset Distribution")
ax_hist.legend(framealpha=1.0, edgecolor="none")

# 4. Force vs DC offset relationship
ax_scatter = fig.add_subplot(gs[1, 1])
ax_scatter.scatter(dc_offset_vals, force_achieved, c=force_errors, s=50, alpha=0.6)
ax_scatter.scatter(
    best_trial.params["dc_offset"],
    best_trial.user_attrs["force_achieved"],
    s=200,
    marker="*",
    color=colors[3],
    label="Best Trial",
)
ax_scatter.axhline(
    TARGET_FORCE__N, linestyle="--", color=colors[1], alpha=0.5, label="Target Force"
)
ax_scatter.set_xlabel("DC Offset (Hz)")
ax_scatter.set_ylabel("Force (N)")
ax_scatter.set_title("Force-DC Offset Relationship")
ax_scatter.legend(framealpha=1.0, edgecolor="none")

plt.tight_layout()
plt.savefig(
    RESULTS_DIR / "optimization_history_dc_offset.png",
    dpi=150,
    bbox_inches="tight",
)
plt.show()

print("\n[DONE] DC offset optimization complete!")
print(f"Best DC offset: {best_trial.user_attrs['dc_offset__Hz']:.2f} Hz")
print(f"Original Watanabe: 65 Hz constant → 58 Hz + oscillation (ratio: {58 / 65:.3f})")
print(
    f"This implementation: {REFERENCE_DRIVE__HZ:.1f} Hz constant → {best_trial.user_attrs['dc_offset__Hz']:.1f} Hz + oscillation (ratio: {best_trial.user_attrs['dc_offset__Hz'] / REFERENCE_DRIVE__HZ:.3f})"
)

Total running time of the script: ( 0 minutes 0.000 seconds)

Download Python source code: 02_optimize_oscillating_dc.py

Gallery generated by mkdocs-gallery