Jaxley

Jaxley is a differentiable simulator for biophysical neuron models built on the JAX library. It is designed for gradient-based optimization of large parameter sets and supports multicompartment neurons. Key technical features include:

• Automatic Differentiation: Enables gradient-based optimization of thousands of parameters.
• Hardware Agnostic: Runs on CPU, GPU, or TPU without code changes.
• JIT-Compilation: Leverages JAX's jit for high-performance execution.
• Parameter Sharing: Provides mechanisms for elegant parameter management across models.

Jaxley is designed for biophysical, Hodgkin-Huxley-type models and their networks. It provides high flexibility and accuracy for:

• Single-compartment (point neuron) Hodgkin-Huxley models
• Multi-compartment Hodgkin-Huxley models

Key capabilities for biophysical models include:

• Adding new channel models.
• Using various synapse types (e.g., conductance-based, tanh).
• Inserting different channel types into specific branches or compartments within a single cell.
• Simulating complex ion dynamics such as diffusion and pumps.

Jaxley implements a backward-Euler solver to ensure stable numerical solutions for multi-compartment neurons, similar to the NEURON simulator but built using JAX.

Jaxley supports the simulation of simplified neuron models and their networks, including:

• Leaky-integrate-and-fire (LIF) neurons
• Izhikevich neurons
• Rate-based neurons

Limitation Note: For LIF and Izhikevich models, Jaxley does not currently support surrogate gradient descent. This means these specific models cannot be used for efficient training that relies on backpropagating through discontinuous spikes.

Jaxley provides a series of step-by-step tutorials designed to guide you through the core concepts and usage of the library. It is recommended to follow these tutorials in order to build a foundational understanding of the framework.

Recommended Learning Path:

1. Morph Neurons: Learn about neuron morphology.
2. Small Network: Understand how to construct and manage small-scale neural networks.
3. JIT and Vmap: Learn how to leverage JAX transformations like jit and vmap within the Jaxley ecosystem.
4. Channel and Synapse Models: Explore detailed modeling of ion channels and synaptic connections.
5. Gradient Descent: Understand how to apply optimization techniques to your models.

Install the stable version of Jaxley with CPU support using pip.

If you require GPU support, you must install JAX with GPU support separately following the instructions in the JAX GitHub repository. For NVIDIA GPUs, you can install the CUDA-enabled version of JAX using the command below after installing Jaxley.

pip install jaxley

# For NVIDIA GPU support (install JAX with CUDA)
pip install -U "jax[cuda13]"

To build the Sphinx documentation, use make html followed by a local HTTP server to view the output. This process executes all Jupyter notebooks found in the repository, collects their outputs, and incorporates them into the generated documentation.

make html && python -m http.server --directory _build/html

The official documentation for Jaxley has moved from this site to Read the Docs. For the most up-to-date guides, API references, and tutorials, visit the new documentation site.

https://jaxley.readthedocs.io/en/latest/

Install jaxley via PyPI for default CPU support. If you require GPU support, you must first install JAX with the appropriate CUDA support for your hardware (e.g., NVIDIA GPUs) before or alongside installing jaxley.

# Install Jaxley with CPU support
pip install jaxley

# Example: Install JAX with CUDA 13 support for NVIDIA GPUs
pip install -U "jax[cuda13]"

To preview the documentation locally, navigate to the mkdocs/ directory and run the jupyter nbconvert command to transform tutorials into markdown, followed by the mkdocs serve command to start a local web server.

jupyter nbconvert --to markdown ../docs/tutorials/*.ipynb --output-dir docs/tutorial/
mkdocs serve

To update the hosted documentation on GitHub, navigate to the mkdocs/ directory and run the jupyter nbconvert command to process tutorials, then use mkdocs gh-deploy to push the built site to GitHub Pages.

jupyter nbconvert --to markdown ../docs/tutorials/*.ipynb --output-dir docs/tutorial/
mkdocs gh-deploy

To install Jaxley in editable mode from the source repository, clone the repository and use pip install -e ..

Requirement: You must use pip >= 21.3 to support editable installations using pyproject.toml.

git clone https://github.com/jaxleyverse/jaxley.git
cd jaxley
pip install -e .