Welcome to Jaxley!

Welcome to Jaxley!#

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, check out our material:

🧠 Tutorials

Step-by-step introductions.

Tutorials
⚙️ Advanced tutorials

In-depth guides for power-users.

Advanced tutorials
🧩 How-to guides

Practical recipes for common tasks.

How-to guide
📚 API Reference

Full documentation of modules and functions.

API reference

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.

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}
}
Navigation