GPU-native jaxon simulator accelerates peripheral nerve modeling

Researchers David Lung and Max Haberbusch have introduced jaxon, an open-source software package that reimplements canonical peripheral-nerve fiber models using the JAX machine-learning framework. The tool allows scientists to simulate entire populations of nerve fibers on a single graphics processing unit while maintaining fidelity to established biophysical models.

The standard approach for modeling peripheral nerve stimulation couples a volume-conductor description of tissue, typically solved by the finite-element method, with multi-compartment cable models of individual axons. The reference implementation for the fiber side of this pipeline is the NEURON software, often accessed through the PyFibers wrapper. While this reference stack is accurate, it faces two primary structural limitations for modern neural-interface design. First, its throughput is CPU-bound, as a NEURON run integrates one fiber at a time. Although independent fibers can be run in parallel on multi-core nodes, scoring an optimized stimulus over an entire fiber population remains computationally expensive.

Because the per-fiber cost is high and the population must be re-scored at every optimization step, models almost universally reduce each fascicle to one or a few representative fibers. Whether this reduction biases the selectivity predicted by the models has been effectively untestable due to the computational burden. The second limitation is that the standard stack is non-differentiable. Because there is no gradient of the fiber response with respect to the stimulus, stimulation parameters must be tuned by grid search or gradient-free optimization, methods that scale poorly as the number of contacts and waveform degrees of freedom grow.

jaxon was built to remove both limitations while staying faithful to the published biophysics. It reimplements the canonical peripheral-nerve fiber models, including the myelinated MRG and Sweeney axons and the unmyelinated Sundt and Rattay C-fibers, in JAX on top of the Jaxley differentiable cable simulator. The software adds a custom backward-Euler solver for the coupled intracellular and periaxonal double-cable equation, which serves as the principled equivalent of NEURON’s extracellular mechanism. Because the whole forward model is built in JAX, it is simultaneously vectorized, allowing whole populations to run in parallel, and differentiable, exposing exact gradients of the fiber response with respect to the stimulus.

The software architecture organizes the model as a differentiable forward model wrapped in an optimization and analysis layer. At its center is a vectorized integrator of the coupled cable equations that maps a set of per-contact stimulus amplitudes to the membrane response of an arbitrary number of fibers at once. The package is structured by concern, with modules for ion-channel mechanisms, morphology builders, stimulus helpers, and optimization loops. The coupled solver integrates the double-cable equation using backward Euler, where each time step is a 2×2 block-tridiagonal system solved by a JAX-compatible block-Thomas sweep.

Validation of jaxon was conducted against PyFibers-wrapped NEURON across 943 activation-threshold configurations. Each model was tested over its diameter range, with eight waveform shapes and six pulse widths. The results showed that 99.6% of configurations agreed within a 1% threshold error, with a median absolute error of 0.08% and an R-squared value of 0.99999. The myelinated MRG and Sweeney axons agreed to a median of 0.04% and 0.02%, respectively, and reproduced conduction velocity to machine precision.

For the unmyelinated Sundt and Rattay C-fibers, the agreement was a median of 0.34% and 0.07%, with conduction velocity within 0.24%. The only four configurations exceeding a 1% error were short-pulse charge-balanced-biphasic Rattay cases at the smallest diameters, where the bisection threshold search is most sensitive. These validation results confirm that jaxon maintains high fidelity to the reference biophysics while enabling new computational capabilities.

Performance benchmarks compared the jax.vmap-batched forward integration of a rectangular-pulse response over N fibers on a single NVIDIA A100 GPU against single-process PyFibers/NEURON on a dual-socket AMD EPYC 7742 node. The GPU times reported are steady-state run times, excluding the one-time JAX compilation cost. Across five model variants, the batched forward pass was a geometric-mean 255 times faster at N=1000 fibers and 819 times faster at N=100,000 fibers. The range of speedups at the larger population size was between 527 and 1377 times.

The comparison against the CPU baseline highlights that while a multi-core node can parallelize the serial baseline almost linearly, the GPU throughput keeps scaling with population size. This makes whole-nerve re-scoring cheap enough to run inside an optimization loop. Furthermore, the same forward pass is differentiable end-to-end, allowing the stimulus itself to be an optimization variable. Per-contact amplitudes, arbitrary waveforms, and electrode positions can be optimized by gradient descent through the exact cable equation, rather than by grid search or a learned surrogate.

Unlike surrogate models, which are fits to the biophysics, jaxon provides the gradient of the biophysics itself. This allows an optimized stimulus to be validated in the same tool that produced it. The software drops into existing pipelines with minimal friction, as it is validated against the community-standard PyFibers/NEURON reference and released under a permissive MIT license. Its per-contact-lead-field interface connects it to FEM tooling such as golgi, positioning it as an interoperable component in peripheral-nerve modeling pipelines.

