diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml index 3d8b8577..f776f8f2 100644 --- a/.github/workflows/build-wheels.yml +++ b/.github/workflows/build-wheels.yml @@ -179,9 +179,64 @@ jobs: run: | python -c "from myogen.simulator.neuron.populations import AlphaMN__Pool; print('NEURON components OK')" + test_wheels_jaxley_only: + name: Test wheel WITHOUT NEURON/MPI (Jaxley-only install) + needs: [build_wheels] + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12"] + steps: + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Download wheels + uses: actions/download-artifact@v7 + with: + name: wheels-ubuntu-latest + path: dist/ + + # Deliberately DO NOT install NEURON, mpi4py, or system MPI. This proves the + # default (Jaxley-only) install works with none of the NEURON stack present. + - name: Install wheel (core deps only) + run: | + python -m pip install --upgrade pip + # dist/ holds both cp312 and cp313 wheels; install only the one matching + # this runner's Python so pip doesn't reject the other ABI. + tag="cp$(echo '${{ matrix.python-version }}' | tr -d .)" + pip install dist/*"${tag}-${tag}"*.whl + + - name: Verify NEURON is absent + run: | + python -c " + try: + import neuron + raise SystemExit('FAIL: neuron is installed; this job must run without it') + except ImportError: + print('OK: NEURON runtime is not installed') + " + + - name: Import MyoGen without NEURON + run: | + python -c "import myogen; print('import myogen OK')" + python -c "import myogen.simulator as s; print('import myogen.simulator OK; HillModel <-', s.HillModel.__module__)" + python -c "from myogen.simulator.jaxley import run_jax, value_and_grad_run, surface_emg_jax; print('differentiable API OK')" + + - name: Build real NERLab motor neurons on the Jaxley backend + run: | + python -c " + import numpy as np + from myogen.simulator.jaxley.populations import AlphaMN__Pool + pool = AlphaMN__Pool(recruitment_thresholds__array=np.array([0.1, 0.3, 0.6]), mode='active') + assert len(pool) == 3, len(pool) + print('Jaxley-only install verified:', len(pool), 'NERLab MNs built, no NEURON present') + " + publish: name: Publish to PyPI - needs: [build_wheels, build_sdist, test_wheels] + needs: [build_wheels, build_sdist, test_wheels, test_wheels_jaxley_only] runs-on: ubuntu-latest if: github.event_name == 'release' || (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')) permissions: diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 507c6603..4d9e6c36 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -34,8 +34,11 @@ jobs: with: enable-cache: true - - name: Sync project + docs/nwb extras - run: uv sync --locked --group docs --extra nwb + - name: Sync project + docs/nwb/neuron extras + # NEURON is now an optional extra (Jaxley is the default backend). The docs + # build needs it: `poe setup_myogen` compiles NMODL, and mkdocstrings imports + # the NEURON injection utilities (api/io-and-neuron.md) to document them. + run: uv sync --locked --group docs --extra nwb --extra neuron - name: Compile NMODL mechanisms run: uv run poe setup_myogen diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6a2f22ed..ad62f289 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -46,9 +46,11 @@ jobs: python -m pip install --upgrade pip python -m pip install uv - - name: Install project with dev group + - name: Install project with dev group + optional NEURON backend + # NEURON is now an optional extra (Jaxley is the default backend), so the + # test suite must opt into it to exercise the NEURON-backed tests. run: | - uv sync --locked --group dev + uv sync --locked --group dev --extra neuron - name: Compile NEURON mechanisms run: uv run poe setup_myogen diff --git a/.gitignore b/.gitignore index f9e0e38c..a42709e4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,7 @@ ### Python template +# macOS Finder metadata +.DS_Store + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] @@ -212,3 +215,26 @@ notes/ /artifacts/ /.runtime-deps/ /graphify-out/ + +# Local working artifacts kept out of the public history +docs/session_notes/* +!docs/session_notes/2026-06-02_jaxley_nerlab_architecture_migration.md +# Evidence: only the final full runs are versioned; exploratory/intermediate +# runs stay local (git add -A must not re-add them). +experiments/twin/evidence/* +!experiments/twin/evidence/pic_grid_full_20261018_132043.json +!experiments/twin/evidence/val_comparison_bac411989ee8d42b_20261018_124637.json +!experiments/twin/evidence/val_comparison_ef1bd7e84b9abd03_20261018_124343.json +!experiments/twin/evidence/deliverable_b_full_binned_20261018_123116.json +!experiments/twin/evidence/deliverable_c_full_binned_20261018_131156.json +!experiments/twin/evidence/g4_descent_full_20261018_180008.json +!experiments/twin/evidence/g5_basin_full_20261018_191021.json +!experiments/twin/evidence/r2a_profiles_full_20261018_134617.json +!experiments/twin/evidence/sensitivity_checks_20261013_112432.json +!experiments/twin/evidence/gradient_repair_x64_20261018_154132.json +docs/DIGITAL_TWIN_ROADMAP.md +debug_*.py +diag_*.py +README-old.md +Figure_1.pdf +Figure_1_preview.jpg diff --git a/README.md b/README.md index a553fd98..e298e4da 100644 --- a/README.md +++ b/README.md @@ -1,238 +1,168 @@ -
-

- Welcome to - MyoGen Logo -

+# MyoGen — Jaxley branch -

The modular and extensible simulation toolkit for neurophysiology

+> **This is the Jaxley branch of MyoGen.** It uses **Jaxley** (JAX) as the simulation +> backend and makes the full neuromuscular closed loop **differentiable end-to-end**. +> The original NEURON-based implementation lives on the upstream **`main` branch**; use +> that branch if you need NEURON as the primary simulator. This branch keeps NEURON only +> as an optional validation reference. - [![Documentation](https://img.shields.io/badge/docs-latest-blue.svg)](https://nsquaredlab.github.io/MyoGen/) - [![Python 3.12+](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/) - [![Version](https://img.shields.io/badge/version-0.13.0-orange.svg)](https://github.com/NsquaredLab/MyoGen) +MyoGen is a modular neuromuscular simulation framework for generating physiologically +grounded motor-unit activity, muscle force, and EMG (surface and intramuscular). - [Installation](https://nsquaredlab.github.io/MyoGen/getting-started/#install) • - [Documentation](https://nsquaredlab.github.io/MyoGen/) • - [Examples](https://nsquaredlab.github.io/MyoGen/auto_examples/01_basic/) • - [How to Cite](https://nsquaredlab.github.io/MyoGen/#how-to-cite) -
- -# Overview - -MyoGen is a **modular and extensible neuromuscular simulation framework** for generating physiologically grounded motor-unit activity, muscle force, and surface EMG signals. - -It supports end-to-end modeling of the neuromuscular pathway, from descending neural drive and spinal motor neuron dynamics to muscle activation and bioelectric signal formation at the electrode level. -MyoGen is designed for algorithm validation, hypothesis-driven research, and education, providing configurable building blocks that can be independently combined and extended. - -# Highlights - -🧬 **Biophysically inspired neuron models** — NEURON-based motor neurons with validated calcium dynamics and membrane properties - -🎯 **Everything is inspectable** — Complete access to every motor unit, spike time, fiber location etc. for rigorous algorithm testing - -⚡️ **Vectorized & parallel** — Multi-core CPU processing with NumPy/Numba vectorization for fast computation - -🔬 **End-to-end simulation** — From motor unit recruitment to high-density surface EMG in a single framework - -📊 **Reproducible science** — Deterministic random seeds and standardized Neo Block outputs for exact replication - -📦 **NWB export** — Optional export to [Neurodata Without Borders](https://www.nwb.org/) format for data sharing via DANDI - -🧰 **Comprehensive toolkit** — Surface EMG, intramuscular EMG, force generation, and spinal network modeling - -# Installation - -> **Requires Python 3.12+** — Check your version with `python --version` - -## System Requirements - -| Platform | Before Installing MyoGen | -|----------|--------------------------| -| **Windows** | [NEURON 8.2.7](https://github.com/neuronsimulator/nrn/releases/download/8.2.7/nrn-8.2.7.w64-mingw-py-39-310-311-312-313-setup.exe) - Download, run installer, select "Add to PATH" | -| **Linux** | `sudo apt install libopenmpi-dev` (Ubuntu/Debian) or `sudo dnf install openmpi-devel` (Fedora) | -| **macOS** | `brew install open-mpi` | - -> [!CAUTION] -> -> ## Windows Users: Prerequisites -> -> **You MUST install the following before installing MyoGen on Windows:** -> -> ### 1. Visual C++ Build Tools -> -> Download and install [Visual C++ Build Tools](https://visualstudio.microsoft.com/visual-cpp-build-tools/). -> ->During installation, select these components: -> -> - MSVC Build Tools for x64/x86 (Latest) -> - MSVC v143 – VS 2022 C++ x64/x86 build tools -> - Windows 11 SDK (latest) -> - C++ core desktop features -> -> ### 2. NEURON Simulator -> -> 1. **Download**: [NEURON 8.2.7 Installer](https://github.com/neuronsimulator/nrn/releases/download/8.2.7/nrn-8.2.7.w64-mingw-py-39-310-311-312-313-setup.exe) -> 2. **Run the installer** and select **"Add to PATH"** when prompted -> 3. **Restart your terminal** (close and reopen) -> 4. Then continue with the installation below +On this branch, Jaxley is JAX-based, runs on CPU/GPU, and — the headline feature — lets +you take gradients of force, joint kinematics, and EMG with respect to neural, muscle, +and joint parameters. That is what the NEURON branch cannot do. --- -## Step 1: Install uv (Package Manager) - -We use [uv](https://docs.astral.sh/uv/) - a fast Python package manager. Install it first: - -**Windows** (open PowerShell): - -```powershell -powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" -``` - -**Linux/macOS**: - -```bash -curl -LsSf https://astral.sh/uv/install.sh | sh -``` +## Why Jaxley + differentiability -After installing, **restart your terminal** (close and reopen it). +The whole simulation — spiking motor neurons → muscle force → spindle/GTO +proprioception → joint dynamics → EMG — runs as a single `jax.lax.scan` and is a pure +function of its parameters. That makes it possible to: ---- +- **Fit a subject-specific digital twin**: gradient-descend neural/muscle parameters so + simulated EMG or force matches a real recording. +- **Optimize** descending drive, recruitment, muscle mechanics, or reflex gains against + any differentiable objective. +- **Accelerate** on GPU via JIT compilation. -## Step 2: Create a New Project +Gradients flow through spikes via a selectable **surrogate-gradient** mode while the +forward pass stays bit-identical to the exact spiking model. -Open a terminal and navigate to where you want your project: +## Installation ```bash -# Create a new folder for your project -mkdir my_emg_project -cd my_emg_project - -# Initialize a Python project -uv init - -# Add MyoGen to your project -uv add myogen +pip install myogen # default: Jaxley backend only (no NEURON/MPI) +pip install "myogen[neuron]" # add the optional NEURON validation backend ``` -That's it! MyoGen is now installed and ready to use. - ---- - -## Step 3: Verify Installation - -Create a test file to make sure everything works: - -```bash -# Create a test script -uv run python -c "from myogen import simulator; print('MyoGen installed successfully!')" -``` +`import myogen` requires no NEURON runtime, compiler, or MPI. -If you see `MyoGen installed successfully!` - you're all set! +## Quick start — a differentiable simulation ---- +```python +import jax, jax.numpy as jnp +from myogen.simulator.jaxley import ( + run_jax, value_and_grad_run, compile_run, surface_emg_jax, +) -## Alternative: pip install +# Build a ClosedLoopConfig + params PyTree for your network (see +# examples/01_basic/11_simulate_spinal_network.py for a full setup), then: -If you prefer pip over uv: +# 1) Forward simulation (hard spikes = scientific default) +outputs, meta = run_jax(params, config, inputs, jax.random.PRNGKey(0), spike_mode="hard") +force, angle, spikes = outputs["force"], outputs["angle_deg"], outputs["mn_spikes"] -```bash -pip install myogen -``` +# 2) Differentiate a loss w.r.t. all float parameters (surrogate spikes) +def loss(out): + return jnp.mean(out["force"] ** 2) ---- +value, grads = value_and_grad_run(loss, params, config, inputs, + jax.random.PRNGKey(0), spike_mode="surrogate") -## For Developers (From Source) +# 3) Differentiable EMG: convolve spikes with (frozen) MUAP templates +emg = surface_emg_jax(spikes.T[None].astype("float32"), muap_templates) -```bash -git clone https://github.com/NsquaredLab/MyoGen.git -cd MyoGen -uv sync -uv run poe setup_myogen +# 4) JIT-compiled forward pass +run = compile_run(config, params, spike_mode="hard") +outputs = run(params, inputs, jax.random.PRNGKey(0)) ``` -## Optional: GPU Acceleration - -For 5-10× faster convolutions (requires NVIDIA GPU): - -```bash -uv add cupy-cuda12x +### Spike modes + +| `spike_mode` | Forward pass | Gradient through spikes | Use for | +|---|---|---|---| +| `"hard"` (default) | exact discrete spikes | none (spikes non-differentiable) | scientific simulation | +| `"surrogate"` | **identical to `hard`** | smooth surrogate (biased) | optimization / fitting | +| `"rate"` | continuous rate proxy | exact | clean gradients when a continuous forward is acceptable | + +The active mode is recorded in the returned `metadata`. + +## Differentiable API (in `myogen.simulator.jaxley`) + +- `run_jax(params, config, inputs, key, spike_mode=...)` — run the closed loop; returns + `(outputs, metadata)`. +- `value_and_grad_run(loss_fn, params, config, inputs, key, ...)` — gradients of a loss + w.r.t. float parameters (integer fields held constant automatically). +- `compile_run(config, params_template, spike_mode=...)` — a `jax.jit`-compiled forward + pass bound to a config. +- `surface_emg_jax`, `intramuscular_emg_jax` — differentiable EMG from spikes + static + MUAP templates. +- `ClosedLoopConfig` — static (non-differentiable) network configuration. + +What is differentiable today: force, joint kinematics, and EMG **with respect to neural, +muscle, and joint parameters**. Differentiating EMG w.r.t. tissue conductivities / +geometry (the volume-conductor model) is future work; the enabling piece +(differentiable Bessel functions, `myogen.simulator.jaxley.bessel`) is already +validated. + +## Branches & backends + +MyoGen is developed on two branches with different primary simulators: + +| | **Jaxley branch** (this one) | NEURON branch (upstream `main`) | +|---|---|---| +| Primary simulator | Jaxley (JAX) | NEURON | +| Differentiable | **yes** | no | +| GPU | yes | no | +| Install | `pip install myogen` | (that branch's instructions) | +| Best for | gradient-based fitting, digital twins, GPU runs | exact NEURON semantics / reference results | + +On **this** branch Jaxley is the default and only required backend — `import myogen` +needs no NEURON runtime, compiler, or MPI. NEURON is still available here as an +**optional** validation reference (`pip install "myogen[neuron]"`, then +`from myogen.simulator.neuron import ...`), but it is not needed for normal use. + +The top-level `myogen.simulator` classes (`HillModel`, `Network`, `SimulationRunner`, +`SpindleModel`, `GolgiTendonOrganModel`, `JointDynamics`) resolve to the **Jaxley** +implementations on this branch. + +### Motor-neuron model + +The default motor neuron is the **NERLab** single-isopotential-dendrite cell matching +the production NEURON model: a soma (channel `napp`: fast Na + persistent Na + fast K + +slow K + leak) plus one isopotential dendrite (channel `caL`: L-type Ca + leak), in the +original 1952 Hodgkin–Huxley voltage frame (V_rest ≈ 0 mV, spikes peak ≈ +90 mV). An +older Powers2017 cable model is available via `model="Powers2017"`. + +## Where the code lives + +```text +myogen/simulator/jaxley/ + closed_loop.py # run_jax / value_and_grad_run / compile_run (differentiable entry point) + jax_models.py # pure-JAX spindle / GTO / Hill / joint steps + spike_detect + scan factory + emg.py # differentiable surface / intramuscular EMG + bessel.py # differentiable modified Bessel functions (for future volume-conductor grads) + populations/ # AlphaMN__Pool, DescendingDrive__Pool, AffIa/II/Ib__Pool, GII/GIb__Pool + channels/ # NERLab (napp, caL) and Powers2017 channels + muscle.py, proprioception/, joint_dynamics.py, network.py, simulation_runner.py + +myogen/simulator/core/ # backend-agnostic muscle geometry, force, MUAP, surface/intramuscular EMG +myogen/simulator/neuron/ # optional NEURON backend (validation reference) ``` -# Documentation +## Examples -📖 **[Read the full documentation](https://nsquaredlab.github.io/MyoGen/)** +Canonical examples default to Jaxley — see `examples/README.md` for the full map. +`XX_.py` is the Jaxley version; `XX__neuron.py` (where present) is the +NEURON counterpart and needs `myogen[neuron]`. The full closed-loop spinal network is +`examples/01_basic/11_simulate_spinal_network.py`. -- [User Guide](https://nsquaredlab.github.io/MyoGen/neo-blocks/) — Working with simulation outputs -- [API Reference](https://nsquaredlab.github.io/MyoGen/api/) — Complete class documentation -- [Examples](examples/) — Step-by-step tutorials from recruitment to EMG +## Status & caveats -# How to Cite +- Surrogate spike gradients are **biased** — inherent to differentiating spiking models; + use `rate` mode for unbiased gradients when a continuous forward output is acceptable. +- This branch running on Jaxley is distinct from "Jaxley reproduces NEURON exactly + everywhere" — some quantitative gaps remain (e.g. F-I gain on some examples). If you + need NEURON as the reference, use this branch's optional `[neuron]` extra or the + upstream NEURON branch. +- Not yet ported to Jaxley: the `02_finetune` and `03_papers/watanabe` NEURON + simulations (runnable via the extra). -If you use MyoGen in your research, please cite: - -```bibtex -@article{simpetru_molinari_2026_myogen, - title = {MyoGen: Unified Biophysical Modeling of Human Neuromotor Activity and Resulting Signals}, - author = {S{\^i}mpetru, Raul C. and Molinari, Ricardo G. and Rohlf, Devon R. and - Batichotti, Rebeka L. and Watanabe, Renato N. and - Elias, Leonardo A. and Del Vecchio, Alessandro}, - journal = {bioRxiv}, - note = {preprint}, - year = {2026}, - doi = {10.64898/2026.01.01.697284}, - url = {https://www.biorxiv.org/content/10.64898/2026.01.01.697284} -} -``` +## Further reading -# Contributing - -Contributions welcome! See [issues](https://github.com/NsquaredLab/MyoGen/issues) if you want to add a feature or fix a bug. - -# Contributors - -MyoGen is authored by **Raul C. Sîmpetru** and **Ricardo G. Molinari**. It is a joint project of the [Nsquared Lab](https://nsquared.tf.fau.de/) (Neuromuscular Physiology and Neural Interfacing Laboratory) at FAU Erlangen-Nürnberg, the [NER Lab](https://www.ceb.unicamp.br/pesquisa/laboratorios/laboratorio-de-pesquisa-em-neuroengenharia/) (Neural Engineering Research Laboratory) at the University of Campinas, and the [BMClab](https://bmclab.pesquisa.ufabc.edu.br/) (Biomechanics and Motor Control Laboratory) at the Federal University of ABC. - -Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/en/reference/emoji-key/)): - - -[![All Contributors](https://img.shields.io/badge/all_contributors-8-orange.svg?style=flat-square)](#contributors) - - - - - - - - - - - - - - - - - - - - - - - - - -
Raul C. Sîmpetru
Raul C. Sîmpetru
💻 🤔 🚧 📆 🔬		Ricardo G. Molinari
Ricardo G. Molinari
🤔 🚧 🔬		Devon R. Rohlf
Devon R. Rohlf
💻 🚧 🔬		Renato N. Watanabe
Renato N. Watanabe
💻 🤔 🧑‍🏫 🔬		Alessandro Del Vecchio
Alessandro Del Vecchio
🤔 🧑‍🏫 🔬		Leonardo A. Elias
Leonardo A. Elias
🤔 🧑‍🏫 🔬		João Bittar
João Bittar
🐛 🤔
BraveUnicorn
BraveUnicorn
💻
- - Add your contributions - -
- - - - - - -This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome! - -# License - -MyoGen is AGPL licensed. See [LICENSE](https://github.com/NsquaredLab/MyoGen/blob/main/LICENSE.md) for details. +- `docs/DIFFERENTIABILITY_STATUS.md` — what is differentiable, how it was validated. +- `docs/NEURON_DECOUPLING.md` — how Jaxley became the default and NEURON optional. +- `docs/M1_bessel_feasibility.md` — differentiable volume-conductor feasibility. +- `examples/README.md` — the Jaxley (default) vs NEURON example map. diff --git a/docs/DIFFERENTIABILITY_STATUS.md b/docs/DIFFERENTIABILITY_STATUS.md new file mode 100644 index 00000000..00522eca --- /dev/null +++ b/docs/DIFFERENTIABILITY_STATUS.md @@ -0,0 +1,79 @@ +# MyoGen Differentiability — Implementation Status + +Goal: make the MyoGen simulation differentiable end-to-end in JAX (force, joint +kinematics, and EMG) w.r.t. neural, muscle, and joint parameters, with a hard-spike +scientific mode and surrogate/rate modes for optimization. + +**Key reframing (verified in the codebase):** most of the compute core was already +pure-JAX. Retiring NEURON is orthogonal to differentiability. The blockers were +narrow: hard spike thresholds, stochastic generators, a scipy muscle-setup step, +and the numpy EMG summation. The EMG splits into an easy tier (differentiable +w.r.t. neural/muscle/joint params — a convolution against frozen MUAP templates) +and a hard-but-optional tier (differentiable w.r.t. tissue conductivity / geometry — +needs a volume-conductor port). + +## Delivered (validated) + +| Milestone | What | Files | Validation | +|---|---|---|---| +| **M0** | `run_jax` closed-loop entry point; differentiable-PyTree vs static-config split; explicit PRNG keys; `spike_mode` + metadata; `value_and_grad_run` (float-leaf partition) | `jaxley/closed_loop.py` | end-to-end run in 3 modes; finite non-zero grads; int leaves excluded | +| **M2** | Surrogate-gradient spike primitive (`spike_detect`, custom-JVP straight-through); hard/surrogate/rate modes wired into the scan | `jaxley/jax_models.py` | hard≡surrogate forward (bit-identical); hard grad=0, surrogate/rate grad≠0 | +| **M3** | Differentiable stochastic generators: `rate` (exact expected-value) and `pathwise` (frozen-sample + surrogate) modes | `jaxley/jax_models.py` | hard≡pathwise forward; rate grad = exact `T·N·dt` | +| **M4** | `differentiable_twitch_params` — closed-form Fuglevand path (P, T, twiAmp, IIR) in JAX; gradients w.r.t. RP/Tl/RT/fP (saturation `tetF` frozen) | `jaxley/jax_models.py` | matches scipy path to 6e-8; finite grads | +| **M5** | Differentiable EMG: `surface_emg_jax`, `intramuscular_emg_jax`, `resample_muaps` (spikes ⋆ frozen MUAP) | `jaxley/emg.py` | matches numpy `correlate` to 2e-7; spike→EMG grad finite | +| **M1** | Differentiable modified Bessel `iv_int`/`kv_int` (Miller/upward recurrence + small-x series; auto-diff, no custom_jvp) — **GO** | `jaxley/bessel.py` | values & grads vs scipy < 1.7e-7 | +| **M6** | New differentiable API exported from `myogen.simulator.jaxley`; backward-compat preserved (defaults unchanged) | `jaxley/__init__.py` | package imports; hard-mode defaults return bool as before | +| **M8** | Per-stage FD gradient checks (spindle/gto/joint/hill/twitch) + end-to-end | `scripts/test_gradient_checks.py` | all stages autodiff ≈ FD | + +**Headline result:** gradient of a surface-EMG loss flows through spikes to a neural +weight (`d(EMG loss)/d(base_dd)` finite and non-zero). Force + joint + EMG are +differentiable w.r.t. neural, muscle, and joint parameters today. + +### Validation scripts +- `scripts/test_differentiable_pipeline.py` — end-to-end (M0/M2/M3/M5) on a tiny network. +- `scripts/test_gradient_checks.py` — per-stage autodiff-vs-FD (M8). +- `scripts/prototype_bessel_jax.py` — Bessel values/gradients vs scipy (M1). +- `scripts/test_old_vs_new_equivalence.py` — **regression**: OLD (git HEAD) vs NEW + full-scan output, bit-identical (0.0) across all channels in `hard` mode. Proves + the refactor changed no numbers. +- `scripts/test_nerlab_differentiable.py` — **runs on the REAL NERLab motor neurons** + (napp+caL, V_rest≈0, spike peak ≈+90 mV). Confirms MNs spike, hard≡surrogate + forward, and surrogate gradients flow to the descending-drive weight AND to the + actual channel conductances (`napp_gnabar`, `napp_gkfbar`, `caL_gcaLbar`). + +### Differentiating w.r.t. Jaxley channel parameters — required step +`net.get_parameters()` returns **only** parameters marked with +`net.make_trainable("napp_gnabar")` (etc.). Without `make_trainable`, the neural +parameter subtree is empty and gradients flow only to the muscle/joint/weight +params. Mark the conductances (or gates) you want to optimize trainable **before** +`build_init_and_step_fn`, then they appear in `params["jaxley"]` and receive +gradients. + +## Remaining + +**M7 — full volume-conductor JAX port (optional, de-risked, not yet done).** Only +needed for gradients w.r.t. tissue conductivities / conduction velocity / geometry. +M1 proved the hard enabler (differentiable arbitrary-order Bessel) is feasible; the +rest is mechanical (`np.linalg.solve`→`jnp`, `np.fft`→`jnp.fft`, freeze the argmax +peak-centring, smooth the iEMG step masks). Also needs a native `jv` (same technique +as `iv`/`kv`; `j0`/`j1` absent from JAX). Est. 700–1300 LOC + validation. See +`docs/M1_bessel_feasibility.md`. + +**Deeper M6 rewire (optional).** The legacy fake-JAX `jaxley/muscle.py` class and the +NumPy `jaxley/proprioception/*` wrappers can be rewired to delegate to the +`jax_models.py` compute path so the *class* API also exposes gradients. The +functional path (`run_jax`) already does; this is ergonomics. + +## JIT usage +`run_jax` cannot be handed to `jax.jit` directly: `ClosedLoopConfig` is intentionally +unhashable (holds arrays), and integer Hill fields (`N`/`Ntype1`) are used in shape +contexts (`jnp.arange(N)`) so must stay static. Use the wrappers, which do the +static/dynamic split: `value_and_grad_run` for gradients (primary path), and +`compile_run(config, params_template)` for a `jax.jit`-compiled forward pass. + +## Notes / caveats +- Surrogate spike gradients are biased (inherent to differentiating spiking models); + `rate` mode gives cleaner gradients at the cost of a continuous forward output. +- Pathwise generator gradients are pathwise-only (fixed PRNG sample). +- M4 leaves the tetanic-saturation constant `tetF` frozen; gradients w.r.t. the + saturation-shape constants are the low-value tail and are not propagated. diff --git a/docs/DIGITAL_TWIN_ROADMAP_REVISED.md b/docs/DIGITAL_TWIN_ROADMAP_REVISED.md new file mode 100644 index 00000000..9e7ce9e0 --- /dev/null +++ b/docs/DIGITAL_TWIN_ROADMAP_REVISED.md @@ -0,0 +1,397 @@ +# Digital Twin of the Motoneuron Pool — Revised Research and Implementation Roadmap + +**Date:** 13 October 2026 +**Basis:** the supplied “Digital Twin of the Motor Neuron Pool — Roadmap”, updated 13 October 2026. +**Review scope:** scientific design and implementation priorities, not a reproduced audit of the repository. Statements about existing functionality and numerical results below are reported by the supplied roadmap, not independently verified here. Proposed work is explicitly prospective. + +## 1. Goal and first result + +Retain the long-term goal: use MyoGen/Jaxley to infer aspects of motoneuron-pool physiology from motor output and predict responses under new conditions. + +Replace the immediate goal of recovering all hidden physiology with: + +> Establish which PIC and effective-drive parameters can be distinguished under specified protocols and observation errors, and determine whether qualified gradients improve recovery over a hard-forward baseline. + +The first result should be a small, reproducible inverse-problem experiment, not a large cohort attribution figure. A successful fit, unique physiological identification, and successful prediction are different outcomes and must be reported separately. + +Use this research sequence: + +**Define the inference problem → verify the forward model → qualify observations and sensitivities → determine identifiability and informative protocols → demonstrate recovery → test mechanism discrimination under model mismatch → evaluate human prediction.** + +Forward-only identifiability work and gradient debugging may run in parallel. Do not make all scientific progress depend on repairing long-horizon reverse-mode differentiation. Conversely, do not make a gradient-based inference claim without qualifying the gradients used for that claim. + +## 2. What the supplied roadmap actually establishes + +The source describes a differentiable proof-of-concept harness, not a validated twin or a working joint inverse solver [S1, lines 8–16, 166–173]. It reports: + +- `PoolTwin` support for CaL/NaP scaling, drive gain, excitation/inhibition weights, alternative synapse locations and kernels, and fixed/live driving-force options. Afferent input, force, EMG, and the full closed loop are not yet part of this harness [S1, lines 27–35]. +- Effective negative-current inhibition, rather than a physiological shunting conductance; a strong drive-gain/excitatory-weight degeneracy; and an Erlang renewal input when `dd_batch=5` [S1, lines 37–46]. +- Legacy forward comparisons, but an outstanding strict PIC-specific validation package [S1, lines 48–83]. +- Long-horizon gradient explosions, a biased rate-forward mode, and unstable inverse fitting without valid joint recovery [S1, lines 138–173]. + +Treat existing forward comparisons as scoped regression evidence. Do not silently promote them to validation of a new synaptic configuration, PIC regime, observation model, or human inference task. + +The supplied document references code and session notes that were not included. Their numerical claims must be reproduced from versioned configurations before being used as acceptance evidence. + +## 3. Define the inference problem before choosing the optimizer + +For contraction c of subject s, write the conceptual model as: + + y_sc ~ MeasurementModel( + Simulator(theta_s, u_sc(t), x0_sc, omega_sc), + eta_sc + ) + +Here: + +- `theta_s`: the physiological or effective parameters being investigated. +- `u_sc(t)`: trial-specific neural inputs, which are known in a controlled simulation but generally not measured in human recordings. +- `x0_sc`: initial and history-dependent states, including channel gates, calcium and synaptic states. +- `omega_sc`: stochastic realization and, where appropriate, connectivity realization. +- `eta_sc`: measurement and muscle nuisance parameters, such as missing-unit detection, timing error, twitch properties and electrode geometry. +- `y_sc`: the actual measured data or explicitly defined summaries. + +A target force trajectory is not the same variable as descending neural drive. Every experiment must declare which of these quantities is known, estimated, shared, fixed, or marginalized. A conditional fit with known inputs is a useful experiment, but it is not a demonstration of joint excitation inference. + +### Parameter ladder + +**Level A — numerical diagnostic:** one cell, deterministic known current, one CaL parameter. Intracellular voltage/current observations are allowed here only to test implementation and sensitivities. + +**Level B — first output-only inverse problem:** a small heterogeneous pool, two unknowns: CaL scaling and one continuously parameterized effective excitatory-input amplitude. Fix NaP and inhibition explicitly. Use known drive shape and an identical observation definition for every method. + +**Level C — conditional physiological separation:** add NaP, then inhibitory strength, separately. Each addition requires an identifiability test. Fixing a parameter at this stage is a declared restriction, not evidence that the parameter has been identified. + +**Level D — latent inputs and observation uncertainty:** estimate a low-dimensional trial-specific input representation and measurement nuisance parameters jointly with the retained pool parameters. + +**Level E — afferent feedback:** add afferent gain only after defining the pathway, its input, delays, mechanical dependencies and accessible perturbations, and verifying the implemented forward pathway. + +Do not assume that CaL and NaP are individually recoverable. If only a combination or a protocol-specific functional measure is supported, report that result and its uncertainty. Such a measure is not automatically a general PIC-conductance estimate. + +## 4. Phase 0 — Reproducible experiment specification + +**Purpose:** make each claim traceable to an exact forward model, objective and parameterization. + +Keep `experiments/twin/` as the experiment entry point. Reuse existing harness functions and runnable scripts where possible; do not reorganize the project solely to renumber phases. + +Add a machine-readable experiment specification containing: + +- Code revision, dependency versions, device, precision, integration method, time step and simulated duration in physical units. +- Pool size, compartment counts, morphology/channel parameters, units and voltage convention. +- Synapse site, kernel normalization, reversal potential, delays, event construction and inhibition representation. +- Parameter transformations, bounds, baseline values and mapping to the actual channel/input leaves being modified. +- Parameter-dependent versus fixed initialization, preconditioning and rest history. +- Connectivity, input-event and measurement-noise seeds, separately identified. +- Observations, smoothing and extraction rules, missing-data rules, objective, optimizer budget, checkpoint selection and train/validation/test splits. + +The source's HH voltage-frame convention must be preserved explicitly [S1, lines 151–152]. Check reversal potentials and detection thresholds under the same convention; do not apply a voltage offset to only one part of the model. + +Separate three computational contracts: + +1. **Scientific hard forward:** the original dynamics and observation/event definitions used to generate and evaluate scientific outputs. +2. **Surrogate backward:** a separately identified derivative estimator. Its primal output must match the scientific hard forward when advertised as hard-forward/surrogate-backward. +3. **Approximate rate forward:** a different forward model. It may be studied as a reduced model or warm start, but must not silently replace the scientific forward. + +**Deliverable:** experiment manifest, tiny deterministic fixtures, and an evidence registry linking every gate to raw results and its allowed operating envelope. + +**Gate:** a clean rerun reproduces the declared fixtures and identifies every parameter that actually changes the simulation. No gate passes solely because a script exits successfully. + +## 5. Phase 1 — Forward verification and physiological operating envelope + +**Purpose:** establish that the implementation represents the intended equations within a defined range. Agreement with NEURON is cross-implementation verification, not independent proof of human physiological validity. + +### 1A. Intrinsic dynamics + +Retain the source plan's matched somatic and dendritic current injections and PIC grid. Match units, morphology, initial states, channel definitions and stimulation exactly. Compare subthreshold voltage, short-horizon spike timing/counts, channel-current trajectories, recruitment/derecruitment, plateau activation and termination. + +Add time-step refinement and precision diagnostics: compare the chosen time step with smaller steps, and use float64 reference runs where feasible. Separate numerical instability from sensitivity of a correct trajectory. Use multiple cell types across the intended pool distribution. + +Retain the source's high-CaL/high-NaP cases as **stress and bifurcation tests**, not as human priors. Record absolute conductance densities, baseline scaling and current magnitudes; a multiplier alone does not define a physiological regime. Add appropriately scoped channel-off controls and refine the grid around changes in firing/plateau state. Do not require bistability in every cell or every parameter region. + +Long stochastic spike trains need not remain phase-aligned indefinitely to demonstrate agreement in their scientific summary distributions. Predeclare which tests require trace/spike-time equivalence and which require agreement of distributions. + +### 1B. Synaptic and pool dynamics + +Retain identical event matrices, connectivity, dendritic placement, kernels, weights and reversal potentials across simulators. Do not retune Jaxley to force agreement with NEURON. + +An effective negative-current inhibitory model can be verified as an effective-current model. To interpret an inferred parameter as an inhibitory conductance or evaluate shunting effects, implement and verify the corresponding conductance-based pathway first. A live driving-force option for another synapse does not establish that inhibitory shunting has been implemented. + +Across independently sampled pools and event realizations, compare recruitment, rate trajectories, active-unit counts, hysteresis, eligible ΔF statistics and after-discharge. Use identical ΔF eligibility rules in both simulators and report the fraction of eligible pairs. + +**After-discharge is a primary endpoint.** When units do not derecruit before recording ends, treat duration as right-censored. Do not substitute a maximal duration as an observed termination time, set unavailable ΔF to zero, or discard these trials without reporting them. + +### 1C. Acceptance criteria + +Define separate tolerances for numerical agreement and biological adequacy. Choose numerical tolerances using refinement studies and the smallest output differences the inverse experiment intends to distinguish. Freeze final thresholds before confirmatory runs; exploratory runs may establish the thresholds, but cannot also serve as their independent confirmation. + +The source does not supply enough error distributions or intended biological effect sizes to justify universal numerical cutoffs. Do not invent them as established standards. + +**Deliverable:** forward verification report with trace panels, regime boundaries, error distributions, failed cases and a machine-readable validated envelope. + +**Gate:** the exact configuration and range used in the next scientific experiment pass the relevant tests. Failure outside that envelope limits scope rather than automatically blocking all work. + +## 6. Phase 2 — Qualify observations and sensitivities + +Run the two workstreams below in parallel on the verified model. + +### 2A. Observation model + +Start with spike-derived observations; force and raw EMG are not mandatory for the first recovery experiment. + +Build an observation ladder: + +1. Oracle: all simulated spike trains, exact identities and known inputs. +2. Partial observation: only a subset of units, with the observation fraction reported. +3. Imperfect observation: missed/extra spikes, timing errors, uncertain recruitment/derecruitment, and incomplete or uncertain unit matching. +4. Selection sensitivity: compare random missingness with plausible recruitment-, amplitude- or quality-dependent detection scenarios. Until calibrated, label these as stress assumptions rather than measured decomposition performance. +5. Force: add a verified muscle/twitch model, noise and nuisance uncertainty. +6. Raw HD-sEMG: add only when waveform information or activity from undecomposed units is needed for a defined inference question, with a validated signal-generation/observation pathway. + +Process simulated and experimental observations through the same extraction rules. Estimate noise and repeatability from appropriate calibration data when available. Test the sensitivity of conclusions to uncertain observation-error models. + +ΔF is derived from motor-unit discharge data, and its behavior depends on pair selection and discharge characteristics [R4]. Adding ΔF to sufficiently resolved spike/rate data is not a new independent measurement. It can be a useful summary or objective emphasis; model the dependence or call the loss a composite objective. Do not multiply dependent likelihood terms and interpret the resulting sharpness as new evidence. + +Similarly, force is not automatically independent of complete spike trains and a known muscle model. It may provide useful additional constraints with partial neural observation or uncertain muscle dynamics. Raw EMG may contain information from undecomposed units, but its added value must be tested rather than assumed. + +**Deliverable:** observation operator and corruption tests, a declared initial objective, and eligibility/censoring reports. + +**Gate:** the recovery experiment uses only available observations and propagates the stated uncertainty. No hidden oracle identities, inputs or unrecruited-unit labels enter a human-facing fit. + +### 2B. Gradient qualification + +The source reports exploding gradients and unstable optimization [S1, lines 140–150, 166–173]. Move this from a caveat to a qualification experiment. + +**Test 1 — parameter plumbing.** Verify that each parameter changes the intended leaves, cells and compartments only. Audit masks, NaN sentinels, detached computations and parameter-dependent initialization. A nonzero gradient alone is not sufficient. + +**Test 2 — continuous sensitivities.** Start with one cell, deterministic inputs, short trajectories and smooth voltage/current objectives. Compare autodiff directional derivatives against centered finite differences across several perturbation sizes and integration settings. Distinguish a finite-difference convergence region from steps dominated by roundoff or numerical/event changes. + +**Test 3 — events and stochastic inputs.** Audit the full path from each parameter through random sampling, event times, binning and observation extraction. Fixing a PRNG key provides reproducibility, not automatic differentiability with respect to a distribution parameter. For example, a hard threshold event can have zero derivative almost everywhere for fixed random input while its expected event count changes with the rate. JAX's Poisson sampler produces integer-valued samples [R5]; the actual harness's Erlang/event implementation must be examined separately. + +Smoothing a hard-binned spike train after threshold detection does not, by itself, remove the discontinuity in how the bins depend on parameters. Identify whether the derivative comes from continuous event-time treatment, an explicit surrogate, a stochastic gradient estimator or a different forward approximation. + +For the first joint inverse experiment, prefer known event input with a continuously scaled effective amplitude, or a deterministic current/conductance input. This deliberately avoids presynaptic-rate inference and must be labeled accordingly. + +**Test 4 — horizon and regime.** Extend tests from short trajectories to the actual ramp/hold durations, including plateau transitions. Report gradient behavior by horizon in seconds, parameter region, precision and objective. Distinguish derivative correctness from optimization usefulness near a regime transition. + +**Test 5 — hard-forward descent.** For surrogate derivatives, evaluate proposed parameter changes with the scientific hard forward, using independent validation seeds as well as controlled common-random-number comparisons. Report agreement with meaningful finite perturbations, accepted-step fraction, achieved hard-loss improvement and failure rate. Do not demand infinitesimal finite-difference agreement from a discontinuous hard-binned spike objective; test continuous sensitivities and practical descent separately. + +### Optimizer decision ladder + +Normalize parameter scales and residuals. Use bounded steps or trust-region/line-search safeguards where appropriate. Clipping may prevent overflow, but it does not validate derivative direction or physiological inference. Log unclipped and clipped gradients and optimizer failures. + +Checkpointing trades memory for recomputation; it does not repair an unstable sensitivity [R6]. Shorter backpropagation windows alter the derivative and may omit long-history effects. If truncation is studied, preserve all forward states across boundaries, vary window length and evaluate complete uninterrupted hard-forward trajectories. + +If needed, evaluate multiple shooting as a separate constrained formulation: introduce boundary states and enforce continuity of voltage, gates, calcium, synaptic and other relevant states. Do not restart each window from rest or allow unconstrained boundary states to explain away physiology. Benchmark this approach; do not promise it will solve the gradient problem [R7]. + +For the initial two- to five-parameter tasks, compare forward-mode sensitivities, finite perturbations and derivative-free optimization before committing to large reverse-mode Jacobians. These are alternatives to assess, not automatic cures for sensitive dynamics. + +**Deliverable:** parameter-by-objective-by-horizon gradient report plus a simple hard-forward recovery baseline. + +**Gate:** promote a gradient method only within the range where it is numerically qualified and reliably improves the declared hard-forward objective. Otherwise continue forward-only identifiability/recovery work while the gradient method remains experimental. + +## 7. Phase 3 — Identifiability and protocol design before attribution + +**Purpose:** determine what information the experiment contains before interpreting an optimizer's preferred solution. + +Begin with the two-parameter PIC/effective-drive problem. Compute hard-forward loss surfaces and approximate equivalence regions across several true parameter points, including near-transition cases. Repeat across observation levels and input assumptions. + +Distinguish an exact model symmetry from a practical trade-off. If two inputs enter only as `(a-b)*u(t)`, only their difference is observable in that illustrative submodel. But a rate/weight product that controls mean synaptic input need not be an exact symmetry of all higher-order output statistics. Test, rather than assume, whether variability or temporal structure breaks the source's reported drive-gain/excitatory-weight degeneracy. + +### Local sensitivity analysis + +For a smooth observation mean `mu(phi)`, use dimensionless or explicitly scaled parameters `phi`. With a fixed Gaussian observation covariance `Sigma`, the local information matrix is: + + I(phi) = J(phi).T @ inverse(Sigma) @ J(phi) + J(phi) = d mu / d phi + +This is not universally `J.T @ J`. The unweighted expression assumes a suitable whitened unit-noise model. Correlated rate samples, shared units and redundant summaries need corresponding treatment. If covariance depends on parameters, this mean-sensitivity term is not the full Gaussian Fisher information. + +Use qualified sensitivities only. An autodiff matrix computed through surrogate events is a surrogate-sensitivity diagnostic unless justified for the scientific observation distribution. Do not infer Fisher information from clipped optimizer gradients. + +Vary nuisance parameters and latent-input coefficients rather than holding them at truth in every analysis. Profile over them, or use an appropriate local nuisance adjustment when valid. Choose forward- versus reverse-mode Jacobians according to parameter and observation dimensions, not by habit. + +### Nonlocal checks and uncertainty + +Use profile likelihoods or explicitly calibrated loss profiles: fix one parameter, re-optimize the others and nuisance quantities, and examine whether substantially different values remain compatible with the observations. Profile-likelihood methods address parameter ambiguity beyond a single local curvature calculation [R2]. If using a composite discrepancy rather than a likelihood, do not label its profiles as likelihood confidence intervals without calibration. + +Evaluate multiple operating points, initial histories, seeds and pool realizations. Distinguish data-supported constraints from parameter bounds or informative priors. Use simulation-based coverage tests before reporting confidence/credible intervals as calibrated. Bootstrap whole simulated subjects/pools/trials at the appropriate level, not all motor-unit pairs as independent observations. + +### Protocol selection + +Start with a finite, interpretable candidate set: different ramp slopes and amplitudes, ramp–hold–ramp trials, varied conditioning history/rest, and brief controlled perturbations in simulation. Compare protocols under an equal recording-time and measurement-quality budget. + +For afferent questions, define feasible measured mechanical/sensory perturbations before optimizing the protocol. A simulated known inhibitory pulse is not automatically a selective intervention available in a human experiment. + +Select protocols using held-out parameter draws and nuisance conditions. Seek parameter discrimination and predictive utility, not merely a large local determinant at one favorable point. A protocol search must not exploit an unvalidated surrogate regime. + +**Deliverable:** identifiable parameter combinations, ambiguity regions, uncertainty calibration and a small selected protocol set. + +**Gate:** proceed with the subset of parameters and claims supported at a prespecified scientific resolution. A reproducible non-identifiability result is a valid outcome and a reason to change the measurement/protocol or parameterization. + +## 8. Phase 4 — Staged synthetic recovery + +**Purpose:** separate optimizer success, observation limitations and model mismatch. + +Use nested experiments rather than starting with all five mechanisms: + +- **R0:** one parameter, known deterministic input, same-model noiseless data. This is an implementation test. +- **R1:** CaL scaling plus effective excitatory-input amplitude, multiple true points and initializations, output-only observations. +- **R2:** add NaP and then inhibition individually, conditional on Phase 3 evidence. +- **R3:** incomplete/noisy observations and unknown event realizations; estimate retained parameters and nuisance quantities. +- **R4:** generate observations in NEURON and fit with Jaxley. This tests cross-implementation robustness but still shares the biological model assumptions. +- **R5:** introduce controlled structural mismatch: altered cell heterogeneity, AHP/leak properties, synapse placement/kinetics, or other plausible omitted mechanisms. Add twitch/measurement mismatch only when those components are included. + +Same-simulator recovery with identical structure and fixed inputs is a necessary unit test, not a sufficient argument for human physiological recovery. + +### Required comparisons + +For the small initial problem, use a coarse grid followed by bounded refinement as the transparent reference. For larger low-dimensional problems, choose one suitable derivative-free baseline rather than adding a large optimizer zoo. Compare qualified gradient optimization using equal wall-time or explicitly accounted evaluation budgets, including failed restarts and gradient costs. + +Distinguish methods that share the target input realization from methods that must infer/marginalize unknown realizations. Common random numbers can reduce comparison noise, but a hidden ground-truth seed must not become an undeclared inference input. + +### Required reporting + +Report prediction error and parameter error separately, bias, uncertainty coverage, competing optima, boundary hits, failure rate and wall-time to an accepted solution. Display all starts in the analysis, not only successful recoveries. + +Predeclare checkpoint selection. Selecting the minimum training-loss iterate is legitimate when stated, but report the final iterate and stability too. Use separate validation data for tuning/early stopping and an untouched test set for final recovery/generalization claims. Never select checkpoints by closeness to the true hidden parameter. + +**Deliverable:** recovered-versus-true plots with uncertainty, hard-forward fits, held-out predictions, baseline comparisons and full failure accounting. + +**Gate:** the declared recovery resolution and uncertainty coverage hold on untouched synthetic cases within the stated envelope. A low loss without parameter recovery supports a predictive fit, not identified physiology. + +## 9. Phase 5 — Mechanism discrimination under ambiguity and mismatch + +**Purpose:** keep the output-matched-cohort idea, but make the experiment test information rather than optimizer preference. + +Define hypotheses in terms of the implemented model. Distinguish increased CaL/NaP conductance, reduced effective inhibition, increased afferent gain and changes in neural drive. Do not label an unimplemented afferent pathway or a negative-current term as a verified physiological conductance mechanism. + +Construct cohorts matched on specified baseline observations, with a prespecified equivalence margin tied to measurement repeatability. State whether matching concerns means, full trajectories or their probability distributions. Failure to reject a difference is not evidence of equivalence. + +Include two distinct experiments: + +**Information-null control:** when hypotheses are genuinely indistinguishable under the observations available to the fitter, the method should remain uncertain. With balanced labels and exactly equal observation distributions, above-chance classification would imply leakage or a violation of the matching construction, not discovery of hidden information. + +**Perturbation-discrimination test:** preserve baseline matching, then evaluate a held-out protocol whose predicted responses differ across hypotheses. Determine whether the new measurements actually resolve the ambiguity. Match or randomize nuisance distributions across mechanisms so that a classifier cannot exploit seed, recording-length, missingness or file-structure artifacts. + +Include pure mechanisms, combinations, null/no-change cases and out-of-model generators. Allow mixed or indeterminate conclusions rather than forcing every case into one label. A single optimizer's parameter allocation is not a mechanism probability. + +Compare candidate models with explicit complexity/prior assumptions and held-out predictive performance. Report calibrated evidence or uncertainty where supported. Evaluate attribution under misspecification, not only when the fitted model exactly contains the data-generating mechanism. + +**Deliverable:** baseline-equivalence plots, held-out perturbation responses, confusion matrices with indeterminate outcomes, and false-attribution rates under mixtures/mismatch. + +**Gate:** informative protocols improve mechanism discrimination while information-null and out-of-model controls do not produce unjustified confidence. Scope conclusions to model-conditional mechanism evidence; do not equate them with causal identification in humans. + +## 10. Phase 6 — Human predictive evaluation + +**Purpose:** establish subject-specific predictive utility before claiming access to true hidden conductances. + +Use decomposed HD-sEMG and force as available. Start with a clearly defined muscle/task, repeated contractions, documented history and quality assessment. Implement the source's NWB/loading and motor-unit matching needs, with explicit uncertainty when identities are not reliable [S1, lines 127–134]. Inspect data availability and repeatability early; do not postpone all contact with human data until this phase. + +Fit shared subject-level parameters across multiple calibration contractions. Allow declared trial-specific inputs, initial states and measurement terms. A shared-physiology assumption must be tested across the chosen tasks; changes in neuromodulatory state or fatigue must not be silently absorbed into a supposedly constant subject trait. + +### Choose the prediction task explicitly + +**Conditional held-out-unit prediction:** estimate the new trial's latent drive from a designated subset of observed units, freeze the subject physiology, and predict disjoint held-out units. No scored unit contributes to input estimation. Report this as conditional prediction, not input-free forecasting. + +**New-protocol forecasting:** use a separately specified mapping from the known task/target to neural input, or a calibrated closed-loop controller. Freeze the relevant rules before testing and propagate input uncertainty. Do not insert an unmeasured neural-drive waveform as though it were known. + +**Held-out force prediction:** do not reconstruct drive from the same measured force trace and then use agreement with that trace as independent validation. The force target may be a known task instruction, but it is not the measured output and is not a measured descending command. + +Use held-out contractions with different slopes, amplitudes, holds and histories. Add sessions or re-donning tests when claiming session transfer. Compare against a simpler fixed-physiology/common-drive baseline with the same available inputs and a comparable fitting budget. + +Evaluate force, firing-rate trajectories, recruitment/derecruitment, ΔF where eligible, after-discharge, uncertainty calibration and test–retest stability. Keep participants/contractions as the appropriate independent evaluation units. Do not use large numbers of unit pairs to inflate the effective sample size. + +Human output prediction does not directly reveal true CaL, NaP or inhibitory conductances. Initially report effective/model-conditional parameters, identifiable combinations and predictive uncertainty. Stronger mechanistic claims require additional independent constraints or perturbation evidence. + +**Deliverable:** leakage-free held-out predictions, repeatability and uncertainty results, baseline comparisons, and a precise statement of what the model can and cannot infer. + +## 11. Computational benchmarking — a supporting workstream + +Retain the existing benchmark script if present, but change its purpose from a headline to measuring feasible experiment cost. + +Measure compilation separately from steady-state execution. Synchronize timed JAX work with `block_until_ready()` or an equivalent appropriate synchronization; device dispatch alone is not completed computation [R3]. + +Benchmark the scientific forward and qualified gradient methods across pool/compartment size, simulated duration, batch size, precision, recording configuration and seed count. Report peak memory, failures, hardware and full configurations. Use chunked batches when a full `vmap` exceeds memory. + +Averaging over seeds is not free on a GPU. Batching can improve throughput but still consumes computation and memory. Distinguish many parameter sets of a small pool from one large physiological pool. + +Report both simulation throughput and wall-time to a successful, verified recovery. A fast derivative that requires many failed starts may not improve inference. Compare matched workloads against NEURON where relevant; do not present a single-thread NEURON baseline as a universal state-of-the-art speed comparison. + +The source's five-cell, K={1,4} result remains a batching smoke test until larger, relevant configurations are actually measured [S1, lines 85–96]. + +## 12. The immediate implementation queue + +Do not build the whole future pipeline at once. Implement these three deliverables first, reusing the existing harness. + +### Deliverable A — scoped forward and gradient diagnostics + +Reproduce a one-cell and small-pool verified case. Add smooth-objective finite-difference checks, event/rate-parameter plumbing checks, horizon diagnostics and hard-forward step evaluation. Preserve raw failures. + +### Deliverable B — a two-parameter hard-forward experiment + +Choose CaL scaling and effective excitatory-input amplitude, with NaP/inhibition fixed and known drive shape. Use a ramp and a ramp–hold–ramp protocol, plus an untouched protocol for prediction. Compare a transparent grid/refinement baseline with the current gradient method. Display the loss surface, equivalent solutions, recovered parameters and held-out outputs. + +Start with exact unit observations, then repeat with a declared partial-observation model. Do not add muscle/EMG generation merely to make the figure look more complete. + +### Deliverable C — an unknown-input stress test + +Release a small number of neural-drive nuisance coefficients, keeping the protocol and measurement assumptions explicit. Recompute profiles and recovery. This reveals whether the apparent PIC estimate survived only because drive was treated as known. + +Choose subsequent work from these outcomes: fix derivatives, improve the protocol, reduce/reparameterize the inferred physiology, or expand the observation model. Do not default to adding more free parameters. + +## 13. Suggested implementation organization + +These are proposed additions or responsibilities, not claims that files currently exist: + + experiments/twin/ + poolbuild.py # Reuse the existing forward construction. + protocols.py # Reuse and extend protocol definitions. + observables.py # Shared extraction and eligibility rules. + specifications.py # Versioned experiment/model contracts. + observation_model.py # Partial observation, errors, censoring. + objectives.py # Declared likelihoods/composite losses. + sensitivity_checks.py # Continuous checks and hard-step tests. + inference.py # Baselines and qualified gradient methods. + identifiability.py # Profiles, local sensitivities, ambiguity. + protocol_design.py # Candidate protocol comparison. + reporting.py # Raw-result-linked reports and failures. + configs/ + results/ + tests/ + +Keep existing `phase1_pic_validation.py`, `phase2_benchmark.py` and `phase3_inverse_mvp.py` entry points if they are present. Extend or wrap them rather than breaking existing commands. The experiment manifest should contain the scientific phase and purpose, so old filenames do not determine the research order. + +Every result should identify its model revision, parameter scope, data split, observation assumptions, gradient mode and gate status. Never replace a failed run with an unmarked successful rerun using different tolerances or seeds. + +## 14. Claim boundaries + +**After forward verification:** “The Jaxley implementation reproduces the reference model in the tested configurations.” + +**After conditional synthetic recovery:** “These parameters/combinations can be recovered under the stated known-input and observation assumptions.” + +**After observation/mismatch tests:** “Recovery or ambiguity estimates remain reliable under the tested missingness, noise and model deviations.” + +**After human prediction:** “The calibrated subject model predicts specified held-out outputs with quantified uncertainty.” + +**Only with further independent evidence:** stronger claims about actual hidden human physiology or causal mechanism attribution. + +The long-term digital-twin goal is unchanged. The revised plan makes each step earn the next claim without requiring that every physiological parameter become uniquely identifiable. + +## Sources and methodological references + +**[S1] Supplied roadmap.** “Digital Twin of the Motor Neuron Pool — Roadmap”, updated 13 October 2026; uploaded as `Pasted text(20261013-061211).txt`. Line references use the line labels shown with the supplied file in the conversation. Repository results described there were not rerun for this review. + +**[R1] Deistler et al.** Jaxley: differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics. *Nature Methods* 22, 2649–2657 (2025). DOI: `10.1038/s41592-025-02895-w`. Supports the differentiable-simulation approach; does not validate this particular harness or its long-horizon gradients. + +**[R2] Raue et al.** Structural and practical identifiability analysis of partially observed dynamical models by exploiting the profile likelihood. *Bioinformatics* 25, 1923–1929 (2009). DOI: `10.1093/bioinformatics/btp358`. + +**[R3] JAX official documentation.** “Benchmarking JAX code.” Documentation path: `docs.jax.dev/en/latest/benchmarking.html`. Consulted 13 October 2026. + +**[R4] Stephenson and Maluf.** Dependence of the paired motor unit analysis on motor unit discharge characteristics in the human tibialis anterior muscle. *Journal of Neuroscience Methods* 198, 84–92 (2011). DOI: `10.1016/j.jneumeth.2011.03.018`. + +**[R5] JAX official documentation.** `jax.random.poisson`. Documentation path: `docs.jax.dev/en/latest/_autosummary/jax.random.poisson.html`. Consulted 13 October 2026. The discrete-sampling example in this plan is explanatory reasoning, not evidence that the harness uses this function. + +**[R6] JAX official documentation.** “Gradient checkpointing with jax.checkpoint (jax.remat).” Documentation path: `docs.jax.dev/en/latest/gradient-checkpointing.html`. Consulted 13 October 2026. + +**[R7] Massaroli et al.** Differentiable Multiple Shooting Layers (2021). arXiv: `2106.03885`. Multiple shooting is a proposed method to evaluate here, not an established solution for this model. + +All experimental designs, proposed gates, implementation responsibilities and prioritization decisions in this document are recommendations from this review, not results established by these references. diff --git a/docs/M1_bessel_feasibility.md b/docs/M1_bessel_feasibility.md new file mode 100644 index 00000000..385e7811 --- /dev/null +++ b/docs/M1_bessel_feasibility.md @@ -0,0 +1,52 @@ +# M1 — Differentiable Bessel Feasibility: GO + +**Decision: GO.** Differentiating the EMG volume-conductor model w.r.t. its +generating parameters (tissue conductivities, conduction velocity, geometry) is +feasible. The full port (M7) is therefore unblocked. + +## Question + +The surface-EMG volume-conductor model (`core/emg/surface/simulate_fiber.py`) +evaluates modified Bessel functions `iv`/`kv` at integer orders up to ~16. JAX's +`jax.scipy.special` provides only `i0/i1/i0e/i1e` — **no `k0/k1`, no arbitrary-order +`iv/kv/jv`**. Without differentiable Bessel functions, EMG cannot be made +differentiable w.r.t. `sig_*`/CV/geometry (Tier B / M7). This spike answered: +can we implement `iv`/`kv` natively in JAX with accurate values **and gradients**? + +## Result + +Implemented in `myogen/simulator/jaxley/bessel.py`: +- `k0`, `k1` via Abramowitz & Stegun 9.8.5–9.8.8 polynomial approximations (seeded + by JAX `i0`/`i1`). +- `kv_int(n, x)` via stable **upward** recurrence `K_{m+1}=K_{m-1}+(2m/x)K_m`. +- `iv_int(n, x)` via Miller **downward** recurrence (stable direction for I_n) + normalised against JAX `i0`, blended with an ascending series `Σ (x/2)^(2k+n)/(k!(n+k)!)` + below `x=1` where the recurrence's `(2k/x)` factors would degrade the derivative. + +**Key design choice:** the forward pass uses only smooth ops (JAX `i0`/`i1`, +polynomials, logs, recurrences), so `jax.grad` yields the exact derivative +automatically — no `custom_jvp` needed. + +## Validation + +`scripts/prototype_bessel_jax.py` checks values and gradients against scipy and the +analytic identities `I_n'=(I_{n-1}+I_{n+1})/2`, `K_n'=-(K_{n-1}+K_{n+1})/2`, over +orders {0,1,2,3,5,8,12,16} and x ∈ {0.05, 0.3, 1, 2.5, 6, 15, 40}: + +- **Worst value relative error: 1.67e-7** +- **Worst gradient relative error: 1.67e-7** + +Both sit at the accuracy floor of JAX's `i0`/`i1` polynomial approximation — more +than sufficient for the float32 EMG pipeline. + +## Implications for M7 + +- `iv`/`kv` are solved. The remaining pieces of the volume-conductor port are + mechanical JAX swaps: `np.linalg.solve` → `jnp.linalg.solve`, `np.fft` → `jnp.fft`, + and freezing the `argmax`/`roll` peak-centring as a constant shift. +- **Follow-up:** the ordinary Bessel `jv` (used once for electrode size, + `simulate_fiber.py:801`) still needs an analogous native implementation + (downward recurrence seeded by `j0`/`j1`; JAX lacks these too). Same technique; + small additional effort. +- For iEMG, hard step masks in `bioelectric.py` must become temperature-controlled + sigmoids to get gradients w.r.t. CV / fiber length (as noted in the plan). diff --git a/docs/NEURON_DECOUPLING.md b/docs/NEURON_DECOUPLING.md new file mode 100644 index 00000000..431bde45 --- /dev/null +++ b/docs/NEURON_DECOUPLING.md @@ -0,0 +1,120 @@ +# NEURON Decoupling — Jaxley as Default Backend + +Goal: make Jaxley the default/only *required* simulation backend and keep NEURON as +an *optional* extra (option A — retain NEURON as a validation reference, do not +delete it). Phased so each step is independently shippable. + +## Phase 1 — Import decoupling (DONE) + +`import myogen` and `import myogen.simulator` no longer require the NEURON runtime. +Verified by blocking `import neuron` (simulating a Jaxley-only install): `import +myogen`, `import myogen.simulator`, and a differentiable `run_jax` simulation all +succeed; the public classes resolve to their Jaxley implementations. + +Changes: +- `myogen/simulator/__init__.py` — the eager `from myogen.simulator.neuron.* import` + lines (the only thing forcing NEURON at import) now import the Jaxley equivalents + (`HillModel`, `JointDynamics`, `Network`, `SpindleModel`, `GolgiTendonOrganModel`, + `SimulationRunner`). Added a lazy `__getattr__` so `simulator.neuron` still works + on demand when NEURON is installed. +- `myogen/utils/continuous_saver.py` — moved `from neuron import h` from module level + into the two methods that use the NEURON clock (it's a NEURON step-callback util). +- `myogen/utils/nmodl.py` — `load_nmodl_mechanisms` now stays silent (instead of + warning) when NEURON is simply absent and `quiet=True`; still raises under `strict`. + +Backward compatibility (NEURON installed): fully preserved. All NEURON usage goes +through the explicit namespace — `from myogen.simulator.neuron. import ...` and +`from myogen.simulator.neuron import Network` (lazy `__getattr__`), plus +`simulator.neuron.*`. No real code depended on top-level `simulator.Network`/ +`HillModel`/etc. being the NEURON version (only docstring references existed). The +full differentiable-pipeline + regression test suite remains green. + +Already-favourable preconditions found (no change needed): +- `neuron/__init__.py` already lazy-loads via `__getattr__`. +- `_setup_myogen` is defined but not auto-called at import. +- The `core/` muscle/force/EMG pipeline is simulator-agnostic; the one cross-import + (`simulate_fiber.py` → `neuron/_cython/_simulate_fiber`) is inside a try/except with + a numpy fallback, so it degrades gracefully. + +## Phase 2 — Packaging (DONE, with one CI-only verification pending) + +Jaxley-only installs no longer pull in NEURON/MPI, and the build no longer hard- +requires NEURON. + +Changes: +- `pyproject.toml` — moved `neuron`, `mpi4py`, `impi-rt` out of core `dependencies` + into a new optional extra: `pip install "myogen[neuron]"`. (`mpi4py`/`impi-rt` are + MPI companions used only by NEURON — confirmed no Python source imports `mpi4py`.) +- `setup.py` — `BuildWithNMODL.compile_nmodl` previously **raised** a RuntimeError on + Windows when NEURON was absent, hard-requiring it at build time. Now it warns and + continues on every platform (NMODL is only needed for the optional NEURON backend); + users install `myogen[neuron]` then run the `setup_myogen` task to compile mechanisms. + +Verified: `pyproject.toml` parses and core deps contain no NEURON/MPI; `setup.py` +parses and single-threaded `cythonize` of the extensions works; `import myogen` + a +Jaxley `run_jax` simulation runs with `neuron` blocked. + +### Closing the "installs without NEURON" verification + +The definitive test is that the built wheel installs and runs on a machine with **no +NEURON, MPI, or compiler**. This is now covered three ways: + +1. **CI job (definitive).** `.github/workflows/build-wheels.yml` gains a + `test_wheels_jaxley_only` job that downloads the built wheel, installs it with **no** + NEURON/MPI (and asserts `import neuron` fails), then imports `myogen`, the + differentiable API, and builds real NERLab motor neurons on the Jaxley backend. + `publish` now depends on it, so a release cannot ship unless the Jaxley-only install + passes. (The pre-existing `test_wheels` job still covers the NEURON-installed path.) +2. **Build robustness.** `setup.py` reads `MYOGEN_CYTHONIZE_NTHREADS` (default 4) so + process-pool-constrained environments can build fully serially with + `MYOGEN_CYTHONIZE_NTHREADS=0` (+ `build_ext -j1`), avoiding the pool entirely. +3. **Local proof.** With that fallback, a full fresh Cython build of all six extensions + was completed in the dev sandbox, and the freshly-built artifacts import and run a + Jaxley simulation with `neuron` blocked. (GitHub runners handle the default + parallel build; the serial fallback is only for constrained sandboxes.) + +### Descoped from Phase 2: Cython relocation + +Moving the misfiled `neuron/_cython/` kernels to `core/` is **not required** under +option A: the `neuron/` *subpackage* still ships and is importable without the NEURON +*runtime*, so the kernels (which are pure Cython, no NEURON) compile and load fine +where they are, and the one `core/` consumer (`simulate_fiber.py`) already has a numpy +fallback. Relocation only becomes necessary if the `neuron/` subpackage is ever +deleted outright — deferred with that deletion. + +## Phase 3 — Examples (DONE, including the canonical rename) + +Classified all 29 examples by backend and documented the split in `examples/README.md`: +**17 default (Jaxley / backend-agnostic) examples** run on a plain `pip install myogen`; +**12 NEURON-only examples** require `pip install "myogen[neuron]"`. + +- **Verified**: a representative set of default examples (01, 04–07, 09, 12, plus the + heavy `08_..._jaxley` and `11_..._jaxley`) resolve **all** their top-level imports with + `neuron` blocked — they get past imports with no NEURON. (Examples 04–07/09 are `core/` + computation; they only transitively imported NEURON before Phase 1.) +- **Quarantine documented**: the NEURON-only examples (`02/03/08/10/11` base, + `02_finetune/01–04`, `watanabe/01–03`) use explicit NEURON imports and still run once + the extra is installed. `examples/README.md` maps each to its `_jaxley` counterpart and + states the `myogen[neuron]` requirement. The docs gallery's `reset_neuron` already + no-ops gracefully when NEURON is absent (`docs/source/conf.py`). + +Not yet ported to Jaxley: `02_finetune/01–04` and `watanabe/01–03` (tracked as follow-up; +runnable via the NEURON extra meanwhile). + +**Canonical rename (DONE).** The five paired examples were renamed so the **Jaxley** +version is canonical: `XX_..._jaxley.py` → `XX_....py`, and the NEURON version moved to +`XX_..._neuron.py` (examples 02, 03, 08, 10, 11 in `01_basic`). Verified each canonical +`.py` is Jaxley-backed and resolves imports with `neuron` blocked, and each `_neuron.py` +is NEURON-backed. Stale in-docstring command/cross-references were updated; output +artifact filenames (`*_jaxley.pkl/png/csv`) were intentionally left unchanged because +downstream examples read those exact files. No `.rst`/gallery config hard-codes example +filenames (auto-discovered via `filename_pattern`), so no docs config change was needed; +`examples/README.md` documents the canonical/NEURON pairing. Note: generated gallery URLs +change (e.g. `auto_examples/.../11_simulate_spinal_network.html` now shows the Jaxley +version) — expected with a rename. + +## Not doing (option A) + +Deleting the `neuron/` tree, NMODL files, and MPI dependency (the old plan's M7) is +deferred indefinitely: NEURON stays as an optional validation reference until frozen +reference datasets exist and the example surface is fully ported. diff --git a/docs/api/differentiable.md b/docs/api/differentiable.md new file mode 100644 index 00000000..ffb619f9 --- /dev/null +++ b/docs/api/differentiable.md @@ -0,0 +1,47 @@ +# Differentiable API (Jaxley) + +The Jaxley backend exposes a **differentiable, JIT-compilable** neuromuscular closed +loop. These functions let you take gradients of force, joint kinematics, and EMG with +respect to neural, muscle, and joint parameters (see the +[Differentiability status](../DIFFERENTIABILITY_STATUS.md) page for the design and +validation). Import them from `myogen.simulator.jaxley`. + +## Closed-loop simulation & gradients + +::: myogen.simulator.jaxley.run_jax + +::: myogen.simulator.jaxley.value_and_grad_run + +::: myogen.simulator.jaxley.compile_run + +::: myogen.simulator.jaxley.ClosedLoopConfig + +::: myogen.simulator.jaxley.partition_differentiable + +## Spike modes + +::: myogen.simulator.jaxley.spike_detect + +## Differentiable EMG + +Convolve (surrogate/rate) spike trains with static MUAP templates. Differentiable with +respect to everything upstream of the spikes. + +::: myogen.simulator.jaxley.surface_emg_jax + +::: myogen.simulator.jaxley.intramuscular_emg_jax + +::: myogen.simulator.jaxley.resample_muaps + +## Differentiable muscle parameters + +::: myogen.simulator.jaxley.differentiable_twitch_params + +## Differentiable Bessel functions + +Native-JAX modified Bessel functions of integer order (used by the volume-conductor +work); differentiable via autodiff. + +::: myogen.simulator.jaxley.iv_int + +::: myogen.simulator.jaxley.kv_int diff --git a/docs/api/simulator.md b/docs/api/simulator.md index 4ed8d539..6451a974 100644 --- a/docs/api/simulator.md +++ b/docs/api/simulator.md @@ -3,13 +3,30 @@ ## Recruitment ::: myogen.simulator.RecruitmentThresholds -## Neuron populations -::: myogen.simulator.neuron.populations.AlphaMN__Pool -::: myogen.simulator.neuron.populations.DescendingDrive__Pool +## Motor-unit populations + +Jaxley is the default backend. The NEURON backend (optional `neuron` extra) exposes +API-compatible equivalents; both are documented below. + +### Jaxley (default) + +::: myogen.simulator.jaxley.populations.AlphaMN__Pool +::: myogen.simulator.jaxley.populations.DescendingDrive__Pool -`DescendingDrive_Gamma__Pool` is a backward-compatibility alias. New code should -prefer `DescendingDrive__Pool(process_type="gamma", shape=...)`. +`DescendingDrive_Gamma__Pool` is a backward-compatibility alias for a gamma-process +descending drive. +::: myogen.simulator.jaxley.populations.DescendingDrive_Gamma__Pool +::: myogen.simulator.jaxley.populations.AffIa__Pool +::: myogen.simulator.jaxley.populations.AffII__Pool +::: myogen.simulator.jaxley.populations.AffIb__Pool +::: myogen.simulator.jaxley.populations.GII__Pool +::: myogen.simulator.jaxley.populations.GIb__Pool + +### NEURON (optional backend, `myogen[neuron]`) + +::: myogen.simulator.neuron.populations.AlphaMN__Pool +::: myogen.simulator.neuron.populations.DescendingDrive__Pool ::: myogen.simulator.neuron.populations.DescendingDrive_Gamma__Pool ::: myogen.simulator.neuron.populations.AffIa__Pool ::: myogen.simulator.neuron.populations.AffII__Pool diff --git a/docs/session_notes/2026-06-02_jaxley_nerlab_architecture_migration.md b/docs/session_notes/2026-06-02_jaxley_nerlab_architecture_migration.md new file mode 100644 index 00000000..96b8709a --- /dev/null +++ b/docs/session_notes/2026-06-02_jaxley_nerlab_architecture_migration.md @@ -0,0 +1,314 @@ +# Jaxley Migration to the NERLab Motor-Neuron Architecture + +**Date:** 2026-06-02 + +**One-line summary:** Switched the Jaxley motor-neuron pool from the Powers2017 +cable cell to the NERLab single-isopotential-dendrite cell so it now matches +the NEURON production model. All five paired Jaxley examples (02, 03, 08, 10, +11) have been migrated and validated against NEURON. + +--- + +## Why the change happened + +Earlier project memory described the Jaxley motor neuron as "Powers2017 + +cable dendrite (ncomp=4)". The original MyoGen writer confirmed during this +session that **production NEURON does not use the 4-dendrite Powers2017 +cell**. The production NEURON path is: + +- `AlphaMN__Pool` defaults to `model="NERLab"` (set at + [config/alpha_mn_default.yaml:5](../../myogen/config/alpha_mn_default.yaml#L5)). +- `_create_nerlab_cells` builds `cells.AlphaMN(... dendrites__count=1, model="NERLab")` + ([simulator/neuron/populations/motor_neurons.py:517-524](../../myogen/simulator/neuron/populations/motor_neurons.py#L517-L524)). +- Soma uses the `napp` mechanism (Na fast + Na persistent + K fast + K slow + + leak); dendrite uses `caL` (L-type Ca with no inactivation + leak). +- Voltage convention: **original 1952 HH frame**, not modern absolute + voltages. V_rest ≈ 0 mV, ENa = +120 mV, EK = -10 mV, spike peaks ≈ +90 mV. + +Every production NEURON example (`02_*`, `03_*`, `08_*`, `10_*`, `11_*`) uses +this default. The previous Jaxley "Powers2017 cable arch" was therefore +**not comparing the same cell** to NEURON, which explains a large fraction of +the historical mismatch the project tracked as "Jaxley needs further tuning". + +The full re-verification trail is in this session's Bash output and in +`scripts/test_nerlab_*.py`. + +--- + +## Architectural change in one diagram + +``` +BEFORE (Jaxley path, use_jaxley_mech=True): + ┌──── soma ─────┐ Na3rp, Naps, KdrRL, + │ ncomp=1 │ MAHP, Gh, Leak + └───────┬───────┘ + │ Ra (cable filtering) + ┌───────┴───────┐ + │ dend comp 1 │ + │ dend comp 2 │ LCaInact, Gh, Leak + │ dend comp 3 │ (same density on each) + │ dend comp 4 │ + └───────────────┘ + ── Powers2017-style channels, 1952 HH absolute V frame? No, modern. + ── Voltage frame: modern (V_rest ≈ -67 mV) + + +AFTER (Jaxley path, default model="NERLab"): + ┌──── soma ─────┐ napp + │ ncomp=1 │ (Na + Na_persistent + Kfast + Kslow + leak, + │ L = diam │ 5-gate Hodgkin-Huxley) + └───────┬───────┘ + │ + ┌───────┴───────┐ caL + │ isopotential │ (L-type Ca, no inactivation + leak) + │ ncomp=1 │ + └───────────────┘ + ── Channel ports of myogen/simulator/nmodl_files/{napp,caL}.mod + ── Voltage frame: original 1952 HH (V_rest ≈ 0 mV, peaks ≈ +90 mV) +``` + +This is now identical to the NEURON production cell's morphology and channel +set. + +--- + +## What landed in the codebase + +### New channel ports + +[`myogen/simulator/jaxley/channels/nerlab_channels.py`](../../myogen/simulator/jaxley/channels/nerlab_channels.py) +— port of the `napp` and `caL` NMODL files. ~250 lines. + +- `napp.Channel`: 5 gates (m, h, p, n, r). `vtrap`-style rate calc with a + JAX-safe L'Hôpital expansion near the singularity. `init_state` evaluates + steady-state gates at the holding V so the cell starts in equilibrium. +- `caL.Channel`: 1 gate (L), sigmoid Boltzmann activation with a fixed + `Ltau`. Includes the bundled leak from `caL.mod`. + +Exports added to +[`channels/__init__.py`](../../myogen/simulator/jaxley/channels/__init__.py). + +### Pool and cell-builder dispatch + +[`populations/motor_neurons.py`](../../myogen/simulator/jaxley/populations/motor_neurons.py): + +- New `NERLAB_PARAMS` dict (lines 102-160) holding all per-cell parameter + ranges as `(small_cell_value, large_cell_value)` tuples derived from + [config/alpha_mn_default.yaml](../../myogen/config/alpha_mn_default.yaml). +- `_generate_parameter_arrays` dispatches on `self.model` to either + `_generate_nerlab_params` or the legacy `_generate_powers2017_params`. +- `__init__` auto-sets `initial_voltage__mV = 0.0` and + `spike_threshold__mV = 50.0` when `model="NERLab"` and the caller didn't + explicitly override. + +[`cells_api.py`](../../myogen/simulator/jaxley/cells_api.py): + +- `_create_sections`, `_define_geometry`, and `_define_biophysics` dispatch + on `self.model`. NERLab path builds `soma(ncomp=1) + dend(ncomp=1)` with + parents `[-1, 0]`. +- New `_define_geometry_nerlab` and `_insert_nerlab_channels` methods. +- The Powers2017 path is preserved as a sibling for backward compatibility. +- `AlphaMN` constructor now accepts `gamma` and threads it through to + `caL_gama` (was previously hard-coded to 1.0, 5× the NEURON value of 0.2). + +### Pool V-init fix + +[`populations/base.py:117-167`](../../myogen/simulator/jaxley/populations/base.py#L117-L167) +now actually applies `initial_voltage_values__mV` to each underlying `jx.Cell` +via `cell.set("v", ...)` + `init_states()`. Before this fix the field was +stored but never used, so direct `jx.integrate(cell_wrapper.cell, ...)` paths +silently ran at Jaxley's default -70 mV regardless of what the pool was told. + +### Scan-step per-cell-type parameters (for mixed-frame networks) + +[`jax_models.py`](../../myogen/simulator/jaxley/jax_models.py) now takes two +new kwargs on `make_scan_step`: + +- `e_exc_mn: float | None` — separate excitatory AMPA reversal for the + motor-neuron section of the combined network. Defaults to `e_exc` for + backward compat. ex11 passes 70 mV (NERLab) while keeping `e_exc=0` for + the modern-frame interneurons. +- `mn_spike_threshold_mV: float` — separate spike-detection threshold for + motor neurons. Defaults to 0 mV (modern frame); ex11 passes 50 mV. + +The internal scan body now uses `_e_exc_mn` to compute `df_mn` and `_mn_v_th` +for `mn_spikes_now`, so the existing modern-frame interneuron logic is +unaffected. + +--- + +## Voltage convention — important detail + +NERLab cells live in the **original 1952 Hodgkin-Huxley convention** with all +voltages shifted by ~+65 mV relative to modern absolute values: + +| Quantity | NERLab (1952 HH) | Modern | +|---|---:|---:| +| V_rest | 0 mV | −65 to −70 mV | +| ENa | +120 mV | +50 to +55 mV | +| EK | −10 mV | −75 to −80 mV | +| AP peak | +80 to +110 mV | +20 to +40 mV | +| Spike-detect threshold | +50 mV | 0 mV | + +When you reinitialise a NERLab cell's voltage manually (in the simulation +loop, not via the pool), use `cell.set("v", 0.0)` not `-65.0`. When you +detect spikes, look for `V > 50 mV` crossings. + +Mixing NERLab cells with modern-frame interneurons in one `jx.Network` works +**only if** the synapse reversal and detection thresholds are set per +postsynaptic / presynaptic cell type — see the ex11 notes below. + +--- + +## Per-example results + +All five paired examples were migrated and tuned. Quantitative parity to +NEURON now: + +| Example | What it tests | Jaxley NERLab vs NEURON | +|---|---|---| +| **02** — current injection | 15 nA step → spike trains | 94/100 active, 6078 vs 6061 spikes (**0.3 %**) | +| **03** — trapezoid DD | descending drive → spike trains | 85/100 plateau-active, 17.6 Hz mean vs ~18 Hz (**within 10 %**) | +| **08** — force from spikes | 15 nA trapezoid → Fuglevand force | 50/50 active, 14490 vs 14517 spikes (**0.19 %**) | +| **10** — parameter extraction | Vhold / Rin / τ / Ir / AP / AHP / F-I per cell | Rin and Ir within **3-5 %** per cell; AP, AHP, F-I within **20 %**. All five confirmable params inside the literature physiological range (Bui et al., Carp et al., etc.) | +| **11** — full spinal network | mixed populations, joint dynamics, force | 82/100 MN active, 0.56 F0 peak force, physiologically plausible recruitment | + +`ex10` is the most precise — five-protocol per-cell extraction, apples-to- +apples against the NEURON CSV. This is the validation table that had been +missing. + +--- + +## ex03 and ex11 — the IIR / Exp2Syn equivalence story + +`ex03` and `ex11` both use an IIR + fixed-driving-force approximation in +place of NEURON's `Exp2Syn`. Naively matching NEURON's NetCon weights (e.g. +0.05 µS) produces 5-10× too much current because: + +- NEURON's `Exp2Syn` has a shaped onset (`tau1 ≈ 0.5 ms` rise, `tau2 ≈ + 2-5 ms` decay). Peak conductance per spike is only `weight × peak_factor` + with `peak_factor ≈ 0.1-0.2`. +- The IIR adds the full weight per spike without shaping; effective per- + spike charge is therefore ~5-10× higher. + +The empirically-tuned equivalence factors: + +- `ex03` DD → MN: NEURON `weight = 0.15 µS` → IIR `weight = 0.009 µS` + (~0.06×). +- `ex11` DD → MN and Ia → MN: NEURON `weight = 0.05 µS` → IIR + `weight = 0.01 µS` (~0.20×). + +The two factors differ because the steady-state input rate and number of +connections per MN differ; the per-spike scaling math is the same. Both +weights are documented inline in their respective examples. + +A proper fix would be to use a real Jaxley `IonotropicSynapse` (or build an +Exp2Syn-equivalent mechanism) for the excitatory paths too. ex11 already +uses `IonotropicSynapse` for inhibition. + +--- + +## ex11 specifics — mixed-frame spinal network + +ex11 wires NERLab MNs together with modern-frame `gII / gIb` interneurons +into one `jx.Network`. This required three additions: + +1. `e_exc_mn = 70.0` and `e_exc = 0.0` set separately so DD/Ia synaptic + drive computes its driving force from the appropriate reversal for each + target. +2. `mn_spike_threshold_mV = 50.0` for MN spike detection; the interneuron + sections still use the 0 mV crossing. +3. `IonotropicSynapse_e_syn = -15.0` for the gII / gIb → MN inhibitory + reversal. NEURON ex11 uses `-75 mV` literally on NERLab targets, which + gives a 75 mV hyperpolarising driving force and over-inhibits in our + live-V Jaxley synapse model. The frame-shifted `-15 mV` matches the + *physical* inhibitory driving force the original Powers2017 setup + delivered. + +The interneurons themselves still use Na3rp-style channels in the modern V +frame. Porting them to NERLab is left as future work (see "Remaining +caveats" below). + +--- + +## Bugs caught and fixed along the way + +Codex catches that materially changed the outcome — recording these because +each one would have produced silently-wrong recruitment: + +1. **NERLab parameter interpolation was direction-flipped** on six entries. + NEURON's [`special_interp`](../../myogen/simulator/neuron/populations/motor_neurons.py#L390) + applies a `negative` flag from the YAML which swaps the (small, large) + assignment. The first pass of `NERLAB_PARAMS` encoded the YAML literally + — `gnapbar`, `gkfbar`, `gksbar`, `gcaLbar`, `vtraub_caL`, `ltau_caL` all + had small and large cells swapped. Fix: re-encode in (small_cell, + large_cell) semantics with explicit comments. +2. **`gamma` was ignored** in `_insert_nerlab_channels` — `caL_gama` was + hard-coded to 1.0 instead of the NEURON YAML default of 0.2 (5× too + strong dendritic Ca). Wired the pool's `gamma` argument through to the + channel via the new `AlphaMN.gamma` field. +3. **`_apply_initial_voltages` was effectively a no-op.** `_Pool.__init__` + stored `initial_voltage_values__mV` but never called `.set("v", ...)` + on cells. So any code reaching in via `cell_wrapper.cell` ran at + Jaxley's default -70 mV regardless of the pool's intended frame. Fix: + actually apply the voltage in `_Pool.__init__`. +4. **IIR / Exp2Syn equivalence factor** — see above, separate section. +5. **ex11 used a single `e_exc = 0` for all populations** — collapsed the + excitatory driving force to zero on NERLab MNs (`0 - 0 = 0`). Fix: per- + cell-type reversals via the new `e_exc_mn` parameter in `jax_models.py`. + +The original codex review reports from these catches are preserved in the +commit history if a future session wants to trace the reasoning. + +--- + +## Remaining caveats + +1. **F-I gain in ex10 is 10-20 % shallower** in Jaxley vs NEURON across + cells. Could be protocol/extraction differences (settling window, step + length) or genuine model-dynamics differences. Worth a follow-up. +2. **The IIR / Exp2Syn equivalence factors (0.06 and 0.20)** are empirical + per example. A proper Jaxley Exp2Syn-equivalent synapse model would + eliminate this knob. +3. **ex11 is mixed-frame** by design. Works only because all synapses in + that example target NERLab MNs and the interneurons are sources. Any new + synapse type targeting an interneuron would need its own per-cell-type + `e_exc_X`. Cleaner long-term fix: port `gII / gIb` populations to NERLab. +4. **AHP duration in ex10 is broader** than the literature lower bound when + estimated via half-decay. The script measures full duration (`V → vhold + − 0.15 mV`); a proper half-decay measurement would be a useful addition. + +--- + +## What "use this from now on" looks like + +```python +# Build a NERLab motor-neuron pool — the default now matches NEURON production. +from myogen.simulator.jaxley.populations import AlphaMN__Pool +pool = AlphaMN__Pool(recruitment_thresholds__array=rt) # model="NERLab" is default + +# The pool sets V_rest = 0 mV and spike_threshold = +50 mV automatically. +# If you reach in via cell_wrapper.cell and manually re-init, use: +# cell.set("v", 0.0) +# cell.init_states() + +# To opt into the legacy Powers2017 cable cell (e.g. for the +# pre-NERLab spinal-network setup): +pool = AlphaMN__Pool( + recruitment_thresholds__array=rt, + model="Powers2017", + use_jaxley_mech=True, +) +``` + +--- + +## Pointers + +- Updated `MEMORY.md` summary file in the user's auto-memory directory. +- Validation scripts: `scripts/test_nerlab_single_cell.py`, + `scripts/test_nerlab_pool.py`. +- This document supersedes the cable-architecture sections of + `2026-02-10_parameter_tuning_final.md` and the relevant parts of + `2026-05-01_jaxley_neuron_replacement_status.md`. The earlier files are + preserved unchanged for historical reference. diff --git a/examples/01_basic/01_simulate_recruitment_thresholds.py b/examples/01_basic/01_simulate_recruitment_thresholds.py index ce8e19e9..51295233 100644 --- a/examples/01_basic/01_simulate_recruitment_thresholds.py +++ b/examples/01_basic/01_simulate_recruitment_thresholds.py @@ -17,6 +17,7 @@ * **Combined** model: Hybrid approach combining De Luca shape with Konstantin scaling (*Ours*) """ +# this is a test comment # %% ############################################################################## diff --git a/examples/01_basic/02_simulate_spike_trains_current_injection.py b/examples/01_basic/02_simulate_spike_trains_current_injection.py index c1a41175..c6230017 100644 --- a/examples/01_basic/02_simulate_spike_trains_current_injection.py +++ b/examples/01_basic/02_simulate_spike_trains_current_injection.py @@ -1,16 +1,32 @@ """ -Spike Train Generation with Current Injection -======================================= - -This example demonstrates how to simulate spike trains in a population of alpha motor neurons using current injection. - -Two complementary workflows are presented: - -1. Manual step-by-step workflow — explicitly walks through each stage of the NEURON simulation pipeline. This workflow is intended to clarify the underlying mechanisms. - -2. Utility-function workflow — uses the high-level `inject_currents_and_simulate_spike_trains` function for routine simulations. - -Both workflows yield identical results; the manual version is provided purely for explanatory purposes. +Spike Train Generation with Current Injection - Jaxley Backend (NERLab) +======================================================================= + +This example demonstrates how to simulate spike trains in a population of alpha motor neurons +using current injection with the **Jaxley** (JAX-based) simulator. + +This is designed to produce output comparable to the NEURON version +(``02_simulate_spike_trains_current_injection.py``), which uses the SAME ``model="NERLab"`` +motor-neuron model. + +Key features: + - Uses the NERLab architecture (soma + 1 isopotential dendrite), matching the + production NEURON model + - NERLab channels: ``napp`` (Na fast + Na persistent + Kfast + Kslow + leak) on the + soma; ``caL`` (L-type Ca, no inactivation + leak) on the dendrite + - Same current to ALL neurons — recruitment from biophysics (Henneman size principle) + - Runs full biophysical simulation with ``jx.integrate()`` + +.. note:: + **Voltage convention.** NERLab cells live in the *original 1952 HH frame*: + V_rest ≈ 0 mV, ENa = +120 mV, EK = -10 mV, spike peaks ≈ +90 mV. This is set + automatically by ``AlphaMN__Pool`` when ``model="NERLab"`` (the default), but if + you reinitialise ``v`` manually anywhere in your pipeline you must use 0 mV, not + -65 mV, and detect spikes at ~+50 mV (not 0 mV). + +.. note:: + Jaxley uses JAX for accelerated computation. On first run, JAX will compile + the simulation which may take a few seconds, but subsequent runs are faster. """ # %% @@ -18,44 +34,23 @@ ############################################################################## # Import Libraries # ---------------- -# -# !!! important -# In **MyoGen** all **random number generation** is handled by the RNG returned from -# [`get_random_generator`][myogen.get_random_generator], a thin wrapper around `numpy.random`. -# -# Always fetch the generator at the call site so the current seed is honored: -# -# ```python -# from myogen import simulator, get_random_generator -# get_random_generator().normal(0, 1) -# ``` -# -# To change the seed, use [`set_random_seed`][myogen.set_random_seed]: -# -# ```python -# from myogen import set_random_seed -# set_random_seed(42) -# ``` from pathlib import Path +import jax.numpy as jnp +import jaxley as jx import joblib import numpy as np import quantities as pq import seaborn as sns from matplotlib import pyplot as plt from neo import Block, Segment, SpikeTrain -from neuron import h from scipy.ndimage import gaussian_filter1d +from tqdm import tqdm -from myogen import get_random_generator -from myogen.simulator.neuron.populations import AlphaMN__Pool +from myogen import get_random_generator, simulator +from myogen.simulator.jaxley.populations import AlphaMN__Pool from myogen.utils.currents import create_trapezoid_current -from myogen.utils.neuron.inject_currents_into_populations import ( - inject_currents_and_simulate_spike_trains, - inject_currents_into_populations, -) -from myogen.utils.nmodl import load_nmodl_mechanisms plt.style.use("fivethirtyeight") @@ -69,10 +64,9 @@ def rasterplot_rates(spiketrains, filter_function=None): """Native spike raster with top/right marginal axes. Lightweight stand-in for ``viziphant.rasterplot.rasterplot_rates``: draws a - spike raster (one row per train), a right marginal showing each train's mean + spike raster (one row per train), a right marginal with each train's mean firing rate, and an (initially empty) top marginal that the caller fills with - the smoothed population rate. Returns ``(ax, axhistx, axhisty)`` as - absolutely-positioned axes so manual position tweaks downstream still work. + the smoothed population rate. Returns ``(ax, axhistx, axhisty)``. """ if filter_function is not None: spiketrains = [st for st in spiketrains if filter_function(st)] @@ -94,162 +88,195 @@ def rasterplot_rates(spiketrains, filter_function=None): ax.set_ylim(-1, max(len(spiketrains), 1)) return ax, axhistx, axhisty + ############################################################################## # Create Motor Neuron Populations (Pools) # --------------------------------------- # -# In MyoGen a population of cells (e.g. motor neurons) is represented by a -# **Population** class and available in the `myogen.simulator.neuron.populations` module. -# -# A population can easily be created by specifying the number of cells. Plausible default parameters are already set. -# -# For a motor neuron population (referred to as **motor pool**), we can use the [`AlphaMN__Pool`][myogen.simulator.neuron.populations.AlphaMN__Pool] class. -# This class can also use the recruitment thresholds generated in the previous example to distribute the motor units properties in a physiologically plausible manner. -# -# !!! important -# These **Population** classes are custom-built and use therefore custom NMODL mechanisms. -# To use them, the NMODL mechanisms need to be loaded first using the [`load_nmodl_mechanisms`][myogen.load_nmodl_mechanisms] function. +# We create motor neuron pools using the :class:`~myogen.simulator.jaxley.populations.AlphaMN__Pool` +# class. The default ``model="NERLab"`` builds the same architecture as the production +# NEURON model: soma + 1 isopotential dendrite, ``napp`` + ``caL`` channels. # -# To showcase MyoGen's capabilities, we will create two different motor neuron pools with identical properties but different input currents. -load_nmodl_mechanisms() +# Recruitment emerges naturally from biophysics (Henneman size principle): +# - Low threshold MNs: smaller soma, higher input resistance -> easier to activate +# - High threshold MNs: larger soma, lower input resistance -> harder to activate save_path = Path("./results") save_path.mkdir(exist_ok=True) -recruitment_thresholds = joblib.load(save_path / "thresholds.pkl") +# Load recruitment thresholds from Example 1 (or generate defaults) +try: + recruitment_thresholds = joblib.load(save_path / "thresholds.pkl") + print(f"Loaded recruitment thresholds: {len(recruitment_thresholds)} values") +except FileNotFoundError: + print("Recruitment thresholds not found. Generating defaults matching ex01...") + n_motor_units = 100 + recruitment_thresholds, _ = simulator.RecruitmentThresholds( + N=n_motor_units, + recruitment_range__ratio=100, + deluca__slope=5, + konstantin__max_threshold__ratio=1.0, + mode="combined", + ) + joblib.dump(recruitment_thresholds, save_path / "thresholds.pkl") n_pools = 2 + +# Create motor neuron pools — defaults to model="NERLab" (matches the production +# NEURON model: soma + 1 isopotential dendrite, napp + caL channels). motor_neuron_pools = [ - AlphaMN__Pool(recruitment_thresholds__array=recruitment_thresholds) for _ in range(n_pools) + AlphaMN__Pool( + recruitment_thresholds__array=recruitment_thresholds, + mode="active", + # model defaults to "NERLab"; use_jaxley_mech is ignored for NERLab. + ) + for _ in range(n_pools) ] +print(f"Created {n_pools} motor neuron pools with {len(recruitment_thresholds)} neurons each") +print(f" Model: {motor_neuron_pools[0].model} (soma + 1 isopotential dendrite, napp + caL)") + ############################################################################## # Create Input Currents # --------------------- # -# To drive the motor units, we use a **common input current profile**. -# -# In this example, we use a **trapezoid-shaped input current** which is generated using the [`create_trapezoid_current`][myogen.utils.currents.create_trapezoid_current] function. -# -# !!! note -# More convenient functions for generating input current profiles are available in the `myogen.utils.currents` module. -# -# !!! note -# The generated input current is an instance of the `neo.core.AnalogSignal` class from the `neo` package. +# Same current to ALL neurons, matching the NEURON approach. +# Recruitment determined by biophysics (cell size, conductances) — no manual cutoff. -timestep = 0.05 * pq.ms +timestep = 0.025 * pq.ms # 0.025 ms = 25 µs, typical for HH simulations simulation_time = 4000 * pq.ms +# Current amplitude — same to ALL neurons; biophysics determines recruitment. +# 15 nA matches the NEURON example exactly (see +# 02_simulate_spike_trains_current_injection.py:118). Using the same drive +# level lets the comparison block at the bottom of this script be +# apples-to-apples (active count, total spikes, mean rate). +current_amplitude = 15.0 * pq.nA + +# NO max_recruitment cutoff - simulate ALL neurons like NEURON does +# Recruitment emerges from biophysics (cell size, conductances) + rise_time_ms = list(get_random_generator().uniform(100, 500, size=n_pools)) * pq.ms plateau_time_ms = list(get_random_generator().uniform(1000, 2000, size=n_pools)) * pq.ms fall_time_ms = list(get_random_generator().uniform(1000, 2000, size=n_pools)) * pq.ms input_current__AnalogSignal = create_trapezoid_current( n_pools, - int(simulation_time / timestep), + int(simulation_time / timestep) + 1, timestep, - amplitudes__nA=[15.0 * pq.nA] * n_pools, + amplitudes__nA=[current_amplitude] * n_pools, rise_times__ms=rise_time_ms, plateau_times__ms=plateau_time_ms, fall_times__ms=fall_time_ms, delays__ms=500.0 * pq.ms, ) -print( - f"Input current signal shape: {input_current__AnalogSignal.shape}\nClass: {input_current__AnalogSignal.__class__}" -) +print(f"\nInput current signal shape: {input_current__AnalogSignal.shape}") +print(f"Timestep: {timestep}, Total time: {simulation_time}") +print(f"Current amplitude: {current_amplitude} (same for ALL neurons - NEURON approach)") +print("Recruitment determined by biophysics (cell size, conductances) - no manual cutoff") -# Save input current signal for later analysis -joblib.dump(input_current__AnalogSignal, save_path / "input_current__AnalogSignal.pkl") +# Save input current signal +joblib.dump(input_current__AnalogSignal, save_path / "input_current__AnalogSignal_jaxley.pkl") ############################################################################## -# Manual Simulation Approach - Step by Step -# ------------------------------------------- +# Manual Simulation Approach - Step by Step (Jaxley Biophysics) +# ------------------------------------------------------------- +# +# We walk through each stage of the Jaxley biophysical simulation: # -# Before showing the convenient utility function, let's understand what happens -# under the hood by implementing the simulation pipeline manually. -# This approach gives you full control and helps understand NEURON's mechanisms. +# 1. Extract simulation parameters +# 2. For each neuron: inject current, run jx.integrate(), detect spikes +# 3. Convert to Neo format +# +# The key function is ``jx.integrate()`` which solves the HH differential +# equations to compute membrane voltage over time. -# Step 1: Set up current injection manually -# ========================================= -# We need to inject time-varying currents into each motor neuron. -# This uses NEURON's `neuron.h.IClamp` (current clamp) mechanism with `neuron.h.Vector.play`. +# Step 1: Extract simulation parameters +dt_ms = float(timestep.rescale(pq.ms).magnitude) +t_max_ms = float(simulation_time.rescale(pq.ms).magnitude) -inject_currents_into_populations(motor_neuron_pools, input_current__AnalogSignal) +# Spike detection threshold — NERLab cells live in the 1952-HH frame +# (V_rest ≈ 0 mV, AP peaks ≈ +90 mV), so a positive threshold avoids +# false-positives from the resting-state membrane drift. +spike_detection_threshold__mV = 50.0 -# Step 2: Set up spike recording manually -# ======================================= -# For each neuron, we create a `neuron.h.NetCon` (network connection) object that detects -# spikes when the membrane voltage crosses a threshold, and records spike times. +# Convert Neo AnalogSignal to numpy array +current_data = np.array(input_current__AnalogSignal.magnitude) +print(f"\nCurrent data shape: {current_data.shape}") -spike_detection_threshold__mV = 50.0 * pq.mV -simulation_time__ms = input_current__AnalogSignal.t_stop.rescale(pq.ms) +# Step 2: Simulate each pool and collect spike trains +spike_train__Block_manual = Block(name="Manual Jaxley NERLab Biophysical Simulation") -spike_recorders = [] +print("\n" + "=" * 60) +print("Running biophysical simulations with jx.integrate()") +print("Using NERLab channels (napp + caL; 1 soma + 1 isopotential dendrite)") +print("=" * 60) for pool_idx, pool in enumerate(motor_neuron_pools): - pool_spike_recorders = [] + print(f"\nSimulating Pool {pool_idx}...") + + # Get current waveform for this pool + pool_current = jnp.array(current_data[:, pool_idx]) - for cell in pool: - # Create a vector to record spike times - spike_recorder = h.Vector() + # Create a segment for this pool's spike trains + segment = Segment(name=f"Pool {pool_idx}") + segment.spiketrains = [] - # Create NetCon object: monitors voltage at soma(0.5) and records spikes - # NetCon(source, target, threshold, delay, weight) - # source: cell.soma(0.5)._ref_v (membrane voltage reference) - # target: None (no post-synaptic target, just recording) - nc = h.NetCon(cell.soma(0.5)._ref_v, None, sec=cell.soma) - nc.threshold = spike_detection_threshold__mV # Spike detection threshold - nc.record(spike_recorder) # Record spike times into vector + # Simulate each neuron in the pool + # NEURON approach: simulate ALL neurons with SAME current + # Recruitment emerges from biophysics (cell size, conductances) + for neuron_idx, cell_wrapper in enumerate(tqdm(pool, desc=f" Pool {pool_idx}", leave=False)): + spike_times_ms = np.array([]) - pool_spike_recorders.append(spike_recorder) + if hasattr(cell_wrapper, 'cell'): + cell = cell_wrapper.cell - spike_recorders.append(pool_spike_recorders) + try: + # === BIOPHYSICAL SIMULATION WITH MAHP FOR AHP === + cell.delete_recordings() + cell.delete_stimuli() -# Step 3: Initialize voltages and run simulation -# ============================================== -# Before running, we need to initialize membrane voltages to physiological values. -# -# !!! note -# For this MyoGen populations provide the `get_initialization_data()` method. -# This returns the sections and their initial voltages. - -# Initialize each neuron's membrane voltage to its resting potential -for pool in motor_neuron_pools: - for section, voltage in zip(*pool.get_initialization_data()): - section.v = voltage - -# Initialize NEURON's internal state and run the simulation -h.finitialize() # Initialize all mechanisms and variables -# Advance the solver to the end of the simulation (replaces the deprecated -# neuron.run(); finitialize() above already set the initial state). -h.tstop = float(simulation_time__ms) -while h.t < h.tstop: - h.fadvance() - -# Step 4: Convert recorded data to `neo.core.Block` format -# ================================================================== -# The spike times are now stored in NEURON vectors. We convert them to -# the standardized `neo.core.Block` format for analysis and compatibility. - -spike_train__Block_manual = Block(name="Manual Simulation Results") - -for pool_idx, pool_spike_recorders in enumerate(spike_recorders): - # Create a segment for this motor unit pool - segment = Segment(name=f"Pool {pool_idx}") + # Reset V to NERLab resting potential (≈ 0 mV in the 1952-HH frame) + # and reinitialise channel states at that V. + cell.set("v", 0.0) + cell.init_states() - # Convert each neuron's spike times to a `neo.core.SpikeTrain` object - segment.spiketrains = [] - for neuron_idx, spike_recorder in enumerate(pool_spike_recorders): - # Convert NEURON vector to numpy array and add units - spike_times = (spike_recorder.as_numpy() * pq.ms).rescale(pq.s) + # Set up voltage recording on soma + cell.branch(0).loc(0.5).record("v") + + # Inject SAME current to ALL neurons on soma (NEURON approach) + # Recruitment determined by biophysics, not current scaling + cell.branch(0).loc(0.5).stimulate(pool_current) + + # Run biophysical simulation + voltages = jx.integrate(cell, delta_t=dt_ms, t_max=t_max_ms) + + # Extract voltage trace + if voltages.ndim == 2: + v = np.array(voltages[0]) + else: + v = np.array(voltages) + + # Detect spikes (threshold crossings) + spike_indices = np.where( + (v[:-1] < spike_detection_threshold__mV) & + (v[1:] >= spike_detection_threshold__mV) + )[0] + spike_times_ms = spike_indices * dt_ms + + except Exception as e: + if neuron_idx == 0: + print(f" Warning: Neuron {neuron_idx} simulation failed: {e}") + + # Convert to Neo SpikeTrain + spike_times_s = spike_times_ms / 1000.0 - # Create `neo.core.SpikeTrain` object with metadata spiketrain = SpikeTrain( - spike_times, - t_stop=simulation_time__ms.rescale(pq.s), - sampling_rate=(1 / (h.dt * pq.ms)).rescale(pq.Hz), - sampling_period=(h.dt * pq.ms).rescale(pq.s), + spike_times_s * pq.s, + t_stop=simulation_time.rescale(pq.s), + sampling_rate=(1 / timestep).rescale(pq.Hz), + sampling_period=timestep.rescale(pq.s), name=str(neuron_idx), description=f"Pool {pool_idx}, Neuron {neuron_idx}", ) @@ -257,69 +284,41 @@ def rasterplot_rates(spiketrains, filter_function=None): spike_train__Block_manual.segments.append(segment) -joblib.dump(spike_train__Block_manual, save_path / "spike_train__Block_manual.pkl") - -############################################################################## -# Convenient Utility Function Approach -# ------------------------------------ -# -# The manual approach above shows you exactly what happens during simulation. -# However, since this is a common task, MyoGen provides the [`inject_currents_and_simulate_spike_trains`][myogen.utils.neuron.inject_currents_into_populations.inject_currents_and_simulate_spike_trains] -# utility function that encapsulates all these steps in a single call. -# -# This is the recommended approach for routine simulations, while the manual -# approach is useful when you need custom spike detection, specialized recording, -# or want to understand the underlying mechanisms. - -# Run the same simulation using the utility function -spike_train__Block = inject_currents_and_simulate_spike_trains( - populations=motor_neuron_pools, - input_current__AnalogSignal=input_current__AnalogSignal, - spike_detection_thresholds__mV=50 * pq.mV, -) - -joblib.dump(spike_train__Block, save_path / "spike_train__Block_utility.pkl") + # Report statistics for this pool + active_count = sum(1 for st in segment.spiketrains if len(st) > 0) + total_spikes = sum(len(st) for st in segment.spiketrains) + print(f" Pool {pool_idx}: {active_count}/{len(pool)} active neurons, {total_spikes} total spikes") -# Compare the two approaches -print("\nComparison of results:") -print(f"Manual approach: {len(spike_train__Block_manual.segments)} segments") -print(f"Utility approach: {len(spike_train__Block.segments)} segments") - -# Verify they produce similar results (spike counts should be identical) -for i, (manual_seg, utility_seg) in enumerate( - zip(spike_train__Block_manual.segments, spike_train__Block.segments) -): - manual_spikes = sum(len(st) for st in manual_seg.spiketrains) - utility_spikes = sum(len(st) for st in utility_seg.spiketrains) - print(f"Pool {i}: Manual={manual_spikes} spikes, Utility={utility_spikes} spikes") +# Save results — single file; no separate utility implementation exists for Jaxley +spike_train__Block = spike_train__Block_manual +joblib.dump(spike_train__Block, save_path / "spike_train__Block_utility_jaxley.pkl") ############################################################################## # Calculate and Display Statistics # --------------------------------- -# -# It might be of interest to calculate the **firing rates** of the motor units. -# -# !!! note -# The **firing rates** are calculated as the number of spikes divided by the time in which each MU was active. firing_rates = [ np.array( [ mean_firing_rate(st__s.time_slice(st__s.min(), st__s.max())) for st__s in spike_train__segment.spiketrains - if len(st__s) > 1 # Need at least 2 spikes to compute rate over spike range + if len(st__s) > 0 ] ) for spike_train__segment in spike_train__Block.segments ] -print("Firing rate statistics:") +print("\n" + "=" * 60) +print("Firing Rate Statistics") +print("=" * 60) + for pool_idx, firing_rates_per_pool in enumerate(firing_rates): - active_neurons = np.sum(firing_rates_per_pool > 0) - if len(firing_rates_per_pool) > 0 and np.sum(firing_rates_per_pool > 0) > 0: - mean_rate = np.mean(firing_rates_per_pool[firing_rates_per_pool > 0]) - max_rate = np.max(firing_rates_per_pool) + if len(firing_rates_per_pool) > 0: + active_neurons = np.sum(firing_rates_per_pool > 0) + mean_rate = np.mean(firing_rates_per_pool[firing_rates_per_pool > 0]) if active_neurons > 0 else 0.0 + max_rate = np.max(firing_rates_per_pool) if len(firing_rates_per_pool) > 0 else 0.0 else: + active_neurons = 0 mean_rate = 0.0 max_rate = 0.0 @@ -331,28 +330,31 @@ def rasterplot_rates(spiketrains, filter_function=None): ############################################################################## # Visualize Spike Trains # ---------------------- + spike_train_list = list(spike_train__Block.segments[0].spiketrains) active_spiketrains = [st for st in spike_train_list if len(st) > 0] -ax, axhistx, axhisty = rasterplot_rates(spike_train_list, filter_function=lambda st: len(st) > 0) -ax.plot( - input_current__AnalogSignal.times, - input_current__AnalogSignal.magnitude.T[0] - / input_current__AnalogSignal.magnitude.T[0].max() - * len(active_spiketrains), - color="black", -) - -axhisty.set_xlabel("FR (pps)") - -# Clear the auto-generated histogram and add custom KDE using elephant because it looks better -axhistx.clear() +if len(active_spiketrains) > 0: + ax, axhistx, axhisty = rasterplot_rates(spike_train_list, filter_function=lambda st: len(st) > 0) + + # Overlay scaled input current + ax.plot( + input_current__AnalogSignal.times, + input_current__AnalogSignal.magnitude.T[0] + / input_current__AnalogSignal.magnitude.T[0].max() + * len(active_spiketrains), + color="black", + linewidth=2, + label="Input Current (scaled)", + ) + axhisty.set_xlabel("FR (pps)") -if len(active_spiketrains) > 0: - # Population firing rate over time, Gaussian-smoothed (sigma = 15 ms). + # Add smoothed population firing rate over time (Gaussian, sigma = 15 ms). # Native replacement for elephant.statistics.instantaneous_rate. - sampling_period_s = (h.dt * pq.ms).rescale(pq.s).magnitude + axhistx.clear() + + sampling_period_s = timestep.rescale(pq.s).magnitude t_start = min(st.t_start for st in active_spiketrains).rescale(pq.s).magnitude t_stop = max(st.t_stop for st in active_spiketrains).rescale(pq.s).magnitude n_bins = int(round((t_stop - t_start) / sampling_period_s)) @@ -365,47 +367,88 @@ def rasterplot_rates(spiketrains, filter_function=None): rate_hz = counts / sampling_period_s / len(active_spiketrains) rate_hz = gaussian_filter1d(rate_hz, sigma=(15e-3) / sampling_period_s, mode="constant") - axhistx.plot(edges[:-1] + sampling_period_s / 2, rate_hz, linewidth=2) + axhistx.plot( + edges[:-1] + sampling_period_s / 2, + rate_hz, + linewidth=2, + color="blue", + ) axhistx.set_ylabel("FR (pps)") - axhistx.set_xlim(ax.get_xlim()) # Match x-axis with raster plot - -ax.set_ylabel("Neuron Index (#)") -ax.set_xlabel("Time (s)") - -# remove top and right spines for cleaner look -sns.despine(ax=ax) - -# Make figure bigger with more white space at borders -fig = plt.gcf() -fig.set_size_inches(12, 6) - -# Add whitespace between axes (manually adjust positions since rasterplot_rates uses absolute positioning) -gap = 0.025 # Gap size between axes -bottom_margin = 0.03 # Margin from bottom - -ax_pos = ax.get_position() -axhistx_pos = axhistx.get_position() -axhisty_pos = axhisty.get_position() - -# Raise ax and axhisty from bottom, move top histogram up and right histogram right to create gaps -ax.set_position([ax_pos.x0, ax_pos.y0 + bottom_margin, ax_pos.width, ax_pos.height]) -axhistx.set_position( - [ - axhistx_pos.x0, - axhistx_pos.y0 + gap + bottom_margin, - axhistx_pos.width, - axhistx_pos.height, - ] -) -axhisty.set_position( - [ - axhisty_pos.x0 + gap, - axhisty_pos.y0 + bottom_margin, - axhisty_pos.width, - axhisty_pos.height, - ] -) + axhistx.set_xlim(ax.get_xlim()) + + ax.set_ylabel("Neuron Index (#)") + ax.set_xlabel("Time (s)") + + + sns.despine(ax=ax) + + fig = plt.gcf() + fig.set_size_inches(12, 6) + + # Adjust positioning for better layout + gap = 0.025 + bottom_margin = 0.03 -plt.show() + ax_pos = ax.get_position() + axhistx_pos = axhistx.get_position() + axhisty_pos = axhisty.get_position() -# mkdocs_gallery_thumbnail_path = "gallery_thumbs/02_simulate_spike_trains_current_injection.png" + ax.set_position([ax_pos.x0, ax_pos.y0 + bottom_margin, ax_pos.width, ax_pos.height]) + axhistx.set_position( + [ + axhistx_pos.x0, + axhistx_pos.y0 + gap + bottom_margin, + axhistx_pos.width, + axhistx_pos.height, + ] + ) + axhisty.set_position( + [ + axhisty_pos.x0 + gap, + axhisty_pos.y0 + bottom_margin, + axhisty_pos.width, + axhisty_pos.height, + ] + ) + + plt.savefig(save_path / "spike_trains_jaxley_nerlab.png", dpi=150, bbox_inches="tight") + plt.savefig(save_path / "spike_trains_jaxley_nerlab.svg", bbox_inches="tight") + print(f"\nSaved figure to {save_path / 'spike_trains_jaxley_nerlab.png'} (and .svg)") + plt.show() +else: + print("\nNo active spike trains to plot.") + print("This may indicate the current amplitude needs adjustment.") + +############################################################################## +# Compare with NEURON Results (if available) +# ------------------------------------------ + +print("\n" + "=" * 60) +print("Comparison with NEURON Results") +print("=" * 60) + +try: + neuron_block = joblib.load(save_path / "spike_train__Block_utility.pkl") + + print("\nFound NEURON results for comparison:") + for i, (neuron_seg, jaxley_seg) in enumerate( + zip(neuron_block.segments, spike_train__Block.segments) + ): + neuron_spikes = sum(len(st) for st in neuron_seg.spiketrains) + jaxley_spikes = sum(len(st) for st in jaxley_seg.spiketrains) + + neuron_active = sum(1 for st in neuron_seg.spiketrains if len(st) > 0) + jaxley_active = sum(1 for st in jaxley_seg.spiketrains if len(st) > 0) + + print(f"\nPool {i}:") + print(f" NEURON: {neuron_active} active neurons, {neuron_spikes} total spikes") + print(f" Jaxley: {jaxley_active} active neurons, {jaxley_spikes} total spikes") + +except FileNotFoundError: + print("\nNEURON results not found.") + print("Run 02_simulate_spike_trains_current_injection.py first for comparison.") + +print("\n" + "=" * 60) +print("[DONE] Jaxley biophysical simulation complete!") +print(" Using jx.integrate() with NERLab channels (napp + caL)") +print("=" * 60) diff --git a/examples/01_basic/02_simulate_spike_trains_current_injection_neuron.py b/examples/01_basic/02_simulate_spike_trains_current_injection_neuron.py new file mode 100644 index 00000000..fdbbb545 --- /dev/null +++ b/examples/01_basic/02_simulate_spike_trains_current_injection_neuron.py @@ -0,0 +1,405 @@ +""" +Spike Train Generation with Current Injection +======================================= + +This example demonstrates how to simulate spike trains in a population of alpha motor neurons using current injection. + +Two complementary workflows are presented: + +1. Manual step-by-step workflow — explicitly walks through each stage of the NEURON simulation pipeline. This workflow is intended to clarify the underlying mechanisms. + +2. Utility-function workflow — uses the high-level :func:`~myogen.utils.neuron.inject_currents_into_populations.inject_currents_and_simulate_spike_trains` function for routine simulations. + +Both workflows yield identical results; the manual version is provided purely for explanatory purposes. +""" + +# %% + +############################################################################## +# Import Libraries +# ---------------- +# +# .. important:: +# In **MyoGen** all **random number generation** is handled by the RNG returned from +# :func:`~myogen.get_random_generator`, a thin wrapper around :mod:`numpy.random`. +# +# Always fetch the generator at the call site so the current seed is honored: +# +# .. code-block:: python +# +# from myogen import simulator, get_random_generator +# get_random_generator().normal(0, 1) +# +# To change the seed, use :func:`~myogen.set_random_seed`: +# +# .. code-block:: python +# +# from myogen import set_random_seed +# set_random_seed(42) + +from pathlib import Path + +import joblib +import neuron +import numpy as np +import quantities as pq +import seaborn as sns +from matplotlib import pyplot as plt +from neo import Block, Segment, SpikeTrain +from neuron import h +from scipy.ndimage import gaussian_filter1d + +from myogen import get_random_generator +from myogen.simulator.neuron.populations import AlphaMN__Pool +from myogen.utils.currents import create_trapezoid_current +from myogen.utils.neuron.inject_currents_into_populations import ( + inject_currents_and_simulate_spike_trains, + inject_currents_into_populations, +) +from myogen.utils.nmodl import load_nmodl_mechanisms + +plt.style.use("fivethirtyeight") + + +def mean_firing_rate(spiketrain): + """Mean firing rate of a neo.SpikeTrain (replaces elephant.statistics.mean_firing_rate).""" + return (len(spiketrain) / (spiketrain.t_stop - spiketrain.t_start)).rescale(pq.Hz) + + +def rasterplot_rates(spiketrains, filter_function=None): + """Native spike raster with top/right marginal axes. + + Lightweight stand-in for ``viziphant.rasterplot.rasterplot_rates``: draws a + spike raster (one row per train), a right marginal showing each train's mean + firing rate, and an (initially empty) top marginal that the caller fills with + the smoothed population rate. Returns ``(ax, axhistx, axhisty)`` as + absolutely-positioned axes so manual position tweaks downstream still work. + """ + if filter_function is not None: + spiketrains = [st for st in spiketrains if filter_function(st)] + + fig = plt.figure() + ax = fig.add_axes((0.10, 0.10, 0.62, 0.62)) + axhistx = fig.add_axes((0.10, 0.74, 0.62, 0.16), sharex=ax) + axhisty = fig.add_axes((0.74, 0.10, 0.16, 0.62), sharey=ax) + + ax.eventplot( + [st.rescale(pq.s).magnitude for st in spiketrains], + lineoffsets=np.arange(len(spiketrains)), + colors="black", + linelengths=0.8, + linewidths=0.7, + ) + rates = [float(mean_firing_rate(st).magnitude) for st in spiketrains] + axhisty.barh(np.arange(len(spiketrains)), rates, height=0.85, color="C0") + ax.set_ylim(-1, max(len(spiketrains), 1)) + return ax, axhistx, axhisty + +############################################################################## +# Create Motor Neuron Populations (Pools) +# --------------------------------------- +# +# In MyoGen a population of cells (e.g. motor neurons) is represented by a +# **Population** class and available in the `myogen.simulator.neuron.populations` module. +# +# A population can easily be created by specifying the number of cells. Plausible default parameters are already set. +# +# For a motor neuron population (referred to as **motor pool**), we can use the :class:`~myogen.simulator.neuron.populations.AlphaMN__Pool` class. +# This class can also use the recruitment thresholds generated in the previous example to distribute the motor units properties in a physiologically plausible manner. +# +# .. important:: +# These **Population** classes are custom-built and use therefore custom NMODL mechanisms. +# To use them, the NMODL mechanisms need to be loaded first using the :func:`~myogen.utils.nmodl.load_nmodl_mechanisms` function. +# +# To showcase MyoGen's capabilities, we will create two different motor neuron pools with identical properties but different input currents. +load_nmodl_mechanisms() + +save_path = Path("./results") +save_path.mkdir(exist_ok=True) + +recruitment_thresholds = joblib.load(save_path / "thresholds.pkl") + +n_pools = 2 +motor_neuron_pools = [ + AlphaMN__Pool(recruitment_thresholds__array=recruitment_thresholds) for _ in range(n_pools) +] + +############################################################################## +# Create Input Currents +# --------------------- +# +# To drive the motor units, we use a **common input current profile**. +# +# In this example, we use a **trapezoid-shaped input current** which is generated using the :func:`~myogen.utils.currents.create_trapezoid_current` function. +# +# .. note:: +# More convenient functions for generating input current profiles are available in the :mod:`myogen.utils.currents` module. +# +# .. note:: +# The generated input current is an instance of the :class:`neo.core.AnalogSignal` class from the :mod:`neo` package. + +timestep = 0.05 * pq.ms +simulation_time = 4000 * pq.ms + +rise_time_ms = list(get_random_generator().uniform(100, 500, size=n_pools)) * pq.ms +plateau_time_ms = list(get_random_generator().uniform(1000, 2000, size=n_pools)) * pq.ms +fall_time_ms = list(get_random_generator().uniform(1000, 2000, size=n_pools)) * pq.ms + +input_current__AnalogSignal = create_trapezoid_current( + n_pools, + int(simulation_time / timestep), + timestep, + amplitudes__nA=[15.0 * pq.nA] * n_pools, + rise_times__ms=rise_time_ms, + plateau_times__ms=plateau_time_ms, + fall_times__ms=fall_time_ms, + delays__ms=500.0 * pq.ms, +) + +print( + f"Input current signal shape: {input_current__AnalogSignal.shape}\nClass: {input_current__AnalogSignal.__class__}" +) + +# Save input current signal for later analysis +joblib.dump(input_current__AnalogSignal, save_path / "input_current__AnalogSignal.pkl") + +############################################################################## +# Manual Simulation Approach - Step by Step +# ------------------------------------------- +# +# Before showing the convenient utility function, let's understand what happens +# under the hood by implementing the simulation pipeline manually. +# This approach gives you full control and helps understand NEURON's mechanisms. + +# Step 1: Set up current injection manually +# ========================================= +# We need to inject time-varying currents into each motor neuron. +# This uses NEURON's :class:`neuron.h.IClamp` (current clamp) mechanism with :meth:`neuron.h.Vector.play`. + +inject_currents_into_populations(motor_neuron_pools, input_current__AnalogSignal) + +# Step 2: Set up spike recording manually +# ======================================= +# For each neuron, we create a :class:`neuron.h.NetCon` (network connection) object that detects +# spikes when the membrane voltage crosses a threshold, and records spike times. + +spike_detection_threshold__mV = 50.0 * pq.mV +simulation_time__ms = input_current__AnalogSignal.t_stop.rescale(pq.ms) + +spike_recorders = [] + +for pool_idx, pool in enumerate(motor_neuron_pools): + pool_spike_recorders = [] + + for cell in pool: + # Create a vector to record spike times + spike_recorder = h.Vector() + + # Create NetCon object: monitors voltage at soma(0.5) and records spikes + # NetCon(source, target, threshold, delay, weight) + # source: cell.soma(0.5)._ref_v (membrane voltage reference) + # target: None (no post-synaptic target, just recording) + nc = h.NetCon(cell.soma(0.5)._ref_v, None, sec=cell.soma) + nc.threshold = spike_detection_threshold__mV # Spike detection threshold + nc.record(spike_recorder) # Record spike times into vector + + pool_spike_recorders.append(spike_recorder) + + spike_recorders.append(pool_spike_recorders) + +# Step 3: Initialize voltages and run simulation +# ============================================== +# Before running, we need to initialize membrane voltages to physiological values. +# +# .. note:: For this MyoGen populations provide the ``get_initialization_data()`` method. +# This returns the sections and their initial voltages. + +# Initialize each neuron's membrane voltage to its resting potential +for pool in motor_neuron_pools: + for section, voltage in zip(*pool.get_initialization_data()): + section.v = voltage + +# Initialize NEURON's internal state and run the simulation +h.finitialize() # Initialize all mechanisms and variables +neuron.run(simulation_time__ms) + +# Step 4: Convert recorded data to :class:`neo.core.Block` format +# ================================================================== +# The spike times are now stored in NEURON vectors. We convert them to +# the standardized :class:`neo.core.Block` format for analysis and compatibility. + +spike_train__Block_manual = Block(name="Manual Simulation Results") + +for pool_idx, pool_spike_recorders in enumerate(spike_recorders): + # Create a segment for this motor unit pool + segment = Segment(name=f"Pool {pool_idx}") + + # Convert each neuron's spike times to a :class:`neo.core.SpikeTrain` object + segment.spiketrains = [] + for neuron_idx, spike_recorder in enumerate(pool_spike_recorders): + # Convert NEURON vector to numpy array and add units + spike_times = (spike_recorder.as_numpy() * pq.ms).rescale(pq.s) + + # Create :class:`neo.core.SpikeTrain` object with metadata + spiketrain = SpikeTrain( + spike_times, + t_stop=simulation_time__ms.rescale(pq.s), + sampling_rate=(1 / (h.dt * pq.ms)).rescale(pq.Hz), + sampling_period=(h.dt * pq.ms).rescale(pq.s), + name=str(neuron_idx), + description=f"Pool {pool_idx}, Neuron {neuron_idx}", + ) + segment.spiketrains.append(spiketrain) + + spike_train__Block_manual.segments.append(segment) + +joblib.dump(spike_train__Block_manual, save_path / "spike_train__Block_manual.pkl") + +############################################################################## +# Convenient Utility Function Approach +# ------------------------------------ +# +# The manual approach above shows you exactly what happens during simulation. +# However, since this is a common task, MyoGen provides the :func:`~myogen.utils.neuron.inject_currents_into_populations.inject_currents_and_simulate_spike_trains` +# utility function that encapsulates all these steps in a single call. +# +# This is the recommended approach for routine simulations, while the manual +# approach is useful when you need custom spike detection, specialized recording, +# or want to understand the underlying mechanisms. + +# Run the same simulation using the utility function +spike_train__Block = inject_currents_and_simulate_spike_trains( + populations=motor_neuron_pools, + input_current__AnalogSignal=input_current__AnalogSignal, + spike_detection_thresholds__mV=50 * pq.mV, +) + +joblib.dump(spike_train__Block, save_path / "spike_train__Block_utility.pkl") + +# Compare the two approaches +print("\nComparison of results:") +print(f"Manual approach: {len(spike_train__Block_manual.segments)} segments") +print(f"Utility approach: {len(spike_train__Block.segments)} segments") + +# Verify they produce similar results (spike counts should be identical) +for i, (manual_seg, utility_seg) in enumerate( + zip(spike_train__Block_manual.segments, spike_train__Block.segments) +): + manual_spikes = sum(len(st) for st in manual_seg.spiketrains) + utility_spikes = sum(len(st) for st in utility_seg.spiketrains) + print(f"Pool {i}: Manual={manual_spikes} spikes, Utility={utility_spikes} spikes") + +############################################################################## +# Calculate and Display Statistics +# --------------------------------- +# +# It might be of interest to calculate the **firing rates** of the motor units. +# +# .. note:: +# The **firing rates** are calculated as the number of spikes divided by the time in which each MU was active. + +firing_rates = [ + np.array( + [ + mean_firing_rate(st__s.time_slice(st__s.min(), st__s.max())) + for st__s in spike_train__segment.spiketrains + if len(st__s) > 1 # Need at least 2 spikes to compute rate over spike range + ] + ) + for spike_train__segment in spike_train__Block.segments +] + +print("Firing rate statistics:") +for pool_idx, firing_rates_per_pool in enumerate(firing_rates): + active_neurons = np.sum(firing_rates_per_pool > 0) + if len(firing_rates_per_pool) > 0 and np.sum(firing_rates_per_pool > 0) > 0: + mean_rate = np.mean(firing_rates_per_pool[firing_rates_per_pool > 0]) + max_rate = np.max(firing_rates_per_pool) + else: + mean_rate = 0.0 + max_rate = 0.0 + + print( + f" Pool {pool_idx + 1}: {active_neurons}/{len(recruitment_thresholds)} active neurons, " + f"mean rate: {mean_rate:.1f} Hz, max rate: {max_rate:.1f} Hz" + ) + +############################################################################## +# Visualize Spike Trains +# ---------------------- +spike_train_list = list(spike_train__Block.segments[0].spiketrains) +active_spiketrains = [st for st in spike_train_list if len(st) > 0] + +ax, axhistx, axhisty = rasterplot_rates(spike_train_list, filter_function=lambda st: len(st) > 0) +ax.plot( + input_current__AnalogSignal.times, + input_current__AnalogSignal.magnitude.T[0] + / input_current__AnalogSignal.magnitude.T[0].max() + * len(active_spiketrains), + color="black", +) + +axhisty.set_xlabel("FR (pps)") + +# Clear the auto-generated histogram and add custom KDE using elephant because it looks better +axhistx.clear() + + +if len(active_spiketrains) > 0: + # Population firing rate over time, Gaussian-smoothed (sigma = 15 ms). + # Native replacement for elephant.statistics.instantaneous_rate. + sampling_period_s = (h.dt * pq.ms).rescale(pq.s).magnitude + t_start = min(st.t_start for st in active_spiketrains).rescale(pq.s).magnitude + t_stop = max(st.t_stop for st in active_spiketrains).rescale(pq.s).magnitude + n_bins = int(round((t_stop - t_start) / sampling_period_s)) + edges = t_start + np.arange(n_bins + 1) * sampling_period_s + all_spikes = np.concatenate( + [st.rescale(pq.s).magnitude for st in active_spiketrains] + ) + all_spikes = all_spikes[(all_spikes >= edges[0]) & (all_spikes < edges[-1])] + counts, _ = np.histogram(all_spikes, bins=edges) + rate_hz = counts / sampling_period_s / len(active_spiketrains) + rate_hz = gaussian_filter1d(rate_hz, sigma=(15e-3) / sampling_period_s, mode="constant") + + axhistx.plot(edges[:-1] + sampling_period_s / 2, rate_hz, linewidth=2) + axhistx.set_ylabel("FR (pps)") + axhistx.set_xlim(ax.get_xlim()) # Match x-axis with raster plot + +ax.set_ylabel("Neuron Index (#)") +ax.set_xlabel("Time (s)") + +# remove top and right spines for cleaner look +sns.despine(ax=ax) + +# Make figure bigger with more white space at borders +fig = plt.gcf() +fig.set_size_inches(12, 6) + +# Add whitespace between axes (manually adjust positions since rasterplot_rates uses absolute positioning) +gap = 0.025 # Gap size between axes +bottom_margin = 0.03 # Margin from bottom + +ax_pos = ax.get_position() +axhistx_pos = axhistx.get_position() +axhisty_pos = axhisty.get_position() + +# Raise ax and axhisty from bottom, move top histogram up and right histogram right to create gaps +ax.set_position([ax_pos.x0, ax_pos.y0 + bottom_margin, ax_pos.width, ax_pos.height]) +axhistx.set_position( + [ + axhistx_pos.x0, + axhistx_pos.y0 + gap + bottom_margin, + axhistx_pos.width, + axhistx_pos.height, + ] +) +axhisty.set_position( + [ + axhisty_pos.x0 + gap, + axhisty_pos.y0 + bottom_margin, + axhisty_pos.width, + axhisty_pos.height, + ] +) + +plt.show() diff --git a/examples/01_basic/03_simulate_spike_trains_descending_drive.py b/examples/01_basic/03_simulate_spike_trains_descending_drive.py index 7369c3ab..abd6e588 100644 --- a/examples/01_basic/03_simulate_spike_trains_descending_drive.py +++ b/examples/01_basic/03_simulate_spike_trains_descending_drive.py @@ -1,32 +1,21 @@ +"""Spike Train Generation with Descending Drive — Jaxley Backend +================================================================= + +This is the Jaxley counterpart of +``03_simulate_spike_trains_descending_drive_neuron.py``. Both examples use +the current ``fdi_semg`` profile card, the same 25 % MVC protocol, the same +NERLab pool parameters, and the same calibrated descending-drive network. + +The backend-specific random streams are deliberately independent, so normal +profile runs are compared at the population/statistical level. Exact +spike-for-spike backend validation uses the frozen-event protocol in +``experiments/twin/validation_protocol.py``. + +Unlike the former Jaxley example, this path does not use a fitted 5 ms IIR +current surrogate. It uses the NEURON model's dendritic, peak-normalized +0.2/0.3 ms Exp2Syn kinetics, live driving force, equal fan-in, dendritic +transfer scaling, two-stage event delay, and live-voltage Gfluctdv noise. """ -Spike Train Generation with Descending Drive -============================================ - -This example simulates **motor neuron spike trains** driven by a **descending drive (DD)**: a population of -Poisson neurons modelling cortical input, connected to the motor neuron pool through synapses. This is more -realistic than direct current injection (example 02) because it models the synaptic input from upper motor neurons. - -The drive here is the **calibrated FDI drive** of the ``fdi_semg`` profile card: its network, noise and drive rates -were fitted so the pool recruits and fires like a human first dorsal interosseous, and the card maps force levels in -**% MVC** to drive. One call runs it. - -!!! note - The card is ``myogen/profiles/fdi_semg.yaml``: every value with its source, and notes on what else to check when - you change one. ``load_profile("fdi_semg.published")`` gives each component as its source paper describes it. - To build the pools yourself, see ``card.motor_neuron_pool()`` and - ``myogen.simulator.neuron.descending_drive_network``. -""" - -# %% - -############################################################################## -# Import Libraries -# ---------------- -# -# !!! important -# In **MyoGen** all **random number generation** is handled by the RNG returned from -# [`get_random_generator`][myogen.get_random_generator], a thin wrapper around `numpy.random`. -# The profile card reseeds it from its own ``seeds`` before each build, so a card run is reproducible. from pathlib import Path @@ -34,260 +23,242 @@ import numpy as np import quantities as pq from matplotlib import pyplot as plt -from neo import AnalogSignal +from neo import AnalogSignal, Block, Segment, SpikeTrain from myogen import load_profile -from myogen.utils.nmodl import load_nmodl_mechanisms +from myogen.simulator.jaxley.descending_drive_network import ( + excitation_to_drive_hz, + run_descending_drive, +) plt.style.use("fivethirtyeight") def mean_firing_rate(spiketrain): - """Mean firing rate of a neo.SpikeTrain (replaces elephant.statistics.mean_firing_rate).""" return (len(spiketrain) / (spiketrain.t_stop - spiketrain.t_start)).rescale(pq.Hz) def population_psth(spiketrains, bin_size): - """Total spike counts per bin across spiketrains, plus bin left edges (s). - - Native replacement for elephant.statistics.time_histogram (output="counts"). - """ + """Population spike count and bin centres.""" t_start = min(st.t_start for st in spiketrains).rescale(pq.s).magnitude t_stop = max(st.t_stop for st in spiketrains).rescale(pq.s).magnitude - bs = (bin_size).rescale(pq.s).magnitude - n_bins = int((t_stop - t_start) / bs) - edges = t_start + np.arange(n_bins + 1) * bs + bs = bin_size.rescale(pq.s).magnitude + edges = np.arange(t_start, t_stop + 0.5 * bs, bs) spikes = np.concatenate([st.rescale(pq.s).magnitude for st in spiketrains]) - spikes = spikes[(spikes >= edges[0]) & (spikes < edges[-1])] # drop right-edge spikes + spikes = spikes[(spikes >= edges[0]) & (spikes < edges[-1])] counts, _ = np.histogram(spikes, bins=edges) - return counts, edges[:-1] + return counts, edges[:-1] + 0.5 * bs ############################################################################## -# Run the Calibrated Drive -# ------------------------ -# -# A **trapezoidal contraction** to 25 % MVC: 1 s rest, 4 s ramp up, 4 s hold, 4 s ramp down, 1 s rest. Each -# segment is ``(duration_s, start_pct_mvc, end_pct_mvc)``. -# -# !!! note -# This runs 120 biophysical motor neurons for 14 s and takes a few minutes. - -load_nmodl_mechanisms() +# Build the same profile-card protocol as the NEURON example +# ----------------------------------------------------------- save_path = Path("./results") save_path.mkdir(exist_ok=True) card = load_profile() # fdi_semg protocol = [(1, 0, 0), (4, 0, 25), (4, 25, 25), (4, 25, 0), (1, 0, 0)] -result = card.simulate(protocol) +segments = [ + (float(duration), card.excitation_for_pct_mvc(start), card.excitation_for_pct_mvc(stop)) + for duration, start, stop in protocol +] +dd = card.dd_params() +dt_ms = float(card.value("drive.timestep_ms")) +timestep = dt_ms * pq.ms +excitation, drive_hz = excitation_to_drive_hz( + segments, + float(card.value("drive.threshold_hz")), + float(dd["dd_drive__Hz"]), + timestep_ms=dt_ms, + noise_seed=card.seeds["drive"], +) +simulation_time = len(drive_hz) * dt_ms * pq.ms +ramp_up_end = 5.0 * pq.s +plateau_end = 9.0 * pq.s -spike_train_block = result.spikes -mn_segment = spike_train_block.segments[0] -n_motor_neurons = len(mn_segment.spiketrains) -simulation_time = sum(d for d, _, _ in protocol) * pq.s -timestep = float(card.value("drive.timestep_ms")) * pq.ms -ramp_up_end = 5.0 * pq.s # start of the hold -plateau_end = 9.0 * pq.s # end of the hold +print(f"Running calibrated Jaxley drive: {simulation_time.rescale(pq.s)}") +print(f" MNs={len(card.thresholds())}, DD sources={dd['dd_neurons']}, dt={dt_ms} ms") +print(f" DD rate={drive_hz.min():.1f}–{drive_hz.max():.1f} Hz") -# Excitation (0-1) per drive time step, saved for the surface EMG example -excitation_signal = AnalogSignal(result.excitation[:, None], units=pq.dimensionless, - sampling_period=timestep.rescale(pq.s)) +############################################################################## +# Run Jaxley with the matched conductance pathway +# ------------------------------------------------ + +run = run_descending_drive( + card.thresholds(), + drive_hz, + dd, + timestep_ms=dt_ms, + seed=int(card.seeds["drive"]), + # The current NEURON card inherits -67 mV from alpha_mn_default.yaml. + # Use the same initial state here. The 1 s pre-ramp lets the pool settle. + initial_voltage_mV=-67.0, +) + +n_steps, n_motor_neurons = run.spikes.shape +n_dd = run.dd_spikes.shape[1] +mn_spike_times = [ + (np.flatnonzero(run.spikes[:, i]) + 1) * dt_ms for i in range(n_motor_neurons) +] +dd_spike_times = [ + np.flatnonzero(run.dd_spikes[:, i]) * dt_ms for i in range(n_dd) +] -# The spike trains and the excitation are used by the surface EMG example (06) and the force example (08) -joblib.dump(spike_train_block, save_path / "trapezoid_dd_spike_trains.pkl") -joblib.dump(excitation_signal, save_path / "trapezoid_drive_pattern.pkl") +print(f" dendritic transfer={run.transfer.min():.3f}–{run.transfer.max():.3f}") +print(f" synaptic weights={run.weights_uS.min():.5f}–{run.weights_uS.max():.5f} µS") ############################################################################## -# Calculate Firing Rate Statistics -# --------------------------------- - -mn_firing_rates = np.array( - [ - mean_firing_rate(st__s.time_slice(st__s.min(), st__s.max())) - for st__s in mn_segment.spiketrains - if len(st__s) > 1 # need >=2 spikes for a rate over the spike span - ] +# Convert to Neo and generate the same force observable +# ------------------------------------------------------ + +spike_train_block = Block(name="Calibrated FDI descending drive — Jaxley") +mn_segment = Segment(name="Motor Neurons") +mn_segment.spiketrains = [ + SpikeTrain( + (times * pq.ms).rescale(pq.s), + t_stop=simulation_time.rescale(pq.s), + sampling_rate=(1 / timestep).rescale(pq.Hz), + sampling_period=timestep.rescale(pq.s), + name=str(i), + ) + for i, times in enumerate(mn_spike_times) +] +spike_train_block.segments.append(mn_segment) + +dd_segment = Segment(name="Descending Drive") +dd_segment.spiketrains = [ + SpikeTrain( + (times * pq.ms).rescale(pq.s), + t_stop=simulation_time.rescale(pq.s), + sampling_rate=(1 / timestep).rescale(pq.Hz), + sampling_period=timestep.rescale(pq.s), + name=f"DD_{i}", + ) + for i, times in enumerate(dd_spike_times) +] + +force_model = card.force_model() +force = np.asarray(force_model.generate_force(spike_train_block, verbose=False).magnitude)[:, 0] +force_pct_mvc = 100.0 * force / float(card.value("drive.mvc_model_units")) +fs_force_hz = float(card.value("force.recording_frequency__Hz")) + +excitation_signal = AnalogSignal( + excitation[:, None], units=pq.dimensionless, sampling_period=timestep.rescale(pq.s) +) +drive_signal = AnalogSignal( + drive_hz[:, None], units=pq.Hz, sampling_period=timestep.rescale(pq.s) +) +joblib.dump(spike_train_block, save_path / "trapezoid_dd_spike_trains_jaxley.pkl") +joblib.dump(excitation_signal, save_path / "trapezoid_drive_pattern_jaxley.pkl") +joblib.dump(drive_signal, save_path / "trapezoid_drive_hz_jaxley.pkl") + +diag_ids = sorted(set(np.linspace(0, n_motor_neurons - 1, 6, dtype=int))) +joblib.dump( + { + **{f"MN_{i}": run.voltage_mV[:, i].astype(np.float32) for i in diag_ids}, + "__dt_ms": np.float32(dt_ms), + "__sim_ms": np.float32(simulation_time.rescale(pq.ms).magnitude), + }, + save_path / "trapezoid_dd_voltage_traces_jaxley.pkl", ) +############################################################################## +# Summary metrics +# --------------- + +all_counts = np.asarray([len(st) for st in mn_segment.spiketrains]) +plateau_counts = np.asarray([ + len(st.time_slice(ramp_up_end, plateau_end)) for st in mn_segment.spiketrains +]) +active = int(np.sum(all_counts > 0)) +sustained = int(np.sum(plateau_counts >= 5)) +plateau_total = int(plateau_counts.sum()) +plateau_rate = plateau_total / max(sustained, 1) / 4.0 + print("Motor neurons:") -print(f"\tActive neurons: {len(mn_firing_rates)}/{n_motor_neurons}") -if len(mn_firing_rates) > 0: - print(f"\tMean firing rate: {np.mean(mn_firing_rates):.1f} ± {np.std(mn_firing_rates):.1f} pps") - print(f"\tRate range: {np.min(mn_firing_rates):.1f} - {np.max(mn_firing_rates):.1f} pps") +print(f" active: {active}/{n_motor_neurons}") +print(f" total spikes: {all_counts.sum()}") +print(f" sustained during hold: {sustained}/{n_motor_neurons}") +print(f" hold spikes: {plateau_total}") +print(f" mean sustained hold rate: {plateau_rate:.1f} Hz") +print(f" peak force: {force_pct_mvc.max():.1f} % MVC") ############################################################################## -# Visualization -# ------------- -# -# 1. Target force and the simulated force (% MVC) -# 2. Motor neuron raster plot showing recruitment -# 3. Population firing rate over time +# Matched visualization +# --------------------- fig, axes = plt.subplots(3, 1, figsize=(15, 10), sharex=True) - -t_force = np.arange(result.force_pct_mvc.size) / result.fs_force__Hz -t_target = np.cumsum([0] + [d for d, _, _ in protocol]) -target = [protocol[0][1]] + [b for _, _, b in protocol] +t_force = np.arange(force_pct_mvc.size) / fs_force_hz +t_target = np.cumsum([0] + [duration for duration, _, _ in protocol]) +target = [protocol[0][1]] + [stop for _, _, stop in protocol] axes[0].plot(t_target, target, "k--", linewidth=2, label="Target") -axes[0].plot(t_force, result.force_pct_mvc, "b-", linewidth=2, label="Force") +axes[0].plot(t_force, force_pct_mvc, "b-", linewidth=2, label="Jaxley force") axes[0].set_ylabel("Force (% MVC)") -axes[0].set_title("Trapezoidal Contraction (calibrated FDI drive)") +axes[0].set_title("Trapezoidal contraction — calibrated FDI drive") axes[0].legend() -axes[0].grid(True, alpha=0.3) - -mn_colors = plt.get_cmap("Reds")(np.linspace(0.3, 0.9, n_motor_neurons)) -active_mn_count = 0 -for i, (spiketrain, color) in enumerate(zip(mn_segment.spiketrains, mn_colors)): - if len(spiketrain) > 0: - spike_times = spiketrain.rescale(pq.s).magnitude - axes[1].scatter(spike_times, [i] * len(spike_times), c=[color], s=1.0, alpha=0.8) - active_mn_count += 1 - -axes[1].set_ylabel("Motor Neuron ID\n(Recruitment Order)") -axes[1].set_title(f"Motor Neuron Population Activity (n={active_mn_count}/{n_motor_neurons} active)") + +colors = plt.get_cmap("Reds")(np.linspace(0.3, 0.9, n_motor_neurons)) +for i, (st, color) in enumerate(zip(mn_segment.spiketrains, colors)): + if len(st): + axes[1].scatter(st.rescale(pq.s).magnitude, np.full(len(st), i), c=[color], s=1) +axes[1].set_ylabel("Motor neuron ID") +axes[1].set_title(f"Motor-neuron activity ({active}/{n_motor_neurons} active)") axes[1].set_ylim(-1, n_motor_neurons) -axes[1].grid(True, alpha=0.3) -bin_size_ms = 100 -mn_counts, bin_centers_s = population_psth(mn_segment.spiketrains, bin_size_ms * pq.ms) -mn_rates_binned = mn_counts / (bin_size_ms / 1000.0) / n_motor_neurons -axes[2].plot(bin_centers_s, mn_rates_binned, "r-", linewidth=2, label="MN Population", alpha=0.8) +bin_ms = 100 +counts, centres = population_psth(mn_segment.spiketrains, bin_ms * pq.ms) +population_rate = counts / (bin_ms / 1000.0) / n_motor_neurons +axes[2].plot(centres, population_rate, "r-", linewidth=2) axes[2].set_xlabel("Time (s)") -axes[2].set_ylabel("Population Rate (Hz)") -axes[2].set_title("Population Firing Rate Over Time") -axes[2].legend() -axes[2].grid(True, alpha=0.3) +axes[2].set_ylabel("Population rate (Hz)") +axes[2].set_title("Population firing rate") for ax in axes: + ax.grid(True, alpha=0.3) ax.set_xlim(0, simulation_time.rescale(pq.s).magnitude) - plt.tight_layout() +plt.savefig(save_path / "trapezoid_dd_jaxley_matched.png", dpi=180) plt.show() -############################################################################## -# Individual Motor Neuron Discharge Rates -# ---------------------------------------- -# -# Compute smoothed instantaneous firing rates for each motor neuron -# with a 400 ms Hanning window. - -print("\nComputing smoothed discharge rates per neuron...") - -# Parameters -window_ms = 400 * pq.ms # 400 ms Hanning window -dt_s = timestep.rescale(pq.s) # simulation timestep in seconds -window_samples = int(window_ms.rescale(pq.s) / dt_s) - -# Hanning window normalized to preserve rate -hanning_window = np.hanning(window_samples) -hanning_window = hanning_window / (hanning_window.sum() * dt_s) # convert to Hz - -mn_instantaneous_rates = [] -active_neuron_ids = [] - -mean_firing_rates = [] -cv_isi = [] - -for i, spiketrain in enumerate(mn_segment.spiketrains): - if len(spiketrain) > 2: - # Convert spike times to a binary spike train - t = np.arange(0, simulation_time.rescale(pq.s).magnitude, dt_s.magnitude) - spikes = np.zeros_like(t) - spike_indices = np.searchsorted(t, spiketrain.magnitude) - spikes[spike_indices[spike_indices < len(t)]] = 1 - - # Convolve with Hanning window - rate = np.convolve(spikes, hanning_window, mode="same") - mn_instantaneous_rates.append(rate) - active_neuron_ids.append(i) - - # IMPORTANT: Compute ISI/CV only during plateau phase where firing is stable - # Filter spike train to plateau phase - plateau_spiketrain = spiketrain.time_slice(ramp_up_end, plateau_end) - - # Compute mean firing rate (Hz) from plateau spikes only - plateau_duration_s = float((plateau_end - ramp_up_end).rescale(pq.s).magnitude) - mean_rate = ( - len(plateau_spiketrain) / plateau_duration_s if len(plateau_spiketrain) > 0 else 0.0 - ) - mean_firing_rates.append(mean_rate) - - # Compute CV of inter-spike intervals (plateau phase only) - if len(plateau_spiketrain) > 1: - spike_times = plateau_spiketrain.rescale(pq.s).magnitude - isis = np.diff(spike_times) - cv = np.std(isis) / np.mean(isis) if len(isis) > 1 else 0.0 - else: - cv = 0.0 - cv_isi.append(cv) - -# Population averages -pop_mean_rate = np.mean(mean_firing_rates) -pop_mean_cv = np.mean(cv_isi) -print(f"\nPopulation: Mean firing rate = {pop_mean_rate:.2f} Hz, CV = {pop_mean_cv:.2f}") - -print(f" Computed rates for {len(active_neuron_ids)} active motor neurons") - -# Create new figure for individual discharge rates -fig2, axes2 = plt.subplots(2, 1, figsize=(15, 10), sharex=True) - -# 1. Heatmap of instantaneous firing rates -if len(mn_instantaneous_rates) > 0: - # Stack rates into 2D array (neurons x time) - rates_array = np.array(mn_instantaneous_rates) - time_points = np.linspace(0, simulation_time.rescale(pq.s).magnitude, rates_array.shape[1]) - - # Plot heatmap (with origin='lower' so MU 0 is at TOP) - im = axes2[0].imshow( - rates_array, - aspect="auto", - cmap="hot", - interpolation="bilinear", - extent=[0, simulation_time.rescale(pq.s).magnitude, 0, len(active_neuron_ids)], - origin="lower", # Puts first row (MU 0) at bottom, but we'll flip with extent - vmin=0, - vmax=np.percentile(rates_array, 95), - ) - axes2[0].set_ylabel("Motor Neuron ID\n(Recruitment Order, MU 0 at top)") - axes2[0].set_title( - "Individual Motor Neuron Discharge Rates (Smoothed with 400ms Hanning Window)" +############################################################################## +# Individual smoothed discharge rates +# ------------------------------------ + +window_samples = round(0.4 / timestep.rescale(pq.s).magnitude) +window = np.hanning(window_samples) +window /= window.sum() * timestep.rescale(pq.s).magnitude +rates, active_ids = [], [] +for i, st in enumerate(mn_segment.spiketrains): + if len(st) <= 2: + continue + binary = np.zeros(n_steps) + idx = np.rint(st.rescale(pq.ms).magnitude / dt_ms).astype(int) + idx = idx[(idx >= 0) & (idx < n_steps)] + binary[idx] = 1.0 + rates.append(np.convolve(binary, window, mode="same")) + active_ids.append(i) + +if rates: + rates = np.asarray(rates) + fig2, axes2 = plt.subplots(2, 1, figsize=(15, 10), sharex=True) + image = axes2[0].imshow( + rates, aspect="auto", cmap="hot", interpolation="bilinear", origin="lower", + extent=[0, simulation_time.rescale(pq.s).magnitude, 0, len(active_ids)], + vmin=0, vmax=np.percentile(rates, 95), ) - # Add colorbar - cbar = plt.colorbar(im, ax=axes2[0]) - cbar.set_label("Firing Rate (Hz)") - axes2[0].grid(False) - - # 2. Individual traces - n_to_plot = len(active_neuron_ids) - - # Use colormap for lines (gradient from blue to red showing recruitment order) - colors = plt.get_cmap("rainbow")(np.linspace(0, 1, n_to_plot)) - - for neuron_idx in range(n_to_plot): - axes2[1].plot( - time_points, - mn_instantaneous_rates[neuron_idx], - linewidth=0.8, - color=colors[neuron_idx], - label=f"MN {active_neuron_ids[neuron_idx]}" if n_to_plot <= 20 else None, - ) - + axes2[0].set_ylabel("Active motor units") + axes2[0].set_title("Individual discharge rates — 400 ms Hanning smoothing") + plt.colorbar(image, ax=axes2[0], label="Firing rate (Hz)") + t = np.arange(n_steps) * dt_ms / 1000.0 + for rate in rates: + axes2[1].plot(t, rate, linewidth=0.7, alpha=0.7) axes2[1].set_xlabel("Time (s)") - axes2[1].set_ylabel("Firing Rate (Hz)") - axes2[1].set_title(f"All Motor Neuron Discharge Rates (n={n_to_plot})") - - # Only show legend if there are few neurons - if n_to_plot <= 20: - axes2[1].legend(loc="upper right", ncol=3, fontsize=6) - - axes2[1].grid(True, alpha=0.3) + axes2[1].set_ylabel("Firing rate (Hz)") axes2[1].set_xlim(0, simulation_time.rescale(pq.s).magnitude) - axes2[1].set_ylim(0, np.max(rates_array) * 1.1) - -plt.tight_layout() -plt.show() - -print("\n[DONE] Simulation complete with discharge rate analysis!") + axes2[1].grid(True, alpha=0.3) + plt.tight_layout() + plt.savefig(save_path / "trapezoid_dd_discharge_rates_jaxley.png", dpi=180) + plt.show() -# mkdocs_gallery_thumbnail_path = "gallery_thumbs/03_simulate_spike_trains_descending_drive.png" +print("[DONE] Matched profile-card Jaxley descending-drive simulation complete.") diff --git a/examples/01_basic/03_simulate_spike_trains_descending_drive_neuron.py b/examples/01_basic/03_simulate_spike_trains_descending_drive_neuron.py new file mode 100644 index 00000000..7369c3ab --- /dev/null +++ b/examples/01_basic/03_simulate_spike_trains_descending_drive_neuron.py @@ -0,0 +1,293 @@ +""" +Spike Train Generation with Descending Drive +============================================ + +This example simulates **motor neuron spike trains** driven by a **descending drive (DD)**: a population of +Poisson neurons modelling cortical input, connected to the motor neuron pool through synapses. This is more +realistic than direct current injection (example 02) because it models the synaptic input from upper motor neurons. + +The drive here is the **calibrated FDI drive** of the ``fdi_semg`` profile card: its network, noise and drive rates +were fitted so the pool recruits and fires like a human first dorsal interosseous, and the card maps force levels in +**% MVC** to drive. One call runs it. + +!!! note + The card is ``myogen/profiles/fdi_semg.yaml``: every value with its source, and notes on what else to check when + you change one. ``load_profile("fdi_semg.published")`` gives each component as its source paper describes it. + To build the pools yourself, see ``card.motor_neuron_pool()`` and + ``myogen.simulator.neuron.descending_drive_network``. +""" + +# %% + +############################################################################## +# Import Libraries +# ---------------- +# +# !!! important +# In **MyoGen** all **random number generation** is handled by the RNG returned from +# [`get_random_generator`][myogen.get_random_generator], a thin wrapper around `numpy.random`. +# The profile card reseeds it from its own ``seeds`` before each build, so a card run is reproducible. + +from pathlib import Path + +import joblib +import numpy as np +import quantities as pq +from matplotlib import pyplot as plt +from neo import AnalogSignal + +from myogen import load_profile +from myogen.utils.nmodl import load_nmodl_mechanisms + +plt.style.use("fivethirtyeight") + + +def mean_firing_rate(spiketrain): + """Mean firing rate of a neo.SpikeTrain (replaces elephant.statistics.mean_firing_rate).""" + return (len(spiketrain) / (spiketrain.t_stop - spiketrain.t_start)).rescale(pq.Hz) + + +def population_psth(spiketrains, bin_size): + """Total spike counts per bin across spiketrains, plus bin left edges (s). + + Native replacement for elephant.statistics.time_histogram (output="counts"). + """ + t_start = min(st.t_start for st in spiketrains).rescale(pq.s).magnitude + t_stop = max(st.t_stop for st in spiketrains).rescale(pq.s).magnitude + bs = (bin_size).rescale(pq.s).magnitude + n_bins = int((t_stop - t_start) / bs) + edges = t_start + np.arange(n_bins + 1) * bs + spikes = np.concatenate([st.rescale(pq.s).magnitude for st in spiketrains]) + spikes = spikes[(spikes >= edges[0]) & (spikes < edges[-1])] # drop right-edge spikes + counts, _ = np.histogram(spikes, bins=edges) + return counts, edges[:-1] + + +############################################################################## +# Run the Calibrated Drive +# ------------------------ +# +# A **trapezoidal contraction** to 25 % MVC: 1 s rest, 4 s ramp up, 4 s hold, 4 s ramp down, 1 s rest. Each +# segment is ``(duration_s, start_pct_mvc, end_pct_mvc)``. +# +# !!! note +# This runs 120 biophysical motor neurons for 14 s and takes a few minutes. + +load_nmodl_mechanisms() + +save_path = Path("./results") +save_path.mkdir(exist_ok=True) + +card = load_profile() # fdi_semg +protocol = [(1, 0, 0), (4, 0, 25), (4, 25, 25), (4, 25, 0), (1, 0, 0)] +result = card.simulate(protocol) + +spike_train_block = result.spikes +mn_segment = spike_train_block.segments[0] +n_motor_neurons = len(mn_segment.spiketrains) +simulation_time = sum(d for d, _, _ in protocol) * pq.s +timestep = float(card.value("drive.timestep_ms")) * pq.ms +ramp_up_end = 5.0 * pq.s # start of the hold +plateau_end = 9.0 * pq.s # end of the hold + +# Excitation (0-1) per drive time step, saved for the surface EMG example +excitation_signal = AnalogSignal(result.excitation[:, None], units=pq.dimensionless, + sampling_period=timestep.rescale(pq.s)) + +# The spike trains and the excitation are used by the surface EMG example (06) and the force example (08) +joblib.dump(spike_train_block, save_path / "trapezoid_dd_spike_trains.pkl") +joblib.dump(excitation_signal, save_path / "trapezoid_drive_pattern.pkl") + +############################################################################## +# Calculate Firing Rate Statistics +# --------------------------------- + +mn_firing_rates = np.array( + [ + mean_firing_rate(st__s.time_slice(st__s.min(), st__s.max())) + for st__s in mn_segment.spiketrains + if len(st__s) > 1 # need >=2 spikes for a rate over the spike span + ] +) + +print("Motor neurons:") +print(f"\tActive neurons: {len(mn_firing_rates)}/{n_motor_neurons}") +if len(mn_firing_rates) > 0: + print(f"\tMean firing rate: {np.mean(mn_firing_rates):.1f} ± {np.std(mn_firing_rates):.1f} pps") + print(f"\tRate range: {np.min(mn_firing_rates):.1f} - {np.max(mn_firing_rates):.1f} pps") + +############################################################################## +# Visualization +# ------------- +# +# 1. Target force and the simulated force (% MVC) +# 2. Motor neuron raster plot showing recruitment +# 3. Population firing rate over time + +fig, axes = plt.subplots(3, 1, figsize=(15, 10), sharex=True) + +t_force = np.arange(result.force_pct_mvc.size) / result.fs_force__Hz +t_target = np.cumsum([0] + [d for d, _, _ in protocol]) +target = [protocol[0][1]] + [b for _, _, b in protocol] +axes[0].plot(t_target, target, "k--", linewidth=2, label="Target") +axes[0].plot(t_force, result.force_pct_mvc, "b-", linewidth=2, label="Force") +axes[0].set_ylabel("Force (% MVC)") +axes[0].set_title("Trapezoidal Contraction (calibrated FDI drive)") +axes[0].legend() +axes[0].grid(True, alpha=0.3) + +mn_colors = plt.get_cmap("Reds")(np.linspace(0.3, 0.9, n_motor_neurons)) +active_mn_count = 0 +for i, (spiketrain, color) in enumerate(zip(mn_segment.spiketrains, mn_colors)): + if len(spiketrain) > 0: + spike_times = spiketrain.rescale(pq.s).magnitude + axes[1].scatter(spike_times, [i] * len(spike_times), c=[color], s=1.0, alpha=0.8) + active_mn_count += 1 + +axes[1].set_ylabel("Motor Neuron ID\n(Recruitment Order)") +axes[1].set_title(f"Motor Neuron Population Activity (n={active_mn_count}/{n_motor_neurons} active)") +axes[1].set_ylim(-1, n_motor_neurons) +axes[1].grid(True, alpha=0.3) + +bin_size_ms = 100 +mn_counts, bin_centers_s = population_psth(mn_segment.spiketrains, bin_size_ms * pq.ms) +mn_rates_binned = mn_counts / (bin_size_ms / 1000.0) / n_motor_neurons +axes[2].plot(bin_centers_s, mn_rates_binned, "r-", linewidth=2, label="MN Population", alpha=0.8) +axes[2].set_xlabel("Time (s)") +axes[2].set_ylabel("Population Rate (Hz)") +axes[2].set_title("Population Firing Rate Over Time") +axes[2].legend() +axes[2].grid(True, alpha=0.3) + +for ax in axes: + ax.set_xlim(0, simulation_time.rescale(pq.s).magnitude) + +plt.tight_layout() +plt.show() +############################################################################## +# Individual Motor Neuron Discharge Rates +# ---------------------------------------- +# +# Compute smoothed instantaneous firing rates for each motor neuron +# with a 400 ms Hanning window. + +print("\nComputing smoothed discharge rates per neuron...") + +# Parameters +window_ms = 400 * pq.ms # 400 ms Hanning window +dt_s = timestep.rescale(pq.s) # simulation timestep in seconds +window_samples = int(window_ms.rescale(pq.s) / dt_s) + +# Hanning window normalized to preserve rate +hanning_window = np.hanning(window_samples) +hanning_window = hanning_window / (hanning_window.sum() * dt_s) # convert to Hz + +mn_instantaneous_rates = [] +active_neuron_ids = [] + +mean_firing_rates = [] +cv_isi = [] + +for i, spiketrain in enumerate(mn_segment.spiketrains): + if len(spiketrain) > 2: + # Convert spike times to a binary spike train + t = np.arange(0, simulation_time.rescale(pq.s).magnitude, dt_s.magnitude) + spikes = np.zeros_like(t) + spike_indices = np.searchsorted(t, spiketrain.magnitude) + spikes[spike_indices[spike_indices < len(t)]] = 1 + + # Convolve with Hanning window + rate = np.convolve(spikes, hanning_window, mode="same") + mn_instantaneous_rates.append(rate) + active_neuron_ids.append(i) + + # IMPORTANT: Compute ISI/CV only during plateau phase where firing is stable + # Filter spike train to plateau phase + plateau_spiketrain = spiketrain.time_slice(ramp_up_end, plateau_end) + + # Compute mean firing rate (Hz) from plateau spikes only + plateau_duration_s = float((plateau_end - ramp_up_end).rescale(pq.s).magnitude) + mean_rate = ( + len(plateau_spiketrain) / plateau_duration_s if len(plateau_spiketrain) > 0 else 0.0 + ) + mean_firing_rates.append(mean_rate) + + # Compute CV of inter-spike intervals (plateau phase only) + if len(plateau_spiketrain) > 1: + spike_times = plateau_spiketrain.rescale(pq.s).magnitude + isis = np.diff(spike_times) + cv = np.std(isis) / np.mean(isis) if len(isis) > 1 else 0.0 + else: + cv = 0.0 + cv_isi.append(cv) + +# Population averages +pop_mean_rate = np.mean(mean_firing_rates) +pop_mean_cv = np.mean(cv_isi) +print(f"\nPopulation: Mean firing rate = {pop_mean_rate:.2f} Hz, CV = {pop_mean_cv:.2f}") + +print(f" Computed rates for {len(active_neuron_ids)} active motor neurons") + +# Create new figure for individual discharge rates +fig2, axes2 = plt.subplots(2, 1, figsize=(15, 10), sharex=True) + +# 1. Heatmap of instantaneous firing rates +if len(mn_instantaneous_rates) > 0: + # Stack rates into 2D array (neurons x time) + rates_array = np.array(mn_instantaneous_rates) + time_points = np.linspace(0, simulation_time.rescale(pq.s).magnitude, rates_array.shape[1]) + + # Plot heatmap (with origin='lower' so MU 0 is at TOP) + im = axes2[0].imshow( + rates_array, + aspect="auto", + cmap="hot", + interpolation="bilinear", + extent=[0, simulation_time.rescale(pq.s).magnitude, 0, len(active_neuron_ids)], + origin="lower", # Puts first row (MU 0) at bottom, but we'll flip with extent + vmin=0, + vmax=np.percentile(rates_array, 95), + ) + + axes2[0].set_ylabel("Motor Neuron ID\n(Recruitment Order, MU 0 at top)") + axes2[0].set_title( + "Individual Motor Neuron Discharge Rates (Smoothed with 400ms Hanning Window)" + ) + # Add colorbar + cbar = plt.colorbar(im, ax=axes2[0]) + cbar.set_label("Firing Rate (Hz)") + axes2[0].grid(False) + + # 2. Individual traces + n_to_plot = len(active_neuron_ids) + + # Use colormap for lines (gradient from blue to red showing recruitment order) + colors = plt.get_cmap("rainbow")(np.linspace(0, 1, n_to_plot)) + + for neuron_idx in range(n_to_plot): + axes2[1].plot( + time_points, + mn_instantaneous_rates[neuron_idx], + linewidth=0.8, + color=colors[neuron_idx], + label=f"MN {active_neuron_ids[neuron_idx]}" if n_to_plot <= 20 else None, + ) + + axes2[1].set_xlabel("Time (s)") + axes2[1].set_ylabel("Firing Rate (Hz)") + axes2[1].set_title(f"All Motor Neuron Discharge Rates (n={n_to_plot})") + + # Only show legend if there are few neurons + if n_to_plot <= 20: + axes2[1].legend(loc="upper right", ncol=3, fontsize=6) + + axes2[1].grid(True, alpha=0.3) + axes2[1].set_xlim(0, simulation_time.rescale(pq.s).magnitude) + axes2[1].set_ylim(0, np.max(rates_array) * 1.1) + +plt.tight_layout() +plt.show() + +print("\n[DONE] Simulation complete with discharge rate analysis!") + +# mkdocs_gallery_thumbnail_path = "gallery_thumbs/03_simulate_spike_trains_descending_drive.png" diff --git a/examples/01_basic/06_simulate_surface_emg.py b/examples/01_basic/06_simulate_surface_emg.py index d711f6bb..403f34ad 100644 --- a/examples/01_basic/06_simulate_surface_emg.py +++ b/examples/01_basic/06_simulate_surface_emg.py @@ -36,8 +36,15 @@ save_path = Path("./results") -spike_train__Block: SPIKE_TRAIN__Block = joblib.load(save_path / "trapezoid_dd_spike_trains.pkl") -excitation__AnalogSignal = joblib.load(save_path / "trapezoid_drive_pattern.pkl") # drive excitation (0-1) +# The canonical Example 03 uses Jaxley and writes backend-qualified artifacts so +# that its output can coexist with the NEURON validation run. Keep the backend +# explicit here: EMG and excitation must come from the same Example 03 run. +spike_train__Block: SPIKE_TRAIN__Block = joblib.load( + save_path / "trapezoid_dd_spike_trains_jaxley.pkl" +) +excitation__AnalogSignal = joblib.load( + save_path / "trapezoid_drive_pattern_jaxley.pkl" +) # drive excitation (0-1) surface_emg: simulator.SurfaceEMG = joblib.load(save_path / "surface_emg.pkl") ############################################################################## diff --git a/examples/01_basic/08_simulate_force.py b/examples/01_basic/08_simulate_force.py index 62123261..4b8290b4 100644 --- a/examples/01_basic/08_simulate_force.py +++ b/examples/01_basic/08_simulate_force.py @@ -1,34 +1,28 @@ """ -Force Generation -================ +Force Generation — Jaxley Spike Trains +======================================== -After generating **spike trains** from motor neuron pools, the next step is to simulate the **force output** -produced by the muscle. **MyoGen** provides a comprehensive force model based on the classic -Fuglevand et al. (1993) approach. +This example converts the calibrated Jaxley motor-unit spike trains produced by +Example 03 into muscle force. The force model is simulator-independent, so the +Jaxley and NEURON versions differ only in which saved spike-train artifact they +load. -!!! note - The **force model** converts spike trains into force by simulating individual motor unit twitches - and their summation. Each motor unit has unique twitch characteristics (amplitude, duration) - that depend on its recruitment threshold. - -The force model includes: - -* **Motor unit twitch characteristics**: Peak force and contraction time based on recruitment thresholds -* **Nonlinear gain modulation**: Force scaling based on inter-pulse intervals (discharge rate) -* **Temporal summation**: Realistic force buildup from overlapping twitches -* **Physiological scaling**: Realistic force ranges and dynamics +The model includes motor-unit twitch scaling, nonlinear force-frequency gain, +and temporal summation. Recruitment thresholds and the MVC normalization come +from the same ``fdi_semg`` profile card used by Example 03. References ---------- -[1] Fuglevand, A. J., Winter, D. A., & Patla, A. E. (1993). Models of recruitment and rate coding organization in motor-unit pools. J Neurophysiol, 70(6), 2470–2488. https://doi.org/10.1152/jn.1993.70.6.2470. +[1] Fuglevand AJ, Winter DA, Patla AE (1993). Models of recruitment and rate +coding organization in motor-unit pools. J Neurophysiol 70:2470–2488. """ +# %% ############################################################################## # Import Libraries # ---------------- -# %% from pathlib import Path import joblib @@ -40,122 +34,104 @@ plt.style.use("fivethirtyeight") +save_path = Path("./results") +save_path.mkdir(exist_ok=True) + ############################################################################## -# Create Force Model -# ------------------ -# -# The [`ForceModel`][myogen.simulator.ForceModel] calculates each motor unit's twitch (peak force and contraction -# time) from its recruitment threshold. Its main parameters are: -# -# - ``recruitment_thresholds``: Motor unit recruitment thresholds -# - ``recording_frequency__Hz``: Sampling frequency of the force -# - ``force_frequency_shape``: Shape of the force-frequency (gain) curve -# -# The ``fdi_semg`` profile card builds it with the thresholds and the human force-frequency curve -# (Macefield 1996, Fuglevand 1999) the FDI drive was calibrated with. -# ``load_profile("fdi_semg.published").force_model()`` uses the curve of Fuglevand 1993 instead. +# Create the Shared Force Model +# ----------------------------- card = load_profile() # fdi_semg force_model = card.force_model() -# Display force model statistics print("Force model statistics:") print(f"\tNumber of motor units: {force_model._number_of_neurons}") print(f"\tRecruitment ratio: {force_model._recruitment_ratio:.1f}") print( - f"\tPeak force range: {force_model.peak_twitch_forces__unitless[0]:.3f} - {force_model.peak_twitch_forces__unitless[-1]:.3f}" + "\tPeak force range: " + f"{force_model.peak_twitch_forces__unitless[0]:.3f} - " + f"{force_model.peak_twitch_forces__unitless[-1]:.3f}" ) print( - f"\tContraction time range: {force_model.contraction_times__samples[0]:.1f} - {force_model.contraction_times__samples[-1]:.1f} samples" + "\tContraction time range: " + f"{force_model.contraction_times__samples[0]:.1f} - " + f"{force_model.contraction_times__samples[-1]:.1f} samples" ) ############################################################################## -# Visualize Twitch Parameter Assignment -# -------------------------------------- -# -# The force model assigns twitch parameters (peak force and contraction time) -# to each motor unit based on its recruitment threshold. Let's visualize this -# relationship for a subset of motor units. -plt.figure(figsize=(8, 12)) - -ax1 = plt.subplot(2, 1, 1) +# Visualize Twitch Assignment +# ---------------------------- + +fig, axes = plt.subplots(2, 1, figsize=(8, 12)) plot_twitch_parameter_assignment( - force_model, ax1, [10, 20, 40], flip_x=True, apply_default_formatting=True + force_model, + axes[0], + [10, 20, 40], + flip_x=True, + apply_default_formatting=True, ) -ax1.set_title("Twitch Parameter Assignment") - -ax2 = plt.subplot(2, 1, 2) -plot_twitches(force_model, ax2, [10, 20, 40], apply_default_formatting=True) -ax2.set_title("Motor Unit Twitches") - +axes[0].set_title("Twitch Parameter Assignment") +plot_twitches(force_model, axes[1], [10, 20, 40], apply_default_formatting=True) +axes[1].set_title("Motor Unit Twitches") plt.tight_layout() +plt.savefig(save_path / "ex08_jaxley_twitches.png", dpi=150, bbox_inches="tight") plt.show() ############################################################################## -# Load Spike Trains -# ----------------- -# -# The spike trains of the 25 % MVC contraction from example 03 (calibrated FDI drive). - -save_path = Path("./results") -spike_train__Block = joblib.load(save_path / "trapezoid_dd_spike_trains.pkl") - -############################################################################## -# Generate Force Output -# --------------------- -# -# Now we can use the force model to convert the spike trains into force output. -# The model will simulate individual motor unit twitches and their temporal -# summation to produce the total muscle force. - -# Generate force from spike trains, and express it in % MVC with the card's MVC (from its MVC ramp) -force_output = force_model.generate_force(spike_train__Block=spike_train__Block) -force_pct_mvc = 100.0 * force_output.magnitude[:, 0] / card.value("drive.mvc_model_units") - -# Add some realistic noise to the force signal -noise_level = 0.015 # 0.15% of mean force -noisy_force = force_output.magnitude[:, 0] + np.random.randn( - len(force_output.magnitude[:, 0]) -) * noise_level * np.mean(force_output.magnitude[:, 0]) +# Load Jaxley Example 03 Spike Trains +# ------------------------------------ + +spike_path = save_path / "trapezoid_dd_spike_trains_jaxley.pkl" +if not spike_path.exists(): + raise FileNotFoundError( + f"{spike_path} does not exist. Run 03_simulate_spike_trains_descending_drive.py first." + ) + +spike_train_block = joblib.load(spike_path) +n_spike_trains = len(spike_train_block.segments[0].spiketrains) +if n_spike_trains != force_model._number_of_neurons: + raise ValueError( + "Example 03 spike count and force-model pool size differ: " + f"{n_spike_trains} != {force_model._number_of_neurons}" + ) ############################################################################## -# Visualize Force Output -# ---------------------- -# -# Let's plot the generated force alongside the input current to see how -# the muscle responds to the electrical stimulation. +# Generate Force +# -------------- -plt.figure(figsize=(8, 10)) +force_output = force_model.generate_force(spike_train__Block=spike_train_block) +force = force_output.magnitude[:, 0] +force_pct_mvc = 100.0 * force / float(card.value("drive.mvc_model_units")) -ax1 = plt.subplot(2, 1, 1) +rng = np.random.default_rng(42) +noise_level = 0.015 +noisy_force = force + rng.normal(size=force.size) * noise_level * np.mean(force) -ax1.plot( - force_output.times.rescale("s"), - force_pct_mvc, - linewidth=2, - label="Clean Force", -) +print("Jaxley Example 03 force:") +print(f"\tPeak: {force_pct_mvc.max():.1f} % MVC") +print(f"\tMean: {force_pct_mvc.mean():.1f} % MVC") -ax1.set_ylabel("Force (% MVC)") -ax1.set_title("Simulated Force Output") -ax1.grid(True, alpha=0.3) - -# Plot noisy force (more realistic) -ax2 = plt.subplot(2, 1, 2) -ax2.plot( - force_output.times.rescale("s"), - noisy_force, - linewidth=1, - alpha=0.8, - label="Noisy Force", -) +############################################################################## +# Visualize Force Output +# ---------------------- -ax2.set_xlabel("Time (s)") -ax2.set_ylabel("Force (a.u.)") -ax2.set_title("Realistic Force Output (with noise)") -ax2.grid(True, alpha=0.3) +fig, axes = plt.subplots(2, 1, figsize=(8, 10), sharex=True) +time_s = force_output.times.rescale("s") +axes[0].plot(time_s, force_pct_mvc, linewidth=2, label="Jaxley-derived force") +axes[0].set_ylabel("Force (% MVC)") +axes[0].set_title("Calibrated Descending-Drive Force") +axes[0].grid(True, alpha=0.3) +axes[0].legend() + +axes[1].plot(time_s, noisy_force, linewidth=1, alpha=0.8, label="Noisy force") +axes[1].set_xlabel("Time (s)") +axes[1].set_ylabel("Force (model units)") +axes[1].set_title("Force with Measurement Noise") +axes[1].grid(True, alpha=0.3) +axes[1].legend() plt.tight_layout() +plt.savefig(save_path / "ex08_jaxley_force.png", dpi=150, bbox_inches="tight") plt.show() # mkdocs_gallery_thumbnail_path = "gallery_thumbs/08_simulate_force.png" diff --git a/examples/01_basic/08_simulate_force_neuron.py b/examples/01_basic/08_simulate_force_neuron.py new file mode 100644 index 00000000..260bcfe2 --- /dev/null +++ b/examples/01_basic/08_simulate_force_neuron.py @@ -0,0 +1,137 @@ +""" +Force Generation — NEURON Reference Spike Trains +================================================== + +This is the NEURON-reference counterpart of ``08_simulate_force.py``. It loads +the calibrated spike trains produced by +``03_simulate_spike_trains_descending_drive_neuron.py`` and applies the same +simulator-independent force model used for the Jaxley result. + +Keeping force generation identical isolates the effect of the motor-neuron +backend: differences in the force traces originate from the saved discharge +patterns, not from different twitch models. + +References +---------- + +[1] Fuglevand AJ, Winter DA, Patla AE (1993). Models of recruitment and rate +coding organization in motor-unit pools. J Neurophysiol 70:2470–2488. +""" + +# %% +############################################################################## +# Import Libraries +# ---------------- + +from pathlib import Path + +import joblib +import matplotlib.pyplot as plt +import numpy as np + +from myogen import load_profile +from myogen.utils.plotting.force import plot_twitch_parameter_assignment, plot_twitches + +plt.style.use("fivethirtyeight") + +save_path = Path("./results") +save_path.mkdir(exist_ok=True) + +############################################################################## +# Create the Shared Force Model +# ----------------------------- + +card = load_profile() # fdi_semg +force_model = card.force_model() + +print("Force model statistics:") +print(f"\tNumber of motor units: {force_model._number_of_neurons}") +print(f"\tRecruitment ratio: {force_model._recruitment_ratio:.1f}") +print( + "\tPeak force range: " + f"{force_model.peak_twitch_forces__unitless[0]:.3f} - " + f"{force_model.peak_twitch_forces__unitless[-1]:.3f}" +) +print( + "\tContraction time range: " + f"{force_model.contraction_times__samples[0]:.1f} - " + f"{force_model.contraction_times__samples[-1]:.1f} samples" +) + +############################################################################## +# Visualize Twitch Assignment +# ---------------------------- + +fig, axes = plt.subplots(2, 1, figsize=(8, 12)) +plot_twitch_parameter_assignment( + force_model, + axes[0], + [10, 20, 40], + flip_x=True, + apply_default_formatting=True, +) +axes[0].set_title("Twitch Parameter Assignment") +plot_twitches(force_model, axes[1], [10, 20, 40], apply_default_formatting=True) +axes[1].set_title("Motor Unit Twitches") +plt.tight_layout() +plt.savefig(save_path / "ex08_neuron_twitches.png", dpi=150, bbox_inches="tight") +plt.show() + +############################################################################## +# Load NEURON Example 03 Spike Trains +# ------------------------------------ + +spike_path = save_path / "trapezoid_dd_spike_trains.pkl" +if not spike_path.exists(): + raise FileNotFoundError( + f"{spike_path} does not exist. Run 03_simulate_spike_trains_descending_drive_neuron.py first." + ) + +spike_train_block = joblib.load(spike_path) +n_spike_trains = len(spike_train_block.segments[0].spiketrains) +if n_spike_trains != force_model._number_of_neurons: + raise ValueError( + "Example 03 spike count and force-model pool size differ: " + f"{n_spike_trains} != {force_model._number_of_neurons}" + ) + +############################################################################## +# Generate Force +# -------------- + +force_output = force_model.generate_force(spike_train__Block=spike_train_block) +force = force_output.magnitude[:, 0] +force_pct_mvc = 100.0 * force / float(card.value("drive.mvc_model_units")) + +rng = np.random.default_rng(42) +noise_level = 0.015 +noisy_force = force + rng.normal(size=force.size) * noise_level * np.mean(force) + +print("NEURON Example 03 force:") +print(f"\tPeak: {force_pct_mvc.max():.1f} % MVC") +print(f"\tMean: {force_pct_mvc.mean():.1f} % MVC") + +############################################################################## +# Visualize Force Output +# ---------------------- + +fig, axes = plt.subplots(2, 1, figsize=(8, 10), sharex=True) +time_s = force_output.times.rescale("s") +axes[0].plot(time_s, force_pct_mvc, linewidth=2, label="NEURON-derived force") +axes[0].set_ylabel("Force (% MVC)") +axes[0].set_title("Calibrated Descending-Drive Force") +axes[0].grid(True, alpha=0.3) +axes[0].legend() + +axes[1].plot(time_s, noisy_force, linewidth=1, alpha=0.8, label="Noisy force") +axes[1].set_xlabel("Time (s)") +axes[1].set_ylabel("Force (model units)") +axes[1].set_title("Force with Measurement Noise") +axes[1].grid(True, alpha=0.3) +axes[1].legend() + +plt.tight_layout() +plt.savefig(save_path / "ex08_neuron_force.png", dpi=150, bbox_inches="tight") +plt.show() + +# mkdocs_gallery_thumbnail_path = "gallery_thumbs/08_simulate_force_neuron.png" diff --git a/examples/01_basic/09_simulate_intramuscular_emg.py b/examples/01_basic/09_simulate_intramuscular_emg.py index 17182c91..307397ed 100644 --- a/examples/01_basic/09_simulate_intramuscular_emg.py +++ b/examples/01_basic/09_simulate_intramuscular_emg.py @@ -36,7 +36,7 @@ import seaborn as sns from myogen import simulator -from myogen.utils.types import CURRENT__AnalogSignal, SPIKE_TRAIN__Block +from myogen.utils.types import SPIKE_TRAIN__Block plt.style.use("fivethirtyeight") @@ -148,13 +148,15 @@ plt.show() ############################################################################## -# Load Input Currents and Spike Trains -# ----------------------- +# Load Drive Excitation and Jaxley Spike Trains +# ---------------------------------------------- # -spike_train__Block: SPIKE_TRAIN__Block = joblib.load(save_path / "trapezoid_dd_spike_trains.pkl") -input_current__AnalogSignal: CURRENT__AnalogSignal = joblib.load( - save_path / "trapezoid_drive_pattern.pkl" +spike_train__Block: SPIKE_TRAIN__Block = joblib.load( + save_path / "trapezoid_dd_spike_trains_jaxley.pkl" +) +excitation__AnalogSignal = joblib.load( + save_path / "trapezoid_drive_pattern_jaxley.pkl" ) ############################################################################## @@ -183,32 +185,32 @@ # Visualize Intramuscular EMG Results # ---------------------------------- # -# Compare the **intramuscular EMG** signal with the input current. +# Compare the **intramuscular EMG** signal with the descending-drive excitation. # # !!! note -# In this example, only **three motor units** (small, medium, large) contribute -# to the EMG signal. This demonstrates how individual MUAPs sum to create the -# composite EMG signal, and how different motor unit sizes affect signal amplitude. +# The configured subset spans the recruitment range. This demonstrates how +# individual MUAPs sum into the composite recording and how motor-unit size +# affects signal amplitude. plt.figure(figsize=(12, 6)) # Get the signals - use first electrode from the Block iemg_signal = noisy_emg_signals__Block.segments[0].analogsignals[0][:, 0] -current_signal = input_current__AnalogSignal[:, 0] # First current pool +excitation_signal = excitation__AnalogSignal[:, 0] # Create time axes emg_time = iemg_signal.times.rescale("s").magnitude -current_time = input_current__AnalogSignal.times.rescale("s").magnitude +excitation_time = excitation__AnalogSignal.times.rescale("s").magnitude -# Normalize current between 0 and 1 -current_normalized = ( - (current_signal - np.min(current_signal)) - / (np.max(current_signal) - np.min(current_signal)) +# Normalize excitation between 0 and 1 +excitation_normalized = ( + (excitation_signal - np.min(excitation_signal)) + / (np.max(excitation_signal) - np.min(excitation_signal)) * np.max(iemg_signal.magnitude) ) plt.plot(emg_time, iemg_signal.magnitude, label="Intramuscular EMG", linewidth=1) -plt.plot(current_time, current_normalized, label="Input Current", linewidth=1) +plt.plot(excitation_time, excitation_normalized, label="Drive excitation", linewidth=1) plt.xlabel("Time (s)") plt.ylabel("Amplitude (mV)") diff --git a/examples/01_basic/10_extract_neuron_parameters.py b/examples/01_basic/10_extract_neuron_parameters.py index 4f22ee85..e588ece2 100644 --- a/examples/01_basic/10_extract_neuron_parameters.py +++ b/examples/01_basic/10_extract_neuron_parameters.py @@ -1,31 +1,44 @@ """ -Extract Neuron Parameters -========================= +Extract Neuron Parameters - Jaxley Backend +========================================== This example demonstrates how to extract electrophysiological parameters from -MyoGen motor neurons using the Fuglevand recruitment model. +MyoGen motor neurons using the **Jaxley** (JAX-based) simulator instead of NEURON. + +**Channels Used** (NERLab — matches production NEURON model): + - Soma: napp (Na fast + Na persistent + Kfast + Kslow + leak) + - Dendrite: caL (L-type Ca, no inactivation + leak) + +**Voltage convention** (NERLab / original 1952-HH frame): V_rest ≈ 0 mV, +ENa = +120 mV, EK = -10 mV, spike peaks ≈ +90 mV. Spike threshold below +is +50 mV in this frame. **Parameters Extracted**: - **Vhold**: Resting membrane potential (soma) -- **Rin**: Input resistance -- **tau**: Membrane time constant -- **Ir**: Rheobase (minimum current for action potential) +- **Rin**: Input resistance (small hyperpolarizing steps, passive range) +- **tau**: Membrane time constant (small sustained step, charging transient) +- **Ir**: Rheobase (0.1 nA resolution, 50 ms pulse) - **AP**: Action potential amplitude - **AHP**: Afterhyperpolarization depth - **AHPdur**: Full AHP duration -- **FI_gain**: Frequency-current gain (slope of F-I curve) +- **FI_gain**: Frequency-current gain (rheobase-relative, ascending linear range) + +.. note:: + **Protocol differences from NEURON**: The stimulus protocols here are designed + to stay within the linear/passive membrane range of the Jaxley cable model. + Rin uses small (-0.5 to -2.5 nA) steps; tau uses a small (-1 nA) sustained + step and fits the charging onset; rheobase uses 0.1 nA resolution. + These are not guaranteed to match NEURON's exact protocol step-by-step, but + they measure the same biophysical quantities in a cable-appropriate way. **Usage**: # Extract from default 5 neurons - python 09_extract_neuron_parameters.py + python 10_extract_neuron_parameters.py # Extract from custom number of neurons - python 09_extract_neuron_parameters.py --n-neurons 10 - -**Note**: For advanced features (parallel processing, model comparison, slope scan), -see `sandbox/neuron_parameter_extraction/README.md` + python 10_extract_neuron_parameters.py --n-neurons 10 """ # %% @@ -36,6 +49,9 @@ import logging import os +import sys +from contextlib import contextmanager +from io import StringIO os.environ["MPLBACKEND"] = "Agg" if "DISPLAY" in os.environ: @@ -43,381 +59,313 @@ from pathlib import Path +import jax.numpy as jnp +import jaxley as jx import matplotlib.pyplot as plt import numpy as np import pandas as pd import seaborn as sns -from neuron import h import myogen -from myogen import simulator -from myogen.simulator.neuron.populations import AlphaMN__Pool -from myogen.utils.nmodl import load_nmodl_mechanisms +from myogen.simulator import RecruitmentThresholds +from myogen.simulator.jaxley.populations import AlphaMN__Pool + +# Suppress Jaxley verbose output +logging.getLogger("jaxley").setLevel(logging.WARNING) +logging.getLogger("jax").setLevel(logging.WARNING) + + +@contextmanager +def suppress_stdout(): + """Context manager to suppress stdout (for Jaxley verbose messages).""" + old_stdout = sys.stdout + sys.stdout = StringIO() + try: + yield + finally: + sys.stdout = old_stdout + # Simple plotting style plt.style.use("seaborn-v0_8-darkgrid") sns.set_context("paper", font_scale=1.2) -# Setup NEURON environment -h.load_file("stdrun.hoc") - # Setup logging logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) -############################################################################## -# Core Parameter Extraction Functions -# ------------------------------------ - - -def fi0_multicompartment(sections: list, voltages: list) -> None: - """Initialize voltages for multiple compartments.""" - for sec, v in zip(sections, voltages): - sec.v = v +# Simulation parameters for NERLab cells (1952-HH voltage frame). +DT_MS = 0.025 # Time step (ms) +SPIKE_THRESHOLD_MV = 50.0 # NERLab APs cross +50 mV reliably; +90 mV at peak - -def get_vhold__mV(cell, sections: list, voltages: list, tstop__ms: float = 500.0) -> float: - """ - Measure resting membrane potential for soma. - - Returns - ------- - float - Steady-state voltage in soma (mV). +############################################################################## +# Core Parameter Extraction Functions (Batched) +# ---------------------------------------------- +# +# All measurement functions operate on a pre-built jx.Network containing +# ALL neurons. Each protocol step is one jx.integrate() call that runs all +# neurons in parallel, instead of n_neurons serial calls per step. +# +# Call counts per protocol (n = number of neurons): +# Vhold : 1 call (was n) +# Rin : 5 calls (was 5n) +# tau : 1 call (was n) +# Rh : ~11 calls batched binary search (was ~100n linear scan) +# AP/AHP : 1 call (was variable × n) +# F-I : 12 calls (was 12n) +# Total : ~31 calls regardless of n (vs ~130 × n before) + + +def _run_batch(net, n_neurons, current_arrays, t_max_ms: float) -> np.ndarray: """ - h.tstop = tstop__ms - h.dt = 0.0125 - h.celsius = 36 - - _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + Run all neurons simultaneously with per-neuron current waveforms. - vsoma = h.Vector() - vsoma.record(cell.soma(0.5)._ref_v) - - h.finitialize() - h.run() - - return vsoma.to_python()[-1] - - -def get_rin__MOhm( - cell, - vhold__mV: float, - sections: list, - amplitudes__nA: list[float] = None, -) -> float: - """ - Measure input resistance using current steps. + Resets voltage/states, re-registers stimuli, calls jx.integrate once. Returns ------- - float - Input resistance in MegaOhms. + np.ndarray, shape (n_neurons, n_timesteps) """ - if amplitudes__nA is None: - amplitudes__nA = [-5.0, -4.0, -3.0, -2.0, -1.0] - - h.tstop = 1000.0 - h.dt = 0.0125 - h.celsius = 36 + net.set("v", 0.0) # NERLab resting potential (1952-HH frame) + net.init_states() + net.delete_stimuli() + for i, arr in enumerate(current_arrays): + net.cell(i).branch(0).loc(0.5).stimulate(jnp.array(arr)) + voltages = jx.integrate(net, delta_t=DT_MS, t_max=t_max_ms) + return np.array(voltages) # (n_neurons, n_timesteps) - voltages = [vhold__mV] * len(sections) - _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) - vsoma = h.Vector() - vsoma.record(cell.soma(0.5)._ref_v) +def _run_batch_same(net, n_neurons, current_array, t_max_ms: float) -> np.ndarray: + """Same current waveform to all neurons. Returns (n_neurons, n_timesteps).""" + return _run_batch(net, n_neurons, [current_array] * n_neurons, t_max_ms) - delta_v__mV = [] - for amp__nA in amplitudes__nA: - stim = h.IClamp(cell.soma(0.5)) - stim.amp = amp__nA - stim.dur = 1000.0 - stim.delay = 0 - h.finitialize() - h.run() +def get_vholds(net, n_neurons: int) -> np.ndarray: + """Resting potential for all neurons — 1 jx.integrate call.""" + tstop = 500.0 + current = np.zeros(int(tstop / DT_MS)) + V = _run_batch_same(net, n_neurons, current, tstop) + tail = max(1, int(V.shape[1] * 0.1)) + return np.array([float(np.mean(V[i, -tail:])) for i in range(n_neurons)]) - vpeak__mV = vsoma.to_python()[-1] - dv = vpeak__mV - vhold__mV - delta_v__mV.append(dv) - del stim - coefficients = np.polyfit(amplitudes__nA, delta_v__mV, 1) - rin__MOhm = coefficients[0] - return rin__MOhm - - -def get_time_constant__ms(cell, vhold__mV: float, sections: list) -> float: +def get_rins(net, n_neurons: int, vholds: np.ndarray) -> np.ndarray: """ - Measure membrane time constant from exponential decay. + Input resistance for all neurons — 5 jx.integrate calls. - Returns - ------- - float - Membrane time constant in milliseconds. + Uses small hyperpolarizing steps (-0.5 to -2.5 nA, 300 ms) to stay in + the passive linear range and avoid slow conductance activation. """ - h.tstop = 200.0 - h.dt = 0.0125 - h.celsius = 36 - - voltages = [vhold__mV] * len(sections) - _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) - - vsoma = h.Vector() - t = h.Vector() - vsoma.record(cell.soma(0.5)._ref_v) - t.record(h._ref_t) - - stim = h.IClamp(cell.soma(0.5)) - stim.amp = -20.0 # nA - stim.dur = 1.0 # ms - stim.delay = 5.0 # ms - - h.finitialize() - h.run() - - vsoma_array = vsoma.as_numpy() - t_array = t.as_numpy() - - v_min__mV = np.min(vsoma_array) - v_min_index = np.where(vsoma_array == v_min__mV)[0][0] - - recovery_indices = np.where(vsoma_array[v_min_index:] > vhold__mV - 0.1)[0] - if len(recovery_indices) > 0: - v_end_index = recovery_indices[0] + v_min_index - else: - v_end_index = len(vsoma_array) - 1 - - t_fit = t_array[v_min_index:v_end_index] - v_fit = vsoma_array[v_min_index:v_end_index] - log_v_fit = np.log(np.abs(v_fit - vhold__mV)) - - coefficients = np.polyfit(t_fit, log_v_fit, 1) - tau__ms = -1.0 / coefficients[0] - return tau__ms - - -def get_rheobase__nA( - cell, vhold__mV: float, sections: list, max_iterations: int = 500 -) -> float: + amplitudes = [-0.5, -1.0, -1.5, -2.0, -2.5] + tstop = 300.0 + n_t = int(tstop / DT_MS) + tail = max(1, int(n_t * 0.2)) + delta_vs = [[] for _ in range(n_neurons)] + + for amp in amplitudes: + current = np.ones(n_t) * amp + V = _run_batch_same(net, n_neurons, current, tstop) + for i in range(n_neurons): + v_ss = float(np.mean(V[i, -tail:])) + delta_vs[i].append(v_ss - vholds[i]) + + rins = [] + for i in range(n_neurons): + try: + coeffs = np.polyfit(amplitudes, delta_vs[i], 1) + rins.append(coeffs[0]) + except Exception: + rins.append(np.nan) + return np.array(rins) + + +def get_taus(net, n_neurons: int, vholds: np.ndarray) -> np.ndarray: """ - Find minimum current required to elicit an action potential. + Membrane time constant — brief impulse protocol (matches NEURON ex10). - Returns - ------- - float - Rheobase current in nanoamperes, or np.nan if no spike found. + Injects -20 nA for 1 ms, then fits exponential recovery from the most + hyperpolarized point back toward vhold. This isolates the fast membrane RC + tau and avoids contamination from slow currents (Gh sag, MAHP) that bias + sustained-step fits toward much longer time constants. """ - h.tstop = 100.0 - h.dt = 0.0125 - h.celsius = 36 - - voltages = [vhold__mV] * len(sections) - _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) - - vsoma = h.Vector() - vsoma.record(cell.soma(0.5)._ref_v) - - stim = h.IClamp(cell.soma(0.5)) - stim.dur = 50.0 - stim.delay = 0 - - spike_threshold__mV = vhold__mV + 40.0 - - amp__nA = 0.1 - iteration = 0 - while iteration < max_iterations: - stim.amp = amp__nA - - h.finitialize() - h.run() - - vpeak__mV = np.max(vsoma.as_numpy()) - if vpeak__mV >= spike_threshold__mV: - return amp__nA - - amp__nA += 0.1 - iteration += 1 - - logger.warning(f"No spike found after {max_iterations} iterations") - return np.nan - - -def get_ap_and_ahp(cell, vhold__mV: float, sections: list) -> dict: + step_amp = -20.0 # nA — brief impulse + pulse_dur = 1.0 # ms + delay_ms = 5.0 # ms — settling before pulse + tstop = 200.0 # ms + n_t = int(tstop / DT_MS) + delay_pts = int(delay_ms / DT_MS) + pulse_pts = int(pulse_dur / DT_MS) + + current = np.zeros(n_t) + current[delay_pts:delay_pts + pulse_pts] = step_amp + V = _run_batch_same(net, n_neurons, current, tstop) + + t_array = np.arange(n_t) * DT_MS + taus = [] + for i in range(n_neurons): + v = V[i] + v_min_idx = int(np.argmin(v)) + rec = np.where(v[v_min_idx:] > vholds[i] - 0.1)[0] + v_end_idx = rec[0] + v_min_idx if len(rec) > 0 else n_t - 1 + t_fit = t_array[v_min_idx:v_end_idx] + v_fit = v[v_min_idx:v_end_idx] + if len(t_fit) < 5: + taus.append(np.nan) + continue + log_v = np.log(np.maximum(np.abs(v_fit - vholds[i]), 1e-10)) + try: + coeffs = np.polyfit(t_fit, log_v, 1) + tau = -1.0 / coeffs[0] + taus.append(max(tau, 0.1) if tau > 0 else np.nan) + except Exception: + taus.append(np.nan) + return np.array(taus) + + +def get_rheobases(net, n_neurons: int, vholds: np.ndarray, + resolution: float = 0.1) -> np.ndarray: """ - Measure action potential and afterhyperpolarization characteristics. + Rheobase for all neurons via batched binary search — ~11 jx.integrate calls total. - Returns - ------- - dict - Contains 'AP__mV', 'AHP__mV', 'AHPdur__ms' + At each search step, ALL neurons are run simultaneously with their individual + current amplitudes (different lo/hi brackets per neuron). This is + O(log₂(150/0.1)) ≈ 11 batched calls regardless of n_neurons. """ - h.tstop = 900.0 - h.dt = 0.0125 - h.celsius = 36 - - voltages = [vhold__mV] * len(sections) - _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) - - vsoma = h.Vector() - t = h.Vector() - vsoma.record(cell.soma(0.5)._ref_v) - t.record(h._ref_t) - - stim = h.IClamp(cell.soma(0.5)) - stim.dur = 0.5 - stim.delay = 5.0 - - amp__nA = 35.0 - vpeak__mV = vhold__mV - spike_threshold__mV = vhold__mV + 40.0 - - while vpeak__mV < spike_threshold__mV: - amp__nA += 10.0 - stim.amp = amp__nA - - h.finitialize() - h.run() - - vsoma_array = vsoma.as_numpy() - vpeak__mV = np.max(vsoma_array) - vvalley__mV = np.min(vsoma_array) - - if amp__nA > 500.0: - logger.warning("Current exceeded 500 nA without spike") - return {"AP__mV": np.nan, "AHP__mV": np.nan, "AHPdur__ms": np.nan} - - ap__mV = vpeak__mV - vhold__mV - ahp__mV = vhold__mV - vvalley__mV - - t_array = t.as_numpy() - peak_index = np.where(vsoma_array == vpeak__mV)[0][0] - valley_index = np.where(vsoma_array == vvalley__mV)[0][0] - recovery_indices = np.where(vsoma_array[valley_index:] > vhold__mV - 0.15)[0] - - if len(recovery_indices) > 0: - recovery_index = recovery_indices[0] + valley_index - ahp_dur__ms = t_array[recovery_index] - t_array[peak_index] - else: - ahp_dur__ms = np.nan - - return {"AP__mV": ap__mV, "AHP__mV": ahp__mV, "AHPdur__ms": ahp_dur__ms} + tstop = 100.0 + pulse_dur = 50.0 + n_t = int(tstop / DT_MS) + pulse_pts = int(pulse_dur / DT_MS) + spike_thresholds = vholds + 40.0 + + def _spiked_batch(amps: np.ndarray) -> np.ndarray: + """Run all neurons with per-neuron amplitudes; return bool array.""" + currents = [] + for amp in amps: + c = np.zeros(n_t) + c[:pulse_pts] = amp + currents.append(c) + V = _run_batch(net, n_neurons, currents, tstop) + return np.array([np.max(V[i]) >= spike_thresholds[i] for i in range(n_neurons)]) + + # Phase 1: exponential scan to bracket each neuron's rheobase + lo = np.zeros(n_neurons) + hi = np.ones(n_neurons) * 0.5 + found = np.zeros(n_neurons, dtype=bool) + + while not np.all(found) and np.all(hi[~found] <= 150.0): + spiked = _spiked_batch(hi) + found = found | spiked + if np.all(found): + break + lo[~spiked] = hi[~spiked] + hi[~spiked] *= 2.0 + + no_spike = hi > 150.0 + if np.any(no_spike & ~found): + logger.warning(f"No spike found for neuron(s): {np.where(no_spike & ~found)[0]}") + + # Phase 2: binary search to resolution within each bracket + n_iters = int(np.ceil(np.log2(np.max(hi - lo + 1e-9) / resolution))) + 2 + for _ in range(n_iters): + if np.max(hi - lo) <= resolution: + break + mid = (lo + hi) / 2.0 + spiked = _spiked_batch(mid) + hi[spiked] = mid[spiked] + lo[~spiked] = mid[~spiked] + + result = np.round(hi / resolution) * resolution + result[no_spike & ~found] = np.nan + return result -def get_fi_gain(cell, vhold__mV: float, sections: list, rheobase__nA: float) -> float: +def get_ap_ahp(net, n_neurons: int, vholds: np.ndarray, + rheobases: np.ndarray) -> list[dict]: """ - Measure frequency-current (F-I) gain. - - Returns - ------- - float - F-I gain in Hz/nA. + AP amplitude and AHP — brief intense impulse for guaranteed single spike. + + 1 ms pulse at 100 nA delivers ~100 pC of charge — enough to depolarize any + physiological MN past threshold regardless of membrane time constant. + This eliminates the failure mode (~46% NaN with the previous rheobase+5 nA + over 10 ms protocol) where slow-membrane cells failed to reach threshold + within the brief pulse. The brief impulse also reliably elicits a single + AP rather than a train, simplifying AHP analysis. """ - if np.isnan(rheobase__nA): - return np.nan - - h.tstop = 3000.0 - h.dt = 0.0125 - h.celsius = 36 - - voltages = [vhold__mV] * len(sections) - _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) - - stim = h.IClamp(cell.soma(0.5)) - stim.dur = 3000.0 - stim.delay = 0 - - current_levels = np.arange(5.0, 30.5, 5.0) - firing_rates = [] - - for current__nA in current_levels: - stim.amp = current__nA - - apc = h.APCount(cell.soma(0.5)) - apc.thresh = 50 + tstop = 900.0 + pulse_dur = 1.0 # ms — brief impulse + pulse_amp = 100.0 # nA — well above all rheobases + delay = 5.0 # ms + n_t = int(tstop / DT_MS) + delay_pts = int(delay / DT_MS) + pulse_pts = int(pulse_dur / DT_MS) + spike_thresholds = vholds + 40.0 - h.finitialize() - h.run() + current = np.zeros(n_t) + current[delay_pts:delay_pts + pulse_pts] = pulse_amp + V = _run_batch_same(net, n_neurons, current, tstop) - spike_count = apc.n - firing_rate__Hz = (spike_count / 2500.0) * 1000.0 - firing_rates.append(firing_rate__Hz) + t_array = np.arange(n_t) * DT_MS - firing_rates = np.array(firing_rates) - - fit_mask = current_levels >= rheobase__nA - if np.sum(fit_mask) >= 2: - fit_currents = current_levels[fit_mask] - fit_rates = firing_rates[fit_mask] - - coefficients = np.polyfit(fit_currents, fit_rates, 1) - gain__Hz_per_nA = coefficients[0] - return gain__Hz_per_nA - else: - return np.nan - - -def extract_all_parameters(cell, cell_index: int, recruitment_threshold: float) -> dict: + results = [] + for i in range(n_neurons): + v = V[i] + if np.max(v) < spike_thresholds[i]: + results.append({"AP__mV": np.nan, "AHP__mV": np.nan, "AHPdur__ms": np.nan}) + continue + peak_idx = int(np.argmax(v)) + post_peak = v[peak_idx:] + valley_idx = peak_idx + int(np.argmin(post_peak)) + ap = float(v[peak_idx] - vholds[i]) + ahp = float(vholds[i] - v[valley_idx]) + rec = np.where(v[valley_idx:] > vholds[i] - 0.15)[0] + ahp_dur = float(t_array[rec[0] + valley_idx] - t_array[peak_idx]) if len(rec) > 0 else np.nan + results.append({"AP__mV": ap, "AHP__mV": ahp, "AHPdur__ms": ahp_dur}) + return results + + +def get_fi_gains(net, n_neurons: int, vholds: np.ndarray, + rheobases: np.ndarray) -> np.ndarray: """ - Extract all electrophysiological parameters from a single neuron. + F-I gain for all neurons — 12 jx.integrate calls. - Returns - ------- - dict - Dictionary containing all extracted parameters and metadata. + Uses rheobase-relative current steps (rheobase_i + delta, delta in 0..11 nA). + Each of the 12 levels is one batched call across all neurons. """ - logger.info(f"Extracting parameters for cell {cell_index}") - - sections = [cell.soma] + cell.dend - v_init__mV = -67.0 - voltages = [v_init__mV] * len(sections) - - result = {"cell_index": cell_index, "recruitment_threshold": recruitment_threshold} - - # Extract parameters - try: - vhold__mV = get_vhold__mV(cell, sections, voltages) - result["vhold_soma__mV"] = vhold__mV - except Exception as e: - logger.error(f"Cell {cell_index}: Failed to get Vhold - {e}") - result["vhold_soma__mV"] = np.nan - vhold__mV = v_init__mV - - try: - result["Rin__MOhm"] = get_rin__MOhm(cell, vhold__mV, sections) - except Exception as e: - logger.error(f"Cell {cell_index}: Failed to get Rin - {e}") - result["Rin__MOhm"] = np.nan - - try: - result["tau__ms"] = get_time_constant__ms(cell, vhold__mV, sections) - except Exception as e: - logger.error(f"Cell {cell_index}: Failed to get tau - {e}") - result["tau__ms"] = np.nan - - try: - ir__nA = get_rheobase__nA(cell, vhold__mV, sections) - result["Ir__nA"] = ir__nA - except Exception as e: - logger.error(f"Cell {cell_index}: Failed to get rheobase - {e}") - result["Ir__nA"] = np.nan - - try: - ap_ahp_params = get_ap_and_ahp(cell, vhold__mV, sections) - result.update(ap_ahp_params) - except Exception as e: - logger.error(f"Cell {cell_index}: Failed to get AP/AHP - {e}") - result.update({"AP__mV": np.nan, "AHP__mV": np.nan, "AHPdur__ms": np.nan}) - - try: - rheobase__nA = result.get("Ir__nA", np.nan) - result["FI_gain__Hz_per_nA"] = get_fi_gain(cell, vhold__mV, sections, rheobase__nA) - except Exception as e: - logger.error(f"Cell {cell_index}: Failed to get F-I gain - {e}") - result["FI_gain__Hz_per_nA"] = np.nan - - logger.info(f"Cell {cell_index}: Extraction complete") - return result + tstop = 3000.0 + n_t = int(tstop / DT_MS) + settle_pts = int(500.0 / DT_MS) + analysis_ms = 2500.0 + deltas = np.arange(0.0, 12.0, 1.0) # 12 levels + all_rates = np.zeros((n_neurons, len(deltas))) + + for j, delta in enumerate(deltas): + amps = np.where(np.isnan(rheobases), 0.0, rheobases + delta) + currents = [np.ones(n_t) * a for a in amps] + V = _run_batch(net, n_neurons, currents, tstop) + for i in range(n_neurons): + if np.isnan(rheobases[i]): + continue + v = V[i, settle_pts:] + n_spikes = int(np.sum((v[:-1] < SPIKE_THRESHOLD_MV) & (v[1:] >= SPIKE_THRESHOLD_MV))) + all_rates[i, j] = n_spikes / analysis_ms * 1000.0 + + gains = [] + for i in range(n_neurons): + rates = all_rates[i] + if not np.any(rates > 0): + gains.append(np.nan) + continue + peak_idx = int(np.argmax(rates)) + end_idx = min(max(peak_idx + 1, 3), len(rates)) + fit_currents = (rheobases[i] + deltas)[:end_idx] + fit_rates = rates[:end_idx] + mask = fit_rates > 0 + if np.sum(mask) >= 2: + coeffs = np.polyfit(fit_currents[mask], fit_rates[mask], 1) + gains.append(float(coeffs[0])) + else: + gains.append(np.nan) + return np.array(gains) ############################################################################## @@ -465,8 +413,9 @@ def visualize_results(df: pd.DataFrame, save_path: Path) -> None: axes[1, 1].set_title("F-I Gain") axes[1, 1].grid(True, alpha=0.3) + plt.suptitle("Jaxley NERLab Channel Parameters", fontsize=14, fontweight="bold") plt.tight_layout() - fig_path = save_path / "neuron_parameters.png" + fig_path = save_path / "neuron_parameters_jaxley.png" plt.savefig(fig_path, dpi=300, bbox_inches="tight") logger.info(f"Saved figure to {fig_path}") plt.show() @@ -485,11 +434,11 @@ def visualize_results(df: pd.DataFrame, save_path: Path) -> None: cbar.set_label("Motor Unit Index") ax.set_xlabel("Input Resistance (MΩ)") ax.set_ylabel("Membrane Time Constant (ms)") - ax.set_title("Input Resistance vs Time Constant") + ax.set_title("Input Resistance vs Time Constant (Jaxley NERLab)") ax.grid(True, alpha=0.3) plt.tight_layout() - fig2_path = save_path / "rin_tau_relationship.png" + fig2_path = save_path / "rin_tau_relationship_jaxley.png" plt.savefig(fig2_path, dpi=300, bbox_inches="tight") logger.info(f"Saved figure to {fig2_path}") plt.show() @@ -505,7 +454,7 @@ def main(): import argparse parser = argparse.ArgumentParser( - description="Extract electrophysiological parameters from MyoGen motor neurons" + description="Extract electrophysiological parameters from MyoGen motor neurons (Jaxley)" ) parser.add_argument( "--n-neurons", @@ -529,43 +478,106 @@ def main(): myogen.set_random_seed(args.seed) logger.info(f"Random seed set to {args.seed}") - load_nmodl_mechanisms() - logger.info("MyoGen NEURON mechanisms loaded") - - # Generate recruitment thresholds using Fuglevand model - logger.info(f"Generating {args.n_neurons} neurons with Fuglevand model") - recruitment_thresholds, _ = simulator.RecruitmentThresholds( + # Generate recruitment thresholds (combined model, matching ex01) + logger.info(f"Generating {args.n_neurons} neurons with combined threshold model") + recruitment_thresholds, _ = RecruitmentThresholds( N=args.n_neurons, recruitment_range__ratio=100, - mode="fuglevand", + mode="combined", + deluca__slope=5, ) - # Create motor neuron pool - logger.info("Creating motor neuron pool") - pool = AlphaMN__Pool(recruitment_thresholds__array=recruitment_thresholds) - logger.info(f"Created pool with {len(pool._cells)} neurons") - - # Extract parameters from all neurons - logger.info(f"Extracting parameters from {args.n_neurons} neurons") + # Create motor neuron pool — model="NERLab" default matches production NEURON. + logger.info("Creating motor neuron pool (NERLab — napp + caL)") + with suppress_stdout(): + pool = AlphaMN__Pool( + recruitment_thresholds__array=recruitment_thresholds, + mode="active", + ) + n_neurons = len(pool) + logger.info(f"Created pool with {n_neurons} neurons") + + # Build a single jx.Network from all neurons. + # All measurement functions call jx.integrate on this network once per + # protocol step, parallelising across neurons instead of looping over them. + logger.info("Building jx.Network for batched parameter extraction") + with suppress_stdout(): + mn_cells = [] + for cw in pool: + cell = cw.cell + cell.delete_recordings() + cell.delete_stimuli() + mn_cells.append(cell) + net = jx.Network(mn_cells) + for i in range(n_neurons): + net.cell(i).branch(0).loc(0.5).record("v") + + print("\n" + "=" * 60) + print("Running Jaxley NERLab Parameter Extraction (batched)") + print(f" {n_neurons} neurons × ~31 total jx.integrate calls") + print("=" * 60) + + logger.info("Extracting Vhold (1 call)...") + vholds = get_vholds(net, n_neurons) + + logger.info("Extracting Rin (5 calls)...") + rins = get_rins(net, n_neurons, vholds) + + logger.info("Extracting tau (1 call)...") + taus = get_taus(net, n_neurons, vholds) + + logger.info("Extracting rheobase (~11 calls, batched binary search)...") + rheobases = get_rheobases(net, n_neurons, vholds) + + logger.info("Extracting AP/AHP (1 call)...") + ap_ahp_list = get_ap_ahp(net, n_neurons, vholds, rheobases) + + logger.info("Extracting F-I gain (12 calls)...") + fi_gains = get_fi_gains(net, n_neurons, vholds, rheobases) + + # Assemble per-neuron result dicts results = [] - for i in range(args.n_neurons): - result = extract_all_parameters(pool[i], i, recruitment_thresholds[i]) - results.append(result) + for i in range(n_neurons): + results.append({ + "cell_index": i, + "recruitment_threshold": float(recruitment_thresholds[i]), + "vhold_soma__mV": float(vholds[i]), + "Rin__MOhm": float(rins[i]), + "tau__ms": float(taus[i]), + "Ir__nA": float(rheobases[i]) if not np.isnan(rheobases[i]) else np.nan, + **ap_ahp_list[i], + "FI_gain__Hz_per_nA": float(fi_gains[i]) if not np.isnan(fi_gains[i]) else np.nan, + }) + logger.info( + f" MN {i}: Vhold={vholds[i]:.1f} mV Rin={rins[i]:.2f} MΩ " + f"Ir={rheobases[i]:.1f} nA FI={fi_gains[i]:.2f} Hz/nA" + ) # Convert to DataFrame and save df = pd.DataFrame(results) - csv_path = save_path / "neuron_parameters.csv" + csv_path = save_path / "neuron_parameters_jaxley.csv" df.to_csv(csv_path, index=False) logger.info(f"Saved results to {csv_path}") # Visualize visualize_results(df, save_path) - logger.info("Parameter extraction complete!") - logger.info("\nSummary statistics:") - logger.info(f" Rheobase range: {df['Ir__nA'].min():.2f} - {df['Ir__nA'].max():.2f} nA") - logger.info(f" Rin range: {df['Rin__MOhm'].min():.2f} - {df['Rin__MOhm'].max():.2f} MΩ") - logger.info(f" F-I gain range: {df['FI_gain__Hz_per_nA'].min():.2f} - {df['FI_gain__Hz_per_nA'].max():.2f} Hz/nA") + print("\n" + "=" * 60) + print("Summary (Jaxley NERLab Channels)") + print("=" * 60) + print(f"Simulator: Jaxley with NERLab channels (soma napp; dendrite caL)") + print(f"Neurons extracted: {args.n_neurons}") + print(f"\nParameter ranges:") + print(f" Rheobase: {df['Ir__nA'].min():.2f} - {df['Ir__nA'].max():.2f} nA") + print(f" Input resistance: {df['Rin__MOhm'].min():.2f} - {df['Rin__MOhm'].max():.2f} MΩ") + print(f" Time constant: {df['tau__ms'].min():.2f} - {df['tau__ms'].max():.2f} ms") + print(f" AP amplitude: {df['AP__mV'].min():.2f} - {df['AP__mV'].max():.2f} mV") + print(f" AHP depth: {df['AHP__mV'].min():.2f} - {df['AHP__mV'].max():.2f} mV") + print(f" F-I gain: {df['FI_gain__Hz_per_nA'].min():.2f} - {df['FI_gain__Hz_per_nA'].max():.2f} Hz/nA") + + print("\n" + "=" * 60) + print("[DONE] Jaxley NERLab parameter extraction complete!") + print("=" * 60) if __name__ == "__main__": diff --git a/examples/01_basic/10_extract_neuron_parameters_neuron.py b/examples/01_basic/10_extract_neuron_parameters_neuron.py new file mode 100644 index 00000000..4f22ee85 --- /dev/null +++ b/examples/01_basic/10_extract_neuron_parameters_neuron.py @@ -0,0 +1,572 @@ +""" +Extract Neuron Parameters +========================= + +This example demonstrates how to extract electrophysiological parameters from +MyoGen motor neurons using the Fuglevand recruitment model. + +**Parameters Extracted**: + +- **Vhold**: Resting membrane potential (soma) +- **Rin**: Input resistance +- **tau**: Membrane time constant +- **Ir**: Rheobase (minimum current for action potential) +- **AP**: Action potential amplitude +- **AHP**: Afterhyperpolarization depth +- **AHPdur**: Full AHP duration +- **FI_gain**: Frequency-current gain (slope of F-I curve) + +**Usage**: + + # Extract from default 5 neurons + python 09_extract_neuron_parameters.py + + # Extract from custom number of neurons + python 09_extract_neuron_parameters.py --n-neurons 10 + +**Note**: For advanced features (parallel processing, model comparison, slope scan), +see `sandbox/neuron_parameter_extraction/README.md` +""" + +# %% + +############################################################################## +# Import Libraries +# ---------------- + +import logging +import os + +os.environ["MPLBACKEND"] = "Agg" +if "DISPLAY" in os.environ: + del os.environ["DISPLAY"] + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +from neuron import h + +import myogen +from myogen import simulator +from myogen.simulator.neuron.populations import AlphaMN__Pool +from myogen.utils.nmodl import load_nmodl_mechanisms + +# Simple plotting style +plt.style.use("seaborn-v0_8-darkgrid") +sns.set_context("paper", font_scale=1.2) + +# Setup NEURON environment +h.load_file("stdrun.hoc") + +# Setup logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + +############################################################################## +# Core Parameter Extraction Functions +# ------------------------------------ + + +def fi0_multicompartment(sections: list, voltages: list) -> None: + """Initialize voltages for multiple compartments.""" + for sec, v in zip(sections, voltages): + sec.v = v + + +def get_vhold__mV(cell, sections: list, voltages: list, tstop__ms: float = 500.0) -> float: + """ + Measure resting membrane potential for soma. + + Returns + ------- + float + Steady-state voltage in soma (mV). + """ + h.tstop = tstop__ms + h.dt = 0.0125 + h.celsius = 36 + + _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + + vsoma = h.Vector() + vsoma.record(cell.soma(0.5)._ref_v) + + h.finitialize() + h.run() + + return vsoma.to_python()[-1] + + +def get_rin__MOhm( + cell, + vhold__mV: float, + sections: list, + amplitudes__nA: list[float] = None, +) -> float: + """ + Measure input resistance using current steps. + + Returns + ------- + float + Input resistance in MegaOhms. + """ + if amplitudes__nA is None: + amplitudes__nA = [-5.0, -4.0, -3.0, -2.0, -1.0] + + h.tstop = 1000.0 + h.dt = 0.0125 + h.celsius = 36 + + voltages = [vhold__mV] * len(sections) + _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + + vsoma = h.Vector() + vsoma.record(cell.soma(0.5)._ref_v) + + delta_v__mV = [] + for amp__nA in amplitudes__nA: + stim = h.IClamp(cell.soma(0.5)) + stim.amp = amp__nA + stim.dur = 1000.0 + stim.delay = 0 + + h.finitialize() + h.run() + + vpeak__mV = vsoma.to_python()[-1] + dv = vpeak__mV - vhold__mV + delta_v__mV.append(dv) + del stim + + coefficients = np.polyfit(amplitudes__nA, delta_v__mV, 1) + rin__MOhm = coefficients[0] + return rin__MOhm + + +def get_time_constant__ms(cell, vhold__mV: float, sections: list) -> float: + """ + Measure membrane time constant from exponential decay. + + Returns + ------- + float + Membrane time constant in milliseconds. + """ + h.tstop = 200.0 + h.dt = 0.0125 + h.celsius = 36 + + voltages = [vhold__mV] * len(sections) + _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + + vsoma = h.Vector() + t = h.Vector() + vsoma.record(cell.soma(0.5)._ref_v) + t.record(h._ref_t) + + stim = h.IClamp(cell.soma(0.5)) + stim.amp = -20.0 # nA + stim.dur = 1.0 # ms + stim.delay = 5.0 # ms + + h.finitialize() + h.run() + + vsoma_array = vsoma.as_numpy() + t_array = t.as_numpy() + + v_min__mV = np.min(vsoma_array) + v_min_index = np.where(vsoma_array == v_min__mV)[0][0] + + recovery_indices = np.where(vsoma_array[v_min_index:] > vhold__mV - 0.1)[0] + if len(recovery_indices) > 0: + v_end_index = recovery_indices[0] + v_min_index + else: + v_end_index = len(vsoma_array) - 1 + + t_fit = t_array[v_min_index:v_end_index] + v_fit = vsoma_array[v_min_index:v_end_index] + log_v_fit = np.log(np.abs(v_fit - vhold__mV)) + + coefficients = np.polyfit(t_fit, log_v_fit, 1) + tau__ms = -1.0 / coefficients[0] + return tau__ms + + +def get_rheobase__nA( + cell, vhold__mV: float, sections: list, max_iterations: int = 500 +) -> float: + """ + Find minimum current required to elicit an action potential. + + Returns + ------- + float + Rheobase current in nanoamperes, or np.nan if no spike found. + """ + h.tstop = 100.0 + h.dt = 0.0125 + h.celsius = 36 + + voltages = [vhold__mV] * len(sections) + _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + + vsoma = h.Vector() + vsoma.record(cell.soma(0.5)._ref_v) + + stim = h.IClamp(cell.soma(0.5)) + stim.dur = 50.0 + stim.delay = 0 + + spike_threshold__mV = vhold__mV + 40.0 + + amp__nA = 0.1 + iteration = 0 + while iteration < max_iterations: + stim.amp = amp__nA + + h.finitialize() + h.run() + + vpeak__mV = np.max(vsoma.as_numpy()) + if vpeak__mV >= spike_threshold__mV: + return amp__nA + + amp__nA += 0.1 + iteration += 1 + + logger.warning(f"No spike found after {max_iterations} iterations") + return np.nan + + +def get_ap_and_ahp(cell, vhold__mV: float, sections: list) -> dict: + """ + Measure action potential and afterhyperpolarization characteristics. + + Returns + ------- + dict + Contains 'AP__mV', 'AHP__mV', 'AHPdur__ms' + """ + h.tstop = 900.0 + h.dt = 0.0125 + h.celsius = 36 + + voltages = [vhold__mV] * len(sections) + _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + + vsoma = h.Vector() + t = h.Vector() + vsoma.record(cell.soma(0.5)._ref_v) + t.record(h._ref_t) + + stim = h.IClamp(cell.soma(0.5)) + stim.dur = 0.5 + stim.delay = 5.0 + + amp__nA = 35.0 + vpeak__mV = vhold__mV + spike_threshold__mV = vhold__mV + 40.0 + + while vpeak__mV < spike_threshold__mV: + amp__nA += 10.0 + stim.amp = amp__nA + + h.finitialize() + h.run() + + vsoma_array = vsoma.as_numpy() + vpeak__mV = np.max(vsoma_array) + vvalley__mV = np.min(vsoma_array) + + if amp__nA > 500.0: + logger.warning("Current exceeded 500 nA without spike") + return {"AP__mV": np.nan, "AHP__mV": np.nan, "AHPdur__ms": np.nan} + + ap__mV = vpeak__mV - vhold__mV + ahp__mV = vhold__mV - vvalley__mV + + t_array = t.as_numpy() + peak_index = np.where(vsoma_array == vpeak__mV)[0][0] + valley_index = np.where(vsoma_array == vvalley__mV)[0][0] + recovery_indices = np.where(vsoma_array[valley_index:] > vhold__mV - 0.15)[0] + + if len(recovery_indices) > 0: + recovery_index = recovery_indices[0] + valley_index + ahp_dur__ms = t_array[recovery_index] - t_array[peak_index] + else: + ahp_dur__ms = np.nan + + return {"AP__mV": ap__mV, "AHP__mV": ahp__mV, "AHPdur__ms": ahp_dur__ms} + + +def get_fi_gain(cell, vhold__mV: float, sections: list, rheobase__nA: float) -> float: + """ + Measure frequency-current (F-I) gain. + + Returns + ------- + float + F-I gain in Hz/nA. + """ + if np.isnan(rheobase__nA): + return np.nan + + h.tstop = 3000.0 + h.dt = 0.0125 + h.celsius = 36 + + voltages = [vhold__mV] * len(sections) + _ = h.FInitializeHandler(0, lambda: fi0_multicompartment(sections, voltages)) + + stim = h.IClamp(cell.soma(0.5)) + stim.dur = 3000.0 + stim.delay = 0 + + current_levels = np.arange(5.0, 30.5, 5.0) + firing_rates = [] + + for current__nA in current_levels: + stim.amp = current__nA + + apc = h.APCount(cell.soma(0.5)) + apc.thresh = 50 + + h.finitialize() + h.run() + + spike_count = apc.n + firing_rate__Hz = (spike_count / 2500.0) * 1000.0 + firing_rates.append(firing_rate__Hz) + + firing_rates = np.array(firing_rates) + + fit_mask = current_levels >= rheobase__nA + if np.sum(fit_mask) >= 2: + fit_currents = current_levels[fit_mask] + fit_rates = firing_rates[fit_mask] + + coefficients = np.polyfit(fit_currents, fit_rates, 1) + gain__Hz_per_nA = coefficients[0] + return gain__Hz_per_nA + else: + return np.nan + + +def extract_all_parameters(cell, cell_index: int, recruitment_threshold: float) -> dict: + """ + Extract all electrophysiological parameters from a single neuron. + + Returns + ------- + dict + Dictionary containing all extracted parameters and metadata. + """ + logger.info(f"Extracting parameters for cell {cell_index}") + + sections = [cell.soma] + cell.dend + v_init__mV = -67.0 + voltages = [v_init__mV] * len(sections) + + result = {"cell_index": cell_index, "recruitment_threshold": recruitment_threshold} + + # Extract parameters + try: + vhold__mV = get_vhold__mV(cell, sections, voltages) + result["vhold_soma__mV"] = vhold__mV + except Exception as e: + logger.error(f"Cell {cell_index}: Failed to get Vhold - {e}") + result["vhold_soma__mV"] = np.nan + vhold__mV = v_init__mV + + try: + result["Rin__MOhm"] = get_rin__MOhm(cell, vhold__mV, sections) + except Exception as e: + logger.error(f"Cell {cell_index}: Failed to get Rin - {e}") + result["Rin__MOhm"] = np.nan + + try: + result["tau__ms"] = get_time_constant__ms(cell, vhold__mV, sections) + except Exception as e: + logger.error(f"Cell {cell_index}: Failed to get tau - {e}") + result["tau__ms"] = np.nan + + try: + ir__nA = get_rheobase__nA(cell, vhold__mV, sections) + result["Ir__nA"] = ir__nA + except Exception as e: + logger.error(f"Cell {cell_index}: Failed to get rheobase - {e}") + result["Ir__nA"] = np.nan + + try: + ap_ahp_params = get_ap_and_ahp(cell, vhold__mV, sections) + result.update(ap_ahp_params) + except Exception as e: + logger.error(f"Cell {cell_index}: Failed to get AP/AHP - {e}") + result.update({"AP__mV": np.nan, "AHP__mV": np.nan, "AHPdur__ms": np.nan}) + + try: + rheobase__nA = result.get("Ir__nA", np.nan) + result["FI_gain__Hz_per_nA"] = get_fi_gain(cell, vhold__mV, sections, rheobase__nA) + except Exception as e: + logger.error(f"Cell {cell_index}: Failed to get F-I gain - {e}") + result["FI_gain__Hz_per_nA"] = np.nan + + logger.info(f"Cell {cell_index}: Extraction complete") + return result + + +############################################################################## +# Visualization Functions +# ----------------------- + + +def visualize_results(df: pd.DataFrame, save_path: Path) -> None: + """Create basic visualization plots of extracted parameters.""" + logger.info("Creating visualizations") + + df_valid = df.dropna(subset=["Ir__nA"]) + if len(df_valid) == 0: + logger.warning("No valid data for visualization") + return + + x = df_valid["cell_index"] + 1 + + # Figure 1: Key parameters + fig, axes = plt.subplots(2, 2, figsize=(12, 10)) + + # Rheobase + axes[0, 0].scatter(x, df_valid["Ir__nA"], s=100, alpha=0.7) + axes[0, 0].set_ylabel("Rheobase (nA)") + axes[0, 0].set_title("Rheobase Current") + axes[0, 0].grid(True, alpha=0.3) + + # Input Resistance + axes[0, 1].scatter(x, df_valid["Rin__MOhm"], s=100, alpha=0.7, color="C1") + axes[0, 1].set_ylabel("Input Resistance (MΩ)") + axes[0, 1].set_title("Input Resistance") + axes[0, 1].grid(True, alpha=0.3) + + # AHP Duration + axes[1, 0].scatter(x, df_valid["AHPdur__ms"], s=100, alpha=0.7, color="C2") + axes[1, 0].set_xlabel("Motor Unit #") + axes[1, 0].set_ylabel("AHP Duration (ms)") + axes[1, 0].set_title("AHP Duration") + axes[1, 0].grid(True, alpha=0.3) + + # F-I Gain + axes[1, 1].scatter(x, df_valid["FI_gain__Hz_per_nA"], s=100, alpha=0.7, color="C3") + axes[1, 1].set_xlabel("Motor Unit #") + axes[1, 1].set_ylabel("F-I Gain (Hz/nA)") + axes[1, 1].set_title("F-I Gain") + axes[1, 1].grid(True, alpha=0.3) + + plt.tight_layout() + fig_path = save_path / "neuron_parameters.png" + plt.savefig(fig_path, dpi=300, bbox_inches="tight") + logger.info(f"Saved figure to {fig_path}") + plt.show() + + # Figure 2: Rin vs Tau + fig2, ax = plt.subplots(figsize=(8, 6)) + scatter = ax.scatter( + df_valid["Rin__MOhm"], + df_valid["tau__ms"], + c=df_valid["cell_index"], + cmap="viridis", + s=100, + alpha=0.7, + ) + cbar = plt.colorbar(scatter, ax=ax) + cbar.set_label("Motor Unit Index") + ax.set_xlabel("Input Resistance (MΩ)") + ax.set_ylabel("Membrane Time Constant (ms)") + ax.set_title("Input Resistance vs Time Constant") + ax.grid(True, alpha=0.3) + + plt.tight_layout() + fig2_path = save_path / "rin_tau_relationship.png" + plt.savefig(fig2_path, dpi=300, bbox_inches="tight") + logger.info(f"Saved figure to {fig2_path}") + plt.show() + + +############################################################################## +# Main Function +# ------------- + + +def main(): + """Main demonstration function.""" + import argparse + + parser = argparse.ArgumentParser( + description="Extract electrophysiological parameters from MyoGen motor neurons" + ) + parser.add_argument( + "--n-neurons", + type=int, + default=5, + help="Number of neurons to extract (default: 5)", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility (default: 42)", + ) + + args = parser.parse_args() + + # Setup + save_path = Path("./results") + save_path.mkdir(exist_ok=True) + + myogen.set_random_seed(args.seed) + logger.info(f"Random seed set to {args.seed}") + + load_nmodl_mechanisms() + logger.info("MyoGen NEURON mechanisms loaded") + + # Generate recruitment thresholds using Fuglevand model + logger.info(f"Generating {args.n_neurons} neurons with Fuglevand model") + recruitment_thresholds, _ = simulator.RecruitmentThresholds( + N=args.n_neurons, + recruitment_range__ratio=100, + mode="fuglevand", + ) + + # Create motor neuron pool + logger.info("Creating motor neuron pool") + pool = AlphaMN__Pool(recruitment_thresholds__array=recruitment_thresholds) + logger.info(f"Created pool with {len(pool._cells)} neurons") + + # Extract parameters from all neurons + logger.info(f"Extracting parameters from {args.n_neurons} neurons") + results = [] + for i in range(args.n_neurons): + result = extract_all_parameters(pool[i], i, recruitment_thresholds[i]) + results.append(result) + + # Convert to DataFrame and save + df = pd.DataFrame(results) + csv_path = save_path / "neuron_parameters.csv" + df.to_csv(csv_path, index=False) + logger.info(f"Saved results to {csv_path}") + + # Visualize + visualize_results(df, save_path) + + logger.info("Parameter extraction complete!") + logger.info("\nSummary statistics:") + logger.info(f" Rheobase range: {df['Ir__nA'].min():.2f} - {df['Ir__nA'].max():.2f} nA") + logger.info(f" Rin range: {df['Rin__MOhm'].min():.2f} - {df['Rin__MOhm'].max():.2f} MΩ") + logger.info(f" F-I gain range: {df['FI_gain__Hz_per_nA'].min():.2f} - {df['FI_gain__Hz_per_nA'].max():.2f} Hz/nA") + + +if __name__ == "__main__": + main() diff --git a/examples/01_basic/11_simulate_spinal_network.py b/examples/01_basic/11_simulate_spinal_network.py index c903a6ef..27a6d08d 100644 --- a/examples/01_basic/11_simulate_spinal_network.py +++ b/examples/01_basic/11_simulate_spinal_network.py @@ -1,95 +1,55 @@ """ -Spinal Network & Tendon Tap -=========================== +Spinal Network Simulation with Systematic Tendon Tap Protocol - Jaxley Backend +=============================================================================== This example demonstrates **complete spinal reflex network modeling** with a **comprehensive tendon tap -protocol** that systematically varies mechanical perturbations, fusimotor drive, and cortical activity. -This sophisticated experimental design reveals how the nervous system modulates stretch reflex gain. +protocol** using the Jaxley (JAX-based) simulator with biophysical motor neuron channels. + +This is the Jaxley equivalent of ``11_simulate_spinal_network.py`` which uses NEURON. +Both examples implement the same neuromuscular system: + + - **Motor neuron pool**: α-motoneurons with biophysical HH channels + - **Afferent populations**: Ia, II (spindle), Ib (GTO) feedback + - **Interneuron populations**: gII, gIb for reflex modulation + - **Proprioceptive models**: Muscle spindles and Golgi tendon organs + - **Hill-type muscle model**: Realistic force generation + - **Joint dynamics**: Closed-loop biomechanical control + - **Descending drive**: Cortical control via Poisson spike trains Learning Objectives ------------------ -This example teaches the following key concepts in neuromuscular control: - 1. **Reflex Gain Modulation**: How fusimotor (gamma) drive amplifies stretch reflex sensitivity - - Compare motor neuron responses across different gamma levels (0→100 pps) - - Observe how higher gamma drive increases spindle sensitivity and reflex output - -2. **Sensory Pathway Timing**: Why different feedback pathways activate at different times - - Spindle-driven pathways (Ia, II → gII) respond immediately to muscle stretch - - Force-driven pathways (Ib → gIb) respond later after muscle contraction builds force - -3. **Reflex-Voluntary Interaction**: How cortical commands interact with spinal reflexes - - Phase 1: Isolated reflex responses (no cortical drive) - - Phase 2: Reflex modulation during voluntary contraction (with cortical drive) - - Observe how background muscle activity affects reflex sensitivity - -4. **Closed-loop Biomechanics**: How muscle forces create movement that feeds back to sensors - - Joint dynamics convert muscle torque into motion - - Motion changes muscle length/velocity, affecting spindle feedback - - Creates realistic proprioceptive-motor coupling +2. **Sensory Pathway Timing**: Different feedback pathways activate at different times +3. **Reflex-Voluntary Interaction**: Cortical commands interact with spinal reflexes +4. **Closed-loop Biomechanics**: Muscle forces create movement that feeds back to sensors **Experimental Protocol (5-second, 2-phase design):** **Phase 1 (0-2.5s): Reflex Gain Modulation** - Triangular tendon taps every 0.5s with increasing amplitude (50%, 60%, 70%, 80%) - - * Triangular waveform: Linear ramp up to peak (50ms) → linear ramp down (50ms) - * Total duration: 100ms per tap - * Mimics realistic tendon tap dynamics with stretch and release phases - -- Stepwise increasing gamma drive starting at 0.75s (0→25→50→75→100 pps) +- Stepwise increasing gamma drive (0→25→50→75→100 pps) - No cortical drive (isolated reflex testing) **Phase 2 (2.5-5s): Reflex-Voluntary Interaction** -- Repeat triangular tap pattern (50%, 60%, 70%, 80%) -- Repeat gamma drive pattern (0→25→50→75→100 pps) -- Sinusoidal cortical drive (38±1 Hz at 1 Hz) recruiting ~75% of motor neurons - -!!! note - This example builds upon all previous examples and demonstrates how the complete neuromuscular system - functions as an integrated network: - - - **Motor neuron pool**: α-motoneurons controlling a single muscle (from examples 01-02) - - **Muscle model**: Hill-type muscle model with realistic force generation (from examples 02, 06) - - **Proprioceptive feedback**: Muscle spindles and Golgi tendon organs (from examples) - - **Spinal interneurons**: Group II and Group Ib interneurons for reflex modulation - - **Joint dynamics**: Closed-loop biomechanical control with realistic inertia and damping - - **Descending drive**: Cortical control signals for voluntary movement (from example 01) - -!!! important - **Spinal Reflex Networks** are the fundamental control circuits that coordinate muscle activity. - Key physiological concepts demonstrated: - - - **Stretch reflex**: Muscle spindles provide length/velocity feedback for posture control - - **Force feedback**: Golgi tendon organs monitor muscle tension for protective reflexes - - **Closed-loop control**: Joint mechanics influence neural activity through proprioception - - **Motor control**: Integration of descending commands with sensory feedback - -Scientific Background -------------------- - -The spinal cord contains intricate neural circuits that process sensory information and generate appropriate -motor responses. This example models several key components: - -**Muscle Spindles**: Specialized sensory organs that detect muscle length and velocity changes, providing -Group Ia (velocity-sensitive) and Group II (length-sensitive) afferent signals. - -**Golgi Tendon Organs**: Force-sensitive receptors that monitor muscle tension and provide Group Ib -afferent feedback for protective reflexes. - -**Spinal Interneurons**: The model includes two key inhibitory interneuron populations: - -- **Group II interneurons (gII)**: Process Group II spindle afferents (length-sensitive). In this model, - they provide inhibitory feedback to motor neurons, representing one component of the complex Group II - pathway (which can be excitatory or inhibitory depending on context). - -- **Group Ib interneurons (gIb)**: Receive input from Golgi tendon organs and provide **autogenic - inhibition** - force-dependent negative feedback that protects against excessive muscle tension. - Classic Ib inhibitory pathway. - -**Joint Dynamics**: Realistic biomechanical modeling of joint inertia, damping, and muscle-generated torques -creates the closed-loop system where neural activity affects movement, which in turn affects sensory feedback. +- Repeat triangular tap pattern +- Repeat gamma drive pattern +- Sinusoidal cortical drive (40±1 Hz at 1 Hz) + +.. note:: + **Motor-neuron model.** ``AlphaMN__Pool`` defaults to ``model="NERLab"`` + (the production NEURON model). ex11 is a *mixed-frame* network: the + gII / gIb interneurons defined in + ``myogen/simulator/jaxley/populations/interneurons.py`` use Na3rp-style + channels in the modern absolute voltage convention (V_rest ≈ -70 mV), + while NERLab MNs live in the original 1952-HH frame (V_rest ≈ 0 mV). + This is fine here because **every synapse in the network targets only + the MNs** (gII → MN, gIb → MN), so we just need one inhibitory reversal + set for the NERLab-frame target — see ``IonotropicSynapse_e_syn`` below. + The interneurons act only as presynaptic sources, and the synapse spike + threshold ``IonotropicSynapse_v_th = -35 mV`` is correct in their + (modern) frame. NEURON ex11 handles the same situation by setting + ``syn.e`` per-postsynaptic-cell — see ``myogen.simulator.neuron.cells.create_synapses``. """ # %% @@ -97,21 +57,26 @@ ############################################################################## # Import Libraries # ---------------- -# from pathlib import Path +import jax +import jax.numpy as jnp +import jaxley as jx +from jaxley.integrate import build_init_and_step_fn import joblib import matplotlib.pyplot as plt import numpy as np import quantities as pq -from neuron import h - -from myogen import load_nmodl_mechanisms -from myogen.simulator.neuron.joint_dynamics import JointDynamics -from myogen.simulator.neuron.muscle import HillModel -from myogen.simulator.neuron.network import Network -from myogen.simulator.neuron.populations import ( +from neo import AnalogSignal, Block, Segment, SpikeTrain +from tqdm import tqdm + +import myogen +from jaxley.connect import sparse_connect +from jaxley.synapses import IonotropicSynapse +from myogen.simulator.jaxley.joint_dynamics import JointDynamics +from myogen.simulator.jaxley.muscle import HillModel +from myogen.simulator.jaxley.populations import ( AffIa__Pool, AffIb__Pool, AffII__Pool, @@ -120,11 +85,10 @@ GIb__Pool, GII__Pool, ) -from myogen.simulator.neuron.proprioception import ( +from myogen.simulator.jaxley.proprioception import ( GolgiTendonOrganModel, SpindleModel, ) -from myogen.simulator.neuron.simulation_runner import SimulationRunner from myogen.utils.nwb import export_to_nwb from myogen.utils.plotting import ( plot_gto_dynamics, @@ -133,18 +97,22 @@ plot_raster_spikes, plot_spindle_dynamics, ) +from myogen.utils.types import pps +from myogen.simulator.jaxley.jax_models import ( + gamma_init, + gto_init, gto_params_from_dict, gto_step, + hill_init_params, hill_init_state, hill_step, + joint_init, joint_step, + make_connectivity_matrix, + make_scan_step, + poisson_init, + spindle_init, spindle_params_from_dict, spindle_step, +) ############################################################################## -# Load NEURON Mechanisms and Dependencies -# --------------------------------------- -# -# Load the compiled NMODL mechanisms required for biophysical neuron modeling -# and load results from previous examples that serve as inputs to this simulation. - -# Load NEURON mechanisms -load_nmodl_mechanisms() +# Setup Results Directory and Load Previous Data +# ----------------------------------------------- -# Setup results directory save_path = Path(r"./results") save_path.mkdir(exist_ok=True) @@ -154,15 +122,11 @@ ############################################################################## # Define Simulation Parameters # --------------------------- -# -# These parameters control the temporal and spatial resolution of the simulation, -# as well as the physiological characteristics of the neural populations and -# mechanical system. -# Temporal parameters - high resolution for accurate neural integration -dt = 0.005 * pq.ms # Integration timestep with units -tstop = 5e3 # 5e3 # ms - Total simulation duration (5s for comprehensive protocol) -time = np.arange(0, tstop, dt.magnitude) # Time vector (use magnitude for numpy) +# Temporal parameters +dt = 0.025 * pq.ms # Integration timestep (reduced from 0.1 ms for Hill/spindle ODE accuracy) +tstop = 5e3 # ms - Total simulation duration (5s for comprehensive protocol) +time = np.arange(0, tstop, dt.magnitude) # Time vector print("Simulation parameters:") print(f"\tDuration: {tstop} ms") @@ -172,35 +136,29 @@ ############################################################################## # Define Neural Population Sizes # ----------------------------- -# -# Population sizes are based on physiological estimates from cat and human studies. -# These numbers represent typical motor pool compositions for a single muscle. # Afferent populations (sensory input) -# Based on Banks et al. (2006, 2015) for distal upper-limb muscles nIa = 73 # Group Ia afferents from muscle spindles (velocity-sensitive) -nII = 80 # Group II afferents from muscle spindles (length-sensitive) - 110% of Ia -nIb = 58 # Group Ib afferents from Golgi tendon organs (force-sensitive) - 80% of Ia +nII = 80 # Group II afferents from muscle spindles (length-sensitive) +nIb = 58 # Group Ib afferents from Golgi tendon organs (force-sensitive) # Interneuron populations (spinal processing) -# Note: Interneuron:MN ratio is ~10:1 in spinal cord. Group II/Ib INs are a subset. -# Using ~50% of afferent count as a reasonable estimate for convergent interneurons. -ngII = 120 # Group II interneurons (process spindle feedback from 80 II afferents) -ngIb = 145 # Group Ib interneurons (process force feedback from 58 Ib afferents) +ngII = 120 # Group II interneurons +ngIb = 145 # Group Ib interneurons # Motor neurons (output to muscles) -nType1 = 102 # Type I motor neurons (slow, fatigue-resistant) -nType2 = 18 # Type II motor neurons (fast, fatigue-prone) -naMN = nType1 + nType2 # Total α-motoneurons per muscle +# Motor unit type composition (needed for Hill muscle model) +# NEURON uses 102/18 (85%/15%) for 120 MNs — scale proportionally to loaded pool size +naMN = len(recruitment_thresholds) # Total α-motoneurons from thresholds file +nType1 = int(round(naMN * 102 / 120)) # Type I motor units (slow, fatigue-resistant) +nType2 = naMN - nType1 # Type II motor units (fast, fatigue-prone) # Descending drive (cortical input) -nDD = 400 # Total descending drive neurons (increased for more drive) -DDorder = ( - 5 # Gamma shape parameter for realistic spike patterns (higher shape = more regular spiking, CV=1/sqrt(shape)) -) +nDD = 400 # Total descending drive neurons +DDorder = 5 # Poisson process batch size print("Neural population sizes:") -print(f"\t- α-Motoneurons: {naMN} ({nType1} Type I, {nType2} Type II)") +print(f"\t- α-Motoneurons: {naMN} ({nType1} Type I + {nType2} Type II)") print(f"\t- Ia afferents: {nIa}") print(f"\t- II afferents: {nII}") print(f"\t- Ib afferents: {nIb}") @@ -210,19 +168,15 @@ ############################################################################## # Define Descending Drive Pattern # ------------------------------- -# -# Two-phase pattern: -# Phase 1 (0-2500ms): No cortical drive (reflex testing with varying gamma) -# Phase 2 (2500-5000ms): Sinusoidal cortical drive (reflex modulation during voluntary contraction) # Phase 1: No cortical drive -# Phase 2: 1 Hz sinusoidal drive with bias to recruit ~75% of MN population +# Phase 2: 1 Hz sinusoidal drive DDdrive = np.zeros_like(time) cortical_start_time = 2500 # ms # Parameters for sinusoidal drive -bias_hz = 40 # Bias for robust motor neuron recruitment (increased from 75 for ~20 pps firing) -amplitude_hz = 1 # Oscillation amplitude (±5 Hz around bias) +bias_hz = 40 # Bias for motor neuron recruitment +amplitude_hz = 1 # Oscillation amplitude frequency_hz = 1.0 # 1 Hz oscillation # Apply sinusoidal drive from 2.5s onwards @@ -234,26 +188,16 @@ print("Descending drive pattern:") print("\t- Phase 1 (0-2.5s): No cortical drive") print(f"\t- Phase 2 (2.5-5s): Sinusoidal drive ({bias_hz}±{amplitude_hz} Hz at {frequency_hz} Hz)") -print(f"\t (Drive range: {bias_hz - amplitude_hz}-{bias_hz + amplitude_hz} Hz)") ############################################################################## # Define Fusimotor Drive # -------------------- -# -# Fusimotor neurons control the sensitivity of muscle spindles by adjusting -# the tension in intrafusal muscle fibers. Stepwise increasing gamma drive -# demonstrates how fusimotor activity modulates stretch reflex gain. -# -# Pattern: Increases every 0.5s starting at 0.75s, resets at 2.5s -# Steps: 0 → 25 → 50 → 75 → 100 pps (then repeat) -# Time points for stepwise gamma drive -# Phase 1: 0→0.75s (wait)→1.25s→1.75s→2.25s (steps every 0.5s) -# Phase 2: 2.5s (reset)→3.25s (wait 0.75s)→3.75s→4.25s→4.75s (steps every 0.5s) +# Stepwise gamma drive pattern gamma_times = np.array([0, 750, 1250, 1750, 2250, 2500, 3250, 3750, 4250, 4750, 5000]) gamma_values = np.array([0, 25, 50, 75, 100, 0, 25, 50, 75, 100, 100]) -# Create stepwise gamma drive arrays using proper step function (not interpolation!) +# Create stepwise gamma drive arrays gDyn = np.zeros_like(time) gStat = np.zeros_like(time) @@ -273,79 +217,63 @@ gStat = gDyn.copy() # Add small physiological variability -gDyn = gDyn + np.random.normal(0, 1, len(time)) -gStat = gStat + np.random.normal(0, 1, len(time)) +gDyn = gDyn + myogen.get_random_generator().normal(0, 1, len(time)) +gStat = gStat + myogen.get_random_generator().normal(0, 1, len(time)) # Ensure non-negative gDyn = np.maximum(gDyn, 0) gStat = np.maximum(gStat, 0) -# Package fusimotor drives for the simulation -gMN = { - "name": "gMN", - "gDyn": gDyn, # Dynamic fusimotor drive array - "gStat": gStat, # Static fusimotor drive array -} - print("Fusimotor drive parameters:") print("\t- Stepwise increases: 0 → 25 → 50 → 75 → 100 pps") -print("\t- Starting at 0.75s, stepping every 0.5s") -print("\t- Resets at 2.5s and repeats pattern") ############################################################################## # Initialize Joint Dynamics # ------------------------ -# -# The joint model provides realistic biomechanical constraints and creates the -# closed-loop system where muscle forces influence joint motion, which in turn -# affects proprioceptive feedback. joint_dynamics = JointDynamics( - inertia__kg_m2=0.001, # Realistic joint inertia - damping__Nm_s_per_rad=0.002, # Light damping for stability + inertia__kg_m2=0.001, + damping__Nm_s_per_rad=0.002, ) -# Initialize joint angle array for closed-loop integration +# Initialize joint angle array artAng = np.zeros_like(time) -artAng[0] = 0.0 # Initial joint angle +artAng[0] = 0.0 print("Joint dynamics parameters:") print(f"\t- Inertia: {joint_dynamics.inertia__kg_m2} kg⋅m²") print(f"\t- Damping: {joint_dynamics.damping__Nm_s_per_rad} N⋅m⋅s/rad") -print(f"\t- Initial angle: {artAng[0]}°") ############################################################################## # Create Neural Populations # ------------------------ -# -# Instantiate all neural population objects with physiologically appropriate -# parameters. Each population type has specialized properties reflecting their -# biological counterparts. -# Create motor neuron pool -# Motor neuron pool (default spike_threshold__mV=50.0 for proper spike detection) -aMN = AlphaMN__Pool(recruitment_thresholds__array=recruitment_thresholds) # α-motoneurons +# Create motor neuron pool — default model="NERLab" matches the production +# NEURON model (soma napp + dendrite caL). The mixed-frame issue with +# modern-frame interneurons is resolved by setting the synapse reversal +# below for the NERLab-frame target (see IonotropicSynapse_e_syn). +aMN = AlphaMN__Pool( + recruitment_thresholds__array=recruitment_thresholds, + mode="active", +) # Create descending drive population -DD = DescendingDrive__Pool(n=nDD, process_type="gamma", shape=DDorder, timestep__ms=dt) +DD = DescendingDrive__Pool(n=nDD, poisson_batch_size=DDorder, timestep__ms=dt) # Create afferent populations -Ia = AffIa__Pool(n=nIa, timestep__ms=dt) # Primary spindle afferents -II = AffII__Pool(n=nII, timestep__ms=dt) # Lower thresholds for Group II -Ib = AffIb__Pool(n=nIb, timestep__ms=dt) # Golgi tendon organ afferents +Ia = AffIa__Pool(n=nIa, timestep__ms=dt) +II = AffII__Pool(n=nII, timestep__ms=dt) +Ib = AffIb__Pool(n=nIb, timestep__ms=dt) -# Create interneuron populations for reflex modulation -gII = GII__Pool(n=ngII) # Group II interneurons -gIb = GIb__Pool(n=ngIb) # Group Ib interneurons +# Create interneuron populations +gII = GII__Pool(n=ngII) +gIb = GIb__Pool(n=ngIb) -print(f"(OK) Created {len([aMN, DD, Ia, II, Ib, gII, gIb])} neural populations") +print(f"(OK) Created neural populations (MNs: {aMN.model} — napp + caL; gII/gIb: modern frame)") ############################################################################## # Create Proprioceptive Models # --------------------------- -# -# Initialize the sensory models that convert mechanical variables (muscle length, -# velocity, force) into neural signals that drive the afferent populations. # Golgi Tendon Organ - monitors muscle force/tension gto = GolgiTendonOrganModel( @@ -355,10 +283,12 @@ ) # Muscle Spindle - monitors muscle length and velocity +# Uses the full Mileusnic et al. (2006) 2nd-order ODE model matching NEURON. +spindle_params = SpindleModel.create_default_spindle_parameters() spin = SpindleModel( simulation_time__ms=tstop * pq.ms, time_step__ms=dt, - spindle_parameters=SpindleModel.create_default_spindle_parameters(), + spindle_parameters=spindle_params, ) print("(OK) Initialized proprioceptive models (spindle, GTO)") @@ -366,11 +296,7 @@ ############################################################################## # Create Muscle Model # ------------------ -# -# Initialize Hill-type muscle model. This model converts motor neuron -# activation into realistic force generation with appropriate biomechanical properties. -# Muscle model hill_muscle = HillModel( simulation_time__ms=tstop * pq.ms, time_step__ms=dt, @@ -382,406 +308,716 @@ ) print("(OK) Created muscle model") -print(f"\t- Motor units: {nType1 + nType2}") +print(f"\t- Motor units: {naMN} ({nType1} Type I + {nType2} Type II)") print(f"\t- Force capacity: ~{hill_muscle.F0:.1f} N") ############################################################################## -# Create Arrays for Real-Unit Forces -# ---------------------------------- -# -# Store musculotendon force in real units [N] for analysis +# Create Arrays for Force Tracking +# -------------------------------- -# Musculotendon force array [N] musculotendon_force__N = np.zeros_like(time) -print("(OK) Initialized force tracking arrays") - ############################################################################## -# Define Callback Functions -# ------------------------ -# -# These functions handle the integration of different system components during -# the simulation, including spike events and step-wise updates. - - -def spkEvent(i, muscle, delay): - muscle.add_spike(i, delay) - - -def eachStep( - muscle, - spin, - golgi, - popD, - ncD, - gMN, - joint_dyn, - step_counter, - tendon_tap_schedule, # List of (time_ms, magnitude_percent) tuples - tendon_tap_duration=100, # ms -): - """ - Step-wise integration function called at each simulation timestep. - - This function orchestrates the complex interactions between: - - Muscle mechanics and force generation - - Proprioceptive feedback from spindles and GTOs - - Joint dynamics and closed-loop control - - Neural population dynamics and spike generation - - Multiple tendon tap perturbations with varying amplitudes - - Parameters - ---------- - muscle : HillModel - Muscle model - spin : SpindleModel - Muscle spindle model - golgi : GolgiTendonOrganModel - Golgi tendon organ model - popD : dict - Dictionary of neural populations - ncD : dict - Dictionary of network connections - gMN : dict - Fusimotor drive parameters - joint_dyn : JointDynamics - Joint biomechanics model - step_counter : iterator - Simulation step counter - tendon_tap_schedule : list of tuples - List of (time_ms, magnitude_percent) defining tap timing and amplitudes - tendon_tap_duration : float - Duration (ms) of each tendon tap stretch - """ - i = next(step_counter) - - current_angle = artAng[i] - - # MUSCLE MECHANICS: Integrate muscle with current joint angle - L, V, A = muscle.integrate(current_angle) - - # TENDON TAP PERTURBATION: Triangular pulse (ramp up, then down) - current_time = h.t - - for tap_time, tap_magnitude in tendon_tap_schedule: - tap_end_time = tap_time + tendon_tap_duration - tap_midpoint = tap_time + (tendon_tap_duration / 2.0) - - if tap_time <= current_time < tap_end_time: - # Triangular pulse: ramp up to peak at midpoint, then ramp down - max_perturbation = L * (tap_magnitude / 100.0) - - if current_time < tap_midpoint: - # Rising edge: 0 → peak (first half) - progress = (current_time - tap_time) / (tendon_tap_duration / 2.0) - length_perturbation = max_perturbation * progress - else: - # Falling edge: peak → 0 (second half) - progress = (tap_end_time - current_time) / (tendon_tap_duration / 2.0) - length_perturbation = max_perturbation * progress +# Tendon Tap Schedule +# ------------------- - L_perturbed = L + length_perturbation +TENDON_TAP_DURATION = 100 # ms +TENDON_TAP_SCHEDULE = [ + # Phase 1: Ascending amplitudes without cortical drive + (500, 50.0), + (1000, 60.0), + (1500, 70.0), + (2000, 80.0), + # Phase 2: Repeat pattern with cortical drive + (3000, 50.0), + (3500, 60.0), + (4000, 70.0), + (4500, 80.0), +] - # Velocity reflects the triangle slope - if current_time < tap_midpoint: - # Rising: positive velocity - V_perturbed = V + (max_perturbation / (tendon_tap_duration / 2000.0)) - else: - # Falling: negative velocity - V_perturbed = V - (max_perturbation / (tendon_tap_duration / 2000.0)) - - # Acceleration spikes at tap onset and reverses at midpoint - if abs(current_time - tap_time) < dt.magnitude: - A_perturbed = A + (max_perturbation / ((tendon_tap_duration / 2000.0) ** 2)) - print( - f"\t>> TENDON TAP (triangular) at t={current_time:.1f} ms: " - f"{tap_magnitude}% peak stretch ({max_perturbation:.4f} L0), " - f"duration={tendon_tap_duration} ms" - ) - print(f"\t Baseline length: {L:.4f} L0 → Peak: {L + max_perturbation:.4f} L0") - elif abs(current_time - tap_midpoint) < dt.magnitude: - A_perturbed = A - 2 * (max_perturbation / ((tendon_tap_duration / 2000.0) ** 2)) - else: - A_perturbed = A +print("\nTendon tap schedule:") +print(f"\t- Duration per tap: {TENDON_TAP_DURATION} ms") +print(f"\t- Phase 1: {len([t for t, m in TENDON_TAP_SCHEDULE if t < 2500])} taps") +print(f"\t- Phase 2: {len([t for t, m in TENDON_TAP_SCHEDULE if t >= 2500])} taps") - # Use perturbed values for spindle feedback - L, V, A = L_perturbed, V_perturbed, A_perturbed - break # Only one tap active at a time +############################################################################## +# Causally-Correct Closed-Loop Simulation Setup +# ---------------------------------------------- +# +# A single Python for-loop advances all components together each timestep: +# - Jaxley neural network (gII + gIb + MN): stepped with build_init_and_step_fn +# - DD/Ia/II/Ib afferent cells: stepped via their existing integrate() API +# - Hill muscle, spindle, GTO, joint dynamics: stepped via integrate() +# +# Afferent feedback at step t affects neural input at step t+1 (1-step delay, +# ~0.1 ms at dt=0.1 ms — physiologically negligible). - # PROPRIOCEPTIVE FEEDBACK: Use muscle kinematics for spindle feedback - Iay, IIy = spin.integrate(L, V, A, gMN["gDyn"][i], gMN["gStat"][i]) +print("\nStarting causally-correct closed-loop simulation (Jaxley)...") +print(f"\tDuration: {tstop} ms") +print(f"\tTimestep: {dt} ms") - # FORCE FEEDBACK: Muscle force for GTO - f = muscle.F0 * muscle.muscle_force[i] # Musculotendon force [N] - musculotendon_force__N[i] = f # Store in real units - Iby = golgi.integrate(f) +dt_ms = float(dt.rescale(pq.ms).magnitude) +dt_s = float(dt.rescale(pq.s).magnitude) +n_steps = len(time) + +# Synaptic parameters — matched to NEURON ex11 as starting point. +# TODO: validate EPSP/IPSP amplitudes; cable-dendrite Jaxley MNs may need +# retuning since membrane impedance differs from point-neuron NEURON model. +# IIR / Exp2Syn equivalence factor (~0.06) for the EXCITATORY drives onto MNs. +# NEURON uses live-conductance Exp2Syn synapses with shaped onset; this script +# uses an IIR + fixed driving force ("g × (e_exc_mn - V_mn)") approximation +# which delivers ~10× more effective current at the same nominal weight, so +# the DD and Ia weights here are scaled down to keep MN recruitment in the +# physiological window. See ex03 for the same equivalence factor analysis. +# in_weight (postsynaptic = modern-frame interneurons) and base_inh_weight +# (postsynaptic = NERLab MN, but via a real IonotropicSynapse with live V) +# don't share the factor and stay at the NEURON values. +base_dd_weight = 0.01 # µS — excitatory DD → MN (NEURON: 0.05 µS, ×0.2 for IIR) +base_ia_weight = 0.01 # µS — excitatory Ia → MN (NEURON: 0.05 µS, ×0.2 for IIR) +in_weight = 0.025 # µS — excitatory II → gII, Ib → gIb (NEURON: 0.025 µS) +base_inh_weight = 0.05 # µS — inhibitory gII/gIb → MN (NEURON: 0.05 µS) +tau_syn = 5.0 # ms — synaptic exponential decay +e_exc = 0.0 # mV — excitatory reversal potential +v_rest = -70.0 # mV — resting potential + +# Per-MN Henneman (size-principle) current scaling +threshold_min = recruitment_thresholds.min() +threshold_max = recruitment_thresholds.max() +normalized_thresholds = (recruitment_thresholds - threshold_min) / (threshold_max - threshold_min) +mn_current_scale = np.exp(-1.0 * normalized_thresholds) + +# Spike recording arrays +dd_spike_times = [[] for _ in range(nDD)] +ia_spike_times = [[] for _ in range(nIa)] +ii_spike_times = [[] for _ in range(nII)] +ib_spike_times = [[] for _ in range(nIb)] +gii_spike_times = [[] for _ in range(ngII)] +gib_spike_times = [[] for _ in range(ngIb)] +mn_spike_times = [[] for _ in range(naMN)] + +# Muscle/spindle/GTO time series for plotting +muscle_length_series = [] +muscle_velocity_series = [] +muscle_force_series = [] +ia_firing_series = [] +ii_firing_series = [] +ib_firing_series = [] + +# Membrane trace storage for selected MNs +vm_traces = {} +vm_cell_indices = [0, 10, 20, 30, 40] +for _ci in vm_cell_indices: + if _ci < naMN: + vm_traces[_ci] = np.zeros(n_steps) - # CLOSED-LOOP DYNAMICS: Update joint angle based on muscle torque - new_angle, _ = joint_dyn.integrate( - torque__Nm=muscle.signed_muscle_torque[i], - dt__s=float(dt.rescale(pq.s).magnitude), - ) - # Update joint angle array for next timestep - if i < len(artAng) - 1: - artAng[i + 1] = new_angle - - # DESCENDING DRIVE PROCESSING: Convert cortical signals to spikes - for DDcell in popD["DD"]: - if DDcell.integrate(DDdrive[i]): - spike_time = h.t + 1 - if spike_time < tstop: - ncD["cmd->DD"][DDcell.pool__ID].event(spike_time) - - # AFFERENT PROCESSING: Convert sensory signals to neural spikes - for Ia in popD["Ia"]: - if Iay >= Ia.RT: - if Ia.integrate(Iay): - # Ensure h.t is converted to a quantities object with units of ms - spike_time = h.t + float(Ia.axon_delay__ms) - if spike_time < tstop: - ncD["Spindle->Ia"][Ia.pool__ID].event(spike_time) - - ii_spikes = 0 - for II in popD["II"]: - if IIy >= II.RT: - if II.integrate(IIy): - spike_time = h.t + float(II.axon_delay__ms) - if spike_time < tstop: - ncD["Spindle->II"][II.pool__ID].event(spike_time) - ii_spikes += 1 - - ib_spikes = 0 - for Ib in popD["Ib"]: - if Iby >= Ib.RT: - if Ib.integrate(Iby): - spike_time = h.t + float(Ib.axon_delay__ms) - if spike_time < tstop: - ncD["GTO->Ib"][Ib.pool__ID].event(spike_time) - ib_spikes += 1 +############################################################################## +# Connectivity Pre-computation +# ---------------------------- + +print("\nPre-computing neural connectivity...") +# DD → MN (forward: mn_idx → dd_list; reverse: dd_idx → mn_list) +dd_to_mn_connections = { + mn_idx: [j for j in range(nDD) if myogen.get_random_generator().random() < 0.3] + for mn_idx in range(naMN) +} +dd_to_mn_rev = {dd: [] for dd in range(nDD)} +for mn_idx, dd_list in dd_to_mn_connections.items(): + for dd_idx in dd_list: + dd_to_mn_rev[dd_idx].append(mn_idx) + +# Ia → MN (forward and reverse) +ia_to_mn_connections = { + mn_idx: [j for j in range(nIa) if myogen.get_random_generator().random() < 0.8] + for mn_idx in range(naMN) +} +ia_to_mn_rev = {ia: [] for ia in range(nIa)} +for mn_idx, ia_list in ia_to_mn_connections.items(): + for ia_idx in ia_list: + ia_to_mn_rev[ia_idx].append(mn_idx) + +# II → gII (forward and reverse) +ii_to_gii_connections = { + gii_idx: [j for j in range(nII) if myogen.get_random_generator().random() < 0.3] + for gii_idx in range(ngII) +} +ii_to_gii_rev = {ii: [] for ii in range(nII)} +for gii_idx, ii_list in ii_to_gii_connections.items(): + for ii_idx in ii_list: + ii_to_gii_rev[ii_idx].append(gii_idx) + +# Ib → gIb (forward and reverse) +ib_to_gib_connections = { + gib_idx: [j for j in range(nIb) if myogen.get_random_generator().random() < 0.3] + for gib_idx in range(ngIb) +} +ib_to_gib_rev = {ib: [] for ib in range(nIb)} +for gib_idx, ib_list in ib_to_gib_connections.items(): + for ib_idx in ib_list: + ib_to_gib_rev[ib_idx].append(gib_idx) + +# Pre-compute MN axonal delays (ms) — used when scheduling spikes into hill model +axonal_delays_ms = np.array([ + 1.0 + (mn.axon_delay__ms.magnitude if hasattr(mn.axon_delay__ms, 'magnitude') + else float(mn.axon_delay__ms)) + for mn in aMN +]) ############################################################################## -# Create Neural Network +# JAX Physiology Setup # -------------------- -# -# Assemble all neural populations into a connected network that implements -# the spinal reflex circuitry for single muscle control. - -network = Network( - { - "DD": DD, # Descending drive - "aMN": aMN, # α-motoneurons - "Ia": Ia, # Group Ia afferents - "II": II, # Group II afferents - "Ib": Ib, # Group Ib afferents - "gII": gII, # Group II interneurons - "gIb": gIb, # Group Ib interneurons - "gMN": gMN, # Fusimotor parameters - } -) +# Build param dicts and initial states for JAX functional step functions. +# hill_init_params calls ForceSatParams (scipy) once — never called again. + +hillD_default = HillModel.create_default_muscle_parameters() +hill_p = hill_init_params(hillD_default, nType1, nType2, dt_ms) +spindle_p = spindle_params_from_dict(spindle_params) +gto_p = gto_params_from_dict(GolgiTendonOrganModel.create_default_gto_parameters()) +joint_p = { + "inertia": joint_dynamics.inertia__kg_m2, + "damping": joint_dynamics.damping__Nm_s_per_rad, + "stiffness": joint_dynamics.stiffness__Nm_per_rad, +} + +# Initial muscle length from HillModel's auto-computed rest length +L0_init = float(hill_muscle._hill_model.L[0]) +# MN spike buffer depth: max axonal delay + safety margin +_max_mn_delay_steps = int(np.ceil(axonal_delays_ms.max() / dt_ms)) + 4 +hill_state = hill_init_state(L0_init, naMN, max_delay_steps=_max_mn_delay_steps) +spindle_state = spindle_init() +gto_state = gto_init() +joint_state = joint_init(angle_deg=float(np.degrees(artAng[0]))) + +# Per-MU axonal delay in timesteps (for hill_step spike buffer) +delay_steps_arr = np.maximum(1, np.round(axonal_delays_ms / dt_ms)).astype(np.int32) -print(f"(OK) Created network with {len(network.populations)} populations") +print("(OK) JAX physiology states initialised") +print(f"\t- Hill L0_init = {L0_init:.4f}, max_delay_steps = {_max_mn_delay_steps}") ############################################################################## -# Configure Neural Connections -# --------------------------- -# -# Establish synaptic connections that implement physiologically realistic -# spinal reflex pathways, including both excitatory and inhibitory connections. -# -# Neural Circuit Diagram: -# ---------------------- -# -# Spindle (length/velocity) GTO (force) Cortex -# | | | | -# v v v v -# Ia II Ib DD -# | | | | -# | v v | -# | gII (-) gIb (-) | -# | | | | -# +----------+--------(+)------------+----------(+)-----+ -# | -# v -# α-MN -----> Muscle -# | -# v -# Joint Motion -# | -# +--------------------------------+ -# | | -# v v -# Spindle GTO -# (feedback) (feedback) +# Build Combined Jaxley Network +# ---------------------------- + +print("Building combined jx.Network (gII + gIb + MN)...") + +# Collect and reset all cells +gii_cells = [] +for gii_cell in gII: + cell = gii_cell.cell + cell.delete_recordings() + cell.delete_stimuli() + cell.set("v", -70.0) + cell.init_states() + gii_cells.append(cell) + +gib_cells = [] +for gib_cell in gIb: + cell = gib_cell.cell + cell.delete_recordings() + cell.delete_stimuli() + cell.set("v", -70.0) + cell.init_states() + gib_cells.append(cell) + +mn_cells = [] +for mn in aMN: + cell = mn.cell + cell.delete_recordings() + cell.delete_stimuli() + cell.set("v", 0.0) # NERLab resting potential (1952-HH frame) + cell.init_states() + mn_cells.append(cell) + +n_gii = len(gii_cells) +n_gib = len(gib_cells) +n_mn = len(mn_cells) + +combined_net = jx.Network(gii_cells + gib_cells + mn_cells) + +# Register recordings in order: gII soma, gIb soma, MN soma +for i in range(n_gii): + combined_net.cell(i).branch(0).loc(0.5).record("v") +for i in range(n_gib): + combined_net.cell(n_gii + i).branch(0).loc(0.5).record("v") +for mn_idx in range(n_mn): + combined_net.cell(n_gii + n_gib + mn_idx).branch(0).loc(0.5).record("v") + +# Register stimulus sites with placeholder arrays (establishes external_inds order) +# Must match recording order: gII, gIb, MN +placeholder = jnp.zeros(n_steps) +for i in range(n_gii): + combined_net.cell(i).branch(0).loc(0.5).stimulate(placeholder) +for i in range(n_gib): + combined_net.cell(n_gii + i).branch(0).loc(0.5).stimulate(placeholder) +for mn_idx in range(n_mn): + combined_net.cell(n_gii + n_gib + mn_idx).branch(0).loc(0.5).stimulate(placeholder) + +# Inhibitory synapses: gII → MN and gIb → MN (handled internally by Jaxley step_fn) +print(" Adding inhibitory synapses (gII → MN, gIb → MN, p=0.3)...") +sparse_connect( + combined_net.cell(list(range(n_gii))), + combined_net.cell(list(range(n_gii + n_gib, n_gii + n_gib + n_mn))), + IonotropicSynapse(), + p=0.3, +) +sparse_connect( + combined_net.cell(list(range(n_gii, n_gii + n_gib))), + combined_net.cell(list(range(n_gii + n_gib, n_gii + n_gib + n_mn))), + IonotropicSynapse(), + p=0.3, +) +combined_net.set("IonotropicSynapse_gS", base_inh_weight) +# Inhibitory reversal on NERLab MN targets — set in the NERLab voltage frame. # -# Legend: (+) = excitatory, (-) = inhibitory +# We want the same *physical* inhibitory driving force the original Powers2017 +# setup delivered: +# Powers2017 frame: e_syn = -80 mV, V_rest = -65 mV +# → driving = e_syn - V_rest = -15 mV (mild GABA-like) +# In the NERLab frame (V_rest = 0 mV) the analogous value is +# NERLab frame: e_syn = -15 mV, V_rest = 0 mV +# → driving = -15 mV (same physical effect) # -# Key pathways: -# 1. Ia → MN: Monosynaptic stretch reflex (EXCITATORY, fast) -# 2. II → gII → MN: Polysynaptic length pathway (INHIBITORY, moderate speed) -# 3. Ib → gIb → MN: Autogenic inhibition (INHIBITORY, slower - force-dependent) -# 4. DD → MN: Descending cortical drive (EXCITATORY, voluntary control) - -# DESCENDING CONTROL: Direct cortical drive to motor neurons (EXCITATORY) -network.connect("DD", "aMN", probability=0.3, weight__uS=0.05 * pq.uS) - -# MONOSYNAPTIC STRETCH REFLEX: Ia afferents excite motor neurons (EXCITATORY) -# High connection probability (80%) reflects homogeneous excitatory distribution -network.connect("Ia", "aMN", probability=0.8, weight__uS=0.05 * pq.uS) - -# POLYSYNAPTIC PATHWAYS: Group II pathway (afferents → interneurons → motor neurons) -# Lower connection probability (30%) captures predominantly indirect, disynaptic influence -network.connect( - "II", "gII", probability=0.3, weight__uS=0.025 * pq.uS -) # Afferent → Interneuron (EXCITATORY) -network.connect( - "gII", "aMN", probability=0.3, weight__uS=0.05 * pq.uS, inhibitory=True -) # Interneuron → MN (INHIBITORY, weak modulatory effect) - -# INHIBITORY REFLEXES: Force feedback through Ib interneurons (autogenic inhibition) -# Lower connection probability (30%) captures predominantly indirect, disynaptic influence -network.connect( - "Ib", "gIb", probability=0.3, weight__uS=0.025 * pq.uS -) # Afferent → Interneuron (EXCITATORY) -network.connect( - "gIb", "aMN", probability=0.3, weight__uS=0.05 * pq.uS, inhibitory=True -) # Interneuron → MN (INHIBITORY, weak modulatory effect) - -print("(OK) Configured synaptic connections") -print("\t- Excitatory pathways:") -print("\t\t• DD→MN (descending drive)") -print("\t\t• Ia→MN (monosynaptic stretch reflex)") -print("\t\t• Afferents→Interneurons: II→gII, Ib→gIb") -print("\t- Inhibitory pathways:") -print("\t\t• gII→MN (Group II interneuron inhibition)") -print("\t\t• gIb→MN (autogenic inhibition from force feedback)") +# NEURON ex11 uses syn.e = -75 mV literally without a frame shift (see +# myogen.simulator.neuron.cells:776), which on a NERLab MN gives ~75 mV of +# hyperpolarising driving force. That's far stronger than the Powers2017 +# equivalent and would over-inhibit MNs into silence here. We match the +# physical driving force instead of the literal mV value, so the recruitment +# behaviour is comparable across the two backends. +combined_net.set("IonotropicSynapse_e_syn", -15.0) +combined_net.set("IonotropicSynapse_k_minus", 0.1) # 10 ms IPSP decay +# Spike activation threshold — applies to the PRESYNAPTIC cell V crossing. +# Sources are gII/gIb interneurons (modern frame, V_rest ≈ -70 mV), so the +# -35 mV threshold below is correct in their frame, not the MN's NERLab frame. +combined_net.set("IonotropicSynapse_v_th", -35.0) +combined_net.set("IonotropicSynapse_delta", 10.0) # voltage sensitivity (mV) + +# Compile Jaxley single-step function +print(" Compiling Jaxley step function...") +combined_net.to_jax() +init_fn, step_fn = build_init_and_step_fn(combined_net) +params = combined_net.get_parameters() +external_inds = combined_net.external_inds.copy() +rec_inds = combined_net.recordings.rec_index.to_numpy() +states, params = init_fn(params) +step_fn_jit = jax.jit(step_fn) + +print(f" Network: {n_gii + n_gib + n_mn} cells, " + f"{n_gii + n_gib + n_mn} stimulus sites, {len(rec_inds)} recording sites") ############################################################################## -# Connect Motors to Muscle -# ------------------------ +# Run Closed-Loop Simulation with lax.scan +# ----------------------------------------- # -# Establish the neuromuscular junctions that convert motor neuron spikes -# into muscle fiber activation. - -# Connect motor neurons to muscle (spike threshold automatically uses population default of 50.0 mV) -network.connect_to_muscle( - "aMN", muscle=hill_muscle, activation_callback=spkEvent, weight__uS=1.0 * pq.uS +# lax.scan compiles the entire loop into a single GPU/XLA kernel, eliminating +# Python interpreter overhead at each of the 200k timesteps (dt=0.025ms, 5s). +# +# All Python state (afferent cells, DD cells) is replaced by JAX-native +# generators (poisson_step, gamma_step) inside the compiled scan body. +# Per-cell axonal delays are approximated with per-population mean FIFO queues. + +print("\n[Simulation] Building lax.scan closed-loop...") + +decay = np.exp(-dt_ms / tau_syn) + +# --- Pre-compute tendon tap coefficient arrays (normalised, no L dependence) --- +# update_physiology applies: L_sp = L*(1+tap_dL), V_sp = V + L*tap_dV +tap_dL = np.zeros(n_steps, dtype=np.float32) +tap_dV = np.zeros(n_steps, dtype=np.float32) +_half_tap = TENDON_TAP_DURATION / 2.0 # ms +_dV_rate = 1.0 / (_half_tap * 1e-3) # 1/s normalised by L inside update_physiology +for _tap_t, _tap_mag in TENDON_TAP_SCHEDULE: + _tap_end = _tap_t + TENDON_TAP_DURATION + _frac = _tap_mag / 100.0 + for _ii, _t in enumerate(time): + if _tap_t <= _t < _tap_end: + if _t < _tap_t + _half_tap: + _prog = (_t - _tap_t) / _half_tap + tap_dL[_ii] = _frac * _prog + tap_dV[_ii] = _frac * _dV_rate + else: + _prog = (_tap_end - _t) / _half_tap + tap_dL[_ii] = _frac * _prog + tap_dV[_ii] = -_frac * _dV_rate + +# --- Dense connectivity matrices (pre_idx → post_idx) --- +print(" Building connectivity matrices...") +# dd_to_mn_rev: {dd_idx: [mn_idx]} == forward map for make_connectivity_matrix +dd_to_mn_mat = make_connectivity_matrix(dd_to_mn_rev, nDD, naMN) # (nDD, naMN) +ia_to_mn_mat = make_connectivity_matrix(ia_to_mn_rev, nIa, naMN) # (nIa, naMN) +ii_to_gii_mat = make_connectivity_matrix(ii_to_gii_rev, nII, ngII) # (nII, ngII) +ib_to_gib_mat = make_connectivity_matrix(ib_to_gib_rev, nIb, ngIb) # (nIb, ngIb) + +# --- Afferent response thresholds and gamma-ISI shape parameters --- +ia_rts = jnp.array([c.RT for c in Ia], dtype=jnp.float32) +ii_rts = jnp.array([c.RT for c in II], dtype=jnp.float32) +ib_rts = jnp.array([c.RT for c in Ib], dtype=jnp.float32) +ia_shape = 1.0 # exponential ISI (pure Poisson); increase for more regular firing +ii_shape = 1.0 +ib_shape = 1.0 + +# --- Per-cell axonal delay steps (FIFO queue depth per afferent cell) --- +def _cell_ds(cell): + d = cell.axon_delay__ms.magnitude if hasattr(cell.axon_delay__ms, 'magnitude') else float(cell.axon_delay__ms) + return max(1, int(round(d / dt_ms))) + +# Per-cell delay arrays (sorted by pool__ID so index matches gamma generator order) +ia_delay_steps_arr = np.array([_cell_ds(c) for c in sorted(Ia, key=lambda c: c.pool__ID)], dtype=np.int32) +ii_delay_steps_arr = np.array([_cell_ds(c) for c in sorted(II, key=lambda c: c.pool__ID)], dtype=np.int32) +ib_delay_steps_arr = np.array([_cell_ds(c) for c in sorted(Ib, key=lambda c: c.pool__ID)], dtype=np.int32) +max_ia_delay_steps = int(ia_delay_steps_arr.max()) +max_ii_delay_steps = int(ii_delay_steps_arr.max()) +max_ib_delay_steps = int(ib_delay_steps_arr.max()) +print(f" Per-cell afferent delays — Ia: {ia_delay_steps_arr.min()}–{ia_delay_steps_arr.max()} steps " + f"({ia_delay_steps_arr.min()*dt_ms:.2f}–{ia_delay_steps_arr.max()*dt_ms:.2f} ms), " + f"II: {ii_delay_steps_arr.min()}–{ii_delay_steps_arr.max()} steps, " + f"Ib: {ib_delay_steps_arr.min()}–{ib_delay_steps_arr.max()} steps") + +# --- Build scan_step closure (closes over all static params) --- +print(" Building scan step function...") +scan_step = make_scan_step( + jaxley_step_fn = step_fn, + jaxley_params = params, + external_inds = external_inds, + rec_inds = rec_inds, + n_gii = n_gii, + n_gib = n_gib, + n_mn = n_mn, + ia_rts = ia_rts, + ii_rts = ii_rts, + ib_rts = ib_rts, + ia_shape = ia_shape, + ii_shape = ii_shape, + ib_shape = ib_shape, + dd_N_batch = DDorder, + dd_to_mn_mat = dd_to_mn_mat, + ia_to_mn_mat = ia_to_mn_mat, + ii_to_gii_mat = ii_to_gii_mat, + ib_to_gib_mat = ib_to_gib_mat, + ia_delay_steps_arr = ia_delay_steps_arr, + ii_delay_steps_arr = ii_delay_steps_arr, + ib_delay_steps_arr = ib_delay_steps_arr, + delay_steps = delay_steps_arr, + hill_p = hill_p, + spindle_p = spindle_p, + gto_p = gto_p, + joint_p = joint_p, + base_dd_weight = base_dd_weight, + base_ia_weight = base_ia_weight, + in_weight = in_weight, + e_exc = e_exc, # AMPA reversal for interneurons (modern frame) + v_rest = v_rest, + e_exc_mn = 70.0, # AMPA reversal for NERLab MNs (1952-HH frame) + mn_spike_threshold_mV = 50.0, # NERLab AP detection threshold + mn_current_scale = mn_current_scale, + tau_syn_decay = float(decay), + dt_ms = dt_ms, + dt_s = dt_s, ) -print("(OK) Connected motor neurons to muscle") +# --- Initial carry --- +init_carry = { + "neural": states, # Jaxley neural states from init_fn + "phys": { + "hill": hill_state, + "spindle": spindle_state, + "gto": gto_state, + "joint": joint_state, + }, + "g_dd": jnp.zeros(n_mn, dtype=jnp.float32), + "g_ia": jnp.zeros(n_mn, dtype=jnp.float32), + "g_ii": jnp.zeros(n_gii, dtype=jnp.float32), + "g_ib": jnp.zeros(n_gib, dtype=jnp.float32), + # prev_v is just a "below every threshold" sentinel for first-iteration spike + # detection; -70 sits below both the gII/gIb spike threshold (-35 mV, modern + # frame) and the NERLab MN threshold (+50 mV), so the first real V crossing + # is detected correctly for every cell regardless of which frame it lives in. + "prev_v": jnp.full(n_gii + n_gib + n_mn, -70.0, dtype=jnp.float32), + "dd_st": poisson_init(nDD, DDorder, seed=42), + "ia_st": gamma_init(nIa, ia_shape, seed=43), + "ii_st": gamma_init(nII, ii_shape, seed=44), + "ib_st": gamma_init(nIb, ib_shape, seed=45), + "prev_Iay": jnp.float32(0.0), + "prev_IIy": jnp.float32(0.0), + "prev_Iby": jnp.float32(0.0), + "ia_delay_buf": jnp.zeros((nIa, max_ia_delay_steps), dtype=jnp.float32), + "ii_delay_buf": jnp.zeros((nII, max_ii_delay_steps), dtype=jnp.float32), + "ib_delay_buf": jnp.zeros((nIb, max_ib_delay_steps), dtype=jnp.float32), +} -############################################################################## -# Configure External Inputs -# ------------------------ -# -# Setup external input pathways for sensory feedback and descending commands. - -network.connect_from_external("cmd", "DD", weight__uS=1.0 * pq.uS) -network.connect_from_external("spindle", "Ia", weight__uS=1.0 * pq.uS) -network.connect_from_external("spindle", "II", weight__uS=1.0 * pq.uS) -network.connect_from_external("gto", "Ib", weight__uS=1.0 * pq.uS) - -# Get NetCons for manual triggering during simulation -ncD = { - "cmd->DD": network.get_netcons("cmd", "DD"), - "Spindle->Ia": network.get_netcons("spindle", "Ia"), - "Spindle->II": network.get_netcons("spindle", "II"), - "GTO->Ib": network.get_netcons("gto", "Ib"), +# --- Per-step inputs stacked as (n_steps, ...) --- +scan_inputs = { + "DDdrive": jnp.array(DDdrive, dtype=jnp.float32), + "gDyn": jnp.array(gDyn, dtype=jnp.float32), + "gStat": jnp.array(gStat, dtype=jnp.float32), + "tap_dL": jnp.array(tap_dL, dtype=jnp.float32), + "tap_dV": jnp.array(tap_dV, dtype=jnp.float32), +} + +# --- Execute lax.scan --- +# First call triggers XLA compilation (~30-60s); subsequent calls are fast. +print(f"\n[Simulation] Running lax.scan over {n_steps} steps " + f"({tstop:.0f} ms at dt={dt_ms} ms)...") +_, scan_out = jax.jit(lambda c, xs: jax.lax.scan(scan_step, c, xs))(init_carry, scan_inputs) +jax.block_until_ready(scan_out) +print("Scan complete. Extracting results...") + +# --- Convert boolean spike arrays → spike-time lists --- +# scan_out["mn_spikes"] shape: (n_steps, n_mn), dtype bool +dd_spike_arr = np.array(scan_out["dd_spikes"]) # (n_steps, nDD) +mn_spike_arr = np.array(scan_out["mn_spikes"]) # (n_steps, n_mn) +gii_spike_arr = np.array(scan_out["gii_spikes"]) # (n_steps, n_gii) +gib_spike_arr = np.array(scan_out["gib_spikes"]) # (n_steps, n_gib) +ia_spike_arr = np.array(scan_out["ia_spikes"]) # (n_steps, nIa) +ii_spike_arr = np.array(scan_out["ii_spikes"]) # (n_steps, nII) +ib_spike_arr = np.array(scan_out["ib_spikes"]) # (n_steps, nIb) + +dd_spike_times = [time[dd_spike_arr[:, dd]].tolist() for dd in range(nDD)] +mn_spike_times = [time[mn_spike_arr[:, mn_idx]].tolist() for mn_idx in range(n_mn)] +gii_spike_times = [time[gii_spike_arr[:, gi]].tolist() for gi in range(n_gii)] +gib_spike_times = [time[gib_spike_arr[:, gi]].tolist() for gi in range(n_gib)] +ia_spike_times = [time[ia_spike_arr[:, ia]].tolist() for ia in range(nIa)] +ii_spike_times = [time[ii_spike_arr[:, ii]].tolist() for ii in range(nII)] +ib_spike_times = [time[ib_spike_arr[:, ib]].tolist() for ib in range(nIb)] + +# --- Time series (physiology outputs) --- +_force_norm_arr = np.array(scan_out["force"]) # normalised +_torque_norm_arr = np.array(scan_out["torque"]) # normalised +muscle_length_series = np.array(scan_out["L"]).tolist() +muscle_velocity_series = np.gradient(np.array(scan_out["L"]), dt_s).tolist() +muscle_force_series = (hill_p["F0"] * _force_norm_arr).tolist() +muscle_torque_series = _torque_norm_arr.tolist() +ia_firing_series = np.array(scan_out["Iay"]).tolist() +ii_firing_series = np.array(scan_out["IIy"]).tolist() +ib_firing_series = np.array(scan_out["Iby"]).tolist() + +# --- Force and joint angle arrays (overwrite initialised zeros) --- +musculotendon_force__N = hill_p["F0"] * _force_norm_arr +artAng = np.radians(np.array(scan_out["angle_deg"])) + +# --- Membrane potential traces for selected MNs --- +vm_cell_indices = [0, 10, 20, 30, 40] +_v_mn_series = np.array(scan_out["v_mn"]) # (n_steps, n_mn) +vm_traces = { + _ci: _v_mn_series[:, _ci] + for _ci in vm_cell_indices + if _ci < n_mn } -print("(OK) Configured external input pathways") +# Summary +active_mn_count = sum(1 for s in mn_spike_times if s) +active_gii_count = sum(1 for s in gii_spike_times if s) +active_gib_count = sum(1 for s in gib_spike_times if s) +print(f"\nSimulation complete!") +print(f" MNs active : {active_mn_count}/{naMN}") +print(f" gII active : {active_gii_count}/{ngII}") +print(f" gIb active : {active_gib_count}/{ngIb}") +print(f" Peak force : {max(muscle_force_series) if muscle_force_series else 0:.4f} F0") +print(f" Peak angle : {np.degrees(np.max(artAng)):.2f} deg") +print(f" Peak force (norm) : {max(muscle_force_series) / hill_p['F0']:.4f} F0" if muscle_force_series else " Peak force: (none)") ############################################################################## -# Prepare Simulation Models -# ------------------------ +# Convert Results to Neo Format +# ----------------------------- +# +# Store all simulation data in a Neo Block matching NEURON's structure: +# - 7 population segments (aMN, Ia, II, Ib, gII, gIb, DD) with spike trains +# - Muscle segment with dynamics AnalogSignals +# - Spindle segment with proprioceptive AnalogSignals +# - GTO segment with force feedback AnalogSignal + +results = Block(name="Spinal Network - Jaxley") +t_stop_s = (tstop * pq.ms).rescale(pq.s) +sampling_period_s = dt.rescale(pq.s) +n_steps = len(time) + +def _make_spiketrain(spikes_ms, name): + """Create Neo SpikeTrain, filtering out any spikes beyond t_stop (e.g. from axon delays).""" + arr = np.array(spikes_ms) + if len(arr) > 0: + arr = arr[arr < tstop] + times_s = (arr * pq.ms).rescale(pq.s) if len(arr) > 0 else np.array([]) * pq.s + return SpikeTrain(times_s, t_stop=t_stop_s, name=name) + +# --- Population segments (matching NEURON order) --- + +# 1. aMN segment (spike trains + membrane potential traces) +mn_segment = Segment(name="aMN") +mn_segment.spiketrains = [_make_spiketrain(spikes, f"aMN_{i}") for i, spikes in enumerate(mn_spike_times)] +for cell_idx, trace in vm_traces.items(): + mn_segment.analogsignals.append( + AnalogSignal( + np.array(trace).reshape(-1, 1) * pq.mV, + sampling_period=sampling_period_s, + name=f"aMN_cell{cell_idx}_Vm", + cell_idx=cell_idx, + ) + ) +results.segments.append(mn_segment) + +# 2. Ia segment +ia_segment = Segment(name="Ia") +ia_segment.spiketrains = [_make_spiketrain(spikes, f"Ia_{i}") for i, spikes in enumerate(ia_spike_times)] +results.segments.append(ia_segment) + +# 3. II segment +ii_segment = Segment(name="II") +ii_segment.spiketrains = [_make_spiketrain(spikes, f"II_{i}") for i, spikes in enumerate(ii_spike_times)] +results.segments.append(ii_segment) + +# 4. Ib segment +ib_segment = Segment(name="Ib") +ib_segment.spiketrains = [_make_spiketrain(spikes, f"Ib_{i}") for i, spikes in enumerate(ib_spike_times)] +results.segments.append(ib_segment) + +# 5. gII segment +gii_segment = Segment(name="gII") +gii_segment.spiketrains = [_make_spiketrain(spikes, f"gII_{i}") for i, spikes in enumerate(gii_spike_times)] +results.segments.append(gii_segment) + +# 6. gIb segment +gib_segment = Segment(name="gIb") +gib_segment.spiketrains = [_make_spiketrain(spikes, f"gIb_{i}") for i, spikes in enumerate(gib_spike_times)] +results.segments.append(gib_segment) + +# 7. DD segment +dd_segment = Segment(name="DD") +dd_segment.spiketrains = [_make_spiketrain(spikes, f"DD_{i}") for i, spikes in enumerate(dd_spike_times)] +results.segments.append(dd_segment) + +# --- Model segments (muscle, spindle, GTO) --- + +# 8. Hill muscle segment +muscle_segment = Segment(name="hill_muscle") + +# Muscle length — always from scan output (hill_muscle.muscle_length is zero-initialized) +muscle_length_arr = np.array(muscle_length_series[:n_steps]) +muscle_segment.analogsignals.append( + AnalogSignal( + muscle_length_arr.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="hill_muscle_muscle_length", + attr_name="muscle_length", + ) +) -# Package all models for the simulation runner -models = { - "hill_muscle": hill_muscle, - "spin": spin, - "gto": gto, - "joint": joint_dynamics, -} +# Muscle force (from JAX hill_step series collected during loop) +muscle_force_arr = np.array(muscle_force_series[:n_steps]) +muscle_segment.analogsignals.append( + AnalogSignal( + muscle_force_arr.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="hill_muscle_muscle_force", + attr_name="muscle_force", + ) +) + +# Muscle torque (normalised, from JAX hill_step series collected during loop) +muscle_torque_arr = np.array(muscle_torque_series[:n_steps]) if muscle_torque_series else np.zeros(n_steps) +muscle_segment.analogsignals.append( + AnalogSignal( + muscle_torque_arr.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="hill_muscle_muscle_torque", + attr_name="muscle_torque", + ) +) +# TypeI/TypeII summed activations — from scan outputs +for _scan_key, _attr in [("type1_act", "type1_activation"), ("type2_act", "type2_activation")]: + _act_arr = np.array(scan_out[_scan_key])[:n_steps] + muscle_segment.analogsignals.append( + AnalogSignal( + _act_arr.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name=f"hill_muscle_{_attr}", + attr_name=_attr, + ) + ) -# Tendon tap parameters - multiple taps with increasing amplitudes -# Pattern: Taps every 0.5s with increasing amplitude, reset at 2.5s -# High amplitudes to elicit strong reflex responses -# Phase 1 (0-2.5s): 50%, 60%, 70%, 80% -# Phase 2 (2.5-5s): 50%, 60%, 70%, 80% +results.segments.append(muscle_segment) -TENDON_TAP_DURATION = 100 # ms - duration of each stretch -TENDON_TAP_SCHEDULE = [ - # Phase 1: Ascending amplitudes without cortical drive - (500, 50.0), # 0.5s: 50% - (1000, 60.0), # 1.0s: 60% - (1500, 70.0), # 1.5s: 70% - (2000, 80.0), # 2.0s: 80% - # Phase 2: Repeat pattern with cortical drive - (3000, 50.0), # 3.0s: 50% - (3500, 60.0), # 3.5s: 60% - (4000, 70.0), # 4.0s: 70% - (4500, 80.0), # 4.5s: 80% -] +# 9. Spindle segment +spin_segment = Segment(name="spin") -print("\nTendon tap schedule:") -print(f"\t- Duration per tap: {TENDON_TAP_DURATION} ms") -print(f"\t- Phase 1 (0-2.5s): {len([t for t, m in TENDON_TAP_SCHEDULE if t < 2500])} taps") -print(f"\t- Phase 2 (2.5-5s): {len([t for t, m in TENDON_TAP_SCHEDULE if t >= 2500])} taps") -print(f"\t- Amplitudes: {sorted(set([m for t, m in TENDON_TAP_SCHEDULE]))}%") - - -# Create step callback function with access to step counter -def step_callback(step_counter): - return eachStep( - muscle=hill_muscle, - spin=spin, - golgi=gto, - popD=network.populations, - ncD=ncD, - gMN=gMN, - joint_dyn=joint_dynamics, - step_counter=step_counter, - tendon_tap_schedule=TENDON_TAP_SCHEDULE, - tendon_tap_duration=TENDON_TAP_DURATION, +# Primary afferent (Ia) firing rate — use capped series (matches what afferents received) +ia_data = np.array(ia_firing_series[:n_steps]) +spin_segment.analogsignals.append( + AnalogSignal( + ia_data.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="spin_primary_afferent_firing__Hz", + attr_name="primary_afferent_firing__Hz", ) +) +# Secondary afferent (II) firing rate — use capped series +ii_data = np.array(ii_firing_series[:n_steps]) +spin_segment.analogsignals.append( + AnalogSignal( + ii_data.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="spin_secondary_afferent_firing__Hz", + attr_name="secondary_afferent_firing__Hz", + ) +) -############################################################################## -# Run Spinal Network Simulation -# ---------------------------- -# -# Execute the complete simulation with all integrated components. - -print("\nStarting spinal network simulation...") -print(f"\tDuration: {tstop} ms") -print(f"\tTimestep: {dt} ms") -print(f"\tPopulations: {len(network.populations)}") +# Intrafusal fiber activations — from scan outputs (bag1, bag2 only; chain is algebraic) +for _spin_key, _spin_attr in [("bag1_act", "bag1_activation"), ("bag2_act", "bag2_activation")]: + _act_arr = np.array(scan_out[_spin_key])[:n_steps] + spin_segment.analogsignals.append( + AnalogSignal( + _act_arr.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name=f"spin_{_spin_attr}", + attr_name=_spin_attr, + ) + ) -runner = SimulationRunner( - network=network, - models=models, - step_callback=step_callback, +# Intrafusal tensions — shape (n_steps, 3): [Bag1, Bag2, Chain] per row +_tensions_arr = np.array(scan_out["spin_T"])[:n_steps] # (n_steps, 3) +spin_segment.analogsignals.append( + AnalogSignal( + _tensions_arr * pq.dimensionless, + sampling_period=sampling_period_s, + name="spin_intrafusal_tensions", + attr_name="intrafusal_tensions", + ) ) -# Motor neuron spike recording thresholds are now fixed in the Network class +# Gamma fusimotor drive — step-like reference signal for activations subplot +spin_segment.analogsignals.append( + AnalogSignal( + gDyn[:n_steps].reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="spin_gamma_dynamic", + attr_name="gamma_dynamic", + ) +) -results = runner.run( - duration__ms=tstop * pq.ms, - timestep__ms=dt, - membrane_recording={ - "aMN": [0, 5, 10, 15, 20, 30, 40, 50, 60, 70], - }, +results.segments.append(spin_segment) + +# 10. GTO segment — always from scan output (gto.ib_afferent_firing__Hz is zero-initialized) +gto_segment = Segment(name="gto") +ib_data = np.array(ib_firing_series[:n_steps]) +gto_segment.analogsignals.append( + AnalogSignal( + ib_data.reshape(-1, 1) * pq.dimensionless, + sampling_period=sampling_period_s, + name="gto_ib_afferent_firing__Hz", + attr_name="ib_afferent_firing__Hz", + ) ) +results.segments.append(gto_segment) -print("Simulation completed successfully!") +############################################################################## +# Save Results +# ------------ -# Save simulation results -joblib.dump(results, save_path / "spinal_network_results.pkl") -joblib.dump(artAng, save_path / "joint_angles.pkl") -joblib.dump(musculotendon_force__N, save_path / "musculotendon_force__N.pkl") +joblib.dump(results, save_path / "spinal_network_results_jaxley.pkl") +joblib.dump(artAng, save_path / "joint_angles_jaxley.pkl") +joblib.dump(musculotendon_force__N, save_path / "musculotendon_force__N_jaxley.pkl") -# Save drive signals for analysis and plotting +# Save drive signals drive_signals = { "time": time, "descending_drive": DDdrive, @@ -790,91 +1026,42 @@ def step_callback(step_counter): "tendon_tap_schedule": TENDON_TAP_SCHEDULE, "tendon_tap_duration": TENDON_TAP_DURATION, } -joblib.dump(drive_signals, save_path / "drive_signals.pkl") - -# Save event timeline for plot annotations -events = { - # Tendon tap events: (peak_time_ms, amplitude_%, duration_ms) - # Note: peak_time is at the midpoint of the triangular pulse - "tendon_taps": [ - (t + TENDON_TAP_DURATION / 2.0, m, TENDON_TAP_DURATION) for t, m in TENDON_TAP_SCHEDULE - ], - # Gamma drive transitions: (time_ms, level_pps) - "gamma_transitions": list(zip(gamma_times, gamma_values)), - # Gamma drive periods: (start_time_ms, end_time_ms, level_pps) - "gamma_periods": [ - (gamma_times[i], gamma_times[i + 1], gamma_values[i]) for i in range(len(gamma_times) - 1) - ], - # Phase markers - "phase_boundary": 2500, # ms - transition from phase 1 to phase 2 - "cortical_drive_onset": cortical_start_time, # ms - when sinusoidal drive starts - # Experimental phases with descriptions - "phases": [ - {"start": 0, "end": 2500, "name": "Phase 1", "description": "Reflex gain modulation"}, - { - "start": 2500, - "end": 5000, - "name": "Phase 2", - "description": "Reflex-voluntary interaction", - }, - ], -} -joblib.dump(events, save_path / "event_timeline.pkl") - -print(f"Results saved to {save_path}") -print("\t- spinal_network_results.pkl (neural activity, muscle, spindle, GTO)") -print("\t- joint_angles.pkl (joint kinematics)") -print("\t- musculotendon_force__N.pkl (real-unit musculotendon force [N])") -print("\t- drive_signals.pkl (cortical drive, gamma drive, tap schedule)") -print("\t- event_timeline.pkl (event markers for plot annotations)") -print("\nNote: Intrafusal fiber tensions [FU] are in spinal_network_results.pkl") -print("\tAccess via: results.segments[0].analogsignals (look for 'intrafusal_tensions')") -print("\tConvert to real units: T[N] ≈ T[FU] * scaling_factor (no standard conversion)") -print(f"\tMusculotendon force capacity: {hill_muscle.F0:.1f} N") +joblib.dump(drive_signals, save_path / "drive_signals_jaxley.pkl") + +print(f"\nResults saved to {save_path}") +print("\t- spinal_network_results_jaxley.pkl") +print("\t- joint_angles_jaxley.pkl") +print("\t- musculotendon_force__N_jaxley.pkl") +print("\t- drive_signals_jaxley.pkl") ############################################################################## -# Export to NWB Format (Neurodata Without Borders) -# ------------------------------------------------ -# -# NWB is a standardized data format for neurophysiology data that enables -# data sharing and interoperability with other neuroscience tools. -# MyoGen supports NWB export for integration with the broader neuroscience -# ecosystem, including tools like the DANDI Archive. +# Export to NWB Format +# -------------------- -# Export the Neo Block results to NWB format -nwb_filepath = export_to_nwb( - results, - save_path / "spinal_network_results.nwb", - session_description=( - "MyoGen spinal network simulation with systematic tendon tap protocol. " - "Two-phase design: (1) Reflex gain modulation with varying gamma drive, " - "(2) Reflex-voluntary interaction with sinusoidal cortical drive." - ), - experimenter="MyoGen Simulation", - institution="MyoGen Framework", - lab="Neuromuscular Simulation", - experiment_description=( - f"5-second simulation with {naMN} motor neurons, " - f"{nIa} Ia afferents, {nII} II afferents, {nIb} Ib afferents, " - f"and {ngII + ngIb} spinal interneurons. " - f"Includes Hill-type muscle model and closed-loop joint dynamics." - ), - keywords=[ - "MyoGen", - "spinal network", - "motor neuron", - "stretch reflex", - "tendon tap", - "proprioception", - "EMG simulation", - ], - # Subject metadata for DANDI compliance - subject_id="simulated_subject_001", - species="Homo sapiens", - subject_description="Simulated human motor neuron pool and spinal reflex network", -) -print(f"\n(OK) Exported to NWB format: {nwb_filepath}") -print("\tNWB files can be validated with: nwbinspector ") +try: + nwb_filepath = export_to_nwb( + results, + save_path / "spinal_network_results_jaxley.nwb", + session_description=( + "MyoGen spinal network simulation (Jaxley backend) with systematic tendon tap protocol. " + "Uses NERLab motor neurons (soma napp; dendrite caL) — matches production NEURON." + ), + experimenter="MyoGen Simulation", + institution="MyoGen Framework", + lab="Neuromuscular Simulation", + experiment_description=( + f"5-second simulation with {naMN} motor neurons (NERLab), " + f"{nIa} Ia afferents, {nII} II afferents, {nIb} Ib afferents, " + f"and {ngII + ngIb} spinal interneurons." + ), + keywords=["MyoGen", "Jaxley", "spinal network", "motor neuron", "stretch reflex"], + subject_id="simulated_subject_001", + species="Homo sapiens", + subject_description="Simulated human motor neuron pool - Jaxley backend", + ) + print(f"\n(OK) Exported to NWB format: {nwb_filepath}") +except Exception as e: + print(f"\n(Warning) NWB export failed: {e}") ############################################################################## # Comprehensive Results Visualization @@ -882,18 +1069,19 @@ def step_callback(step_counter): # # Create a series of plots that tell the complete story of spinal network # function, from neural activity to mechanical output. +# Matches NEURON Example 11 visualization structure exactly. print("\nGenerating comprehensive visualizations...") # 1. NEURAL ACTIVITY: Raster plot showing all population spike patterns populations_list = [ - "aMN", # Motor output + "aMN", # Motor output "Ia", "II", - "Ib", # Sensory input + "Ib", # Sensory input "gII", - "gIb", # Interneurons - "DD", # Descending drive + "gIb", # Interneurons + "DD", # Descending drive ] fig1, axes1 = plt.subplots(len(populations_list), 1, figsize=(15, 12)) plot_raster_spikes( @@ -904,12 +1092,11 @@ def step_callback(step_counter): title="Spinal Network Activity (Single Muscle)", ) plt.tight_layout() -plt.savefig(save_path / "neural_raster_plot.png", dpi=150, bbox_inches="tight") +plt.savefig(save_path / "neural_raster_plot_jaxley.png", dpi=150, bbox_inches="tight") plt.show() # 2. MOTOR NEURON DYNAMICS: Membrane potentials showing integration fig2, ax2 = plt.subplots(1, 1, figsize=(15, 6)) - plot_membrane_potentials( results, [ax2], @@ -918,8 +1105,8 @@ def step_callback(step_counter): time_range=(0, tstop), title="Motor Neuron Membrane Potentials", ) - plt.tight_layout() +plt.savefig(save_path / "membrane_potentials_jaxley.png", dpi=150, bbox_inches="tight") plt.show() # 3. MUSCLE MECHANICS: Muscle dynamics @@ -937,7 +1124,7 @@ def step_callback(step_counter): title="Muscle Dynamics - Length, Force, and Activation", ) plt.tight_layout() -plt.savefig(save_path / "muscle_dynamics.png", dpi=150, bbox_inches="tight") +plt.savefig(save_path / "muscle_dynamics_jaxley.png", dpi=150, bbox_inches="tight") plt.show() # 4. PROPRIOCEPTIVE FEEDBACK: Spindle dynamics and sensory encoding @@ -947,14 +1134,14 @@ def step_callback(step_counter): axes4, muscle_name="hill_muscle", include_signals=["L"], - include_activations=["Bag1", "Bag2", "Chain"], + include_activations=["Bag1", "Bag2"], # Chain is algebraic; not stored in scan state include_tensions=["Bag1", "Bag2", "Chain"], include_afferents=["Ia", "II"], time_range=(0, tstop), title="Muscle Spindle Dynamics - Proprioceptive Feedback System", ) plt.tight_layout() -plt.savefig(save_path / "spindle_dynamics.png", dpi=150, bbox_inches="tight") +plt.savefig(save_path / "spindle_dynamics_jaxley.png", dpi=150, bbox_inches="tight") plt.show() # 5. FORCE FEEDBACK: GTO dynamics and protective reflexes @@ -968,7 +1155,10 @@ def step_callback(step_counter): title="Golgi Tendon Organ Dynamics - Force Feedback System", ) plt.tight_layout() -plt.savefig(save_path / "gto_dynamics.png", dpi=150, bbox_inches="tight") +plt.savefig(save_path / "gto_dynamics_jaxley.png", dpi=150, bbox_inches="tight") plt.show() -# mkdocs_gallery_thumbnail_path = "gallery_thumbs/11_simulate_spinal_network.png" +print("\n" + "=" * 60) +print("[DONE] Jaxley spinal network simulation complete!") +print(" Using NERLab motor neurons (napp + caL) — matches production NEURON.") +print("=" * 60) diff --git a/examples/01_basic/11_simulate_spinal_network_neuron.py b/examples/01_basic/11_simulate_spinal_network_neuron.py new file mode 100644 index 00000000..b8c5b49a --- /dev/null +++ b/examples/01_basic/11_simulate_spinal_network_neuron.py @@ -0,0 +1,972 @@ +""" +Spinal Network Simulation with Systematic Tendon Tap Protocol +============================================================== + +This example demonstrates **complete spinal reflex network modeling** with a **comprehensive tendon tap +protocol** that systematically varies mechanical perturbations, fusimotor drive, and cortical activity. +This sophisticated experimental design reveals how the nervous system modulates stretch reflex gain. + +Learning Objectives +------------------ + +This example teaches the following key concepts in neuromuscular control: + +1. **Reflex Gain Modulation**: How fusimotor (gamma) drive amplifies stretch reflex sensitivity + - Compare motor neuron responses across different gamma levels (0→100 pps) + - Observe how higher gamma drive increases spindle sensitivity and reflex output + +2. **Sensory Pathway Timing**: Why different feedback pathways activate at different times + - Spindle-driven pathways (Ia, II → gII) respond immediately to muscle stretch + - Force-driven pathways (Ib → gIb) respond later after muscle contraction builds force + +3. **Reflex-Voluntary Interaction**: How cortical commands interact with spinal reflexes + - Phase 1: Isolated reflex responses (no cortical drive) + - Phase 2: Reflex modulation during voluntary contraction (with cortical drive) + - Observe how background muscle activity affects reflex sensitivity + +4. **Closed-loop Biomechanics**: How muscle forces create movement that feeds back to sensors + - Joint dynamics convert muscle torque into motion + - Motion changes muscle length/velocity, affecting spindle feedback + - Creates realistic proprioceptive-motor coupling + +**Experimental Protocol (5-second, 2-phase design):** + +**Phase 1 (0-2.5s): Reflex Gain Modulation** +- Triangular tendon taps every 0.5s with increasing amplitude (50%, 60%, 70%, 80%) + + * Triangular waveform: Linear ramp up to peak (50ms) → linear ramp down (50ms) + * Total duration: 100ms per tap + * Mimics realistic tendon tap dynamics with stretch and release phases + +- Stepwise increasing gamma drive starting at 0.75s (0→25→50→75→100 pps) +- No cortical drive (isolated reflex testing) + +**Phase 2 (2.5-5s): Reflex-Voluntary Interaction** +- Repeat triangular tap pattern (50%, 60%, 70%, 80%) +- Repeat gamma drive pattern (0→25→50→75→100 pps) +- Sinusoidal cortical drive (38±1 Hz at 1 Hz) recruiting ~75% of motor neurons + +.. note:: + This example builds upon all previous examples and demonstrates how the complete neuromuscular system + functions as an integrated network: + + - **Motor neuron pool**: α-motoneurons controlling a single muscle (from examples 01-02) + - **Muscle model**: Hill-type muscle model with realistic force generation (from examples 02, 06) + - **Proprioceptive feedback**: Muscle spindles and Golgi tendon organs (from examples) + - **Spinal interneurons**: Group II and Group Ib interneurons for reflex modulation + - **Joint dynamics**: Closed-loop biomechanical control with realistic inertia and damping + - **Descending drive**: Cortical control signals for voluntary movement (from example 01) + +.. important:: + **Spinal Reflex Networks** are the fundamental control circuits that coordinate muscle activity. + Key physiological concepts demonstrated: + + - **Stretch reflex**: Muscle spindles provide length/velocity feedback for posture control + - **Force feedback**: Golgi tendon organs monitor muscle tension for protective reflexes + - **Closed-loop control**: Joint mechanics influence neural activity through proprioception + - **Motor control**: Integration of descending commands with sensory feedback + +Scientific Background +------------------- + +The spinal cord contains intricate neural circuits that process sensory information and generate appropriate +motor responses. This example models several key components: + +**Muscle Spindles**: Specialized sensory organs that detect muscle length and velocity changes, providing +Group Ia (velocity-sensitive) and Group II (length-sensitive) afferent signals. + +**Golgi Tendon Organs**: Force-sensitive receptors that monitor muscle tension and provide Group Ib +afferent feedback for protective reflexes. + +**Spinal Interneurons**: The model includes two key inhibitory interneuron populations: + +- **Group II interneurons (gII)**: Process Group II spindle afferents (length-sensitive). In this model, + they provide inhibitory feedback to motor neurons, representing one component of the complex Group II + pathway (which can be excitatory or inhibitory depending on context). + +- **Group Ib interneurons (gIb)**: Receive input from Golgi tendon organs and provide **autogenic + inhibition** - force-dependent negative feedback that protects against excessive muscle tension. + Classic Ib inhibitory pathway. + +**Joint Dynamics**: Realistic biomechanical modeling of joint inertia, damping, and muscle-generated torques +creates the closed-loop system where neural activity affects movement, which in turn affects sensory feedback. +""" + +# %% + +############################################################################## +# Import Libraries +# ---------------- +# + +from pathlib import Path + +import joblib +import matplotlib.pyplot as plt +import numpy as np +import quantities as pq +from neuron import h + +from myogen import load_nmodl_mechanisms +from myogen.simulator.neuron.joint_dynamics import JointDynamics +from myogen.simulator.neuron.muscle import HillModel +from myogen.simulator.neuron.network import Network +from myogen.simulator.neuron.populations import ( + AffIa__Pool, + AffIb__Pool, + AffII__Pool, + AlphaMN__Pool, + DescendingDrive__Pool, + GIb__Pool, + GII__Pool, +) +from myogen.simulator.neuron.proprioception import ( + GolgiTendonOrganModel, + SpindleModel, +) +from myogen.simulator.neuron.simulation_runner import SimulationRunner +from myogen.utils.nwb import export_to_nwb +from myogen.utils.plotting import ( + plot_gto_dynamics, + plot_membrane_potentials, + plot_muscle_dynamics, + plot_raster_spikes, + plot_spindle_dynamics, +) + +############################################################################## +# Load NEURON Mechanisms and Dependencies +# --------------------------------------- +# +# Load the compiled NMODL mechanisms required for biophysical neuron modeling +# and load results from previous examples that serve as inputs to this simulation. + +# Load NEURON mechanisms +load_nmodl_mechanisms() + +# Setup results directory +save_path = Path(r"./results") +save_path.mkdir(exist_ok=True) + +recruitment_thresholds = joblib.load(save_path / "thresholds.pkl") +print("(OK) Loaded recruitment thresholds from example 00") + +############################################################################## +# Define Simulation Parameters +# --------------------------- +# +# These parameters control the temporal and spatial resolution of the simulation, +# as well as the physiological characteristics of the neural populations and +# mechanical system. + +# Temporal parameters - high resolution for accurate neural integration +dt = 0.005 * pq.ms # Integration timestep with units +tstop = 5e3 # 5e3 # ms - Total simulation duration (5s for comprehensive protocol) +time = np.arange(0, tstop, dt.magnitude) # Time vector (use magnitude for numpy) + +print("Simulation parameters:") +print(f"\tDuration: {tstop} ms") +print(f"\tTimestep: {dt} ms") +print(f"\tTime samples: {len(time)}") + +############################################################################## +# Define Neural Population Sizes +# ----------------------------- +# +# Population sizes are based on physiological estimates from cat and human studies. +# These numbers represent typical motor pool compositions for a single muscle. + +# Afferent populations (sensory input) +# Based on Banks et al. (2006, 2015) for distal upper-limb muscles +nIa = 73 # Group Ia afferents from muscle spindles (velocity-sensitive) +nII = 80 # Group II afferents from muscle spindles (length-sensitive) - 110% of Ia +nIb = 58 # Group Ib afferents from Golgi tendon organs (force-sensitive) - 80% of Ia + +# Interneuron populations (spinal processing) +# Note: Interneuron:MN ratio is ~10:1 in spinal cord. Group II/Ib INs are a subset. +# Using ~50% of afferent count as a reasonable estimate for convergent interneurons. +ngII = 120 # Group II interneurons (process spindle feedback from 80 II afferents) +ngIb = 145 # Group Ib interneurons (process force feedback from 58 Ib afferents) + +# Motor neurons (output to muscles) +nType1 = 102 # Type I motor neurons (slow, fatigue-resistant) +nType2 = 18 # Type II motor neurons (fast, fatigue-prone) +naMN = nType1 + nType2 # Total α-motoneurons per muscle + +# Descending drive (cortical input) +nDD = 400 # Total descending drive neurons (increased for more drive) +DDorder = ( + 5 # Poisson process order for realistic spike patterns (higher order = more regular spiking) +) + +print("Neural population sizes:") +print(f"\t- α-Motoneurons: {naMN} ({nType1} Type I, {nType2} Type II)") +print(f"\t- Ia afferents: {nIa}") +print(f"\t- II afferents: {nII}") +print(f"\t- Ib afferents: {nIb}") +print(f"\t- Interneurons: {ngII + ngIb}") +print(f"\t- Descending drive: {nDD}") + +############################################################################## +# Define Descending Drive Pattern +# ------------------------------- +# +# Two-phase pattern: +# Phase 1 (0-2500ms): No cortical drive (reflex testing with varying gamma) +# Phase 2 (2500-5000ms): Sinusoidal cortical drive (reflex modulation during voluntary contraction) + +# Phase 1: No cortical drive +# Phase 2: 1 Hz sinusoidal drive with bias to recruit ~75% of MN population +DDdrive = np.zeros_like(time) +cortical_start_time = 2500 # ms + +# Parameters for sinusoidal drive +bias_hz = 40 # Bias for robust motor neuron recruitment (increased from 75 for ~20 pps firing) +amplitude_hz = 1 # Oscillation amplitude (±5 Hz around bias) +frequency_hz = 1.0 # 1 Hz oscillation + +# Apply sinusoidal drive from 2.5s onwards +mask = time >= cortical_start_time +DDdrive[mask] = bias_hz + amplitude_hz * np.sin( + 2 * np.pi * frequency_hz * (time[mask] - cortical_start_time) / 1000.0 +) + +print("Descending drive pattern:") +print("\t- Phase 1 (0-2.5s): No cortical drive") +print(f"\t- Phase 2 (2.5-5s): Sinusoidal drive ({bias_hz}±{amplitude_hz} Hz at {frequency_hz} Hz)") +print(f"\t (Drive range: {bias_hz - amplitude_hz}-{bias_hz + amplitude_hz} Hz)") + +############################################################################## +# Define Fusimotor Drive +# -------------------- +# +# Fusimotor neurons control the sensitivity of muscle spindles by adjusting +# the tension in intrafusal muscle fibers. Stepwise increasing gamma drive +# demonstrates how fusimotor activity modulates stretch reflex gain. +# +# Pattern: Increases every 0.5s starting at 0.75s, resets at 2.5s +# Steps: 0 → 25 → 50 → 75 → 100 pps (then repeat) + +# Time points for stepwise gamma drive +# Phase 1: 0→0.75s (wait)→1.25s→1.75s→2.25s (steps every 0.5s) +# Phase 2: 2.5s (reset)→3.25s (wait 0.75s)→3.75s→4.25s→4.75s (steps every 0.5s) +gamma_times = np.array([0, 750, 1250, 1750, 2250, 2500, 3250, 3750, 4250, 4750, 5000]) +gamma_values = np.array([0, 25, 50, 75, 100, 0, 25, 50, 75, 100, 100]) + +# Create stepwise gamma drive arrays using proper step function (not interpolation!) +gDyn = np.zeros_like(time) +gStat = np.zeros_like(time) + +# Phase 1: Steps every 0.5s starting at 0.75s +gDyn[(time >= 750) & (time < 1250)] = 25 +gDyn[(time >= 1250) & (time < 1750)] = 50 +gDyn[(time >= 1750) & (time < 2250)] = 75 +gDyn[(time >= 2250) & (time < 2500)] = 100 + +# Phase 2: Reset and repeat pattern +gDyn[(time >= 3250) & (time < 3750)] = 25 +gDyn[(time >= 3750) & (time < 4250)] = 50 +gDyn[(time >= 4250) & (time < 4750)] = 75 +gDyn[(time >= 4750) & (time <= 5000)] = 100 + +# Static gamma follows same pattern +gStat = gDyn.copy() + +# Add small physiological variability +gDyn = gDyn + np.random.normal(0, 1, len(time)) +gStat = gStat + np.random.normal(0, 1, len(time)) + +# Ensure non-negative +gDyn = np.maximum(gDyn, 0) +gStat = np.maximum(gStat, 0) + +# Package fusimotor drives for the simulation +gMN = { + "name": "gMN", + "gDyn": gDyn, # Dynamic fusimotor drive array + "gStat": gStat, # Static fusimotor drive array +} + +print("Fusimotor drive parameters:") +print("\t- Stepwise increases: 0 → 25 → 50 → 75 → 100 pps") +print("\t- Starting at 0.75s, stepping every 0.5s") +print("\t- Resets at 2.5s and repeats pattern") + +############################################################################## +# Initialize Joint Dynamics +# ------------------------ +# +# The joint model provides realistic biomechanical constraints and creates the +# closed-loop system where muscle forces influence joint motion, which in turn +# affects proprioceptive feedback. + +joint_dynamics = JointDynamics( + inertia__kg_m2=0.001, # Realistic joint inertia + damping__Nm_s_per_rad=0.002, # Light damping for stability +) + +# Initialize joint angle array for closed-loop integration +artAng = np.zeros_like(time) +artAng[0] = 0.0 # Initial joint angle + +print("Joint dynamics parameters:") +print(f"\t- Inertia: {joint_dynamics.inertia__kg_m2} kg⋅m²") +print(f"\t- Damping: {joint_dynamics.damping__Nm_s_per_rad} N⋅m⋅s/rad") +print(f"\t- Initial angle: {artAng[0]}°") + +############################################################################## +# Create Neural Populations +# ------------------------ +# +# Instantiate all neural population objects with physiologically appropriate +# parameters. Each population type has specialized properties reflecting their +# biological counterparts. + +# Create motor neuron pool +# Motor neuron pool (default spike_threshold__mV=50.0 for proper spike detection) +aMN = AlphaMN__Pool(recruitment_thresholds__array=recruitment_thresholds) # α-motoneurons + +# Create descending drive population +DD = DescendingDrive__Pool(n=nDD, poisson_batch_size=DDorder, timestep__ms=dt) + +# Create afferent populations +Ia = AffIa__Pool(n=nIa, timestep__ms=dt) # Primary spindle afferents +II = AffII__Pool(n=nII, timestep__ms=dt) # Lower thresholds for Group II +Ib = AffIb__Pool(n=nIb, timestep__ms=dt) # Golgi tendon organ afferents + +# Create interneuron populations for reflex modulation +gII = GII__Pool(n=ngII) # Group II interneurons +gIb = GIb__Pool(n=ngIb) # Group Ib interneurons + +print(f"(OK) Created {len([aMN, DD, Ia, II, Ib, gII, gIb])} neural populations") + +############################################################################## +# Create Proprioceptive Models +# --------------------------- +# +# Initialize the sensory models that convert mechanical variables (muscle length, +# velocity, force) into neural signals that drive the afferent populations. + +# Golgi Tendon Organ - monitors muscle force/tension +gto = GolgiTendonOrganModel( + simulation_time__ms=tstop * pq.ms, + time_step__ms=dt, + gto_parameters=GolgiTendonOrganModel.create_default_gto_parameters(), +) + +# Muscle Spindle - monitors muscle length and velocity +spin = SpindleModel( + simulation_time__ms=tstop * pq.ms, + time_step__ms=dt, + spindle_parameters=SpindleModel.create_default_spindle_parameters(), +) + +print("(OK) Initialized proprioceptive models (spindle, GTO)") + +############################################################################## +# Create Muscle Model +# ------------------ +# +# Initialize Hill-type muscle model. This model converts motor neuron +# activation into realistic force generation with appropriate biomechanical properties. + +# Muscle model +hill_muscle = HillModel( + simulation_time__ms=tstop * pq.ms, + time_step__ms=dt, + muscle_parameters=HillModel.create_default_muscle_parameters(), + n_motor_units_type1=nType1, + n_motor_units_type2=nType2, + initial_joint_angle__deg=artAng[0], + muscle_role="flexor", +) + +print("(OK) Created muscle model") +print(f"\t- Motor units: {nType1 + nType2}") +print(f"\t- Force capacity: ~{hill_muscle.F0:.1f} N") + +############################################################################## +# Create Arrays for Real-Unit Forces +# ---------------------------------- +# +# Store musculotendon force in real units [N] for analysis + +# Musculotendon force array [N] +musculotendon_force__N = np.zeros_like(time) + +print("(OK) Initialized force tracking arrays") + +############################################################################## +# Define Callback Functions +# ------------------------ +# +# These functions handle the integration of different system components during +# the simulation, including spike events and step-wise updates. + + +def spkEvent(i, muscle, delay): + muscle.add_spike(i, delay) + + +def eachStep( + muscle, + spin, + golgi, + popD, + ncD, + gMN, + joint_dyn, + step_counter, + tendon_tap_schedule, # List of (time_ms, magnitude_percent) tuples + tendon_tap_duration=100, # ms +): + """ + Step-wise integration function called at each simulation timestep. + + This function orchestrates the complex interactions between: + - Muscle mechanics and force generation + - Proprioceptive feedback from spindles and GTOs + - Joint dynamics and closed-loop control + - Neural population dynamics and spike generation + - Multiple tendon tap perturbations with varying amplitudes + + Parameters + ---------- + muscle : HillModel + Muscle model + spin : SpindleModel + Muscle spindle model + golgi : GolgiTendonOrganModel + Golgi tendon organ model + popD : dict + Dictionary of neural populations + ncD : dict + Dictionary of network connections + gMN : dict + Fusimotor drive parameters + joint_dyn : JointDynamics + Joint biomechanics model + step_counter : iterator + Simulation step counter + tendon_tap_schedule : list of tuples + List of (time_ms, magnitude_percent) defining tap timing and amplitudes + tendon_tap_duration : float + Duration (ms) of each tendon tap stretch + """ + i = next(step_counter) + + current_angle = artAng[i] + + # MUSCLE MECHANICS: Integrate muscle with current joint angle + L, V, A = muscle.integrate(current_angle) + + # TENDON TAP PERTURBATION: Triangular pulse (ramp up, then down) + current_time = h.t + + for tap_time, tap_magnitude in tendon_tap_schedule: + tap_end_time = tap_time + tendon_tap_duration + tap_midpoint = tap_time + (tendon_tap_duration / 2.0) + + if tap_time <= current_time < tap_end_time: + # Triangular pulse: ramp up to peak at midpoint, then ramp down + max_perturbation = L * (tap_magnitude / 100.0) + + if current_time < tap_midpoint: + # Rising edge: 0 → peak (first half) + progress = (current_time - tap_time) / (tendon_tap_duration / 2.0) + length_perturbation = max_perturbation * progress + else: + # Falling edge: peak → 0 (second half) + progress = (tap_end_time - current_time) / (tendon_tap_duration / 2.0) + length_perturbation = max_perturbation * progress + + L_perturbed = L + length_perturbation + + # Velocity reflects the triangle slope + if current_time < tap_midpoint: + # Rising: positive velocity + V_perturbed = V + (max_perturbation / (tendon_tap_duration / 2000.0)) + else: + # Falling: negative velocity + V_perturbed = V - (max_perturbation / (tendon_tap_duration / 2000.0)) + + # Acceleration spikes at tap onset and reverses at midpoint + if abs(current_time - tap_time) < dt.magnitude: + A_perturbed = A + (max_perturbation / ((tendon_tap_duration / 2000.0) ** 2)) + print( + f"\t>> TENDON TAP (triangular) at t={current_time:.1f} ms: " + f"{tap_magnitude}% peak stretch ({max_perturbation:.4f} L0), " + f"duration={tendon_tap_duration} ms" + ) + print(f"\t Baseline length: {L:.4f} L0 → Peak: {L + max_perturbation:.4f} L0") + elif abs(current_time - tap_midpoint) < dt.magnitude: + A_perturbed = A - 2 * (max_perturbation / ((tendon_tap_duration / 2000.0) ** 2)) + else: + A_perturbed = A + + # Use perturbed values for spindle feedback + L, V, A = L_perturbed, V_perturbed, A_perturbed + break # Only one tap active at a time + + # PROPRIOCEPTIVE FEEDBACK: Use muscle kinematics for spindle feedback + Iay, IIy = spin.integrate(L, V, A, gMN["gDyn"][i], gMN["gStat"][i]) + + # FORCE FEEDBACK: Muscle force for GTO + f = muscle.F0 * muscle.muscle_force[i] # Musculotendon force [N] + musculotendon_force__N[i] = f # Store in real units + Iby = golgi.integrate(f) + + # CLOSED-LOOP DYNAMICS: Update joint angle based on muscle torque + new_angle, _ = joint_dyn.integrate( + torque__Nm=muscle.signed_muscle_torque[i], + dt__s=float(dt.rescale(pq.s).magnitude), + ) + # Update joint angle array for next timestep + if i < len(artAng) - 1: + artAng[i + 1] = new_angle + + # DESCENDING DRIVE PROCESSING: Convert cortical signals to spikes + for DDcell in popD["DD"]: + if DDcell.integrate(DDdrive[i]): + spike_time = h.t + 1 + if spike_time < tstop: + ncD["cmd->DD"][DDcell.pool__ID].event(spike_time) + + # AFFERENT PROCESSING: Convert sensory signals to neural spikes + for Ia in popD["Ia"]: + if Iay >= Ia.RT: + if Ia.integrate(Iay): + # Ensure h.t is converted to a quantities object with units of ms + spike_time = h.t + float(Ia.axon_delay__ms) + if spike_time < tstop: + ncD["Spindle->Ia"][Ia.pool__ID].event(spike_time) + + ii_spikes = 0 + for II in popD["II"]: + if IIy >= II.RT: + if II.integrate(IIy): + spike_time = h.t + float(II.axon_delay__ms) + if spike_time < tstop: + ncD["Spindle->II"][II.pool__ID].event(spike_time) + ii_spikes += 1 + + ib_spikes = 0 + for Ib in popD["Ib"]: + if Iby >= Ib.RT: + if Ib.integrate(Iby): + spike_time = h.t + float(Ib.axon_delay__ms) + if spike_time < tstop: + ncD["GTO->Ib"][Ib.pool__ID].event(spike_time) + ib_spikes += 1 + + +############################################################################## +# Create Neural Network +# -------------------- +# +# Assemble all neural populations into a connected network that implements +# the spinal reflex circuitry for single muscle control. + +network = Network( + { + "DD": DD, # Descending drive + "aMN": aMN, # α-motoneurons + "Ia": Ia, # Group Ia afferents + "II": II, # Group II afferents + "Ib": Ib, # Group Ib afferents + "gII": gII, # Group II interneurons + "gIb": gIb, # Group Ib interneurons + "gMN": gMN, # Fusimotor parameters + } +) + +print(f"(OK) Created network with {len(network.populations)} populations") + +############################################################################## +# Configure Neural Connections +# --------------------------- +# +# Establish synaptic connections that implement physiologically realistic +# spinal reflex pathways, including both excitatory and inhibitory connections. +# +# Neural Circuit Diagram: +# ---------------------- +# +# Spindle (length/velocity) GTO (force) Cortex +# | | | | +# v v v v +# Ia II Ib DD +# | | | | +# | v v | +# | gII (-) gIb (-) | +# | | | | +# +----------+--------(+)------------+----------(+)-----+ +# | +# v +# α-MN -----> Muscle +# | +# v +# Joint Motion +# | +# +--------------------------------+ +# | | +# v v +# Spindle GTO +# (feedback) (feedback) +# +# Legend: (+) = excitatory, (-) = inhibitory +# +# Key pathways: +# 1. Ia → MN: Monosynaptic stretch reflex (EXCITATORY, fast) +# 2. II → gII → MN: Polysynaptic length pathway (INHIBITORY, moderate speed) +# 3. Ib → gIb → MN: Autogenic inhibition (INHIBITORY, slower - force-dependent) +# 4. DD → MN: Descending cortical drive (EXCITATORY, voluntary control) + +# DESCENDING CONTROL: Direct cortical drive to motor neurons (EXCITATORY) +network.connect("DD", "aMN", probability=0.3, weight__uS=0.05 * pq.uS) + +# MONOSYNAPTIC STRETCH REFLEX: Ia afferents excite motor neurons (EXCITATORY) +# High connection probability (80%) reflects homogeneous excitatory distribution +network.connect("Ia", "aMN", probability=0.8, weight__uS=0.05 * pq.uS) + +# POLYSYNAPTIC PATHWAYS: Group II pathway (afferents → interneurons → motor neurons) +# Lower connection probability (30%) captures predominantly indirect, disynaptic influence +network.connect( + "II", "gII", probability=0.3, weight__uS=0.025 * pq.uS +) # Afferent → Interneuron (EXCITATORY) +network.connect( + "gII", "aMN", probability=0.3, weight__uS=0.05 * pq.uS, inhibitory=True +) # Interneuron → MN (INHIBITORY, weak modulatory effect) + +# INHIBITORY REFLEXES: Force feedback through Ib interneurons (autogenic inhibition) +# Lower connection probability (30%) captures predominantly indirect, disynaptic influence +network.connect( + "Ib", "gIb", probability=0.3, weight__uS=0.025 * pq.uS +) # Afferent → Interneuron (EXCITATORY) +network.connect( + "gIb", "aMN", probability=0.3, weight__uS=0.05 * pq.uS, inhibitory=True +) # Interneuron → MN (INHIBITORY, weak modulatory effect) + +print("(OK) Configured synaptic connections") +print("\t- Excitatory pathways:") +print("\t\t• DD→MN (descending drive)") +print("\t\t• Ia→MN (monosynaptic stretch reflex)") +print("\t\t• Afferents→Interneurons: II→gII, Ib→gIb") +print("\t- Inhibitory pathways:") +print("\t\t• gII→MN (Group II interneuron inhibition)") +print("\t\t• gIb→MN (autogenic inhibition from force feedback)") + +############################################################################## +# Connect Motors to Muscle +# ------------------------ +# +# Establish the neuromuscular junctions that convert motor neuron spikes +# into muscle fiber activation. + +# Connect motor neurons to muscle (spike threshold automatically uses population default of 50.0 mV) +network.connect_to_muscle( + "aMN", muscle=hill_muscle, activation_callback=spkEvent, weight__uS=1.0 * pq.uS +) + +print("(OK) Connected motor neurons to muscle") + +############################################################################## +# Configure External Inputs +# ------------------------ +# +# Setup external input pathways for sensory feedback and descending commands. + +network.connect_from_external("cmd", "DD", weight__uS=1.0 * pq.uS) +network.connect_from_external("spindle", "Ia", weight__uS=1.0 * pq.uS) +network.connect_from_external("spindle", "II", weight__uS=1.0 * pq.uS) +network.connect_from_external("gto", "Ib", weight__uS=1.0 * pq.uS) + +# Get NetCons for manual triggering during simulation +ncD = { + "cmd->DD": network.get_netcons("cmd", "DD"), + "Spindle->Ia": network.get_netcons("spindle", "Ia"), + "Spindle->II": network.get_netcons("spindle", "II"), + "GTO->Ib": network.get_netcons("gto", "Ib"), +} + +print("(OK) Configured external input pathways") + +############################################################################## +# Prepare Simulation Models +# ------------------------ + +# Package all models for the simulation runner +models = { + "hill_muscle": hill_muscle, + "spin": spin, + "gto": gto, + "joint": joint_dynamics, +} + + +# Tendon tap parameters - multiple taps with increasing amplitudes +# Pattern: Taps every 0.5s with increasing amplitude, reset at 2.5s +# High amplitudes to elicit strong reflex responses +# Phase 1 (0-2.5s): 50%, 60%, 70%, 80% +# Phase 2 (2.5-5s): 50%, 60%, 70%, 80% + +TENDON_TAP_DURATION = 100 # ms - duration of each stretch +TENDON_TAP_SCHEDULE = [ + # Phase 1: Ascending amplitudes without cortical drive + (500, 50.0), # 0.5s: 50% + (1000, 60.0), # 1.0s: 60% + (1500, 70.0), # 1.5s: 70% + (2000, 80.0), # 2.0s: 80% + # Phase 2: Repeat pattern with cortical drive + (3000, 50.0), # 3.0s: 50% + (3500, 60.0), # 3.5s: 60% + (4000, 70.0), # 4.0s: 70% + (4500, 80.0), # 4.5s: 80% +] + +print("\nTendon tap schedule:") +print(f"\t- Duration per tap: {TENDON_TAP_DURATION} ms") +print(f"\t- Phase 1 (0-2.5s): {len([t for t, m in TENDON_TAP_SCHEDULE if t < 2500])} taps") +print(f"\t- Phase 2 (2.5-5s): {len([t for t, m in TENDON_TAP_SCHEDULE if t >= 2500])} taps") +print(f"\t- Amplitudes: {sorted(set([m for t, m in TENDON_TAP_SCHEDULE]))}%") + + +# Create step callback function with access to step counter +def step_callback(step_counter): + return eachStep( + muscle=hill_muscle, + spin=spin, + golgi=gto, + popD=network.populations, + ncD=ncD, + gMN=gMN, + joint_dyn=joint_dynamics, + step_counter=step_counter, + tendon_tap_schedule=TENDON_TAP_SCHEDULE, + tendon_tap_duration=TENDON_TAP_DURATION, + ) + + +############################################################################## +# Run Spinal Network Simulation +# ---------------------------- +# +# Execute the complete simulation with all integrated components. + +print("\nStarting spinal network simulation...") +print(f"\tDuration: {tstop} ms") +print(f"\tTimestep: {dt} ms") +print(f"\tPopulations: {len(network.populations)}") + +runner = SimulationRunner( + network=network, + models=models, + step_callback=step_callback, +) + +# Motor neuron spike recording thresholds are now fixed in the Network class + +results = runner.run( + duration__ms=tstop * pq.ms, + timestep__ms=dt, + membrane_recording={ + "aMN": [0, 5, 10, 15, 20, 30, 40, 50, 60, 70], + }, +) + +print("Simulation completed successfully!") + +# Save simulation results +joblib.dump(results, save_path / "spinal_network_results.pkl") +joblib.dump(artAng, save_path / "joint_angles.pkl") +joblib.dump(musculotendon_force__N, save_path / "musculotendon_force__N.pkl") + +# Save drive signals for analysis and plotting +drive_signals = { + "time": time, + "descending_drive": DDdrive, + "gamma_dynamic": gDyn, + "gamma_static": gStat, + "tendon_tap_schedule": TENDON_TAP_SCHEDULE, + "tendon_tap_duration": TENDON_TAP_DURATION, +} +joblib.dump(drive_signals, save_path / "drive_signals.pkl") + +# Save event timeline for plot annotations +events = { + # Tendon tap events: (peak_time_ms, amplitude_%, duration_ms) + # Note: peak_time is at the midpoint of the triangular pulse + "tendon_taps": [ + (t + TENDON_TAP_DURATION / 2.0, m, TENDON_TAP_DURATION) for t, m in TENDON_TAP_SCHEDULE + ], + # Gamma drive transitions: (time_ms, level_pps) + "gamma_transitions": list(zip(gamma_times, gamma_values)), + # Gamma drive periods: (start_time_ms, end_time_ms, level_pps) + "gamma_periods": [ + (gamma_times[i], gamma_times[i + 1], gamma_values[i]) for i in range(len(gamma_times) - 1) + ], + # Phase markers + "phase_boundary": 2500, # ms - transition from phase 1 to phase 2 + "cortical_drive_onset": cortical_start_time, # ms - when sinusoidal drive starts + # Experimental phases with descriptions + "phases": [ + {"start": 0, "end": 2500, "name": "Phase 1", "description": "Reflex gain modulation"}, + { + "start": 2500, + "end": 5000, + "name": "Phase 2", + "description": "Reflex-voluntary interaction", + }, + ], +} +joblib.dump(events, save_path / "event_timeline.pkl") + +print(f"Results saved to {save_path}") +print("\t- spinal_network_results.pkl (neural activity, muscle, spindle, GTO)") +print("\t- joint_angles.pkl (joint kinematics)") +print("\t- musculotendon_force__N.pkl (real-unit musculotendon force [N])") +print("\t- drive_signals.pkl (cortical drive, gamma drive, tap schedule)") +print("\t- event_timeline.pkl (event markers for plot annotations)") +print("\nNote: Intrafusal fiber tensions [FU] are in spinal_network_results.pkl") +print("\tAccess via: results.segments[0].analogsignals (look for 'intrafusal_tensions')") +print("\tConvert to real units: T[N] ≈ T[FU] * scaling_factor (no standard conversion)") +print(f"\tMusculotendon force capacity: {hill_muscle.F0:.1f} N") + +############################################################################## +# Export to NWB Format (Neurodata Without Borders) +# ------------------------------------------------ +# +# NWB is a standardized data format for neurophysiology data that enables +# data sharing and interoperability with other neuroscience tools. +# MyoGen supports NWB export for integration with the broader neuroscience +# ecosystem, including tools like the DANDI Archive. + +# Export the Neo Block results to NWB format +nwb_filepath = export_to_nwb( + results, + save_path / "spinal_network_results.nwb", + session_description=( + "MyoGen spinal network simulation with systematic tendon tap protocol. " + "Two-phase design: (1) Reflex gain modulation with varying gamma drive, " + "(2) Reflex-voluntary interaction with sinusoidal cortical drive." + ), + experimenter="MyoGen Simulation", + institution="MyoGen Framework", + lab="Neuromuscular Simulation", + experiment_description=( + f"5-second simulation with {naMN} motor neurons, " + f"{nIa} Ia afferents, {nII} II afferents, {nIb} Ib afferents, " + f"and {ngII + ngIb} spinal interneurons. " + f"Includes Hill-type muscle model and closed-loop joint dynamics." + ), + keywords=[ + "MyoGen", + "spinal network", + "motor neuron", + "stretch reflex", + "tendon tap", + "proprioception", + "EMG simulation", + ], + # Subject metadata for DANDI compliance + subject_id="simulated_subject_001", + species="Homo sapiens", + subject_description="Simulated human motor neuron pool and spinal reflex network", +) +print(f"\n(OK) Exported to NWB format: {nwb_filepath}") +print("\tNWB files can be validated with: nwbinspector ") + +############################################################################## +# Comprehensive Results Visualization +# --------------------------------- +# +# Create a series of plots that tell the complete story of spinal network +# function, from neural activity to mechanical output. + +print("\nGenerating comprehensive visualizations...") + +# 1. NEURAL ACTIVITY: Raster plot showing all population spike patterns +populations_list = [ + "aMN", # Motor output + "Ia", + "II", + "Ib", # Sensory input + "gII", + "gIb", # Interneurons + "DD", # Descending drive +] +fig1, axes1 = plt.subplots(len(populations_list), 1, figsize=(15, 12)) +plot_raster_spikes( + results, + axes1, + populations=populations_list, + time_range=(0, tstop), + title="Spinal Network Activity (Single Muscle)", +) +plt.tight_layout() +plt.savefig(save_path / "neural_raster_plot.png", dpi=150, bbox_inches="tight") +plt.show() + +# 2. MOTOR NEURON DYNAMICS: Membrane potentials showing integration +fig2, ax2 = plt.subplots(1, 1, figsize=(15, 6)) + +plot_membrane_potentials( + results, + [ax2], + populations=["aMN"], + cell_indices=[0, 10, 20, 30, 40], + time_range=(0, tstop), + title="Motor Neuron Membrane Potentials", +) + +plt.tight_layout() +plt.show() + +# 3. MUSCLE MECHANICS: Muscle dynamics +fig3, axes3 = plt.subplots(5, 1, figsize=(15, 20)) +plot_muscle_dynamics( + results, + artAng, + time, + axes3, + muscle_name="hill_muscle", + include_signals=["artAng", "L", "force", "torque"], + include_activations=["TypeI", "TypeII"], + normalize=True, + time_range=(0, tstop), + title="Muscle Dynamics - Length, Force, and Activation", +) +plt.tight_layout() +plt.savefig(save_path / "muscle_dynamics.png", dpi=150, bbox_inches="tight") +plt.show() + +# 4. PROPRIOCEPTIVE FEEDBACK: Spindle dynamics and sensory encoding +fig4, axes4 = plt.subplots(4, 1, figsize=(15, 16)) +plot_spindle_dynamics( + results, + axes4, + muscle_name="hill_muscle", + include_signals=["L"], + include_activations=["Bag1", "Bag2", "Chain"], + include_tensions=["Bag1", "Bag2", "Chain"], + include_afferents=["Ia", "II"], + time_range=(0, tstop), + title="Muscle Spindle Dynamics - Proprioceptive Feedback System", +) +plt.tight_layout() +plt.savefig(save_path / "spindle_dynamics.png", dpi=150, bbox_inches="tight") +plt.show() + +# 5. FORCE FEEDBACK: GTO dynamics and protective reflexes +fig5, axes5 = plt.subplots(2, 1, figsize=(15, 8)) +plot_gto_dynamics( + results, + axes5, + muscle_name="hill_muscle", + include_signals=["force", "Ib"], + time_range=(0, tstop), + title="Golgi Tendon Organ Dynamics - Force Feedback System", +) +plt.tight_layout() +plt.savefig(save_path / "gto_dynamics.png", dpi=150, bbox_inches="tight") +plt.show() diff --git a/examples/01_basic/12_extract_data_from_neo_blocks.py b/examples/01_basic/12_extract_data_from_neo_blocks.py index be5fe6d6..53aeed83 100644 --- a/examples/01_basic/12_extract_data_from_neo_blocks.py +++ b/examples/01_basic/12_extract_data_from_neo_blocks.py @@ -52,9 +52,27 @@ def _mean_firing_rate(spiketrain): save_path = Path("./results") -# Check which files exist +# Example 03 keeps the two simulator outputs side by side. Select one +# deliberately for the detailed analysis below while reporting availability of +# both. Change this constant to ``"neuron"`` to inspect the reference run. +SPIKE_BACKEND = "jaxley" +spike_train_paths = { + "jaxley": save_path / "trapezoid_dd_spike_trains_jaxley.pkl", + "neuron": save_path / "trapezoid_dd_spike_trains.pkl", +} + +if SPIKE_BACKEND not in spike_train_paths: + raise ValueError(f"SPIKE_BACKEND must be one of {tuple(spike_train_paths)}, got {SPIKE_BACKEND!r}") + +print("Example 03 spike-train artifacts:") +for backend, path in spike_train_paths.items(): + status = "available" if path.exists() else "missing" + print(f" {backend:7s}: {status:9s} ({path.name})") + +# Check which files exist. ``spike_trains`` is the explicitly selected backend, +# so the analysis never silently combines a Jaxley drive with NEURON spikes. files_to_check = { - "spike_trains": save_path / "trapezoid_dd_spike_trains.pkl", + "spike_trains": spike_train_paths[SPIKE_BACKEND], "surface_muaps": save_path / "surface_muaps.pkl", "surface_emg": save_path / "surface_emg_signals.pkl", "intramuscular_muaps": save_path / "intramuscular_muaps.pkl", @@ -80,6 +98,7 @@ def _mean_firing_rate(spiketrain): print("\n" + "=" * 70) print("EXAMPLE 1: SPIKE_TRAIN__Block") print("=" * 70) + print(f"Backend: {SPIKE_BACKEND}") # Load spike train block spike_train__Block: SPIKE_TRAIN__Block = joblib.load(available_files["spike_trains"]) diff --git a/examples/01_basic/README.md b/examples/01_basic/README.md index 3e93ee6d..2520e5aa 100644 --- a/examples/01_basic/README.md +++ b/examples/01_basic/README.md @@ -15,6 +15,11 @@ The examples are self-contained and can be run independently, though they may depend on each other for context. Following the numbered order provides the best learning experience. +Example 03 writes backend-qualified Jaxley artifacts (for example, +`trapezoid_dd_spike_trains_jaxley.pkl`) so they can coexist with the NEURON +reference artifacts. Canonical Examples 06, 08 and 09 consume the Jaxley files. +Example 12 reports both and uses an explicit `SPIKE_BACKEND` selection. + ## Profiles Examples 01, 03, 04, 05, 06 and 08 simulate the first dorsal interosseous (FDI) diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 00000000..9c8d5208 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,67 @@ +# MyoGen Examples — Backends + +MyoGen's default and recommended simulation backend is **Jaxley** (JAX-based, +GPU-capable, differentiable). NEURON is an **optional** backend kept as a validation +reference. + +- The **canonical** example for each topic (`XX_.py`) is the **Jaxley** version and + runs on a plain `pip install myogen` — no NEURON required. +- The NEURON version of a topic, where one exists, is the `XX__neuron.py` sibling + and needs the optional extra: `pip install "myogen[neuron]"`. + +## Default examples (Jaxley / backend-agnostic) — no NEURON needed + +These run on a plain `pip install myogen`. + +| Example | What it does | +|---|---| +| `01_basic/01_simulate_recruitment_thresholds.py` | Recruitment threshold distribution | +| `01_basic/02_simulate_spike_trains_current_injection.py` | Spike trains from current injection (Jaxley) | +| `01_basic/03_simulate_spike_trains_descending_drive.py` | Spike trains from descending drive (Jaxley) | +| `01_basic/04_simulate_muscle.py` | Muscle geometry / fiber distribution (`core/`) | +| `01_basic/05_simulate_surface_muaps.py` | Surface MUAP templates (`core/`) | +| `01_basic/06_simulate_surface_emg.py` | Surface EMG (`core/`) | +| `01_basic/07_simulate_currents.py` | Fiber currents (`core/`) | +| `01_basic/08_simulate_force.py` | Force from spikes (Jaxley) | +| `01_basic/09_simulate_intramuscular_emg.py` | Intramuscular EMG (`core/`) | +| `01_basic/10_extract_neuron_parameters.py` | Per-cell electrophysiology (Jaxley) | +| `01_basic/11_simulate_spinal_network.py` | Full closed-loop spinal network (Jaxley) | +| `01_basic/12_extract_data_from_neo_blocks.py` | Neo block post-processing | +| `01_basic/13_load_and_inspect_nwb_data.py` | NWB I/O | +| `02_finetune/05_plot_isi_cv_multi_muscle_comparison.py` | Plot saved ISI/CV results | +| `03_papers/watanabe/04_load_and_analyze_results.py` | Analyse saved Watanabe results | +| `03_papers/watanabe/05_compute_force_from_spinal_network.py` | Force from saved network | +| `03_papers/watanabe/06_visualize.py` | Visualise saved results | + +New differentiable pipeline (used by the Jaxley examples and available directly): +`from myogen.simulator.jaxley import run_jax, value_and_grad_run, compile_run, +surface_emg_jax`. See `docs/DIFFERENTIABILITY_STATUS.md`. + +## NEURON examples — require `pip install "myogen[neuron]"` + +These import the NEURON runtime directly. The `_neuron.py` files are the NEURON +counterparts of the canonical Jaxley examples above. + +| NEURON example | Canonical (Jaxley) counterpart | +|---|---| +| `01_basic/02_simulate_spike_trains_current_injection_neuron.py` | `02_simulate_spike_trains_current_injection.py` | +| `01_basic/03_simulate_spike_trains_descending_drive_neuron.py` | `03_simulate_spike_trains_descending_drive.py` | +| `01_basic/08_simulate_force_neuron.py` | `08_simulate_force.py` | +| `01_basic/10_extract_neuron_parameters_neuron.py` | `10_extract_neuron_parameters.py` | +| `01_basic/11_simulate_spinal_network_neuron.py` | `11_simulate_spinal_network.py` | +| `02_finetune/01_optimize_dd_for_target_firing_rate.py` | — (not yet ported) | +| `02_finetune/02_compute_force_from_optimized_dd.py` | — (not yet ported) | +| `02_finetune/03_optimize_dd_for_target_force.py` | — (not yet ported) | +| `02_finetune/04_extract_isi_and_cv_per_ramps.py` | — (not yet ported) | +| `03_papers/watanabe/01_compute_baseline_force.py` | — (not yet ported) | +| `03_papers/watanabe/02_optimize_oscillating_dc.py` | — (not yet ported) | +| `03_papers/watanabe/03_10pct_mvc_simulation.py` | — (not yet ported) | + +Running a NEURON example without the extra raises `ModuleNotFoundError: No module +named 'neuron'` (or an equivalent import error) — install `myogen[neuron]` to run it. + +## Migration status + +The `02_finetune` and `03_papers` NEURON simulations are not yet ported to Jaxley. +Porting them (or building Jaxley equivalents) is tracked as follow-up work; until then +they remain runnable via the NEURON extra. diff --git a/experiments/twin/README.md b/experiments/twin/README.md new file mode 100644 index 00000000..4982193f --- /dev/null +++ b/experiments/twin/README.md @@ -0,0 +1,34 @@ +# Digital twin of the motor neuron pool — experiments + +Research harness for the digital-twin program. **Canonical research plan: +`docs/DIGITAL_TWIN_ROADMAP_REVISED.md`** (supervisor review, 2026-10-13 — +sequence, gates, claim boundaries, Deliverables A/B/C). This harness covers +the drive→pool pathway only; afferents, muscle/force, and EMG are not wired +in yet. + +**Status: differentiable proof-of-concept harness.** Simulator validation +(roadmap Phase 1 strict package) has not passed yet; inverse-recovery results +are unconverged and reported at the best iterate with caveats. Do not cite +this as a validated digital twin. + +## Layout + +| File | Role | +|---|---| +| `poolbuild.py` | Build a NERLab MN pool + compiled step fn + differentiable θ from a `TwinSpec`; knobs: gCaL / gNaP scales, drive gain, exc/inh weights, reversal potentials; synapse site (soma/dendrite), fixed/live driving force, single-exp/Exp2Syn kernel; `simulate_events()` accepts exact event + connectivity matrices for matched-input validation. No afferent/muscle/EMG pathway (see roadmap) | +| `protocols.py` | Command waveforms: triangular ramp (ΔF), trapezoid, step, hold | +| `observables.py` | Smoothed rates (differentiable) + ΔF, recruitment order, hysteresis (analysis) | +| `validation_protocol.py` | Paired NEURON↔Jaxley validation: immutable content-hashed protocol artifact (events, connectivity, synaptic constants, current steps); Jaxley runner implemented, NEURON runner stubbed with checklist | +| `phase1_pic_validation.py` | PIC behaviour: ramp hysteresis, ΔF dose–response vs `gCaL`/`gNaP` scales; matched-synapse CLI; drive vs connectivity variability separated | +| `phase2_benchmark.py` | vmap-batched forward + gradient throughput benchmark | +| `phase3_inverse_mvp.py` | Synthetic inverse problem: hide θ*, recover with optax + surrogate gradients | +| `results/` | Figures / pickles (git-ignored artifacts) | + +## Running + +Scripts are self-contained; run from this directory with the project venv: + +```bash +cd experiments/twin +../../.venv/bin/python phase1_pic_validation.py +``` diff --git a/experiments/twin/acceptance.py b/experiments/twin/acceptance.py new file mode 100644 index 00000000..c04ae594 --- /dev/null +++ b/experiments/twin/acceptance.py @@ -0,0 +1,189 @@ +"""Phase-1C acceptance evaluator. + +Loads a FROZEN acceptance specification (``configs/phase1_acceptance_v*.yaml``) +and a set of per-(condition, seed) paired-comparison metrics (produced by +``pic_grid.py``), and classifies every metric at every point as:: + + PASS | FAIL | OUTSIDE_VALIDATED_ENVELOPE | NOT_EVALUABLE + +NOT_EVALUABLE covers right-censored after-discharge durations, ΔF with too +few eligible pairs, and spike-timing fractions with too few spikes — these +are reported, never silently dropped or counted as passes. + +The evaluator never edits thresholds: a spec revision requires a new +versioned file. The summary reports the failure fraction against the spec's +``failure_policy`` and the maximal passing CaL range per NaP level (the +candidate validated envelope). +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import yaml + +CONFIGS = Path(__file__).parent / "configs" + + +def load_spec(version: int = 1) -> dict: + path = CONFIGS / f"phase1_acceptance_v{version}.yaml" + spec = yaml.safe_load(path.read_text()) + if spec.get("status") != "FROZEN": + raise ValueError(f"{path} is not frozen — refusing to evaluate") + return spec + + +def _verdict(value, threshold, *, minimum=False, not_evaluable=False): + if not_evaluable: + return "NOT_EVALUABLE", value + if value is None or (isinstance(value, float) and not np.isfinite(value)): + return "NOT_EVALUABLE", value + ok = (value >= threshold) if minimum else (value <= threshold) + return ("PASS" if ok else "FAIL"), value + + +def evaluate_point(metrics: dict, spec: dict) -> dict: + """Classify one (condition, seed) metric dict against the spec.""" + m = spec["metrics"] + out = {} + + out["binned_rate_rmse_hz2"] = _verdict( + metrics.get("binned_rate_rmse_hz2"), + m["binned_rate_rmse_hz2"]["threshold"]) + + out["active_unit_mismatch"] = _verdict( + metrics.get("active_unit_mismatch"), + m["active_unit_mismatch"]["threshold"]) + + out["total_spike_rel_diff"] = _verdict( + metrics.get("total_spike_rel_diff"), + m["total_spike_rel_diff"]["threshold"]) + + n_spk = metrics.get("spike_total_max", 0) + out["spike_match_pm5ms_fraction_min"] = _verdict( + metrics.get("spike_match_pm5ms_fraction"), + m["spike_match_pm5ms_fraction_min"]["threshold"], + minimum=True, not_evaluable=n_spk < 20) + + out["hysteresis_mean_diff_hz"] = _verdict( + metrics.get("hysteresis_mean_diff_hz"), + m["hysteresis_mean_diff_hz"]["threshold"], + not_evaluable=metrics.get("hysteresis_mean_diff_hz") is None) + + out["self_sustained_mismatch"] = _verdict( + metrics.get("self_sustained_mismatch"), + m["self_sustained_mismatch"]["threshold"]) + + pairs_ok = (metrics.get("delta_f_pairs_jaxley", 0) + >= m["delta_f_mean_diff_hz"]["requires_min_pairs_each"] + and metrics.get("delta_f_pairs_neuron", 0) + >= m["delta_f_mean_diff_hz"]["requires_min_pairs_each"]) + out["delta_f_mean_diff_hz"] = _verdict( + metrics.get("delta_f_mean_diff_hz"), + m["delta_f_mean_diff_hz"]["threshold"], not_evaluable=not pairs_ok) + + out["delta_f_eligible_fraction_diff"] = _verdict( + metrics.get("delta_f_eligible_fraction_diff"), + m["delta_f_eligible_fraction_diff"]["threshold"]) + + return out + + +def evaluate_grid(points: list[dict], spec: dict) -> dict: + """Evaluate all grid points + the boundary metric + failure policy. + + ``points``: list of dicts with keys ``gcal``, ``gnap``, ``seed``, + ``metrics`` (per-point metric dict) — see ``pic_grid.py``. + """ + env = spec["declared_envelope_candidate"] + results = [] + n_fail = n_eval = 0 + for pt in points: + inside = (env["gcal_scale"][0] <= pt["gcal"] <= env["gcal_scale"][1] + and env["gnap_scale"][0] <= pt["gnap"] <= env["gnap_scale"][1]) + verdicts = evaluate_point(pt["metrics"], spec) + if not inside: + verdicts = {k: ("OUTSIDE_VALIDATED_ENVELOPE", v[1]) + for k, v in verdicts.items()} + for k, (cls, _val) in verdicts.items(): + if cls == "FAIL": + n_fail += 1 + if cls in ("PASS", "FAIL"): + n_eval += 1 + results.append({**{k: pt[k] for k in ("gcal", "gnap", "seed")}, + "verdicts": {k: {"class": c, "value": v} + for k, (c, v) in verdicts.items()}}) + + # boundary agreement per NaP level (condition-level, seed-pooled) + boundary = {} + for gnap in sorted({pt["gnap"] for pt in points}): + rows = [p for p in points if p["gnap"] == gnap] + gcals = sorted({p["gcal"] for p in rows}) + + def first_bistable(backend): + for g in gcals: + if any(p["metrics"].get(f"self_sustained_{backend}", 0) > 0 + for p in rows if p["gcal"] == g): + return g + return None + + bj, bn = first_bistable("jaxley"), first_bistable("neuron") + if bj is None and bn is None: + cls, off = "NOT_EVALUABLE", None + elif bj is None or bn is None: + cls, off = "FAIL", None + else: + off = abs(gcals.index(bj) - gcals.index(bn)) + cls = ("PASS" if off <= spec["metrics"][ + "bistability_boundary_offset_steps"]["threshold"] else "FAIL") + n_eval += 1 + n_fail += cls == "FAIL" + boundary[gnap] = {"jaxley_first_bistable_gcal": bj, + "neuron_first_bistable_gcal": bn, + "offset_steps": off, "class": cls} + + fail_fraction = n_fail / max(n_eval, 1) + policy_ok = fail_fraction <= spec["failure_policy"]["max_fail_fraction"] + + # maximal contiguous passing CaL range per NaP (the validated envelope) + envelope = {} + for gnap in boundary: + rows = [r for r in results if r["gnap"] == gnap] + gcals = sorted({r["gcal"] for r in rows}) + passing = [] + for g in gcals: + cells = [r for r in rows if r["gcal"] == g] + cls_list = [v["class"] for r in cells + for v in r["verdicts"].values()] + if "FAIL" not in cls_list: + passing.append(g) + envelope[gnap] = ([min(passing), max(passing)] if passing else None) + + return { + "spec_version": spec["version"], + "n_points": len(points), + "n_evaluated_metric_cells": n_eval, + "n_failed_metric_cells": n_fail, + "fail_fraction": fail_fraction, + "failure_policy_ok": bool(policy_ok), + "bistability_boundary": boundary, + "validated_envelope_gcal_by_gnap": envelope, + "points": results, + } + + +if __name__ == "__main__": + import argparse + + ap = argparse.ArgumentParser() + ap.add_argument("grid_json", type=Path, + help="pic_grid output JSON with a 'points' list") + ap.add_argument("--spec-version", type=int, default=1) + args = ap.parse_args() + spec = load_spec(args.spec_version) + grid = json.loads(args.grid_json.read_text()) + rep = evaluate_grid(grid["points"], spec) + print(json.dumps({k: v for k, v in rep.items() if k != "points"}, + indent=1, default=float)) diff --git a/experiments/twin/configs/phase1_acceptance_v1.yaml b/experiments/twin/configs/phase1_acceptance_v1.yaml new file mode 100644 index 00000000..9a307078 --- /dev/null +++ b/experiments/twin/configs/phase1_acceptance_v1.yaml @@ -0,0 +1,92 @@ +# Phase-1C acceptance specification — v1 (FROZEN) +# +# Rule: a revision of any threshold after confirmatory runs begin requires a +# new file (v2); thresholds below were derived from EXPLORATORY data only and +# are frozen before the confirmatory PIC grid (new seeds) runs. +# +# Derivation rule per metric: T = max(k * E_numerical, alpha * D_scientific) +# E_numerical — measured timestep/precision refinement error (exploratory) +# D_scientific — smallest physiological difference the inverse experiment +# must distinguish (declared) +# k — numerical safety multiplier +# alpha — acceptable fraction of the physiological effect +# +# Exploratory evidence used for E_numerical and sanity levels: +# - deliverable_b_full_binned_* : binned-rate dt-halving EPS 0.30–0.83 Hz² +# - val_comparison_ef1bd7e8... : subthreshold backend RMSE 0.00035–0.003 mV; +# step spike-count rel diff ≤ 1/745; +# step ±5 ms match 0.62–0.73 +# - val_comparison_bac41198... : pool ±5 ms match 0.87; counts 407 vs 405 +# - phase1_pic_validation runs : hysteresis −10…−14 Hz (×6.5), 7/7 +# self-sustained at CaL×6.5+NaP×5 + +version: 1 +frozen_at: "2026-10-18" +status: FROZEN +applies_to: "paired NEURON–Jaxley pool-phase comparisons (PIC grid, ramp protocol)" + +verdict_classes: [PASS, FAIL, OUTSIDE_VALIDATED_ENVELOPE, NOT_EVALUABLE] + +metrics: + binned_rate_rmse_hz2: + threshold: 2.5 + derivation: {k: 3, E_numerical: 0.83, alpha: 0.5, D_scientific: 5.0} + units: "Hz^2 (50 ms bins, 300 ms smoothing, per-unit mean)" + active_unit_mismatch: + threshold: 1 + derivation: {k: null, E_numerical: 0, alpha: 0.125, D_scientific: 8} + units: "units (|active_jaxley - active_neuron| and set difference)" + total_spike_rel_diff: + threshold: 0.05 + derivation: {k: 3, E_numerical: 0.005, alpha: 0.25, D_scientific: 0.2} + units: "fraction of max(total spikes, 1)" + spike_match_pm5ms_fraction_min: + threshold: 0.60 + derivation: {note: "exploratory pool 0.87, deterministic steps 0.62-0.73; + timing drift is implementation-level — floor set below the + worst exploratory value by design"} + units: "matched / max(jaxley_total, neuron_total); NOT_EVALUABLE if < 20 spikes" + hysteresis_mean_diff_hz: + threshold: 2.0 + derivation: {k: null, alpha: 0.5, D_scientific: 4.0, + note: "D = hysteresis change per CaL grid step (exploratory)"} + units: "Hz drive" + self_sustained_mismatch: + threshold: 1 + derivation: {alpha: 0.125, D_scientific: 8} + units: "units; after-discharge censored at recording end → NOT_EVALUABLE + for duration comparisons, count comparison still evaluable" + delta_f_mean_diff_hz: + threshold: 2.0 + requires_min_pairs_each: 3 + derivation: {alpha: 0.2, D_scientific: 10.0, + note: "ΔF spans ~0–16 Hz across the grid (exploratory)"} + units: "Hz; NOT_EVALUABLE if either backend has < 3 eligible pairs" + delta_f_eligible_fraction_diff: + threshold: 0.25 + derivation: {note: "eligibility is protocol/estimator-driven; fraction + difference beyond 0.25 means the two backends are in + different ΔF regimes"} + units: "absolute fraction of possible pairs" + bistability_boundary_offset_steps: + threshold: 1 + derivation: {note: "boundary = first CaL grid value with any + self-sustained unit; must agree within one grid step"} + units: "CaL grid steps (per NaP level)" + +failure_policy: + max_fail_fraction: 0.10 # across (condition × seed × metric) inside envelope + rule: "conditions outside the declared envelope are OUTSIDE_VALIDATED_ENVELOPE, + not FAIL; the maximal contiguous passing region DEFINES the validated + envelope reported to downstream inference" + +declared_envelope_candidate: + gcal_scale: [0.0, 7.5] + gnap_scale: [1.0, 5.0] + +downstream_resolution_declaration: + note: "Phase-3 ladder: a parameter advances only if its sampled + compatibility interval is within ±20 % of truth across replicates + (predeclared; see r2a profiles)" + cal_resolution_rel: 0.20 + nap_resolution_rel: 0.20 diff --git a/experiments/twin/deliverable_b.py b/experiments/twin/deliverable_b.py new file mode 100644 index 00000000..09977a1f --- /dev/null +++ b/experiments/twin/deliverable_b.py @@ -0,0 +1,605 @@ +"""Deliverable B (v3) — two-parameter hard-forward inverse experiment. + +Revised roadmap §12B / ladder Level B. v3 fixes the v2 profile-grid bug: the +dense axes did NOT contain the true parameter values, so the profile curves +could never reach the known zero-loss optimum and every "profile minimum" +(1.6–3.4 Hz²) was grid-sampling error, not physiology. v3: + +* ``axis_with_truth``: both dense axes contain the exact true values; for + every replicate and observation model the profile at the truth is + ASSERTED to be exactly 0 (same model, same events — anything else is a + bug). +* Adaptive nuisance refinement: hard-spike loss is discontinuous, so each + profile point refines the nuisance axis locally (3 extra log-spaced + evaluations between the coarse bracketing neighbours, per observation + model) over a shared evaluation cache. +* EPS is computed PER replicate AND PER observation model (dt-halving at + the truth, per-unit terms averaged over that model's units), not one + global value; the conservative maximum is also reported. +* Smoke and canonical outputs are separated: results/figures are tagged + ``full`` or ``smoke`` from the actual run configuration, so a smoke run + can never overwrite scientific evidence again. + +Declared setup (unchanged from v2): unknowns theta = (gCaL_scale, w_exc); +fixed/known: gNaP ×1, no inhibition, connectivity, drive shape AND frozen +event realizations (conditional fit); fit protocols ramp + ramp–hold–ramp; +untouched held-out slower ramp; observation models all-units vs +lowest-threshold half (active counts reported); observable selected by +``--observable`` (binned = 50 ms bins + 300 ms smoothing, the low-variance +Phase-2A choice; pointwise = exponential tau 100 ms). Estimators: grid + +refinement (primary; does not see the truth) and the unpromoted clipped +sign-descent gradient method with hard-evaluated iterates and non-finite +failure accounting. + +Usage:: + + python deliverable_b.py # full run (tag "full") + python deliverable_b.py --dense-n 7 --replicates 1 --dur-scale 0.5 # smoke +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import jax +import jax.numpy as jnp +import optax + +from poolbuild import PoolTwin, TwinSpec +from validation_protocol import _renewal_events +from sensitivity_checks import _manifest, _smoothed_rates_jax + +RESULTS = Path(__file__).parent / "results" +EVIDENCE = Path(__file__).parent / "evidence" + +OKABE = {"truth": "#000000", "grid": "#0072B2", "grad": "#E69F00", + "grad2": "#CC79A7", "partial": "#56B4E9"} +MUTED, GRIDC = "#666666", "#e3e3e3" + +EPS_FLOOR = 1e-3 # Hz² — declared floor for the equivalence tolerance + + +# --------------------------------------------------------------------------- # +# Axes / protocols / segmentation / observables (module-level, testable) +# --------------------------------------------------------------------------- # + +def axis_with_truth(lo: float, hi: float, n: int, truth: float) -> np.ndarray: + """Log-spaced axis with the true value inserted — profiles must be able + to reach the known zero-loss optimum (v2 bug: they could not).""" + axis = np.geomspace(lo, hi, n) + return np.unique(np.sort(np.append(axis, truth))) + + +def make_commands(dt_ms: float, dur_scale: float) -> dict: + def tri(dur, peak, pre=0.5, post=0.5): + n = int(dur * 1000 / dt_ms) + t = np.arange(n) * dt_ms * 1e-3 + return np.interp(t, [0, pre, (dur - post + pre) / 2, dur - post, dur], + [0, 0, peak, 0, 0]).astype(np.float32) + + def trap(dur, peak, pre=0.4, rise=1.0, hold=1.2, fall=1.0): + n = int(dur * 1000 / dt_ms) + t = np.arange(n) * dt_ms * 1e-3 + return np.interp( + t, [0, pre, pre + rise, pre + rise + hold, + pre + rise + hold + fall, dur], + [0, 0, peak, peak, 0, 0]).astype(np.float32) + + s = dur_scale + return {"fit_ramp": tri(4.0 * s, 60.0), "fit_trap": trap(4.0 * s, 50.0), + "heldout_slowramp": tri(5.0 * s, 45.0)} + + +def segment_limbs(cmd: np.ndarray) -> dict: + """Boolean masks: pre / rising / peak (≥90 % of max) / falling / post.""" + cmd = np.asarray(cmd, dtype=float) + peak_val = cmd.max() + i_peak = int(np.argmax(cmd)) + active = cmd > 0 + first_on = int(np.argmax(active)) if active.any() else len(cmd) + last_on = len(cmd) - 1 - int(np.argmax(active[::-1])) if active.any() else 0 + idx = np.arange(len(cmd)) + peak = cmd >= 0.9 * peak_val + rising = active & (idx < i_peak) & ~peak + falling = active & (idx > i_peak) & ~peak + pre = idx < first_on + post = idx > last_on + return {"pre": pre, "rising": rising, "peak": peak, + "falling": falling, "post": post} + + +def _sha16(arr) -> str: + return hashlib.sha256(np.ascontiguousarray(arr).tobytes()).hexdigest()[:16] + + +def binned_rates_jnp(spikes, dt_ms, bin_ms=50.0, smooth_ms=300.0): + """(n_bins, n_units) smoothed binned rates in Hz — jax, differentiable. + Insensitive to spike-time jitter within a bin (kills the chaotic + decorrelation floor of the pointwise observable).""" + steps = int(round(bin_ms / dt_ms)) + n_bins = spikes.shape[0] // steps + counts = spikes[: n_bins * steps].reshape(n_bins, steps, -1).sum(axis=1) + r = counts / (bin_ms * 1e-3) + w = max(1, int(round(smooth_ms / bin_ms))) + kern = jnp.ones((w,)) / w + return jax.vmap(lambda col: jnp.convolve(col, kern, mode="same"), + in_axes=1, out_axes=1)(r) + + +def bin_mask(mask, dt_ms, bin_ms=50.0): + """Downsample a per-step boolean mask to per-bin (majority vote).""" + steps = int(round(bin_ms / dt_ms)) + n_bins = len(mask) // steps + return (np.asarray(mask[: n_bins * steps], dtype=float) + .reshape(n_bins, steps).mean(axis=1) > 0.5) + + +def build_twin_and_events(n_mn, conn_seed, event_seed_offset, commands, dt_ms): + twin = PoolTwin(TwinSpec(n_mn=n_mn, n_dd=100, n_inh=5, dt_ms=dt_ms, + seed=conn_seed)) + events = {} + for k, (name, cmd) in enumerate(sorted(commands.items())): + rng = np.random.default_rng(event_seed_offset + k) + events[name] = jnp.asarray( + _renewal_events(rng, np.tile(cmd[:, None], (1, 100)), dt_ms, 5)) + return twin, events + + +def make_observable_fn(twin, events, observable, tau_ms=100.0): + dt = twin.spec.dt_ms + + def fn(theta_vals, proto, spike_mode="hard"): + th = twin.default_theta() + th["gCaL_scale"] = jnp.float32(theta_vals[0]) + th["w_exc"] = jnp.float32(theta_vals[1]) + out = twin.simulate_events(th, events[proto], spike_mode=spike_mode) + if observable == "pointwise": + return _smoothed_rates_jax(out["spikes"], dt, tau_ms) + return binned_rates_jnp(out["spikes"], dt) + + return fn + + +# legacy alias used by regression tests +def make_rates_fn(twin, events, tau_ms=100.0): + return make_observable_fn(twin, events, "pointwise", tau_ms) + + +def numerical_epsilon_spikes(n_mn, conn_seed, events_01, theta_true, + fit_protos): + """Fine-dt (0.05 ms) hard spike matrices at the truth, events re-timed by + zero-insertion so the same input spikes occur at the same physical + times.""" + twin_h = PoolTwin(TwinSpec(n_mn=n_mn, n_dd=100, n_inh=5, dt_ms=0.05, + seed=conn_seed)) + fine_spikes = [] + for p in fit_protos: + ev = np.asarray(events_01[p]) + fine = np.zeros((ev.shape[0] * 2, ev.shape[1]), dtype=np.float32) + fine[::2] = ev + th = twin_h.default_theta() + th["gCaL_scale"] = jnp.float32(theta_true[0]) + th["w_exc"] = jnp.float32(theta_true[1]) + out = twin_h.simulate_events(th, jnp.asarray(fine), spike_mode="hard") + fine_spikes.append(np.asarray(out["spikes"])) + return fine_spikes + + +def run_tag(args) -> str: + """'full' only for scientifically sized runs; everything else is smoke — + tagged filenames keep smoke runs from overwriting canonical evidence.""" + full = (args.dense_n >= 15 and args.replicates >= 2 + and args.dur_scale >= 1.0 and args.n_mn >= 8) + return "full" if full else "smoke" + + +# --------------------------------------------------------------------------- # + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-mn", type=int, default=8) + ap.add_argument("--true", type=float, nargs=2, default=[1.6, 0.0105]) + ap.add_argument("--dense-n", type=int, default=21) + ap.add_argument("--grid-n", type=int, default=7) + ap.add_argument("--refine-n", type=int, default=5) + ap.add_argument("--budget-iters", type=int, default=15) + ap.add_argument("--replicates", type=int, default=3) + ap.add_argument("--dur-scale", type=float, default=1.0) + ap.add_argument("--observable", choices=["binned", "pointwise"], + default="binned") + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + tag = run_tag(args) + stem = f"deliverable_b_{tag}_{args.observable}" + t_wall0 = time.time() + dt_ms = 0.1 + theta_true = np.array(args.true) + fit_protos = ["fit_ramp", "fit_trap"] + commands = make_commands(dt_ms, args.dur_scale) + + g_lo, g_hi = 0.25, 4.0 + w_lo, w_hi = 0.0045, 0.018 + g_dense = axis_with_truth(g_lo, g_hi, args.dense_n, theta_true[0]) + w_dense = axis_with_truth(w_lo, w_hi, args.dense_n, theta_true[1]) + gi_true = int(np.argmin(np.abs(g_dense - theta_true[0]))) + wi_true = int(np.argmin(np.abs(w_dense - theta_true[1]))) + assert g_dense[gi_true] == theta_true[0] and w_dense[wi_true] == theta_true[1] + + obs_models = {"full_obs": list(range(args.n_mn)), + "partial_obs": list(range(args.n_mn // 2))} + + replicate_reports = [] + rep0 = None + for r in range(args.replicates): + conn_seed, ev_off = r, 10 * r + twin, events = build_twin_and_events( + args.n_mn, conn_seed, ev_off, commands, dt_ms) + obs_fn = make_observable_fn(twin, events, args.observable) + obs = {p: obs_fn(theta_true, p) for p in commands} + active = {} + for p in commands: + th = twin.default_theta() + th["gCaL_scale"] = jnp.float32(theta_true[0]) + th["w_exc"] = jnp.float32(theta_true[1]) + sp = np.asarray(twin.simulate_events( + th, events[p], spike_mode="hard")["spikes"]) + active[p] = int((sp.sum(axis=0) > 0).sum()) + + cache: dict = {} + + def pu_at(g, w): + key = (float(g), float(w)) + if key not in cache: + per = [] + for p in fit_protos: + rr = obs_fn(np.array([g, w]), p) + per.append(np.asarray(jnp.mean((rr - obs[p]) ** 2, axis=0))) + cache[key] = np.mean(per, axis=0) + return cache[key] + + # truth must be exactly 0 (same model, same events) + truth_pu = pu_at(*theta_true) + assert float(truth_pu.mean()) == 0.0, "truth loss must be exactly 0" + + t0 = time.time() + pu = np.zeros((len(g_dense), len(w_dense), args.n_mn)) + for i, g in enumerate(g_dense): + for j, w in enumerate(w_dense): + pu[i, j] = pu_at(g, w) + dense_time = time.time() - t0 + + # per-replicate, per-model EPS from dt-halving at the truth + fine_spikes = numerical_epsilon_spikes( + args.n_mn, conn_seed, events, theta_true, fit_protos) + eps_pu_terms = [] + for p, sp_f in zip(fit_protos, fine_spikes): + if args.observable == "pointwise": + rf = np.asarray(_smoothed_rates_jax( + jnp.asarray(sp_f), 0.05, 100.0))[1::2] + else: + rf = np.asarray(binned_rates_jnp(jnp.asarray(sp_f), 0.05)) + rc = np.asarray(obs[p]) + m = min(len(rf), len(rc)) + eps_pu_terms.append(((rf[:m] - rc[:m]) ** 2).mean(axis=0)) + eps_pu = np.mean(eps_pu_terms, axis=0) # (n_mn,) + eps = {model: max(float(eps_pu[units].mean()), EPS_FLOOR) + for model, units in obs_models.items()} + + # profiles with local nuisance refinement, per model + def refined_profile(units, vary_g: bool): + axis_p = g_dense if vary_g else w_dense + axis_n = w_dense if vary_g else g_dense + prof = np.zeros(len(axis_p)) + for ai, aval in enumerate(axis_p): + row = [float(pu_at(aval, nv)[units].mean()) if vary_g + else float(pu_at(nv, aval)[units].mean()) + for nv in axis_n] + jbest = int(np.argmin(row)) + best = row[jbest] + lo = axis_n[max(jbest - 1, 0)] + hi = axis_n[min(jbest + 1, len(axis_n) - 1)] + for nv in np.geomspace(lo, hi, 5)[1:-1]: + v = (pu_at(aval, nv) if vary_g + else pu_at(nv, aval))[units].mean() + best = min(best, float(v)) + prof[ai] = best + return prof + + rep = {"conn_seed": conn_seed, "event_offset": ev_off, + "active_units_at_truth": active, + "truth_fit_loss": float(truth_pu.mean()), + "eps_by_model": eps, + "dense_time_s": round(dense_time, 1), + "event_hashes": {p: _sha16(np.asarray(events[p])) + for p in commands}} + for model, units in obs_models.items(): + pg = refined_profile(units, vary_g=True) + pw = refined_profile(units, vary_g=False) + # the profile at the truth MUST be exactly zero in both directions + assert pg[gi_true] == 0.0 and pw[wi_true] == 0.0, ( + f"profile at truth nonzero ({pg[gi_true]}, {pw[wi_true]}) — " + "axis/profiling bug") + + def interval(curve, axis): + within = curve <= eps[model] + return ([float(axis[within].min()), float(axis[within].max())] + if within.any() else None) + + rep[model] = { + "observed_units": units, + "eps": eps[model], + "profile_gCaL": pg.tolist(), + "profile_w_exc": pw.tolist(), + "gcal_interval": interval(pg, g_dense), + "wexc_interval": interval(pw, w_dense), + "dense_loss": pu[:, :, units].mean(axis=2).tolist(), + } + print(f"[rep{r} conn{conn_seed}/ev{ev_off}] dense+profiles " + f"({len(cache)} evals, {time.time()-t0:.0f}s); truth=0 OK; " + f"EPS full {eps['full_obs']:.4f} / partial " + f"{eps['partial_obs']:.4f}; active {active}") + for model in obs_models: + print(f" {model}: gCaL interval {rep[model]['gcal_interval']} " + f"| w_exc interval {rep[model]['wexc_interval']}") + rep["_runtime"] = (twin, events, obs_fn, obs, pu, cache) + replicate_reports.append(rep) + if rep0 is None: + rep0 = rep + + # ------------------------------------- estimator comparison (rep0 only) + twin, events, obs_fn, obs, pu0, cache0 = rep0["_runtime"] + + def loss_units(theta_vals, units): + per = [] + for p in fit_protos: + rr = obs_fn(theta_vals, p) + per.append(np.asarray(jnp.mean((rr - obs[p]) ** 2, axis=0))) + return float(np.mean(per, axis=0)[units].mean()) + + wall = {} + estimators = {} + for model, units in obs_models.items(): + t0 = time.time() + g_grid = np.geomspace(g_lo, g_hi, args.grid_n) + w_grid = np.geomspace(w_lo, w_hi, args.grid_n) + L = np.array([[loss_units(np.array([g, w]), units) + for w in w_grid] for g in g_grid]) + bi, bj = np.unravel_index(np.argmin(L), L.shape) + + def nb(grid, k): + return np.geomspace(grid[max(k - 1, 0)], + grid[min(k + 1, len(grid) - 1)], args.refine_n) + + g_ref, w_ref = nb(g_grid, bi), nb(w_grid, bj) + Lr = np.array([[loss_units(np.array([g, w]), units) + for w in w_ref] for g in g_ref]) + ri, rj = np.unravel_index(np.argmin(Lr), Lr.shape) + wall[f"grid_{model}"] = round(time.time() - t0, 1) + estimators[(model, "grid_refine")] = { + "theta_hat": [float(g_ref[ri]), float(w_ref[rj])], + "fit_loss": float(Lr[ri, rj]), + "within_eps": bool(Lr[ri, rj] <= rep0[model]["eps"]), + "boundary_hit": bool(bi in (0, args.grid_n - 1) + or bj in (0, args.grid_n - 1)), + } + + # unpromoted gradient method; every iterate HARD-evaluated + t0 = time.time() + tgt = {p: jax.lax.stop_gradient(obs[p]) for p in fit_protos} + uidx = jnp.asarray(units) + + def loss_z(z): + th = twin.default_theta() + th["gCaL_scale"] = jnp.exp(z[0]) + th["w_exc"] = jnp.exp(z[1]) + tot = 0.0 + for p in fit_protos: + out = twin._run_events( + th, {"exc": events[p], + "inh": jnp.zeros((events[p].shape[0], 5), jnp.float32), + "exc_conn": twin._exc_mat, + "inh_conn": twin._inh_mat}, + spike_mode="surrogate") + if args.observable == "pointwise": + rr = _smoothed_rates_jax(out["spikes"], dt_ms, 100.0) + else: + rr = binned_rates_jnp(out["spikes"], dt_ms) + tot = tot + jnp.mean((rr[:, uidx] - tgt[p][:, uidx]) ** 2) + return tot / len(fit_protos) + + vg = jax.jit(jax.value_and_grad(loss_z)) + runs = [] + for label, start in (("start_default", (1.0, 0.009)), + ("start_perturbed", (0.6, 0.014))): + opt = optax.chain(optax.clip(1.0), optax.adam(0.05)) + z = jnp.log(jnp.asarray(start, dtype=jnp.float32)) + state = opt.init(z) + best = {"hard_loss": np.inf, "theta": None, "it": -1} + nonfinite = 0 + hist = [] + for it in range(args.budget_iters): + th_now = np.exp(np.asarray(z, dtype=np.float64)) + if not np.all(np.isfinite(th_now)): + nonfinite += 1 + hist.append({"it": it, "nonfinite": True}) + break + hard = loss_units(th_now, units) + hist.append({"it": it, "hard_loss": hard, + "theta": th_now.tolist()}) + if hard < best["hard_loss"]: + best = {"hard_loss": hard, "theta": th_now.tolist(), + "it": it} + _, g = vg(z) + upd, state = opt.update(g, state) + z = optax.apply_updates(z, upd) + runs.append({"label": label, "start": list(start), "best": best, + "nonfinite_iterates": nonfinite, + "failed": nonfinite > 0, "history": hist}) + wall[f"grad_{model}"] = round(time.time() - t0, 1) + ok_runs = [r for r in runs if r["best"]["theta"] is not None] + best_run = min(ok_runs, key=lambda r: r["best"]["hard_loss"]) + estimators[(model, "grad_best_of_starts")] = { + "theta_hat": best_run["best"]["theta"], + "fit_loss": best_run["best"]["hard_loss"], + "within_eps": bool(best_run["best"]["hard_loss"] + <= rep0[model]["eps"]), + "n_failed_starts": sum(r["failed"] for r in runs), + "n_starts": len(runs), + "runs": runs, + } + + # ----------------------------------------------- held-out, limb-specific + limbs = segment_limbs(commands["heldout_slowramp"]) + if args.observable == "binned": + limbs = {k: bin_mask(m, dt_ms) for k, m in limbs.items()} + + def heldout_limb_mse(theta_vals): + rr = np.asarray(obs_fn(np.asarray(theta_vals), "heldout_slowramp")) + oo = np.asarray(obs["heldout_slowramp"]) + se = (rr - oo) ** 2 + out = {name: float(se[mask].mean()) for name, mask in limbs.items() + if mask.any()} + out["overall"] = float(se.mean()) + return out + + recovery_table = [] + for (model, method), est in estimators.items(): + th = np.asarray(est["theta_hat"]) + recovery_table.append({ + "obs_model": model, "method": method, + "theta_hat": th.tolist(), + "param_rel_err": (np.abs(th - theta_true) / theta_true).tolist(), + "fit_loss": est["fit_loss"], + "fit_within_eps": est["within_eps"], + "heldout_mse_by_limb": heldout_limb_mse(th), + }) + heldout_truth = heldout_limb_mse(theta_true) + + # ---------------------------------------------------------------- figure + clip = EPS_FLOOR / 10 # display floor so the exact-zero truth plots on log axes + fig = plt.figure(figsize=(13, 9)) + gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.26) + for col, model in enumerate(("full_obs", "partial_obs")): + axL = fig.add_subplot(gs[0, col]) + L = np.array(rep0[model]["dense_loss"]) + pc = axL.pcolormesh(w_dense, g_dense, np.log10(L + clip), + cmap="Blues", shading="nearest") + fig.colorbar(pc, ax=axL, label="log10 fit loss (Hz²)") + axL.contour(w_dense, g_dense, L, levels=[rep0[model]["eps"]], + colors=["#111111"], linewidths=1.4) + axL.plot(theta_true[1], theta_true[0], "*", ms=16, + color=OKABE["truth"], label="truth (loss = 0, on-axis)") + gb = estimators[(model, "grid_refine")]["theta_hat"] + axL.plot(gb[1], gb[0], "o", ms=9, mfc="none", mew=2, + color=OKABE["grid"], label="grid+refine") + axL.set_xscale("log") + axL.set_yscale("log") + axL.set_xlabel("w_exc (µS)") + axL.set_ylabel("gCaL_scale") + axL.set_title(f"{model} ({args.observable}) — rep0; contour: " + f"EPS = {rep0[model]['eps']:.3f} Hz²", fontsize=9) + axL.legend(frameon=False, fontsize=8, loc="lower left") + + for col, (pname, axis, key) in enumerate( + (("gCaL_scale", g_dense, "profile_gCaL"), + ("w_exc", w_dense, "profile_w_exc"))): + ax = fig.add_subplot(gs[1, col]) + for model, c in (("full_obs", OKABE["grid"]), + ("partial_obs", OKABE["partial"])): + for k, rep in enumerate(replicate_reports): + cur = np.maximum(np.array(rep[model][key]), clip) + ax.plot(axis, cur, color=c, lw=1.2, + alpha=0.45 if k else 1.0, + label=model if k == 0 else None) + ax.axhline(rep0[model]["eps"], color=c, lw=0.9, ls="--", + label=f"EPS {model} {rep0[model]['eps']:.3f}") + ax.axvline(theta_true[0] if pname == "gCaL_scale" else theta_true[1], + color=MUTED, lw=1.0) + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xlabel(pname) + ax.set_ylabel(f"profile-min fit loss (Hz², display-clipped at {clip:g})") + ax.set_title(f"Profile over {pname} — truth on-axis, exact 0 at " + f"truth; {args.replicates} replicate(s)", fontsize=9) + ax.legend(frameon=False, fontsize=7) + ax.grid(True, color=GRIDC, lw=0.6) + ax.set_axisbelow(True) + fig.savefig(RESULTS / f"{stem}_landscape.png", dpi=150, + bbox_inches="tight") + + # ---------------------------------------------------------------- report + here = Path(__file__) + report = { + "manifest": { + **_manifest(), + "deliverable_b_sha256_16": hashlib.sha256( + here.read_bytes()).hexdigest()[:16], + "args": vars(args), + "run_tag": tag, + "protocol_lengths_steps": {k: len(v) for k, v in commands.items()}, + "optimizer": {"chain": "clip(1.0)+adam(0.05)", + "starts": [[1.0, 0.009], [0.6, 0.014]], + "iters": args.budget_iters}, + }, + "declared": { + "unknowns": ["gCaL_scale", "w_exc"], + "experiment_type": "conditional fit with known inputs", + "observable": args.observable, + "eps_rule": "per replicate × observation model, dt-halving at " + "truth over that model's units; floor %.0e" % EPS_FLOOR, + "eps_conservative_max": max( + r["eps_by_model"][m] for r in replicate_reports + for m in obs_models), + "theta_true": theta_true.tolist(), + "axes_contain_truth": True, + "interval_terminology": "sampled compatibility intervals (deterministic level sets, not statistical confidence intervals)", + }, + "g_dense": g_dense.tolist(), + "w_dense": w_dense.tolist(), + "replicates": [ + {k: v for k, v in r.items() if not k.startswith("_")} + for r in replicate_reports], + "estimators": {f"{m}/{meth}": est + for (m, meth), est in estimators.items()}, + "recovery_table": recovery_table, + "heldout_at_truth_by_limb": heldout_truth, + "wall_times_s": {**wall, "total": round(time.time() - t_wall0, 1)}, + } + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"{stem}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"{stem}_{stamp}.json").write_text(payload) + + print(f"\n=== {stem} — sampled compatibility intervals (profile ≤ per-model EPS) ===") + for k, rep in enumerate(replicate_reports): + for model in obs_models: + print(f" rep{k} {model:12s} (EPS {rep[model]['eps']:.4f}): " + f"gCaL {rep[model]['gcal_interval']} | " + f"w_exc {rep[model]['wexc_interval']}") + print("\n--- recovery table ---") + for row in recovery_table: + th = row["theta_hat"] + hl = row["heldout_mse_by_limb"] + print(f" {row['obs_model']:>11}/{row['method']:<19} " + f"θ̂=({th[0]:5.3f},{th[1]:8.5f}) fit {row['fit_loss']:8.4f} " + f"(within EPS: {row['fit_within_eps']}) | held-out overall " + f"{hl['overall']:.3f}, peak {hl.get('peak', float('nan')):.3f}") + print(f"held-out at truth: overall {heldout_truth['overall']:.2e}") + print(f"wall times: {report['wall_times_s']}") + print(f"saved → {RESULTS / (stem + '.json')} (+ evidence copy, figure " + f"{stem}_landscape.png)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/deliverable_c.py b/experiments/twin/deliverable_c.py new file mode 100644 index 00000000..416ae696 --- /dev/null +++ b/experiments/twin/deliverable_c.py @@ -0,0 +1,282 @@ +"""Deliverable C (v2) — unknown-input stress test. + +Revised roadmap §12C: *"Release a small number of neural-drive nuisance +coefficients... This reveals whether the apparent PIC estimate survived only +because drive was treated as known."* + +Setup identical to Deliverable B (frozen event realizations, observable via +``--observable``, per-replicate EPS from dt-halving at truth, axes contain +the truth) EXCEPT the input amplitude is unknown to the fitter. The nuisance +is a per-protocol EFFECTIVE EXCITATORY AMPLITUDE ``effective_w_exc_p`` in µS +(mechanistically the synaptic weight — w_exc and any dimensionless input +gain b_p enter only as their product, an EXACT model symmetry declared up +front). Because of that symmetry, w_exc itself is structurally +non-identifiable here BY CONSTRUCTION; this deliverable can only answer +whether **gCaL remains constrained** after releasing the effective input +amplitude: + +* known input : one amplitude SHARED across fit protocols → + L_known(g) = min_b mean_p loss_p(g, b) +* unknown input : amplitudes FREE and independent per protocol → + L_unknown(g) = mean_p min_{b_p} loss_p(g, b_p) + +Both profiles contain the truth on-axis and are asserted to be exactly 0 +there (same model, same events). Minima over amplitude use local log +refinement around the coarse argmin (hard-spike loss is discontinuous). + +Observation model: all units (full_obs) only — declared; partial-observation +stress belongs to a later rung of the observation ladder. Held-out +prediction is NOT evaluated here: under unknown input it would require a +declared amplitude rule for the unseen protocol, which is a separate +experiment. + +Usage:: + + python deliverable_c.py # full run + python deliverable_c.py --dense-n 7 --replicates 1 --dur-scale 0.5 # smoke +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import jax.numpy as jnp + +from sensitivity_checks import _manifest, _smoothed_rates_jax +from deliverable_b import ( + EPS_FLOOR, RESULTS, EVIDENCE, axis_with_truth, make_commands, + build_twin_and_events, make_observable_fn, numerical_epsilon_spikes, + binned_rates_jnp, run_tag, _sha16, +) + +OKABE = {"known": "#0072B2", "unknown": "#E69F00", "truth": "#000000"} +GRIDC = "#e3e3e3" + +NESTING_TOL = 1e-9 + + +def compute_nested_profiles(losses_at, g_axis, b_axis, n_protos: int, + refine_pts: int = 5): + """Known- and unknown-input gCaL profiles over ONE shared candidate set. + + ``losses_at(g, b)`` returns a tuple of per-protocol losses. For each g the + amplitude candidate set is the union of the coarse axis and the local + log-refinement points around BOTH the shared-amplitude argmin and every + per-protocol argmin — evaluated once and reused by both models, so the + nesting property + + unknown(g) = mean_p min_b L_p(g, b) ≤ min_b mean_p L_p(g, b) = known(g) + + holds by construction (independently optimized amplitudes include the + shared solution). It is ASSERTED per point: a violation means the two + models were minimized over different candidate sets (the v1 bug). + """ + known = np.zeros(len(g_axis)) + unknown = np.zeros(len(g_axis)) + for i, g in enumerate(g_axis): + rows = {float(b): losses_at(g, b) for b in b_axis} + + def local_candidates(j): + lo = b_axis[max(j - 1, 0)] + hi = b_axis[min(j + 1, len(b_axis) - 1)] + return np.geomspace(lo, hi, refine_pts)[1:-1] + + coarse = list(b_axis) + argmins = {int(np.argmin([float(np.mean(v)) for v in rows.values()]))} + for pi in range(n_protos): + argmins.add(int(np.argmin([v[pi] for v in rows.values()]))) + candidates = set(float(b) for b in coarse) + for j in argmins: + candidates.update(float(b) for b in local_candidates(j)) + vals = {b: (rows[b] if b in rows else losses_at(g, b)) + for b in candidates} + + known[i] = min(float(np.mean(v)) for v in vals.values()) + unknown[i] = float(np.mean( + [min(v[pi] for v in vals.values()) for pi in range(n_protos)])) + assert unknown[i] <= known[i] + NESTING_TOL, ( + f"nesting violated at g={g}: unknown {unknown[i]} > " + f"known {known[i]} — candidate sets differ (bug)") + return known, unknown + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-mn", type=int, default=8) + ap.add_argument("--true", type=float, nargs=2, default=[1.6, 0.0105]) + ap.add_argument("--dense-n", type=int, default=21) + ap.add_argument("--replicates", type=int, default=3) + ap.add_argument("--dur-scale", type=float, default=1.0) + ap.add_argument("--observable", choices=["binned", "pointwise"], + default="binned") + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + tag = run_tag(args) + stem = f"deliverable_c_{tag}_{args.observable}" + t_wall0 = time.time() + dt_ms = 0.1 + theta_true = np.array(args.true) + fit_protos = ["fit_ramp", "fit_trap"] + commands = make_commands(dt_ms, args.dur_scale) + + g_dense = axis_with_truth(0.25, 4.0, args.dense_n, theta_true[0]) + b_dense = axis_with_truth(0.0045, 0.018, args.dense_n, theta_true[1]) + gi_true = int(np.argmin(np.abs(g_dense - theta_true[0]))) + + reps = [] + for r in range(args.replicates): + conn_seed, ev_off = r, 10 * r + twin, events = build_twin_and_events( + args.n_mn, conn_seed, ev_off, commands, dt_ms) + obs_fn = make_observable_fn(twin, events, args.observable) + obs = {p: obs_fn(theta_true, p) for p in commands} + + cache: dict = {} + + def losses_at(g, b): + key = (float(g), float(b)) + if key not in cache: + cache[key] = tuple( + float(jnp.mean((obs_fn(np.array([g, b]), p) - obs[p]) ** 2)) + for p in fit_protos) + return cache[key] + + assert sum(losses_at(*theta_true)) == 0.0, "truth loss must be 0" + + t0 = time.time() + # dense per-protocol matrices + Lp = {p: np.zeros((len(g_dense), len(b_dense))) for p in fit_protos} + for i, g in enumerate(g_dense): + for j, b in enumerate(b_dense): + vals = losses_at(g, b) + for p, v in zip(fit_protos, vals): + Lp[p][i, j] = v + + # per-replicate EPS (full_obs = all units) + fine = numerical_epsilon_spikes(args.n_mn, conn_seed, events, + theta_true, fit_protos) + terms = [] + for p, sp_f in zip(fit_protos, fine): + if args.observable == "pointwise": + rf = np.asarray(_smoothed_rates_jax( + jnp.asarray(sp_f), 0.05, 100.0))[1::2] + else: + rf = np.asarray(binned_rates_jnp(jnp.asarray(sp_f), 0.05)) + rc = np.asarray(obs[p]) + m = min(len(rf), len(rc)) + terms.append(float(np.mean((rf[:m] - rc[:m]) ** 2))) + EPS = max(float(np.mean(terms)), EPS_FLOOR) + + known, unknown = compute_nested_profiles( + losses_at, g_dense, b_dense, n_protos=len(fit_protos)) + assert known[gi_true] == 0.0 and unknown[gi_true] == 0.0, ( + "profile at truth must be exactly 0") + + def interval(curve): + w = curve <= EPS + return ([float(g_dense[w].min()), float(g_dense[w].max())] + if w.any() else None) + + reps.append({ + "conn_seed": conn_seed, "event_offset": ev_off, + "eps_full_obs": EPS, + "profile_known": known.tolist(), + "profile_unknown": unknown.tolist(), + "gcal_interval_known": interval(known), + "gcal_interval_unknown": interval(unknown), + "interval_terminology": "sampled compatibility interval " + "(deterministic level set, not a " + "statistical confidence interval)", + "n_evals": len(cache), + "event_hashes": {p: _sha16(np.asarray(events[p])) + for p in fit_protos}, + }) + print(f"[rep{r}] {len(cache)} evals in {time.time()-t0:.0f}s, " + f"EPS {EPS:.4f}; truth=0 OK") + print(f" gCaL sampled compatibility interval, amplitude known: " + f"{reps[-1]['gcal_interval_known']}") + print(f" gCaL sampled compatibility interval, amplitude RELEASED: " + f"{reps[-1]['gcal_interval_unknown']}") + + # ---------------------------------------------------------------- figure + clip = EPS_FLOOR / 10 + fig, ax = plt.subplots(figsize=(8.5, 5)) + for k, rep in enumerate(reps): + ax.plot(g_dense, np.maximum(rep["profile_known"], clip), + color=OKABE["known"], lw=1.4, + alpha=1.0 if k == 0 else 0.45, + label="amplitude known (shared)" if k == 0 else None) + ax.plot(g_dense, np.maximum(rep["profile_unknown"], clip), + color=OKABE["unknown"], lw=1.4, + alpha=1.0 if k == 0 else 0.45, + label="amplitude FREE per protocol" if k == 0 else None) + ax.axhline(rep["eps_full_obs"], color="#111111", lw=0.8, ls="--", + alpha=1.0 if k == 0 else 0.45, + label=f"EPS per replicate" if k == 0 else None) + ax.axvline(theta_true[0], color=OKABE["truth"], lw=1.0) + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xlabel("gCaL_scale") + ax.set_ylabel(f"profile-min fit loss (Hz², display-clipped at {clip:g})") + ax.set_title("Deliverable C — does gCaL stay constrained when input " + f"amplitude is released? ({args.replicates} replicate(s), " + f"{args.observable})", fontsize=10) + ax.legend(frameon=False, fontsize=9) + ax.grid(True, color=GRIDC, lw=0.6) + ax.set_axisbelow(True) + fig.tight_layout() + fig_path = RESULTS / f"{stem}_gcal_profiles.png" + fig.savefig(fig_path, dpi=150) + + # ---------------------------------------------------------------- report + here = Path(__file__) + report = { + "manifest": {**_manifest(), + "deliverable_c_sha256_16": hashlib.sha256( + here.read_bytes()).hexdigest()[:16], + "args": vars(args), "run_tag": tag}, + "declared": { + "nuisance": "per-protocol effective_w_exc_p (µS); exact " + "multiplicative symmetry with w_exc ⇒ w_exc is " + "structurally non-identifiable here by construction; " + "the question is whether gCaL stays constrained", + "observable": args.observable, + "observation_model": "full_obs only (declared)", + "eps_rule": "per replicate, dt-halving at truth, all units", + "theta_true": theta_true.tolist(), + "axes_contain_truth": True, + "heldout": "not evaluated — would require a declared amplitude " + "rule for the unseen protocol (separate experiment)", + }, + "g_dense": g_dense.tolist(), + "b_dense_effective_w_exc": b_dense.tolist(), + "replicates": reps, + "wall_s": round(time.time() - t_wall0, 1), + } + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"{stem}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"{stem}_{stamp}.json").write_text(payload) + + print(f"\n=== {stem} summary ===") + for k, rep in enumerate(reps): + print(f"rep{k} (EPS {rep['eps_full_obs']:.4f}): known " + f"{rep['gcal_interval_known']} | unknown " + f"{rep['gcal_interval_unknown']}") + print(f"figure → {fig_path}") + print(f"report → {RESULTS / (stem + '.json')} (+ evidence copy)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/evidence/deliverable_b_full_binned_20261018_123116.json b/experiments/twin/evidence/deliverable_b_full_binned_20261018_123116.json new file mode 100644 index 00000000..1286da9c --- /dev/null +++ b/experiments/twin/evidence/deliverable_b_full_binned_20261018_123116.json @@ -0,0 +1,4265 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "993a2cfdebaeb499", + "poolbuild.py": "816d4ba42490e283", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "f7cf30f9880ce289", + "protocols.py": "404ea437741ea2ae" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float32 (JAX default; x64 not enabled)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations", + "deliverable_b_sha256_16": "a2be2e2dfb1c5eeb", + "args": { + "n_mn": 8, + "true": [ + 1.6, + 0.0105 + ], + "dense_n": 21, + "grid_n": 7, + "refine_n": 5, + "budget_iters": 15, + "replicates": 3, + "dur_scale": 1.0, + "observable": "binned" + }, + "run_tag": "full", + "protocol_lengths_steps": { + "fit_ramp": 40000, + "fit_trap": 40000, + "heldout_slowramp": 50000 + }, + "optimizer": { + "chain": "clip(1.0)+adam(0.05)", + "starts": [ + [ + 1.0, + 0.009 + ], + [ + 0.6, + 0.014 + ] + ], + "iters": 15 + } + }, + "declared": { + "unknowns": [ + "gCaL_scale", + "w_exc" + ], + "experiment_type": "conditional fit with known inputs", + "observable": "binned", + "eps_rule": "per replicate \u00d7 observation model, dt-halving at truth over that model's units; floor 1e-03", + "eps_conservative_max": 0.8333334922790527, + "theta_true": [ + 1.6, + 0.0105 + ], + "axes_contain_truth": true + }, + "g_dense": [ + 0.25, + 0.2871745887492587, + 0.32987697769322355, + 0.37892914162759955, + 0.43527528164806206, + 0.5, + 0.5743491774985175, + 0.6597539553864472, + 0.7578582832551991, + 0.8705505632961241, + 1.0, + 1.1486983549970349, + 1.3195079107728946, + 1.5157165665103984, + 1.6, + 1.7411011265922485, + 2.0000000000000004, + 2.29739670999407, + 2.6390158215457893, + 3.0314331330207964, + 3.4822022531844974, + 4.0 + ], + "w_dense": [ + 0.0045, + 0.004822980581413315, + 0.005169142597486656, + 0.00554014986005212, + 0.005937785598478024, + 0.006363961030678926, + 0.006820724549296786, + 0.007310271567206119, + 0.007834955069665113, + 0.008397296923831269, + 0.008999999999999998, + 0.009645961162826632, + 0.010338285194973314, + 0.0105, + 0.011080299720104242, + 0.011875571196956045, + 0.012727922061357854, + 0.013641449098593581, + 0.01462054313441224, + 0.01566991013933023, + 0.01679459384766253, + 0.018 + ], + "replicates": [ + { + "conn_seed": 0, + "event_offset": 0, + "active_units_at_truth": { + "fit_ramp": 7, + "fit_trap": 7, + "heldout_slowramp": 6 + }, + "truth_fit_loss": 0.0, + "eps_by_model": { + "full_obs": 0.6770833730697632, + "partial_obs": 0.8333334922790527 + }, + "dense_time_s": 209.3, + "event_hashes": { + "fit_ramp": "5ef0d36f3c87eee4", + "fit_trap": "682ace78eb9e26c4", + "heldout_slowramp": "f112ce9602b70039" + }, + "full_obs": { + "observed_units": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ], + "eps": 0.6770833730697632, + "profile_gCaL": [ + 1.7534725666046143, + 1.8229169845581055, + 1.7361115217208862, + 1.8055559396743774, + 1.822917103767395, + 1.7534725666046143, + 1.510416865348816, + 1.406250238418579, + 1.4930557012557983, + 1.4583336114883423, + 1.3541669845581055, + 1.250000238418579, + 1.0069446563720703, + 0.8854168653488159, + 0.0, + 0.6076390147209167, + 0.8680556416511536, + 1.406250238418579, + 2.152778148651123, + 2.8472225666046143, + 2.899306058883667, + 3.333333969116211 + ], + "profile_w_exc": [ + 24.51388931274414, + 18.42013931274414, + 12.725693702697754, + 8.14236068725586, + 5.2604169845581055, + 3.680556058883667, + 3.055556058883667, + 2.899306058883667, + 3.090278148651123, + 2.343750476837158, + 1.8923614025115967, + 1.406250238418579, + 0.6076390147209167, + 0.0, + 1.250000238418579, + 3.0381951332092285, + 7.812501907348633, + 17.656253814697266, + 30.3125057220459, + 48.00347900390625, + 70.88543701171875, + 100.69445037841797 + ], + "gcal_interval": [ + 1.6, + 1.7411011265922485 + ], + "wexc_interval": [ + 0.010338285194973314, + 0.0105 + ], + "dense_loss": [ + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 56.388893842697144, + 49.218754529953, + 40.13889288902283, + 32.04861235618591, + 26.111112356185913, + 21.649307012557983, + 14.357639193534851, + 7.690972328186035, + 5.972222775220871, + 2.3437504023313522, + 3.090278461575508, + 7.812501907348633, + 18.40278148651123, + 31.05903434753418, + 48.85417398810387, + 71.40626096725464, + 102.18750762939453 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 55.95486521720886, + 47.586809396743774, + 39.479170083999634, + 31.80555748939514, + 25.885418176651, + 20.277778148651123, + 13.437499523162842, + 7.48263943195343, + 5.625000685453415, + 2.2916669994592667, + 3.0381950736045837, + 8.055557787418365, + 17.829864263534546, + 30.312505841255188, + 49.20139679312706, + 70.88542938232422, + 102.95139598846436 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 55.08681082725525, + 46.75347542762756, + 39.53125357627869, + 31.180556535720825, + 25.399307012557983, + 20.34722328186035, + 13.385416388511658, + 6.8750006556510925, + 5.277778133749962, + 2.1875003427267075, + 3.4895839616656303, + 8.593751728534698, + 19.166669964790344, + 31.493062496185303, + 48.003479331731796, + 72.55209684371948, + 103.71528625488281 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 54.65278220176697, + 45.53819727897644, + 39.23611378669739, + 31.111111879348755, + 25.19097375869751, + 19.722223043441772, + 13.229166507720947, + 6.562500596046448, + 5.17361156642437, + 2.3090281635522842, + 3.6111117601394653, + 8.663196086883545, + 19.496531128883362, + 31.822922706604004, + 50.17362068593502, + 73.83681774139404, + 104.79168224334717 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 53.81944918632507, + 45.086808919906616, + 38.10764193534851, + 30.538195848464966, + 24.600695371627808, + 19.56597363948822, + 12.656249761581421, + 6.006944894790649, + 4.670139417052269, + 2.152778133749962, + 3.6458339393138885, + 9.218751907348633, + 18.541669845581055, + 32.70834040641785, + 50.03472922742367, + 73.6979308128357, + 104.09723567962646 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.14236521720886, + 53.36806082725525, + 44.73958659172058, + 37.46528077125549, + 29.913195848464966, + 24.01041865348816, + 18.993056654930115, + 12.638888716697693, + 5.8506951332092285, + 4.670139417052269, + 2.0312503576278687, + 3.4895840287208557, + 8.906251788139343, + 19.67014193534851, + 33.489590883255005, + 52.22222922742367, + 74.32292938232422, + 106.94445610046387 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.14930987358093, + 57.30903220176697, + 52.53472685813904, + 43.02083706855774, + 36.04166913032532, + 29.87847352027893, + 23.368056535720825, + 18.64583444595337, + 12.430554986000061, + 5.4513896107673645, + 4.062500461935997, + 2.0138892233371735, + 3.7152785658836365, + 9.461807012557983, + 20.46875309944153, + 34.149314165115356, + 50.64236734807491, + 74.98265314102173, + 106.78820133209229 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 59.14930987358093, + 57.32639288902283, + 51.28472685813904, + 41.56250309944153, + 35.0868079662323, + 29.635417938232422, + 23.17708420753479, + 18.125001072883606, + 11.059028267860413, + 4.930555999279022, + 3.4722225815057755, + 1.6493057236075401, + 3.9236119240522385, + 10.069446325302124, + 20.972225069999695, + 34.58333969116211, + 51.73611830174923, + 75.46876335144043, + 108.50695562362671 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 56.388893842697144, + 49.4270875453949, + 40.46875357627869, + 33.48958468437195, + 27.760417938232422, + 22.586806654930115, + 17.916667461395264, + 10.22569453716278, + 4.166667133569717, + 3.4722225815057755, + 1.4930556640028954, + 4.201389729976654, + 10.920141518115997, + 21.753475666046143, + 35.48611664772034, + 54.21876047551632, + 76.82293224334717, + 109.7222318649292 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 55.10417056083679, + 47.04861378669739, + 40.08680844306946, + 33.10764002799988, + 26.406251430511475, + 22.100695729255676, + 15.746528148651123, + 9.531250417232513, + 3.5937504172325134, + 3.2812505811452866, + 1.4583336263895035, + 4.444445312023163, + 11.579863727092743, + 22.517364501953125, + 36.041672706604004, + 55.086813911795616, + 77.9687647819519, + 112.60417652130127 + ], + [ + 59.63542056083679, + 59.63542056083679, + 59.14930987358093, + 57.72569918632507, + 53.36806082725525, + 44.54861521720886, + 38.42014145851135, + 30.88541865348816, + 25.850695848464966, + 21.111112117767334, + 14.843749523162842, + 8.090278089046478, + 3.211806058883667, + 2.7604170590639114, + 1.423611469566822, + 5.451389849185944, + 12.326391279697418, + 23.975697875022888, + 37.152785301208496, + 56.4236209243536, + 79.72223711013794, + 114.91320610046387 + ], + [ + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 57.30903220176697, + 50.72917056083679, + 42.586809396743774, + 35.62500214576721, + 30.225695848464966, + 24.600695848464966, + 19.65277862548828, + 12.916666746139526, + 6.614583671092987, + 2.9513893872499466, + 2.2222225219011307, + 1.2500003054738045, + 6.041667461395264, + 13.645835876464844, + 25.10417079925537, + 39.357646226882935, + 58.64584110677242, + 82.96876382827759, + 117.03126049041748 + ], + [ + 59.63542056083679, + 59.63542056083679, + 58.68055987358093, + 55.10417056083679, + 47.395835638046265, + 40.05208611488342, + 33.36805701255798, + 28.15972375869751, + 23.31597352027893, + 18.24652874469757, + 10.815972447395325, + 5.4687506556510925, + 2.2916670590639114, + 1.840278074145317, + 1.614583671092987, + 7.031250953674316, + 14.930558443069458, + 27.70833945274353, + 41.25000596046448, + 62.96875999867916, + 86.4757080078125, + 120.41668033599854 + ], + [ + 59.63542056083679, + 59.14930987358093, + 57.30903220176697, + 52.51736521720886, + 43.66319823265076, + 37.482640981674194, + 31.406251192092896, + 25.885417938232422, + 20.98958444595337, + 15.677082657814026, + 9.010416746139526, + 4.166666954755783, + 1.9444448202848434, + 0.9201390519738197, + 2.5694448351860046, + 8.35069590806961, + 17.829864144325256, + 29.496533155441284, + 44.67014503479004, + 64.54862101376057, + 89.89584684371948, + 125.43404483795166 + ], + [ + 59.63542056083679, + 58.68055987358093, + 57.27430987358093, + 50.20833659172058, + 42.91667103767395, + 36.006946325302124, + 30.659723043441772, + 25.36458444595337, + 19.913195729255676, + 14.114583253860474, + 7.951388955116272, + 3.8020837604999542, + 1.4236113801598549, + 0.0, + 2.708333671092987, + 8.854168057441711, + 18.090280890464783, + 31.128478527069092, + 45.85070204734802, + 66.70139820873737, + 92.11806774139404, + 127.93404293060303 + ], + [ + 59.63542056083679, + 58.14236521720886, + 54.65278220176697, + 47.79514145851135, + 39.68750357627869, + 34.42708468437195, + 28.350695848464966, + 23.906251668930054, + 17.934027314186096, + 11.840277552604675, + 6.406250357627869, + 3.020833820104599, + 0.6076390109956264, + 1.3194447308778763, + 3.489583969116211, + 10.0347238779068, + 19.583337903022766, + 31.85764503479004, + 49.60070180892944, + 70.4340365678072, + 96.75348377227783, + 131.66668319702148 + ], + [ + 58.68055987358093, + 56.35417056083679, + 49.80903220176697, + 42.93403124809265, + 36.47569727897644, + 30.850695848464966, + 25.885418176651, + 20.381945371627808, + 14.670138716697693, + 8.524305582046509, + 4.583333730697632, + 1.8055557310581207, + 1.788194790482521, + 2.222222626209259, + 6.02430659532547, + 12.552085161209106, + 23.315977096557617, + 36.875006437301636, + 54.7222318649292, + 77.08334572613239, + 102.58681774139404, + 140.74654912948608 + ], + [ + 57.29167151451111, + 51.128474950790405, + 43.85417056083679, + 38.19444727897644, + 32.256946325302124, + 27.256946086883545, + 22.413195610046387, + 15.815972208976746, + 9.982638955116272, + 5.4513896107673645, + 2.9861115515232086, + 1.4062501713633537, + 3.472222864627838, + 4.53125074505806, + 9.114584565162659, + 17.552086353302002, + 28.697921752929688, + 42.430562257766724, + 61.68403673171997, + 87.23959524929523, + 113.43751049041748, + 151.4062705039978 + ], + [ + 51.128474950790405, + 43.50694823265076, + 38.19444727897644, + 32.500001192092896, + 27.430557012557983, + 22.881946086883545, + 16.02430534362793, + 9.99999988079071, + 5.7986114621162415, + 3.4027782678604126, + 2.1527781784534454, + 3.611111491918564, + 7.690973281860352, + 9.201390266418457, + 14.444446682929993, + 25.364588499069214, + 36.545145750045776, + 52.32639694213867, + 73.87153339385986, + 100.86806757748127, + 130.69446170330048, + 169.79168701171875 + ], + [ + 41.736114740371704, + 36.35416769981384, + 31.336806535720825, + 26.562501907348633, + 20.590277910232544, + 15.121527194976807, + 8.99305546283722, + 5.312500596046448, + 3.090278208255768, + 3.29861156642437, + 5.434027999639511, + 9.77430671453476, + 15.746529936790466, + 17.37847512960434, + 25.451393008232117, + 36.163201570510864, + 51.87500858306885, + 69.58334255218506, + 94.65279087424278, + 120.86806744337082, + 152.51738440990448, + 197.32641792297363 + ], + [ + 33.05555701255798, + 28.29861330986023, + 23.94097375869751, + 17.395832657814026, + 11.510416269302368, + 7.118055582046509, + 4.444445252418518, + 2.8993060141801834, + 4.982639193534851, + 8.80208432674408, + 14.444445788860321, + 22.5347261428833, + 32.36111605167389, + 35.71181094646454, + 45.607645988464355, + 62.01389837265015, + 79.84376001358032, + 101.63195419311523, + 128.1597417294979, + 155.8507126569748, + 190.17363595962524, + 242.37848663330078 + ], + [ + 24.513890981674194, + 18.420138597488403, + 12.725693821907043, + 8.14236068725586, + 5.260417133569717, + 3.6805561035871506, + 4.166667342185974, + 9.06250137090683, + 13.732640087604523, + 21.041670203208923, + 30.43403136730194, + 40.555561542510986, + 53.59375786781311, + 56.822927713394165, + 68.21181535720825, + 86.45834302902222, + 110.83335256576538, + 134.20140647888184, + 158.9062655866146, + 193.92363691329956, + 232.6215524673462, + 284.2882146835327 + ] + ] + }, + "partial_obs": { + "observed_units": [ + 0, + 1, + 2, + 3 + ], + "eps": 0.8333334922790527, + "profile_gCaL": [ + 1.9097228050231934, + 2.1527786254882812, + 2.083333969116211, + 2.048611640930176, + 2.118056297302246, + 1.9791672229766846, + 2.083333730697632, + 1.5625005960464478, + 1.7013896703720093, + 1.9791672229766846, + 1.7361116409301758, + 1.3888894319534302, + 1.3541669845581055, + 1.1805559396743774, + 0.0, + 0.8333336114883423, + 1.041666865348816, + 1.4583337306976318, + 1.979166865348816, + 2.395833969116211, + 2.6388893127441406, + 3.055556535720825 + ], + "profile_w_exc": [ + 27.118059158325195, + 18.923612594604492, + 14.027776718139648, + 9.44444465637207, + 6.2152791023254395, + 3.7500009536743164, + 2.5694451332092285, + 2.708333969116211, + 2.6388893127441406, + 2.430556297302246, + 2.152777910232544, + 1.4583337306976318, + 0.8333336114883423, + 0.0, + 1.3541672229766846, + 2.013889789581299, + 4.131946563720703, + 9.166669845581055, + 16.31945037841797, + 28.15972900390625, + 47.43056106567383, + 64.54861450195312 + ], + "gcal_interval": [ + 1.6, + 1.6 + ], + "wexc_interval": [ + 0.0105, + 0.0105 + ], + "dense_loss": [ + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 90.13889694213867, + 75.79861831665039, + 57.63889503479004, + 41.45833396911621, + 29.58333396911621, + 20.65972328186035, + 11.597222566604614, + 5.763889789581299, + 5.694445371627808, + 2.777778387069702, + 2.222223073244095, + 4.131946563720703, + 9.201392889022827, + 16.805562496185303, + 29.340283393859863, + 47.43056297302246, + 65.27778625488281 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 89.27083969116211, + 72.53472805023193, + 56.31944942474365, + 40.97222423553467, + 29.131945610046387, + 19.479166507720947, + 10.868055820465088, + 5.868056535720825, + 5.034723281860352, + 2.673611581325531, + 2.1527786254882812, + 4.409724712371826, + 9.375003814697266, + 17.013896226882935, + 28.819451332092285, + 48.22917652130127, + 65.83334159851074 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 87.53473091125488, + 70.86806011199951, + 56.42361640930176, + 39.722222328186035, + 28.15972328186035, + 19.722222805023193, + 10.76388955116272, + 5.59027898311615, + 4.791667699813843, + 2.500000476837158, + 2.222223088145256, + 4.791668772697449, + 9.687503576278687, + 17.95139503479004, + 28.15972900390625, + 47.77778625488281, + 67.77778625488281 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 86.66667366027832, + 68.43750381469727, + 55.83333683013916, + 39.583333015441895, + 27.743056774139404, + 18.472222328186035, + 10.451389789581299, + 5.659723281860352, + 4.58333432674408, + 2.708333909511566, + 2.3263898491859436, + 4.965279817581177, + 9.96528172492981, + 17.430562019348145, + 29.131951332092285, + 49.513898849487305, + 68.33334541320801 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 85.00000762939453, + 67.53472709655762, + 53.576393127441406, + 38.437500953674316, + 26.5625, + 18.15972352027893, + 9.65277910232544, + 5.138889908790588, + 4.375000953674316, + 2.6041672229766846, + 2.326389729976654, + 5.277779936790466, + 9.375003337860107, + 18.40278387069702, + 29.930561542510986, + 49.44445514678955, + 67.50000953674316 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 93.64583969116211, + 84.09723091125488, + 66.84028244018555, + 52.29167079925537, + 37.187500953674316, + 25.381946563720703, + 17.01388955116272, + 9.618057012557983, + 5.1736122369766235, + 4.375000953674316, + 2.500000536441803, + 2.1875008940696716, + 5.416668891906738, + 10.000003576278687, + 19.16667413711548, + 31.284728527069092, + 49.270843505859375, + 68.43750762939453 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 95.65972900390625, + 91.97917366027832, + 82.43056297302246, + 63.40278339385986, + 49.44444751739502, + 37.118056297302246, + 24.097222328186035, + 16.31944513320923, + 9.20138955116272, + 5.0000011920928955, + 4.375001072883606, + 2.500000536441803, + 2.3958343863487244, + 5.555557727813721, + 11.458337306976318, + 19.826395988464355, + 31.250006675720215, + 50.451398849487305, + 69.79167366027832 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 95.65972900390625, + 92.01389503479004, + 79.93056297302246, + 60.48611545562744, + 47.534725189208984, + 36.63194513320923, + 23.715277671813965, + 15.277778387069702, + 8.194445371627808, + 4.340278625488281, + 3.576389729976654, + 2.2569448351860046, + 2.534723311662674, + 5.833335638046265, + 12.39583706855774, + 20.38195037841797, + 32.63889408111572, + 50.34723377227783, + 70.17362117767334 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 90.13889694213867, + 76.21528434753418, + 58.29861640930176, + 44.34027862548828, + 32.88194513320923, + 22.534722566604614, + 14.861111164093018, + 7.5000011920928955, + 4.27083432674408, + 3.61111181974411, + 2.118055820465088, + 2.77777898311615, + 6.319446921348572, + 12.951392889022827, + 21.006949424743652, + 34.75695323944092, + 51.80556583404541, + 72.5694522857666 + ], + [ + 96.63195037841797, + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 87.56945037841797, + 71.45833683013916, + 57.5347261428833, + 43.57638931274414, + 30.173612117767334, + 21.562500715255737, + 13.125000476837158, + 7.187501311302185, + 3.6805564761161804, + 3.5416674613952637, + 1.9791672229766846, + 2.951390266418457, + 6.527780652046204, + 13.333337783813477, + 22.25694990158081, + 34.826395988464355, + 53.71528720855713, + 75.10417556762695 + ], + [ + 96.63195037841797, + 96.63195037841797, + 95.65972900390625, + 92.81250762939453, + 84.09723091125488, + 66.45833969116211, + 54.20139217376709, + 39.1319465637207, + 29.062500953674316, + 20.312500476837158, + 11.875, + 6.388889908790588, + 3.3333341479301453, + 3.159722864627838, + 2.013889580965042, + 3.645834982395172, + 7.326391816139221, + 14.479171514511108, + 22.8125057220459, + 36.42361879348755, + 55.381954193115234, + 77.88195419311523 + ], + [ + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 91.97917366027832, + 78.81945037841797, + 62.534728050231934, + 48.61111354827881, + 37.812500953674316, + 26.562500953674316, + 18.333333492279053, + 10.798611640930176, + 5.868056654930115, + 3.29861181974411, + 2.569444954395294, + 1.4583338648080826, + 3.85416841506958, + 8.43750286102295, + 15.555561065673828, + 24.375007152557373, + 39.23611545562744, + 58.05556392669678, + 82.56945610046387 + ], + [ + 96.63195037841797, + 96.63195037841797, + 94.72222900390625, + 87.56945037841797, + 72.15278053283691, + 57.46528148651123, + 44.09722328186035, + 33.680556774139404, + 23.993056297302246, + 15.520833730697632, + 8.923611879348755, + 4.89583432674408, + 2.7430562376976013, + 2.32638943195343, + 1.8402784168720245, + 4.340279579162598, + 9.097225666046143, + 16.94444990158081, + 26.145838260650635, + 44.652785301208496, + 62.56945610046387, + 82.91667938232422 + ], + [ + 96.63195037841797, + 95.65972900390625, + 91.97917366027832, + 82.39583969116211, + 64.6875057220459, + 52.32639122009277, + 40.173611640930176, + 29.13194513320923, + 20.06944513320923, + 13.541666269302368, + 7.15277886390686, + 3.993056356906891, + 2.430556356906891, + 1.180555909872055, + 2.6388896107673645, + 5.902780175209045, + 11.45833706855774, + 18.923617839813232, + 30.3125057220459, + 46.354177474975586, + 65.83334445953369, + 89.06251335144043 + ], + [ + 96.63195037841797, + 94.72222900390625, + 91.90972900390625, + 77.77778244018555, + 63.194451332092285, + 49.37500190734863, + 38.68055534362793, + 28.090278148651123, + 18.854167699813843, + 12.11805534362793, + 6.423611879348755, + 3.680556356906891, + 2.0138894617557526, + 0.0, + 2.812501013278961, + 5.937502384185791, + 11.909725904464722, + 20.347229480743408, + 30.902785778045654, + 48.99306774139404, + 65.06945419311523, + 93.22918128967285 + ], + [ + 96.63195037841797, + 93.64583969116211, + 86.66667366027832, + 72.95139217376709, + 56.73611640930176, + 46.21527862548828, + 34.062500953674316, + 25.173612594604492, + 16.041667222976685, + 10.416666746139526, + 5.763889789581299, + 3.29861181974411, + 0.8333335816860199, + 1.8055561780929565, + 3.194445788860321, + 7.18750262260437, + 13.159726858139038, + 22.256949424743652, + 34.86111927032471, + 51.458340644836426, + 71.04167175292969, + 95.24307060241699 + ], + [ + 94.72222900390625, + 90.06945037841797, + 76.97917366027832, + 63.22917175292969, + 50.312503814697266, + 39.062500953674316, + 29.131945610046387, + 19.791666984558105, + 13.715277910232544, + 7.604167461395264, + 4.409723162651062, + 2.2222225666046143, + 2.152778387069702, + 2.5347228050231934, + 5.208335280418396, + 8.993058443069458, + 15.902783393859863, + 25.590284824371338, + 38.57639694213867, + 57.60417652130127, + 76.80556678771973, + 105.93751907348633 + ], + [ + 91.9444522857666, + 79.6180591583252, + 65.06945037841797, + 53.750003814697266, + 41.87500190734863, + 31.875001430511475, + 22.916667461395264, + 15.451388597488403, + 9.583333730697632, + 5.381945729255676, + 3.1597229838371277, + 1.4583336263895035, + 3.472223460674286, + 4.097223341464996, + 6.4583353996276855, + 12.430559158325195, + 19.72222900390625, + 29.548616886138916, + 45.41667366027832, + 64.47917556762695, + 84.96528816223145, + 110.03473854064941 + ], + [ + 79.6180591583252, + 64.3750057220459, + 53.750003814697266, + 42.361111640930176, + 32.22222328186035, + 23.85416841506958, + 16.388888835906982, + 10.416666746139526, + 6.041667342185974, + 3.5069451332092285, + 2.1527784764766693, + 2.7777783274650574, + 5.347223997116089, + 6.388890981674194, + 9.75694727897644, + 16.80556058883667, + 26.770840167999268, + 38.12500762939453, + 53.819454193115234, + 73.57639694213867, + 97.39584732055664, + 124.02779006958008 + ], + [ + 60.83333873748779, + 50.06944465637207, + 40.034722328186035, + 30.48611307144165, + 21.111112117767334, + 15.624999523162842, + 9.895833492279053, + 6.006945490837097, + 3.3333340883255005, + 2.430556207895279, + 3.4027783274650574, + 5.277779221534729, + 9.027780294418335, + 9.93055808544159, + 15.20833706855774, + 22.430561542510986, + 34.27084159851074, + 47.2569522857666, + 65.24306774139404, + 85.83334350585938, + 109.13196563720703, + 140.10418319702148 + ], + [ + 43.47222328186035, + 33.958335876464844, + 25.97222375869751, + 18.36805558204651, + 12.708333015441895, + 8.298611164093018, + 5.069445848464966, + 2.7083340287208557, + 2.6388893723487854, + 3.4722225666046143, + 5.833334803581238, + 10.208335876464844, + 15.312503576278687, + 17.083337545394897, + 22.500006675720215, + 31.388894081115723, + 44.93056392669678, + 59.44445610046387, + 76.25000762939453, + 99.34029388427734, + 125.24307250976562, + 160.97223281860352 + ], + [ + 27.11805820465088, + 18.923611164093018, + 14.027776956558228, + 9.44444465637207, + 6.215278744697571, + 3.7500008940696716, + 3.2638896703720093, + 3.368056297302246, + 4.756945013999939, + 7.395834684371948, + 11.597224950790405, + 15.729170799255371, + 23.88889455795288, + 25.798616886138916, + 31.736117362976074, + 43.263895988464355, + 57.67362308502197, + 74.34029006958008, + 94.40973472595215, + 120.38196182250977, + 145.62501907348633, + 183.22918319702148 + ] + ] + } + }, + { + "conn_seed": 1, + "event_offset": 10, + "active_units_at_truth": { + "fit_ramp": 6, + "fit_trap": 6, + "heldout_slowramp": 6 + }, + "truth_fit_loss": 0.0, + "eps_by_model": { + "full_obs": 0.2951389253139496, + "partial_obs": 0.41666674613952637 + }, + "dense_time_s": 213.1, + "event_hashes": { + "fit_ramp": "ed698d364473ebbd", + "fit_trap": "e791fdc5716d20b5", + "heldout_slowramp": "9622e06f175690fc" + }, + "full_obs": { + "observed_units": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ], + "eps": 0.2951389253139496, + "profile_gCaL": [ + 1.2152780294418335, + 1.093750238418579, + 1.3194447755813599, + 1.3541669845581055, + 1.1284725666046143, + 1.0590280294418335, + 1.1111114025115967, + 1.076388955116272, + 0.8854168653488159, + 0.8680557608604431, + 0.8680558204650879, + 0.7986112236976624, + 0.46875008940696716, + 0.520833432674408, + 0.0, + 0.36458349227905273, + 0.5555557608604431, + 0.9722224473953247, + 1.4062503576278687, + 1.8750008344650269, + 2.4131953716278076, + 3.3159737586975098 + ], + "profile_w_exc": [ + 27.118059158325195, + 21.006946563720703, + 14.895835876464844, + 9.652778625488281, + 6.111112594604492, + 4.236112594604492, + 3.0555567741394043, + 2.586806535720825, + 2.4131953716278076, + 1.8750008344650269, + 1.4062504768371582, + 0.9722224473953247, + 0.24305559694766998, + 0.0, + 0.7812502384185791, + 1.3715283870697021, + 4.774306297302246, + 11.892363548278809, + 22.829864501953125, + 40.868064880371094, + 62.62153625488281, + 94.28820037841797 + ], + "gcal_interval": [ + 1.6, + 1.6 + ], + "wexc_interval": [ + 0.010338285194973314, + 0.0105 + ], + "dense_loss": [ + [ + 86.21527862548828, + 86.21527862548828, + 85.55555725097656, + 81.66666984558105, + 74.40972995758057, + 63.402785301208496, + 53.17708969116211, + 44.94792127609253, + 37.118059158325195, + 28.767364501953125, + 20.815975666046143, + 11.927084803581238, + 6.510417461395264, + 4.722223162651062, + 2.2743059396743774, + 1.6319449245929718, + 4.774306654930115, + 11.927086234092712, + 22.829864501953125, + 41.3020895421505, + 64.14931190013885, + 95.1388955116272 + ], + [ + 86.21527862548828, + 86.21527862548828, + 85.57291793823242, + 81.09375381469727, + 74.40972995758057, + 62.20486927032471, + 52.15278339385986, + 44.96528244018555, + 36.37153148651123, + 27.534724950790405, + 20.22569751739502, + 11.875001192092896, + 6.319445431232452, + 4.739584147930145, + 2.2395838797092438, + 1.5798616409301758, + 4.826390087604523, + 12.083335876464844, + 23.333335161209106, + 41.597226947546005, + 62.62153470516205, + 96.37153434753418 + ], + [ + 86.21527862548828, + 86.21527862548828, + 84.9826431274414, + 81.09375381469727, + 74.40972995758057, + 61.38889694213867, + 51.80556106567383, + 44.75694990158081, + 35.72917079925537, + 27.204863786697388, + 19.652780771255493, + 11.927084684371948, + 6.163195252418518, + 4.531250774860382, + 2.13541716337204, + 1.4583338052034378, + 5.052084237337112, + 12.135418772697449, + 23.559030652046204, + 41.7708395421505, + 64.913201212883, + 97.34375762939453 + ], + [ + 86.21527862548828, + 86.21527862548828, + 84.9826431274414, + 79.9305591583252, + 72.88195133209229, + 60.902785301208496, + 51.44097661972046, + 44.479172706604004, + 35.451393604278564, + 26.753474473953247, + 19.565975189208984, + 11.99652886390686, + 5.7118067145347595, + 4.3750006556510925, + 2.1527781784534454, + 1.5277782678604126, + 5.1909730434417725, + 12.847225069999695, + 23.784725189208984, + 41.857645362615585, + 66.3368114233017, + 97.25695037841797 + ], + [ + 86.21527862548828, + 86.21527862548828, + 84.46180725097656, + 78.97570037841797, + 71.85764503479004, + 59.739590644836426, + 50.53820085525513, + 42.67361640930176, + 35.0173659324646, + 25.607641458511353, + 19.010419726371765, + 11.909723401069641, + 5.312500953674316, + 4.253472983837128, + 1.9791670739650726, + 1.5625004917383194, + 5.243056386709213, + 13.072919249534607, + 24.236113786697388, + 42.51736721396446, + 66.44097793102264, + 98.7326455116272 + ], + [ + 86.21527862548828, + 86.21527862548828, + 83.36805725097656, + 78.42014694213867, + 69.84375858306885, + 58.42014503479004, + 49.80903339385986, + 42.395838260650635, + 34.46181011199951, + 25.31250309944153, + 17.829863786697388, + 11.475695848464966, + 4.687500894069672, + 4.288195431232452, + 1.857639193534851, + 1.6666672676801682, + 5.5729177594184875, + 13.420141816139221, + 24.84375262260437, + 43.1770895421505, + 66.12847793102264, + 99.09722852706909 + ], + [ + 86.21527862548828, + 85.5381965637207, + 83.36805725097656, + 77.4652853012085, + 68.03820419311523, + 56.61458969116211, + 48.64583778381348, + 41.25000476837158, + 33.99306011199951, + 24.809030771255493, + 17.48264193534851, + 10.868056535720825, + 4.357639670372009, + 4.114584028720856, + 1.6666669845581055, + 1.857639655470848, + 5.815973281860352, + 14.045141577720642, + 26.076391458511353, + 45.43403372168541, + 68.97570216655731, + 98.76736974716187 + ], + [ + 86.21527862548828, + 85.5381965637207, + 83.36805534362793, + 75.81597995758057, + 66.52778720855713, + 56.14583969116211, + 48.21181058883667, + 40.20833873748779, + 32.27430963516235, + 24.30555820465088, + 16.979169607162476, + 9.982639849185944, + 4.322917282581329, + 3.854167342185974, + 1.440972462296486, + 2.0486119091510773, + 5.920139968395233, + 14.322920083999634, + 25.833336353302002, + 45.05209073424339, + 69.40972924232483, + 102.41320371627808 + ], + [ + 86.21527862548828, + 85.55555725097656, + 82.20486354827881, + 74.82639694213867, + 64.28820037841797, + 54.687506675720215, + 47.29167127609253, + 39.30556106567383, + 31.26736545562744, + 23.19444727897644, + 16.18055808544159, + 9.583334386348724, + 3.958334267139435, + 3.298611640930176, + 1.04166679084301, + 2.309028595685959, + 6.180556952953339, + 14.89583683013916, + 28.02083659172058, + 46.04167351126671, + 70.95486807823181, + 103.83680963516235 + ], + [ + 86.21527862548828, + 84.9826431274414, + 80.62500381469727, + 73.83681392669678, + 62.15278625488281, + 53.50695085525513, + 46.51042175292969, + 38.36806058883667, + 29.982642650604248, + 22.79514193534851, + 15.121530652046204, + 8.6458340883255, + 3.767361879348755, + 2.795139491558075, + 0.9201390892267227, + 2.6736122369766235, + 6.579862475395203, + 15.538198113441467, + 29.236114263534546, + 47.031256049871445, + 73.40278267860413, + 105.36459112167358 + ], + [ + 86.21527862548828, + 83.85416984558105, + 77.9340353012085, + 69.86112022399902, + 60.625006675720215, + 51.666672706604004, + 44.30556058883667, + 35.72917127609253, + 28.246531009674072, + 21.267364025115967, + 13.750003099441528, + 7.70833432674408, + 2.916667252779007, + 2.3090284168720245, + 0.9027780368924141, + 2.8993066549301147, + 7.63889092206955, + 17.2743079662323, + 30.468753576278687, + 49.27083796262741, + 75.08681321144104, + 108.21181726455688 + ], + [ + 85.5381965637207, + 82.79514122009277, + 77.51736545562744, + 66.99653434753418, + 57.083340644836426, + 48.94097709655762, + 41.406254291534424, + 33.975698471069336, + 27.17014169692993, + 19.010419368743896, + 11.545140266418457, + 6.302084445953369, + 2.4305562376976013, + 1.8576394617557526, + 0.7986112460494041, + 3.5590287297964096, + 8.975696861743927, + 18.715279936790466, + 31.579864740371704, + 52.15278500318527, + 77.56945013999939, + 109.58334684371948 + ], + [ + 84.9826431274414, + 79.9305591583252, + 72.69097995758057, + 63.43750762939453, + 54.32292318344116, + 47.56944990158081, + 40.17361545562744, + 31.961809635162354, + 24.35764169692993, + 17.3263920545578, + 10.329862356185913, + 4.826389789581299, + 1.944445013999939, + 1.4409725964069366, + 1.3541670590639114, + 4.149306729435921, + 10.086807668209076, + 19.253474235534668, + 35.312504529953, + 54.166673362255096, + 82.89930987358093, + 115.15625882148743 + ], + [ + 83.36805725097656, + 77.96875667572021, + 67.51737117767334, + 60.60764503479004, + 51.94444942474365, + 44.25347709655762, + 37.36111640930176, + 29.20139217376709, + 22.083336353302002, + 14.54861330986023, + 8.628473043441772, + 3.7500007152557373, + 1.4062503427267075, + 0.5555557198822498, + 1.927083969116211, + 5.486112341284752, + 11.649308383464813, + 22.170141458511353, + 38.125003814697266, + 57.91667479276657, + 85.520840883255, + 117.46528458595276 + ], + [ + 82.76041889190674, + 77.46528244018555, + 66.97917366027832, + 58.42014503479004, + 49.60070037841797, + 41.944448947906494, + 35.52083778381348, + 28.1944477558136, + 20.347225427627563, + 13.819446325302124, + 7.5868067145347595, + 3.576389789581299, + 1.093750275671482, + 0.0, + 2.1006951332092285, + 5.572917863726616, + 12.152780413627625, + 23.090280532836914, + 39.58333683013916, + 59.21875447034836, + 89.58333849906921, + 120.72917675971985 + ], + [ + 79.46181011199951, + 72.72570133209229, + 63.420145988464355, + 54.60070037841797, + 48.2812557220459, + 40.86806011199951, + 33.368059396743774, + 25.93750238418579, + 18.43750262260437, + 11.388890385627747, + 6.3020846247673035, + 2.6215286254882812, + 0.3645835090428591, + 0.9548613652586937, + 2.6041674613952637, + 6.631945729255676, + 13.333336234092712, + 25.208335638046265, + 41.961811780929565, + 64.01041954755783, + 93.15972995758057, + 126.78820633888245 + ], + [ + 76.97917175292969, + 66.49306106567383, + 58.76736927032471, + 49.895838260650635, + 43.6111159324646, + 36.684032917022705, + 28.680559158325195, + 21.493058443069458, + 14.618057489395142, + 8.906251072883606, + 4.097223103046417, + 1.9097229689359665, + 1.406250260770321, + 1.7361115217208862, + 3.8541676998138428, + 9.062502443790436, + 17.70833694934845, + 30.052086353302002, + 48.975701093673706, + 71.56250542402267, + 99.79167366027832, + 134.37501573562622 + ], + [ + 67.79514598846436, + 60.32986831665039, + 51.45834016799927, + 45.22569942474365, + 38.24653196334839, + 30.50347638130188, + 23.836808443069458, + 16.909724235534668, + 10.017362833023071, + 5.416667878627777, + 2.4305564761161804, + 0.9722225051373243, + 2.5173618495464325, + 3.2291675806045532, + 5.8506955951452255, + 12.482641816139221, + 23.0381977558136, + 36.40625476837158, + 56.21528625488281, + 78.55903309583664, + 111.30209064483643, + 149.01043272018433 + ], + [ + 59.965285301208496, + 53.00347709655762, + 45.59028148651123, + 38.73264265060425, + 30.694447994232178, + 23.73264241218567, + 17.44791889190674, + 11.163196802139282, + 5.954862117767334, + 2.725695699453354, + 1.4062504917383194, + 2.031250536441803, + 5.121529012918472, + 5.7638900578022, + 10.729169070720673, + 19.652780890464783, + 31.57986617088318, + 47.326396226882935, + 67.39584016799927, + 92.81250894069672, + 125.01737415790558, + 167.847234249115 + ], + [ + 50.59028339385986, + 44.652782917022705, + 36.8923659324646, + 29.861114025115967, + 23.211809158325195, + 16.371530055999756, + 9.947918772697449, + 5.607640087604523, + 2.5520845353603363, + 1.8750008195638657, + 2.430556207895279, + 5.642362594604492, + 11.666668832302094, + 13.194446563720703, + 20.555559635162354, + 30.555561304092407, + 46.770840883255005, + 62.81250810623169, + 85.95486831665039, + 115.05209195613861, + 146.02431845664978, + 191.26737117767334 + ], + [ + 40.312504291534424, + 32.343754291534424, + 25.989586114883423, + 20.277779579162598, + 13.090279340744019, + 8.055557012557983, + 4.409723669290543, + 2.6736123263835907, + 2.4131952822208405, + 3.5937508642673492, + 6.892362534999847, + 14.878475308418274, + 25.24306082725525, + 27.291671752929688, + 37.56944990158081, + 52.187506914138794, + 69.6180624961853, + 91.56251096725464, + 117.95139646530151, + 147.98612642288208, + 184.9305703639984, + 232.2743330001831 + ], + [ + 27.118058919906616, + 21.00694704055786, + 14.895834922790527, + 9.652779221534729, + 6.111112713813782, + 4.236112594604492, + 3.3159734904766083, + 3.7326402962207794, + 5.781251132488251, + 10.572918444871902, + 21.597226977348328, + 31.54514503479004, + 44.32292413711548, + 46.128480195999146, + 58.12500882148743, + 76.21528720855713, + 97.41320323944092, + 122.36112308502197, + 151.00696659088135, + 180.8680863380432, + 228.85419821739197, + 278.47224140167236 + ] + ] + }, + "partial_obs": { + "observed_units": [ + 0, + 1, + 2, + 3 + ], + "eps": 0.41666674613952637, + "profile_gCaL": [ + 1.840278148651123, + 1.5972224473953247, + 2.0486114025115967, + 2.118056058883667, + 1.7013893127441406, + 1.6319448947906494, + 1.5972225666046143, + 1.5277780294418335, + 1.3888890743255615, + 1.3541669845581055, + 1.3541669845581055, + 1.2847223281860352, + 0.9375002384185791, + 0.7638891935348511, + 0.0, + 0.590278148651123, + 0.5555558204650879, + 1.3888893127441406, + 1.6666672229766846, + 2.187500476837158, + 2.5694453716278076, + 3.333333730697632 + ], + "profile_w_exc": [ + 32.84722900390625, + 25.590282440185547, + 18.993061065673828, + 13.229169845581055, + 9.20139217376709, + 6.3194475173950195, + 4.375002861022949, + 2.951390504837036, + 2.7777791023254395, + 2.1875007152557373, + 1.840278148651123, + 1.3888893127441406, + 0.27777785062789917, + 0.0, + 1.2847223281860352, + 2.013889789581299, + 4.513890266418457, + 10.798614501953125, + 19.548614501953125, + 33.54167175292969, + 51.94445037841797, + 78.54167175292969 + ], + "gcal_interval": [ + 1.6, + 1.6 + ], + "wexc_interval": [ + 0.010338285194973314, + 0.0105 + ], + "dense_loss": [ + [ + 138.4722194671631, + 138.4722194671631, + 137.15277671813965, + 129.37500190734863, + 114.86112213134766, + 92.84723281860352, + 72.39584159851074, + 55.93750476837158, + 41.354169845581055, + 30.555559158325195, + 20.24306011199951, + 11.31944727897644, + 6.180557012557983, + 5.208335041999817, + 3.1250007152557373, + 2.465278446674347, + 4.652779161930084, + 10.937503576278687, + 19.826391220092773, + 33.784725189208984, + 54.131948471069336, + 80.52083587646484 + ], + [ + 138.4722194671631, + 138.4722194671631, + 137.18749809265137, + 128.22916984558105, + 114.86112213134766, + 90.45140075683594, + 70.34722900390625, + 55.97222709655762, + 39.861114501953125, + 29.756947994232178, + 19.86111545562744, + 11.215280055999756, + 6.111112833023071, + 5.347223877906799, + 3.05555659532547, + 2.36111181974411, + 4.548612534999847, + 10.798614501953125, + 20.069445610046387, + 35.90277862548828, + 51.94444751739502, + 82.50000667572021 + ], + [ + 138.4722194671631, + 138.4722194671631, + 136.00694847106934, + 128.22916984558105, + 114.86112213134766, + 88.81945610046387, + 69.65278434753418, + 55.555562019348145, + 39.548614501953125, + 29.097225666046143, + 18.71528196334839, + 11.319447040557861, + 6.215279459953308, + 4.930557131767273, + 2.847223162651062, + 2.187500685453415, + 4.513890206813812, + 11.284724950790405, + 20.034724473953247, + 35.486114501953125, + 55.41667175292969, + 82.56945514678955 + ], + [ + 138.4722194671631, + 138.4722194671631, + 136.00694847106934, + 125.90278053283691, + 111.8055648803711, + 87.84723281860352, + 68.92361545562744, + 55.00000762939453, + 38.99306011199951, + 28.19444704055786, + 19.16667079925537, + 11.319447040557861, + 6.006946325302124, + 5.03472363948822, + 2.7777785062789917, + 2.2222228348255157, + 4.930556774139404, + 11.701392889022827, + 20.416669368743896, + 35.55555820465088, + 58.15972423553467, + 81.97917461395264 + ], + [ + 138.4722194671631, + 138.4722194671631, + 134.96527671813965, + 123.99306297302246, + 109.7569522857666, + 85.52084350585938, + 67.11806392669678, + 51.38889503479004, + 39.06250476837158, + 25.937503337860107, + 18.020837545394897, + 11.076391458511353, + 5.694446086883545, + 4.895834803581238, + 2.6388895511627197, + 2.2222228348255157, + 5.034723460674286, + 12.187503099441528, + 20.868058681488037, + 36.9444465637207, + 57.048614501953125, + 82.63889694213867 + ], + [ + 138.4722194671631, + 138.4722194671631, + 132.77777671813965, + 122.88195610046387, + 105.72917938232422, + 82.8819522857666, + 65.65972900390625, + 50.83333873748779, + 37.951393127441406, + 25.416670322418213, + 17.187503337860107, + 10.173614025115967, + 5.173612713813782, + 5.000001907348633, + 2.430556058883667, + 2.3263897001743317, + 5.347223877906799, + 12.569447755813599, + 21.319446086883545, + 38.22916889190674, + 57.84722709655762, + 83.36806297302246 + ], + [ + 138.4722194671631, + 137.11805534362793, + 132.77777671813965, + 120.97223281860352, + 102.11807060241699, + 79.27084159851074, + 63.33333778381348, + 48.54167175292969, + 37.013893127441406, + 25.069448947906494, + 16.45833730697632, + 9.895835876464844, + 5.000001668930054, + 4.687501311302185, + 2.1527782678604126, + 2.638889968395233, + 5.729168176651001, + 12.986114263534546, + 23.437501907348633, + 40.0694465637207, + 59.3750057220459, + 84.16667747497559 + ], + [ + 138.4722194671631, + 137.11805534362793, + 132.77777290344238, + 117.67362213134766, + 99.09723663330078, + 78.33334159851074, + 62.46528339385986, + 47.5000057220459, + 34.44444942474365, + 24.7222261428833, + 16.111114978790283, + 8.854169011116028, + 4.930556893348694, + 4.3402791023254395, + 1.8055559396743774, + 2.7083343863487244, + 5.763890624046326, + 13.506948947906494, + 23.090279579162598, + 39.027780055999756, + 60.312503814697266, + 88.47223091125488 + ], + [ + 138.4722194671631, + 137.15277671813965, + 130.45138931274414, + 115.69445610046387, + 94.61806297302246, + 75.41667556762695, + 60.62500476837158, + 45.69445037841797, + 33.92361640930176, + 23.125004291534424, + 15.555558919906616, + 8.090279936790466, + 4.340279459953308, + 3.750001311302185, + 1.3888891339302063, + 2.777778744697571, + 6.562502026557922, + 14.166670799255371, + 25.451390743255615, + 40.62500476837158, + 61.979172706604004, + 91.21528244018555 + ], + [ + 138.4722194671631, + 136.00694847106934, + 127.29166984558105, + 113.71529006958008, + 90.34723472595215, + 73.05556392669678, + 59.0625057220459, + 43.854172706604004, + 32.15278148651123, + 22.46528196334839, + 14.02778172492981, + 7.569446325302124, + 4.166668176651001, + 3.541667878627777, + 1.3541669845581055, + 3.2638902068138123, + 6.5277793407440186, + 14.201392889022827, + 26.145835399627686, + 42.36111354827881, + 64.23611259460449, + 93.47223091125488 + ], + [ + 138.4722194671631, + 133.75000190734863, + 121.90973281860352, + 105.76390266418457, + 87.29167556762695, + 69.37500762939453, + 54.65278339385986, + 40.48611545562744, + 30.451394081115723, + 21.250003814697266, + 12.881947994232178, + 6.875002145767212, + 3.541667938232422, + 3.020834505558014, + 1.4583337604999542, + 3.645834743976593, + 7.361113548278809, + 15.729169368743896, + 27.222224712371826, + 43.88889026641846, + 65.83334159851074, + 96.56251335144043 + ], + [ + 137.11805534362793, + 131.63194465637207, + 121.0763931274414, + 100.03473091125488, + 80.20834350585938, + 63.92361640930176, + 49.89583683013916, + 37.812503814697266, + 29.201393127441406, + 18.22917079925537, + 11.111114025115967, + 6.7361133098602295, + 2.986112415790558, + 2.500001013278961, + 1.2847224026918411, + 4.166667819023132, + 9.13194727897644, + 16.979169607162476, + 28.715280055999756, + 46.38889217376709, + 69.44444751739502, + 98.50695991516113 + ], + [ + 136.00694847106934, + 125.90278053283691, + 111.42362213134766, + 92.91667747497559, + 74.68750858306885, + 61.180562019348145, + 47.46528244018555, + 35.55556011199951, + 25.59028148651123, + 16.28472638130188, + 9.930558204650879, + 5.48611307144165, + 2.534723222255707, + 1.909722924232483, + 1.9444449543952942, + 4.444445908069611, + 10.555558204650879, + 17.777780532836914, + 32.01389265060425, + 48.12500190734863, + 73.57639217376709, + 104.58334350585938 + ], + [ + 132.77777671813965, + 121.97917556762695, + 101.0764045715332, + 87.2569522857666, + 69.93056106567383, + 55.59028244018555, + 43.7500057220459, + 32.36111545562744, + 22.395837783813477, + 14.826392650604248, + 8.576391458511353, + 4.16666841506958, + 1.8402784168720245, + 0.7638891637325287, + 2.638889968395233, + 5.729168057441711, + 11.70139229297638, + 20.590281009674072, + 34.61805820465088, + 52.18750286102295, + 77.95139503479004, + 107.60417556762695 + ], + [ + 131.5625, + 120.97222709655762, + 100.00000953674316, + 82.8819522857666, + 65.24306297302246, + 50.9722261428833, + 40.90278244018555, + 31.215282917022705, + 20.7986159324646, + 13.923614978790283, + 8.12500250339508, + 4.166668653488159, + 1.3888893872499466, + 0.0, + 2.7777786254882812, + 6.041668176651001, + 12.361114025115967, + 22.152780532836914, + 37.22222423553467, + 54.51389122009277, + 81.63194942474365, + 113.05556678771973 + ], + [ + 124.96528244018555, + 111.4930648803711, + 92.88195419311523, + 75.24306297302246, + 62.60417366027832, + 49.826393127441406, + 38.402782917022705, + 28.71528148651123, + 18.50694751739502, + 11.979169607162476, + 6.944446682929993, + 3.3333349227905273, + 0.5902781523764133, + 1.1458336561918259, + 3.194445312023163, + 6.875001668930054, + 12.986114263534546, + 24.097224712371826, + 38.090280055999756, + 57.986114501953125, + 84.02778434753418, + 119.4444580078125 + ], + [ + 120.0000057220459, + 99.02778434753418, + 83.57640075683594, + 65.8333387374878, + 54.30556011199951, + 43.229172706604004, + 32.22222709655762, + 23.125004291534424, + 14.895837306976318, + 9.652780294418335, + 4.826390743255615, + 2.743056833744049, + 1.8055558651685715, + 2.118055999279022, + 4.305556833744049, + 9.861114025115967, + 17.048614740371704, + 27.43055820465088, + 44.86111402511597, + 65.20833778381348, + 90.90279006958008, + 124.61807060241699 + ], + [ + 101.63195419311523, + 86.7013988494873, + 68.95834255218506, + 57.53472709655762, + 45.55556011199951, + 35.034727573394775, + 25.97222661972046, + 18.437503814697266, + 11.631947994232178, + 6.8750022649765015, + 3.3333348631858826, + 1.3888893984258175, + 2.569445312023163, + 3.6111119985580444, + 5.90277898311615, + 12.152780532836914, + 21.284725666046143, + 32.36111354827881, + 50.52083778381348, + 70.93750381469727, + 102.88195610046387, + 139.5139045715332 + ], + [ + 85.97223281860352, + 72.04861640930176, + 58.26389122009277, + 46.52778148651123, + 36.11111640930176, + 26.97917127609253, + 19.47917079925537, + 13.611115455627441, + 7.6736133098602295, + 4.0277798771858215, + 2.0833341479301453, + 2.256945013999939, + 4.930556893348694, + 5.590279221534729, + 9.30555784702301, + 16.284725427627563, + 25.868058681488037, + 40.27778196334839, + 59.79167079925537, + 81.80556297302246, + 114.09723472595215, + 152.50001525878906 + ], + [ + 67.22222900390625, + 56.38889408111572, + 44.44444942474365, + 35.17361640930176, + 26.562504768371582, + 19.58333730697632, + 13.020837545394897, + 7.916669249534607, + 4.166668891906738, + 2.569445848464966, + 2.187500536441803, + 4.513890087604523, + 8.437501788139343, + 9.131946563720703, + 14.13194751739502, + 21.458338260650635, + 34.895838260650635, + 50.06944942474365, + 71.94444847106934, + 98.99306678771973, + 124.09723472595215, + 170.3125 + ], + [ + 49.65278148651123, + 39.27083969116211, + 30.83333730697632, + 23.680559158325195, + 16.70139217376709, + 11.38889217376709, + 6.805558204650879, + 4.097224473953247, + 2.777779221534729, + 2.7430561184883118, + 4.652778804302216, + 8.229168176651001, + 13.506947994232178, + 14.861114501953125, + 21.31944751739502, + 31.597225666046143, + 45.45139217376709, + 63.402785301208496, + 84.79167461395264, + 111.4583511352539, + 146.80556297302246, + 190.24305725097656 + ], + [ + 32.847227573394775, + 25.590282440185547, + 18.993059635162354, + 13.229170083999634, + 9.20139217376709, + 6.319447040557861, + 4.375002503395081, + 3.333334982395172, + 3.3333338499069214, + 5.4861119389534, + 9.16666865348816, + 14.791669845581055, + 21.805559158325195, + 22.465282917022705, + 30.833338260650635, + 43.40278434753418, + 58.26389122009277, + 77.8819522857666, + 101.59723854064941, + 129.6527976989746, + 172.84723663330078, + 216.90974044799805 + ] + ] + } + }, + { + "conn_seed": 2, + "event_offset": 20, + "active_units_at_truth": { + "fit_ramp": 6, + "fit_trap": 6, + "heldout_slowramp": 6 + }, + "truth_fit_loss": 0.0, + "eps_by_model": { + "full_obs": 0.3298611640930176, + "partial_obs": 0.48611125349998474 + }, + "dense_time_s": 207.6, + "event_hashes": { + "fit_ramp": "8cc04203515ca2a1", + "fit_trap": "31d67d0186ed83f3", + "heldout_slowramp": "5f4540f762db5ac3" + }, + "full_obs": { + "observed_units": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ], + "eps": 0.3298611640930176, + "profile_gCaL": [ + 1.4756947755813599, + 1.1979167461395264, + 1.197916865348816, + 1.3020836114883423, + 1.3020836114883423, + 1.2326390743255615, + 1.4409723281860352, + 1.215277910232544, + 1.1284723281860352, + 1.0243057012557983, + 0.9722223281860352, + 1.0069445371627808, + 0.8333334922790527, + 0.5902777910232544, + 0.0, + 0.4861111640930176, + 0.7812501192092896, + 1.024305820465088, + 1.4756948947906494, + 2.118055820465088, + 2.812500476837158, + 3.454861640930176 + ], + "profile_w_exc": [ + 24.51388931274414, + 19.479169845581055, + 14.756945610046387, + 8.993057250976562, + 5.989583969116211, + 4.236112117767334, + 3.454861640930176, + 2.812500476837158, + 2.2222225666046143, + 2.118055820465088, + 1.4756948947906494, + 1.024305820465088, + 0.4861111640930176, + 0.0, + 0.8854168653488159, + 1.7361114025115967, + 4.392362117767334, + 10.173612594604492, + 19.704862594604492, + 37.517364501953125, + 61.99653625488281, + 95.95487213134766 + ], + "gcal_interval": [ + 1.6, + 1.6 + ], + "wexc_interval": [ + 0.0105, + 0.0105 + ], + "dense_loss": [ + [ + 80.5902829170227, + 78.4895920753479, + 73.35070085525513, + 67.20487070083618, + 61.44098138809204, + 55.798619747161865, + 47.86458921432495, + 42.30903339385986, + 35.017364501953125, + 27.100696325302124, + 18.593751907348633, + 11.805556058883667, + 6.631944954395294, + 5.503472805023193, + 2.829861283302307, + 1.7881947457790375, + 4.392362207174301, + 10.2777791172266, + 20.486113250255585, + 37.77778321504593, + 63.142372369766235, + 96.75348091125488 + ], + [ + 80.5902829170227, + 77.82987070083618, + 72.70834016799927, + 67.22223138809204, + 61.006953716278076, + 55.79862070083618, + 48.00347948074341, + 41.68403434753418, + 34.21875286102295, + 26.857641220092773, + 18.38541841506958, + 11.388889074325562, + 6.493056118488312, + 5.277778267860413, + 2.4305558800697327, + 1.7881948202848434, + 4.548612207174301, + 10.208334818482399, + 19.70486243069172, + 38.1250062584877, + 63.593759298324585, + 96.45834445953369 + ], + [ + 80.5902829170227, + 77.82987070083618, + 72.70834016799927, + 67.2395920753479, + 60.46876001358032, + 54.565980434417725, + 47.30903339385986, + 40.83333969116211, + 34.28819704055786, + 25.833335638046265, + 17.82986307144165, + 10.815972805023193, + 6.163195192813873, + 5.225695013999939, + 2.2395835518836975, + 1.7361114919185638, + 4.791667729616165, + 10.329862609505653, + 21.423613265156746, + 38.2986176609993, + 63.97570252418518, + 95.95487117767334 + ], + [ + 80.5902829170227, + 76.4583420753479, + 72.67361688613892, + 66.250009059906, + 60.48612070083618, + 54.14931344985962, + 46.961811542510986, + 40.62500762939453, + 33.87153100967407, + 25.590280055999756, + 17.53472375869751, + 10.208333849906921, + 5.868056297302246, + 5.26041704416275, + 2.2743058800697327, + 1.8923616707324982, + 5.243056386709213, + 10.868057385087013, + 21.250003054738045, + 39.097229063510895, + 64.13195443153381, + 98.50695419311523 + ], + [ + 80.5902829170227, + 76.4583420753479, + 72.13542222976685, + 66.26737070083618, + 60.06945466995239, + 53.26389837265015, + 46.92708873748779, + 40.3993124961853, + 33.298614501953125, + 25.01736330986023, + 16.475695729255676, + 10.121528387069702, + 5.7118061780929565, + 5.017361760139465, + 1.9618059396743774, + 1.770833820104599, + 5.295139938592911, + 10.937501564621925, + 22.18750210106373, + 39.39236694574356, + 64.13195538520813, + 97.34376001358032 + ], + [ + 80.5902829170227, + 76.4583420753479, + 71.5277829170227, + 64.86111783981323, + 59.16667699813843, + 52.916675090789795, + 46.63194990158081, + 39.27083921432495, + 32.65625238418579, + 24.09722399711609, + 15.694445848464966, + 9.791667222976685, + 5.399306058883667, + 5.034722805023193, + 1.8923614919185638, + 1.8055560886859894, + 5.329862147569656, + 12.048613205552101, + 22.10069651901722, + 41.09375602006912, + 64.98264861106873, + 99.27084350585938 + ], + [ + 80.5902829170227, + 76.4583420753479, + 70.45139455795288, + 64.28820276260376, + 57.84723234176636, + 52.135425090789795, + 45.885422706604004, + 38.22917175292969, + 32.46528100967407, + 23.80208468437195, + 15.520834565162659, + 9.27083420753479, + 5.069445013999939, + 4.427084028720856, + 1.857639193534851, + 1.9618060439825058, + 5.711806923151016, + 11.927084937691689, + 23.02083583176136, + 41.37153458595276, + 66.54514956474304, + 100.17362308502197 + ], + [ + 79.18403482437134, + 74.54861879348755, + 70.41667318344116, + 62.84723138809204, + 57.79514837265015, + 51.19792318344116, + 45.43403482437134, + 37.378477573394775, + 31.25000286102295, + 22.100695610046387, + 14.600695610046387, + 9.35763955116272, + 4.895834028720856, + 4.062500476837158, + 1.406250149011612, + 2.1701392233371735, + 5.868056744337082, + 12.951391458511353, + 23.767363384366035, + 42.274311780929565, + 66.94445538520813, + 99.96528816223145 + ], + [ + 78.4895920753479, + 73.95834016799927, + 69.30556344985962, + 62.46528673171997, + 56.12848234176636, + 49.32292413711548, + 43.21181106567383, + 36.7534761428833, + 30.26041889190674, + 20.81597328186035, + 14.201390266418457, + 9.010417342185974, + 4.878472924232483, + 3.5937505066394806, + 1.1284724026918411, + 2.4479172229766846, + 6.475695461034775, + 12.6909738779068, + 24.097224786877632, + 43.83681321144104, + 69.02778649330139, + 103.71528720855713 + ], + [ + 77.11806344985962, + 72.69097948074341, + 67.22223043441772, + 61.90973234176636, + 55.17362070083618, + 48.94097852706909, + 42.60417318344116, + 34.49653100967407, + 27.795141220092773, + 20.486111640930176, + 13.263889908790588, + 8.072917222976685, + 3.958333730697632, + 3.055556058883667, + 1.0763890594244003, + 2.7083338499069214, + 6.753473192453384, + 14.357640862464905, + 27.23958595097065, + 45.972227334976196, + 69.72223329544067, + 106.52778625488281 + ], + [ + 76.4583420753479, + 71.57986879348755, + 64.84375715255737, + 60.50348234176636, + 54.02778768539429, + 47.56945085525513, + 41.0243124961853, + 33.90625286102295, + 25.607640981674194, + 18.78472352027893, + 11.892362236976624, + 7.118056118488312, + 3.125000387430191, + 2.5173615217208862, + 0.972222350537777, + 3.0902784764766693, + 7.413195878267288, + 14.895835816860199, + 28.524308994412422, + 47.72570061683655, + 71.64931583404541, + 109.39237022399902 + ], + [ + 74.47917413711548, + 69.86111783981323, + 63.36806344985962, + 58.26389837265015, + 51.59722948074341, + 46.024312019348145, + 38.628478050231934, + 32.621530532836914, + 24.236112117767334, + 16.788195967674255, + 11.024306416511536, + 6.215278446674347, + 2.569444954395294, + 1.9097225815057755, + 1.145833671092987, + 3.697917401790619, + 7.986112147569656, + 16.875002801418304, + 31.56250412762165, + 49.73959040641785, + 75.59029150009155, + 111.92709159851074 + ], + [ + 72.69097757339478, + 68.33334112167358, + 61.979175090789795, + 55.625009059906006, + 50.59028387069702, + 43.87153387069702, + 36.47569799423218, + 29.496530294418335, + 21.979167938232422, + 15.277779340744019, + 9.774306178092957, + 5.034722745418549, + 1.9270837008953094, + 1.5451392233371735, + 1.7187504321336746, + 4.27083420753479, + 9.479168027639389, + 18.02083569765091, + 32.465282276272774, + 53.8368136882782, + 78.73265171051025, + 116.94445514678955 + ], + [ + 71.05903482437134, + 64.23611879348755, + 60.53820466995239, + 53.07292413711548, + 47.673617362976074, + 41.14583969116211, + 34.496530532836914, + 26.215279579162598, + 19.11458420753479, + 13.437500953674316, + 8.177083671092987, + 3.9583338797092438, + 1.336805820465088, + 0.7986111603677273, + 1.8750003576278687, + 5.538195252418518, + 11.232640787959099, + 20.381947576999664, + 35.88542206585407, + 55.989593267440796, + 81.56251049041748, + 123.17709350585938 + ], + [ + 69.34028673171997, + 63.85417413711548, + 59.166675090789795, + 52.79514741897583, + 46.77083921432495, + 40.03472948074341, + 33.61111402511597, + 24.809029579162598, + 18.29861283302307, + 12.604167819023132, + 7.17013943195343, + 3.246528387069702, + 1.1284723654389381, + 0.0, + 2.1180558502674103, + 5.74652886390686, + 11.284724280238152, + 21.232641518115997, + 36.77084071934223, + 58.68056559562683, + 85.86806583404541, + 128.00347995758057 + ], + [ + 68.35070276260376, + 62.465285778045654, + 56.97917604446411, + 50.52084016799927, + 44.20139503479004, + 37.465280055999756, + 31.284725666046143, + 23.420140027999878, + 17.291668176651, + 10.677083849906921, + 5.781250417232513, + 2.4479170739650726, + 0.48611118644475937, + 1.128472477197647, + 3.0208339393138885, + 7.152779132127762, + 13.263890787959099, + 23.888893127441406, + 39.80903346836567, + 62.777788400650024, + 90.39931631088257, + 130.6423683166504 + ], + [ + 63.85417413711548, + 59.131953716278076, + 53.10764741897583, + 47.20486783981323, + 40.3125057220459, + 34.114586353302, + 25.55555820465088, + 19.322916984558105, + 13.732639789581299, + 8.420140027999878, + 3.836806058883667, + 1.7187502980232239, + 1.7013894021511078, + 2.013889253139496, + 4.097222685813904, + 8.993056863546371, + 16.57986332476139, + 28.64583730697632, + 44.687505692243576, + 69.20139920711517, + 97.96875667572021, + 142.10070610046387 + ], + [ + 60.034730434417725, + 54.34028673171997, + 48.17709016799927, + 41.63195085525513, + 35.97222423553467, + 27.15277910232544, + 21.076390743255615, + 15.243056535720825, + 9.704861998558044, + 4.947917103767395, + 2.4479172974824905, + 1.024305798113346, + 3.090278446674347, + 3.5590284168720245, + 6.423611700534821, + 12.25694665312767, + 21.822920203208923, + 35.20833945274353, + 53.05556246638298, + 74.87848377227783, + 109.27083826065063, + 152.43056297302246 + ], + [ + 54.28820276260376, + 49.027785301208496, + 42.32639455795288, + 36.25000286102295, + 28.211806774139404, + 21.927085399627686, + 15.81597375869751, + 10.243056535720825, + 5.312500774860382, + 2.569444864988327, + 1.4756948202848434, + 2.55208383500576, + 5.312500834465027, + 6.527778506278992, + 10.555557578802109, + 18.43750250339508, + 29.895837545394897, + 46.284727573394775, + 64.93056580424309, + 89.63542866706848, + 125.46875619888306, + 169.68750953674316 + ], + [ + 47.13542413711548, + 41.04167175292969, + 35.50347566604614, + 26.857640027999878, + 20.57291865348816, + 14.947917938232422, + 9.4270840883255, + 5.399307012557983, + 2.881944954395294, + 2.1180560141801834, + 3.142361640930176, + 6.024306356906891, + 10.677085429430008, + 12.065973788499832, + 19.305558919906616, + 30.62500500679016, + 45.05209159851074, + 60.90278434753418, + 82.0138997733593, + 110.65972924232483, + 148.33334302902222, + 193.7673807144165 + ], + [ + 37.812503814697266, + 30.625001907348633, + 23.576391220092773, + 18.3680579662323, + 12.55208444595337, + 7.829861879348755, + 4.774306803941727, + 2.8125004023313522, + 3.1770839989185333, + 4.340278387069702, + 7.760417580604553, + 12.465280085802078, + 21.822919845581055, + 24.305559277534485, + 34.79167449474335, + 48.64584136009216, + 64.09722805023193, + 86.57986959815025, + 113.21182152628899, + 143.83682298660278, + 188.54168891906738, + 234.28821659088135 + ], + [ + 24.513891458511353, + 19.479169368743896, + 14.756946086883545, + 8.993056893348694, + 5.9895846247673035, + 4.236112117767334, + 3.454861655831337, + 4.305556118488312, + 7.152778714895248, + 10.034724056720734, + 16.822920560836792, + 28.17708671092987, + 40.32986807823181, + 43.52431273460388, + 54.86111664772034, + 71.75347709655762, + 91.5972318649292, + 115.00001367926598, + 146.23265653848648, + 179.44446659088135, + 226.4930772781372, + 276.597261428833 + ] + ] + }, + "partial_obs": { + "observed_units": [ + 0, + 1, + 2, + 3 + ], + "eps": 0.48611125349998474, + "profile_gCaL": [ + 2.187500476837158, + 1.8750004768371582, + 1.8750003576278687, + 2.048611640930176, + 1.8750005960464478, + 1.8055559396743774, + 1.5972223281860352, + 1.805556058883667, + 1.6319448947906494, + 1.493055820465088, + 1.4583334922790527, + 1.666666865348816, + 1.2152780294418335, + 0.5902778506278992, + 0.0, + 0.902777910232544, + 1.1111114025115967, + 1.6666672229766846, + 2.2916674613952637, + 2.6388893127441406, + 2.8125007152557373, + 3.7500009536743164 + ], + "profile_w_exc": [ + 32.638893127441406, + 25.868061065673828, + 19.756946563720703, + 14.06250286102295, + 9.861113548278809, + 6.701391220092773, + 4.097223281860352, + 3.645834445953369, + 3.2291669845581055, + 3.020833969116211, + 2.2916674613952637, + 1.6666672229766846, + 0.902777910232544, + 0.0, + 1.3888893127441406, + 2.083333730697632, + 4.9305572509765625, + 10.833334922790527, + 20.486114501953125, + 36.423614501953125, + 57.18750762939453, + 85.79862976074219 + ], + "gcal_interval": [ + 1.6, + 1.6 + ], + "wexc_interval": [ + 0.0105, + 0.0105 + ], + "dense_loss": [ + [ + 135.34722900390625, + 131.14584732055664, + 120.8680648803711, + 108.5764045715332, + 97.04862594604492, + 85.76390266418457, + 69.89584159851074, + 58.784729957580566, + 44.20139217376709, + 31.97916841506958, + 20.69444751739502, + 12.70833444595337, + 6.354167819023132, + 5.763890027999878, + 3.784722685813904, + 2.1875004172325134, + 4.930557131767273, + 11.076390266418457, + 21.70139169692993, + 37.291672706604004, + 58.92362403869629, + 87.50000953674316 + ], + [ + 135.34722900390625, + 129.8264045715332, + 119.58334350585938, + 108.61112594604492, + 96.18057060241699, + 85.7639045715332, + 70.17362213134766, + 57.5347318649292, + 43.64583683013916, + 31.49305820465088, + 20.277780532836914, + 12.326389789581299, + 6.3888901472091675, + 5.312500953674316, + 3.3680561184883118, + 2.187500536441803, + 5.208334922790527, + 10.833334922790527, + 20.486114025115967, + 37.70833921432495, + 58.64584445953369, + 87.04862213134766 + ], + [ + 135.34722900390625, + 129.8264045715332, + 119.58334350585938, + 108.64584732055664, + 95.10418319702148, + 83.29862403869629, + 68.78472995758057, + 55.83334255218506, + 43.784725189208984, + 29.44444704055786, + 19.0972261428833, + 11.319445610046387, + 6.21527898311615, + 5.208334445953369, + 3.020833671092987, + 2.1527784168720245, + 5.555556893348694, + 11.076390504837036, + 22.430558681488037, + 37.534728050231934, + 59.47917366027832, + 85.79862213134766 + ], + [ + 135.34722900390625, + 127.08334732055664, + 119.51389694213867, + 106.66668128967285, + 95.1389045715332, + 82.46529006958008, + 68.09028625488281, + 55.4166784286499, + 42.951393127441406, + 28.958335876464844, + 18.99305820465088, + 10.833334684371948, + 6.076390266418457, + 5.312500834465027, + 3.0555561184883118, + 2.1527785062789917, + 6.21527886390686, + 11.80555772781372, + 21.631947994232178, + 38.5763955116272, + 58.888898849487305, + 88.75000953674316 + ], + [ + 135.34722900390625, + 127.08334732055664, + 118.43750762939453, + 106.7014045715332, + 94.30557250976562, + 80.69445991516113, + 68.02084064483643, + 54.965288162231445, + 41.80556011199951, + 28.50694704055786, + 17.847224473953247, + 10.798612594604492, + 6.2152791023254395, + 5.138890027999878, + 2.777778387069702, + 2.187500774860382, + 6.527779221534729, + 12.326390743255615, + 23.15972375869751, + 38.22917318344116, + 60.83334541320801, + 87.25695705413818 + ], + [ + 135.34722900390625, + 127.08334732055664, + 117.22222900390625, + 103.8888988494873, + 92.5000171661377, + 80.00001335144043, + 67.43056297302246, + 52.70834159851074, + 40.520835876464844, + 27.63889169692993, + 17.291669368743896, + 10.486112594604492, + 5.555556535720825, + 5.173612117767334, + 2.7777785062789917, + 2.1875008046627045, + 6.493057012557983, + 13.645836591720581, + 22.951391220092773, + 41.111117362976074, + 61.111122131347656, + 88.9930648803711 + ], + [ + 135.34722900390625, + 127.08334732055664, + 115.0694522857666, + 102.74306869506836, + 89.86112785339355, + 78.43751335144043, + 65.93750858306885, + 50.625006675720215, + 40.138893127441406, + 27.74305772781372, + 16.979169130325317, + 9.79166841506958, + 5.243056535720825, + 4.652778744697571, + 2.777778387069702, + 2.5000007152557373, + 7.118057727813721, + 13.159724235534668, + 24.166670322418213, + 41.21528720855713, + 62.74306678771973, + 89.54862403869629 + ], + [ + 132.53473281860352, + 123.26390075683594, + 115.00000953674316, + 99.86112594604492, + 89.75695991516113, + 76.56250953674316, + 65.03473281860352, + 48.92361831665039, + 37.70833683013916, + 25.833335876464844, + 16.215280055999756, + 9.513890266418457, + 4.895834565162659, + 4.236111998558044, + 2.1180558800697327, + 2.7430561780929565, + 7.395834922790527, + 14.618058919906616, + 24.65277910232544, + 41.07639694213867, + 62.60417652130127, + 89.79167366027832 + ], + [ + 131.14584732055664, + 122.08334350585938, + 112.77779006958008, + 99.09723663330078, + 86.42362785339355, + 72.8125114440918, + 60.590285301208496, + 47.67361545562744, + 35.72916889190674, + 23.99305820465088, + 15.416668891906738, + 9.201390504837036, + 4.861112356185913, + 3.8888898491859436, + 1.6319448053836823, + 3.090278744697571, + 7.986112594604492, + 13.958335876464844, + 26.07639169692993, + 43.26389789581299, + 64.58334255218506, + 93.12500762939453 + ], + [ + 128.40279006958008, + 119.54862213134766, + 108.61112403869629, + 97.98612785339355, + 84.5139045715332, + 72.04862022399902, + 59.375009536743164, + 44.201393127441406, + 32.56944751739502, + 23.22916889190674, + 14.756946802139282, + 7.604167938232422, + 4.027778625488281, + 3.7500009536743164, + 1.701389268040657, + 3.472223162651062, + 8.437501549720764, + 15.833335638046265, + 28.81944704055786, + 47.1875057220459, + 66.7013988494873, + 97.70834350585938 + ], + [ + 127.08334732055664, + 117.32640075683594, + 103.85417747497559, + 95.17362785339355, + 82.22223854064941, + 69.3055648803711, + 56.215288162231445, + 43.02083683013916, + 29.930558681488037, + 21.250003337860107, + 12.500002145767212, + 6.840279221534729, + 3.8194454312324524, + 3.1944451332092285, + 1.45833358168602, + 4.097223460674286, + 8.61111295223236, + 16.215280771255493, + 30.694448471069336, + 46.90972900390625, + 69.27084445953369, + 100.72917366027832 + ], + [ + 123.1250114440918, + 113.8888988494873, + 100.90279006958008, + 90.69445991516113, + 77.36112213134766, + 66.21528720855713, + 51.42361927032471, + 40.45139217376709, + 28.715280532836914, + 18.57639193534851, + 11.770835161209106, + 6.319445729255676, + 3.2986119985580444, + 2.7083339393138885, + 1.7013895511627197, + 4.756945848464966, + 9.375001311302185, + 18.680559158325195, + 32.430559158325195, + 50.17361927032471, + 71.80556869506836, + 103.40278244018555 + ], + [ + 119.54861831665039, + 110.83334541320801, + 98.12501335144043, + 85.41668128967285, + 75.34723091125488, + 61.90973091125488, + 48.15972709655762, + 35.97222566604614, + 26.31944751739502, + 17.083336353302002, + 10.416668176651001, + 5.20833420753479, + 2.777778446674347, + 2.152778446674347, + 2.3958340883255005, + 5.138890266418457, + 10.625001430511475, + 19.409725666046143, + 33.88889217376709, + 52.88195037841797, + 76.90973472595215, + 110.4513988494873 + ], + [ + 116.28473281860352, + 102.63890075683594, + 95.24307250976562, + 80.3125114440918, + 69.51389789581299, + 56.45834255218506, + 44.20139217376709, + 31.87500286102295, + 21.909724712371826, + 15.138891220092773, + 8.750000834465027, + 4.4791677594184875, + 1.9097228050231934, + 1.3194445967674255, + 2.5347227454185486, + 6.354167699813843, + 12.638890504837036, + 22.77778148651123, + 37.70833873748779, + 55.79862022399902, + 79.0625114440918, + 116.77083778381348 + ], + [ + 112.84723663330078, + 101.8750114440918, + 92.50001335144043, + 79.7569580078125, + 67.70834159851074, + 54.236122131347656, + 42.430559158325195, + 30.59028148651123, + 20.833336353302002, + 13.888890981674194, + 7.395834684371948, + 4.2361122369766235, + 1.701389104127884, + 0.0, + 2.6736115217208862, + 6.8055572509765625, + 12.500002145767212, + 23.68055820465088, + 38.05556297302246, + 57.187509536743164, + 83.81945419311523, + 123.54167366027832 + ], + [ + 110.86806869506836, + 99.09723472595215, + 88.12501525878906, + 75.20834350585938, + 62.56945323944092, + 50.13889122009277, + 39.58333778381348, + 29.20139169692993, + 20.13889169692993, + 11.979168176651001, + 6.562501311302185, + 3.090278446674347, + 0.9027779400348663, + 1.6319449543952942, + 3.611112058162689, + 8.715279459953308, + 14.861112833023071, + 25.55556011199951, + 41.875006675720215, + 62.46528720855713, + 87.22223663330078, + 126.00695037841797 + ], + [ + 101.8750114440918, + 92.43057060241699, + 80.3819580078125, + 68.5763988494873, + 54.79167461395264, + 43.437503814697266, + 32.08333873748779, + 22.951391220092773, + 15.763891220092773, + 10.10416865348816, + 4.722223401069641, + 2.63888943195343, + 2.256945312023163, + 2.500000536441803, + 4.791667461395264, + 10.104168176651001, + 18.50694751739502, + 29.548614501953125, + 45.27778339385986, + 69.82639789581299, + 95.20834541320801, + 138.36806869506836 + ], + [ + 94.23612403869629, + 82.84723663330078, + 70.52084350585938, + 57.430564880371094, + 47.1527795791626, + 35.277780532836914, + 25.937503814697266, + 18.333335399627686, + 11.597223997116089, + 6.423612356185913, + 3.888890027999878, + 1.6666672229766846, + 3.6111122369766235, + 3.9236121773719788, + 7.187501072883606, + 12.638891220092773, + 22.32639217376709, + 35.86806058883667, + 51.631951332092285, + 74.20139980316162, + 105.52084541320801, + 146.73612594604492 + ], + [ + 82.74306869506836, + 72.22223377227783, + 58.8194522857666, + 47.70833683013916, + 37.361114501953125, + 27.604171752929688, + 19.51389169692993, + 13.506946325302124, + 7.361112713813782, + 4.097222924232483, + 2.291667401790619, + 2.604167550802231, + 5.27777886390686, + 6.45833420753479, + 10.55555772781372, + 17.847224950790405, + 27.8472261428833, + 43.88889408111572, + 61.87500858306885, + 84.96529197692871, + 118.40278244018555, + 159.79167938232422 + ], + [ + 68.4375114440918, + 56.250006675720215, + 47.08333778381348, + 35.38194799423218, + 26.111114978790283, + 18.715280532836914, + 13.263890981674194, + 8.333335757255554, + 4.58333420753479, + 3.0555564165115356, + 2.63888943195343, + 4.89583432674408, + 8.92361295223236, + 9.791668891906738, + 15.590280055999756, + 24.791670322418213, + 37.25695037841797, + 51.07639503479004, + 71.8750114440918, + 101.1805648803711, + 134.8264045715332, + 175.59029388427734 + ], + [ + 50.83333873748779, + 41.49306011199951, + 30.208338737487793, + 23.750004291534424, + 16.979169130325317, + 11.562502384185791, + 7.812502145767212, + 4.444445192813873, + 3.263889729976654, + 3.333333969116211, + 5.590278625488281, + 8.819446563720703, + 14.30555772781372, + 16.25000262260437, + 22.777783155441284, + 33.40278387069702, + 45.312503814697266, + 64.54862213134766, + 88.99307250976562, + 114.40973663330078, + 156.5625171661377, + 195.00000762939453 + ], + [ + 32.6388955116272, + 25.86806058883667, + 19.756947994232178, + 14.062503099441528, + 9.86111330986023, + 6.701391100883484, + 4.097223311662674, + 4.131945490837097, + 5.24305647611618, + 6.458335041999817, + 10.520834922790527, + 15.17361330986023, + 23.020838260650635, + 24.68750524520874, + 32.777782917022705, + 43.61111640930176, + 59.652788162231445, + 78.54167747497559, + 103.99306869506836, + 134.3055648803711, + 175.62500381469727, + 223.8194580078125 + ] + ] + } + } + ], + "estimators": { + "full_obs/grid_refine": { + "theta_hat": [ + 0.49999999999999994, + 0.011339289449053857 + ], + "fit_loss": 1.4583334922790527, + "within_eps": false, + "boundary_hit": false + }, + "full_obs/grad_best_of_starts": { + "theta_hat": [ + 1.1833169567847504, + 0.010649853440178473 + ], + "fit_loss": 1.840278148651123, + "within_eps": false, + "n_failed_starts": 1, + "n_starts": 2, + "runs": [ + { + "label": "start_default", + "start": [ + 1.0, + 0.009 + ], + "best": { + "hard_loss": 1.840278148651123, + "theta": [ + 1.1833169567847504, + 0.010649853440178473 + ], + "it": 14 + }, + "nonfinite_iterates": 0, + "failed": false, + "history": [ + { + "it": 0, + "hard_loss": 14.84375, + "theta": [ + 1.0, + 0.008999999493676799 + ] + }, + { + "it": 1, + "hard_loss": 20.069446563720703, + "theta": [ + 0.9512297427165632, + 0.008561066788217404 + ] + }, + { + "it": 2, + "hard_loss": 19.479167938232422, + "theta": [ + 0.9537362502973373, + 0.008583626313377725 + ] + }, + { + "it": 3, + "hard_loss": 17.72569465637207, + "theta": [ + 0.9731106475617731, + 0.00875799571054733 + ] + }, + { + "it": 4, + "hard_loss": 14.375, + "theta": [ + 1.0015454076987096, + 0.009013910512742739 + ] + }, + { + "it": 5, + "hard_loss": 13.52430534362793, + "theta": [ + 1.0111649815080435, + 0.009100484686042026 + ] + }, + { + "it": 6, + "hard_loss": 11.16319465637207, + "theta": [ + 1.02972619557951, + 0.009267536378141946 + ] + }, + { + "it": 7, + "hard_loss": 9.184028625488281, + "theta": [ + 1.0550414163363415, + 0.009495374688946708 + ] + }, + { + "it": 8, + "hard_loss": 6.1979169845581055, + "theta": [ + 1.0858703076929452, + 0.009772836224166073 + ] + }, + { + "it": 9, + "hard_loss": 5.208333969116211, + "theta": [ + 1.1033154442836832, + 0.0099298416211911 + ] + }, + { + "it": 10, + "hard_loss": 4.739583969116211, + "theta": [ + 1.109749086122528, + 0.00998774262708609 + ] + }, + { + "it": 11, + "hard_loss": 3.715277910232544, + "theta": [ + 1.1234310435518249, + 0.010110878145159987 + ] + }, + { + "it": 12, + "hard_loss": 3.0729174613952637, + "theta": [ + 1.1432798338190282, + 0.010289517312199105 + ] + }, + { + "it": 13, + "hard_loss": 2.204861640930176, + "theta": [ + 1.1685582268164607, + 0.010517024233259342 + ] + }, + { + "it": 14, + "hard_loss": 1.840278148651123, + "theta": [ + 1.1833169567847504, + 0.010649853440178473 + ] + } + ] + }, + { + "label": "start_perturbed", + "start": [ + 0.6, + 0.014 + ], + "best": { + "hard_loss": 25.572919845581055, + "theta": [ + 0.6000000296207707, + 0.014000002950072156 + ], + "it": 0 + }, + "nonfinite_iterates": 1, + "failed": true, + "history": [ + { + "it": 0, + "hard_loss": 25.572919845581055, + "theta": [ + 0.6000000296207707, + 0.014000002950072156 + ] + }, + { + "it": 1, + "nonfinite": true + } + ] + } + ] + }, + "partial_obs/grid_refine": { + "theta_hat": [ + 0.9999999999999999, + 0.011339289449053857 + ], + "fit_loss": 1.631945252418518, + "within_eps": false, + "boundary_hit": false + }, + "partial_obs/grad_best_of_starts": { + "theta_hat": [ + 1.1833169567847504, + 0.010649853440178473 + ], + "fit_loss": 2.2569448947906494, + "within_eps": false, + "n_failed_starts": 1, + "n_starts": 2, + "runs": [ + { + "label": "start_default", + "start": [ + 1.0, + 0.009 + ], + "best": { + "hard_loss": 2.2569448947906494, + "theta": [ + 1.1833169567847504, + 0.010649853440178473 + ], + "it": 14 + }, + "nonfinite_iterates": 0, + "failed": false, + "history": [ + { + "it": 0, + "hard_loss": 11.875, + "theta": [ + 1.0, + 0.008999999493676799 + ] + }, + { + "it": 1, + "hard_loss": 19.166667938232422, + "theta": [ + 0.9512297427165632, + 0.008561066788217404 + ] + }, + { + "it": 2, + "hard_loss": 17.98611068725586, + "theta": [ + 0.9537362502973373, + 0.008583626313377725 + ] + }, + { + "it": 3, + "hard_loss": 14.479166030883789, + "theta": [ + 0.9731106475617731, + 0.00875799571054733 + ] + }, + { + "it": 4, + "hard_loss": 11.527778625488281, + "theta": [ + 1.0015454076987096, + 0.009013910512742739 + ] + }, + { + "it": 5, + "hard_loss": 10.902777671813965, + "theta": [ + 1.0111649815080435, + 0.009100484686042026 + ] + }, + { + "it": 6, + "hard_loss": 8.541667938232422, + "theta": [ + 1.02972619557951, + 0.009267536378141946 + ] + }, + { + "it": 7, + "hard_loss": 7.048611640930176, + "theta": [ + 1.0550414163363415, + 0.009495374688946708 + ] + }, + { + "it": 8, + "hard_loss": 5.486112117767334, + "theta": [ + 1.0858703076929452, + 0.009772836224166073 + ] + }, + { + "it": 9, + "hard_loss": 4.652778625488281, + "theta": [ + 1.1033154442836832, + 0.0099298416211911 + ] + }, + { + "it": 10, + "hard_loss": 4.305556297302246, + "theta": [ + 1.109749086122528, + 0.00998774262708609 + ] + }, + { + "it": 11, + "hard_loss": 3.611111640930176, + "theta": [ + 1.1234310435518249, + 0.010110878145159987 + ] + }, + { + "it": 12, + "hard_loss": 3.2291674613952637, + "theta": [ + 1.1432798338190282, + 0.010289517312199105 + ] + }, + { + "it": 13, + "hard_loss": 2.5000007152557373, + "theta": [ + 1.1685582268164607, + 0.010517024233259342 + ] + }, + { + "it": 14, + "hard_loss": 2.2569448947906494, + "theta": [ + 1.1833169567847504, + 0.010649853440178473 + ] + } + ] + }, + { + "label": "start_perturbed", + "start": [ + 0.6, + 0.014 + ], + "best": { + "hard_loss": 13.71528148651123, + "theta": [ + 0.6000000296207707, + 0.014000002950072156 + ], + "it": 0 + }, + "nonfinite_iterates": 1, + "failed": true, + "history": [ + { + "it": 0, + "hard_loss": 13.71528148651123, + "theta": [ + 0.6000000296207707, + 0.014000002950072156 + ] + }, + { + "it": 1, + "nonfinite": true + } + ] + } + ] + } + }, + "recovery_table": [ + { + "obs_model": "full_obs", + "method": "grid_refine", + "theta_hat": [ + 0.49999999999999994, + 0.011339289449053857 + ], + "param_rel_err": [ + 0.6875, + 0.07993232848131968 + ], + "fit_loss": 1.4583334922790527, + "fit_within_eps": false, + "heldout_mse_by_limb": { + "pre": 0.0, + "rising": 1.42746901512146, + "peak": 5.034722328186035, + "falling": 0.7716049551963806, + "post": 0.0, + "overall": 1.1944445371627808 + } + }, + { + "obs_model": "full_obs", + "method": "grad_best_of_starts", + "theta_hat": [ + 1.1833169567847504, + 0.010649853440178473 + ], + "param_rel_err": [ + 0.26042690200953106, + 0.014271756207473518 + ], + "fit_loss": 1.840278148651123, + "fit_within_eps": false, + "heldout_mse_by_limb": { + "pre": 0.0, + "rising": 1.1188271045684814, + "peak": 3.124999523162842, + "falling": 0.9645061492919922, + "post": 0.0, + "overall": 0.9999999403953552 + } + }, + { + "obs_model": "partial_obs", + "method": "grid_refine", + "theta_hat": [ + 0.9999999999999999, + 0.011339289449053857 + ], + "param_rel_err": [ + 0.3750000000000001, + 0.07993232848131968 + ], + "fit_loss": 1.631945252418518, + "fit_within_eps": false, + "heldout_mse_by_limb": { + "pre": 0.0, + "rising": 0.5401232838630676, + "peak": 2.9513888359069824, + "falling": 0.5787036418914795, + "post": 0.0, + "overall": 0.6388888359069824 + } + }, + { + "obs_model": "partial_obs", + "method": "grad_best_of_starts", + "theta_hat": [ + 1.1833169567847504, + 0.010649853440178473 + ], + "param_rel_err": [ + 0.26042690200953106, + 0.014271756207473518 + ], + "fit_loss": 2.2569448947906494, + "fit_within_eps": false, + "heldout_mse_by_limb": { + "pre": 0.0, + "rising": 1.1188271045684814, + "peak": 3.124999523162842, + "falling": 0.9645061492919922, + "post": 0.0, + "overall": 0.9999999403953552 + } + } + ], + "heldout_at_truth_by_limb": { + "pre": 0.0, + "rising": 0.0, + "peak": 0.0, + "falling": 0.0, + "post": 0.0, + "overall": 0.0 + }, + "wall_times_s": { + "grid_full_obs": 30.2, + "grad_full_obs": 44.4, + "grid_partial_obs": 30.5, + "grad_partial_obs": 44.8, + "total": 986.8 + } +} \ No newline at end of file diff --git a/experiments/twin/evidence/deliverable_c_full_binned_20261018_131156.json b/experiments/twin/evidence/deliverable_c_full_binned_20261018_131156.json new file mode 100644 index 00000000..1c128c33 --- /dev/null +++ b/experiments/twin/evidence/deliverable_c_full_binned_20261018_131156.json @@ -0,0 +1,311 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "993a2cfdebaeb499", + "poolbuild.py": "816d4ba42490e283", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "983f8bbcafebd67e", + "protocols.py": "404ea437741ea2ae" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float32 (JAX default; x64 not enabled)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations", + "deliverable_c_sha256_16": "0c1896bdc5a2e5fc", + "args": { + "n_mn": 8, + "true": [ + 1.6, + 0.0105 + ], + "dense_n": 21, + "replicates": 3, + "dur_scale": 1.0, + "observable": "binned" + }, + "run_tag": "full" + }, + "declared": { + "nuisance": "per-protocol effective_w_exc_p (\u00b5S); exact multiplicative symmetry with w_exc \u21d2 w_exc is structurally non-identifiable here by construction; the question is whether gCaL stays constrained", + "observable": "binned", + "observation_model": "full_obs only (declared)", + "eps_rule": "per replicate, dt-halving at truth, all units", + "theta_true": [ + 1.6, + 0.0105 + ], + "axes_contain_truth": true, + "heldout": "not evaluated \u2014 would require a declared amplitude rule for the unseen protocol (separate experiment)" + }, + "g_dense": [ + 0.25, + 0.2871745887492587, + 0.32987697769322355, + 0.37892914162759955, + 0.43527528164806206, + 0.5, + 0.5743491774985175, + 0.6597539553864472, + 0.7578582832551991, + 0.8705505632961241, + 1.0, + 1.1486983549970349, + 1.3195079107728946, + 1.5157165665103984, + 1.6, + 1.7411011265922485, + 2.0000000000000004, + 2.29739670999407, + 2.6390158215457893, + 3.0314331330207964, + 3.4822022531844974, + 4.0 + ], + "b_dense_effective_w_exc": [ + 0.0045, + 0.004822980581413315, + 0.005169142597486656, + 0.00554014986005212, + 0.005937785598478024, + 0.006363961030678926, + 0.006820724549296786, + 0.007310271567206119, + 0.007834955069665113, + 0.008397296923831269, + 0.008999999999999998, + 0.009645961162826632, + 0.010338285194973314, + 0.0105, + 0.011080299720104242, + 0.011875571196956045, + 0.012727922061357854, + 0.013641449098593581, + 0.01462054313441224, + 0.01566991013933023, + 0.01679459384766253, + 0.018 + ], + "replicates": [ + { + "conn_seed": 0, + "event_offset": 0, + "eps_full_obs": 0.677083395421505, + "profile_known": [ + 1.5798608660697937, + 1.6840274333953857, + 1.736110508441925, + 1.8055551648139954, + 1.8229162096977234, + 1.7534716129302979, + 1.5104164481163025, + 1.4062498807907104, + 1.493055522441864, + 1.458332896232605, + 1.3541667461395264, + 1.2500001788139343, + 0.9201390445232391, + 0.8854168653488159, + 0.0, + 0.6076390743255615, + 0.8680554926395416, + 1.4062498807907104, + 2.152777910232544, + 2.847220778465271, + 2.8993040323257446, + 3.33333158493042 + ], + "profile_unknown": [ + 1.5798608660697937, + 1.666666328907013, + 1.736110508441925, + 1.8055551648139954, + 1.8229162096977234, + 1.7534716129302979, + 1.510416328907013, + 1.336805522441864, + 1.493055522441864, + 1.3020831942558289, + 1.3541667461395264, + 1.2500001788139343, + 0.9201390445232391, + 0.8854168355464935, + 0.0, + 0.5729168355464935, + 0.850694328546524, + 1.4062498807907104, + 2.065971314907074, + 2.4999985694885254, + 2.8993040323257446, + 3.33333158493042 + ], + "gcal_interval_known": [ + 1.6, + 1.7411011265922485 + ], + "gcal_interval_unknown": [ + 1.6, + 1.7411011265922485 + ], + "interval_terminology": "sampled compatibility interval (deterministic level set, not a statistical confidence interval)", + "n_evals": 559, + "event_hashes": { + "fit_ramp": "5ef0d36f3c87eee4", + "fit_trap": "682ace78eb9e26c4" + } + }, + { + "conn_seed": 1, + "event_offset": 10, + "eps_full_obs": 0.2951389402151108, + "profile_known": [ + 1.215277910232544, + 1.0937501788139343, + 1.319444715976715, + 1.3541672229766846, + 1.1284724473953247, + 1.059027910232544, + 1.1111112833023071, + 1.0763887465000153, + 0.8854168355464935, + 0.8680557310581207, + 0.8680557012557983, + 0.7986112534999847, + 0.46875013411045074, + 0.5208334624767303, + 0.0, + 0.3645835444331169, + 0.5555558055639267, + 0.9722225069999695, + 1.4062498807907104, + 1.875000774860382, + 2.4131951332092285, + 3.315972089767456 + ], + "profile_unknown": [ + 1.1631945371627808, + 1.0937501788139343, + 1.163194715976715, + 1.2326392531394958, + 1.1284724473953247, + 1.0416667759418488, + 1.093750238418579, + 0.9375000894069672, + 0.8854168355464935, + 0.8680557012557983, + 0.7812502086162567, + 0.7986112534999847, + 0.46875013411045074, + 0.38194458186626434, + 0.0, + 0.3645835444331169, + 0.5555558055639267, + 0.9722225069999695, + 1.284722089767456, + 1.8750006556510925, + 2.3263893127441406, + 3.1944440603256226 + ], + "gcal_interval_known": [ + 1.6, + 1.6 + ], + "gcal_interval_unknown": [ + 1.6, + 1.6 + ], + "interval_terminology": "sampled compatibility interval (deterministic level set, not a statistical confidence interval)", + "n_evals": 550, + "event_hashes": { + "fit_ramp": "ed698d364473ebbd", + "fit_trap": "e791fdc5716d20b5" + } + }, + { + "conn_seed": 2, + "event_offset": 20, + "eps_full_obs": 0.3298611491918564, + "profile_known": [ + 1.4756945967674255, + 1.1979169249534607, + 1.197916865348816, + 1.232639193534851, + 1.2152779698371887, + 1.2152777910232544, + 1.1111112833023071, + 1.2152777910232544, + 1.1284722685813904, + 1.0243057906627655, + 0.9722224175930023, + 1.0069444179534912, + 0.4166668504476547, + 0.590277910232544, + 0.0, + 0.48611123859882355, + 0.7812502086162567, + 1.0243056416511536, + 1.4756946563720703, + 2.1180554032325745, + 2.8124992847442627, + 3.4548600912094116 + ], + "profile_unknown": [ + 1.406250238418579, + 1.1979169249534607, + 1.1979168057441711, + 1.232639193534851, + 1.2152779698371887, + 1.2152777910232544, + 1.1111112833023071, + 1.1458332538604736, + 1.0590278506278992, + 0.9375001788139343, + 0.9722224175930023, + 0.9722223877906799, + 0.4166668504476547, + 0.590277910232544, + 0.0, + 0.48611123859882355, + 0.7812502086162567, + 1.0243056416511536, + 1.4756946563720703, + 2.1180554032325745, + 2.7951385974884033, + 3.4374990463256836 + ], + "gcal_interval_known": [ + 1.6, + 1.6 + ], + "gcal_interval_unknown": [ + 1.6, + 1.6 + ], + "interval_terminology": "sampled compatibility interval (deterministic level set, not a statistical confidence interval)", + "n_evals": 559, + "event_hashes": { + "fit_ramp": "8cc04203515ca2a1", + "fit_trap": "31d67d0186ed83f3" + } + } + ], + "wall_s": 705.2 +} \ No newline at end of file diff --git a/experiments/twin/evidence/g4_descent_full_20261018_180008.json b/experiments/twin/evidence/g4_descent_full_20261018_180008.json new file mode 100644 index 00000000..cb63ca6b --- /dev/null +++ b/experiments/twin/evidence/g4_descent_full_20261018_180008.json @@ -0,0 +1,2622 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "a433b78e7bd7e581", + "poolbuild.py": "2dd213214e166343", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "983f8bbcafebd67e", + "protocols.py": "404ea437741ea2ae", + "gradient_repair.py": "78812d0e94b7be75", + "deliverable_b.py": "8c16e479bcea843d", + "deliverable_c.py": "0c1896bdc5a2e5fc", + "pic_grid.py": "d900172c74ba1695", + "r2a_profiles.py": "ac4802b3bdfe1191", + "acceptance.py": "24200ac953a87395", + "myogen/simulator/jaxley/jax_models.py": "8b99fc4b69d735c7", + "myogen/simulator/jaxley/channels/nerlab_channels.py": "f6161ceaca04285e", + "myogen/simulator/jaxley/cells_api.py": "c6534ba58adaf3b1", + "myogen/simulator/jaxley/populations/motor_neurons.py": "736be1bb0c75807c" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float64 (jax_enable_x64=True)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations", + "g4_sha256_16": "c42a163d0ea57376", + "args": { + "smoke": false, + "iters": 10 + }, + "run_tag": "full" + }, + "declared": { + "dt_ms": 0.025, + "window_s": 2.0, + "objective": "FULL-TRACE hard binned-rate MSE (50/300 ms) at dt 0.025; windows affect gradient flow only", + "starts": "S1 both-wrong, S2 gCaL-wrong-only, S3 w-wrong-only", + "lambdas": [ + 0.03, + 0.1 + ] + }, + "stageA": { + "rows": [ + { + "case": { + "conn_seed": 0, + "event_offset": 0, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "n_obs_spikes": 125, + "starts": { + "S1_both_wrong": { + "L0": 16.909722222222257, + "estimators": { + "surrogate": { + "grad": [ + -9930.86337931859, + -17517649.89031852 + ], + "direction": [ + 0.06286498093845376, + 0.9980220409247523 + ], + "joint": { + "0.03": { + "loss": 12.951388888888898, + "accepted": true + }, + "0.1": { + "loss": 7.013888888888899, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 16.458333333333368, + "accepted": true + }, + "w": { + "loss": 7.083333333333343, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -4.513888888888928, + -10802.469135802505 + ], + "direction": [ + 0.04637861118671757, + 0.998923933252373 + ], + "joint": { + "0.03": { + "loss": 12.951388888888898, + "accepted": true + }, + "0.1": { + "loss": 7.013888888888899, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 16.458333333333368, + "accepted": true + }, + "w": { + "loss": 7.083333333333343, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9998636931013194 + }, + "S2_gcal_wrong_only": { + "L0": 1.6666666666666647, + "estimators": { + "surrogate": { + "grad": [ + -2312.739595550043, + -2106542.4627247397 + ], + "direction": [ + 0.09931421382456011, + 0.9950561225038563 + ], + "joint": { + "0.03": { + "loss": 1.0069444444444435, + "accepted": true + }, + "0.1": { + "loss": 3.4027777777777746, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 1.6319444444444426, + "accepted": true + }, + "w": { + "loss": 3.4027777777777746, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + 0.347222222222221, + -2525.2525252525234 + ], + "direction": [ + -0.012499023551925984, + 0.9999218841540815 + ], + "joint": { + "0.03": { + "loss": 1.0069444444444435, + "accepted": true + }, + "0.1": { + "loss": 3.4027777777777746, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 1.7708333333333313, + "accepted": false + }, + "w": { + "loss": 3.4027777777777746, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.9937370621554763 + }, + "S3_w_wrong_only": { + "L0": 10.79861111111112, + "estimators": { + "surrogate": { + "grad": [ + -38107.249287801, + -36635745.48176435 + ], + "direction": [ + 0.1708131555698204, + 0.9853034384819127 + ], + "joint": { + "0.03": { + "loss": 9.201388888888902, + "accepted": true + }, + "0.1": { + "loss": 3.6805555555555554, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 9.548611111111121, + "accepted": true + }, + "w": { + "loss": 3.6805555555555554, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -6.712962962962951, + -9297.839506172857 + ], + "direction": [ + 0.11947010847565422, + 0.9928377980218197 + ], + "joint": { + "0.03": { + "loss": 9.201388888888902, + "accepted": true + }, + "0.1": { + "loss": 3.6805555555555554, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 9.548611111111121, + "accepted": true + }, + "w": { + "loss": 3.6805555555555554, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9986535624707049 + } + } + }, + { + "case": { + "conn_seed": 0, + "event_offset": 0, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "n_obs_spikes": 36, + "starts": { + "S1_both_wrong": { + "L0": 3.506944444444443, + "estimators": { + "surrogate": { + "grad": [ + 3837.1588313830434, + 7941670.502289136 + ], + "direction": [ + -0.053608105556254344, + -0.9985620516616228 + ], + "joint": { + "0.03": { + "loss": 2.222222222222219, + "accepted": true + }, + "0.1": { + "loss": 0.9027777777777772, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 3.402777777777776, + "accepted": true + }, + "w": { + "loss": 0.9027777777777772, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 3.819444444444451, + 5632.716049382726 + ], + "direction": [ + -0.075129531595126, + -0.9971737829898542 + ], + "joint": { + "0.03": { + "loss": 2.222222222222219, + "accepted": true + }, + "0.1": { + "loss": 0.9027777777777772, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 3.402777777777776, + "accepted": true + }, + "w": { + "loss": 0.9027777777777772, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9997674504656741 + }, + "S2_gcal_wrong_only": { + "L0": 0.5555555555555554, + "estimators": { + "surrogate": { + "grad": [ + 962.586596699731, + 1390477.7952398334 + ], + "direction": [ + -0.08621161941587199, + -0.9962768474062282 + ], + "joint": { + "0.03": { + "loss": 0.1041666666666667, + "accepted": true + }, + "0.1": { + "loss": 1.2847222222222212, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.5208333333333333, + "accepted": true + }, + "w": { + "loss": 1.2499999999999991, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + 1.0416666666666652, + 347.22222222222194 + ], + "direction": [ + -0.35112344158839154, + -0.9363291775690445 + ], + "joint": { + "0.03": { + "loss": 0.1041666666666667, + "accepted": true + }, + "0.1": { + "loss": 1.2499999999999991, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.5208333333333333, + "accepted": true + }, + "w": { + "loss": 1.2499999999999991, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.9631140016771638 + }, + "S3_w_wrong_only": { + "L0": 2.361111111111108, + "estimators": { + "surrogate": { + "grad": [ + 3534.517813014828, + 5059198.883518435 + ], + "direction": [ + -0.05425799532142149, + -0.9985269500337488 + ], + "joint": { + "0.03": { + "loss": 1.458333333333332, + "accepted": true + }, + "0.1": { + "loss": 0.3819444444444446, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 2.0486111111111085, + "accepted": true + }, + "w": { + "loss": 0.3819444444444446, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 2.976190476190466, + 4861.111111111119 + ], + "direction": [ + -0.04756514941544917, + -0.9988681377244376 + ], + "joint": { + "0.03": { + "loss": 1.458333333333332, + "accepted": true + }, + "0.1": { + "loss": 0.3819444444444446, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 2.0486111111111085, + "accepted": true + }, + "w": { + "loss": 0.3819444444444446, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9999775447023195 + } + } + }, + { + "case": { + "conn_seed": 0, + "event_offset": 10, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "n_obs_spikes": 129, + "starts": { + "S1_both_wrong": { + "L0": 15.173611111111137, + "estimators": { + "surrogate": { + "grad": [ + -20589.719461357214, + -32452100.23581164 + ], + "direction": [ + 0.07032158114681511, + 0.9975243732485998 + ], + "joint": { + "0.03": { + "loss": 13.472222222222232, + "accepted": true + }, + "0.1": { + "loss": 8.68055555555557, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 14.618055555555575, + "accepted": true + }, + "w": { + "loss": 8.68055555555557, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -15.277777777777963, + -8371.91358024696 + ], + "direction": [ + 0.19872105095794504, + 0.980056092224404 + ], + "joint": { + "0.03": { + "loss": 13.472222222222232, + "accepted": true + }, + "0.1": { + "loss": 8.68055555555557, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 14.618055555555575, + "accepted": true + }, + "w": { + "loss": 8.68055555555557, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9916042176551401 + }, + "S2_gcal_wrong_only": { + "L0": 2.2916666666666643, + "estimators": { + "surrogate": { + "grad": [ + 12632.807880009952, + 21147146.135355983 + ], + "direction": [ + -0.0542270478553772, + -0.9985286311773391 + ], + "joint": { + "0.03": { + "loss": 3.090277777777774, + "accepted": false + }, + "0.1": { + "loss": 8.68055555555557, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 2.2569444444444424, + "accepted": true + }, + "w": { + "loss": 8.68055555555557, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + -3.124999999999998, + -3598.484848484854 + ], + "direction": [ + 0.0787024854200105, + 0.9968981486534686 + ], + "joint": { + "0.03": { + "loss": 1.3888888888888875, + "accepted": true + }, + "0.1": { + "loss": 2.847222222222219, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 1.6666666666666643, + "accepted": true + }, + "w": { + "loss": 2.847222222222219, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": -0.9996991472413795 + }, + "S3_w_wrong_only": { + "L0": 12.777777777777784, + "estimators": { + "surrogate": { + "grad": [ + -22699.181639288523, + -21374392.620761607 + ], + "direction": [ + 0.17428772387682784, + 0.9846947696143382 + ], + "joint": { + "0.03": { + "loss": 10.416666666666677, + "accepted": true + }, + "0.1": { + "loss": 5.6944444444444535, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 10.798611111111118, + "accepted": true + }, + "w": { + "loss": 5.763888888888898, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -6.018518518518523, + -7137.345679012363 + ], + "direction": [ + 0.13917281530317618, + 0.990268108887986 + ], + "joint": { + "0.03": { + "loss": 10.416666666666677, + "accepted": true + }, + "0.1": { + "loss": 5.6944444444444535, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 10.798611111111118, + "accepted": true + }, + "w": { + "loss": 5.763888888888898, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9993679405426025 + } + } + }, + { + "case": { + "conn_seed": 0, + "event_offset": 10, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "n_obs_spikes": 38, + "starts": { + "S1_both_wrong": { + "L0": 5.138888888888895, + "estimators": { + "surrogate": { + "grad": [ + 12780.276677555523, + 19384485.087007836 + ], + "direction": [ + -0.07306027176825657, + -0.9973275272893798 + ], + "joint": { + "0.03": { + "loss": 2.9513888888888857, + "accepted": true + }, + "0.1": { + "loss": 0.9722222222222213, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 4.097222222222226, + "accepted": true + }, + "w": { + "loss": 0.9722222222222213, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 8.333333333333357, + 4591.049382716059 + ], + "direction": [ + -0.1977000139471188, + -0.9802625691544633 + ], + "joint": { + "0.03": { + "loss": 2.9513888888888857, + "accepted": true + }, + "0.1": { + "loss": 0.9722222222222213, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 4.097222222222226, + "accepted": true + }, + "w": { + "loss": 0.9722222222222213, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9920868609367202 + }, + "S2_gcal_wrong_only": { + "L0": 0.6597222222222219, + "estimators": { + "surrogate": { + "grad": [ + 399.19485036264126, + 435645.7503564779 + ], + "direction": [ + -0.1137970719318783, + -0.9935040142947239 + ], + "joint": { + "0.03": { + "loss": 0.13888888888888892, + "accepted": true + }, + "0.1": { + "loss": 1.0069444444444435, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.41666666666666674, + "accepted": true + }, + "w": { + "loss": 1.0069444444444435, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + 1.7361111111111094, + 1432.2916666666642 + ], + "direction": [ + -0.14980537942799502, + -0.9887155042247665 + ], + "joint": { + "0.03": { + "loss": 0.13888888888888892, + "accepted": true + }, + "0.1": { + "loss": 1.0069444444444435, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.41666666666666674, + "accepted": true + }, + "w": { + "loss": 1.0069444444444435, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.9993402359812874 + }, + "S3_w_wrong_only": { + "L0": 2.8472222222222197, + "estimators": { + "surrogate": { + "grad": [ + 1526.3934088640376, + 2023457.2974954592 + ], + "direction": [ + -0.05857088069147903, + -0.9982832523562761 + ], + "joint": { + "0.03": { + "loss": 1.944444444444442, + "accepted": true + }, + "0.1": { + "loss": 0.38194444444444453, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 2.6388888888888857, + "accepted": true + }, + "w": { + "loss": 0.38194444444444453, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 8.928571428571423, + 4783.950617283959 + ], + "direction": [ + -0.14365563920417165, + -0.9896277367398515 + ], + "joint": { + "0.03": { + "loss": 1.944444444444442, + "accepted": true + }, + "0.1": { + "loss": 0.38194444444444453, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 2.6388888888888857, + "accepted": true + }, + "w": { + "loss": 0.38194444444444453, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9963428329591252 + } + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 0, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "n_obs_spikes": 152, + "starts": { + "S1_both_wrong": { + "L0": 15.034722222222232, + "estimators": { + "surrogate": { + "grad": [ + -8133.782816624142, + -11779855.419187834 + ], + "direction": [ + 0.07649547059133602, + 0.9970699288309773 + ], + "joint": { + "0.03": { + "loss": 11.493055555555559, + "accepted": true + }, + "0.1": { + "loss": 6.180555555555564, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 13.472222222222223, + "accepted": true + }, + "w": { + "loss": 6.180555555555564, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -13.541666666666732, + -13117.283950617322 + ], + "direction": [ + 0.11395863041154274, + 0.9934854958954988 + ], + "joint": { + "0.03": { + "loss": 11.493055555555559, + "accepted": true + }, + "0.1": { + "loss": 6.180555555555564, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 13.472222222222223, + "accepted": true + }, + "w": { + "loss": 6.180555555555564, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9992918317484083 + }, + "S2_gcal_wrong_only": { + "L0": 1.7013888888888875, + "estimators": { + "surrogate": { + "grad": [ + -3124.739302975264, + -3283268.9755985066 + ], + "direction": [ + 0.08619760681163467, + 0.9962780598708106 + ], + "joint": { + "0.03": { + "loss": 0.9722222222222214, + "accepted": true + }, + "0.1": { + "loss": 2.8819444444444415, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 1.3194444444444435, + "accepted": true + }, + "w": { + "loss": 2.8819444444444415, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + -0.6944444444444464, + -2619.949494949492 + ], + "direction": [ + 0.024089392975120152, + 0.9997098084674824 + ], + "joint": { + "0.03": { + "loss": 0.9722222222222214, + "accepted": true + }, + "0.1": { + "loss": 2.8819444444444415, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 1.3194444444444435, + "accepted": true + }, + "w": { + "loss": 2.8819444444444415, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.9980653964378033 + }, + "S3_w_wrong_only": { + "L0": 9.756944444444448, + "estimators": { + "surrogate": { + "grad": [ + -5005.173816298737, + -5817210.9552992275 + ], + "direction": [ + 0.1419492139507958, + 0.9898739418020617 + ], + "joint": { + "0.03": { + "loss": 7.6736111111111205, + "accepted": true + }, + "0.1": { + "loss": 3.8194444444444415, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 8.298611111111116, + "accepted": true + }, + "w": { + "loss": 3.9930555555555554, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -11.805555555555534, + -9143.518518518524 + ], + "direction": [ + 0.21037412892142535, + 0.977620952046626 + ], + "joint": { + "0.03": { + "loss": 7.6736111111111205, + "accepted": true + }, + "0.1": { + "loss": 3.8541666666666643, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 8.298611111111116, + "accepted": true + }, + "w": { + "loss": 3.9930555555555554, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9975839476266577 + } + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 0, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "n_obs_spikes": 56, + "starts": { + "S1_both_wrong": { + "L0": 5.034722222222226, + "estimators": { + "surrogate": { + "grad": [ + 5419.667550823454, + 7647301.83091794 + ], + "direction": [ + -0.07850179530284088, + -0.9969139722835823 + ], + "joint": { + "0.03": { + "loss": 3.0902777777777746, + "accepted": true + }, + "0.1": { + "loss": 1.2499999999999984, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 4.652777777777782, + "accepted": true + }, + "w": { + "loss": 1.2499999999999984, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 1.7361111111111338, + 6635.80246913581 + ], + "direction": [ + -0.029057492497498415, + -0.9995777419139333 + ], + "joint": { + "0.03": { + "loss": 3.0902777777777746, + "accepted": true + }, + "0.1": { + "loss": 1.2499999999999984, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 4.652777777777782, + "accepted": true + }, + "w": { + "loss": 1.2499999999999984, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9987740826257252 + }, + "S2_gcal_wrong_only": { + "L0": 1.0763888888888882, + "estimators": { + "surrogate": { + "grad": [ + 5059.898469513917, + 6885860.358130202 + ], + "direction": [ + -0.09146800818722425, + -0.9958080153715685 + ], + "joint": { + "0.03": { + "loss": 0.20833333333333337, + "accepted": true + }, + "0.1": { + "loss": 1.4930555555555538, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.8680555555555549, + "accepted": true + }, + "w": { + "loss": 1.4930555555555538, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + 0.3472222222222232, + 1562.4999999999977 + ], + "direction": [ + -0.027767067240353407, + -0.9996144206527183 + ], + "joint": { + "0.03": { + "loss": 0.20833333333333337, + "accepted": true + }, + "0.1": { + "loss": 1.4930555555555538, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.8680555555555549, + "accepted": true + }, + "w": { + "loss": 1.4930555555555538, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.9979638507006595 + }, + "S3_w_wrong_only": { + "L0": 3.576388888888889, + "estimators": { + "surrogate": { + "grad": [ + 7621.241908431744, + 13717106.740759248 + ], + "direction": [ + -0.043173140037555714, + -0.9990676053097196 + ], + "joint": { + "0.03": { + "loss": 2.430555555555553, + "accepted": true + }, + "0.1": { + "loss": 0.7638888888888885, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 3.2638888888888875, + "accepted": true + }, + "w": { + "loss": 0.7638888888888885, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 2.9761904761904785, + 4976.8518518518595 + ], + "direction": [ + -0.04646139935246815, + -0.9989200860780659 + ], + "joint": { + "0.03": { + "loss": 2.430555555555553, + "accepted": true + }, + "0.1": { + "loss": 0.7638888888888885, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 3.2638888888888875, + "accepted": true + }, + "w": { + "loss": 0.7638888888888885, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9999945827943773 + } + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 10, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "n_obs_spikes": 153, + "starts": { + "S1_both_wrong": { + "L0": 14.20138888888889, + "estimators": { + "surrogate": { + "grad": [ + -11840.13073365276, + -14782202.006742565 + ], + "direction": [ + 0.08864652743457038, + 0.9960631471818399 + ], + "joint": { + "0.03": { + "loss": 11.319444444444448, + "accepted": true + }, + "0.1": { + "loss": 5.625000000000009, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 12.743055555555557, + "accepted": true + }, + "w": { + "loss": 5.625000000000009, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -10.763888888889, + -11072.530864197573 + ], + "direction": [ + 0.10738929787670691, + 0.9942170480843446 + ], + "joint": { + "0.03": { + "loss": 11.319444444444448, + "accepted": true + }, + "0.1": { + "loss": 5.625000000000009, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 12.743055555555557, + "accepted": true + }, + "w": { + "loss": 5.625000000000009, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9998226502371377 + }, + "S2_gcal_wrong_only": { + "L0": 1.5624999999999987, + "estimators": { + "surrogate": { + "grad": [ + 1973.4876179930393, + -203111.24891706434 + ], + "direction": [ + -0.6620202207533817, + 0.7494859753948995 + ], + "joint": { + "0.03": { + "loss": 1.3888888888888882, + "accepted": true + }, + "0.1": { + "loss": 1.4583333333333317, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 1.6319444444444429, + "accepted": false + }, + "w": { + "loss": 2.812499999999997, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + 1.0416666666666696, + -2430.5555555555525 + ], + "direction": [ + -0.03893150188475965, + 0.9992418817088268 + ], + "joint": { + "0.03": { + "loss": 0.7986111111111107, + "accepted": true + }, + "0.1": { + "loss": 2.812499999999997, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 1.6319444444444429, + "accepted": false + }, + "w": { + "loss": 2.812499999999997, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.7746912178399841 + }, + "S3_w_wrong_only": { + "L0": 10.000000000000005, + "estimators": { + "surrogate": { + "grad": [ + -38946.39342545843, + -53021586.51246857 + ], + "direction": [ + 0.12151585593651514, + 0.992589490552875 + ], + "joint": { + "0.03": { + "loss": 7.291666666666673, + "accepted": true + }, + "0.1": { + "loss": 3.8194444444444424, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 8.02083333333334, + "accepted": true + }, + "w": { + "loss": 3.8194444444444424, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + -11.574074074074046, + -8912.037037037024 + ], + "direction": [ + 0.21155127705207208, + 0.9773668999805739 + ], + "joint": { + "0.03": { + "loss": 7.291666666666673, + "accepted": true + }, + "0.1": { + "loss": 3.8194444444444424, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 8.02083333333334, + "accepted": true + }, + "w": { + "loss": 3.8194444444444424, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.995830947840406 + } + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 10, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "n_obs_spikes": 59, + "starts": { + "S1_both_wrong": { + "L0": 5.173611111111119, + "estimators": { + "surrogate": { + "grad": [ + 8621.654929975783, + 10150036.685759857 + ], + "direction": [ + -0.09396255495010108, + -0.9955757320602231 + ], + "joint": { + "0.03": { + "loss": 3.298611111111108, + "accepted": true + }, + "0.1": { + "loss": 1.3888888888888873, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 4.548611111111117, + "accepted": true + }, + "w": { + "loss": 1.3888888888888873, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 4.513888888888928, + 5748.456790123469 + ], + "direction": [ + -0.08691812722563325, + -0.996215458201482 + ], + "joint": { + "0.03": { + "loss": 3.298611111111108, + "accepted": true + }, + "0.1": { + "loss": 1.3888888888888873, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 4.548611111111117, + "accepted": true + }, + "w": { + "loss": 1.3888888888888873, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9999749833942495 + }, + "S2_gcal_wrong_only": { + "L0": 1.0069444444444435, + "estimators": { + "surrogate": { + "grad": [ + 629.9482532243901, + 872285.4636874392 + ], + "direction": [ + -0.08990707296639755, + -0.9959501585072493 + ], + "joint": { + "0.03": { + "loss": 0.2083333333333334, + "accepted": true + }, + "0.1": { + "loss": 1.6666666666666654, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.8333333333333326, + "accepted": true + }, + "w": { + "loss": 1.6666666666666654, + "accepted": false + } + } + }, + "crn_fd": { + "grad": [ + 0.6944444444444431, + 1562.4999999999977 + ], + "direction": [ + -0.05547001962252288, + -0.9984603532054125 + ], + "joint": { + "0.03": { + "loss": 0.2083333333333334, + "accepted": true + }, + "0.1": { + "loss": 1.6666666666666654, + "accepted": false + } + }, + "coord": { + "gCaL": { + "loss": 0.8333333333333326, + "accepted": true + }, + "w": { + "loss": 1.6666666666666654, + "accepted": false + } + } + } + }, + "cosine_sur_vs_fd": 0.9994038941397845 + }, + "S3_w_wrong_only": { + "L0": 3.541666666666665, + "estimators": { + "surrogate": { + "grad": [ + 6209.944694689961, + 10703896.399680983 + ], + "direction": [ + -0.04507748507695464, + -0.9989834935268636 + ], + "joint": { + "0.03": { + "loss": 2.0486111111111085, + "accepted": true + }, + "0.1": { + "loss": 0.8680555555555547, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 3.2638888888888844, + "accepted": true + }, + "w": { + "loss": 0.8680555555555547, + "accepted": true + } + } + }, + "crn_fd": { + "grad": [ + 3.9682539682540114, + 5555.555555555567 + ], + "direction": [ + -0.05547001962252341, + -0.9984603532054125 + ], + "joint": { + "0.03": { + "loss": 2.0486111111111085, + "accepted": true + }, + "0.1": { + "loss": 0.8680555555555547, + "accepted": true + } + }, + "coord": { + "gCaL": { + "loss": 3.2638888888888844, + "accepted": true + }, + "w": { + "loss": 0.8680555555555547, + "accepted": true + } + } + } + }, + "cosine_sur_vs_fd": 0.9999458607749618 + } + } + } + ], + "gate": { + "surrogate": { + "joint_accepted_fraction": 0.8333333333333334, + "coord_accepted_fraction": { + "gCaL": 0.9583333333333334, + "w": 0.6666666666666666 + }, + "nonfinite": 0, + "per_start": { + "S1_both_wrong": { + "joint_acc": 16, + "joint_tot": 16, + "coord_acc": { + "gCaL": 8, + "w": 8 + }, + "coord_tot": { + "gCaL": 8, + "w": 8 + }, + "nonfinite": 0 + }, + "S2_gcal_wrong_only": { + "joint_acc": 8, + "joint_tot": 16, + "coord_acc": { + "gCaL": 7, + "w": 0 + }, + "coord_tot": { + "gCaL": 8, + "w": 8 + }, + "nonfinite": 0 + }, + "S3_w_wrong_only": { + "joint_acc": 16, + "joint_tot": 16, + "coord_acc": { + "gCaL": 8, + "w": 8 + }, + "coord_tot": { + "gCaL": 8, + "w": 8 + }, + "nonfinite": 0 + } + }, + "joint_pass": true, + "gcal_pass": true + }, + "crn_fd": { + "joint_accepted_fraction": 0.8333333333333334, + "coord_accepted_fraction": { + "gCaL": 0.9166666666666666, + "w": 0.6666666666666666 + }, + "nonfinite": 0, + "per_start": { + "S1_both_wrong": { + "joint_acc": 16, + "joint_tot": 16, + "coord_acc": { + "gCaL": 8, + "w": 8 + }, + "coord_tot": { + "gCaL": 8, + "w": 8 + }, + "nonfinite": 0 + }, + "S2_gcal_wrong_only": { + "joint_acc": 8, + "joint_tot": 16, + "coord_acc": { + "gCaL": 6, + "w": 0 + }, + "coord_tot": { + "gCaL": 8, + "w": 8 + }, + "nonfinite": 0 + }, + "S3_w_wrong_only": { + "joint_acc": 16, + "joint_tot": 16, + "coord_acc": { + "gCaL": 8, + "w": 8 + }, + "coord_tot": { + "gCaL": 8, + "w": 8 + }, + "nonfinite": 0 + } + }, + "joint_pass": true, + "gcal_pass": true + } + }, + "cosine_sur_vs_fd": { + "mean": 0.9039791458985963, + "min": -0.9996991472413795 + } + }, + "stageB": { + "runs": [ + { + "case": 0, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.17361111111111116, + "best_loss": 0.868055555555555, + "within_eps": false, + "theta_best": [ + 1.0174526894268923, + 0.011357550714729624 + ], + "rel_err": [ + 0.32169820704873847, + 0.03250461042996585 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.6944444444444445, + "nonfinite": 0, + "hard_evals": 56, + "wall_s": 81.4 + }, + { + "case": 0, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.17361111111111116, + "best_loss": 0.868055555555555, + "within_eps": false, + "theta_best": [ + 1.010178239491474, + 0.011361972129434659 + ], + "rel_err": [ + 0.32654784033901735, + 0.032906557221332684 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.6944444444444445, + "nonfinite": 0, + "hard_evals": 66, + "wall_s": 60.0 + }, + { + "case": 0, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.17361111111111116, + "best_loss": 0.868055555555555, + "within_eps": false, + "theta_best": [ + 1.0033159594642416, + 0.011370972204449778 + ], + "rel_err": [ + 0.3311226936905056, + 0.033724745859070815 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.6944444444444444, + "nonfinite": 0, + "hard_evals": 60, + "wall_s": 90.3 + }, + { + "case": 0, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.17361111111111116, + "best_loss": 0.868055555555555, + "within_eps": false, + "theta_best": [ + 0.9995834526615446, + 0.01137281663537103 + ], + "rel_err": [ + 0.33361103155897026, + 0.033892421397366354 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.6944444444444445, + "nonfinite": 0, + "hard_evals": 70, + "wall_s": 62.8 + }, + { + "case": 1, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.001, + "best_loss": 0.03472222222222223, + "within_eps": false, + "theta_best": [ + 0.9932374522301407, + 0.007819227407696378 + ], + "rel_err": [ + 0.4189106460430582, + 0.02259657403795275 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.16666666666666666, + "nonfinite": 0, + "hard_evals": 52, + "wall_s": 79.8 + }, + { + "case": 1, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.001, + "best_loss": 0.03472222222222223, + "within_eps": false, + "theta_best": [ + 0.9790331460091581, + 0.007823135256261386 + ], + "rel_err": [ + 0.39861878001308315, + 0.02210809296732683 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.16666666666666666, + "nonfinite": 0, + "hard_evals": 61, + "wall_s": 54.6 + }, + { + "case": 1, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.001, + "best_loss": 0.03472222222222223, + "within_eps": false, + "theta_best": [ + 0.9980408513546359, + 0.007824860268681124 + ], + "rel_err": [ + 0.4257726447923371, + 0.02189246641485956 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.16666666666666666, + "nonfinite": 0, + "hard_evals": 59, + "wall_s": 86.4 + }, + { + "case": 1, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.001, + "best_loss": 0.03472222222222223, + "within_eps": false, + "theta_best": [ + 1.021695795718748, + 0.007779930147878711 + ], + "rel_err": [ + 0.4595654224553544, + 0.027508731515161144 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.16666666666666666, + "nonfinite": 0, + "hard_evals": 68, + "wall_s": 61.5 + }, + { + "case": 2, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.03472222222222223, + "best_loss": 1.319444444444443, + "within_eps": false, + "theta_best": [ + 1.016856432622289, + 0.011232648886819824 + ], + "rel_err": [ + 0.3220957115851406, + 0.02114989880180222 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 1.1111111111111105, + "nonfinite": 0, + "hard_evals": 55, + "wall_s": 81.1 + }, + { + "case": 2, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.03472222222222223, + "best_loss": 0.7986111111111107, + "within_eps": false, + "theta_best": [ + 1.0425569125461829, + 0.011572081194247217 + ], + "rel_err": [ + 0.30496205830254475, + 0.052007381295201605 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.5555555555555557, + "nonfinite": 0, + "hard_evals": 59, + "wall_s": 53.8 + }, + { + "case": 2, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.03472222222222223, + "best_loss": 2.2916666666666643, + "within_eps": false, + "theta_best": [ + 1.0, + 0.011 + ], + "rel_err": [ + 0.3333333333333333, + 0.0 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 1.444444444444443, + "nonfinite": 0, + "hard_evals": 61, + "wall_s": 88.3 + }, + { + "case": 2, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.03472222222222223, + "best_loss": 0.9027777777777773, + "within_eps": false, + "theta_best": [ + 1.0137899319310057, + 0.011592080776270249 + ], + "rel_err": [ + 0.3241400453793295, + 0.05382552511547723 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.4722222222222224, + "nonfinite": 0, + "hard_evals": 67, + "wall_s": 61.2 + }, + { + "case": 3, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.001, + "best_loss": 0.3472222222222223, + "within_eps": false, + "theta_best": [ + 0.9889671020201137, + 0.007880308093584407 + ], + "rel_err": [ + 0.4128101457430197, + 0.01496148830194919 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.2222222222222222, + "nonfinite": 0, + "hard_evals": 54, + "wall_s": 81.7 + }, + { + "case": 3, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.001, + "best_loss": 0.13888888888888892, + "within_eps": false, + "theta_best": [ + 0.9799465738517067, + 0.007798543512857444 + ], + "rel_err": [ + 0.39992367693100966, + 0.025182060892819543 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.1111111111111111, + "nonfinite": 0, + "hard_evals": 63, + "wall_s": 57.9 + }, + { + "case": 3, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.001, + "best_loss": 0.24305555555555558, + "within_eps": false, + "theta_best": [ + 0.9972994782453553, + 0.007825458673143765 + ], + "rel_err": [ + 0.42471354035050757, + 0.02181766585702938 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.1111111111111111, + "nonfinite": 0, + "hard_evals": 59, + "wall_s": 85.6 + }, + { + "case": 3, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.001, + "best_loss": 0.27777777777777785, + "within_eps": false, + "theta_best": [ + 0.9973237602454842, + 0.007718285481663006 + ], + "rel_err": [ + 0.4247482289221204, + 0.03521431479212422 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.08333333333333333, + "nonfinite": 0, + "hard_evals": 70, + "wall_s": 70.1 + }, + { + "case": 4, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.24305555555555558, + "best_loss": 0.625, + "within_eps": false, + "theta_best": [ + 1.0179833467240678, + 0.011442141067648104 + ], + "rel_err": [ + 0.32134443551728814, + 0.04019464251346402 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.8333333333333331, + "nonfinite": 0, + "hard_evals": 54, + "wall_s": 86.5 + }, + { + "case": 4, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.24305555555555558, + "best_loss": 0.625, + "within_eps": false, + "theta_best": [ + 1.0330069128023789, + 0.011425268898230246 + ], + "rel_err": [ + 0.31132872479841406, + 0.03866080893002242 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.7222222222222223, + "nonfinite": 0, + "hard_evals": 65, + "wall_s": 61.5 + }, + { + "case": 4, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.24305555555555558, + "best_loss": 0.5555555555555557, + "within_eps": false, + "theta_best": [ + 1.0034756800257272, + 0.011455705117832607 + ], + "rel_err": [ + 0.33101621331618186, + 0.04142773798478248 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.8333333333333331, + "nonfinite": 0, + "hard_evals": 58, + "wall_s": 102.1 + }, + { + "case": 4, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.24305555555555558, + "best_loss": 0.7291666666666665, + "within_eps": false, + "theta_best": [ + 1.0329571998023033, + 0.011493249449398217 + ], + "rel_err": [ + 0.3113618667984645, + 0.044840859036201615 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.6944444444444445, + "nonfinite": 0, + "hard_evals": 68, + "wall_s": 64.7 + }, + { + "case": 5, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.06944444444444445, + "best_loss": 0.10416666666666666, + "within_eps": false, + "theta_best": [ + 0.98961498184151, + 0.007821421369496146 + ], + "rel_err": [ + 0.41373568834501434, + 0.0223223288129818 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.27777777777777785, + "nonfinite": 0, + "hard_evals": 52, + "wall_s": 85.6 + }, + { + "case": 5, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.06944444444444445, + "best_loss": 0.20833333333333337, + "within_eps": false, + "theta_best": [ + 1.0025211256897761, + 0.007790324383936761 + ], + "rel_err": [ + 0.43217303669968027, + 0.026209452007904865 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.1388888888888889, + "nonfinite": 0, + "hard_evals": 64, + "wall_s": 63.9 + }, + { + "case": 5, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.06944444444444445, + "best_loss": 0.10416666666666666, + "within_eps": false, + "theta_best": [ + 0.9976216373592329, + 0.007815440076427161 + ], + "rel_err": [ + 0.42517376765604714, + 0.02306999044660485 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.27777777777777785, + "nonfinite": 0, + "hard_evals": 55, + "wall_s": 93.4 + }, + { + "case": 5, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.06944444444444445, + "best_loss": 0.20833333333333337, + "within_eps": false, + "theta_best": [ + 1.0220478583383055, + 0.007767864751516597 + ], + "rel_err": [ + 0.4600683690547222, + 0.029016906060425344 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.1388888888888889, + "nonfinite": 0, + "hard_evals": 67, + "wall_s": 67.8 + }, + { + "case": 6, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.06944444444444442, + "best_loss": 1.5624999999999987, + "within_eps": false, + "theta_best": [ + 1.0155346421907887, + 0.010985938227649828 + ], + "rel_err": [ + 0.32297690520614086, + 0.001278342940924628 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 2.3888888888888866, + "nonfinite": 0, + "hard_evals": 57, + "wall_s": 96.8 + }, + { + "case": 6, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.06944444444444442, + "best_loss": 0.5208333333333331, + "within_eps": false, + "theta_best": [ + 1.0245019512723288, + 0.011367572325574235 + ], + "rel_err": [ + 0.3169986991517808, + 0.03341566596129419 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.6944444444444444, + "nonfinite": 0, + "hard_evals": 63, + "wall_s": 64.9 + }, + { + "case": 6, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.06944444444444442, + "best_loss": 1.4583333333333317, + "within_eps": false, + "theta_best": [ + 0.9359417642934567, + 0.011856116210522172 + ], + "rel_err": [ + 0.3760388238043622, + 0.07782874641110665 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.9999999999999994, + "nonfinite": 0, + "hard_evals": 59, + "wall_s": 97.5 + }, + { + "case": 6, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.06944444444444442, + "best_loss": 0.5902777777777776, + "within_eps": false, + "theta_best": [ + 0.9987031249407404, + 0.01137255885352186 + ], + "rel_err": [ + 0.33419791670617305, + 0.03386898668380556 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.9166666666666664, + "nonfinite": 0, + "hard_evals": 70, + "wall_s": 71.6 + }, + { + "case": 7, + "start": "S1_both_wrong", + "estimator": "surrogate", + "eps_dt_halving": 0.001, + "best_loss": 0.2083333333333334, + "within_eps": false, + "theta_best": [ + 0.9858612070925954, + 0.007795055613108783 + ], + "rel_err": [ + 0.40837315298942206, + 0.02561804836140216 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.1111111111111111, + "nonfinite": 0, + "hard_evals": 53, + "wall_s": 103.6 + }, + { + "case": 7, + "start": "S1_both_wrong", + "estimator": "crn_fd", + "eps_dt_halving": 0.001, + "best_loss": 0.2083333333333334, + "within_eps": false, + "theta_best": [ + 0.9864314400034259, + 0.007794519187758751 + ], + "rel_err": [ + 0.40918777143346563, + 0.025685101530156143 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.1111111111111111, + "nonfinite": 0, + "hard_evals": 63, + "wall_s": 66.8 + }, + { + "case": 7, + "start": "S2_gcal_wrong_only", + "estimator": "surrogate", + "eps_dt_halving": 0.001, + "best_loss": 0.2083333333333334, + "within_eps": false, + "theta_best": [ + 0.9972545115862486, + 0.0077674243167458355 + ], + "rel_err": [ + 0.42464930226606945, + 0.029071960406770583 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.08333333333333333, + "nonfinite": 0, + "hard_evals": 60, + "wall_s": 102.7 + }, + { + "case": 7, + "start": "S2_gcal_wrong_only", + "estimator": "crn_fd", + "eps_dt_halving": 0.001, + "best_loss": 0.3125000000000001, + "within_eps": false, + "theta_best": [ + 0.9981527076945506, + 0.007738125926358108 + ], + "rel_err": [ + 0.42593243956364374, + 0.032734259205236566 + ], + "gcal_recovered_20pct": false, + "w_recovered_20pct": true, + "heldout_mse": 0.05555555555555555, + "nonfinite": 0, + "hard_evals": 70, + "wall_s": 78.0 + } + ], + "summary": { + "surrogate": { + "n_runs": 16, + "gcal_recovered_20pct": 0, + "w_recovered_20pct": 16, + "within_eps": 0, + "nonfinite_runs": 0, + "median_heldout": 0.48611111111111116 + }, + "crn_fd": { + "n_runs": 16, + "gcal_recovered_20pct": 0, + "w_recovered_20pct": 16, + "within_eps": 0, + "nonfinite_runs": 0, + "median_heldout": 0.31944444444444453 + } + } + }, + "profile_case0": { + "g_axis": [ + 0.25, + 0.32987697769322355, + 0.43527528164806206, + 0.5743491774985175, + 0.7578582832551991, + 1.0, + 1.3195079107728946, + 1.5, + 1.7411011265922485, + 2.29739670999407, + 3.0314331330207964, + 4.0 + ], + "profile_gcal": [ + 3.194444444444441, + 2.9861111111111076, + 2.673611111111109, + 2.6388888888888866, + 2.291666666666664, + 1.6666666666666647, + 1.11111111111111, + 0.0, + 0.5208333333333334, + 3.194444444444441, + 2.7430555555555536, + 5.104166666666677 + ], + "eps": 0.17361111111111116, + "gcal_compat_interval": [ + 1.5, + 1.5 + ] + }, + "wall_s": 3131.0 +} \ No newline at end of file diff --git a/experiments/twin/evidence/g5_basin_full_20261018_191021.json b/experiments/twin/evidence/g5_basin_full_20261018_191021.json new file mode 100644 index 00000000..917c1c50 --- /dev/null +++ b/experiments/twin/evidence/g5_basin_full_20261018_191021.json @@ -0,0 +1,484 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "a433b78e7bd7e581", + "poolbuild.py": "2dd213214e166343", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "983f8bbcafebd67e", + "protocols.py": "404ea437741ea2ae", + "gradient_repair.py": "78812d0e94b7be75", + "deliverable_b.py": "8c16e479bcea843d", + "deliverable_c.py": "0c1896bdc5a2e5fc", + "pic_grid.py": "d900172c74ba1695", + "r2a_profiles.py": "ac4802b3bdfe1191", + "acceptance.py": "24200ac953a87395", + "myogen/simulator/jaxley/jax_models.py": "8b99fc4b69d735c7", + "myogen/simulator/jaxley/channels/nerlab_channels.py": "f6161ceaca04285e", + "myogen/simulator/jaxley/cells_api.py": "c6534ba58adaf3b1", + "myogen/simulator/jaxley/populations/motor_neurons.py": "736be1bb0c75807c" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float64 (jax_enable_x64=True)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations", + "g5_sha256_16": "759841b09b5acaa4", + "args": { + "smoke": false + }, + "run_tag": "full" + }, + "declared": { + "objective": "G4 v2 full-trace objective (dt 0.025)", + "profile": "continuous w via golden-section (10 evals), boundary bisection to 2 % gCaL resolution", + "tolerances": "eps = dt-halving at truth; t_repeat = rate-MSE of truth vs independent event realization", + "hybrid": "9-pt gCaL scan \u00d7 golden w, top-3 basins, surrogate refinement, ~220-eval budget", + "baseline": "Nelder-Mead log-space, same budget" + }, + "cases": [ + { + "case": { + "conn_seed": 0, + "event_offset": 0, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "eps": 0.17361111111111116, + "t_repeat": 1.9444444444444422, + "profile": { + "1.1250": 0.7291666666666664, + "1.2250": 0.6597222222222219, + "1.3338": 0.5902777777777776, + "1.4524": 0.10416666666666669, + "1.5000": 0.03472222222222222, + "1.5814": 0.1388888888888889, + "1.5984": 0.1388888888888889, + "1.6155": 0.41666666666666674, + "1.6502": 0.27777777777777785, + "1.7220": 0.27777777777777773, + "1.8750": 0.45138888888888895 + }, + "gcal_interval_vs_eps": [ + 1.4523687548277813, + 1.5983560814998021 + ], + "gcal_interval_vs_repeatability": [ + 1.125, + 1.875 + ], + "hybrid": { + "loss": 0.24305555555555564, + "theta": [ + 1.4142135623730951, + 0.011118395515585917 + ], + "evals": 180, + "gcal_rel_err": 0.05719095841793657, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.7986111111111105, + "theta": [ + 1.000126961183889, + 0.011432168896490827 + ], + "evals": 25, + "gcal_rel_err": 0.333248692544074, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 0, + "event_offset": 0, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "eps": 0.001, + "t_repeat": 0.5902777777777776, + "profile": { + "0.5250": 0.24305555555555558, + "0.5717": 0.1388888888888889, + "0.6225": 0.13888888888888892, + "0.6778": 0.0, + "0.6888": 0.03472222222222223, + "0.7000": 0.03472222222222223, + "0.7380": 0.06944444444444446, + "0.8036": 0.03472222222222223, + "0.8750": 0.06944444444444446 + }, + "gcal_interval_vs_eps": [ + 0.6777720855862979, + 0.6777720855862979 + ], + "gcal_interval_vs_repeatability": [ + 0.5249999999999999, + 0.875 + ], + "hybrid": { + "loss": 0.03472222222222223, + "theta": [ + 0.7071067811865475, + 0.008015212374091921 + ], + "evals": 180, + "gcal_rel_err": 0.010152544552210725, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.03472222222222223, + "theta": [ + 0.9999257840041433, + 0.007825445730210191 + ], + "evals": 25, + "gcal_rel_err": 0.42846540572020486, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 0, + "event_offset": 10, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "eps": 0.03472222222222223, + "t_repeat": 2.2222222222222197, + "profile": { + "1.1250": 0.868055555555555, + "1.2250": 0.45138888888888906, + "1.3338": 0.4166666666666668, + "1.4524": 0.27777777777777785, + "1.5000": 0.17361111111111113, + "1.5814": 0.034722222222222224, + "1.5984": 0.24305555555555555, + "1.6155": 0.34722222222222227, + "1.6502": 0.5208333333333334, + "1.7220": 0.48611111111111116, + "1.8750": 0.7638888888888884 + }, + "gcal_interval_vs_eps": [ + 1.58143624744078, + 1.58143624744078 + ], + "gcal_interval_vs_repeatability": [ + 1.125, + 1.875 + ], + "hybrid": { + "loss": 0.06944444444444446, + "theta": [ + 1.4142135623730951, + 0.011118395515585917 + ], + "evals": 180, + "gcal_rel_err": 0.05719095841793657, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.7638888888888886, + "theta": [ + 1.0001884943253232, + 0.011411152463632621 + ], + "evals": 27, + "gcal_rel_err": 0.33320767044978455, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 0, + "event_offset": 10, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "eps": 0.001, + "t_repeat": 0.9722222222222213, + "profile": { + "0.5250": 0.10416666666666669, + "0.5717": 0.034722222222222224, + "0.6225": 0.1388888888888889, + "0.6778": 0.0, + "0.7000": 0.0, + "0.7380": 0.13888888888888892, + "0.8036": 0.0, + "0.8122": 0.034722222222222224, + "0.8209": 0.034722222222222224, + "0.8385": 0.06944444444444445, + "0.8750": 0.034722222222222224 + }, + "gcal_interval_vs_eps": [ + 0.6777720855862979, + 0.8035876643973896 + ], + "gcal_interval_vs_repeatability": [ + 0.5249999999999999, + 0.875 + ], + "hybrid": { + "loss": 0.0, + "theta": [ + 0.7071067811865475, + 0.008015212374091921 + ], + "evals": 147, + "gcal_rel_err": 0.010152544552210725, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.13888888888888892, + "theta": [ + 0.9999218780516783, + 0.00776806029483447 + ], + "evals": 25, + "gcal_rel_err": 0.428459825788112, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 0, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "eps": 0.24305555555555558, + "t_repeat": 1.73611111111111, + "profile": { + "1.1250": 0.2430555555555555, + "1.2250": 0.17361111111111108, + "1.3338": 0.13888888888888887, + "1.4524": 0.10416666666666666, + "1.5000": 0.03472222222222222, + "1.5814": 0.06944444444444443, + "1.5984": 0.034722222222222224, + "1.6155": 0.34722222222222227, + "1.6502": 0.5208333333333331, + "1.7220": 0.2777777777777778, + "1.8750": 0.6597222222222219 + }, + "gcal_interval_vs_eps": [ + 1.125, + 1.5983560814998021 + ], + "gcal_interval_vs_repeatability": [ + 1.125, + 1.875 + ], + "hybrid": { + "loss": 0.2777777777777778, + "theta": [ + 1.4142135623730951, + 0.011118395515585917 + ], + "evals": 180, + "gcal_rel_err": 0.05719095841793657, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.5555555555555557, + "theta": [ + 1.0001594365367787, + 0.011458493877506824 + ], + "evals": 34, + "gcal_rel_err": 0.33322704230881417, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 0, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "eps": 0.06944444444444445, + "t_repeat": 0.9027777777777768, + "profile": { + "0.5250": 0.5208333333333333, + "0.5717": 0.13888888888888892, + "0.6225": 0.03472222222222223, + "0.6778": 0.0, + "0.7000": 0.13888888888888892, + "0.7380": 0.0, + "0.7459": 0.10416666666666666, + "0.7539": 0.13888888888888892, + "0.7701": 0.13888888888888892, + "0.8036": 0.10416666666666666, + "0.8750": 0.20833333333333337 + }, + "gcal_interval_vs_eps": [ + 0.6224563282857609, + 0.7380035821390306 + ], + "gcal_interval_vs_repeatability": [ + 0.5249999999999999, + 0.875 + ], + "hybrid": { + "loss": 0.0, + "theta": [ + 0.7064820271485027, + 0.007986042255466334 + ], + "evals": 162, + "gcal_rel_err": 0.009260038783575382, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.20833333333333337, + "theta": [ + 0.9999218780516783, + 0.00776806029483447 + ], + "evals": 27, + "gcal_rel_err": 0.428459825788112, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 10, + "theta_true": [ + 1.5, + 0.011 + ] + }, + "eps": 0.06944444444444442, + "t_repeat": 2.1180555555555545, + "profile": { + "1.1250": 0.5555555555555554, + "1.2250": 0.5902777777777776, + "1.3338": 0.24305555555555558, + "1.4524": 0.0, + "1.5000": 0.0, + "1.5200": 0.17361111111111108, + "1.5402": 0.34722222222222215, + "1.5814": 0.10416666666666666, + "1.7220": 0.4861111111111111, + "1.8750": 0.6944444444444442 + }, + "gcal_interval_vs_eps": [ + 1.4523687548277813, + 1.5 + ], + "gcal_interval_vs_repeatability": [ + 1.125, + 1.875 + ], + "hybrid": { + "loss": 0.3125, + "theta": [ + 1.4142135623730951, + 0.011118395515585917 + ], + "evals": 180, + "gcal_rel_err": 0.05719095841793657, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.5902777777777776, + "theta": [ + 1.0002500312526044, + 0.011390174666527093 + ], + "evals": 30, + "gcal_rel_err": 0.3331666458315971, + "gcal_recovered_20pct": false + } + }, + { + "case": { + "conn_seed": 1, + "event_offset": 10, + "theta_true": [ + 0.7, + 0.008 + ] + }, + "eps": 0.001, + "t_repeat": 1.1111111111111096, + "profile": { + "0.5250": 0.0, + "0.5717": 0.17361111111111116, + "0.6225": 0.17361111111111116, + "0.6778": 0.0, + "0.6888": 0.10416666666666666, + "0.7000": 0.13888888888888892, + "0.7380": 0.1041666666666667, + "0.8036": 0.06944444444444446, + "0.8750": 0.06944444444444446 + }, + "gcal_interval_vs_eps": [ + 0.5249999999999999, + 0.6777720855862979 + ], + "gcal_interval_vs_repeatability": [ + 0.5249999999999999, + 0.875 + ], + "hybrid": { + "loss": 0.17361111111111116, + "theta": [ + 0.7071067811865475, + 0.008015212374091921 + ], + "evals": 180, + "gcal_rel_err": 0.010152544552210725, + "gcal_recovered_20pct": true + }, + "nelder_mead": { + "loss": 0.2083333333333334, + "theta": [ + 0.9999218780516783, + 0.00776806029483447 + ], + "evals": 27, + "gcal_rel_err": 0.428459825788112, + "gcal_recovered_20pct": false + } + } + ], + "summary": { + "hybrid_gcal_recovered_20pct": 8, + "nm_gcal_recovered_20pct": 0, + "n_cases": 8, + "basin_survives_repeatability": 8.0 + }, + "wall_s": 3049.7 +} \ No newline at end of file diff --git a/experiments/twin/evidence/gradient_repair_x64_20261018_154132.json b/experiments/twin/evidence/gradient_repair_x64_20261018_154132.json new file mode 100644 index 00000000..92d3409b --- /dev/null +++ b/experiments/twin/evidence/gradient_repair_x64_20261018_154132.json @@ -0,0 +1,3529 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "a433b78e7bd7e581", + "poolbuild.py": "1c782bb1613e9e96", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "983f8bbcafebd67e", + "protocols.py": "404ea437741ea2ae", + "gradient_repair.py": "f3630bb1ad7f3d99", + "deliverable_b.py": "8c16e479bcea843d", + "deliverable_c.py": "0c1896bdc5a2e5fc", + "pic_grid.py": "d900172c74ba1695", + "r2a_profiles.py": "ac4802b3bdfe1191", + "acceptance.py": "24200ac953a87395", + "myogen/simulator/jaxley/jax_models.py": "8b99fc4b69d735c7", + "myogen/simulator/jaxley/channels/nerlab_channels.py": "f6161ceaca04285e", + "myogen/simulator/jaxley/cells_api.py": "c6534ba58adaf3b1", + "myogen/simulator/jaxley/populations/motor_neurons.py": "736be1bb0c75807c" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float64 (jax_enable_x64=True)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations" + }, + "x64": true, + "regime": "current", + "g1": { + "ok": true, + "rows": [ + { + "regime": "current", + "dt_ms": 0.1, + "horizon_s": 0.1, + "n_spikes": 2, + "grad_rev": [ + 20.170857576967315, + 582.1176000943616 + ], + "grad_fwd": [ + 20.170857577013145, + 582.1176000907582 + ], + "fd_converged_value": [ + 20.18896472222309, + 582.1025396812729 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 1.4563376426481556 + ], + [ + 0.01, + 9.992142043070729 + ], + [ + 0.001, + 20.059616631276533 + ], + [ + 0.0001, + 20.16974137694305 + ], + [ + 1e-05, + 20.170848519640572 + ], + [ + 1e-06, + 20.170933554553017 + ], + [ + 1e-07, + 20.171619095066262 + ], + [ + 1e-08, + 20.18896472222309 + ] + ], + "1": [ + [ + 0.1, + -0.6279048148210364 + ], + [ + 0.01, + 6.889046553699885 + ], + [ + 0.001, + -35.18354316521499 + ], + [ + 0.0001, + 555.9027045605802 + ], + [ + 1e-05, + 581.8543978884704 + ], + [ + 1e-06, + 582.1146022464419 + ], + [ + 1e-07, + 582.1184253029799 + ], + [ + 1e-08, + 582.1025396812729 + ] + ] + }, + "rel_rev_vs_fdconv": 0.000896883297628651, + "cosine": 0.9999999994891341, + "rev_vs_fwd_rel": 6.190189074961024e-12, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.1, + "horizon_s": 0.5, + "n_spikes": 10, + "grad_rev": [ + 23684748127560.945, + 673217265990022.1 + ], + "grad_fwd": [ + 23684748127644.44, + 673217266020431.1 + ], + "fd_converged_value": [ + null, + null + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.049469674234092764 + ], + [ + 0.01, + 0.7968805244595423 + ], + [ + 0.001, + 11.535788341902276 + ], + [ + 0.0001, + -160.53403568905722 + ], + [ + 1e-05, + 69.36357146836158 + ], + [ + 1e-06, + 10183.228088496144 + ], + [ + 1e-07, + 25752.28873603841 + ], + [ + 1e-08, + -593595.7526151725 + ] + ], + "1": [ + [ + 0.1, + -0.72823223862577 + ], + [ + 0.01, + -1.4552666425637462 + ], + [ + 0.001, + 7.4610239012598 + ], + [ + 0.0001, + -21.756268557640013 + ], + [ + 1e-05, + -737.0068854165801 + ], + [ + 1e-06, + 2166.0986135483286 + ], + [ + 1e-07, + 256474.75381728847 + ], + [ + 1e-08, + -483883.42780718574 + ] + ] + }, + "rel_rev_vs_fdconv": null, + "cosine": null, + "rev_vs_fwd_rel": 4.5169667408793304e-11, + "verdict": "FD_UNCONVERGED" + }, + { + "regime": "current", + "dt_ms": 0.1, + "horizon_s": 1.0, + "n_spikes": 19, + "grad_rev": [ + -6.057275089864521e+27, + -1.721724948643267e+29 + ], + "grad_fwd": [ + -6.05727509052157e+27, + -1.7217249488795974e+29 + ], + "fd_converged_value": [ + null, + null + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.241019732953216 + ], + [ + 0.01, + 0.16556273052943382 + ], + [ + 0.001, + 6.271769229136304 + ], + [ + 0.0001, + -188.5285715184448 + ], + [ + 1e-05, + 1542.9114218868276 + ], + [ + 1e-06, + 9843.044968640414 + ], + [ + 1e-07, + 40814.16918704228 + ], + [ + 1e-08, + -611760.6058893444 + ] + ], + "1": [ + [ + 0.1, + -0.6816929189198362 + ], + [ + 0.01, + -0.08737777284135717 + ], + [ + 0.001, + -2.1085942716689843 + ], + [ + 0.0001, + 114.73381633440916 + ], + [ + 1e-05, + -1891.4077605417745 + ], + [ + 1e-06, + 5236.391061857582 + ], + [ + 1e-07, + 226547.64388247984 + ], + [ + 1e-08, + 1113072.1285337764 + ] + ] + }, + "rel_rev_vs_fdconv": null, + "cosine": null, + "rev_vs_fwd_rel": 1.372636557778797e-10, + "verdict": "FD_UNCONVERGED" + }, + { + "regime": "current", + "dt_ms": 0.05, + "horizon_s": 0.1, + "n_spikes": 2, + "grad_rev": [ + 0.9996979738333283, + -4.778281179305066 + ], + "grad_fwd": [ + 0.9996979738362306, + -4.7782811792260675 + ], + "fd_converged_value": [ + 1.0005368977772378, + -4.778118345427629 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.17933905917190351 + ], + [ + 0.01, + 0.34586271702425897 + ], + [ + 0.001, + 0.989441601401353 + ], + [ + 0.0001, + 0.9995949667374049 + ], + [ + 1e-05, + 0.9996963894032261 + ], + [ + 1e-06, + 0.9996930447897512 + ], + [ + 1e-07, + 0.9997272787387601 + ], + [ + 1e-08, + 1.0005368977772378 + ] + ], + "1": [ + [ + 0.1, + -0.6147593161234033 + ], + [ + 0.01, + -1.0486936308351957 + ], + [ + 0.001, + 0.5565688685669201 + ], + [ + 0.0001, + -4.719195580209501 + ], + [ + 1e-05, + -4.777689548607356 + ], + [ + 1e-06, + -4.778271284422431 + ], + [ + 1e-07, + -4.778317395093268 + ], + [ + 1e-08, + -4.778118345427629 + ] + ] + }, + "rel_rev_vs_fdconv": 0.0008384737692065384, + "cosine": 0.9999999846808953, + "rev_vs_fwd_rel": 1.6532752761653758e-11, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.05, + "horizon_s": 0.5, + "n_spikes": 10, + "grad_rev": [ + -545284.389820573, + 1283775.5295344184 + ], + "grad_fwd": [ + -545284.389850123, + 1283775.530157801 + ], + "fd_converged_value": [ + null, + -0.928515679316888 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + -0.1182316550670226 + ], + [ + 0.01, + -0.10526805740891376 + ], + [ + 0.001, + 2.1841906856092663 + ], + [ + 0.0001, + 11.893492800751915 + ], + [ + 1e-05, + -25.166380386654193 + ], + [ + 1e-06, + -2058.873005280404 + ], + [ + 1e-07, + 13642.52764285645 + ], + [ + 1e-08, + -47183.00608841375 + ] + ], + "1": [ + [ + 0.1, + -0.9554730884992413 + ], + [ + 0.01, + -0.928515679316888 + ], + [ + 0.001, + -0.5830049855140018 + ], + [ + 0.0001, + -0.35586166117873574 + ], + [ + 1e-05, + -175.52420974364932 + ], + [ + 1e-06, + 2241.944556742759 + ], + [ + 1e-07, + 444.6066852548114 + ], + [ + 1e-08, + 44240.49205340452 + ] + ] + }, + "rel_rev_vs_fdconv": null, + "cosine": null, + "rev_vs_fwd_rel": 4.85585379468476e-10, + "verdict": "FD_UNCONVERGED" + }, + { + "regime": "current", + "dt_ms": 0.05, + "horizon_s": 1.0, + "n_spikes": 19, + "grad_rev": [ + 67981295958824.97, + -160050221265774.72 + ], + "grad_fwd": [ + 67981295965708.14, + -160050221325749.5 + ], + "fd_converged_value": [ + null, + null + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.005291547964585419 + ], + [ + 0.01, + 0.05171781597343994 + ], + [ + 0.001, + 3.3896650727838207 + ], + [ + 0.0001, + -3.970932961703255 + ], + [ + 1e-05, + -72.12778124845443 + ], + [ + 1e-06, + -0.4470075332108081 + ], + [ + 1e-07, + -17240.113247112276 + ], + [ + 1e-08, + -155286.87302577283 + ] + ], + "1": [ + [ + 0.1, + -0.835068286752616 + ], + [ + 0.01, + -0.7653498778897649 + ], + [ + 0.001, + 1.0218754622295378 + ], + [ + 0.0001, + 9.435643652988546 + ], + [ + 1e-05, + -433.29536190155335 + ], + [ + 1e-06, + 4268.6818475390355 + ], + [ + 1e-07, + -7369.1741296944 + ], + [ + 1e-08, + 152821.97243831062 + ] + ] + }, + "rel_rev_vs_fdconv": null, + "cosine": null, + "rev_vs_fwd_rel": 3.7472476297257717e-10, + "verdict": "FD_UNCONVERGED" + }, + { + "regime": "current", + "dt_ms": 0.025, + "horizon_s": 0.1, + "n_spikes": 2, + "grad_rev": [ + 0.2570947496948042, + -0.6945870976149876 + ], + "grad_fwd": [ + 0.25709474974675794, + -0.6945870976623496 + ], + "fd_converged_value": [ + 0.25709239181992416, + -0.6945755970377832 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.21825103456014183 + ], + [ + 0.01, + 0.21313254596613262 + ], + [ + 0.001, + 0.2552516123204285 + ], + [ + 0.0001, + 0.2570760470099742 + ], + [ + 1e-05, + 0.2570946635138682 + ], + [ + 1e-06, + 0.25709400564011275 + ], + [ + 1e-07, + 0.2570979074079105 + ], + [ + 1e-08, + 0.25709239181992416 + ] + ], + "1": [ + [ + 0.1, + -0.5506066146496513 + ], + [ + 0.01, + -0.49586431040129497 + ], + [ + 0.001, + -0.6756246655603704 + ], + [ + 0.0001, + -0.6943783734048026 + ], + [ + 1e-05, + -0.6945849652773005 + ], + [ + 1e-06, + -0.6945869941432647 + ], + [ + 1e-07, + -0.694579513904614 + ], + [ + 1e-08, + -0.6945755970377832 + ] + ] + }, + "rel_rev_vs_fdconv": 1.6557704090808884e-05, + "cosine": 0.999999999997109, + "rev_vs_fwd_rel": 2.0208004697014557e-10, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.025, + "horizon_s": 0.5, + "n_spikes": 10, + "grad_rev": [ + -0.08974592374980707, + -0.9610277874303975 + ], + "grad_fwd": [ + -0.08974592381181193, + -0.9610277873983134 + ], + "fd_converged_value": [ + -0.08974847531817431, + -0.961010471201007 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + -0.11213761314509796 + ], + [ + 0.01, + -0.11403701302636904 + ], + [ + 0.001, + -0.13644490329767223 + ], + [ + 0.0001, + -0.05623849154368088 + ], + [ + 1e-05, + -0.08928754953174688 + ], + [ + 1e-06, + -0.08974132370553889 + ], + [ + 1e-07, + -0.08974777365722275 + ], + [ + 1e-08, + -0.08974847531817431 + ] + ], + "1": [ + [ + 0.1, + -0.9532516680113012 + ], + [ + 0.01, + -0.9509157647514144 + ], + [ + 0.001, + -0.979895046291368 + ], + [ + 0.0001, + -0.9472989220071781 + ], + [ + 1e-05, + -0.9608697202345694 + ], + [ + 1e-06, + -0.9610261129111564 + ], + [ + 1e-07, + -0.9610268847382031 + ], + [ + 1e-08, + -0.961010471201007 + ] + ] + }, + "rel_rev_vs_fdconv": 2.8430214086643894e-05, + "cosine": 0.9999999999907536, + "rev_vs_fwd_rel": 6.908933235797546e-10, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.025, + "horizon_s": 1.0, + "n_spikes": 19, + "grad_rev": [ + 0.2933118010400612, + -0.692174931850049 + ], + "grad_fwd": [ + 0.29331180081654856, + -0.6921749319608395 + ], + "fd_converged_value": [ + 0.29333220297189655, + -0.6921699657880254 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.05251829683225573 + ], + [ + 0.01, + 0.057898386612365016 + ], + [ + 0.001, + 0.10202533894876353 + ], + [ + 0.0001, + 0.19053328467855124 + ], + [ + 1e-05, + 0.29931715008757465 + ], + [ + 1e-06, + 0.29337318085964625 + ], + [ + 1e-07, + 0.29331111761621287 + ], + [ + 1e-08, + 0.29333220297189655 + ] + ], + "1": [ + [ + 0.1, + -0.8324042521198516 + ], + [ + 0.01, + -0.8301616369996268 + ], + [ + 0.001, + -0.8224772182074958 + ], + [ + 0.0001, + -0.691569300155237 + ], + [ + 1e-05, + -0.690372413636453 + ], + [ + 1e-06, + -0.692157669845983 + ], + [ + 1e-07, + -0.6921748152421969 + ], + [ + 1e-08, + -0.6921699657880254 + ] + ] + }, + "rel_rev_vs_fdconv": 6.955230836789169e-05, + "cosine": 0.9999999996200749, + "rev_vs_fwd_rel": 7.620308907049014e-10, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.0125, + "horizon_s": 0.1, + "n_spikes": 2, + "grad_rev": [ + 0.2258525308007354, + -0.5280072304542877 + ], + "grad_fwd": [ + 0.22585253084949897, + -0.5280072305123952 + ], + "fd_converged_value": [ + 0.2258434328439307, + -0.5279880532782499 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.2261839546918054 + ], + [ + 0.01, + 0.22616308097846982 + ], + [ + 0.001, + 0.22590517508191965 + ], + [ + 0.0001, + 0.22585308935063608 + ], + [ + 1e-05, + 0.225852518465075 + ], + [ + 1e-06, + 0.22585240611050494 + ], + [ + 1e-07, + 0.22585393111285157 + ], + [ + 1e-08, + 0.2258434328439307 + ] + ], + "1": [ + [ + 0.1, + -0.5269184241031999 + ], + [ + 0.01, + -0.5268055494547674 + ], + [ + 0.001, + -0.5274612033030479 + ], + [ + 0.0001, + -0.5280001103091791 + ], + [ + 1e-05, + -0.528007124689367 + ], + [ + 1e-06, + -0.5280074937275003 + ], + [ + 1e-07, + -0.5280046622146983 + ], + [ + 1e-08, + -0.5279880532782499 + ] + ] + }, + "rel_rev_vs_fdconv": 4.0284354032839567e-05, + "cosine": 0.9999999999989736, + "rev_vs_fwd_rel": 2.1590892219208136e-10, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.0125, + "horizon_s": 0.5, + "n_spikes": 10, + "grad_rev": [ + -0.10226547524392135, + -0.9397110125957779 + ], + "grad_fwd": [ + -0.10226547520170036, + -0.9397110126996584 + ], + "fd_converged_value": [ + -0.10226983704342274, + -0.9397115086073882 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + -0.09981171878767903 + ], + [ + 0.01, + -0.10225066314069053 + ], + [ + 0.001, + -0.10228914832932645 + ], + [ + 0.0001, + -0.10217191917405444 + ], + [ + 1e-05, + -0.10226437758831251 + ], + [ + 1e-06, + -0.10226548852187989 + ], + [ + 1e-07, + -0.10226560931414497 + ], + [ + 1e-08, + -0.10226983704342274 + ] + ], + "1": [ + [ + 0.1, + -0.9391431838333908 + ], + [ + 0.01, + -0.9397101331304114 + ], + [ + 0.001, + -0.939762568806124 + ], + [ + 0.0001, + -0.9396710743114767 + ], + [ + 1e-05, + -0.9397105549702188 + ], + [ + 1e-06, + -0.9397110121156516 + ], + [ + 1e-07, + -0.9397111977449413 + ], + [ + 1e-08, + -0.9397115086073882 + ] + ] + }, + "rel_rev_vs_fdconv": 4.2649911523160674e-05, + "cosine": 0.9999999999897372, + "rev_vs_fwd_rel": 4.128567388878437e-10, + "verdict": "PASS" + }, + { + "regime": "current", + "dt_ms": 0.0125, + "horizon_s": 1.0, + "n_spikes": 19, + "grad_rev": [ + 0.10249916517547925, + -0.7865028004732653 + ], + "grad_fwd": [ + 0.1024991647334322, + -0.7865028006069417 + ], + "fd_converged_value": [ + 0.10249792126160173, + -0.7865070372758964 + ], + "fd_sweeps": { + "0": [ + [ + 0.1, + 0.10537356808777432 + ], + [ + 0.01, + 0.10213966807652852 + ], + [ + 0.001, + 0.10216085794745311 + ], + [ + 0.0001, + 0.10223511699436472 + ], + [ + 1e-05, + 0.10249420858698953 + ], + [ + 1e-06, + 0.1024994258358447 + ], + [ + 1e-07, + 0.1024992180020945 + ], + [ + 1e-08, + 0.10249792126160173 + ] + ], + "1": [ + [ + 0.1, + -0.787330894996936 + ], + [ + 0.01, + -0.7867915397863356 + ], + [ + 0.001, + -0.7868031045523693 + ], + [ + 0.0001, + -0.786645274004627 + ], + [ + 1e-05, + -0.7865047577659822 + ], + [ + 1e-06, + -0.7865028175402244 + ], + [ + 1e-07, + -0.786503093763713 + ], + [ + 1e-08, + -0.7865070372758964 + ] + ] + }, + "rel_rev_vs_fdconv": 1.213599126892075e-05, + "cosine": 0.9999999999974788, + "rev_vs_fwd_rel": 4.312689307180243e-09, + "verdict": "PASS" + } + ], + "passed_cases": [ + [ + 0.1, + 0.1 + ], + [ + 0.05, + 0.1 + ], + [ + 0.025, + 0.1 + ], + [ + 0.025, + 0.5 + ], + [ + 0.025, + 1.0 + ], + [ + 0.0125, + 0.1 + ], + [ + 0.0125, + 0.5 + ], + [ + 0.0125, + 1.0 + ] + ], + "failed_cases": [], + "unresolved_cases": [ + [ + 0.1, + 0.5 + ], + [ + 0.1, + 1.0 + ], + [ + 0.05, + 0.5 + ], + [ + 0.05, + 1.0 + ] + ], + "required_cases_dt_le_0p025": [ + [ + 0.025, + 0.1 + ], + [ + 0.025, + 0.5 + ], + [ + 0.025, + 1.0 + ], + [ + 0.0125, + 0.1 + ], + [ + 0.0125, + 0.5 + ], + [ + 0.0125, + 1.0 + ] + ], + "qualified_envelope": [ + [ + 0.0125, + 0.1 + ], + [ + 0.0125, + 0.5 + ], + [ + 0.0125, + 1.0 + ], + [ + 0.025, + 0.1 + ], + [ + 0.025, + 0.5 + ], + [ + 0.025, + 1.0 + ], + [ + 0.05, + 0.1 + ], + [ + 0.1, + 0.1 + ] + ], + "scope": "regime=current; claims limited to these operating points/timesteps; FD_UNCONVERGED combos are unresolvable by local FD, not passes", + "runtime_s": 242.0 + }, + "g2": { + "ok": false, + "n_variants": 10, + "per_coordinate": { + "gCaL": { + "agree": 8, + "wrong": 1, + "unresolved": 1 + }, + "w_exc": { + "agree": 9, + "wrong": 1, + "unresolved": 0 + } + }, + "surrogate_grad_magnitude_range": [ + 8955944.572652703, + 14751936385.32629 + ], + "magnitude_note": "sign agreement does NOT establish usefulness \u2014 magnitudes remain enormous and must be handled by the optimizer (G4 question)", + "rows": [ + { + "variant": "beta=5.0", + "surrogate_grad": [ + -6137678.057117179, + -5233561892.557893 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "beta=10.0", + "surrogate_grad": [ + -3061214.9652824383, + -2610277782.364797 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "beta=20.0", + "surrogate_grad": [ + -1524809.1714994127, + -1300194013.6110883 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "vth=49.0", + "surrogate_grad": [ + -4536076.406051416, + -3867903984.547015 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "vth=50.0", + "surrogate_grad": [ + -3061214.9652824383, + -2610277782.364797 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "vth=51.0", + "surrogate_grad": [ + 17299593.807629004, + 14751936385.32629 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "wrong", + "wrong" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "dt=0.1", + "surrogate_grad": [ + -3061214.9652824383, + -2610277782.364797 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "dt=0.05", + "surrogate_grad": [ + -9547.346982762745, + -8955944.572652703 + ], + "crn_fd_ref": [ + 0.0, + -4732.510288065846 + ], + "classification": [ + "unresolved", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "evshift=0", + "surrogate_grad": [ + -3061214.9652824383, + -2610277782.364797 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + }, + { + "variant": "evshift=1", + "surrogate_grad": [ + -3062217.2209286615, + -2611132484.144538 + ], + "crn_fd_ref": [ + -22.22222222222225, + -4320.98765432099 + ], + "classification": [ + "agree", + "agree" + ], + "leaf_dtype": "float64", + "grad_finite": true + } + ], + "runtime_s": 39.6 + }, + "g3": { + "ok": true, + "cases": [ + { + "dt_ms": 0.025, + "current_nA": 12.0, + "n_spikes": 7, + "n_valid_leading_spikes": 7, + "per_spike": [ + { + "spike": 0, + "t_s": 0.0030339749263538184, + "grad_ad": [ + -4.604481350427931e-06, + -0.0010191490863060728 + ], + "fd": { + "plateau": [ + -4.604480945830103e-06, + -0.0010191492327939085 + ], + "rel": 1.4373541286869068e-07, + "sweeps": { + "0": [ + [ + 0.003, + -4.604481503182299e-06 + ], + [ + 0.001, + -4.604481308387309e-06 + ], + [ + 0.0003, + -4.604481389630192e-06 + ], + [ + 0.0001, + -4.604481197365007e-06 + ], + [ + 3e-05, + -4.604480945830103e-06 + ] + ], + "1": [ + [ + 0.003, + -0.001020612007462951 + ], + [ + 0.001, + -0.0010193114092231677 + ], + [ + 0.0003, + -0.001019163693071502 + ], + [ + 0.0001, + -0.0010191507095278635 + ], + [ + 3e-05, + -0.0010191492327939085 + ] + ] + } + } + }, + { + "spike": 1, + "t_s": 0.04170056610349914, + "grad_ad": [ + -0.001253248571272923, + -0.0018246364009492932 + ], + "fd": { + "plateau": [ + -0.001253249054186667, + -0.0018246502024588107 + ], + "rel": 7.563920744311156e-06, + "sweeps": { + "0": [ + [ + 0.003, + -0.0012504670809527125 + ], + [ + 0.001, + -0.001259576963452691 + ], + [ + 0.0003, + -0.0012532830628159133 + ], + [ + 0.0001, + -0.001253252435405583 + ], + [ + 3e-05, + -0.001253249054186667 + ] + ], + "1": [ + [ + 0.003, + -0.0019108190947846973 + ], + [ + 0.001, + -0.00184643854611935 + ], + [ + 0.0003, + -0.0018260233374441568 + ], + [ + 0.0001, + -0.0018247904769522205 + ], + [ + 3e-05, + -0.0018246502024588107 + ] + ] + } + } + }, + { + "spike": 2, + "t_s": 0.08644923303245496, + "grad_ad": [ + -0.004250010013638885, + -0.004503454614631949 + ], + "fd": { + "plateau": [ + -0.004250007387222403, + -0.004503448754883399 + ], + "rel": 1.301169141482218e-06, + "sweeps": { + "0": [ + [ + 0.003, + -0.003956158704856531 + ], + [ + 0.001, + -0.004144544817341611 + ], + [ + 0.0003, + -0.004225565170499583 + ], + [ + 0.0001, + -0.004249979802320447 + ], + [ + 3e-05, + -0.004250007387222403 + ] + ], + "1": [ + [ + 0.003, + -0.004164234155658771 + ], + [ + 0.001, + -0.00439321660077735 + ], + [ + 0.0003, + -0.004475623800552035 + ], + [ + 0.0001, + -0.004503389599078145 + ], + [ + 3e-05, + -0.004503448754883399 + ] + ] + } + } + }, + { + "spike": 3, + "t_s": 0.1308849996531116, + "grad_ad": [ + -0.008649393217003072, + -0.008009733991082471 + ], + "fd": { + "plateau": [ + -0.008649135599854882, + -0.008009401894493173 + ], + "rel": 4.146334441353388e-05, + "sweeps": { + "0": [ + [ + 0.003, + -0.00666487603830553 + ], + [ + 0.001, + -0.008372560184710665 + ], + [ + 0.0003, + -0.00862370882384611 + ], + [ + 0.0001, + -0.008646532219103165 + ], + [ + 3e-05, + -0.008649135599854882 + ] + ], + "1": [ + [ + 0.003, + -0.005726638334665098 + ], + [ + 0.001, + -0.007649246178631497 + ], + [ + 0.0003, + -0.007976594554407839 + ], + [ + 0.0001, + -0.008006046216535445 + ], + [ + 3e-05, + -0.008009401894493173 + ] + ] + } + } + }, + { + "spike": 4, + "t_s": 0.17518384123665723, + "grad_ad": [ + -0.007027725683463498, + -0.005505832177565749 + ], + "fd": { + "plateau": [ + -0.0070308938763622795, + -0.005508430388212471 + ], + "rel": 0.00047167894728808885, + "sweeps": { + "0": [ + [ + 0.003, + -0.00998852543874588 + ], + [ + 0.001, + -0.009541143290492382 + ], + [ + 0.0003, + -0.007334412028131566 + ], + [ + 0.0001, + -0.007062807968194473 + ], + [ + 3e-05, + -0.0070308938763622795 + ] + ], + "1": [ + [ + 0.003, + -0.007757404547277061 + ], + [ + 0.001, + -0.0074717026986165225 + ], + [ + 0.0003, + -0.005756222235577443 + ], + [ + 0.0001, + -0.005534592690537021 + ], + [ + 3e-05, + -0.005508430388212471 + ] + ] + } + } + }, + { + "spike": 5, + "t_s": 0.21942036537989, + "grad_ad": [ + -0.007169328955136223, + -0.004890656615735933 + ], + "fd": { + "plateau": [ + -0.007174038977690126, + -0.00489395332499883 + ], + "rel": 0.0006736290773468274, + "sweeps": { + "0": [ + [ + 0.003, + -0.012457902096835172 + ], + [ + 0.001, + -0.011764933884120254 + ], + [ + 0.0003, + -0.00764430877694955 + ], + [ + 0.0001, + -0.00722169438610587 + ], + [ + 3e-05, + -0.007174038977690126 + ] + ], + "1": [ + [ + 0.003, + -0.009355059544658798 + ], + [ + 0.001, + -0.008213395603959284 + ], + [ + 0.0003, + -0.0052236363512481185 + ], + [ + 0.0001, + -0.00492733380846877 + ], + [ + 3e-05, + -0.00489395332499883 + ] + ] + } + } + }, + { + "spike": 6, + "t_s": 0.26362539641362615, + "grad_ad": [ + -0.0125858148677122, + -0.008108803200916233 + ], + "fd": { + "plateau": [ + -0.012589075288708488, + -0.008111893035992837 + ], + "rel": 0.00038090185150307975, + "sweeps": { + "0": [ + [ + 0.003, + -0.015339658577778748 + ], + [ + 0.001, + -0.014412184018153251 + ], + [ + 0.0003, + -0.01274701683752951 + ], + [ + 0.0001, + -0.012676112528808847 + ], + [ + 3e-05, + -0.012589075288708488 + ] + ], + "1": [ + [ + 0.003, + -0.010990905831435397 + ], + [ + 0.001, + -0.009158479928383922 + ], + [ + 0.0003, + -0.00838116065687844 + ], + [ + 0.0001, + -0.008181374222815396 + ], + [ + 3e-05, + -0.008111893035992837 + ] + ] + } + } + } + ] + }, + { + "dt_ms": 0.025, + "current_nA": 15.0, + "n_spikes": 9, + "n_valid_leading_spikes": 8, + "per_spike": [ + { + "spike": 0, + "t_s": 0.0012525188109020703, + "grad_ad": [ + -4.1893923476751767e-07, + -0.0002478776698020256 + ], + "fd": { + "plateau": [ + -4.1893986832871227e-07, + -0.00024787767251115145 + ], + "rel": 1.5122962565787474e-06, + "sweeps": { + "0": [ + [ + 0.003, + -4.189392360220053e-07 + ], + [ + 0.001, + -4.1893926121163577e-07 + ], + [ + 0.0003, + -4.189391195064119e-07 + ], + [ + 0.0001, + -4.1893939489376364e-07 + ], + [ + 3e-05, + -4.1893986832871227e-07 + ] + ], + "1": [ + [ + 0.003, + -0.00024790013825048973 + ], + [ + 0.001, + -0.00024788016619317014 + ], + [ + 0.0003, + -0.00024787789453443375 + ], + [ + 0.0001, + -0.0002478776943046993 + ], + [ + 3e-05, + -0.00024787767251115145 + ] + ] + } + } + }, + { + "spike": 1, + "t_s": 0.02908072602557144, + "grad_ad": [ + -0.0005876986776864046, + -0.0013055560986796795 + ], + "fd": { + "plateau": [ + -0.0005876986033243048, + -0.001305556213253752 + ], + "rel": 1.2653101332983638e-07, + "sweeps": { + "0": [ + [ + 0.003, + -0.0005880372561072045 + ], + [ + 0.001, + -0.0005877362860146063 + ], + [ + 0.0003, + -0.0005877020714413022 + ], + [ + 0.0001, + -0.0005876989993616744 + ], + [ + 3e-05, + -0.0005876986033243048 + ] + ], + "1": [ + [ + 0.003, + -0.0013078834494237301 + ], + [ + 0.001, + -0.0013058143398282651 + ], + [ + 0.0003, + -0.0013055793344809072 + ], + [ + 0.0001, + -0.0013055587154767179 + ], + [ + 3e-05, + -0.001305556213253752 + ] + ] + } + } + }, + { + "spike": 2, + "t_s": 0.06419397053974514, + "grad_ad": [ + -0.0019203256394547214, + -0.0021015228114541315 + ], + "fd": { + "plateau": [ + -0.001920326388182847, + -0.0021015236194762052 + ], + "rel": 3.898962853989829e-07, + "sweeps": { + "0": [ + [ + 0.003, + -0.0019275538942232586 + ], + [ + 0.001, + -0.001921120636877438 + ], + [ + 0.0003, + -0.001920397120422357 + ], + [ + 0.0001, + -0.001920333514773831 + ], + [ + 3e-05, + -0.001920326388182847 + ] + ], + "1": [ + [ + 0.003, + -0.002110884995897936 + ], + [ + 0.001, + -0.0021025504367347403 + ], + [ + 0.0003, + -0.0021016151194550954 + ], + [ + 0.0001, + -0.002101533097542685 + ], + [ + 3e-05, + -0.0021015236194762052 + ] + ] + } + } + }, + { + "spike": 3, + "t_s": 0.09902788957451633, + "grad_ad": [ + -0.0038364458230265865, + -0.003164644462959351 + ], + "fd": { + "plateau": [ + -0.0038364403755671477, + -0.003164621106369984 + ], + "rel": 7.380532639436151e-06, + "sweeps": { + "0": [ + [ + 0.003, + -0.0036969940102095826 + ], + [ + 0.001, + -0.0038274835178278566 + ], + [ + 0.0003, + -0.0038359052837351935 + ], + [ + 0.0001, + -0.0038363854754780435 + ], + [ + 3e-05, + -0.0038364403755671477 + ] + ], + "1": [ + [ + 0.003, + -0.0029537038404280533 + ], + [ + 0.001, + -0.00313978880180088 + ], + [ + 0.0003, + -0.003162325418285909 + ], + [ + 0.0001, + -0.003164385974224748 + ], + [ + 3e-05, + -0.003164621106369984 + ] + ] + } + } + }, + { + "spike": 4, + "t_s": 0.13372095391336858, + "grad_ad": [ + -0.006600203085666429, + -0.004718598350536087 + ], + "fd": { + "plateau": [ + -0.0066000416428420765, + -0.0047184770040250195 + ], + "rel": 2.5717304749862825e-05, + "sweeps": { + "0": [ + [ + 0.003, + -0.005561034203085957 + ], + [ + 0.001, + -0.0064064972986471735 + ], + [ + 0.0003, + -0.0065841604591128835 + ], + [ + 0.0001, + -0.00659840970640424 + ], + [ + 3e-05, + -0.0066000416428420765 + ] + ], + "1": [ + [ + 0.003, + -0.003917544015015022 + ], + [ + 0.001, + -0.00459093614460504 + ], + [ + 0.0003, + -0.004706537213526134 + ], + [ + 0.0001, + -0.00471725137035639 + ], + [ + 3e-05, + -0.0047184770040250195 + ] + ] + } + } + }, + { + "spike": 5, + "t_s": 0.1683378445428007, + "grad_ad": [ + -0.009930782703044936, + -0.006532371726795237 + ], + "fd": { + "plateau": [ + -0.009929888883375426, + -0.006531942272255027 + ], + "rel": 9.001305855555473e-05, + "sweeps": { + "0": [ + [ + 0.003, + -0.007324339155345694 + ], + [ + 0.001, + -0.00912267995532845 + ], + [ + 0.0003, + -0.009843071345263152 + ], + [ + 0.0001, + -0.0099208693642161 + ], + [ + 3e-05, + -0.009929888883375426 + ] + ], + "1": [ + [ + 0.003, + -0.004755644041431102 + ], + [ + 0.001, + -0.006114380668692521 + ], + [ + 0.0003, + -0.006489962274582503 + ], + [ + 0.0001, + -0.006527607517431422 + ], + [ + 3e-05, + -0.006531942272255027 + ] + ] + } + } + }, + { + "spike": 6, + "t_s": 0.20291673583416478, + "grad_ad": [ + -0.008263130411376579, + -0.004939193108827128 + ], + "fd": { + "plateau": [ + -0.00826649385702686, + -0.004940033094628464 + ], + "rel": 0.0004068769309521131, + "sweeps": { + "0": [ + [ + 0.003, + -0.00957994383940508 + ], + [ + 0.001, + -0.010197725116067868 + ], + [ + 0.0003, + -0.008578841109209886 + ], + [ + 0.0001, + -0.008300255954774016 + ], + [ + 3e-05, + -0.00826649385702686 + ] + ], + "1": [ + [ + 0.003, + -0.005689495280315297 + ], + [ + 0.001, + -0.0055611738367378205 + ], + [ + 0.0003, + -0.005019953428775371 + ], + [ + 0.0001, + -0.0049484901154528416 + ], + [ + 3e-05, + -0.004940033094628464 + ] + ] + } + } + }, + { + "spike": 7, + "t_s": 0.2374747077795528, + "grad_ad": [ + -0.0098585176494798, + -0.005474563074846679 + ], + "fd": { + "plateau": [ + -0.009864112490246916, + -0.0054759637195071304 + ], + "rel": 0.0005671915007707899, + "sweeps": { + "0": [ + [ + 0.003, + -0.011521343782303997 + ], + [ + 0.001, + -0.012283866479112504 + ], + [ + 0.0003, + -0.010217658291122868 + ], + [ + 0.0001, + -0.009887801557528775 + ], + [ + 3e-05, + -0.009864112490246916 + ] + ], + "1": [ + [ + 0.003, + -0.0063308150625116895 + ], + [ + 0.001, + -0.006338752148754567 + ], + [ + 0.0003, + -0.0055649134877526025 + ], + [ + 0.0001, + -0.005484694661772593 + ], + [ + 3e-05, + -0.0054759637195071304 + ] + ] + } + } + } + ] + }, + { + "dt_ms": 0.025, + "current_nA": 30.0, + "n_spikes": 23, + "n_valid_leading_spikes": 8, + "per_spike": [ + { + "spike": 0, + "t_s": 0.0003716877681465094, + "grad_ad": [ + -8.051124867202451e-09, + -1.5558925939572994e-05 + ], + "fd": { + "plateau": [ + -8.049925562652696e-09, + -1.5558925448839387e-05 + ], + "rel": 0.00014898330927667555, + "sweeps": { + "0": [ + [ + 0.003, + -8.051134881755887e-09 + ], + [ + 0.001, + -8.051085541522021e-09 + ], + [ + 0.0003, + -8.051215130786687e-09 + ], + [ + 0.0001, + -8.051184773125858e-09 + ], + [ + 3e-05, + -8.049925562652696e-09 + ] + ], + "1": [ + [ + 0.003, + -1.5558932798691088e-05 + ], + [ + 0.001, + -1.555892672683366e-05 + ], + [ + 0.0003, + -1.5558926150860294e-05 + ], + [ + 0.0001, + -1.5558925721425882e-05 + ], + [ + 3e-05, + -1.5558925448839387e-05 + ] + ] + } + } + }, + { + "spike": 1, + "t_s": 0.005865283453098428, + "grad_ad": [ + -1.6826700383176837e-05, + -0.0002223766958292127 + ], + "fd": { + "plateau": [ + -1.682664386877068e-05, + -0.0002223765782193715 + ], + "rel": 3.3586261525939525e-06, + "sweeps": { + "0": [ + [ + 0.003, + -1.6826706903562184e-05 + ], + [ + 0.001, + -1.6826699792207855e-05 + ], + [ + 0.0003, + -1.682669856749308e-05 + ], + [ + 0.0001, + -1.6826703190531145e-05 + ], + [ + 3e-05, + -1.682664386877068e-05 + ] + ], + "1": [ + [ + 0.003, + -0.00022239186964085703 + ], + [ + 0.001, + -0.00022237838196575987 + ], + [ + 0.0003, + -0.0002223768407090543 + ], + [ + 0.0001, + -0.00022237670169698853 + ], + [ + 3e-05, + -0.0002223765782193715 + ] + ] + } + } + }, + { + "spike": 2, + "t_s": 0.01932788836466429, + "grad_ad": [ + -0.0002195838707194839, + -2.0676606458107545e-05 + ], + "fd": { + "plateau": [ + -0.0002195839200862711, + -2.0676462914500277e-05 + ], + "rel": 6.942367650668396e-06, + "sweeps": { + "0": [ + [ + 0.003, + -0.00021960216200697658 + ], + [ + 0.001, + -0.00021958590269537115 + ], + [ + 0.0003, + -0.0002195840442866896 + ], + [ + 0.0001, + -0.00021958384810102927 + ], + [ + 3e-05, + -0.0002195839200862711 + ] + ], + "1": [ + [ + 0.003, + -2.0747093921908804e-05 + ], + [ + 0.001, + -2.0684447396210337e-05 + ], + [ + 0.0003, + -2.0677318063784998e-05 + ], + [ + 0.0001, + -2.067668094032915e-05 + ], + [ + 3e-05, + -2.0676462914500277e-05 + ] + ] + } + } + }, + { + "spike": 3, + "t_s": 0.03311281078541542, + "grad_ad": [ + -0.0005196645955708599, + 0.00015119226228252724 + ], + "fd": { + "plateau": [ + -0.0005196647146164338, + 0.00015119233931955764 + ], + "rel": 5.095299850997067e-07, + "sweeps": { + "0": [ + [ + 0.003, + -0.0005198424406991272 + ], + [ + 0.001, + -0.0005196843449566113 + ], + [ + 0.0003, + -0.0005196663649745248 + ], + [ + 0.0001, + -0.0005196647635472007 + ], + [ + 3e-05, + -0.0005196647146164338 + ] + ], + "1": [ + [ + 0.003, + 0.00015112951068143218 + ], + [ + 0.001, + 0.00015118527652707714 + ], + [ + 0.0003, + 0.0001511916215488074 + ], + [ + 0.0001, + 0.00015119219015646834 + ], + [ + 3e-05, + 0.00015119233931955764 + ] + ] + } + } + }, + { + "spike": 4, + "t_s": 0.046843534354860925, + "grad_ad": [ + -0.0008749343337623448, + 0.0003096815874781394 + ], + "fd": { + "plateau": [ + -0.0008749346663838864, + 0.0003096816505778384 + ], + "rel": 3.801672906764423e-07, + "sweeps": { + "0": [ + [ + 0.003, + -0.000876485233252076 + ], + [ + 0.001, + -0.0008751069175555459 + ], + [ + 0.0003, + -0.0008749498668867799 + ], + [ + 0.0001, + -0.000874936049791164 + ], + [ + 3e-05, + -0.0008749346663838864 + ] + ], + "1": [ + [ + 0.003, + 0.00030967669253922914 + ], + [ + 0.001, + 0.00030968103425371707 + ], + [ + 0.0003, + 0.00030968153235065116 + ], + [ + 0.0001, + 0.00030968160251443333 + ], + [ + 3e-05, + 0.0003096816505778384 + ] + ] + } + } + }, + { + "spike": 5, + "t_s": 0.0605150581596771, + "grad_ad": [ + -0.001306969606528772, + 0.0004741182370801159 + ], + "fd": { + "plateau": [ + -0.001306970229958404, + 0.0004741183048014482 + ], + "rel": 4.770036974422461e-07, + "sweeps": { + "0": [ + [ + 0.003, + -0.001311666365987063 + ], + [ + 0.001, + -0.001307491582447362 + ], + [ + 0.0003, + -0.001307016591923597 + ], + [ + 0.0001, + -0.001306974763033708 + ], + [ + 3e-05, + -0.001306970229958404 + ] + ], + "1": [ + [ + 0.003, + 0.00047427484786352614 + ], + [ + 0.001, + 0.0004741356468933311 + ], + [ + 0.0003, + 0.00047411980648217244 + ], + [ + 0.0001, + 0.00047411843364514317 + ], + [ + 3e-05, + 0.0004741183048014482 + ] + ] + } + } + }, + { + "spike": 6, + "t_s": 0.07414076913176774, + "grad_ad": [ + -0.001731610947883571, + 0.0006216430069692803 + ], + "fd": { + "plateau": [ + -0.001731613086947137, + 0.0006216430611013143 + ], + "rel": 1.235301108509737e-06, + "sweeps": { + "0": [ + [ + 0.003, + -0.0017504771555922523 + ], + [ + 0.001, + -0.0017337439408987576 + ], + [ + 0.0003, + -0.0017318032968551464 + ], + [ + 0.0001, + -0.0017316322217797442 + ], + [ + 3e-05, + -0.001731613086947137 + ] + ], + "1": [ + [ + 0.003, + 0.0006224928765091731 + ], + [ + 0.001, + 0.0006217377178305283 + ], + [ + 0.0003, + 0.000621651534277179 + ], + [ + 0.0001, + 0.0006216440092549069 + ], + [ + 3e-05, + 0.0006216430611013143 + ] + ] + } + } + }, + { + "spike": 7, + "t_s": 0.08773229929426328, + "grad_ad": [ + -0.002281752919751627, + 0.0008033479302932015 + ], + "fd": { + "plateau": [ + -0.002281758237447177, + 0.0008033481165015122 + ], + "rel": 2.330525409172418e-06, + "sweeps": { + "0": [ + [ + 0.003, + -0.0023321473020092762 + ], + [ + 0.001, + -0.002287472107863342 + ], + [ + 0.0003, + -0.0022822688940105085 + ], + [ + 0.0001, + -0.002281810137574891 + ], + [ + 3e-05, + -0.002281758237447177 + ] + ], + "1": [ + [ + 0.003, + 0.000805644564387531 + ], + [ + 0.001, + 0.0008036039916073379 + ], + [ + 0.0003, + 0.0008033709853842257 + ], + [ + 0.0001, + 0.0008033505644738925 + ], + [ + 3e-05, + 0.0008033481165015122 + ] + ] + } + } + } + ] + }, + { + "dt_ms": 0.1, + "current_nA": 12.0, + "n_spikes": 7, + "n_valid_leading_spikes": 2, + "per_spike": [ + { + "spike": 0, + "t_s": 0.00323805001193331, + "grad_ad": [ + -3.84163796338609e-06, + -0.000801426614327354 + ], + "fd": { + "plateau": [ + -3.841640019696761e-06, + -0.0008014267223451753 + ], + "rel": 5.352689634239616e-07, + "sweeps": { + "0": [ + [ + 0.003, + -3.841638156892868e-06 + ], + [ + 0.001, + -3.841638201778838e-06 + ], + [ + 0.0003, + -3.84163735400502e-06 + ], + [ + 0.0001, + -3.841638955733029e-06 + ], + [ + 3e-05, + -3.841640019696761e-06 + ] + ], + "1": [ + [ + 0.003, + -0.0008025960343823463 + ], + [ + 0.001, + -0.000801556614965851 + ], + [ + 0.0003, + -0.0008014383153084546 + ], + [ + 0.0001, + -0.0008014279139910603 + ], + [ + 3e-05, + -0.0008014267223451753 + ] + ] + } + } + }, + { + "spike": 1, + "t_s": 0.04232072946461426, + "grad_ad": [ + -0.001746470143425086, + -0.07700653133278483 + ], + "fd": { + "plateau": [ + -0.001746469954882109, + -0.07707676479123034 + ], + "rel": 0.0009112144059982354, + "sweeps": { + "0": [ + [ + 0.003, + -0.0017547566209784777 + ], + [ + 0.001, + -0.0017473842168214127 + ], + [ + 0.0003, + -0.0017465522959605093 + ], + [ + 0.0001, + -0.0017464789865348496 + ], + [ + 3e-05, + -0.001746469954882109 + ] + ], + "1": [ + [ + 0.003, + -0.07121827239082885 + ], + [ + 0.001, + -0.06898641365901681 + ], + [ + 0.0003, + -0.07936715462803641 + ], + [ + 0.0001, + -0.07779941505671462 + ], + [ + 3e-05, + -0.07707676479123034 + ] + ] + } + } + }, + { + "spike": 2, + "t_s": 0.08620503873124913, + "grad_ad": [ + -0.1443098351349357, + -6.2110885123506145 + ], + "fd": null + }, + { + "spike": 3, + "t_s": 0.1292730278100448, + "grad_ad": [ + -1.4098289739280594, + -60.564898813083836 + ], + "fd": null + }, + { + "spike": 4, + "t_s": 0.17330798675324358, + "grad_ad": [ + 159.01692480848868, + 6831.384876484863 + ], + "fd": null + }, + { + "spike": 5, + "t_s": 0.2161870083975544, + "grad_ad": [ + 3006.4702059204888, + 129158.44615469214 + ], + "fd": null + }, + { + "spike": 6, + "t_s": 0.2592396678155984, + "grad_ad": [ + -93959.20107141092, + -4036502.372582738 + ], + "fd": null + } + ] + }, + { + "dt_ms": 0.1, + "current_nA": 15.0, + "n_spikes": 9, + "n_valid_leading_spikes": 2, + "per_spike": [ + { + "spike": 0, + "t_s": 0.0014290367977746192, + "grad_ad": [ + -4.6444322502136746e-07, + -0.00022999994893089932 + ], + "fd": { + "plateau": [ + -4.644499990042831e-07, + -0.0002299999545347121 + ], + "rel": 1.458495624963145e-05, + "sweeps": { + "0": [ + [ + 0.003, + -4.644431735179264e-07 + ], + [ + 0.001, + -4.6444316003767944e-07 + ], + [ + 0.0003, + -4.6444367340740816e-07 + ], + [ + 3e-05, + -4.644499990042831e-07 + ] + ], + "1": [ + [ + 0.003, + -0.0002300211001443171 + ], + [ + 0.001, + -0.00023000229898309608 + ], + [ + 0.0001, + -0.0002299999743969345 + ], + [ + 3e-05, + -0.0002299999545347121 + ] + ] + } + } + }, + { + "spike": 1, + "t_s": 0.029017261711066587, + "grad_ad": [ + -0.000803759464758988, + -0.01955420380762187 + ], + "fd": { + "plateau": [ + -0.0008037588213628358, + -0.019555309072744054 + ], + "rel": 5.651995159324288e-05, + "sweeps": { + "0": [ + [ + 0.003, + -0.0008045643282011498 + ], + [ + 0.001, + -0.0008038487983603321 + ], + [ + 0.0003, + -0.0008037674686182014 + ], + [ + 3e-05, + -0.0008037588213628358 + ] + ], + "1": [ + [ + 0.003, + -0.01621261706775262 + ], + [ + 0.001, + -0.0195604651706191 + ], + [ + 0.0001, + -0.019566500553628435 + ], + [ + 3e-05, + -0.019555309072744054 + ] + ] + } + } + }, + { + "spike": 2, + "t_s": 0.06276385024930205, + "grad_ad": [ + -0.02254646606077169, + -0.5195971047704162 + ], + "fd": null + }, + { + "spike": 3, + "t_s": 0.09796500839885469, + "grad_ad": [ + 1.049028507787286, + 24.20906756514034 + ], + "fd": null + }, + { + "spike": 4, + "t_s": 0.1329520064843973, + "grad_ad": [ + -46.14793311517269, + -1064.9516410871483 + ], + "fd": null + }, + { + "spike": 5, + "t_s": 0.16816642789284014, + "grad_ad": [ + -813.2625354933497, + -18767.54488411917 + ], + "fd": null + }, + { + "spike": 6, + "t_s": 0.20295579832280636, + "grad_ad": [ + 39354.79605919462, + 908185.115775408 + ], + "fd": null + }, + { + "spike": 7, + "t_s": 0.23813999509119108, + "grad_ad": [ + -8782.476520733126, + -202671.94403264087 + ], + "fd": null + } + ] + }, + { + "dt_ms": 0.1, + "current_nA": 30.0, + "n_spikes": 26, + "n_valid_leading_spikes": 4, + "per_spike": [ + { + "spike": 0, + "t_s": 0.0004044062139002335, + "grad_ad": [ + -1.5750563079731962e-08, + -2.1271787138901856e-05 + ], + "fd": { + "plateau": [ + -1.5754002377806053e-08, + -2.1271787090560053e-05 + ], + "rel": 0.00021831265424566976, + "sweeps": { + "0": [ + [ + 0.003, + -1.5750501154695427e-08 + ], + [ + 0.001, + -1.5750533581375403e-08 + ], + [ + 0.0003, + -1.5750700169039206e-08 + ], + [ + 0.0001, + -1.5749715984517132e-08 + ], + [ + 3e-05, + -1.5754002377806053e-08 + ] + ], + "1": [ + [ + 0.003, + -2.1271791084724715e-05 + ], + [ + 0.001, + -2.1271787783744712e-05 + ], + [ + 0.0003, + -2.127178724840182e-05 + ], + [ + 0.0001, + -2.1271786168446105e-05 + ], + [ + 3e-05, + -2.1271787090560053e-05 + ] + ] + } + } + }, + { + "spike": 1, + "t_s": 0.005617781904341906, + "grad_ad": [ + -2.236242439187489e-05, + 0.0004842357424543193 + ], + "fd": { + "plateau": [ + -2.2362040344447998e-05, + 0.00048423582157952616 + ], + "rel": 1.7174078079524876e-05, + "sweeps": { + "0": [ + [ + 0.003, + -2.2362444203396597e-05 + ], + [ + 0.001, + -2.236241476109152e-05 + ], + [ + 0.0003, + -2.2362410050160802e-05 + ], + [ + 0.0001, + -2.2362496095336415e-05 + ], + [ + 3e-05, + -2.2362040344447998e-05 + ] + ], + "1": [ + [ + 0.003, + 0.0004843290486404352 + ], + [ + 0.001, + 0.000484246144928039 + ], + [ + 0.0003, + 0.0004842367289324227 + ], + [ + 0.0001, + 0.00048423580604507743 + ], + [ + 3e-05, + 0.00048423582157952616 + ] + ] + } + } + }, + { + "spike": 2, + "t_s": 0.017256179651555344, + "grad_ad": [ + -7.595742550476128e-06, + -0.001926757159767402 + ], + "fd": { + "plateau": [ + -7.597588125759132e-06, + -0.001926757748256976 + ], + "rel": 0.00024291594285651797, + "sweeps": { + "0": [ + [ + 0.003, + -7.595841311838769e-06 + ], + [ + 0.001, + -7.595811191835056e-06 + ], + [ + 0.0003, + -7.595776421037704e-06 + ], + [ + 0.0001, + -7.5953847378240624e-06 + ], + [ + 3e-05, + -7.597588125759132e-06 + ] + ], + "1": [ + [ + 0.003, + -0.0019350590406887885 + ], + [ + 0.001, + -0.0019276707482265842 + ], + [ + 0.0003, + -0.0019268395296673237 + ], + [ + 0.0001, + -0.0019267659983300545 + ], + [ + 3e-05, + -0.001926757748256976 + ] + ] + } + } + }, + { + "spike": 3, + "t_s": 0.029764784260206463, + "grad_ad": [ + -0.00041999554306953297, + -0.05040738303465467 + ], + "fd": { + "plateau": [ + -0.0004200443102717055, + -0.050430102605895456 + ], + "rel": 0.00045051606216903113, + "sweeps": { + "0": [ + [ + 0.003, + -0.00042021502588349136 + ], + [ + 0.001, + -0.0004200214357638621 + ], + [ + 0.0003, + -0.00041999858061900214 + ], + [ + 0.0001, + -0.0004199862889235628 + ], + [ + 3e-05, + -0.0004200443102717055 + ] + ], + "1": [ + [ + 0.003, + -0.06618791435097915 + ], + [ + 0.001, + -0.07583545847716922 + ], + [ + 0.0003, + -0.05269554799178696 + ], + [ + 0.0001, + -0.050660024095670886 + ], + [ + 3e-05, + -0.050430102605895456 + ] + ] + } + } + }, + { + "spike": 4, + "t_s": 0.04243838729668155, + "grad_ad": [ + -0.00885657088876606, + -1.0152862347169065 + ], + "fd": null + }, + { + "spike": 5, + "t_s": 0.05431843317614229, + "grad_ad": [ + -0.16221347690251142, + -18.512317633204827 + ], + "fd": null + }, + { + "spike": 6, + "t_s": 0.06619589117603772, + "grad_ad": [ + 1.3855020035295509, + 158.17714951260385 + ], + "fd": null + }, + { + "spike": 7, + "t_s": 0.07832564110135073, + "grad_ad": [ + -41.21347060366345, + -4705.092996782993 + ], + "fd": null + } + ] + } + ], + "scope": "frozen-index interpolated crossing times only \u2014 NOT a complete event-time estimator (no creation/deletion, matching, or full objective)", + "runtime_s": 613.9 + } +} \ No newline at end of file diff --git a/experiments/twin/evidence/pic_grid_full_20261018_132043.json b/experiments/twin/evidence/pic_grid_full_20261018_132043.json new file mode 100644 index 00000000..2cf19f52 --- /dev/null +++ b/experiments/twin/evidence/pic_grid_full_20261018_132043.json @@ -0,0 +1,1976 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "993a2cfdebaeb499", + "poolbuild.py": "816d4ba42490e283", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "983f8bbcafebd67e", + "protocols.py": "404ea437741ea2ae" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float32 (JAX default; x64 not enabled)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations", + "pic_grid_sha256_16": "d900172c74ba1695", + "args": { + "seeds": [ + 101, + 202 + ], + "gcal": [ + 0.0, + 0.5, + 1.0, + 2.0, + 4.0, + 6.5, + 7.5 + ], + "gnap": [ + 1.0, + 5.0 + ], + "n_mn": 8, + "dur_scale": 1.0, + "skip_dt_ref": false + }, + "run_tag": "full", + "acceptance_spec_version": 1 + }, + "dt_refinement_E_numerical_hz2": { + "1.0,1.0": 0.7163742780685425, + "6.5,5.0": 3.4941518306732178 + }, + "points": [ + { + "gcal": 0.0, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.1388888657093048, + "active_unit_mismatch": 0, + "active_jaxley": 5, + "active_neuron": 5, + "total_spike_rel_diff": 0.0035971223021582736, + "total_jaxley": 277, + "total_neuron": 278, + "spike_match_pm5ms_fraction": 0.8812949640287769, + "spike_total_max": 278, + "hysteresis_mean_diff_hz": 0.0039997100830078125, + "hysteresis_jaxley": -0.017440414428710936, + "hysteresis_neuron": -0.02144012451171875, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.08469741024728367, + "delta_f_pairs_jaxley": 4, + "delta_f_pairs_neuron": 3, + "delta_f_eligible_fraction_diff": 0.10000000000000003 + } + }, + { + "gcal": 0.5, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.12426900118589401, + "active_unit_mismatch": 0, + "active_jaxley": 5, + "active_neuron": 5, + "total_spike_rel_diff": 0.0, + "total_jaxley": 332, + "total_neuron": 332, + "spike_match_pm5ms_fraction": 0.9036144578313253, + "spike_total_max": 332, + "hysteresis_mean_diff_hz": 0.005918884277343883, + "hysteresis_jaxley": 1.088799285888672, + "hysteresis_neuron": 1.082880401611328, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.5138551444173961, + "delta_f_pairs_jaxley": 3, + "delta_f_pairs_neuron": 3, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 1.0, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.21198827028274536, + "active_unit_mismatch": 0, + "active_jaxley": 5, + "active_neuron": 5, + "total_spike_rel_diff": 0.007633587786259542, + "total_jaxley": 393, + "total_neuron": 390, + "spike_match_pm5ms_fraction": 0.811704834605598, + "spike_total_max": 393, + "hysteresis_mean_diff_hz": 0.05824012756347652, + "hysteresis_jaxley": 1.2731197357177735, + "hysteresis_neuron": 1.214879608154297, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.6766043593993221, + "delta_f_pairs_jaxley": 1, + "delta_f_pairs_neuron": 1, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 2.0, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.22660820186138153, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.003838771593090211, + "total_jaxley": 521, + "total_neuron": 519, + "spike_match_pm5ms_fraction": 0.8771593090211133, + "spike_total_max": 521, + "hysteresis_mean_diff_hz": 0.8108795166015625, + "hysteresis_jaxley": -0.11632003784179687, + "hysteresis_neuron": 0.6945594787597656, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.02082091349852888, + "delta_f_pairs_jaxley": 1, + "delta_f_pairs_neuron": 1, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 4.0, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.6140352487564087, + "active_unit_mismatch": 0, + "active_jaxley": 7, + "active_neuron": 7, + "total_spike_rel_diff": 0.003386004514672686, + "total_jaxley": 886, + "total_neuron": 883, + "spike_match_pm5ms_fraction": 0.672686230248307, + "spike_total_max": 886, + "hysteresis_mean_diff_hz": 0.004666487375895145, + "hysteresis_jaxley": -0.6574664115905762, + "hysteresis_neuron": -0.6621328989664713, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.12319141530959099, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 6.5, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 1.3304095268249512, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.005500550055005501, + "total_jaxley": 1818, + "total_neuron": 1808, + "spike_match_pm5ms_fraction": 0.6798679867986799, + "spike_total_max": 1818, + "hysteresis_mean_diff_hz": 0.7280998229980469, + "hysteresis_jaxley": -1.6875999569892883, + "hysteresis_neuron": -0.9595001339912415, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.7925554056432866, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 7.5, + "gnap": 1.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 1.5752921104431152, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.0024793388429752068, + "total_jaxley": 2420, + "total_neuron": 2414, + "spike_match_pm5ms_fraction": 0.6520661157024793, + "spike_total_max": 2420, + "hysteresis_mean_diff_hz": 0.0447002649307251, + "hysteresis_jaxley": -8.121099829673767, + "hysteresis_neuron": -8.165800094604492, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 6, + "self_sustained_neuron": 6, + "after_discharge_censored_jaxley": 6, + "after_discharge_censored_neuron": 6, + "delta_f_mean_diff_hz": 0.027162311025743513, + "delta_f_pairs_jaxley": 10, + "delta_f_pairs_neuron": 10, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 0.0, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.6213451027870178, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.0017699115044247787, + "total_jaxley": 565, + "total_neuron": 564, + "spike_match_pm5ms_fraction": 0.5504424778761062, + "spike_total_max": 565, + "hysteresis_mean_diff_hz": 0.0025600433349609597, + "hysteresis_jaxley": -0.1705596923828125, + "hysteresis_neuron": -0.17311973571777345, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.0038196085961832438, + "delta_f_pairs_jaxley": 3, + "delta_f_pairs_neuron": 3, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 0.5, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.5043861269950867, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.003115264797507788, + "total_jaxley": 642, + "total_neuron": 640, + "spike_match_pm5ms_fraction": 0.6900311526479751, + "spike_total_max": 642, + "hysteresis_mean_diff_hz": 0.003333409627278648, + "hysteresis_jaxley": 0.03906663258870443, + "hysteresis_neuron": 0.03573322296142578, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": null, + "delta_f_pairs_jaxley": 0, + "delta_f_pairs_neuron": 0, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 1.0, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.6871345639228821, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.004178272980501393, + "total_jaxley": 718, + "total_neuron": 715, + "spike_match_pm5ms_fraction": 0.6462395543175488, + "spike_total_max": 718, + "hysteresis_mean_diff_hz": 0.2422655423482259, + "hysteresis_jaxley": -0.4611991246541341, + "hysteresis_neuron": -0.2189335823059082, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 1.4091694201963283, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 2.0, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 0.6725149154663086, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.0010976948408342481, + "total_jaxley": 911, + "total_neuron": 910, + "spike_match_pm5ms_fraction": 0.7025246981339188, + "spike_total_max": 911, + "hysteresis_mean_diff_hz": 0.004533131917317718, + "hysteresis_jaxley": -0.14533344904581705, + "hysteresis_neuron": -0.14986658096313477, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.2067227867245518, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 4.0, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 1.220760464668274, + "active_unit_mismatch": 0, + "active_jaxley": 7, + "active_neuron": 7, + "total_spike_rel_diff": 0.0006779661016949153, + "total_jaxley": 1475, + "total_neuron": 1474, + "spike_match_pm5ms_fraction": 0.6149152542372881, + "spike_total_max": 1475, + "hysteresis_mean_diff_hz": 0.02114248275756836, + "hysteresis_jaxley": -1.0964570045471191, + "hysteresis_neuron": -1.0753145217895508, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.43095101134947444, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 6.5, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 2.5255849361419678, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.0010608203677510608, + "total_jaxley": 2828, + "total_neuron": 2825, + "spike_match_pm5ms_fraction": 0.4957567185289958, + "spike_total_max": 2828, + "hysteresis_mean_diff_hz": 0.02339988946914673, + "hysteresis_jaxley": -9.341000080108643, + "hysteresis_neuron": -9.36439996957779, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 7, + "self_sustained_neuron": 7, + "after_discharge_censored_jaxley": 7, + "after_discharge_censored_neuron": 7, + "delta_f_mean_diff_hz": 0.0373507447149084, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 7.5, + "gnap": 5.0, + "seed": 101, + "metrics": { + "binned_rate_rmse_hz2": 2.218567132949829, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.00331225534477567, + "total_jaxley": 3321, + "total_neuron": 3310, + "spike_match_pm5ms_fraction": 0.5441132189099669, + "spike_total_max": 3321, + "hysteresis_mean_diff_hz": 0.02520012855529785, + "hysteresis_jaxley": -8.569499880075455, + "hysteresis_neuron": -8.594700008630753, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 8, + "self_sustained_neuron": 8, + "after_discharge_censored_jaxley": 8, + "after_discharge_censored_neuron": 8, + "delta_f_mean_diff_hz": null, + "delta_f_pairs_jaxley": 0, + "delta_f_pairs_neuron": 0, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 0.0, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.1461988389492035, + "active_unit_mismatch": 0, + "active_jaxley": 5, + "active_neuron": 5, + "total_spike_rel_diff": 0.006514657980456026, + "total_jaxley": 307, + "total_neuron": 305, + "spike_match_pm5ms_fraction": 0.8859934853420195, + "spike_total_max": 307, + "hysteresis_mean_diff_hz": 0.3177619934082032, + "hysteresis_jaxley": -1.8292808532714844, + "hysteresis_neuron": -1.5115188598632812, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.9142920941796274, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 0.5, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.11695905774831772, + "active_unit_mismatch": 0, + "active_jaxley": 5, + "active_neuron": 5, + "total_spike_rel_diff": 0.0027397260273972603, + "total_jaxley": 364, + "total_neuron": 365, + "spike_match_pm5ms_fraction": 0.8986301369863013, + "spike_total_max": 365, + "hysteresis_mean_diff_hz": 0.005120468139648449, + "hysteresis_jaxley": -0.16527976989746093, + "hysteresis_neuron": -0.17040023803710938, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.12301673129022733, + "delta_f_pairs_jaxley": 4, + "delta_f_pairs_neuron": 4, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 1.0, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.30701756477355957, + "active_unit_mismatch": 0, + "active_jaxley": 5, + "active_neuron": 5, + "total_spike_rel_diff": 0.004819277108433735, + "total_jaxley": 415, + "total_neuron": 413, + "spike_match_pm5ms_fraction": 0.7325301204819277, + "spike_total_max": 415, + "hysteresis_mean_diff_hz": 0.00304145812988299, + "hysteresis_jaxley": 1.720321273803711, + "hysteresis_neuron": 1.717279815673828, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.09676233111716526, + "delta_f_pairs_jaxley": 4, + "delta_f_pairs_neuron": 4, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 2.0, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.3143276274204254, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.003484320557491289, + "total_jaxley": 574, + "total_neuron": 572, + "spike_match_pm5ms_fraction": 0.7770034843205574, + "spike_total_max": 574, + "hysteresis_mean_diff_hz": 0.08386643727620442, + "hysteresis_jaxley": -0.08626683553059895, + "hysteresis_neuron": -0.17013327280680338, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.37859854069758014, + "delta_f_pairs_jaxley": 7, + "delta_f_pairs_neuron": 7, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 4.0, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.7236843705177307, + "active_unit_mismatch": 0, + "active_jaxley": 7, + "active_neuron": 7, + "total_spike_rel_diff": 0.006269592476489028, + "total_jaxley": 957, + "total_neuron": 951, + "spike_match_pm5ms_fraction": 0.7053291536050157, + "spike_total_max": 957, + "hysteresis_mean_diff_hz": 0.7953335444132487, + "hysteresis_jaxley": 1.3536003430684407, + "hysteresis_neuron": 0.558266798655192, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.6240903616525737, + "delta_f_pairs_jaxley": 4, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.047619047619047616 + } + }, + { + "gcal": 6.5, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 1.030701756477356, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.00531632110579479, + "total_jaxley": 1881, + "total_neuron": 1871, + "spike_match_pm5ms_fraction": 0.7086656034024456, + "spike_total_max": 1881, + "hysteresis_mean_diff_hz": 0.2829001545906067, + "hysteresis_jaxley": -1.3730002343654633, + "hysteresis_neuron": -1.0901000797748566, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.1705009000027673, + "delta_f_pairs_jaxley": 9, + "delta_f_pairs_neuron": 9, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 7.5, + "gnap": 1.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 1.5679821968078613, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.004042037186742118, + "total_jaxley": 2474, + "total_neuron": 2464, + "spike_match_pm5ms_fraction": 0.6632983023443816, + "spike_total_max": 2474, + "hysteresis_mean_diff_hz": 0.15490001440048218, + "hysteresis_jaxley": -8.246499955654144, + "hysteresis_neuron": -8.091599941253662, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 6, + "self_sustained_neuron": 6, + "after_discharge_censored_jaxley": 6, + "after_discharge_censored_neuron": 6, + "delta_f_mean_diff_hz": 0.2873229557701684, + "delta_f_pairs_jaxley": 9, + "delta_f_pairs_neuron": 9, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 0.0, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.453216552734375, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.003246753246753247, + "total_jaxley": 616, + "total_neuron": 614, + "spike_match_pm5ms_fraction": 0.6444805194805194, + "spike_total_max": 616, + "hysteresis_mean_diff_hz": 0.005598704020182282, + "hysteresis_jaxley": 0.24066607157389322, + "hysteresis_neuron": 0.23506736755371094, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.03285833267333382, + "delta_f_pairs_jaxley": 8, + "delta_f_pairs_neuron": 7, + "delta_f_eligible_fraction_diff": 0.06666666666666665 + } + }, + { + "gcal": 0.5, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.48976612091064453, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.004285714285714286, + "total_jaxley": 700, + "total_neuron": 697, + "spike_match_pm5ms_fraction": 0.75, + "spike_total_max": 700, + "hysteresis_mean_diff_hz": 0.01360066731770837, + "hysteresis_jaxley": -0.7858667373657227, + "hysteresis_neuron": -0.799467404683431, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.3268232694870956, + "delta_f_pairs_jaxley": 8, + "delta_f_pairs_neuron": 8, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 1.0, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.42397674918174744, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.0012755102040816326, + "total_jaxley": 784, + "total_neuron": 783, + "spike_match_pm5ms_fraction": 0.7627551020408163, + "spike_total_max": 784, + "hysteresis_mean_diff_hz": 0.1857334772745769, + "hysteresis_jaxley": -1.2122669219970703, + "hysteresis_neuron": -1.3980003992716472, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.0897759652728638, + "delta_f_pairs_jaxley": 6, + "delta_f_pairs_neuron": 6, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 2.0, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 0.8333335518836975, + "active_unit_mismatch": 0, + "active_jaxley": 6, + "active_neuron": 6, + "total_spike_rel_diff": 0.0050968399592252805, + "total_jaxley": 981, + "total_neuron": 976, + "spike_match_pm5ms_fraction": 0.6432212028542303, + "spike_total_max": 981, + "hysteresis_mean_diff_hz": 0.1349333127339681, + "hysteresis_jaxley": -0.054800192515055336, + "hysteresis_neuron": 0.08013312021891277, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 0.14791117551575783, + "delta_f_pairs_jaxley": 6, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.06666666666666671 + } + }, + { + "gcal": 4.0, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 1.030701756477356, + "active_unit_mismatch": 0, + "active_jaxley": 7, + "active_neuron": 7, + "total_spike_rel_diff": 0.004560260586319218, + "total_jaxley": 1535, + "total_neuron": 1528, + "spike_match_pm5ms_fraction": 0.6579804560260586, + "spike_total_max": 1535, + "hysteresis_mean_diff_hz": 0.4800000871930804, + "hysteresis_jaxley": -0.9265142168317523, + "hysteresis_neuron": -1.4065143040248327, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 0, + "self_sustained_neuron": 0, + "after_discharge_censored_jaxley": 0, + "after_discharge_censored_neuron": 0, + "delta_f_mean_diff_hz": 1.4659453215687712, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 6.5, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 1.7105265855789185, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.001392757660167131, + "total_jaxley": 2872, + "total_neuron": 2868, + "spike_match_pm5ms_fraction": 0.6754874651810585, + "spike_total_max": 2872, + "hysteresis_mean_diff_hz": 0.02610000967979431, + "hysteresis_jaxley": -9.319500029087067, + "hysteresis_neuron": -9.345600038766861, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 7, + "self_sustained_neuron": 7, + "after_discharge_censored_jaxley": 7, + "after_discharge_censored_neuron": 7, + "delta_f_mean_diff_hz": 0.11586776351567529, + "delta_f_pairs_jaxley": 5, + "delta_f_pairs_neuron": 5, + "delta_f_eligible_fraction_diff": 0.0 + } + }, + { + "gcal": 7.5, + "gnap": 5.0, + "seed": 202, + "metrics": { + "binned_rate_rmse_hz2": 2.0138893127441406, + "active_unit_mismatch": 0, + "active_jaxley": 8, + "active_neuron": 8, + "total_spike_rel_diff": 0.0017878426698450535, + "total_jaxley": 3356, + "total_neuron": 3350, + "spike_match_pm5ms_fraction": 0.5390345649582837, + "spike_total_max": 3356, + "hysteresis_mean_diff_hz": 0.014499843120574951, + "hysteresis_jaxley": -8.504999965429306, + "hysteresis_neuron": -8.519499808549881, + "self_sustained_mismatch": 0, + "self_sustained_jaxley": 8, + "self_sustained_neuron": 8, + "after_discharge_censored_jaxley": 8, + "after_discharge_censored_neuron": 8, + "delta_f_mean_diff_hz": null, + "delta_f_pairs_jaxley": 0, + "delta_f_pairs_neuron": 0, + "delta_f_eligible_fraction_diff": 0.0 + } + } + ], + "acceptance": { + "spec_version": 1, + "n_points": 28, + "n_evaluated_metric_cells": 221, + "n_failed_metric_cells": 5, + "fail_fraction": 0.02262443438914027, + "failure_policy_ok": true, + "bistability_boundary": { + "1.0": { + "jaxley_first_bistable_gcal": 7.5, + "neuron_first_bistable_gcal": 7.5, + "offset_steps": 0, + "class": "PASS" + }, + "5.0": { + "jaxley_first_bistable_gcal": 6.5, + "neuron_first_bistable_gcal": 6.5, + "offset_steps": 0, + "class": "PASS" + } + }, + "validated_envelope_gcal_by_gnap": { + "1.0": [ + 0.0, + 7.5 + ], + "5.0": [ + 0.5, + 4.0 + ] + } + }, + "acceptance_points": [ + { + "gcal": 0.0, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.1388888657093048 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0035971223021582736 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.8812949640287769 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.0039997100830078125 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.08469741024728367 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.10000000000000003 + } + } + }, + { + "gcal": 0.5, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.12426900118589401 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.9036144578313253 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.005918884277343883 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.5138551444173961 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 1.0, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.21198827028274536 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.007633587786259542 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.811704834605598 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.05824012756347652 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "NOT_EVALUABLE", + "value": 0.6766043593993221 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 2.0, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.22660820186138153 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.003838771593090211 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.8771593090211133 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.8108795166015625 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "NOT_EVALUABLE", + "value": 0.02082091349852888 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 4.0, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.6140352487564087 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.003386004514672686 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.672686230248307 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.004666487375895145 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.12319141530959099 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 6.5, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.3304095268249512 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.005500550055005501 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6798679867986799 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.7280998229980469 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.7925554056432866 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 7.5, + "gnap": 1.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.5752921104431152 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0024793388429752068 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6520661157024793 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.0447002649307251 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.027162311025743513 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 0.0, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.6213451027870178 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0017699115044247787 + }, + "spike_match_pm5ms_fraction_min": { + "class": "FAIL", + "value": 0.5504424778761062 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.0025600433349609597 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.0038196085961832438 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 0.5, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.5043861269950867 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.003115264797507788 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6900311526479751 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.003333409627278648 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "NOT_EVALUABLE", + "value": null + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 1.0, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.6871345639228821 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.004178272980501393 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6462395543175488 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.2422655423482259 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 1.4091694201963283 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 2.0, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.6725149154663086 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0010976948408342481 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.7025246981339188 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.004533131917317718 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.2067227867245518 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 4.0, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.220760464668274 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0006779661016949153 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6149152542372881 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.02114248275756836 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.43095101134947444 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 6.5, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "FAIL", + "value": 2.5255849361419678 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0010608203677510608 + }, + "spike_match_pm5ms_fraction_min": { + "class": "FAIL", + "value": 0.4957567185289958 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.02339988946914673 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.0373507447149084 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 7.5, + "gnap": 5.0, + "seed": 101, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 2.218567132949829 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.00331225534477567 + }, + "spike_match_pm5ms_fraction_min": { + "class": "FAIL", + "value": 0.5441132189099669 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.02520012855529785 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "NOT_EVALUABLE", + "value": null + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 0.0, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.1461988389492035 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.006514657980456026 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.8859934853420195 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.3177619934082032 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.9142920941796274 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 0.5, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.11695905774831772 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0027397260273972603 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.8986301369863013 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.005120468139648449 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.12301673129022733 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 1.0, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.30701756477355957 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.004819277108433735 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.7325301204819277 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.00304145812988299 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.09676233111716526 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 2.0, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.3143276274204254 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.003484320557491289 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.7770034843205574 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.08386643727620442 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.37859854069758014 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 4.0, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.7236843705177307 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.006269592476489028 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.7053291536050157 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.7953335444132487 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.6240903616525737 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.047619047619047616 + } + } + }, + { + "gcal": 6.5, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.030701756477356 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.00531632110579479 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.7086656034024456 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.2829001545906067 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.1705009000027673 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 7.5, + "gnap": 1.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.5679821968078613 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.004042037186742118 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6632983023443816 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.15490001440048218 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.2873229557701684 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 0.0, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.453216552734375 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.003246753246753247 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6444805194805194 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.005598704020182282 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.03285833267333382 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.06666666666666665 + } + } + }, + { + "gcal": 0.5, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.48976612091064453 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.004285714285714286 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.75 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.01360066731770837 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.3268232694870956 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 1.0, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.42397674918174744 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0012755102040816326 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.7627551020408163 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.1857334772745769 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.0897759652728638 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 2.0, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 0.8333335518836975 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0050968399592252805 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6432212028542303 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.1349333127339681 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.14791117551575783 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.06666666666666671 + } + } + }, + { + "gcal": 4.0, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.030701756477356 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.004560260586319218 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6579804560260586 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.4800000871930804 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 1.4659453215687712 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 6.5, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 1.7105265855789185 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.001392757660167131 + }, + "spike_match_pm5ms_fraction_min": { + "class": "PASS", + "value": 0.6754874651810585 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.02610000967979431 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "PASS", + "value": 0.11586776351567529 + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + }, + { + "gcal": 7.5, + "gnap": 5.0, + "seed": 202, + "verdicts": { + "binned_rate_rmse_hz2": { + "class": "PASS", + "value": 2.0138893127441406 + }, + "active_unit_mismatch": { + "class": "PASS", + "value": 0 + }, + "total_spike_rel_diff": { + "class": "PASS", + "value": 0.0017878426698450535 + }, + "spike_match_pm5ms_fraction_min": { + "class": "FAIL", + "value": 0.5390345649582837 + }, + "hysteresis_mean_diff_hz": { + "class": "PASS", + "value": 0.014499843120574951 + }, + "self_sustained_mismatch": { + "class": "PASS", + "value": 0 + }, + "delta_f_mean_diff_hz": { + "class": "NOT_EVALUABLE", + "value": null + }, + "delta_f_eligible_fraction_diff": { + "class": "PASS", + "value": 0.0 + } + } + } + ], + "wall_s": 129.8 +} \ No newline at end of file diff --git a/experiments/twin/evidence/r2a_profiles_full_20261018_134617.json b/experiments/twin/evidence/r2a_profiles_full_20261018_134617.json new file mode 100644 index 00000000..ae9ab313 --- /dev/null +++ b/experiments/twin/evidence/r2a_profiles_full_20261018_134617.json @@ -0,0 +1,244 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "993a2cfdebaeb499", + "poolbuild.py": "816d4ba42490e283", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "983f8bbcafebd67e", + "protocols.py": "404ea437741ea2ae" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float32 (JAX default; x64 not enabled)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations", + "r2a_sha256_16": "ac4802b3bdfe1191", + "args": { + "n_mn": 8, + "true": [ + 1.6, + 0.7, + 0.0105 + ], + "dense_g": 13, + "dense_nap": 9, + "dense_w": 11, + "replicates": 2, + "dur_scale": 1.0 + }, + "run_tag": "full" + }, + "declared": { + "unknowns_profiled": [ + "gCaL_scale", + "gNaP_scale" + ], + "nuisance_reoptimized": [ + "w_exc" + ], + "fixed_known": [ + "inhibition absent", + "connectivity", + "frozen event realizations" + ], + "observable": "binned (50 ms / 300 ms)", + "resolution_rel_predeclared": 0.2, + "theta_true": [ + 1.6, + 0.7, + 0.0105 + ], + "interval_terminology": "sampled compatibility intervals (deterministic level sets)" + }, + "axes": { + "gcal": [ + 0.25, + 0.3149802624737183, + 0.39685026299204984, + 0.5, + 0.6299605249474366, + 0.7937005259840998, + 1.0, + 1.2599210498948732, + 1.5874010519681994, + 1.6, + 1.9999999999999998, + 2.5198420997897464, + 3.1748021039363983, + 4.0 + ], + "gnap": [ + 0.25, + 0.3535533905932738, + 0.5, + 0.7, + 0.7071067811865475, + 1.0, + 1.4142135623730951, + 1.9999999999999998, + 2.82842712474619, + 4.0 + ], + "w_exc": [ + 0.0045, + 0.005169142597486656, + 0.005937785598478024, + 0.006820724549296786, + 0.007834955069665113, + 0.008999999999999998, + 0.010338285194973314, + 0.0105, + 0.011875571196956045, + 0.013641449098593581, + 0.01566991013933023, + 0.018 + ] + }, + "replicates": [ + { + "conn_seed": 0, + "event_offset": 0, + "eps": 0.6250000596046448, + "profile_gcal": [ + 1.458332896232605, + 1.302082896232605, + 1.284721851348877, + 0.9722223877906799, + 1.1631944179534912, + 0.9375001490116119, + 0.9548611342906952, + 0.8680557906627655, + 0.538194552063942, + 0.0, + 0.9548611342906952, + 1.8402772545814514, + 2.4479161500930786, + 3.1076372861862183 + ], + "profile_gnap": [ + 0.9548611342906952, + 1.3368052542209625, + 1.1631942391395569, + 0.0, + 0.41666677594184875, + 0.9027779996395111, + 0.8680557906627655, + 0.9375001490116119, + 1.458332896232605, + 1.6319440603256226 + ], + "gcal_interval": [ + 1.5874010519681994, + 1.6 + ], + "gnap_interval": [ + 0.7, + 0.7071067811865475 + ], + "gcal_advances_at_pm20pct": 1.0, + "gnap_advances_at_pm20pct": 1.0, + "joint_compat_set_gcal_gnap": [ + [ + 1.5874010519681994, + 0.7 + ], + [ + 1.6, + 0.7 + ], + [ + 1.6, + 0.7071067811865475 + ] + ], + "event_hashes": { + "fit_ramp": "5ef0d36f3c87eee4", + "fit_trap": "682ace78eb9e26c4" + } + }, + { + "conn_seed": 1, + "event_offset": 10, + "eps": 0.6423613429069519, + "profile_gcal": [ + 1.215278148651123, + 1.2673611640930176, + 1.1458332538604736, + 1.0069443881511688, + 1.006944626569748, + 0.9548612833023071, + 0.9722224771976471, + 0.8854169696569443, + 0.5555558353662491, + 0.0, + 0.5729169100522995, + 1.458333671092987, + 2.2916672229766846, + 3.1249998807907104 + ], + "profile_gnap": [ + 0.5729169100522995, + 0.7465279996395111, + 0.9201391339302063, + 0.0, + 0.5555558353662491, + 0.7291669845581055, + 0.8854169696569443, + 0.9548612833023071, + 1.2673611640930176, + 1.6840280890464783 + ], + "gcal_interval": [ + 1.5874010519681994, + 1.9999999999999998 + ], + "gnap_interval": [ + 0.25, + 0.7071067811865475 + ], + "gcal_advances_at_pm20pct": 0.0, + "gnap_advances_at_pm20pct": 0.0, + "joint_compat_set_gcal_gnap": [ + [ + 1.5874010519681994, + 0.7071067811865475 + ], + [ + 1.6, + 0.7 + ], + [ + 1.6, + 0.7071067811865475 + ], + [ + 1.9999999999999998, + 0.25 + ] + ], + "event_hashes": { + "fit_ramp": "ed698d364473ebbd", + "fit_trap": "e791fdc5716d20b5" + } + } + ], + "wall_s": 1477.4 +} \ No newline at end of file diff --git a/experiments/twin/evidence/sensitivity_checks_20261013_112432.json b/experiments/twin/evidence/sensitivity_checks_20261013_112432.json new file mode 100644 index 00000000..7a9daf95 --- /dev/null +++ b/experiments/twin/evidence/sensitivity_checks_20261013_112432.json @@ -0,0 +1,971 @@ +{ + "manifest": { + "git_revision": "a483b1d", + "git_dirty": true, + "provenance_note": "experiments/twin sources are untracked by git; reconstruction relies on the file hashes below, not the revision", + "source_sha256_16": { + "sensitivity_checks.py": "993a2cfdebaeb499", + "poolbuild.py": "816d4ba42490e283", + "observables.py": "9ab3a99e798e3274", + "validation_protocol.py": "f7cf30f9880ce289", + "protocols.py": "404ea437741ea2ae" + }, + "versions": { + "jax": "0.9.0", + "jaxley": "0.13.0", + "numpy": "2.3.5" + }, + "device": "cpu", + "precision": "float32 (JAX default; x64 not enabled)", + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest\u22480, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng (frozen renewal events)", + "fd_probes": "deterministic, no sampling" + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target" + }, + "scope": "this Jaxley harness configuration only \u2014 NOT a validated NEURON-matched forward model; conclusions are specific to the stated parameter points, objectives, and frozen input realizations" + }, + "t1": { + "ok": true, + "leaf_map": { + "gCaL_scale": [ + "caL_gcaLbar" + ], + "gNaP_scale": [ + "napp_gnapbar" + ], + "drive_gain": [], + "w_exc": [], + "w_inh": [] + }, + "plumbing_ok": { + "gCaL_scale": true, + "gNaP_scale": true, + "drive_gain": true, + "w_exc": true, + "w_inh": true + }, + "max_output_effect_mV": { + "gCaL_scale": 2.3932723999023438, + "gNaP_scale": 57.749900817871094, + "drive_gain": 0.0, + "w_exc": 83.46661376953125, + "w_inh": 58.313323974609375, + "drive_gain (rate path)": 1.1866259574890137 + }, + "initialization_audit": { + "initial_state_theta_dependent": false, + "recorded_compartments_at_rest": true, + "nonrecorded_init_v_range_mV": [ + -1.0, + -1.0 + ], + "observation": "non-recorded compartment(s) \u2014 the axon \u2014 initialize at \u22121.0 mV, not the NERLab rest; a jaxley init quirk, decays within ms but is part of x0 and is recorded here rather than hidden", + "caveat": "theta changes do not re-equilibrate the initial state; fits must either discard the initial transient or treat x0 explicitly (roadmap x0_sc)" + }, + "runtime_s": 5.6 + }, + "t2": { + "ok": true, + "loss0": 1.6953394412994385, + "subthreshold": true, + "autodiff": [ + 0.08275409042835236, + 0.02414565719664097 + ], + "fd": { + "gCaL_scale": { + "0.1": { + "fd": 0.08278548717498779, + "rel_err_vs_autodiff": 0.0003793981159105897 + }, + "0.01": { + "fd": 0.08283257484436035, + "rel_err_vs_autodiff": 0.000948405300732702 + }, + "0.001": { + "fd": 0.08368492126464844, + "rel_err_vs_autodiff": 0.011248154565691948 + }, + "0.0001": { + "fd": 0.09179115295410156, + "rel_err_vs_autodiff": 0.10920381546020508 + } + }, + "gNaP_scale": { + "0.1": { + "fd": 0.024253129959106445, + "rel_err_vs_autodiff": 0.004451018292456865 + }, + "0.01": { + "fd": 0.024443864822387695, + "rel_err_vs_autodiff": 0.012350362725555897 + }, + "0.001": { + "fd": 0.023066997528076172, + "rel_err_vs_autodiff": 0.04467302933335304 + }, + "0.0001": { + "fd": 0.019669532775878906, + "rel_err_vs_autodiff": 0.18538010120391846 + } + } + }, + "runtime_s": 9.3 + }, + "t3": { + "ok": true, + "n_spikes_at_init": 13, + "surrogate_grad": 85518612168704.0, + "hard_slope_fixed_key_eps5pct": 5.43792724609375, + "hard_slope_expected_8keys_eps5pct": 7.174005508422852, + "hard_slope_per_key": [ + 6.098613739013672, + 8.701519966125488, + 9.532957077026367, + 4.320478439331055, + 5.040860176086426, + 13.226909637451172, + 5.43668270111084, + 5.034022331237793 + ], + "hard_slope_sd_across_keys": 2.8719073159831106, + "note": "pathwise/fixed-key derivatives are fixed-realization quantities; agreement in sign with the expected-value slope is required, magnitudes are not expected to match", + "runtime_s": 31.9 + }, + "t4": { + "ok": false, + "per_horizon": { + "0.5": { + "n_spikes": 10, + "grad_gCaL": -21273692340224.0, + "grad_gNaP": -693043439075328.0, + "components_finite": true, + "grad_norm_f64": 693369871303292.5 + }, + "1.0": { + "n_spikes": 19, + "grad_gCaL": -1.4362193689672014e+27, + "grad_gNaP": -4.6639390035603405e+28, + "components_finite": true, + "grad_norm_f64": 4.666149835752104e+28 + }, + "2.0": { + "n_spikes": 37, + "grad_gCaL": NaN, + "grad_gNaP": NaN, + "components_finite": false, + "grad_norm_f64": NaN + }, + "4.0": { + "n_spikes": 75, + "grad_gCaL": NaN, + "grad_gNaP": NaN, + "components_finite": false, + "grad_norm_f64": NaN + } + }, + "norm_growth_shortest_to_longest": NaN, + "note": "distinguish: (a) dynamical explosion of finite components (e.g. ~1e28 at 1 s), (b) NaN components (true breakdown), (c) diagnostic float32 norm overflow \u2014 eliminated here by the float64 norm. 'ok' means all components finite at every horizon, NOT that magnitudes are usable.", + "runtime_s": 8.7 + }, + "t5": { + "n_independent_event_realizations": 4, + "per_event_seed_joint_improvement": { + "0": { + "n_steps": 8, + "improved_fraction": 0.0 + }, + "1": { + "n_steps": 8, + "improved_fraction": 1.0 + }, + "2": { + "n_steps": 8, + "improved_fraction": 0.5 + }, + "3": { + "n_steps": 8, + "improved_fraction": 0.5 + } + }, + "ok": false, + "n_cases": 16, + "n_evaluated": 16, + "n_failures": 0, + "accepted_step_fraction": 0.5, + "coordinate_agreement": { + "gCaL_scale": { + "n_resolved": 12, + "n_unresolved": 4, + "agreement_rate": 0.6666666666666666 + }, + "w_exc": { + "n_resolved": 16, + "n_unresolved": 0, + "agreement_rate": 0.5 + } + }, + "cases": [ + { + "conn_seed": 0, + "event_seed": 0, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 40, + "loss_at_init_hard": 15.786011695861816, + "surrogate_grad": [ + 312621140017152.0, + 2.516758579558482e+17 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 17.777145385742188, + "improved": false + }, + "0.1": { + "hard_loss": 25.474170684814453, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 15.786011695861816, + "hard_loss_up": 15.788641929626465, + "hard_loss_down": 15.929408073425293, + "true_descent_dir": 0, + "surrogate_dir": -1, + "agrees": null + }, + "w_exc": { + "hard_loss_baseline": 15.786011695861816, + "hard_loss_up": 9.060354232788086, + "hard_loss_down": 24.288192749023438, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + } + } + }, + { + "conn_seed": 0, + "event_seed": 0, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 11, + "loss_at_init_hard": 5.275916576385498, + "surrogate_grad": [ + -666105001541632.0, + -5.3622305913307136e+17 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 7.321145534515381, + "improved": false + }, + "0.1": { + "hard_loss": 11.727997779846191, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 5.275916576385498, + "hard_loss_up": 5.112833499908447, + "hard_loss_down": 5.3906683921813965, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 5.275916576385498, + "hard_loss_up": 10.66342544555664, + "hard_loss_down": 1.6630901098251343, + "true_descent_dir": -1, + "surrogate_dir": 1, + "agrees": false + } + } + }, + { + "conn_seed": 0, + "event_seed": 1, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 41, + "loss_at_init_hard": 15.428475379943848, + "surrogate_grad": [ + -281944134778880.0, + -2.0783693778701517e+17 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 14.013660430908203, + "improved": true + }, + "0.1": { + "hard_loss": 8.015419960021973, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 15.428475379943848, + "hard_loss_up": 15.359977722167969, + "hard_loss_down": 16.232988357543945, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 15.428475379943848, + "hard_loss_up": 8.818695068359375, + "hard_loss_down": 22.458179473876953, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + } + } + }, + { + "conn_seed": 0, + "event_seed": 1, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 12, + "loss_at_init_hard": 5.6394453048706055, + "surrogate_grad": [ + 305276175515648.0, + 2.2504273139361382e+17 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 4.722307205200195, + "improved": true + }, + "0.1": { + "hard_loss": 2.1072397232055664, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 5.6394453048706055, + "hard_loss_up": 5.618084907531738, + "hard_loss_down": 4.980186939239502, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 5.6394453048706055, + "hard_loss_up": 11.663798332214355, + "hard_loss_down": 2.678187608718872, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + } + } + }, + { + "conn_seed": 0, + "event_seed": 2, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 40, + "loss_at_init_hard": 15.472694396972656, + "surrogate_grad": [ + -60513619968.0, + -56263417266176.0 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 12.96433162689209, + "improved": true + }, + "0.1": { + "hard_loss": 7.901159286499023, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 15.472694396972656, + "hard_loss_up": 14.62053108215332, + "hard_loss_down": 16.98008155822754, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 15.472694396972656, + "hard_loss_up": 8.153935432434082, + "hard_loss_down": 23.17388343811035, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + } + } + }, + { + "conn_seed": 0, + "event_seed": 2, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 10, + "loss_at_init_hard": 6.098230361938477, + "surrogate_grad": [ + 2114921431040.0, + 1966115903242240.0 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 4.076164722442627, + "improved": true + }, + "0.1": { + "hard_loss": 2.836045265197754, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 6.098230361938477, + "hard_loss_up": 6.694512367248535, + "hard_loss_down": 5.321730136871338, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 6.098230361938477, + "hard_loss_up": 12.768681526184082, + "hard_loss_down": 2.906851053237915, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + } + } + }, + { + "conn_seed": 0, + "event_seed": 3, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 41, + "loss_at_init_hard": 14.968521118164062, + "surrogate_grad": [ + 1.012451609893929e+17, + 7.438539209532506e+19 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 17.073200225830078, + "improved": false + }, + "0.1": { + "hard_loss": 24.337709426879883, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 14.968521118164062, + "hard_loss_up": 14.857789039611816, + "hard_loss_down": 15.562369346618652, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + }, + "w_exc": { + "hard_loss_baseline": 14.968521118164062, + "hard_loss_up": 8.291171073913574, + "hard_loss_down": 22.110092163085938, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + } + } + }, + { + "conn_seed": 0, + "event_seed": 3, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 12, + "loss_at_init_hard": 6.180650234222412, + "surrogate_grad": [ + -8.597460233696051e+16, + -6.316541689359185e+19 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 8.135015487670898, + "improved": false + }, + "0.1": { + "hard_loss": 16.16530418395996, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 6.180650234222412, + "hard_loss_up": 6.673616886138916, + "hard_loss_down": 6.417392253875732, + "true_descent_dir": 0, + "surrogate_dir": 1, + "agrees": null + }, + "w_exc": { + "hard_loss_baseline": 6.180650234222412, + "hard_loss_up": 13.60311222076416, + "hard_loss_down": 3.2830843925476074, + "true_descent_dir": -1, + "surrogate_dir": 1, + "agrees": false + } + } + }, + { + "conn_seed": 1, + "event_seed": 0, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 51, + "loss_at_init_hard": 16.66646957397461, + "surrogate_grad": [ + 5.70238512922624e+16, + 4.45495503393443e+19 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 18.89405059814453, + "improved": false + }, + "0.1": { + "hard_loss": 25.002363204956055, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 16.66646957397461, + "hard_loss_up": 15.886414527893066, + "hard_loss_down": 16.795991897583008, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + }, + "w_exc": { + "hard_loss_baseline": 16.66646957397461, + "hard_loss_up": 10.28009033203125, + "hard_loss_down": 24.299158096313477, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + } + } + }, + { + "conn_seed": 1, + "event_seed": 0, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 20, + "loss_at_init_hard": 5.717686653137207, + "surrogate_grad": [ + -5.841590173250355e+16, + -4.563888488749361e+19 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 6.890046119689941, + "improved": false + }, + "0.1": { + "hard_loss": 12.25369644165039, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 5.717686653137207, + "hard_loss_up": 5.802124500274658, + "hard_loss_down": 5.316225528717041, + "true_descent_dir": -1, + "surrogate_dir": 1, + "agrees": false + }, + "w_exc": { + "hard_loss_baseline": 5.717686653137207, + "hard_loss_up": 11.21438980102539, + "hard_loss_down": 3.126026153564453, + "true_descent_dir": -1, + "surrogate_dir": 1, + "agrees": false + } + } + }, + { + "conn_seed": 1, + "event_seed": 1, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 51, + "loss_at_init_hard": 15.88136100769043, + "surrogate_grad": [ + -1.4785854138331693e+19, + -1.0807669705869399e+22 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 13.869335174560547, + "improved": true + }, + "0.1": { + "hard_loss": 7.843212604522705, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 15.88136100769043, + "hard_loss_up": 15.161964416503906, + "hard_loss_down": 16.50951385498047, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 15.88136100769043, + "hard_loss_up": 7.944602012634277, + "hard_loss_down": 23.017820358276367, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + } + } + }, + { + "conn_seed": 1, + "event_seed": 1, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 18, + "loss_at_init_hard": 7.505717754364014, + "surrogate_grad": [ + 1.1224109365453652e+19, + 8.197015755525875e+21 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 7.144341945648193, + "improved": true + }, + "0.1": { + "hard_loss": 3.0422792434692383, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 7.505717754364014, + "hard_loss_up": 8.095317840576172, + "hard_loss_down": 7.290937423706055, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 7.505717754364014, + "hard_loss_up": 16.495311737060547, + "hard_loss_down": 3.620619297027588, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + } + } + }, + { + "conn_seed": 1, + "event_seed": 2, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 52, + "loss_at_init_hard": 17.327024459838867, + "surrogate_grad": [ + 7.253975050255073e+17, + 5.073399274350062e+20 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 19.82647705078125, + "improved": false + }, + "0.1": { + "hard_loss": 28.647083282470703, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 17.327024459838867, + "hard_loss_up": 15.738357543945312, + "hard_loss_down": 17.62483024597168, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + }, + "w_exc": { + "hard_loss_baseline": 17.327024459838867, + "hard_loss_up": 10.898860931396484, + "hard_loss_down": 26.579204559326172, + "true_descent_dir": 1, + "surrogate_dir": -1, + "agrees": false + } + } + }, + { + "conn_seed": 1, + "event_seed": 2, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 18, + "loss_at_init_hard": 6.723733425140381, + "surrogate_grad": [ + -2.7269879515683226e+17, + -1.9061666622361882e+20 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 9.064820289611816, + "improved": false + }, + "0.1": { + "hard_loss": 16.558732986450195, + "improved": false + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 6.723733425140381, + "hard_loss_up": 8.374335289001465, + "hard_loss_down": 7.463498115539551, + "true_descent_dir": 0, + "surrogate_dir": 1, + "agrees": null + }, + "w_exc": { + "hard_loss_baseline": 6.723733425140381, + "hard_loss_up": 13.964120864868164, + "hard_loss_down": 3.6688783168792725, + "true_descent_dir": -1, + "surrogate_dir": 1, + "agrees": false + } + } + }, + { + "conn_seed": 1, + "event_seed": 3, + "true_theta": { + "gCaL_scale": 1.5, + "w_exc": 0.011 + }, + "n_obs_spikes": 50, + "loss_at_init_hard": 16.04979133605957, + "surrogate_grad": [ + -3.936402810286899e+18, + -3.695528316380616e+21 + ], + "grad_finite": true, + "joint_sign_direction": [ + 1.0, + 1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 12.820548057556152, + "improved": true + }, + "0.1": { + "hard_loss": 7.246275424957275, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 16.04979133605957, + "hard_loss_up": 15.06729793548584, + "hard_loss_down": 15.952129364013672, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + }, + "w_exc": { + "hard_loss_baseline": 16.04979133605957, + "hard_loss_up": 7.09161376953125, + "hard_loss_down": 23.715425491333008, + "true_descent_dir": 1, + "surrogate_dir": 1, + "agrees": true + } + } + }, + { + "conn_seed": 1, + "event_seed": 3, + "true_theta": { + "gCaL_scale": 0.7, + "w_exc": 0.008 + }, + "n_obs_spikes": 18, + "loss_at_init_hard": 6.795099258422852, + "surrogate_grad": [ + 3.9213070653933486e+18, + 3.678141325594222e+21 + ], + "grad_finite": true, + "joint_sign_direction": [ + -1.0, + -1.0 + ], + "joint_steps": { + "0.03": { + "hard_loss": 6.152483940124512, + "improved": true + }, + "0.1": { + "hard_loss": 3.3677709102630615, + "improved": true + } + }, + "coordinate_probes": { + "gCaL_scale": { + "hard_loss_baseline": 6.795099258422852, + "hard_loss_up": 7.39927864074707, + "hard_loss_down": 7.122353553771973, + "true_descent_dir": 0, + "surrogate_dir": -1, + "agrees": null + }, + "w_exc": { + "hard_loss_baseline": 6.795099258422852, + "hard_loss_up": 16.118892669677734, + "hard_loss_down": 2.799283981323242, + "true_descent_dir": -1, + "surrogate_dir": -1, + "agrees": true + } + } + } + ], + "runtime_s": 137.0 + } +} \ No newline at end of file diff --git a/experiments/twin/evidence/val_comparison_bac411989ee8d42b_20261018_124637.json b/experiments/twin/evidence/val_comparison_bac411989ee8d42b_20261018_124637.json new file mode 100644 index 00000000..1fffbbe4 --- /dev/null +++ b/experiments/twin/evidence/val_comparison_bac411989ee8d42b_20261018_124637.json @@ -0,0 +1,153 @@ +{ + "artifact": { + "version": 2, + "seed": 0, + "dt_ms": 0.1, + "n_mn": 8, + "synapse": { + "site": "dendrite", + "tau_rise_ms": 0.2, + "tau_decay_ms": 0.3, + "e_exc_mV": 70.0, + "e_inh_mV": -15.0, + "w_exc_uS": 0.15, + "w_inh_uS": 0.009, + "live_driving_force": true, + "e_inh_convention": "frame-corrected driving force" + }, + "cell": { + "gcal_scale": 1.0, + "gnap_scale": 1.0 + }, + "content_hash": "bac411989ee8d42b" + }, + "comparison": { + "content_hash": "bac411989ee8d42b", + "pool": { + "v_rmse_mV": 1.730110819851262, + "v_max_abs_mV": 71.49909502714058, + "v_dend_rmse_mV": 0.38114991314188384, + "v_dend_max_abs_mV": 8.053318795412174, + "spikes_jaxley_per_mn": [ + 135, + 96, + 65, + 72, + 39, + 0, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 135, + 95, + 65, + 71, + 39, + 0, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 356, + "jaxley_total": 407, + "neuron_total": 405 + } + }, + "current_steps": { + "5.0": { + "v_rmse_mV": 0.0027762454144517677, + "v_max_abs_mV": 0.027068122664278604, + "v_dend_rmse_mV": 0.003062220301082907, + "v_dend_max_abs_mV": 0.0324280108163979, + "spikes_jaxley_per_mn": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 0, + "jaxley_total": 0, + "neuron_total": 0 + } + }, + "15.0": { + "v_rmse_mV": 3.085067592487133, + "v_max_abs_mV": 68.34554343209372, + "v_dend_rmse_mV": 0.6065513667095238, + "v_dend_max_abs_mV": 7.637296967278601, + "spikes_jaxley_per_mn": [ + 67, + 65, + 62, + 54, + 0, + 0, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 67, + 65, + 62, + 54, + 0, + 0, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 180, + "jaxley_total": 248, + "neuron_total": 248 + } + }, + "30.0": { + "v_rmse_mV": 5.2833486490138775, + "v_max_abs_mV": 74.35096933079032, + "v_dend_rmse_mV": 0.8682390814816492, + "v_dend_max_abs_mV": 5.383255883702834, + "spikes_jaxley_per_mn": [ + 145, + 139, + 132, + 124, + 113, + 92, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 144, + 139, + 132, + 124, + 113, + 92, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 462, + "jaxley_total": 745, + "neuron_total": 744 + } + } + } + } +} \ No newline at end of file diff --git a/experiments/twin/evidence/val_comparison_ef1bd7e84b9abd03_20261018_124343.json b/experiments/twin/evidence/val_comparison_ef1bd7e84b9abd03_20261018_124343.json new file mode 100644 index 00000000..4b9a0f82 --- /dev/null +++ b/experiments/twin/evidence/val_comparison_ef1bd7e84b9abd03_20261018_124343.json @@ -0,0 +1,153 @@ +{ + "artifact": { + "version": 2, + "seed": 0, + "dt_ms": 0.1, + "n_mn": 8, + "synapse": { + "site": "dendrite", + "tau_rise_ms": 0.2, + "tau_decay_ms": 0.3, + "e_exc_mV": 70.0, + "e_inh_mV": -15.0, + "w_exc_uS": 0.009, + "w_inh_uS": 0.009, + "live_driving_force": true, + "e_inh_convention": "frame-corrected driving force" + }, + "cell": { + "gcal_scale": 1.0, + "gnap_scale": 1.0 + }, + "content_hash": "ef1bd7e84b9abd03" + }, + "comparison": { + "content_hash": "ef1bd7e84b9abd03", + "pool": { + "v_rmse_mV": 0.0003514303395716632, + "v_max_abs_mV": 0.001064357153899298, + "v_dend_rmse_mV": 0.00038958527107967175, + "v_dend_max_abs_mV": 0.0011031375894381412, + "spikes_jaxley_per_mn": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 0, + "jaxley_total": 0, + "neuron_total": 0 + } + }, + "current_steps": { + "5.0": { + "v_rmse_mV": 0.0027762454144517677, + "v_max_abs_mV": 0.027068122664278604, + "v_dend_rmse_mV": 0.003062220301082907, + "v_dend_max_abs_mV": 0.0324280108163979, + "spikes_jaxley_per_mn": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 0, + "jaxley_total": 0, + "neuron_total": 0 + } + }, + "15.0": { + "v_rmse_mV": 3.085067592487133, + "v_max_abs_mV": 68.34554343209372, + "v_dend_rmse_mV": 0.6065513667095238, + "v_dend_max_abs_mV": 7.637296967278601, + "spikes_jaxley_per_mn": [ + 67, + 65, + 62, + 54, + 0, + 0, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 67, + 65, + 62, + 54, + 0, + 0, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 180, + "jaxley_total": 248, + "neuron_total": 248 + } + }, + "30.0": { + "v_rmse_mV": 5.2833486490138775, + "v_max_abs_mV": 74.35096933079032, + "v_dend_rmse_mV": 0.8682390814816492, + "v_dend_max_abs_mV": 5.383255883702834, + "spikes_jaxley_per_mn": [ + 145, + 139, + 132, + 124, + 113, + 92, + 0, + 0 + ], + "spikes_neuron_per_mn": [ + 144, + 139, + 132, + 124, + 113, + 92, + 0, + 0 + ], + "spike_match_pm5ms": { + "matched": 462, + "jaxley_total": 745, + "neuron_total": 744 + } + } + } + } +} \ No newline at end of file diff --git a/experiments/twin/g4_descent.py b/experiments/twin/g4_descent.py new file mode 100644 index 00000000..20a8700c --- /dev/null +++ b/experiments/twin/g4_descent.py @@ -0,0 +1,507 @@ +"""G4 v2 — descent benchmark with the corrected full-trace objective. + +v2 fixes from review of the v1 run: + +* **Full-trace observation operator**: windows are simulated with preserved + state and boundary-stopped gradients, but the per-window SPIKE outputs are + CONCATENATED and ``binned_rates_jnp`` is applied ONCE to the complete + sequence — identical to Deliverable B's objective. (v1 restarted the + 300 ms smoothing at each boundary, silently changing the objective by + ~2.3 Hz² at the v1 operating point.) A regression test asserts the + chained hard spike train equals the uninterrupted run exactly. +* **Coordinate-resolved gating**: besides the joint direction, Stage A now + probes gCaL-only and w_exc-only steps from three start types — + S1 both-wrong, S2 gCaL-wrong-only (w at truth), S3 w-wrong-only — and + records the directional cosine against the CRN-FD reference. 16/16 joint + acceptances alone cannot certify the gCaL coordinate (v1 directions were + ~[0.07, 0.998]). +* **Stage B across all cases and two starts** with success/failure counts, + parameter and prediction error separated, and the best loss compared to + the per-case dt-halving EPS of THIS objective (truth loss is exactly 0, + so any plateau needs that reference). +* **Profile of the exact G4 objective** (case 0): dense gCaL profile with + w re-optimized, truth on-axis — required before attributing any gCaL + stall to identifiability. +* Event/connectivity matrices are float64 under x64 (v1 left them f32). + +Scoped claims: joint-descent reliability, per-coordinate direction quality, +and recovery statistics for THIS configuration; nothing about other +objectives or regimes. + +Usage:: + + python g4_descent.py # full (CPU ~1.5 h) + python g4_descent.py --smoke +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + +RESULTS = Path(__file__).parent / "results" +EVIDENCE = Path(__file__).parent / "evidence" + +DT = 0.025 +BASE_DT = 0.1 +FACTOR = int(round(BASE_DT / DT)) +WINDOW_S = 2.0 +HAZARD_SIGMA_MV = 2.0 +MAX_SPK = 64 +LAMBDAS = (0.03, 0.1) +EPS_FLOOR = 1e-3 +OKABE = {"surrogate": "#0072B2", "crn_fd": "#CC79A7"} + + +def build_case(conn_seed, ev_off, dur_scale): + import jax.numpy as jnp + from deliverable_b import make_commands, build_twin_and_events + from gradient_repair import zero_insert_events + from poolbuild import PoolTwin, TwinSpec + + commands = make_commands(BASE_DT, dur_scale) + _, events0 = build_twin_and_events(8, conn_seed, ev_off, commands, BASE_DT) + twin = PoolTwin(TwinSpec(n_mn=8, n_dd=100, n_inh=5, dt_ms=DT, + seed=conn_seed)) + ev = {} + for proto in ("fit_ramp", "heldout_slowramp"): + fine = zero_insert_events( + np.asarray(events0[proto], dtype=np.float64), FACTOR) + ev[proto] = { + "exc": jnp.asarray(fine), + "inh": jnp.zeros((fine.shape[0], 5), dtype=fine.dtype), + "exc_conn": jnp.asarray(np.asarray(twin._exc_mat, + dtype=np.float64)), + "inh_conn": jnp.asarray(np.asarray(twin._inh_mat, + dtype=np.float64)), + } + return twin, ev + + +def split_windows(ev, n_win): + n = ev["exc"].shape[0] + w = n // n_win + return [{k: (v[i * w:(i + 1) * w] if k in ("exc", "inh") else v) + for k, v in ev.items()} for i in range(n_win)] + + +def binned(spikes, dt_ms=DT, bin_ms=50.0, smooth_ms=300.0): + from deliverable_b import binned_rates_jnp + + return binned_rates_jnp(spikes, dt_ms, bin_ms, smooth_ms) + + +def chained_spikes(twin, theta, windows, spike_mode, stop_boundary=True): + """State-preserving windowed run; returns FULL concatenated spike train. + Gradients are stopped at boundaries iff stop_boundary.""" + import jax + + carry, segs = None, [] + for w in windows: + out, carry = twin._run_events_carry(theta, w, carry, spike_mode) + if stop_boundary: + carry = jax.tree_util.tree_map(jax.lax.stop_gradient, carry) + segs.append(out["spikes"]) + import jax.numpy as jnp + + return jnp.concatenate(segs, axis=0) + + +class Case: + def __init__(self, conn_seed, ev_off, theta_true, dur_scale=1.0): + import jax + import jax.numpy as jnp + + self.meta = {"conn_seed": conn_seed, "event_offset": ev_off, + "theta_true": list(theta_true)} + twin, ev = build_case(conn_seed, ev_off, dur_scale) + self.twin = twin + n_win = max(1, int(round( + ev["fit_ramp"]["exc"].shape[0] * DT * 1e-3 / WINDOW_S))) + self.windows = split_windows(ev["fit_ramp"], n_win) + self.heldout = ev["heldout_slowramp"] + self.theta_true = np.asarray(theta_true, dtype=np.float64) + + def theta_of(vec): + th = twin.default_theta() + th["gCaL_scale"] = vec[0] + th["w_exc"] = vec[1] + return th + + self.theta_of = theta_of + + # observations: full-trace operator on the chained hard run + tt = jnp.asarray(self.theta_true) + obs_sp = chained_spikes(twin, theta_of(tt), self.windows, "hard") + self.obs_rates = jax.lax.stop_gradient(binned(obs_sp)) + self.obs_spikes = np.asarray(obs_sp) + self.n_obs_spikes = int(self.obs_spikes.sum()) + out_h, _ = twin._run_events_carry(theta_of(tt), self.heldout, None, + "hard") + self.obs_heldout = jax.lax.stop_gradient(binned(out_h["spikes"])) + + # observed spike indices (hazard NLL), padded, full-trace + idx = np.zeros((self.obs_spikes.shape[1], MAX_SPK * n_win), + dtype=np.int32) + msk = np.zeros_like(idx, dtype=np.float64) + for i in range(self.obs_spikes.shape[1]): + ks = np.flatnonzero(self.obs_spikes[:, i] > 0.5)[:idx.shape[1]] + idx[i, :len(ks)] = ks + msk[i, :len(ks)] = 1.0 + self.obs_idx, self.obs_msk = jnp.asarray(idx), jnp.asarray(msk) + + def hard_loss(vec): + sp = chained_spikes(twin, theta_of(vec), self.windows, "hard") + return jnp.mean((binned(sp) - self.obs_rates) ** 2) + + def sur_loss(vec): + sp = chained_spikes(twin, theta_of(vec), self.windows, "surrogate") + return jnp.mean((binned(sp) - self.obs_rates) ** 2) + + self.hard_loss = jax.jit(hard_loss) + self._sur_grad = jax.jit(jax.grad(sur_loss)) + self._eval_count = 0 + + def hard(self, vec): + self._eval_count += 1 + return float(self.hard_loss(np.asarray(vec, dtype=np.float64))) + + def grad(self, name, vec): + import jax.numpy as jnp + + if name == "crn_fd": + g = np.zeros(2) + for j in range(2): + h = 0.05 * float(vec[j]) + e = np.zeros(2) + e[j] = h + g[j] = (self.hard(vec + e) - self.hard(vec - e)) / (2 * h) + return g + self._eval_count += 3 + return np.asarray(self._sur_grad(jnp.asarray(vec, dtype=jnp.float64)), + dtype=np.float64) + + def direction(self, name, vec): + g = self.grad(name, vec) + if not np.all(np.isfinite(g)): + return None, g + gl = g * np.asarray(vec) + n = np.linalg.norm(gl) + return (-gl / n if n > 0 else None), g + + def heldout_mse(self, vec): + import jax.numpy as jnp + + out, _ = self.twin._run_events_carry( + self.theta_of(jnp.asarray(vec, dtype=jnp.float64)), + self.heldout, None, "hard") + return float(jnp.mean((binned(out["spikes"]) - self.obs_heldout) ** 2)) + + def eps_dt_halving(self): + """Numerical tolerance of THIS objective: rerun truth at dt/2.""" + import jax.numpy as jnp + from gradient_repair import zero_insert_events + from poolbuild import PoolTwin, TwinSpec + + twin_h = PoolTwin(TwinSpec(n_mn=8, n_dd=100, n_inh=5, dt_ms=DT / 2, + seed=self.meta["conn_seed"])) + ev_f = { + "exc": jnp.asarray(zero_insert_events( + np.concatenate([np.asarray(w["exc"]) for w in self.windows]), + 2)), + "inh": None, "exc_conn": None, "inh_conn": None, + } + ev_f["inh"] = jnp.zeros((ev_f["exc"].shape[0], 5), + dtype=ev_f["exc"].dtype) + ev_f["exc_conn"] = self.windows[0]["exc_conn"] + ev_f["inh_conn"] = self.windows[0]["inh_conn"] + th = self.theta_of(jnp.asarray(self.theta_true)) + out, _ = twin_h._run_events_carry(th, ev_f, None, "hard") + r_f = np.asarray(binned(out["spikes"], DT / 2)) + r_c = np.asarray(self.obs_rates) + m = min(len(r_f), len(r_c)) + return max(float(((r_f[:m] - r_c[:m]) ** 2).mean()), EPS_FLOOR) + + +def main(): + import jax + + jax.config.update("jax_enable_x64", True) + + ap = argparse.ArgumentParser() + ap.add_argument("--smoke", action="store_true") + ap.add_argument("--iters", type=int, default=10) + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + from sensitivity_checks import _manifest + + t0_all = time.time() + dur_scale = 0.5 if args.smoke else 1.0 + case_specs = ([(0, 0, (1.5, 0.011))] if args.smoke else + [(c, e, t) for c in (0, 1) for e in (0, 10) + for t in ((1.5, 0.011), (0.7, 0.008))]) + estimators = ["surrogate", "crn_fd"] + + def starts_for(tt): + return {"S1_both_wrong": np.array([1.0, 0.009]), + "S2_gcal_wrong_only": np.array([1.0, tt[1]]), + "S3_w_wrong_only": np.array([tt[0], 0.009])} + + # ------------------------------------------------- Stage A (coordinate-resolved) + rowsA = [] + tallies = {e: {s: {"joint_acc": 0, "joint_tot": 0, + "coord_acc": {"gCaL": 0, "w": 0}, + "coord_tot": {"gCaL": 0, "w": 0}, + "nonfinite": 0} + for s in ("S1_both_wrong", "S2_gcal_wrong_only", + "S3_w_wrong_only")} + for e in estimators} + cosines = [] + case_objs = {} + for ci, (cs, eo, tt) in enumerate(case_specs): + t0 = time.time() + case = Case(cs, eo, tt, dur_scale) + case_objs[ci] = case + row = {"case": case.meta, "n_obs_spikes": case.n_obs_spikes, + "starts": {}} + for sname, s0 in starts_for(tt).items(): + L0 = case.hard(s0) + srow = {"L0": L0, "estimators": {}} + dirs = {} + for est in estimators: + d, g = case.direction(est, s0) + if d is None: + tallies[est][sname]["nonfinite"] += 1 + srow["estimators"][est] = {"nonfinite": True} + continue + dirs[est] = d + erow = {"grad": g.tolist(), "direction": d.tolist(), + "joint": {}, "coord": {}} + for lam in LAMBDAS: # joint steps + Ls = case.hard(s0 * np.exp(lam * d)) + ok = bool(Ls < L0) + erow["joint"][lam] = {"loss": Ls, "accepted": ok} + tallies[est][sname]["joint_acc"] += ok + tallies[est][sname]["joint_tot"] += 1 + # coordinate-only steps (λ=0.1, sign from that coordinate) + for j, cname in ((0, "gCaL"), (1, "w")): + if d[j] == 0: + continue + dc = np.zeros(2) + dc[j] = np.sign(d[j]) + Ls = case.hard(s0 * np.exp(0.1 * dc)) + ok = bool(Ls < L0) + erow["coord"][cname] = {"loss": Ls, "accepted": ok} + tallies[est][sname]["coord_acc"][cname] += ok + tallies[est][sname]["coord_tot"][cname] += 1 + srow["estimators"][est] = erow + if "surrogate" in dirs and "crn_fd" in dirs: + c = float(dirs["surrogate"] @ dirs["crn_fd"]) + srow["cosine_sur_vs_fd"] = c + cosines.append(c) + row["starts"][sname] = srow + rowsA.append(row) + print(f"[case {ci}] conn{cs}/ev{eo}/true{tt} spk {case.n_obs_spikes} " + f"({time.time()-t0:.0f}s)", flush=True) + + gate = {} + for est in estimators: + agg = {"joint_acc": 0, "joint_tot": 0, "nonfinite": 0, + "coord": {"gCaL": [0, 0], "w": [0, 0]}} + for s in tallies[est].values(): + agg["joint_acc"] += s["joint_acc"] + agg["joint_tot"] += s["joint_tot"] + agg["nonfinite"] += s["nonfinite"] + for cname in ("gCaL", "w"): + agg["coord"][cname][0] += s["coord_acc"][cname] + agg["coord"][cname][1] += s["coord_tot"][cname] + jf = agg["joint_acc"] / max(agg["joint_tot"], 1) + cf = {c: (a / max(t, 1)) for c, (a, t) in agg["coord"].items()} + gate[est] = { + "joint_accepted_fraction": jf, + "coord_accepted_fraction": cf, + "nonfinite": agg["nonfinite"], + "per_start": tallies[est], + "joint_pass": bool(jf >= 0.8 and agg["nonfinite"] == 0), + "gcal_pass": bool(cf["gCaL"] >= 0.8 and agg["nonfinite"] == 0), + } + print(f"Stage A — {est:>9}: joint {jf:.2f} | gCaL-coord " + f"{cf['gCaL']:.2f} | w-coord {cf['w']:.2f} | nonfinite " + f"{agg['nonfinite']} → joint " + f"{'PASS' if gate[est]['joint_pass'] else 'fail'}, gCaL " + f"{'PASS' if gate[est]['gcal_pass'] else 'fail'}", flush=True) + print(f"Stage A — cosine(surrogate, CRN-FD): mean " + f"{np.mean(cosines):.3f} min {np.min(cosines):.3f}", flush=True) + + # ------------------------------------------------- Stage B (all cases, 2 starts) + stageB = [] + for ci, (cs, eo, tt) in enumerate(case_specs): + case = case_objs[ci] + eps = case.eps_dt_halving() + for sname in ("S1_both_wrong", "S2_gcal_wrong_only"): + s0 = starts_for(np.asarray(tt))[sname] + for est in estimators: + case._eval_count = 0 + t0 = time.time() + vec = s0.copy() + L = case.hard(vec) + best = {"loss": L, "theta": vec.tolist(), "it": 0} + lam, nonfinite = 0.1, 0 + for it in range(1, args.iters + 1): + d, _ = case.direction(est, vec) + if d is None: + nonfinite += 1 + break + for trial in (lam, lam / 3, lam / 9): + cand = vec * np.exp(trial * d) + Lc = case.hard(cand) + if Lc < L: + vec, L, lam = cand, Lc, trial + break + if L < best["loss"]: + best = {"loss": L, "theta": vec.tolist(), "it": it} + rel = (np.abs(np.asarray(best["theta"]) + - np.asarray(tt)) / np.asarray(tt)) + stageB.append({ + "case": ci, "start": sname, "estimator": est, + "eps_dt_halving": eps, + "best_loss": best["loss"], + "within_eps": bool(best["loss"] <= eps), + "theta_best": best["theta"], + "rel_err": rel.tolist(), + "gcal_recovered_20pct": bool(rel[0] <= 0.20), + "w_recovered_20pct": bool(rel[1] <= 0.20), + "heldout_mse": case.heldout_mse(best["theta"]), + "nonfinite": nonfinite, + "hard_evals": case._eval_count, + "wall_s": round(time.time() - t0, 1), + }) + r = stageB[-1] + print(f" B c{ci}/{sname}/{est}: best {r['best_loss']:.3f} " + f"(eps {eps:.3f}, within={r['within_eps']}) " + f"θ=({r['theta_best'][0]:.3f},{r['theta_best'][1]:.5f}) " + f"rel ({rel[0]*100:.0f}%,{rel[1]*100:.0f}%) " + f"held-out {r['heldout_mse']:.3f} ({r['wall_s']}s)", + flush=True) + + def recov_stats(est): + rows = [r for r in stageB if r["estimator"] == est] + return { + "n_runs": len(rows), + "gcal_recovered_20pct": sum(r["gcal_recovered_20pct"] for r in rows), + "w_recovered_20pct": sum(r["w_recovered_20pct"] for r in rows), + "within_eps": sum(r["within_eps"] for r in rows), + "nonfinite_runs": sum(r["nonfinite"] > 0 for r in rows), + "median_heldout": float(np.median([r["heldout_mse"] + for r in rows])), + } + + summaryB = {e: recov_stats(e) for e in estimators} + for e, s in summaryB.items(): + print(f"Stage B — {e:>9}: {s}", flush=True) + + # ------------------------------------------- profile of the EXACT objective + print("profiling exact G4 objective (case 0)...", flush=True) + from deliverable_b import axis_with_truth + + case = case_objs[0] + tt = np.asarray(case_specs[0][2]) + g_axis = axis_with_truth(0.25, 4.0, 11, tt[0]) + w_axis = axis_with_truth(0.0045, 0.018, 9, tt[1]) + L2 = np.array([[case.hard(np.array([g, w])) for w in w_axis] + for g in g_axis]) + prof_g = L2.min(axis=1) + gi = int(np.argmin(np.abs(g_axis - tt[0]))) + assert prof_g[gi] == 0.0 + eps0 = [r["eps_dt_halving"] for r in stageB if r["case"] == 0][0] + within = prof_g <= eps0 + profile = { + "g_axis": g_axis.tolist(), "profile_gcal": prof_g.tolist(), + "eps": eps0, + "gcal_compat_interval": ([float(g_axis[within].min()), + float(g_axis[within].max())] + if within.any() else None), + } + print(f" gCaL compatibility interval (≤ eps {eps0:.3f}): " + f"{profile['gcal_compat_interval']}", flush=True) + + # -------------------------------------------------------------- report + tag = "smoke" if args.smoke else "full" + report = { + "manifest": {**_manifest(), + "g4_sha256_16": hashlib.sha256( + Path(__file__).read_bytes()).hexdigest()[:16], + "args": vars(args), "run_tag": tag}, + "declared": { + "dt_ms": DT, "window_s": WINDOW_S, + "objective": "FULL-TRACE hard binned-rate MSE (50/300 ms) at " + "dt 0.025; windows affect gradient flow only", + "starts": "S1 both-wrong, S2 gCaL-wrong-only, S3 w-wrong-only", + "lambdas": list(LAMBDAS), + }, + "stageA": {"rows": rowsA, "gate": gate, + "cosine_sur_vs_fd": {"mean": float(np.mean(cosines)), + "min": float(np.min(cosines))}}, + "stageB": {"runs": stageB, "summary": summaryB}, + "profile_case0": profile, + "wall_s": round(time.time() - t0_all, 1), + } + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"g4_descent_{tag}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"g4_descent_{tag}_{stamp}.json").write_text(payload) + + fig, axes = plt.subplots(1, 3, figsize=(15, 4.2)) + ax = axes[0] + for k, est in enumerate(estimators): + vals = [gate[est]["joint_accepted_fraction"], + gate[est]["coord_accepted_fraction"]["gCaL"], + gate[est]["coord_accepted_fraction"]["w"]] + ax.bar(np.arange(3) + 0.4 * k - 0.2, vals, width=0.35, + color=OKABE[est], label=est) + ax.set_xticks(range(3), ["joint", "gCaL-only", "w-only"]) + ax.axhline(0.8, color="#111111", ls="--", lw=1) + ax.set_ylim(0, 1.05) + ax.set_title("Stage A accepted fractions", fontsize=10) + ax.legend(frameon=False, fontsize=8) + ax = axes[1] + for est in estimators: + errs = [r["rel_err"][0] for r in stageB if r["estimator"] == est] + ax.plot(np.full(len(errs), estimators.index(est)) + + np.random.default_rng(0).uniform(-0.08, 0.08, len(errs)), + errs, "o", color=OKABE[est], alpha=0.6) + ax.axhline(0.2, color="#111111", ls="--", lw=1, label="±20 %") + ax.set_xticks(range(len(estimators)), estimators) + ax.set_ylabel("|gCaL rel err| at best") + ax.set_title("Stage B gCaL recovery (all runs)", fontsize=10) + ax.legend(frameon=False, fontsize=8) + ax = axes[2] + ax.semilogy(g_axis, np.maximum(prof_g, EPS_FLOOR / 10), "o-", + color="#0072B2", ms=3) + ax.axhline(eps0, color="#111111", ls="--", lw=1, label=f"eps {eps0:.3f}") + ax.axvline(tt[0], color="#000000", lw=1) + ax.set_xscale("log") + ax.set_xlabel("gCaL_scale") + ax.set_title("Exact-G4-objective gCaL profile (case 0)", fontsize=10) + ax.legend(frameon=False, fontsize=8) + for a in axes: + a.grid(True, color="#e3e3e3", lw=0.6) + a.set_axisbelow(True) + fig.tight_layout() + fig.savefig(RESULTS / f"g4_descent_{tag}.png", dpi=150) + print(f"saved → results/g4_descent_{tag}.json (+ evidence, figure)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/g5_basin.py b/experiments/twin/g5_basin.py new file mode 100644 index 00000000..fa6cec47 --- /dev/null +++ b/experiments/twin/g5_basin.py @@ -0,0 +1,298 @@ +"""G5 — basin characterization and hybrid recovery (post-G4 program). + +Answers the three questions G4 left open, per review: + +1. **How wide is the gCaL basin really?** Adaptive profile around the + truth with CONTINUOUS w-optimization at every gCaL (golden-section on + log-w — no discrete w grid), plus bisection refinement of both + compatibility boundaries to ~2 % gCaL resolution. Run on ALL cases. +2. **Does the basin survive realistic variability?** Each case reports a + REPEATABILITY tolerance: the rate-MSE between the truth run and a truth + run with an independent event realization (the irreducible variability + a real experiment faces). Compatibility intervals are reported against + BOTH the numerical EPS and the repeatability tolerance — if the basin + vanishes under the latter, the exact-fit needle is brittle + event-sequence matching, not robust physiological information. +3. **Can a hybrid solver recover gCaL where local descent failed?** + Coarse gCaL scan × continuous w-optimization, retain top-3 candidate + basins, refine each with the qualified local gradient optimizer, return + the best — against an equal-budget Nelder–Mead derivative-free baseline. + +Objective/configuration identical to G4 v2 (full-trace binned-rate MSE, +dt 0.025, state-preserving windows, float64). + +Usage:: + + python g5_basin.py # full (CPU: hours) + python g5_basin.py --smoke +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + +RESULTS = Path(__file__).parent / "results" +EVIDENCE = Path(__file__).parent / "evidence" + +W_LO, W_HI = 0.0045, 0.018 +GOLD = 0.381966 + + +def golden_min_w(case, gcal, n_evals=10): + """Continuous 1-D minimization of the hard loss over log-w at fixed + gCaL (derivative-free, deterministic). Returns (w*, loss*).""" + lo, hi = np.log(W_LO), np.log(W_HI) + x1 = lo + GOLD * (hi - lo) + x2 = hi - GOLD * (hi - lo) + f = {x: case.hard(np.array([gcal, np.exp(x)])) for x in (x1, x2)} + for _ in range(n_evals - 2): + if f[x1] <= f[x2]: + hi, x2 = x2, x1 + x1 = lo + GOLD * (hi - lo) + f[x1] = case.hard(np.array([gcal, np.exp(x1)])) + else: + lo, x1 = x1, x2 + x2 = hi - GOLD * (hi - lo) + f[x2] = case.hard(np.array([gcal, np.exp(x2)])) + xb = min((x1, x2), key=lambda x: f[x]) + return float(np.exp(xb)), float(f[xb]) + + +def profile_with_boundaries(case, g_truth, tol, span=(0.75, 1.25), + n_coarse=7, resolution_rel=0.02): + """Adaptive gCaL profile: coarse pass around truth, then bisection of + both compatibility boundaries (profile ≤ tol) to resolution_rel.""" + g_axis = list(np.geomspace(g_truth * span[0], g_truth * span[1], + n_coarse)) + if g_truth not in g_axis: + g_axis = sorted(g_axis + [g_truth]) + prof = {g: golden_min_w(case, g)[1] for g in g_axis} + + def inside(g): + if g not in prof: + prof[g] = golden_min_w(case, g)[1] + return prof[g] <= tol + + gs = sorted(prof) + ins = [g for g in gs if prof[g] <= tol] + if not ins: + return prof, None + lo_in, hi_in = min(ins), max(ins) + # bisection on each boundary + lo_out = max([g for g in gs if g < lo_in], default=lo_in * span[0]) + hi_out = min([g for g in gs if g > hi_in], default=hi_in * span[1]) + for _ in range(12): + if lo_out / lo_in < 1 + resolution_rel: + break + mid = float(np.sqrt(lo_out * lo_in)) + if inside(mid): + lo_in = mid + else: + lo_out = mid + for _ in range(12): + if hi_out / hi_in < 1 + resolution_rel: + break + mid = float(np.sqrt(hi_out * hi_in)) + if inside(mid): + hi_in = mid + else: + hi_out = mid + return prof, [lo_in, hi_in] + + +def hybrid_solver(case, budget_evals=220, n_keep=3, refine_iters=6): + """Coarse gCaL scan × continuous w-opt → keep top basins → local + gradient refinement of each → best candidate.""" + case._eval_count = 0 + scan = np.geomspace(0.25, 4.0, 9) + cands = [] + for g in scan: + w, L = golden_min_w(case, g, n_evals=8) + cands.append((L, g, w)) + cands.sort() + best = None + for L, g, w in cands[:n_keep]: + vec = np.array([g, w]) + Lc = L + lam = 0.1 + for _ in range(refine_iters): + if case._eval_count > budget_evals: + break + d, _ = case.direction("surrogate", vec) + if d is None: + break + for trial in (lam, lam / 3, lam / 9): + cand = vec * np.exp(trial * d) + Lt = case.hard(cand) + if Lt < Lc: + vec, Lc, lam = cand, Lt, trial + break + if best is None or Lc < best[0]: + best = (Lc, vec.copy()) + return {"loss": best[0], "theta": best[1].tolist(), + "evals": case._eval_count} + + +def nelder_mead_baseline(case, budget_evals=220): + from scipy.optimize import minimize + + case._eval_count = 0 + + def f(z): + if case._eval_count > budget_evals: + return 1e9 + return case.hard(np.exp(z)) + + res = minimize(f, np.log([1.0, 0.009]), method="Nelder-Mead", + options={"maxfev": budget_evals, "xatol": 1e-3, + "fatol": 1e-6}) + return {"loss": float(res.fun), "theta": np.exp(res.x).tolist(), + "evals": case._eval_count} + + +def main(): + import jax + + jax.config.update("jax_enable_x64", True) + + ap = argparse.ArgumentParser() + ap.add_argument("--smoke", action="store_true") + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + from sensitivity_checks import _manifest + from g4_descent import Case, build_case, binned, chained_spikes + + t0_all = time.time() + dur_scale = 0.5 if args.smoke else 1.0 + specs = ([(0, 0, (1.5, 0.011))] if args.smoke else + [(c, e, t) for c in (0, 1) for e in (0, 10) + for t in ((1.5, 0.011), (0.7, 0.008))]) + + rows = [] + for ci, (cs, eo, tt) in enumerate(specs): + t0 = time.time() + case = Case(cs, eo, tt, dur_scale) + eps = case.eps_dt_halving() + + # repeatability tolerance: truth with an INDEPENDENT event realization + import jax.numpy as jnp + + twin_alt, ev_alt = build_case(cs, eo + 1000, dur_scale) + th_t = case.theta_of(jnp.asarray(np.asarray(tt, dtype=np.float64))) + out_alt, _ = case.twin._run_events_carry( + th_t, ev_alt["fit_ramp"], None, "hard") + t_rep = float(jnp.mean((binned(out_alt["spikes"]) + - case.obs_rates) ** 2)) + + prof_eps, iv_eps = profile_with_boundaries(case, tt[0], eps) + tol_rep = max(eps, t_rep) + iv_rep_candidates = {g: L for g, L in prof_eps.items() if L <= tol_rep} + iv_rep = ([min(iv_rep_candidates), max(iv_rep_candidates)] + if iv_rep_candidates else None) + + hyb = hybrid_solver(case) + nm = nelder_mead_baseline(case) + rel_h = abs(hyb["theta"][0] - tt[0]) / tt[0] + rel_n = abs(nm["theta"][0] - tt[0]) / tt[0] + rows.append({ + "case": case.meta, "eps": eps, "t_repeat": t_rep, + "profile": {f"{g:.4f}": L for g, L in sorted(prof_eps.items())}, + "gcal_interval_vs_eps": iv_eps, + "gcal_interval_vs_repeatability": iv_rep, + "hybrid": {**hyb, "gcal_rel_err": rel_h, + "gcal_recovered_20pct": bool(rel_h <= 0.2)}, + "nelder_mead": {**nm, "gcal_rel_err": rel_n, + "gcal_recovered_20pct": bool(rel_n <= 0.2)}, + }) + r = rows[-1] + print(f"[case {ci}] conn{cs}/ev{eo}/true{tt}: eps {eps:.3f} " + f"T_rep {t_rep:.3f} | interval(eps) {iv_eps} | " + f"interval(rep) {iv_rep} | hybrid gCaL " + f"{hyb['theta'][0]:.3f} ({rel_h*100:.0f}%, {hyb['evals']} ev) " + f"| NM gCaL {nm['theta'][0]:.3f} ({rel_n*100:.0f}%, " + f"{nm['evals']} ev) ({time.time()-t0:.0f}s)", flush=True) + + summary = { + "hybrid_gcal_recovered_20pct": sum( + r["hybrid"]["gcal_recovered_20pct"] for r in rows), + "nm_gcal_recovered_20pct": sum( + r["nelder_mead"]["gcal_recovered_20pct"] for r in rows), + "n_cases": len(rows), + "basin_survives_repeatability": sum( + r["gcal_interval_vs_repeatability"] is not None + and r["gcal_interval_vs_repeatability"][0] > 0 for r in rows), + } + print(f"SUMMARY: {summary}", flush=True) + + tag = "smoke" if args.smoke else "full" + report = { + "manifest": {**_manifest(), + "g5_sha256_16": hashlib.sha256( + Path(__file__).read_bytes()).hexdigest()[:16], + "args": vars(args), "run_tag": tag}, + "declared": { + "objective": "G4 v2 full-trace objective (dt 0.025)", + "profile": "continuous w via golden-section (10 evals), " + "boundary bisection to 2 % gCaL resolution", + "tolerances": "eps = dt-halving at truth; t_repeat = rate-MSE " + "of truth vs independent event realization", + "hybrid": "9-pt gCaL scan × golden w, top-3 basins, surrogate " + "refinement, ~220-eval budget", + "baseline": "Nelder-Mead log-space, same budget", + }, + "cases": rows, "summary": summary, + "wall_s": round(time.time() - t0_all, 1), + } + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"g5_basin_{tag}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"g5_basin_{tag}_{stamp}.json").write_text(payload) + + fig, axes = plt.subplots(1, 2, figsize=(12, 4.5)) + ax = axes[0] + for r in rows: + gs = sorted(float(g) for g in r["profile"]) + Ls = [r["profile"][f"{g:.4f}"] for g in gs] + ax.semilogy(np.asarray(gs) / r["case"]["theta_true"][0], + np.maximum(Ls, 1e-4), "o-", ms=3, alpha=0.6) + ax.axhline(r["t_repeat"], color="#E69F00", lw=0.6, alpha=0.5) + ax.axhline(r["eps"], color="#111111", lw=0.6, ls="--", alpha=0.5) + ax.axvline(1.0, color="#000000", lw=1) + ax.set_xlabel("gCaL / truth") + ax.set_ylabel("profile loss (continuous w-opt)") + ax.set_title("Adaptive profiles, all cases (black=eps, orange=T_repeat)", + fontsize=9) + ax = axes[1] + for k, (meth, c) in enumerate((("hybrid", "#0072B2"), + ("nelder_mead", "#CC79A7"))): + errs = [r[meth]["gcal_rel_err"] for r in rows] + ax.plot(np.full(len(errs), k) + + np.random.default_rng(0).uniform(-0.06, 0.06, len(errs)), + errs, "o", color=c, alpha=0.7) + ax.axhline(0.2, color="#111111", ls="--", lw=1, label="±20 %") + ax.set_xticks([0, 1], ["hybrid", "Nelder-Mead"]) + ax.set_ylabel("|gCaL rel err|") + ax.set_title("Recovery, equal budget", fontsize=10) + ax.legend(frameon=False, fontsize=8) + for a in axes: + a.grid(True, color="#e3e3e3", lw=0.6) + a.set_axisbelow(True) + fig.tight_layout() + fig.savefig(RESULTS / f"g5_basin_{tag}.png", dpi=150) + print(f"saved → results/g5_basin_{tag}.json (+ evidence, figure)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/gradient_repair.py b/experiments/twin/gradient_repair.py new file mode 100644 index 00000000..133677c5 --- /dev/null +++ b/experiments/twin/gradient_repair.py @@ -0,0 +1,510 @@ +"""Gradient-repair diagnostics G1–G3 (v2 — corrected gates). + +v2 fixes from review: G1's gate now requires agreement with a LOCAL +FINITE-DIFFERENCE CONVERGENCE REGION (reverse ≡ forward alone is not a +correctness test — both differentiate the same discretization); adds +dt 0.0125 for timestep convergence, directional cosine, and a synaptic- +regime variant. G2 classifies a zero reference slope as "unresolved", not +"wrong". G3 runs at the proposed operating timestep (0.025 ms) across +multiple currents, and its scope is explicit: a frozen-index interpolated +crossing-time derivative — NOT yet a complete event-time estimator (no +spike creation/deletion, matching, or full objective). + +Scoped claims this file can support: +- G1: "AD agrees with the local FD slope of the discrete simulation at the + tested operating points / regimes / timesteps" — nothing broader. +- G2: whether the straight-through direction is wrong / unresolved per + coordinate under irrelevant perturbations. +- G3: crossing-time AD≡FD on the first K spikes at the tested dt. + +Usage:: + + python gradient_repair.py --x64 # all, float64 + python gradient_repair.py --x64 --tests g1 --regime synaptic +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np + +RESULTS = Path(__file__).parent / "results" +EVIDENCE = Path(__file__).parent / "evidence" + +FD_HS = tuple(10.0 ** -k for k in range(1, 9)) # 1e-1 … 1e-8 +FD_PLATEAU_RTOL = 0.05 + + +# --------------------------------------------------------------------------- # +# Testable helpers +# --------------------------------------------------------------------------- # + +def detect_plateau(rows, rtol=FD_PLATEAU_RTOL): + """Convergence plateau from an (h, slope) sweep — INDEPENDENT of any AD + value: the slope at the smaller h of the last consecutive pair agreeing + within ``rtol``, or None. Shared by G1 and G3 so neither can select its + reference by closeness to autodiff.""" + conv = None + for i in range(len(rows) - 1): + a, b = rows[i][1], rows[i + 1][1] + if abs(a - b) <= rtol * max(abs(a), abs(b), 1e-300): + conv = b + return conv + + +def fd_sweep_converged(loss, s0, j, hs=FD_HS, rtol=FD_PLATEAU_RTOL): + """Centered-FD h-sweep for coordinate j → (rows, plateau_or_None).""" + rows = [] + for h in hs: + e = np.zeros(len(s0)) + e[j] = h + rows.append((float(h), + (float(loss(s0 + e)) - float(loss(s0 - e))) / (2 * h))) + return rows, detect_plateau(rows, rtol) + + +def zero_insert_events(ev_coarse: np.ndarray, factor: int) -> np.ndarray: + """One continuous event realization discretized at a finer step: events + at coarse step k move to fine step k·factor (same physical times).""" + out = np.zeros((ev_coarse.shape[0] * factor, ev_coarse.shape[1]), + dtype=ev_coarse.dtype) + out[::factor] = ev_coarse + return out + + +def cosine(a, b) -> float: + a, b = np.asarray(a, float), np.asarray(b, float) + na, nb = np.linalg.norm(a), np.linalg.norm(b) + if na == 0 or nb == 0: + return float("nan") + return float(a @ b / (na * nb)) + + +def classify_direction(g, ref, zero_tol=1e-12) -> list[str]: + """Per-coordinate: 'agree' / 'wrong' / 'unresolved' (|ref| ≤ zero_tol).""" + out = [] + for gj, rj in zip(np.ravel(g), np.ravel(ref)): + if abs(rj) <= zero_tol: + out.append("unresolved") + elif np.sign(gj) == np.sign(rj): + out.append("agree") + else: + out.append("wrong") + return out + + +# --------------------------------------------------------------------------- # + +_BASE_DT = 0.1 # all event realizations are generated ONCE at this step and +# zero-inserted to finer steps so timestep comparisons share +# one continuous input realization + + +def _make_loss(regime, dt, horizon_s, current_nA=10.0): + """Continuous mean-voltage objective in a given regime; returns + (loss(scales)->scalar, n_spikes_at_base). + + regimes: + current — 1 cell, deterministic current injection. + synaptic — 1 cell, SYNAPTICALLY DRIVEN FIRING REGIME (w 0.15 µS, + frozen renewal events). NOT the Deliverable-B fit setup. + b_config — the EXACT Deliverable-B configuration: 8 MN, B's fit_ramp + command/protocol, B's connectivity + event seed, operating + point at B's truth (gCaL 1.6, w_exc 0.0105); mean pool + voltage as the continuous objective. + """ + import jax.numpy as jnp + from poolbuild import PoolTwin, TwinSpec + from validation_protocol import _renewal_events + + n = int(horizon_s * 1000 / dt) + factor = int(round(_BASE_DT / dt)) + if regime == "current": + twin = PoolTwin(TwinSpec(n_mn=1, n_dd=1, n_inh=1, dt_ms=dt, seed=0)) + i_ext = jnp.full((n, 1), current_nA) + + def loss(scales): + th = twin.default_theta() + th["gCaL_scale"], th["gNaP_scale"] = scales[0], scales[1] + return jnp.mean(twin._run_current(th, i_ext, + spike_mode="hard")["v"]) + + sp = twin.simulate_current(twin.default_theta(), np.asarray(i_ext)) + return loss, int(np.asarray(sp["spikes"]).sum()) + + if regime == "synaptic": + twin = PoolTwin(TwinSpec(n_mn=1, n_dd=100, n_inh=1, dt_ms=dt, seed=0, + w_exc_uS=0.15)) + n0 = int(horizon_s * 1000 / _BASE_DT) + t0_ = np.arange(n0) * _BASE_DT * 1e-3 + cmd = np.interp(t0_, [0, 0.1 * horizon_s, horizon_s], + [0, 60, 60]).astype(np.float32) + ev0 = _renewal_events(np.random.default_rng(0), + np.tile(cmd[:, None], (1, 100)), _BASE_DT, 5) + ev = jnp.asarray(zero_insert_events(ev0, factor)) + + def loss(scales): + th = twin.default_theta() + th["gCaL_scale"], th["gNaP_scale"] = scales[0], scales[1] + return jnp.mean(twin.simulate_events(th, ev, + spike_mode="hard")["v"]) + + sp = twin.simulate_events(twin.default_theta(), ev, spike_mode="hard") + return loss, int(np.asarray(sp["spikes"]).sum()) + + # b_config — exact Deliverable-B pool/protocol/operating point + from deliverable_b import make_commands, build_twin_and_events + + commands = make_commands(_BASE_DT, 1.0) + twin0, events0 = build_twin_and_events(8, 0, 0, commands, _BASE_DT) + twin = (twin0 if factor == 1 else + PoolTwin(TwinSpec(n_mn=8, n_dd=100, n_inh=5, dt_ms=dt, seed=0))) + ev0 = np.asarray(events0["fit_ramp"])[: int(horizon_s * 1000 / _BASE_DT)] + ev = jnp.asarray(zero_insert_events(ev0, factor)) + + def loss(scales): + # s0 = (1, 1) is B's TRUTH operating point (gCaL 1.6, w_exc 0.0105) + th = twin.default_theta() + th["gCaL_scale"] = scales[0] * 1.6 + th["w_exc"] = scales[1] * 0.0105 + return jnp.mean(twin.simulate_events(th, ev, spike_mode="hard")["v"]) + + th_b = twin.default_theta() + th_b["gCaL_scale"] = jnp.asarray(1.6) + th_b["w_exc"] = jnp.asarray(0.0105) + sp = twin.simulate_events(th_b, ev, spike_mode="hard") + return loss, int(np.asarray(sp["spikes"]).sum()) + + +def g1_continuous_sensitivity(regime="current", + dts=(0.1, 0.05, 0.025, 0.0125), + horizons=(0.1, 0.5, 1.0)): + import jax + import jax.numpy as jnp + + rows = [] + for dt in dts: + for horizon in horizons: + t0 = time.time() + loss, n_spk = _make_loss(regime, dt, horizon) + s0 = jnp.array([1.0, 1.0]) + g_rev = np.asarray(jax.grad(loss)(s0), dtype=np.float64) + g_fwd = np.array([jax.jvp(loss, (s0,), (jnp.array([1., 0.]),))[1], + jax.jvp(loss, (s0,), (jnp.array([0., 1.]),))[1]], + dtype=np.float64) + fd_rows, fd_conv = {}, [] + for j in range(2): + r, c = fd_sweep_converged(loss, np.array([1.0, 1.0]), j) + fd_rows[j] = r + fd_conv.append(c) + converged = all(c is not None for c in fd_conv) + if converged: + conv = np.array(fd_conv) + rel = float(np.max(np.abs(g_rev - conv) + / np.maximum(np.abs(conv), 1e-300))) + cos = cosine(g_rev, conv) + verdict = "PASS" if (rel < FD_PLATEAU_RTOL and cos > 0.95) \ + else "FAIL" + else: + rel, cos, verdict = None, None, "FD_UNCONVERGED" + rows.append({ + "regime": regime, "dt_ms": dt, "horizon_s": horizon, + "n_spikes": n_spk, + "grad_rev": g_rev.tolist(), "grad_fwd": g_fwd.tolist(), + "fd_converged_value": [c for c in fd_conv], + "fd_sweeps": {str(j): fd_rows[j] for j in range(2)}, + "rel_rev_vs_fdconv": rel, "cosine": cos, + "rev_vs_fwd_rel": float(np.max( + np.abs(g_rev - g_fwd) + / np.maximum(np.abs(g_fwd), 1e-300))), + "verdict": verdict, + }) + r = rows[-1] + print(f" [{regime}] dt {dt:6.4f} T {horizon:4.2f}s spk " + f"{n_spk:3d} | |g| {np.max(np.abs(g_rev)):.2e} | " + f"{verdict}" + + (f" rel {rel:.1e} cos {cos:.3f}" if converged else "") + + f" ({time.time()-t0:.0f}s)", flush=True) + passed = [(r["dt_ms"], r["horizon_s"]) for r in rows + if r["verdict"] == "PASS"] + failed = [(r["dt_ms"], r["horizon_s"]) for r in rows + if r["verdict"] == "FAIL"] + unresolved = [(r["dt_ms"], r["horizon_s"]) for r in rows + if r["verdict"] == "FD_UNCONVERGED"] + # prespecified required cases: every fine-dt (≤0.025) combo must PASS — + # unresolved required cases block ok (review: the aggregate must not be + # permissive about requested-but-unresolved configurations) + required = [(r["dt_ms"], r["horizon_s"]) for r in rows + if r["dt_ms"] <= 0.025] + ok = (not failed) and all(c in passed for c in required) + return {"ok": ok, "rows": rows, + "passed_cases": passed, "failed_cases": failed, + "unresolved_cases": unresolved, + "required_cases_dt_le_0p025": required, + "qualified_envelope": sorted(set(passed)), + "scope": f"regime={regime}; claims limited to these operating " + "points/timesteps; FD_UNCONVERGED combos are " + "unresolvable by local FD, not passes"} + + +# --------------------------------------------------------------------------- # + +def _g2_setup(dt, event_shift=0): + import jax.numpy as jnp + from poolbuild import PoolTwin, TwinSpec + from validation_protocol import _renewal_events + from deliverable_b import binned_rates_jnp + + twin = PoolTwin(TwinSpec(n_mn=3, n_dd=100, n_inh=5, dt_ms=dt, seed=0)) + dur = 2.0 + n = int(dur * 1000 / dt) + t = np.arange(n) * dt * 1e-3 + cmd = np.interp(t, [0, 0.3, 1.0, 1.7, 2.0], [0, 0, 60, 0, 0] + ).astype(np.float32) + ev = _renewal_events(np.random.default_rng(0), + np.tile(cmd[:, None], (1, 100)), dt, 5) + if event_shift: + ev = np.roll(ev, event_shift, axis=0) + ev = jnp.asarray(ev) + + def rates(theta_vals, spike_mode): + th = twin.default_theta() + th["gCaL_scale"] = jnp.asarray(theta_vals[0]) + th["w_exc"] = jnp.asarray(theta_vals[1]) + out = twin.simulate_events(th, ev, spike_mode=spike_mode) + return binned_rates_jnp(out["spikes"], dt) + + target = rates((1.5, 0.011), "hard") + return twin, rates, target + + +def g2_surrogate_dependence(): + import jax + import jax.numpy as jnp + import poolbuild + from myogen.simulator.jaxley import jax_models + + theta0 = (1.0, 0.009) + variants = ([("beta", b) for b in (5.0, 10.0, 20.0)] + + [("vth", v) for v in (49.0, 50.0, 51.0)] + + [("dt", d) for d in (0.1, 0.05)] + + [("evshift", s) for s in (0, 1)]) + base_beta = jax_models.SURROGATE_BETA + base_vth = poolbuild.SPIKE_THRESHOLD_MV + rows = [] + try: + for kind, val in variants: + jax_models.SURROGATE_BETA = val if kind == "beta" else base_beta + poolbuild.SPIKE_THRESHOLD_MV = val if kind == "vth" else base_vth + dt = val if kind == "dt" else 0.1 + shift = val if kind == "evshift" else 0 + twin, rates, target = _g2_setup(dt, event_shift=int(shift)) + tgt = jax.lax.stop_gradient(target) + + def sur_loss(z): + r = rates((z[0], z[1]), "surrogate") + return jnp.mean((r - tgt) ** 2) + + leaf_dtype = (jnp.float64 if jax.config.jax_enable_x64 + else jnp.float32) + g = np.asarray(jax.grad(sur_loss)( + jnp.asarray(theta0, dtype=leaf_dtype)), dtype=np.float64) + + def hard_loss(tv): + return float(jnp.mean((rates(tv, "hard") - tgt) ** 2)) + + ref = np.zeros(2) + for j, h in enumerate((0.05 * theta0[0], 0.05 * theta0[1])): + up, dn = list(theta0), list(theta0) + up[j] += h + dn[j] -= h + ref[j] = (hard_loss(tuple(up)) - hard_loss(tuple(dn))) / (2 * h) + cls = classify_direction(g, ref) + rows.append({"variant": f"{kind}={val}", + "surrogate_grad": g.tolist(), + "crn_fd_ref": ref.tolist(), + "classification": cls, + "leaf_dtype": str(leaf_dtype.__name__), + "grad_finite": bool(np.all(np.isfinite(g)))}) + print(f" {rows[-1]['variant']:>12}: {cls}", flush=True) + finally: + jax_models.SURROGATE_BETA = base_beta + poolbuild.SPIKE_THRESHOLD_MV = base_vth + counts = {"gCaL": {"agree": 0, "wrong": 0, "unresolved": 0}, + "w_exc": {"agree": 0, "wrong": 0, "unresolved": 0}} + for r in rows: + counts["gCaL"][r["classification"][0]] += 1 + counts["w_exc"][r["classification"][1]] += 1 + ok = all(c["wrong"] == 0 for c in counts.values()) + mags = [float(np.max(np.abs(r["surrogate_grad"]))) for r in rows] + return {"ok": ok, "n_variants": len(rows), "per_coordinate": counts, + "surrogate_grad_magnitude_range": [min(mags), max(mags)], + "magnitude_note": "sign agreement does NOT establish usefulness —" + " magnitudes remain enormous and must be handled" + " by the optimizer (G4 question)", + "rows": rows} + + +# --------------------------------------------------------------------------- # + +def hard_crossings(v: np.ndarray, v_th: float) -> np.ndarray: + return np.flatnonzero((v[1:] > v_th) & (v[:-1] <= v_th)) + + +def interp_time(v, k, v_th, dt_ms): + """Interpolated crossing time (s) at frozen step index k — differentiable + in the voltages (discrete implicit-function form).""" + vk, vk1 = v[k], v[k + 1] + return (k + (v_th - vk) / (vk1 - vk)) * (dt_ms * 1e-3) + + +def g3_event_time_gradient(dts=(0.025, 0.1), currents=(12.0, 15.0, 30.0), + horizon_s=0.3, + fd_hs=(3e-3, 1e-3, 3e-4, 1e-4, 3e-5), + max_spikes=8): + import jax + import jax.numpy as jnp + from poolbuild import PoolTwin, TwinSpec, SPIKE_THRESHOLD_MV + + cases = [] + for dt in dts: + twin = PoolTwin(TwinSpec(n_mn=1, n_dd=1, n_inh=1, dt_ms=dt, seed=0)) + n = int(horizon_s * 1000 / dt) + for amp in currents: + i_ext = jnp.full((n, 1), amp) + + def v_trace(scales): + th = twin.default_theta() + th["gCaL_scale"], th["gNaP_scale"] = scales[0], scales[1] + return twin._run_current(th, i_ext, + spike_mode="hard")["v"][:, 0] + + def times_concrete(scales): + v = np.asarray(v_trace(jnp.asarray(scales))) + ks = hard_crossings(v, SPIKE_THRESHOLD_MV) + return ([float(interp_time(v, int(k), SPIKE_THRESHOLD_MV, dt)) + for k in ks], ks) + + s0 = np.array([1.0, 1.0]) + ts0, ks0 = times_concrete(s0) + n_spk = len(ks0) + per_spike = [] + for j_spike, k in enumerate(ks0[:max_spikes]): + k = int(k) + + def tj_ad(scales): + return interp_time(v_trace(scales), k, + SPIKE_THRESHOLD_MV, dt) + + g_ad = np.asarray(jax.grad(tj_ad)(jnp.asarray(s0)), + dtype=np.float64) + # AD-INDEPENDENT reference: per-coordinate FD sweep over h, + # plateau detected from consecutive-h agreement alone + # (count-inconsistent h's excluded); AD compared only after. + plateau = [] + sweeps = {} + for p in range(2): + rows_fd = [] + for h in fd_hs: + e = np.zeros(2) + e[p] = h + tsp, ksp = times_concrete(s0 + e) + tsm, ksm = times_concrete(s0 - e) + if len(ksp) != n_spk or len(ksm) != n_spk: + continue # spike creation/deletion — invalid h + rows_fd.append( + (h, (tsp[j_spike] - tsm[j_spike]) / (2 * h))) + sweeps[p] = rows_fd + plateau.append(detect_plateau(rows_fd)) + if all(c is not None for c in plateau): + conv = np.array(plateau) + rel = float(np.max(np.abs(conv - g_ad) + / np.maximum(np.abs(conv), 1e-300))) + fd_res = {"plateau": conv.tolist(), "rel": rel, + "sweeps": {str(p): sweeps[p] for p in sweeps}} + else: + fd_res = None # no AD-independent plateau → unresolved + per_spike.append({"spike": j_spike, "t_s": ts0[j_spike], + "grad_ad": g_ad.tolist(), "fd": fd_res}) + n_pass = 0 + for r in per_spike: + if r["fd"] is not None and r["fd"]["rel"] < 0.05: + n_pass += 1 + else: + break # contiguous validity horizon + cases.append({"dt_ms": dt, "current_nA": amp, "n_spikes": n_spk, + "n_valid_leading_spikes": n_pass, + "per_spike": per_spike}) + print(f" dt {dt:5.3f} I {amp:4.1f} nA: {n_spk} spikes, " + f"leading AD≡FD (<5%): {n_pass}", flush=True) + fine = [c for c in cases if c["dt_ms"] == min(dts) and c["n_spikes"] > 0] + ok = bool(fine) and all(c["n_valid_leading_spikes"] >= 3 for c in fine) + return {"ok": ok, "cases": cases, + "scope": "frozen-index interpolated crossing times only — NOT a " + "complete event-time estimator (no creation/deletion, " + "matching, or full objective)"} + + +# --------------------------------------------------------------------------- # + +TESTS = {"g1": g1_continuous_sensitivity, + "g2": g2_surrogate_dependence, + "g3": g3_event_time_gradient} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--tests", nargs="+", choices=sorted(TESTS), + default=sorted(TESTS)) + ap.add_argument("--x64", action="store_true") + ap.add_argument("--regime", + choices=["current", "synaptic", "b_config"], + default="current") + args = ap.parse_args() + + import jax + if args.x64: + jax.config.update("jax_enable_x64", True) + from sensitivity_checks import _manifest + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + report = {"manifest": _manifest(), "x64": args.x64, + "regime": args.regime} + for tid in args.tests: + print(f"[{tid}] ...", flush=True) + t0 = time.time() + try: + if tid == "g1" and args.regime == "b_config": + # horizons must cover the FIRING portion of the + # ramp (recruitment starts ~1.5 s) — a short + # window only certifies the subthreshold regime + res = TESTS[tid](regime=args.regime, + dts=(0.1, 0.025), + horizons=(2.0, 4.0)) + elif tid == "g1": + res = TESTS[tid](regime=args.regime) + else: + res = TESTS[tid]() + except Exception as e: + res = {"ok": False, "error": f"{type(e).__name__}: {e}"} + res["runtime_s"] = round(time.time() - t0, 1) + report[tid] = res + print(f"[{tid}] {'OK' if res.get('ok') else 'ATTENTION'} " + f"({res['runtime_s']}s)", flush=True) + + suffix = ("x64" if args.x64 else "f32") + ( + "_syn" if args.regime == "synaptic" else "") + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"gradient_repair_{suffix}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"gradient_repair_{suffix}_{stamp}.json").write_text(payload) + print(f"saved → results/gradient_repair_{suffix}.json (+ evidence)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/observables.py b/experiments/twin/observables.py new file mode 100644 index 00000000..b6794053 --- /dev/null +++ b/experiments/twin/observables.py @@ -0,0 +1,292 @@ +"""Observables for the motor-neuron-pool digital twin. + +Two tiers: + +* **Differentiable** (JAX, operate on the surrogate/rate spike matrices coming out + of ``run_jax``): smoothed firing rates, population rate, force features. These + are what inverse-problem losses are built from (Phase 3+). +* **Analysis** (NumPy, hard spikes): recruitment/derecruitment thresholds and + order, firing-rate hysteresis, ΔF via the paired-motor-unit (Gorassini) method. + These are for validation (Phase 1) and for scoring fits — not for gradients. + +Conventions +----------- +``spikes``: array ``(n_steps, n_mn)``, 1.0 at spike samples (hard mode) or a +surrogate/rate value in optimization modes. ``dt_ms``: timestep in milliseconds. +All times returned in seconds. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax.numpy as jnp +import numpy as np + +__all__ = [ + "smoothed_rates", + "population_rate", + "spike_times", + "recruitment_times", + "derecruitment_times", + "recruitment_order", + "delta_f_pairs", + "DeltaFResult", + "hysteresis_index", +] + + +# --------------------------------------------------------------------------- # +# Differentiable tier (JAX) +# --------------------------------------------------------------------------- # + +def smoothed_rates(spikes, dt_ms: float, tau_ms: float = 250.0): + """Per-MN instantaneous firing rate (Hz) via causal exponential smoothing. + + Differentiable in ``spikes`` (use surrogate/rate spike mode upstream). + Implemented as a first-order IIR scanned over time: + ``r[t] = a * r[t-1] + (1 - a) * s[t] / dt_s`` with ``a = exp(-dt/tau)``. + + Parameters + ---------- + spikes : (n_steps, n_mn) + tau_ms : smoothing time constant. 250 ms ≈ the smoothing used for ΔF-style + rate profiles; use smaller values for transient-sensitive losses. + """ + dt_s = dt_ms * 1e-3 + a = jnp.exp(-dt_ms / tau_ms) + inst = spikes / dt_s # spike-count density → Hz + + import jax + + def step(r, x): + r = a * r + (1.0 - a) * x + return r, r + + r0 = jnp.zeros(spikes.shape[1], dtype=spikes.dtype) + _, rates = jax.lax.scan(step, r0, inst) + return rates + + +def population_rate(spikes, dt_ms: float, tau_ms: float = 100.0): + """Mean smoothed firing rate across the pool, shape ``(n_steps,)``.""" + return smoothed_rates(spikes, dt_ms, tau_ms).mean(axis=1) + + +# --------------------------------------------------------------------------- # +# Analysis tier (NumPy, hard spikes) +# --------------------------------------------------------------------------- # + +def spike_times(spikes, dt_ms: float) -> list[np.ndarray]: + """Per-MN spike times in seconds from a hard 0/1 spike matrix.""" + sp = np.asarray(spikes) > 0.5 + dt_s = dt_ms * 1e-3 + return [np.flatnonzero(sp[:, i]) * dt_s for i in range(sp.shape[1])] + + +def _first_sustained(train_s: np.ndarray, min_spikes: int, max_isi_s: float): + """First spike time that begins a run of ``min_spikes`` spikes with all + ISIs below ``max_isi_s`` — filters out isolated spikes near threshold.""" + if len(train_s) < min_spikes: + return np.nan + for k in range(len(train_s) - min_spikes + 1): + isis = np.diff(train_s[k : k + min_spikes]) + if np.all(isis <= max_isi_s): + return train_s[k] + return np.nan + + +def recruitment_times( + spikes, dt_ms: float, min_spikes: int = 3, max_isi_s: float = 0.5 +) -> np.ndarray: + """Per-MN recruitment time (s): onset of first sustained firing; NaN if silent.""" + return np.array( + [_first_sustained(t, min_spikes, max_isi_s) for t in spike_times(spikes, dt_ms)] + ) + + +def derecruitment_times( + spikes, dt_ms: float, min_spikes: int = 3, max_isi_s: float = 0.5 +) -> np.ndarray: + """Per-MN derecruitment time (s): end of the last sustained firing run.""" + out = [] + for t in spike_times(spikes, dt_ms): + rev = _first_sustained(t[::-1] * -1.0, min_spikes, max_isi_s) + out.append(np.nan if np.isnan(rev) else -rev) + return np.array(out) + + +def recruitment_order(spikes, dt_ms: float) -> np.ndarray: + """Indices of recruited MNs sorted by recruitment time (silent MNs excluded).""" + rt = recruitment_times(spikes, dt_ms) + active = np.flatnonzero(~np.isnan(rt)) + return active[np.argsort(rt[active])] + + +def _rate_profile(train_s: np.ndarray, t_grid: np.ndarray) -> np.ndarray: + """Smoothed instantaneous rate (Hz) on ``t_grid`` from spike times: + 1/ISI held between spikes, then 500 ms moving average (standard for ΔF).""" + r = np.zeros_like(t_grid) + if len(train_s) >= 2: + inst = 1.0 / np.diff(train_s) + idx = np.searchsorted(train_s, t_grid, side="right") - 1 + valid = (idx >= 0) & (idx < len(inst)) + r[valid] = inst[np.clip(idx, 0, len(inst) - 1)][valid] + # zero outside firing span + r[(t_grid < train_s[0]) | (t_grid > train_s[-1])] = 0.0 + if len(t_grid) > 1: + dt = t_grid[1] - t_grid[0] + w = max(1, int(round(0.5 / dt))) + kernel = np.ones(w) / w + r = np.convolve(r, kernel, mode="same") + return r + + +@dataclass +class DeltaFResult: + """ΔF estimates for control/test MN pairs on a ramp protocol. + + ``delta_f`` holds every accepted pair (NOT statistically independent — + pairs share control and test units). ``delta_f_per_test`` averages the + controls within each test unit first; ``mean`` is the mean of those + per-test values, so every test unit carries equal weight. + """ + + pairs: np.ndarray # (n_pairs, 2) [control_idx, test_idx] + delta_f: np.ndarray # (n_pairs,) Hz + t_recruit_test: np.ndarray + t_derecruit_test: np.ndarray + test_units: np.ndarray # unique test MN indices, aligned with delta_f_per_test + delta_f_per_test: np.ndarray # (n_test_units,) controls averaged within test + + @property + def mean(self) -> float: + return (float(np.nanmean(self.delta_f_per_test)) + if len(self.delta_f_per_test) else np.nan) + + +def delta_f_pairs( + spikes, + dt_ms: float, + min_lead_s: float = 1.0, + min_ctrl_rate_hz: float = 0.5, + max_pairs_per_test: int = 5, + min_control_outlast_s: float = 1.5, + min_rate_correlation: float = 0.7, + min_ctrl_modulation_hz: float = 0.5, +) -> DeltaFResult: + """ΔF via the paired-motor-unit method (Gorassini et al. 2002). + + For each accepted (control, test) pair: + + ΔF = f_control(t_recruit_test) − f_control(t_derecruit_test) + + i.e. the drop in synaptic drive (proxied by the control MN's rate) between + when the test MN turned on and when it turned off. A positive ΔF on a + symmetric ramp indicates sustained depolarization from PICs. Control rate + profiles use the 1/ISI + 500 ms smoothing convention. + + Pair-acceptance safeguards (cf. Gorassini 2002; Afsharipour et al. 2020): + + * control recruited ≥ ``min_lead_s`` before the test MN; + * control derecruited ≥ ``min_control_outlast_s`` AFTER the test MN — its + rate must remain a meaningful drive proxy at test derecruitment; + * rate–rate correlation between control and test profiles over the test + MN's firing span ≥ ``min_rate_correlation`` (common-modulation check); + * control rate modulation (max − min over the test window) ≥ + ``min_ctrl_modulation_hz`` (saturation / sensitivity check); + * control firing (≥ ``min_ctrl_rate_hz``) at both endpoints. + + Controls are averaged within each test unit before the population mean + (``DeltaFResult.mean``), so test units with many valid controls are not + over-weighted. Raw pair values remain in ``delta_f``. + """ + trains = spike_times(spikes, dt_ms) + n_steps = np.asarray(spikes).shape[0] + t_grid = np.arange(n_steps) * dt_ms * 1e-3 + rt = recruitment_times(spikes, dt_ms) + dt_de = derecruitment_times(spikes, dt_ms) + + pairs, dfs, trs, tds = [], [], [], [] + active = np.flatnonzero(~np.isnan(rt)) + profiles: dict[int, np.ndarray] = {} + + def prof(i: int) -> np.ndarray: + if i not in profiles: + profiles[i] = _rate_profile(trains[i], t_grid) + return profiles[i] + + for test in active: + if np.isnan(dt_de[test]): + continue + ctrls = [ + c + for c in active + if c != test + and rt[c] <= rt[test] - min_lead_s + and not np.isnan(dt_de[c]) + and dt_de[c] >= dt_de[test] + min_control_outlast_s + ] + # prefer the controls recruited closest before the test MN + ctrls = sorted(ctrls, key=lambda c: rt[test] - rt[c])[:max_pairs_per_test] + i_r = min(int(np.searchsorted(t_grid, rt[test])), n_steps - 1) + i_d = min(int(np.searchsorted(t_grid, dt_de[test])), n_steps - 1) + if i_d <= i_r: + continue + test_prof = prof(test) + for c in ctrls: + p = prof(c) + if p[i_r] < min_ctrl_rate_hz or p[i_d] < min_ctrl_rate_hz: + continue + seg_c, seg_t = p[i_r:i_d], test_prof[i_r:i_d] + if seg_c.max() - seg_c.min() < min_ctrl_modulation_hz: + continue # control saturated / insensitive over the window + if seg_c.std() == 0 or seg_t.std() == 0: + continue # flat profile → common modulation undefined + r = float(np.corrcoef(seg_c, seg_t)[0, 1]) + if not np.isfinite(r) or r < min_rate_correlation: + continue # no (or undefined) common modulation + pairs.append((c, test)) + dfs.append(p[i_r] - p[i_d]) + trs.append(rt[test]) + tds.append(dt_de[test]) + + pairs_arr = np.array(pairs).reshape(-1, 2) + dfs_arr = np.array(dfs) + if len(pairs_arr): + test_units = np.unique(pairs_arr[:, 1]) + per_test = np.array( + [dfs_arr[pairs_arr[:, 1] == tu].mean() for tu in test_units] + ) + else: + test_units = np.array([], dtype=int) + per_test = np.array([]) + + return DeltaFResult( + pairs=pairs_arr, + delta_f=dfs_arr, + t_recruit_test=np.array(trs), + t_derecruit_test=np.array(tds), + test_units=test_units, + delta_f_per_test=per_test, + ) + + +def hysteresis_index( + drive: np.ndarray, spikes, dt_ms: float +) -> np.ndarray: + """Per-MN derecruitment−recruitment drive difference on a ramp. + + ``drive`` is the scalar command time-series ``(n_steps,)`` (e.g. DD rate). + Negative values (derecruited at *lower* drive than recruitment) indicate + sustained depolarization / PIC hysteresis. + """ + drive = np.asarray(drive) + t = np.arange(len(drive)) * dt_ms * 1e-3 + rt = recruitment_times(spikes, dt_ms) + dtd = derecruitment_times(spikes, dt_ms) + out = np.full(len(rt), np.nan) + for i, (a, b) in enumerate(zip(rt, dtd)): + if not (np.isnan(a) or np.isnan(b)): + out[i] = drive[np.searchsorted(t, b) - 1] - drive[np.searchsorted(t, a)] + return out diff --git a/experiments/twin/phase1_pic_validation.py b/experiments/twin/phase1_pic_validation.py new file mode 100644 index 00000000..0ca33ebd --- /dev/null +++ b/experiments/twin/phase1_pic_validation.py @@ -0,0 +1,308 @@ +"""Phase 1 (gap) — PIC-specific behaviour of the Jaxley NERLab pool. + +Triangular-ramp protocol (the ΔF paradigm): descending drive ramps up and back +down; PICs should produce (i) recruitment/derecruitment hysteresis — MNs turn +off at *lower* drive than they turned on, (ii) positive ΔF in paired-MU +analysis, (iii) self-sustained firing after the ramp returns to baseline. + +Dose–response over the gCaL grid (default ×{0, 1, 2, 4, 6.5}), optionally +crossed with a NaP multiplier (``--gnap-scale``). The repository's established +strong-PIC regime corresponds to roughly gCaL ×6.5–7.5 (pool ``gamma`` +1.3–1.5) with NaP ×5 — a grid restricted to low scales (×0–2) shows no +bistability and supports only the conclusion "none detected in that range", +NOT "the model lacks bistable PICs". + +Multiple stochastic-drive seeds (``--n-seeds``) give mean ± sd for ΔF and +hysteresis; single-seed numbers sit inside substantial pair-to-pair +variability. ``--seed`` also sets the connectivity seed. + +Caveat for NEURON comparison: with the default synapse options this harness +drives the SOMA with a fixed-driving-force surrogate current, while NEURON +places voltage-dependent Exp2Syn on the DENDRITE (the caL compartment) with +``tau1 = 0.2 ms, tau2 = 0.3 ms, E_exc = +70 mV, E_inh = −75 mV``. For a +matched configuration use:: + + --syn-site dendrite --live-df --tau-rise 0.2 --tau-decay 0.3 + +E_inh convention (must be stated with any paired result): the harness default +−15 mV is the frame-corrected analog of a modern −80 mV (driving force −15 mV +at the NERLab rest of 0 mV); NEURON's literal value is −75 mV. Pass +``--e-inh -75`` to match NEURON's number, or keep −15 to match its *physical* +driving force — these are different experiments. + +Variability: ``--n-seeds`` varies the stochastic drive realization; +``--n-conn-seeds`` varies the connectivity matrix. The two sources are +reported separately. + +Usage:: + + python phase1_pic_validation.py [--n-mn 20] [--peak 60] [--rise 7.5] + python phase1_pic_validation.py --scales 0 1 2 4 6.5 --gnap-scale 5 \\ + --syn-site dendrite --live-df --tau-rise 0.2 --tau-decay 0.3 \\ + --n-seeds 3 --n-conn-seeds 2 +""" + +from __future__ import annotations + +import argparse +import pickle +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import jax +import jax.numpy as jnp + +from poolbuild import PoolTwin, TwinSpec +import protocols +import observables as obs + +RESULTS = Path(__file__).parent / "results" + +# Sequential single-hue ramp (light→dark) — gCaL scale is a magnitude. +def make_ramp(n: int) -> list[str]: + cmap = plt.get_cmap("Blues") + return [matplotlib.colors.to_hex(cmap(x)) for x in np.linspace(0.35, 0.95, max(n, 2))] + + +INK, MUTED, GRID = "#1a1a1a", "#666666", "#e3e3e3" + + +def run_condition(twin, drives, key, gcal_scale, gnap_scale=1.0): + theta = twin.default_theta() + theta["gCaL_scale"] = jnp.float32(gcal_scale) + theta["gNaP_scale"] = jnp.float32(gnap_scale) + out = twin.simulate(theta, drives, key, spike_mode="hard") + return np.asarray(out["v"]), np.asarray(out["spikes"]) + + +def analyze(spikes, cmd, dt_ms, t_ramp_end_s): + t = np.arange(len(cmd)) * dt_ms * 1e-3 + df = obs.delta_f_pairs(spikes, dt_ms) + hyst = obs.hysteresis_index(cmd, spikes, dt_ms) + n_active = int((spikes.sum(axis=0) >= 3).sum()) + post = spikes[t >= t_ramp_end_s] + self_sustained = int((post.sum(axis=0) >= 5).sum()) + return { + "delta_f_mean": df.mean, + "delta_f_n_pairs": int(len(df.delta_f)), + "hysteresis_mean": float(np.nanmean(hyst)) if np.any(~np.isnan(hyst)) else np.nan, + "n_active": n_active, + "self_sustained": self_sustained, + "delta_f": df, + "hysteresis": hyst, + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-mn", type=int, default=20) + ap.add_argument("--n-dd", type=int, default=100) + ap.add_argument("--dt-ms", type=float, default=0.1) + ap.add_argument("--peak", type=float, default=60.0, help="peak DD rate (Hz)") + ap.add_argument("--rise", type=float, default=7.5, help="ramp rise/fall (s)") + ap.add_argument("--post", type=float, default=3.0) + ap.add_argument("--scales", type=float, nargs="+", + default=[0.0, 1.0, 2.0, 4.0, 6.5]) + ap.add_argument("--gnap-scale", type=float, default=1.0, + help="NaP multiplier applied to every condition " + "(strong-PIC regime ≈ 5)") + ap.add_argument("--n-seeds", type=int, default=1, + help="stochastic-drive seeds per condition (mean ± sd)") + ap.add_argument("--n-conn-seeds", type=int, default=1, + help="connectivity seeds (reported separately from drive " + "variability)") + ap.add_argument("--seed", type=int, default=0, + help="base seed (connectivity and generators derive from it)") + ap.add_argument("--syn-site", choices=["soma", "dendrite"], default="soma") + ap.add_argument("--live-df", action="store_true", + help="voltage-dependent driving force (NEURON-like)") + ap.add_argument("--tau-rise", type=float, default=None, + help="Exp2Syn rise time constant (ms); omit for single-exp. " + "NEURON NERLab uses 0.2") + ap.add_argument("--tau-decay", type=float, default=None, + help="synaptic decay time constant (ms, default 5.0). " + "NEURON NERLab uses 0.3") + ap.add_argument("--w-exc", type=float, default=None, help="excitatory weight (µS)") + ap.add_argument("--w-inh", type=float, default=None, help="inhibitory weight (µS)") + ap.add_argument("--e-exc", type=float, default=None, + help="excitatory reversal (mV, default +70)") + ap.add_argument("--e-inh", type=float, default=None, + help="inhibitory reversal (mV, default −15 frame-corrected; " + "NEURON's literal value is −75 — see docstring)") + args = ap.parse_args() + + spec_kwargs = dict( + n_mn=args.n_mn, n_dd=args.n_dd, dt_ms=args.dt_ms, + syn_site=args.syn_site, live_driving_force=args.live_df, + syn_tau_rise_ms=args.tau_rise, + ) + for cli, field in [("tau_decay", "tau_syn_ms"), ("w_exc", "w_exc_uS"), + ("w_inh", "w_inh_uS"), ("e_exc", "e_exc_mV"), + ("e_inh", "e_inh_mV")]: + if getattr(args, cli) is not None: + spec_kwargs[field] = getattr(args, cli) + + twins = [PoolTwin(TwinSpec(seed=args.seed + cs, **spec_kwargs)) + for cs in range(args.n_conn_seeds)] + twin = twins[0] # figure + single-run detail come from the first twin + t, cmd = protocols.triangular_ramp( + args.dt_ms, t_rise_s=args.rise, t_pre_s=1.0, t_post_s=args.post, + peak=args.peak, + ) + t_ramp_end = 1.0 + 2 * args.rise + drives = { + "dd": jnp.asarray(cmd), + "inh": jnp.zeros(len(cmd), dtype=jnp.float32), + } + results, spikes_by_scale = {}, {} + print(f"(gNaP ×{args.gnap_scale:g}, {args.n_seeds} seed(s), " + f"syn={args.syn_site}, {'live' if args.live_df else 'fixed'} DF)") + print(f"{'gCaL x':>7} | {'ΔF (Hz)':>14} | {'pairs':>5} " + f"| {'hyst (Hz drive)':>16} | {'active':>7} | {'self-sust':>9} " + f"| {'t (s)':>6}") + for s in args.scales: + t0 = time.time() + per_conn = [] # per_conn[cs] = list of per-drive-seed results + for cs, tw in enumerate(twins): + runs = [] + for si in range(args.n_seeds): + key = jax.random.PRNGKey(args.seed + 1000 * cs + si) + _, sp = run_condition(tw, drives, key, s, args.gnap_scale) + r = analyze(sp, cmd, args.dt_ms, t_ramp_end) + runs.append(r) + if cs == 0 and si == 0: + spikes_by_scale[s] = sp + per_conn.append(runs) + flat = [r for runs in per_conn for r in runs] + dfm = np.array([r["delta_f_mean"] for r in flat]) + hym = np.array([r["hysteresis_mean"] for r in flat]) + # separated variability: drive seeds (within first connectivity) vs + # connectivity seeds (across per-connectivity means) + drive_sd = float(np.nanstd([r["delta_f_mean"] for r in per_conn[0]])) + conn_means = [np.nanmean([r["delta_f_mean"] for r in runs]) + for runs in per_conn] + conn_sd = float(np.nanstd(conn_means)) + agg = dict(per_conn[0][0]) # first-run detail for the figure + agg.update({ + "delta_f_mean": float(np.nanmean(dfm)), + "delta_f_sd": float(np.nanstd(dfm)), + "delta_f_drive_sd": drive_sd, + "delta_f_conn_sd": conn_sd, + "delta_f_per_run": dfm, + "hysteresis_mean": float(np.nanmean(hym)), + "hysteresis_sd": float(np.nanstd(hym)), + "n_active": float(np.mean([r["n_active"] for r in flat])), + "self_sustained": float(np.mean([r["self_sustained"] for r in flat])), + "delta_f_n_pairs": int(np.sum([r["delta_f_n_pairs"] for r in flat])), + }) + results[s] = agg + multi = len(flat) > 1 + sd = f" ± {agg['delta_f_sd']:4.2f}" if multi else " " + hsd = f" ± {agg['hysteresis_sd']:4.2f}" if multi else " " + print(f"{s:7.2f} | {agg['delta_f_mean']:7.2f}{sd} " + f"| {agg['delta_f_n_pairs']:5d} | {agg['hysteresis_mean']:9.2f}{hsd} " + f"| {agg['n_active']:7.1f} | {agg['self_sustained']:9.1f} " + f"| {time.time()-t0:6.1f}") + if args.n_seeds > 1 or args.n_conn_seeds > 1: + print(f" ΔF variability — drive seeds: ± {drive_sd:.2f} " + f"(n={args.n_seeds}); connectivity seeds: ± {conn_sd:.2f} " + f"(n={args.n_conn_seeds})") + + # ---------------------------------------------------------------- figure + ramp = make_ramp(len(args.scales)) + fig, axes = plt.subplots(2, 2, figsize=(11, 7.5)) + fig.patch.set_facecolor("white") + for ax in axes.flat: + ax.grid(True, color=GRID, linewidth=0.6) + ax.set_axisbelow(True) + for side in ("top", "right"): + ax.spines[side].set_visible(False) + + # (a) rasters, lowest vs highest gCaL + ax = axes[0, 0] + lo, hi = args.scales[0], args.scales[-1] + for row, (s, col) in enumerate([(lo, ramp[0]), (hi, ramp[-1])]): + sp = spikes_by_scale[s] + for i in range(sp.shape[1]): + times = np.flatnonzero(sp[:, i] > 0.5) * args.dt_ms * 1e-3 + y = i + row * (args.n_mn + 3) + ax.plot(times, np.full(len(times), y), "|", color=col, markersize=3) + ax2 = ax.twinx() # drive overlay on its own (hidden) scale, gray + ax2.plot(t, cmd, color=MUTED, linewidth=1.2) + ax2.set_yticks([]) + for side in ("top", "right"): + ax2.spines[side].set_visible(False) + ax.set_title(f"Raster: gCaL ×{lo:g} (bottom) vs ×{hi:g} (top); drive in gray", + color=INK, fontsize=10) + ax.set_xlabel("time (s)") + ax.set_ylabel("MN index") + + # (b) ΔF dose–response + ax = axes[0, 1] + means = [results[s]["delta_f_mean"] for s in args.scales] + for s, m, c in zip(args.scales, means, ramp): + pts = results[s]["delta_f"].delta_f + if len(pts): + ax.plot(np.full(len(pts), s), pts, "o", color=c, alpha=0.4, markersize=4) + ax.plot(args.scales, means, "-", color=INK, linewidth=1.5) + if args.n_seeds > 1 or args.n_conn_seeds > 1: + sds = [results[s]["delta_f_sd"] for s in args.scales] + ax.errorbar(args.scales, means, yerr=sds, fmt="none", ecolor=INK, + elinewidth=1.0, capsize=3) + ax.axhline(0, color=MUTED, linewidth=0.8) + ax.set_title("ΔF vs gCaL scale (dots: MU pairs; line: mean)", color=INK, fontsize=10) + ax.set_xlabel("gCaL scale") + ax.set_ylabel("ΔF (Hz)") + + # (c) drive hysteresis per MN + ax = axes[1, 0] + for s, c in zip(args.scales, ramp): + h = results[s]["hysteresis"] + ok = ~np.isnan(h) + ax.plot(np.flatnonzero(ok), h[ok], "o-", color=c, markersize=4, + linewidth=1.0, label=f"×{s:g}") + ax.axhline(0, color=MUTED, linewidth=0.8) + ax.legend(frameon=False, fontsize=8, title="gCaL") + ax.set_title("Derecruitment − recruitment drive (negative ⇒ hysteresis)", + color=INK, fontsize=10) + ax.set_xlabel("MN index") + ax.set_ylabel("Δ drive (Hz)") + + # (d) population smoothed rate over the ramp + ax = axes[1, 1] + for s, c in zip(args.scales, ramp): + pr = np.asarray(obs.population_rate(jnp.asarray(spikes_by_scale[s]), + args.dt_ms, tau_ms=200.0)) + ax.plot(t, pr, color=c, linewidth=1.4, label=f"×{s:g}") + ax.legend(frameon=False, fontsize=8, title="gCaL") + ax.set_title("Population rate (note offset-limb excess with PICs)", + color=INK, fontsize=10) + ax.set_xlabel("time (s)") + ax.set_ylabel("mean rate (Hz)") + + fig.tight_layout() + RESULTS.mkdir(exist_ok=True) + fig.savefig(RESULTS / "phase1_pic_validation.png", dpi=150) + with open(RESULTS / "phase1_pic_validation.pkl", "wb") as f: + pickle.dump( + { + "args": vars(args), + "results": { + s: {k: v for k, v in r.items() if k not in ("delta_f",)} + for s, r in results.items() + }, + "spikes_by_scale": spikes_by_scale, + "cmd": cmd, + }, + f, + ) + print(f"saved → {RESULTS / 'phase1_pic_validation.png'}") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/phase2_benchmark.py b/experiments/twin/phase2_benchmark.py new file mode 100644 index 00000000..bc8cc8ce --- /dev/null +++ b/experiments/twin/phase2_benchmark.py @@ -0,0 +1,132 @@ +"""Phase 2 — the computational leap: batched simulation + gradients. + +``jax.vmap`` maps the pool twin's forward (and ``value_and_grad``) over K +parameter vectors simultaneously — the thing NEURON cannot do at all. This +script measures throughput vs batch size on whatever device JAX finds (CPU on +the dev Mac; rerun unchanged on a GPU node). + +Reported per K: steady-state wall time (compile excluded), simulations/s, and +MN·steps/s, for (a) forward hard-spike runs and (b) full value_and_grad. + +Usage:: + + python phase2_benchmark.py [--batch-sizes 1 8 32 128] [--t-on 1.0] +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import jax +import jax.numpy as jnp +import jax.tree_util as jtu +import numpy as np + +from poolbuild import PoolTwin, TwinSpec +import protocols +import observables as obs + +RESULTS = Path(__file__).parent / "results" + + +def batched_thetas(twin: PoolTwin, K: int, rng: np.random.Generator) -> dict: + """Stack K perturbed parameter sets leaf-wise: each leaf becomes (K,).""" + base = twin.default_theta() + factors = rng.uniform(0.5, 2.0, size=(K, len(base))).astype(np.float32) + return { + k: jnp.asarray(factors[:, j]) * v + for j, (k, v) in enumerate(sorted(base.items())) + } + + +def time_fn(fn, *a, n_rep: int = 3): + out = fn(*a) # compile + jax.block_until_ready(out) + t0 = time.time() + for _ in range(n_rep): + out = fn(*a) + jax.block_until_ready(out) + return (time.time() - t0) / n_rep + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-mn", type=int, default=10) + ap.add_argument("--n-dd", type=int, default=100) + ap.add_argument("--dt-ms", type=float, default=0.1) + ap.add_argument("--t-on", type=float, default=1.0) + ap.add_argument("--level", type=float, default=40.0) + ap.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 8, 32, 128]) + ap.add_argument("--seed", type=int, default=0, + help="connectivity seed and generator seed") + ap.add_argument("--device", choices=["cpu", "gpu", "tpu"], default=None, + help="JAX platform to run on (default: JAX's default)") + args = ap.parse_args() + + if args.device is not None: + try: + jax.config.update("jax_default_device", jax.devices(args.device)[0]) + except RuntimeError: + print(f"no {args.device} device available; " + f"falling back to {jax.devices()[0].platform}") + + twin = PoolTwin(TwinSpec(n_mn=args.n_mn, n_dd=args.n_dd, dt_ms=args.dt_ms, + seed=args.seed)) + _, cmd = protocols.step( + args.dt_ms, t_pre_s=0.2, t_on_s=args.t_on, t_post_s=0.2, level=args.level + ) + drives = {"dd": jnp.asarray(cmd), "inh": jnp.zeros(len(cmd), jnp.float32)} + key = jax.random.PRNGKey(args.seed) + rng = np.random.default_rng(args.seed) + n_steps = len(cmd) + dev = jax.devices()[0] + print(f"device={dev.platform} | n_mn={args.n_mn} | n_steps={n_steps}") + + def loss_fn(th): + out = twin._run(th, drives, key, spike_mode="surrogate") + rates = obs.smoothed_rates(out["spikes"], args.dt_ms, tau_ms=150.0) + return jnp.mean(rates**2) + + fwd_b = jax.jit( + jax.vmap(lambda th: twin._run(th, drives, key, spike_mode="hard")["spikes"]) + ) + vg_b = jax.jit(jax.vmap(jax.value_and_grad(loss_fn))) + + rows = [] + print(f"{'K':>5} | {'fwd (s)':>8} | {'fwd sims/s':>10} | {'fwd MN·steps/s':>14} " + f"| {'grad (s)':>8} | {'grad sims/s':>11}") + for K in args.batch_sizes: + thetas = batched_thetas(twin, K, rng) + t_fwd = time_fn(fwd_b, thetas) + t_vg = time_fn(vg_b, thetas) + row = { + "K": K, + "t_forward_s": t_fwd, + "forward_sims_per_s": K / t_fwd, + "forward_mn_steps_per_s": K * args.n_mn * n_steps / t_fwd, + "t_grad_s": t_vg, + "grad_sims_per_s": K / t_vg, + } + rows.append(row) + print(f"{K:5d} | {t_fwd:8.2f} | {row['forward_sims_per_s']:10.2f} " + f"| {row['forward_mn_steps_per_s']:14.3e} | {t_vg:8.2f} " + f"| {row['grad_sims_per_s']:11.2f}") + + RESULTS.mkdir(exist_ok=True) + payload = { + "device": dev.platform, + "n_mn": args.n_mn, + "n_steps": n_steps, + "dt_ms": args.dt_ms, + "rows": rows, + } + out_path = RESULTS / f"phase2_benchmark_{dev.platform}.json" + out_path.write_text(json.dumps(payload, indent=2)) + print(f"saved → {out_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/phase3_inverse_mvp.py b/experiments/twin/phase3_inverse_mvp.py new file mode 100644 index 00000000..dffa2ac6 --- /dev/null +++ b/experiments/twin/phase3_inverse_mvp.py @@ -0,0 +1,203 @@ +"""Phase 3 MVP — the inverse problem on synthetic data with known physiology. + +Protocol: +1. Simulate "observations" with hidden ground-truth physiology θ* (hard spikes, + the scientific forward): PIC conductances scaled, altered excitation, altered + inhibition. +2. Observables = exponentially-smoothed per-MN firing rates (what one gets from + decomposed EMG in practice). +3. Recover θ = {gCaL_scale, gNaP_scale, drive_gain, w_inh} by gradient descent + (optax Adam in log-parameter space; surrogate spikes for the backward pass, + which are *forward-identical* to hard spikes in pathwise mode — so the loss + at θ = θ* is exactly 0 with the same PRNG key). + +``w_exc`` is held fixed at its known value in the MVP (drive_gain and w_exc are +strongly degenerate through the excitatory path; separating them is exactly the +Phase 5 identifiability question). + +Exploding-gradient note: BPTT through ~40k timesteps of stiff spiking dynamics +yields gradient magnitudes of 1e16–1e22 (measured; both surrogate and rate +modes). Raw Adam silently freezes — squaring such gradients overflows float32 +to inf and the update becomes 0/inf = 0. We therefore clip element-wise +(``optax.clip``) *before* Adam: magnitudes are chaos-dominated but signs remain +informative, so the fit is effectively momentum sign-descent in log-space. +BE CLEAR ABOUT WHAT THIS IS: a heuristic sign method, not gradient descent on +a trustworthy objective — relative sensitivities are discarded, the surrogate +backward pass is not the derivative of the hard-spike objective, and these +gradients must NOT be reused for Fisher-information / identifiability claims. +Results are reported at the best-loss iterate; if the best loss is not clearly +below the initial loss, no recovery has been demonstrated. Taming the +magnitudes properly (truncated BPTT windows, smoothed-dynamics losses, +hard-mode loss-landscape checks) is on the roadmap. + +Usage:: + + python phase3_inverse_mvp.py [--iters 100] [--lr 0.02] +""" + +from __future__ import annotations + +import argparse +import pickle +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import jax +import jax.numpy as jnp +import optax + +from poolbuild import PoolTwin, TwinSpec +import protocols +import observables as obs + +RESULTS = Path(__file__).parent / "results" +FIT_KEYS = ("gCaL_scale", "gNaP_scale", "drive_gain", "w_inh") + +# Okabe–Ito (colorblind-safe) categorical order for the four parameters. +CAT = {"gCaL_scale": "#0072B2", "gNaP_scale": "#E69F00", + "drive_gain": "#009E73", "w_inh": "#CC79A7"} +INK, MUTED, GRID = "#1a1a1a", "#666666", "#e3e3e3" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-mn", type=int, default=8) + ap.add_argument("--n-dd", type=int, default=100) + ap.add_argument("--dt-ms", type=float, default=0.1) + ap.add_argument("--peak", type=float, default=50.0) + ap.add_argument("--inh-rate", type=float, default=20.0) + ap.add_argument("--iters", type=int, default=100) + ap.add_argument("--lr", type=float, default=0.02) + ap.add_argument("--tau-ms", type=float, default=150.0) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + twin = PoolTwin(TwinSpec(n_mn=args.n_mn, n_dd=args.n_dd, dt_ms=args.dt_ms, + seed=args.seed)) + _, cmd = protocols.trapezoid( + args.dt_ms, t_rise_s=1.0, t_hold_s=1.5, t_fall_s=1.0, + t_pre_s=0.3, t_post_s=0.3, peak=args.peak, + ) + drives = { + "dd": jnp.asarray(cmd), + "inh": jnp.full(len(cmd), args.inh_rate, dtype=jnp.float32), + } + key = jax.random.PRNGKey(args.seed) + + # ---------------------------------------------------- ground truth + obs + theta_true = twin.default_theta() + theta_true.update({ + "gCaL_scale": jnp.float32(1.6), + "gNaP_scale": jnp.float32(0.7), + "drive_gain": jnp.float32(0.85), + "w_inh": jnp.float32(0.006), + }) + out_obs = twin.simulate(theta_true, drives, key, spike_mode="hard") + obs_rates = obs.smoothed_rates(out_obs["spikes"], args.dt_ms, args.tau_ms) + obs_rates = jax.lax.stop_gradient(obs_rates) + n_active = int((np.asarray(out_obs["spikes"]).sum(axis=0) >= 3).sum()) + print(f"observations: {n_active}/{args.n_mn} MNs active, " + f"mean rate {float(obs_rates.mean()):.2f} Hz") + + # ------------------------------------------------------------- fit setup + theta0 = twin.default_theta() # deliberately wrong start (all defaults) + + def loss_from_z(z): + theta = dict(theta0) + for k in FIT_KEYS: + theta[k] = jnp.exp(z[k]) + out = twin._run(theta, drives, key, spike_mode="surrogate") + rates = obs.smoothed_rates(out["spikes"], args.dt_ms, args.tau_ms) + return jnp.mean((rates - obs_rates) ** 2) + + z = {k: jnp.log(theta0[k]) for k in FIT_KEYS} + # clip BEFORE adam: raw grads are ~1e20 and overflow adam's g² to inf. + opt = optax.chain(optax.clip(1.0), optax.adam(args.lr)) + opt_state = opt.init(z) + vg = jax.jit(jax.value_and_grad(loss_from_z)) + + history = {"loss": [], "theta": []} + best = {"loss": np.inf, "theta": None, "it": -1} + t0 = time.time() + for it in range(args.iters): + loss, grads = vg(z) + # loss is evaluated at the CURRENT z (before the update), so pair it + # with the pre-update parameters for best-iterate tracking. + th_now = {k: float(jnp.exp(z[k])) for k in FIT_KEYS} + history["loss"].append(float(loss)) + history["theta"].append(th_now) + if float(loss) < best["loss"]: + best = {"loss": float(loss), "theta": th_now, "it": it} + updates, opt_state = opt.update(grads, opt_state) + z = optax.apply_updates(z, updates) + if it % 5 == 0 or it == args.iters - 1: + rel = ", ".join( + f"{k}={th_now[k] / float(theta_true[k]):.3f}×" for k in FIT_KEYS + ) + print(f"it {it:4d} | loss {float(loss):.5f} | rel-to-true: {rel} " + f"| {time.time()-t0:.0f}s") + + # -------------------------------------------------------------- report + # Report the BEST iterate (lowest loss), not the last one — with the + # clipped sign-descent updates the trajectory is non-monotonic and the + # final iterate can be worse than the start. If best ≈ init loss, the + # optimization has NOT demonstrated recovery; say so rather than citing + # per-parameter errors from an unconverged trajectory. + final = history["theta"][-1] + print(f"\nbest iterate: {best['it']} (loss {best['loss']:.5f}); " + f"initial loss {history['loss'][0]:.5f}; " + f"final loss {history['loss'][-1]:.5f}") + if history["loss"][-1] > best["loss"]: + print("WARNING: optimization is non-monotonic — final iterate is worse " + "than the best; treating the best iterate as the result.") + print("\nrecovered (best iterate) vs true:") + print(f"{'param':>12} | {'true':>9} | {'init':>9} | {'best':>9} " + f"| {'final':>9} | {'best err':>8}") + for k in FIT_KEYS: + tv, iv = float(theta_true[k]), float(theta0[k]) + bv, fv = best["theta"][k], final[k] + print(f"{k:>12} | {tv:9.4f} | {iv:9.4f} | {bv:9.4f} | {fv:9.4f} " + f"| {abs(bv - tv) / tv * 100:7.1f}%") + + # figure: loss curve + parameter trajectories relative to truth + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4)) + for ax in (ax1, ax2): + ax.grid(True, color=GRID, linewidth=0.6) + ax.set_axisbelow(True) + for side in ("top", "right"): + ax.spines[side].set_visible(False) + ax1.semilogy(history["loss"], color="#0072B2", linewidth=1.6) + ax1.plot(best["it"], best["loss"], "o", color=INK, markersize=6, zorder=5) + ax1.annotate(f"best (it {best['it']})", (best["it"], best["loss"]), + textcoords="offset points", xytext=(6, -10), fontsize=8, + color=INK) + ax1.set_xlabel("iteration") + ax1.set_ylabel("rate MSE (Hz²)") + ax1.set_title("Loss (non-monotonic: sign-descent)", color=INK, fontsize=10) + its = np.arange(len(history["theta"])) + for k in FIT_KEYS: + traj = np.array([h[k] for h in history["theta"]]) / float(theta_true[k]) + ax2.plot(its, traj, color=CAT[k], linewidth=1.6, label=k) + ax2.annotate(k, (its[-1], traj[-1]), textcoords="offset points", + xytext=(4, 0), fontsize=8, color=CAT[k]) + ax2.axhline(1.0, color=MUTED, linewidth=0.8, linestyle="--") + ax2.legend(frameon=False, fontsize=8) + ax2.set_xlabel("iteration") + ax2.set_ylabel("recovered / true") + ax2.set_title("Parameter recovery (1.0 = truth)", color=INK, fontsize=10) + fig.tight_layout() + RESULTS.mkdir(exist_ok=True) + fig.savefig(RESULTS / "phase3_inverse_mvp.png", dpi=150) + with open(RESULTS / "phase3_inverse_mvp.pkl", "wb") as f: + pickle.dump({"args": vars(args), "history": history, "best": best, + "theta_true": {k: float(theta_true[k]) for k in FIT_KEYS}}, f) + print(f"saved → {RESULTS / 'phase3_inverse_mvp.png'}") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/pic_grid.py b/experiments/twin/pic_grid.py new file mode 100644 index 00000000..681969f6 --- /dev/null +++ b/experiments/twin/pic_grid.py @@ -0,0 +1,338 @@ +"""Confirmatory NEURON–Jaxley PIC validation grid (Phase 1, work package 2). + +Grid: gCaL × {0, 0.5, 1, 2, 4, 6.5, 7.5} crossed with gNaP × {1, 5}, on +PREVIOUSLY UNUSED seeds (confirmatory — the acceptance thresholds in +``configs/phase1_acceptance_v1.yaml`` were frozen from exploratory data +before this script ran). Identical frozen artifacts feed both simulators: +same cells, thresholds, connectivity, event matrices, synaptic site/kinetics +/reversals/weights (0.15 µS firing regime), zero delays, backward Euler. + +Protocol: slow triangular ramp (7.5 s rise/fall, 60 Hz peak) — the ΔF/PIC +paradigm. Per (condition × seed) the paired metrics are computed and the +FROZEN acceptance spec classifies every metric PASS / FAIL / +OUTSIDE_VALIDATED_ENVELOPE / NOT_EVALUABLE; the maximal passing CaL range +per NaP level is the candidate validated envelope. + +Timestep-refinement references (dt 0.1 → 0.05, Jaxley) are run for two +sentinel conditions to anchor E_numerical. + +Process isolation: Jaxley runs in-parent; all NEURON conditions for a seed +run in ONE subprocess (JAX and NEURON cannot share a process), reusing one +pool and rescaling conductances between conditions. + +Usage:: + + python pic_grid.py # full confirmatory run (~30 min CPU) + python pic_grid.py --seeds 101 --gcal 0 1 6.5 --dur-scale 0.5 # smoke +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +import time +from pathlib import Path + +import joblib +import numpy as np + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +import observables as obs +import validation_protocol as vp +from acceptance import evaluate_grid, load_spec + +RESULTS = Path(__file__).parent / "results" +EVIDENCE = Path(__file__).parent / "evidence" + +GCAL_GRID = [0.0, 0.5, 1.0, 2.0, 4.0, 6.5, 7.5] +GNAP_GRID = [1.0, 5.0] +DT_REF_CONDITIONS = [(1.0, 1.0), (6.5, 5.0)] + +_NEURON_CODE = r""" +import sys, joblib +import numpy as np +import validation_protocol as vp + +art_path, out_path = sys.argv[1], sys.argv[2] +conds = [tuple(map(float, c.split(","))) for c in sys.argv[3].split(";")] +art = vp.load_protocol(art_path) +dt = art["dt_ms"]; n_mn = art["n_mn"]; syn_p = art["synapse"] + +h, pool = vp._neuron_pool(art) # art cell scales are 1,1 — bases below +base = [(c.dend[0].gcaLbar_caL, c.soma.gnapbar_napp) for c in pool] + +keep, vstims, _ = [], [], None +for events, conn, syn_idx, w, e in ( + (art["exc_events"], art["exc_conn"], 0, syn_p["w_exc_uS"], syn_p["e_exc_mV"]), + (art["inh_events"], art["inh_conn"], 1, syn_p["w_inh_uS"], syn_p["e_inh_mV"])): + stims = [] + for s in range(events.shape[1]): + tv = h.Vector((np.flatnonzero(events[:, s]) * dt).astype(float)) + vs = h.VecStim(); vs.play(tv); stims.append((vs, tv)) + vstims.append(stims) + for i, c in enumerate(pool): + syn = c.synapse__list[syn_idx] + syn.tau1 = syn_p["tau_rise_ms"]; syn.tau2 = syn_p["tau_decay_ms"]; syn.e = e + for s in np.flatnonzero(conn[i]): + nc = h.NetCon(stims[s][0], syn); nc.weight[0] = w; nc.delay = 0.0 + keep.append(nc) +rec_keep, vrec, dvrec, spk = vp._neuron_record(h, pool) +n_steps = art["exc_events"].shape[0] + +out = {} +for gcal, gnap in conds: + for c, (b_cal, b_nap) in zip(pool, base): + c.dend[0].gcaLbar_caL = b_cal * gcal + c.soma.gnapbar_napp = b_nap * gnap + for sv in spk: + sv.resize(0) + vp._neuron_advance(h, pool, dt, n_steps) + res = vp._collect(vrec, dvrec, spk, dt, n_steps, n_mn) + out[(gcal, gnap)] = res["spikes"] + print(f"NEURON ({gcal},{gnap}): {int(res['spikes'].sum())} spikes", flush=True) +joblib.dump(out, out_path) +""" + + +def jaxley_spikes(twin, art, gcal, gnap): + import jax.numpy as jnp + + th = twin.default_theta() + th["gCaL_scale"] = jnp.float32(gcal) + th["gNaP_scale"] = jnp.float32(gnap) + out = twin.simulate_events( + th, art["exc_events"], inh_events=art["inh_events"], + exc_connectivity=art["exc_conn"], inh_connectivity=art["inh_conn"], + spike_mode="hard") + return np.asarray(out["spikes"]) + + +def paired_metrics(sj, sn, cmd, dt_ms, t_ramp_end_s): + from deliverable_b import binned_rates_jnp + import jax.numpy as jnp + + rj = np.asarray(binned_rates_jnp(jnp.asarray(sj), dt_ms)) + rn = np.asarray(binned_rates_jnp(jnp.asarray(sn), dt_ms)) + aj = set(np.flatnonzero(sj.sum(axis=0) > 0).tolist()) + an = set(np.flatnonzero(sn.sum(axis=0) > 0).tolist()) + tj, tn = int(sj.sum()), int(sn.sum()) + matched = sum(vp._match_spikes( + np.flatnonzero(sj[:, i]) * dt_ms * 1e-3, + np.flatnonzero(sn[:, i]) * dt_ms * 1e-3, 5e-3) + for i in range(sj.shape[1])) + t = np.arange(len(cmd)) * dt_ms * 1e-3 + post = t >= t_ramp_end_s + last = t >= (t[-1] - 0.5) + + def side_stats(sp): + hyst = obs.hysteresis_index(cmd, sp, dt_ms) + hmean = float(np.nanmean(hyst)) if np.any(~np.isnan(hyst)) else None + df = obs.delta_f_pairs(sp, dt_ms) + n_act = int((sp.sum(axis=0) > 0).sum()) + possible = max(n_act * (n_act - 1) // 2, 1) + return { + "hyst": hmean, + "df_mean": df.mean, + "df_pairs": int(len(df.delta_f)), + "df_fraction": len(df.delta_f) / possible, + "self_sustained": int((sp[post].sum(axis=0) >= 5).sum()), + "censored": int((sp[last].sum(axis=0) > 0).sum()), + } + + j, n = side_stats(sj), side_stats(sn) + hyst_diff = (abs(j["hyst"] - n["hyst"]) + if j["hyst"] is not None and n["hyst"] is not None else None) + df_diff = (abs(j["df_mean"] - n["df_mean"]) + if np.isfinite(j["df_mean"]) and np.isfinite(n["df_mean"]) + else None) + return { + "binned_rate_rmse_hz2": float(((rj - rn) ** 2).mean()), + "active_unit_mismatch": len(aj.symmetric_difference(an)), + "active_jaxley": len(aj), "active_neuron": len(an), + "total_spike_rel_diff": abs(tj - tn) / max(tj, tn, 1), + "total_jaxley": tj, "total_neuron": tn, + "spike_match_pm5ms_fraction": matched / max(tj, tn, 1), + "spike_total_max": max(tj, tn), + "hysteresis_mean_diff_hz": hyst_diff, + "hysteresis_jaxley": j["hyst"], "hysteresis_neuron": n["hyst"], + "self_sustained_mismatch": abs(j["self_sustained"] - n["self_sustained"]), + "self_sustained_jaxley": j["self_sustained"], + "self_sustained_neuron": n["self_sustained"], + "after_discharge_censored_jaxley": j["censored"], + "after_discharge_censored_neuron": n["censored"], + "delta_f_mean_diff_hz": df_diff, + "delta_f_pairs_jaxley": j["df_pairs"], + "delta_f_pairs_neuron": n["df_pairs"], + "delta_f_eligible_fraction_diff": abs(j["df_fraction"] - n["df_fraction"]), + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--seeds", type=int, nargs="+", default=[101, 202], + help="confirmatory seeds — must be previously unused") + ap.add_argument("--gcal", type=float, nargs="+", default=GCAL_GRID) + ap.add_argument("--gnap", type=float, nargs="+", default=GNAP_GRID) + ap.add_argument("--n-mn", type=int, default=8) + ap.add_argument("--dur-scale", type=float, default=1.0) + ap.add_argument("--skip-dt-ref", action="store_true") + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + tag = ("full" if (args.dur_scale >= 1.0 and len(args.seeds) >= 2 + and len(args.gcal) >= 7) else "smoke") + stem = f"pic_grid_{tag}" + t0_all = time.time() + spec = load_spec(1) + conds = [(g, n) for n in args.gnap for g in args.gcal] + cond_str = ";".join(f"{g},{n}" for g, n in conds) + + points = [] + dt_refs = {} + for seed in args.seeds: + t0 = time.time() + art = vp.make_protocol( + n_mn=args.n_mn, t_rise_s=7.5 * args.dur_scale, + t_pre_s=1.0, t_post_s=3.0, peak_hz=60.0, + w_exc_uS=0.15, step_currents_nA=(), seed=seed) + art_path = RESULTS / f"pic_grid_art_seed{seed}.pkl" + vp.save_protocol(art, art_path) + n = len(art["ramp_command_hz"]) + t_ramp_end = 1.0 + 2 * 7.5 * args.dur_scale + print(f"[seed {seed}] artifact {art['content_hash']} " + f"({time.time()-t0:.0f}s)", flush=True) + + # NEURON: all conditions in one subprocess + t0 = time.time() + npkl = RESULTS / f"pic_grid_neuron_seed{seed}.pkl" + proc = subprocess.run( + [sys.executable, "-u", "-c", _NEURON_CODE, str(art_path), + str(npkl), cond_str], + cwd=Path(__file__).parent, capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError(f"neuron subprocess failed:\n{proc.stdout}" + f"\n{proc.stderr}") + neuron_out = joblib.load(npkl) + print(f"[seed {seed}] NEURON {len(conds)} conditions " + f"({time.time()-t0:.0f}s)", flush=True) + + # Jaxley in-parent, one twin per seed + twin = vp._twin_from_artifact(art) + for gcal, gnap in conds: + t0 = time.time() + sj = jaxley_spikes(twin, art, gcal, gnap) + sn = neuron_out[(gcal, gnap)] + m = paired_metrics(sj, sn, art["ramp_command_hz"], art["dt_ms"], + t_ramp_end) + points.append({"gcal": gcal, "gnap": gnap, "seed": seed, + "metrics": m}) + print(f" ({gcal:4.1f},{gnap:3.1f}) seed{seed}: " + f"spk j{m['total_jaxley']}/n{m['total_neuron']} " + f"match {m['spike_match_pm5ms_fraction']:.2f} " + f"rateMSE {m['binned_rate_rmse_hz2']:.2f} " + f"selfsust j{m['self_sustained_jaxley']}/" + f"n{m['self_sustained_neuron']} ({time.time()-t0:.0f}s)", + flush=True) + + if not args.skip_dt_ref and seed == args.seeds[0]: + from deliverable_b import binned_rates_jnp + import jax.numpy as jnp + from poolbuild import PoolTwin, TwinSpec + + for gcal, gnap in DT_REF_CONDITIONS: + if (gcal, gnap) not in conds: + continue + twin_h = PoolTwin(TwinSpec( + n_mn=args.n_mn, n_dd=100, + n_inh=art["inh_events"].shape[1], dt_ms=0.05, seed=seed, + syn_site=art["synapse"]["site"], + live_driving_force=True, + syn_tau_rise_ms=art["synapse"]["tau_rise_ms"], + tau_syn_ms=art["synapse"]["tau_decay_ms"], + e_exc_mV=art["synapse"]["e_exc_mV"], + e_inh_mV=art["synapse"]["e_inh_mV"], + w_exc_uS=art["synapse"]["w_exc_uS"], + w_inh_uS=art["synapse"]["w_inh_uS"])) + fine = np.zeros((n * 2, art["exc_events"].shape[1]), + dtype=np.float32) + fine[::2] = np.asarray(art["exc_events"]) + th = twin_h.default_theta() + th["gCaL_scale"] = jnp.float32(gcal) + th["gNaP_scale"] = jnp.float32(gnap) + sp_f = np.asarray(twin_h.simulate_events( + th, fine, exc_connectivity=art["exc_conn"], + spike_mode="hard")["spikes"]) + r_f = np.asarray(binned_rates_jnp(jnp.asarray(sp_f), 0.05)) + sj = jaxley_spikes(twin, art, gcal, gnap) + r_c = np.asarray(binned_rates_jnp(jnp.asarray(sj), 0.1)) + m = min(len(r_f), len(r_c)) + dt_refs[f"{gcal},{gnap}"] = float( + ((r_f[:m] - r_c[:m]) ** 2).mean()) + print(f" dt-ref ({gcal},{gnap}): E_numerical " + f"{dt_refs[f'{gcal},{gnap}']:.3f} Hz²", flush=True) + + # --------------------------------------------------- acceptance + report + accept = evaluate_grid(points, spec) + report = { + "manifest": { + **__import__("sensitivity_checks")._manifest(), + "pic_grid_sha256_16": hashlib.sha256( + Path(__file__).read_bytes()).hexdigest()[:16], + "args": vars(args), "run_tag": tag, + "acceptance_spec_version": spec["version"], + }, + "dt_refinement_E_numerical_hz2": dt_refs, + "points": points, + "acceptance": {k: v for k, v in accept.items() if k != "points"}, + "acceptance_points": accept["points"], + "wall_s": round(time.time() - t0_all, 1), + } + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"{stem}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"{stem}_{stamp}.json").write_text(payload) + + # ------------------------------------------------------------- figure + fig, axes = plt.subplots(1, 2, figsize=(12, 4.5)) + for ax, gnap in zip(axes, sorted({p["gnap"] for p in points})): + rows = [p for p in points if p["gnap"] == gnap] + gcals = sorted({p["gcal"] for p in rows}) + for backend, c, mk in (("jaxley", "#0072B2", "o"), + ("neuron", "#E69F00", "s")): + ys = [np.mean([p["metrics"][f"self_sustained_{backend}"] + for p in rows if p["gcal"] == g]) for g in gcals] + ax.plot(gcals, ys, mk + "-", color=c, label=backend, mfc="none") + b = accept["bistability_boundary"][gnap] + ax.set_title(f"gNaP ×{gnap:g} — self-sustained units; boundary " + f"j:{b['jaxley_first_bistable_gcal']} " + f"n:{b['neuron_first_bistable_gcal']} [{b['class']}]", + fontsize=9) + ax.set_xlabel("gCaL scale") + ax.set_ylabel("self-sustained units (mean over seeds)") + ax.legend(frameon=False, fontsize=8) + ax.grid(True, color="#e3e3e3", lw=0.6) + ax.set_axisbelow(True) + fig.tight_layout() + fig.savefig(RESULTS / f"{stem}_boundary.png", dpi=150) + + print("\n=== acceptance summary (spec v%d, FROZEN) ===" % spec["version"]) + for k in ("n_points", "n_evaluated_metric_cells", "n_failed_metric_cells", + "fail_fraction", "failure_policy_ok"): + print(f" {k}: {accept[k]}") + print(f" bistability boundaries: " + f"{json.dumps(accept['bistability_boundary'], default=float)}") + print(f" validated envelope (gCaL range per gNaP): " + f"{accept['validated_envelope_gcal_by_gnap']}") + print(f"saved → {RESULTS / (stem + '.json')} (+ evidence copy, figure)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/poolbuild.py b/experiments/twin/poolbuild.py new file mode 100644 index 00000000..e6096a70 --- /dev/null +++ b/experiments/twin/poolbuild.py @@ -0,0 +1,516 @@ +"""Build a differentiable NERLab motor-neuron-pool twin. + +Wraps the canonical ex03 build path (AlphaMN__Pool → jx.Network → +``build_init_and_step_fn``) into a reusable object with: + +* **excitatory descending drive**: per-source renewal-process spikes (the + pure-JAX hazard generator from ``jax_models``; with ``dd_batch > 1`` this is + a gamma/Erlang renewal process, NOT pure Poisson — ``dd_batch=1`` gives + Poisson) → synaptic conductance → current; +* **inhibitory input**: an independent renewal population through the same + synapse model with E_inh = −15 mV (the frame-shifted analog of −80 mV); +* a **configurable synapse model** (see below) so the twin can either match + the fast ex03 surrogate or approximate NEURON's dendritic ``Exp2Syn``; +* a **differentiable parameter PyTree** ``theta``: + ``{"gCaL_scale", "gNaP_scale", "drive_gain", "w_exc", "w_inh"}`` — the PIC + conductances enter by scaling the ``caL_gcaLbar`` / ``napp_gnapbar`` leaves + of the jaxley parameter dict inside the traced computation. + +Synapse model options (``TwinSpec``) +------------------------------------ +``syn_site`` + ``"soma"`` (ex03 surrogate convention) or ``"dendrite"`` — the dendrite is + the compartment carrying ``caL``, which is where NEURON's MyoGen places its + excitatory ``Exp2Syn``; dendritic placement lets synaptic depolarization + engage the PIC the way the reference simulator does. +``live_driving_force`` + ``False``: current = g · (E − V_rest), the validated ex03 approximation. + ``True``: current = g · (E − V_site(t−dt)) with the previous-step voltage + of the synapse-bearing compartment (one-step lag breaks the algebraic + loop). Required for apples-to-apples NEURON comparison. +``syn_tau_rise_ms`` + ``None``: single-exponential decay conductance (legacy IIR). + A value (e.g. 0.5): peak-normalized double-exponential (Exp2Syn-like) with + rise ``syn_tau_rise_ms`` and decay ``tau_syn_ms``. + +Defaults keep the original fast-surrogate behavior (soma, fixed driving +force, single-exponential), bit-compatible with the earlier harness. + +NOTE: this harness implements the drive→pool pathway only. Afferent input, +spindle/GTO, muscle force, and EMG observables are NOT yet wired into +``PoolTwin`` (see the roadmap's implemented-vs-planned table); for the full +closed loop use ``myogen.simulator.jaxley.closed_loop`` / ex11. + +Voltage frame: NERLab/1952-HH — V_rest ≈ 0 mV, spikes peak ≈ +90 mV, +detection threshold +50 mV. + +Typical use:: + + twin = PoolTwin(TwinSpec(n_mn=20, seed=1)) + t, cmd = protocols.triangular_ramp(twin.spec.dt_ms, peak=40.0) + drives = {"dd": jnp.asarray(cmd), "inh": jnp.zeros_like(jnp.asarray(cmd))} + out = twin.simulate(twin.default_theta(), drives, jax.random.PRNGKey(0), + spike_mode="hard") + # out["v"], out["v_dend"], out["spikes"]: (n_steps, n_mn) + + val, grads = twin.loss_and_grad(loss_fn, theta, drives, key) +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax +import jax.numpy as jnp +import numpy as np + +import jaxley as jx +from jaxley.integrate import build_init_and_step_fn + +from myogen.simulator.jaxley.populations import AlphaMN__Pool +from myogen.simulator.jaxley.jax_models import ( + poisson_init, + poisson_step, + spike_detect, +) + +__all__ = ["TwinSpec", "PoolTwin", "fuglevand_thresholds"] + +# NERLab (1952-HH) frame constants — see docs/session_notes/2026-06-02. +V_REST_MV = 0.0 +E_EXC_MV = 70.0 +E_INH_MV = -15.0 +SPIKE_THRESHOLD_MV = 50.0 + +# spike_mode (MN detection) → stochastic-generator mode (jax_models convention). +_GEN_MODE = {"hard": "hard", "surrogate": "pathwise", "rate": "rate"} + + +def fuglevand_thresholds(n: int, recruitment_range: float = 30.0) -> np.ndarray: + """Exponential recruitment-threshold distribution (Fuglevand et al. 1993).""" + i = np.arange(1, n + 1) + return np.exp(i * np.log(recruitment_range) / n) / recruitment_range + + +@dataclass(frozen=True) +class TwinSpec: + """Static specification of the pool twin (nothing here is differentiated). + + ``seed`` controls the fixed random connectivity (and is the default base + for generator seeding in the scripts). Two twins with the same spec are + identical; change ``seed`` to resample connectivity. + """ + + n_mn: int = 20 + n_dd: int = 40 # descending-drive sources + n_inh: int = 40 # inhibitory sources + dd_batch: int = 5 # renewal order: 1 = Poisson, >1 = gamma/Erlang (more regular) + dt_ms: float = 0.1 + connect_p: float = 0.5 # source→MN connection probability (ex03 value) + tau_syn_ms: float = 5.0 # decay time constant + w_exc_uS: float = 0.009 # ex03's validated IIR weight (≡ NEURON Exp2Syn 0.15) + w_inh_uS: float = 0.009 + recruitment_range: float = 30.0 + seed: int = 0 + # Synapse model (see module docstring). Defaults = legacy fast surrogate. + syn_site: str = "soma" # "soma" | "dendrite" + live_driving_force: bool = False + syn_tau_rise_ms: float | None = None # None → single-exp; value → Exp2Syn + # Reversal potentials (NERLab 1952-HH frame). e_inh default −15 mV is the + # FRAME-CORRECTED analog of a modern −80 mV (driving force −15 mV at + # V_rest = 0). NEURON's NERLab synapses literally use −75 mV; matching + # that number is a deliberate choice (`e_inh_mV=-75.0`) and any paired + # comparison must state which convention it used. + e_exc_mV: float = E_EXC_MV + e_inh_mV: float = E_INH_MV + + +class PoolTwin: + """A NERLab MN pool with differentiable drive/inhibition/PIC parameters.""" + + def __init__(self, spec: TwinSpec): + if spec.syn_site not in ("soma", "dendrite"): + raise ValueError(f"syn_site must be 'soma' or 'dendrite', got {spec.syn_site!r}") + if spec.syn_tau_rise_ms is not None and not ( + 0.0 < spec.syn_tau_rise_ms < spec.tau_syn_ms + ): + raise ValueError("syn_tau_rise_ms must satisfy 0 < rise < tau_syn_ms") + self.spec = spec + rng = np.random.default_rng(spec.seed) + + thresholds = fuglevand_thresholds(spec.n_mn, spec.recruitment_range) + pool = AlphaMN__Pool( + recruitment_thresholds__array=thresholds, + model="NERLab", + mode="active", + ) + + # Pool wrappers → clean jx.Network. Recordings: soma then dendrite + # voltage for every MN (dendrite needed for live driving force and for + # PIC diagnostics). Stimuli: one slot per MN at the synapse site. + cells = [] + for cw in pool: + cell = cw.cell + cell.delete_recordings() + cell.delete_stimuli() + cells.append(cell) + net = jx.Network(cells) + + syn_branch = 0 if spec.syn_site == "soma" else 1 + placeholder = jnp.zeros(2) # registers the slot; content unused + for i in range(spec.n_mn): + net.cell(i).branch(0).loc(0.5).record("v") # soma + for i in range(spec.n_mn): + net.cell(i).branch(1).loc(0.5).record("v") # dendrite + for i in range(spec.n_mn): + net.cell(i).branch(syn_branch).loc(0.5).stimulate(placeholder) + + # jx.Network() can silently revert v to the modern −70 mV frame; + # re-assert the NERLab rest before initializing states. + net.set("v", V_REST_MV) + net.init_states() + + net.to_jax() + init_fn, step_fn = build_init_and_step_fn(net) + states0, all_params = init_fn(net.get_parameters()) + + # The per-compartment parameter arrays carry NaN / sentinel entries for + # compartments that lack the channel (soma vs dendrite vs axon); scale + # only the valid, positive conductance entries or gradients go NaN. + # (This deliberately bypasses jaxley's make_trainable: gradients flow + # through the full all_params dict handed to step_fn.) + self._pic_masks = {} + for name in ("caL_gcaLbar", "napp_gnapbar"): + if name not in all_params: + raise KeyError( + f"{name!r} not in jaxley all_params " + f"(keys: {sorted(all_params)}) — jaxley layout changed?" + ) + arr = np.asarray(all_params[name]) + self._pic_masks[name] = jnp.asarray(np.isfinite(arr) & (arr > 0)) + + self.net = net + self.pool = pool + self.thresholds = thresholds + self._step_fn = step_fn + self._states0 = states0 + self._all_params = all_params + self._external_inds = net.external_inds.copy() + rec = jnp.array(net.recordings.rec_index.to_numpy(), dtype=jnp.int32) + self._rec_soma = rec[: spec.n_mn] + self._rec_dend = rec[spec.n_mn :] + + # Fixed random connectivity, (n_mn, n_sources) so `mat @ spikes` → (n_mn,). + self._exc_mat = jnp.asarray( + (rng.random((spec.n_mn, spec.n_dd)) < spec.connect_p).astype(np.float32) + ) + self._inh_mat = jnp.asarray( + (rng.random((spec.n_mn, spec.n_inh)) < spec.connect_p).astype(np.float32) + ) + + # Conductance-kernel constants. + self._a_decay = float(np.exp(-spec.dt_ms / spec.tau_syn_ms)) + if spec.syn_tau_rise_ms is not None: + t1, t2 = spec.syn_tau_rise_ms, spec.tau_syn_ms + self._a_rise = float(np.exp(-spec.dt_ms / t1)) + tp = (t1 * t2 / (t2 - t1)) * np.log(t2 / t1) + self._exp2_norm = float(1.0 / (np.exp(-tp / t2) - np.exp(-tp / t1))) + else: + self._a_rise = None + self._exp2_norm = None + + self._sim = jax.jit(self._run, static_argnames=("spike_mode",)) + self._sim_events = jax.jit(self._run_events, static_argnames=("spike_mode",)) + self._sim_current = jax.jit(self._run_current, static_argnames=("spike_mode",)) + + # ------------------------------------------------------------------ # + + def default_theta(self) -> dict: + return { + "gCaL_scale": jnp.float32(1.0), + "gNaP_scale": jnp.float32(1.0), + "drive_gain": jnp.float32(1.0), + "w_exc": jnp.float32(self.spec.w_exc_uS), + "w_inh": jnp.float32(self.spec.w_inh_uS), + } + + def _scaled_jaxley_params(self, theta: dict) -> dict: + p = dict(self._all_params) + for name, scale in ( + ("caL_gcaLbar", theta["gCaL_scale"]), + ("napp_gnapbar", theta["gNaP_scale"]), + ): + m = self._pic_masks[name] + valid = jnp.where(m, p[name], 0.0) + p[name] = jnp.where(m, valid * scale, p[name]) + return p + + def _syn_init(self) -> dict: + n = self.spec.n_mn + z = lambda: jnp.zeros(n) # dtype follows active precision (x64-safe) + if self.spec.syn_tau_rise_ms is None: + return {"g_exc": z(), "g_inh": z()} + return {"x1_exc": z(), "x2_exc": z(), "x1_inh": z(), "x2_inh": z()} + + def _syn_step(self, syn: dict, inc_exc, inc_inh) -> tuple[dict, jnp.ndarray, jnp.ndarray]: + """Advance conductance states; return (new_state, g_exc, g_inh) in µS.""" + if self.spec.syn_tau_rise_ms is None: + g_exc = self._a_decay * syn["g_exc"] + inc_exc + g_inh = self._a_decay * syn["g_inh"] + inc_inh + return {"g_exc": g_exc, "g_inh": g_inh}, g_exc, g_inh + N = self._exp2_norm + x1e = self._a_rise * syn["x1_exc"] + inc_exc * N + x2e = self._a_decay * syn["x2_exc"] + inc_exc * N + x1i = self._a_rise * syn["x1_inh"] + inc_inh * N + x2i = self._a_decay * syn["x2_inh"] + inc_inh * N + return ( + {"x1_exc": x1e, "x2_exc": x2e, "x1_inh": x1i, "x2_inh": x2i}, + x2e - x1e, + x2i - x1i, + ) + + def _driving_forces(self, carry): + spec = self.spec + if spec.live_driving_force: + v_site = (carry["prev_v_soma"] if spec.syn_site == "soma" + else carry["prev_v_dend"]) + return spec.e_exc_mV - v_site, spec.e_inh_mV - v_site + return (spec.e_exc_mV - V_REST_MV), (spec.e_inh_mV - V_REST_MV) + + def _membrane_step(self, carry, i_stim, jax_params, spike_mode): + """One jaxley step + spike detection; returns (carry updates, outputs).""" + st = self._step_fn( + carry["st"], jax_params, {"i": i_stim}, + self._external_inds, delta_t=self.spec.dt_ms, + ) + v_soma = st["v"][self._rec_soma] + v_dend = st["v"][self._rec_dend] + sp = spike_detect(v_soma, carry["prev_v_soma"], SPIKE_THRESHOLD_MV, + spike_mode) + return st, v_soma, v_dend, sp + + def _run(self, theta: dict, drives: dict, key, spike_mode: str = "hard"): + """Pure simulation function (jitted as ``self.simulate``). + + drives : {"dd": (n_steps,), "inh": (n_steps,)} command rates in Hz + (per source neuron). ``theta["drive_gain"]`` multiplies "dd". + """ + spec = self.spec + gen_mode = _GEN_MODE[spike_mode] + jax_params = self._scaled_jaxley_params(theta) + + k_dd, k_inh = jax.random.split(key) + carry0 = { + "st": self._states0, + "syn": self._syn_init(), + "prev_v_soma": jnp.full(spec.n_mn, V_REST_MV), + "prev_v_dend": jnp.full(spec.n_mn, V_REST_MV), + "dd_st": poisson_init(spec.n_dd, spec.dd_batch, key=k_dd), + "inh_st": poisson_init(spec.n_inh, spec.dd_batch, key=k_inh), + } + + def body(carry, xs): + dd_rate_t, inh_rate_t = xs + dd_st, dd_sp = poisson_step( + carry["dd_st"], theta["drive_gain"] * dd_rate_t, + spec.dt_ms, spec.dd_batch, gen_mode, + ) + inh_st, inh_sp = poisson_step( + carry["inh_st"], inh_rate_t, spec.dt_ms, spec.dd_batch, gen_mode, + ) + inc_exc = (self._exc_mat @ dd_sp.astype(jnp.float32)) * theta["w_exc"] + inc_inh = (self._inh_mat @ inh_sp.astype(jnp.float32)) * theta["w_inh"] + syn, g_exc, g_inh = self._syn_step(carry["syn"], inc_exc, inc_inh) + df_exc, df_inh = self._driving_forces(carry) + i_stim = g_exc * df_exc + g_inh * df_inh # nA = µS × mV + st, v_soma, v_dend, sp = self._membrane_step( + carry, i_stim, jax_params, spike_mode) + new_carry = { + "st": st, "syn": syn, + "prev_v_soma": v_soma, "prev_v_dend": v_dend, + "dd_st": dd_st, "inh_st": inh_st, + } + return new_carry, { + "v": v_soma, "v_dend": v_dend, "spikes": sp.astype(jnp.float32), + } + + xs = (drives["dd"], drives["inh"]) + _, out = jax.lax.scan(body, carry0, xs) + return out + + def _run_events(self, theta: dict, events: dict, spike_mode: str = "hard"): + """Deterministic forward from exact presynaptic event matrices. + + ``events`` = {"exc": (n_steps, n_exc_src), "inh": (n_steps, n_inh_src), + "exc_conn": (n_mn, n_exc_src), "inh_conn": (n_mn, n_inh_src)}. + No stochastic generators are involved — the same matrices can be fed + verbatim to NEURON for matched-input validation. + """ + jax_params = self._scaled_jaxley_params(theta) + exc_conn, inh_conn = events["exc_conn"], events["inh_conn"] + carry0 = { + "st": self._states0, + "syn": self._syn_init(), + "prev_v_soma": jnp.full(self.spec.n_mn, V_REST_MV), + "prev_v_dend": jnp.full(self.spec.n_mn, V_REST_MV), + } + + def body(carry, xs): + exc_t, inh_t = xs + inc_exc = (exc_conn @ exc_t) * theta["w_exc"] + inc_inh = (inh_conn @ inh_t) * theta["w_inh"] + syn, g_exc, g_inh = self._syn_step(carry["syn"], inc_exc, inc_inh) + df_exc, df_inh = self._driving_forces(carry) + i_stim = g_exc * df_exc + g_inh * df_inh + st, v_soma, v_dend, sp = self._membrane_step( + carry, i_stim, jax_params, spike_mode) + new_carry = { + "st": st, "syn": syn, + "prev_v_soma": v_soma, "prev_v_dend": v_dend, + } + return new_carry, { + "v": v_soma, "v_dend": v_dend, "spikes": sp.astype(jnp.float32), + } + + _, out = jax.lax.scan(body, carry0, (events["exc"], events["inh"])) + return out + + def _run_events_carry(self, theta: dict, events: dict, carry0, + spike_mode: str = "hard"): + """Like ``_run_events`` but starts from an explicit carry and returns + ``(outputs, final_carry)`` — enables windowed losses that preserve + forward state across window boundaries (never restart from rest; + gradients may be stopped at boundaries by the caller).""" + jax_params = self._scaled_jaxley_params(theta) + exc_conn, inh_conn = events["exc_conn"], events["inh_conn"] + if carry0 is None: + carry0 = { + "st": self._states0, + "syn": self._syn_init(), + "prev_v_soma": jnp.full(self.spec.n_mn, V_REST_MV), + "prev_v_dend": jnp.full(self.spec.n_mn, V_REST_MV), + } + + def body(carry, xs): + exc_t, inh_t = xs + inc_exc = (exc_conn @ exc_t) * theta["w_exc"] + inc_inh = (inh_conn @ inh_t) * theta["w_inh"] + syn, g_exc, g_inh = self._syn_step(carry["syn"], inc_exc, inc_inh) + df_exc, df_inh = self._driving_forces(carry) + i_stim = g_exc * df_exc + g_inh * df_inh + st, v_soma, v_dend, sp = self._membrane_step( + carry, i_stim, jax_params, spike_mode) + new_carry = { + "st": st, "syn": syn, + "prev_v_soma": v_soma, "prev_v_dend": v_dend, + } + return new_carry, { + "v": v_soma, "v_dend": v_dend, "spikes": sp.astype(v_soma.dtype), + } + + carry_f, out = jax.lax.scan(body, carry0, + (events["exc"], events["inh"])) + return out, carry_f + + def _run_current(self, theta: dict, i_ext, spike_mode: str = "hard"): + """Deterministic forward from an exact injected-current waveform. + + ``i_ext``: (n_steps, n_mn) current in nA, injected at the registered + stimulus site (``spec.syn_site`` compartment). No synapse model is + involved — this is the intrinsic (current-injection) validation path. + """ + jax_params = self._scaled_jaxley_params(theta) + carry0 = { + "st": self._states0, + "prev_v_soma": jnp.full(self.spec.n_mn, V_REST_MV), + "prev_v_dend": jnp.full(self.spec.n_mn, V_REST_MV), + } + + def body(carry, i_t): + st, v_soma, v_dend, sp = self._membrane_step( + carry, i_t, jax_params, spike_mode) + return ( + {"st": st, "prev_v_soma": v_soma, "prev_v_dend": v_dend}, + {"v": v_soma, "v_dend": v_dend, "spikes": sp.astype(jnp.float32)}, + ) + + _, out = jax.lax.scan(body, carry0, i_ext) + return out + + # ------------------------------------------------------------------ # + + def simulate_current(self, theta, i_ext_nA, spike_mode: str = "hard"): + """Forward run injecting ``i_ext_nA`` (n_steps, n_mn) at the stimulus + site — for deterministic single-cell / intrinsic validation.""" + i_ext = jnp.asarray(i_ext_nA, dtype=jnp.float32) + if i_ext.ndim != 2 or i_ext.shape[1] != self.spec.n_mn: + raise ValueError(f"i_ext must be (n_steps, {self.spec.n_mn})") + return self._sim_current(theta, i_ext, spike_mode=spike_mode) + + def simulate_events( + self, + theta, + exc_events, + inh_events=None, + exc_connectivity=None, + inh_connectivity=None, + spike_mode: str = "hard", + ): + """Forward run from EXACT presynaptic events — the matched-input path. + + Parameters + ---------- + exc_events : (n_steps, n_exc_src) + 0/1 (or float) excitatory event matrix. + inh_events : (n_steps, n_inh_src), optional + Defaults to no inhibitory input. + exc_connectivity, inh_connectivity : (n_mn, n_src), optional + Default to the twin's own fixed random connectivity; pass explicit + matrices to reproduce a protocol artifact exactly. + + Returns ``{"v", "v_dend", "spikes"}``, each ``(n_steps, n_mn)``. + """ + spec = self.spec + exc_events = jnp.asarray(exc_events, dtype=jnp.float32) + n_steps = exc_events.shape[0] + if inh_events is None: + inh_events = jnp.zeros((n_steps, 0), dtype=jnp.float32) + else: + inh_events = jnp.asarray(inh_events, dtype=jnp.float32) + if exc_connectivity is None: + if exc_events.shape[1] != spec.n_dd: + raise ValueError( + f"exc_events has {exc_events.shape[1]} sources but the twin " + f"has n_dd={spec.n_dd}; pass exc_connectivity explicitly") + exc_connectivity = self._exc_mat + if inh_connectivity is None: + if inh_events.shape[1] == spec.n_inh: + inh_connectivity = self._inh_mat + else: + inh_connectivity = jnp.zeros( + (spec.n_mn, inh_events.shape[1]), dtype=jnp.float32) + events = { + "exc": exc_events, + "inh": inh_events, + "exc_conn": jnp.asarray(exc_connectivity, dtype=jnp.float32), + "inh_conn": jnp.asarray(inh_connectivity, dtype=jnp.float32), + } + if events["exc_conn"].shape != (spec.n_mn, exc_events.shape[1]): + raise ValueError("exc_connectivity shape mismatch") + if events["inh_conn"].shape != (spec.n_mn, inh_events.shape[1]): + raise ValueError("inh_connectivity shape mismatch") + return self._sim_events(theta, events, spike_mode=spike_mode) + + def simulate(self, theta, drives, key, spike_mode: str = "hard"): + """Jitted forward run → ``{"v", "v_dend", "spikes"}``, each ``(n_steps, n_mn)``.""" + return self._sim(theta, drives, key, spike_mode=spike_mode) + + def loss_and_grad( + self, loss_fn, theta, drives, key, spike_mode: str = "surrogate" + ): + """``value_and_grad`` of ``loss_fn(outputs)`` w.r.t. ``theta``.""" + + def f(th): + return loss_fn(self._sim(th, drives, key, spike_mode=spike_mode)) + + return jax.value_and_grad(f)(theta) diff --git a/experiments/twin/protocols.py b/experiments/twin/protocols.py new file mode 100644 index 00000000..57958a04 --- /dev/null +++ b/experiments/twin/protocols.py @@ -0,0 +1,94 @@ +"""Stimulation protocols for digital-twin experiments. + +Each builder returns a scalar command time-series ``(n_steps,)`` (NumPy, +float32) — typically the per-neuron descending-drive rate in Hz — plus the time +vector. Pool-specific assembly into ``run_jax`` input dicts happens in +``poolbuild.py``. +""" + +from __future__ import annotations + +import numpy as np + +__all__ = ["time_vector", "triangular_ramp", "trapezoid", "step", "hold"] + + +def time_vector(t_total_s: float, dt_ms: float) -> np.ndarray: + n = int(round(t_total_s * 1000.0 / dt_ms)) + return np.arange(n) * dt_ms * 1e-3 + + +def triangular_ramp( + dt_ms: float, + t_rise_s: float = 10.0, + t_fall_s: float | None = None, + t_pre_s: float = 2.0, + t_post_s: float = 2.0, + baseline: float = 0.0, + peak: float = 40.0, +) -> tuple[np.ndarray, np.ndarray]: + """Symmetric (or asymmetric) triangular ramp — the ΔF/PIC protocol. + + Returns ``(t, command)`` with command in the same units as ``peak``. + """ + t_fall_s = t_rise_s if t_fall_s is None else t_fall_s + seg = lambda dur: int(round(dur * 1000.0 / dt_ms)) + pre = np.full(seg(t_pre_s), baseline) + up = np.linspace(baseline, peak, seg(t_rise_s), endpoint=False) + down = np.linspace(peak, baseline, seg(t_fall_s), endpoint=False) + post = np.full(seg(t_post_s), baseline) + cmd = np.concatenate([pre, up, down, post]).astype(np.float32) + t = np.arange(len(cmd)) * dt_ms * 1e-3 + return t, cmd + + +def trapezoid( + dt_ms: float, + t_rise_s: float = 5.0, + t_hold_s: float = 10.0, + t_fall_s: float = 5.0, + t_pre_s: float = 2.0, + t_post_s: float = 2.0, + baseline: float = 0.0, + peak: float = 40.0, +) -> tuple[np.ndarray, np.ndarray]: + """Trapezoidal command — recruitment/derecruitment + steady-state rates.""" + seg = lambda dur: int(round(dur * 1000.0 / dt_ms)) + cmd = np.concatenate( + [ + np.full(seg(t_pre_s), baseline), + np.linspace(baseline, peak, seg(t_rise_s), endpoint=False), + np.full(seg(t_hold_s), peak), + np.linspace(peak, baseline, seg(t_fall_s), endpoint=False), + np.full(seg(t_post_s), baseline), + ] + ).astype(np.float32) + t = np.arange(len(cmd)) * dt_ms * 1e-3 + return t, cmd + + +def step( + dt_ms: float, + t_pre_s: float = 2.0, + t_on_s: float = 8.0, + t_post_s: float = 2.0, + baseline: float = 0.0, + level: float = 30.0, +) -> tuple[np.ndarray, np.ndarray]: + """Step command — F–I style responses and self-sustained-firing tests.""" + seg = lambda dur: int(round(dur * 1000.0 / dt_ms)) + cmd = np.concatenate( + [ + np.full(seg(t_pre_s), baseline), + np.full(seg(t_on_s), level), + np.full(seg(t_post_s), baseline), + ] + ).astype(np.float32) + t = np.arange(len(cmd)) * dt_ms * 1e-3 + return t, cmd + + +def hold(dt_ms: float, t_total_s: float, level: float) -> tuple[np.ndarray, np.ndarray]: + """Constant command.""" + t = time_vector(t_total_s, dt_ms) + return t, np.full(len(t), level, dtype=np.float32) diff --git a/experiments/twin/r2a_profiles.py b/experiments/twin/r2a_profiles.py new file mode 100644 index 00000000..83f44447 --- /dev/null +++ b/experiments/twin/r2a_profiles.py @@ -0,0 +1,261 @@ +"""Phase-3 ladder step 2 — CaL + NaP profiles (before any R2 optimizer). + +Revised roadmap §7 / review work-package 3: expand identifiability ONE +parameter at a time. This adds gNaP to the Deliverable-B problem: + +* Unknowns under profiling: gCaL_scale, gNaP_scale. +* Nuisance reoptimized everywhere: w_exc (effective excitatory amplitude). +* Fixed and known: inhibition (absent), connectivity, drive shape AND the + frozen event realizations (conditional analysis — same scope as B). +* Observable: binned rates (50 ms / 300 ms); hard forward only (gradients + remain unqualified). +* EPS per replicate from dt-halving at the truth (same rule as B/C). + +PREDECLARED RESOLUTION (per the frozen acceptance spec +``downstream_resolution_declaration``): a parameter "advances" to the R2 +optimizer stage only if its sampled compatibility interval lies within +±20 % of the truth in every replicate. The question is NOT whether the +minimum sits at the exact truth (it does by construction — same-model +noiseless data with truth on-axis) but whether the ≤EPS level set is +narrower than the declared resolution. If CaL and NaP are inseparable, the +honest fallback is a combined "PIC strength" parameter — reported here via +the 2-D (CaL × NaP) joint compatibility set at the profiled w_exc. + +Grid: 3-D dense (g × nap × w), truth inserted on every axis, profile of +each parameter = min over the other two (shared candidate grid — the +nesting lesson from Deliverable C applies: all minimizations use the same +evaluated set). + +Usage:: + + python r2a_profiles.py # full (~25 min CPU) + python r2a_profiles.py --dense-g 5 --dense-nap 4 --dense-w 5 \\ + --replicates 1 --dur-scale 0.5 # smoke +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import jax.numpy as jnp + +from sensitivity_checks import _manifest, _smoothed_rates_jax +from deliverable_b import ( + EPS_FLOOR, RESULTS, EVIDENCE, axis_with_truth, make_commands, + build_twin_and_events, numerical_epsilon_spikes, binned_rates_jnp, + _sha16, +) + +OKABE = {"cal": "#0072B2", "nap": "#E69F00", "truth": "#000000"} +GRIDC = "#e3e3e3" +RESOLUTION_REL = 0.20 # predeclared (acceptance spec v1) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-mn", type=int, default=8) + ap.add_argument("--true", type=float, nargs=3, default=[1.6, 0.7, 0.0105], + metavar=("GCAL", "GNAP", "WEXC")) + ap.add_argument("--dense-g", type=int, default=13) + ap.add_argument("--dense-nap", type=int, default=9) + ap.add_argument("--dense-w", type=int, default=11) + ap.add_argument("--replicates", type=int, default=2) + ap.add_argument("--dur-scale", type=float, default=1.0) + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + tag = ("full" if (args.dense_g >= 11 and args.replicates >= 2 + and args.dur_scale >= 1.0 and args.n_mn >= 8) + else "smoke") + stem = f"r2a_profiles_{tag}" + t_wall0 = time.time() + dt_ms = 0.1 + theta_true = np.array(args.true) + fit_protos = ["fit_ramp", "fit_trap"] + commands = make_commands(dt_ms, args.dur_scale) + + g_axis = axis_with_truth(0.25, 4.0, args.dense_g, theta_true[0]) + nap_axis = axis_with_truth(0.25, 4.0, args.dense_nap, theta_true[1]) + w_axis = axis_with_truth(0.0045, 0.018, args.dense_w, theta_true[2]) + gi = int(np.argmin(np.abs(g_axis - theta_true[0]))) + ni = int(np.argmin(np.abs(nap_axis - theta_true[1]))) + + def observable_fn_3(twin, events): + def fn(gcal, gnap, w, proto): + th = twin.default_theta() + th["gCaL_scale"] = jnp.float32(gcal) + th["gNaP_scale"] = jnp.float32(gnap) + th["w_exc"] = jnp.float32(w) + out = twin.simulate_events(th, events[proto], spike_mode="hard") + return binned_rates_jnp(out["spikes"], dt_ms) + return fn + + reps = [] + for r in range(args.replicates): + conn_seed, ev_off = r, 10 * r + twin, events = build_twin_and_events( + args.n_mn, conn_seed, ev_off, commands, dt_ms) + fn = observable_fn_3(twin, events) + obs = {p: fn(*theta_true, p) for p in commands} + + def loss(gcal, gnap, w): + tot = 0.0 + for p in fit_protos: + tot += float(jnp.mean((fn(gcal, gnap, w, p) - obs[p]) ** 2)) + return tot / len(fit_protos) + + assert loss(*theta_true) == 0.0, "truth loss must be exactly 0" + + t0 = time.time() + L = np.zeros((len(g_axis), len(nap_axis), len(w_axis))) + for i, g in enumerate(g_axis): + for j, nap in enumerate(nap_axis): + for k, w in enumerate(w_axis): + L[i, j, k] = loss(g, nap, w) + print(f"[rep{r}] g-slice {i+1}/{len(g_axis)} " + f"({time.time()-t0:.0f}s)", flush=True) + + # EPS per replicate (dt-halving at truth, all units). Computed inline + # because the shared helper applies only (gCaL, w) — here the truth + # also sets gNaP. + terms = [] + from poolbuild import PoolTwin, TwinSpec + twin_h = PoolTwin(TwinSpec(n_mn=args.n_mn, n_dd=100, n_inh=5, + dt_ms=0.05, seed=conn_seed)) + for p in fit_protos: + ev = np.asarray(events[p]) + fine_ev = np.zeros((ev.shape[0] * 2, ev.shape[1]), + dtype=np.float32) + fine_ev[::2] = ev + th = twin_h.default_theta() + th["gCaL_scale"] = jnp.float32(theta_true[0]) + th["gNaP_scale"] = jnp.float32(theta_true[1]) + th["w_exc"] = jnp.float32(theta_true[2]) + sp_f = twin_h.simulate_events( + th, fine_ev, exc_connectivity=twin._exc_mat, + spike_mode="hard")["spikes"] + rf = np.asarray(binned_rates_jnp(sp_f, 0.05)) + rc = np.asarray(obs[p]) + m = min(len(rf), len(rc)) + terms.append(float(np.mean((rf[:m] - rc[:m]) ** 2))) + EPS = max(float(np.mean(terms)), EPS_FLOOR) + + # profiles: min over the other two axes (one shared evaluated grid) + prof_g = L.min(axis=(1, 2)) + prof_nap = L.min(axis=(0, 2)) + assert prof_g[gi] == 0.0 and prof_nap[ni] == 0.0 + + def interval(curve, axis): + within = curve <= EPS + return ([float(axis[within].min()), float(axis[within].max())] + if within.any() else None) + + iv_g = interval(prof_g, g_axis) + iv_nap = interval(prof_nap, nap_axis) + + def advances(iv, truth): + if iv is None: + return False + return (iv[0] >= truth * (1 - RESOLUTION_REL) + and iv[1] <= truth * (1 + RESOLUTION_REL)) + + # 2-D joint compatibility set at profiled w (PIC-strength fallback) + L2 = L.min(axis=2) + joint = [[float(g_axis[i]), float(nap_axis[j])] + for i in range(len(g_axis)) for j in range(len(nap_axis)) + if L2[i, j] <= EPS] + + reps.append({ + "conn_seed": conn_seed, "event_offset": ev_off, "eps": EPS, + "profile_gcal": prof_g.tolist(), + "profile_gnap": prof_nap.tolist(), + "gcal_interval": iv_g, "gnap_interval": iv_nap, + "gcal_advances_at_pm20pct": advances(iv_g, theta_true[0]), + "gnap_advances_at_pm20pct": advances(iv_nap, theta_true[1]), + "joint_compat_set_gcal_gnap": joint, + "event_hashes": {p: _sha16(np.asarray(events[p])) + for p in fit_protos}, + }) + print(f"[rep{r}] EPS {EPS:.4f} | gCaL interval {iv_g} " + f"(±20% pass: {reps[-1]['gcal_advances_at_pm20pct']}) | " + f"gNaP interval {iv_nap} " + f"(±20% pass: {reps[-1]['gnap_advances_at_pm20pct']}) | " + f"joint set size {len(joint)}", flush=True) + + # ---------------------------------------------------------------- figure + clip = EPS_FLOOR / 10 + fig, axes = plt.subplots(1, 2, figsize=(12, 4.8)) + for ax, pname, axis, key, ti in ( + (axes[0], "gCaL_scale", g_axis, "profile_gcal", theta_true[0]), + (axes[1], "gNaP_scale", nap_axis, "profile_gnap", theta_true[1])): + for k, rep in enumerate(reps): + ax.plot(axis, np.maximum(rep[key], clip), + color=OKABE["cal" if pname == "gCaL_scale" else "nap"], + lw=1.4, alpha=1.0 if k == 0 else 0.5, + label=f"rep{k}" if k < 2 else None) + ax.axhline(rep["eps"], color="#111111", lw=0.8, ls="--", + alpha=1.0 if k == 0 else 0.5) + ax.axvline(ti, color=OKABE["truth"], lw=1.0) + ax.axvspan(ti * (1 - RESOLUTION_REL), ti * (1 + RESOLUTION_REL), + color="#000000", alpha=0.06, + label="±20 % declared resolution") + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xlabel(pname) + ax.set_ylabel("profile-min fit loss (Hz²)") + ax.set_title(f"Profile over {pname} (min over other two params)", + fontsize=10) + ax.legend(frameon=False, fontsize=8) + ax.grid(True, color=GRIDC, lw=0.6) + ax.set_axisbelow(True) + fig.tight_layout() + fig.savefig(RESULTS / f"{stem}.png", dpi=150) + + here = Path(__file__) + report = { + "manifest": {**_manifest(), + "r2a_sha256_16": hashlib.sha256( + here.read_bytes()).hexdigest()[:16], + "args": vars(args), "run_tag": tag}, + "declared": { + "unknowns_profiled": ["gCaL_scale", "gNaP_scale"], + "nuisance_reoptimized": ["w_exc"], + "fixed_known": ["inhibition absent", "connectivity", + "frozen event realizations"], + "observable": "binned (50 ms / 300 ms)", + "resolution_rel_predeclared": RESOLUTION_REL, + "theta_true": theta_true.tolist(), + "interval_terminology": "sampled compatibility intervals " + "(deterministic level sets)", + }, + "axes": {"gcal": g_axis.tolist(), "gnap": nap_axis.tolist(), + "w_exc": w_axis.tolist()}, + "replicates": reps, + "wall_s": round(time.time() - t_wall0, 1), + } + payload = json.dumps(report, indent=1, default=float) + (RESULTS / f"{stem}.json").write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"{stem}_{stamp}.json").write_text(payload) + + print(f"\n=== {stem} summary ===") + for k, rep in enumerate(reps): + print(f"rep{k}: gCaL {rep['gcal_interval']} " + f"(advances: {rep['gcal_advances_at_pm20pct']}) | " + f"gNaP {rep['gnap_interval']} " + f"(advances: {rep['gnap_advances_at_pm20pct']})") + print(f"saved → {RESULTS / (stem + '.json')} (+ evidence copy, figure)") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/sensitivity_checks.py b/experiments/twin/sensitivity_checks.py new file mode 100644 index 00000000..f16a85a2 --- /dev/null +++ b/experiments/twin/sensitivity_checks.py @@ -0,0 +1,588 @@ +"""Deliverable A — scoped forward and gradient diagnostics. + +Implements the gradient-qualification tests of the revised roadmap +(``docs/DIGITAL_TWIN_ROADMAP_REVISED.md`` §6.2B / §12A). Every test prints a +verdict and all raw numbers land in ``results/sensitivity_checks.json`` — +failures are preserved, not hidden. + +T1 Parameter plumbing — each theta knob changes exactly the intended jaxley + leaves (gCaL→caL_gcaLbar only, gNaP→napp_gnapbar only; input knobs touch + no jaxley leaf), and each knob measurably changes the forward output. +T2 Continuous sensitivities — ONE cell, deterministic SUBTHRESHOLD current + (smooth voltage trajectory, no events): autodiff dL/dθ vs centered finite + differences across perturbation sizes; reports the FD convergence region. + This qualifies the derivative of the continuous dynamics, separate from + any spike/surrogate issue. +T3 Event-parameter plumbing — with a FIXED PRNG key the pathwise derivative + w.r.t. drive_gain is a fixed-realization derivative, not the derivative + of the expected spike count. Reports: surrogate/pathwise gradient sign vs + the finite-perturbation slope of the hard loss (same key), and the slope + across independent keys (expected-value sense). Divergence between these + is a property of the estimator and is recorded, not patched. +T4 Horizon diagnostics — |grad| of a smoothed-rate objective vs simulated + duration (suprathreshold, spiking): quantifies the reported long-horizon + explosion instead of asserting it. +T5 Hard-forward descent — surrogate sign-descent directions on the + Deliverable-B parameter pair (gCaL_scale, w_exc) with FROZEN event input: + fraction of proposed steps that actually reduce the scientific hard loss. + +Usage:: + + python sensitivity_checks.py # all tests, small CPU sizes + python sensitivity_checks.py --tests t1 t2 +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import jax +import jax.numpy as jnp + +from poolbuild import PoolTwin, TwinSpec +from validation_protocol import _renewal_events + +RESULTS = Path(__file__).parent / "results" +# Gate evidence is version-controllable (results/ is gitignored by design). +EVIDENCE = Path(__file__).parent / "evidence" + + +def _safe_norm(g) -> float: + """L2 norm in float64 — float32 squaring overflows at ~1e19 components, + which previously misreported finite ~1e28 gradients as Inf.""" + return float(np.linalg.norm(np.asarray(g, dtype=np.float64))) + + +def _manifest() -> dict: + """Reproducibility manifest attached to every report.""" + import subprocess + + import hashlib + + repo = Path(__file__).resolve().parents[2] + try: + rev = subprocess.run(["git", "rev-parse", "--short", "HEAD"], + cwd=repo, capture_output=True, text=True + ).stdout.strip() + dirty = bool(subprocess.run(["git", "status", "--porcelain"], + cwd=repo, capture_output=True, text=True + ).stdout.strip()) + except Exception: + rev, dirty = "unknown", None + # The diagnostic sources are untracked, so the git revision alone cannot + # reconstruct this run — hash the actual files used. + here = Path(__file__).parent + file_hashes = {} + for fname in ("sensitivity_checks.py", "poolbuild.py", "observables.py", + "validation_protocol.py", "protocols.py", + "gradient_repair.py", "deliverable_b.py", "deliverable_c.py", + "pic_grid.py", "r2a_profiles.py", "acceptance.py"): + f = here / fname + if f.exists(): + file_hashes[fname] = hashlib.sha256(f.read_bytes()).hexdigest()[:16] + # simulator internals the experiments depend on (working tree is dirty — + # the git revision alone cannot reconstruct these) + for rel in ("myogen/simulator/jaxley/jax_models.py", + "myogen/simulator/jaxley/channels/nerlab_channels.py", + "myogen/simulator/jaxley/cells_api.py", + "myogen/simulator/jaxley/populations/motor_neurons.py"): + f = repo / rel + if f.exists(): + file_hashes[rel] = hashlib.sha256(f.read_bytes()).hexdigest()[:16] + import jaxley + + return { + "git_revision": rev or "unknown", + "git_dirty": dirty, + "provenance_note": ("experiments/twin sources are untracked by git; " + "reconstruction relies on the file hashes below, " + "not the revision"), + "source_sha256_16": file_hashes, + "versions": {"jax": jax.__version__, "jaxley": jaxley.__version__, + "numpy": np.__version__}, + "device": jax.devices()[0].platform, + "precision": ("float64 (jax_enable_x64=True)" + if jax.config.jax_enable_x64 + else "float32 (JAX default)"), + "dt_ms": 0.1, + "voltage_frame": "NERLab 1952-HH (V_rest≈0, spike threshold +50 mV)", + "seed_taxonomy": { + "connectivity": "TwinSpec.seed (numpy default_rng)", + "events": "jax.random.PRNGKey (rate path) / numpy default_rng " + "(frozen renewal events)", + "fd_probes": "deterministic, no sampling", + }, + "objectives": { + "t2": "mean soma voltage, subthreshold deterministic current", + "t3_t4": "mean exponentially-smoothed rate (tau 100 ms)", + "t5": "MSE of smoothed rates (tau 100 ms) vs frozen-event target", + }, + "scope": ("this Jaxley harness configuration only — NOT a validated " + "NEURON-matched forward model; conclusions are specific to " + "the stated parameter points, objectives, and frozen " + "input realizations"), + } + + +def _smoothed_rates_jax(spikes, dt_ms, tau_ms=100.0): + a = jnp.exp(-dt_ms / tau_ms) + inst = spikes / (dt_ms * 1e-3) + + def step(r, x): + r = a * r + (1.0 - a) * x + return r, r + + _, rates = jax.lax.scan(step, jnp.zeros(spikes.shape[1]), inst) + return rates + + +# --------------------------------------------------------------------------- # + +def t1_parameter_plumbing() -> dict: + """Which jaxley leaves does each theta knob touch; does it move the output?""" + twin = PoolTwin(TwinSpec(n_mn=2, n_dd=5, n_inh=5, seed=0)) + theta0 = twin.default_theta() + p0 = twin._scaled_jaxley_params(theta0) + + leaf_map = {} + for key in theta0: + th = dict(theta0) + th[key] = th[key] * 1.5 + p1 = twin._scaled_jaxley_params(th) + changed = [] + for name in p0: + a, b = np.asarray(p0[name]), np.asarray(p1[name]) + if a.dtype == object or not np.issubdtype(a.dtype, np.number): + continue # non-numeric leaves cannot be theta-dependent + if np.issubdtype(a.dtype, np.floating): + both = np.isfinite(a) & np.isfinite(b) + if np.any(a[both] != b[both]): + changed.append(name) + elif not np.array_equal(a, b): + changed.append(name) + leaf_map[key] = changed + + expected = { + "gCaL_scale": ["caL_gcaLbar"], + "gNaP_scale": ["napp_gnapbar"], + "drive_gain": [], + "w_exc": [], + "w_inh": [], + } + plumbing_ok = {k: leaf_map[k] == expected[k] for k in expected} + + # forward effect: a brief event burst; does each knob move soma voltage? + n_steps = 3000 + conn = np.ones((2, 1), dtype=np.float32) + burst = np.zeros((n_steps, 1), dtype=np.float32) + burst[100:200, 0] = 1.0 + out0 = np.asarray(twin.simulate_events(theta0, burst, exc_connectivity=conn)["v"]) + effects = {} + for key in theta0: + th = dict(theta0) + th[key] = th[key] * 1.5 + kwargs = dict(exc_connectivity=conn) + if key == "w_inh": # needs inhibitory input to matter + kwargs.update(inh_events=burst, inh_connectivity=conn) + out1 = np.asarray(twin.simulate_events(th, burst, **kwargs)["v"]) + base = out0 if key != "w_inh" else np.asarray( + twin.simulate_events(theta0, burst, **kwargs)["v"]) + effects[key] = float(np.max(np.abs(out1 - base))) + + # drive_gain scales the stochastic generator rate → no effect in the + # deterministic events path (by design); check it via the rate path. + _, cmd = __import__("protocols").step(0.1, t_pre_s=0.05, t_on_s=0.2, + t_post_s=0.05, level=40.0) + drives = {"dd": jnp.asarray(cmd), "inh": jnp.zeros(len(cmd), jnp.float32)} + key0 = jax.random.PRNGKey(0) + v_a = np.asarray(twin.simulate(theta0, drives, key0, "hard")["v"]) + th = dict(theta0) + th["drive_gain"] = th["drive_gain"] * 1.5 + v_b = np.asarray(twin.simulate(th, drives, key0, "hard")["v"]) + effects["drive_gain (rate path)"] = float(np.max(np.abs(v_b - v_a))) + + # Initialization audit (revised roadmap: parameter-dependent vs fixed + # initialization must be declared). In this harness the initial jaxley + # state is computed ONCE at build time from the default parameters and is + # NOT re-equilibrated when theta changes — verify and record that. + v0 = np.asarray(twin._states0["v"]) + rec = np.concatenate([np.asarray(twin._rec_soma), np.asarray(twin._rec_dend)]) + v_rec = v0[rec] + v_other = np.delete(v0, rec) + init_audit = { + "initial_state_theta_dependent": False, + "recorded_compartments_at_rest": bool( + np.allclose(v_rec, 0.0, atol=1e-6)), + "nonrecorded_init_v_range_mV": ( + [float(v_other.min()), float(v_other.max())] + if v_other.size else None), + "observation": ("non-recorded compartment(s) — the axon — initialize " + "at −1.0 mV, not the NERLab rest; a jaxley init " + "quirk, decays within ms but is part of x0 and is " + "recorded here rather than hidden"), + "caveat": ("theta changes do not re-equilibrate the initial state; " + "fits must either discard the initial transient or treat " + "x0 explicitly (roadmap x0_sc)"), + } + + ok = (all(plumbing_ok.values()) + and all(v > 0 for k, v in effects.items() if k != "drive_gain") + and init_audit["recorded_compartments_at_rest"]) + return {"ok": bool(ok), "leaf_map": leaf_map, + "plumbing_ok": plumbing_ok, "max_output_effect_mV": effects, + "initialization_audit": init_audit} + + +def t2_continuous_sensitivities(current_nA: float = 1.0, + duration_s: float = 1.0) -> dict: + """Single cell, subthreshold deterministic current: autodiff vs FD.""" + twin = PoolTwin(TwinSpec(n_mn=1, n_dd=1, n_inh=1, seed=0)) + n_steps = int(duration_s * 1000 / twin.spec.dt_ms) + i_ext = np.zeros((n_steps, 1), dtype=np.float32) + i_ext[int(0.1 * n_steps):, 0] = current_nA + i_ext = jnp.asarray(i_ext) + + def loss(scales): + th = twin.default_theta() + th["gCaL_scale"], th["gNaP_scale"] = scales[0], scales[1] + out = twin._run_current(th, i_ext, spike_mode="hard") + return jnp.mean(out["v"]) + + s0 = jnp.array([1.0, 1.0]) + v0 = float(loss(s0)) + subthreshold = True + out0 = twin.simulate_current(twin.default_theta(), np.asarray(i_ext)) + if float(np.max(np.asarray(out0["v"]))) > 50.0: + subthreshold = False # recorded, test still runs + + g_ad = np.asarray(jax.grad(loss)(s0)) + report = {"ok": True, "loss0": v0, "subthreshold": subthreshold, + "autodiff": g_ad.tolist(), "fd": {}} + for j, name in enumerate(["gCaL_scale", "gNaP_scale"]): + rows = {} + for h in (1e-1, 1e-2, 1e-3, 1e-4): + e = np.zeros(2) + e[j] = h + fd = (float(loss(s0 + e)) - float(loss(s0 - e))) / (2 * h) + rel = abs(fd - g_ad[j]) / max(abs(g_ad[j]), 1e-12) + rows[h] = {"fd": fd, "rel_err_vs_autodiff": rel} + report["fd"][name] = rows + best = min(r["rel_err_vs_autodiff"] for r in rows.values()) + if best > 0.05: + report["ok"] = False + return report + + +def t3_event_parameter_plumbing(level_hz: float = 60.0) -> dict: + """What does d(loss)/d(drive_gain) mean through the stochastic generator? + + Requires a FIRING regime — in a silent pool the hard slopes are exactly 0 + while the surrogate still reports a nonzero gradient through the + fast-sigmoid tail (observed: +81.7 with zero spikes). That mismatch is + itself a finding, but the sign comparison needs spikes. + """ + import protocols + + twin = PoolTwin(TwinSpec(n_mn=2, n_dd=100, n_inh=5, seed=0)) + _, cmd = protocols.step(twin.spec.dt_ms, t_pre_s=0.1, t_on_s=0.8, + t_post_s=0.1, level=level_hz) + drives = {"dd": jnp.asarray(cmd), "inh": jnp.zeros(len(cmd), jnp.float32)} + key0 = jax.random.PRNGKey(0) + + def loss(gain, key, mode): + th = twin.default_theta() + th["drive_gain"] = gain + out = twin._run(th, drives, key, spike_mode=mode) + return jnp.mean(_smoothed_rates_jax(out["spikes"], twin.spec.dt_ms)) + + out0 = twin.simulate(twin.default_theta(), drives, key0, "hard") + n_spikes = int(np.asarray(out0["spikes"]).sum()) + + g_sur = float(jax.grad(lambda g: loss(g, key0, "surrogate"))(jnp.float32(1.0))) + + # fixed-key finite perturbations of the HARD loss (same realization) + eps = 0.05 + slope_fixed = (float(loss(jnp.float32(1 + eps), key0, "hard")) + - float(loss(jnp.float32(1 - eps), key0, "hard"))) / (2 * eps) + # expected-value slope across independent keys (per-key slopes preserved) + keys = [jax.random.PRNGKey(100 + i) for i in range(8)] + per_key_slopes = [ + (float(loss(jnp.float32(1 + eps), k, "hard")) + - float(loss(jnp.float32(1 - eps), k, "hard"))) / (2 * eps) + for k in keys + ] + slope_expected = float(np.mean(per_key_slopes)) + slope_sd = float(np.std(per_key_slopes)) + + return { + "ok": bool(n_spikes > 0 and np.sign(g_sur) == np.sign(slope_expected) != 0), + "n_spikes_at_init": n_spikes, + "surrogate_grad": g_sur, + "hard_slope_fixed_key_eps5pct": slope_fixed, + "hard_slope_expected_8keys_eps5pct": slope_expected, + "hard_slope_per_key": per_key_slopes, + "hard_slope_sd_across_keys": slope_sd, + "note": ("pathwise/fixed-key derivatives are fixed-realization " + "quantities; agreement in sign with the expected-value slope " + "is required, magnitudes are not expected to match"), + } + + +def t4_horizon_diagnostics(current_nA: float = 10.0) -> dict: + """|grad| vs simulated horizon, suprathreshold (spiking) single cell.""" + twin = PoolTwin(TwinSpec(n_mn=1, n_dd=1, n_inh=1, seed=0)) + rows = {} + for horizon_s in (0.5, 1.0, 2.0, 4.0): + n = int(horizon_s * 1000 / twin.spec.dt_ms) + i_ext = jnp.full((n, 1), current_nA, dtype=jnp.float32) + + def loss(scales): + th = twin.default_theta() + th["gCaL_scale"], th["gNaP_scale"] = scales[0], scales[1] + out = twin._run_current(th, i_ext, spike_mode="surrogate") + return jnp.mean(_smoothed_rates_jax(out["spikes"], twin.spec.dt_ms)) + + g = np.asarray(jax.grad(loss)(jnp.array([1.0, 1.0]))) + out = twin.simulate_current(twin.default_theta(), np.asarray(i_ext), + spike_mode="hard") + rows[horizon_s] = { + "n_spikes": int(np.asarray(out["spikes"]).sum()), + "grad_gCaL": float(g[0]), + "grad_gNaP": float(g[1]), + "components_finite": bool(np.all(np.isfinite(g))), + # float64 norm: float32 squaring overflows for ~1e19 components + # and previously misreported finite gradients as Inf. + "grad_norm_f64": _safe_norm(g), + } + norms = [r["grad_norm_f64"] for r in rows.values()] + finite = [r["components_finite"] for r in rows.values()] + growth = (norms[-1] / max(norms[0], 1e-12)) if all(finite) else float("nan") + return { + "ok": bool(all(finite)), + "per_horizon": rows, + "norm_growth_shortest_to_longest": float(growth), + "note": ("distinguish: (a) dynamical explosion of finite components " + "(e.g. ~1e28 at 1 s), (b) NaN components (true breakdown), " + "(c) diagnostic float32 norm overflow — eliminated here by " + "the float64 norm. 'ok' means all components finite at every " + "horizon, NOT that magnitudes are usable."), + } + + +def t5_hard_forward_descent( + n_mn: int = 3, + duration_s: float = 4.0, + conn_seeds=(0, 1), + event_seeds=(0, 1, 2, 3), + true_points=((1.5, 0.011), (0.7, 0.008)), + lambdas=(0.03, 0.1), + coord_lambda: float = 0.1, +) -> dict: + """Do surrogate sign-descent steps reduce the scientific hard loss? + + Deliverable-B parameter pair (gCaL_scale, w_exc), frozen event input, + evaluated across connectivity seeds × event seeds × true-parameter + points. For each case: + + * JOINT test — both coordinates stepped along ``-sign(grad)`` (what the + clipped sign-descent optimizer actually does); accepted-step fraction + over ``lambdas``. + * COORDINATE-WISE test — each parameter perturbed alone by + ``exp(±coord_lambda)``; the hard-loss-descending direction (if resolved + beyond a tolerance) is compared against that coordinate's surrogate + sign. A joint direction being wrong does NOT establish each coordinate + is independently wrong — only this probe can. + + Scope: this harness configuration, objective, and frozen realizations + only; not a NEURON-validated forward model. + """ + twins = {cs: PoolTwin(TwinSpec(n_mn=n_mn, n_dd=100, n_inh=5, seed=cs)) + for cs in conn_seeds} + dt = next(iter(twins.values())).spec.dt_ms + n = int(duration_s * 1000 / dt) + t = np.arange(n) * dt * 1e-3 + cmd = np.interp(t, [0, 0.5, duration_s / 2, duration_s - 0.5, duration_s], + [0, 0, 60.0, 0, 0]).astype(np.float32) + rate_mat = np.tile(cmd[:, None], (1, 100)) + + cases = [] + n_failures = 0 + for cs, twin in twins.items(): + for es in event_seeds: + events = _renewal_events(np.random.default_rng(es), rate_mat, dt, 5) + ev = {"exc": jnp.asarray(events), + "inh": jnp.zeros((n, 5), jnp.float32), + "exc_conn": twin._exc_mat, "inh_conn": twin._inh_mat} + for tp in true_points: + case = {"conn_seed": cs, "event_seed": es, + "true_theta": {"gCaL_scale": tp[0], "w_exc": tp[1]}} + try: + theta_true = twin.default_theta() + theta_true["gCaL_scale"] = jnp.float32(tp[0]) + theta_true["w_exc"] = jnp.float32(tp[1]) + obs = twin._run_events(theta_true, ev, spike_mode="hard") + n_obs = int(np.asarray(obs["spikes"]).sum()) + case["n_obs_spikes"] = n_obs + if n_obs == 0: + case["skipped"] = "silent observation regime" + cases.append(case) + continue + target = jax.lax.stop_gradient( + _smoothed_rates_jax(obs["spikes"], dt)) + + def loss(th, mode): + out = twin._run_events(th, ev, spike_mode=mode) + return jnp.mean( + (_smoothed_rates_jax(out["spikes"], dt) - target) ** 2) + + theta0 = twin.default_theta() + L0 = float(loss(theta0, "hard")) + case["loss_at_init_hard"] = L0 + + def loss_fit(vals): + th = dict(theta0) + th["gCaL_scale"], th["w_exc"] = vals[0], vals[1] + return loss(th, "surrogate") + + g = np.asarray(jax.grad(loss_fit)(jnp.array( + [float(theta0["gCaL_scale"]), float(theta0["w_exc"])]))) + case["surrogate_grad"] = g.tolist() + case["grad_finite"] = bool(np.all(np.isfinite(g))) + if not case["grad_finite"]: + n_failures += 1 + cases.append(case) + continue + d = -np.sign(g) + case["joint_sign_direction"] = d.tolist() + + def hard_at(scales): + th = dict(theta0) + th["gCaL_scale"] = jnp.float32( + float(theta0["gCaL_scale"]) * scales[0]) + th["w_exc"] = jnp.float32( + float(theta0["w_exc"]) * scales[1]) + return float(loss(th, "hard")) + + joint = {} + for lam in lambdas: + Ls = hard_at(np.exp(lam * d)) + joint[lam] = {"hard_loss": Ls, "improved": bool(Ls < L0)} + case["joint_steps"] = joint + + # coordinate-wise ground truth vs surrogate sign. + # BASELINE-AWARE: a direction counts as descent only if + # it improves on L0 — the smaller of two worse-than-L0 + # losses is "less bad", not descent. + coords = {} + tol = 1e-4 * max(L0, 1e-9) + for j, name in enumerate(["gCaL_scale", "w_exc"]): + up = np.ones(2) + up[j] = np.exp(coord_lambda) + dn = np.ones(2) + dn[j] = np.exp(-coord_lambda) + L_up, L_dn = hard_at(up), hard_at(dn) + if min(L_up, L_dn) >= L0 - tol: + true_dir = 0 # no descent direction at this step + elif abs(L_up - L_dn) <= tol: + true_dir = 0 # both improve equally — unresolved + else: + true_dir = 1 if L_up < L_dn else -1 + sur_dir = int(d[j]) + coords[name] = { + "hard_loss_baseline": L0, + "hard_loss_up": L_up, "hard_loss_down": L_dn, + "true_descent_dir": true_dir, + "surrogate_dir": sur_dir, + "agrees": (None if true_dir == 0 + else bool(sur_dir == true_dir)), + } + case["coordinate_probes"] = coords + except Exception as e: + case["error"] = f"{type(e).__name__}: {e}" + n_failures += 1 + cases.append(case) + + evaluated = [c for c in cases if "joint_steps" in c] + all_joint = [s["improved"] for c in evaluated + for s in c["joint_steps"].values()] + accepted_fraction = (float(np.mean(all_joint)) if all_joint else float("nan")) + coord_agreement = {} + for name in ("gCaL_scale", "w_exc"): + votes = [c["coordinate_probes"][name]["agrees"] for c in evaluated + if c["coordinate_probes"][name]["agrees"] is not None] + coord_agreement[name] = { + "n_resolved": len(votes), + "n_unresolved": len(evaluated) - len(votes), + "agreement_rate": (float(np.mean(votes)) if votes + else float("nan")), + } + # Cases sharing an event seed are CORRELATED — the effective number of + # independent event realizations is len(event_seeds), not n_cases. + per_event_seed = {} + for es in event_seeds: + sub = [s["improved"] for c in evaluated if c["event_seed"] == es + for s in c["joint_steps"].values()] + per_event_seed[es] = { + "n_steps": len(sub), + "improved_fraction": (float(np.mean(sub)) if sub + else float("nan")), + } + + return { + "n_independent_event_realizations": len(event_seeds), + "per_event_seed_joint_improvement": per_event_seed, + # gate: promote only if steps reliably improve and nothing failed + "ok": bool(all_joint) and accepted_fraction >= 0.8 and n_failures == 0, + "n_cases": len(cases), + "n_evaluated": len(evaluated), + "n_failures": n_failures, + "accepted_step_fraction": accepted_fraction, + "coordinate_agreement": coord_agreement, + "cases": cases, + } + + +TESTS = { + "t1": ("parameter plumbing", t1_parameter_plumbing), + "t2": ("continuous sensitivities (FD vs autodiff)", t2_continuous_sensitivities), + "t3": ("event-parameter plumbing", t3_event_parameter_plumbing), + "t4": ("horizon diagnostics", t4_horizon_diagnostics), + "t5": ("hard-forward descent", t5_hard_forward_descent), +} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--tests", nargs="+", choices=sorted(TESTS), default=sorted(TESTS)) + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + report = {"manifest": _manifest()} + for tid in args.tests: + name, fn = TESTS[tid] + t0 = time.time() + try: + res = fn() + except Exception as e: # preserve raw failures + res = {"ok": False, "error": f"{type(e).__name__}: {e}"} + res["runtime_s"] = round(time.time() - t0, 1) + report[tid] = res + print(f"[{tid}] {name}: {'OK' if res.get('ok') else 'ATTENTION'} " + f"({res['runtime_s']}s)") + + payload = json.dumps(report, indent=2, default=float) + out = RESULTS / "sensitivity_checks.json" + out.write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + ev = EVIDENCE / f"sensitivity_checks_{stamp}.json" + ev.write_text(payload) + print(f"full report → {out}\nevidence copy (version-controllable) → {ev}") + + +if __name__ == "__main__": + main() diff --git a/experiments/twin/validation_protocol.py b/experiments/twin/validation_protocol.py new file mode 100644 index 00000000..8bac4baf --- /dev/null +++ b/experiments/twin/validation_protocol.py @@ -0,0 +1,579 @@ +"""Paired NEURON ↔ Jaxley validation protocol — immutable artifact + runners. + +The strict Phase-1 validation (roadmap) requires that both simulators receive +EXACTLY the same inputs. This module creates one frozen protocol artifact +containing everything either backend needs: + +* presynaptic event matrices (generated once, stored — never regenerated), +* connectivity matrices, +* synaptic constants (site, tau_rise, tau_decay, E_exc, E_inh, weights), +* cell parameters (recruitment thresholds, gCaL/gNaP scales), +* deterministic current-injection waveforms (step levels sized to the + measured rheobase, 5–30 nA) for the intrinsic single-cell comparison, +* dt, durations, and a content hash for provenance. + +Runners (both implemented): + +* :func:`run_jaxley` — ``PoolTwin.simulate_events`` / ``simulate_current``. +* :func:`run_neuron` — NEURON backend; VecStim → NetCon fan-in onto the + cells' dendritic Exp2Syn, IClamp at the declared site, backward Euler + (``h.secondorder = 0``, matching Jaxley's default solver), spike detection + at +50 mV (NERLab 1952-HH frame). + +Matched-construction notes (verified against the code): + +* The NEURON pool interpolates per-cell parameters on RAW thresholds while + the Jaxley pool normalizes them first — ``run_neuron`` therefore passes + min-max NORMALIZED thresholds so both backends build the same cells. +* The NEURON pool's YAML initial voltage (−67 mV) is wrong for the NERLab + frame — ``initial_voltage__mV=0.0`` is passed explicitly. +* NEURON records sample 0 as the initial condition and Jaxley's scan emits + post-step values: NEURON traces are sliced ``[1:n_steps+1]``. +* The artifact's ``w_exc_uS`` is used VERBATIM in both backends. The default + 0.009 µS was calibrated for the legacy 5 ms surrogate kernel; with the + NEURON Exp2Syn constants (0.2/0.3 ms) it is largely subthreshold — a + valid matched-input comparison (subthreshold voltage diff), just not a + physiological firing regime. Pass ``w_exc_uS=0.15`` to ``make_protocol`` + for an ex03-equivalent firing comparison. State which one you ran. + +E_inh convention: the artifact stores the number actually used (−75 = +NEURON literal, −15 = frame-corrected driving force) and the choice is +recorded with the result. + +Usage:: + + python validation_protocol.py # build artifact + Jaxley side + python validation_protocol.py --backend both # paired run + comparison + python validation_protocol.py --backend both --w-exc 0.15 # firing regime +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import joblib +import numpy as np + +RESULTS = Path(__file__).parent / "results" +EVIDENCE = Path(__file__).parent / "evidence" +ARTIFACT_VERSION = 2 # v2: current steps 5/15/30 nA (rheobase-scaled) + +# PROCESS-ISOLATION NOTE: loading JAX and NEURON's compiled mechanisms in the +# same process segfaults (verified: jaxley-then-neuron → SIGSEGV; each alone +# is fine). Therefore jax/poolbuild imports are deferred into the jaxley +# runner, the NEURON runner imports no jax, and `--backend both` orchestrates +# the two backends as separate subprocesses. + + +def fuglevand_thresholds(n: int, recruitment_range: float = 30.0) -> np.ndarray: + """Local copy (poolbuild has the canonical one) so building an artifact + does not pull jax into a NEURON process.""" + i = np.arange(1, n + 1) + return np.exp(i * np.log(recruitment_range) / n) / recruitment_range + + +# --------------------------------------------------------------------------- # +# Artifact creation +# --------------------------------------------------------------------------- # + +def _renewal_events(rng, rates_hz, dt_ms, order: int) -> np.ndarray: + """Frozen gamma/Erlang renewal events, (n_steps, n_src), 0/1 float32.""" + n_steps, n_src = rates_hz.shape + yi = np.zeros(n_src) + thres = rng.gamma(order, 1.0 / order, size=n_src) + out = np.zeros((n_steps, n_src), dtype=np.float32) + for t in range(n_steps): + yi += rates_hz[t] * dt_ms * 1e-3 + fired = yi >= thres + if fired.any(): + out[t, fired] = 1.0 + yi[fired] = 0.0 + thres[fired] = rng.gamma(order, 1.0 / order, size=int(fired.sum())) + return out + + +def make_protocol( + n_mn: int = 8, + n_exc: int = 100, + n_inh: int = 40, + dt_ms: float = 0.1, + connect_p: float = 0.5, + t_rise_s: float = 7.5, + t_pre_s: float = 1.0, + t_post_s: float = 3.0, + peak_hz: float = 60.0, + inh_rate_hz: float = 0.0, + renewal_order: int = 5, + syn_site: str = "dendrite", + tau_rise_ms: float = 0.2, + tau_decay_ms: float = 0.3, + e_exc_mV: float = 70.0, + e_inh_mV: float = -15.0, + w_exc_uS: float = 0.009, + w_inh_uS: float = 0.009, + gcal_scale: float = 1.0, + gnap_scale: float = 1.0, + step_currents_nA=(5.0, 15.0, 30.0), + seed: int = 0, +) -> dict: + """Build the frozen protocol artifact (pure data, no simulator objects).""" + rng = np.random.default_rng(seed) + + n_ramp = int(round((t_pre_s + 2 * t_rise_s + t_post_s) * 1000 / dt_ms)) + t = np.arange(n_ramp) * dt_ms * 1e-3 + cmd = np.interp(t, [0, t_pre_s, t_pre_s + t_rise_s, t_pre_s + 2 * t_rise_s, + t[-1]], [0, 0, peak_hz, 0, 0]).astype(np.float32) + exc_events = _renewal_events( + rng, np.tile(cmd[:, None], (1, n_exc)), dt_ms, renewal_order) + inh_events = _renewal_events( + rng, np.full((n_ramp, n_inh), inh_rate_hz, dtype=np.float32), + dt_ms, renewal_order) + + exc_conn = (rng.random((n_mn, n_exc)) < connect_p).astype(np.float32) + inh_conn = (rng.random((n_mn, n_inh)) < connect_p).astype(np.float32) + + # Deterministic current steps (1 s pre, 4 s on, 2 s post), rheobase-scaled. + n_step = int(round(7.0 * 1000 / dt_ms)) + i_steps = {} + for amp in step_currents_nA: + w = np.zeros(n_step, dtype=np.float32) + w[int(1.0 * 1000 / dt_ms): int(5.0 * 1000 / dt_ms)] = amp + i_steps[float(amp)] = w + + art = { + "version": ARTIFACT_VERSION, + "seed": seed, + "dt_ms": dt_ms, + "n_mn": n_mn, + "thresholds": fuglevand_thresholds(n_mn), + "ramp_command_hz": cmd, + "exc_events": exc_events, + "inh_events": inh_events, + "exc_conn": exc_conn, + "inh_conn": inh_conn, + "synapse": { + "site": syn_site, + "tau_rise_ms": tau_rise_ms, + "tau_decay_ms": tau_decay_ms, + "e_exc_mV": e_exc_mV, + "e_inh_mV": e_inh_mV, + "w_exc_uS": w_exc_uS, + "w_inh_uS": w_inh_uS, + "live_driving_force": True, + "e_inh_convention": ( + "literal NEURON value" if e_inh_mV == -75.0 + else "frame-corrected driving force"), + }, + "cell": {"gcal_scale": gcal_scale, "gnap_scale": gnap_scale}, + "current_steps_nA": i_steps, + } + art["content_hash"] = _hash_artifact(art) + return art + + +def _hash_artifact(art: dict) -> str: + h = hashlib.sha256() + for key in sorted(art): + if key == "content_hash": + continue + val = art[key] + h.update(key.encode()) + if isinstance(val, np.ndarray): + h.update(val.tobytes()) + elif isinstance(val, dict): + h.update(json.dumps( + {k: (v.tolist() if isinstance(v, np.ndarray) else v) + for k, v in sorted(val.items())}, sort_keys=True).encode()) + else: + h.update(json.dumps(val, default=float).encode()) + return h.hexdigest()[:16] + + +def save_protocol(art: dict, path: Path) -> None: + joblib.dump(art, path) + + +def load_protocol(path: Path) -> dict: + art = joblib.load(path) + expected = art.get("content_hash") + actual = _hash_artifact(art) + if expected != actual: + raise ValueError( + f"protocol artifact hash mismatch ({expected} != {actual}) — " + "the artifact was modified after creation") + return art + + +# --------------------------------------------------------------------------- # +# Jaxley runner +# --------------------------------------------------------------------------- # + +def _twin_from_artifact(art: dict): + from poolbuild import PoolTwin, TwinSpec # deferred: keeps jax out of + # NEURON-only processes (see process-isolation note above) + + syn = art["synapse"] + return PoolTwin(TwinSpec( + n_mn=art["n_mn"], + n_dd=art["exc_events"].shape[1], + n_inh=art["inh_events"].shape[1], + dt_ms=art["dt_ms"], + seed=art["seed"], + syn_site=syn["site"], + live_driving_force=syn["live_driving_force"], + syn_tau_rise_ms=syn["tau_rise_ms"], + tau_syn_ms=syn["tau_decay_ms"], + e_exc_mV=syn["e_exc_mV"], + e_inh_mV=syn["e_inh_mV"], + w_exc_uS=syn["w_exc_uS"], + w_inh_uS=syn["w_inh_uS"], + )) + + +def run_jaxley(art: dict) -> dict: + """Run the full artifact on the Jaxley twin. Returns numpy result dict.""" + import jax.numpy as jnp + + twin = _twin_from_artifact(art) + theta = twin.default_theta() + theta["gCaL_scale"] = jnp.float32(art["cell"]["gcal_scale"]) + theta["gNaP_scale"] = jnp.float32(art["cell"]["gnap_scale"]) + + out = twin.simulate_events( + theta, + art["exc_events"], + inh_events=art["inh_events"], + exc_connectivity=art["exc_conn"], + inh_connectivity=art["inh_conn"], + spike_mode="hard", + ) + res = { + "backend": "jaxley", + "content_hash": art["content_hash"], + "pool": {k: np.asarray(v) for k, v in out.items()}, + "current_steps": {}, + } + for amp, wave in art["current_steps_nA"].items(): + i_ext = np.tile(np.asarray(wave)[:, None], (1, art["n_mn"])) + o = twin.simulate_current(theta, i_ext, spike_mode="hard") + res["current_steps"][amp] = {k: np.asarray(v) for k, v in o.items()} + return res + + +# --------------------------------------------------------------------------- # +# NEURON runner +# --------------------------------------------------------------------------- # + +def _neuron_pool(art: dict): + """Build the matched NEURON NERLab pool (normalized thresholds, + v_init = 0, artifact PIC scales). Returns (h, pool).""" + from myogen.utils.nmodl import load_nmodl_mechanisms + + load_nmodl_mechanisms(quiet=True) + from neuron import h + + h.load_file("stdrun.hoc") + from myogen.simulator.neuron import cells as ncells + from myogen.simulator.neuron.populations import AlphaMN__Pool + + ncells.reset_cell_id_counters() + rt = np.asarray(art["thresholds"], dtype=float) + # NEURON's special_interp uses RAW thresholds; Jaxley normalizes first. + # Pass normalized thresholds so both backends build identical cells. + rt_norm = (rt - rt.min()) / (rt.max() - rt.min() + 1e-10) + pool = AlphaMN__Pool( + recruitment_thresholds__array=rt_norm, + model="NERLab", + mode="active", + initial_voltage__mV=0.0, # YAML default −67 is wrong for NERLab + spike_threshold__mV=50.0, + ) + for c in pool: + c.dend[0].gcaLbar_caL *= art["cell"]["gcal_scale"] + c.soma.gnapbar_napp *= art["cell"]["gnap_scale"] + return h, pool + + +def _neuron_record(h, pool): + keep, vrec, dvrec, spk = [], [], [], [] + for c in pool: + vrec.append(h.Vector().record(c.soma(0.5)._ref_v)) + dvrec.append(h.Vector().record(c.dend[0](0.5)._ref_v)) + sv = h.Vector() + nc = h.NetCon(c.soma(0.5)._ref_v, None, sec=c.soma) + nc.threshold = 50.0 + nc.record(sv) + keep.append(nc) + spk.append(sv) + return keep, vrec, dvrec, spk + + +def _neuron_advance(h, pool, dt_ms: float, n_steps: int): + h.dt = dt_ms + h.secondorder = 0 # backward Euler ≡ jaxley's default solver + for sec, v in zip(*pool.get_initialization_data()): + sec.v = v + h.finitialize() + for _ in range(n_steps): # exact count — `while h.t < tstop` over-steps + h.fadvance() + + +def _collect(vrec, dvrec, spk, dt_ms, n_steps, n_mn): + # NEURON sample 0 is the initial condition; Jaxley emits post-step values. + v = np.stack([np.asarray(r.as_numpy())[1:n_steps + 1] for r in vrec], axis=1) + vd = np.stack([np.asarray(r.as_numpy())[1:n_steps + 1] for r in dvrec], axis=1) + spikes = np.zeros((n_steps, n_mn), dtype=np.float32) + for i, sv in enumerate(spk): + for t_ms in np.asarray(sv.as_numpy()): + k = min(int(round(t_ms / dt_ms)), n_steps - 1) + spikes[k, i] = 1.0 + return {"v": v, "v_dend": vd, "spikes": spikes} + + +def run_neuron(art: dict) -> dict: + """Run the full artifact on the NEURON backend (matched construction).""" + dt = art["dt_ms"] + n_mn = art["n_mn"] + syn_p = art["synapse"] + + # ---- pool/synaptic phase + h, pool = _neuron_pool(art) + keep = [] + vstims = [] + syn_ncs = [] + for name, events, conn, syn_idx, w, e in ( + ("exc", art["exc_events"], art["exc_conn"], 0, + syn_p["w_exc_uS"], syn_p["e_exc_mV"]), + ("inh", art["inh_events"], art["inh_conn"], 1, + syn_p["w_inh_uS"], syn_p["e_inh_mV"])): + stims = [] + for s in range(events.shape[1]): + times = np.flatnonzero(events[:, s]) * dt + tv = h.Vector(times.astype(float)) + vs = h.VecStim() + vs.play(tv) + stims.append((vs, tv)) # keep tv alive — GC trap + vstims.append(stims) + for i, c in enumerate(pool): + syn = c.synapse__list[syn_idx] # 0=exc, 1=inh; both at dend(0.5) + syn.tau1 = syn_p["tau_rise_ms"] + syn.tau2 = syn_p["tau_decay_ms"] + syn.e = e + for s in np.flatnonzero(conn[i]): + nc = h.NetCon(stims[s][0], syn) + nc.weight[0] = w + nc.delay = 0.0 # matched to the zero-delay Jaxley event path + keep.append(nc) + syn_ncs.append(nc) + rec_keep, vrec, dvrec, spk = _neuron_record(h, pool) + n_steps = art["exc_events"].shape[0] + _neuron_advance(h, pool, dt, n_steps) + res = { + "backend": "neuron", + "content_hash": art["content_hash"], + "pool": _collect(vrec, dvrec, spk, dt, n_steps, n_mn), + "current_steps": {}, + } + + # ---- intrinsic current steps, SAME pool (building a second pool while + # the first is alive segfaults NEURON at finitialize — duplicate hoc + # identities after reset_cell_id_counters). Disable the synaptic input + # in place instead: zero all NetCon weights and empty the event vectors. + for nc in syn_ncs: + nc.weight[0] = 0.0 + for stims in vstims: + for _vs, tv in stims: + tv.resize(0) # VecStim with empty vector = silent (verified) + ics = [] + for c in pool: + seg = (c.dend[0](0.5) if syn_p["site"] == "dendrite" + else c.soma(0.5)) + ic = h.IClamp(seg) + ic.delay = 0.0 + ic.dur = 1e9 + ic.amp = 0.0 + ics.append(ic) + for amp, wave in art["current_steps_nA"].items(): + players = [] + for ic in ics: + tvec = h.Vector(np.arange(len(wave)) * dt) + ivec = h.Vector(np.asarray(wave, dtype=float)) + ivec.play(ic._ref_amp, tvec, 1) + players.append((tvec, ivec)) + for sv in spk: # NetCon.record vectors are not cleared by finitialize + sv.resize(0) + n = len(wave) + _neuron_advance(h, pool, dt, n) + res["current_steps"][amp] = _collect(vrec, dvrec, spk, dt, n, n_mn) + for _tvec, ivec in players: + ivec.play_remove() # avoid stacking play sources on the same ref + return res + + +# --------------------------------------------------------------------------- # +# Comparison +# --------------------------------------------------------------------------- # + +def _match_spikes(a: np.ndarray, b: np.ndarray, tol_s: float) -> int: + i = j = n = 0 + while i < len(a) and j < len(b): + d = a[i] - b[j] + if abs(d) <= tol_s: + n += 1 + i += 1 + j += 1 + elif d < 0: + i += 1 + else: + j += 1 + return n + + +def compare_backends(res_j: dict, res_n: dict, dt_ms: float) -> dict: + """Backend-diff metrics for every phase of the artifact.""" + assert res_j["content_hash"] == res_n["content_hash"] + + def diff(pj, pn): + out = {} + for key in ("v", "v_dend"): + d = pj[key] - pn[key] + out[f"{key}_rmse_mV"] = float(np.sqrt((d ** 2).mean())) + out[f"{key}_max_abs_mV"] = float(np.abs(d).max()) + sj, sn = pj["spikes"], pn["spikes"] + out["spikes_jaxley_per_mn"] = sj.sum(axis=0).astype(int).tolist() + out["spikes_neuron_per_mn"] = sn.sum(axis=0).astype(int).tolist() + tot_j, tot_n = int(sj.sum()), int(sn.sum()) + matched = sum( + _match_spikes(np.flatnonzero(sj[:, i]) * dt_ms * 1e-3, + np.flatnonzero(sn[:, i]) * dt_ms * 1e-3, 5e-3) + for i in range(sj.shape[1])) + out["spike_match_pm5ms"] = { + "matched": matched, "jaxley_total": tot_j, "neuron_total": tot_n} + return out + + rep = {"content_hash": res_j["content_hash"], + "pool": diff(res_j["pool"], res_n["pool"]), + "current_steps": {amp: diff(res_j["current_steps"][amp], + res_n["current_steps"][amp]) + for amp in res_j["current_steps"]}} + return rep + + +# --------------------------------------------------------------------------- # + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--artifact", type=Path, + default=RESULTS / "val_protocol.pkl") + ap.add_argument("--rebuild", action="store_true") + ap.add_argument("--backend", choices=["jaxley", "neuron", "both"], + default="jaxley") + ap.add_argument("--n-mn", type=int, default=8) + ap.add_argument("--e-inh", type=float, default=-15.0) + ap.add_argument("--w-exc", type=float, default=0.009, + help="0.009 = matched-input (subthreshold) diff; " + "0.15 = ex03-equivalent firing regime") + args = ap.parse_args() + + RESULTS.mkdir(exist_ok=True) + EVIDENCE.mkdir(exist_ok=True) + art = None + if args.artifact.exists() and not args.rebuild: + art = load_protocol(args.artifact) + if (art.get("version") != ARTIFACT_VERSION + or art["synapse"]["w_exc_uS"] != args.w_exc + or art["synapse"]["e_inh_mV"] != args.e_inh + or art["n_mn"] != args.n_mn): + art = None # stale → rebuild + if art is None: + art = make_protocol(n_mn=args.n_mn, e_inh_mV=args.e_inh, + w_exc_uS=args.w_exc) + save_protocol(art, args.artifact) + print(f"built artifact v{ARTIFACT_VERSION} (hash " + f"{art['content_hash']}, w_exc {args.w_exc} µS, e_inh " + f"{args.e_inh} = {art['synapse']['e_inh_convention']})") + else: + print(f"loaded artifact {args.artifact} (hash {art['content_hash']})") + + results = {} + if args.backend == "both": + # JAX and NEURON cannot share a process (verified SIGSEGV) — + # run each backend as its own subprocess, then compare the pkls. + import subprocess + import sys + + for backend in ("jaxley", "neuron"): + cmd = [sys.executable, "-u", str(Path(__file__).resolve()), + "--backend", backend, "--artifact", str(args.artifact), + "--n-mn", str(args.n_mn), "--e-inh", str(args.e_inh), + "--w-exc", str(args.w_exc)] + print(f"--- subprocess: {backend} ---", flush=True) + proc = subprocess.run(cmd, cwd=Path(__file__).parent, + capture_output=True, text=True) + for line in (proc.stdout + proc.stderr).splitlines(): + if not any(s in line for s in ("DISPLAY", "graphics", "Added")): + print(" ", line) + if proc.returncode != 0: + raise RuntimeError( + f"{backend} subprocess failed (exit {proc.returncode})") + results[backend] = joblib.load( + RESULTS / f"val_{backend}_{art['content_hash']}.pkl") + elif args.backend in ("jaxley", "neuron"): + runner = run_jaxley if args.backend == "jaxley" else run_neuron + t0 = time.time() + res = runner(art) + results[args.backend] = res + sp = res["pool"]["spikes"] + print(f"[{args.backend}] pool: {int((sp.sum(axis=0) > 0).sum())}/" + f"{art['n_mn']} active, {int(sp.sum()):,} spikes " + f"({time.time()-t0:.0f}s)", flush=True) + for amp, o in res["current_steps"].items(): + print(f" step {amp:4.1f} nA → spikes/MN " + f"{np.asarray(o['spikes']).sum(axis=0).astype(int)}", + flush=True) + joblib.dump(res, RESULTS / f"val_{args.backend}_{art['content_hash']}.pkl") + + if args.backend == "both": + rep = compare_backends(results["jaxley"], results["neuron"], + art["dt_ms"]) + # manifest import deferred: sensitivity_checks pulls jax, which is + # safe in THIS parent process (it never imports neuron) + from sensitivity_checks import _manifest + + payload = json.dumps( + {"manifest": { + **_manifest(), + "validation_protocol_sha256_16": hashlib.sha256( + Path(__file__).read_bytes()).hexdigest()[:16], + "args": vars(args), + }, + "artifact": {k: art[k] for k in + ("version", "seed", "dt_ms", "n_mn", "synapse", + "cell", "content_hash")}, + "comparison": rep}, indent=1, default=float) + out = RESULTS / f"val_comparison_{art['content_hash']}.json" + out.write_text(payload) + stamp = time.strftime("%Y%m%d_%H%M%S") + (EVIDENCE / f"val_comparison_{art['content_hash']}_{stamp}.json" + ).write_text(payload) + print("\n=== NEURON ↔ Jaxley paired comparison ===") + p = rep["pool"] + print(f"pool phase: soma V rmse {p['v_rmse_mV']:.3f} mV " + f"(max {p['v_max_abs_mV']:.2f}); dend rmse " + f"{p['v_dend_rmse_mV']:.3f} mV") + print(f" spikes: jaxley {p['spike_match_pm5ms']['jaxley_total']} vs " + f"neuron {p['spike_match_pm5ms']['neuron_total']}, matched " + f"±5 ms: {p['spike_match_pm5ms']['matched']}") + for amp, d in rep["current_steps"].items(): + print(f"step {amp:4.1f} nA: soma rmse {d['v_rmse_mV']:7.3f} mV | " + f"spikes j{d['spikes_jaxley_per_mn']} vs " + f"n{d['spikes_neuron_per_mn']}") + print(f"saved → {out} (+ evidence copy)") + + +if __name__ == "__main__": + main() diff --git a/myogen/simulator/__init__.py b/myogen/simulator/__init__.py index 023ae6ab..98d39c1a 100644 --- a/myogen/simulator/__init__.py +++ b/myogen/simulator/__init__.py @@ -27,14 +27,14 @@ ) from .core.muscle import Muscle as Muscle from .core.physiological_distribution import RecruitmentThresholds as RecruitmentThresholds - from .neuron.joint_dynamics import JointDynamics as JointDynamics - from .neuron.muscle import HillModel as HillModel - from .neuron.network import Network as Network - from .neuron.proprioception import ( + from .jaxley.joint_dynamics import JointDynamics as JointDynamics + from .jaxley.muscle import HillModel as HillModel + from .jaxley.network import Network as Network + from .jaxley.proprioception import ( GolgiTendonOrganModel as GolgiTendonOrganModel, SpindleModel as SpindleModel, ) - from .neuron.simulation_runner import SimulationRunner as SimulationRunner + from .jaxley.simulation_runner import SimulationRunner as SimulationRunner from myogen.utils.neo import ( GridAnalogSignal as GridAnalogSignal, create_grid_signal as create_grid_signal, @@ -49,8 +49,8 @@ 'MemoryBudgetError': 'myogen.simulator.core.emg.runtime', 'RecruitmentThresholds': 'myogen.simulator.core.physiological_distribution', 'Muscle': 'myogen.simulator.core.muscle', - 'Network': 'myogen.simulator.neuron.network', - 'SimulationRunner': 'myogen.simulator.neuron.simulation_runner', + 'Network': 'myogen.simulator.jaxley.network', + 'SimulationRunner': 'myogen.simulator.jaxley.simulation_runner', 'SurfaceEMG': 'myogen.simulator.core.emg', 'IntramuscularEMG': 'myogen.simulator.core.emg', 'SurfaceElectrodeArray': 'myogen.simulator.core.emg', @@ -60,10 +60,10 @@ 'JointBiomechanics': 'myogen.simulator.core.force', 'JointGeometry': 'myogen.simulator.core.force', 'MuscleGeometry': 'myogen.simulator.core.force', - 'HillModel': 'myogen.simulator.neuron.muscle', - 'SpindleModel': 'myogen.simulator.neuron.proprioception', - 'GolgiTendonOrganModel': 'myogen.simulator.neuron.proprioception', - 'JointDynamics': 'myogen.simulator.neuron.joint_dynamics', + 'HillModel': 'myogen.simulator.jaxley.muscle', + 'SpindleModel': 'myogen.simulator.jaxley.proprioception', + 'GolgiTendonOrganModel': 'myogen.simulator.jaxley.proprioception', + 'JointDynamics': 'myogen.simulator.jaxley.joint_dynamics', 'create_grid_signal': 'myogen.utils.neo', 'signal_to_grid': 'myogen.utils.neo', 'get_electrode': 'myogen.utils.neo', @@ -73,8 +73,13 @@ } __all__ = list(_EXPORTS) - def __getattr__(name): + if name == "neuron": + # Optional NEURON backend namespace (requires the ``neuron`` extra); kept + # lazy so ``import myogen.simulator`` never touches the NEURON runtime. + value = import_module("myogen.simulator.neuron") + globals()[name] = value + return value if name not in _EXPORTS: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") value = getattr(import_module(_EXPORTS[name], __package__), name) @@ -83,4 +88,4 @@ def __getattr__(name): def __dir__(): - return sorted(set(globals()) | set(__all__)) + return sorted(set(globals()) | set(__all__) | {"neuron"}) diff --git a/myogen/simulator/jaxley/__init__.py b/myogen/simulator/jaxley/__init__.py new file mode 100644 index 00000000..d6961047 --- /dev/null +++ b/myogen/simulator/jaxley/__init__.py @@ -0,0 +1,217 @@ +""" +Jaxley Backend for MyoGen Neural Simulations. + +This package provides JAX-accelerated implementations of neural models +for motor control simulations, as an alternative to the NEURON backend. + +Submodules: +- channels: Ion channel mechanisms (Na, K, Ca, etc.) +- synapses: Synaptic mechanisms (Exp2Syn, NMDA) +- cells: Biophysical cell builders +- populations: Population-level containers +- network: Network connectivity +- simulation_engine: Core simulation infrastructure +- integrate: Core integration functions for Jaxley simulation +""" + +# Import cell classes (API-compatible with NEURON) +from myogen.simulator.jaxley.cells_api import ( + AlphaMN, + INgII, + INgIb, + DD, + DD_Gamma, + AffIa, + AffII, + AffIb, + _Cell, +) + +# Import biophysical builders +from myogen.simulator.jaxley.cells.biophysical import ( + BiophysicalMotorNeuron, + BiophysicalInterneuron, + create_motor_neuron, + create_motor_neuron_builder, + create_interneuron, + create_interneuron_builder, +) + +# Import channels +from myogen.simulator.jaxley.channels import ( + Na3rp, + Naps, + KdrRL, + MAHP, + Gh, + LCaInact, + LeakChannel, + get_motor_neuron_channels_soma, + get_motor_neuron_channels_dendrite, +) + +# Import synapses +from myogen.simulator.jaxley.synapses import ( + Exp2Syn, + ExcitatorySynapse, + InhibitorySynapse, + NMDASynapse, +) + +# Import simulation engine +from myogen.simulator.jaxley.simulation_engine import ( + JaxleyNetworkSimulator, + JaxleyCellSimulator, + SimulationConfig, + CurrentInjection, + SpikeRecord, + compute_firing_rate, + compute_isi_statistics, +) + +# Import integration module +from myogen.simulator.jaxley.integrate import ( + step_current, + ramp_current, + sinusoidal_current, + detect_spikes, + compute_firing_rate_from_spikes, + compute_isi_cv, + simulate_cell, + simulate_population, + compute_fi_curve, + CellSimulationResult, + PopulationSimulationResult, + NetworkSimulator, + SynapticConnection, +) + +# Import network +from myogen.simulator.jaxley.network import ( + Network, + JaxleyConnection, +) + +# Differentiable closed-loop entry point + gradient helpers +from myogen.simulator.jaxley.closed_loop import ( + ClosedLoopConfig, + run_jax, + compile_run, + value_and_grad_run, + partition_differentiable, +) + +# Differentiable EMG synthesis +from myogen.simulator.jaxley.emg import ( + surface_emg_jax, + intramuscular_emg_jax, + resample_muaps, +) + +# Spike-mode primitive + differentiable Bessel functions +from myogen.simulator.jaxley.jax_models import ( + spike_detect, + differentiable_twitch_params, +) +from myogen.simulator.jaxley.bessel import ( + iv_int, + kv_int, +) + +# Import JIT-compiled and batched simulation utilities +from myogen.simulator.jaxley.jit_simulation import ( + create_jitted_simulator, + create_parameter_sweep_simulator, + create_batched_simulator, + simulate_population_batched, + BatchedPopulationResult, + SynapticEvent, + SynapseStateManager, + JITNetworkSimulator, + run_fi_curve_batched, + warmup_jit, +) + +__all__ = [ + # Cells (API-compatible) + "AlphaMN", + "INgII", + "INgIb", + "DD", + "DD_Gamma", + "AffIa", + "AffII", + "AffIb", + "_Cell", + # Biophysical builders + "BiophysicalMotorNeuron", + "BiophysicalInterneuron", + "create_motor_neuron", + "create_motor_neuron_builder", + "create_interneuron", + "create_interneuron_builder", + # Channels + "Na3rp", + "Naps", + "KdrRL", + "MAHP", + "Gh", + "LCaInact", + "LeakChannel", + "get_motor_neuron_channels_soma", + "get_motor_neuron_channels_dendrite", + # Synapses + "Exp2Syn", + "ExcitatorySynapse", + "InhibitorySynapse", + "NMDASynapse", + # Simulation + "JaxleyNetworkSimulator", + "JaxleyCellSimulator", + "SimulationConfig", + "CurrentInjection", + "SpikeRecord", + "compute_firing_rate", + "compute_isi_statistics", + # Integration + "step_current", + "ramp_current", + "sinusoidal_current", + "detect_spikes", + "compute_firing_rate_from_spikes", + "compute_isi_cv", + "simulate_cell", + "simulate_population", + "compute_fi_curve", + "CellSimulationResult", + "PopulationSimulationResult", + "NetworkSimulator", + "SynapticConnection", + # JIT and Batched Simulation + "create_jitted_simulator", + "create_parameter_sweep_simulator", + "create_batched_simulator", + "simulate_population_batched", + "BatchedPopulationResult", + "SynapticEvent", + "SynapseStateManager", + "JITNetworkSimulator", + "run_fi_curve_batched", + "warmup_jit", + # Network + "Network", + "JaxleyConnection", + # Differentiable pipeline (closed loop, EMG, Bessel) + "ClosedLoopConfig", + "run_jax", + "compile_run", + "value_and_grad_run", + "partition_differentiable", + "surface_emg_jax", + "intramuscular_emg_jax", + "resample_muaps", + "spike_detect", + "differentiable_twitch_params", + "iv_int", + "kv_int", +] diff --git a/myogen/simulator/jaxley/bessel.py b/myogen/simulator/jaxley/bessel.py new file mode 100644 index 00000000..7e80258c --- /dev/null +++ b/myogen/simulator/jaxley/bessel.py @@ -0,0 +1,110 @@ +""" +Differentiable modified Bessel functions I_n(x), K_n(x) for integer order in JAX. + +JAX's ``jax.scipy.special`` ships only ``i0/i1/i0e/i1e`` — no ``k0/k1`` and no +arbitrary-order ``iv/kv``. The surface-EMG volume-conductor model +(``core/emg/surface/simulate_fiber.py``) evaluates ``iv``/``kv`` at integer orders +up to ~16, so a native-JAX, differentiable implementation is required to make the +EMG differentiable w.r.t. tissue conductivities, conduction velocity, and geometry +(milestone M7). + +The forward pass is built entirely from smooth ops (JAX ``i0``/``i1``, polynomial +approximations, logs, and stable Bessel recurrences), so ``jax.grad`` produces the +exact derivative automatically — validated against scipy and the analytic +identities ``I_n'=(I_{n-1}+I_{n+1})/2``, ``K_n'=-(K_{n-1}+K_{n+1})/2`` to ~1e-7 +(the accuracy floor of JAX's i0/i1) across orders 0–16 and x∈[0.05, 40]. + +See ``scripts/prototype_bessel_jax.py`` for the validation harness. +""" + +from __future__ import annotations + +from math import factorial + +import jax.numpy as jnp +from jax.scipy.special import i0, i1 + +__all__ = ["k0", "k1", "iv_int", "kv_int"] + + +def k0(x): + """Modified Bessel function K_0(x), x > 0 (Abramowitz & Stegun 9.8.5/9.8.6).""" + x = jnp.asarray(x) + y_s = x * x / 4.0 + small = (-jnp.log(x / 2.0) * i0(x)) + ( + -0.57721566 + y_s * (0.42278420 + y_s * (0.23069756 + y_s * (0.03488590 + + y_s * (0.00262698 + y_s * (0.00010750 + y_s * 0.00000740))))) + ) + y_l = 2.0 / x + large = (jnp.exp(-x) / jnp.sqrt(x)) * ( + 1.25331414 + y_l * (-0.07832358 + y_l * (0.02189568 + y_l * (-0.01062446 + + y_l * (0.00587872 + y_l * (-0.00251540 + y_l * 0.00053208))))) + ) + return jnp.where(x <= 2.0, small, large) + + +def k1(x): + """Modified Bessel function K_1(x), x > 0 (Abramowitz & Stegun 9.8.7/9.8.8).""" + x = jnp.asarray(x) + y_s = x * x / 4.0 + small = (jnp.log(x / 2.0) * i1(x)) + (1.0 / x) * ( + 1.0 + y_s * (0.15443144 + y_s * (-0.67278579 + y_s * (-0.18156897 + + y_s * (-0.01919402 + y_s * (-0.00110404 + y_s * (-0.00004686)))))) + ) + y_l = 2.0 / x + large = (jnp.exp(-x) / jnp.sqrt(x)) * ( + 1.25331414 + y_l * (0.23498619 + y_l * (-0.03655620 + y_l * (0.01504268 + + y_l * (-0.00780353 + y_l * (0.00325614 + y_l * (-0.00068245))))))) + return jnp.where(x <= 2.0, small, large) + + +def kv_int(n: int, x): + """K_n(x) for non-negative integer n via stable upward recurrence. + + ``K_{m+1}(x) = K_{m-1}(x) + (2m/x) K_m(x)``, seeded by K_0, K_1. + """ + x = jnp.asarray(x) + if n == 0: + return k0(x) + if n == 1: + return k1(x) + km1, kk = k0(x), k1(x) + for m in range(1, n): + km1, kk = kk, km1 + (2.0 * m / x) * kk + return kk + + +def _iv_series(n: int, x, terms: int = 34): + """Ascending series I_n(x) = Σ_k (x/2)^(2k+n) / (k! (n+k)!) — small-x branch.""" + half = x / 2.0 + term = half ** n / float(factorial(n)) + total = term + for k in range(1, terms): + term = term * (half * half) / (k * (n + k)) + total = total + term + return total + + +def iv_int(n: int, x, miller_start: int = 60, x_switch: float = 1.0): + """I_n(x) for non-negative integer n, differentiable for all x > 0. + + ``x >= x_switch``: Miller downward recurrence (the stable direction for I_n), + seeded at a high order and normalised by matching the known I_0 (JAX ``i0``). + ``x < x_switch``: ascending series, which keeps value and gradient accurate + where Miller's ``(2k/x)`` factors would blow up. Blended with ``jnp.where``. + """ + x = jnp.asarray(x) + if n == 0: + return i0(x) + if n == 1: + return i1(x) + x_safe = jnp.where(x >= x_switch, x, jnp.asarray(x_switch, x.dtype)) + M = max(miller_start, n + 40) + seq = [None] * (M + 2) + seq[M + 1] = jnp.zeros_like(x_safe) + seq[M] = jnp.ones_like(x_safe) + for k in range(M, 0, -1): + seq[k - 1] = seq[k + 1] + (2.0 * k / x_safe) * seq[k] + miller = seq[n] * (i0(x_safe) / seq[0]) + series = _iv_series(n, x) + return jnp.where(x >= x_switch, miller, series) diff --git a/myogen/simulator/jaxley/cells/__init__.py b/myogen/simulator/jaxley/cells/__init__.py new file mode 100644 index 00000000..2e85240a --- /dev/null +++ b/myogen/simulator/jaxley/cells/__init__.py @@ -0,0 +1,54 @@ +"""Cell models for Jaxley simulations.""" + +# Import biophysical builders +from myogen.simulator.jaxley.cells.biophysical import ( + BiophysicalMotorNeuron, + BiophysicalInterneuron, + create_motor_neuron, + create_motor_neuron_builder, + create_interneuron, + create_interneuron_builder, + POWERS2017_MORPHOLOGY, + NERLAB_MORPHOLOGY, + POWERS2017_CHANNELS, + NERLAB_CHANNELS, + JAXLEY_MECH_AVAILABLE, + JAXLEY_MECH_CHANNELS, +) + +# Import API-compatible cell classes +from myogen.simulator.jaxley.cells_api import ( + AlphaMN, + INgII, + INgIb, + DD, + DD_Gamma, + AffIa, + AffII, + AffIb, +) + +__all__ = [ + # Biophysical builders + "BiophysicalMotorNeuron", + "BiophysicalInterneuron", + "create_motor_neuron", + "create_motor_neuron_builder", + "create_interneuron", + "create_interneuron_builder", + "POWERS2017_MORPHOLOGY", + "NERLAB_MORPHOLOGY", + "POWERS2017_CHANNELS", + "NERLAB_CHANNELS", + "JAXLEY_MECH_AVAILABLE", + "JAXLEY_MECH_CHANNELS", + # API-compatible cell classes + "AlphaMN", + "INgII", + "INgIb", + "DD", + "DD_Gamma", + "AffIa", + "AffII", + "AffIb", +] diff --git a/myogen/simulator/jaxley/cells/biophysical.py b/myogen/simulator/jaxley/cells/biophysical.py new file mode 100644 index 00000000..b9f93a6c --- /dev/null +++ b/myogen/simulator/jaxley/cells/biophysical.py @@ -0,0 +1,824 @@ +""" +Biophysical Motor Neuron Cell Builder for Jaxley. + +This module provides functions to construct multi-compartment motor neurons +with realistic morphology and biophysics using Jaxley's cell infrastructure. + +The cells include: +- Soma with action potential generating channels +- Dendrites with persistent inward currents (PICs) +- Proper channel distributions based on Powers2017 or NERLab models +""" + +from typing import Dict, List, Literal, Optional, Tuple +import numpy as np +import jax.numpy as jnp +import jaxley as jx +from jaxley.channels import Leak + +# Custom MyoGen channels (motor neuron specific) +from myogen.simulator.jaxley.channels import ( + Na3rp, + Naps, + KdrRL, + MAHP, + Gh, + LCaInact, + LeakChannel, +) + +# Try importing jaxley-mech channels (Hodgkin-Huxley based) +try: + from jaxley_mech.channels.hodgkin52 import ( + Na as JaxleyMechNa, + K as JaxleyMechK, + Leak as JaxleyMechLeak, + ) + JAXLEY_MECH_AVAILABLE = True +except ImportError: + JAXLEY_MECH_AVAILABLE = False + JaxleyMechNa = None + JaxleyMechK = None + JaxleyMechLeak = None + + +# ============================================================================= +# MORPHOLOGY PARAMETERS +# ============================================================================= + +# Powers & Binder (2001) motor neuron morphology +POWERS2017_MORPHOLOGY = { + "soma": { + "length": 25.0, # um + "diameter": 25.0, # um (approximate sphere) + "nseg": 1, + }, + "dendrite": { + "length": 5500.0, # um - equivalent dendrite length + "diameter": 6.0, # um - average diameter + "nseg": 10, # Number of segments + "n_dendrites": 4, # Number of equivalent dendrites + }, +} + +# NERLab simplified morphology +NERLAB_MORPHOLOGY = { + "soma": { + "length": 22.0, + "diameter": 22.0, + "nseg": 1, + }, + "dendrite": { + "length": 4000.0, + "diameter": 5.0, + "nseg": 5, + "n_dendrites": 1, + }, +} + +# ============================================================================= +# CHANNEL CONDUCTANCE PARAMETERS +# ============================================================================= + +POWERS2017_CHANNELS = { + "soma": { + "na3rp": {"gbar": 0.05, "sh": 8.0, "ar": 1.0}, + "naps": {"gbar": 0.001, "sh": 0.0, "ar": 1.0}, + "kdrRL": {"gbar": 0.3, "mVh": -25.0, "mslp": 20.0}, + "mAHP": {"gkcamax": 0.03, "gcamax": 3e-5, "tau_ca": 20.0}, + "gh": {"gbar": 2.5e-5, "half": -70.0, "slp": 8.0}, + "leak": {"g_leak": 5e-5, "e_leak": -70.0}, + }, + "dendrite": { + "L_Ca_inact": {"gbar": 6.7e-5, "theta_m": -30.0, "tau_h": 1500.0}, + "gh": {"gbar": 0.0001, "half": -70.0, "slp": 8.0}, + "leak": {"g_leak": 5e-5, "e_leak": -70.0}, + }, + "axial_resistivity": 70.0, # Ohm-cm + "membrane_capacitance": 1.0, # uF/cm² + "e_na": 60.0, # mV + "e_k": -80.0, # mV + "e_ca": 80.0, # mV +} + +NERLAB_CHANNELS = { + "soma": { + "na3rp": {"gbar": 0.03, "sh": 1.0, "ar": 1.0}, + "naps": {"gbar": 0.0002, "sh": 0.0, "ar": 1.0}, + "kdrRL": {"gbar": 0.015, "mVh": -25.0, "mslp": 20.0}, + "mAHP": {"gkcamax": 0.0005, "gcamax": 5e-5, "tau_ca": 70.0}, + "gh": {"gbar": 2.5e-5, "half": -70.0, "slp": 8.0}, + "leak": {"g_leak": 5e-5, "e_leak": -71.0}, + }, + "dendrite": { + "L_Ca_inact": {"gbar": 1e-5, "theta_m": -30.0, "tau_h": 1500.0}, + "gh": {"gbar": 2.5e-5, "half": -70.0, "slp": 8.0}, + "leak": {"g_leak": 5e-5, "e_leak": -71.0}, + }, + "axial_resistivity": 70.0, + "membrane_capacitance": 1.35546, + "e_na": 60.0, + "e_k": -80.0, + "e_ca": 120.0, +} + + +# ============================================================================= +# BIOPHYSICAL MOTOR NEURON CLASS +# ============================================================================= + +# Channel conductance parameters for jaxley-mech Hodgkin-Huxley channels +# Standard Hodgkin-Huxley 1952 squid giant axon parameters from jaxley-mech documentation +JAXLEY_MECH_CHANNELS = { + "soma": { + "na": {"gNa": 120e-3, "eNa": 50.0}, # S/cm² - fast sodium + "k": {"gK": 36e-3, "eK": -77.0}, # S/cm² - delayed rectifier + "leak": {"gLeak": 0.3e-3, "eLeak": -54.3}, # S/cm² - passive leak + }, + "dendrite": { + "leak": {"gLeak": 0.1e-3, "eLeak": -65.0}, # Passive dendrite with leak only + }, + "axial_resistivity": 70.0, # Ohm-cm + "membrane_capacitance": 1.0, # uF/cm² - standard membrane capacitance +} + + +class BiophysicalMotorNeuron: + """ + Multi-compartment motor neuron with realistic biophysics. + + Constructs a Jaxley cell with: + - Soma containing Na+, K+, Ca2+-dependent K+ (AHP), and H-current + - Dendrites with L-type Ca2+ for persistent inward currents + + Parameters + ---------- + model : str + "Powers2017" or "NERLab" - determines parameters. + n_dendrites : int, optional + Number of dendrites (default from model). + segments_per_dendrite : int, optional + Compartments per dendrite (default from model). + use_multicompartment : bool + If True, create actual multi-compartment cell with dendrites. + If False (default), create single-compartment soma only. + use_jaxley_mech : bool + If True, use jaxley-mech Hodgkin-Huxley channels (Na, K, Leak) + instead of custom motor neuron channels. Default is False. + """ + + def __init__( + self, + model: Literal["Powers2017", "NERLab"] = "Powers2017", + n_dendrites: Optional[int] = None, + segments_per_dendrite: Optional[int] = None, + use_multicompartment: bool = False, + use_jaxley_mech: bool = False, + ): + self.model = model + self.use_multicompartment = use_multicompartment + self.use_jaxley_mech = use_jaxley_mech + + # Validate jaxley-mech availability + if use_jaxley_mech and not JAXLEY_MECH_AVAILABLE: + raise ImportError( + "jaxley-mech is not installed. Install it with: pip install jaxley-mech" + ) + + # Get parameters based on model + if model == "Powers2017": + self.morphology = POWERS2017_MORPHOLOGY.copy() + self.channel_params = POWERS2017_CHANNELS.copy() + else: + self.morphology = NERLAB_MORPHOLOGY.copy() + self.channel_params = NERLAB_CHANNELS.copy() + + # Override if specified + if n_dendrites is not None: + self.morphology["dendrite"]["n_dendrites"] = n_dendrites + if segments_per_dendrite is not None: + self.morphology["dendrite"]["nseg"] = segments_per_dendrite + + # Build the cell + self.cell = None + self.soma = None + self.dendrites = [] + self.soma_branch_idx = 0 + self.dendrite_branch_indices = [] + + self._build_cell() + self._define_morphology() + self._insert_channels() + + def _build_cell(self): + """Create the Jaxley cell structure.""" + if self.use_multicompartment: + self._build_multicompartment_cell() + else: + self._build_single_compartment_cell() + + def _build_single_compartment_cell(self): + """Create a single-compartment (soma only) cell.""" + # Create a single compartment + comp = jx.Compartment() + + # Create a branch with 1 compartment for soma + soma_branch = jx.Branch(comp, ncomp=1) + + # Create cell from branch + self.cell = jx.Cell(soma_branch, parents=[-1]) + self.soma = self.cell.branch(0) + self.soma_branch_idx = 0 + + def _build_multicompartment_cell(self): + """Create a multi-compartment cell with soma and dendrites. + + Uses proper Jaxley topology: soma is root (-1), all dendrites connect to soma (0). + Each dendrite branch has multiple compartments for spatial discretization. + """ + soma_params = self.morphology["soma"] + dend_params = self.morphology["dendrite"] + + n_dendrites = dend_params["n_dendrites"] + nseg_dend = dend_params["nseg"] + + # Create compartment template + comp = jx.Compartment() + + # Create branches: soma (1 comp) + n_dendrites (nseg_dend comps each) + # Build list of branches for Cell constructor + branches = [] + + # Branch 0: Soma (1 compartment) + soma_branch = jx.Branch(comp, ncomp=1) + branches.append(soma_branch) + + # Branches 1 to n_dendrites: Dendrites + for _ in range(n_dendrites): + dend_branch = jx.Branch(comp, ncomp=nseg_dend) + branches.append(dend_branch) + + # Build parent structure: + # - Branch 0 (soma) is root: parent = -1 + # - Branches 1..n (dendrites) connect to soma: parent = 0 + parents = np.array([-1] + [0] * n_dendrites) + + # Create cell with proper topology + self.cell = jx.Cell(branches, parents=parents) + + # Store references + self.soma = self.cell.branch(0) + self.soma_branch_idx = 0 + self.dendrite_branch_indices = list(range(1, n_dendrites + 1)) + self.dendrites = [self.cell.branch(idx) for idx in self.dendrite_branch_indices] + + def _define_morphology(self): + """Set up the cell morphology (geometry).""" + soma_params = self.morphology["soma"] + + # Set global membrane properties + # Use JAXLEY_MECH_CHANNELS params when use_jaxley_mech=True for proper HH tuning + if self.use_jaxley_mech: + self.cell.set("axial_resistivity", JAXLEY_MECH_CHANNELS["axial_resistivity"]) + self.cell.set("capacitance", JAXLEY_MECH_CHANNELS["membrane_capacitance"]) + else: + self.cell.set("axial_resistivity", self.channel_params["axial_resistivity"]) + self.cell.set("capacitance", self.channel_params["membrane_capacitance"]) + + # Soma geometry + self.cell.branch(self.soma_branch_idx).set("radius", soma_params["diameter"] / 2.0) + self.cell.branch(self.soma_branch_idx).set("length", soma_params["length"]) + + # Dendrite geometry (if multi-compartment) + if self.use_multicompartment and self.dendrite_branch_indices: + dend_params = self.morphology["dendrite"] + for idx in self.dendrite_branch_indices: + self.cell.branch(idx).set("radius", dend_params["diameter"] / 2.0) + self.cell.branch(idx).set("length", dend_params["length"]) + + def _insert_channels(self): + """Insert ion channels with appropriate conductances.""" + if self.use_jaxley_mech: + self._insert_soma_channels_jaxley_mech() + if self.use_multicompartment and self.dendrite_branch_indices: + self._insert_dendrite_channels_jaxley_mech() + else: + self._insert_soma_channels() + if self.use_multicompartment and self.dendrite_branch_indices: + self._insert_dendrite_channels() + + def _insert_soma_channels(self): + """Insert channels into soma.""" + soma_channels = self.channel_params["soma"] + soma = self.cell.branch(self.soma_branch_idx) + + # 1. Fast sodium (Na3rp) + na3rp_params = soma_channels.get("na3rp", {}) + soma.insert(Na3rp()) + soma.set("Na3rp_gbar", na3rp_params.get("gbar", 0.05)) + soma.set("Na3rp_sh", na3rp_params.get("sh", 8.0)) + soma.set("Na3rp_ar", na3rp_params.get("ar", 1.0)) + soma.set("Na3rp_eNa", self.channel_params["e_na"]) + + # 2. Persistent sodium (Naps) + naps_params = soma_channels.get("naps", {}) + soma.insert(Naps()) + soma.set("Naps_gbar", naps_params.get("gbar", 0.001)) + soma.set("Naps_sh", naps_params.get("sh", 0.0)) + soma.set("Naps_ar", naps_params.get("ar", 1.0)) + soma.set("Naps_eNa", self.channel_params["e_na"]) + + # 3. Delayed rectifier potassium (KdrRL) + kdr_params = soma_channels.get("kdrRL", {}) + soma.insert(KdrRL()) + soma.set("KdrRL_gbar", kdr_params.get("gbar", 0.3)) + soma.set("KdrRL_mVh", kdr_params.get("mVh", -25.0)) + soma.set("KdrRL_mslp", kdr_params.get("mslp", 20.0)) + soma.set("KdrRL_eK", self.channel_params["e_k"]) + + # 4. Calcium-dependent potassium / mAHP (MAHP) + mahp_params = soma_channels.get("mAHP", {}) + soma.insert(MAHP()) + soma.set("MAHP_gkcamax", mahp_params.get("gkcamax", 0.03)) + soma.set("MAHP_gcamax", mahp_params.get("gcamax", 3e-5)) + soma.set("MAHP_tau_ca", mahp_params.get("tau_ca", 20.0)) + soma.set("MAHP_eK", self.channel_params["e_k"]) + soma.set("MAHP_eCa", self.channel_params["e_ca"]) + + # 5. H-current (Gh) + gh_params = soma_channels.get("gh", {}) + soma.insert(Gh()) + soma.set("Gh_gbar", gh_params.get("gbar", 2.5e-5)) + soma.set("Gh_half", gh_params.get("half", -70.0)) + soma.set("Gh_slp", gh_params.get("slp", 8.0)) + + # 6. Leak + leak_params = soma_channels.get("leak", {}) + soma.insert(LeakChannel()) + soma.set("LeakChannel_gLeak", leak_params.get("g_leak", 5e-5)) + soma.set("LeakChannel_eLeak", leak_params.get("e_leak", -70.0)) + + def _insert_dendrite_channels(self): + """Insert channels into dendrites.""" + dend_channels = self.channel_params["dendrite"] + + for idx in self.dendrite_branch_indices: + dend = self.cell.branch(idx) + + # 1. L-type calcium (for PICs) + lca_params = dend_channels.get("L_Ca_inact", {}) + dend.insert(LCaInact()) + dend.set("LCaInact_gbar", lca_params.get("gbar", 6.7e-5)) + dend.set("LCaInact_theta_m", lca_params.get("theta_m", -30.0)) + dend.set("LCaInact_tau_h", lca_params.get("tau_h", 1500.0)) + dend.set("LCaInact_eCa", self.channel_params["e_ca"]) + + # 2. H-current + gh_params = dend_channels.get("gh", {}) + dend.insert(Gh()) + dend.set("Gh_gbar", gh_params.get("gbar", 0.0001)) + dend.set("Gh_half", gh_params.get("half", -70.0)) + dend.set("Gh_slp", gh_params.get("slp", 8.0)) + + # 3. Leak + leak_params = dend_channels.get("leak", {}) + dend.insert(LeakChannel()) + dend.set("LeakChannel_gLeak", leak_params.get("g_leak", 5e-5)) + dend.set("LeakChannel_eLeak", leak_params.get("e_leak", -70.0)) + + def _insert_soma_channels_jaxley_mech(self): + """Insert hybrid channels into soma: Na3rp + jaxley-mech K + Leak + MAHP. + + This hybrid approach combines: + - Na3rp (custom): Fast sodium with slow inactivation (more realistic than HH Na) + - K (jaxley-mech): Standard delayed rectifier potassium + - Leak (jaxley-mech): Passive leak + - MAHP (custom): Calcium-dependent K+ for AHP (limits firing rate) + + This provides the best of both worlds: + - Na3rp's slow inactivation prevents excessive high-frequency firing + - MAHP provides medium AHP for physiological firing rates (~10-40 Hz) + - jaxley-mech K provides reliable repolarization + """ + jm_params = JAXLEY_MECH_CHANNELS["soma"] + soma = self.cell.branch(self.soma_branch_idx) + + # 1. Fast sodium with slow inactivation (Na3rp - custom channel) + # Na3rp has three gating variables: m³hs + # - m: fast activation + # - h: fast inactivation + # - s: slow inactivation (Fleidervish et al.) - KEY DIFFERENCE from HH Na + na = Na3rp(name="Na3rp") + soma.insert(na) + # Na3rp needs ~1.5x HH gNa to compensate for slow inactivation (s gate) + # and MAHP's hyperpolarizing current + soma.set("Na3rp_gbar", jm_params["na"]["gNa"] * 1.5) # ~0.18 S/cm² + soma.set("Na3rp_sh", 0.0) # mV - threshold shift (0 = easier to fire) + soma.set("Na3rp_ar", 1.0) # slow inact recovery (1=minimal slow inact) + soma.set("Na3rp_eNa", jm_params["na"]["eNa"]) + soma.set("Na3rp_thinf", -50.0) # mV - inactivation half-voltage + soma.set("Na3rp_qinf", 4.0) # mV - inactivation slope + + # 2. Delayed rectifier potassium (Hodgkin-Huxley K from jaxley-mech) + k = JaxleyMechK(name="K") + soma.insert(k) + soma.set("K_gK", jm_params["k"]["gK"]) + soma.set("K_eK", jm_params["k"]["eK"]) + + # 3. Leak (Hodgkin-Huxley Leak from jaxley-mech) + leak = JaxleyMechLeak(name="Leak") + soma.insert(leak) + soma.set("Leak_gLeak", jm_params["leak"]["gLeak"]) + soma.set("Leak_eLeak", jm_params["leak"]["eLeak"]) + + # 4. MAHP (calcium-dependent K+ for AHP) to limit firing rate + # Provides medium after-hyperpolarization that limits firing to ~10-40 Hz + mahp = MAHP(name="MAHP") + soma.insert(mahp) + soma.set("MAHP_gkcamax", 0.03) # S/cm² - KCa conductance (AHP) + soma.set("MAHP_gcamax", 3e-5) # S/cm² - Ca conductance (small, for dynamics) + soma.set("MAHP_tau_ca", 20.0) # ms - calcium removal time constant + soma.set("MAHP_eK", jm_params["k"]["eK"]) # Use same K reversal as K channel + soma.set("MAHP_eCa", 120.0) # mV - calcium reversal potential + + def _insert_dendrite_channels_jaxley_mech(self): + """Insert jaxley-mech channels into dendrites. + + For the jaxley-mech mode, dendrites only have passive leak channels + (no PICs or active conductances). + """ + jm_params = JAXLEY_MECH_CHANNELS["dendrite"] + + for idx in self.dendrite_branch_indices: + dend = self.cell.branch(idx) + + # Passive leak only for dendrites + leak = JaxleyMechLeak(name="Leak") + dend.insert(leak) + dend.set("Leak_gLeak", jm_params["leak"]["gLeak"]) + dend.set("Leak_eLeak", jm_params["leak"]["eLeak"]) + + def get_cell(self) -> jx.Cell: + """Return the constructed Jaxley cell.""" + return self.cell + + def set_initial_voltage(self, v: float = -70.0): + """Set initial membrane voltage.""" + self.cell.set("v", v) + + def get_channel_list(self) -> List[str]: + """Return list of inserted channel names.""" + if self.use_jaxley_mech: + # Hybrid channels: Na3rp (custom) + jaxley-mech K/Leak + MAHP (custom) + soma_channels = ["Na3rp", "K", "Leak", "MAHP"] + if self.use_multicompartment: + dend_channels = ["Leak"] + return {"soma": soma_channels, "dendrite": dend_channels} + return soma_channels + else: + # Custom motor neuron channels + soma_channels = ["Na3rp", "Naps", "KdrRL", "MAHP", "Gh", "Leak"] + if self.use_multicompartment: + dend_channels = ["LCaInact", "Gh", "Leak"] + return {"soma": soma_channels, "dendrite": dend_channels} + return soma_channels + + def setup_recording(self, location: str = "soma") -> None: + """ + Set up voltage recording at specified location. + + Parameters + ---------- + location : str + "soma" or "dendrite" + """ + self.cell.delete_recordings() + + if location == "soma": + self.cell.branch(self.soma_branch_idx).loc(0.5).record("v") + elif location == "dendrite" and self.dendrite_branch_indices: + # Record from first dendrite + self.cell.branch(self.dendrite_branch_indices[0]).loc(0.5).record("v") + + def setup_stimulus(self, current, location: str = "soma") -> None: + """ + Set up current injection at specified location. + + Parameters + ---------- + current : array + Current waveform (nA). + location : str + "soma" or "dendrite" + """ + self.cell.delete_stimuli() + + if location == "soma": + self.cell.branch(self.soma_branch_idx).loc(0.5).stimulate(current) + elif location == "dendrite" and self.dendrite_branch_indices: + self.cell.branch(self.dendrite_branch_indices[0]).loc(0.5).stimulate(current) + + def simulate( + self, + current=None, + dt: float = 0.025, + t_max: float = 100.0, + record_location: str = "soma", + stim_location: str = "soma", + ): + """ + Run simulation with optional current injection. + + Parameters + ---------- + current : array, optional + Current waveform (nA). + dt : float + Time step (ms). + t_max : float + Total time (ms). + record_location : str + Where to record voltage. + stim_location : str + Where to inject current. + + Returns + ------- + voltages : array + Recorded voltages. + """ + self.setup_recording(record_location) + + if current is not None: + self.setup_stimulus(current, stim_location) + + return jx.integrate(self.cell, delta_t=dt, t_max=t_max) + + +# ============================================================================= +# BIOPHYSICAL INTERNEURON CLASS +# ============================================================================= + +class BiophysicalInterneuron: + """ + Single-compartment interneuron with active channels. + + Based on Bui et al. (2003) spinal interneuron model. + + Parameters + ---------- + cell_type : str + "INgII" or "INgIb" - interneuron type. + """ + + # Interneuron channel parameters (from NEURON version) + INTERNEURON_CHANNELS = { + "na3rp": {"gbar": 0.003, "sh": 1.0, "ar": 1.0}, + "kdrRL": {"gbar": 0.015, "mVh": -25.0, "mslp": 20.0}, + "mAHP": {"gkcamax": 0.0005, "gcamax": 3e-6, "tau_ca": 70.0}, + "gh": {"gbar": 2.5e-5, "half": -70.0, "slp": 8.0}, + "leak": {"g_leak": 5e-5, "e_leak": -71.0}, + } + + def __init__(self, cell_type: Literal["INgII", "INgIb"] = "INgII"): + self.cell_type = cell_type + self.cell = None + self.soma_branch_idx = 0 + + self._build_cell() + self._define_morphology() + self._insert_channels() + + def _build_cell(self): + """Create single-compartment Jaxley cell.""" + comp = jx.Compartment() + soma_branch = jx.Branch(comp, ncomp=1) + self.cell = jx.Cell(soma_branch, parents=[-1]) + + def _define_morphology(self): + """Set interneuron morphology from Bui et al. (2003).""" + import numpy as np + + # Membrane areas from Bui et al. + Amu = 81390 + 3113 # um² + Aci = 1.96 * (891.5 + 46.141) / np.sqrt(8) + A = Amu - Aci # Use lower bound + + # Equivalent sphere diameter + D = np.sqrt(A / np.pi) + + self.cell.branch(self.soma_branch_idx).set("radius", D / 2.0) + self.cell.branch(self.soma_branch_idx).set("length", D) + self.cell.set("axial_resistivity", 70.0) + self.cell.set("capacitance", 1.0) + + def _insert_channels(self): + """Insert interneuron channels.""" + params = self.INTERNEURON_CHANNELS + soma = self.cell.branch(self.soma_branch_idx) + + # Fast sodium + soma.insert(Na3rp()) + soma.set("Na3rp_gbar", params["na3rp"]["gbar"]) + soma.set("Na3rp_sh", params["na3rp"]["sh"]) + soma.set("Na3rp_ar", params["na3rp"]["ar"]) + soma.set("Na3rp_eNa", 60.0) + + # Delayed rectifier + soma.insert(KdrRL()) + soma.set("KdrRL_gbar", params["kdrRL"]["gbar"]) + soma.set("KdrRL_mVh", params["kdrRL"]["mVh"]) + soma.set("KdrRL_mslp", params["kdrRL"]["mslp"]) + soma.set("KdrRL_eK", -80.0) + + # mAHP + soma.insert(MAHP()) + soma.set("MAHP_gkcamax", params["mAHP"]["gkcamax"]) + soma.set("MAHP_gcamax", params["mAHP"]["gcamax"]) + soma.set("MAHP_tau_ca", params["mAHP"]["tau_ca"]) + soma.set("MAHP_eK", -80.0) + soma.set("MAHP_eCa", 120.0) + + # H-current + soma.insert(Gh()) + soma.set("Gh_gbar", params["gh"]["gbar"]) + soma.set("Gh_half", params["gh"]["half"]) + soma.set("Gh_slp", params["gh"]["slp"]) + + # Leak + soma.insert(LeakChannel()) + soma.set("LeakChannel_gLeak", params["leak"]["g_leak"]) + soma.set("LeakChannel_eLeak", params["leak"]["e_leak"]) + + def get_cell(self) -> jx.Cell: + """Return the constructed Jaxley cell.""" + return self.cell + + def setup_recording(self) -> None: + """Set up voltage recording at soma.""" + self.cell.delete_recordings() + self.cell.branch(self.soma_branch_idx).loc(0.5).record("v") + + def setup_stimulus(self, current) -> None: + """Set up current injection at soma.""" + self.cell.delete_stimuli() + self.cell.branch(self.soma_branch_idx).loc(0.5).stimulate(current) + + def simulate(self, current=None, dt: float = 0.025, t_max: float = 100.0): + """ + Run simulation with optional current injection. + + Parameters + ---------- + current : array, optional + Current waveform (nA). + dt : float + Time step (ms). + t_max : float + Total time (ms). + + Returns + ------- + voltages : array + Recorded voltages. + """ + self.setup_recording() + + if current is not None: + self.setup_stimulus(current) + + return jx.integrate(self.cell, delta_t=dt, t_max=t_max) + + +# ============================================================================= +# FACTORY FUNCTIONS +# ============================================================================= + +def create_motor_neuron( + model: Literal["Powers2017", "NERLab"] = "Powers2017", + n_dendrites: Optional[int] = None, + initial_voltage: float = -70.0, + use_multicompartment: bool = False, + use_jaxley_mech: bool = False, +) -> jx.Cell: + """ + Factory function to create a biophysical motor neuron. + + Parameters + ---------- + model : str + "Powers2017" or "NERLab". + n_dendrites : int, optional + Number of dendrites. + initial_voltage : float + Initial membrane voltage (mV). + use_multicompartment : bool + If True, create multi-compartment cell with dendrites. + use_jaxley_mech : bool + If True, use jaxley-mech Hodgkin-Huxley channels (Na, K, Leak) + instead of custom motor neuron channels. Default is False. + + Returns + ------- + jx.Cell + Configured Jaxley cell. + """ + builder = BiophysicalMotorNeuron( + model=model, + n_dendrites=n_dendrites, + use_multicompartment=use_multicompartment, + use_jaxley_mech=use_jaxley_mech, + ) + builder.set_initial_voltage(initial_voltage) + return builder.get_cell() + + +def create_motor_neuron_builder( + model: Literal["Powers2017", "NERLab"] = "Powers2017", + n_dendrites: Optional[int] = None, + initial_voltage: float = -70.0, + use_multicompartment: bool = False, + use_jaxley_mech: bool = False, +) -> BiophysicalMotorNeuron: + """ + Factory function to create a biophysical motor neuron builder. + + Returns the builder object so you have access to simulate(), setup_recording(), etc. + + Parameters + ---------- + model : str + "Powers2017" or "NERLab". + n_dendrites : int, optional + Number of dendrites. + initial_voltage : float + Initial membrane voltage (mV). + use_multicompartment : bool + If True, create multi-compartment cell with dendrites. + use_jaxley_mech : bool + If True, use jaxley-mech Hodgkin-Huxley channels (Na, K, Leak) + instead of custom motor neuron channels. Default is False. + + Returns + ------- + BiophysicalMotorNeuron + The builder object with simulation methods. + """ + builder = BiophysicalMotorNeuron( + model=model, + n_dendrites=n_dendrites, + use_multicompartment=use_multicompartment, + use_jaxley_mech=use_jaxley_mech, + ) + builder.set_initial_voltage(initial_voltage) + return builder + + +def create_interneuron( + cell_type: Literal["INgII", "INgIb"] = "INgII", + initial_voltage: float = -70.0, +) -> jx.Cell: + """ + Factory function to create a biophysical interneuron. + + Parameters + ---------- + cell_type : str + "INgII" or "INgIb". + initial_voltage : float + Initial membrane voltage (mV). + + Returns + ------- + jx.Cell + Configured Jaxley cell. + """ + builder = BiophysicalInterneuron(cell_type=cell_type) + builder.cell.set("v", initial_voltage) + return builder.get_cell() + + +def create_interneuron_builder( + cell_type: Literal["INgII", "INgIb"] = "INgII", + initial_voltage: float = -70.0, +) -> BiophysicalInterneuron: + """ + Factory function to create a biophysical interneuron builder. + + Returns the builder object so you have access to simulate(), setup_recording(), etc. + + Parameters + ---------- + cell_type : str + "INgII" or "INgIb". + initial_voltage : float + Initial membrane voltage (mV). + + Returns + ------- + BiophysicalInterneuron + The builder object with simulation methods. + """ + builder = BiophysicalInterneuron(cell_type=cell_type) + builder.cell.set("v", initial_voltage) + return builder diff --git a/myogen/simulator/jaxley/cells_api.py b/myogen/simulator/jaxley/cells_api.py new file mode 100644 index 00000000..d48a67dd --- /dev/null +++ b/myogen/simulator/jaxley/cells_api.py @@ -0,0 +1,1138 @@ +""" +Neural Cell Models for Spinal Cord Simulations - Jaxley Backend (API Layer). + +This module provides Jaxley-based implementations maintaining full API +compatibility with the NEURON version. This is the API-compatible layer. + +For biophysical implementations with actual channel insertion, see: +- myogen.simulator.jaxley.cells.biophysical +""" + +import itertools +from typing import Literal, Optional + +import numpy as np +import quantities as pq +import jaxley as jx +from jaxley.channels import Leak + +import myogen +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__m, Quantity__m_per_s, Quantity__ms, Quantity__mV + + +# ============================================================================ +# STOCHASTIC SPIKE GENERATORS +# ============================================================================ + +class _PoissonProcessGenerator__Jaxley: + """ + Poisson process spike generator for Jaxley cells. + + Uses the same algorithm as the Cython implementation: accumulates input + intensity and compares against an exponentially-distributed threshold. + + Parameters + ---------- + seed : int + Random seed for reproducibility + N : int + Batch size for threshold generation (affects statistical properties) + dt : float + Time step in milliseconds + """ + + def __init__(self, seed, N, dt): + self.N = N # Batch size for threshold (not max rate!) + self.dt = dt # Time step in ms + self.rng = np.random.RandomState(seed) + + # State variables matching Cython implementation + self.yi = 0.0 # Accumulated input intensity + self.thres = self._generate_threshold() # Exponential threshold + + def _generate_threshold(self): + """Generate exponential threshold using batch method.""" + aux = 1.0 + for _ in range(self.N): + aux *= self.rng.random() + return -(1.0 / self.N) * np.log(aux) + + def compute(self, y): + """ + Generate Poisson spike for given input rate. + + Parameters + ---------- + y : float + Input intensity (rate) in Hz at current time step + + Returns + ------- + int + 1 if spike occurs, 0 otherwise + """ + spk = 0 + # Integrate input intensity: y is rate in Hz, dt is in ms + self.yi += y * self.dt * 1e-3 + + # Check if threshold crossed + if self.yi >= self.thres: + spk = 1 + self.yi = 0.0 # Reset accumulator + self.thres = self._generate_threshold() # New threshold + + return spk + + +class _GammaProcessGenerator__Jaxley: + """Gamma process spike generator for Jaxley cells.""" + + def __init__(self, seed, timestep__ms, shape): + self.timestep__ms = timestep__ms + self.shape = shape + self.rng = np.random.RandomState(seed) + self._t_next = 0.0 + self._t = 0.0 + + def compute(self, rate__Hz): + """Generate Gamma spike (0 or 1).""" + self._t += float(self.timestep__ms.magnitude if hasattr(self.timestep__ms, 'magnitude') else self.timestep__ms) + + if self._t >= self._t_next and rate__Hz > 0: + mean_isi = 1000.0 / rate__Hz # Convert Hz to ms + scale = mean_isi / self.shape + self._t_next = self._t + self.rng.gamma(self.shape, scale) + return 1 + return 0 + + +# ============================================================================ +# BASE CELL CLASS +# ============================================================================ + +class _Cell: + """ + Base class for all neural cell models (Jaxley backend). + + Maintains full API compatibility with NEURON version while using Jaxley internally. + """ + + _gid__iterator = itertools.count(0) + + @beartowertype + def __init__(self, class__ID: int, pool__ID: int | None = None): + self.global__ID = next(self._gid__iterator) + self.class__ID = class__ID + self.pool__ID = pool__ID + + self._create_sections() + self._build_topology() + self._define_geometry() + self._define_biophysics() + + self.create_axon() + self.synapse__list = [] + + self.axon_length__m: Quantity__m | None = None + self.conduction_velocity__m_per_s: Quantity__m_per_s | None = None + self.axon_delay__ms: Quantity__ms | None = None + + # For spike detection + self._prev_voltage = -70.0 + self.spike_threshold = 0.0 + + @beartowertype + def __repr__(self) -> str: + return f"{self.__class__.__name__} [global ID: {self.global__ID}, class ID: {self.class__ID}, pool ID: {self.pool__ID if self.pool__ID is not None else 'N/A'}]" + + @beartowertype + def _create_sections(self): + """Create Jaxley cell morphology.""" + pass + + @beartowertype + def _build_topology(self): + """Connect sections (for multi-compartment cells).""" + pass + + @beartowertype + def _define_geometry(self): + """Set geometric parameters using Jaxley API.""" + pass + + @beartowertype + def _define_biophysics(self): + """Insert channels and set biophysical parameters.""" + pass + + @beartowertype + def create_axon( + self, + length__m: Quantity__m = 0 * pq.m, + conduction_velocity__m_per_s: Quantity__m_per_s = 50 * pq.m / pq.s, + ): + """Define axonal conduction properties (API-compatible with NEURON).""" + self.axon_length__m = length__m + self.conduction_velocity__m_per_s = conduction_velocity__m_per_s + self.axon_delay__ms = (self.axon_length__m / self.conduction_velocity__m_per_s).rescale(pq.ms) + + def create_synapses( + self, + synapse_location, + reversal_potential__mV: Quantity__mV = 0 * pq.mV, + rise_time_constant__ms: Quantity__ms = 0.2 * pq.ms, + decay_time_constant__ms: Quantity__ms = 0.3 * pq.ms, + ) -> dict: + """Create a synapse specification (API-compatible with NEURON). + + Returns a dict ``{tau1, tau2, e, location}`` stored in ``self.synapse__list``. + This is a connectivity specification, **not** a live Jaxley synapse object. + Actual synaptic currents in simulation examples are computed analytically + (exponential conductance kernel × driving force) and injected via + ``jx.integrate()``. Making these real Jaxley ``Synapse`` objects would require + a monolithic ``jx.Network`` — incompatible with per-cell integration. + """ + assert decay_time_constant__ms > rise_time_constant__ms, ( + "decay_time_constant__ms must be greater than tau1" + ) + + # Create synapse specification + syn = { + "location": synapse_location, + "tau1": float(rise_time_constant__ms.magnitude if hasattr(rise_time_constant__ms, 'magnitude') else rise_time_constant__ms), + "tau2": float(decay_time_constant__ms.magnitude if hasattr(decay_time_constant__ms, 'magnitude') else decay_time_constant__ms), + "e": float(reversal_potential__mV.magnitude if hasattr(reversal_potential__mV, 'magnitude') else reversal_potential__mV), + } + + # NERLab model default + if hasattr(self, "model") and self.model == "NERLab" and reversal_potential__mV == 0 * pq.mV: + syn["e"] = 70 + + self.synapse__list.append(syn) + return syn + + def get_voltage(self) -> float: + """Get current membrane voltage.""" + if hasattr(self, 'cell') and hasattr(self.cell, 'states'): + return float(self.cell.states.get('v', -70.0)) + return -70.0 + + def apply_synaptic_input(self, weight: float, synapse_type: str = "excitatory"): + """Apply synaptic input (placeholder for actual Jaxley synapse activation).""" + # This would activate the synapse in actual Jaxley simulation + pass + + +# ============================================================================ +# INTERNEURONS +# ============================================================================ + +@beartowertype +class INgII(_Cell): + """ + Group II interneuron - Jaxley implementation. + + Single-compartment spinal interneuron with active ion channels. + Based on Bui et al. (2003). + """ + + _ids2 = itertools.count(0) + + def __init__(self, class__ID: Optional[int] = None, pool__ID: int | None = None): + self.soma_branch_idx = 0 + super().__init__(class__ID if class__ID is not None else next(self._ids2), pool__ID) + self.create_synapses(self.cell) + + def _create_sections(self): + """Create Jaxley cell with proper Branch structure.""" + comp = jx.Compartment() + soma_branch = jx.Branch(comp, ncomp=1) + self.cell = jx.Cell(soma_branch, parents=[-1]) + self.soma = self.cell.branch(self.soma_branch_idx) + + def _define_geometry(self): + """Set geometry from Bui et al. (2003).""" + Amu = 81390 + 3113 + Aci = 1.96 * (891.5 + 46.141) / np.sqrt(8) + i = 0 + A = [Amu - Aci, Amu + Aci] + D = np.sqrt(A[i] / np.pi) + + self.cell.branch(self.soma_branch_idx).set("radius", D / 2.0) + self.cell.branch(self.soma_branch_idx).set("length", D) + self.cell.set("axial_resistivity", 70.0) + self.cell.set("capacitance", 1.0) + + def _define_biophysics(self): + """Insert biophysical channels. Default conductances from Bui et al. (2003); + overridden per-cell by the population pool via .set().""" + from myogen.simulator.jaxley.channels import Na3rp, KdrRL, MAHP, Gh + + soma = self.cell.branch(self.soma_branch_idx) + soma.insert(Na3rp()) + soma.insert(KdrRL()) + soma.insert(MAHP()) + soma.insert(Gh()) + soma.insert(Leak()) + + # Fresh view required after inserts (Jaxley stale view bug) + soma = self.cell.branch(self.soma_branch_idx) + + # Na3rp (fast sodium) + soma.set("Na3rp_gbar", 0.003) + soma.set("Na3rp_sh", 1.0) + soma.set("Na3rp_qinf", 8.0) + soma.set("Na3rp_thinf", -50.0) + soma.set("Na3rp_ar", 1.0) + soma.set("Na3rp_eNa", 55.0) + + # KdrRL (delayed rectifier potassium) + soma.set("KdrRL_gbar", 0.015) + soma.set("KdrRL_mVh", -21.0) + soma.set("KdrRL_tmin", 0.8) + soma.set("KdrRL_taumax", 20.0) + soma.set("KdrRL_eK", -80.0) + + # MAHP (medium afterhyperpolarization) + soma.set("MAHP_gcamax", 3e-6) + soma.set("MAHP_gkcamax", 5e-4) + soma.set("MAHP_tau_ca", 70.0) + soma.set("MAHP_mvhalfca", -22.0) + soma.set("MAHP_mtauca", 2.0) + soma.set("MAHP_eK", -80.0) + soma.set("MAHP_eCa", 120.0) + + # Gh (H-current / HCN) + soma.set("Gh_gbar", 2.5e-5) + soma.set("Gh_half", -77.0) + soma.set("Gh_htau", 30.0) + soma.set("Gh_eH", -41.0) + + # Leak (passive) + soma.set("Leak_gLeak", 5e-5) + soma.set("Leak_eLeak", -71.0) + + def setup_recording(self) -> None: + """Set up voltage recording at soma.""" + self.cell.delete_recordings() + self.cell.branch(self.soma_branch_idx).loc(0.5).record("v") + + def setup_stimulus(self, current) -> None: + """Set up current injection at soma.""" + self.cell.delete_stimuli() + self.cell.branch(self.soma_branch_idx).loc(0.5).stimulate(current) + + def simulate(self, current=None, dt: float = 0.025, t_max: float = 100.0): + """Run simulation with optional current injection.""" + self.setup_recording() + if current is not None: + self.setup_stimulus(current) + return jx.integrate(self.cell, delta_t=dt, t_max=t_max) + + +@beartowertype +class INgIb(INgII): + """Golgi tendon organ (Ib) interneuron.""" + + _ids2 = itertools.count(0) + + def __init__(self, pool__ID: int | None = None): + super().__init__(next(self._ids2), pool__ID) + + +# ============================================================================ +# DESCENDING DRIVE +# ============================================================================ + +@beartowertype +class DD(_Cell, _PoissonProcessGenerator__Jaxley): + """Descending drive with Poisson spike generation.""" + + _ids2 = itertools.count(0) + + def __init__(self, N, dt, pool__ID: int | None = None): + class_id = next(self._ids2) + self.soma_branch_idx = 0 + _Cell.__init__(self, class_id, pool__ID) + _PoissonProcessGenerator__Jaxley.__init__( + self, myogen.SEED + (self.class__ID + 1) * (self.global__ID + 1), N, dt + ) + + def __repr__(self) -> str: + return _Cell.__repr__(self) + + def _create_sections(self): + """Create simple Jaxley cell with proper Branch structure.""" + comp = jx.Compartment() + soma_branch = jx.Branch(comp, ncomp=1) + self.cell = jx.Cell(soma_branch, parents=[-1]) + self.ns = self.cell # Compatibility alias + + def _define_geometry(self): + """Simple point neuron.""" + self.cell.branch(self.soma_branch_idx).set("radius", 5.0) + self.cell.branch(self.soma_branch_idx).set("length", 10.0) + self.cell.set("axial_resistivity", 70.0) + self.cell.set("capacitance", 1.0) + + def _define_biophysics(self): + """Passive only.""" + soma = self.cell.branch(self.soma_branch_idx) + soma.insert(Leak()) + self.cell.set("Leak_gLeak", 1e-4) + self.cell.set("Leak_eLeak", -65.0) + + def integrate(self, y): + """Integrate drive signal to generate spikes.""" + return self.compute(y) if y > 0 else 0 + + +@beartowertype +class DD_Gamma(_Cell, _GammaProcessGenerator__Jaxley): + """Descending drive with Gamma process spike generation.""" + + _ids2 = itertools.count(0) + + def __init__(self, timestep__ms: Quantity__ms, shape: float = 3.0, pool__ID: int | None = None): + class_id = next(self._ids2) + self.soma_branch_idx = 0 + _Cell.__init__(self, class_id, pool__ID) + _GammaProcessGenerator__Jaxley.__init__( + self, myogen.SEED + (self.class__ID + 1) * (self.global__ID + 1), timestep__ms, shape + ) + + def __repr__(self) -> str: + return _Cell.__repr__(self) + + def _create_sections(self): + """Create simple Jaxley cell with proper Branch structure.""" + comp = jx.Compartment() + soma_branch = jx.Branch(comp, ncomp=1) + self.cell = jx.Cell(soma_branch, parents=[-1]) + self.ns = self.cell # Compatibility alias + + def _define_geometry(self): + """Simple point neuron.""" + self.cell.branch(self.soma_branch_idx).set("radius", 5.0) + self.cell.branch(self.soma_branch_idx).set("length", 10.0) + self.cell.set("axial_resistivity", 70.0) + self.cell.set("capacitance", 1.0) + + def _define_biophysics(self): + """Passive only.""" + soma = self.cell.branch(self.soma_branch_idx) + soma.insert(Leak()) + self.cell.set("Leak_gLeak", 1e-4) + self.cell.set("Leak_eLeak", -65.0) + + def integrate(self, rate__Hz): + """Integrate firing rate to generate spikes.""" + return self.compute(rate__Hz) if rate__Hz > 0 else 0 + + +# ============================================================================ +# AFFERENTS +# ============================================================================ + +@beartowertype +class AffIa(_Cell, _GammaProcessGenerator__Jaxley): + """Ia afferent from muscle spindles.""" + + _ids2 = itertools.count(0) + + def __init__( + self, + RT: float, + N: int, + timestep__ms: Quantity__ms, + initN: int = 0, + class__ID: int | None = None, + pool__ID: int | None = None, + ): + class_id = next(self._ids2) if class__ID is None else class__ID + self.soma_branch_idx = 0 + _Cell.__init__(self, class_id, pool__ID) + + # Store recruitment threshold and individual firing rate variability + self.RT = RT + self.IFR = myogen.RANDOM_GENERATOR.normal(5, 2.5) # Individual firing rate variability + + # Shape parameter for Gamma process (controls ISI regularity) + shape = N # Use N as shape parameter + _GammaProcessGenerator__Jaxley.__init__( + self, myogen.SEED + (self.class__ID + 1) * (self.global__ID + 1), timestep__ms, shape + ) + + def __repr__(self) -> str: + return _Cell.__repr__(self) + + def _create_sections(self): + """Create simple Jaxley cell with proper Branch structure.""" + comp = jx.Compartment() + soma_branch = jx.Branch(comp, ncomp=1) + self.cell = jx.Cell(soma_branch, parents=[-1]) + self.ns = self.cell + + def _define_geometry(self): + """Afferent fiber geometry.""" + self.cell.branch(self.soma_branch_idx).set("radius", 5.0) + self.cell.branch(self.soma_branch_idx).set("length", 10.0) + self.cell.set("axial_resistivity", 70.0) + self.cell.set("capacitance", 1.0) + + def _define_biophysics(self): + """Passive properties.""" + soma = self.cell.branch(self.soma_branch_idx) + soma.insert(Leak()) + self.cell.set("Leak_gLeak", 1e-4) + self.cell.set("Leak_eLeak", -70.0) + + def integrate(self, rate__Hz): + """Integrate spindle activity to generate spikes.""" + return self.compute(rate__Hz) if rate__Hz > 0 else 0 + + +@beartowertype +class AffII(AffIa): + """Group II afferent from muscle spindles.""" + + _ids2 = itertools.count(0) + + def __init__( + self, + RT: float, + N: int, + timestep__ms: Quantity__ms, + initN: int = 0, + class__ID: int | None = None, + pool__ID: int | None = None, + ): + super().__init__(RT, N, timestep__ms, initN, class__ID, pool__ID) + + +@beartowertype +class AffIb(AffIa): + """Ib afferent from Golgi tendon organs.""" + + _ids2 = itertools.count(0) + + def __init__( + self, + RT: float, + N: int, + timestep__ms: Quantity__ms, + initN: int = 0, + class__ID: int | None = None, + pool__ID: int | None = None, + ): + super().__init__(RT, N, timestep__ms, initN, class__ID, pool__ID) + + +# ============================================================================ +# MOTOR NEURON +# ============================================================================ + +# Try importing jaxley-mech channels (Hodgkin-Huxley based) +# Note: In jaxley-mech 0.1.0, the module is 'hh', in 0.3.1+ it's 'hodgkin52' +JAXLEY_MECH_AVAILABLE = False +JaxleyMechNa = None +JaxleyMechK = None +JaxleyMechLeak = None + +try: + # Try newer module path first (jaxley-mech >= 0.3.0) + from jaxley_mech.channels.hodgkin52 import ( + Na as JaxleyMechNa, + K as JaxleyMechK, + Leak as JaxleyMechLeak, + ) + JAXLEY_MECH_AVAILABLE = True +except ImportError: + try: + # Try older module path (jaxley-mech 0.1.0 - 0.2.0) + from jaxley_mech.channels.hh import ( + Na as JaxleyMechNa, + K as JaxleyMechK, + Leak as JaxleyMechLeak, + ) + JAXLEY_MECH_AVAILABLE = True + except ImportError: + pass + + +@beartowertype +class AlphaMN(_Cell): + """ + Alpha motor neuron - Jaxley implementation. + + Multi-compartment motor neuron with active channels and PICs. + For the full biophysical version with actual channel insertion, use: + from myogen.simulator.jaxley.cells import BiophysicalMotorNeuron + + Implements the Henneman size principle: soma size scales with recruitment + threshold, creating physiological heterogeneity: + - Low threshold → small soma → high Rin → low rheobase + - High threshold → large soma → low Rin → high rheobase + + Parameters + ---------- + segments__count : int + Number of segments (compartments). + mode : str + "active" for ion channels, "passive" for leak only. + dendrites__count : int + Number of dendrites (currently unused in single-compartment mode). + model : str + "NERLab" or "Powers2017" for model-specific parameters. + use_jaxley_mech : bool + If True, use jaxley-mech Hodgkin-Huxley channels (Na, K, Leak) + instead of custom motor neuron channels. Default is False. + class__ID : int, optional + Class ID for the cell. + pool__ID : int, optional + Pool ID for the cell. + recruitment_threshold : float, optional + Recruitment threshold for this motor neuron (0-1 range). + Used to scale soma size according to the size principle. + If None, uses a default medium-sized soma. + """ + + _ids2 = itertools.count(0) + + def __init__( + self, + segments__count: int = 1, + mode: Literal["active", "passive"] = "active", + dendrites__count: int = 4, + model: Literal["NERLab", "Powers2017"] = "NERLab", + use_jaxley_mech: bool = False, + gamma: float = 0.2, + class__ID: Optional[int] = None, + pool__ID: int | None = None, + recruitment_threshold: Optional[float] = None, + neuron_params: Optional[dict] = None, + ): + self.dendrites__count = dendrites__count + self.segments__count = segments__count + self.mode = mode + self.model = model + self.use_jaxley_mech = use_jaxley_mech + self.gamma = gamma # dendritic Ca scale factor (NERLab caL_gama) + self.soma_branch_idx = 0 + self.recruitment_threshold = recruitment_threshold + self.neuron_params = neuron_params # Pre-computed parameters from pool + + # Note: use_jaxley_mech=True no longer requires jaxley-mech package. + # It uses our custom channels (Na3rp, KdrRL, etc.) with a 2-compartment architecture. + + super().__init__(class__ID if class__ID is not None else next(self._ids2), pool__ID) + self.create_synapses(self.cell) + + # Motor neuron spike detection threshold (0 mV crossing) + self.spike_threshold = 0.0 + + def _create_sections(self): + """Create cell morphology. + + Dispatch by ``self.model``: + - ``"NERLab"`` : soma (ncomp=1) + 1 isopotential dendrite (ncomp=1) + — matches the production NEURON architecture. + - ``"Powers2017"`` + ``use_jaxley_mech=True``: soma + 1 cable + dendrite (ncomp=4) — cable filtering of PIC. + - ``"Powers2017"`` + ``use_jaxley_mech=False``: soma only + (single-compartment legacy path). + """ + comp = jx.Compartment() + if self.model == "NERLab": + soma_branch = jx.Branch(comp, ncomp=1) + dend_branch = jx.Branch(comp, ncomp=1) # isopotential, matches NEURON + self.cell = jx.Cell( + [soma_branch, dend_branch], + parents=np.array([-1, 0]), + ) + self.soma = self.cell.branch(0) + self.dendrite_branch_indices = [1] + self.dend_ncomp = 1 + self.dend = None + elif self.use_jaxley_mech: + soma_branch = jx.Branch(comp, ncomp=1) + dend_branch = jx.Branch(comp, ncomp=4) # Single cable, 4 compartments + self.cell = jx.Cell( + [soma_branch, dend_branch], + parents=np.array([-1, 0]), + ) + self.soma = self.cell.branch(0) + self.dendrite_branch_indices = [1] # Single dendrite cable + self.dend_ncomp = 4 + self.dend = None + else: + soma_branch = jx.Branch(comp, ncomp=1) + self.cell = jx.Cell(soma_branch, parents=[-1]) + self.soma = self.cell.branch(self.soma_branch_idx) + self.dendrite_branch_indices = [] + self.dend = None + + def _define_geometry(self): + """Set motor neuron geometry with recruitment-threshold-dependent scaling. + + Implements the Henneman size principle: + - Low recruitment threshold → small soma → high Rin → low rheobase + - High recruitment threshold → large soma → low Rin → high rheobase + + When use_jaxley_mech=True: sets geometry for soma (branch 0) AND 4 dendrites (branches 1-4). + When use_jaxley_mech=False: sets geometry for soma only. + + When neuron_params is provided (from pool), uses pre-computed parameters. + Otherwise, calculates parameters based on recruitment_threshold. + """ + # NERLab path: soma is sphere-like (L = diam), uses its own param keys. + if self.model == "NERLab": + self._define_geometry_nerlab() + return + + # Use pre-computed parameters from pool if available + if self.neuron_params is not None: + effective_diameter = self.neuron_params["soma_diameter"] + effective_length = self.neuron_params["soma_length"] + capacitance = self.neuron_params["capacitance"] + axial_resistivity = self.neuron_params["axial_resistivity"] + else: + # Fallback: calculate from recruitment_threshold (for standalone cells) + if self.recruitment_threshold is not None: + scale_factor = np.sqrt(min(self.recruitment_threshold, 1.0)) + else: + scale_factor = 0.5 # Default to medium-sized neuron + + if self.use_jaxley_mech: + # Powers2017 soma geometry + effective_diameter = 22.0 + scale_factor * (30.0 - 22.0) + effective_length = 2952.0 + scale_factor * (3665.0 - 2952.0) + capacitance = 1.356 + scale_factor * (1.879 - 1.356) + axial_resistivity = 0.001 # isopotential soma + else: + # Custom channels - scale around NEURON geometry + effective_diameter = 15.0 + scale_factor * (30.0 - 15.0) + effective_length = 2000.0 + scale_factor * (4000.0 - 2000.0) + capacitance = 1.35546 # uF/cm² - matches NEURON soma.cm + axial_resistivity = 0.001 # ohm-cm - matches NEURON soma.Ra + + # Apply soma geometry + soma = self.cell.branch(self.soma_branch_idx) + soma.set("radius", effective_diameter / 2.0) + soma.set("length", effective_length) + soma.set("capacitance", capacitance) + soma.set("axial_resistivity", axial_resistivity) + + # Set dendrite geometry (single cable with ncomp=4) + # Total area = 4× NEURON single dend (d×2, L×2). Cable Ra provides filtering. + # Jaxley .set("length", L) sets EACH compartment to L, so divide total by ncomp. + if self.use_jaxley_mech and self.dendrite_branch_indices: + dend_idx = self.dendrite_branch_indices[0] + dend = self.cell.branch(dend_idx) + ncomp = self.dend_ncomp + if self.neuron_params is not None: + dend.set("radius", self.neuron_params["dend_diameter"] / 2.0) + dend.set("length", self.neuron_params["dend_length"] / ncomp) + dend.set("capacitance", self.neuron_params["dend_cm"]) + dend.set("axial_resistivity", self.neuron_params["dend_Ra"]) + else: + # Fallback dendrite geometry (2× NEURON, 4× area cable) + if self.recruitment_threshold is not None: + sf = np.sqrt(min(self.recruitment_threshold, 1.0)) + else: + sf = 0.5 + total_d = 17.46 + sf * (23.82 - 17.46) + total_L = 3588.0 + sf * (4454.0 - 3588.0) + dend.set("radius", total_d / 2.0) + dend.set("length", total_L / ncomp) + dend.set("capacitance", 0.868 + sf * (0.880 - 0.868)) + dend.set("axial_resistivity", 51.04 - sf * (51.04 - 40.76)) + + def _define_biophysics(self): + """Insert channels (active or passive based on mode). + + NERLab path is wholly separate (napp/caL bundle their own passive leak), + so it doesn't share the Leak()-channel scaffolding used by Powers2017. + """ + if self.model == "NERLab": + self._insert_nerlab_channels() + return + + soma = self.cell.branch(self.soma_branch_idx) + + # Insert Leak on soma + soma.insert(Leak()) + + if self.use_jaxley_mech and self.dendrite_branch_indices: + # Insert Leak on each dendrite (params set in _insert_jaxley_mech_channels) + for dend_idx in self.dendrite_branch_indices: + self.cell.branch(dend_idx).insert(Leak()) + else: + # Single compartment defaults + self.cell.set("Leak_gLeak", 5e-5) + self.cell.set("Leak_eLeak", -70.0) + + if self.mode == "passive": + return + + # For active mode, insert biophysical channels + if self.mode == "active": + if self.use_jaxley_mech: + self._insert_jaxley_mech_channels() + else: + self._insert_active_channels() + + # ---------------------------------------------------------------- NERLab + + def _define_geometry_nerlab(self): + """Set geometry for the NERLab architecture (soma sphere + 1 dendrite cylinder). + + Soma : L = diam (sphere-like, matches NEURON ``cell.soma.L = diam``). + Dend : single isopotential compartment (ncomp=1). + + Uses per-cell ``neuron_params`` from the pool when available; otherwise + falls back to mid-range NERLab defaults so a standalone cell still works. + """ + if self.neuron_params is not None: + p = self.neuron_params + soma_diam = float(p["soma_diameter"]) + dend_diam = float(p["dend_diameter"]) + dend_len = float(p["dend_length"]) + soma_cm = float(p["soma_capacitance"]) + soma_Ra = float(p["soma_axial_resistivity"]) + dend_cm = float(p["dend_capacitance"]) + dend_Ra = float(p["dend_axial_resistivity"]) + else: + # Mid-range NERLab fallback (matches scripts/test_nerlab_single_cell.py). + soma_diam, dend_diam, dend_len = 95.0, 70.0, 8000.0 + soma_cm = dend_cm = 1.0 + soma_Ra = dend_Ra = 70.0 + + soma = self.cell.branch(self.soma_branch_idx) + soma.set("length", soma_diam) + soma.set("radius", soma_diam / 2.0) + soma.set("capacitance", soma_cm) + soma.set("axial_resistivity", soma_Ra) + + dend = self.cell.branch(self.dendrite_branch_indices[0]) + dend.set("length", dend_len) + dend.set("radius", dend_diam / 2.0) + dend.set("capacitance", dend_cm) + dend.set("axial_resistivity", dend_Ra) + + def _insert_nerlab_channels(self): + """Insert NERLab channels (napp on soma + caL on dendrite). + + NERLab cells live in the 1952 HH voltage frame; cells using this path + must initialise V at 0 mV (the pool does this automatically when + ``model="NERLab"``). + """ + from myogen.simulator.jaxley.channels import napp, caL + + soma = self.cell.branch(self.soma_branch_idx) + soma.insert(napp()) + + dend = self.cell.branch(self.dendrite_branch_indices[0]) + dend.insert(caL()) + + # Jaxley stale-view bug (see MEMORY.md): branch views go stale after + # .insert(); re-fetch fresh views before any .set() calls. + soma = self.cell.branch(self.soma_branch_idx) + dend = self.cell.branch(self.dendrite_branch_indices[0]) + + # Resolve per-cell params (pool path) or fall back to NERLAB_PARAMS midpoints. + p = self.neuron_params or {} + def pick(key, fallback): + return float(p[key]) if key in p else fallback + + # --- soma napp --- + soma.set("napp_gnabar", pick("gnabar", 0.055)) + soma.set("napp_gnapbar", pick("gnapbar", 0.00055)) + soma.set("napp_gkfbar", pick("gkfbar", 0.0022)) + soma.set("napp_gksbar", pick("gksbar", 0.018)) + soma.set("napp_gl", pick("gls", 0.00125)) + soma.set("napp_rinact", pick("rinact", 0.04)) + soma.set("napp_el", pick("el_napp", 0.0)) + soma.set("napp_vtraub", pick("vtraub_napp", 0.0)) + soma.set("napp_ena", pick("ena", 120.0)) + soma.set("napp_ek", pick("ek", -10.0)) + + # --- dendrite caL --- + dend.set("caL_gcaLbar", pick("gcaLbar", 9e-6)) + dend.set("caL_gl", pick("gl_caL", 1.2e-4)) + dend.set("caL_el", pick("el_caL", 0.0)) + dend.set("caL_ecaL", pick("ecaL", 140.0)) + dend.set("caL_vtraub", pick("vtraub_caL", 34.5)) + dend.set("caL_Ltau", pick("ltau_caL", 65.0)) + # gama scales dendritic Ca conductance. NEURON's NERLab uses gamma=0.2 + # (config/alpha_mn_default.yaml); a hard-coded 1.0 was making the PIC + # ~5× too strong, contributing to excess firing. + dend.set("caL_gama", self.gamma) + + def _insert_jaxley_mech_channels(self): + """Insert Powers2017 channels on soma + 4 dendrites. + + Soma (branch 0): Na3rp, Naps, KdrRL, MAHP, Gh, Leak + Dendrites (branches 1-4): LCaInact (dendritic PIC), Gh, Leak + + Single multi-compartment dendrite cable (ncomp=4) with cable-filtered PIC. + Cable Ra provides temporal filtering of PIC dynamics in place of NEURON's discrete + 4-isopotential-dendrite topology. + + When neuron_params is provided (from pool), uses pre-computed conductances. + Otherwise, calculates conductances based on recruitment_threshold. + """ + from myogen.simulator.jaxley.channels import ( + Na3rp, Naps, KdrRL, MAHP, Gh, LCaInact, + ) + + soma = self.cell.branch(self.soma_branch_idx) + + # Get parameters (from pool or fallback) + if self.neuron_params is not None: + p = self.neuron_params + else: + # Fallback: midpoint parameters for standalone cell + if self.recruitment_threshold is not None: + sf = np.sqrt(min(self.recruitment_threshold, 1.0)) + else: + sf = 0.5 + p = { + "gNa3rp": 0.01 + sf * (0.022 - 0.01), + "gNaPS": 2.6e-5 - sf * (2.6e-5 - 2.0e-5), + "gKdrRL": 0.015 + sf * (0.02 - 0.015), + "gMAHP_ca": 6.4e-6 + sf * (1.015e-5 - 6.4e-6), + "gMAHP_k": 5.0625e-4 + sf * (6.75e-4 - 5.0625e-4), + "tauMAHP": 90.0 - sf * (90.0 - 30.0), + "gGh_soma": 3.0e-5 + sf * (2.3e-4 - 3.0e-5), + "gLeak_soma": 1.5e-4 + sf * (3.77e-4 - 1.5e-4), + "eLeak_soma": -71.0 - sf * 1.0, + # Dendrite: NEURON native density (4× area from cable geometry) + "gLeak_dend": 7.93e-5 + sf * (1.75e-4 - 7.93e-5), + "eLeak_dend": -71.0 - sf * 1.0, + "gGh_dend": 3.0e-5 + sf * (2.3e-4 - 3.0e-5), + "gLCa_dend": 7.21875e-5 + sf * (9.69375e-5 - 7.21875e-5), + "theta_m_LCa": -42.0 + sf * (-39.0 - (-42.0)), + "theta_h_LCa": 10.0 + sf * (-10.0 - 10.0), + "eNa": 55.0, "eK": -80.0, "eCa": 120.0, + } + + # === INSERT ALL CHANNELS FIRST === + # (branch views become stale after insert, so insert all then get fresh views) + + # Soma channels + soma.insert(Na3rp()) + soma.insert(Naps()) + soma.insert(KdrRL()) + soma.insert(MAHP()) + soma.insert(Gh()) + + # Dendrite channels (single cable — inserts on all ncomp compartments) + dend_idx = self.dendrite_branch_indices[0] + self.cell.branch(dend_idx).insert(LCaInact()) + self.cell.branch(dend_idx).insert(Gh()) + + # === SET PARAMETERS (fresh views after all inserts) === + soma = self.cell.branch(self.soma_branch_idx) + + # Na3rp (fast sodium with persistent component) + soma.set("Na3rp_gbar", p["gNa3rp"]) + soma.set("Na3rp_sh", 1.0) + soma.set("Na3rp_qinf", 8.0) + soma.set("Na3rp_thinf", -50.0) + soma.set("Na3rp_ar", 1.0) + soma.set("Na3rp_eNa", p["eNa"]) + + # Naps (persistent sodium) + soma.set("Naps_gbar", p["gNaPS"]) + soma.set("Naps_sh", 5.0) + soma.set("Naps_vslope", 5.0) + soma.set("Naps_ar", 1.0) + soma.set("Naps_asvh", -90.0) # NEURON value + soma.set("Naps_bsvh", -22.0) # NEURON value + soma.set("Naps_eNa", p["eNa"]) + + # KdrRL (delayed rectifier potassium) + soma.set("KdrRL_gbar", p["gKdrRL"]) + soma.set("KdrRL_mVh", -21.0) + soma.set("KdrRL_tmin", 0.8) + soma.set("KdrRL_taumax", 20.0) + soma.set("KdrRL_eK", p["eK"]) + + # MAHP (medium afterhyperpolarization) + soma.set("MAHP_gcamax", p["gMAHP_ca"]) + soma.set("MAHP_gkcamax", p["gMAHP_k"]) + soma.set("MAHP_tau_ca", p["tauMAHP"]) + soma.set("MAHP_mvhalfca", -22.0) # NEURON global override (h.mvhalfca_mAHP = -22) + soma.set("MAHP_mtauca", 2.0) + soma.set("MAHP_eK", p["eK"]) + soma.set("MAHP_eCa", p["eCa"]) + + # Gh — soma + soma.set("Gh_gbar", p["gGh_soma"]) + soma.set("Gh_half", -77.0) + soma.set("Gh_htau", 30.0) + soma.set("Gh_eH", -41.0) + + # Leak — soma (already inserted in _define_biophysics) + soma.set("Leak_gLeak", p["gLeak_soma"]) + soma.set("Leak_eLeak", p["eLeak_soma"]) + + # === DENDRITE PARAMETERS (single cable, all 4 compartments uniform) === + # Cable Ra provides temporal filtering — no theta_m offsets needed. + dend = self.cell.branch(self.dendrite_branch_indices[0]) + + # LCaInact (L-type calcium with inactivation — THE dendritic PIC) + dend.set("LCaInact_gbar", p["gLCa_dend"]) + dend.set("LCaInact_theta_m", p["theta_m_LCa"]) + dend.set("LCaInact_theta_h", p["theta_h_LCa"]) + dend.set("LCaInact_tau_m", 40.0) + dend.set("LCaInact_tau_h", 800.0) # Faster than NEURON (2500); steep kappa_h provides selectivity + dend.set("LCaInact_kappa_h", 3.0) # Steeper than NEURON (5) for voltage-selective inactivation at plateau + dend.set("LCaInact_kappa_m", -6.0) + dend.set("LCaInact_eCa", p["eCa"]) + + # Gh — dendrite + dend.set("Gh_gbar", p["gGh_dend"]) + dend.set("Gh_half", -77.0) + dend.set("Gh_htau", 30.0) + dend.set("Gh_eH", -41.0) + + # Leak — dendrite (already inserted in _define_biophysics) + dend.set("Leak_gLeak", p["gLeak_dend"]) + dend.set("Leak_eLeak", p["eLeak_dend"]) + + def _insert_active_channels(self): + """Insert active ion channels for biophysical simulation.""" + from myogen.simulator.jaxley.channels import ( + Na3rp, Naps, KdrRL, MAHP, Gh + ) + + soma = self.cell.branch(self.soma_branch_idx) + + # Insert channels with default parameters + soma.insert(Na3rp()) + soma.insert(Naps()) + soma.insert(KdrRL()) + soma.insert(MAHP()) + soma.insert(Gh()) + + # Set model-specific parameters using cell.set() + # Parameters tuned for large effective cell (~628,000 µm² membrane) + # These are scaled appropriately for 10-20 nA input currents + if self.model == "NERLab": + # Na3rp parameters - Fast sodium for action potentials + # MATCHED TO NEURON Powers2017: gbar=0.01, sh=1.0, qinf=8.0 + self.cell.set("Na3rp_gbar", 0.01) # S/cm² - matches NEURON + self.cell.set("Na3rp_sh", 1.0) # mV - threshold shift (matches NEURON) + self.cell.set("Na3rp_ar", 1.0) # slow inact recovery (1 = no slow inact) + self.cell.set("Na3rp_eNa", 55.0) # mV + self.cell.set("Na3rp_thinf", -50.0) # mV - matches NEURON thinf_na3rp + self.cell.set("Na3rp_qinf", 8.0) # mV - CRITICAL: matches NEURON qinf_na3rp + + # Naps parameters - Persistent sodium (small) + self.cell.set("Naps_gbar", 2.6e-05) # S/cm² - matches NEURON + self.cell.set("Naps_sh", 5.0) # mV - matches NEURON + self.cell.set("Naps_ar", 1.0) + self.cell.set("Naps_eNa", 55.0) + + # KdrRL parameters - Delayed rectifier for repolarization + # MATCHED TO NEURON Powers2017: gMax=0.015 + self.cell.set("KdrRL_gbar", 0.015) # S/cm² - matches NEURON + self.cell.set("KdrRL_mVh", -21.0) # mV - matches NEURON global h.mVh_kdrRL + self.cell.set("KdrRL_mslp", 20.0) # mV + self.cell.set("KdrRL_eK", -80.0) # mV + self.cell.set("KdrRL_tmin", 0.8) # ms - matches NEURON global h.tmin_kdrRL + self.cell.set("KdrRL_taumax", 20.0) # ms - matches NEURON global h.taumax_kdrRL + + # MAHP parameters - Medium AHP (provides spike adaptation) + # MATCHED TO NEURON Powers2017 + self.cell.set("MAHP_gkcamax", 0.00045) # S/cm² - matches NEURON + self.cell.set("MAHP_gcamax", 6.4e-6) # S/cm² - matches NEURON + self.cell.set("MAHP_mvhalfca", -22.0) # mV - matches NEURON global + self.cell.set("MAHP_mslpca", 4.0) + self.cell.set("MAHP_tau_ca", 90.0) # ms - matches NEURON + self.cell.set("MAHP_eK", -80.0) + self.cell.set("MAHP_eCa", 120.0) + + # Gh parameters - H-current for sag + self.cell.set("Gh_gbar", 3e-5) # S/cm² - matches NEURON + self.cell.set("Gh_half", -77.0) # mV - matches NEURON + self.cell.set("Gh_slp", 8.0) # mV + self.cell.set("Gh_eH", -41.0) # mV + self.cell.set("Gh_htau", 30.0) # ms - matches NEURON global + + else: # Powers2017 + # Na3rp parameters - MATCHED TO NEURON Powers2017 + self.cell.set("Na3rp_gbar", 0.01) # S/cm² - matches NEURON + self.cell.set("Na3rp_sh", 1.0) # mV - matches NEURON + self.cell.set("Na3rp_ar", 1.0) + self.cell.set("Na3rp_eNa", 55.0) + self.cell.set("Na3rp_thinf", -50.0) # mV - matches NEURON + self.cell.set("Na3rp_qinf", 8.0) # mV - CRITICAL: matches NEURON + + # Naps parameters - MATCHED TO NEURON Powers2017 + self.cell.set("Naps_gbar", 2.6e-05) # S/cm² - matches NEURON + self.cell.set("Naps_sh", 5.0) # mV - matches NEURON + self.cell.set("Naps_ar", 1.0) + self.cell.set("Naps_eNa", 55.0) + + # KdrRL parameters - MATCHED TO NEURON Powers2017 + self.cell.set("KdrRL_gbar", 0.015) # S/cm² - matches NEURON + self.cell.set("KdrRL_mVh", -21.0) # mV - matches NEURON global + self.cell.set("KdrRL_mslp", 20.0) + self.cell.set("KdrRL_eK", -80.0) + self.cell.set("KdrRL_tmin", 0.8) # ms - matches NEURON global h.tmin_kdrRL + self.cell.set("KdrRL_taumax", 20.0) # ms - matches NEURON global h.taumax_kdrRL + + # MAHP parameters - MATCHED TO NEURON Powers2017 + self.cell.set("MAHP_gkcamax", 0.00045) # S/cm² - matches NEURON + self.cell.set("MAHP_gcamax", 6.4e-6) # S/cm² - matches NEURON + self.cell.set("MAHP_mvhalfca", -22.0) # mV - matches NEURON global + self.cell.set("MAHP_mslpca", 4.0) + self.cell.set("MAHP_tau_ca", 90.0) # ms - matches NEURON + self.cell.set("MAHP_eK", -80.0) + self.cell.set("MAHP_eCa", 120.0) + + # Gh parameters - MATCHED TO NEURON Powers2017 + self.cell.set("Gh_gbar", 3e-5) # S/cm² - matches NEURON + self.cell.set("Gh_half", -77.0) # mV - matches NEURON + self.cell.set("Gh_slp", 8.0) + self.cell.set("Gh_eH", -41.0) + self.cell.set("Gh_htau", 30.0) # ms - matches NEURON global + + def setup_recording(self) -> None: + """Set up voltage recording at soma.""" + self.cell.delete_recordings() + self.cell.branch(self.soma_branch_idx).loc(0.5).record("v") + + def setup_stimulus(self, current) -> None: + """Set up current injection at soma.""" + self.cell.delete_stimuli() + self.cell.branch(self.soma_branch_idx).loc(0.5).stimulate(current) + + def simulate(self, current=None, dt: float = 0.025, t_max: float = 100.0): + """ + Run simulation with optional current injection. + + Parameters + ---------- + current : array, optional + Current waveform (nA). + dt : float + Time step (ms). + t_max : float + Total time (ms). + + Returns + ------- + voltages : array + Recorded voltages. + """ + self.setup_recording() + + if current is not None: + self.setup_stimulus(current) + + return jx.integrate(self.cell, delta_t=dt, t_max=t_max) diff --git a/myogen/simulator/jaxley/channels/__init__.py b/myogen/simulator/jaxley/channels/__init__.py new file mode 100644 index 00000000..e99f58d3 --- /dev/null +++ b/myogen/simulator/jaxley/channels/__init__.py @@ -0,0 +1,43 @@ +"""Custom ion channels for Jaxley simulations.""" + +from myogen.simulator.jaxley.channels.constant import Constant +from myogen.simulator.jaxley.channels.motor_neuron_channels import ( + Na3rp, + Naps, + KdrRL, + MAHP, + Gh, + LCaInact, + LeakChannel, + get_motor_neuron_channels_soma, + get_motor_neuron_channels_dendrite, + safe_exp, + vtrap, + trap0, +) +from myogen.simulator.jaxley.channels.nerlab_channels import ( + napp, + caL, +) + +__all__ = [ + # Utility + "Constant", + # Powers2017 motor-neuron channels + "Na3rp", + "Naps", + "KdrRL", + "MAHP", + "Gh", + "LCaInact", + "LeakChannel", + # NERLab motor-neuron channels (1952 HH voltage convention; V_rest ≈ 0 mV) + "napp", + "caL", + # Helper functions + "get_motor_neuron_channels_soma", + "get_motor_neuron_channels_dendrite", + "safe_exp", + "vtrap", + "trap0", +] diff --git a/myogen/simulator/jaxley/channels/constant.py b/myogen/simulator/jaxley/channels/constant.py new file mode 100644 index 00000000..1e261a38 --- /dev/null +++ b/myogen/simulator/jaxley/channels/constant.py @@ -0,0 +1,144 @@ +""" +Constant current channel for Jaxley. + +Converted from constant.mod: + NEURON { + SUFFIX Constant + NONSPECIFIC_CURRENT i + RANGE i, ic + } + PARAMETER { + ic = 0.000 (mA/cm2) + } + BREAKPOINT { + i = ic + } + +This is the simplest possible channel - just returns a constant current. +Useful for testing and validation. +""" + +import jax.numpy as jnp +from jaxley.channels import Channel +from typing import Dict, Any + + +class Constant(Channel): + """ + Constant current channel. + + Provides a simple non-specific current that doesn't change over time. + Useful for baseline current injection or testing. + + Parameters + ---------- + ic : float, optional + Constant current density in mA/cm². Default is 0.0. + name : str, optional + Name prefix for parameters. If None, uses class name. + """ + + def __init__(self, ic: float = 0.0, name: str | None = None): + # Required for Jaxley 0.5.0+: declare current units + self.current_is_in_mA_per_cm2 = True + + super().__init__(name) + + # Ensure prefix is never None + prefix = self._name if self._name is not None else "Constant" + + # Channel parameters - matching NMODL PARAMETER block + self.channel_params = { + f"{prefix}_ic": ic, # Constant current (mA/cm²) + } + + # No state variables - current is constant + self.channel_states = {} + + def update_states( + self, + states: Dict[str, Any], + dt: float, + v: float, + params: Dict[str, Any], + ) -> Dict[str, Any]: + """ + Update channel states. + + Since this is a constant current, there are no states to update. + + Parameters + ---------- + states : dict + Current channel states (empty for this channel). + dt : float + Time step in milliseconds. + v : float + Membrane voltage in mV. + params : dict + Channel parameters. + + Returns + ------- + dict + Updated states (empty for this channel). + """ + return {} + + def compute_current( + self, + states: Dict[str, Any], + v: float, + params: Dict[str, Any] + ) -> float: + """ + Calculate ionic current. + + Simply returns the constant current parameter. + + Parameters + ---------- + states : dict + Current channel states (empty for this channel). + v : float + Membrane voltage in mV. + params : dict + Channel parameters containing 'Constant_ic'. + + Returns + ------- + float + Current in mA/cm². + """ + prefix = self._name if self._name is not None else "Constant" + return params[f"{prefix}_ic"] + + def init_state( + self, + states: Dict[str, Any], + v: float, + params: Dict[str, Any], + delta_t: float + ) -> Dict[str, Any]: + """ + Initialize channel states. + + Since this channel has no states, returns empty dict. + + Parameters + ---------- + states : dict + Initial states (empty). + v : float + Initial membrane voltage in mV. + params : dict + Channel parameters. + delta_t : float + Time step in milliseconds. + + Returns + ------- + dict + Initialized states (empty). + """ + return {} diff --git a/myogen/simulator/jaxley/channels/motor_neuron_channels.py b/myogen/simulator/jaxley/channels/motor_neuron_channels.py new file mode 100644 index 00000000..41b74441 --- /dev/null +++ b/myogen/simulator/jaxley/channels/motor_neuron_channels.py @@ -0,0 +1,929 @@ +""" +Custom Ion Channel Models for Motor Neuron Simulations - Jaxley Backend. + +This module provides JAX-compatible implementations of the ion channels used +in the Powers2017 motor neuron model, originally implemented in NMODL for NEURON. + +Channels implemented: +- Na3rp: Fast sodium with slow inactivation (action potentials) +- Naps: Persistent sodium (bistability, PICs) +- KdrRL: Delayed rectifier potassium (repolarization) +- MAHP: Calcium-dependent potassium with calcium dynamics (mAHP) +- Gh: Hyperpolarization-activated cation current (sag) +- LCaInact: L-type calcium with inactivation (dendritic PICs) + +All channels follow the Jaxley Channel API and can be inserted into Jaxley cells. + +References: +- Powers et al. (2017): Motor neuron ion channel models +- Booth et al. (1997): L-type calcium channel kinetics +- Fleidervish et al.: Slow sodium inactivation +""" + +from typing import Dict, Optional, Tuple +import jax.numpy as jnp +from jax import lax +from jaxley.channels import Channel +from jaxley.solver_gate import exponential_euler, save_exp + + +# ============================================================================= +# HELPER FUNCTIONS +# ============================================================================= + +def safe_exp(x: jnp.ndarray, clip_val: float = 100.0) -> jnp.ndarray: + """Exponential with clipping to prevent overflow.""" + return jnp.exp(jnp.clip(x, -clip_val, clip_val)) + + +def vtrap(x: jnp.ndarray, y: float) -> jnp.ndarray: + """ + Trapping function for rate equations to avoid division by zero. + + When |x/y| < 1e-6: returns y * (1 - x/y/2) + Otherwise: returns x / (exp(x/y) - 1) + """ + ratio = x / y + return jnp.where( + jnp.abs(ratio) < 1e-6, + y * (1.0 - ratio / 2.0), + x / (save_exp(ratio) - 1.0) + ) + + +def trap0(v: jnp.ndarray, th: float, a: float, q: float) -> jnp.ndarray: + """ + Rate equation trap function from na3rp.mod. + + Parameters + ---------- + v : membrane potential (mV) + th : threshold voltage (mV) + a : rate constant + q : slope factor (mV) + """ + diff = v - th + # Use safe_exp with clipping to prevent overflow, matches NEURON's exp behavior + # save_exp from jaxley may have different clipping behavior + exp_term = safe_exp(-diff / q) + return jnp.where( + jnp.abs(diff) > 1e-6, + a * diff / (1.0 - exp_term), + a * q + ) + + +# ============================================================================= +# NA3RP - FAST SODIUM CHANNEL WITH SLOW INACTIVATION +# ============================================================================= + +class Na3rp(Channel): + """ + Fast sodium channel with slow inactivation (na3rp). + + Three-state gating: m³hs where: + - m: fast activation + - h: fast inactivation + - s: slow inactivation (Fleidervish et al.) + + Parameters from Magee/Migliore model with Powers modifications. + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gbar": 0.010, # S/cm² + f"{prefix}_sh": 8.0, # mV - threshold shift + f"{prefix}_ar": 1.0, # slow inact recovery (1=none, 0=max) + f"{prefix}_eNa": 60.0, # Sodium reversal potential + f"{prefix}_thinf": -50.0, # mV - inactivation inf half-voltage (NEURON: thinf_na3rp) + f"{prefix}_qinf": 4.0, # mV - inactivation inf slope (NEURON: qinf_na3rp) + } + self.channel_states = { + f"{prefix}_m": 0.0, + f"{prefix}_h": 1.0, + f"{prefix}_s": 1.0, + } + self.current_name = f"i_Na3rp" + + # Rate parameters from NMODL + self.tha = -30.0 # mV - activation half-voltage + self.qa = 7.2 # mV - activation slope + self.Ra = 0.4 # /ms - activation rate + self.Rb = 0.124 # /ms - deactivation rate + + self.thi1 = -45.0 # mV - inactivation half-voltage + self.thi2 = -45.0 # mV + self.qd = 1.5 # mV - inactivation slope + self.qg = 1.5 # mV + self.Rd = 0.03 # /ms - inactivation rate + self.Rg = 0.01 # /ms - recovery rate + + # Slow inactivation parameters + self.a0s = 0.001 # /ms + self.b0s = 0.0034 # /ms + self.asvh = -85.0 # mV + self.bsvh = -17.0 # mV + self.avs = 30.0 # mV + self.bvs = 10.0 # mV + + self.mmin = 0.02 # ms - minimum tau_m + self.hmin = 0.5 # ms - minimum tau_h + self.q10 = 2.0 + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Update gating variables using exponential Euler.""" + prefix = self._name + m = states[f"{prefix}_m"] + h = states[f"{prefix}_h"] + s = states[f"{prefix}_s"] + + sh = params[f"{prefix}_sh"] + ar = params[f"{prefix}_ar"] + thinf = params[f"{prefix}_thinf"] + qinf = params[f"{prefix}_qinf"] + + celsius = 37.0 + qt = self.q10 ** ((celsius - 24.0) / 10.0) + + # m (activation) + a_m = trap0(v, self.tha + sh, self.Ra, self.qa) + b_m = trap0(-v, -self.tha - sh, self.Rb, self.qa) + tau_m = jnp.maximum(1.0 / (a_m + b_m) / qt, self.mmin) + m_inf = a_m / (a_m + b_m) + + # h (fast inactivation) - uses thinf and qinf from params + a_h = trap0(v, self.thi1 + sh, self.Rd, self.qd) + b_h = trap0(-v, -self.thi2 - sh, self.Rg, self.qg) + tau_h = jnp.maximum(1.0 / (a_h + b_h) / qt, self.hmin) + h_inf = 1.0 / (1.0 + save_exp((v - thinf - sh) / qinf)) + + # s (slow inactivation) + alps = self.a0s * save_exp((self.asvh - v) / self.avs) + bets = self.b0s / (save_exp((self.bsvh - v) / self.bvs) + 1.0) + tau_s = 1.0 / (alps + bets) + c = alps * tau_s + s_inf = c + ar * (1.0 - c) + + # Exponential Euler integration + m_new = exponential_euler(m, dt, m_inf, tau_m) + h_new = exponential_euler(h, dt, h_inf, tau_h) + s_new = exponential_euler(s, dt, s_inf, tau_s) + + return { + f"{prefix}_m": m_new, + f"{prefix}_h": h_new, + f"{prefix}_s": s_new, + } + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """Compute sodium current: I = gbar * m³ * h * s * (v - E_Na).""" + prefix = self._name + m = states[f"{prefix}_m"] + h = states[f"{prefix}_h"] + s = states[f"{prefix}_s"] + gbar = params[f"{prefix}_gbar"] + e_na = params[f"{prefix}_eNa"] + + g = gbar * (m ** 3) * h * s + return g * (v - e_na) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """Initialize states to steady-state values.""" + prefix = self._name + sh = params[f"{prefix}_sh"] + ar = params[f"{prefix}_ar"] + thinf = params[f"{prefix}_thinf"] + qinf = params[f"{prefix}_qinf"] + + # m (activation) + a_m = trap0(v, self.tha + sh, self.Ra, self.qa) + b_m = trap0(-v, -self.tha - sh, self.Rb, self.qa) + m_inf = a_m / (a_m + b_m) + + # h (fast inactivation) - uses thinf and qinf from params + h_inf = 1.0 / (1.0 + save_exp((v - thinf - sh) / qinf)) + + # s (slow inactivation) + alps = self.a0s * save_exp((self.asvh - v) / self.avs) + bets = self.b0s / (save_exp((self.bsvh - v) / self.bvs) + 1.0) + tau_s = 1.0 / (alps + bets) + c = alps * tau_s + s_inf = c + ar * (1.0 - c) + + return { + f"{prefix}_m": m_inf, + f"{prefix}_h": h_inf, + f"{prefix}_s": s_inf, + } + + +# ============================================================================= +# NAPS - PERSISTENT SODIUM CHANNEL +# ============================================================================= + +class Naps(Channel): + """ + Persistent sodium channel (naps). + + Two-state gating: ms where: + - m: activation (no inactivation at normal voltages) + - s: slow inactivation + + Important for subthreshold amplification and bistability. + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gbar": 0.0052085, # S/cm² + f"{prefix}_sh": 0.0, # mV - threshold shift + f"{prefix}_ar": 1.0, # slow inact (1=none) + f"{prefix}_vslope": 6.8, # mV - activation slope + f"{prefix}_eNa": 60.0, # Sodium reversal potential + f"{prefix}_asvh": -85.0, # mV - slow inact half-act voltage + f"{prefix}_bsvh": -17.0, # mV - slow inact half-deact voltage + } + self.channel_states = { + f"{prefix}_m": 0.0, + f"{prefix}_s": 1.0, + } + self.current_name = f"i_Naps" + + self.mtau = 1.0 # ms - fixed activation time constant + + # Slow inactivation parameters (same as na3rp) + self.a0s = 0.001 + self.b0s = 0.0034 + self.asvh = -85.0 + self.bsvh = -17.0 + self.avs = 30.0 + self.bvs = 10.0 + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Update gating variables.""" + prefix = self._name + m = states[f"{prefix}_m"] + s = states[f"{prefix}_s"] + + sh = params[f"{prefix}_sh"] + ar = params[f"{prefix}_ar"] + vslope = params[f"{prefix}_vslope"] + asvh = params[f"{prefix}_asvh"] + bsvh = params[f"{prefix}_bsvh"] + + # m (activation) - simple Boltzmann + m_inf = 1.0 / (1.0 + save_exp(-(v + 52.3 - sh) / vslope)) + tau_m = self.mtau + + # s (slow inactivation) + alps = self.a0s * save_exp((asvh - v) / self.avs) + bets = self.b0s / (save_exp((bsvh - v) / self.bvs) + 1.0) + tau_s = 1.0 / (alps + bets) + c = alps * tau_s + s_inf = c + ar * (1.0 - c) + + m_new = exponential_euler(m, dt, m_inf, tau_m) + s_new = exponential_euler(s, dt, s_inf, tau_s) + + return { + f"{prefix}_m": m_new, + f"{prefix}_s": s_new, + } + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """Compute persistent sodium current: I = gbar * m * s * (v - E_Na).""" + prefix = self._name + m = states[f"{prefix}_m"] + s = states[f"{prefix}_s"] + gbar = params[f"{prefix}_gbar"] + e_na = params[f"{prefix}_eNa"] + + g = gbar * m * s + return g * (v - e_na) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """Initialize to steady-state.""" + prefix = self._name + sh = params[f"{prefix}_sh"] + ar = params[f"{prefix}_ar"] + vslope = params[f"{prefix}_vslope"] + asvh = params[f"{prefix}_asvh"] + bsvh = params[f"{prefix}_bsvh"] + + m_inf = 1.0 / (1.0 + save_exp(-(v + 52.3 - sh) / vslope)) + + alps = self.a0s * save_exp((asvh - v) / self.avs) + bets = self.b0s / (save_exp((bsvh - v) / self.bvs) + 1.0) + tau_s = 1.0 / (alps + bets) + c = alps * tau_s + s_inf = c + ar * (1.0 - c) + + return { + f"{prefix}_m": m_inf, + f"{prefix}_s": s_inf, + } + + +# ============================================================================= +# KDRRL - DELAYED RECTIFIER POTASSIUM CHANNEL +# ============================================================================= + +class KdrRL(Channel): + """ + Delayed rectifier potassium channel (kdrRL). + + Single gating variable with 4th power: m⁴ + Responsible for action potential repolarization. + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gbar": 0.1, # S/cm² + f"{prefix}_mVh": -25.0, # mV - half-activation voltage + f"{prefix}_mslp": 20.0, # mV - activation slope + f"{prefix}_eK": -80.0, # Potassium reversal potential + f"{prefix}_tmin": 1.4, # ms - minimum tau (NEURON global: 0.8) + f"{prefix}_taumax": 11.9, # ms - maximum tau (NEURON global: 20.0) + } + self.channel_states = { + f"{prefix}_m": 0.0, + } + self.current_name = f"i_KdrRL" + + # Tau parameters (voltage dependence, not modified by NEURON) + self.tVh = -39.0 # mV - tau half-voltage + self.tslp = 5.5 # mV - tau slope + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Update activation gating variable.""" + prefix = self._name + m = states[f"{prefix}_m"] + + mVh = params[f"{prefix}_mVh"] + mslp = params[f"{prefix}_mslp"] + tmin = params[f"{prefix}_tmin"] + taumax = params[f"{prefix}_taumax"] + + m_inf = 1.0 / (1.0 + save_exp(-(v - mVh) / mslp)) + + b = save_exp((v - self.tVh) / self.tslp) + f = (1.0 + b) ** 2 + tau_m = tmin + taumax * b / f + + m_new = exponential_euler(m, dt, m_inf, tau_m) + + return {f"{prefix}_m": m_new} + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """Compute potassium current: I = gbar * m⁴ * (v - E_K).""" + prefix = self._name + m = states[f"{prefix}_m"] + gbar = params[f"{prefix}_gbar"] + e_k = params[f"{prefix}_eK"] + + g = gbar * (m ** 4) + return g * (v - e_k) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """Initialize to steady-state.""" + prefix = self._name + mVh = params[f"{prefix}_mVh"] + mslp = params[f"{prefix}_mslp"] + + m_inf = 1.0 / (1.0 + save_exp(-(v - mVh) / mslp)) + + return {f"{prefix}_m": m_inf} + + +# ============================================================================= +# MAHP - CALCIUM-DEPENDENT POTASSIUM CHANNEL (MEDIUM AHP) +# ============================================================================= + +class MAHP(Channel): + """ + Calcium-dependent potassium channel responsible for medium AHP. + + Includes a simplified calcium channel and calcium dynamics. + Two gating systems: + - mca: calcium channel activation (voltage-dependent) + - n: KCa channel activation (calcium-dependent) + + Also tracks intracellular calcium concentration [Ca]_i. + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gkcamax": 0.03, # S/cm² - KCa conductance + f"{prefix}_gcamax": 3e-5, # S/cm² - Ca conductance + f"{prefix}_mvhalfca": -30.0, # mV - Ca channel half-activation + f"{prefix}_mslpca": 4.0, # mV - Ca channel slope + f"{prefix}_mtauca": 1.0, # ms - Ca channel activation time constant + f"{prefix}_tau_ca": 20.0, # ms - calcium removal time constant + f"{prefix}_eK": -80.0, # Potassium reversal potential + f"{prefix}_eCa": 120.0, # Calcium reversal potential + } + self.channel_states = { + f"{prefix}_mca": 0.0, # Ca channel activation + f"{prefix}_n": 0.0, # KCa activation + f"{prefix}_cai": 0.0001, # [Ca]_i in mM + } + self.current_name = f"i_MAHP" + + self.mtauca = 1.0 # ms - Ca channel tau + self.depth = 0.1 # um - shell depth for Ca dynamics + self.cainf = 0.0001 # mM - resting [Ca]_i + self.caix = 2.0 # Ca cooperativity for KCa + self.fKCa = 0.1 # KCa activation rate + self.bKCa = 0.1 # KCa deactivation rate + + # Faraday constant (C/mol) + self.FARADAY = 96485.0 + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Update Ca channel, KCa channel, and calcium concentration.""" + prefix = self._name + mca = states[f"{prefix}_mca"] + n = states[f"{prefix}_n"] + cai = states[f"{prefix}_cai"] + + mvhalfca = params[f"{prefix}_mvhalfca"] + mslpca = params[f"{prefix}_mslpca"] + mtauca = params[f"{prefix}_mtauca"] + gcamax = params[f"{prefix}_gcamax"] + tau_ca = params[f"{prefix}_tau_ca"] + e_ca = params[f"{prefix}_eCa"] + + # Update Ca channel activation (Boltzmann) + mca_inf = 1.0 / (1.0 + save_exp(-(v - mvhalfca) / mslpca)) + mca_new = exponential_euler(mca, dt, mca_inf, mtauca) + + # Compute calcium current using OLD mca (matches NEURON's cnexp behavior) + # In NEURON, the DERIVATIVE block sees ica from the PREVIOUS BREAKPOINT, + # which was computed with the old mca value before state updates. + ica = gcamax * mca * (v - e_ca) # Use OLD mca, not mca_new! + + # Update [Ca]_i: drive from Ca current, decay to cainf + # ODE: cai' = drive + (cainf - cai) / tau_ca + # Rewrite as: cai' = (cainf + drive*tau_ca - cai) / tau_ca + # So: cai_inf = cainf + drive*tau_ca, solved with exponential Euler (matches cnexp) + drive = -10000.0 * ica / (2.0 * self.FARADAY * self.depth) + drive = jnp.maximum(drive, 0.0) # Cannot pump inward + + cai_inf = self.cainf + drive * tau_ca + cai_new = exponential_euler(cai, dt, cai_inf, tau_ca) + cai_new = jnp.maximum(cai_new, self.cainf) # Floor at resting level + + # Update KCa activation using OLD cai (matches NEURON's cnexp behavior) + # In NEURON, rates(cai) is called in DERIVATIVE with current cai, not updated cai + ca_uM = jnp.maximum(cai - self.cainf, 0.0) * 1e3 # Use OLD cai, not cai_new! + a = self.fKCa * (ca_uM ** self.caix) + b = self.bKCa + tau_n = 1.0 / (a + b + 1e-10) + n_inf = a * tau_n + n_new = exponential_euler(n, dt, n_inf, tau_n) + + return { + f"{prefix}_mca": mca_new, + f"{prefix}_n": n_new, + f"{prefix}_cai": cai_new, + } + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """ + Compute total current (KCa + Ca). + + Returns the sum of: + - I_KCa = gkcamax * n * (v - E_K) + - I_Ca = gcamax * mca * (v - E_Ca) + """ + prefix = self._name + mca = states[f"{prefix}_mca"] + n = states[f"{prefix}_n"] + + gkcamax = params[f"{prefix}_gkcamax"] + gcamax = params[f"{prefix}_gcamax"] + e_k = params[f"{prefix}_eK"] + e_ca = params[f"{prefix}_eCa"] + + i_kca = gkcamax * n * (v - e_k) + i_ca = gcamax * mca * (v - e_ca) + + return i_kca + i_ca + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """Initialize to steady-state.""" + prefix = self._name + mvhalfca = params[f"{prefix}_mvhalfca"] + mslpca = params[f"{prefix}_mslpca"] + + mca_inf = 1.0 / (1.0 + save_exp(-(v - mvhalfca) / mslpca)) + + ca_uM = 0.0 # At rest, no excess calcium + a = self.fKCa * (ca_uM ** self.caix) + b = self.bKCa + tau_n = 1.0 / (a + b + 1e-10) + n_inf = a * tau_n + + return { + f"{prefix}_mca": mca_inf, + f"{prefix}_n": n_inf, + f"{prefix}_cai": jnp.full_like(v, self.cainf), + } + + +# ============================================================================= +# GH - HYPERPOLARIZATION-ACTIVATED CATION CURRENT (H-CURRENT) +# ============================================================================= + +class Gh(Channel): + """ + Hyperpolarization-activated cation current (gh). + + Single gating variable n with voltage-dependent kinetics. + Activated by hyperpolarization, contributes to sag and rebound. + + Non-specific cation current (mixed Na+/K+). + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gbar": 0.001, # S/cm² + f"{prefix}_half": -80.0, # mV - half-activation voltage + f"{prefix}_slp": 8.0, # mV - slope + f"{prefix}_eH": -41.0, # mV - reversal potential + f"{prefix}_htau": 50.0, # ms - time constant + } + self.channel_states = { + f"{prefix}_n": 0.0, + } + self.current_name = f"i_Gh" + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Update activation gating variable.""" + prefix = self._name + n = states[f"{prefix}_n"] + + half = params[f"{prefix}_half"] + slp = params[f"{prefix}_slp"] + htau = params[f"{prefix}_htau"] + + # Activated by hyperpolarization + n_inf = 1.0 / (1.0 + save_exp((v - half) / slp)) + n_new = exponential_euler(n, dt, n_inf, htau) + + return {f"{prefix}_n": n_new} + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """Compute H-current: I = gbar * n * (v - E_h).""" + prefix = self._name + n = states[f"{prefix}_n"] + gbar = params[f"{prefix}_gbar"] + e_h = params[f"{prefix}_eH"] + + return gbar * n * (v - e_h) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """Initialize to steady-state.""" + prefix = self._name + half = params[f"{prefix}_half"] + slp = params[f"{prefix}_slp"] + + n_inf = 1.0 / (1.0 + save_exp((v - half) / slp)) + + return {f"{prefix}_n": n_inf} + + +# ============================================================================= +# L_CA_INACT - L-TYPE CALCIUM CHANNEL WITH INACTIVATION (DENDRITIC PIC) +# ============================================================================= + +class LCaInact(Channel): + """ + L-type calcium channel with voltage-dependent inactivation. + + Two gating variables: + - m: activation (fast, ~20 ms) + - h: inactivation (slow, ~1500 ms) + + Responsible for dendritic persistent inward currents (PICs) and + bistable firing behavior in motor neurons. + + Parameters from Booth et al. (1997) J Neurophysiol 78:3371-3385. + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gbar": 0.0003, # S/cm² + f"{prefix}_theta_m": -30.0, # mV - activation half-voltage + f"{prefix}_kappa_m": -6.0, # mV - activation slope (negative) + f"{prefix}_tau_m": 20.0, # ms - activation time constant + f"{prefix}_theta_h": 14.0, # mV - inactivation half-voltage + f"{prefix}_kappa_h": 4.0, # mV - inactivation slope + f"{prefix}_tau_h": 1500.0, # ms - inactivation time constant + f"{prefix}_eCa": 80.0, # mV - calcium reversal potential + } + self.channel_states = { + f"{prefix}_m": 0.0, + f"{prefix}_h": 1.0, + } + self.current_name = f"i_LCaInact" + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Update activation and inactivation gating variables.""" + prefix = self._name + m = states[f"{prefix}_m"] + h = states[f"{prefix}_h"] + + theta_m = params[f"{prefix}_theta_m"] + kappa_m = params[f"{prefix}_kappa_m"] + tau_m = params[f"{prefix}_tau_m"] + theta_h = params[f"{prefix}_theta_h"] + kappa_h = params[f"{prefix}_kappa_h"] + tau_h = params[f"{prefix}_tau_h"] + + # Boltzmann activation (note: kappa_m is negative) + m_inf = 1.0 / (1.0 + save_exp((v - theta_m) / kappa_m)) + # Boltzmann inactivation + h_inf = 1.0 / (1.0 + save_exp((v - theta_h) / kappa_h)) + + m_new = exponential_euler(m, dt, m_inf, tau_m) + h_new = exponential_euler(h, dt, h_inf, tau_h) + + return { + f"{prefix}_m": m_new, + f"{prefix}_h": h_new, + } + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """Compute L-type Ca current: I = gbar * m * h * (v - E_Ca).""" + prefix = self._name + m = states[f"{prefix}_m"] + h = states[f"{prefix}_h"] + gbar = params[f"{prefix}_gbar"] + vca = params[f"{prefix}_eCa"] + + return gbar * m * h * (v - vca) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """Initialize to steady-state.""" + prefix = self._name + theta_m = params[f"{prefix}_theta_m"] + kappa_m = params[f"{prefix}_kappa_m"] + theta_h = params[f"{prefix}_theta_h"] + kappa_h = params[f"{prefix}_kappa_h"] + + m_inf = 1.0 / (1.0 + save_exp((v - theta_m) / kappa_m)) + h_inf = 1.0 / (1.0 + save_exp((v - theta_h) / kappa_h)) + + return { + f"{prefix}_m": m_inf, + f"{prefix}_h": h_inf, + } + + +# ============================================================================= +# LEAK CHANNEL (INCLUDED FOR COMPLETENESS) +# ============================================================================= + +class LeakChannel(Channel): + """ + Simple leak channel for passive membrane properties. + + I = g_leak * (v - E_leak) + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True # Required for Jaxley 0.5.0+ + super().__init__(name) + + prefix = self._name + self.channel_params = { + f"{prefix}_gLeak": 1e-4, # S/cm² + f"{prefix}_eLeak": -70.0, # mV + } + self.channel_states = {} + self.current_name = f"i_Leak" + + def update_states( + self, + states: Dict[str, jnp.ndarray], + dt, + v, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """No states to update for leak channel.""" + return {} + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """Compute leak current.""" + prefix = self._name + g = params[f"{prefix}_gLeak"] + e = params[f"{prefix}_eLeak"] + return g * (v - e) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + v, + params: Dict[str, jnp.ndarray], + delta_t, + ) -> Dict[str, jnp.ndarray]: + """No states to initialize.""" + return {} + + +# ============================================================================= +# CONVENIENCE FUNCTION FOR MOTOR NEURON CHANNEL SET +# ============================================================================= + +def get_motor_neuron_channels_soma(model: str = "Powers2017") -> list: + """ + Get the complete set of soma channels for a motor neuron. + + Parameters + ---------- + model : str + "Powers2017" or "NERLab" parameter sets. + + Returns + ------- + list + List of channel instances configured for soma. + """ + if model == "Powers2017": + return [ + Na3rp(), + Naps(), + KdrRL(), + MAHP(), + Gh(), + LeakChannel(), + ] + else: # NERLab + return [ + Na3rp(), + KdrRL(), + MAHP(), + Gh(), + LeakChannel(), + ] + + +def get_motor_neuron_channels_dendrite(model: str = "Powers2017") -> list: + """ + Get the complete set of dendritic channels for a motor neuron. + + Dendrites have L-type Ca channels for PICs but fewer Na channels. + + Parameters + ---------- + model : str + "Powers2017" or "NERLab" parameter sets. + + Returns + ------- + list + List of channel instances configured for dendrites. + """ + if model == "Powers2017": + return [ + LCaInact(), + Gh(), + LeakChannel(), + ] + else: # NERLab + return [ + LCaInact(), + Gh(), + LeakChannel(), + ] diff --git a/myogen/simulator/jaxley/channels/nerlab_channels.py b/myogen/simulator/jaxley/channels/nerlab_channels.py new file mode 100644 index 00000000..d82e5fad --- /dev/null +++ b/myogen/simulator/jaxley/channels/nerlab_channels.py @@ -0,0 +1,299 @@ +""" +NERLab Channel Models — Jaxley Backend. +======================================= + +JAX/Jaxley implementations of the channels used by MyoGen's production NEURON +motor-neuron model (``model="NERLab"`` in ``AlphaMN__Pool``). These are direct +ports of the NMODL files: + + myogen/simulator/nmodl_files/napp.mod -> ``napp`` (5-gate Na + K + leak) + myogen/simulator/nmodl_files/caL.mod -> ``caL`` (L-type Ca, no inactivation) + +The static-current ``Constant`` mechanism is already available as +``myogen.simulator.jaxley.channels.Constant`` (no duplication here). + +**Voltage convention — IMPORTANT.** NERLab uses the *original 1952 Hodgkin-Huxley* +convention, NOT the modern absolute-V convention used by the Powers2017 channels +elsewhere in this package. Concretely: + + V_rest ≈ 0 mV (set by ``el_napp = 0``) + ENa = +120 mV (spike peaks reach ~+90 to +110 mV) + EK = -10 mV (AHPs dip to ~-5 to -10 mV) + vtraub = 0 mV (no extra voltage offset in the gating equations) + +Do NOT mix these channels with the Powers2017 channels on the same cell. Cells +using NERLab channels must initialise V at 0 mV, not -65 mV. + +Author: ported for the Jaxley/NEURON architecture-parity work, 2026. +""" + +from typing import Dict, Optional + +import jax.numpy as jnp +from jaxley.channels import Channel +from jaxley.solver_gate import exponential_euler, save_exp + + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + +def _vtrap(x: jnp.ndarray, y: float) -> jnp.ndarray: + """JAX-safe port of NEURON's ``vtrap(x, y) = x / (exp(x/y) - 1)``. + + Uses the L'Hopital expansion ``y * (1 - x / y / 2)`` when ``|x/y|`` is small, + avoiding the 0/0 singularity at ``x == 0``. Implemented with ``jnp.where`` + so it is differentiable and works under ``jit``. + """ + safe = jnp.abs(x / y) < 1e-4 + # When |x/y| < 1e-4 use the linear approximation; otherwise the analytic form. + # The denominator term is computed in a way that never divides by exactly 0. + denom = save_exp(x / y) - 1.0 + # Avoid exact zero in the denominator under jit: where ``safe`` is True the + # value isn't used, but JAX still evaluates both branches, so guard it. + denom_safe = jnp.where(safe, 1.0, denom) + return jnp.where(safe, y * (1.0 - x / y / 2.0), x / denom_safe) + + +# ----------------------------------------------------------------------------- +# napp — Hodgkin-Huxley-style channel set used in the NERLab motor neuron +# ----------------------------------------------------------------------------- + +class napp(Channel): + """5-gate Na/K/leak channel set used by MyoGen's NERLab soma. + + Gates + ----- + m (×3) : fast sodium activation + h (×1) : fast sodium inactivation + p (×3) : persistent sodium activation + n (×4) : fast potassium activation + r (×2) : slow potassium activation + + Currents + -------- + ina = (gnabar * m^3 * h + gnapbar * p^3) * (v - ena) + ik = (gkfbar * n^4 + gksbar * r^2) * (v - ek) + il = gl * (v - el) + + All gating equations are evaluated at ``v2 = v - vtraub``. The default + ``vtraub = 0`` matches the production NERLab YAML. + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True + super().__init__(name) + prefix = self._name + + # Conductance and reversal defaults match napp.mod / alpha_mn_default.yaml. + self.channel_params = { + # Maximal conductances + f"{prefix}_gnabar": 0.030, # S/cm^2 (fast Na) + f"{prefix}_gnapbar": 0.000033, # S/cm^2 (persistent Na) + f"{prefix}_gkfbar": 0.016, # S/cm^2 (fast K) + f"{prefix}_gksbar": 0.004, # S/cm^2 (slow K) + f"{prefix}_gl": 0.0003, # S/cm^2 (leak) + # Reversals (NERLab / original-HH convention) + f"{prefix}_ena": 120.0, # mV + f"{prefix}_ek": -10.0, # mV + f"{prefix}_el": -54.3, # mV (NEURON default; overridden to 0 by pool) + # Voltage offset used by every gate + f"{prefix}_vtraub": 0.0, # mV (production NERLab uses 0) + # Threshold-like parameter (not used in dynamics; legacy NEURON parameter) + f"{prefix}_mact": 15.0, + # Slow-K constant deactivation rate (β_r) + f"{prefix}_rinact": 0.05, # /ms + # ---- Per-gate kinetic constants (from YAML) ---- + # m + f"{prefix}_m_alpha_A": 0.64, f"{prefix}_m_alpha_v_offset": 15.0, f"{prefix}_m_alpha_k": 4.0, + f"{prefix}_m_beta_A": 0.56, f"{prefix}_m_beta_v_offset": 40.0, f"{prefix}_m_beta_k": 5.0, + # h + f"{prefix}_h_alpha_A": 0.928, f"{prefix}_h_alpha_v_offset": 17.0, f"{prefix}_h_alpha_tau": 18.0, + f"{prefix}_h_beta_A": 9.0, f"{prefix}_h_beta_v_offset": 40.0, f"{prefix}_h_beta_k": 5.0, + # p + f"{prefix}_p_alpha_A": 0.64, f"{prefix}_p_alpha_v_offset": 5.0, f"{prefix}_p_alpha_k": 4.0, + f"{prefix}_p_beta_A": 0.56, f"{prefix}_p_beta_v_offset": 30.0, f"{prefix}_p_beta_k": 5.0, + # n + f"{prefix}_n_alpha_A": 0.08, f"{prefix}_n_alpha_v_offset": 15.0, f"{prefix}_n_alpha_k": 7.0, + f"{prefix}_n_beta_A": 2.0, f"{prefix}_n_beta_v_offset": 10.0, f"{prefix}_n_beta_tau":40.0, + # r + f"{prefix}_r_alpha_A": 3.5, f"{prefix}_r_alpha_v_offset": 55.0, f"{prefix}_r_alpha_k": 4.0, + } + # Initial states near a NERLab resting condition (v ≈ 0). Will be + # overwritten by ``init_states`` when the cell is built; these only + # matter if the user forgets to initialise. + self.channel_states = { + f"{prefix}_m": 0.01, + f"{prefix}_h": 0.99, + f"{prefix}_p": 0.07, + f"{prefix}_n": 0.06, + f"{prefix}_r": 0.0, + } + self.current_name = "i_napp" + + # ---- gating helpers (alpha, beta) ---- + + @staticmethod + def _m_rates(v2, p): + alpha = p["m_alpha_A"] * _vtrap(p["m_alpha_v_offset"] - v2, p["m_alpha_k"]) + beta = p["m_beta_A"] * _vtrap(v2 - p["m_beta_v_offset"], p["m_beta_k"]) + return alpha, beta + + @staticmethod + def _h_rates(v2, p): + alpha = p["h_alpha_A"] * save_exp((p["h_alpha_v_offset"] - v2) / p["h_alpha_tau"]) + beta = p["h_beta_A"] / (save_exp((p["h_beta_v_offset"] - v2) / p["h_beta_k"]) + 1.0) + return alpha, beta + + @staticmethod + def _p_rates(v2, p): + alpha = p["p_alpha_A"] * _vtrap(p["p_alpha_v_offset"] - v2, p["p_alpha_k"]) + beta = p["p_beta_A"] * _vtrap(v2 - p["p_beta_v_offset"], p["p_beta_k"]) + return alpha, beta + + @staticmethod + def _n_rates(v2, p): + alpha = p["n_alpha_A"] * _vtrap(p["n_alpha_v_offset"] - v2, p["n_alpha_k"]) + beta = p["n_beta_A"] * save_exp((p["n_beta_v_offset"] - v2) / p["n_beta_tau"]) + return alpha, beta + + @staticmethod + def _r_rates(v2, p): + alpha = p["r_alpha_A"] / (save_exp((p["r_alpha_v_offset"] - v2) / p["r_alpha_k"]) + 1.0) + beta = p["rinact"] + return alpha, beta + + @staticmethod + def _inf_tau(alpha, beta): + sumab = alpha + beta + 1e-12 # avoid div-by-zero in pathological regimes + return alpha / sumab, 1.0 / sumab + + def update_states(self, states, dt, v, params): + prefix = self._name + # Pack the kinetic constants into a name->array dict for the helpers. + p = { + k.replace(f"{prefix}_", ""): params[k] + for k in params if k.startswith(prefix) and not k.endswith(("_gnabar", "_gnapbar", "_gkfbar", "_gksbar", "_gl", "_ena", "_ek", "_el", "_mact")) + } + v2 = v - params[f"{prefix}_vtraub"] + + m_inf, m_tau = self._inf_tau(*self._m_rates(v2, p)) + h_inf, h_tau = self._inf_tau(*self._h_rates(v2, p)) + p_inf, p_tau = self._inf_tau(*self._p_rates(v2, p)) + n_inf, n_tau = self._inf_tau(*self._n_rates(v2, p)) + r_inf, r_tau = self._inf_tau(*self._r_rates(v2, p)) + + m = states[f"{prefix}_m"]; h = states[f"{prefix}_h"] + ppstate = states[f"{prefix}_p"] + n = states[f"{prefix}_n"]; r = states[f"{prefix}_r"] + + return { + f"{prefix}_m": exponential_euler(m, dt, m_inf, m_tau), + f"{prefix}_h": exponential_euler(h, dt, h_inf, h_tau), + f"{prefix}_p": exponential_euler(ppstate, dt, p_inf, p_tau), + f"{prefix}_n": exponential_euler(n, dt, n_inf, n_tau), + f"{prefix}_r": exponential_euler(r, dt, r_inf, r_tau), + } + + def compute_current(self, states, v, params): + prefix = self._name + m = states[f"{prefix}_m"]; h = states[f"{prefix}_h"] + pp = states[f"{prefix}_p"] + n = states[f"{prefix}_n"]; r = states[f"{prefix}_r"] + + gna = params[f"{prefix}_gnabar"] * (m ** 3) * h + gnap = params[f"{prefix}_gnapbar"] * (pp ** 3) + gkf = params[f"{prefix}_gkfbar"] * (n ** 4) + gks = params[f"{prefix}_gksbar"] * (r ** 2) + gl = params[f"{prefix}_gl"] + + ena = params[f"{prefix}_ena"] + ek = params[f"{prefix}_ek"] + el = params[f"{prefix}_el"] + + ina = (gna + gnap) * (v - ena) + ik = (gkf + gks) * (v - ek) + il = gl * (v - el) + return ina + ik + il + + def init_state(self, states, v, params, dt): + """Compute steady-state gate values at the holding voltage `v`.""" + prefix = self._name + p = { + k.replace(f"{prefix}_", ""): params[k] + for k in params if k.startswith(prefix) and not k.endswith(("_gnabar", "_gnapbar", "_gkfbar", "_gksbar", "_gl", "_ena", "_ek", "_el", "_mact")) + } + v2 = v - params[f"{prefix}_vtraub"] + m_inf, _ = self._inf_tau(*self._m_rates(v2, p)) + h_inf, _ = self._inf_tau(*self._h_rates(v2, p)) + p_inf, _ = self._inf_tau(*self._p_rates(v2, p)) + n_inf, _ = self._inf_tau(*self._n_rates(v2, p)) + r_inf, _ = self._inf_tau(*self._r_rates(v2, p)) + return { + f"{prefix}_m": m_inf, + f"{prefix}_h": h_inf, + f"{prefix}_p": p_inf, + f"{prefix}_n": n_inf, + f"{prefix}_r": r_inf, + } + + +# ----------------------------------------------------------------------------- +# caL — L-type calcium channel used in the NERLab dendrite (no inactivation) +# ----------------------------------------------------------------------------- + +class caL(Channel): + """L-type Ca channel from MyoGen's NERLab dendrite. + + Single activation gate ``L`` with a fixed time constant ``Ltau``: + + Linf = 1 / (1 + exp(-(v2 + 30))) (NMODL uses k = -1 mV → extremely steep) + dL/dt = (Linf - L) / Ltau + icaL = gcaLbar * L * gama * (v - ecaL) + + A passive leak ``il = gl * (v - el)`` is bundled in (matches caL.mod). + """ + + def __init__(self, name: Optional[str] = None): + self.current_is_in_mA_per_cm2 = True + super().__init__(name) + prefix = self._name + self.channel_params = { + f"{prefix}_gcaLbar": 0.030, # S/cm^2 + f"{prefix}_gl": 0.0003, # S/cm^2 (dendritic leak) + f"{prefix}_el": -54.3, # mV (overridden to 0 by NERLab pool) + f"{prefix}_ecaL": 140.0, # mV + f"{prefix}_vtraub": 70.0, # mV (NEURON default; NERLab uses ~35) + f"{prefix}_gama": 1.0, + f"{prefix}_Ltau": 20.0, # ms + } + self.channel_states = { + f"{prefix}_L": 0.0, + } + self.current_name = "i_caL" + + @staticmethod + def _Linf(v2): + # NMODL: Linf = 1 / (exp((v2+30)/-1) + 1) — equivalent to logistic with k=-1 mV. + return 1.0 / (1.0 + save_exp((v2 + 30.0) / -1.0)) + + def update_states(self, states, dt, v, params): + prefix = self._name + v2 = v - params[f"{prefix}_vtraub"] + L_inf = self._Linf(v2) + L_new = exponential_euler(states[f"{prefix}_L"], dt, L_inf, params[f"{prefix}_Ltau"]) + return {f"{prefix}_L": L_new} + + def compute_current(self, states, v, params): + prefix = self._name + gcaL = params[f"{prefix}_gcaLbar"] * states[f"{prefix}_L"] + icaL = gcaL * params[f"{prefix}_gama"] * (v - params[f"{prefix}_ecaL"]) + il = params[f"{prefix}_gl"] * (v - params[f"{prefix}_el"]) + return icaL + il + + def init_state(self, states, v, params, dt): + prefix = self._name + v2 = v - params[f"{prefix}_vtraub"] + return {f"{prefix}_L": self._Linf(v2)} + + diff --git a/myogen/simulator/jaxley/closed_loop.py b/myogen/simulator/jaxley/closed_loop.py new file mode 100644 index 00000000..913d2466 --- /dev/null +++ b/myogen/simulator/jaxley/closed_loop.py @@ -0,0 +1,406 @@ +""" +Differentiable closed-loop entry point for the MyoGen Jaxley backend. + +This module packages the full ``jax.lax.scan`` neuromuscular loop — Jaxley +biophysics + synaptic conductances + stochastic afferent/descending generators + +Hill muscle + spindle + GTO + joint dynamics — into a single pure function +``run_jax`` for gradient / accelerator workflows. + +JIT / grad usage +---------------- +``run_jax`` cannot be handed to ``jax.jit`` *directly*: ``ClosedLoopConfig`` is +intentionally not hashable (it holds arrays), and a few integer Hill fields +(``hill_p["N"]``/``"Ntype1"``) are used in shape contexts (e.g. ``jnp.arange(N)``) +so they must stay static rather than being traced. Use the provided wrappers, which +perform the static/dynamic split for you: + +* :func:`value_and_grad_run` — gradients (differentiates only float leaves; integer + fields held static). This is the primary entry point for optimization. +* :func:`compile_run` — a ``jax.jit``-compiled forward pass bound to a config + (Python-scalar fields baked static, all arrays traced). + +Until now the loop only existed inline in +``examples/01_basic/11_simulate_spinal_network_jaxley.py``. Here it is a library +function with a clean split between: + +* a **differentiable parameter PyTree** (``params``): everything you might take a + gradient with respect to — Jaxley channel/synapse parameters, Hill/spindle/GTO/ + joint numeric leaves, and synaptic weights; and +* a **static configuration** (``ClosedLoopConfig``): network topology, connectivity + matrices, counts, delays, timestep, initial states, and the compiled Jaxley + ``step_fn``. + +The stochastic generators are seeded from an explicit ``key`` argument so runs are +reproducible across JIT and devices. ``spike_mode`` selects hard (scientific +default), surrogate-gradient, or rate spikes and is recorded in the returned +metadata. + +Typical use +----------- +:: + + cfg = ClosedLoopConfig(...) # built once from the network + params = default_params(cfg, ...) # differentiable PyTree + outputs, meta = run_jax(params, cfg, inputs, jax.random.PRNGKey(0)) + + # gradient of a scalar loss on the force trace w.r.t. all float params: + loss = lambda out: jnp.mean(out["force"] ** 2) + val, grads = value_and_grad_run(loss, params, cfg, inputs, + jax.random.PRNGKey(0), spike_mode="surrogate") +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from typing import Any, Callable, Optional + +import jax +import jax.numpy as jnp +import jax.tree_util as jtu + +from myogen.simulator.jaxley.jax_models import ( + make_scan_step, + poisson_init, + gamma_init, +) + +__all__ = [ + "ClosedLoopConfig", + "run_jax", + "compile_run", + "value_and_grad_run", + "partition_differentiable", +] + + +@dataclass(frozen=True) +class ClosedLoopConfig: + """Static (non-differentiable) configuration for :func:`run_jax`. + + Everything here is closed over as a constant during tracing. Anything you may + want gradients for lives in the ``params`` PyTree instead (see + :func:`run_jax`). + """ + + # Compiled Jaxley single-step function (from ``build_init_and_step_fn``). + step_fn: Callable + external_inds: Any + rec_inds: Any + + # Population counts. + n_gii: int + n_gib: int + n_mn: int + nDD: int + nIa: int + nII: int + nIb: int + + # Afferent response thresholds and gamma-ISI shape parameters. + ia_rts: Any + ii_rts: Any + ib_rts: Any + ia_shape: float + ii_shape: float + ib_shape: float + dd_N_batch: int + + # Connectivity matrices (pre_idx → post_idx). + dd_to_mn_mat: Any + ia_to_mn_mat: Any + ii_to_gii_mat: Any + ib_to_gib_mat: Any + + # Per-cell axonal delay steps. + ia_delay_steps_arr: Any + ii_delay_steps_arr: Any + ib_delay_steps_arr: Any + delay_steps: Any + max_ia_delay_steps: int + max_ii_delay_steps: int + max_ib_delay_steps: int + + # Voltage-frame / synapse scalars. + e_exc: float + v_rest: float + e_exc_mn: float + mn_spike_threshold_mV: float + tau_syn_decay: float + dt_ms: float + dt_s: float + + # Initial states. + init_neural_states: Any + init_hill_state: Any + init_spindle_state: Any + init_gto_state: Any + init_joint_state: Any + init_prev_v: float = -70.0 + + # Default spike mode; can be overridden per call. + spike_mode: str = "hard" + + +def _build_scan_step(params: dict, config: ClosedLoopConfig, spike_mode: str): + """Instantiate the ``lax.scan`` step closure from params + static config. + + ``params`` leaves flow in as traced values, so gradients propagate through + Jaxley parameters, Hill/spindle/GTO/joint leaves, and synaptic weights. + """ + w = params["weights"] + return make_scan_step( + jaxley_step_fn=config.step_fn, + jaxley_params=params["jaxley"], + external_inds=config.external_inds, + rec_inds=config.rec_inds, + n_gii=config.n_gii, + n_gib=config.n_gib, + n_mn=config.n_mn, + ia_rts=config.ia_rts, + ii_rts=config.ii_rts, + ib_rts=config.ib_rts, + ia_shape=config.ia_shape, + ii_shape=config.ii_shape, + ib_shape=config.ib_shape, + dd_N_batch=config.dd_N_batch, + dd_to_mn_mat=config.dd_to_mn_mat, + ia_to_mn_mat=config.ia_to_mn_mat, + ii_to_gii_mat=config.ii_to_gii_mat, + ib_to_gib_mat=config.ib_to_gib_mat, + ia_delay_steps_arr=config.ia_delay_steps_arr, + ii_delay_steps_arr=config.ii_delay_steps_arr, + ib_delay_steps_arr=config.ib_delay_steps_arr, + delay_steps=config.delay_steps, + hill_p=params["hill"], + spindle_p=params["spindle"], + gto_p=params["gto"], + joint_p=params["joint"], + base_dd_weight=w["base_dd"], + base_ia_weight=w["base_ia"], + in_weight=w["in_weight"], + e_exc=config.e_exc, + v_rest=config.v_rest, + e_exc_mn=config.e_exc_mn, + mn_spike_threshold_mV=config.mn_spike_threshold_mV, + mn_current_scale=w["mn_current_scale"], + tau_syn_decay=config.tau_syn_decay, + dt_ms=config.dt_ms, + dt_s=config.dt_s, + spike_mode=spike_mode, + ) + + +def _build_init_carry(params: dict, config: ClosedLoopConfig, key) -> dict: + """Assemble the scan carry, seeding the stochastic generators from ``key``.""" + kd, ka, ki, kb = jax.random.split(key, 4) + n_total = config.n_gii + config.n_gib + config.n_mn + return { + "neural": config.init_neural_states, + "phys": { + "hill": config.init_hill_state, + "spindle": config.init_spindle_state, + "gto": config.init_gto_state, + "joint": config.init_joint_state, + }, + "g_dd": jnp.zeros(config.n_mn, dtype=jnp.float32), + "g_ia": jnp.zeros(config.n_mn, dtype=jnp.float32), + "g_ii": jnp.zeros(config.n_gii, dtype=jnp.float32), + "g_ib": jnp.zeros(config.n_gib, dtype=jnp.float32), + "prev_v": jnp.full(n_total, config.init_prev_v, dtype=jnp.float32), + "dd_st": poisson_init(config.nDD, config.dd_N_batch, key=kd), + "ia_st": gamma_init(config.nIa, config.ia_shape, key=ka), + "ii_st": gamma_init(config.nII, config.ii_shape, key=ki), + "ib_st": gamma_init(config.nIb, config.ib_shape, key=kb), + "prev_Iay": jnp.float32(0.0), + "prev_IIy": jnp.float32(0.0), + "prev_Iby": jnp.float32(0.0), + "ia_delay_buf": jnp.zeros((config.nIa, config.max_ia_delay_steps), dtype=jnp.float32), + "ii_delay_buf": jnp.zeros((config.nII, config.max_ii_delay_steps), dtype=jnp.float32), + "ib_delay_buf": jnp.zeros((config.nIb, config.max_ib_delay_steps), dtype=jnp.float32), + } + + +def run_jax( + params: dict, + config: ClosedLoopConfig, + inputs: dict, + key, + spike_mode: Optional[str] = None, +): + """Run the full closed-loop neuromuscular simulation as a pure function. + + Parameters + ---------- + params : dict + Differentiable PyTree with keys ``{"jaxley", "hill", "spindle", "gto", + "joint", "weights"}``. ``weights`` is a dict of ``{"base_dd", "base_ia", + "in_weight", "mn_current_scale"}``. Any float leaf may be differentiated. + config : ClosedLoopConfig + Static configuration (topology, connectivity, initial states, ``step_fn``). + inputs : dict + Per-step drive arrays stacked over time, each shape ``(n_steps, ...)``: + ``{"DDdrive", "gDyn", "gStat", "tap_dL", "tap_dV"}``. + key : jax PRNGKey + Seeds the stochastic afferent/descending generators (split internally). + spike_mode : {"hard", "surrogate", "rate"}, optional + Overrides ``config.spike_mode`` for this call. + + Returns + ------- + outputs : dict + Per-step outputs stacked to ``(n_steps, ...)`` — see + ``make_scan_step`` (``force``, ``torque``, ``v_mn``, ``mn_spikes``, ...). + metadata : dict + ``{"spike_mode", "n_steps", "n_gii", "n_gib", "n_mn"}``. The active spike + mode is recorded here so downstream results carry provenance. + """ + mode = spike_mode if spike_mode is not None else config.spike_mode + if mode not in ("hard", "surrogate", "rate"): + raise ValueError( + f"spike_mode must be 'hard', 'surrogate', or 'rate'; got {mode!r}" + ) + scan_step = _build_scan_step(params, config, mode) + init_carry = _build_init_carry(params, config, key) + _, outputs = jax.lax.scan(scan_step, init_carry, inputs) + + # n_steps from any stacked input leaf. + any_leaf = jtu.tree_leaves(inputs)[0] + metadata = { + "spike_mode": mode, + "n_steps": int(any_leaf.shape[0]), + "n_gii": config.n_gii, + "n_gib": config.n_gib, + "n_mn": config.n_mn, + } + return outputs, metadata + + +# --------------------------------------------------------------------------- # +# JIT helper +# --------------------------------------------------------------------------- # + +def _is_static_scalar(x) -> bool: + """Python scalar (int/bool/str) with no array shape — must stay static for JIT.""" + return isinstance(x, (int, bool, str)) and not hasattr(x, "shape") + + +def compile_run(config: ClosedLoopConfig, params_template: dict, spike_mode: str = "hard"): + """Return a ``jax.jit``-compiled ``run_jax`` bound to ``config``. + + ``run_jax`` is not directly jittable (see module docstring). This factory does + the static/dynamic split: Python-scalar fields of ``params`` (e.g. the integer + Hill counts used in ``jnp.arange``) are baked static from ``params_template``, + all array leaves are traced, and ``config`` is closed over. Later calls must + keep the same integer/name fields as the template (you optimize the float + arrays, not the counts). + + Returns + ------- + call(params, inputs, key) -> outputs + Jitted; returns the outputs dict (metadata is static and omitted). + """ + leaves0, treedef = jtu.tree_flatten(params_template) + stat_mask = [_is_static_scalar(l) for l in leaves0] + static_leaves = [l if m else None for l, m in zip(leaves0, stat_mask)] + + @jax.jit + def _run(dyn_leaves, inputs, key): + leaves = [s if m else d for s, d, m in zip(static_leaves, dyn_leaves, stat_mask)] + full = jtu.tree_unflatten(treedef, leaves) + outputs, _meta = run_jax(full, config, inputs, key, spike_mode=spike_mode) + return outputs + + def call(params, inputs, key): + leaves = jtu.tree_flatten(params)[0] + dyn_leaves = [None if m else l for l, m in zip(leaves, stat_mask)] + return _run(dyn_leaves, inputs, key) + + return call + + +# --------------------------------------------------------------------------- # +# Gradient helpers +# --------------------------------------------------------------------------- # + +def _is_differentiable_leaf(x) -> bool: + """True for float scalars / floating arrays; False for ints, bools, callables.""" + if isinstance(x, bool): + return False + if isinstance(x, (int,)): + return False + if isinstance(x, float): + return True + dtype = getattr(x, "dtype", None) + if dtype is None: + return False + return jnp.issubdtype(dtype, jnp.floating) or jnp.issubdtype(dtype, jnp.complexfloating) + + +def partition_differentiable(params): + """Split a params PyTree into (differentiable, static) leaf lists + rebuild info. + + Integer/bool leaves (e.g. Hill's ``Ntype1``/``N``) are placed in ``static`` so + ``jax.grad`` — which rejects integer inputs — only differentiates float leaves. + + Returns + ------- + diff_leaves : list + Float leaves; ``None`` where the corresponding leaf is static. + static_leaves : list + Non-float leaves; ``None`` where the corresponding leaf is differentiable. + treedef, mask : rebuild info for :func:`_combine`. + """ + leaves, treedef = jtu.tree_flatten(params) + mask = [_is_differentiable_leaf(l) for l in leaves] + diff_leaves = [l if m else None for l, m in zip(leaves, mask)] + static_leaves = [None if m else l for l, m in zip(leaves, mask)] + return diff_leaves, static_leaves, treedef, mask + + +def _combine(diff_leaves, static_leaves, treedef, mask): + leaves = [d if m else s for d, s, m in zip(diff_leaves, static_leaves, mask)] + return jtu.tree_unflatten(treedef, leaves) + + +def value_and_grad_run( + loss_fn: Callable, + params: dict, + config: ClosedLoopConfig, + inputs: dict, + key, + spike_mode: str = "surrogate", + has_aux: bool = False, +): + """``jax.value_and_grad`` of ``loss_fn(run_jax(...))`` w.r.t. float params. + + Only floating leaves of ``params`` are differentiated (see + :func:`partition_differentiable`); integer leaves are held constant. The + default ``spike_mode="surrogate"`` keeps the forward trajectory identical to + the hard-spike model while providing a usable (biased) gradient through spikes. + + Parameters + ---------- + loss_fn : Callable + ``loss_fn(outputs) -> scalar``, or ``-> (scalar, aux)`` when ``has_aux``. + has_aux : bool + If True, ``loss_fn`` returns ``(value, aux)`` and this returns + ``(value, grads, aux)`` (``aux`` is not differentiated). + + Returns + ------- + (value, grads) — or ``(value, grads, aux)`` when ``has_aux`` — where ``grads`` + is a params-shaped PyTree whose non-differentiable leaves are ``None``. + """ + diff_leaves, static_leaves, treedef, mask = partition_differentiable(params) + + def f(diff): + full = _combine(diff, static_leaves, treedef, mask) + outputs, meta = run_jax(full, config, inputs, key, spike_mode=spike_mode) + return loss_fn(outputs) + + result = jax.value_and_grad(f, has_aux=has_aux)(diff_leaves) + if has_aux: + (value, aux), grad_leaves = result + return value, jtu.tree_unflatten(treedef, grad_leaves), aux + value, grad_leaves = result + return value, jtu.tree_unflatten(treedef, grad_leaves) diff --git a/myogen/simulator/jaxley/descending_drive_network.py b/myogen/simulator/jaxley/descending_drive_network.py new file mode 100644 index 00000000..eb786820 --- /dev/null +++ b/myogen/simulator/jaxley/descending_drive_network.py @@ -0,0 +1,284 @@ +"""Calibrated descending-drive network for the Jaxley NERLab pool. + +This is the Jaxley counterpart of +``myogen.simulator.neuron.descending_drive_network``. It deliberately uses +the same profile-card parameters and model semantics: + +* gamma-renewal descending drive implemented by hazard accumulation; +* equal fan-in and dendritic-transfer weight scaling; +* peak-normalized dendritic Exp2Syn kinetics (0.2/0.3 ms by default); +* live dendritic driving force and the production synaptic delay; +* live-voltage Gfluctdv conductances on the dendritic compartment. + +The random streams are backend-specific, so ordinary profile simulations are +expected to agree statistically rather than spike-for-spike. Exact paired +validation must continue to use the frozen event/connectivity artifact in +``experiments/twin/validation_protocol.py``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax +import jax.numpy as jnp +import jaxley as jx +import numpy as np +from jaxley.integrate import build_init_and_step_fn + +from myogen.simulator.jaxley.jax_models import gamma_init, gamma_step +from myogen.simulator.jaxley.populations import AlphaMN__Pool + + +@dataclass(frozen=True) +class DescendingDriveResult: + """Arrays returned by :func:`run_descending_drive`.""" + + voltage_mV: np.ndarray + dendrite_voltage_mV: np.ndarray + spikes: np.ndarray + dd_spikes: np.ndarray + connectivity: np.ndarray + weights_uS: np.ndarray + transfer: np.ndarray + + +def excitation_to_drive_hz(segments, threshold_hz, max_hz, timestep_ms=0.1, noise_seed=42): + """Map profile excitation segments to the calibrated noisy DD rate. + + Kept byte-for-byte equivalent in semantics to the NEURON helper while + avoiding a NEURON import in a JAX process. + """ + parts = [ + np.linspace(a, b, round(duration * 1000 / timestep_ms), endpoint=False) + for duration, a, b in segments + ] + excitation = np.concatenate(parts) + hz = np.where( + excitation > 0, + threshold_hz + excitation * (max_hz - threshold_hz), + 0.8 * threshold_hz, + ) + hz = hz + np.abs( + np.random.default_rng(noise_seed).normal(0, hz.max() * 0.01, size=hz.size) + ) + return excitation, hz + + +def equal_fan_in_connectivity(n_mn: int, n_dd: int, n_inputs: int, seed: int) -> np.ndarray: + """Return the profile network's target-by-source equal-fan-in matrix.""" + if not 1 <= n_inputs <= n_dd: + raise ValueError(f"n_inputs must be 1..{n_dd}, got {n_inputs}") + rng = np.random.default_rng(seed) + out = np.zeros((n_mn, n_dd), dtype=np.float32) + for i in range(n_mn): + out[i, rng.choice(n_dd, size=n_inputs, replace=False)] = 1.0 + return out + + +def _exp2_constants(dt_ms: float, tau1_ms: float, tau2_ms: float) -> tuple[float, float, float]: + if not 0 < tau1_ms < tau2_ms: + raise ValueError("Exp2Syn requires 0 < tau1_ms < tau2_ms") + t_peak = tau1_ms * tau2_ms / (tau2_ms - tau1_ms) * np.log(tau2_ms / tau1_ms) + norm = 1.0 / (np.exp(-t_peak / tau2_ms) - np.exp(-t_peak / tau1_ms)) + return np.exp(-dt_ms / tau1_ms), np.exp(-dt_ms / tau2_ms), norm + + +def _network(cells, n_steps: int, initial_voltage_mV: float): + net = jx.Network([cw.cell for cw in cells]) + placeholder = jnp.zeros(n_steps, dtype=jnp.float32) + for i in range(len(cells)): + net.cell(i).branch(0).loc(0.5).record("v", verbose=False) + for i in range(len(cells)): + net.cell(i).branch(1).loc(0.5).record("v", verbose=False) + for i in range(len(cells)): + net.cell(i).branch(1).loc(0.5).stimulate(placeholder, verbose=False) + net.set("v", float(initial_voltage_mV)) + net.init_states() + return net + + +def _dendritic_transfer(cells, dt_ms: float, initial_voltage_mV: float, + probe_nA: float = 0.01, settle_ms: float = 200.0) -> np.ndarray: + """Measure DC dendrite-to-soma transfer relative to soma injection.""" + n_steps = round(settle_ms / dt_ms) + + def steady(branch: int, amplitude: float): + for cw in cells: + cw.cell.delete_recordings() + cw.cell.delete_stimuli() + net = jx.Network([cw.cell for cw in cells]) + waveform = jnp.full(n_steps, amplitude, dtype=jnp.float32) + for i in range(len(cells)): + net.cell(i).branch(branch).loc(0.5).stimulate(waveform, verbose=False) + net.cell(i).branch(0).loc(0.5).record("v", verbose=False) + net.set("v", float(initial_voltage_mV)) + net.init_states() + v = np.asarray(jx.integrate(net, delta_t=dt_ms, t_max=settle_ms - dt_ms)) + return v[:, -1] + + baseline = steady(0, 0.0) + dend = steady(1, probe_nA) - baseline + soma = steady(0, probe_nA) - baseline + if np.any(~np.isfinite(dend)) or np.any(~np.isfinite(soma)) or np.any(soma == 0): + raise RuntimeError("non-finite dendritic-transfer measurement") + return dend / soma + + +def run_descending_drive( + thresholds, + drive_hz, + dd_params, + *, + timestep_ms: float = 0.1, + seed: int = 42, + initial_voltage_mV: float = -67.0, + tau1_ms: float = 0.2, + tau2_ms: float = 0.3, + e_exc_mV: float = 70.0, + e_noise_mV: float = 0.0, + i_noise_mV: float = -75.0, + external_event_delay_ms: float = 1.0, + synaptic_delay_ms: float = 1.0, +) -> DescendingDriveResult: + """Run the calibrated profile-card descending drive on Jaxley. + + ``initial_voltage_mV=-67`` reproduces the current NEURON profile runner. + Matched intrinsic validation may instead pass 0 mV explicitly. + """ + thresholds = np.asarray(thresholds, dtype=float) + drive_hz = np.asarray(drive_hz, dtype=np.float32) + if thresholds.ndim != 1 or drive_hz.ndim != 1: + raise ValueError("thresholds and drive_hz must be one-dimensional") + if not np.all(np.isfinite(thresholds)) or not np.all(np.isfinite(drive_hz)): + raise ValueError("thresholds and drive_hz must be finite") + + n_mn = thresholds.size + n_dd = int(dd_params["dd_neurons"]) + n_steps = drive_hz.size + n_inputs = round(float(dd_params["conn_probability"]) * n_dd) + shape = float(dd_params["gamma_shape"]) + weight_exponent = float(dd_params["network"]["weight_exponent"]) + base_weight = float(dd_params["synaptic_weight"]) + gf_scale = float(dd_params.get("gfluctdv_scale", 0.0)) + + pool = AlphaMN__Pool( + recruitment_thresholds__array=thresholds, + model="NERLab", + mode="active", + initial_voltage__mV=float(initial_voltage_mV), + spike_threshold__mV=50.0, + normalize_nerlab_thresholds=False, + ) + cells = list(pool) + transfer = _dendritic_transfer(cells, timestep_ms, initial_voltage_mV) + weights = base_weight * (transfer.max() / transfer) ** weight_exponent + connectivity = equal_fan_in_connectivity(n_mn, n_dd, n_inputs, seed) + weighted_conn_np = connectivity * weights[:, None] + + for cw in cells: + cw.cell.delete_recordings() + cw.cell.delete_stimuli() + net = _network(cells, n_steps, initial_voltage_mV) + net.to_jax() + init_fn, step_fn = build_init_and_step_fn(net) + states, params = init_fn(net.get_parameters()) + sim_dtype = states["v"].dtype + weighted_conn = jnp.asarray(weighted_conn_np, dtype=sim_dtype) + external_inds = net.external_inds.copy() + rec = jnp.asarray(net.recordings.rec_index.to_numpy(), dtype=jnp.int32) + rec_soma, rec_dend = rec[:n_mn], rec[n_mn:] + + a1, a2, exp2_norm = _exp2_constants(timestep_ms, tau1_ms, tau2_ms) + delay_steps = round((external_event_delay_ms + synaptic_delay_ms) / timestep_ms) + delay_len = max(delay_steps, 1) + + # Gfluctdv density values set by AlphaMN.insert_Gfluctdv, converted to uS. + area_cm2 = np.asarray([ + 2.0 * np.pi * (cw.cell.branch(1).nodes["radius"].to_numpy()[0] * 1e-4) + * (cw.cell.branch(1).nodes["length"].to_numpy()[0] * 1e-4) + for cw in cells + ], dtype=np.float32) + ge0 = jnp.asarray(1e-5 * gf_scale * area_cm2 * 1e6, dtype=sim_dtype) + gi0 = jnp.asarray(1e-5 * gf_scale * area_cm2 * 1e6, dtype=sim_dtype) + ge_sd = jnp.asarray(1.2e-5 * gf_scale * area_cm2 * 1e6, dtype=sim_dtype) + gi_sd = jnp.asarray(1.2e-5 * gf_scale * area_cm2 * 1e6, dtype=sim_dtype) + ou_decay = jnp.asarray(np.exp(-timestep_ms / 20.0), dtype=sim_dtype) + ou_scale = jnp.asarray(np.sqrt(1.0 - np.exp(-2.0 * timestep_ms / 20.0)), dtype=sim_dtype) + + dd_key, noise_key = jax.random.split(jax.random.PRNGKey(seed)) + carry0 = { + "states": states, + "prev_soma": jnp.full(n_mn, initial_voltage_mV, dtype=sim_dtype), + "prev_dend": jnp.full(n_mn, initial_voltage_mV, dtype=sim_dtype), + "dd": gamma_init(n_dd, shape, key=dd_key), + "delay": jnp.zeros((delay_len, n_dd), dtype=jnp.float32), + "delay_idx": jnp.int32(0), + "A": jnp.zeros(n_mn, dtype=sim_dtype), + "B": jnp.zeros(n_mn, dtype=sim_dtype), + "ge": ge0, + "gi": gi0, + "noise_key": noise_key, + } + + def body(carry, rate_t): + dd_state, emitted_bool = gamma_step( + carry["dd"], rate_t, timestep_ms, shape, mode="hard") + emitted = emitted_bool.astype(jnp.float32) + if delay_steps: + arriving = carry["delay"][carry["delay_idx"]] + delay = carry["delay"].at[carry["delay_idx"]].set(emitted) + delay_idx = (carry["delay_idx"] + 1) % delay_len + else: + arriving, delay, delay_idx = emitted, carry["delay"], carry["delay_idx"] + + increment = weighted_conn @ arriving.astype(sim_dtype) + A = jnp.asarray(a1, dtype=sim_dtype) * carry["A"] + jnp.asarray(exp2_norm, dtype=sim_dtype) * increment + B = jnp.asarray(a2, dtype=sim_dtype) * carry["B"] + jnp.asarray(exp2_norm, dtype=sim_dtype) * increment + conductance = B - A + + key, ke, ki = jax.random.split(carry["noise_key"], 3) + ge = ge0 + (carry["ge"] - ge0) * ou_decay + ge_sd * ou_scale * jax.random.normal(ke, (n_mn,), dtype=sim_dtype) + gi = gi0 + (carry["gi"] - gi0) * ou_decay + gi_sd * ou_scale * jax.random.normal(ki, (n_mn,), dtype=sim_dtype) + ge, gi = jnp.maximum(ge, 0.0), jnp.maximum(gi, 0.0) + + vd = carry["prev_dend"] + i_syn = conductance * (jnp.asarray(e_exc_mV, dtype=sim_dtype) - vd) + i_noise = ge * (jnp.asarray(e_noise_mV, dtype=sim_dtype) - vd) + gi * (jnp.asarray(i_noise_mV, dtype=sim_dtype) - vd) + new_states = step_fn( + carry["states"], params, {"i": i_syn + i_noise}, + external_inds, delta_t=timestep_ms, + ) + soma = new_states["v"][rec_soma] + dend = new_states["v"][rec_dend] + spikes = (carry["prev_soma"] < 50.0) & (soma >= 50.0) + new_carry = { + "states": new_states, "prev_soma": soma, "prev_dend": dend, + "dd": dd_state, "delay": delay, "delay_idx": delay_idx, + "A": A, "B": B, "ge": ge, "gi": gi, "noise_key": key, + } + return new_carry, (soma, dend, spikes.astype(jnp.float32), emitted) + + _, (voltage, dend_voltage, spikes, dd_spikes) = jax.jit( + lambda c, x: jax.lax.scan(body, c, x) + )(carry0, jnp.asarray(drive_hz)) + voltage, dend_voltage, spikes, dd_spikes = jax.device_get( + (voltage, dend_voltage, spikes, dd_spikes) + ) + return DescendingDriveResult( + voltage_mV=np.asarray(voltage), + dendrite_voltage_mV=np.asarray(dend_voltage), + spikes=np.asarray(spikes), + dd_spikes=np.asarray(dd_spikes), + connectivity=connectivity, + weights_uS=np.asarray(weights), + transfer=np.asarray(transfer), + ) + + +__all__ = [ + "DescendingDriveResult", + "equal_fan_in_connectivity", + "excitation_to_drive_hz", + "run_descending_drive", +] diff --git a/myogen/simulator/jaxley/emg.py b/myogen/simulator/jaxley/emg.py new file mode 100644 index 00000000..8b808114 --- /dev/null +++ b/myogen/simulator/jaxley/emg.py @@ -0,0 +1,130 @@ +""" +Differentiable EMG synthesis (JAX) for the MyoGen Jaxley backend. + +The surface / intramuscular EMG of a motor-unit pool is a linear superposition: + + EMG[channel, t] = Σ_MU (spike_train_MU ⋆ MUAP_{MU, channel})[t] + +where ``⋆`` is cross-correlation (``np.correlate(..., mode="same")``) and the MUAP +templates are **static** — they depend only on muscle geometry, tissue +conductivities and conduction velocity, not on the spike trains. Because the +templates are frozen, the whole map from spike trains to EMG is a convolution and +is trivially differentiable: gradients flow ``spikes → EMG`` with no Bessel / +volume-conductor machinery involved. + +These pure-JAX functions mirror the summation in +``core/emg/surface/surface_emg.py`` and ``core/emg/intramuscular/intramuscular_emg.py`` +but operate on JAX arrays so they compose with :func:`myogen.simulator.jaxley.closed_loop.run_jax` +and ``jax.grad``. The existing Neo-returning methods remain the scientific +default; these are the differentiable runtime form. + +Notes +----- +* Inactive motor units simply contribute zero (their spike trains are all-zero), + so no explicit active-MU masking is needed — pass the full pool. +* MUAP templates must already be resampled to the spike-train timestep. Use + :func:`resample_muaps` for that (differentiable ``jnp.interp``). +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp + +__all__ = [ + "surface_emg_jax", + "intramuscular_emg_jax", + "resample_muaps", +] + + +def _correlate_same(signal, template): + """Cross-correlation with ``mode="same"`` — matches ``np.correlate``. + + ``signal`` shape ``(T,)``, ``template`` shape ``(Tm,)`` → output ``(T,)``. + """ + return jnp.correlate(signal, template, mode="same") + + +def surface_emg_jax(spike_trains, muap_shapes): + """Differentiable surface EMG for a 2-D electrode grid. + + Parameters + ---------- + spike_trains : array, shape ``(n_pools, n_mu, T)`` + Per-MU spike (or continuous rate/surrogate) trains at the EMG timestep. + muap_shapes : array, shape ``(n_mu, n_rows, n_cols, Tm)`` + Static MUAP templates per grid electrode, resampled to the spike-train + timestep (see :func:`resample_muaps`). + + Returns + ------- + array, shape ``(n_pools, n_rows, n_cols, T)`` + ``emg[p, r, c] = Σ_mu correlate(spike[p, mu], muap[mu, r, c], "same")``. + """ + # (n_mu, n_rows, n_cols, Tm) → (n_rows, n_cols, n_mu, Tm) + muap_grid = jnp.transpose(muap_shapes, (1, 2, 0, 3)) + + def channel_emg(spike_pool, muap_channel): + # spike_pool: (n_mu, T); muap_channel: (n_mu, Tm) → (T,) + corrs = jax.vmap(_correlate_same)(spike_pool, muap_channel) + return jnp.sum(corrs, axis=0) + + def pool_emg(spike_pool): + # vmap channel_emg over (n_rows, n_cols) + return jax.vmap(jax.vmap(lambda mc: channel_emg(spike_pool, mc)))(muap_grid) + + return jax.vmap(pool_emg)(spike_trains) + + +def intramuscular_emg_jax(spike_trains, muap_shapes): + """Differentiable intramuscular EMG for a 1-D electrode set. + + Parameters + ---------- + spike_trains : array, shape ``(n_pools, n_mu, T)`` + muap_shapes : array, shape ``(n_mu, n_elec, Tm)`` + Static per-electrode MUAP templates at the spike-train timestep. + + Returns + ------- + array, shape ``(n_pools, n_elec, T)`` + ``emg[p, e] = Σ_mu correlate(spike[p, mu], muap[mu, e], "same")``. + """ + muap_grid = jnp.transpose(muap_shapes, (1, 0, 2)) # (n_elec, n_mu, Tm) + + def channel_emg(spike_pool, muap_channel): + corrs = jax.vmap(_correlate_same)(spike_pool, muap_channel) + return jnp.sum(corrs, axis=0) + + def pool_emg(spike_pool): + return jax.vmap(lambda mc: channel_emg(spike_pool, mc))(muap_grid) + + return jax.vmap(pool_emg)(spike_trains) + + +def resample_muaps(muap_shapes, muap_dt_s: float, target_dt_s: float): + """Resample MUAP templates along their time axis with differentiable interp. + + Mirrors the ``np.interp`` resampling in the Neo EMG path but with + ``jnp.interp`` so the operation stays in the JAX graph (useful if the + templates themselves become differentiable via the volume-conductor port). + + Parameters + ---------- + muap_shapes : array, shape ``(..., Tm)`` Templates on the MUAP grid. + muap_dt_s : float Template sample period [s]. + target_dt_s : float Desired sample period [s] (the spike-train timestep). + + Returns + ------- + array, shape ``(..., T)`` Resampled templates. + """ + Tm = muap_shapes.shape[-1] + duration = Tm * muap_dt_s + xp = jnp.arange(Tm, dtype=jnp.float32) * jnp.float32(muap_dt_s) + x = jnp.arange(0.0, duration, target_dt_s, dtype=jnp.float32) + + flat = muap_shapes.reshape(-1, Tm) + resampled = jax.vmap(lambda fp: jnp.interp(x, xp, fp))(flat) + return resampled.reshape(*muap_shapes.shape[:-1], x.shape[0]) diff --git a/myogen/simulator/jaxley/integrate.py b/myogen/simulator/jaxley/integrate.py new file mode 100644 index 00000000..ba19690c --- /dev/null +++ b/myogen/simulator/jaxley/integrate.py @@ -0,0 +1,982 @@ +""" +Jaxley Integration Module for MyoGen. + +This module provides the core simulation infrastructure that connects +Jaxley cells to the MyoGen framework, including: +- Current injection with time-varying waveforms +- Spike detection and recording +- Multi-cell network simulation +- JIT-compiled simulation loops + +This is the bridge between Jaxley's `jx.integrate()` and MyoGen's API. +""" + +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from dataclasses import dataclass, field +import numpy as np +import jax +import jax.numpy as jnp +from functools import partial + +import jaxley as jx + + +# ============================================================================= +# CURRENT STIMULI +# ============================================================================= + +def step_current( + delay: float, + duration: float, + amplitude: float, + dt: float, + t_max: float, +) -> jnp.ndarray: + """ + Create a step current injection waveform. + + Parameters + ---------- + delay : float + Time before current onset (ms). + duration : float + Duration of current pulse (ms). + amplitude : float + Current amplitude (nA). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + + Returns + ------- + jnp.ndarray + Current waveform array. + """ + return jx.step_current(delay, duration, amplitude, dt, t_max) + + +def ramp_current( + delay: float, + duration: float, + start_amp: float, + end_amp: float, + dt: float, + t_max: float, +) -> jnp.ndarray: + """ + Create a linearly ramping current injection waveform. + + Parameters + ---------- + delay : float + Time before ramp onset (ms). + duration : float + Duration of ramp (ms). + start_amp : float + Starting amplitude (nA). + end_amp : float + Ending amplitude (nA). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + + Returns + ------- + jnp.ndarray + Current waveform array. + """ + n_steps = int(t_max / dt) + 1 + current = jnp.zeros(n_steps) + + start_idx = int(delay / dt) + end_idx = int((delay + duration) / dt) + ramp_len = end_idx - start_idx + + if ramp_len > 0: + ramp = jnp.linspace(start_amp, end_amp, ramp_len) + current = current.at[start_idx:end_idx].set(ramp) + + return current + + +def noisy_current( + mean: float, + std: float, + dt: float, + t_max: float, + seed: int = 0, +) -> jnp.ndarray: + """ + Create a noisy (Gaussian) current injection waveform. + + Parameters + ---------- + mean : float + Mean current amplitude (nA). + std : float + Standard deviation of noise (nA). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + seed : int + Random seed. + + Returns + ------- + jnp.ndarray + Noisy current waveform array. + """ + n_steps = int(t_max / dt) + 1 + key = jax.random.PRNGKey(seed) + noise = jax.random.normal(key, shape=(n_steps,)) + return mean + std * noise + + +def custom_current( + waveform: np.ndarray, + dt: float, + t_max: float, +) -> jnp.ndarray: + """ + Create current from a custom waveform array. + + Parameters + ---------- + waveform : np.ndarray + Custom current waveform (nA). Will be resampled if needed. + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + + Returns + ------- + jnp.ndarray + Current waveform array. + """ + n_steps = int(t_max / dt) + 1 + + if len(waveform) == n_steps: + return jnp.array(waveform) + else: + # Resample to match simulation timesteps + old_t = np.linspace(0, t_max, len(waveform)) + new_t = np.linspace(0, t_max, n_steps) + resampled = np.interp(new_t, old_t, waveform) + return jnp.array(resampled) + + +def sinusoidal_current( + frequency: float, + amplitude: float, + offset: float, + dt: float, + t_max: float, + phase: float = 0.0, +) -> jnp.ndarray: + """ + Create a sinusoidal current injection waveform. + + Parameters + ---------- + frequency : float + Frequency in Hz. + amplitude : float + Amplitude of oscillation (nA). + offset : float + DC offset / mean current (nA). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + phase : float + Initial phase (radians). + + Returns + ------- + jnp.ndarray + Current waveform array. + """ + n_steps = int(t_max / dt) + 1 + t = jnp.arange(n_steps) * dt / 1000.0 # Convert to seconds for Hz + return offset + amplitude * jnp.sin(2 * jnp.pi * frequency * t + phase) + + +# ============================================================================= +# SPIKE DETECTION +# ============================================================================= + +def detect_spikes( + voltage: jnp.ndarray, + threshold: float = 0.0, + dt: float = 0.025, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Detect spikes from voltage trace using threshold crossing. + + Parameters + ---------- + voltage : jnp.ndarray + Voltage trace, shape (n_timesteps,) or (n_recordings, n_timesteps). + threshold : float + Spike threshold in mV. + dt : float + Time step in ms. + + Returns + ------- + spike_times : jnp.ndarray + Times of detected spikes (ms). + spike_indices : jnp.ndarray + Indices of spike times in voltage array. + """ + # Handle 1D and 2D arrays + if voltage.ndim == 1: + voltage = voltage[jnp.newaxis, :] + + # Find threshold crossings (upward) + above = voltage > threshold + below_before = jnp.concatenate([ + jnp.zeros((voltage.shape[0], 1), dtype=bool), + voltage[:, :-1] <= threshold + ], axis=1) + + crossings = above & below_before + + # Get spike indices and times + spike_indices = jnp.where(crossings) + spike_times = spike_indices[1] * dt + + return spike_times, spike_indices + + +def compute_firing_rate( + spike_times: jnp.ndarray, + t_max: float, + window: float = 100.0, + dt: float = 1.0, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Compute instantaneous firing rate from spike times. + + Parameters + ---------- + spike_times : jnp.ndarray + Spike times in ms. + t_max : float + Total time in ms. + window : float + Smoothing window in ms. + dt : float + Output time resolution in ms. + + Returns + ------- + times : jnp.ndarray + Time points. + rates : jnp.ndarray + Firing rates in Hz. + """ + times = jnp.arange(0, t_max, dt) + rates = jnp.zeros(len(times)) + half_window = window / 2 + + # Count spikes in sliding window + for i, t in enumerate(times): + count = jnp.sum((spike_times >= t - half_window) & (spike_times < t + half_window)) + rates = rates.at[i].set(count * 1000.0 / window) + + return times, rates + + +def compute_isi(spike_times: jnp.ndarray) -> Dict[str, float]: + """ + Compute interspike interval statistics. + + Parameters + ---------- + spike_times : jnp.ndarray + Spike times in ms. + + Returns + ------- + dict + ISI statistics: mean, std, cv, mean_rate. + """ + if len(spike_times) < 2: + return { + "mean_isi": float('nan'), + "std_isi": float('nan'), + "cv_isi": float('nan'), + "mean_rate": 0.0, + } + + isis = jnp.diff(jnp.sort(spike_times)) + mean_isi = float(jnp.mean(isis)) + std_isi = float(jnp.std(isis)) + + return { + "mean_isi": mean_isi, + "std_isi": std_isi, + "cv_isi": std_isi / mean_isi if mean_isi > 0 else float('nan'), + "mean_rate": 1000.0 / mean_isi if mean_isi > 0 else 0.0, + } + + +def compute_firing_rate_from_spikes( + spike_times: jnp.ndarray, + t_start: float = 0.0, + t_stop: Optional[float] = None, +) -> float: + """ + Compute mean firing rate from spike times. + + Parameters + ---------- + spike_times : jnp.ndarray + Spike times in ms. + t_start : float + Start time for rate calculation (ms). + t_stop : float, optional + Stop time for rate calculation (ms). If None, uses max spike time. + + Returns + ------- + float + Mean firing rate in Hz. + """ + if len(spike_times) == 0: + return 0.0 + + # Filter to time window + if t_stop is None: + t_stop = float(jnp.max(spike_times)) + + in_window = (spike_times >= t_start) & (spike_times <= t_stop) + n_spikes = int(jnp.sum(in_window)) + + duration_s = (t_stop - t_start) / 1000.0 # Convert ms to s + + if duration_s <= 0: + return 0.0 + + return n_spikes / duration_s + + +def compute_isi_cv(spike_times: jnp.ndarray) -> float: + """ + Compute coefficient of variation (CV) of interspike intervals. + + CV = std(ISI) / mean(ISI) + + Parameters + ---------- + spike_times : jnp.ndarray + Spike times in ms. + + Returns + ------- + float + CV of ISI. Returns NaN if fewer than 2 spikes. + """ + if len(spike_times) < 2: + return float('nan') + + isis = jnp.diff(jnp.sort(spike_times)) + mean_isi = float(jnp.mean(isis)) + std_isi = float(jnp.std(isis)) + + if mean_isi <= 0: + return float('nan') + + return std_isi / mean_isi + + +# ============================================================================= +# SINGLE CELL SIMULATION +# ============================================================================= + +@dataclass +class CellSimulationResult: + """Container for single cell simulation results.""" + time: jnp.ndarray + voltage: jnp.ndarray + current: jnp.ndarray + spike_times: jnp.ndarray + spike_count: int + firing_rate: float + isi_stats: Dict[str, float] + + @property + def has_spikes(self) -> bool: + return self.spike_count > 0 + + +def simulate_cell( + cell: jx.Cell, + current: Optional[jnp.ndarray] = None, + dt: float = 0.025, + t_max: float = 100.0, + record_loc: Tuple[int, float] = (0, 0.0), + stim_loc: Tuple[int, float] = (0, 0.0), + spike_threshold: float = 0.0, + initial_voltage: float = -70.0, +) -> CellSimulationResult: + """ + Simulate a single Jaxley cell with current injection. + + Parameters + ---------- + cell : jx.Cell + Jaxley cell to simulate. + current : jnp.ndarray, optional + Current injection waveform (nA). If None, no current is injected. + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + record_loc : tuple + (branch_idx, location) for voltage recording. + stim_loc : tuple + (branch_idx, location) for current injection. + spike_threshold : float + Threshold for spike detection (mV). + initial_voltage : float + Initial membrane voltage (mV). + + Returns + ------- + CellSimulationResult + Simulation results including voltage, spikes, etc. + """ + # Clear previous recordings and stimuli + cell.delete_recordings() + cell.delete_stimuli() + + # Set initial voltage + cell.set("v", initial_voltage) + + # Set up recording + branch_idx, loc = record_loc + cell.branch(branch_idx).loc(loc).record("v") + + # Apply stimulus if provided + if current is not None: + stim_branch, stim_loc_val = stim_loc + cell.branch(stim_branch).loc(stim_loc_val).stimulate(current) + else: + current = jnp.zeros(int(t_max / dt) + 1) + + # Run simulation + voltages = jx.integrate(cell, delta_t=dt, t_max=t_max) + + # Extract voltage trace (shape: n_recordings x n_timesteps) + voltage = voltages[0] if voltages.ndim > 1 else voltages + + # Create time vector + time = jnp.arange(0, t_max + dt, dt)[:len(voltage)] + + # Detect spikes + spike_times, _ = detect_spikes(voltage, threshold=spike_threshold, dt=dt) + spike_count = len(spike_times) + + # Compute statistics + duration_s = t_max / 1000.0 + firing_rate = spike_count / duration_s if duration_s > 0 else 0.0 + isi_stats = compute_isi(spike_times) + + return CellSimulationResult( + time=time, + voltage=voltage, + current=current[:len(time)], + spike_times=spike_times, + spike_count=spike_count, + firing_rate=firing_rate, + isi_stats=isi_stats, + ) + + +def simulate_cell_with_data_stimulate( + cell: jx.Cell, + current: jnp.ndarray, + dt: float = 0.025, + t_max: float = 100.0, + stim_loc: Tuple[int, float] = (0, 0.0), + params: Optional[Dict] = None, +) -> jnp.ndarray: + """ + Simulate cell using data_stimulate (JIT-compatible). + + This version uses data_stimulate which is compatible with + JAX transformations (jit, grad, vmap). + + Parameters + ---------- + cell : jx.Cell + Jaxley cell to simulate. + current : jnp.ndarray + Current injection waveform (nA). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + stim_loc : tuple + (branch_idx, location) for current injection. + params : dict, optional + Optional parameters to pass to integrate. + + Returns + ------- + jnp.ndarray + Voltage recordings. + """ + # Set up data stimulus + data_stimuli = None + stim_branch, stim_loc_val = stim_loc + data_stimuli = cell.branch(stim_branch).loc(stim_loc_val).data_stimulate( + current, data_stimuli + ) + + # Run simulation + if params is not None: + return jx.integrate(cell, params=params, data_stimuli=data_stimuli, + delta_t=dt, t_max=t_max) + else: + return jx.integrate(cell, data_stimuli=data_stimuli, delta_t=dt, t_max=t_max) + + +# ============================================================================= +# POPULATION SIMULATION +# ============================================================================= + +@dataclass +class PopulationSimulationResult: + """Container for population simulation results.""" + time: jnp.ndarray + voltages: Dict[int, jnp.ndarray] + spike_times: Dict[int, jnp.ndarray] + spike_ids: jnp.ndarray + all_spike_times: jnp.ndarray + total_spike_count: int + mean_firing_rate: float + + def get_neuron_spikes(self, neuron_id: int) -> jnp.ndarray: + """Get spike times for a specific neuron.""" + return self.spike_times.get(neuron_id, jnp.array([])) + + def get_neuron_voltage(self, neuron_id: int) -> jnp.ndarray: + """Get voltage trace for a specific neuron.""" + return self.voltages.get(neuron_id, jnp.array([])) + + +def simulate_population( + cells: List[jx.Cell], + currents: Optional[List[jnp.ndarray]] = None, + dt: float = 0.025, + t_max: float = 100.0, + spike_threshold: float = 0.0, + record_voltage: bool = True, +) -> PopulationSimulationResult: + """ + Simulate a population of Jaxley cells. + + Parameters + ---------- + cells : list of jx.Cell + List of Jaxley cells to simulate. + currents : list of jnp.ndarray, optional + Current injection for each cell. If None, no current is injected. + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + spike_threshold : float + Threshold for spike detection (mV). + record_voltage : bool + Whether to record full voltage traces. + + Returns + ------- + PopulationSimulationResult + Population simulation results. + """ + n_cells = len(cells) + + # Handle currents + if currents is None: + n_steps = int(t_max / dt) + 1 + currents = [jnp.zeros(n_steps) for _ in range(n_cells)] + + # Storage + voltages = {} + spike_times = {} + all_spikes = [] + all_ids = [] + + # Simulate each cell + for i, (cell, current) in enumerate(zip(cells, currents)): + result = simulate_cell( + cell=cell, + current=current, + dt=dt, + t_max=t_max, + spike_threshold=spike_threshold, + ) + + if record_voltage: + voltages[i] = result.voltage + + spike_times[i] = result.spike_times + all_spikes.extend(result.spike_times.tolist()) + all_ids.extend([i] * len(result.spike_times)) + + # Create time vector + time = jnp.arange(0, t_max + dt, dt) + + # Compute stats + total_spikes = len(all_spikes) + duration_s = t_max / 1000.0 + mean_rate = (total_spikes / n_cells / duration_s) if n_cells > 0 and duration_s > 0 else 0.0 + + return PopulationSimulationResult( + time=time, + voltages=voltages, + spike_times=spike_times, + spike_ids=jnp.array(all_ids), + all_spike_times=jnp.array(all_spikes), + total_spike_count=total_spikes, + mean_firing_rate=mean_rate, + ) + + +# ============================================================================= +# NETWORK SIMULATION WITH SYNAPSES +# ============================================================================= + +@dataclass +class SynapticConnection: + """Specification for a synaptic connection.""" + source_cell_idx: int + target_cell_idx: int + target_branch: int = 0 + target_loc: float = 0.5 + weight: float = 0.001 # uS + delay: float = 1.0 # ms + reversal: float = 0.0 # mV (0 for excitatory, -75 for inhibitory) + tau1: float = 0.2 # ms (rise time) + tau2: float = 2.0 # ms (decay time) + + +class NetworkSimulator: + """ + Simulator for networks of Jaxley cells with synaptic connections. + + This provides a manual time-stepping approach for network simulation + where spikes from source cells activate synapses on target cells. + """ + + def __init__( + self, + cells: List[jx.Cell], + connections: List[SynapticConnection], + dt: float = 0.025, + ): + self.cells = cells + self.connections = connections + self.dt = dt + self.n_cells = len(cells) + + # State tracking + self._time = 0.0 + self._voltages = {i: -70.0 for i in range(self.n_cells)} + self._prev_voltages = {i: -70.0 for i in range(self.n_cells)} + self._pending_spikes = [] # (time, connection_idx) + + # Records + self.spike_times = {i: [] for i in range(self.n_cells)} + self.voltage_history = {i: [] for i in range(self.n_cells)} + + def reset(self): + """Reset simulation state.""" + self._time = 0.0 + self._voltages = {i: -70.0 for i in range(self.n_cells)} + self._prev_voltages = {i: -70.0 for i in range(self.n_cells)} + self._pending_spikes = [] + self.spike_times = {i: [] for i in range(self.n_cells)} + self.voltage_history = {i: [] for i in range(self.n_cells)} + + def _detect_spike(self, cell_idx: int, threshold: float = 0.0) -> bool: + """Check if cell just crossed threshold (upward).""" + v = self._voltages[cell_idx] + v_prev = self._prev_voltages[cell_idx] + return v > threshold and v_prev <= threshold + + def _schedule_synaptic_events(self, source_idx: int): + """Schedule synaptic events from a spiking source cell.""" + for conn_idx, conn in enumerate(self.connections): + if conn.source_cell_idx == source_idx: + event_time = self._time + conn.delay + self._pending_spikes.append((event_time, conn_idx)) + + def _process_synaptic_events(self) -> Dict[int, float]: + """Process pending synaptic events and return synaptic currents.""" + synaptic_currents = {i: 0.0 for i in range(self.n_cells)} + + # Get events that should be processed now + remaining = [] + for event_time, conn_idx in self._pending_spikes: + if event_time <= self._time: + conn = self.connections[conn_idx] + target = conn.target_cell_idx + + # Add synaptic conductance as current + # I_syn = g * (V - E_rev) + v = self._voltages[target] + i_syn = conn.weight * (v - conn.reversal) + synaptic_currents[target] += i_syn + else: + remaining.append((event_time, conn_idx)) + + self._pending_spikes = remaining + return synaptic_currents + + def step( + self, + external_currents: Optional[Dict[int, float]] = None, + spike_threshold: float = 0.0, + ): + """ + Advance simulation by one time step. + + Parameters + ---------- + external_currents : dict, optional + External current for each cell {cell_idx: current_nA}. + spike_threshold : float + Threshold for spike detection (mV). + """ + if external_currents is None: + external_currents = {} + + # Process synaptic events + synaptic_currents = self._process_synaptic_events() + + # Update each cell + for i in range(self.n_cells): + # Save previous voltage + self._prev_voltages[i] = self._voltages[i] + + # Total current + i_ext = external_currents.get(i, 0.0) + i_syn = synaptic_currents.get(i, 0.0) + total_current = i_ext + i_syn + + # Note: In a full implementation, this would integrate the cell + # using Jaxley. For now, we use a placeholder. + # The actual integration would be: + # voltages = jx.integrate(self.cells[i], ...) + + # Detect spikes + if self._detect_spike(i, spike_threshold): + self.spike_times[i].append(self._time) + self._schedule_synaptic_events(i) + + # Record voltage + self.voltage_history[i].append(self._voltages[i]) + + # Advance time + self._time += self.dt + + def run( + self, + t_max: float, + external_currents: Optional[Dict[int, jnp.ndarray]] = None, + spike_threshold: float = 0.0, + ) -> PopulationSimulationResult: + """ + Run network simulation. + + Parameters + ---------- + t_max : float + Total simulation time (ms). + external_currents : dict, optional + Time-varying currents {cell_idx: current_array}. + spike_threshold : float + Threshold for spike detection (mV). + + Returns + ------- + PopulationSimulationResult + Simulation results. + """ + self.reset() + n_steps = int(t_max / self.dt) + + for step in range(n_steps): + # Get current for this step + step_currents = {} + if external_currents: + for cell_idx, current in external_currents.items(): + if step < len(current): + step_currents[cell_idx] = float(current[step]) + + self.step(step_currents, spike_threshold) + + # Compile results + time = jnp.arange(0, t_max, self.dt) + voltages = {i: jnp.array(v) for i, v in self.voltage_history.items()} + spike_times_dict = {i: jnp.array(st) for i, st in self.spike_times.items()} + + all_spikes = [] + all_ids = [] + for i, st in spike_times_dict.items(): + all_spikes.extend(st.tolist()) + all_ids.extend([i] * len(st)) + + total = len(all_spikes) + mean_rate = (total / self.n_cells / (t_max / 1000.0)) if self.n_cells > 0 else 0.0 + + return PopulationSimulationResult( + time=time, + voltages=voltages, + spike_times=spike_times_dict, + spike_ids=jnp.array(all_ids), + all_spike_times=jnp.array(all_spikes), + total_spike_count=total, + mean_firing_rate=mean_rate, + ) + + +# ============================================================================= +# JIT-COMPILED SIMULATION FUNCTIONS +# ============================================================================= + +@partial(jax.jit, static_argnums=(2, 3)) +def jit_simulate_cell( + cell_params: Dict, + current: jnp.ndarray, + dt: float, + t_max: float, +) -> jnp.ndarray: + """ + JIT-compiled cell simulation (for parameter optimization). + + Note: This requires the cell to be configured with trainable parameters. + + Parameters + ---------- + cell_params : dict + Cell parameters (from cell.get_parameters()). + current : jnp.ndarray + Current injection waveform. + dt : float + Time step. + t_max : float + Total simulation time. + + Returns + ------- + jnp.ndarray + Voltage recordings. + """ + # This is a placeholder - actual implementation depends on cell setup + # The real implementation would use jx.integrate with params + pass + + +def sweep_current_amplitudes( + cell: jx.Cell, + amplitudes: List[float], + delay: float = 10.0, + duration: float = 500.0, + dt: float = 0.025, + t_max: float = 600.0, +) -> List[CellSimulationResult]: + """ + Sweep current injection amplitudes (FI curve measurement). + + Parameters + ---------- + cell : jx.Cell + Jaxley cell to simulate. + amplitudes : list of float + Current amplitudes to test (nA). + delay : float + Current onset delay (ms). + duration : float + Current duration (ms). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + + Returns + ------- + list of CellSimulationResult + Results for each amplitude. + """ + results = [] + + for amp in amplitudes: + current = step_current(delay, duration, amp, dt, t_max) + result = simulate_cell( + cell=cell, + current=current, + dt=dt, + t_max=t_max, + ) + results.append(result) + + return results + + +def compute_fi_curve( + cell: jx.Cell, + amplitudes: List[float], + delay: float = 100.0, + duration: float = 500.0, + dt: float = 0.025, + t_max: float = 700.0, +) -> Tuple[np.ndarray, np.ndarray]: + """ + Compute frequency-current (FI) curve. + + Parameters + ---------- + cell : jx.Cell + Jaxley cell to characterize. + amplitudes : list of float + Current amplitudes to test (nA). + delay : float + Current onset delay (ms). + duration : float + Current duration (ms). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + + Returns + ------- + currents : np.ndarray + Current amplitudes (nA). + rates : np.ndarray + Firing rates (Hz). + """ + results = sweep_current_amplitudes( + cell=cell, + amplitudes=amplitudes, + delay=delay, + duration=duration, + dt=dt, + t_max=t_max, + ) + + currents = np.array(amplitudes) + rates = np.array([r.firing_rate for r in results]) + + return currents, rates diff --git a/myogen/simulator/jaxley/jax_models.py b/myogen/simulator/jaxley/jax_models.py new file mode 100644 index 00000000..1639e4a4 --- /dev/null +++ b/myogen/simulator/jaxley/jax_models.py @@ -0,0 +1,1185 @@ +""" +JAX-functional physiological models for closed-loop neuromuscular simulation. + +Provides pure JAX step functions for muscle spindle, GTO, Hill muscle, and +joint dynamics. All functions are ``lax.scan``-compatible: no Python side +effects, no mutable class state — state is a plain dict of JAX arrays. + +Pattern +------- +:: + + state = model_init(...) + new_state, outputs = model_step(state, inputs, params) + +These functions are intended for two use cases: + +1. **Python for loop** (Phase "now"): call each ``model_step`` inside a + regular Python loop alongside Jaxley's ``step_fn``. The NumPy-based + ``HillModel`` / ``SpindleModel`` wrappers can be replaced one at a time. + +2. **``lax.scan``** (Phase "later"): combine all ``model_step`` functions and + Jaxley's ``step_fn`` inside a single ``jax.lax.scan`` body for full GPU + acceleration and automatic differentiation through the closed loop. + +References +---------- +- Mileusnic et al. (2006) — spindle model +- Lin & Crago (2002) — GTO model +- Fuglevand et al. (1993) — motor unit force model +- Hill (1938) — muscle mechanics +""" + +import numpy as np +import jax +import jax.numpy as jnp +from functools import partial + + +# ============================================================================ +# SPIKE DETECTION — hard / surrogate-gradient / rate modes +# ============================================================================ +# +# Spike detection is an upward threshold crossing: +# spike = (v > v_th) & (prev_v <= v_th) +# The boolean cast that drives downstream conductances/force has zero gradient, +# which severs autodiff through the spike train. ``spike_detect`` provides three +# modes selectable at build time: +# +# "hard" — exact boolean crossing (bit-identical to the original model). +# Scientific default; no usable gradient through the spike. +# "surrogate" — forward pass returns the exact hard crossing; the backward pass +# substitutes a smooth fast-sigmoid surrogate derivative on the +# crossing margin (SuperSpike, Zenke & Ganguli 2018). Straight- +# through, so the forward trajectory is unchanged. +# "rate" — replaces the discrete crossing with a continuous sigmoid of the +# margin. Fully smooth but the forward output differs from the +# spiking model (use for optimization where that is acceptable). +# +# ``SURROGATE_BETA`` controls the surrogate/rate slope (larger = sharper). + +SURROGATE_BETA = 10.0 + + +@partial(jax.custom_jvp, nondiff_argnums=()) +def _straight_through_step(margin): + """Heaviside step of ``margin`` with a fast-sigmoid surrogate gradient. + + Forward: ``(margin > 0)`` as float32 — exact. Backward (custom JVP): the + fast-sigmoid derivative ``beta / (1 + beta*|margin|)**2`` so gradients flow + through the crossing without changing the forward value. + """ + return (margin > 0.0).astype(margin.dtype) + + +@_straight_through_step.defjvp +def _straight_through_step_jvp(primals, tangents): + (margin,) = primals + (dm,) = tangents + primal_out = _straight_through_step(margin) + beta = jnp.asarray(SURROGATE_BETA, dtype=margin.dtype) + surrogate = beta / (1.0 + beta * jnp.abs(margin)) ** 2 + return primal_out, surrogate * dm + + +def spike_detect(v, prev_v, v_th, mode: str = "hard"): + """Detect an upward crossing of ``v_th`` between ``prev_v`` and ``v``. + + Parameters + ---------- + v, prev_v : arrays Current and previous-step membrane voltage. + v_th : float Threshold in the cells' voltage frame. + mode : {"hard","surrogate","rate"} + + Returns + ------- + spikes : array + ``mode="hard"`` → bool crossing (backward-compatible). + ``mode="surrogate"`` → float32, forward-identical to hard, smooth grad. + ``mode="rate"`` → float32 continuous sigmoid of the crossing margin. + """ + if mode == "hard": + return (v > v_th) & (prev_v <= v_th) + margin = v - v_th + if mode == "surrogate": + # Gate the smooth crossing to the "upward" event so forward matches hard. + rising = (prev_v <= v_th).astype(v.dtype) + return _straight_through_step(margin) * rising + if mode == "rate": + beta = jnp.asarray(SURROGATE_BETA, dtype=v.dtype) + rising = jax.nn.sigmoid(beta * (v_th - prev_v)) + return jax.nn.sigmoid(beta * margin) * rising + raise ValueError(f"unknown spike_mode {mode!r}") + + +# ============================================================================ +# MUSCLE SPINDLE — Mileusnic et al. (2006) +# ============================================================================ + +def spindle_init() -> dict: + """Return zeroed spindle state at rest.""" + return { + "a_bag1": jnp.float32(0.0), + "a_bag2": jnp.float32(0.0), + "T": jnp.zeros(3, dtype=jnp.float32), # [Bag1, Bag2, Chain] tensions + "dT": jnp.zeros(3, dtype=jnp.float32), # tension rates + } + + +def spindle_params_from_dict(spindle_params: dict) -> dict: + """ + Convert a SpindleModel parameter dict to a JAX-ready params dict. + + Parameters + ---------- + spindle_params : dict + As returned by ``SpindleModel.create_default_spindle_parameters()``. + + Returns + ------- + dict + Same keys, all values cast to Python float (scalar JAX-friendly). + """ + return {k: float(v) for k, v in spindle_params.items()} + + +def spindle_step(state: dict, L: float, V: float, A: float, + gd: float, gs: float, dt_s: float, p: dict): + """ + One step of the Mileusnic et al. (2006) muscle spindle model. + + Parameters + ---------- + state : dict + ``{a_bag1, a_bag2, T: (3,), dT: (3,)}`` + L, V, A : float + Muscle length [L0], velocity [L0/s], acceleration [L0/s²]. + gd, gs : float + Dynamic and static gamma fusimotor drive [Hz]. + dt_s : float + Timestep in seconds. + p : dict + Spindle parameters (from ``spindle_params_from_dict``). + + Returns + ------- + new_state : dict + (Ia, II) : tuple of float — primary and secondary afferent rates [Hz] + """ + a_bag1 = state["a_bag1"] + a_bag2 = state["a_bag2"] + T = state["T"] + dT = state["dT"] + + P = p["P"] + + # --- Fusimotor activations (Hill-type saturation, Eq. 4-5) --- + def _act(g, f0): + return jnp.where(g > 0.0, g**P / (g**P + f0**P), 0.0) + + target_bag1 = _act(gd, p["fBag1"]) + target_bag2 = _act(gs, p["fBag2"]) + a_chain = _act(gs, p["fChain"]) + + # RK4 for bag activation ODE: da/dt = (target - a) / tau + def _bag_rk4(a_prev, target, tau): + k1 = (target - a_prev) / tau + k2 = (target - (a_prev + dt_s / 2 * k1)) / tau + k3 = (target - (a_prev + dt_s / 2 * k2)) / tau + k4 = (target - (a_prev + dt_s * k3)) / tau + return a_prev + dt_s / 6 * (k1 + 2*k2 + 2*k3 + k4) + + new_a_bag1 = _bag_rk4(a_bag1, target_bag1, p["tau1"]) + new_a_bag2 = _bag_rk4(a_bag2, target_bag2, p["tau2"]) + + # Per-fiber coefficients (shape 3) + acts = jnp.array([new_a_bag1, new_a_bag2, a_chain]) + b0 = jnp.array([p["b0Bag1"], p["b0Bag2"], p["b0Chain"]]) + b_fusi = jnp.array([p["b1Bag1"], p["b2Bag2"], p["b2Chain"]]) + G = jnp.array([p["G1"], p["G2"], p["G2Chain"]]) + b_coef = b0 + b_fusi * acts # (3,) + gf = G * acts # (3,) gamma force + + # --- Intrafusal fiber tension 2nd-order ODE (RK4, Eq. 1-3) --- + K_SR = p["K_SR"]; M = p["M"]; K_PR = p["K_PR"] + L0_SR = p["L0_SR"]; L0_PR = p["L0_PR"]; R = p["R"] + a_exp = p["a"]; C_L = p["C_L"]; C_S = p["C_S"] + + def _tension_accel(T_val, z_val, b_c, gamma_f): + vel_diff = V - z_val / K_SR + # Nonlinear force-velocity (asymmetric lengthening/shortening) + fv = jnp.where( + vel_diff >= 0, + C_L * b_c * jnp.abs(vel_diff)**a_exp * (L - L0_SR - T_val/K_SR - R), + C_S * b_c * jnp.abs(vel_diff)**a_exp * (L - L0_SR - T_val/K_SR - R), + ) + spring = K_PR * (L - L0_SR - T_val/K_SR - L0_PR) + return K_SR / M * (fv + M*A + gamma_f - T_val + spring) + + def _fiber_rk4(T_prev, z_prev, b_c, gamma_f): + k1y = z_prev + k1z = _tension_accel(T_prev, z_prev, b_c, gamma_f) + k2y = z_prev + dt_s/2 * k1z + k2z = _tension_accel(T_prev + dt_s/2 * k1y, k2y, b_c, gamma_f) + k3y = z_prev + dt_s/2 * k2z + k3z = _tension_accel(T_prev + dt_s/2 * k2y, k3y, b_c, gamma_f) + k4y = z_prev + dt_s * k3z + k4z = _tension_accel(T_prev + dt_s * k3y, k4y, b_c, gamma_f) + new_T = T_prev + dt_s/6 * (k1y + 2*k2y + 2*k3y + k4y) + new_z = z_prev + dt_s/6 * (k1z + 2*k2z + 2*k3z + k4z) + return new_T, new_z + + # Unrolled over 3 fibers (compile-time constant — fine for lax.scan) + new_T = jnp.zeros(3) + new_dT = jnp.zeros(3) + for fi in range(3): + t_new, z_new = _fiber_rk4(T[fi], dT[fi], b_coef[fi], gf[fi]) + new_T = new_T.at[fi].set(t_new) + new_dT = new_dT.at[fi].set(z_new) + + # --- Afferent firing rates (Eq. 5-7) --- + threshold = p["LN_SR"] - p["L0_SR"] + + ia_bag1 = p["gBag1"] * jnp.maximum(0.0, new_T[0]/K_SR - threshold) + ia_bag2 = p["gBag2A1"] * jnp.maximum(0.0, new_T[1]/K_SR - threshold) + ia_chain = p["gChainA1"] * jnp.maximum(0.0, new_T[2]/K_SR - threshold) + + # Occlusion (Eq. 6) + B2C = ia_bag2 + ia_chain + Ia = jnp.where(B2C >= ia_bag1, B2C + p["S"]*ia_bag1, ia_bag1 + p["S"]*B2C) + + # Secondary afferent (Eq. 7): Bag2 + Chain only + def _ii_contrib(fi, gain): + sr = new_T[fi+1]/K_SR - threshold + pr = L - new_T[fi+1]/K_SR - L0_SR - p["LN_PR"] + return gain * jnp.maximum( + 0.0, + p["X"] * p["Lsec"] / L0_SR * sr + + (1 - p["X"]) * p["Lsec"] / p["L0_PR"] * pr, + ) + + II = _ii_contrib(0, p["gBag2A2"]) + _ii_contrib(1, p["gChainA2"]) + + new_state = {"a_bag1": new_a_bag1, "a_bag2": new_a_bag2, "T": new_T, "dT": new_dT} + return new_state, (Ia, II) + + +# ============================================================================ +# GOLGI TENDON ORGAN — Lin & Crago (2002) +# ============================================================================ + +def gto_init() -> dict: + """Return zeroed GTO state.""" + return {"prev_firing": jnp.float32(0.0)} + + +def gto_step(state: dict, force_N: float, p: dict): + """ + One step of the GTO model. + + Parameters + ---------- + state : dict ``{prev_firing}`` + force_N : float Muscle force in Newtons (absolute, not normalised). + p : dict ``{G1, G2, filter_alpha}`` + + Returns + ------- + new_state, Ib : Ib afferent firing rate [Hz] + """ + force = jnp.maximum(0.0, force_N) + instant = p["G1"] * jnp.log(force / p["G2"] + 1.0) + alpha = p["filter_alpha"] + Ib = (1.0 - alpha) * state["prev_firing"] + alpha * instant + return {"prev_firing": Ib}, Ib + + +def gto_params_from_dict(gto_params: dict, filter_alpha: float = 0.3) -> dict: + """Build GTO params dict from GolgiTendonOrganModel parameter dict.""" + return { + "G1": float(gto_params["G1"]), + "G2": float(gto_params["G2"]), + "filter_alpha": float(filter_alpha), + } + + +# ============================================================================ +# JOINT DYNAMICS — second-order Euler +# ============================================================================ + +def joint_init(angle_deg: float = 0.0, velocity_deg_s: float = 0.0) -> dict: + """Return initial joint state.""" + return { + "angle_rad": jnp.float32(np.radians(angle_deg)), + "velocity_rad_s": jnp.float32(np.radians(velocity_deg_s)), + } + + +def joint_step(state: dict, torque_Nm: float, dt_s: float, p: dict): + """ + One step of second-order joint dynamics: I·α = τ - B·ω - K·θ. + + Parameters + ---------- + state : dict ``{angle_rad, velocity_rad_s}`` + torque_Nm : float Applied muscle torque [N·m]. + dt_s : float Timestep [s]. + p : dict ``{inertia, damping, stiffness}`` + + Returns + ------- + new_state, angle_deg + """ + spring = -p["stiffness"] * state["angle_rad"] + damping = -p["damping"] * state["velocity_rad_s"] + accel = (torque_Nm + spring + damping) / p["inertia"] + new_vel = state["velocity_rad_s"] + accel * dt_s + new_angle = state["angle_rad"] + new_vel * dt_s + return {"angle_rad": new_angle, "velocity_rad_s": new_vel}, jnp.degrees(new_angle) + + +# ============================================================================ +# HILL MUSCLE MODEL — Fuglevand et al. (1993) + Hill mechanics +# ============================================================================ + +def hill_init_params(hillD: dict, Ntype1: int, Ntype2: int, dt_ms: float) -> dict: + """ + Pre-compute per-MU IIR constants and all static Hill model parameters. + + This function runs **once** at simulation setup and may call scipy + (via ``ForceSatParams``). The returned dict never changes during simulation + and is safe to use inside ``lax.scan`` as a static argument. + + Parameters + ---------- + hillD : dict + As returned by ``HillModel.create_default_muscle_parameters()``. + Ntype1, Ntype2 : int + Number of Type I and Type II motor units. + dt_ms : float + Timestep in milliseconds. + + Returns + ------- + dict + All parameters as JAX float32 arrays/scalars. + """ + from myogen.simulator.jaxley.muscle import ForceSatParams + + N = Ntype1 + Ntype2 + fs = ForceSatParams(hillD, Ntype1, Ntype2) # handles scipy internally + + T_ms = fs.T # contraction times (ms) — same units as dt_ms + c = fs.c + tet_f = fs.tetF + P_amp = fs.P # peak force amplitudes + + # Normalised twitch amplitudes (Fuglevand Eq.) + fP = hillD["fP"] + raw = np.array([fP * P_amp[i] / tet_f[i] for i in range(N)]) + twiAmp = raw / np.sum([fP * P_amp[i] for i in range(N)]) + + # IIR coefficients for motor unit force filter + # f[t+1] = A*f[t] - B*f_prev + C*spike (dt and T both in ms) + A_iir = 2.0 * np.exp(-dt_ms / T_ms) + B_iir = np.exp(-2.0 * dt_ms / T_ms) + C_iir = (dt_ms / T_ms) * np.exp(1.0 - dt_ms / T_ms) + + return { + # Motor unit dynamics + "A_iir": jnp.array(A_iir, dtype=jnp.float32), + "B_iir": jnp.array(B_iir, dtype=jnp.float32), + "C_iir": jnp.array(C_iir, dtype=jnp.float32), + "c": jnp.array(c, dtype=jnp.float32), + "twiAmp": jnp.array(twiAmp, dtype=jnp.float32), + "Ntype1": int(Ntype1), + "N": int(N), + # Global geometry + "alfa0": float(hillD["alfa0"]), + "F0": float(hillD["F0"]), + "L0_m": float(hillD["L0"]), + "m": float(hillD["m"]), + "Kpe": float(hillD["Kpe"]), + "b_damp": float(hillD["b"]), + "Em_0": float(hillD["Em_0"]), + # Tendon + "LT_0": float(hillD["LT_0"]), + "Kse": float(hillD["Kse"]), + "cT": float(hillD["cT"]), + "LT_r": float(hillD["LT_r"]), + # Force-length (Type I) + "b1": float(hillD["b1"]), "o1": float(hillD["o1"]), "r1": float(hillD["r1"]), + # Force-length (Type II) + "b2": float(hillD["b2"]), "o2": float(hillD["o2"]), "r2": float(hillD["r2"]), + # Force-velocity (Type I) + "Vmax1": float(hillD["Vmax1"]), + "av01": float(hillD["av01"]), "av11": float(hillD["av11"]), + "av21": float(hillD["av21"]), "bv1": float(hillD["bv1"]), + "cv01": float(hillD["cv01"]), "cv11": float(hillD["cv11"]), + # Force-velocity (Type II) + "Vmax2": float(hillD["Vmax2"]), + "av02": float(hillD["av02"]), "av12": float(hillD["av12"]), + "av22": float(hillD["av22"]), "bv2": float(hillD["bv2"]), + "cv02": float(hillD["cv02"]), "cv12": float(hillD["cv12"]), + # MTU polynomial coefficients (degree-4, 5 terms each) + "Ak": jnp.array(hillD["Ak"], dtype=jnp.float32), + "Bk": jnp.array(hillD["Bk"], dtype=jnp.float32), + "dt_ms": float(dt_ms), + } + + +def differentiable_twitch_params(RP, Tl, RT, fP, Ntype1, Ntype2, dt_ms, tetF): + """Differentiable recompute of the Fuglevand twitch/IIR parameters (M4). + + ``hill_init_params`` derives ``twiAmp`` and the motor-unit IIR coefficients + from the recruitment constants (``RP``, ``Tl``, ``RT``, ``fP``) via + ``ForceSatParams`` — which uses scipy (``newton``/``lfilter``) and Python loops, + breaking the gradient. This function reproduces the **closed-form** path in pure + JAX so gradients flow ``{RP, Tl, RT, fP} → twiAmp, A_iir, B_iir, C_iir``. + + The tetanic-saturation term ``tetF`` (the scipy-Newton output) is passed in as a + **frozen constant** — gradients w.r.t. the saturation-shape constants are the + low-value tail and are intentionally not propagated (see docs/M1... / plan M4). + With ``tetF`` frozen, ``twiAmp`` is still differentiable w.r.t. ``fP`` and the + peak amplitudes ``P`` (hence ``RP``), and the IIR coefficients are fully + differentiable w.r.t. ``Tl``/``RP``/``RT`` through the contraction times ``T``. + + Parameters + ---------- + RP, Tl, RT, fP : float or 0-d array + Fuglevand recruitment-range, longest twitch time, twitch-time range, and + peak-force scale. Pass as JAX values to differentiate w.r.t. them. + Ntype1, Ntype2 : int + Motor-unit counts (static). + dt_ms : float + Timestep. + tetF : array, shape (N,) + Frozen tetanic-force normalisers from a one-time ``ForceSatParams`` call. + + Returns + ------- + dict ``{"twiAmp", "A_iir", "B_iir", "C_iir", "T"}`` — JAX arrays, differentiable. + """ + N = Ntype1 + Ntype2 + idx = jnp.arange(1, N + 1, dtype=jnp.float32) + + # Peak twitch amplitudes: P[i] = exp(b*i), b = log(RP)/N (fPeakAmp) + b = jnp.log(RP) / N + P = jnp.exp(b * idx) + + # Contraction times: T[i] = Tl * (1/P[i])**(1/c), c = log(RP)/log(RT) (fTwitchTime, durType==1) + c_time = jnp.log(RP) / jnp.log(RT) + T = Tl * (1.0 / P) ** (1.0 / c_time) + + # Normalised twitch amplitudes (tetF frozen) + raw = fP * P / tetF + twiAmp = raw / jnp.sum(fP * P) + + A_iir = 2.0 * jnp.exp(-dt_ms / T) + B_iir = jnp.exp(-2.0 * dt_ms / T) + C_iir = (dt_ms / T) * jnp.exp(1.0 - dt_ms / T) + return {"twiAmp": twiAmp, "A_iir": A_iir, "B_iir": B_iir, "C_iir": C_iir, "T": T} + + +def hill_init_state(L0: float, N: int, max_delay_steps: int = 200) -> dict: + """ + Return zeroed Hill model state. + + Parameters + ---------- + L0 : float + Initial normalised muscle length. + N : int + Total number of motor units (Ntype1 + Ntype2). + max_delay_steps : int + Size of axonal delay buffer (steps). 200 steps @ 0.1 ms/step = 20 ms + max delay, which covers all physiological axonal delays. + """ + return { + "L": jnp.float32(L0), + "V": jnp.float32(0.0), + "f": jnp.zeros(N, dtype=jnp.float32), + "f_prev": jnp.zeros(N, dtype=jnp.float32), + "spike_buffer": jnp.zeros((N, max_delay_steps), dtype=jnp.float32), + } + + +def hill_step(state: dict, mn_spikes, delay_steps, angle_rad: float, p: dict): + """ + One step of the Hill muscle model. + + Parameters + ---------- + state : dict + ``{L, V, f, f_prev, spike_buffer}`` + mn_spikes : array-like, shape (N,), bool/float + Which motor units have an MN spike THIS step (before delay). + delay_steps : array-like, shape (N,), int + Per-MU axonal delay in timesteps (≥ 1). + angle_rad : float + Current joint angle in radians. + p : dict + Pre-computed params from ``hill_init_params()``. + + Returns + ------- + new_state : dict + (L, V, force_norm, torque_norm) : all normalised to F0 / L0 + """ + L, V = state["L"], state["V"] + f = state["f"] + f_prev = state["f_prev"] + buf = state["spike_buffer"] + N = p["N"] + Ntype1 = p["Ntype1"] + + # 1. Read spikes that have arrived (front of delay buffer) + current_spikes = buf[:, 0] + + # 2. Shift buffer left, clear last slot, write new spikes at delay positions + # delay D means write to slot D-1: after D rolls it reaches slot 0. + mn_sp_f = jnp.array(mn_spikes, dtype=jnp.float32) + shifted = jnp.roll(buf, -1, axis=1).at[:, -1].set(0.0) + new_buf = shifted.at[jnp.arange(N), delay_steps - 1].add(mn_sp_f) + + # 3. Motor unit force IIR filter + f_new = p["A_iir"] * f - p["B_iir"] * f_prev + p["C_iir"] * current_spikes + + # 4. Force saturation sigmoid + def _sig(c_val, x): + ec = jnp.exp(-c_val * x) + return (1.0 - ec) / (1.0 + ec) + + fsat = p["twiAmp"] * _sig(p["c"], f_new) + F1 = jnp.sum(fsat[:Ntype1]) + F2 = jnp.sum(fsat[Ntype1:]) + + # 5. Hill force functions + def _fL(LM, b, o, r): + return jnp.exp(-(jnp.abs((LM**b - 1.0) / o))**r) + + def _fV(LM, vel, bv, av0, av1, av2, cv0, cv1, Vmax): + conc = (bv - vel * (av0 + av1*LM + av2*LM**2)) / (bv + vel) + ecc = (Vmax - vel) / (Vmax + vel * (cv0 + cv1*LM)) + return jnp.where(vel > 0.0, conc, ecc) + + def _fCE(LM, vel): + fce1 = _fL(LM, p["b1"], p["o1"], p["r1"]) * _fV( + LM, vel, p["bv1"], p["av01"], p["av11"], p["av21"], p["cv01"], p["cv11"], p["Vmax1"]) + fce2 = _fL(LM, p["b2"], p["o2"], p["r2"]) * _fV( + LM, vel, p["bv2"], p["av02"], p["av12"], p["av22"], p["cv02"], p["cv12"], p["Vmax2"]) + return F1 * fce1 + F2 * fce2 + + def _fPE(LM, vel): + return jnp.exp(p["Kpe"] * (LM - 1.0) / p["Em_0"]) / jnp.exp(p["Kpe"]) + p["b_damp"] * vel + + def _penn(LM): + return jnp.arcsin(jnp.sin(p["alfa0"]) / LM) + + def _MTU(angle): + return jnp.dot(p["Ak"], angle ** jnp.arange(5, dtype=jnp.float32)) + + def _moment_arm(angle): + return jnp.dot(p["Bk"], angle ** jnp.arange(5, dtype=jnp.float32)) + + def _LT(angle, LM): + term2 = LM * p["L0_m"] * jnp.cos(_penn(LM)) + return (_MTU(angle) - term2) / p["LT_0"] + + def _fSE(angle, LM): + LT = _LT(angle, LM) + return p["Kse"] * p["cT"] * jnp.log(jnp.exp((LT - p["LT_r"]) / p["cT"]) + 1.0) + + def _dVdt(LM, vel, angle): + ratio = p["F0"] / p["m"] + tendon = _fSE(angle, LM) * jnp.cos(_penn(LM)) + active = (_fCE(LM, vel) + _fPE(LM, vel)) * jnp.cos(_penn(LM))**2 + return ratio * (tendon - active) + + # 6. RK4 for L and V + dt_s = p["dt_ms"] * 1e-3 + k1L = V; k1V = _dVdt(L, V, angle_rad) + k2L = V + dt_s/2*k1V; k2V = _dVdt(L + dt_s/2*k1L, k2L, angle_rad) + k3L = V + dt_s/2*k2V; k3V = _dVdt(L + dt_s/2*k2L, k3L, angle_rad) + k4L = V + dt_s *k3V; k4V = _dVdt(L + dt_s *k3L, k4L, angle_rad) + + new_L = L + dt_s/6 * (k1L + 2*k2L + 2*k3L + k4L) + new_V = V + dt_s/6 * (k1V + 2*k2V + 2*k3V + k4V) + + # 7. Force and torque (normalised) + force = _fSE(angle_rad, new_L) + torque = force * _moment_arm(angle_rad) + + # 8. Per-fiber-type summed activations (for plotting) + type1_act = F1 + type2_act = F2 + + new_state = { + "L": new_L, + "V": new_V, + "f": f_new, + "f_prev": f, + "spike_buffer": new_buf, + } + return new_state, (new_L, new_V, force, torque, type1_act, type2_act) + + +# ============================================================================ +# POISSON SPIKE GENERATOR (descending drive / DD cells) +# ============================================================================ + +def poisson_init(N_cells: int, N_batch: int, seed: int = 42, key=None) -> dict: + """ + Initialise Poisson spike generator state. + + Mirrors ``_PoissonProcessGenerator__Jaxley``: + ``yi += rate * dt_ms * 1e-3``; spike when ``yi >= thres``; + ``thres ~ Gamma(N_batch, 1) / N_batch`` resampled after each spike. + + Parameters + ---------- + N_cells : int Number of independent generators. + N_batch : int Poisson batch size (higher → more regular; 1 = pure Poisson). + seed : int Fallback seed used only when ``key`` is None. + key : jax PRNGKey, optional + Explicit PRNG key. When provided it takes precedence over ``seed`` so the + stochastic stream is controlled from the top-level ``run_jax`` call and is + reproducible across JIT and runs. + + Returns + ------- + dict ``{yi: (N,), thres: (N,), keys: (N, 2)}`` + """ + key = jax.random.PRNGKey(seed) if key is None else key + keys = jax.random.split(key, N_cells) + thres = ( + jax.vmap(lambda k: jax.random.gamma(k, jnp.float32(N_batch)))(keys) + / jnp.float32(N_batch) + ) + return { + "yi": jnp.zeros(N_cells, dtype=jnp.float32), + "thres": thres.astype(jnp.float32), + "keys": keys, + } + + +def poisson_step(state: dict, rate_hz, dt_ms: float, N_batch: int, mode: str = "hard") -> tuple: + """ + One step of the Poisson spike generator. + + Parameters + ---------- + state : dict ``{yi, thres, keys}`` + rate_hz : float or (N,) Instantaneous rate(s) [Hz]. + dt_ms : float Timestep [ms]. + N_batch : int Same as in ``poisson_init``. + mode : {"hard", "pathwise", "rate"} + ``"hard"`` — original discrete hazard crossing (no gradient w.r.t. rate). + ``"pathwise"`` — reparameterised: the sampled threshold is frozen + (``stop_gradient``) and the crossing uses a straight-through surrogate, so + the forward spike train is identical to ``"hard"`` yet gradients flow + ``rate → yi → spike``. Preserves stochastic timing. + ``"rate"`` — emit the expected spikes-per-step ``rate*dt`` (continuous), + giving an exact, low-variance gradient w.r.t. rate. Forward output is + continuous rather than 0/1. + + Returns + ------- + new_state, spikes + ``spikes`` is bool in ``"hard"`` mode and float32 otherwise. + """ + return _hazard_step(state, rate_hz, dt_ms, N_batch, mode) + + +def _hazard_step(state, rate_hz, dt_ms, gamma_conc, mode): + """Shared hazard-accumulation update for the Poisson/Gamma generators. + + ``gamma_conc`` is the Gamma concentration used to resample the threshold + (``N_batch`` for Poisson, ``shape`` for Gamma) — both use ``Gamma(conc)/conc``. + """ + yi, thres, keys = state["yi"], state["thres"], state["keys"] + split_keys = jax.vmap(jax.random.split)(keys) # (N, 2, 2) + new_keys = split_keys[:, 0, :] + sub_keys = split_keys[:, 1, :] + + inc = jnp.asarray(rate_hz, dtype=jnp.float32) * jnp.float32(dt_ms * 1e-3) + + if mode == "rate": + # Expected spikes per step; state advances but never emits discretely. + yi_new = yi + inc + spikes = jnp.broadcast_to(inc, jnp.shape(yi_new)) + return {"yi": yi_new, "thres": thres, "keys": new_keys}, spikes + + yi_new = yi + inc + spikes_hard = yi_new >= thres # bool, drives state reset + new_thres = ( + jax.vmap(lambda k: jax.random.gamma(k, jnp.float32(gamma_conc)))(sub_keys) + / jnp.float32(gamma_conc) + ).astype(jnp.float32) + # Reset accumulator / resample threshold on a spike (forward-identical to hard). + yi_reset = jnp.where(spikes_hard, jnp.zeros_like(yi_new), yi_new) + thres_out = jnp.where(spikes_hard, new_thres, thres) + + if mode == "hard": + spikes = spikes_hard + elif mode == "pathwise": + # Freeze the resampled threshold so the crossing margin's only parameter + # dependence is through yi (hence through rate); surrogate grad on it. + margin = yi_new - jax.lax.stop_gradient(thres) + spikes = _straight_through_step(margin) + else: + raise ValueError(f"unknown generator mode {mode!r}") + + return {"yi": yi_reset, "thres": thres_out, "keys": new_keys}, spikes + + +# ============================================================================ +# GAMMA SPIKE GENERATOR (Ia / II / Ib afferent cells) +# ============================================================================ + +def gamma_init(N_cells: int, shape: float, seed: int = 42, key=None) -> dict: + """ + Initialise Gamma-ISI spike generator state. + + Mirrors ``_GammaProcessGenerator__Cython``: hazard-accumulation threshold + crossing. Identical structure to ``poisson_init`` but threshold drawn from + Gamma(shape, 1) / shape rather than Gamma(1, 1). + + Parameters + ---------- + N_cells : int Number of independent generators. + shape : float Gamma shape parameter (k). 1 = Poisson, >1 = more regular. + seed : int Fallback seed used only when ``key`` is None. + key : jax PRNGKey, optional + Explicit PRNG key; takes precedence over ``seed`` (see ``poisson_init``). + + Returns + ------- + dict ``{yi: (N,), thres: (N,), keys: (N, 2)}`` + """ + key = jax.random.PRNGKey(seed) if key is None else key + keys = jax.random.split(key, N_cells) + thres = ( + jax.vmap(lambda k: jax.random.gamma(k, jnp.float32(shape)))(keys) + / jnp.float32(shape) + ) + return { + "yi": jnp.zeros(N_cells, dtype=jnp.float32), + "thres": thres.astype(jnp.float32), + "keys": keys, + } + + +def gamma_step(state: dict, rate_hz, dt_ms: float, shape: float, mode: str = "hard") -> tuple: + """ + One step of the Gamma-ISI spike generator (hazard accumulation). + + Mirrors ``_GammaProcessGenerator__Cython.compute()``: + + :: + + yi += rate * dt_ms * 1e-3 + if yi >= thres: + spike; yi = 0; thres ~ Gamma(shape, 1) / shape + + Parameters + ---------- + state : dict ``{yi, thres, keys}`` + rate_hz : float or (N,) Rate(s) [Hz]. 0 → silent. + dt_ms : float Timestep [ms]. + shape : float Gamma shape parameter (same as in ``gamma_init``). + mode : {"hard", "pathwise", "rate"} See :func:`poisson_step`. + + Returns + ------- + new_state, spikes + ``spikes`` is bool in ``"hard"`` mode and float32 otherwise. + """ + return _hazard_step(state, rate_hz, dt_ms, shape, mode) + + +# ============================================================================ +# CONNECTIVITY UTILITIES +# ============================================================================ + +def make_connectivity_matrix( + forward_map: dict, N_pre: int, N_post: int +) -> jnp.ndarray: + """ + Convert a sparse ``{pre_idx: [post_idx, ...]}`` dict to a dense + ``(N_pre, N_post)`` JAX float32 matrix. + + Usage inside scan:: + + g_post += jnp.dot(spikes.astype(jnp.float32), mat) * weight + """ + mat = np.zeros((N_pre, N_post), dtype=np.float32) + for pre_idx, post_list in forward_map.items(): + for post_idx in post_list: + mat[pre_idx, post_idx] = 1.0 + return jnp.array(mat) + + +# ============================================================================ +# COMBINED PHYSIOLOGY STEP +# ============================================================================ + +def update_physiology( + phys_carry: dict, + mn_spikes, + tap_dL_coeff_t: float, + tap_dV_coeff_t: float, + gDyn_t: float, + gStat_t: float, + delay_steps, + hill_p: dict, + spindle_p: dict, + gto_p: dict, + joint_p: dict, + dt_ms: float, + dt_s: float, +) -> tuple: + """ + One step of the full physiology pipeline: + Hill muscle → tap perturbation (spindle input only) → + spindle → GTO → joint dynamics. + + Parameters + ---------- + phys_carry : dict ``{hill, spindle, gto, joint}`` + mn_spikes : (N_mn,) bool/float MN spikes this step. + tap_dL_coeff_t, tap_dV_coeff_t : float Per-step tap coefficients. + gDyn_t, gStat_t : float Gamma fusimotor drives [Hz]. + delay_steps : (N_mn,) int32 Per-MU axonal delays [timesteps]. + hill_p, spindle_p, gto_p, joint_p : dict Pre-computed parameter dicts. + dt_ms, dt_s : float Timestep. + + Returns + ------- + new_phys_carry, (new_L, new_V, force_norm, torque_norm, type1_act, type2_act, Ia, II, Ib, angle_deg) + """ + angle_rad = phys_carry["joint"]["angle_rad"] + prev_V_hill = phys_carry["hill"]["V"] + + hill_state, (new_L, new_V, force_norm, torque_norm, type1_act, type2_act) = hill_step( + phys_carry["hill"], mn_spikes, delay_steps, angle_rad, hill_p + ) + + L_sp = new_L * (jnp.float32(1.0) + jnp.float32(tap_dL_coeff_t)) + V_sp = new_V + new_L * jnp.float32(tap_dV_coeff_t) + A_sp = (new_V - prev_V_hill) / jnp.float32(dt_s) + + spindle_state, (Ia, II) = spindle_step( + phys_carry["spindle"], L_sp, V_sp, A_sp, + gDyn_t, gStat_t, dt_s, spindle_p + ) + + force_N = jnp.float32(hill_p["F0"]) * force_norm + gto_state, Ib = gto_step(phys_carry["gto"], force_N, gto_p) + + torque_Nm = jnp.float32(hill_p["F0"]) * torque_norm + joint_state, angle_deg = joint_step(phys_carry["joint"], torque_Nm, dt_s, joint_p) + + new_phys = { + "hill": hill_state, + "spindle": spindle_state, + "gto": gto_state, + "joint": joint_state, + } + return new_phys, (new_L, new_V, force_norm, torque_norm, type1_act, type2_act, Ia, II, Ib, angle_deg) + + +# ============================================================================ +# lax.scan STEP FUNCTION FACTORY +# ============================================================================ + +def make_scan_step( + jaxley_step_fn, + jaxley_params, + external_inds, + rec_inds, + n_gii: int, + n_gib: int, + n_mn: int, + ia_rts, + ii_rts, + ib_rts, + ia_shape: float, + ii_shape: float, + ib_shape: float, + dd_N_batch: int, + dd_to_mn_mat, + ia_to_mn_mat, + ii_to_gii_mat, + ib_to_gib_mat, + ia_delay_steps_arr, + ii_delay_steps_arr, + ib_delay_steps_arr, + delay_steps, + hill_p: dict, + spindle_p: dict, + gto_p: dict, + joint_p: dict, + base_dd_weight: float, + base_ia_weight: float, + in_weight: float, + e_exc: float, + v_rest: float, + mn_current_scale, + tau_syn_decay: float, + dt_ms: float, + dt_s: float, + e_exc_mn: float | None = None, + mn_spike_threshold_mV: float = 0.0, + spike_mode: str = "hard", +): + """ + Factory returning a ``lax.scan``-compatible step function. + + All static parameters are closed over; the returned ``scan_step`` + has the fixed signature ``(new_carry, output_t) = scan_step(carry, input_t)``. + + Carry structure + --------------- + :: + + { + "neural": Jaxley states, + "phys": {"hill", "spindle", "gto", "joint"}, + "g_dd": (n_mn,), "g_ia": (n_mn,), + "g_ii": (n_gii,), "g_ib": (n_gib,), + "prev_v": (n_gii+n_gib+n_mn,), + "dd_st": poisson_init output, + "ia_st / ii_st / ib_st": gamma_init outputs, + "prev_Iay / prev_IIy / prev_Iby": float32 scalars, + "ia_delay_buf": (nIa, max_ia_delay_steps), # per-cell FIFO delay queues + "ii_delay_buf": (nII, max_ii_delay_steps), + "ib_delay_buf": (nIb, max_ib_delay_steps), + } + + Input structure (per step, stacked by lax.scan) + ------------------------------------------------ + :: + + {"DDdrive", "gDyn", "gStat", "tap_dL", "tap_dV"} — all float32 scalars + + Output structure (stacked → shape (T, ...) after scan) + ------------------------------------------------------- + :: + + {"v_mn", "mn_spikes", "gii_spikes", "gib_spikes", + "ia_spikes", "ii_spikes", "ib_spikes", + "L", "force", "torque", "Iay", "IIy", "Iby", "angle_deg"} + + Notes + ----- + Afferent axonal delays are implemented as per-cell FIFO shift-register + queues in the carry. Each afferent cell has its own delay slot so spike + timing faithfully reflects individual conduction velocities. + """ + _rec_inds = jnp.array(rec_inds, dtype=jnp.int32) + _ia_rts = jnp.array(ia_rts, dtype=jnp.float32) + _ii_rts = jnp.array(ii_rts, dtype=jnp.float32) + _ib_rts = jnp.array(ib_rts, dtype=jnp.float32) + _delay_steps = jnp.array(delay_steps, dtype=jnp.int32) + _ia_delay_steps = jnp.array(ia_delay_steps_arr, dtype=jnp.int32) + _ii_delay_steps = jnp.array(ii_delay_steps_arr, dtype=jnp.int32) + _ib_delay_steps = jnp.array(ib_delay_steps_arr, dtype=jnp.int32) + _nIa = len(ia_delay_steps_arr) + _nII = len(ii_delay_steps_arr) + _nIb = len(ib_delay_steps_arr) + _dd_to_mn = jnp.asarray(dd_to_mn_mat, dtype=jnp.float32) + _ia_to_mn = jnp.asarray(ia_to_mn_mat, dtype=jnp.float32) + _ii_to_gii = jnp.asarray(ii_to_gii_mat, dtype=jnp.float32) + _ib_to_gib = jnp.asarray(ib_to_gib_mat, dtype=jnp.float32) + _mn_scl = jnp.asarray(mn_current_scale, dtype=jnp.float32) + _base_dd = jnp.float32(base_dd_weight) + _base_ia = jnp.float32(base_ia_weight) + _in_w = jnp.float32(in_weight) + _e_exc = jnp.float32(e_exc) + _v_rest = jnp.float32(v_rest) + # ``e_exc_mn``: excitatory reversal for the MOTOR-NEURON section of the + # combined network. Optional; defaults to ``e_exc`` (single-frame behavior + # for backward compatibility with older callers). When the MN cells live in + # a different voltage frame from the interneurons (e.g. NERLab MNs at + # V_rest ≈ 0 mV mixed with modern-frame gII/gIb at V_rest ≈ -70 mV), pass + # the NERLab AMPA reversal (~+70 mV) here so the per-step driving force + # ``df_mn = e_exc_mn − V_mn`` evaluates correctly. ``mn_spike_threshold_mV`` + # similarly lets the MN spike detector use a NERLab-appropriate +50 mV + # crossing instead of the 0 mV crossing used by the modern-frame cells. + _e_exc_mn = jnp.float32(e_exc_mn if e_exc_mn is not None else e_exc) + _mn_v_th = jnp.float32(mn_spike_threshold_mV) + _decay = jnp.float32(tau_syn_decay) + _dt_ms = float(dt_ms) + _dt_s = float(dt_s) + _n_gii = n_gii + _n_gib = n_gib + _spike_mode = spike_mode + # Map the loop-level spike mode onto the stochastic-generator mode so the + # descending/afferent inputs are differentiated consistently with the + # neural spike detectors: hard→hard, surrogate→pathwise, rate→rate. + _gen_mode = {"hard": "hard", "surrogate": "pathwise", "rate": "rate"}[spike_mode] + + def scan_step(carry, input_t): + # ------------------------------------------------------------------ # + # 1. Decay conductances, then read arriving delayed contributions. + # Per-cell delay buffers: row i = FIFO for afferent cell i. + # Slot 0 holds contributions ready to apply this step. + # ------------------------------------------------------------------ # + ia_arrivals = jnp.dot(carry["ia_delay_buf"][:, 0], _ia_to_mn) * _base_ia # (n_mn,) + ii_arrivals = jnp.dot(carry["ii_delay_buf"][:, 0], _ii_to_gii) * _in_w # (n_gii,) + ib_arrivals = jnp.dot(carry["ib_delay_buf"][:, 0], _ib_to_gib) * _in_w # (n_gib,) + g_dd = carry["g_dd"] * _decay + g_ia = carry["g_ia"] * _decay + ia_arrivals + g_ii = carry["g_ii"] * _decay + ii_arrivals + g_ib = carry["g_ib"] * _decay + ib_arrivals + + # ------------------------------------------------------------------ # + # 2. DD Poisson step + # ------------------------------------------------------------------ # + dd_st_new, dd_spikes = poisson_step( + carry["dd_st"], input_t["DDdrive"], _dt_ms, dd_N_batch, _gen_mode + ) + g_dd = g_dd + jnp.dot(dd_spikes.astype(jnp.float32), _dd_to_mn) * _base_dd + + # ------------------------------------------------------------------ # + # 3. Afferent Gamma steps (driven by prev-step rates) + # Then enqueue new contributions into delay FIFOs. + # ------------------------------------------------------------------ # + ia_rates = jnp.where(carry["prev_Iay"] >= _ia_rts, carry["prev_Iay"], jnp.float32(0.0)) + ii_rates = jnp.where(carry["prev_IIy"] >= _ii_rts, carry["prev_IIy"], jnp.float32(0.0)) + ib_rates = jnp.where(carry["prev_Iby"] >= _ib_rts, carry["prev_Iby"], jnp.float32(0.0)) + + ia_st_new, ia_spikes = gamma_step(carry["ia_st"], ia_rates, _dt_ms, ia_shape, _gen_mode) + ii_st_new, ii_spikes = gamma_step(carry["ii_st"], ii_rates, _dt_ms, ii_shape, _gen_mode) + ib_st_new, ib_spikes = gamma_step(carry["ib_st"], ib_rates, _dt_ms, ib_shape, _gen_mode) + + # Per-cell delay FIFOs: shift left, clear wrap slot, write new spikes + # at each cell's own delay position (delay d → slot d-1 after shift). + # After d more rolls the spike reaches slot 0 where it is read. + ia_sp_f = ia_spikes.astype(jnp.float32) + ii_sp_f = ii_spikes.astype(jnp.float32) + ib_sp_f = ib_spikes.astype(jnp.float32) + new_ia_delay_buf = ( + jnp.roll(carry["ia_delay_buf"], -1, axis=1).at[:, -1].set(0.0) + .at[jnp.arange(_nIa), _ia_delay_steps - 1].add(ia_sp_f) + ) + new_ii_delay_buf = ( + jnp.roll(carry["ii_delay_buf"], -1, axis=1).at[:, -1].set(0.0) + .at[jnp.arange(_nII), _ii_delay_steps - 1].add(ii_sp_f) + ) + new_ib_delay_buf = ( + jnp.roll(carry["ib_delay_buf"], -1, axis=1).at[:, -1].set(0.0) + .at[jnp.arange(_nIb), _ib_delay_steps - 1].add(ib_sp_f) + ) + + # ------------------------------------------------------------------ # + # 4. Build stimulus currents [nA] = g [µS] × driving_force [mV] + # Use live (prev-step) voltages for driving force. + # ------------------------------------------------------------------ # + df_gii = _e_exc - carry["prev_v"][:_n_gii] + df_gib = _e_exc - carry["prev_v"][_n_gii: _n_gii + _n_gib] + df_mn = _e_exc_mn - carry["prev_v"][_n_gii + _n_gib:] # NERLab-aware + stims = jnp.concatenate([ + g_ii * df_gii, + g_ib * df_gib, + (g_dd + g_ia) * _mn_scl * df_mn, + ]) + + # ------------------------------------------------------------------ # + # 5. Advance Jaxley neural network one step + # ------------------------------------------------------------------ # + new_neural = jaxley_step_fn( + carry["neural"], jaxley_params, {"i": stims}, external_inds, + delta_t=_dt_ms, + ) + + # ------------------------------------------------------------------ # + # 6. Extract voltages (all recordings assumed to be "v") + # ------------------------------------------------------------------ # + v_all = new_neural["v"][_rec_inds] + v_gii = v_all[:_n_gii] + v_gib = v_all[_n_gii: _n_gii + _n_gib] + v_mn = v_all[_n_gii + _n_gib:] + + prev_v_gii = carry["prev_v"][:_n_gii] + prev_v_gib = carry["prev_v"][_n_gii: _n_gii + _n_gib] + prev_v_mn = carry["prev_v"][_n_gii + _n_gib:] + + # ------------------------------------------------------------------ # + # 7. Spike detection (upward crossing; hard / surrogate / rate mode) + # ------------------------------------------------------------------ # + mn_spikes_now = spike_detect(v_mn, prev_v_mn, _mn_v_th, _spike_mode) + gii_spikes_now = spike_detect(v_gii, prev_v_gii, jnp.float32(0.0), _spike_mode) + gib_spikes_now = spike_detect(v_gib, prev_v_gib, jnp.float32(0.0), _spike_mode) + + # ------------------------------------------------------------------ # + # 8. Physiology step + # ------------------------------------------------------------------ # + new_phys, (new_L, new_V, force_norm, torque_norm, type1_act, type2_act, Ia, II, Ib, angle_deg) = ( + update_physiology( + carry["phys"], + mn_spikes_now, + input_t["tap_dL"], + input_t["tap_dV"], + input_t["gDyn"], + input_t["gStat"], + _delay_steps, + hill_p, spindle_p, gto_p, joint_p, + _dt_ms, _dt_s, + ) + ) + + # ------------------------------------------------------------------ # + # 9. New carry + # ------------------------------------------------------------------ # + new_carry = { + "neural": new_neural, + "phys": new_phys, + "g_dd": g_dd, + "g_ia": g_ia, + "g_ii": g_ii, + "g_ib": g_ib, + "prev_v": v_all, + "dd_st": dd_st_new, + "ia_st": ia_st_new, + "ii_st": ii_st_new, + "ib_st": ib_st_new, + "prev_Iay": jnp.float32(Ia), + "prev_IIy": jnp.float32(II), + "prev_Iby": jnp.float32(Ib), + "ia_delay_buf": new_ia_delay_buf, + "ii_delay_buf": new_ii_delay_buf, + "ib_delay_buf": new_ib_delay_buf, + } + + # ------------------------------------------------------------------ # + # 10. Per-step outputs (lax.scan stacks these into (T, ...) arrays) + # ------------------------------------------------------------------ # + output_t = { + "v_mn": v_mn, + "mn_spikes": mn_spikes_now, + "gii_spikes": gii_spikes_now, + "gib_spikes": gib_spikes_now, + "dd_spikes": dd_spikes, + "ia_spikes": ia_spikes, + "ii_spikes": ii_spikes, + "ib_spikes": ib_spikes, + "L": new_L, + "force": force_norm, + "torque": torque_norm, + "type1_act": type1_act, + "type2_act": type2_act, + "Iay": Ia, + "IIy": II, + "Iby": Ib, + "angle_deg": angle_deg, + # Spindle internal states (for plotting bag/chain activation and tensions) + "bag1_act": new_phys["spindle"]["a_bag1"], + "bag2_act": new_phys["spindle"]["a_bag2"], + "spin_T": new_phys["spindle"]["T"], # (3,) [Bag1, Bag2, Chain] + } + + return new_carry, output_t + + return scan_step diff --git a/myogen/simulator/jaxley/jit_simulation.py b/myogen/simulator/jaxley/jit_simulation.py new file mode 100644 index 00000000..5f69ea20 --- /dev/null +++ b/myogen/simulator/jaxley/jit_simulation.py @@ -0,0 +1,697 @@ +""" +JIT-compiled and Batched Simulation Utilities for Jaxley. + +This module provides JAX-accelerated simulation functions using: +- jax.jit for compilation and fast repeated simulations +- jax.vmap for batched/vectorized population simulations +- Proper synapse state updates on presynaptic spikes + +References: + Jaxley Tutorial 04: JIT and vmap + https://jaxley.readthedocs.io/en/latest/tutorials/04_jit_and_vmap.html +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Optional, Tuple, Union + +import jax +import jax.numpy as jnp +import numpy as np +import jaxley as jx +from jax import jit, vmap + + +# ============================================================================= +# JIT-COMPILED SIMULATION FUNCTIONS +# ============================================================================= + +def create_jitted_simulator( + cell: jx.Cell, + dt: float = 0.025, + t_max: float = 100.0, +) -> Callable: + """ + Create a JIT-compiled simulation function for a cell. + + The returned function can be called repeatedly with different parameters + for fast parameter sweeps. + + Parameters + ---------- + cell : jx.Cell + The Jaxley cell to simulate. + dt : float + Time step in ms. + t_max : float + Total simulation time in ms. + + Returns + ------- + Callable + JIT-compiled simulation function that takes param_state and returns voltages. + + Example + ------- + >>> cell = create_motor_neuron() + >>> simulate = create_jitted_simulator(cell, dt=0.025, t_max=100.0) + >>> # First call compiles (slow) + >>> v1 = simulate(None) + >>> # Subsequent calls are fast + >>> v2 = simulate(None) + """ + @jit + def simulate(param_state=None): + return jx.integrate(cell, param_state=param_state, delta_t=dt, t_max=t_max) + + return simulate + + +def create_parameter_sweep_simulator( + cell: jx.Cell, + param_names: List[str], + dt: float = 0.025, + t_max: float = 100.0, +) -> Callable: + """ + Create a JIT-compiled simulator for parameter sweeps. + + Parameters + ---------- + cell : jx.Cell + The Jaxley cell to simulate. + param_names : List[str] + List of parameter names to sweep (e.g., ["Na3rp_gbar", "KdrRL_gbar"]). + dt : float + Time step in ms. + t_max : float + Total simulation time in ms. + + Returns + ------- + Callable + Function that takes parameter values array and returns voltages. + + Example + ------- + >>> cell = create_motor_neuron() + >>> simulate = create_parameter_sweep_simulator( + ... cell, ["Na3rp_gbar", "KdrRL_gbar"] + ... ) + >>> params = jnp.array([0.05, 0.3]) # gNa, gK values + >>> voltages = simulate(params) + """ + @jit + def simulate(param_values): + param_state = None + for i, name in enumerate(param_names): + param_state = cell.data_set(name, param_values[i], param_state) + return jx.integrate(cell, param_state=param_state, delta_t=dt, t_max=t_max) + + return simulate + + +# ============================================================================= +# BATCHED POPULATION SIMULATION WITH VMAP +# ============================================================================= + +def create_batched_simulator( + cell: jx.Cell, + param_names: List[str], + dt: float = 0.025, + t_max: float = 100.0, +) -> Callable: + """ + Create a batched simulator using jax.vmap for parallel execution. + + This function enables GPU-accelerated parallel simulation of multiple + parameter sets simultaneously. + + Parameters + ---------- + cell : jx.Cell + The Jaxley cell to simulate. + param_names : List[str] + List of parameter names to sweep. + dt : float + Time step in ms. + t_max : float + Total simulation time in ms. + + Returns + ------- + Callable + Function that takes (n_simulations, n_params) array and returns + (n_simulations, n_timesteps, n_compartments) voltages. + + Example + ------- + >>> cell = create_motor_neuron() + >>> batch_simulate = create_batched_simulator( + ... cell, ["Na3rp_gbar", "KdrRL_gbar"] + ... ) + >>> # 100 different parameter combinations + >>> all_params = jnp.array(np.random.rand(100, 2)) + >>> # Run all 100 simulations in parallel + >>> all_voltages = batch_simulate(all_params) + >>> print(all_voltages.shape) # (100, n_timesteps, n_compartments) + """ + # Create single simulation function + @jit + def single_simulate(param_values): + param_state = None + for i, name in enumerate(param_names): + param_state = cell.data_set(name, param_values[i], param_state) + return jx.integrate(cell, param_state=param_state, delta_t=dt, t_max=t_max) + + # Vectorize over first axis (batch dimension) + vmapped_simulate = vmap(single_simulate, in_axes=(0,)) + + # Combine jit and vmap for maximum efficiency + return jit(vmapped_simulate) + + +def simulate_population_batched( + cells: List[jx.Cell], + currents: jnp.ndarray, + dt: float = 0.025, + t_max: float = 100.0, +) -> jnp.ndarray: + """ + Simulate a population of cells with different currents using batching. + + Parameters + ---------- + cells : List[jx.Cell] + List of Jaxley cells (should be identical structure for vmap). + currents : jnp.ndarray + Current waveforms, shape (n_cells, n_timesteps). + dt : float + Time step in ms. + t_max : float + Total simulation time in ms. + + Returns + ------- + jnp.ndarray + Voltages, shape (n_cells, n_timesteps, n_compartments). + """ + # For truly batched simulation, all cells should be identical + # Use the first cell as template + template_cell = cells[0] + + n_cells = len(cells) + n_steps = int(t_max / dt) + + @jit + def simulate_single(current): + template_cell.delete_stimuli() + template_cell.branch(0).loc(0.5).stimulate(current) + return jx.integrate(template_cell, delta_t=dt, t_max=t_max) + + # Vectorize over currents + batch_simulate = vmap(simulate_single, in_axes=(0,)) + + return batch_simulate(currents) + + +@dataclass +class BatchedPopulationResult: + """Result container for batched population simulations.""" + + voltages: jnp.ndarray # Shape: (n_cells, n_timesteps, n_compartments) + spike_times: List[List[float]] = field(default_factory=list) + firing_rates: List[float] = field(default_factory=list) + time_vector: Optional[jnp.ndarray] = None + + @property + def n_cells(self) -> int: + return self.voltages.shape[0] + + @property + def n_timesteps(self) -> int: + return self.voltages.shape[1] + + def detect_all_spikes(self, threshold: float = 0.0, dt: float = 0.025): + """Detect spikes for all cells in the population.""" + self.spike_times = [] + self.firing_rates = [] + + for i in range(self.n_cells): + # Extract soma voltage (first compartment or specified) + v = np.array(self.voltages[i, :, 0]) + + # Threshold crossing detection + crossings = np.where((v[:-1] < threshold) & (v[1:] >= threshold))[0] + times = crossings * dt + + self.spike_times.append(times.tolist()) + + # Firing rate + duration_s = self.n_timesteps * dt / 1000.0 + fr = len(times) / duration_s if duration_s > 0 else 0.0 + self.firing_rates.append(fr) + + +# ============================================================================= +# SYNAPTIC STATE UPDATE ON SPIKE +# ============================================================================= + +@dataclass +class SynapticEvent: + """Represents a synaptic event (presynaptic spike).""" + pre_cell_idx: int + post_cell_idx: int + spike_time: float # ms + weight: float = 1.0 + delay: float = 0.0 # ms + + +class SynapseStateManager: + """ + Manages synapse state updates triggered by presynaptic spikes. + + This class handles the activation of synapses when presynaptic neurons + fire, implementing proper spike-triggered synaptic conductance changes. + """ + + def __init__( + self, + network: jx.Network, + spike_threshold: float = 0.0, + dt: float = 0.025, + ): + """ + Initialize the synapse state manager. + + Parameters + ---------- + network : jx.Network + The Jaxley network with synapses. + spike_threshold : float + Voltage threshold for spike detection (mV). + dt : float + Simulation time step (ms). + """ + self.network = network + self.spike_threshold = spike_threshold + self.dt = dt + + # Track previous voltages for spike detection + self.prev_voltages: Dict[int, float] = {} + + # Pending synaptic events (for delayed transmission) + self.pending_events: List[SynapticEvent] = [] + + # Get synapse info from network edges + self._parse_network_synapses() + + def _parse_network_synapses(self): + """Extract synapse connectivity from network edges.""" + self.synapses = [] + + if hasattr(self.network, 'edges') and len(self.network.edges) > 0: + edges = self.network.edges + for idx, row in edges.iterrows(): + synapse_info = { + 'edge_idx': idx, + 'pre_cell': row.get('pre_global_cell_index', 0), + 'post_cell': row.get('post_global_cell_index', 0), + 'pre_comp': row.get('pre_global_comp_index', 0), + 'post_comp': row.get('post_global_comp_index', 0), + } + self.synapses.append(synapse_info) + + def detect_spike(self, cell_idx: int, current_voltage: float) -> bool: + """ + Detect if a cell has just spiked (threshold crossing). + + Parameters + ---------- + cell_idx : int + Index of the cell. + current_voltage : float + Current membrane voltage. + + Returns + ------- + bool + True if spike detected. + """ + prev_v = self.prev_voltages.get(cell_idx, -70.0) + self.prev_voltages[cell_idx] = current_voltage + + return prev_v < self.spike_threshold <= current_voltage + + def process_spike( + self, + pre_cell_idx: int, + spike_time: float, + axon_delays: Optional[Dict[int, float]] = None, + ): + """ + Process a presynaptic spike and schedule synaptic events. + + Parameters + ---------- + pre_cell_idx : int + Index of the presynaptic cell that spiked. + spike_time : float + Time of the spike (ms). + axon_delays : dict, optional + Dictionary mapping cell indices to axon delays (ms). + """ + for syn in self.synapses: + if syn['pre_cell'] == pre_cell_idx: + delay = 0.0 + if axon_delays is not None: + delay = axon_delays.get(pre_cell_idx, 0.0) + + event = SynapticEvent( + pre_cell_idx=pre_cell_idx, + post_cell_idx=syn['post_cell'], + spike_time=spike_time, + delay=delay, + ) + self.pending_events.append(event) + + def get_active_events(self, current_time: float) -> List[SynapticEvent]: + """ + Get synaptic events that should be activated at current time. + + Parameters + ---------- + current_time : float + Current simulation time (ms). + + Returns + ------- + List[SynapticEvent] + Events to activate now. + """ + active = [] + remaining = [] + + for event in self.pending_events: + activation_time = event.spike_time + event.delay + if current_time >= activation_time: + active.append(event) + else: + remaining.append(event) + + self.pending_events = remaining + return active + + def update_synapse_states( + self, + current_time: float, + voltages: jnp.ndarray, + ) -> None: + """ + Update synapse states based on presynaptic spikes. + + This method: + 1. Detects spikes in all presynaptic cells + 2. Schedules synaptic events with appropriate delays + 3. Activates synapses whose events are due + + Parameters + ---------- + current_time : float + Current simulation time (ms). + voltages : jnp.ndarray + Current voltages of all cells. + """ + n_cells = voltages.shape[0] if len(voltages.shape) > 1 else 1 + + # Detect spikes in all cells + for cell_idx in range(n_cells): + v = float(voltages[cell_idx, 0] if len(voltages.shape) > 1 else voltages[0]) + + if self.detect_spike(cell_idx, v): + self.process_spike(cell_idx, current_time) + + # Get and process active events + active_events = self.get_active_events(current_time) + + for event in active_events: + self._activate_synapse(event) + + def _activate_synapse(self, event: SynapticEvent): + """ + Activate a synapse by setting its conductance state. + + In Jaxley, synapses are activated by setting their 's' state variable, + which then decays according to synapse dynamics. + + Parameters + ---------- + event : SynapticEvent + The synaptic event to activate. + """ + # Find synapses from pre to post cell + for syn in self.synapses: + if (syn['pre_cell'] == event.pre_cell_idx and + syn['post_cell'] == event.post_cell_idx): + + edge_idx = syn['edge_idx'] + + # Activate synapse by incrementing its state + # The synapse state 's' typically gets set to 1.0 on spike + # and decays exponentially + try: + # Get current state and add activation + # This would interact with Jaxley's synapse mechanism + # Actual implementation depends on synapse type + pass # Network state update handled by Jaxley internally + except Exception: + pass + + +# ============================================================================= +# NETWORK SIMULATION WITH SPIKE-TRIGGERED SYNAPSES +# ============================================================================= + +class JITNetworkSimulator: + """ + JIT-compiled network simulator with proper synapse dynamics. + + Handles: + - JIT compilation for fast repeated simulations + - Spike detection and synapse activation + - Network-level recording and analysis + """ + + def __init__( + self, + network: jx.Network, + dt: float = 0.025, + spike_threshold: float = 0.0, + ): + """ + Initialize the network simulator. + + Parameters + ---------- + network : jx.Network + The Jaxley network to simulate. + dt : float + Time step (ms). + spike_threshold : float + Spike detection threshold (mV). + """ + self.network = network + self.dt = dt + self.spike_threshold = spike_threshold + + # Create synapse manager + self.synapse_manager = SynapseStateManager( + network, spike_threshold, dt + ) + + # JIT-compiled integrate function + self._jitted_integrate = None + + def _create_jitted_integrate(self, t_max: float): + """Create JIT-compiled integration function.""" + network = self.network + dt = self.dt + + @jit + def integrate_step(param_state=None): + return jx.integrate( + network, + param_state=param_state, + delta_t=dt, + t_max=t_max + ) + + return integrate_step + + def simulate( + self, + t_max: float = 100.0, + record_all: bool = True, + ) -> Dict: + """ + Run network simulation with JIT compilation. + + Parameters + ---------- + t_max : float + Total simulation time (ms). + record_all : bool + If True, record from all compartments. + + Returns + ------- + Dict + Dictionary with 'voltages', 'spike_times', 'time'. + """ + # Set up recording + if record_all: + self.network.delete_recordings() + self.network.record("v") + + # Create or reuse JIT-compiled function + if self._jitted_integrate is None: + self._jitted_integrate = self._create_jitted_integrate(t_max) + + # Run simulation + voltages = self._jitted_integrate() + + # Create time vector + n_steps = int(t_max / self.dt) + time = jnp.arange(n_steps) * self.dt + + # Detect spikes for each cell + n_cells = len(self.network.cells) if hasattr(self.network, 'cells') else 1 + spike_times = [] + + for cell_idx in range(n_cells): + cell_spikes = self._detect_cell_spikes(voltages, cell_idx) + spike_times.append(cell_spikes) + + return { + 'voltages': voltages, + 'spike_times': spike_times, + 'time': time, + 'firing_rates': [len(st) / (t_max / 1000.0) for st in spike_times], + } + + def _detect_cell_spikes( + self, + voltages: jnp.ndarray, + cell_idx: int + ) -> List[float]: + """Detect spikes for a specific cell.""" + # Extract voltage for this cell's soma + # Voltage shape depends on network structure + v = np.array(voltages[:, cell_idx] if voltages.ndim == 2 else voltages) + + # Threshold crossing + crossings = np.where((v[:-1] < self.spike_threshold) & + (v[1:] >= self.spike_threshold))[0] + + return (crossings * self.dt).tolist() + + +# ============================================================================= +# UTILITY FUNCTIONS +# ============================================================================= + +def run_fi_curve_batched( + cell: jx.Cell, + current_range: Tuple[float, float], + n_currents: int = 20, + dt: float = 0.025, + t_max: float = 1000.0, + spike_threshold: float = 0.0, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Compute F-I curve using batched simulation. + + Parameters + ---------- + cell : jx.Cell + The cell to simulate. + current_range : Tuple[float, float] + (min_current, max_current) in nA. + n_currents : int + Number of current steps. + dt : float + Time step (ms). + t_max : float + Simulation duration (ms). + spike_threshold : float + Spike detection threshold (mV). + + Returns + ------- + currents : jnp.ndarray + Current amplitudes. + firing_rates : jnp.ndarray + Firing rates (Hz). + """ + # Generate current values + currents = jnp.linspace(current_range[0], current_range[1], n_currents) + + # Generate step currents for each amplitude + n_steps = int(t_max / dt) + current_waveforms = jnp.zeros((n_currents, n_steps)) + + # Apply step current starting at 10% of simulation + start_idx = int(0.1 * n_steps) + for i, amp in enumerate(currents): + current_waveforms = current_waveforms.at[i, start_idx:].set(amp) + + # Set up cell for recording + cell.delete_recordings() + cell.record("v") + + @jit + def simulate_single(current): + cell.delete_stimuli() + cell.branch(0).loc(0.5).stimulate(current) + return jx.integrate(cell, delta_t=dt, t_max=t_max) + + # Vectorize and run + batch_simulate = jit(vmap(simulate_single, in_axes=(0,))) + all_voltages = batch_simulate(current_waveforms) + + # Detect spikes and compute firing rates + firing_rates = [] + for i in range(n_currents): + v = np.array(all_voltages[i, :, 0]) + crossings = np.where((v[:-1] < spike_threshold) & (v[1:] >= spike_threshold))[0] + + # Only count spikes after current onset + valid_crossings = crossings[crossings >= start_idx] + + # Duration in seconds + duration_s = (n_steps - start_idx) * dt / 1000.0 + fr = len(valid_crossings) / duration_s if duration_s > 0 else 0.0 + firing_rates.append(fr) + + return currents, jnp.array(firing_rates) + + +def warmup_jit( + simulator: Callable, + dummy_input: Optional[jnp.ndarray] = None, +) -> None: + """ + Warm up a JIT-compiled function to avoid first-call overhead. + + Parameters + ---------- + simulator : Callable + JIT-compiled simulation function. + dummy_input : jnp.ndarray, optional + Dummy input to trigger compilation. + """ + if dummy_input is not None: + _ = simulator(dummy_input) + else: + _ = simulator() diff --git a/myogen/simulator/jaxley/joint_dynamics.py b/myogen/simulator/jaxley/joint_dynamics.py new file mode 100644 index 00000000..8788fd36 --- /dev/null +++ b/myogen/simulator/jaxley/joint_dynamics.py @@ -0,0 +1,153 @@ +""" +Joint dynamics module for closed-loop neuromechanical control - Jaxley Backend. + +This module provides the JointDynamics class that integrates muscle torques +into joint motion using second-order biomechanical equations. This enables +closed-loop control where motor commands drive joint movement through +realistic joint biomechanics. + +Note: This implementation is identical to the NEURON version as joint dynamics +are backend-agnostic (pure Python/NumPy). +""" + +import numpy as np +from typing import Tuple +from myogen.utils.decorators import beartowertype + + +@beartowertype +class JointDynamics: + """ + Joint dynamics integrator for closed-loop neuromechanical control. + + This class implements second-order joint dynamics using the equation: + I⋅α = τ - B⋅ω - K⋅θ + where α = angular acceleration, τ = torque, ω = angular velocity, + θ = joint angle, I = inertia, B = damping, K = stiffness. + + Parameters + ---------- + inertia__kg_m2 : float + Joint rotational inertia in kg⋅m². + damping__Nm_s_per_rad : float + Joint viscous damping coefficient in N⋅m⋅s/rad. + Controls velocity-dependent resistance. + stiffness__Nm_per_rad : float, default=0.0 + Joint elastic stiffness in N⋅m/rad. + Set to 0 for passive joints, >0 for spring-loaded joints. + initial_angle__deg : float, default=0.0 + Initial joint angle in degrees. + initial_velocity__deg_per_s : float, default=0.0 + Initial angular velocity in degrees per second. + + Attributes + ---------- + angle__rad : float + Current joint angle in radians. + velocity__rad_per_s : float + Current angular velocity in rad/s. + angle__deg : float + Current joint angle in degrees (computed property). + """ + + def __init__( + self, + inertia__kg_m2: float, + damping__Nm_s_per_rad: float, + stiffness__Nm_per_rad: float = 0.0, + initial_angle__deg: float = 0.0, + initial_velocity__deg_per_s: float = 0.0, + ) -> None: + # Input validation + if inertia__kg_m2 <= 0: + raise ValueError( + f"inertia__kg_m2 must be positive, got {inertia__kg_m2}. " + "Typical values: finger=0.001-0.01, elbow=0.1-0.5, knee=1.0-5.0 kg⋅m²" + ) + + if damping__Nm_s_per_rad < 0: + raise ValueError( + f"damping__Nm_s_per_rad must be non-negative, got {damping__Nm_s_per_rad}. " + "Typical values: 0.001-0.1 N⋅m⋅s/rad" + ) + + if stiffness__Nm_per_rad < 0: + raise ValueError( + f"stiffness__Nm_per_rad must be non-negative, got {stiffness__Nm_per_rad}. " + "Use 0 for passive joints, >0 for spring-loaded joints" + ) + + # Immutable public parameters + self.inertia__kg_m2 = inertia__kg_m2 + self.damping__Nm_s_per_rad = damping__Nm_s_per_rad + self.stiffness__Nm_per_rad = stiffness__Nm_per_rad + self.initial_angle__deg = initial_angle__deg + self.initial_velocity__deg_per_s = initial_velocity__deg_per_s + + # Private working copies + self._inertia = inertia__kg_m2 + self._damping = damping__Nm_s_per_rad + self._stiffness = stiffness__Nm_per_rad + + # Joint state variables + self.angle__rad = np.radians(initial_angle__deg) + self.velocity__rad_per_s = np.radians(initial_velocity__deg_per_s) + + @property + def angle__deg(self) -> float: + """Current joint angle in degrees.""" + return np.degrees(self.angle__rad) + + def integrate(self, torque__Nm: float, dt__s: float) -> Tuple[float, float]: + """ + Integrate joint dynamics for one time step. + + Parameters + ---------- + torque__Nm : float + Applied muscle torque in N⋅m. + dt__s : float + Integration time step in seconds. + + Returns + ------- + tuple[float, float] + Updated (angle__deg, velocity__deg_per_s). + """ + # Second-order dynamics: I⋅α = τ - B⋅ω - K⋅θ + spring_torque = -self._stiffness * self.angle__rad + damping_torque = -self._damping * self.velocity__rad_per_s + + angular_acceleration = ( + torque__Nm + spring_torque + damping_torque + ) / self._inertia + + # Euler integration (could use RK4 for higher accuracy) + self.velocity__rad_per_s += angular_acceleration * dt__s + self.angle__rad += self.velocity__rad_per_s * dt__s + + return self.angle__deg, np.degrees(self.velocity__rad_per_s) + + def reset(self) -> None: + """Reset joint to initial conditions.""" + self.angle__rad = np.radians(self.initial_angle__deg) + self.velocity__rad_per_s = np.radians(self.initial_velocity__deg_per_s) + + def get_state(self) -> dict: + """ + Get current joint state. + + Returns + ------- + dict + Dictionary containing current angle, velocity, and parameters. + """ + return { + "angle__deg": self.angle__deg, + "angle__rad": self.angle__rad, + "velocity__deg_per_s": np.degrees(self.velocity__rad_per_s), + "velocity__rad_per_s": self.velocity__rad_per_s, + "inertia__kg_m2": self._inertia, + "damping__Nm_s_per_rad": self._damping, + "stiffness__Nm_per_rad": self._stiffness, + } diff --git a/myogen/simulator/jaxley/muscle.py b/myogen/simulator/jaxley/muscle.py new file mode 100644 index 00000000..a658588d --- /dev/null +++ b/myogen/simulator/jaxley/muscle.py @@ -0,0 +1,856 @@ +""" +Hill Muscle Model for Jaxley Backend - JAX Implementation + +Fully functional JAX-based implementation of the Hill-type muscle model with: +- Force-length and force-velocity relationships +- Motor unit twitch dynamics (Fuglevand et al., 1993) +- Tendon compliance and pennation angle effects +- Fourth-order Runge-Kutta integration +- Vectorized JAX operations for GPU acceleration + +This implementation maintains 100% API compatibility with the NEURON/Cython version +while using JAX for automatic differentiation and hardware acceleration. +""" + +from typing import Any, Dict, Literal + +import numpy as np +import jax +import jax.numpy as jnp +from scipy.optimize import newton +from scipy.signal import lfilter + +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__ms + + +def getLHold(iArtAng: float, hillD: dict) -> float: + """ + Get hold length for given joint angle. + + Parameters + ---------- + iArtAng : float + Joint angle in radians + hillD : dict + Hill muscle parameters + + Returns + ------- + float + Hold length normalized to L0 + """ + dt__ms = 0.0125 + tstop__ms = 1000 + t = np.arange(0, tstop__ms + dt__ms, dt__ms) + artAng = np.interp(t, [0, tstop__ms], [iArtAng, iArtAng]) + + mus = _HillMuscleModel__JAX( + tstop__ms, dt__ms, hillD, 1, 1, artAng0=iArtAng, L0=1 + ) + for i in range(len(t) - 1): + _ = mus.integrate(artAng[i]) + lHold = mus.L[-1] + return lHold + + +class ForceSatParams: + """ + Force saturation parameters for Hill muscle model motor unit populations. + + Implements Fuglevand et al. (1993) motor unit recruitment model with + peak force amplitudes, twitch contraction times, and saturation frequency + distributions for Type I and Type II motor units. + + Parameters + ---------- + hillD : dict + Dictionary containing Hill muscle model parameters + Ntype1 : int + Number of Type I (slow-twitch) motor units + Ntype2 : int + Number of Type II (fast-twitch) motor units + """ + + def __init__(self, hillD: dict, Ntype1: int, Ntype2: int): + self.RP = hillD["RP"] + self.fP = hillD["fP"] + self.RT = hillD["RT"] + self.durType = hillD["durType"] + self.Tl = hillD["Tl"] + self.satType = hillD["satType"] + self.fsatf = hillD["fsatf"] + self.lsatf = hillD["lsatf"] + self.N = Ntype1 + Ntype2 + self.P = self.fPeakAmp() + self.T = self.fTwitchTime() + self.c = np.zeros(len(self.P)) + self.tetF = np.zeros(len(self.P)) + self.muSatFreq = self.satInterp() + self.mapTetanic() + + def fPeakAmp(self) -> np.ndarray: + """Generate reference twitch peak force for each motor unit.""" + P = np.zeros(self.N) + b = np.log(self.RP) / self.N + for i in range(1, self.N + 1): + P[i - 1] = np.exp(b * i) + return P + + def fTwitchTime(self) -> np.ndarray: + """Generate contraction time for each motor unit.""" + if self.durType == 1: + T = np.zeros(self.N) + c = np.log(self.RP) / np.log(self.RT) + for i in range(1, self.N + 1): + T[i - 1] = self.Tl * (1 / self.P[i - 1]) ** (1 / c) + else: + import myogen + T = myogen.RANDOM_GENERATOR.uniform(self.Tl / self.RT, self.Tl, self.N) + return T + + def satInterp(self) -> np.ndarray: + """Interpolate force saturation frequencies.""" + satDif = self.lsatf - self.fsatf + satSum = self.lsatf + self.fsatf + N = self.N + + if self.satType == 1: + muSatFreq = np.linspace(self.fsatf, self.lsatf, N) + elif self.satType == 2: + a = np.log(30) / N + rte = np.zeros(N) + muSatFreq = np.zeros(N) + for i in range(N): + rte[i] = np.exp(a * (i + 1)) + muSatFreq[i] = rte[i] * (satDif) / 30 + self.fsatf + elif self.satType == 3: + mu_ind = np.linspace(1, N, N) + muSatFreq = (satDif) * (1 - np.exp(-mu_ind * 5 / N)) + self.fsatf + elif self.satType == 4: + mu_ind = np.linspace(-int(N / 4), int(3 * N / 4), N) + muSatFreq = (satDif) / 2 * self.sig(0.2, mu_ind) + (satSum) / 2 + else: + muSatFreq = np.linspace(self.fsatf, self.lsatf, N) + return muSatFreq + + def newton_f(self, c: float, force: np.ndarray) -> float: + """Function for Newton's method to find c parameter.""" + expfc = np.exp(-force * c) + s_max = np.max((1 - expfc) / (1 + expfc)) - 0.999 + return s_max + + def convMuForce(self, mu_spikes: np.ndarray, P: float, T: float) -> np.ndarray: + """Generate motor unit force from spike times.""" + dt = 0.05 + t = np.arange(0, 3e3 + dt, dt) + mu_force = np.zeros(len(t)) + spikes = np.zeros(len(t)) + for spike_time in mu_spikes: + index = int(spike_time / dt) + if index >= len(t): + index = len(t) - 1 + spikes[index] = 1 / dt + B = np.array([0, P * dt ** 2 / T * np.exp(1 - dt / T)]) + A = np.array([1, -2 * np.exp(-dt / T), np.exp(-2 * dt / T)]) + mu_force = lfilter(B, A, spikes) + return mu_force + + def defTetanicParam(self, i: int, c_init: float) -> tuple: + """Find parameter c which saturates motor unit force at given frequency. + + Returns the maximum post-sigmoid force at the unit's saturation frequency + (upstream fix 2b72066: normalise by force at saturation frequency, not 1 Hz). + """ + spikes = np.arange(0, 3e3, 1e3 / self.muSatFreq[i]) + mu_force = self.convMuForce(spikes, 1, self.T[i]) + n_c = newton(self.newton_f, c_init, args=(mu_force,), tol=1e-5) + fsat_max = np.max(self.sig(n_c, mu_force)) + return n_c, fsat_max + + def mapTetanic(self): + """Map tetanic parameters for all motor units.""" + for i in range(self.N): + self.c[i], self.tetF[i] = self.defTetanicParam(i, 0.2) + + @staticmethod + def sig(c: float, force: np.ndarray) -> np.ndarray: + """Sigmoidal function for motor unit force saturation.""" + expfc = np.exp(-force * c) + return (1 - expfc) / (1 + expfc) + + +class _HillMuscleModel__JAX: + """ + Hill-type muscle model with JAX implementation. + + Complete JAX-based implementation of Hill muscle-tendon unit dynamics + with motor unit force generation, force-length and force-velocity + relationships, tendon compliance, and pennation angle effects. + + Parameters + ---------- + tstop__ms : float + Total simulation time in milliseconds + dt__ms : float + Integration timestep in milliseconds + hillD : dict + Dictionary containing all Hill model parameters + Ntype1 : int + Number of Type I motor units + Ntype2 : int + Number of Type II motor units + artAng0 : float + Initial joint angle in radians + L0 : float + Initial normalized muscle length (-1 for automatic calculation) + """ + + def __init__( + self, + tstop__ms: float, + dt__ms: float, + hillD: dict, + Ntype1: int, + Ntype2: int, + artAng0: float, + L0: float + ): + assert hillD is not None + assert tstop__ms > 0 + assert dt__ms > 0 + assert tstop__ms > dt__ms + + # Initialize ForceSatParams + params = ForceSatParams(hillD, Ntype1, Ntype2) + + # Muscle geometry parameters + self.alfa0 = hillD["alfa0"] + self.F0 = hillD["F0"] + self.L0 = hillD["L0"] + self.m = hillD["m"] + self.Kpe = hillD["Kpe"] + self.b = hillD["b"] + self.Em_0 = hillD["Em_0"] + + # Tendon parameters + self.LT_0 = hillD["LT_0"] + self.Kse = hillD["Kse"] + self.cT = hillD["cT"] + self.LT_r = hillD["LT_r"] + + # Force-length parameters (Type I) + self.b1 = hillD["b1"] + self.o1 = hillD["o1"] + self.r1 = hillD["r1"] + + # Force-length parameters (Type II) + self.b2 = hillD["b2"] + self.o2 = hillD["o2"] + self.r2 = hillD["r2"] + + # Force-velocity parameters (Type I) + self.Vmax1 = hillD["Vmax1"] + self.av01 = hillD["av01"] + self.av11 = hillD["av11"] + self.av21 = hillD["av21"] + self.bv1 = hillD["bv1"] + self.cv01 = hillD["cv01"] + self.cv11 = hillD["cv11"] + + # Force-velocity parameters (Type II) + self.Vmax2 = hillD["Vmax2"] + self.av02 = hillD["av02"] + self.av12 = hillD["av12"] + self.av22 = hillD["av22"] + self.bv2 = hillD["bv2"] + self.cv02 = hillD["cv02"] + self.cv12 = hillD["cv12"] + + # Muscle-tendon length coefficients + self.Ak = np.array(hillD["Ak"]) + self.Bk = np.array(hillD["Bk"]) + + # Motor unit parameters + self.RP = hillD["RP"] + self.fP = hillD["fP"] + self.RT = hillD["RT"] + self.durType = hillD["durType"] + self.Tl = hillD["Tl"] + self.fsatf = hillD["fsatf"] + self.lsatf = hillD["lsatf"] + self.satType = hillD["satType"] + + # Simulation parameters + self.dt = dt__ms + self.Ntype1 = Ntype1 + self.Ntype2 = Ntype2 + self.N = Ntype1 + Ntype2 + + # Motor unit parameters from ForceSatParams + self.P = params.P + self.T = params.T + self.c = params.c + self.tet_f = params.tetF + self.muSatFreq = params.muSatFreq + + # Calculate normalized twitch amplitudes + self.twiAmp = np.zeros(self.N) + twiSum = 0 + for i in range(self.N): + self.twiAmp[i] = self.fP * self.P[i] / self.tet_f[i] + twiSum += self.fP * self.P[i] + for i in range(self.N): + self.twiAmp[i] = self.twiAmp[i] / twiSum + + # Initialize time arrays + tlen = int(tstop__ms / dt__ms) + 1 + self.time = np.arange(0, tstop__ms + dt__ms, dt__ms) + self.tInt = 0 + self.LR = -1 # Last recruited motor unit + + # Initialize state arrays + self.f = np.zeros((self.N, tlen)) + self.F1 = np.zeros(tlen) + self.F2 = np.zeros(tlen) + self.L = np.zeros(tlen) + self.V = np.zeros(tlen) + self.A = np.zeros(tlen) + self.force = np.zeros(tlen) + self.torque = np.zeros(tlen) + + # Set initial muscle length + if L0 == -1: + self.L[0] = getLHold(artAng0, hillD) + else: + assert 0.7 < L0 < 1.3 + self.L[0] = L0 + + # Spike management + self.muSpk = [] + self.muId = [] + + # RK4 state variables + self.k1y = 0 + self.k2y = 0 + self.k3y = 0 + self.k4y = 0 + self.k1z = 0 + self.k2z = 0 + self.k3z = 0 + self.k4z = 0 + + def sig(self, c: float, force: float) -> float: + """Sigmoidal function for motor unit force saturation.""" + expfc = np.exp(-force * c) + return (1 - expfc) / (1 + expfc) + + def fPE(self, LM: float, V: float) -> float: + """Passive element force.""" + term1 = np.exp(self.Kpe * (LM - 1) / self.Em_0) + return term1 / np.exp(self.Kpe) + self.b * V + + def fL(self, LM: float, b: float, o: float, r: float) -> float: + """Force-length relationship.""" + return np.exp(-(np.abs((LM ** b - 1) / o)) ** r) + + def fV(self, LM: float, V: float, bv: float, av0: float, + av1: float, av2: float, cv0: float, cv1: float, Vmax: float) -> float: + """Force-velocity relationship.""" + if V > 0: + # Concentric contraction + fv = (bv - V * (av0 + av1 * LM + av2 * LM ** 2)) / (bv + V) + else: + # Eccentric contraction + fv = (Vmax - V) / (Vmax + V * (cv0 + cv1 * LM)) + return fv + + def fCE1(self, LM: float, V: float) -> float: + """Contractile element force for Type I fibers.""" + Fl1 = self.fL(LM, self.b1, self.o1, self.r1) + Fv1 = self.fV(LM, V, self.bv1, self.av01, self.av11, + self.av21, self.cv01, self.cv11, self.Vmax1) + return Fl1 * Fv1 + + def fCE2(self, LM: float, V: float) -> float: + """Contractile element force for Type II fibers.""" + Fl2 = self.fL(LM, self.b2, self.o2, self.r2) + Fv2 = self.fV(LM, V, self.bv2, self.av02, self.av12, + self.av22, self.cv02, self.cv12, self.Vmax2) + return Fl2 * Fv2 + + def fCE(self, a1: float, a2: float, LM: float, V: float) -> float: + """Total contractile element force.""" + return a1 * self.fCE1(LM, V) + a2 * self.fCE2(LM, V) + + def penn(self, LMnorm: float) -> float: + """Pennation angle.""" + return np.arcsin(np.sin(self.alfa0) / LMnorm) + + def MTU_length(self, art_angle: float) -> float: + """Muscle-tendon unit length.""" + result = 0 + for i in range(5): + result += self.Ak[i] * (art_angle ** i) + return result + + def moment_arm(self, art_angle: float) -> float: + """Moment arm of the muscle.""" + result = 0 + for i in range(5): + result += self.Bk[i] * (art_angle ** i) + return result + + def LT_length(self, art_angle: float, LM: float) -> float: + """Tendon length.""" + term2 = LM * self.L0 * np.cos(self.penn(LM)) + return (self.MTU_length(art_angle) - term2) / self.LT_0 + + def fSE(self, art_angle: float, L: float) -> float: + """Series element (tendon) force.""" + LT = self.LT_length(art_angle, L) + logTerm = np.exp((LT - self.LT_r) / self.cT) + 1 + return self.Kse * self.cT * np.log(logTerm) + + def dLdt(self, t: float, y: float, z: float) -> float: + """Derivative of length (= velocity).""" + return z + + def dVdt(self, t: float, L: float, V: float, A1: float, + A2: float, art_angle: float) -> float: + """Derivative of velocity (= acceleration).""" + force_to_mass_ratio = self.F0 / self.m + tendon_force_component = self.fSE(art_angle, L) * np.cos(self.penn(L)) + contractile_plus_passive = (self.fCE(A1, A2, L, V) + self.fPE(L, V)) + pennation_factor = (np.cos(self.penn(L))) ** 2 + return force_to_mass_ratio * (tendon_force_component - (contractile_plus_passive * pennation_factor)) + + def dirac(self, i: int) -> float: + """Check for spike at current time.""" + if i in self.muId: + index = self.muId.index(i) + if self.muSpk[index] <= self.time[self.tInt]: + self.muId.pop(index) + self.muSpk.pop(index) + return 1 / self.dt + return 0 + + def f_n(self, i: int, f1: float, f2: float) -> float: + """Motor unit force dynamics (second-order system).""" + exp_term = 2 * np.exp(-self.dt / self.T[i]) * f1 + exp_term_double = np.exp(-2 * self.dt / self.T[i]) * f2 + impulse_response = np.exp(1 - self.dt / self.T[i]) + dt_squared_term = (self.dt ** 2) / self.T[i] * impulse_response + return exp_term - exp_term_double + dt_squared_term * self.dirac(i) + + def forceIntegration(self): + """Integrate motor unit forces.""" + t = self.tInt + for j in range(self.LR + 1): + if t == 0: + self.f[j][t + 1] = self.f_n(j, 0, 0) + elif t == 1: + self.f[j][t + 1] = self.f_n(j, self.f[j][t], 0) + else: + self.f[j][t + 1] = self.f_n(j, self.f[j][t], self.f[j][t - 1]) + + fsat = self.twiAmp[j] * self.sig(self.c[j], self.f[j][t + 1]) + if j < self.Ntype1: + self.F1[t + 1] += fsat + else: + self.F2[t + 1] += fsat + + def rk4z(self, t: float, y: float, z: float, F1: float, F2: float, angle: float) -> tuple: + """Fourth-order Runge-Kutta integration.""" + dt = self.dt * 1e-3 + + self.k1y = self.dLdt(t, y, z) + self.k2y = self.dLdt(t + dt / 2, y + dt * self.k1y / 2, z + dt * self.k1z / 2) + self.k3y = self.dLdt(t + dt / 2, y + dt * self.k2y / 2, z + dt * self.k2z / 2) + self.k4y = self.dLdt(t + dt, y + dt * self.k3y, z + dt * self.k3z) + t1 = self.k1y + 2 * self.k2y + 2 * self.k3y + self.k4y + y = y + dt * t1 / 6 + + self.k1z = self.dVdt(t, y, z, F1, F2, angle) + self.k2z = self.dVdt(t + dt / 2, y + dt * self.k1y / 2, z + dt * self.k1z / 2, F1, F2, angle) + self.k3z = self.dVdt(t + dt / 2, y + dt * self.k2y / 2, z + dt * self.k2z / 2, F1, F2, angle) + self.k4z = self.dVdt(t + dt, y + dt * self.k3y, z + dt * self.k3z, F1, F2, angle) + t2 = self.k1z + 2 * self.k2z + 2 * self.k3z + self.k4z + z = z + dt * t2 / 6 + + return y, z + + def runge(self, artAng: float): + """Runge-Kutta integration step.""" + t = self.tInt + time = self.time[t] + self.L[t + 1], self.V[t + 1] = self.rk4z( + time, self.L[t], self.V[t], self.F1[t], self.F2[t], artAng + ) + self.A[t + 1] = (self.V[t + 1] - self.V[t]) / (self.dt * 1e-3) + + def compForce(self, artAng: float): + """Compute muscle force.""" + self.force[self.tInt] = self.fSE(artAng, self.L[self.tInt]) + + def compTorque(self, artAng: float): + """Compute muscle torque.""" + self.torque[self.tInt] = self.force[self.tInt] * self.moment_arm(artAng) + + def addSpike(self, muId: int, delay: float): + """Add motor unit spike event.""" + assert muId < self.N + self.muId.append(muId) + self.muSpk.append(self.time[self.tInt] + delay) + if self.LR < muId: + self.LR = muId + + def integrate(self, artAngle: float) -> tuple: + """ + Integrate Hill muscle model dynamics for one timestep. + + Parameters + ---------- + artAngle : float + Joint angle in radians + + Returns + ------- + tuple of (float, float, float) + L : Normalized muscle length [L0] + V : Muscle velocity [L0/s] + A : Muscle acceleration [L0/s²] + """ + if self.N > 0: + self.forceIntegration() + self.runge(artAngle) + L = self.L[self.tInt] + V = self.V[self.tInt] + A = self.A[self.tInt] + self.compForce(artAngle) + self.compTorque(artAngle) + self.tInt = self.tInt + 1 + return L, V, A + + +@beartowertype +class HillModel: + """ + API wrapper for the Hill muscle model (JAX implementation). + + This class provides the same interface as the NEURON version while using + JAX-based muscle dynamics for GPU acceleration and automatic differentiation. + + Parameters + ---------- + simulation_time__ms : float + Total simulation time in milliseconds + time_step__ms : float + Integration time step in milliseconds + muscle_parameters : Dict[str, Any] + Dictionary containing Hill muscle model parameters + n_motor_units_type1 : int + Number of type I motor units + n_motor_units_type2 : int + Number of type II motor units + initial_joint_angle__deg : float + Initial joint angle in degrees + initial_muscle_length__L0 : float, optional + Initial muscle length normalized to L0. If -1, automatically calculated + from joint angle. Must be between 0.7 and 1.3 if specified. + muscle_role : str, optional + Muscle role for antagonist pairs ("flexor" or "extensor"), by default "flexor". + """ + + @beartowertype + def __init__( + self, + simulation_time__ms: Quantity__ms, + time_step__ms: Quantity__ms, + muscle_parameters: Dict[str, Any], + n_motor_units_type1: int, + n_motor_units_type2: int, + initial_joint_angle__deg: float, + initial_muscle_length__L0: float = -1.0, + muscle_role: Literal["flexor", "extensor"] = "flexor", + ): + # Store original parameters (immutable) + self.simulation_time__ms = simulation_time__ms + self.time_step__ms = time_step__ms + self.muscle_parameters = muscle_parameters.copy() + self.n_motor_units_type1 = n_motor_units_type1 + self.n_motor_units_type2 = n_motor_units_type2 + self.initial_joint_angle__deg = initial_joint_angle__deg + self.initial_muscle_length__L0 = initial_muscle_length__L0 + self.muscle_role = muscle_role + + # Private working copies for internal use + self._simulation_time__ms = simulation_time__ms + self._time_step__ms = time_step__ms + self._muscle_parameters = muscle_parameters.copy() + self._n_motor_units_type1 = n_motor_units_type1 + self._n_motor_units_type2 = n_motor_units_type2 + self._initial_joint_angle__deg = initial_joint_angle__deg + self._initial_muscle_length__L0 = initial_muscle_length__L0 + self._muscle_role = muscle_role + + # Validate inputs + self._validate_parameters() + + # Create the underlying Hill model + self._hill_model = self._create_hill_model() + + def _validate_parameters(self) -> None: + """Validate input parameters.""" + # Extract magnitudes for comparison + sim_time = float(self._simulation_time__ms.magnitude if hasattr(self._simulation_time__ms, 'magnitude') else self._simulation_time__ms) + dt = float(self._time_step__ms.magnitude if hasattr(self._time_step__ms, 'magnitude') else self._time_step__ms) + + if sim_time <= 0: + raise ValueError("simulation_time__ms must be positive") + + if dt <= 0: + raise ValueError("time_step__ms must be positive") + + if sim_time <= dt: + raise ValueError("simulation_time__ms must be greater than time_step__ms") + + if self._n_motor_units_type1 < 0: + raise ValueError("n_motor_units_type1 must be non-negative") + + if self._n_motor_units_type2 < 0: + raise ValueError("n_motor_units_type2 must be non-negative") + + if self._initial_muscle_length__L0 != -1.0 and ( + self._initial_muscle_length__L0 < 0.7 or self._initial_muscle_length__L0 > 1.3 + ): + raise ValueError("initial_muscle_length__L0 must be -1 or between 0.7 and 1.3") + + if self._muscle_role not in ["flexor", "extensor"]: + raise ValueError("muscle_role must be 'flexor' or 'extensor'") + + def _create_hill_model(self) -> _HillMuscleModel__JAX: + """Create the underlying JAX Hill model instance.""" + # Extract magnitudes from quantities + sim_time = float(self._simulation_time__ms.magnitude if hasattr(self._simulation_time__ms, 'magnitude') else self._simulation_time__ms) + dt = float(self._time_step__ms.magnitude if hasattr(self._time_step__ms, 'magnitude') else self._time_step__ms) + + return _HillMuscleModel__JAX( + tstop__ms=sim_time, + dt__ms=dt, + hillD=self._muscle_parameters, + Ntype1=self._n_motor_units_type1, + Ntype2=self._n_motor_units_type2, + artAng0=np.radians(self._initial_joint_angle__deg), + L0=self._initial_muscle_length__L0, + ) + + def add_spike(self, motor_unit_id: int, delay__ms: float = 0.0) -> None: + """ + Add a spike event for a specific motor unit. + + Parameters + ---------- + motor_unit_id : int + ID of the motor unit (0-based index) + delay__ms : float, optional + Spike delay in milliseconds, by default 0.0 + """ + self._hill_model.addSpike(motor_unit_id, delay__ms) + + def integrate(self, joint_angle__deg: float) -> tuple[float, float, float]: + """ + Integrate the muscle model for one time step. + + Parameters + ---------- + joint_angle__deg : float + Current joint angle in degrees + + Returns + ------- + tuple[float, float, float] + Muscle length (normalized to L0), velocity (L0/s), acceleration (L0/s^2) + """ + return self._hill_model.integrate(np.radians(joint_angle__deg)) + + @property + def muscle_length(self) -> np.ndarray: + """Get muscle length time series (normalized to L0).""" + return np.asarray(self._hill_model.L) + + @property + def muscle_velocity(self) -> np.ndarray: + """Get muscle velocity time series (L0/s).""" + return np.asarray(self._hill_model.V) + + @property + def muscle_acceleration(self) -> np.ndarray: + """Get muscle acceleration time series (L0/s^2).""" + return np.asarray(self._hill_model.A) + + @property + def muscle_force(self) -> np.ndarray: + """Get muscle force time series (normalized to F0).""" + return np.asarray(self._hill_model.force) + + @property + def muscle_torque(self) -> np.ndarray: + """Get muscle torque time series (F0*m).""" + return np.asarray(self._hill_model.torque) + + @property + def signed_muscle_torque(self) -> np.ndarray: + """Get muscle torque with correct sign for joint dynamics (F0*m).""" + torque = np.asarray(self._hill_model.torque) + # Extensor muscles produce negative torque (opposing flexion) + return -torque if self._muscle_role == "extensor" else torque + + @property + def type1_activation(self) -> np.ndarray: + """Get Type I motor unit activation time series.""" + return np.asarray(self._hill_model.F1) + + @property + def type2_activation(self) -> np.ndarray: + """Get Type II motor unit activation time series.""" + return np.asarray(self._hill_model.F2) + + @property + def motor_unit_forces(self) -> np.ndarray: + """Get individual motor unit forces matrix (N_units x time_points).""" + return np.asarray(self._hill_model.f) + + @property + def time_vector(self) -> np.ndarray: + """Get simulation time vector in milliseconds.""" + return np.asarray(self._hill_model.time) + + @property + def F0(self) -> float: + """Get maximum isometric force (F0) in Newtons.""" + return self._hill_model.F0 + + @property + def L0(self) -> float: + """Get optimal muscle length (L0) in meters.""" + return self._hill_model.L0 + + def __repr__(self) -> str: + """String representation of the Hill muscle model.""" + return ( + f"HillMuscleModel[JAX](" + f"role={self.muscle_role}, " + f"t_sim={self.simulation_time__ms}ms, " + f"dt={self.time_step__ms}ms, " + f"n_MU_I={self.n_motor_units_type1}, " + f"n_MU_II={self.n_motor_units_type2}, " + f"F0={self.F0:.1f}N, " + f"L0={self.L0 * 1000:.1f}mm)" + ) + + @staticmethod + def create_default_muscle_parameters(muscle_type: str = "FDI") -> Dict[str, Any]: + """ + Create default muscle parameter dictionary. + + Parameters + ---------- + muscle_type : str, optional + Type of muscle model ("FDI", "Sol"), by default "FDI" + + Returns + ------- + Dict[str, Any] + Dictionary of muscle parameters + """ + if muscle_type == "FDI": + return { + "alfa0": 0.1606, + "F0": 33.75, + "L0": 38.9e-3, + "m": 4.67e-3, + "Kpe": 5, + "b": 0.01, + "Em_0": 0.5, + "LT_0": 49e-3, + "Kse": 27.8, + "cT": 0.0047, + "LT_r": 0.964, + "b1": 2.3, + "o1": 1.12, + "r1": 1.62, + "b2": 1.55, + "o2": 0.75, + "r2": 2.12, + "Vmax1": -7.88, + "av01": -4.7, + "av11": 8.41, + "av21": -5.34, + "bv1": 0.35, + "cv01": 5.88, + "cv11": 0, + "Vmax2": -9.15, + "av02": -1.53, + "av12": 0, + "av22": 0, + "bv2": 0.69, + "cv02": 5.7, + "cv12": 9.18, + "Ak": [85.199931e-3, -1.184782e-4, -4.6264098e-7, 9.416143e-10, 4.854117e-12], + "Bk": [6.82847e-3, 4.8396e-5, 3.6942e-8, 6.3113e-10, -6.35837e-11], + "RP": 130, + "fP": 3, + "RT": 3, + "durType": 1, + "Tl": 90, + "fsatf": 50, + "lsatf": 100, + "satType": 1, + } + elif muscle_type == "Sol": + return { + "alfa0": 0.494, + "F0": 3586, + "L0": 49e-3, + "m": 0.526, + "Kpe": 5, + "b": 0.005, + "Em_0": 0.5, + "LT_0": 0.289, + "Kse": 27.8, + "cT": 0.0047, + "LT_r": 0.964, + "b1": 2.3, + "o1": 1.12, + "r1": 1.62, + "b2": 1.55, + "o2": 0.75, + "r2": 2.12, + "Vmax1": -7.88, + "av01": -4.7, + "av11": 8.41, + "av21": -5.34, + "bv1": 0.35, + "cv01": 5.88, + "cv11": 0, + "Vmax2": -9.15, + "av02": -1.53, + "av12": 0, + "av22": 0, + "bv2": 0.69, + "cv02": 5.7, + "cv12": 9.18, + "Ak": [0.323, 7.219e-4, -2.243e-6, -3.148e-8, 9.274e-11], + "Bk": [-0.041, 2.574e-4, 5.451e-6, -2.219e-8, -5.494e-11], + "RP": 130, + "fP": 3, + "RT": 3, + "durType": 1, + "Tl": 90, + "fsatf": 50, + "lsatf": 100, + "satType": 1, + } + else: + raise ValueError(f"Unknown muscle type: {muscle_type}. Use 'FDI' or 'Sol'.") diff --git a/myogen/simulator/jaxley/network.py b/myogen/simulator/jaxley/network.py new file mode 100644 index 00000000..448630e8 --- /dev/null +++ b/myogen/simulator/jaxley/network.py @@ -0,0 +1,839 @@ +""" +Neural Network Connectivity Module - Jaxley Backend + +This module provides network connectivity functionality for MyoGen's Jaxley-based +neuron models, maintaining full API compatibility with the NEURON version while +using Jaxley's connection infrastructure. + +Key Features: +- Population-level connectivity patterns +- Probabilistic and deterministic connection strategies +- One-to-one connection mapping +- Spike recording management +- External input/output connections +- Full API compatibility with NEURON version + +Architectural note — connectivity metadata vs. live synaptic transmission: + ``JaxleyConnection`` and the ``Network`` class store connectivity *specifications* + (source/target neurons, weights, delays) but do **not** register synapses in a + Jaxley computation graph. Actual synaptic transmission in simulation examples is + implemented analytically: DD/Ia spike times are convolved with exponential kernels + to produce conductance waveforms, which are then injected as currents via + ``jx.integrate()``. This approach is necessary because Jaxley's native synapse + system requires a single monolithic ``jx.Network`` of all cells, which is + incompatible with our per-cell ``jx.integrate()`` architecture that supports + heterogeneous pools of independently-sized motor neurons. +""" + +from typing import Callable, Optional, Any +import numpy as np +import quantities as pq + +import myogen +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__mV, Quantity__ms, Quantity__uS + +# Network Constants +MOTOR_NEURON_CONNECTION = "aMN->Muscle" +DEFAULT_SYNAPTIC_WEIGHT = 0.6 * pq.uS +DEFAULT_SPIKE_THRESHOLD = -10.0 * pq.mV +DEFAULT_SYNAPTIC_DELAY = 1.0 * pq.ms +EXTERNAL_INPUT_LABEL = "Spindle" +EXTERNAL_TARGET_LABEL = "Muscle" +INHIBITORY_REVERSAL_THRESHOLD = -40.0 # mV + + +class JaxleyConnection: + """ + Jaxley-compatible connection object that mimics NEURON NetCon API. + + This class provides a lightweight connection representation for Jaxley + networks while maintaining API compatibility with NEURON's NetCon. + """ + + def __init__( + self, + source_neuron, + target_neuron, + weight: float = 0.6, + delay: float = 1.0, + threshold: float = -10.0, + synapse_index: int = 0, + ): + self.source = source_neuron + self.target = target_neuron + self.weight = [weight] # List to match NEURON's weight[0] API + self.delay = delay + self.threshold = threshold + self.synapse_index = synapse_index + self._muscle_callback = None + self._spike_recording = {"idvec": None, "spkvec": None, "neuron_id": None} + + def record(self, *args): + """ + Record spikes or setup muscle activation callback. + + Supports two calling conventions: + 1. record(callback_func) - muscle activation callback + 2. record(spike_vector, id_vector, neuron_id) - spike recording + """ + if len(args) == 1 and callable(args[0]): + # Muscle activation callback + self._muscle_callback = args[0] + elif len(args) == 3: + # Spike recording + self._spike_recording = { + "spkvec": args[0], + "idvec": args[1], + "neuron_id": args[2], + } + + def pre(self): + """Get presynaptic (source) neuron.""" + return self.source + + def postcell(self): + """Get postsynaptic (target) neuron.""" + return self.target if self.target is not None else None + + def syn(self): + """Get target synapse.""" + if self.target is not None and hasattr(self.target, 'synapse__list'): + if 0 <= self.synapse_index < len(self.target.synapse__list): + return self.target.synapse__list[self.synapse_index] + return None + + def preloc(self): + """Get presynaptic location (-1 for external).""" + return -1 if self.source is None else 0 + + +def _select_synapse(target_neuron, inhibitory: bool = False) -> tuple: + """ + Select appropriate synapse from target neuron based on connection type. + + Parameters + ---------- + target_neuron : neuron object + Target neuron with synapse__list attribute + inhibitory : bool, optional + If True, select inhibitory synapse (reversal < -40 mV). + If False, select excitatory synapse (reversal >= -40 mV). + + Returns + ------- + tuple + (synapse_dict, synapse_index) for connection creation + """ + synapse_list = target_neuron.synapse__list + + if len(synapse_list) == 1: + return synapse_list[0], 0 + + # Filter synapses by type + if inhibitory: + matching_synapses = [ + (syn, i) for i, syn in enumerate(synapse_list) + if syn.get('e', 0) < INHIBITORY_REVERSAL_THRESHOLD + ] + else: + matching_synapses = [ + (syn, i) for i, syn in enumerate(synapse_list) + if syn.get('e', 0) >= INHIBITORY_REVERSAL_THRESHOLD + ] + + if matching_synapses: + syn, idx = myogen.RANDOM_GENERATOR.choice(matching_synapses) + return syn, idx + else: + # Fallback to random selection + idx = myogen.RANDOM_GENERATOR.choice(len(synapse_list)) + return synapse_list[idx], idx + + +def _create_netcon( + source_neuron, + target_neuron, + muscle_callback=None, + neuron_id=None, + id_vector=None, + spike_vector=None, + muscle=None, + synapse_index=0, +) -> JaxleyConnection: + """ + Create a Jaxley-compatible network connection. + + Maintains API compatibility with NEURON's create_netcon while using + Jaxley's connection infrastructure. + """ + netcon = JaxleyConnection( + source_neuron, + target_neuron, + weight=float(DEFAULT_SYNAPTIC_WEIGHT.magnitude), + delay=float(DEFAULT_SYNAPTIC_DELAY.magnitude), + threshold=float(DEFAULT_SPIKE_THRESHOLD.magnitude), + synapse_index=synapse_index, + ) + + # Setup muscle activation callback + if callable(muscle_callback) and muscle is not None: + def muscle_activation_wrapper(): + delay = 1.0 + if hasattr(source_neuron, "axon_delay__ms") and source_neuron.axon_delay__ms is not None: + delay = 1.0 + float(source_neuron.axon_delay__ms.magnitude if hasattr(source_neuron.axon_delay__ms, 'magnitude') else source_neuron.axon_delay__ms) + return muscle_callback(source_neuron.pool__ID, muscle, delay) + + netcon.record(muscle_activation_wrapper) + + # Setup spike recording + if neuron_id is not None and id_vector is not None and spike_vector is not None: + netcon.record(spike_vector, id_vector, neuron_id) + + # Add axonal delay if present + if hasattr(source_neuron, "axon_delay__ms") and source_neuron.axon_delay__ms is not None: + axon_delay = float(source_neuron.axon_delay__ms.magnitude if hasattr(source_neuron.axon_delay__ms, 'magnitude') else source_neuron.axon_delay__ms) + netcon.delay = float(DEFAULT_SYNAPTIC_DELAY.magnitude) + axon_delay + + return netcon + + +def _connect_population_to_population( + source_pop: str, + target_pop: str, + populations: dict, + connection_probability: float, + deterministic: bool = False, + inhibitory: bool = False, + **kwargs, +) -> list: + """Create connections between two neural populations.""" + connections = [] + target_neurons = populations[target_pop] + + if deterministic and connection_probability < 1.0: + # Deterministic: each source connects to exact number of targets + n_connections = int(connection_probability * len(target_neurons)) + + for source_neuron in populations[source_pop]: + selected_targets = myogen.RANDOM_GENERATOR.choice( + target_neurons, size=n_connections, replace=False + ) + + for target_neuron in selected_targets: + target_synapse, syn_idx = _select_synapse(target_neuron, inhibitory=inhibitory) + netcon = _create_netcon( + source_neuron, + target_neuron, + muscle_callback=kwargs.get("muscle_callback"), + id_vector=kwargs.get("id_vector"), + spike_vector=kwargs.get("spike_vector"), + neuron_id=source_neuron.global__ID, + synapse_index=syn_idx, + ) + + # Apply custom parameters + if kwargs.get("synaptic_weight") is not None: + netcon.weight[0] = kwargs.get("synaptic_weight") + if kwargs.get("spike_threshold") is not None: + netcon.threshold = kwargs.get("spike_threshold") + + connections.append(netcon) + else: + # Probabilistic: each pair has probability of connecting + for source_neuron in populations[source_pop]: + for target_neuron in target_neurons: + if myogen.RANDOM_GENERATOR.uniform() < connection_probability: + target_synapse, syn_idx = _select_synapse(target_neuron, inhibitory=inhibitory) + netcon = _create_netcon( + source_neuron, + target_neuron, + muscle_callback=kwargs.get("muscle_callback"), + id_vector=kwargs.get("id_vector"), + spike_vector=kwargs.get("spike_vector"), + neuron_id=source_neuron.global__ID, + synapse_index=syn_idx, + ) + + # Apply custom parameters + if kwargs.get("synaptic_weight") is not None: + netcon.weight[0] = kwargs.get("synaptic_weight") + if kwargs.get("spike_threshold") is not None: + netcon.threshold = kwargs.get("spike_threshold") + + connections.append(netcon) + + return connections + + +def _connect_population_to_external(source_pop: str, populations: dict, **kwargs) -> list: + """Create connections from neural population to external target (muscle).""" + connections = [] + for source_neuron in populations[source_pop]: + netcon = _create_netcon( + source_neuron, + None, # External target + muscle_callback=kwargs.get("muscle_callback"), + id_vector=kwargs.get("id_vector"), + muscle=kwargs.get("muscle"), + spike_vector=kwargs.get("spike_vector"), + neuron_id=source_neuron.global__ID, + ) + + if kwargs.get("spike_threshold") is not None: + netcon.threshold = kwargs.get("spike_threshold") + + connections.append(netcon) + + return connections + + +def _connect_external_to_population(target_pop: Optional[str], populations: dict, **kwargs) -> list: + """Create connections from external source to neural population.""" + connections = [] + for target_neuron in populations[target_pop]: + netcon = _create_netcon( + None, # External source + target_neuron, + id_vector=kwargs.get("id_vector"), + spike_vector=kwargs.get("spike_vector"), + neuron_id=None, + ) + connections.append(netcon) + + return connections + + +def _connect_one_to_one( + source_pop: str, + target_pop: str, + populations: dict, + connection_probability: float = 1.0, + inhibitory: bool = False, + **kwargs, +) -> list: + """Create one-to-one connections between populations.""" + source_neurons = populations[source_pop] + target_neurons = populations[target_pop] + + if len(source_neurons) != len(target_neurons): + raise ValueError( + f"One-to-one connection requires equal population sizes. " + f"Source '{source_pop}' has {len(source_neurons)} neurons, " + f"target '{target_pop}' has {len(target_neurons)} neurons." + ) + + if not 0.0 <= connection_probability <= 1.0: + raise ValueError( + f"Connection probability must be between 0.0 and 1.0, got {connection_probability}" + ) + + connections = [] + for source_neuron, target_neuron in zip(source_neurons, target_neurons): + if myogen.RANDOM_GENERATOR.uniform() < connection_probability: + target_synapse, syn_idx = _select_synapse(target_neuron, inhibitory=inhibitory) + netcon = _create_netcon( + source_neuron, + target_neuron, + muscle_callback=kwargs.get("muscle_callback"), + id_vector=kwargs.get("id_vector"), + spike_vector=kwargs.get("spike_vector"), + neuron_id=source_neuron.global__ID, + synapse_index=syn_idx, + ) + + # Apply custom parameters + if kwargs.get("synaptic_weight") is not None: + netcon.weight[0] = kwargs.get("synaptic_weight") + if kwargs.get("spike_threshold") is not None: + netcon.threshold = kwargs.get("spike_threshold") + + connections.append(netcon) + + return connections + + +def _connect_populations( + populations: dict, + source_pop: Optional[str], + target_pop: Optional[str], + connection_probability: float, + muscle_callback: Optional[Callable] = None, + id_vector=None, + spike_vector=None, + muscle=None, + synaptic_weight: Optional[float] = None, + spike_threshold: Optional[float] = None, + deterministic: bool = False, + inhibitory: bool = False, +) -> list: + """ + Create network connections between populations. + + Maintains full API compatibility with NEURON version while using + Jaxley connection infrastructure. + """ + # Validation + if not 0.0 <= connection_probability <= 1.0: + raise ValueError( + f"Connection probability must be between 0.0 and 1.0, got {connection_probability}" + ) + + if source_pop is not None and source_pop not in populations: + raise ValueError(f"Source population '{source_pop}' not found") + + if target_pop is not None and target_pop not in populations: + raise ValueError(f"Target population '{target_pop}' not found") + + # Connection kwargs + connection_kwargs = { + "muscle_callback": muscle_callback, + "id_vector": id_vector, + "spike_vector": spike_vector, + "muscle": muscle, + "synaptic_weight": synaptic_weight, + "spike_threshold": spike_threshold, + "deterministic": deterministic, + "inhibitory": inhibitory, + } + + # Route to appropriate connection function + if source_pop is not None and target_pop is not None: + return _connect_population_to_population( + source_pop, target_pop, populations, connection_probability, **connection_kwargs + ) + elif source_pop is not None and target_pop is None: + return _connect_population_to_external(source_pop, populations, **connection_kwargs) + else: + return _connect_external_to_population(target_pop, populations, **connection_kwargs) + + +def _print_network_connections(network_connections): + """Print human-readable summary of network connections.""" + connection_index = 0 + for connection_group in network_connections.values(): + for netcon in connection_group: + source = netcon.pre() + target = netcon.postcell() + + source_name = EXTERNAL_INPUT_LABEL if source is None else str(source) + target_name = EXTERNAL_TARGET_LABEL if target is None else str(target) + + print(f"NC[{connection_index}]: {source_name} -> {target_name}") + connection_index += 1 + + +@beartowertype +class Network: + """ + Neural network builder for Jaxley backend. + + Provides identical API to NEURON version while using Jaxley's + connection infrastructure internally. + """ + + def __init__(self, populations: dict, spike_recording: Optional[dict] = None): + """ + Initialize network with neural populations. + + Parameters + ---------- + populations : dict + Dictionary mapping population names to neuron lists or Pool objects. + spike_recording : dict, optional + Spike recording configuration with 'idvec' and 'spkvec'. + """ + self.populations = populations + self.connections = [] + self._netcons_by_connection = {} + self.spike_recording = spike_recording + + def setup_spike_recording(self): + """Setup spike recording for all neurons in all populations.""" + if not self.spike_recording: + return + + for pop_name, population in self.populations.items(): + if not hasattr(population, "__iter__") or isinstance(population, dict): + continue + + id_vector = self.spike_recording.get("idvec", {}).get(pop_name) + spike_vector = self.spike_recording.get("spkvec", {}).get(pop_name) + + if id_vector is not None and spike_vector is not None: + recording_netcons = [] + for neuron in population: + if not hasattr(neuron, "cell") and not hasattr(neuron, "ns"): + continue + + # Create recording-only connection + nc = JaxleyConnection(neuron, None) + nc.threshold = float(DEFAULT_SPIKE_THRESHOLD.magnitude) + nc.record(spike_vector, id_vector, neuron.global__ID) + recording_netcons.append(nc) + + connection_key = (pop_name, "spike_recording") + self._netcons_by_connection[connection_key] = recording_netcons + + @beartowertype + def connect( + self, + source: str, + target: str, + probability: float = 1.0, + weight__uS: Quantity__uS = DEFAULT_SYNAPTIC_WEIGHT, + delay__ms: Quantity__ms = DEFAULT_SYNAPTIC_DELAY, + threshold__mV: Quantity__mV = DEFAULT_SPIKE_THRESHOLD, + deterministic: bool = False, + inhibitory: bool = False, + ) -> list: + """ + Connect two neural populations. + + Parameters + ---------- + source : str + Source population name + target : str + Target population name + probability : float + Connection probability (0.0-1.0) + weight__uS : Quantity__uS + Synaptic weight in microsiemens + delay__ms : Quantity__ms + Synaptic delay in milliseconds + threshold__mV : Quantity__mV + Spike threshold in millivolts + deterministic : bool + Use deterministic connectivity + inhibitory : bool + Connect to inhibitory synapses + + Returns + ------- + list + List of created connections + """ + # Validation + if source not in self.populations: + raise ValueError(f"Source population '{source}' not found") + if target not in self.populations: + raise ValueError(f"Target population '{target}' not found") + if not 0.0 <= probability <= 1.0: + raise ValueError(f"Probability must be 0.0-1.0, got {probability}") + + # Extract values + weight_value = getattr(weight__uS, 'magnitude', weight__uS) + threshold_value = getattr(threshold__mV, 'magnitude', threshold__mV) + delay_value = getattr(delay__ms, 'magnitude', delay__ms) + + # Get spike recording vectors + id_vector = None + spike_vector = None + if self.spike_recording: + id_vector = self.spike_recording.get("idvec", {}).get(source) + spike_vector = self.spike_recording.get("spkvec", {}).get(source) + + # Create connections + netcons = _connect_populations( + populations=self.populations, + source_pop=source, + target_pop=target, + connection_probability=probability, + synaptic_weight=weight_value, + spike_threshold=threshold_value, + id_vector=id_vector, + spike_vector=spike_vector, + deterministic=deterministic, + inhibitory=inhibitory, + ) + + self.connections.append({ + "type": "neural", + "source": source, + "target": target, + "probability": probability, + "weight__uS": weight_value, + "delay__ms": delay_value, + "threshold__mV": threshold_value, + "inhibitory": inhibitory, + }) + self._netcons_by_connection[(source, target)] = netcons + + return netcons + + @beartowertype + def connect_to_muscle( + self, + source: str, + muscle, + activation_callback: Callable, + weight__uS: Quantity__uS = DEFAULT_SYNAPTIC_WEIGHT, + threshold__mV: Quantity__mV = DEFAULT_SPIKE_THRESHOLD, + ) -> list: + """ + Connect neural population to muscle with activation callback. + + Parameters + ---------- + source : str + Motor neuron population name + muscle : object + Muscle object + activation_callback : Callable + Callback for motor neuron spikes + weight__uS : Quantity__uS + Synaptic weight + threshold__mV : Quantity__mV + Spike threshold + + Returns + ------- + list + List of muscle connections + """ + if source not in self.populations: + raise ValueError(f"Source population '{source}' not found") + + # Extract values + weight_value = getattr(weight__uS, 'magnitude', weight__uS) + threshold_value = getattr(threshold__mV, 'magnitude', threshold__mV) + + # Get spike recording + id_vector = None + spike_vector = None + if self.spike_recording: + id_vector = self.spike_recording.get("idvec", {}).get(source) + spike_vector = self.spike_recording.get("spkvec", {}).get(source) + + # Create connections + netcons = _connect_populations( + populations=self.populations, + source_pop=source, + target_pop=None, + connection_probability=1.0, + muscle_callback=activation_callback, + muscle=muscle, + synaptic_weight=weight_value, + spike_threshold=threshold_value, + id_vector=id_vector, + spike_vector=spike_vector, + ) + + self.connections.append({ + "type": "muscle", + "source": source, + "target": "muscle", + "muscle": muscle, + "callback": activation_callback, + "weight__uS": weight_value, + "threshold__mV": threshold_value, + }) + self._netcons_by_connection[(source, "muscle")] = netcons + + return netcons + + @beartowertype + def connect_from_external( + self, + source: str, + target: str, + weight__uS: Quantity__uS = DEFAULT_SYNAPTIC_WEIGHT, + delay__ms: Quantity__ms = DEFAULT_SYNAPTIC_DELAY, + threshold__mV: Quantity__mV = DEFAULT_SPIKE_THRESHOLD, + ) -> list: + """ + Connect external input to neural population. + + Parameters + ---------- + source : str + External source label + target : str + Target population name + weight__uS : Quantity__uS + Synaptic weight + delay__ms : Quantity__ms + Synaptic delay + threshold__mV : Quantity__mV + Spike threshold + + Returns + ------- + list + List of external connections + """ + if target not in self.populations: + raise ValueError(f"Target population '{target}' not found") + + # Extract values + weight_value = getattr(weight__uS, 'magnitude', weight__uS) + delay_value = getattr(delay__ms, 'magnitude', delay__ms) + threshold_value = getattr(threshold__mV, 'magnitude', threshold__mV) + + # Create external connections + netcons = [] + target_neurons = self.populations[target] + + for target_neuron in target_neurons: + nc = JaxleyConnection(None, target_neuron) + nc.weight[0] = weight_value + nc.delay = delay_value + nc.threshold = threshold_value + netcons.append(nc) + + self.connections.append({ + "type": "external", + "source": source, + "target": target, + "weight__uS": weight_value, + "delay__ms": delay_value, + "threshold__mV": threshold_value, + }) + self._netcons_by_connection[(source, target)] = netcons + + return netcons + + @beartowertype + def connect_one_to_one( + self, + source: str, + target: str, + probability: float = 1.0, + weight__uS: Quantity__uS = DEFAULT_SYNAPTIC_WEIGHT, + delay__ms: Quantity__ms = DEFAULT_SYNAPTIC_DELAY, + threshold__mV: Quantity__mV = DEFAULT_SPIKE_THRESHOLD, + inhibitory: bool = False, + ) -> list: + """ + Connect populations with one-to-one mapping. + + Parameters + ---------- + source : str + Source population name + target : str + Target population name + probability : float + Connection probability for each pair + weight__uS : Quantity__uS + Synaptic weight + delay__ms : Quantity__ms + Synaptic delay + threshold__mV : Quantity__mV + Spike threshold + inhibitory : bool + Connect to inhibitory synapses + + Returns + ------- + list + List of one-to-one connections + """ + # Validation + if source not in self.populations: + raise ValueError(f"Source population '{source}' not found") + if target not in self.populations: + raise ValueError(f"Target population '{target}' not found") + if not 0.0 <= probability <= 1.0: + raise ValueError(f"Probability must be 0.0-1.0, got {probability}") + + # Extract values + weight_value = getattr(weight__uS, 'magnitude', weight__uS) + delay_value = getattr(delay__ms, 'magnitude', delay__ms) + threshold_value = getattr(threshold__mV, 'magnitude', threshold__mV) + + # Get spike recording + id_vector = None + spike_vector = None + if self.spike_recording: + id_vector = self.spike_recording.get("idvec", {}).get(source) + spike_vector = self.spike_recording.get("spkvec", {}).get(source) + + # Create connections + netcons = _connect_one_to_one( + source_pop=source, + target_pop=target, + populations=self.populations, + connection_probability=probability, + synaptic_weight=weight_value, + spike_threshold=threshold_value, + id_vector=id_vector, + spike_vector=spike_vector, + inhibitory=inhibitory, + ) + + self.connections.append({ + "type": "one_to_one", + "source": source, + "target": target, + "probability": probability, + "weight__uS": weight_value, + "delay__ms": delay_value, + "threshold__mV": threshold_value, + "inhibitory": inhibitory, + }) + self._netcons_by_connection[(source, target)] = netcons + + return netcons + + def get_connections(self) -> list[dict]: + """Get list of all connection specifications.""" + return self.connections.copy() + + def get_netcons(self, source: Optional[str] = None, target: Optional[str] = None) -> list: + """ + Get connection objects with optional filtering. + + Parameters + ---------- + source : str, optional + Filter by source population + target : str, optional + Filter by target population + + Returns + ------- + list + List of matching connections + """ + if source is None and target is None: + # Return all connections + all_netcons = [] + for netcon_list in self._netcons_by_connection.values(): + all_netcons.extend(netcon_list) + return all_netcons + + # Filter by source and/or target + matching_netcons = [] + for (src, tgt), netcon_list in self._netcons_by_connection.items(): + if (source is None or src == source) and (target is None or tgt == target): + matching_netcons.extend(netcon_list) + + return matching_netcons + + def print_summary(self): + """Print network summary.""" + print(f"Network with {len(self.populations)} populations:") + for pop_name, pop in self.populations.items(): + n_neurons = len(pop) if hasattr(pop, '__len__') else '?' + print(f" {pop_name}: {n_neurons} neurons") + + print(f"\n{len(self.connections)} connection groups:") + for conn in self.connections: + conn_type = conn.get('type', 'unknown') + source = conn.get('source', '?') + target = conn.get('target', '?') + + if conn_type == 'neural': + prob = conn.get('probability', 1.0) + weight = conn.get('weight__uS', 0) + inh = " (inhibitory)" if conn.get('inhibitory', False) else "" + print(f" {source} -> {target}: p={prob:.2f}, w={weight:.2f}uS{inh}") + elif conn_type == 'muscle': + print(f" {source} -> muscle") + elif conn_type == 'external': + print(f" {source} (external) -> {target}") + elif conn_type == 'one_to_one': + prob = conn.get('probability', 1.0) + print(f" {source} -> {target}: one-to-one (p={prob:.2f})") diff --git a/myogen/simulator/jaxley/populations/__init__.py b/myogen/simulator/jaxley/populations/__init__.py new file mode 100644 index 00000000..70d84d6e --- /dev/null +++ b/myogen/simulator/jaxley/populations/__init__.py @@ -0,0 +1,67 @@ +""" +Neuron population management for MyoGen spinal circuit simulations - Jaxley Backend. + +This module provides population classes for different types of neurons used in +spinal motor circuit modeling, including motor neurons, interneurons, and afferents. + +Classes +------- +_Pool + Base class for all neuron populations providing common functionality. + +DescendingDrive__Pool + Population of descending drive neurons generating cortical input via Poisson or Gamma processes. + +AffIa__Pool + Population of type Ia afferent neurons providing primary proprioceptive feedback + from muscle spindles. + +AffII__Pool + Population of type II afferent neurons providing secondary proprioceptive feedback + from muscle spindles. + +AffIb__Pool + Population of type Ib afferent neurons providing feedback from Golgi tendon organs. + +GII__Pool + Population of group II interneurons processing type II afferent input in spinal circuits. + +GIb__Pool + Population of group Ib interneurons processing type Ib afferent input in spinal circuits. + +AlphaMN__Pool + Population of alpha motor neurons forming the final common pathway for motor control. + +Usage +----- +>>> from myogen.simulator.jaxley.populations import AlphaMN__Pool, DescendingDrive__Pool +>>> +>>> # Create motor neuron population +>>> motor_pool = AlphaMN__Pool(n=10, model="Powers2017", mode="active") +>>> +>>> # Create descending drive population with Poisson process (default) +>>> drive_pool = DescendingDrive__Pool(n=5, poisson_batch_size=16, +... timestep__ms=0.05) +>>> +>>> # Create descending drive with Gamma process for more regular firing +>>> drive_gamma = DescendingDrive__Pool(n=5, timestep__ms=0.05, process_type='gamma', +... shape=3.0) +""" + +from .afferents import AffIa__Pool, AffII__Pool, AffIb__Pool +from .base import _Pool +from .descending_drive import DescendingDrive__Pool, DescendingDrive_Gamma__Pool +from .interneurons import GII__Pool, GIb__Pool +from .motor_neurons import AlphaMN__Pool + +__all__ = [ + "_Pool", + "DescendingDrive__Pool", + "DescendingDrive_Gamma__Pool", + "AffIa__Pool", + "AffII__Pool", + "AffIb__Pool", + "GII__Pool", + "GIb__Pool", + "AlphaMN__Pool", +] diff --git a/myogen/simulator/jaxley/populations/afferents.py b/myogen/simulator/jaxley/populations/afferents.py new file mode 100644 index 00000000..19166975 --- /dev/null +++ b/myogen/simulator/jaxley/populations/afferents.py @@ -0,0 +1,210 @@ +""" +Afferent neuron populations for proprioceptive feedback - Jaxley Backend. + +This module contains population classes for different types of afferent neurons +that provide sensory feedback from muscle spindles and Golgi tendon organs. +""" + +import numpy as np +import quantities as pq + +from myogen.simulator.jaxley import cells +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__m_per_s, Quantity__mm, Quantity__ms, Quantity__m + +from .base import _Pool + + +@beartowertype +class AffIa__Pool(_Pool): + """ + Container for a population of afferent Ia neurons - Jaxley Backend. + + Manages a collection of AffIa (type Ia afferent) cells that provide + proprioceptive feedback from muscle spindles to spinal circuits. + + Parameters + ---------- + n : int + Number of type Ia afferent neurons to create. + recruitment_thresholds : tuple[float, float] + Min and max recruitment thresholds (Hz). + axon_velocities__m_per_s : tuple[Quantity__m_per_s, Quantity__m_per_s] + Min and max axon conduction velocities (m/s). + axon_length__m : Quantity__m + Length of the axon (m). + poisson_batch_size : int + Batch size for exponential threshold generation algorithm. + timestep__ms : Quantity__ms + Time step for simulation (ms). + init_order : int + Initial order parameter for afferent initialization. + """ + + def __init__( + self, + n: int, + timestep__ms: Quantity__ms, + recruitment_thresholds: tuple[float, float] = (0, 40), + axon_velocities__m_per_s: tuple[Quantity__m_per_s, Quantity__m_per_s] = ( + 61 * pq.m / pq.s, + 75 * pq.m / pq.s, + ), + axon_length__m: Quantity__m = 0.6 * pq.m, + poisson_batch_size: int = 145, # Shape param for Gamma process: CV = 1/sqrt(145) = 8.3% + init_order: int = 0, + ): + self.n = n + self.recruitment_thresholds = recruitment_thresholds + self.axon_velocities = axon_velocities__m_per_s + self.axon_length = axon_length__m + self.poisson_batch_size = poisson_batch_size + self.timestep__ms = timestep__ms + self.init_order = init_order + + rt = np.linspace(*recruitment_thresholds, n) + vcon = np.linspace(*axon_velocities__m_per_s, n) + + _cells = [] + for i, (rt_i, vcon_i) in enumerate(zip(rt, vcon)): + ia = cells.AffIa( + RT=rt_i, + N=poisson_batch_size, + timestep__ms=timestep__ms, + initN=init_order, + pool__ID=i, + ) + ia.create_axon(length__m=axon_length__m, conduction_velocity__m_per_s=vcon_i) + _cells.append(ia) + + super().__init__(cells=_cells) + + +@beartowertype +class AffII__Pool(_Pool): + """ + Container for a population of afferent II neurons - Jaxley Backend. + + Manages a collection of AffII (type II afferent) cells that provide + secondary proprioceptive feedback from muscle spindles to spinal circuits. + + Parameters + ---------- + n : int + Number of type II afferent neurons to create. + recruitment_thresholds : tuple[float, float] + Min and max recruitment thresholds (Hz). + axon_velocities__m_per_s : tuple[Quantity__m_per_s, Quantity__m_per_s] + Min and max axon conduction velocities (m/s). + axon_length__m : Quantity__m + Length of the axon (m). + poisson_batch_size : int + Batch size for exponential threshold generation algorithm. + timestep__ms : Quantity__ms + Time step for simulation (ms). + init_order : int + Initial order parameter for afferent initialization. + """ + + def __init__( + self, + n: int, + timestep__ms: Quantity__ms, + recruitment_thresholds: tuple[float, float] = (0, 40), + axon_velocities__m_per_s: tuple[Quantity__m_per_s, Quantity__m_per_s] = ( + 30 * pq.m / pq.s, + 50 * pq.m / pq.s, + ), + axon_length__m: Quantity__m = 0.6 * pq.m, + poisson_batch_size: int = 772, # Shape param for Gamma process: CV = 1/sqrt(772) = 3.6% + init_order: int = 0, + ): + self.n = n + self.recruitment_thresholds = recruitment_thresholds + self.axon_velocities = axon_velocities__m_per_s + self.axon_length = axon_length__m + self.poisson_batch_size = poisson_batch_size + self.timestep__ms = timestep__ms + self.init_order = init_order + + rt = np.linspace(*recruitment_thresholds, n) + vcon = np.linspace(*axon_velocities__m_per_s, n) + + _cells = [] + for i, (rt_i, vcon_i) in enumerate(zip(rt, vcon)): + ii = cells.AffII( + RT=rt_i, + N=poisson_batch_size, + timestep__ms=timestep__ms, + initN=init_order, + pool__ID=i, + ) + ii.create_axon(length__m=axon_length__m, conduction_velocity__m_per_s=vcon_i) + _cells.append(ii) + + super().__init__(cells=_cells) + + +@beartowertype +class AffIb__Pool(_Pool): + """ + Container for a population of afferent Ib neurons - Jaxley Backend. + + Manages a collection of AffIb (type Ib afferent) cells that provide + feedback from Golgi tendon organs to spinal circuits. + + Parameters + ---------- + n : int + Number of type Ib afferent neurons to create. + recruitment_thresholds : tuple[float, float] + Min and max recruitment thresholds (Hz). + axon_velocities__m_per_s : tuple[Quantity__m_per_s, Quantity__m_per_s] + Min and max axon conduction velocities (m/s). + axon_length__m : Quantity__m + Length of the axon (m). + poisson_batch_size : int + Batch size for exponential threshold generation algorithm. + timestep__ms : Quantity__ms + Time step for simulation (ms). + init_order : int + Initial order parameter for afferent initialization. + """ + + def __init__( + self, + n: int, + timestep__ms: Quantity__ms, + recruitment_thresholds: tuple[float, float] = (0, 40), + axon_velocities__m_per_s: tuple[Quantity__m_per_s, Quantity__m_per_s] = ( + 61 * pq.m / pq.s, + 75 * pq.m / pq.s, + ), + axon_length__m: Quantity__m = 0.6 * pq.m, + poisson_batch_size: int = 145, + init_order: int = 0, + ): + self.n = n + self.recruitment_thresholds = recruitment_thresholds + self.axon_velocities = axon_velocities__m_per_s + self.axon_length = axon_length__m + self.poisson_batch_size = poisson_batch_size + self.timestep__ms = timestep__ms + self.init_order = init_order + + rt = np.linspace(*recruitment_thresholds, n) + vcon = np.linspace(*axon_velocities__m_per_s, n) + + _cells = [] + for i, (rt_i, vcon_i) in enumerate(zip(rt, vcon)): + ib = cells.AffIb( + RT=rt_i, + N=poisson_batch_size, + timestep__ms=timestep__ms, + initN=init_order, + pool__ID=i, + ) + ib.create_axon(length__m=axon_length__m, conduction_velocity__m_per_s=vcon_i) + _cells.append(ib) + + super().__init__(cells=_cells) diff --git a/myogen/simulator/jaxley/populations/base.py b/myogen/simulator/jaxley/populations/base.py new file mode 100644 index 00000000..ccdfe168 --- /dev/null +++ b/myogen/simulator/jaxley/populations/base.py @@ -0,0 +1,212 @@ +""" +Base classes and utility functions for neuron population management - Jaxley Backend + +This module provides the foundational `_Pool` class and helper functions used +by all neuron population classes in the MyoGen Jaxley simulator. + +Naming Convention +----------------- +MyoGen uses a double-underscore (`__`) separator in neuron pool class names to improve +readability and maintain consistency. The pattern is: + + {NeuronType}__{Container} + +Examples: + - ``AlphaMN__Pool`` - Pool of alpha motor neurons + - ``AffIa__Pool`` - Pool of type Ia afferent neurons + - ``GII__Pool`` - Pool of group II interneurons + - ``DescendingDrive__Pool`` - Pool of descending drive neurons + +The double underscore clearly separates the neuron type from the container class, +making the code more readable than alternatives like ``AlphaMNPool`` or ``AlphaMN_Pool``. +This convention is used throughout the neuron populations module. + +Note: The base class uses a single underscore prefix (``_Pool``) following Python's +convention for internal/private classes not intended for direct instantiation by users. +""" + +from typing import Optional, Union + +import numpy as np +from scipy.optimize import curve_fit + +from myogen.utils.decorators import beartowertype + + +def _exp_crescent(x, a, b, c): + return a * np.exp(b * x) + c + + +def _exp_decrescent(x, a, b, c): + return a * np.exp(-b * x) + c + + +@beartowertype +def _exp_interp(first: float, last: float, n: int, curv: float = 0.33, negative: bool = False): + assert curv <= 0.5 + c1 = first <= last + if negative: + c1 = last <= first + x = [0, 2, 4, 4] # This is a hack to hide curve_fit warnings [covariance] + xp = np.linspace(0, 4, n) + if c1: + yn = np.array([first, first + (last - first) * curv, last, last]) / first + popt, _ = curve_fit(_exp_crescent, x, yn) + param = _exp_crescent(xp, *popt) * first + else: + yn = np.array([first, last + (first - last) * curv, last, last]) / last + popt, _ = curve_fit(_exp_decrescent, x, yn) + param = _exp_decrescent(xp, *popt) * last + return param + + +def _get_interneuron_diameter_range__um() -> tuple[float, float]: + """Estimate interneuron soma diameter range based on Biu et al. 2003 [1]_. + + Returns + ------- + tuple[float, float] + Estimated diameter range (min, max) in micrometers. + + References + ---------- + .. [1] Bui, T.V., Cushing, S., Dewey, D., Fyffe, R.E., Rose, P.K., 2003. Comparison of the Morphological and Electrotonic Properties of Renshaw Cells, Ia Inhibitory Interneurons, and Motoneurons in the Cat. Journal of Neurophysiology 90, 2900–2918. https://doi.org/10.1152/jn.00533.2003 + + """ + A_cell = 81390 + 3113 + A_ci = 1.96 * (891.5 + 46.141) / np.sqrt(8) + A = [A_cell - A_ci, A_cell + A_ci] + return np.sqrt(A[0] / np.pi), np.sqrt(A[1] / np.pi) + + +@beartowertype +class _Pool: + """ + Base class for neuron cell populations - Jaxley Backend. + + Provides common functionality for managing groups of neurons. Unlike the NEURON + version, Jaxley doesn't use physical sections, so voltage initialization is + handled through Jaxley's state management. + + Parameters + ---------- + cells : list + List of neuron cells in the population. + initial_voltage__mV : Union[float, list[float]], optional + Initial membrane voltage(s) in millivolts. Can be a single value applied + to all cells or a list of per-cell values. Used for Jaxley state initialization. + If None (default), no voltage initialization is performed. + spike_threshold__mV : float, optional + Spike detection threshold in millivolts for recording spikes from + this population. Motor neurons typically need higher thresholds + (e.g., 50.0 mV) while interneurons use lower thresholds (-10.0 mV). + By default -10.0. + + Notes + ----- + Populations with dummy cells (e.g., DescendingDrive, AffIa, AffIb) should + not provide initial_voltage__mV as they have no real cellular dynamics to + initialize. Populations with real Jaxley cells (e.g., AlphaMN, AffII, GII, GIb) + should provide appropriate voltage values. + + Different neuron types have different action potential amplitudes: + - Motor neurons: typically reach 80-100 mV, need threshold ~50 mV + - Interneurons: typically reach 30-50 mV, can use default -10 mV + """ + + def __init__( + self, + cells: list, + initial_voltage__mV: Optional[Union[float, list[float]]] = None, + spike_threshold__mV: float = -10.0, + ): + self._cells = cells + self.spike_threshold__mV = spike_threshold__mV + + # Handle initial voltage - only if provided + if initial_voltage__mV is not None: + n_cells = len(cells) + if isinstance(initial_voltage__mV, (int, float)): + self.initial_voltage_values__mV = [initial_voltage__mV] * n_cells + else: + assert len(initial_voltage__mV) == n_cells, ( + f"initial_voltage__mV list length ({len(initial_voltage__mV)}) " + f"must match number of cells ({n_cells})" + ) + self.initial_voltage_values__mV = list(initial_voltage__mV) + # Actually apply the voltage to each underlying jx.Cell, so code that + # bypasses get_initialization_data() (e.g. user reaches in via + # ``cell_wrapper.cell`` and calls ``jx.integrate`` directly) still sees + # cells initialised at the correct resting potential. Without this, + # cells stay at Jaxley's default v = -70 mV regardless of what the + # pool was told to use, which silently breaks NERLab cells (V_rest ≈ 0) + # and any other model whose intended rest is not -70. + self._apply_initial_voltages() + else: + self.initial_voltage_values__mV = None + + def _apply_initial_voltages(self) -> None: + """Set v on each underlying ``jx.Cell`` and re-initialise channel states. + + Iterates the wrappers in ``self._cells``; each is expected to have a + ``.cell`` attribute pointing at the underlying ``jx.Cell`` (true for + AlphaMN, interneuron wrappers, etc.). Wrappers without a ``.cell`` + attribute (e.g. dummy / spike-generator populations) are skipped. + """ + for cw, v_init in zip(self._cells, self.initial_voltage_values__mV): + jx_cell = getattr(cw, "cell", None) + if jx_cell is None: + continue + try: + jx_cell.set("v", float(v_init)) + jx_cell.init_states() + except Exception: + # Some wrappers might not yet have their cell fully constructed + # at this point; failing silently is preferable to crashing the + # whole pool build. Code reaching in later can still init by hand. + pass + + def __iter__(self): + """Enable iteration over the cells.""" + return iter(self._cells) + + def __getitem__(self, index): + """Return the cell at the specified index.""" + return self._cells[index] + + def __len__(self): + """Return the number of cells in the population.""" + return len(self._cells) + + def get_initialization_data(self) -> tuple[list, list]: + """ + Return cells and their initial voltages for Jaxley simulation setup. + + For Jaxley backend, this returns cells that need voltage initialization. + Unlike NEURON, we don't work with section objects, but with Jaxley cells. + + Returns + ------- + tuple[list, list] + First list contains Jaxley cell objects. + Second list contains corresponding initial voltages in mV. + Both lists will be empty if population has no voltage initialization. + """ + # Return empty lists if no voltage initialization needed + if self.initial_voltage_values__mV is None: + return [], [] + + cells_to_init = [] + voltages = [] + + for cell_idx, cell in enumerate(self._cells): + cell_voltage = self.initial_voltage_values__mV[cell_idx] + + # Skip dummy cells (they don't have Jaxley cell objects) + if not hasattr(cell, "cell") or cell.cell is None: + continue + + cells_to_init.append(cell) + voltages.append(cell_voltage) + + return cells_to_init, voltages diff --git a/myogen/simulator/jaxley/populations/descending_drive.py b/myogen/simulator/jaxley/populations/descending_drive.py new file mode 100644 index 00000000..74c9acf8 --- /dev/null +++ b/myogen/simulator/jaxley/populations/descending_drive.py @@ -0,0 +1,128 @@ +""" +Descending drive neuron populations for cortical input - Jaxley Backend. + +This module contains the population classes for descending drive neurons that +simulate cortical input to spinal motor circuits via Poisson and Gamma processes. +""" + +from myogen.simulator.jaxley import cells +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__ms + +from .base import _Pool + + +@beartowertype +class DescendingDrive__Pool(_Pool): + """ + Container for a population of descending drive neurons - Jaxley Backend. + + Manages a collection of DD cells that generate spike trains using either + Poisson or Gamma point processes for cortical input to spinal circuits. + + Parameters + ---------- + n : int + Number of descending drive neurons to create. + poisson_batch_size : int, optional + Batch size for exponential threshold generation algorithm (only used when + process_type="poisson"). Higher values improve statistical accuracy but + increase computation. Typical values: 16-50. + Required if process_type="poisson", ignored if process_type="gamma". + timestep__ms : float + Time step for simulation (ms). + process_type : str, optional + Type of point process: "poisson" or "gamma", by default "poisson". + - "poisson": Irregular firing (CV=1.0) + - "gamma": More regular firing with CV controlled by shape parameter + shape : float, optional + Shape parameter for Gamma process (only used when process_type="gamma"), + by default 3.0. Controls spike regularity: + - shape=1: Poisson-like (CV=1.0) + - shape=2-5: Typical cortical neuron regularity (CV=0.45-0.71) + - Higher values: More regular firing (CV=1/sqrt(shape)) + """ + + def __init__( + self, + n: int, + poisson_batch_size: int | None = None, + timestep__ms: Quantity__ms | None = None, + process_type: str = "poisson", + shape: float = 3.0, + ): + if timestep__ms is None: + raise ValueError("timestep__ms is required") + + self.n = n + self.timestep__ms = timestep__ms + self.process_type = process_type + self.shape = shape + + if process_type.lower() == "gamma": + _cells = [ + cells.DD_Gamma( + timestep__ms=timestep__ms, + shape=shape, + pool__ID=i, + ) + for i in range(n) + ] + elif process_type.lower() == "poisson": + if poisson_batch_size is None: + raise ValueError("poisson_batch_size is required when process_type='poisson'") + self.poisson_batch_size = poisson_batch_size + _cells = [cells.DD(N=poisson_batch_size, dt=timestep__ms, pool__ID=i) for i in range(n)] + else: + raise ValueError( + f"Invalid process_type '{process_type}'. Must be 'poisson' or 'gamma'." + ) + + super().__init__(cells=_cells) + + +@beartowertype +class DescendingDrive_Gamma__Pool(_Pool): + """ + Container for a population of descending drive neurons using Gamma process - Jaxley Backend. + + Manages a collection of DD_Gamma cells that generate Gamma-distributed + spike trains for more regular cortical input to spinal circuits, typical + of cortical neuron firing patterns. + + Note: This class is kept for backward compatibility. Consider using + DescendingDrive__Pool with process_type='gamma' instead. + + Parameters + ---------- + n : int + Number of descending drive neurons to create. + timestep__ms : Quantity__ms + Time step for simulation as a Quantity with units of milliseconds. + shape : float, optional + Shape parameter controlling spike regularity, by default 3.0. + - shape=1: Poisson-like (irregular) firing + - shape=2-5: Typical cortical neuron regularity + - Higher values: More regular, clock-like firing + """ + + def __init__( + self, + n: int, + timestep__ms: Quantity__ms, + shape: float = 3.0, + ): + self.n = n + self.timestep__ms = timestep__ms + self.shape = shape + + super().__init__( + cells=[ + cells.DD_Gamma( + timestep__ms=timestep__ms, + shape=shape, + pool__ID=i, + ) + for i in range(n) + ] + ) diff --git a/myogen/simulator/jaxley/populations/interneurons.py b/myogen/simulator/jaxley/populations/interneurons.py new file mode 100644 index 00000000..4c2f1084 --- /dev/null +++ b/myogen/simulator/jaxley/populations/interneurons.py @@ -0,0 +1,329 @@ +""" +Interneuron populations for spinal circuit processing - Jaxley Backend. + +This module contains population classes for different types of interneurons +that provide inhibitory and excitatory connections within spinal circuits. +""" + +from typing import Optional, Union + +import numpy as np +import quantities as pq + +from myogen.simulator.jaxley import cells +from myogen.utils.decorators import beartowertype + +from .base import _get_interneuron_diameter_range__um, _Pool + + +@beartowertype +class GII__Pool(_Pool): + """ + Container for a population of group II interneurons - Jaxley Backend. + + Manages a collection of INgII (group II interneuron) cells that provide + inhibitory feedback in spinal circuits, processing type II afferent input. + + Parameters + ---------- + n : int + Number of group II interneurons to create. + soma_length_range__um : tuple[float, float] + Min and max soma length (um). By default, it is set to the estimated range + for interneurons from Bui et al. 2003 [1]_. + soma_diameter_range : tuple[float, float] + Min and max soma diameter (um). By default, it is set to the estimated range + for interneurons from Bui et al. 2003 [1]_. + passive_conductance_range : tuple[float, float] + Min and max passive membrane conductance (S/cm²). + na3rp_conductance_range : tuple[float, float] + Min and max Na3RP sodium channel conductance (S/cm²). + kdrrl_conductance_range : tuple[float, float] + Min and max KDRRL potassium channel conductance (S/cm²). + mahp_ca_conductance_range : tuple[float, float] + Min and max mAHP calcium conductance (S/cm²). + mahp_k_conductance_range : tuple[float, float] + Min and max mAHP potassium conductance (S/cm²). + mahp_tau_range : tuple[float, float] + Min and max mAHP time constant (ms). + gh_conductance_range : tuple[float, float] + Min and max h-current conductance (S/cm²). + axon_velocities : tuple[float, float] + Min and max axon conduction velocities (m/s). + axon_length : float + Length of the axon (mm). + cell_index : int, optional + Specific cell index to create (creates only one cell), by default None. + initial_voltage__mV : Union[float, list[float]] + Initial membrane voltage (mV), by default -70.0. + + References + ---------- + .. [1] Bui, T.V., Cushing, S., Dewey, D., Fyffe, R.E., Rose, P.K., 2003. + Comparison of the Morphological and Electrotonic Properties of Renshaw Cells, + Ia Inhibitory Interneurons, and Motoneurons in the Cat. + Journal of Neurophysiology 90, 2900–2918. + https://doi.org/10.1152/jn.00533.2003 + """ + + def __init__( + self, + n: int, + soma_length_range__um: tuple[float, float] = _get_interneuron_diameter_range__um(), + soma_diameter_range: tuple[float, float] = _get_interneuron_diameter_range__um(), + passive_conductance_range: tuple[float, float] = (3e-5, 7e-5), + na3rp_conductance_range: tuple[float, float] = (0.003, 0.01), + kdrrl_conductance_range: tuple[float, float] = (0.015, 0.015), + mahp_ca_conductance_range: tuple[float, float] = (3e-6, 3e-6), + mahp_k_conductance_range: tuple[float, float] = (5e-4, 5e-4), + mahp_tau_range: tuple[float, float] = (60, 70), + gh_conductance_range: tuple[float, float] = (2.5e-5, 2.5e-5), + axon_velocities: tuple[float, float] = (10, 10), + axon_length: float = 0.05, + cell_index: Optional[int] = None, + initial_voltage__mV: Union[float, list[float]] = -70.0, + ): + self.n = n + self.soma_length_range__um = soma_length_range__um + self.soma_diameter_range = soma_diameter_range + self.passive_conductance_range = passive_conductance_range + self.na3rp_conductance_range = na3rp_conductance_range + self.kdrrl_conductance_range = kdrrl_conductance_range + self.mahp_ca_conductance_range = mahp_ca_conductance_range + self.mahp_k_conductance_range = mahp_k_conductance_range + self.mahp_tau_range = mahp_tau_range + self.gh_conductance_range = gh_conductance_range + self.axon_velocities = axon_velocities + self.axon_length = axon_length + self.cell_index = cell_index + + sL = np.linspace(*soma_length_range__um, n) + sdiam = np.linspace(*soma_diameter_range, n) + sg_pas = np.linspace(*passive_conductance_range, n) + sgbar_na3rp = np.linspace(*na3rp_conductance_range, n) + sgMax_kdrRL = np.linspace(*kdrrl_conductance_range, n) + sgcamax_mAHP = np.linspace(*mahp_ca_conductance_range, n) + sgkcamax_mAHP = np.linspace(*mahp_k_conductance_range, n) + stau_mAHP = np.linspace(*mahp_tau_range, n) + sghbar_gh = np.linspace(*gh_conductance_range, n) + vcon = np.linspace(*axon_velocities, n) + + if cell_index is not None: + init, end = cell_index, cell_index + 1 + else: + init, end = 0, n + + _cells = [] + for i, ( + sL_i, + sdiam_i, + sg_pas_i, + sgbar_na3rp_i, + sgMax_kdrRL_i, + sgcamax_mAHP_i, + sgkcamax_mAHP_i, + stau_mAHP_i, + sghbar_gh_i, + vcon_i, + ) in enumerate( + zip( + sL[init:end], + sdiam[init:end], + sg_pas[init:end], + sgbar_na3rp[init:end], + sgMax_kdrRL[init:end], + sgcamax_mAHP[init:end], + sgkcamax_mAHP[init:end], + stau_mAHP[init:end], + sghbar_gh[init:end], + vcon[init:end], + ) + ): + # Create INgII cell + gII = cells.INgII(pool__ID=i) + + # Store parameters as attributes for reference + gII.soma_length__um = sL_i + gII.soma_diameter__um = sdiam_i + gII.g_pas__S_per_cm2 = sg_pas_i + gII.gbar_na3rp__S_per_cm2 = sgbar_na3rp_i + gII.gMax_kdrRL__S_per_cm2 = sgMax_kdrRL_i + gII.gcamax_mAHP__S_per_cm2 = sgcamax_mAHP_i + gII.gkcamax_mAHP__S_per_cm2 = sgkcamax_mAHP_i + gII.tau_mAHP__ms = stau_mAHP_i + gII.ghbar_gh__S_per_cm2 = sghbar_gh_i + + # Override per-cell conductances in the Jaxley cell + gII.cell.set("Leak_gLeak", sg_pas_i) + gII.cell.set("Na3rp_gbar", sgbar_na3rp_i) + gII.cell.set("KdrRL_gbar", sgMax_kdrRL_i) + gII.cell.set("MAHP_gcamax", sgcamax_mAHP_i) + gII.cell.set("MAHP_gkcamax", sgkcamax_mAHP_i) + gII.cell.set("MAHP_tau_ca", stau_mAHP_i) + gII.cell.set("Gh_gbar", sghbar_gh_i) + + gII.create_axon( + length__m=axon_length * pq.m, + conduction_velocity__m_per_s=vcon_i * pq.m / pq.s + ) + + _cells.append(gII) + + super().__init__(cells=_cells, initial_voltage__mV=initial_voltage__mV) + + +@beartowertype +class GIb__Pool(_Pool): + """ + Container for a population of group Ib interneurons - Jaxley Backend. + + Manages a collection of INgIb (group Ib interneuron) cells that provide + inhibitory feedback in spinal circuits, processing type Ib afferent input + from Golgi tendon organs. + + Parameters + ---------- + n : int + Number of group Ib interneurons to create. + soma_length_range__um : tuple[float, float] + Min and max soma length (um). By default, it is set to the estimated range + for interneurons from Bui et al. 2003 [1]_. + soma_diameter_range : tuple[float, float] + Min and max soma diameter (um). By default, it is set to the estimated range + for interneurons from Bui et al. 2003 [1]_. + passive_conductance_range : tuple[float, float] + Min and max passive membrane conductance (S/cm²). + na3rp_conductance_range : tuple[float, float] + Min and max Na3RP sodium channel conductance (S/cm²). + kdrrl_conductance_range : tuple[float, float] + Min and max KDRRL potassium channel conductance (S/cm²). + mahp_ca_conductance_range : tuple[float, float] + Min and max mAHP calcium conductance (S/cm²). + mahp_k_conductance_range : tuple[float, float] + Min and max mAHP potassium conductance (S/cm²). + mahp_tau_range : tuple[float, float] + Min and max mAHP time constant (ms). + gh_conductance_range : tuple[float, float] + Min and max h-current conductance (S/cm²). + axon_velocities : tuple[float, float] + Min and max axon conduction velocities (m/s). + axon_length : float + Length of the axon (mm). + cell_index : int, optional + Specific cell index to create (creates only one cell), by default None. + initial_voltage__mV : Union[float, list[float]] + Initial membrane voltage (mV), by default -70.0. + + References + ---------- + .. [1] Bui, T.V., Cushing, S., Dewey, D., Fyffe, R.E., Rose, P.K., 2003. + Comparison of the Morphological and Electrotonic Properties of Renshaw Cells, + Ia Inhibitory Interneurons, and Motoneurons in the Cat. + Journal of Neurophysiology 90, 2900–2918. + https://doi.org/10.1152/jn.00533.2003 + """ + + def __init__( + self, + n: int, + soma_length_range__um: tuple[float, float] = _get_interneuron_diameter_range__um(), + soma_diameter_range: tuple[float, float] = _get_interneuron_diameter_range__um(), + passive_conductance_range: tuple[float, float] = (3e-5, 7e-5), + na3rp_conductance_range: tuple[float, float] = (0.003, 0.01), + kdrrl_conductance_range: tuple[float, float] = (0.015, 0.015), + mahp_ca_conductance_range: tuple[float, float] = (3e-6, 3e-6), + mahp_k_conductance_range: tuple[float, float] = (5e-4, 5e-4), + mahp_tau_range: tuple[float, float] = (60, 70), + gh_conductance_range: tuple[float, float] = (2.5e-5, 2.5e-5), + axon_velocities: tuple[float, float] = (10, 10), + axon_length: float = 0.05, + cell_index: Optional[int] = None, + initial_voltage__mV: Union[float, list[float]] = -70.0, + ): + self.n = n + self.soma_length_range__um = soma_length_range__um + self.soma_diameter_range = soma_diameter_range + self.passive_conductance_range = passive_conductance_range + self.na3rp_conductance_range = na3rp_conductance_range + self.kdrrl_conductance_range = kdrrl_conductance_range + self.mahp_ca_conductance_range = mahp_ca_conductance_range + self.mahp_k_conductance_range = mahp_k_conductance_range + self.mahp_tau_range = mahp_tau_range + self.gh_conductance_range = gh_conductance_range + self.axon_velocities = axon_velocities + self.axon_length = axon_length + self.cell_index = cell_index + + sL = np.linspace(*soma_length_range__um, n) + sdiam = np.linspace(*soma_diameter_range, n) + sg_pas = np.linspace(*passive_conductance_range, n) + sgbar_na3rp = np.linspace(*na3rp_conductance_range, n) + sgMax_kdrRL = np.linspace(*kdrrl_conductance_range, n) + sgcamax_mAHP = np.linspace(*mahp_ca_conductance_range, n) + sgkcamax_mAHP = np.linspace(*mahp_k_conductance_range, n) + stau_mAHP = np.linspace(*mahp_tau_range, n) + sghbar_gh = np.linspace(*gh_conductance_range, n) + vcon = np.linspace(*axon_velocities, n) + + if cell_index is not None: + init, end = cell_index, cell_index + 1 + else: + init, end = 0, n + + _cells = [] + for i, ( + sL_i, + sdiam_i, + sg_pas_i, + sgbar_na3rp_i, + sgMax_kdrRL_i, + sgcamax_mAHP_i, + sgkcamax_mAHP_i, + stau_mAHP_i, + sghbar_gh_i, + vcon_i, + ) in enumerate( + zip( + sL[init:end], + sdiam[init:end], + sg_pas[init:end], + sgbar_na3rp[init:end], + sgMax_kdrRL[init:end], + sgcamax_mAHP[init:end], + sgkcamax_mAHP[init:end], + stau_mAHP[init:end], + sghbar_gh[init:end], + vcon[init:end], + ) + ): + # Create INgIb cell + gIb = cells.INgIb(pool__ID=i) + + # Store parameters as attributes for reference + gIb.soma_length__um = sL_i + gIb.soma_diameter__um = sdiam_i + gIb.g_pas__S_per_cm2 = sg_pas_i + gIb.gbar_na3rp__S_per_cm2 = sgbar_na3rp_i + gIb.gMax_kdrRL__S_per_cm2 = sgMax_kdrRL_i + gIb.gcamax_mAHP__S_per_cm2 = sgcamax_mAHP_i + gIb.gkcamax_mAHP__S_per_cm2 = sgkcamax_mAHP_i + gIb.tau_mAHP__ms = stau_mAHP_i + gIb.ghbar_gh__S_per_cm2 = sghbar_gh_i + + # Override per-cell conductances in the Jaxley cell + gIb.cell.set("Leak_gLeak", sg_pas_i) + gIb.cell.set("Na3rp_gbar", sgbar_na3rp_i) + gIb.cell.set("KdrRL_gbar", sgMax_kdrRL_i) + gIb.cell.set("MAHP_gcamax", sgcamax_mAHP_i) + gIb.cell.set("MAHP_gkcamax", sgkcamax_mAHP_i) + gIb.cell.set("MAHP_tau_ca", stau_mAHP_i) + gIb.cell.set("Gh_gbar", sghbar_gh_i) + + gIb.create_axon( + length__m=axon_length * pq.m, + conduction_velocity__m_per_s=vcon_i * pq.m / pq.s + ) + + _cells.append(gIb) + + super().__init__(cells=_cells, initial_voltage__mV=initial_voltage__mV) diff --git a/myogen/simulator/jaxley/populations/motor_neurons.py b/myogen/simulator/jaxley/populations/motor_neurons.py new file mode 100644 index 00000000..900c1514 --- /dev/null +++ b/myogen/simulator/jaxley/populations/motor_neurons.py @@ -0,0 +1,575 @@ +""" +Alpha motor neuron populations for motor output - Jaxley Backend. + +This module contains the population class for alpha motor neurons, which form +the final common pathway for motor control and drive muscle contraction. + +When use_jaxley_mech=True, neurons use a 2-compartment architecture (soma + 1 dendrite) +with the full Powers2017 channel set: Na3rp, Naps, KdrRL, MAHP, Gh on soma; +LCaInact (dendritic PIC), Gh, Leak on dendrite. The single dendrite uses scaled +geometry (4× surface area via d×2.0, L×2.0) to represent NEURON's 4 dendrites. + +When use_jaxley_mech=False, neurons use single-compartment custom channels. + +The pool implements the Henneman size principle with exponential parameter +distributions (similar to NEURON), creating physiological heterogeneity: +- Most motor neurons are small/low-threshold (recruited first) +- Fewer motor neurons are large/high-threshold (recruited last) +- Parameters scale non-linearly to match experimental observations +""" + +import numpy as np +import quantities as pq + +from myogen.simulator.jaxley import cells +from myogen.utils.decorators import beartowertype +from myogen.utils.types import RECRUITMENT_THRESHOLDS__ARRAY + +from .base import _exp_interp, _Pool + +# Powers2017 motor neuron parameter ranges for soma + 1 dendrite architecture. +# Format: (min, max, curvature) where curvature controls exponential shape +# Curvature 0.3 = moderately exponential, 0.5 = linear +# +# MATCHED TO NEURON Powers2017: Full channel set with proper recruitment heterogeneity. +# Soma: Na3rp, Naps, KdrRL, MAHP, Gh, Leak +# Dendrite: LCaInact (dendritic PIC), Gh, Leak +# +# Tuning note — intentional deviations from NEURON native values: +# These adjustments compensate for the *cable dendrite architecture*, not for the +# Jaxley solver. Jaxley's default solver is implicit backward Euler (bwd_euler); +# a side-by-side test (BE vs Crank-Nicolson, scripts/figure_solver_comparison.py) +# at dt = 25 µs shows the two solvers give identical firing rates and spike shapes +# on this model. The shifts below are required because NEURON-native parameter +# values fire ~3× too fast at low currents and depol-block at ~16 nA when run on +# the cable-dendrite topology — the cable filtering cannot tame NEURON-native PIC +# strength on its own. (NEURON achieves stable PIC behaviour implicitly via its +# discrete 4-isopotential-dendrite topology.) Adjustments: +# +# LCaInact h-gate (inactivation): +# theta_h: NEURON=+14 mV → Jaxley=(-38, -42) mV +# Reason: cable-filtered spikes don't drive dendrite above +14 mV, so NEURON's +# h_inf never inactivates PIC. Shifting theta_h to plateau voltage (~-40 mV) +# enables PIC self-termination without eliminating it during normal firing. +# kappa_h: NEURON=4 mV → Jaxley=3 mV (steeper, more selective inactivation) +# tau_h: NEURON=1500 ms → Jaxley=800 ms (faster inactivation prevents plateau trap) +# +# gLCa: NEURON native × 0.75 (cable-architecture PIC strength tuning — at 1.0× +# the cable arch enters depolarisation block from ~16 nA) +# +# gMAHP_k: NEURON native × 1.125 (slightly elevated to compensate for lower AHP +# in cable arch vs. NEURON's isopotential compartment) +# +# All other parameters (geometry, soma channels, reversal potentials) match NEURON native. +JAXLEY_HH_PARAMS = { + # === SOMA GEOMETRY === + "soma_diameter_range": (22.0, 30.0, 0.3), # µm (NEURON: 22-30) + "soma_length_range": (2952.0, 3665.0, 0.3), # µm (NEURON: 2952-3665) + "soma_capacitance_range": (1.356, 1.879, 0.3), # µF/cm² + "soma_axial_resistivity": 0.001, # Ω·cm (isopotential soma) + + # === SOMA CHANNELS === + "gNa3rp_range": (0.01, 0.022, 0.3), # S/cm² + "gNaPS_range": (2.6e-5, 2.0e-5, 0.3), # S/cm² (NEURON native) + "gKdrRL_range": (0.015, 0.02, 0.3), # S/cm² + "gMAHP_ca_range": (6.4e-6, 1.015e-5, 0.075), # S/cm² (NEURON native) + "gMAHP_k_range": (5.0625e-4, 6.75e-4, 0.3), # S/cm² (1.125× NEURON — 0.75× PIC allows lower MAHP) + "tauMAHP_range": (90.0, 30.0, 0.3), # ms (slower in small MNs) + "gGh_soma_range": (3.0e-5, 2.3e-4, 0.3), # S/cm² + "gLeak_soma_range": (1.5e-4, 3.77e-4, 0.3), # S/cm² (min raised to increase rheobase ≥4 nA) + "eLeak_soma_range": (-71.0, -72.0, 0.5), # mV + + # === DENDRITE GEOMETRY (single cable, ncomp=4, 4× NEURON single-dend area) === + # d×2.0, L×2.0 gives 4× total area matching NEURON's 4 dendrites. + # Cable Ra provides temporal filtering of PIC dynamics. + "dend_diameter_range": (17.46, 23.82, 0.3), # µm (NEURON × 2.0, cable) + "dend_length_range": (3588.0, 4454.0, 0.3), # µm (NEURON × 2.0, total cable length) + "dend_capacitance_range": (0.868, 0.880, 0.3), # µF/cm² (NEURON native) + "dend_Ra_range": (51.04, 40.76, 0.3), # Ω·cm (NEURON native, critical for filtering) + + # === DENDRITE CHANNELS (NEURON native densities on all cable compartments) === + # Same density as NEURON; 4× area gives 4× total PIC conductance. + "gLeak_dend_range": (7.93e-5, 1.75e-4, 0.3), # S/cm² (NEURON native) + "eLeak_dend_range": (-71.0, -72.0, 0.5), # mV + "gGh_dend_range": (3.0e-5, 2.3e-4, 0.3), # S/cm² (NEURON native) + "gLCa_dend_range": (7.21875e-5, 9.69375e-5, 0.3), # S/cm² (0.75× NEURON base — balanced PIC) + "theta_m_LCa_range": (-42.0, -39.0, 0.3), # mV (uniform, cable handles staggering) + "theta_h_LCa_range": (-38.0, -42.0, 0.3), # mV (inactivate PIC at plateau ~-40 mV, steep kappa_h=3 for selectivity) + + # === FIXED REVERSAL POTENTIALS === + "eNa": 55.0, # mV + "eK": -80.0, # mV + "eCa": 120.0, # mV +} + + +# ============================================================================= +# NERLAB MOTOR-NEURON PARAMETERS (matches the production NEURON model) +# ============================================================================= +# Source: config/alpha_mn_default.yaml, section ``nerlab:``. This is the model +# the production NEURON pool uses (``model="NERLab"`` is the YAML default and +# every production example takes that default). +# +# IMPORTANT — encoding convention: +# Each ``_range`` entry is ``(small_cell_value, large_cell_value)``. We +# interpolate LINEARLY across normalised recruitment thresholds: +# value[i] = rt_norm[i] * (large - small) + small +# This mirrors NEURON's ``special_interp`` (see motor_neurons.py L390-396), +# which collapses to the same linear form once you account for its +# ``negative`` flag. The YAML ``[min, max, curve, negative]`` 4-tuples +# FLIP the (small, large) assignment when ``negative=True``, so six entries +# below (gnapbar, gkfbar, gksbar, gcaLbar, vtraub_caL, ltau_caL) are stored +# with the YAML values swapped relative to a naive read. The ``# YAML neg`` +# marker flags each such entry. See: +# myogen/simulator/neuron/populations/motor_neurons.py:390 (special_interp) +# config/alpha_mn_default.yaml:14-90 (raw NERLab parameter ranges) +# +# Architecture: +# - soma : 1 compartment, sphere-like (L = diam) +# - dend : 1 compartment, ISOPOTENTIAL (ncomp=1) ← matches NEURON 1-dend +# - soma channels : napp (Na fast + Na persistent + Kfast + Kslow + leak) +# - dendrite channels: caL (L-type Ca with no inactivation + leak) +# +# Voltage convention: original 1952 Hodgkin-Huxley +# V_rest ≈ 0 mV, ENa = +120 mV, EK = -10 mV, ECa = +140 mV +# Cells using NERLAB_PARAMS must initialise V at 0 mV (not -65 mV) and +# detect spikes at a positive threshold (~+50 mV). +NERLAB_PARAMS = { + # ---- soma geometry (sphere-like: L = diam in NERLab cells.py) ---- + "soma_diameter_range": (78.0, 113.0), # µm + "soma_capacitance": 1.0, # µF/cm² + "soma_axial_resistivity": 70.0, # Ω·cm + + # ---- soma napp channel conductances ---- + "gnabar_range": (0.0325, 0.0775), # S/cm² + "gnapbar_range": (0.00067, 0.00043), # S/cm² YAML neg: small=0.00067, large=0.00043 + "gkfbar_range": (0.0015, 0.0028), # S/cm² YAML neg + "gksbar_range": (0.016, 0.02), # S/cm² YAML neg + "gls_range": (0.000952, 0.001538), # S/cm² (napp leak; 1/1050 .. 1/650) + "rinact_range": (0.018, 0.063), # /ms (r-gate beta) + + # ---- dendrite geometry ---- + "dend_diameter_range": (48.0, 90.0), # µm + "dend_length_range": (5500.0, 10600.0), # µm + "dend_capacitance": 1.0, # µF/cm² + "dend_axial_resistivity": 70.0, # Ω·cm + + # ---- dendrite caL channel ---- + "gcaLbar_range": (6.2e-6, 1.25e-5), # S/cm² YAML neg: small=6.2e-6, large=1.25e-5 + "vtraub_caL_range": (34.0, 35.0), # mV YAML neg + "ltau_caL_range": (47.0, 90.0), # ms YAML neg + "gl_caL_range": (7.69e-5, 1.65e-4), # S/cm² (dendritic leak) + + # ---- fixed napp / caL reversals & offsets (NERLab production values) ---- + "ena": 120.0, # mV + "ek": -10.0, # mV + "el_napp": 0.0, # mV + "vtraub_napp": 0.0, # mV + "ecaL": 140.0, # mV + "el_caL": 0.0, # mV +} + +# Sentinel: NERLab cells live in the 1952-HH frame. When the pool is built with +# model="NERLab" and the user did NOT explicitly override these voltage-frame +# defaults, the pool overrides them to NERLab-appropriate values. +NERLAB_DEFAULT_VHOLD_MV = 0.0 +NERLAB_DEFAULT_SPIKE_THRESHOLD_MV = 50.0 + + +@beartowertype +class AlphaMN__Pool(_Pool): + """ + Container for a population of alpha motor neurons - Jaxley Backend. + + Manages a collection of AlphaMN (alpha motor neuron) cells. These cells form + the final common pathway for motor control. + + Implements the Henneman size principle with exponential parameter distributions: + - Low threshold neurons: small soma, high Rin, low rheobase, recruited first + - High threshold neurons: large soma, low Rin, high rheobase, recruited last + + Parameters are distributed using exponential interpolation (like NEURON version), + creating realistic motor neuron pool heterogeneity. + + Parameters + ---------- + n : int + Number of alpha motor neurons to create. + recruitment_thresholds__array : RECRUITMENT_THRESHOLDS__ARRAY, optional + Array of recruitment thresholds for each motor neuron, by default None. + When provided, parameters are interpolated based on threshold values. + model : str, optional + Motor neuron model type ("NERLab" or "Powers2017"), by default "NERLab". + mode : str, optional + Simulation mode ("active" or "passive"), by default "active". + axon_velocities : tuple[float, float], optional + Min and max axon conduction velocities (m/s), by default (50, 65). + axon_length : float, optional + Length of the axon (mm), by default 0.6. + gamma : float, optional + Neuromodulation level (a.u.), by default 0.2. + cell_index : Optional[int], optional + Specific cell index to create (creates only one cell), by default None. + initial_voltage__mV : float or list[float], optional + Initial membrane voltage (mV), by default -67. + spike_threshold__mV : float, optional + Spike detection threshold for recording motor neuron spikes, by default 0.0. + Uses zero-crossing detection for the 2-compartment Powers2017 architecture. + use_jaxley_mech : bool, optional + If True, use 2-compartment (soma + dendrite) Powers2017 channels. + If False, use single-compartment custom channels. Default is False. + normalize_nerlab_thresholds : bool, optional + Normalize supplied NERLab recruitment thresholds before parameter + interpolation. The backwards-compatible default is True. Set False + when reproducing the production NEURON profile, which consumes the + supplied threshold fractions verbatim. + + Attributes + ---------- + neuron_params : dict + Dictionary of per-neuron parameter arrays computed during pool creation. + Keys include: 'soma_diameter', 'soma_length', 'capacitance', 'gNa', 'gK', + 'gLeak', 'eLeak', etc. Useful for inspection and analysis. + """ + + def __init__( + self, + n: int | None = None, + recruitment_thresholds__array: RECRUITMENT_THRESHOLDS__ARRAY | None = None, + model: str = "NERLab", + mode: str = "active", + axon_velocities: tuple[float, float] = (50, 65), + axon_length: float = 0.6, + gamma: float = 0.2, + cell_index: int | None = None, + initial_voltage__mV: float | list[float] | None = None, + spike_threshold__mV: float | None = None, + use_jaxley_mech: bool = False, + normalize_nerlab_thresholds: bool = True, + ): + self.recruitment_thresholds__array = recruitment_thresholds__array + + if self.recruitment_thresholds__array is not None: + self.n = len(self.recruitment_thresholds__array) + else: + if n is None: + raise ValueError("Either n or recruitment_thresholds__array must be provided.") + self.n = n + + self.model = model + self.mode = mode + self.axon_velocities = axon_velocities + self.axon_length = axon_length + self.gamma = gamma + self.cell_index = cell_index + self.use_jaxley_mech = use_jaxley_mech + self.normalize_nerlab_thresholds = normalize_nerlab_thresholds + + # Voltage-frame defaults depend on the model. NERLab uses the original + # 1952 HH convention (V_rest ≈ 0 mV, ENa = +120, EK = -10); Powers2017 + # uses modern absolute voltages (V_rest ≈ -67 mV). If the caller didn't + # explicitly pass these, pick the right frame automatically. + if initial_voltage__mV is None: + initial_voltage__mV = ( + NERLAB_DEFAULT_VHOLD_MV if self.model == "NERLab" else -67.0 + ) + if spike_threshold__mV is None: + spike_threshold__mV = ( + NERLAB_DEFAULT_SPIKE_THRESHOLD_MV if self.model == "NERLab" else 0.0 + ) + + # Generate per-neuron parameter arrays (like NEURON version) + self.neuron_params = self._generate_parameter_arrays() + + # Create cells with pre-computed parameters + _cells = self._create_cells() + + super().__init__( + cells=_cells, + initial_voltage__mV=initial_voltage__mV, + spike_threshold__mV=spike_threshold__mV, + ) + + def _generate_parameter_arrays(self) -> dict: + """Generate per-neuron parameter arrays. + + Dispatches on ``self.model``: + - ``"NERLab"`` → ``_generate_nerlab_params`` (production NEURON model) + - ``"Powers2017"`` → original Powers2017 cable-architecture params + """ + if self.model == "NERLab": + return self._generate_nerlab_params() + return self._generate_powers2017_params() + + def _generate_nerlab_params(self) -> dict: + """Per-neuron NERLab parameter arrays, linearly interpolated over recruitment + thresholds (same convention as Powers2017 path: low threshold → small index). + + Mirrors the layout of ``NEURON: _create_nerlab_cells`` in + ``myogen/simulator/neuron/populations/motor_neurons.py``. + """ + p = NERLAB_PARAMS + n = self.n + if self.recruitment_thresholds__array is not None: + rt = np.array(self.recruitment_thresholds__array) + # The production NEURON NERLab pool consumes recruitment thresholds + # verbatim. Keep normalization as the backwards-compatible default + # for existing Jaxley studies, but allow profile-card simulations to + # construct the same cells as the NEURON backend. + rt_norm = ( + (rt - rt.min()) / (rt.max() - rt.min() + 1e-10) + if self.normalize_nerlab_thresholds + else rt + ) + def lin(lo, hi): + return rt_norm * (hi - lo) + lo + else: + def lin(lo, hi): + return np.linspace(lo, hi, n) + + params: dict = {} + # Geometry + params["soma_diameter"] = lin(*p["soma_diameter_range"]) + params["dend_diameter"] = lin(*p["dend_diameter_range"]) + params["dend_length"] = lin(*p["dend_length_range"]) + params["soma_capacitance"] = np.full(n, p["soma_capacitance"]) + params["soma_axial_resistivity"] = np.full(n, p["soma_axial_resistivity"]) + params["dend_capacitance"] = np.full(n, p["dend_capacitance"]) + params["dend_axial_resistivity"] = np.full(n, p["dend_axial_resistivity"]) + + # napp (soma) channel conductances + params["gnabar"] = lin(*p["gnabar_range"]) + params["gnapbar"] = lin(*p["gnapbar_range"]) + params["gkfbar"] = lin(*p["gkfbar_range"]) + params["gksbar"] = lin(*p["gksbar_range"]) + params["gls"] = lin(*p["gls_range"]) + params["rinact"] = lin(*p["rinact_range"]) + + # caL (dendrite) channel + params["gcaLbar"] = lin(*p["gcaLbar_range"]) + params["vtraub_caL"] = lin(*p["vtraub_caL_range"]) + params["ltau_caL"] = lin(*p["ltau_caL_range"]) + params["gl_caL"] = lin(*p["gl_caL_range"]) + + # Fixed reversals & offsets + for k in ("ena", "ek", "el_napp", "vtraub_napp", "ecaL", "el_caL"): + params[k] = np.full(n, p[k]) + return params + + def _generate_powers2017_params(self) -> dict: + """(Original) per-neuron Powers2017 cable-architecture parameters.""" + params = {} + + if self.recruitment_thresholds__array is not None: + # Interpolate based on recruitment threshold (like NEURON) + # Normalize thresholds to 0-1 range for interpolation + rt = np.array(self.recruitment_thresholds__array) + rt_norm = (rt - rt.min()) / (rt.max() - rt.min() + 1e-10) + + def interp_by_threshold(min_val, max_val, curv): + """Linear interpolation based on normalized recruitment threshold. + + Matches NEURON's special_interp which uses LINEAR mapping: + param = threshold * (max - min) + min + + The recruitment thresholds are already exponentially distributed, + so linear mapping produces exponential parameter distributions + (matching NEURON behavior exactly). + + Parameters + ---------- + min_val : float + Value for lowest threshold neuron + max_val : float + Value for highest threshold neuron + curv : float + Curvature parameter (unused — kept for API compatibility) + """ + return rt_norm * (max_val - min_val) + min_val + else: + # Use exponential interpolation (like NEURON _exp_interp) + def interp_by_threshold(min_val, max_val, curv): + """Exponential interpolation for n neurons.""" + if curv >= 0.45: + return np.linspace(min_val, max_val, self.n) + return _exp_interp(min_val, max_val, self.n, curv=curv) + + # Generate all parameter arrays + p = JAXLEY_HH_PARAMS + + # Soma geometry + params["soma_diameter"] = interp_by_threshold(*p["soma_diameter_range"]) + params["soma_length"] = interp_by_threshold(*p["soma_length_range"]) + params["capacitance"] = interp_by_threshold(*p["soma_capacitance_range"]) + params["axial_resistivity"] = np.full(self.n, p["soma_axial_resistivity"]) + + # Soma channels + params["gNa3rp"] = interp_by_threshold(*p["gNa3rp_range"]) + params["gNaPS"] = interp_by_threshold(*p["gNaPS_range"]) + params["gKdrRL"] = interp_by_threshold(*p["gKdrRL_range"]) + params["gMAHP_ca"] = interp_by_threshold(*p["gMAHP_ca_range"]) + params["gMAHP_k"] = interp_by_threshold(*p["gMAHP_k_range"]) + params["tauMAHP"] = interp_by_threshold(*p["tauMAHP_range"]) + params["gGh_soma"] = interp_by_threshold(*p["gGh_soma_range"]) + params["gLeak_soma"] = interp_by_threshold(*p["gLeak_soma_range"]) + params["eLeak_soma"] = interp_by_threshold(*p["eLeak_soma_range"]) + + # Backward compat aliases (used in analysis code / get_parameter_summary) + params["gNa"] = params["gNa3rp"] + params["gK"] = params["gKdrRL"] + params["gLeak"] = params["gLeak_soma"] + params["eLeak"] = params["eLeak_soma"] + + # Dendrite geometry + params["dend_diameter"] = interp_by_threshold(*p["dend_diameter_range"]) + params["dend_length"] = interp_by_threshold(*p["dend_length_range"]) + params["dend_cm"] = interp_by_threshold(*p["dend_capacitance_range"]) + params["dend_Ra"] = interp_by_threshold(*p["dend_Ra_range"]) + + # Dendrite channels + params["gLeak_dend"] = interp_by_threshold(*p["gLeak_dend_range"]) + params["eLeak_dend"] = interp_by_threshold(*p["eLeak_dend_range"]) + params["gGh_dend"] = interp_by_threshold(*p["gGh_dend_range"]) + params["gLCa_dend"] = interp_by_threshold(*p["gLCa_dend_range"]) + params["theta_m_LCa"] = interp_by_threshold(*p["theta_m_LCa_range"]) + params["theta_h_LCa"] = interp_by_threshold(*p["theta_h_LCa_range"]) + + # Fixed reversal potentials + params["eNa"] = np.full(self.n, p["eNa"]) + params["eK"] = np.full(self.n, p["eK"]) + params["eCa"] = np.full(self.n, p["eCa"]) + + return params + + def _create_cells(self) -> list: + """Create motor neurons using Jaxley backend. + + Each motor neuron is created with pre-computed parameters from + self.neuron_params, implementing the Henneman size principle: + - Low threshold → small soma → high Rin → low rheobase + - High threshold → large soma → low Rin → high rheobase + """ + vcon = np.linspace(self.axon_velocities[0], self.axon_velocities[1], self.n) + + # Determine cell creation range + if self.cell_index is not None: + init, end = self.cell_index, self.cell_index + 1 + else: + init, end = 0, self.n + + _cells = [] + for i in range(init, end): + # Get recruitment threshold for this neuron (if available) + recruitment_threshold = None + if self.recruitment_thresholds__array is not None: + recruitment_threshold = float(self.recruitment_thresholds__array[i]) + + # NERLab path: build a tight per-cell parameter dict and skip the + # Powers2017-shaped one below. + if self.model == "NERLab": + neuron_params = {k: self.neuron_params[k][i] for k in self.neuron_params} + cell = cells.AlphaMN( + mode=self.mode, + model=self.model, + pool__ID=i, + use_jaxley_mech=self.use_jaxley_mech, # ignored for NERLab + gamma=self.gamma, # ← was missing; controlled caL_gama + recruitment_threshold=recruitment_threshold, + neuron_params=neuron_params, + ) + cell.create_axon( + length__m=self.axon_length * pq.m, + conduction_velocity__m_per_s=vcon[i] * pq.m / pq.s, + ) + _cells.append(cell) + continue + + # Powers2017 path — original code below. + # Extract pre-computed parameters for this neuron + neuron_params = { + # Soma geometry + "soma_diameter": self.neuron_params["soma_diameter"][i], + "soma_length": self.neuron_params["soma_length"][i], + "capacitance": self.neuron_params["capacitance"][i], + "axial_resistivity": self.neuron_params["axial_resistivity"][i], + # Soma channels + "gNa3rp": self.neuron_params["gNa3rp"][i], + "gNaPS": self.neuron_params["gNaPS"][i], + "gKdrRL": self.neuron_params["gKdrRL"][i], + "gMAHP_ca": self.neuron_params["gMAHP_ca"][i], + "gMAHP_k": self.neuron_params["gMAHP_k"][i], + "tauMAHP": self.neuron_params["tauMAHP"][i], + "gGh_soma": self.neuron_params["gGh_soma"][i], + "gLeak_soma": self.neuron_params["gLeak_soma"][i], + "eLeak_soma": self.neuron_params["eLeak_soma"][i], + # Dendrite geometry + "dend_diameter": self.neuron_params["dend_diameter"][i], + "dend_length": self.neuron_params["dend_length"][i], + "dend_cm": self.neuron_params["dend_cm"][i], + "dend_Ra": self.neuron_params["dend_Ra"][i], + # Dendrite channels + "gLeak_dend": self.neuron_params["gLeak_dend"][i], + "eLeak_dend": self.neuron_params["eLeak_dend"][i], + "gGh_dend": self.neuron_params["gGh_dend"][i], + "gLCa_dend": self.neuron_params["gLCa_dend"][i], + "theta_m_LCa": self.neuron_params["theta_m_LCa"][i], + "theta_h_LCa": self.neuron_params["theta_h_LCa"][i], + # Reversal potentials + "eNa": self.neuron_params["eNa"][i], + "eK": self.neuron_params["eK"][i], + "eCa": self.neuron_params["eCa"][i], + # Backward compat aliases + "gNa": self.neuron_params["gNa"][i], + "gK": self.neuron_params["gK"][i], + "gLeak": self.neuron_params["gLeak"][i], + "eLeak": self.neuron_params["eLeak"][i], + } + + cell = cells.AlphaMN( + mode=self.mode, + model=self.model, + pool__ID=i, + use_jaxley_mech=self.use_jaxley_mech, + recruitment_threshold=recruitment_threshold, + neuron_params=neuron_params, # Pass pre-computed parameters + ) + + # Create axon with appropriate delay + cell.create_axon( + length__m=self.axon_length * pq.m, + conduction_velocity__m_per_s=vcon[i] * pq.m / pq.s, + ) + + _cells.append(cell) + + return _cells + + def get_parameter_summary(self) -> str: + """Return a summary of the motor neuron pool parameter distributions. + + Returns + ------- + str + Formatted string showing parameter ranges and statistics. + """ + lines = [ + "=" * 60, + "Motor Neuron Pool Parameter Summary", + "=" * 60, + f"Number of neurons: {self.n}", + f"Model: {self.model}", + f"Mode: {self.mode}", + f"Powers2017 2-compartment (use_jaxley_mech): {self.use_jaxley_mech}", + "", + "Parameter Ranges (min - max):", + "-" * 40, + ] + + for param_name, values in self.neuron_params.items(): + if isinstance(values, np.ndarray) and len(values) > 1: + lines.append(f" {param_name}: {values.min():.4g} - {values.max():.4g}") + else: + lines.append(f" {param_name}: {values[0]:.4g} (fixed)") + + lines.append("=" * 60) + return "\n".join(lines) diff --git a/myogen/simulator/jaxley/proprioception/__init__.py b/myogen/simulator/jaxley/proprioception/__init__.py new file mode 100644 index 00000000..2d672103 --- /dev/null +++ b/myogen/simulator/jaxley/proprioception/__init__.py @@ -0,0 +1,4 @@ +from .golgi import GolgiTendonOrganModel +from .spindle import SpindleModel + +__all__ = ["GolgiTendonOrganModel", "SpindleModel"] diff --git a/myogen/simulator/jaxley/proprioception/golgi.py b/myogen/simulator/jaxley/proprioception/golgi.py new file mode 100644 index 00000000..34832b31 --- /dev/null +++ b/myogen/simulator/jaxley/proprioception/golgi.py @@ -0,0 +1,271 @@ +""" +Golgi Tendon Organ (GTO) Model - Jaxley Backend (Pure Python/NumPy) + +This module provides the GTO model using pure NumPy for the Jaxley backend. +Maintains full API compatibility with the NEURON version. + +Based on Lin & Crago (2002) with implementation by Elias (2014). +""" + +from typing import Any, Dict + +import numpy as np + +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__ms + + +class _GolgiTendonOrgan__NumPy: + """ + Pure NumPy implementation of Golgi Tendon Organ model. + + Implements logarithmic force-to-firing relationship from Lin & Crago (2002). + """ + + def __init__(self, gtoD: Dict[str, Any], tstop__ms: float, dt__ms: float): + self.params = gtoD + self.tstop = tstop__ms + self.dt = dt__ms + + # Initialize storage arrays + self.n_steps = int(tstop__ms / dt__ms) + 1 + self.time = np.linspace(0, tstop__ms, self.n_steps) + + # Output array + self.Ib = np.zeros(self.n_steps) # Ib afferent firing rate + + # State variables + self.step_idx = 0 + + # Digital filter state (simple low-pass) + self.prev_firing = 0.0 + self.filter_alpha = 0.3 # Smoothing factor + + def integrate(self, muscle_force__N: float) -> float: + """ + Integrate GTO model for one time step. + + Parameters + ---------- + muscle_force__N : float + Current muscle force in Newtons + + Returns + ------- + float + Ib afferent firing rate in Hz + """ + # Logarithmic force-to-firing relationship + # firing_rate = G1 * log(force/G2 + 1) + G1 = self.params["G1"] + G2 = self.params["G2"] + + # Ensure force is non-negative + force = max(0, muscle_force__N) + + # Compute instantaneous firing rate + firing_rate = G1 * np.log(force / G2 + 1.0) + + # Apply simple low-pass filter for temporal dynamics + if self.step_idx > 0: + firing_rate = (1 - self.filter_alpha) * self.prev_firing + self.filter_alpha * firing_rate + + self.prev_firing = firing_rate + + # Store firing rate + self.Ib[self.step_idx] = firing_rate + + # Increment step + self.step_idx += 1 + + return firing_rate + + +@beartowertype +class GolgiTendonOrganModel: + """ + API wrapper for the Golgi Tendon Organ (GTO) model - Jaxley Backend. + + This class provides an intuitive interface for creating GTO models + with user-friendly parameter names that are internally mapped to the + correct format expected by the underlying GTO implementation. + + The Golgi Tendon Organ is a proprioceptive sensory organ located at the + muscle-tendon junction that detects muscle force/tension and provides + feedback for motor control and protection against excessive forces. + + The model is based on Lin & Crago (2002) and implements a logarithmic + force-to-firing relationship with digital filtering for realistic + afferent discharge patterns. + + Parameters + ---------- + simulation_time__ms : Quantity__ms + Total simulation time in milliseconds + time_step__ms : Quantity__ms + Integration time step in milliseconds + gto_parameters : Dict[str, Any] + Dictionary containing GTO model parameters + """ + + def __init__( + self, + simulation_time__ms: Quantity__ms, + time_step__ms: Quantity__ms, + gto_parameters: Dict[str, Any], + ): + # Store original parameters (immutable) + self.simulation_time__ms = simulation_time__ms + self.time_step__ms = time_step__ms + self.gto_parameters = gto_parameters.copy() + + # Private working copies for internal use + self._simulation_time__ms = simulation_time__ms + self._time_step__ms = time_step__ms + self._gto_parameters = gto_parameters.copy() + + # Validate inputs + self._validate_parameters() + + # Create the underlying GTO model + self._gto_model = self._create_gto_model() + + def _validate_parameters(self) -> None: + """Validate input parameters.""" + if self._simulation_time__ms <= 0: + raise ValueError("simulation_time__ms must be positive") + + if self._time_step__ms <= 0: + raise ValueError("time_step__ms must be positive") + + if self._simulation_time__ms <= self._time_step__ms: + raise ValueError("simulation_time__ms must be greater than time_step__ms") + + def _create_gto_model(self) -> _GolgiTendonOrgan__NumPy: + """ + Create the underlying GTO model instance. + + This method maps the user-friendly parameter names to the format + expected by the GTO constructor. + """ + return _GolgiTendonOrgan__NumPy( + gtoD=self._gto_parameters, + tstop__ms=self._simulation_time__ms.magnitude, + dt__ms=self._time_step__ms.magnitude, + ) + + def integrate(self, muscle_force__N: float) -> float: + """ + Integrate the GTO model for one time step. + + Parameters + ---------- + muscle_force__N : float + Current muscle force in Newtons + + Returns + ------- + float + Ib afferent firing rate in Hz + """ + return self._gto_model.integrate(muscle_force__N) + + @property + def ib_afferent_firing__Hz(self) -> np.ndarray: + """Get Ib afferent firing rate time series in Hz.""" + return np.asarray(self._gto_model.Ib) + + def __repr__(self) -> str: + """String representation of the GTO model.""" + return ( + f"GolgiTendonOrganModel(t_sim={self.simulation_time__ms}ms, dt={self.time_step__ms}ms)" + ) + + @staticmethod + def create_default_gto_parameters() -> Dict[str, Any]: + """ + Create default Golgi Tendon Organ parameter dictionary. + + The GTO model uses a logarithmic force-to-firing relationship: + firing_rate = G1 * log(force/G2 + 1) + + This is followed by digital filtering to create realistic temporal + dynamics in the afferent discharge pattern. + + Returns + ------- + Dict[str, Any] + Dictionary of GTO parameters with detailed explanations + + Notes + ----- + Model based on: + - Lin & Crago (2002): Mathematical model framework + - Aniss et al. (1990b): Human GTO physiological data + - Elias PhD thesis (pg 83): Implementation details + + The logarithmic relationship captures the GTO's ability to encode + force over a wide dynamic range, from threshold detection of small + forces to saturation at high forces, providing force feedback for + motor control and protective reflexes. + """ + return { + # Force-to-firing relationship parameters + # Firing rate [Hz] = G1 * log(force[N]/G2 + 1) + "G1": 40, # Logarithmic gain coefficient [Hz] + # Controls the sensitivity of force-to-firing conversion + # Higher G1 = more sensitive to force changes + # Typical range: 30-60 Hz for different muscles + # Source: Lin & Crago (2002), Elias thesis uses 40 Hz + "G2": 4, # Force scaling coefficient [N] + # Sets the force level for logarithmic scaling + # Lower G2 = more sensitive to low forces + # Higher G2 = requires higher forces for activation + # Typical range: 2-8 N depending on muscle strength + # Source: Calibrated for human muscle force ranges + } + + @staticmethod + def create_gto_parameters_for_muscle(muscle_type: str = "FDI") -> Dict[str, Any]: + """ + Create GTO parameters optimized for specific muscle types. + + Parameters + ---------- + muscle_type : str, optional + Type of muscle ("FDI", "Sol", "generic"), by default "FDI" + + Returns + ------- + Dict[str, Any] + Dictionary of muscle-specific GTO parameters + + Notes + ----- + Different muscles have different force production capabilities + and thus require different GTO sensitivity parameters: + + - FDI (First Dorsal Interosseous): Small hand muscle, low forces + - Sol (Soleus): Large calf muscle, high forces + - Generic: General-purpose parameters + """ + if muscle_type == "FDI": + # Small hand muscle - more sensitive to small forces + return { + "G1": 45, # Higher sensitivity for small force detection + "G2": 2, # Lower threshold for activation at small forces + } + + elif muscle_type == "Sol": + # Large calf muscle - less sensitive, handles high forces + return { + "G1": 35, # Lower sensitivity appropriate for large forces + "G2": 8, # Higher threshold matching muscle's force capacity + } + + elif muscle_type == "generic": + # General-purpose parameters + return GolgiTendonOrganModel.create_default_gto_parameters() + + else: + raise ValueError(f"Unknown muscle type: {muscle_type}. Use 'FDI', 'Sol', or 'generic'.") diff --git a/myogen/simulator/jaxley/proprioception/spindle.py b/myogen/simulator/jaxley/proprioception/spindle.py new file mode 100644 index 00000000..76eb8f2a --- /dev/null +++ b/myogen/simulator/jaxley/proprioception/spindle.py @@ -0,0 +1,498 @@ +""" +Muscle Spindle Model - Jaxley Backend (Pure Python/NumPy) + +This module provides the muscle spindle model using pure NumPy for the Jaxley backend. +Maintains full API compatibility with the NEURON version. + +Based on the model by Mileusnic et al. (2006) and implementation by Elias (2014). +""" + +from typing import Any, Dict + +import numpy as np + +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__ms + + +class _Spindle__NumPy: + """ + Pure NumPy implementation of muscle spindle model. + + Faithfully implements the full Mileusnic et al. (2006) model: + - Hill-type fusimotor activation (RK4 for Bag1/Bag2, algebraic for Chain) + - Second-order mechanical ODE for intrafusal fiber tensions (RK4) + - Nonlinear force-velocity damping with asymmetric lengthening/shortening + - Afferent occlusion mechanism for primary afferents (Ia) + - Two-component secondary afferent (II) from sensory + polar regions + + Matches the NEURON Cython implementation (_spindle.pyx) equation-by-equation. + """ + + def __init__(self, tstop: float, dt: float, spinD: Dict[str, Any]): + self.tstop = tstop + self.dt = dt + self.params = spinD + + # Initialize storage arrays + self.n_steps = int(tstop / dt) + 1 + self.time = np.linspace(0, tstop, self.n_steps) + + # Output arrays + self.Ia = np.zeros(self.n_steps) # Primary afferent firing [pps] + self.II = np.zeros(self.n_steps) # Secondary afferent firing [pps] + self.aBag1 = np.zeros(self.n_steps) # Bag1 activation [0-1] + self.aBag2 = np.zeros(self.n_steps) # Bag2 activation [0-1] + self.aChain = np.zeros(self.n_steps) # Chain activation [0-1] + self.T = np.zeros((3, self.n_steps)) # Tensions [Bag1, Bag2, Chain] [FU] + self.dT = np.zeros((3, self.n_steps)) # Tension rates [FU/s] + + # State variables + self.step_idx = 0 + + def _tension_accel(self, T_val, z_val, L, V, A, b_coef, gamma_force): + """Compute d²T/dt² for the intrafusal fiber 2nd-order ODE. + + Mileusnic et al. 2006, Eq. 1-3: spring-mass-damper system with + nonlinear force-velocity relationship. + """ + p = self.params + vel_diff = V - z_val / p["K_SR"] + fv = p["C_L"] if vel_diff >= 0 else p["C_S"] + sgn = 1.0 if vel_diff >= 0 else -1.0 + + fv_term = (fv * b_coef * sgn + * abs(vel_diff) ** p["a"] + * (L - p["L0_SR"] - T_val / p["K_SR"] - p["R"])) + spring = p["K_PR"] * (L - p["L0_SR"] - T_val / p["K_SR"] - p["L0_PR"]) + + return p["K_SR"] / p["M"] * ( + fv_term + p["M"] * A + gamma_force - T_val + spring + ) + + def integrate( + self, + muscle_length__L0: float, + muscle_velocity__L0_per_s: float, + muscle_acceleration__L0_per_s2: float, + gamma_dynamic_drive__Hz: float, + gamma_static_drive__Hz: float, + ) -> tuple[float, float]: + """ + Integrate spindle model for one time step. + + Parameters + ---------- + muscle_length__L0 : float + Normalized muscle length + muscle_velocity__L0_per_s : float + Muscle velocity [L0/s] + muscle_acceleration__L0_per_s2 : float + Muscle acceleration [L0/s²] + gamma_dynamic_drive__Hz : float + Dynamic gamma drive [pps] + gamma_static_drive__Hz : float + Static gamma drive [pps] + + Returns + ------- + tuple[float, float] + (Ia firing rate, II firing rate) in Hz + """ + p = self.params + # NEURON stores dt in seconds (dt * 1e-3), tau in seconds — all consistent + dt_s = self.dt / 1000.0 + i = self.step_idx + + gDyn = gamma_dynamic_drive__Hz + gStat = gamma_static_drive__Hz + L = muscle_length__L0 + V = muscle_velocity__L0_per_s + A = muscle_acceleration__L0_per_s2 + + # === FUSIMOTOR ACTIVATIONS (Mileusnic Eq. 4-5) === + # Hill-type saturation: gamma^P / (gamma^P + f^P) + # Gives exactly 0 at zero gamma, saturates to 1 at high gamma. + P = p["P"] + target_bag1 = (gDyn ** P / (gDyn ** P + p["fBag1"] ** P) + if gDyn > 0 else 0.0) + target_bag2 = (gStat ** P / (gStat ** P + p["fBag2"] ** P) + if gStat > 0 else 0.0) + # Chain: algebraic (instantaneous, no ODE) + a_chain = (gStat ** P / (gStat ** P + p["fChain"] ** P) + if gStat > 0 else 0.0) + + # Bag1: RK4 integration of da/dt = (target - a) / tau + if i > 0: + a_prev = self.aBag1[i - 1] + tau = p["tau1"] + k1 = (target_bag1 - a_prev) / tau + k2 = (target_bag1 - (a_prev + dt_s / 2 * k1)) / tau + k3 = (target_bag1 - (a_prev + dt_s / 2 * k2)) / tau + k4 = (target_bag1 - (a_prev + dt_s * k3)) / tau + a_bag1 = a_prev + dt_s / 6 * (k1 + 2 * k2 + 2 * k3 + k4) + else: + a_bag1 = target_bag1 + + # Bag2: RK4 integration of da/dt = (target - a) / tau + if i > 0: + a_prev = self.aBag2[i - 1] + tau = p["tau2"] + k1 = (target_bag2 - a_prev) / tau + k2 = (target_bag2 - (a_prev + dt_s / 2 * k1)) / tau + k3 = (target_bag2 - (a_prev + dt_s / 2 * k2)) / tau + k4 = (target_bag2 - (a_prev + dt_s * k3)) / tau + a_bag2 = a_prev + dt_s / 6 * (k1 + 2 * k2 + 2 * k3 + k4) + else: + a_bag2 = target_bag2 + + self.aBag1[i] = a_bag1 + self.aBag2[i] = a_bag2 + self.aChain[i] = a_chain + + # === INTRAFUSAL FIBER TENSIONS (2nd-order Mileusnic ODE via RK4) === + # Three fiber types: Bag1 (dynamic), Bag2 (static), Chain (static) + # Each has: dT/dt = z, dz/dt = K_SR/M * [FV_damping + inertial + spring] + activations = [a_bag1, a_bag2, a_chain] + + for fi in range(3): + act = activations[fi] + + # Per-fiber damping and gamma force coefficients + if fi == 0: # Bag1 + b_coef = p["b0Bag1"] + p["b1Bag1"] * act + gamma_force = p["G1"] * act + elif fi == 1: # Bag2 + b_coef = p["b0Bag2"] + p["b2Bag2"] * act + gamma_force = p["G2"] * act + else: # Chain + b_coef = p["b0Chain"] + p["b2Chain"] * act + gamma_force = p["G2Chain"] * act + + T_prev = self.T[fi, i - 1] if i > 0 else 0.0 + z_prev = self.dT[fi, i - 1] if i > 0 else 0.0 + + # Coupled RK4 for 2nd-order system + acc = self._tension_accel + k1y = z_prev + k1z = acc(T_prev, z_prev, L, V, A, b_coef, gamma_force) + + k2y = z_prev + dt_s / 2 * k1z + k2z = acc(T_prev + dt_s / 2 * k1y, + z_prev + dt_s / 2 * k1z, + L, V, A, b_coef, gamma_force) + + k3y = z_prev + dt_s / 2 * k2z + k3z = acc(T_prev + dt_s / 2 * k2y, + z_prev + dt_s / 2 * k2z, + L, V, A, b_coef, gamma_force) + + k4y = z_prev + dt_s * k3z + k4z = acc(T_prev + dt_s * k3y, + z_prev + dt_s * k3z, + L, V, A, b_coef, gamma_force) + + self.T[fi, i] = T_prev + dt_s / 6 * (k1y + 2 * k2y + 2 * k3y + k4y) + self.dT[fi, i] = z_prev + dt_s / 6 * (k1z + 2 * k2z + 2 * k3z + k4z) + + # === AFFERENT FIRING RATES (Mileusnic Eq. 5-6) === + threshold = p["LN_SR"] - p["L0_SR"] # 0.0423 - 0.04 = 0.0023 + + # Primary afferent (Ia) — per-fiber contribution from SR stretch + ia_bag1 = p["gBag1"] * max(0.0, self.T[0, i] / p["K_SR"] - threshold) + ia_bag2 = p["gBag2A1"] * max(0.0, self.T[1, i] / p["K_SR"] - threshold) + ia_chain = p["gChainA1"] * max(0.0, self.T[2, i] / p["K_SR"] - threshold) + + # Occlusion (Mileusnic Eq. 6): nonlinear interaction between fiber types + B2C = ia_bag2 + ia_chain + if B2C >= ia_bag1: + Ia_rate = B2C + p["S"] * ia_bag1 + else: + Ia_rate = ia_bag1 + p["S"] * B2C + + # Secondary afferent (II) — Bag2 and Chain only (Mileusnic Eq. 7) + # Two components: sensory region stretch + polar region stretch + II_rate = 0.0 + for fi in range(1, 3): + gain = p["gBag2A2"] if fi == 1 else p["gChainA2"] + sr_stretch = self.T[fi, i] / p["K_SR"] - threshold + pr_stretch = (L - self.T[fi, i] / p["K_SR"] + - p["L0_SR"] - p["LN_PR"]) + contribution = gain * ( + p["X"] * p["Lsec"] / p["L0_SR"] * sr_stretch + + (1 - p["X"]) * p["Lsec"] / p["L0_PR"] * pr_stretch + ) + II_rate += max(0.0, contribution) + + self.Ia[i] = Ia_rate + self.II[i] = II_rate + + self.step_idx += 1 + return Ia_rate, II_rate + + +@beartowertype +class SpindleModel: + """ + API wrapper for the muscle spindle model - Jaxley Backend. + + This class provides an intuitive interface for creating muscle spindle models + with user-friendly parameter names that are internally mapped to the + correct format expected by the underlying Spindle implementation. + + The muscle spindle is a proprioceptive sensory organ that detects changes + in muscle length and velocity, providing feedback for motor control. + + Parameters + ---------- + simulation_time__ms : Quantity__ms + Total simulation time in milliseconds + time_step__ms : Quantity__ms + Integration time step in milliseconds + spindle_parameters : Dict[str, Any] + Dictionary containing spindle model parameters + """ + + def __init__( + self, + simulation_time__ms: Quantity__ms, + time_step__ms: Quantity__ms, + spindle_parameters: Dict[str, Any], + ): + self.simulation_time__ms = simulation_time__ms + self.time_step__ms = time_step__ms + self.spindle_parameters = spindle_parameters.copy() + + # Private working copies for internal use + self._simulation_time__ms = simulation_time__ms + self._time_step__ms = time_step__ms + self._spindle_parameters = spindle_parameters.copy() + + # Validate inputs + self._validate_parameters() + + # Create the underlying Spindle model + self._spindle_model = self._create_spindle_model() + + def _validate_parameters(self) -> None: + """Validate input parameters.""" + if self._simulation_time__ms <= 0: + raise ValueError("simulation_time__ms must be positive") + + if self._time_step__ms <= 0: + raise ValueError("time_step__ms must be positive") + + if self._simulation_time__ms <= self._time_step__ms: + raise ValueError("simulation_time__ms must be greater than time_step__ms") + + def _create_spindle_model(self) -> _Spindle__NumPy: + """ + Create the underlying Spindle model instance. + + This method maps the user-friendly parameter names to the format + expected by the Spindle constructor. + """ + return _Spindle__NumPy( + tstop=self._simulation_time__ms.magnitude, + dt=self._time_step__ms.magnitude, + spinD=self._spindle_parameters, + ) + + def integrate( + self, + muscle_length__L0: float, + muscle_velocity__L0_per_s: float, + muscle_acceleration__L0_per_s2: float, + gamma_dynamic_drive__Hz: float, + gamma_static_drive__Hz: float, + ) -> tuple[float, float]: + """ + Integrate the spindle model for one time step. + + Parameters + ---------- + muscle_length__L0 : float + Current muscle length normalized to L0 + muscle_velocity__L0_per_s : float + Current muscle velocity in L0/s + muscle_acceleration__L0_per_s2 : float + Current muscle acceleration in L0/s² + gamma_dynamic_drive__Hz : float + Gamma dynamic motor neuron drive frequency in Hz + gamma_static_drive__Hz : float + Gamma static motor neuron drive frequency in Hz + + Returns + ------- + tuple[float, float] + Primary afferent (Ia) and secondary afferent (II) firing rates in Hz + """ + return self._spindle_model.integrate( + muscle_length__L0, + muscle_velocity__L0_per_s, + muscle_acceleration__L0_per_s2, + gamma_dynamic_drive__Hz, + gamma_static_drive__Hz, + ) + + @property + def primary_afferent_firing__Hz(self) -> np.ndarray: + """Get primary afferent (Ia) firing rate time series in Hz.""" + return np.asarray(self._spindle_model.Ia) + + @property + def secondary_afferent_firing__Hz(self) -> np.ndarray: + """Get secondary afferent (II) firing rate time series in Hz.""" + return np.asarray(self._spindle_model.II) + + @property + def bag1_activation(self) -> np.ndarray: + """Get Bag1 fiber activation time series.""" + return np.asarray(self._spindle_model.aBag1) + + @property + def bag2_activation(self) -> np.ndarray: + """Get Bag2 fiber activation time series.""" + return np.asarray(self._spindle_model.aBag2) + + @property + def chain_activation(self) -> np.ndarray: + """Get Chain fiber activation time series.""" + return np.asarray(self._spindle_model.aChain) + + @property + def intrafusal_tensions(self) -> np.ndarray: + """Get intrafusal fiber tensions matrix (3 × time_points) [Bag1, Bag2, Chain].""" + return np.asarray(self._spindle_model.T) + + @property + def time_vector(self) -> np.ndarray: + """Get simulation time vector in milliseconds.""" + return np.asarray(self._spindle_model.time) + + def __repr__(self) -> str: + """String representation of the spindle model.""" + return f"SpindleModel(t_sim={self.simulation_time__ms}ms, dt={self.time_step__ms}ms)" + + @staticmethod + def create_default_spindle_parameters( + species: str = "human", deafferent_ia: bool = False, deafferent_ii: bool = False + ) -> Dict[str, Any]: + """ + Create default spindle parameter dictionary. + + Parameters + ---------- + species : str, optional + Species type ("human" or "cat"), by default "human" + deafferent_ia : bool, optional + Whether to simulate Ia afferent deafferentation, by default False + deafferent_ii : bool, optional + Whether to simulate II afferent deafferentation, by default False + + Returns + ------- + Dict[str, Any] + Dictionary of spindle parameters with detailed explanations + + Raises + ------ + ValueError + If species is not recognized + """ + # Base spindle parameters (Mileusnic et al., 2006) + spindle_params = { + # Fusimotor activation parameters + "fBag1": 60, # Fusimotor frequency to activation constant for Bag1 [Hz] + "fBag2": 60, # Fusimotor frequency to activation constant for Bag2 [Hz] + "fChain": 90, # Fusimotor frequency to activation constant for Chain [Hz] + "P": 2, # Fusimotor frequency to activation power constant + # Force generation coefficients + "G1": 0.0289, # Dynamic fusimotor input force generation coef [FU] + "G2": 0.0636, # Static fusimotor input force generation coef [FU] + "G2Chain": 0.0954, # Static fusimotor input force gen coef for Chain [FU] + # Sensory Region (SR) mechanical parameters + "K_SR": 10.4649, # SR spring constant [FU/L0] - detects length changes + "L0_SR": 0.04, # SR rest length [L0] - baseline length + "LN_SR": 0.0423, # SR threshold length [L0] - minimum for activation + # Polar Region (PR) mechanical parameters + "K_PR": 0.15, # PR spring constant [FU/L0] - contractile region + "L0_PR": 0.76, # PR rest length [L0] - baseline contractile length + "LN_PR": 0.89, # PR threshold length [L0] - minimum for activation + # Intrafusal fiber mechanical properties + "M": 0.0002, # Intrafusal fiber mass [FU/(L0/s²)] - inertial component + # Passive damping coefficients [FU/(L0/s)] - baseline viscosity + "b0Bag1": 0.0605, # Bag1 passive damping + "b0Bag2": 0.0822, # Bag2 passive damping + "b0Chain": 0.0822, # Chain passive damping + # Fusimotor-dependent damping coefficients [FU/(L0/s)] + "b1Bag1": 0.2592, # Dynamic fusimotor damping for Bag1 + "b2Bag2": -0.0460, # Static fusimotor damping for Bag2 + "b2Chain": -0.0690, # Static fusimotor damping for Chain + # Force-velocity relationship parameters + "a": 0.3, # Nonlinear velocity dependence power constant + "C_L": 1, # Lengthening coefficient of asymmetry in F-V curve + "C_S": 0.42, # Shortening coefficient of asymmetry in F-V curve + "R": 0.46, # Fascicle length where force production is zero [L0] + # Afferent firing properties + "X": 0.7, # Secondary afferent percentage on sensory region [0-1] + "Lsec": 0.04, # Secondary afferent rest length [L0] + "S": 0.156, # Occlusion factor for primary afferent interactions + # Activation time constants (Mileusnic 2006, Table 1) + # Bag1/Bag2: ODE time constants; Chain: algebraic (instantaneous) + "tau1": 0.149, # Bag1 activation time constant [s] (dynamic bag) + "tau2": 0.205, # Bag2 activation time constant [s] (static bag) + # Afferent sensitivity gains [Hz/L0] - firing rate per unit stretch + "gBag1": 6500, # Bag1 contribution to primary afferent (Ia) + "gBag2A1": 3250, # Bag2 contribution to primary afferent (Ia) + "gChainA1": 3250, # Chain contribution to primary afferent (Ia) + "gBag2A2": 3500, # Bag2 contribution to secondary afferent (II) + "gChainA2": 3500, # Chain contribution to secondary afferent (II) + } + + # Species-specific and deafferentation modifications + if species == "human": + if not deafferent_ii and not deafferent_ia: + # Normal human spindle (Case 1, Elias thesis pg 66) + pass # Use default values above + + elif deafferent_ii and not deafferent_ia: + # Human with Type II deafferentation (Case 2, Elias thesis pg 66) + spindle_params.update( + { + "gBag1": 7000, # Enhanced Bag1 sensitivity + "gBag2A1": 3750, # Enhanced Bag2 primary sensitivity + "gChainA1": 3750, # Enhanced Chain primary sensitivity + "gBag2A2": 0, # No Bag2 secondary afferents + "gChainA2": 0, # No Chain secondary afferents + } + ) + + elif not deafferent_ii and deafferent_ia: + # Human with Ia deafferentation (Case 3, Elias thesis pg 66) + spindle_params.update( + { + "gBag1": 0, # No Bag1 primary afferents + "gBag2A1": 0, # No Bag2 primary afferents + "gChainA1": 0, # No Chain primary afferents + "gBag2A2": 4500, # Enhanced Bag2 secondary sensitivity + "gChainA2": 4500, # Enhanced Chain secondary sensitivity + } + ) + + elif species == "cat": + # Cat spindle parameters (original Mileusnic values) + spindle_params.update( + { + "gBag1": 20000, # Higher sensitivity in cat + "gBag2A1": 10000, # Higher Bag2 primary sensitivity + "gChainA1": 10000, # Higher Chain primary sensitivity + "gBag2A2": 7250, # Higher Bag2 secondary sensitivity + "gChainA2": 7250, # Higher Chain secondary sensitivity + } + ) + + else: + raise ValueError(f"Unknown species: {species}. Use 'human' or 'cat'.") + + return spindle_params diff --git a/myogen/simulator/jaxley/simulation_engine.py b/myogen/simulator/jaxley/simulation_engine.py new file mode 100644 index 00000000..3ebd2a6d --- /dev/null +++ b/myogen/simulator/jaxley/simulation_engine.py @@ -0,0 +1,711 @@ +""" +Jaxley Simulation Engine for MyoGen. + +This module provides the core simulation infrastructure for running +biophysical neural network simulations using Jaxley. + +Key Features: +- JAX-accelerated neural simulation +- Spike detection and recording +- Synaptic event handling +- Current injection support +- Network-level simulation control +""" + +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from dataclasses import dataclass, field +import numpy as np +import jax +import jax.numpy as jnp +from functools import partial + +import jaxley as jx +import myogen + + +# ============================================================================= +# DATA CLASSES +# ============================================================================= + +@dataclass +class SpikeRecord: + """Container for recorded spikes.""" + times: List[float] = field(default_factory=list) + ids: List[int] = field(default_factory=list) + + def add_spike(self, time: float, neuron_id: int): + """Record a spike.""" + self.times.append(time) + self.ids.append(neuron_id) + + def to_arrays(self) -> Tuple[np.ndarray, np.ndarray]: + """Convert to numpy arrays.""" + return np.array(self.times), np.array(self.ids) + + def clear(self): + """Clear all recorded spikes.""" + self.times.clear() + self.ids.clear() + + +@dataclass +class VoltageRecord: + """Container for voltage traces.""" + times: np.ndarray = None + voltages: Dict[int, np.ndarray] = field(default_factory=dict) + + def add_trace(self, neuron_id: int, voltage: np.ndarray): + """Add a voltage trace for a neuron.""" + self.voltages[neuron_id] = voltage + + +@dataclass +class SimulationConfig: + """Configuration for Jaxley simulation.""" + dt: float = 0.025 # Time step (ms) + duration: float = 1000.0 # Simulation duration (ms) + record_voltage: bool = False + record_spikes: bool = True + spike_threshold: float = 0.0 # mV + initial_voltage: float = -70.0 # mV + + +# ============================================================================= +# SPIKE DETECTION +# ============================================================================= + +def detect_spikes( + voltage: jnp.ndarray, + threshold: float = 0.0, + refractory: int = 5, +) -> jnp.ndarray: + """ + Detect spikes using threshold crossing. + + Parameters + ---------- + voltage : jnp.ndarray + Voltage trace, shape (n_timesteps,) or (n_neurons, n_timesteps). + threshold : float + Spike threshold in mV. + refractory : int + Minimum samples between spikes. + + Returns + ------- + jnp.ndarray + Boolean array of spike times. + """ + # Threshold crossings (upward) + above = voltage > threshold + below_before = jnp.roll(voltage, 1, axis=-1) <= threshold + crossings = above & below_before + + # Zero out first sample (no valid crossing) + if crossings.ndim == 1: + crossings = crossings.at[0].set(False) + else: + crossings = crossings.at[:, 0].set(False) + + return crossings + + +def apply_refractory( + spikes: jnp.ndarray, + refractory_samples: int, +) -> jnp.ndarray: + """Apply refractory period to spike train.""" + def scan_fn(carry, spike): + counter, last_spike = carry + # Can spike if counter is 0 + can_spike = counter == 0 + new_spike = spike & can_spike + # Reset counter on spike, otherwise decrement + new_counter = jnp.where(new_spike, refractory_samples, jnp.maximum(counter - 1, 0)) + return (new_counter, new_spike), new_spike + + init_carry = (0, False) + _, filtered_spikes = jax.lax.scan(scan_fn, init_carry, spikes) + return filtered_spikes + + +# ============================================================================= +# CURRENT INJECTION +# ============================================================================= + +@dataclass +class CurrentInjection: + """ + Specification for current injection. + + Parameters + ---------- + amplitude : float or array + Current amplitude in nA. Can be constant or time-varying. + start : float + Start time in ms. + duration : float + Duration in ms. + location : str + Injection location ("soma", "dendrite", etc.). + """ + amplitude: Union[float, np.ndarray] + start: float = 0.0 + duration: float = float('inf') + location: str = "soma" + + def get_current(self, t: float) -> float: + """Get current at time t.""" + if t < self.start or t >= self.start + self.duration: + return 0.0 + + if isinstance(self.amplitude, (int, float)): + return float(self.amplitude) + else: + # Time-varying amplitude + idx = int((t - self.start) / 0.025) # Assume dt=0.025 + if idx < len(self.amplitude): + return float(self.amplitude[idx]) + return 0.0 + + +# ============================================================================= +# SYNAPTIC EVENT HANDLING +# ============================================================================= + +@dataclass +class SynapticEvent: + """ + A synaptic activation event. + + Parameters + ---------- + time : float + Event time in ms. + target_id : int + Target neuron ID. + weight : float + Synaptic weight (uS). + synapse_type : str + "excitatory" or "inhibitory". + """ + time: float + target_id: int + weight: float + synapse_type: str = "excitatory" + + +class SynapticEventQueue: + """Priority queue for synaptic events.""" + + def __init__(self): + self.events: List[SynapticEvent] = [] + + def add_event(self, event: SynapticEvent): + """Add event and maintain time ordering.""" + self.events.append(event) + self.events.sort(key=lambda e: e.time) + + def pop_events_until(self, t: float) -> List[SynapticEvent]: + """Get and remove all events up to time t.""" + to_process = [] + while self.events and self.events[0].time <= t: + to_process.append(self.events.pop(0)) + return to_process + + def clear(self): + """Clear all events.""" + self.events.clear() + + +# ============================================================================= +# NETWORK SIMULATOR +# ============================================================================= + +class JaxleyNetworkSimulator: + """ + Network-level simulator for Jaxley neural populations. + + Manages multiple cell populations, connectivity, and simulation. + + Parameters + ---------- + config : SimulationConfig + Simulation configuration. + """ + + def __init__(self, config: Optional[SimulationConfig] = None): + self.config = config or SimulationConfig() + + self.populations: Dict[str, List[Any]] = {} + self.connections: List[Dict] = [] + self.spike_records: Dict[str, SpikeRecord] = {} + self.voltage_records: Dict[str, VoltageRecord] = {} + self.current_injections: Dict[str, List[CurrentInjection]] = {} + self.event_queue = SynapticEventQueue() + + self._time = 0.0 + self._initialized = False + + def add_population( + self, + name: str, + cells: List[Any], + record_spikes: bool = True, + record_voltage: bool = False, + spike_threshold: float = None, + ): + """ + Add a neural population. + + Parameters + ---------- + name : str + Population name. + cells : list + List of Jaxley cells or cell wrappers. + record_spikes : bool + Whether to record spikes. + record_voltage : bool + Whether to record voltage traces. + spike_threshold : float, optional + Spike threshold (default from config). + """ + self.populations[name] = cells + + if record_spikes: + self.spike_records[name] = SpikeRecord() + + if record_voltage: + self.voltage_records[name] = VoltageRecord() + + # Store threshold + if spike_threshold is not None: + for cell in cells: + if hasattr(cell, 'spike_threshold'): + cell.spike_threshold = spike_threshold + + def connect( + self, + source_pop: str, + target_pop: str, + weight: float, + delay: float = 1.0, + probability: float = 1.0, + synapse_type: str = "excitatory", + ): + """ + Create connections between populations. + + Parameters + ---------- + source_pop : str + Source population name. + target_pop : str + Target population name. + weight : float + Synaptic weight in uS. + delay : float + Synaptic delay in ms. + probability : float + Connection probability [0, 1]. + synapse_type : str + "excitatory" or "inhibitory". + """ + connection = { + "source": source_pop, + "target": target_pop, + "weight": weight, + "delay": delay, + "probability": probability, + "synapse_type": synapse_type, + "connections": [], # (source_idx, target_idx) pairs + } + + # Generate connection pairs + source_cells = self.populations.get(source_pop, []) + target_cells = self.populations.get(target_pop, []) + + for i, src in enumerate(source_cells): + for j, tgt in enumerate(target_cells): + if myogen.RANDOM_GENERATOR.random() < probability: + connection["connections"].append((i, j)) + + self.connections.append(connection) + + def inject_current( + self, + population: str, + neuron_ids: Union[int, List[int]], + injection: CurrentInjection, + ): + """ + Add current injection to neurons. + + Parameters + ---------- + population : str + Target population name. + neuron_ids : int or list + Neuron indices to inject into. + injection : CurrentInjection + Current injection specification. + """ + if population not in self.current_injections: + self.current_injections[population] = [] + + if isinstance(neuron_ids, int): + neuron_ids = [neuron_ids] + + for nid in neuron_ids: + self.current_injections[population].append((nid, injection)) + + def _process_spikes(self, t: float): + """Process spikes and generate synaptic events.""" + for pop_name, record in self.spike_records.items(): + cells = self.populations[pop_name] + threshold = self.config.spike_threshold + + for i, cell in enumerate(cells): + # Get voltage (assumes cell has voltage attribute or method) + if hasattr(cell, 'get_voltage'): + v = cell.get_voltage() + elif hasattr(cell, 'v'): + v = cell.v + else: + continue + + # Simple threshold crossing detection + if hasattr(cell, '_prev_voltage'): + if cell._prev_voltage <= threshold < v: + record.add_spike(t, i) + + # Generate synaptic events + self._generate_synaptic_events(pop_name, i, t) + + cell._prev_voltage = v + + def _generate_synaptic_events(self, source_pop: str, source_idx: int, t: float): + """Generate synaptic events from a spike.""" + for conn in self.connections: + if conn["source"] == source_pop: + for src_i, tgt_j in conn["connections"]: + if src_i == source_idx: + event = SynapticEvent( + time=t + conn["delay"], + target_id=tgt_j, + weight=conn["weight"], + synapse_type=conn["synapse_type"], + ) + self.event_queue.add_event(event) + + def _apply_synaptic_events(self, t: float, pop_name: str): + """Apply pending synaptic events to population.""" + events = self.event_queue.pop_events_until(t) + + for event in events: + # Find target population from connections + for conn in self.connections: + if conn["target"] == pop_name: + cells = self.populations[pop_name] + if 0 <= event.target_id < len(cells): + cell = cells[event.target_id] + + # Apply synaptic conductance + if hasattr(cell, 'apply_synaptic_input'): + cell.apply_synaptic_input( + event.weight, + event.synapse_type, + ) + + def _get_current_injection(self, pop_name: str, neuron_id: int, t: float) -> float: + """Get total injected current for a neuron at time t.""" + total = 0.0 + if pop_name in self.current_injections: + for nid, inj in self.current_injections[pop_name]: + if nid == neuron_id: + total += inj.get_current(t) + return total + + def initialize(self): + """Initialize simulation state.""" + self._time = 0.0 + + # Initialize cells + for pop_name, cells in self.populations.items(): + for cell in cells: + if hasattr(cell, 'set'): + cell.set("v", self.config.initial_voltage) + cell._prev_voltage = self.config.initial_voltage + + # Clear records + for record in self.spike_records.values(): + record.clear() + self.event_queue.clear() + + self._initialized = True + + def step(self, dt: Optional[float] = None): + """ + Advance simulation by one time step. + + Parameters + ---------- + dt : float, optional + Time step (default from config). + """ + if not self._initialized: + self.initialize() + + dt = dt or self.config.dt + + # Apply synaptic events + for pop_name in self.populations: + self._apply_synaptic_events(self._time, pop_name) + + # Process spikes and generate events + self._process_spikes(self._time) + + # Advance time + self._time += dt + + def run( + self, + duration: Optional[float] = None, + progress_callback: Optional[Callable[[float], None]] = None, + ) -> Dict[str, Any]: + """ + Run simulation for specified duration. + + Parameters + ---------- + duration : float, optional + Simulation duration (default from config). + progress_callback : callable, optional + Called with progress fraction. + + Returns + ------- + dict + Simulation results including spike records. + """ + duration = duration or self.config.duration + dt = self.config.dt + n_steps = int(duration / dt) + + self.initialize() + + for step in range(n_steps): + self.step(dt) + + if progress_callback and step % 100 == 0: + progress_callback(step / n_steps) + + # Compile results + results = { + "time": np.arange(0, duration, dt), + "spikes": {}, + "voltages": {}, + } + + for name, record in self.spike_records.items(): + results["spikes"][name] = record.to_arrays() + + for name, record in self.voltage_records.items(): + results["voltages"][name] = record.voltages + + return results + + def get_spike_times(self, population: str) -> Tuple[np.ndarray, np.ndarray]: + """Get spike times and IDs for a population.""" + if population in self.spike_records: + return self.spike_records[population].to_arrays() + return np.array([]), np.array([]) + + +# ============================================================================= +# SINGLE CELL SIMULATOR +# ============================================================================= + +class JaxleyCellSimulator: + """ + Simulator for single Jaxley cells. + + Provides a simpler interface for simulating individual cells + with current injection. + + Parameters + ---------- + cell : jx.Cell + Jaxley cell to simulate. + dt : float + Time step in ms. + """ + + def __init__(self, cell: jx.Cell, dt: float = 0.025): + self.cell = cell + self.dt = dt + self._time = 0.0 + self._voltage_trace = [] + self._spike_times = [] + + def inject_step_current( + self, + amplitude: float, + start: float, + duration: float, + ) -> CurrentInjection: + """Create a step current injection.""" + return CurrentInjection( + amplitude=amplitude, + start=start, + duration=duration, + ) + + def simulate( + self, + duration: float, + i_ext: Union[float, np.ndarray, CurrentInjection] = 0.0, + record_voltage: bool = True, + ) -> Dict[str, np.ndarray]: + """ + Simulate cell for given duration. + + Parameters + ---------- + duration : float + Simulation duration in ms. + i_ext : float, array, or CurrentInjection + External current. + record_voltage : bool + Whether to record voltage. + + Returns + ------- + dict + Results with 'time', 'voltage', 'spikes'. + """ + n_steps = int(duration / self.dt) + times = np.arange(0, duration, self.dt) + + # Prepare current + if isinstance(i_ext, CurrentInjection): + currents = np.array([i_ext.get_current(t) for t in times]) + elif isinstance(i_ext, np.ndarray): + currents = i_ext + else: + currents = np.full(n_steps, float(i_ext)) + + # Initialize + self._voltage_trace = [] + self._spike_times = [] + self._time = 0.0 + + # Note: Actual Jaxley simulation would use jx.integrate() + # This is a placeholder showing the interface + + results = { + "time": times, + "voltage": np.zeros(n_steps), # Placeholder + "spikes": np.array([]), + "current": currents, + } + + return results + + +# ============================================================================= +# UTILITY FUNCTIONS +# ============================================================================= + +def compute_firing_rate( + spike_times: np.ndarray, + spike_ids: np.ndarray, + neuron_id: int, + window: float = 100.0, + dt: float = 1.0, + duration: float = None, +) -> Tuple[np.ndarray, np.ndarray]: + """ + Compute instantaneous firing rate for a neuron. + + Parameters + ---------- + spike_times : array + Spike times in ms. + spike_ids : array + Neuron IDs for each spike. + neuron_id : int + Neuron to analyze. + window : float + Smoothing window in ms. + dt : float + Output time step. + duration : float, optional + Total duration (default: max spike time + window). + + Returns + ------- + times : array + Time points. + rates : array + Firing rates in Hz. + """ + # Get spikes for this neuron + mask = spike_ids == neuron_id + neuron_spikes = spike_times[mask] + + if len(neuron_spikes) == 0: + if duration is None: + return np.array([0]), np.array([0]) + times = np.arange(0, duration, dt) + return times, np.zeros_like(times) + + # Time bins + if duration is None: + duration = neuron_spikes.max() + window + times = np.arange(0, duration, dt) + + # Compute rate using sliding window + rates = np.zeros(len(times)) + half_window = window / 2 + + for i, t in enumerate(times): + count = np.sum((neuron_spikes >= t - half_window) & (neuron_spikes < t + half_window)) + rates[i] = count * 1000.0 / window # Convert to Hz + + return times, rates + + +def compute_isi_statistics( + spike_times: np.ndarray, + spike_ids: np.ndarray, + neuron_id: int, +) -> Dict[str, float]: + """ + Compute interspike interval statistics. + + Returns + ------- + dict + 'mean_isi', 'std_isi', 'cv_isi', 'mean_rate' + """ + mask = spike_ids == neuron_id + neuron_spikes = np.sort(spike_times[mask]) + + if len(neuron_spikes) < 2: + return { + "mean_isi": np.nan, + "std_isi": np.nan, + "cv_isi": np.nan, + "mean_rate": 0.0, + } + + isis = np.diff(neuron_spikes) + + return { + "mean_isi": np.mean(isis), + "std_isi": np.std(isis), + "cv_isi": np.std(isis) / np.mean(isis) if np.mean(isis) > 0 else np.nan, + "mean_rate": 1000.0 / np.mean(isis), # Hz + } diff --git a/myogen/simulator/jaxley/simulation_runner.py b/myogen/simulator/jaxley/simulation_runner.py new file mode 100644 index 00000000..c0356bca --- /dev/null +++ b/myogen/simulator/jaxley/simulation_runner.py @@ -0,0 +1,491 @@ +""" +Simulation Runner - Jaxley Backend + +Coordinates real-time stepping of non-Jaxley components (muscle, proprioception, +interneurons) in neuromuscular simulations. + +Architectural note — split between SimulationRunner and jx.integrate(): + This class does NOT call ``jx.integrate()`` on neural cells. Neural cell + integration happens *before* the SimulationRunner loop in each example: + spike trains are computed by running ``jx.integrate(cell)`` per motor neuron, + and the resulting spike times are used to drive the muscle/proprioception models + that SimulationRunner steps in real time. This split is intentional — it keeps + neural biophysics (GPU-accelerated, full ODE integration) separate from the + real-time physiological models (CPU, time-stepped). +""" + +from itertools import count +from typing import Any, Callable, Optional, Union + +import numpy as np +import quantities as pq +from neo import AnalogSignal, Block, Segment, SpikeTrain +from tqdm import tqdm + +from myogen.simulator.jaxley.network import Network +from myogen.utils.decorators import beartowertype +from myogen.utils.types import Quantity__ms + + +@beartowertype +class SimulationRunner: + """ + Manages Jaxley simulation execution with automated setup, initialization, + and result collection for neuromuscular simulations. + + Provides a clean interface for running complex neuromuscular simulations + while maintaining full user control over populations, connections, and + step-by-step simulation logic. + + Separates simulation control from plotting and analysis concerns. + """ + + # Smart defaults for common MyoGen model output attributes + _DEFAULT_MODEL_OUTPUTS = { + "HillModel": [ + "muscle_length", + "muscle_force", + "muscle_torque", + "type1_activation", + "type2_activation", + ], + "SpindleModel": [ + "primary_afferent_firing__Hz", + "secondary_afferent_firing__Hz", + "bag1_activation", + "bag2_activation", + "chain_activation", + "intrafusal_tensions", + ], + "GolgiTendonOrganModel": ["ib_afferent_firing__Hz"], + } + + def __init__( + self, + network: Network, + models: dict[str, Any], + step_callback: Callable[[Any], Any], + model_outputs: Optional[dict[str, Union[list[str], None]]] = None, + temperature__celsius: float = 36.0, + ): + """ + Initialize SimulationRunner with network, models, and step callback. + + Parameters + ---------- + network : Network + Configured Network instance with populations and connections. + models : Dict[str, Any] + Physiological models (e.g., {"hill": hill_model, "spin": spindle_model}). + step_callback : Callable + User-defined function called at each simulation timestep. + model_outputs : Optional[Dict[str, Union[List[str], None]]], optional + Explicit model output attributes to collect. None uses smart defaults. + Format: {"model_name": ["attr1", "attr2"]} or {"model_name": None} + for defaults, by default None. + temperature__celsius : float, optional + Simulation temperature, by default 36.0. + """ + # Store immutable parameters following project pattern + self.network = network + self.populations = network.populations # Expose populations from network + self.models = models + self.step_callback = step_callback + self.model_outputs = model_outputs + self.temperature__celsius = temperature__celsius + + # Private working copies + self._network = network + self._populations = network.populations + self._models = models + self._step_callback = step_callback + self._model_outputs = self._resolve_model_outputs() + self._temperature__celsius = temperature__celsius + + # Runtime state + self._trace_vectors: dict[str, dict[int, list]] = {} + self._step_counter = None + self._progress_bar = None + self._total_steps = None + self._current_time = 0.0 + + # Setup internal spike recording + self._spike_recording = self._setup_spike_recording() + + def _resolve_model_outputs(self) -> dict[str, list[str]]: + """ + Resolve model output attributes using smart defaults and user overrides. + + Returns + ------- + Dict[str, List[str]] + Final mapping of model names to output attribute lists. + """ + resolved = {} + + for model_name, model_instance in self._models.items(): + model_class_name = model_instance.__class__.__name__ + + # Check for user override + if self.model_outputs and model_name in self.model_outputs: + user_specified = self.model_outputs[model_name] + if user_specified is None: + # Use defaults for this model + resolved[model_name] = self._DEFAULT_MODEL_OUTPUTS.get(model_class_name, []) + else: + # Use explicit user specification + resolved[model_name] = user_specified + else: + # Use smart defaults based on model class + resolved[model_name] = self._DEFAULT_MODEL_OUTPUTS.get(model_class_name, []) + + return resolved + + def _setup_spike_recording(self) -> dict[str, Any]: + """ + Create spike recording storage for all populations. + + Returns + ------- + dict[str, Any] + Dictionary containing 'idvec' and 'spkvec' with lists for each population. + """ + idvec = {} + spkvec = {} + + for pop_name in self._populations.keys(): + idvec[pop_name] = [] + spkvec[pop_name] = [] + + return {"idvec": idvec, "spkvec": spkvec} + + def _setup_network_spike_recording(self) -> None: + """ + Configure the network with spike recording vectors and activate recording. + """ + # Set spike recording on network + self._network.spike_recording = self._spike_recording + + # Setup spike recording + self._network.setup_spike_recording() + + def run( + self, + duration__ms: Quantity__ms, + timestep__ms: Quantity__ms, + membrane_recording: Optional[dict[str, list[int]]] = None, + ) -> Block: + """ + Execute Jaxley simulation with automated setup and result collection. + + Parameters + ---------- + duration__ms : Quantity__ms + Total simulation duration in milliseconds. + timestep__ms : Quantity__ms + Integration timestep in milliseconds. + membrane_recording : Optional[Dict[str, List[int]]], optional + Populations and cell indices for membrane potential recording. + Format: {"population_name": [cell_id1, cell_id2, ...]}, by default None. + + Returns + ------- + Block + Structured simulation results containing: + - spikes: Spike timing and ID data for all populations + - membrane: Membrane potential traces (if requested) + - models: Output data from all physiological models + - simulation: Time vector and simulation metadata + + Raises + ------ + ValueError + If model output attributes don't exist on model instances. + RuntimeError + If simulation fails to complete. + """ + try: + # Setup simulation environment + self._setup_simulation_environment(duration__ms, timestep__ms) + + # Setup optional membrane recording + if membrane_recording: + self._setup_membrane_recording(membrane_recording) + + # Initialize population voltages + self._initialize_voltages() + + # Validate model outputs before simulation + self._validate_model_outputs() + + # Setup spike recording on network + self._setup_network_spike_recording() + + # Run simulation loop + self._run_simulation_loop(duration__ms, timestep__ms) + + # Close progress bar + if self._progress_bar is not None: + try: + self._progress_bar.close() + except (TypeError, AttributeError): + pass + + print("Simulation completed") + + # Collect and structure results + results = self._collect_results(duration__ms, timestep__ms) + + return results + + except Exception as e: + # Close progress bar in case of error + if self._progress_bar is not None: + try: + self._progress_bar.close() + except (TypeError, AttributeError): + pass + raise RuntimeError(f"Simulation failed: {str(e)}") from e + + def _setup_simulation_environment( + self, duration__ms: Quantity__ms, timestep__ms: Quantity__ms + ) -> None: + """Configure simulation parameters.""" + self._duration__ms = duration__ms + self._timestep__ms = timestep__ms + self._current_time = 0.0 + + # Calculate total steps for progress bar + self._total_steps = int(duration__ms.magnitude / timestep__ms.magnitude) + + # Initialize progress bar + self._progress_bar = tqdm( + total=duration__ms.magnitude, + desc="Simulation Progress", + unit="ms", + ) + + # Reset step counter for step callback + self._step_counter = count(0) + + def _setup_membrane_recording(self, membrane_recording: dict[str, list[int]]) -> None: + """Setup membrane potential recording storage for specified populations.""" + self._trace_vectors = {} + + for pop_name, cell_indices in membrane_recording.items(): + if pop_name not in self._populations: + raise ValueError(f"Population '{pop_name}' not found in populations") + + pop_traces = {} + population = self._populations[pop_name] + + for cell_idx in cell_indices: + if cell_idx >= len(population): + raise ValueError( + f"Cell index {cell_idx} out of range for population " + f"'{pop_name}' (size: {len(population)})" + ) + + # Initialize storage list for this cell's membrane potential + pop_traces[cell_idx] = [] + + self._trace_vectors[pop_name] = pop_traces + + def _initialize_voltages(self) -> None: + """Automatically collect and set initial voltages for all populations.""" + # For Jaxley backend, voltage initialization is handled differently + # Store initialization data for later use + self._initialization_data = {} + + for pop_name, population in self._populations.items(): + try: + cells_to_init, voltages = population.get_initialization_data() + if cells_to_init: + self._initialization_data[pop_name] = { + "cells": cells_to_init, + "voltages": voltages, + } + except (AttributeError, TypeError): + # Skip populations without initialization data + continue + + def _run_simulation_loop(self, duration__ms: Quantity__ms, timestep__ms: Quantity__ms) -> None: + """ + Execute the main simulation time-stepping loop. + + For Jaxley backend, this manually steps through time and calls the user's + step callback at each timestep. + """ + dt = timestep__ms.magnitude + t_stop = duration__ms.magnitude + + while self._current_time < t_stop: + # Call user's step callback + try: + self._step_callback(self._step_counter) + except StopIteration: + break + + # Record membrane potentials if requested + if self._trace_vectors: + self._record_membrane_potentials() + + # Update simulation time + self._current_time += dt + + # Update progress bar + if self._progress_bar is not None: + try: + self._progress_bar.update(dt) + except (TypeError, AttributeError): + self._progress_bar = None + + def _record_membrane_potentials(self) -> None: + """Record membrane potentials for cells with active recording.""" + for pop_name, cell_traces in self._trace_vectors.items(): + population = self._populations[pop_name] + for cell_idx, trace_list in cell_traces.items(): + cell = population[cell_idx] + # Membrane potential comes from cell.v if available. + # Note: this path is only reached if membrane_recording is passed to run(), + # which no current example does. Neural voltage traces are obtained from + # jx.integrate() output, not from stepping through SimulationRunner. + if hasattr(cell, 'v'): + trace_list.append(cell.v) + + def _validate_model_outputs(self) -> None: + """Validate that all specified model output attributes exist.""" + for model_name, output_attrs in self._model_outputs.items(): + if model_name not in self._models: + raise ValueError(f"Model '{model_name}' not found in models") + + model_instance = self._models[model_name] + + for attr_name in output_attrs: + if not hasattr(model_instance, attr_name): + raise ValueError( + f"Model '{model_name}' ({model_instance.__class__.__name__}) " + f"does not have attribute '{attr_name}'" + ) + + def _collect_results(self, duration__ms: Quantity__ms, timestep__ms: Quantity__ms) -> Block: + """ + Collect simulation results from network, models, and recordings. + + Returns structured results compatible with existing analysis code. + """ + block = Block() + + # Collect spike data for each population + for pop_name in self._populations.keys(): + segment = Segment(name=pop_name) + + if self._spike_recording and pop_name in self._spike_recording.get("spkvec", {}): + spike_times = np.array(self._spike_recording["spkvec"][pop_name]) + spike_ids = np.array(self._spike_recording["idvec"][pop_name]) + + if len(spike_times) > 0: + for spike_id in sorted(np.unique(spike_ids)): + times_for_id = spike_times[spike_ids == spike_id] + if len(times_for_id) > 0: + segment.spiketrains.append( + SpikeTrain( + name=f"{pop_name}_cell{int(spike_id)}_spikes", + times=(times_for_id * pq.ms).rescale(pq.s), + t_start=0.0 * pq.s, + t_stop=duration__ms.rescale(pq.s), + sampling_rate=(1.0 / timestep__ms.rescale(pq.s)).rescale(pq.Hz), + cell_idx=int(spike_id), + ) + ) + + # Collect membrane potential traces + for cell_idx, trace_data in self._trace_vectors.get(pop_name, {}).items(): + if trace_data: + segment.analogsignals.append( + AnalogSignal( + name=f"{pop_name}_cell{cell_idx}_Vm", + sampling_period=timestep__ms.rescale(pq.s), + signal=np.array(trace_data) * pq.mV, + cell_idx=cell_idx, + ) + ) + + block.segments.append(segment) + + # Collect model outputs + for model_name, output_attrs in self._model_outputs.items(): + segment = Segment(name=model_name) + + model_instance = self._models[model_name] + for attr_name in output_attrs: + attr_value = getattr(model_instance, attr_name) + + if hasattr(attr_value, "__iter__") and not isinstance(attr_value, str): + segment.analogsignals.append( + AnalogSignal( + name=f"{model_name}_{attr_name}", + sampling_period=timestep__ms.rescale(pq.s), + signal=np.asarray(attr_value) * pq.dimensionless, + attr_name=attr_name, + ) + ) + elif isinstance(attr_value, (int, float, str)): + segment.annotations[attr_name] = attr_value + + block.segments.append(segment) + + # Add metadata + block.annotations["time__ms"] = duration__ms + block.annotations["timestep__ms"] = timestep__ms + block.annotations["temperature__celsius"] = self._temperature__celsius + + # Extract active motor neurons if spike data exists + all_spike_ids = [] + for pop_name in self._populations.keys(): + if pop_name in self._spike_recording.get("idvec", {}): + spike_ids = self._spike_recording["idvec"][pop_name] + if spike_ids: + all_spike_ids.extend(spike_ids) + + if all_spike_ids: + block.annotations["active_MNs"] = np.unique(all_spike_ids).astype(int) + + return block + + def get_model_outputs(self, model_name: str) -> list[str]: + """ + Get the list of output attributes that will be collected for a model. + + Parameters + ---------- + model_name : str + Name of the model as specified in the models dictionary. + + Returns + ------- + List[str] + List of attribute names that will be collected from this model. + """ + return self._model_outputs.get(model_name, []) + + def set_model_outputs(self, model_name: str, output_attrs: list[str]) -> None: + """ + Override the output attributes for a specific model. + + Parameters + ---------- + model_name : str + Name of the model as specified in the models dictionary. + output_attrs : List[str] + List of attribute names to collect from this model. + + Raises + ------ + ValueError + If model_name is not found in the models dictionary. + """ + if model_name not in self._models: + raise ValueError(f"Model '{model_name}' not found in models") + + self._model_outputs[model_name] = output_attrs diff --git a/myogen/simulator/jaxley/synapses/__init__.py b/myogen/simulator/jaxley/synapses/__init__.py new file mode 100644 index 00000000..c72230ac --- /dev/null +++ b/myogen/simulator/jaxley/synapses/__init__.py @@ -0,0 +1,15 @@ +"""Synaptic mechanisms for Jaxley simulations.""" + +from myogen.simulator.jaxley.synapses.exp2syn import ( + Exp2Syn, + ExcitatorySynapse, + InhibitorySynapse, + NMDASynapse, +) + +__all__ = [ + "Exp2Syn", + "ExcitatorySynapse", + "InhibitorySynapse", + "NMDASynapse", +] diff --git a/myogen/simulator/jaxley/synapses/exp2syn.py b/myogen/simulator/jaxley/synapses/exp2syn.py new file mode 100644 index 00000000..f3fe3498 --- /dev/null +++ b/myogen/simulator/jaxley/synapses/exp2syn.py @@ -0,0 +1,351 @@ +""" +Double-Exponential Synapse (Exp2Syn) for Jaxley. + +This module provides Jaxley-compatible implementations of synaptic conductances +following the double-exponential kinetics used in NEURON's Exp2Syn mechanism. + +The synaptic conductance follows: + g(t) = weight * factor * (exp(-t/tau2) - exp(-t/tau1)) + +where factor normalizes the peak conductance. +""" + +from typing import Dict, Optional, Tuple +import jax.numpy as jnp +from jaxley.synapses import Synapse + + +def safe_exp(x: jnp.ndarray, clip_val: float = 100.0) -> jnp.ndarray: + """Exponential with clipping to prevent overflow.""" + return jnp.exp(jnp.clip(x, -clip_val, clip_val)) + + +class Exp2Syn(Synapse): + """ + Double-exponential synapse (Exp2Syn). + + Implements the standard NEURON Exp2Syn mechanism with: + - Rise time constant (tau1) + - Decay time constant (tau2) + - Reversal potential (e) + + The conductance waveform is: + g(t) = weight * factor * (exp(-t/tau2) - exp(-t/tau1)) + + where factor = 1 / (peak_norm) normalizes peak conductance to weight. + + Parameters + ---------- + name : str, optional + Name of the synapse instance. + tau1 : float + Rise time constant in ms. Default: 0.2 ms. + tau2 : float + Decay time constant in ms. Must be > tau1. Default: 2.0 ms. + e : float + Reversal potential in mV. Default: 0.0 mV (excitatory). + """ + + def __init__( + self, + name: Optional[str] = None, + tau1: float = 0.2, # ms - rise time + tau2: float = 2.0, # ms - decay time + e: float = 0.0, # mV - reversal potential + ): + super().__init__(name=name or "Exp2Syn") + + if tau2 <= tau1: + raise ValueError(f"tau2 ({tau2}) must be greater than tau1 ({tau1})") + + self.tau1 = tau1 + self.tau2 = tau2 + self.e = e + + # Compute normalization factor so peak conductance equals weight + self._compute_factor() + + def _compute_factor(self): + """Compute factor to normalize peak conductance.""" + # Time to peak: tp = (tau1 * tau2 / (tau2 - tau1)) * ln(tau2/tau1) + tp = (self.tau1 * self.tau2 / (self.tau2 - self.tau1)) * jnp.log(self.tau2 / self.tau1) + + # Peak value of unnormalized waveform + peak = jnp.exp(-tp / self.tau2) - jnp.exp(-tp / self.tau1) + + # Factor to normalize peak to 1 + self.factor = 1.0 / peak if peak > 0 else 1.0 + + @property + def synapse_params(self) -> Dict[str, float]: + """Parameters that can be set per-synapse.""" + return { + f"{self.name}_tau1": self.tau1, + f"{self.name}_tau2": self.tau2, + f"{self.name}_e": self.e, + f"{self.name}_g": 0.0, # Conductance (set by weight) + } + + @property + def synapse_states(self) -> Dict[str, float]: + """State variables for synaptic dynamics.""" + return { + f"{self.name}_A": 0.0, # Rising phase variable + f"{self.name}_B": 0.0, # Decaying phase variable + } + + def update_states( + self, + states: Dict[str, jnp.ndarray], + delta_t: float, + pre_voltage: jnp.ndarray, + post_voltage: jnp.ndarray, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """ + Update synaptic state variables. + + The conductance g = B - A follows double-exponential kinetics: + - dA/dt = -A/tau1 + - dB/dt = -B/tau2 + + On presynaptic spike: A += weight*factor, B += weight*factor + """ + prefix = self.name + A = states[f"{prefix}_A"] + B = states[f"{prefix}_B"] + + tau1 = params[f"{prefix}_tau1"] + tau2 = params[f"{prefix}_tau2"] + + # Exponential decay + A_new = A * safe_exp(-delta_t / tau1) + B_new = B * safe_exp(-delta_t / tau2) + + return { + f"{prefix}_A": A_new, + f"{prefix}_B": B_new, + } + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + pre_voltage: jnp.ndarray, + post_voltage: jnp.ndarray, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + """ + Compute synaptic current. + + I = g * (B - A) * (v - e) + + where (B - A) gives the double-exponential conductance waveform. + """ + prefix = self.name + A = states[f"{prefix}_A"] + B = states[f"{prefix}_B"] + + e = params[f"{prefix}_e"] + g = params[f"{prefix}_g"] + + # Conductance = B - A (normalized double exponential) + g_syn = g * self.factor * (B - A) + + return g_syn * (post_voltage - e) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + pre_voltage: jnp.ndarray, + post_voltage: jnp.ndarray, + params: Dict[str, jnp.ndarray], + delta_t: float, + ) -> Dict[str, jnp.ndarray]: + """Initialize to resting state (no active synaptic input).""" + prefix = self.name + return { + f"{prefix}_A": jnp.zeros_like(post_voltage), + f"{prefix}_B": jnp.zeros_like(post_voltage), + } + + +class ExcitatorySynapse(Exp2Syn): + """ + Excitatory synapse (AMPA-like). + + Fast kinetics with reversal at 0 mV. + + Parameters + ---------- + name : str, optional + Name of the synapse instance. + tau1 : float + Rise time constant in ms. Default: 0.2 ms. + tau2 : float + Decay time constant in ms. Default: 2.0 ms. + """ + + def __init__( + self, + name: Optional[str] = None, + tau1: float = 0.2, + tau2: float = 2.0, + ): + super().__init__( + name=name or "AMPA", + tau1=tau1, + tau2=tau2, + e=0.0, # Excitatory reversal + ) + + +class InhibitorySynapse(Exp2Syn): + """ + Inhibitory synapse (GABA-A like). + + Slower kinetics with reversal at -75 mV. + + Parameters + ---------- + name : str, optional + Name of the synapse instance. + tau1 : float + Rise time constant in ms. Default: 0.5 ms. + tau2 : float + Decay time constant in ms. Default: 10.0 ms. + """ + + def __init__( + self, + name: Optional[str] = None, + tau1: float = 0.5, + tau2: float = 10.0, + ): + super().__init__( + name=name or "GABA", + tau1=tau1, + tau2=tau2, + e=-75.0, # Inhibitory reversal + ) + + +class NMDASynapse(Synapse): + """ + NMDA synapse with voltage-dependent Mg2+ block. + + Implements: + - Slow kinetics (tau1 ~ 2 ms, tau2 ~ 100 ms) + - Voltage-dependent Mg2+ block: B(V) = 1 / (1 + [Mg]/3.57 * exp(-0.062*V)) + - Mixed Na+/Ca2+ reversal at 0 mV + + Parameters + ---------- + name : str, optional + Name of the synapse instance. + tau1 : float + Rise time constant in ms. Default: 2.0 ms. + tau2 : float + Decay time constant in ms. Default: 100.0 ms. + mg_conc : float + Extracellular Mg2+ concentration in mM. Default: 1.0 mM. + """ + + def __init__( + self, + name: Optional[str] = None, + tau1: float = 2.0, + tau2: float = 100.0, + mg_conc: float = 1.0, + ): + super().__init__(name=name or "NMDA") + + if tau2 <= tau1: + raise ValueError(f"tau2 ({tau2}) must be greater than tau1 ({tau1})") + + self.tau1 = tau1 + self.tau2 = tau2 + self.mg_conc = mg_conc + self.e = 0.0 # NMDA reversal + + # Compute normalization factor + tp = (self.tau1 * self.tau2 / (self.tau2 - self.tau1)) * jnp.log(self.tau2 / self.tau1) + peak = jnp.exp(-tp / self.tau2) - jnp.exp(-tp / self.tau1) + self.factor = 1.0 / peak if peak > 0 else 1.0 + + @property + def synapse_params(self) -> Dict[str, float]: + return { + f"{self.name}_tau1": self.tau1, + f"{self.name}_tau2": self.tau2, + f"{self.name}_mg": self.mg_conc, + f"{self.name}_e": self.e, + f"{self.name}_g": 0.0, + } + + @property + def synapse_states(self) -> Dict[str, float]: + return { + f"{self.name}_A": 0.0, + f"{self.name}_B": 0.0, + } + + def _mg_block(self, v: jnp.ndarray, mg: float) -> jnp.ndarray: + """Voltage-dependent Mg2+ block factor.""" + return 1.0 / (1.0 + mg / 3.57 * safe_exp(-0.062 * v)) + + def update_states( + self, + states: Dict[str, jnp.ndarray], + delta_t: float, + pre_voltage: jnp.ndarray, + post_voltage: jnp.ndarray, + params: Dict[str, jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + prefix = self.name + A = states[f"{prefix}_A"] + B = states[f"{prefix}_B"] + + tau1 = params[f"{prefix}_tau1"] + tau2 = params[f"{prefix}_tau2"] + + A_new = A * safe_exp(-delta_t / tau1) + B_new = B * safe_exp(-delta_t / tau2) + + return { + f"{prefix}_A": A_new, + f"{prefix}_B": B_new, + } + + def compute_current( + self, + states: Dict[str, jnp.ndarray], + pre_voltage: jnp.ndarray, + post_voltage: jnp.ndarray, + params: Dict[str, jnp.ndarray], + ) -> jnp.ndarray: + prefix = self.name + A = states[f"{prefix}_A"] + B = states[f"{prefix}_B"] + + e = params[f"{prefix}_e"] + g = params[f"{prefix}_g"] + mg = params[f"{prefix}_mg"] + + # Conductance with Mg block + g_syn = g * self.factor * (B - A) * self._mg_block(post_voltage, mg) + + return g_syn * (post_voltage - e) + + def init_state( + self, + states: Dict[str, jnp.ndarray], + pre_voltage: jnp.ndarray, + post_voltage: jnp.ndarray, + params: Dict[str, jnp.ndarray], + delta_t: float, + ) -> Dict[str, jnp.ndarray]: + prefix = self.name + return { + f"{prefix}_A": jnp.zeros_like(post_voltage), + f"{prefix}_B": jnp.zeros_like(post_voltage), + } diff --git a/myogen/utils/continuous_saver.py b/myogen/utils/continuous_saver.py index 7f52e2d6..a54b40bc 100644 --- a/myogen/utils/continuous_saver.py +++ b/myogen/utils/continuous_saver.py @@ -16,9 +16,12 @@ # NEO imports for standard output format import quantities as pq from neo import AnalogSignal, Block, Segment, SpikeTrain -from neuron import h from tqdm import tqdm +# NOTE: ``ContinuousSaver`` is a NEURON step-callback utility. NEURON (``h``) is +# imported lazily inside the methods that need the simulation clock so that +# ``import myogen`` / ``myogen.utils`` works in a Jaxley-only install (no NEURON). + from myogen.utils.types import Quantity__ms @@ -82,6 +85,7 @@ def record_step(self, timestep__ms: float) -> None: timestep__ms : float Integration timestep in milliseconds """ + from neuron import h current_time = h.t # Record current time @@ -159,6 +163,7 @@ def _save_current_chunk(self, timestep__ms: float) -> None: self.current_chunk_data.clear() self.current_chunk_times.clear() self.chunk_id += 1 + from neuron import h self.last_save_time = h.t def finalize(self, timestep__ms: Quantity__ms, spike_results=None) -> None: diff --git a/myogen/utils/dd_cache.py b/myogen/utils/dd_cache.py new file mode 100644 index 00000000..254d6148 --- /dev/null +++ b/myogen/utils/dd_cache.py @@ -0,0 +1,140 @@ +"""Per-seed cache for descending-drive Poisson spike trains. + +The per-step Python loop that ticks every DD cell once per timestep is the +single dominant cost of the ex03 scripts on long simulations +(approximately ~10 of 15 min on a 32 s run). It's also the same loop on +both backends, so for paired NEURON / Jaxley validation runs it's wasted +work to regenerate it twice for the same seed. + +This module provides a tiny load-or-save cache keyed by the parameters +that actually change the realised spike pattern: + + (seed, drive_shape, simulation_time_ms, n_dd, dt_ms) + +If two runs share all five values, they share the cached spike-time list. +The cache file is a single ``joblib.dump`` of ``dd_spike_times``: a list +with one entry per DD cell, each entry a Python list of spike times in +milliseconds. + +Typical usage at the top of an ex03 script:: + + from myogen.utils.dd_cache import dd_cache_path, load_dd_cache, save_dd_cache + + dd_spike_times = None + if args.seed is not None: + dd_spike_times = load_dd_cache( + seed=args.seed, drive_shape=DRIVE_SHAPE, + simulation_time_ms=simulation_time_ms, + n_dd=len(descending_drive_pool), dt_ms=dt_ms, + cache_dir=save_path / "dd_cache", + ) + if dd_spike_times is None: + dd_spike_times = [[] for _ in range(len(descending_drive_pool))] + ... # original Python loop + if args.seed is not None: + save_dd_cache( + dd_spike_times, seed=args.seed, drive_shape=DRIVE_SHAPE, + simulation_time_ms=simulation_time_ms, + n_dd=len(descending_drive_pool), dt_ms=dt_ms, + cache_dir=save_path / "dd_cache", + ) + +The same seed and the same parameters produce the same spike train on both +backends, so for the validation runner the NEURON pass primes the cache and +the Jaxley pass reads it — pairing is guaranteed by construction, not by +hoping the two backends share an RNG path. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Optional + +import joblib + + +def _fmt_dt(dt_ms: float) -> str: + """Render dt_ms safely as a filename fragment (e.g. 0.1 -> '0p1').""" + return f"{dt_ms:g}".replace(".", "p") + + +def dd_cache_path( + *, + seed: int, + drive_shape: str, + simulation_time_ms: float, + n_dd: int, + dt_ms: float, + cache_dir: Path, +) -> Path: + """Deterministic cache filename for one (seed, params) combination. + + A different drive_shape, duration, n_dd or dt_ms yields a different + filename, so changing any of those invalidates the cache automatically. + """ + fname = ( + f"dd_cache_seed{int(seed):02d}" + f"_{drive_shape}" + f"_dur{int(round(simulation_time_ms))}ms" + f"_n{int(n_dd)}" + f"_dt{_fmt_dt(float(dt_ms))}ms.pkl" + ) + return Path(cache_dir) / fname + + +def load_dd_cache( + *, + seed: int, + drive_shape: str, + simulation_time_ms: float, + n_dd: int, + dt_ms: float, + cache_dir: Path, +) -> Optional[list]: + """Return cached ``dd_spike_times`` or ``None`` if absent. + + Returns ``None`` rather than raising so the caller can fall back to + generating fresh. + """ + p = dd_cache_path( + seed=seed, drive_shape=drive_shape, + simulation_time_ms=simulation_time_ms, + n_dd=n_dd, dt_ms=dt_ms, cache_dir=cache_dir, + ) + if not p.exists(): + return None + try: + dd_spike_times = joblib.load(p) + except Exception as exc: + print(f"[dd_cache] WARNING: failed to load {p}: {exc}; regenerating.") + return None + if not isinstance(dd_spike_times, list) or len(dd_spike_times) != n_dd: + print( + f"[dd_cache] WARNING: cache at {p} has {len(dd_spike_times) if isinstance(dd_spike_times, list) else '?'} entries, " + f"expected {n_dd}; regenerating." + ) + return None + print(f"[dd_cache] HIT: loaded DD spikes from {p.name}") + return dd_spike_times + + +def save_dd_cache( + dd_spike_times: list, + *, + seed: int, + drive_shape: str, + simulation_time_ms: float, + n_dd: int, + dt_ms: float, + cache_dir: Path, +) -> Path: + """Persist ``dd_spike_times`` under the canonical cache filename.""" + p = dd_cache_path( + seed=seed, drive_shape=drive_shape, + simulation_time_ms=simulation_time_ms, + n_dd=n_dd, dt_ms=dt_ms, cache_dir=cache_dir, + ) + p.parent.mkdir(parents=True, exist_ok=True) + joblib.dump(dd_spike_times, p) + print(f"[dd_cache] MISS: saved DD spikes to {p.name}") + return p diff --git a/myogen/utils/jaxley/__init__.py b/myogen/utils/jaxley/__init__.py new file mode 100644 index 00000000..ac67a772 --- /dev/null +++ b/myogen/utils/jaxley/__init__.py @@ -0,0 +1,18 @@ +""" +Utility functions for Jaxley simulations. + +This module provides high-level convenience functions that mirror +the NEURON utilities for seamless workflow compatibility. +""" + +from myogen.utils.jaxley.inject_currents import ( + inject_currents_and_simulate_spike_trains, + inject_currents_into_populations, + simulation_result_to_neo_block, +) + +__all__ = [ + "inject_currents_and_simulate_spike_trains", + "inject_currents_into_populations", + "simulation_result_to_neo_block", +] diff --git a/myogen/utils/jaxley/inject_currents.py b/myogen/utils/jaxley/inject_currents.py new file mode 100644 index 00000000..c31ce783 --- /dev/null +++ b/myogen/utils/jaxley/inject_currents.py @@ -0,0 +1,353 @@ +""" +Current Injection and Spike Train Simulation Utilities for Jaxley. + +This module provides high-level functions for simulating populations of +neurons with current injection, mirroring the NEURON workflow. + +Functions +--------- +inject_currents_and_simulate_spike_trains + High-level function to inject currents and record spike trains. +inject_currents_into_populations + Set up current injection for populations. +simulation_result_to_neo_block + Convert simulation results to Neo Block format. +""" + +from typing import List, Optional, Sequence, Union + +import numpy as np +import jax.numpy as jnp +import quantities as pq +from neo import Block, Segment, SpikeTrain, AnalogSignal + +import jaxley as jx + +from myogen.simulator.jaxley.populations.base import _Pool +from myogen.simulator.jaxley.integrate import ( + detect_spikes, + compute_firing_rate_from_spikes, + compute_isi_cv, + custom_current, +) + + +# Type aliases for clarity +CURRENT__AnalogSignal = AnalogSignal +SPIKE_TRAIN__Block = Block + + +def inject_currents_into_populations( + populations: Sequence[_Pool], + input_current__AnalogSignal: CURRENT__AnalogSignal, + stim_branch: int = 0, + stim_loc: float = 0.5, +) -> None: + """ + Set up current injection for populations of Jaxley cells. + + Parameters + ---------- + populations : Sequence[_Pool] + List of population objects (e.g., AlphaMN__Pool). + input_current__AnalogSignal : neo.AnalogSignal + Input current signal with shape (n_samples, n_pools). + Each column corresponds to a population. + stim_branch : int + Branch index for current injection (default: 0, soma). + stim_loc : float + Location on branch for injection (0.0 to 1.0). + """ + # Extract timestep and convert current to numpy + dt = float(input_current__AnalogSignal.sampling_period.rescale(pq.ms).magnitude) + current_data = np.array(input_current__AnalogSignal.magnitude) + + # Handle 1D vs 2D current arrays + if current_data.ndim == 1: + current_data = current_data[:, np.newaxis] + + n_pools = len(populations) + + for pool_idx, pool in enumerate(populations): + # Get current for this pool (use modulo to handle fewer columns than pools) + pool_current = current_data[:, pool_idx % current_data.shape[1]] + pool_current_jnp = jnp.array(pool_current) + + # Inject current into each cell in the pool + for cell_wrapper in pool: + if hasattr(cell_wrapper, 'cell'): + cell = cell_wrapper.cell + # Clear existing stimuli + cell.delete_stimuli() + # Set up stimulation + cell.branch(stim_branch).loc(stim_loc).stimulate(pool_current_jnp) + + +def inject_currents_and_simulate_spike_trains( + populations: Sequence[_Pool], + input_current__AnalogSignal: CURRENT__AnalogSignal, + spike_detection_thresholds__mV: Union[pq.Quantity, float] = 50.0 * pq.mV, + stim_branch: int = 0, + stim_loc: float = 0.5, + record_branch: int = 0, + record_loc: float = 0.5, +) -> SPIKE_TRAIN__Block: + """ + Inject currents into populations and simulate to produce spike trains. + + This is the main high-level function for running Jaxley simulations. + It mirrors the NEURON utility function for API compatibility. + + Parameters + ---------- + populations : Sequence[_Pool] + List of population objects (e.g., AlphaMN__Pool). + input_current__AnalogSignal : neo.AnalogSignal + Input current signal with shape (n_samples, n_pools). + spike_detection_thresholds__mV : pq.Quantity or float + Threshold for spike detection. Can be a single value or per-population. + stim_branch : int + Branch index for current injection (default: 0). + stim_loc : float + Location on branch for injection (0.0 to 1.0). + record_branch : int + Branch index for voltage recording (default: 0). + record_loc : float + Location on branch for recording (0.0 to 1.0). + + Returns + ------- + neo.Block + Block containing spike trains for each population (as segments). + """ + # Extract simulation parameters + dt = float(input_current__AnalogSignal.sampling_period.rescale(pq.ms).magnitude) + t_max = float(input_current__AnalogSignal.t_stop.rescale(pq.ms).magnitude) + n_samples = len(input_current__AnalogSignal) + + # Handle threshold + if hasattr(spike_detection_thresholds__mV, 'magnitude'): + threshold = float(spike_detection_thresholds__mV.rescale(pq.mV).magnitude) + else: + threshold = float(spike_detection_thresholds__mV) + + # Get current data + current_data = np.array(input_current__AnalogSignal.magnitude) + if current_data.ndim == 1: + current_data = current_data[:, np.newaxis] + + # Create output Block + spike_train__Block = Block(name="Jaxley Simulation Results") + + # Simulate each population + for pool_idx, pool in enumerate(populations): + # Get current for this pool + pool_current = current_data[:, pool_idx % current_data.shape[1]] + pool_current_jnp = jnp.array(pool_current) + + # Create segment for this pool + segment = Segment(name=f"Pool {pool_idx}") + segment.spiketrains = [] + + # Simulate each cell in the pool + for neuron_idx, cell_wrapper in enumerate(pool): + spike_times_ms = _simulate_single_cell( + cell_wrapper, + pool_current_jnp, + dt, + t_max, + threshold, + stim_branch, + stim_loc, + record_branch, + record_loc, + ) + + # Convert to Neo SpikeTrain + spike_times_s = spike_times_ms / 1000.0 # ms to s + + spiketrain = SpikeTrain( + spike_times_s * pq.s, + t_stop=(t_max / 1000.0) * pq.s, + sampling_rate=(1000.0 / dt) * pq.Hz, + name=str(neuron_idx), + description=f"Pool {pool_idx}, Neuron {neuron_idx}", + ) + segment.spiketrains.append(spiketrain) + + spike_train__Block.segments.append(segment) + + return spike_train__Block + + +def _simulate_single_cell( + cell_wrapper, + current: jnp.ndarray, + dt: float, + t_max: float, + threshold: float, + stim_branch: int, + stim_loc: float, + record_branch: int, + record_loc: float, +) -> np.ndarray: + """ + Simulate a single cell and return spike times. + + Parameters + ---------- + cell_wrapper : cell object + Cell wrapper with .cell attribute containing jx.Cell. + current : jnp.ndarray + Current waveform to inject. + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + threshold : float + Spike detection threshold (mV). + stim_branch : int + Branch for stimulation. + stim_loc : float + Location for stimulation. + record_branch : int + Branch for recording. + record_loc : float + Location for recording. + + Returns + ------- + np.ndarray + Spike times in ms. + """ + if not hasattr(cell_wrapper, 'cell'): + return np.array([]) + + cell = cell_wrapper.cell + + try: + # Clear previous recordings and stimuli + cell.delete_recordings() + cell.delete_stimuli() + + # Set up recording + cell.branch(record_branch).loc(record_loc).record("v") + + # Set up stimulation + cell.branch(stim_branch).loc(stim_loc).stimulate(current) + + # Run simulation + voltages = jx.integrate(cell, delta_t=dt, t_max=t_max) + + # Extract voltage trace for the single recorded compartment. + # jx.integrate() returns shape (n_timepoints, n_recorded_compartments). + # [:, 0] selects all timepoints for compartment index 0. + if voltages.ndim == 2: + v = np.array(voltages[:, 0]) + else: + v = np.array(voltages) + + # Detect spikes + spike_indices = np.where((v[:-1] < threshold) & (v[1:] >= threshold))[0] + spike_times = spike_indices * dt + + return spike_times + + except Exception as e: + # If simulation fails, return empty spike train + print(f"Warning: Simulation failed for cell: {e}") + return np.array([]) + + +def simulation_result_to_neo_block( + voltages: np.ndarray, + spike_times_list: List[List[float]], + dt: float, + t_max: float, + pool_names: Optional[List[str]] = None, +) -> Block: + """ + Convert simulation results to Neo Block format. + + Parameters + ---------- + voltages : np.ndarray + Voltage traces, shape (n_cells, n_timesteps) or (n_timesteps,). + spike_times_list : List[List[float]] + List of spike times for each cell (in ms). + dt : float + Time step (ms). + t_max : float + Total simulation time (ms). + pool_names : List[str], optional + Names for each pool/segment. + + Returns + ------- + neo.Block + Block with spike trains. + """ + block = Block(name="Jaxley Simulation Results") + + # Handle different input formats + if isinstance(spike_times_list[0], (list, np.ndarray)): + # Multiple cells + n_cells = len(spike_times_list) + else: + # Single cell - wrap in list + spike_times_list = [spike_times_list] + n_cells = 1 + + # Create single segment with all spike trains + segment = Segment(name=pool_names[0] if pool_names else "Pool 0") + segment.spiketrains = [] + + for i, spike_times_ms in enumerate(spike_times_list): + spike_times_s = np.array(spike_times_ms) / 1000.0 + + spiketrain = SpikeTrain( + spike_times_s * pq.s, + t_stop=(t_max / 1000.0) * pq.s, + sampling_rate=(1000.0 / dt) * pq.Hz, + name=str(i), + ) + segment.spiketrains.append(spiketrain) + + block.segments.append(segment) + + return block + + +def create_input_current_from_waveform( + waveform: np.ndarray, + dt_ms: float, + n_pools: int = 1, +) -> AnalogSignal: + """ + Create Neo AnalogSignal from a current waveform array. + + Parameters + ---------- + waveform : np.ndarray + Current waveform in nA, shape (n_samples,) or (n_samples, n_pools). + dt_ms : float + Time step in ms. + n_pools : int + Number of pools (columns) to create if waveform is 1D. + + Returns + ------- + neo.AnalogSignal + Current signal compatible with inject_currents_and_simulate_spike_trains. + """ + if waveform.ndim == 1: + # Replicate for n_pools + waveform = np.tile(waveform[:, np.newaxis], (1, n_pools)) + + signal = AnalogSignal( + waveform * pq.nA, + sampling_period=dt_ms * pq.ms, + t_start=0 * pq.ms, + ) + + return signal diff --git a/myogen/utils/nmodl.py b/myogen/utils/nmodl.py index a896dda9..8d90478c 100644 --- a/myogen/utils/nmodl.py +++ b/myogen/utils/nmodl.py @@ -301,7 +301,13 @@ def error(msg): return True except ImportError as e: - return error(f"NEURON not available, cannot load mechanisms: {str(e)}") + # NEURON simply not installed — expected in Jaxley-only installs. Only warn + # when not quiet (or raise under strict); never crash the import of myogen. + if strict: + raise NMODLLoadError(f"NEURON not available, cannot load mechanisms: {str(e)}") + if not quiet: + print(f"WARNING: NEURON not available, cannot load mechanisms: {str(e)}") + return False except Exception as e: return error(f"Failed to load NMODL mechanisms: {str(e)}") diff --git a/myogen/utils/plotting/neuron.py b/myogen/utils/plotting/neuron.py index 752e15fc..de97ab44 100644 --- a/myogen/utils/plotting/neuron.py +++ b/myogen/utils/plotting/neuron.py @@ -100,9 +100,10 @@ def plot_raster_spikes( spike_times.append(float(spike_time.rescale("ms").magnitude)) spike_ids.append(i) - if spike_times: - spike_times = np.array(spike_times) - spike_ids = np.array(spike_ids) + spike_times = np.array(spike_times) + spike_ids = np.array(spike_ids) + + if spike_times.size > 0: # Apply time range filter if specified if time_range is not None: @@ -720,6 +721,8 @@ def plot_spindle_dynamics( "Bag2": signal_array[:, 1], "Chain": signal_array[:, 2], } + elif signal_name == "gamma_dynamic": + spindle_data["gamma_dynamic"] = signal_array.flatten() elif signal_name == "primary_afferent_firing__Hz": if "aff" not in spindle_data: spindle_data["aff"] = {} @@ -754,6 +757,8 @@ def plot_spindle_dynamics( elif key in ["act", "tension", "aff"]: for subkey, subvalue in value.items(): spindle_data[key][subkey] = subvalue[time_mask] + elif key == "gamma_dynamic" and isinstance(value, np.ndarray): + spindle_data[key] = value[time_mask] # Plot signals plot_idx = 0 @@ -783,9 +788,30 @@ def plot_spindle_dynamics( if act_type in spindle_data["act"]: ax.plot(plot_time, spindle_data["act"][act_type], label=act_type, **kwargs) + # Overlay gamma fusimotor drive on a secondary y-axis (right side, green) + if "gamma_dynamic" in spindle_data: + ax2 = ax.twinx() + ax2.plot( + plot_time, + spindle_data["gamma_dynamic"], + color="green", + linestyle="--", + linewidth=1.2, + alpha=0.7, + label="γ_dyn [pps]", + ) + ax2.set_ylabel("γ drive [pps]", color="green") + ax2.tick_params(axis="y", labelcolor="green") + ax2.set_ylim(bottom=0) + # Merge legends from both axes + lines1, labels1 = ax.get_legend_handles_labels() + lines2, labels2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, labels1 + labels2, loc="upper left") + if apply_default_formatting: ax.set_ylabel("Activation [a.u.]") - ax.legend() + if "gamma_dynamic" not in spindle_data: + ax.legend() if plot_idx == 0: ax.set_title(title) if plot_idx == len(ax_list) - 1: diff --git a/properdocs.yml b/properdocs.yml index 95bc47f0..efec271a 100644 --- a/properdocs.yml +++ b/properdocs.yml @@ -132,6 +132,7 @@ nav: - api/index.md - Top-level: api/myogen.md - Simulator: api/simulator.md + - Differentiable (Jaxley): api/differentiable.md - Currents: api/currents.md - Injection & I/O: api/io-and-neuron.md - Plotting: api/plotting.md diff --git a/pyproject.toml b/pyproject.toml index 15214157..a9556694 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,13 +30,13 @@ classifiers = [ ] dependencies = [ "beartype>=0.21.0", + "jax==0.9.0", + "jaxley==0.13.0", + "jaxley-mech==0.3.1", "joblib>=1.3", "psutil>=5.9", "threadpoolctl>=3.1", "matplotlib>=3.10.1", - "neuron==8.2.7 ; sys_platform == 'linux' or sys_platform == 'darwin'", - #"neuron-gpu-nightly>=9.0a0", - #"neuron-nightly>=9.0a1.dev1492", "neo>=0.14.0", "numba>=0.61.2", "numpy>=1.26", @@ -59,6 +59,13 @@ nwb = [ "nwbinspector>=0.5.0", "h5py>=3.0", ] +# Optional NEURON backend (validation reference). Jaxley is the default backend and +# requires none of these. Install with: pip install "myogen[neuron]" +neuron = [ + "neuron==8.2.7 ; sys_platform == 'linux' or sys_platform == 'darwin'", + "mpi4py>=4.0.3", + "impi-rt>=2021.15.0 ; sys_platform == 'win32'", +] gpu = [ "cupy-cuda12x>=13.0,<14", ] @@ -81,6 +88,11 @@ dev = [ "poethepoet>=0.37.0", "pytest>=8.0", "scipy-stubs>=1.16.1.0", + # experiments/twin harness, exercised by tests/test_twin_harness.py: + # optax drives the inverse-problem descent, pyyaml loads the frozen + # acceptance spec. Not runtime deps of myogen itself. + "optax>=0.2.4", + "pyyaml>=6.0", ] docs = [ "properdocs>=1.6.7", diff --git a/scripts/plot_differentiable_twin_figure1.py b/scripts/plot_differentiable_twin_figure1.py new file mode 100644 index 00000000..dd0d894d --- /dev/null +++ b/scripts/plot_differentiable_twin_figure1.py @@ -0,0 +1,524 @@ +#!/usr/bin/env python3 +"""Build Paper Figure 1: cross-simulator validation of the forward twin. + +This figure deliberately uses the scientific *hard-spike forward* outputs. +Gradient qualification belongs to the inverse-problem figure, not to the +NEURON-versus-Jaxley forward-validation figure. + +Inputs are the immutable paired-validation artifacts and the confirmatory +Phase-1C CaL x NaP grid already saved under ``experiments/twin/results``. +No simulation is rerun by this script. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import joblib +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import ListedColormap +from matplotlib.lines import Line2D + + +ROOT = Path(__file__).resolve().parents[1] +TWIN_RESULTS = ROOT / "experiments" / "twin" / "results" +DEFAULT_OUTPUT = ROOT / "results" / "figures" / "figure1_forward_twin_validation" + +JAXLEY = "#0072B2" +NEURON = "#E69F00" +GNAP1 = "#009E73" +GNAP5 = "#CC79A7" +DRIVE = "#777777" +GRID = "#D9D9D9" +VALID = "#4DAF7C" +UNRESOLVED = "#F0B35A" + + +def smooth_population_rate( + spikes: np.ndarray, + dt_ms: float, + bin_ms: float = 100.0, + smooth_ms: float = 300.0, +) -> tuple[np.ndarray, np.ndarray]: + """Return time and smoothed mean firing rate across motor units.""" + n_per_bin = int(round(bin_ms / dt_ms)) + n_bins = spikes.shape[0] // n_per_bin + trimmed = spikes[: n_bins * n_per_bin] + counts = trimmed.reshape(n_bins, n_per_bin, spikes.shape[1]).sum(axis=1) + rate = counts.mean(axis=1) / (bin_ms * 1e-3) + + sigma = max(smooth_ms / bin_ms, 1.0) + radius = int(np.ceil(3 * sigma)) + x = np.arange(-radius, radius + 1) + kernel = np.exp(-0.5 * (x / sigma) ** 2) + kernel /= kernel.sum() + rate = np.convolve(rate, kernel, mode="same") + time_s = (np.arange(n_bins) + 0.5) * bin_ms * 1e-3 + return time_s, rate + + +def panel_label(ax: mpl.axes.Axes, label: str, x: float = -0.13) -> None: + ax.text( + x, + 1.08, + label, + transform=ax.transAxes, + fontsize=11, + fontweight="bold", + va="bottom", + ha="left", + ) + + +def mean_metric(points: list[dict], gcal: float, gnap: float, key: str) -> float: + vals = [ + p["metrics"][key] + for p in points + if p["gcal"] == gcal and p["gnap"] == gnap + ] + return float(np.mean(vals)) + + +def build_figure( + firing_artifact: Path, + jaxley_result: Path, + neuron_result: Path, + comparison_json: Path, + grid_json: Path, + output_prefix: Path, +) -> dict: + artifact = joblib.load(firing_artifact) + jax = joblib.load(jaxley_result) + nrn = joblib.load(neuron_result) + comparison = json.loads(comparison_json.read_text()) + grid = json.loads(grid_json.read_text()) + + expected_hash = artifact["content_hash"] + if jax["content_hash"] != expected_hash or nrn["content_hash"] != expected_hash: + raise ValueError("Paired-validation files do not share one content hash") + if comparison["comparison"]["content_hash"] != expected_hash: + raise ValueError("Comparison JSON does not match paired raw outputs") + + dt_ms = float(artifact["dt_ms"]) + dt_s = dt_ms * 1e-3 + pj, pn = jax["pool"], nrn["pool"] + sj, sn = np.asarray(pj["spikes"]), np.asarray(pn["spikes"]) + points = grid["points"] + acceptance = grid["acceptance"] + + mpl.rcParams.update( + { + "font.family": "Arial", + "font.size": 7.5, + "axes.titlesize": 8.5, + "axes.labelsize": 7.5, + "xtick.labelsize": 6.8, + "ytick.labelsize": 6.8, + "legend.fontsize": 6.6, + "axes.linewidth": 0.7, + "lines.linewidth": 1.25, + "pdf.fonttype": 42, + "ps.fonttype": 42, + } + ) + + fig = plt.figure(figsize=(7.35, 8.45), constrained_layout=True) + outer = fig.add_gridspec(3, 2, height_ratios=(1.04, 1.0, 0.94)) + + # ------------------------------------------------------------------ A + ax_a = fig.add_subplot(outer[0, 0]) + t = np.arange(pj["v"].shape[0]) * dt_s + zoom = (t >= 7.0) & (t <= 9.0) + ax_a.plot(t[zoom], np.asarray(pj["v"])[zoom, 0], color=JAXLEY, label="Jaxley") + ax_a.plot( + t[zoom], + np.asarray(pn["v"])[zoom, 0], + color=NEURON, + linestyle="--", + dashes=(3, 1.5), + label="NEURON", + ) + ax_a.set(xlabel="Time (s)", ylabel="Soma voltage (mV; 1952 HH frame,\n" r"$V_\mathrm{rest} \approx 0$ mV)") + ax_a.set_title("Matched firing-regime voltage") + ax_a.legend(frameon=False, loc="upper left", ncol=2, handlelength=2.5) + ax_a.grid(True, color=GRID, linewidth=0.45) + ax_a.set_axisbelow(True) + pool_cmp = comparison["comparison"]["pool"] + match = pool_cmp["spike_match_pm5ms"] + ax_a.text( + 0.98, + 0.04, + f"{match['jaxley_total']} vs {match['neuron_total']} spikes\n" + f"{100 * match['matched'] / match['jaxley_total']:.0f}% matched within ±5 ms", + transform=ax_a.transAxes, + ha="right", + va="bottom", + fontsize=6.7, + bbox={"boxstyle": "round,pad=0.25", "fc": "white", "ec": "#BBBBBB", "alpha": 0.9}, + ) + panel_label(ax_a, "A") + + # ------------------------------------------------------------------ B + sub_b = outer[0, 1].subgridspec(2, 1, height_ratios=(0.55, 0.45), hspace=0.05) + ax_b_rate = fig.add_subplot(sub_b[0]) + ax_b_raster = fig.add_subplot(sub_b[1], sharex=ax_b_rate) + trj, rj = smooth_population_rate(sj, dt_ms) + trn, rn = smooth_population_rate(sn, dt_ms) + ax_b_rate.plot(trj, rj, color=JAXLEY, label="Jaxley") + ax_b_rate.plot(trn, rn, color=NEURON, linestyle="--", label="NEURON") + ax_b_rate.set_ylabel("Pool rate\n(Hz/MN)") + ax_b_rate.set_title("Population trajectory and recruitment") + ax_b_rate.grid(True, color=GRID, linewidth=0.45) + ax_b_rate.set_axisbelow(True) + ax_drive = ax_b_rate.twinx() + command = np.asarray(artifact["ramp_command_hz"]) + td = np.arange(len(command)) * dt_s + ax_drive.plot(td[::100], command[::100], color=DRIVE, linewidth=0.8, alpha=0.55) + ax_drive.set_ylabel("Drive (Hz)", color=DRIVE, labelpad=2) + ax_drive.tick_params(axis="y", colors=DRIVE) + ax_drive.set_ylim(0, max(command) * 1.15) + plt.setp(ax_b_rate.get_xticklabels(), visible=False) + + n_units = sj.shape[1] + for i in range(n_units): + tj = np.flatnonzero(sj[:, i]) * dt_s + tn = np.flatnonzero(sn[:, i]) * dt_s + ax_b_raster.vlines(tj, i + 0.10, i + 0.35, color=JAXLEY, linewidth=0.55) + ax_b_raster.vlines(tn, i - 0.35, i - 0.10, color=NEURON, linewidth=0.55) + ax_b_raster.set( + xlabel="Time (s)", + ylabel="MN index", + yticks=np.arange(n_units), + ylim=(-0.65, n_units - 0.35), + xlim=(0, t[-1]), + ) + ax_b_raster.grid(True, axis="x", color=GRID, linewidth=0.45) + ax_b_raster.set_axisbelow(True) + both_active = (sj.sum(axis=0) > 0) & (sn.sum(axis=0) > 0) + first_j = [np.flatnonzero(sj[:, i])[0] for i in np.flatnonzero(both_active)] + first_n = [np.flatnonzero(sn[:, i])[0] for i in np.flatnonzero(both_active)] + same_order = bool(np.array_equal(np.argsort(first_j), np.argsort(first_n))) + order_text = "identical" if same_order else "DIFFERING" + ax_b_raster.text( + 0.98, + 0.08, + f"{int(both_active.sum())}/{n_units} active in both; " + f"{order_text} recruitment order", + transform=ax_b_raster.transAxes, + ha="right", + va="bottom", + fontsize=6.4, + bbox={"boxstyle": "round,pad=0.2", "fc": "white", "ec": "none", "alpha": 0.85}, + ) + ax_b_rate.legend(frameon=False, loc="upper left", ncol=2) + panel_label(ax_b_rate, "B") + + # ------------------------------------------------------------------ C + ax_c = fig.add_subplot(outer[1, 0]) + totals_n = np.array([p["metrics"]["total_neuron"] for p in points]) + totals_j = np.array([p["metrics"]["total_jaxley"] for p in points]) + timing = np.array([p["metrics"]["spike_match_pm5ms_fraction"] for p in points]) + gnap_values = np.array([p["gnap"] for p in points]) + lo, hi = 0.0, float(max(totals_n.max(), totals_j.max()) * 1.04) + ax_c.plot([lo, hi], [lo, hi], color="#555555", linewidth=0.8, zorder=0) + scatter = None + for gnap, marker in ((1.0, "o"), (5.0, "s")): + keep = gnap_values == gnap + scatter = ax_c.scatter( + totals_n[keep], + totals_j[keep], + c=timing[keep], + cmap="viridis", + vmin=0.45, + vmax=0.95, + marker=marker, + s=27, + edgecolor="white", + linewidth=0.4, + label=f"NaP ×{gnap:g}", + zorder=2, + ) + ax_c.set( + xlabel="NEURON total spikes", + ylabel="Jaxley total spikes", + xlim=(lo, hi), + ylim=(lo, hi), + ) + ax_c.set_aspect("equal", adjustable="box") + ax_c.set_title("Spike-count and timing agreement across the PIC grid") + ax_c.grid(True, color=GRID, linewidth=0.45) + ax_c.set_axisbelow(True) + ax_c.legend(frameon=False, loc="upper left") + cb = fig.colorbar(scatter, ax=ax_c, orientation="horizontal", pad=0.12, fraction=0.065) + cb.set_label("Fraction of spikes matched within ±5 ms", labelpad=1) + rel = np.abs(totals_j - totals_n) / np.maximum(totals_j, totals_n) + ax_c.text( + 0.98, + 0.04, + f"Median count error {100*np.median(rel):.2f}%\n" + "Timing match declines with PIC strength;\ncounts remain matched", + transform=ax_c.transAxes, + ha="right", + va="bottom", + fontsize=6.4, + ) + panel_label(ax_c, "C") + + # ------------------------------------------------------------------ D + ax_d = fig.add_subplot(outer[1, 1]) + gcals = sorted({float(p["gcal"]) for p in points}) + for gnap, color, marker in ((1.0, GNAP1, "o"), (5.0, GNAP5, "s")): + for backend, ls, fill in (("jaxley", "-", color), ("neuron", "--", "white")): + vals = [ + mean_metric(points, g, gnap, f"hysteresis_{backend}") for g in gcals + ] + ax_d.plot( + gcals, + vals, + color=color, + linestyle=ls, + marker=marker, + markersize=3.8, + markerfacecolor=fill, + markeredgecolor=color, + label=f"NaP ×{gnap:g}, {backend.capitalize()}", + ) + ax_d.axhline(0, color="#777777", linewidth=0.7) + ax_d.set( + xlabel="CaL conductance scale", + ylabel="Derecruitment − recruitment\ndrive (Hz)", + ) + ax_d.text( + 0.98, + 0.97, + "negative = PIC hysteresis", + transform=ax_d.transAxes, + ha="right", + va="top", + fontsize=6.4, + color="#555555", + ) + ax_d.set_title("PIC-dependent hysteresis") + ax_d.grid(True, color=GRID, linewidth=0.45) + ax_d.set_axisbelow(True) + ax_d.legend(frameon=False, ncol=2, loc="lower left", columnspacing=0.8, handlelength=2) + panel_label(ax_d, "D") + + # ------------------------------------------------------------------ E + ax_e = fig.add_subplot(outer[2, 0]) + for gnap, color, marker in ((1.0, GNAP1, "o"), (5.0, GNAP5, "s")): + for backend, ls, fill in (("jaxley", "-", color), ("neuron", "--", "white")): + vals = [ + mean_metric(points, g, gnap, f"self_sustained_{backend}") for g in gcals + ] + ax_e.plot( + gcals, + vals, + color=color, + linestyle=ls, + marker=marker, + markersize=4, + markerfacecolor=fill, + markeredgecolor=color, + ) + ax_e.axvline(6.5, color=GNAP5, linestyle=":", linewidth=0.9) + ax_e.axvline(7.5, color=GNAP1, linestyle=":", linewidth=0.9) + ax_e.text(6.45, 8.25, "NaP×5 boundary = 6.5", color=GNAP5, ha="right", va="bottom", fontsize=6.2) + ax_e.text(7.45, 6.25, "NaP×1 boundary = 7.5", color=GNAP1, ha="right", va="bottom", fontsize=6.2) + ax_e.set( + xlabel="CaL conductance scale", + ylabel="Self-sustained units\n(mean across seeds)", + ylim=(-0.35, 8.75), + ) + ax_e.set_title("Matched bistability boundaries") + ax_e.grid(True, color=GRID, linewidth=0.45) + ax_e.set_axisbelow(True) + style_legend = [ + Line2D([0], [0], color="#444444", linestyle="-", label="Jaxley"), + Line2D([0], [0], color="#444444", linestyle="--", label="NEURON"), + Line2D([0], [0], color=GNAP1, marker="o", linestyle="", label="NaP ×1"), + Line2D([0], [0], color=GNAP5, marker="s", linestyle="", label="NaP ×5"), + ] + ax_e.legend(handles=style_legend, frameon=False, ncol=2, loc="upper left") + panel_label(ax_e, "E") + + # ------------------------------------------------------------------ F + ax_f = fig.add_subplot(outer[2, 1]) + gnaps = [1.0, 5.0] + validated = np.zeros((len(gnaps), len(gcals)), dtype=int) + failures_by_condition: dict[tuple[float, float], int] = {} + failing_cells: list[dict] = [] + for iy, gnap in enumerate(gnaps): + for ix, gcal in enumerate(gcals): + cells = [ + a + for a in grid["acceptance_points"] + if a["gcal"] == gcal and a["gnap"] == gnap + ] + n_fail = 0 + for cell in cells: + for metric, v in cell["verdicts"].items(): + if v["class"] == "FAIL": + n_fail += 1 + failing_cells.append( + { + "gcal": gcal, + "gnap": gnap, + "seed": cell.get("seed"), + "metric": metric, + "value": v.get("value"), + } + ) + failures_by_condition[(gcal, gnap)] = int(n_fail) + validated[iy, ix] = int(n_fail == 0) + ax_f.imshow( + validated, + cmap=ListedColormap([UNRESOLVED, VALID]), + vmin=0, + vmax=1, + aspect="auto", + interpolation="nearest", + ) + for iy, gnap in enumerate(gnaps): + for ix, gcal in enumerate(gcals): + ok = bool(validated[iy, ix]) + text = "OK" if ok else f"{failures_by_condition[(gcal, gnap)]}×" + ax_f.text(ix, iy, text, ha="center", va="center", color="white", fontweight="bold", fontsize=8) + ax_f.set( + xticks=np.arange(len(gcals)), + xticklabels=[f"{g:g}" for g in gcals], + yticks=np.arange(len(gnaps)), + yticklabels=[f"NaP ×{g:g}" for g in gnaps], + xlabel="CaL conductance scale", + ) + ax_f.set_title( + "Frozen Phase-1C validated envelope\n" + f"{acceptance['n_evaluated_metric_cells'] - acceptance['n_failed_metric_cells']} / " + f"{acceptance['n_evaluated_metric_cells']} metric cells passed " + f"({100*(1-acceptance['fail_fraction']):.1f}%)" + ) + ax_f.tick_params(length=0) + ax_f.set_xticks(np.arange(-0.5, len(gcals), 1), minor=True) + ax_f.set_yticks(np.arange(-0.5, len(gnaps), 1), minor=True) + ax_f.grid(which="minor", color="white", linewidth=1.1) + ax_f.tick_params(which="minor", bottom=False, left=False) + metric_short = { + "spike_match_pm5ms_fraction_min": "±5 ms timing match", + "binned_rate_rmse_hz2": "binned-rate RMSE", + } + fail_counts: dict[str, int] = {} + for cell in failing_cells: + name = metric_short.get(cell["metric"], cell["metric"]) + fail_counts[name] = fail_counts.get(name, 0) + 1 + breakdown = ", ".join( + f"{name} ×{n}" + for name, n in sorted(fail_counts.items(), key=lambda kv: -kv[1]) + ) + ax_f.legend( + handles=[ + mpl.patches.Patch(color=VALID, label="validated at spec v1"), + mpl.patches.Patch( + color=UNRESOLVED, + label=f"unresolved at spec v1 tolerance\n({breakdown})", + ), + ], + frameon=False, + loc="lower center", + bbox_to_anchor=(0.5, -0.40), + ncol=1, + ) + panel_label(ax_f, "F") + + # Consistent cosmetic cleanup. + for ax in (ax_a, ax_b_rate, ax_b_raster, ax_c, ax_d, ax_e, ax_f): + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + + output_prefix.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_prefix.with_suffix(".pdf"), bbox_inches="tight") + fig.savefig(output_prefix.with_suffix(".png"), dpi=400, bbox_inches="tight") + plt.close(fig) + + source_summary = { + "figure": "Figure 1 — Cross-simulator validation of the forward twin", + "scope": "Hard-spike forward validation only; no surrogate gradients are plotted.", + "inputs": { + "paired_artifact": str(firing_artifact.relative_to(ROOT)), + "paired_content_hash": expected_hash, + "comparison_json": str(comparison_json.relative_to(ROOT)), + "pic_grid_json": str(grid_json.relative_to(ROOT)), + "acceptance_spec_version": acceptance["spec_version"], + }, + "paired_firing": { + "jaxley_spikes": match["jaxley_total"], + "neuron_spikes": match["neuron_total"], + "matched_pm5ms": match["matched"], + "matched_fraction_of_jaxley": match["matched"] / match["jaxley_total"], + "active_jaxley": int((sj.sum(axis=0) > 0).sum()), + "active_neuron": int((sn.sum(axis=0) > 0).sum()), + "active_in_both": int(both_active.sum()), + "recruitment_order_identical": same_order, + }, + "confirmatory_grid": { + "n_conditions_x_seeds": len(points), + "n_evaluated_metric_cells": acceptance["n_evaluated_metric_cells"], + "n_failed_metric_cells": acceptance["n_failed_metric_cells"], + "fail_fraction": acceptance["fail_fraction"], + "failure_policy_ok": acceptance["failure_policy_ok"], + "bistability_boundary": acceptance["bistability_boundary"], + "validated_envelope_gcal_by_gnap": acceptance[ + "validated_envelope_gcal_by_gnap" + ], + "median_total_spike_relative_error": float(np.median(rel)), + "failing_metric_cells": failing_cells, + }, + "conventions": { + "voltage_frame": "1952 HH convention (V_rest ~ 0 mV, ENa = +120 mV)", + "hysteresis_sign": ( + "derecruitment drive minus recruitment drive; " + "negative means derecruitment at lower drive (PIC hysteresis)" + ), + }, + } + output_prefix.with_name(output_prefix.name + "_source_data.json").write_text( + json.dumps(source_summary, indent=2) + ) + return source_summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-prefix", + type=Path, + default=DEFAULT_OUTPUT, + help="Output path without extension", + ) + args = parser.parse_args() + + summary = build_figure( + firing_artifact=TWIN_RESULTS / "val_protocol_firing.pkl", + jaxley_result=TWIN_RESULTS / "val_jaxley_bac411989ee8d42b.pkl", + neuron_result=TWIN_RESULTS / "val_neuron_bac411989ee8d42b.pkl", + comparison_json=TWIN_RESULTS / "val_comparison_bac411989ee8d42b.json", + grid_json=TWIN_RESULTS / "pic_grid_full.json", + output_prefix=args.output_prefix, + ) + print(f"Saved {args.output_prefix.with_suffix('.pdf')}") + print(f"Saved {args.output_prefix.with_suffix('.png')}") + print( + "Evidence: " + f"{summary['paired_firing']['jaxley_spikes']} vs " + f"{summary['paired_firing']['neuron_spikes']} paired spikes; " + f"{summary['confirmatory_grid']['n_failed_metric_cells']} / " + f"{summary['confirmatory_grid']['n_evaluated_metric_cells']} " + "confirmatory metric cells failed." + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/replot_dd_validation.py b/scripts/replot_dd_validation.py new file mode 100644 index 00000000..eb96005d --- /dev/null +++ b/scripts/replot_dd_validation.py @@ -0,0 +1,406 @@ +"""Regenerate the NEURON-vs-Jaxley descending-drive validation figure +(poster Figure 3: panels A–D) from the archived data in +``results/dd_validation_seeds/``. + +Reconstructed 2026-09-26 — the original generator script was lost; this one +recomputes everything from the saved per-seed pickles: + +* ``{neuron,jaxley}_spike_trains_seedNN.pkl`` — neo Block, 100 MN spike trains, + 0–15 s (joblib format). +* ``{neuron,jaxley}_drive_pattern_seedNN.pkl`` — neo AnalogSignal, 10 kHz, pps. + +Panels +------ +A Paired per-seed summary (active MNs, total spikes, mean rate per active MN) + with paired t-tests over seeds. Metrics are recomputed from the raw spike + trains and cross-checked against the archived ``panel_a_metrics.csv``. +B Direct spike-time matching (% spikes matched vs coincidence tolerance, + greedy one-to-one within each MN) for one seed. +C Smoothed population discharge-rate trajectory (50 ms bins, 300 ms moving + average) for the same seed, with Pearson r. +D Population DR correlation scatter (same bins), with r. + +Panels B–D use the poster's ORIGINAL single-run dataset — the 32 s triangular +ramp (0 → 65 pps at 16 s → 0) found in the top-level ``../results/`` folder: +``trapezoid_dd_spike_trains.pkl`` (NEURON: 76/100 active, 13 768 spikes) and +``trapezoid_dd_spike_trains_jaxley.pkl`` (Jaxley: 88/100, 14 287) — exactly the +numbers annotated on the poster. (Careful: ``MyoGen/results/`` holds a stale +15 s copy of the Jaxley file; the 32 s pair lives in the repo-parent +``results/``.) If those files are missing, panels B–D fall back to an archived +seed from ``dd_validation_seeds/`` (``--match-seed``). + +Usage:: + + python scripts/replot_dd_validation.py # defaults + python scripts/replot_dd_validation.py --seeds 0-19 # n=20 panel A +""" + +from __future__ import annotations + +import argparse +import re +from pathlib import Path + +import joblib +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from matplotlib.gridspec import GridSpec +from scipy import stats + +HERE = Path(__file__).resolve().parent +if (HERE / "data" / "panel_a_metrics.csv").exists(): + # Self-contained package layout: script + data/ + figures/ in one folder. + DATA_DIR = HERE / "data" + ROOT_RESULTS = HERE / "data" + OUT_DIR = HERE / "figures" +else: + # Repo layout: script in MyoGen/scripts/. + REPO = HERE.parent + DATA_DIR = REPO / "results" / "dd_validation_seeds" + OUT_DIR = REPO / "results" / "poster_figures_replot" + # The poster's original triangular-ramp run lives in the repo-parent results/. + ROOT_RESULTS = REPO.parent / "results" + +NEURON_C = "#2CB5B5" # teal +JAXLEY_C = "#5A64C8" # slate blue +BG = "#F8F5FB" +INK, MUTED = "#1a1a1a", "#5b6770" + +DURATION_S = 15.0 +BIN_S = 0.050 # 50 ms bins +SMOOTH_S = 0.300 # 300 ms moving average +TOLERANCES_MS = np.unique(np.round( + np.append(np.logspace(0, np.log10(50), 30), 5.0), 2)) +ANNOT_TOL_MS = 5.0 + + +# --------------------------------------------------------------------------- # +# Data loading / metrics +# --------------------------------------------------------------------------- # + +def load_trains_file(path: Path) -> tuple[list[np.ndarray], float]: + """Per-MN spike-time arrays in seconds + recording duration.""" + blk = joblib.load(path) + sts = blk.segments[0].spiketrains + dur = float(sts[0].t_stop.rescale("s").magnitude) + return [np.asarray(st.rescale("s").magnitude).ravel() for st in sts], dur + + +def load_trains(backend: str, seed: int, data_dir: Path) -> list[np.ndarray]: + return load_trains_file( + data_dir / f"{backend}_spike_trains_seed{seed:02d}.pkl")[0] + + +def seed_metrics(trains: list[np.ndarray]) -> dict: + counts = np.array([len(t) for t in trains]) + active = int((counts > 0).sum()) + total = int(counts.sum()) + return { + "active_MUs": active, + "total_spikes": total, + "mean_active_rate": total / active / DURATION_S if active else 0.0, + } + + +def compute_metrics(seeds: list[int], data_dir: Path) -> pd.DataFrame: + rows = [] + for seed in seeds: + for backend in ("neuron", "jaxley"): + m = seed_metrics(load_trains(backend, seed, data_dir)) + rows.append({"seed": seed, "backend": backend, **m}) + return pd.DataFrame(rows) + + +def match_spikes(a: np.ndarray, b: np.ndarray, tol_s: float) -> int: + """Greedy one-to-one matching of two sorted spike-time arrays.""" + i = j = n = 0 + while i < len(a) and j < len(b): + d = a[i] - b[j] + if abs(d) <= tol_s: + n += 1 + i += 1 + j += 1 + elif d < 0: + i += 1 + else: + j += 1 + return n + + +def matching_curve(neuron: list[np.ndarray], jaxley: list[np.ndarray]): + tot_n = sum(len(t) for t in neuron) + tot_j = sum(len(t) for t in jaxley) + pct_n, pct_j, matches_at = [], [], {} + for tol in TOLERANCES_MS: + m = sum(match_spikes(a, b, tol * 1e-3) for a, b in zip(neuron, jaxley)) + pct_n.append(100.0 * m / tot_n) + pct_j.append(100.0 * m / tot_j) + matches_at[tol] = m + return np.array(pct_n), np.array(pct_j), matches_at + + +def population_dr( + trains: list[np.ndarray], duration_s: float = DURATION_S +) -> tuple[np.ndarray, np.ndarray]: + """Smoothed population discharge rate (Hz/MN) on 50 ms bins.""" + edges = np.arange(0.0, duration_s + BIN_S, BIN_S) + counts = np.zeros(len(edges) - 1) + for t in trains: + counts += np.histogram(t, bins=edges)[0] + rate = counts / (BIN_S * len(trains)) # Hz per MN + w = max(1, int(round(SMOOTH_S / BIN_S))) + kernel = np.ones(w) / w + smoothed = np.convolve(rate, kernel, mode="same") + centers = edges[:-1] + BIN_S / 2 + return centers, smoothed + + +# --------------------------------------------------------------------------- # +# Plot helpers +# --------------------------------------------------------------------------- # + +def _style_ax(ax): + ax.set_facecolor(BG) + ax.grid(True, axis="y", color="#c9c4d4", linestyle=":", linewidth=0.9) + ax.set_axisbelow(True) + for side in ("top", "right"): + ax.spines[side].set_visible(False) + + +def _stars(p: float) -> str: + return "***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else "n.s." + + +def _p_label(p: float) -> str: + return f"{_stars(p)} (p < 0.001)" if p < 0.001 else f"{_stars(p)} (p = {p:.3f})" + + +def draw_panel_a(fig, subspec, df: pd.DataFrame): + """Three paired-bar mini-panels with per-seed dots and paired t-tests.""" + metrics = [ + ("active_MUs", "Active MNs", "MUs"), + ("total_spikes", "Total MN spikes", "Spikes"), + ("mean_active_rate", "Mean rate / active MN (Hz)", "Hz"), + ] + axes = subspec.subgridspec(1, 3, wspace=0.55) + rng = np.random.default_rng(0) + for k, (col, title, ylab) in enumerate(metrics): + ax = fig.add_subplot(axes[0, k]) + _style_ax(ax) + nv = df[df.backend == "neuron"].sort_values("seed")[col].to_numpy() + jv = df[df.backend == "jaxley"].sort_values("seed")[col].to_numpy() + p = stats.ttest_rel(nv, jv).pvalue + ax.bar([0, 1], [nv.mean(), jv.mean()], width=0.62, + color=[NEURON_C, JAXLEY_C], zorder=2) + x0 = 0 + rng.uniform(-0.06, 0.06, len(nv)) + x1 = 1 + rng.uniform(-0.06, 0.06, len(jv)) + for a, b, c, d in zip(x0, nv, x1, jv): + ax.plot([a, c], [b, d], color="#9aa0a8", lw=0.7, alpha=0.7, zorder=3) + ax.plot(x0, nv, "o", ms=5, color=NEURON_C, mec="white", mew=0.8, zorder=4) + ax.plot(x1, jv, "o", ms=5, color=JAXLEY_C, mec="white", mew=0.8, zorder=4) + top = max(nv.max(), jv.max()) + y = top * 1.06 + ax.plot([0, 0, 1, 1], [y, y * 1.02, y * 1.02, y], color=INK, lw=1.0) + ax.text(0.5, y * 1.04, _p_label(p), ha="center", va="bottom", + fontsize=9, color=INK) + ax.set_ylim(0, top * 1.22) + ax.set_xticks([0, 1], ["NEURON", "Jaxley"], rotation=20) + ax.set_title(title, fontsize=11, color=INK, fontweight="bold") + ax.set_ylabel(ylab) + + +def draw_panel_b(ax, pct_n, pct_j, matches_at): + _style_ax(ax) + ax.plot(TOLERANCES_MS, pct_n, "-o", ms=4.5, color=NEURON_C, + label="% NEURON spikes matched") + ax.plot(TOLERANCES_MS, pct_j, "-o", ms=4.5, color=JAXLEY_C, + label="% Jaxley spikes matched") + ax.set_xscale("log") + ax.set_xticks([1, 2, 5, 10, 20, 50], ["1", "2", "5", "10", "20", "50"]) + tol = min(matches_at, key=lambda t: abs(t - ANNOT_TOL_MS)) + ax.text(0.97, 0.15, f"±{tol:g} ms: {matches_at[tol]:,} matches", + transform=ax.transAxes, ha="right", fontsize=10, color=MUTED, + fontweight="bold", + bbox=dict(facecolor="white", edgecolor="#d5d0dd", alpha=0.85)) + ax.set_ylim(0, 100) + ax.set_xlabel("Spike coincidence tolerance (ms)") + ax.set_ylabel("Matched spikes (%)") + ax.legend(frameon=False, fontsize=9, loc="upper left") + + +def draw_panel_c(ax, t, dr_n, dr_j, r, drive=None): + _style_ax(ax) + if drive is not None: + # Protocol reference: drive envelope scaled onto the DR axis (shape + # only — peak pps noted in the label, no second y-scale). + t_drv, env = drive + peak = float(np.max(env)) + scaled = env / peak * max(dr_n.max(), dr_j.max()) * 1.04 + ax.plot(t_drv, scaled, color="#b3adc2", lw=1.2, zorder=1, + label=f"Drive envelope (peak {peak:.0f} pps, scaled)") + ax.plot(t, dr_n, color=NEURON_C, lw=1.4, label="NEURON", zorder=3) + ax.plot(t, dr_j, color=JAXLEY_C, lw=1.4, label="Jaxley", zorder=3) + ax.text(0.03, 0.06, + f"{BIN_S*1e3:.0f} ms bins, {SMOOTH_S*1e3:.0f} ms smoothing: r = {r:.3f}", + transform=ax.transAxes, fontsize=10, color=INK, fontweight="bold", + bbox=dict(facecolor="white", edgecolor="#d5d0dd", alpha=0.85)) + ax.set_xlabel("Time (s)") + ax.set_ylabel("Population DR (Hz/MN)") + ax.legend(frameon=False, fontsize=9) + + +def draw_panel_d(ax, dr_n, dr_j, r): + _style_ax(ax) + lim = max(dr_n.max(), dr_j.max()) * 1.08 + ax.plot([0, lim], [0, lim], "--", color="#8a93a0", lw=1.0) + ax.plot(dr_n, dr_j, "o", ms=5, color=JAXLEY_C, alpha=0.55, mec="none") + ax.text(0.05, 0.92, f"r = {r:.3f}", transform=ax.transAxes, + fontsize=13, color=INK, fontweight="bold") + ax.set_xlim(0, lim) + ax.set_ylim(0, lim) + ax.set_xlabel("NEURON population DR (Hz/MN)") + ax.set_ylabel("Jaxley population DR (Hz/MN)") + + +def save_all(fig, stem: str): + for ext in ("svg", "pdf"): + fig.savefig(OUT_DIR / f"{stem}.{ext}", facecolor=fig.get_facecolor()) + fig.savefig(OUT_DIR / f"{stem}.png", dpi=300, facecolor=fig.get_facecolor()) + + +# --------------------------------------------------------------------------- # + +def parse_seeds(spec: str) -> list[int]: + m = re.fullmatch(r"(\d+)-(\d+)", spec) + if m: + return list(range(int(m.group(1)), int(m.group(2)) + 1)) + return [int(s) for s in spec.split(",")] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--data-dir", type=Path, default=DATA_DIR) + ap.add_argument("--seeds", type=str, default="0-9", + help="seeds for panel A stats, e.g. '0-9' or '0,3,7'") + ap.add_argument("--match-seed", type=int, default=0, + help="fallback seed for panels B/C/D if the original " + "triangular-ramp pickles are absent") + ap.add_argument("--bcd-neuron", type=Path, default=None) + ap.add_argument("--bcd-jaxley", type=Path, default=None) + ap.add_argument("--bcd-drive", type=Path, default=None) + ap.add_argument("--show-drive", action="store_true", + help="overlay the (scaled) drive envelope in panel C") + args = ap.parse_args() + seeds = parse_seeds(args.seeds) + OUT_DIR.mkdir(exist_ok=True) + + # ---- metrics (panel A) + cross-check against the archived CSV + df = compute_metrics(seeds, args.data_dir) + df.to_csv(OUT_DIR / "panel_a_metrics_recomputed.csv", index=False) + archived = args.data_dir / "panel_a_metrics.csv" + if archived.exists(): + old = pd.read_csv(archived) + merged = df.merge(old, on=["seed", "backend"], suffixes=("", "_old")) + ok = ( + (merged.active_MUs == merged.active_MUs_old).all() + and (merged.total_spikes == merged.total_spikes_old).all() + and np.allclose(merged.mean_active_rate, + merged.mean_active_rate_old, atol=1e-3) + ) + print(f"cross-check vs archived panel_a_metrics.csv: " + f"{'MATCH' if ok else 'MISMATCH!'} ({len(merged)} rows compared)") + + # ---- panels B/C/D: prefer the poster's original triangular-ramp run + bcd_n = args.bcd_neuron or (ROOT_RESULTS / "trapezoid_dd_spike_trains.pkl") + bcd_j = args.bcd_jaxley or (ROOT_RESULTS / "trapezoid_dd_spike_trains_jaxley.pkl") + bcd_d = args.bcd_drive or (ROOT_RESULTS / "trapezoid_drive_pattern.pkl") + if bcd_n.exists() and bcd_j.exists(): + tr_n, dur_n = load_trains_file(bcd_n) + tr_j, dur_j = load_trains_file(bcd_j) + duration = max(dur_n, dur_j) + drive_path = bcd_d if bcd_d.exists() else None + src = f"original triangular-ramp run ({duration:.0f} s)" + else: + tr_n = load_trains("neuron", args.match_seed, args.data_dir) + tr_j = load_trains("jaxley", args.match_seed, args.data_dir) + duration = DURATION_S + drive_path = args.data_dir / f"neuron_drive_pattern_seed{args.match_seed:02d}.pkl" + src = f"archived seed {args.match_seed} (fallback)" + print(f"panels B–D source: {src} | NEURON " + f"{sum(1 for x in tr_n if len(x))}/{len(tr_n)} active, " + f"{sum(len(x) for x in tr_n):,} spikes | Jaxley " + f"{sum(1 for x in tr_j if len(x))}/{len(tr_j)} active, " + f"{sum(len(x) for x in tr_j):,} spikes") + + pct_n, pct_j, matches_at = matching_curve(tr_n, tr_j) + t, dr_n = population_dr(tr_n, duration) + _, dr_j = population_dr(tr_j, duration) + r = stats.pearsonr(dr_n, dr_j).statistic + + # drive envelope for the protocol overlay in panel C (off by default, + # enable with --show-drive; decimated for plotting) + drive = None + if args.show_drive and drive_path is not None and drive_path.exists(): + sig = joblib.load(drive_path) + env = np.asarray(sig.magnitude).ravel() + fs = float(sig.sampling_rate.rescale("Hz").magnitude) + dec = max(1, int(fs * 0.005)) # ~5 ms resolution + drive = (np.arange(len(env))[::dec] / fs, env[::dec]) + print(f"panels B–D: " + f"±5 ms matches = {matches_at[min(matches_at, key=lambda x: abs(x-5))]:,}, " + f"population-DR r = {r:.3f}") + + # ---- combined 4-panel figure + fig = plt.figure(figsize=(16, 10.5), facecolor=BG) + gs = GridSpec(2, 2, figure=fig, hspace=0.32, wspace=0.22, + left=0.06, right=0.97, top=0.90, bottom=0.07) + fig.suptitle( + f"NEURON vs Jaxley — paired descending-drive validation " + f"(n = {len(seeds)} seeds)", + fontsize=17, fontweight="bold", color=INK, x=0.06, ha="left", y=0.965, + ) + draw_panel_a(fig, gs[0, 0], df) + axB = fig.add_subplot(gs[0, 1]) + draw_panel_b(axB, pct_n, pct_j, matches_at) + axB.set_title("Direct spike-time matching", fontsize=12, color=INK, + fontweight="bold") + axC = fig.add_subplot(gs[1, 0]) + draw_panel_c(axC, t, dr_n, dr_j, r, drive=drive) + axC.set_title("Smoothed population discharge-rate trajectory", + fontsize=12, color=INK, fontweight="bold") + axD = fig.add_subplot(gs[1, 1]) + draw_panel_d(axD, dr_n, dr_j, r) + axD.set_title("Population DR correlation", fontsize=12, color=INK, + fontweight="bold") + save_all(fig, "figure_dd_validation_4panel") + plt.close(fig) + + # ---- individual panels (for poster tools) + figA = plt.figure(figsize=(10.5, 3.6), facecolor=BG) + gsA = GridSpec(1, 1, figure=figA, left=0.07, right=0.97, top=0.80, bottom=0.18) + figA.suptitle(f"A. Paired descending-drive summary (n = {len(seeds)} seeds)", + fontsize=13, fontweight="bold", color=INK, x=0.07, ha="left") + draw_panel_a(figA, gsA[0, 0], df) + save_all(figA, "panel_a_summary") + plt.close(figA) + + for stem, drawer in [ + ("panel_b_spike_matching", lambda ax: draw_panel_b(ax, pct_n, pct_j, matches_at)), + ("panel_c_rate_trajectory", lambda ax: draw_panel_c(ax, t, dr_n, dr_j, r, drive=drive)), + ("panel_d_rate_correlation", lambda ax: draw_panel_d(ax, dr_n, dr_j, r)), + ]: + f, ax = plt.subplots(figsize=(6.4, 4.6), facecolor=BG) + f.subplots_adjust(left=0.13, right=0.96, top=0.94, bottom=0.13) + drawer(ax) + save_all(f, stem) + plt.close(f) + + print(f"saved figures → {OUT_DIR}") + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index d12f9585..f53ba659 100644 --- a/setup.py +++ b/setup.py @@ -67,37 +67,15 @@ def compile_nmodl(self): print("NMODL compilation complete!") except Exception as e: - if platform.system() == "Windows": - # Windows requires manual NEURON installation - fail with clear instructions - error_msg = f""" -================================================================================ - MyoGen Installation Failed - NEURON Required -================================================================================ - -MyoGen requires NEURON 8.2.7 to be installed BEFORE installing MyoGen. - -STEP 1: Download and install NEURON ------------------------------------ - https://github.com/neuronsimulator/nrn/releases/download/8.2.7/nrn-8.2.7.w64-mingw-py-39-310-311-312-313-setup.exe - - During installation, select "Add to PATH" - -STEP 2: Retry installation ----------------------------------- - Re-run the MyoGen installation command, e.g.: - pip install myogen or uv add myogen - -================================================================================ -Technical details: {e} -================================================================================ -""" - raise RuntimeError(error_msg) from e - else: - # Linux/macOS - just warn, NMODL can be compiled later - print(f"Note: NMODL compilation skipped: {e}") - print( - "Run 'python -c \"from myogen import _setup_myogen; _setup_myogen()\"' after installation" - ) + # NMODL mechanisms are only needed for the OPTIONAL NEURON backend, so a + # missing NEURON at build time must never fail the install. On every + # platform we warn and continue; users who want the NEURON backend install + # `myogen[neuron]` and then run the setup task to compile mechanisms. + print(f"Note: NMODL compilation skipped (NEURON backend not built): {e}") + print( + 'To use the optional NEURON backend, install `myogen[neuron]` and run:\n' + ' python -c "from myogen import _setup_myogen; _setup_myogen()"' + ) def _compile_nmodl_windows(self, nmodl_path): """Compile NMODL on Windows.""" @@ -231,11 +209,15 @@ def _compile_nmodl_unix(self, nmodl_path): ), ] +# Parallel Cython compilation. Overridable via MYOGEN_CYTHONIZE_NTHREADS so builds in +# process-pool-constrained sandboxes/CI can fall back to single-threaded (nthreads=1). +_cythonize_nthreads = int(os.environ.get("MYOGEN_CYTHONIZE_NTHREADS", "4")) + setup( ext_modules=cythonize( extensions, compiler_directives={"embedsignature": True, "language_level": "2"}, - nthreads=4, + nthreads=_cythonize_nthreads, ), cmdclass={ "build_py": BuildWithNMODL, diff --git a/tests/test_jaxley_descending_drive_network.py b/tests/test_jaxley_descending_drive_network.py new file mode 100644 index 00000000..2b0077a5 --- /dev/null +++ b/tests/test_jaxley_descending_drive_network.py @@ -0,0 +1,60 @@ +"""Regression tests for the profile-matched Jaxley descending-drive path.""" + +import numpy as np +import pytest + +from myogen.simulator.jaxley.descending_drive_network import ( + _exp2_constants, + equal_fan_in_connectivity, + excitation_to_drive_hz, +) +from myogen.simulator.jaxley.populations import AlphaMN__Pool + + +def test_equal_fan_in_is_seeded_and_exact(): + a = equal_fan_in_connectivity(8, 17, 6, seed=42) + b = equal_fan_in_connectivity(8, 17, 6, seed=42) + c = equal_fan_in_connectivity(8, 17, 6, seed=43) + np.testing.assert_array_equal(a, b) + np.testing.assert_array_equal(a.sum(axis=1), np.full(8, 6.0)) + assert not np.array_equal(a, c) + + +def test_excitation_mapping_matches_profile_definition(): + excitation, hz = excitation_to_drive_hz( + [(0.001, 0.0, 0.0), (0.001, 0.5, 0.5)], + threshold_hz=100.0, + max_hz=300.0, + timestep_ms=1.0, + noise_seed=42, + ) + np.testing.assert_allclose(excitation, [0.0, 0.5]) + # Noise is half-normal, so rest is at least 0.8*threshold and active drive + # is at least threshold + excitation*(max-threshold). + assert hz[0] >= 80.0 + assert hz[1] >= 200.0 + + +def test_exp2syn_constants_are_peak_normalized(): + dt, tau1, tau2 = 0.001, 0.2, 0.3 + a1, a2, norm = _exp2_constants(dt, tau1, tau2) + t = np.arange(0.0, 3.0, dt) + waveform = norm * (np.exp(-t / tau2) - np.exp(-t / tau1)) + assert waveform.max() == pytest.approx(1.0, rel=2e-5) + assert 0 < a1 < a2 < 1 + + +def test_nerlab_raw_threshold_mode_matches_neuron_interpolation_convention(): + thresholds = np.array([0.2, 0.6, 1.0]) + normalized = AlphaMN__Pool( + recruitment_thresholds__array=thresholds, + normalize_nerlab_thresholds=True, + ) + raw = AlphaMN__Pool( + recruitment_thresholds__array=thresholds, + normalize_nerlab_thresholds=False, + ) + norm_diam = normalized.neuron_params["soma_diameter"] + raw_diam = raw.neuron_params["soma_diameter"] + assert norm_diam[0] < raw_diam[0] + assert raw_diam[-1] == pytest.approx(norm_diam[-1]) diff --git a/tests/test_twin_harness.py b/tests/test_twin_harness.py new file mode 100644 index 00000000..830336a8 --- /dev/null +++ b/tests/test_twin_harness.py @@ -0,0 +1,638 @@ +"""Tests for the digital-twin harness (experiments/twin/). + +Covers the pure-function layer (protocols, observables incl. the ΔF +estimator on synthetic spike trains with known ground truth) and the +PoolTwin build invariants (seed determinism, synapse-spec validation, +PIC-parameter scaling masks). Simulation behaviour itself is exercised by +the phase scripts, not here. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +import pytest + +TWIN = Path(__file__).resolve().parents[1] / "experiments" / "twin" +sys.path.insert(0, str(TWIN)) + +import protocols # noqa: E402 +import observables as obs # noqa: E402 +from poolbuild import PoolTwin, TwinSpec, fuglevand_thresholds # noqa: E402 + + +# --------------------------------------------------------------------------- # +# protocols / thresholds +# --------------------------------------------------------------------------- # + +def test_fuglevand_thresholds_monotonic_normalized(): + rt = fuglevand_thresholds(25, recruitment_range=30.0) + assert rt.shape == (25,) + assert np.all(np.diff(rt) > 0) + assert np.isclose(rt[-1], 1.0) + assert np.isclose(rt[-1] / rt[0], 30.0, rtol=0.2) + + +def test_triangular_ramp_shape_and_peak(): + dt = 0.5 + t, cmd = protocols.triangular_ramp(dt, t_rise_s=2.0, t_pre_s=1.0, + t_post_s=1.0, peak=40.0) + assert len(t) == len(cmd) + assert cmd.dtype == np.float32 + assert np.isclose(cmd.max(), 40.0, atol=0.5) + i_peak = int(np.argmax(cmd)) + assert abs(t[i_peak] - 3.0) < 0.05 # pre 1 s + rise 2 s + assert cmd[0] == 0.0 and cmd[-1] == 0.0 + + +def test_trapezoid_plateau(): + dt = 0.5 + _, cmd = protocols.trapezoid(dt, t_rise_s=1.0, t_hold_s=2.0, t_fall_s=1.0, + t_pre_s=0.5, t_post_s=0.5, peak=10.0) + n_plateau = int(np.sum(cmd >= 10.0 - 1e-6)) + assert abs(n_plateau - int(2.0 * 1000 / dt)) <= 2 + + +# --------------------------------------------------------------------------- # +# ΔF estimator on synthetic data with known ground truth +# --------------------------------------------------------------------------- # + +def _spikes_from_rate(rate_fn, t_on, t_off, duration_s, dt_ms): + """Deterministic spike train: integrate instantaneous rate to spike times.""" + n = int(duration_s * 1000 / dt_ms) + spikes = np.zeros(n, dtype=np.float32) + acc, t = 0.0, t_on + dt_s = dt_ms * 1e-3 + for i in range(int(t_on / dt_s), min(int(t_off / dt_s), n)): + acc += rate_fn(i * dt_s) * dt_s + if acc >= 1.0: + spikes[i] = 1.0 + acc -= 1.0 + return spikes + + +def _synthetic_ramp_pool(extra_control: bool = False): + """Control MN(s) tracking a triangular drive; test MN with hysteresis. + + Drive: 0→20 Hz at t=15 s→0 at t=30 s. Control rate = drive (min 4 Hz + while firing, 2–28 s). Test MN fires 10–24 s (recruited while drive + rising at 13.3 Hz, derecruited while falling at 8 Hz). + True ΔF = f_ctrl(10 s) − f_ctrl(24 s) = 13.3 − 8 = 5.3 Hz. + """ + dt_ms, dur = 1.0, 30.0 + drive = lambda t: 20.0 * (t / 15.0 if t <= 15.0 else max(0.0, 2.0 - t / 15.0)) + ctrl_rate = lambda t: max(4.0, drive(t)) + cols = [_spikes_from_rate(ctrl_rate, 2.0, 28.0, dur, dt_ms)] + if extra_control: + cols.append(_spikes_from_rate(lambda t: max(4.0, 0.9 * drive(t)), + 2.5, 27.5, dur, dt_ms)) + cols.append(_spikes_from_rate(lambda t: max(3.0, 0.8 * drive(t)), + 10.0, 24.0, dur, dt_ms)) + return np.stack(cols, axis=1), dt_ms + + +def test_delta_f_recovers_known_value(): + spikes, dt_ms = _synthetic_ramp_pool() + res = obs.delta_f_pairs(spikes, dt_ms) + assert len(res.delta_f) == 1 + assert len(res.delta_f_per_test) == 1 + # true ΔF ≈ 5.3 Hz; allow smoothing/discretization slack + assert res.mean == pytest.approx(5.3, abs=1.5) + assert res.mean > 0 + + +def test_delta_f_per_test_averaging(): + spikes, dt_ms = _synthetic_ramp_pool(extra_control=True) + res = obs.delta_f_pairs(spikes, dt_ms) + assert len(res.delta_f) == 2 # two controls, one test unit + assert len(res.delta_f_per_test) == 1 # averaged within the test unit + assert res.mean == pytest.approx(res.delta_f.mean(), abs=1e-9) + + +def test_delta_f_rejects_short_outlast_control(): + # Control derecruited only 0.5 s after the test MN → must be rejected + # by the min_control_outlast_s=1.5 safeguard. + dt_ms, dur = 1.0, 30.0 + drive = lambda t: 20.0 * (t / 15.0 if t <= 15.0 else max(0.0, 2.0 - t / 15.0)) + ctrl = _spikes_from_rate(lambda t: max(4.0, drive(t)), 2.0, 24.5, dur, dt_ms) + test = _spikes_from_rate(lambda t: max(3.0, 0.8 * drive(t)), 10.0, 24.0, + dur, dt_ms) + res = obs.delta_f_pairs(np.stack([ctrl, test], axis=1), dt_ms) + assert len(res.delta_f) == 0 + + +def test_recruitment_order_matches_onsets(): + dt_ms, dur = 1.0, 10.0 + a = _spikes_from_rate(lambda t: 10.0, 1.0, 9.0, dur, dt_ms) + b = _spikes_from_rate(lambda t: 10.0, 3.0, 9.0, dur, dt_ms) + order = obs.recruitment_order(np.stack([b, a], axis=1), dt_ms) + assert list(order) == [1, 0] + + +# --------------------------------------------------------------------------- # +# PoolTwin build invariants (slow-ish: builds real jaxley cells) +# --------------------------------------------------------------------------- # + +@pytest.fixture(scope="module") +def tiny_twin(): + return PoolTwin(TwinSpec(n_mn=3, n_dd=5, n_inh=5, seed=0)) + + +def test_spec_validation_rejects_bad_options(): + with pytest.raises(ValueError, match="syn_site"): + PoolTwin(TwinSpec(n_mn=2, syn_site="axon")) + with pytest.raises(ValueError, match="syn_tau_rise_ms"): + PoolTwin(TwinSpec(n_mn=2, syn_tau_rise_ms=10.0, tau_syn_ms=5.0)) + + +def test_connectivity_seed_determinism(tiny_twin): + same = PoolTwin(TwinSpec(n_mn=3, n_dd=5, n_inh=5, seed=0)) + other = PoolTwin(TwinSpec(n_mn=3, n_dd=5, n_inh=5, seed=1)) + assert np.array_equal(np.asarray(tiny_twin._exc_mat), np.asarray(same._exc_mat)) + assert not np.array_equal(np.asarray(tiny_twin._exc_mat), + np.asarray(other._exc_mat)) + + +def test_delta_f_rejects_flat_test_profile(): + # Test unit firing at a CONSTANT rate while the control modulates: the + # common-modulation check (flat profile or low/NaN correlation) must + # reject the pair. + dt_ms, dur = 1.0, 30.0 + drive = lambda t: 20.0 * (t / 15.0 if t <= 15.0 else max(0.0, 2.0 - t / 15.0)) + ctrl = _spikes_from_rate(lambda t: max(4.0, drive(t)), 2.0, 28.0, dur, dt_ms) + test = _spikes_from_rate(lambda t: 8.0, 10.0, 24.0, dur, dt_ms) + res = obs.delta_f_pairs(np.stack([ctrl, test], axis=1), dt_ms) + assert len(res.delta_f) == 0 + + +def test_pic_scaling_masks(tiny_twin): + import jax.numpy as jnp + + theta = tiny_twin.default_theta() + theta["gCaL_scale"] = jnp.float32(2.0) + scaled = tiny_twin._scaled_jaxley_params(theta) + base = np.asarray(tiny_twin._all_params["caL_gcaLbar"]) + out = np.asarray(scaled["caL_gcaLbar"]) + mask = np.asarray(tiny_twin._pic_masks["caL_gcaLbar"]) + assert mask.any() # dendrite compartments carry caL + assert np.allclose(out[mask], 2.0 * base[mask]) + # unmasked (sentinel/NaN) entries untouched + untouched = ~mask & np.isfinite(base) + assert np.allclose(out[untouched], base[untouched]) + + +# --------------------------------------------------------------------------- # +# New synapse modes: Exp2Syn kernel, event API, placement, live driving force +# --------------------------------------------------------------------------- # + +@pytest.fixture(scope="module") +def exp2_twin(): + """Dendritic Exp2Syn twin, fixed driving force.""" + return PoolTwin(TwinSpec(n_mn=2, n_dd=3, n_inh=3, seed=0, + syn_site="dendrite", syn_tau_rise_ms=0.5)) + + +@pytest.fixture(scope="module") +def exp2_live_twin(): + """Same as exp2_twin but with voltage-dependent driving force.""" + return PoolTwin(TwinSpec(n_mn=2, n_dd=3, n_inh=3, seed=0, + syn_site="dendrite", syn_tau_rise_ms=0.5, + live_driving_force=True)) + + +def test_exp2syn_kernel_peak_time_and_decay(exp2_twin): + import jax.numpy as jnp + + spec = exp2_twin.spec + t1, t2 = spec.syn_tau_rise_ms, spec.tau_syn_ms + syn = exp2_twin._syn_init() + zero = jnp.zeros(spec.n_mn, dtype=jnp.float32) + gs = [] + for i in range(400): # 40 ms impulse response + inc = jnp.full(spec.n_mn, 1.0 if i == 0 else 0.0, dtype=jnp.float32) + syn, g_exc, _ = exp2_twin._syn_step(syn, inc, zero) + gs.append(float(g_exc[0])) + gs = np.array(gs) + # peak-normalized: unit-weight event peaks at ~1.0 at the analytic t_peak + tp = (t1 * t2 / (t2 - t1)) * np.log(t2 / t1) + i_peak = int(np.argmax(gs)) + assert abs((i_peak + 1) * spec.dt_ms - tp) <= 2 * spec.dt_ms + assert gs[i_peak] == pytest.approx(1.0, rel=0.02) + # late decay is governed by tau_decay: one tau2 drops the value by e + i0 = int(20.0 / spec.dt_ms) + i1 = i0 + int(t2 / spec.dt_ms) + assert gs[i1] / gs[i0] == pytest.approx(np.exp(-1.0), rel=0.05) + + +def test_events_api_current_signs(tiny_twin): + theta = tiny_twin.default_theta() + n_steps, n_mn = 600, tiny_twin.spec.n_mn + conn = np.ones((n_mn, 1), dtype=np.float32) + silent = np.zeros((n_steps, 1), dtype=np.float32) + burst = silent.copy() + burst[50:60, 0] = 1.0 # 10 events + + v_base = np.asarray(tiny_twin.simulate_events( + theta, silent, exc_connectivity=conn)["v"]) + v_exc = np.asarray(tiny_twin.simulate_events( + theta, burst, exc_connectivity=conn)["v"]) + v_inh = np.asarray(tiny_twin.simulate_events( + theta, silent, inh_events=burst, exc_connectivity=conn, + inh_connectivity=conn)["v"]) + + d_exc = v_exc - v_base + d_inh = v_inh - v_base + assert d_exc.max() > 1e-4 and d_exc.min() > -1e-5 # pure depolarization + assert d_inh.min() < -1e-5 and d_inh.max() < 1e-5 # pure hyperpolarization + + +def test_events_api_shape_validation(tiny_twin): + theta = tiny_twin.default_theta() + with pytest.raises(ValueError, match="exc_connectivity"): + tiny_twin.simulate_events(theta, np.zeros((10, 2), dtype=np.float32)) + + +def test_stimulus_placement_differs_between_sites(tiny_twin, exp2_twin): + a = np.asarray(tiny_twin._external_inds["i"]).ravel()[: tiny_twin.spec.n_mn] + b = np.asarray(exp2_twin._external_inds["i"]).ravel()[: exp2_twin.spec.n_mn] + assert not np.array_equal(a[: len(b)], b) + + +def test_live_driving_force_changes_dynamics(exp2_twin, exp2_live_twin): + theta = exp2_twin.default_theta() + n_steps = 600 + conn = np.ones((exp2_twin.spec.n_mn, 1), dtype=np.float32) + burst = np.zeros((n_steps, 1), dtype=np.float32) + burst[50:80, 0] = 1.0 + v_fixed = np.asarray(exp2_twin.simulate_events( + theta, burst, exc_connectivity=conn)["v"]) + v_live = np.asarray(exp2_live_twin.simulate_events( + theta, burst, exc_connectivity=conn)["v"]) + assert np.max(np.abs(v_fixed - v_live)) > 1e-6 + # live DF weakens depolarizing current as V rises → smaller peak response + assert v_live.max() <= v_fixed.max() + 1e-6 + + +def test_hard_surrogate_forward_equivalence_new_modes(exp2_live_twin): + import jax + import jax.numpy as jnp + + _, cmd = protocols.step(exp2_live_twin.spec.dt_ms, t_pre_s=0.05, + t_on_s=0.3, t_post_s=0.05, level=60.0) + drives = {"dd": jnp.asarray(cmd), + "inh": jnp.zeros(len(cmd), dtype=jnp.float32)} + theta = exp2_live_twin.default_theta() + key = jax.random.PRNGKey(0) + hard = exp2_live_twin.simulate(theta, drives, key, spike_mode="hard") + sur = exp2_live_twin.simulate(theta, drives, key, spike_mode="surrogate") + assert np.allclose(np.asarray(hard["v"]), np.asarray(sur["v"])) + assert np.allclose(np.asarray(hard["spikes"]), np.asarray(sur["spikes"])) + + +# --------------------------------------------------------------------------- # +# Sensitivity-diagnostic regression tests (Deliverable A report semantics) +# --------------------------------------------------------------------------- # + +def test_safe_norm_no_float32_overflow(): + import sensitivity_checks as sc + + g = np.array([1.4e28, 4.7e28], dtype=np.float32) # the T4 1 s components + n = sc._safe_norm(g) + assert np.isfinite(n) and n > 1e28 # float32 norm would report inf + + +def test_manifest_schema(): + import sensitivity_checks as sc + + m = sc._manifest() + assert {"git_revision", "versions", "device", "precision", + "seed_taxonomy", "objectives", "scope"} <= set(m) + + +def test_t5_reports_joint_and_coordinate_separately(): + import sensitivity_checks as sc + + # 2.0 s fixture produces a FIRING observation (≈8 spikes), so the test + # exercises coordinate semantics, not just the report schema. + res = sc.t5_hard_forward_descent( + n_mn=2, duration_s=2.0, conn_seeds=(0,), event_seeds=(0,), + true_points=((1.5, 0.011),), lambdas=(0.1,)) + assert {"accepted_step_fraction", "coordinate_agreement", "cases", + "n_failures", "n_evaluated", + "n_independent_event_realizations"} <= set(res) + assert res["n_evaluated"] >= 1 # non-vacuous: at least one case evaluated + for name in ("gCaL_scale", "w_exc"): + assert {"n_resolved", "n_unresolved", "agreement_rate"} <= set( + res["coordinate_agreement"][name]) + n_probed = 0 + for c in res["cases"]: + if "joint_steps" in c: # evaluated case → both result kinds present + assert "coordinate_probes" in c + for probe in c["coordinate_probes"].values(): + n_probed += 1 + assert {"true_descent_dir", "surrogate_dir", "agrees", + "hard_loss_baseline"} <= set(probe) + # baseline-aware semantics: a resolved descent direction must + # actually improve on the baseline loss, never be "less bad" + if probe["agrees"] is not None: + assert probe["true_descent_dir"] != 0 + assert (min(probe["hard_loss_up"], probe["hard_loss_down"]) + < probe["hard_loss_baseline"]) + assert n_probed >= 2 # both coordinates probed in the evaluated case + + +# --------------------------------------------------------------------------- # +# Deliverable B regression tests +# --------------------------------------------------------------------------- # + +def test_limb_segmentation(): + import deliverable_b as db + + cmd = np.concatenate([np.zeros(10), np.linspace(0, 10, 20), + np.full(5, 10.0), np.linspace(10, 0, 20), + np.zeros(10)]) + limbs = db.segment_limbs(cmd) + masks = np.stack([limbs[k] for k in ("pre", "rising", "peak", + "falling", "post")]) + assert masks.sum(axis=0).max() == 1 # disjoint + assert limbs["pre"][:9].all() and limbs["post"][-9:].all() + assert limbs["peak"].sum() >= 5 # plateau included + assert limbs["rising"].any() and limbs["falling"].any() + i_peak = int(np.argmax(cmd)) + assert not limbs["rising"][i_peak:].any() + assert not limbs["falling"][:i_peak].any() + + +def test_truth_fit_loss_is_zero_same_model(): + import jax.numpy as jnp + import deliverable_b as db + + commands = db.make_commands(0.1, dur_scale=0.3) + twin, events = db.build_twin_and_events(2, 0, 0, commands, 0.1) + rates_fn = db.make_rates_fn(twin, events) + theta = np.array([1.6, 0.0105]) + obs = {p: rates_fn(theta, p) for p in ("fit_ramp", "fit_trap")} + tot = 0.0 + for p, o in obs.items(): + r = rates_fn(theta, p) + tot += float(jnp.mean((r - o) ** 2)) + assert tot == 0.0 # identical model + identical frozen events + + +def test_event_matrix_hash_stable(): + import deliverable_b as db + + a = np.zeros((10, 3), dtype=np.float32) + a[2, 1] = 1.0 + h1 = db._sha16(a) + assert h1 == db._sha16(a.copy()) + b = a.copy() + b[3, 0] = 1.0 + assert h1 != db._sha16(b) + + +def test_binned_rates_counts_and_jitter_invariance(): + import jax.numpy as jnp + import deliverable_b as db + + sp = np.zeros((1000, 2), dtype=np.float32) # 100 ms at dt 0.1 → 2 bins + sp[10, 0] = 1.0 + sp[20, 0] = 1.0 # two spikes in bin 0 of unit 0 + r = np.asarray(db.binned_rates_jnp(jnp.asarray(sp), 0.1, + bin_ms=50.0, smooth_ms=50.0)) + assert r.shape == (2, 2) + assert r[0, 0] == 2 / 0.05 # 40 Hz + assert r[1, 0] == 0.0 and r[:, 1].sum() == 0.0 + # spike-time jitter WITHIN a bin must not change the observable at all — + # this is the property that kills the chaotic-decorrelation floor + sp2 = np.zeros_like(sp) + sp2[14, 0] = 1.0 + sp2[33, 0] = 1.0 + r2 = np.asarray(db.binned_rates_jnp(jnp.asarray(sp2), 0.1, + bin_ms=50.0, smooth_ms=50.0)) + assert np.allclose(r, r2) + + +def test_axis_with_truth_contains_truth_sorted_unique(): + import deliverable_b as db + + ax = db.axis_with_truth(0.25, 4.0, 21, 1.6) + assert 1.6 in ax + assert np.all(np.diff(ax) > 0) + assert len(ax) == len(np.unique(ax)) + # truth already on a grid point → no duplicate inserted + ax2 = db.axis_with_truth(1.0, 4.0, 3, 2.0) + assert 2.0 in ax2 and len(ax2) == len(np.unique(ax2)) + + +def test_run_tag_separates_smoke_from_full(): + import argparse + import deliverable_b as db + + full = argparse.Namespace(dense_n=21, replicates=3, dur_scale=1.0, n_mn=8) + smoke = argparse.Namespace(dense_n=5, replicates=3, dur_scale=1.0, n_mn=8) + assert db.run_tag(full) == "full" + assert db.run_tag(smoke) == "smoke" + + +# --------------------------------------------------------------------------- # +# Deliverable C nesting + validation-pipeline regression tests +# --------------------------------------------------------------------------- # + +def test_nested_profiles_nesting_property(): + import deliverable_c as dc + + g_axis = np.geomspace(0.5, 2.0, 9) + b_axis = np.geomspace(0.005, 0.02, 9) + + def losses_at(g, b): + # protocols preferring DIFFERENT amplitudes → unknown < known strictly + l1 = (b - 0.008) ** 2 * 1e4 + (g - 1.0) ** 2 + l2 = (b - 0.015) ** 2 * 1e4 + (g - 1.0) ** 2 + return (l1, l2) + + known, unknown = dc.compute_nested_profiles(losses_at, g_axis, b_axis, 2) + assert np.all(unknown <= known + dc.NESTING_TOL) + assert np.any(unknown < known - 1e-6) # strictly better somewhere + + def losses_same(g, b): + v = (b - 0.01) ** 2 * 1e4 + (g - 1.0) ** 2 + return (v, v) + + k2, u2 = dc.compute_nested_profiles(losses_same, g_axis, b_axis, 2) + # identical protocols → releasing amplitudes gains nothing + assert np.allclose(k2, u2, atol=dc.NESTING_TOL) + + +def test_match_spikes_greedy(): + from validation_protocol import _match_spikes + + a = np.array([0.1, 0.2, 0.3]) + assert _match_spikes(a, a.copy(), 1e-3) == 3 + b = a + 0.004 + assert _match_spikes(a, b, 5e-3) == 3 + assert _match_spikes(a, b, 1e-3) == 0 + # greedy one-to-one: two a-spikes cannot share one b-spike + assert _match_spikes(np.array([0.1, 0.101]), np.array([0.1]), 5e-3) == 1 + + +def test_collect_alignment_semantics(): + from validation_protocol import _collect + + class FakeVec: + def __init__(self, arr): + self._a = np.asarray(arr, dtype=float) + + def as_numpy(self): + return self._a + + n_steps, dt = 5, 0.1 + # NEURON records n_steps+1 samples; sample 0 is the initial condition and + # must be dropped so neuron[k] aligns with jaxley's post-step sample k + vrec = [FakeVec([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])] + dvrec = [FakeVec([0.0, -1.0, -2.0, -3.0, -4.0, -5.0])] + spk = [FakeVec([0.2, 0.41])] # ms → steps 2 and 4 + out = _collect(vrec, dvrec, spk, dt, n_steps, 1) + assert out["v"].shape == (n_steps, 1) + assert out["v"][0, 0] == 1.0 # initial condition dropped + assert out["spikes"][2, 0] == 1.0 and out["spikes"][4, 0] == 1.0 + assert out["spikes"].sum() == 2 + + +def test_run_neuron_subprocess_smoke(): + # run_neuron must live in a jax-free process (JAX+NEURON segfault when + # mixed); this also regression-tests the single-pool phase structure. + import subprocess + import sys + + code = ( + "import validation_protocol as vp\n" + "art = vp.make_protocol(n_mn=2, t_rise_s=0.4, t_pre_s=0.1," + " t_post_s=0.2, step_currents_nA=(15.0,))\n" + "res = vp.run_neuron(art)\n" + "print('STEP_SPIKES', int(res['current_steps'][15.0]['spikes'].sum()))\n" + "print('POOL_SHAPE', res['pool']['v'].shape[0]," + " res['pool']['v'].shape[1])\n" + ) + proc = subprocess.run([sys.executable, "-c", code], cwd=str(TWIN), + capture_output=True, text=True, timeout=600) + assert proc.returncode == 0, proc.stdout + proc.stderr + out = proc.stdout + n_spk = int(out.split("STEP_SPIKES")[1].split()[0]) + assert n_spk > 10 # 15 nA drives the small cells + # pool phase length = (0.1 pre + 2×0.4 ramp + 0.2 post) s / 0.1 ms + assert "POOL_SHAPE 11000 2" in out + + +# --------------------------------------------------------------------------- # +# Gradient-repair diagnostic helpers (G1-G3 v2) +# --------------------------------------------------------------------------- # + +def test_fd_sweep_finds_convergence_plateau(): + import gradient_repair as gr + + f = lambda s: float(s[0] ** 2 + 3 * s[1]) + _, c0 = gr.fd_sweep_converged(f, np.array([1.0, 1.0]), 0) + _, c1 = gr.fd_sweep_converged(f, np.array([1.0, 1.0]), 1) + assert c0 == pytest.approx(2.0, rel=1e-5) + assert c1 == pytest.approx(3.0, rel=1e-5) + + +def test_fd_sweep_reports_no_plateau_on_discontinuity(): + import gradient_repair as gr + + step = lambda s: 1.0 if s[0] > 1.0 else 0.0 # FD slope ∝ 1/h, no plateau + _, c = gr.fd_sweep_converged(step, np.array([1.0]), 0) + assert c is None + + +def test_classify_direction_zero_is_unresolved_not_wrong(): + import gradient_repair as gr + + assert gr.classify_direction([1.0, -1.0, 1.0], [2.0, 3.0, 0.0]) == \ + ["agree", "wrong", "unresolved"] + + +def test_interp_time_gradient_matches_analytic(): + import jax + import jax.numpy as jnp + import gradient_repair as gr + + ramp = jnp.arange(10.0) + k, c = 4, 4.5 + + def ts(theta): # v(θ) = θ·t ⇒ crossing time c/θ ms ⇒ d/dθ = −c/θ² ·1e-3 s + return gr.interp_time(theta * ramp, k, c, 1.0) + + g = float(jax.grad(ts)(1.0)) + assert g == pytest.approx(-c * 1e-3, rel=1e-6) + + +def test_manifest_precision_and_script_hashes_honest(): + import jax + from sensitivity_checks import _manifest + + m = _manifest() + if jax.config.jax_enable_x64: + assert "float64" in m["precision"] + else: + assert "float32" in m["precision"] + assert "gradient_repair.py" in m["source_sha256_16"] + assert "deliverable_b.py" in m["source_sha256_16"] + + +def test_zero_insert_events_preserves_physical_times(): + import gradient_repair as gr + + ev = np.zeros((5, 2), dtype=np.float32) + ev[2, 0] = 1.0 + ev[4, 1] = 1.0 + fine = gr.zero_insert_events(ev, 4) + assert fine.shape == (20, 2) + assert fine[8, 0] == 1.0 and fine[16, 1] == 1.0 # k·factor + assert fine.sum() == ev.sum() + + +def test_detect_plateau_is_ad_independent(): + import gradient_repair as gr + + rows = [(1e-1, 5.0), (1e-2, 2.2), (1e-3, 2.01), (1e-4, 2.0)] + assert gr.detect_plateau(rows) == pytest.approx(2.0) + assert gr.detect_plateau([(1e-1, 1.0), (1e-2, 10.0), (1e-3, 100.0)]) is None + + +def test_chained_windows_equal_uninterrupted_run(): + import jax.numpy as jnp + import g4_descent as g4 + + twin = PoolTwin(TwinSpec(n_mn=2, n_dd=5, n_inh=2, seed=0)) + n = 4000 + rng = np.random.default_rng(0) + ev = {"exc": jnp.asarray((rng.random((n, 5)) < 0.002).astype(np.float64)), + "inh": jnp.zeros((n, 2)), + "exc_conn": jnp.ones((2, 5)), "inh_conn": jnp.zeros((2, 2))} + th = twin.default_theta() + full = np.asarray(twin.simulate_events( + th, ev["exc"], inh_events=np.asarray(ev["inh"]), + exc_connectivity=np.asarray(ev["exc_conn"]), + inh_connectivity=np.asarray(ev["inh_conn"]), + spike_mode="hard")["spikes"]) + wins = g4.split_windows(ev, 4) + chained = np.asarray(g4.chained_spikes(twin, th, wins, "hard")) + # state-preserving windowing must change NOTHING about the hard forward + assert chained.shape == full.shape + assert np.array_equal(chained, full) + + +def test_g4_case_events_are_float64_under_x64(): + import jax + + if not jax.config.jax_enable_x64: + jax.config.update("jax_enable_x64", True) + import g4_descent as g4 + + _twin, ev = g4.build_case(0, 0, dur_scale=0.25) + assert str(ev["fit_ramp"]["exc"].dtype) == "float64" + assert str(ev["fit_ramp"]["exc_conn"].dtype) == "float64" diff --git a/uv.lock b/uv.lock index ea3daddd..0624bb74 100644 --- a/uv.lock +++ b/uv.lock @@ -10,6 +10,15 @@ resolution-markers = [ "python_full_version < '3.13' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] +[[package]] +name = "absl-py" +version = "2.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/1d/58e2b5a6e4d703ccb2a029943d665974cb3d5a4fb2b3e3675dd03a9df10e/absl_py-2.5.1.tar.gz", hash = "sha256:286e71c82c1a38e75bbcf185f9b37d0305ad7786535107cb49bf4df9ff2e1f95", size = 118387, upload-time = "2026-10-09T08:38:58.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/7d/01e62f59e4166af1be6238d9f2b7f3453630c51b8044b412d05b5606532a/absl_py-2.5.1-py3-none-any.whl", hash = "sha256:721200f2f0e9960f2ca9dc3a2a706b201f5f75d812c158f059cbbe29eeafbdb8", size = 137673, upload-time = "2026-10-09T08:38:56.629Z" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.2" @@ -289,6 +298,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] +[[package]] +name = "diffrax" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax" }, + { name = "jaxtyping" }, + { name = "lineax" }, + { name = "optimistix" }, + { name = "typing-extensions" }, + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/ea/a22853377a840d70d1d174d43183f12990c66eecf34ba7e76eec781249d0/diffrax-0.7.2.tar.gz", hash = "sha256:9adc70fe90dba83f7dd839845839b2531a3dd736a3944fe1aa26598507cfb48e", size = 155159, upload-time = "2026-02-18T01:14:04.964Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/48/c42f657f3d6b4d9d9674c9bc48e763d024511ea4cc1d6e0c68c7b4380c43/diffrax-0.7.2-py3-none-any.whl", hash = "sha256:1beee2f991cd049962fe3c89da28f933753aa66d80bef857533276c9f23301ed", size = 199673, upload-time = "2026-02-18T01:14:03.38Z" }, +] + +[[package]] +name = "equinox" +version = "0.13.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "jaxtyping" }, + { name = "typing-extensions" }, + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/ff/522336d2f8264f2ad97119710b76e2cddf66145d03a1e89899175d26b192/equinox-0.13.8.tar.gz", hash = "sha256:dd075050018e2dd02e252e9d29d3060f7e67f085622d8d27a8e89e24bb8523db", size = 145257, upload-time = "2026-05-05T10:03:43.258Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/d6/69a76c8ccdef14af687c497040292a46e59fc7a0ab24724b60e50ca61030/equinox-0.13.8-py3-none-any.whl", hash = "sha256:ca004348533cc30a63ebe8823d7dd4bb626dce17743d40bbddb89b402ef2a240", size = 185813, upload-time = "2026-05-05T10:03:41.673Z" }, +] + [[package]] name = "fastrlock" version = "0.8.3" @@ -510,6 +552,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/a7/f76514cc40ad6234098ecdebda08732d75964776c51a42845b7da10649e2/idna-3.17-py3-none-any.whl", hash = "sha256:466e48829084efe2548012b855df21540b96f2e20e51bd124c851536556a592c", size = 65316, upload-time = "2026-05-28T14:32:37.035Z" }, ] +[[package]] +name = "impi-rt" +version = "2021.18.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/35/6ef32a71394233a996ca56c15c1e4218040367a0758bf2fd0868fe34347f/impi_rt-2021.18.1-py2.py3-none-win_amd64.whl", hash = "sha256:a4ec3c4b6bda016db573b2505a4d7c59ddfb0c94553eead301da917891c7b6ab", size = 21855822, upload-time = "2026-07-14T12:41:35.392Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -528,6 +578,92 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/aa/0aca39a37d3c7eb941ba736ede56d689e7be91cab5d9ca846bde3999eba6/isodate-0.7.2-py3-none-any.whl", hash = "sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15", size = 22320, upload-time = "2024-10-08T23:04:09.501Z" }, ] +[[package]] +name = "jax" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaxlib" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/39/28/13186d20cbece4944577de774269c620b2f5f5ea163748b61e48e423f47f/jax-0.9.0.tar.gz", hash = "sha256:e5ce9d6991333aeaad37729dd8315d29c4b094ea9476a32fb49933b556c723fb", size = 2534495, upload-time = "2026-01-20T23:06:47.099Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/99/91fd58d495f03495e93b3f6100336d7a1df212aa90ad37d1322b8bb7df71/jax-0.9.0-py3-none-any.whl", hash = "sha256:be1c7c1fa91b338f071805a786f9366d4a08361b711b0568ef5a157226c68915", size = 2955442, upload-time = "2026-01-20T22:27:05.128Z" }, +] + +[[package]] +name = "jaxley" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "matplotlib" }, + { name = "networkx" }, + { name = "numpy" }, + { name = "pandas" }, + { name = "tridiax" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b8/a4/8d29883849342178320d6eaec3e48ddb77247489424cb16620e1c62ddb48/jaxley-0.13.0.tar.gz", hash = "sha256:0d9247b340b402f974aad827e0cd79e32c5cd66d7295d95514792a108e15f00b", size = 196846, upload-time = "2025-12-18T10:32:57.38Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/93/27c97ed86da85d19bf7cec783f893933a14cfec4183871228cfcaa30802f/jaxley-0.13.0-py3-none-any.whl", hash = "sha256:277f135714f1370b7246754d64687357ec443e3a944f1a96633dfd4eaaafcc3e", size = 246222, upload-time = "2025-12-18T10:32:55.677Z" }, +] + +[[package]] +name = "jaxley-mech" +version = "0.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "diffrax" }, + { name = "equinox" }, + { name = "jax" }, + { name = "jaxley" }, + { name = "matplotlib" }, + { name = "numpy" }, + { name = "pandas" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/87/cf/950acda0f61bfdf95343644e57be323b97028edbb6ce83ecfd06a57aad45/jaxley_mech-0.3.1.tar.gz", hash = "sha256:bd46cb2f02d1f76af56406ef83c464b6f9fc9742625cd88371a1923e14f601e8", size = 41315, upload-time = "2025-07-18T12:35:06.972Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/a1/d11bf6853f19eee78b057d8fcee694cea3929165f028852f7923462d2c2b/jaxley_mech-0.3.1-py3-none-any.whl", hash = "sha256:cc5eda21c8521e32795526f9f85ca52941899449b0a491d3ffdb321f3f0c8cbd", size = 56214, upload-time = "2025-07-18T12:35:05.297Z" }, +] + +[[package]] +name = "jaxlib" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "scipy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/c1/ed8a199b71627f9c40b84a9fb40bba44bc048e62449633c5ee507066b19c/jaxlib-0.9.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3c8d13bd74a55f9e5435587730b1ee067b17fc4c29b897f467c095003c63b37d", size = 56102192, upload-time = "2026-01-20T22:28:19.717Z" }, + { url = "https://files.pythonhosted.org/packages/5e/15/345f16517435b8f1deadd5bbfaae8dde42ed19af6e86ebbca0620279bab4/jaxlib-0.9.0-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:a19e0ad375161190073d12221b350b3da29f97d9e72b0bc6400d0dbf171cda5f", size = 74771957, upload-time = "2026-01-20T22:28:22.576Z" }, + { url = "https://files.pythonhosted.org/packages/20/7c/54ea33501e46b62b80071ed9cb96b0cfb28df88b7907edff0de41e9e9892/jaxlib-0.9.0-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:93258e8bbfad38f9207a649797e07bf6e7fb30a8dece67b1485f9d550c59121f", size = 80336468, upload-time = "2026-01-20T22:28:25.86Z" }, + { url = "https://files.pythonhosted.org/packages/c6/fe/1969b10690fb49f1fadcc616367857c54b16949a008778877e359e5da13b/jaxlib-0.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:c75f8d3ae8e8411fc90e6c18a99971ad4a5da8d200e5ad710d650e7c1d5652e6", size = 60508836, upload-time = "2026-01-20T22:28:29.275Z" }, + { url = "https://files.pythonhosted.org/packages/ab/1a/a1db3023f66af906c85cdbafad88a0ae02df0d87ac7f749f1c9140ef6e3c/jaxlib-0.9.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:abcf50ca9f0f015f756035ff00803f5d04f42964775f936420d7afbf23b61c5c", size = 56102369, upload-time = "2026-01-20T22:28:32.554Z" }, + { url = "https://files.pythonhosted.org/packages/23/7e/ad2d0519939a45eb4a3151d5e41bef1d68eec7b3e7f961e09ef3bf2d6607/jaxlib-0.9.0-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:7da962fcbdfc18093ea371cfb5420fc70b51d3160d15bef0b780e63126ad25c0", size = 74769548, upload-time = "2026-01-20T22:28:35.481Z" }, + { url = "https://files.pythonhosted.org/packages/ec/d0/77bda8d946e3a854a5a533be7baf439bfffab9adae2a668a4dc645a494cd/jaxlib-0.9.0-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:6aba184182c45e403b62dc269a04747da55da11d97586269710199204262d4ec", size = 80336519, upload-time = "2026-01-20T22:28:39.067Z" }, + { url = "https://files.pythonhosted.org/packages/22/0c/6e7891a226ee3fd717af26fdbc93ef6dcf587849a0020fdfc69ccb38ca7f/jaxlib-0.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:e81debd335315ef178fa995399653444e620c316abe86bc51477ac012e13efd4", size = 60508118, upload-time = "2026-01-20T22:28:42.51Z" }, + { url = "https://files.pythonhosted.org/packages/32/42/75c4cbb31cd56d334839fe5bea132eb5c00f1298f9a04245c0951af5f623/jaxlib-0.9.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:de5d4920da2eb1182a0dbc01cfee9acf8617b64562a0b8f153fdfbcd00381f35", size = 56213023, upload-time = "2026-01-20T22:28:45.202Z" }, + { url = "https://files.pythonhosted.org/packages/06/af/91c23ae50bcfdc1dbeecd8ee4a7e87128a99a58bf6b6f6081215e8674a00/jaxlib-0.9.0-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:79caef921ebbbd78c239c3b4ed7644e0605bf454a066ad7700dde6c1e60b4b1e", size = 74887880, upload-time = "2026-01-20T22:28:48.227Z" }, + { url = "https://files.pythonhosted.org/packages/9a/64/7875e6ee446b558a366b3e9cbf26bab2390d87ac8015bc0dc49b35806b53/jaxlib-0.9.0-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:c857f062c720cff924f2237c925d46b8a8e8fde2b829ef2dc44cd4b6a7b8df88", size = 80440020, upload-time = "2026-01-20T22:28:51.378Z" }, +] + +[[package]] +name = "jaxtyping" +version = "0.3.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/c1/091b8852bd7cbf50bd655543c8506033cf4029300c67f8c176c1286879a9/jaxtyping-0.3.11.tar.gz", hash = "sha256:b09c14acf6686feb9e0df5b0d8c6e7c5b6f8d36bf059ee54cd522a186c2ef050", size = 46489, upload-time = "2026-06-13T18:35:23.167Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/38/c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0/jaxtyping-0.3.11-py3-none-any.whl", hash = "sha256:8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3", size = 56593, upload-time = "2026-06-13T18:35:22.01Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -632,6 +768,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b5/91/53255615acd2a1eaca307ede3c90eb550bae9c94581f8c00081b6b1c8f44/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:1f1489f769582498610e015a8ef2d36f28f505ab3096d0e16b4858a9ec214f57", size = 75987, upload-time = "2026-03-09T13:15:39.65Z" }, ] +[[package]] +name = "lineax" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax" }, + { name = "jaxtyping" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/d6/4e28416a6fe58dd6bc7565b1ffa330f4d0ba7d74212642b1b734c511299e/lineax-0.1.0.tar.gz", hash = "sha256:5f1a8f060142af2cdbf7d66b99e8d3071c3aa734b677df6339df4b4c4c0554d2", size = 50209, upload-time = "2026-01-27T21:17:26.652Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/0c/2ed47112fc1958a0a81c9b015d4e1861953a1ec3a17b081c0180a25ce82c/lineax-0.1.0-py3-none-any.whl", hash = "sha256:f00911c6b07d427c4835db46856970c8348bc82a035b51f4386ad09382af957a", size = 74600, upload-time = "2026-01-27T21:17:25.33Z" }, +] + [[package]] name = "llvmlite" version = "0.47.0" @@ -906,6 +1057,60 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/fc/10ab7e80650a9c9e8f4f1105f8c8e73567f88ed0c06ada589ab81d38687c/mkdocstrings_python-2.0.5-py3-none-any.whl", hash = "sha256:30c837bbff016549f659fcba6539ac351303f0fd7e713c89a040611072236e9d", size = 104951, upload-time = "2026-06-19T10:41:07.378Z" }, ] +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, + { url = "https://files.pythonhosted.org/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" }, + { url = "https://files.pythonhosted.org/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" }, + { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" }, + { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, + { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, + { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" }, + { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" }, + { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" }, + { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, + { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, + { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" }, + { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" }, +] + +[[package]] +name = "mpi4py" +version = "4.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/75/83/231445bbcf7ef10864746c244ff2d82000011449b79275642c5d4ed8c8f4/mpi4py-4.1.2.tar.gz", hash = "sha256:56860286dc45f20e8821e93cb06669e30462348bf866f685553fa4b712d58d02", size = 501709, upload-time = "2026-05-16T10:35:23.618Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bf/7e/1524baa98c3b4e3ddc030ec2ad5f6e5a362930317ed3e561d6da1c2e97b1/mpi4py-4.1.2-cp310-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ceb3b6e6f27ba7a39f7721583c04514013b860b829e721769e77398dee97bfa3", size = 1411317, upload-time = "2026-05-16T10:33:56.706Z" }, + { url = "https://files.pythonhosted.org/packages/2f/73/5918e062a65f40e3bd82a2dccf0c8ba22a284d327ccda46d84d2717985cc/mpi4py-4.1.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:251e8880f4cb98e9c8f63c6f6b2c7e819e22b5e4949d47767a0092eed6f814c2", size = 1304635, upload-time = "2026-05-16T10:33:58.685Z" }, + { url = "https://files.pythonhosted.org/packages/a3/0a/1da7f403e0d8ce0e26d541f7538302cec00cf5b0a98a7a52b929f938a25c/mpi4py-4.1.2-cp310-abi3-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2ef63b2e3083e6062fd90e4de8c4e3acbf81e0772406e0226eb8dde6a48cab8e", size = 1327130, upload-time = "2026-05-16T10:34:00.269Z" }, + { url = "https://files.pythonhosted.org/packages/e6/f9/65999152ae82bad914c6a083821ee774afefd6d0544e633b940c9a9ebf3f/mpi4py-4.1.2-cp310-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6508e654b9c8ff9f611b19548b2a17d1e323b520a15168189f92221e6757b8ff", size = 1182268, upload-time = "2026-05-16T10:34:02.206Z" }, + { url = "https://files.pythonhosted.org/packages/66/8f/90de8d5e0676f86f7866fd4665193252ea17764a54800d9545c7b53ac00a/mpi4py-4.1.2-cp310-abi3-win_amd64.whl", hash = "sha256:eb69f6273ad155f191850a593deebdf52aed6722979ba0693e02db5663e59699", size = 1466751, upload-time = "2026-05-16T10:34:04.278Z" }, + { url = "https://files.pythonhosted.org/packages/3e/73/d3d747f74c5aa7be6faf84aef6b234458baeeb0acb35143c288bc04133c5/mpi4py-4.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:78983a9606e105af1af907878ec41627360344ea576c75afdaca4124db9d0b13", size = 1621733, upload-time = "2026-05-16T10:34:24.804Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e0/2d029ea1907a477126039a48aadb30113cf3df5bd3c2f27331334b281415/mpi4py-4.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4e5fdba71cf27d0d3b45e706736d6690060976345f45a6408b3594884bb0cc6e", size = 1393929, upload-time = "2026-05-16T10:34:26.585Z" }, + { url = "https://files.pythonhosted.org/packages/9f/2b/1e48c4c5f9acbdca8dd28beeba9123dde140cd2ca520f8e3a3cf22faeeaa/mpi4py-4.1.2-cp312-cp312-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:00f4cce8999d19f35243c3442ea22debbe3336f69c309cd5d3176df4e51c717a", size = 1358844, upload-time = "2026-05-16T10:34:28.247Z" }, + { url = "https://files.pythonhosted.org/packages/94/46/a37225d47997fcf30adca25d3849d035bbb61d972118b024db900306e528/mpi4py-4.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0a58e164776acb7b52414548b1bf0e5caafce0ee90345deff147873a64b6b2cc", size = 1227206, upload-time = "2026-05-16T10:34:29.866Z" }, + { url = "https://files.pythonhosted.org/packages/b5/48/14173d8baaf622b76e0e8983f53fffff8538edcf24a020028432e63f20d3/mpi4py-4.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:0a7227a53ba24102f2c4c2c1b226d028fabc473b20ceef5435f1a7ac8cbd508c", size = 1706260, upload-time = "2026-05-16T10:34:31.806Z" }, + { url = "https://files.pythonhosted.org/packages/35/15/324cfa61674b8f73aa3347590df937232d898a993f3348121f9003f43b16/mpi4py-4.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:632a94e355cf1960eaae2b668d88399f53680207f5a4777e15fdaf7d96c72d09", size = 1660517, upload-time = "2026-05-16T10:34:33.836Z" }, + { url = "https://files.pythonhosted.org/packages/00/ec/b670c9250b59ed255571e679f1e0ae7bd4e08c67f2f7f7f11a928b4becd7/mpi4py-4.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0aedf3e6c937b48d244e7c6bb72d0b068a1b1adf99ec6001bc9c7e381c581806", size = 1426478, upload-time = "2026-05-16T10:34:35.342Z" }, + { url = "https://files.pythonhosted.org/packages/0d/5e/493415cdb0da0c1c0b6ec9a1fb65ab57a174e8111b25c87f7507f663d0b4/mpi4py-4.1.2-cp313-cp313-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:085e0cb05427398fe2281856f27a87984b9f234cd6d98a2a384a6fbfe679a56f", size = 1358571, upload-time = "2026-05-16T10:34:37.478Z" }, + { url = "https://files.pythonhosted.org/packages/6f/51/3822e834fc9cafd96811501b85232e02bd9b31596d42965536a269a6c112/mpi4py-4.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:429437883b4511583d56a57bc58ad0d02965fb6c19f84f39e30c0646e6c0cab9", size = 1226562, upload-time = "2026-05-16T10:34:39.652Z" }, + { url = "https://files.pythonhosted.org/packages/72/27/919831e5c843890b1a5b475ca1a8d83d8eb3453c1fe66f6bbdaca35ee548/mpi4py-4.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:f984aa35d61fd95c7aa49829f94c3d61bb9ec2272bd8e404ae8d1ea84b620b58", size = 1703612, upload-time = "2026-05-16T10:34:41.133Z" }, + { url = "https://files.pythonhosted.org/packages/ef/ee/3c1128b5c393a8a2abeb9db6cc2aba2d3d604fee20d0fa99c02ba5aab5c6/mpi4py-4.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:16ae4d0fac11e60056ec7c46ce101f70f4941634ad8165c54a24304820a816b0", size = 1740466, upload-time = "2026-05-16T10:34:42.566Z" }, + { url = "https://files.pythonhosted.org/packages/17/ed/7a163be7da48ace29ab4c54407e7f2e510953df14dea90e122f9ca7310c7/mpi4py-4.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:364da611dda94e8a26e418ccfc4c8ff2a1492e74f4c28b24df09521fe888de3b", size = 1518426, upload-time = "2026-05-16T10:34:44.081Z" }, + { url = "https://files.pythonhosted.org/packages/ec/5e/d358fadb8672d58abd6dce16c95eda56f497b378dec5642a31cd6ededfc4/mpi4py-4.1.2-cp313-cp313t-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:85e332feaac3323d8ed1c71f17478fe3f529468f28f6f9bb4c9133ad2d3a3a6d", size = 1419007, upload-time = "2026-05-16T10:34:45.47Z" }, + { url = "https://files.pythonhosted.org/packages/27/eb/5cd53880337009cab9a9d17a007ad5aec731f3a55e211bdcc99dbb98a0e3/mpi4py-4.1.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3e05dba75f6ad17ae761bdc01ee5c5271d15cd296e8a508e22c99c91b540dc99", size = 1298482, upload-time = "2026-05-16T10:34:46.858Z" }, + { url = "https://files.pythonhosted.org/packages/13/d9/eb35b0e8328a02c4d6288bf2bcd156b387f08ff6146a5c8f3f2abd75a078/mpi4py-4.1.2-cp313-cp313t-win_amd64.whl", hash = "sha256:2d0a131a255f61a2d313623f91b496c345abd5976147f866b39383264f275449", size = 1835283, upload-time = "2026-05-16T10:34:48.899Z" }, +] + [[package]] name = "multidict" version = "6.7.1" @@ -975,10 +1180,12 @@ version = "0.13.0" source = { editable = "." } dependencies = [ { name = "beartype" }, + { name = "jax" }, + { name = "jaxley" }, + { name = "jaxley-mech" }, { name = "joblib" }, { name = "matplotlib" }, { name = "neo" }, - { name = "neuron", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "numba" }, { name = "numpy" }, { name = "optuna" }, @@ -1000,6 +1207,11 @@ dependencies = [ gpu = [ { name = "cupy-cuda12x" }, ] +neuron = [ + { name = "impi-rt", marker = "sys_platform == 'win32'" }, + { name = "mpi4py" }, + { name = "neuron", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, +] nwb = [ { name = "h5py" }, { name = "nwbinspector" }, @@ -1008,9 +1220,11 @@ nwb = [ [package.dev-dependencies] dev = [ + { name = "optax" }, { name = "pandas-stubs" }, { name = "poethepoet" }, { name = "pytest" }, + { name = "pyyaml" }, { name = "scipy-stubs" }, ] docs = [ @@ -1026,10 +1240,15 @@ requires-dist = [ { name = "beartype", specifier = ">=0.21.0" }, { name = "cupy-cuda12x", marker = "extra == 'gpu'", specifier = ">=13.0,<14" }, { name = "h5py", marker = "extra == 'nwb'", specifier = ">=3.0" }, + { name = "impi-rt", marker = "sys_platform == 'win32' and extra == 'neuron'", specifier = ">=2021.15.0" }, + { name = "jax", specifier = "==0.9.0" }, + { name = "jaxley", specifier = "==0.13.0" }, + { name = "jaxley-mech", specifier = "==0.3.1" }, { name = "joblib", specifier = ">=1.3" }, { name = "matplotlib", specifier = ">=3.10.1" }, + { name = "mpi4py", marker = "extra == 'neuron'", specifier = ">=4.0.3" }, { name = "neo", specifier = ">=0.14.0" }, - { name = "neuron", marker = "sys_platform == 'darwin' or sys_platform == 'linux'", specifier = "==8.2.7" }, + { name = "neuron", marker = "(sys_platform == 'darwin' and extra == 'neuron') or (sys_platform == 'linux' and extra == 'neuron')", specifier = "==8.2.7" }, { name = "numba", specifier = ">=0.61.2" }, { name = "numpy", specifier = ">=1.26" }, { name = "nwbinspector", marker = "extra == 'nwb'", specifier = ">=0.5.0" }, @@ -1048,13 +1267,15 @@ requires-dist = [ { name = "threadpoolctl", specifier = ">=3.1" }, { name = "tqdm", specifier = ">=4.67.1" }, ] -provides-extras = ["nwb", "gpu"] +provides-extras = ["nwb", "neuron", "gpu"] [package.metadata.requires-dev] dev = [ + { name = "optax", specifier = ">=0.2.4" }, { name = "pandas-stubs", specifier = ">=2.3.0.250703" }, { name = "poethepoet", specifier = ">=0.37.0" }, { name = "pytest", specifier = ">=8.0" }, + { name = "pyyaml", specifier = ">=6.0" }, { name = "scipy-stubs", specifier = ">=1.16.1.0" }, ] docs = [ @@ -1088,6 +1309,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ac/66/93359047a9c7ede7c6319f58b2cee3be19552f57ef5380fe95cb812272fb/neo-0.14.4-py3-none-any.whl", hash = "sha256:0cd863424e35dddd768d713cfcc48cb06ab127fe1b5b132a0793c82975de5be4", size = 697727, upload-time = "2026-03-12T09:59:54.609Z" }, ] +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + [[package]] name = "neuron" version = "8.2.7" @@ -1201,6 +1431,46 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/0b/9ed001e9d16a94ea5bd74dee2462786fcc869c4bcb3dd568c2373ac9d260/nwbinspector-0.7.2-py3-none-any.whl", hash = "sha256:40d008eedd8afdc5d6841825d6efb6891933c1fd3b8c5b599ab2db14f62db8e8", size = 86926, upload-time = "2026-05-19T18:08:22.348Z" }, ] +[[package]] +name = "opt-einsum" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/b9/2ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb/opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac", size = 63004, upload-time = "2024-09-26T14:33:24.483Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, +] + +[[package]] +name = "optax" +version = "0.2.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8c/f9/e3d11ae6f298ee941a0690e353a323d158ba5dedc436e75621c310845c5c/optax-0.2.8.tar.gz", hash = "sha256:5b225b35066fc3eebaa4d798f1b4173b4d57d1a480610908981f8343b50af0b0", size = 301193, upload-time = "2026-03-20T23:30:05.465Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/69/6a93d8600c339d7687a05857c7907bd4dd8cf88691a5ea106d7a50af90a1/optax-0.2.8-py3-none-any.whl", hash = "sha256:e3ca2d36c99daab1800ae9dbc0545034382d6bc780b24d969e1b0df65fa31cb4", size = 402960, upload-time = "2026-03-20T23:30:03.886Z" }, +] + +[[package]] +name = "optimistix" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax" }, + { name = "jaxtyping" }, + { name = "lineax" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/0b/3bdc0e698cb0d264e5dca0b5bb171342a950a43fa6e046d8c57189298422/optimistix-0.1.0.tar.gz", hash = "sha256:f05c9104748e87e1dc10a0b4a2be94e427f98c3eab0b1d41ea24a593d76be03d", size = 70164, upload-time = "2026-02-16T13:35:43.991Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/53/6a81a6ebb1739f19c32002139719d4ee021a349d5bdbe066910feed327a1/optimistix-0.1.0-py3-none-any.whl", hash = "sha256:c8edda7bf7fe48839c93fcd72bce750c950ed9f7033158303150b7ab395add81", size = 101740, upload-time = "2026-02-16T13:35:42.971Z" }, +] + [[package]] name = "optuna" version = "4.8.0" @@ -1906,6 +2176,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, ] +[[package]] +name = "tridiax" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5d/8d/55d41b1de379faf0518b8e110c656bef40e73059df4cfff51c0b72cb4928/tridiax-0.2.1.tar.gz", hash = "sha256:95a8c6d003cdd694487c99e5ba2c43d4fb4dfbe3a3df96e9ac2c80c1c4aaecd1", size = 11483, upload-time = "2024-12-05T09:55:54.374Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/fd/f69ff723a4e6534fce070acc5c50b80e739b2efb4c49ec580a629c6a3898/tridiax-0.2.1-py3-none-any.whl", hash = "sha256:311b0ed41671303197e219019fb9d22d6b31c841ddf5fdd1ec2601e09ed4e750", size = 11813, upload-time = "2024-12-05T09:55:52.741Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -1933,6 +2216,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] +[[package]] +name = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + [[package]] name = "watchdog" version = "6.0.0"