Biologically Detailed Brain Cell Modeling in JAX

braincell provides a unified interface for modeling single-compartment and multi-compartment Hodgkin-Huxley-styled neuron models. It is built on top of JAX and brainstate, offering a highly parallelized and efficient simulation of biophysically detailed brain cell models.

Quick start

Here is an example to model the single-compartment thalamus neuron model by using the interface of braincell.SingleCompartment:

import braincell
import brainstate
import brainunit as u

class HTC(braincell.SingleCompartment):
    def __init__(self, size, solver: str = 'ind_exp_euler'):
        super().__init__(size, V_initializer=brainstate.init.Constant(-65. * u.mV), V_th=20. * u.mV, solver=solver)

        self.na = braincell.ion.SodiumFixed(size, E=50. * u.mV)
        self.na.add(INa=braincell.channel.INa_Ba2002(size, V_sh=-30 * u.mV))

        self.k = braincell.ion.PotassiumFixed(size, E=-90. * u.mV)
        self.k.add(IKL=braincell.channel.IK_Leak(size, g_max=0.01 * (u.mS / u.cm ** 2)))
        self.k.add(IDR=braincell.channel.IKDR_Ba2002(size, V_sh=-30. * u.mV, phi=0.25))

        self.ca = braincell.ion.CalciumDetailed(size, C_rest=5e-5 * u.mM, tau=10. * u.ms, d=0.5 * u.um)
        self.ca.add(ICaL=braincell.channel.ICaL_IS2008(size, g_max=0.5 * (u.mS / u.cm ** 2)))
        self.ca.add(ICaN=braincell.channel.ICaN_IS2008(size, g_max=0.5 * (u.mS / u.cm ** 2)))
        self.ca.add(ICaT=braincell.channel.ICaT_HM1992(size, g_max=2.1 * (u.mS / u.cm ** 2)))
        self.ca.add(ICaHT=braincell.channel.ICaHT_HM1992(size, g_max=3.0 * (u.mS / u.cm ** 2)))

        self.kca = braincell.MixIons(self.k, self.ca)
        self.kca.add(IAHP=braincell.channel.IAHP_De1994(size, g_max=0.3 * (u.mS / u.cm ** 2)))

        self.Ih = braincell.channel.Ih_HM1992(size, g_max=0.01 * (u.mS / u.cm ** 2), E=-43 * u.mV)
        self.IL = braincell.channel.IL(size, g_max=0.0075 * (u.mS / u.cm ** 2), E=-70 * u.mV)

Here is an example to model the multi-compartment neuron model by using the interface of braincell.MultiCompartment:

import braincell
import brainstate
import brainunit as u


class HTC(braincell.MultiCompartment):
    def __init__(self, size, solver: str = 'staggered'):
        morphology = braincell.Morphology.from_swc(...)
        super().__init__(size, 
                         morphology=morphology,   # the only difference from SingleCompartment
                         V_initializer=brainstate.init.Constant(-65. * u.mV), 
                         V_th=20. * u.mV, 
                         solver=solver)
        
        self.na = braincell.ion.SodiumFixed(size, E=50. * u.mV)
        self.na.add(INa=braincell.channel.INa_Ba2002(size, V_sh=-30 * u.mV))

        self.k = braincell.ion.PotassiumFixed(size, E=-90. * u.mV)
        self.k.add(IDR=braincell.channel.IKDR_Ba2002(size, V_sh=-30. * u.mV, phi=0.25))

Installation

You can install braincell via pip:

pip install braincell --upgrade

Alternatively, you can install BrainX, which bundles braincell with other compatible packages for a comprehensive brain modeling ecosystem:

pip install BrainX -U

Documentation

The official documentation is hosted on Read the Docs: https://braincell.readthedocs.io

See also the brain modeling ecosystem

We are building the brain modeling ecosystem: https://brainmodeling.readthedocs.io/