Overview

Welcome to the world of braincell!

This section provides a brief introduction to some of the key concepts and modeling features of the braincell framework.

braincell is a high-performance computing framework specifically designed for neuron modeling, built on top of JAX and brainstate. It offers a comprehensive toolchain for neuroscientists, computational neuroscientists, and neuromorphic computing engineers to construct, simulate, and optimize electrophysiologically precise models ranging from ion channels to multi-scale neural networks. The framework integrates advanced features such as modern hardware acceleration and automatic differentiation. The following tutorials will walk you through its layered architecture and core functionalities, helping you quickly understand the design logic behind braincell.

import brainunit as u
import brainpy
import braincell
import brainstate
import matplotlib.pyplot as plt

E0812 11:38:41.347369 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 59.46GiB (63839404032 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.348555 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 53.51GiB (57455460352 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.350331 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 48.16GiB (51709911040 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.351318 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 43.34GiB (46538919936 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.353535 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 39.01GiB (41885028352 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.357931 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 35.11GiB (37696524288 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.359740 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 31.60GiB (33926871040 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.362117 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 28.44GiB (30534182912 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.364472 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 25.59GiB (27480764416 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.366473 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 23.03GiB (24732686336 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.370873 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 20.73GiB (22259417088 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.371701 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 18.66GiB (20033474560 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.373213 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 16.79GiB (18030127104 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.373939 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 15.11GiB (16227113984 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.374595 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 13.60GiB (14604402688 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.376286 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 12.24GiB (13143961600 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.377894 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 11.02GiB (11829565440 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.381062 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 9.92GiB (10646608896 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.381780 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 8.92GiB (9581947904 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.383212 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 8.03GiB (8623753216 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
E0812 11:38:41.384493 1801193 cuda_executor.cc:1273] [0] Failed to allocate device memory: INTERNAL: [0] Failed to allocate 7.23GiB (7761377792 bytes) of device memory: : CUDA_ERROR_OUT_OF_MEMORY: out of memory

Hierarchical Architecture

In neuron modeling, models are typically divided into the following major levels based on structural complexity and functional hierarchy:

• Channel

• Ion

• Cell

braincell focuses on precise neuronal dynamics modeling at these three levels.

• Channel: Ion channels on the membrane, controlling the permeability of specific ions and determining the flow of ionic currents.

• Ion: The fundamental charged particles, such as Na⁺, K⁺, and Ca²⁺, that drive changes in membrane potential.

• Cell: The basic modeling unit of a single neuron, integrating multiple channels to produce membrane potential dynamics.

Taking the Hodgkin–Huxley (HH) neuron as an example, the relationship among these three levels can be illustrated as follows:

From this diagram, it is clear that Channel, Ion, and Cell are tightly interconnected. Biologically, these three levels form the fundamental mechanisms underlying neuronal electrical activity:

• A neuron Cell is surrounded by a cell membrane populated with various ion Channels.

• Ion Channels regulate the flow of different Ions across the membrane.

• The movement of Ions changes the potential difference across the membrane, thereby driving neuronal electrical activity.

Together, these levels provide a complete modeling pathway from microscopic mechanisms to macroscopic dynamics of neurons.

Core Features

The main features of braincell include:

• Ion channel modeling: Support for constructing electrophysiologically precise ion channel models using modular components.

• Ion modeling: Support for modular construction of ion dynamics models.

• Neuron modeling: Support for single-compartment and multi-compartment neuron models based on the HH framework.

• Differential equation solvers: High-performance solvers built on JAX.

In the following sections, we will explore how to use and optimize these features, helping you build models ranging from complete neurons down to fine-grained ion mechanisms, and fully leverage the modeling capabilities of braincell.

Single-Compartment Neuron Modeling

We now turn to how braincell can be used to construct neuron models.

Below is an example of building an HH neuron using braincell:

class HH(braincell.SingleCompartment):
    def __init__(self, in_size):
        super().__init__(in_size, C=Cm, solver='ind_exp_euler')

        self.na = braincell.ion.SodiumFixed(in_size, E=50. * u.mV)
        self.na.add(
            INa=braincell.channel.Na_TM1991(in_size, g_max=(100. * u.mS * u.cm **-2) * area, V_sh=-63. * u.mV)
        )

        self.k = braincell.ion.PotassiumFixed(in_size, E=-90 * u.mV)
        self.k.add(
            IK=braincell.channel.K_TM1991(in_size, g_max=(30. * u.mS * u.cm** -2) * area, V_sh=-63. * u.mV)
        )

        self.IL = braincell.channel.IL(
            in_size,
            E=-60. * u.mV,
            g_max=(5. * u.nS * u.cm **-2) * area
        )

We construct the HH neuron model by inheriting from SingleCompartment, which represents a single-compartment neuron. SingleCompartment itself is the base class for all single-compartment neurons and is derived from HHTypedNeuron. Through HHTypedNeuron, both single-compartment neurons (SingleCompartment) and multi-compartment neurons (MultiCompartment) can be modeled.

SingleCompartment provides built-in interfaces for membrane potential while supporting the integration of multiple Channels. For example, INa_TM1991, IK_TM1991, and IL are concrete subclasses of the Channel class. Taking INa_TM1991 as an example, it inherits from the sodium channel base class SodiumChannel, and thus provides the relevant interfaces.

From this example, we can see that modeling a neuron with braincell requires only the composition of the relevant classes at the three levels: Cell, Channel, and Ion. The inheritance relationships within these three levels are deliberately kept simple in braincell.

The neuron Cell regulates the flow of different Ions, and each Ion is controlled by its corresponding Channel. In the HH neuron model above:

• The neuron includes sodium ions (implemented as SodiumFixed) and potassium ions (implemented as PotassiumFixed).

• Sodium ions determine the behavior of the sodium channel INa_TM1991.

• Potassium ions determine the behavior of the potassium channel IK_TM1991.

• In addition, IL represents a leakage channel, which is not associated with any ion.

This design makes neuron model construction both flexible and modular, facilitating extension and modification.

Neural Network Modeling

Building upon single-neuron modeling in braincell, we can further construct neural network models, seamlessly integrating with the network modeling capabilities of BrainState.

Returning to our example, we can use the constructed HH neuron model to build networks for practical applications, such as an excitatory–inhibitory (E–I) network:

V_th = -20. * u.mV
area = 20000 * u.um ** 2
area = area.in_unit(u.cm ** 2)
Cm = (1 * u.uF * u.cm ** -2) * area  # Membrane Capacitance [pF]

class EINet(brainstate.nn.Module):
    def __init__(self):
        super().__init__()
        self.n_exc = 3200
        self.n_inh = 800
        self.num = self.n_exc + self.n_inh
        self.N = HH(self.num)

        self.E = brainpy.state.AlignPostProj(
            comm=brainstate.nn.EventFixedProb(self.n_exc, self.num, conn_num=0.02, conn_weight=6. * u.nS),
            syn=brainpy.state.Expon(self.num, tau=5. * u.ms),
            out=brainpy.state.COBA(E=0. * u.mV),
            post=self.N
        )
        self.I = brainpy.state.AlignPostProj(
            comm=brainstate.nn.EventFixedProb(self.n_inh, self.num, conn_num=0.02, conn_weight=67. * u.nS),
            syn=brainpy.state.Expon(self.num, tau=10. * u.ms),
            out=brainpy.state.COBA(E=-80. * u.mV),
            post=self.N
        )

    def update(self, t):
        with brainstate.environ.context(t=t):
            spk = self.N.spike.value
            self.E(spk[:self.n_exc])
            self.I(spk[self.n_exc:])
            spk = self.N(0. * u.nA)
            return spk

# network
net = EINet()
_ = brainstate.nn.init_all_states(net)

# simulation
with brainstate.environ.context(dt=0.1 * u.ms):
    times = u.math.arange(0. * u.ms, 100. * u.ms, brainstate.environ.get_dt())
    spikes = brainstate.transform.for_loop(net.update, times, pbar=brainstate.transform.ProgressBar(10))

# visualization
t_indices, n_indices = u.math.where(spikes)
plt.scatter(times[t_indices]/u.ms, n_indices, s=1)
plt.xlabel('Time (ms)')
plt.ylabel('Neuron index')
plt.show()

In this way, we have used braincell to model the HH neuron and implemented a complete E–I network!