CPU、GPU或TPU上的可微分神经元模拟

文档 | 入门 | 安装指南 | 参考文档 | 常见问题

什么是Jaxley？

Jaxley 是一个使用 Python 库 JAX 编写的可微分模拟器，用于 生物物理神经元模型。它的主要特点包括

• 自动微分，允许基于梯度的优化数千个参数
• 支持CPU、GPU或TPU，无需更改代码
• jit-编译，使其在完全用Python编写的条件下，速度与其他软件包相当
• 支持多室神经元
• 优雅的参数共享机制

入门指南

Jaxley 允许在CPU、GPU或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.

这里 可以找到在 Jaxley 中可以实现的模型概述。如果您想了解更多信息，我们建议您查看我们关于如何

• 开始使用 Jaxley
• 模拟神经元网络
• 使用GPU和 jit 加速模拟
• 定义自己的通道和突触
• 计算梯度并训练生物物理模型

安装

Jaxley 可在 PyPI 上获得

pip install jaxley

这将安装具有CPU支持的 Jaxley。如果您想安装具有GPU支持的版本，请按照 JAX Github仓库 中的说明安装具有GPU支持的 JAX（除了安装 Jaxley）。例如，对于NVIDIA GPU，请运行

pip install -U "jax[cuda12]"

反馈和贡献

我们欢迎您对我们如何为您的神经元模型工作提供反馈，并乐于接收错误报告、拉取请求和其他反馈（请参阅贡献）。我们希望保持一个积极的社区，请阅读我们的行为准则。

许可证

Apache License Version 2.0 (Apache-2.0)

引用

如果您使用 Jaxley，请考虑引用相应的论文

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