Jaxley Integration

.github/workflows/build-wheels.yml

@@ -169,9 +169,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:

.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

.github/workflows/tests.yml

@@ -43,9 +43,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 --group dev
+          uv sync --group dev --extra neuron
 
       - name: Compile NEURON mechanisms
         run: uv run poe setup_myogen

.gitignore

@@ -1,4 +1,7 @@
 ### Python template
+# macOS Finder metadata
+.DS_Store
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

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.

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).

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.<sub> 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.