Back

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

Lung, D.; Haberbusch, M.

2026-07-31 bioengineering
10.64898/2026.07.30.741846 bioRxiv
Show 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.

Matching journals

The top 2 journals account for 50% of the predicted probability mass.

1
Journal of Neural Engineering
○ 221 papers in training set
Top 0.1%
41.7%
Similar papers in this journal
2
Brain Stimulation
○ 125 papers in training set
Top 0.2%
11.2%
Similar papers in this journal
50% of probability mass above
3
PLOS Computational Biology
● 1863 papers in training set
Top 8%
4.3%
Similar papers in this journal
4
PLOS ONE
● 5266 papers in training set
Top 38%
2.9%
Similar papers in this journal
5
Nature Communications
○ 5641 papers in training set
Top 39%
2.5%
6
Journal of Open Source Software
● 25 papers in training set
Top 0.1%
2.5%
Similar papers in this journal
7
NeuroImage
○ 903 papers in training set
Top 4%
2.0%
8
IEEE Transactions on Biomedical Engineering
● 40 papers in training set
Top 0.7%
1.4%
9
Computer Methods and Programs in Biomedicine
○ 28 papers in training set
Top 0.7%
1.2%
10
Scientific Reports
○ 3612 papers in training set
Top 61%
1.2%
11
Journal of NeuroEngineering and Rehabilitation
○ 36 papers in training set
Top 0.5%
1.2%
12
eneuro
● 439 papers in training set
Top 6%
1.2%
13
Imaging Neuroscience
● 282 papers in training set
Top 3%
1.1%
14
Annals of Biomedical Engineering
○ 37 papers in training set
Top 0.9%
1.0%
15
IEEE Transactions on Neural Systems and Rehabilitation Engineering
● 49 papers in training set
Top 0.8%
1.0%
16
Bioinformatics
◐ 1204 papers in training set
Top 8%
1.0%
17
iScience
○ 1154 papers in training set
Top 33%
0.9%
18
eLife
● 5828 papers in training set
Top 64%
0.9%
19
Bioengineering & Translational Medicine
○ 21 papers in training set
Top 0.5%
0.9%
20
Biomedical Physics & Engineering Express
○ 11 papers in training set
Top 0.3%
0.6%
21
npj Systems Biology and Applications
○ 125 papers in training set
Top 2%
0.6%
22
Nature Biomedical Engineering
○ 47 papers in training set
Top 2%
0.6%
23
Gigabyte
● 62 papers in training set
Top 1%
0.6%
24
Network Neuroscience
● 126 papers in training set
Top 2%
0.6%
25
Frontiers in Physiology
○ 106 papers in training set
Top 3%
0.6%
26
Communications Biology
○ 993 papers in training set
Top 33%
0.6%
27
Neuroinformatics
○ 46 papers in training set
Top 1%
0.5%
28
International Journal for Numerical Methods in Biomedical Engineering
○ 14 papers in training set
Top 0.3%
0.5%
29
Cell Reports
○ 1498 papers in training set
Top 30%
0.5%
30
Advanced Science
○ 286 papers in training set
Top 12%
0.5%

"Similar papers" are the closest papers from that journal in the model's embedding space. They show what the match is built on, but the ranking comes mostly from a classifier over the whole training set, not from these examples alone.