Home

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

Getting started

Jaxley allows to simulate biophysical neuron models on CPU, GPU, or TPU:

import matplotlib.pyplot as plt
from jax import config

import jaxley as jx
from jaxley.channels import HH

config.update("jax_platform_name", "cpu")  # Or "gpu" / "tpu".

cell = jx.Cell()  # Define cell.
cell.insert(HH())  # Insert channels.

current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0)
cell.stimulate(current)  # Stimulate with step current.
cell.record("v")  # Record voltage.

v = jx.integrate(cell)  # Run simulation.
plt.plot(v.T)  # Plot voltage trace.

If you want to learn more, we have tutorials on how to:

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

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.

License

Apache License Version 2.0 (Apache-2.0)

Citation

If you use Jaxley, consider citing the corresponding paper:

@article{deistler2024differentiable,
  doi = {10.1101/2024.08.21.608979},
  year = {2024},
  publisher = {Cold Spring Harbor Laboratory},
  author = {Deistler, Michael and Kadhim, Kyra L. and Pals, Matthijs and Beck, Jonas and Huang, Ziwei and Gloeckler, Manuel and Lappalainen, Janne K. and Schr{\"o}der, Cornelius and Berens, Philipp and Gon{\c c}alves, Pedro J. and Macke, Jakob H.},
  title = {Differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics},
  journal = {bioRxiv}
}