Jaxley 0.1.1

Jaxley is a differentiable simulator for biophysical neuron models in JAX. Its key features are:

• automatic differentiation, allowing gradient-based optimization of thousands of parameters
• support for CPU, GPU, or TPU without any changes to the code
• jit-compilation, making it as fast as other packages while being fully written in python
• backward-Euler solver for stable numerical solution of multicompartment neurons
• elegant mechanisms for parameter sharing

Tutorials

Tutorials are available on our website. We currently have tutorials on how to:

• simulate morphologically detailed neurons
• simulate networks of such neurons
• set parameters of cells and networks
• speed up simulations with jit and vmap
• define your own channels and synapses
• define groups
• read and handle SWC files
• train biophysical models

Units

Jaxley uses the same units as NEURON.

Installation

Jaxley is available on pypi:

pip install jaxley

This will install Jaxley with CPU support. If you want GPU support, follow the instructions on the JAX github repository to install JAX with GPU support (in addition to installing Jaxley). For example, for NVIDIA GPUs, run

pip install -U "jax[cuda12]"

Feedback and Contributions

We welcome any feedback on how Jaxley is working for your neuron models and are happy to receive bug reports, pull requests and other feedback (see contribute). We wish to maintain a positive community, please read our Code of Conduct.

Acknowledgements

We greatly benefited from previous toolboxes for simulating multicompartment neurons, in particular NEURON.

License

Apache License