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

Abstract

jaxon is an open-source, fully differentiable and GPU-native reimplementation of the canonical peripheral-nerve fiber models in JAX/Jaxley: the myelinated McIntyre-Richardson-Grill (MRG) and Sweeney axons and the un-myelinated Sundt and Rattay C-fibers. It reproduces NEURONs extracellular mechanism through a custom backward-Euler coupled intracellular/periaxonal double-cable solver, agreeing with PyFibers-wrapped NEURON on 99.6% of 943 activation-threshold configurations within 1% and matching conduction velocity to machine precision. Because the entire forward model is expressed in JAX, it is both vectorized--simulating whole fiber populations in parallel and reaching a geometric-mean [~]820x speedup at N = 100,000 fibers on a single GPU--and differentiable, so extracellular-stimulation parameters (per-contact amplitudes, waveform shape, and electrode position) can be optimized directly through the cable equation rather than grid-searched. jaxon slots into existing peripheral-nerve modeling pipelines as a gradient-enabled, population-scale replacement for the NEURON forward solver.