The introduction of jaxon changes what is computationally feasible in peripheral-nerve-stimulation modeling. It makes population-scale fiber modeling routine, turning a previously prohibitive whole-nerve re-scoring into a step cheap enough to run inside an optimization loop. This directly enables the central question of a companion study, which investigates whether reducing each fascicle to a representative fiber biases predicted selectivity. That question was previously untestable precisely because the population simulation was too expensive.

By combining validated biophysics with GPU-scale throughput and end-to-end differentiability, jaxon allows neural-interface researchers to ask and answer optimization and population-scale questions that established serial, non-differentiable tools cannot.

The software is distributed under the MIT license and is available through the Python Package Index using the command pip install jaxon-nerve, with the import name jaxon. The current executable version is v1.0.1, and a permanent link to the executables is hosted on Zenodo. The package is platform-independent, supporting both Linux and macOS operating systems. Dependencies are managed through a Conda environment defined in an environment.yml file, or installed directly via pip. The core software stack relies on Python 3.11, JAX 0.6.2, Jaxley 0.13.0, Optax 0.2.8, NumPy 2.2.6, and SciPy 1.15.3. For validation purposes, the package includes NEURON 9.0.1 and PyFibers 0.8.5 as reference tools, while Matplotlib 3.10.9 is used for visualization. An optional NVIDIA GPU with CUDA 12 is required to enable batched GPU execution, which is essential for achieving the reported population-scale throughput.

The package structure is organized by specific functional concerns to facilitate modular use. The jaxon.channels module contains each NMODL ion-channel mechanism hand-translated into a pure-JAX Channel class, with rate constants checked against the published .mod formulas to machine precision. The jaxon.fibers module provides morphology builders for each model, returning a jaxley.Cell and a geometry dataclass that describes the per-compartment static parameters. The jaxon.stim module houses the intracellular and point-source extracellular stimulus helpers, the coupled backward-Euler solver, the vmapped multi-fiber forward pass, and the multi-contact ring-cuff field. Finally, the jaxon.optim module includes the recruitment and activation proxy, the selectivity index, and the Adam optimization loops.

For optimization tasks, jaxon offers distinct gradient strategies depending on the specific stimulus parameters being adjusted. For per-contact rectangular-amplitude optimizations, the software uses a packed finite-difference gradient rather than reverse-mode autodiff. This approach is chosen because the thresholded recruitment objective has an almost-everywhere-zero exact gradient, and the K-contact directional derivative requires only K+1 forward configurations evaluated in a single batched pass. This method avoids the computational overhead of a reverse-mode tape through the ODE scan. Reverse-mode autodiff is retained for arbitrary K×T waveforms, where it remains the only tractable option for differentiating through the complex temporal dynamics. A head-to-head comparison of these two gradient paths is reported in the companion paper, providing detailed insights into their respective efficiencies and accuracy.

The software includes specific functionalities for selectivity scoring and field support. A recruitment proxy and a selectivity index are defined as the difference between the recruited fraction of on-target and off-target fibers. These metrics are usable both as an optimization objective during the design phase and as a post-hoc population score for evaluating existing stimuli. The package supports per-contact lead fields from the golgi platform, which are derived from histology-segmented nerve geometries. It also includes a multi-contact ring cuff with a superposition-based current-steering model. For local demonstrations and testing, a lightweight analytic point-source field is provided.

Three illustrative examples demonstrate the typical usage progression of the software. The first example involves building a single model fiber, placing it in an extracellular field, and computing its activation threshold using the same bisection method used for the NEURON comparison. This unit is the fundamental component of the validation suite. The second example simulates an entire population in a single batched pass using jax.vmap, showcasing the population-scale throughput that enables whole-nerve re-scoring. The third example optimizes a multi-contact stimulus using gradient descent, where the per-contact amplitudes of a cuff are optimized directly against the selectivity objective. These examples are reproducible from the released code, and the REPRODUCE.md file documents the full path from fiber build to every figure and reported number in the companion study.

Subject of Research: Development of a differentiable, GPU-native software package for simulating peripheral nerve fiber models

Article Title: jaxon: a differentiable, GPU-native simulator for peripheral-nerve fiber models

Article References: Lung, D., & Haberbusch, M. (2026). jaxon: a differentiable, GPU-native simulator for peripheral-nerve fiber models. SoftwareX, 36, Article 103196. https://doi.org/10.1016/j.softx.2026.103196

Image Credits: AI Generated

DOI: 10.1016/j.softx.2026.103196

Keywords: peripheral nerve stimulation, GPU computing, differentiable simulation, JAX framework, biophysical modeling, electrode optimization, neural interfaces, open-source software, jaxon, differentiable, GPU-native, simulator