Author Response: BrainPy, a Flexible, Integrative, Efficient, and Extensible Framework for General-Purpose Brain Dynamics Programming
Chaoming Wang,Tianqiu Zhang,Xiaoyu Chen,Simin He,Shangyang Li,Si Wu
DOI: https://doi.org/10.7554/elife.86365.sa2
2023-01-01
Abstract:Full text Figures and data Side by side Abstract Editor's evaluation Introduction Method and results Discussion Appendix 1 Appendix 2 Appendix 3 Appendix 4 Appendix 5 Appendix 6 Appendix 7 Appendix 8 Appendix 9 Appendix 10 Appendix 11 Data availability References Decision letter Author response Article and author information Metrics Abstract Elucidating the intricate neural mechanisms underlying brain functions requires integrative brain dynamics modeling. To facilitate this process, it is crucial to develop a general-purpose programming framework that allows users to freely define neural models across multiple scales, efficiently simulate, train, and analyze model dynamics, and conveniently incorporate new modeling approaches. In response to this need, we present BrainPy. BrainPy leverages the advanced just-in-time (JIT) compilation capabilities of JAX and XLA to provide a powerful infrastructure tailored for brain dynamics programming. It offers an integrated platform for building, simulating, training, and analyzing brain dynamics models. Models defined in BrainPy can be JIT compiled into binary instructions for various devices, including Central Processing Unit, Graphics Processing Unit, and Tensor Processing Unit, which ensures high-running performance comparable to native C or CUDA. Additionally, BrainPy features an extensible architecture that allows for easy expansion of new infrastructure, utilities, and machine-learning approaches. This flexibility enables researchers to incorporate cutting-edge techniques and adapt the framework to their specific needs. Editor's evaluation The paper introduces a new, important framework for neural modelling that promises to offer efficient simulation and analysis tools for a wide range of biologically-realistic neural networks. It provides convincing support for the ease of use, flexibility, and performance of the framework, and features a solid comparison to existing solutions in terms of accuracy. The work is of potential interest to a wide range of computational neuroscientists and researchers working on biologically inspired machine learning applications. https://doi.org/10.7554/eLife.86365.sa0 Decision letter Reviews on Sciety eLife's review process Introduction Brain dynamics modeling, which uses computational models to simulate and elucidate brain functions, is receiving increasing attention from researchers across different disciplines. Recently, gigantic projects in brain science have been initiated worldwide, including the BRAIN Initiative (Jorgenson et al., 2015), Human Brain Project (Amunts et al., 2016), and China Brain Project (Poo et al., 2016), which are continuously producing new data about the structures and activity patterns of neural systems. Computational modeling is a fundamental and indispensable tool for interpreting this vast amount of data. However, to date, we still lack a general-purpose programming framework for brain dynamics modeling. By general purpose, we mean that such a programming framework can implement most brain dynamics models, integrate diverse modeling demands (e.g., simulation, training, and analysis), and accommodate new modeling approaches constantly emerging in the field while maintaining high-running performance. General-purpose programming frameworks are exemplified by TensorFlow (Abadi et al., 2016) and PyTorch (Paszke et al., 2019) in the field of Deep Learning, which provides convenient interfaces for researchers to define various AI models flexibly and efficiently. These frameworks have become essential infrastructure in AI research, and play an indispensable role in this round of the AI revolution (Dean, 2022). Brain dynamics modeling also needs such a general-purpose programming framework urgently (D’Angelo and Jirsa, 2022). To develop a general-purpose programming framework for brain dynamics modeling, we face several challenges. The first challenge comes from the complexity of the brain. The brain is organized modularly, hierarchically, and across multi-scales (Meunier et al., 2010), implying that the framework must support model construction at different levels (e.g., channel, neuron, network) and model composition across multiple scales (e.g., neurons to networks, networks to circuits). Current brain simulators typically focus on only one or two scales, for example, spiking networks (Gewaltig and Diesmann, 2007; Davison et al., 2008; Beyeler et al., 2015; Stimberg et al., 2019) or firing rate models (Sanz Leon et al., 2013; Cakan et al., 2021). Recently, NetPyNE (Dura-Bernal et al., 2019) and BMTK (Dai et al., 2020a) have adopted descriptive languages to expand the modeling scales from channels to neurons and networks, but their modeling interfaces are still limited to predefined scales. The second challenge is the integration of different modeling needs (Ramezanian-Panahi et al., 2022; D’Angelo and Jirsa, 2022). To elucidate brain functions comprehensively with computational models, we need to not only simulate neural activities, but also analyze the underlying mechanisms, and sometimes, we need to train models from data or tasks, implying that a general-purpose programming framework needs to be a platform to integrate multiple modeling demands. Current brain simulators mainly focus on simulation (Brette et al., 2007; Tikidji-Hamburyan et al., 2017; Blundell et al., 2018), and largely ignore training and analysis. The third challenge is achieving high-running performance while maintaining programming convenience (Tikidji-Hamburyan et al., 2017; Blundell et al., 2018), which is particularly true for brain dynamics modeling, as its unique characteristics make it difficult to run efficiently within a convenient Python interface. The current popular approach for solving this challenge is code generation based on descriptive languages (Goodman, 2010; Blundell et al., 2018). However, this approach has intrinsic limitations regarding transparency, flexibility, and extensibility (Tikidji-Hamburyan et al., 2017; Blundell et al., 2018) (Appendix 1). The fourth challenge comes from the rapid development of the field. Brain dynamics modeling is relatively new and developing rapidly. New concepts, models, and mathematical approaches are constantly emerging, implying that a general-purpose programming framework needs to be extensible to take up new advances in the field conveniently. In this paper, we propose BrainPy (‘Brain Dynamics Programming in Python’, Figure 1) as a solution to address all the above challenges. BrainPy provides infrastructure tailored for brain dynamics programming, including mathematical operators, differential equation solvers, universal model-building formats, and object-oriented JIT compilation. Such infrastructure provides the flexibility for users to define brain dynamics models freely and lays the foundation for BrainPy to build an integrative framework for brain dynamics modeling. First, BrainPy introduces a brainpy.DynamicalSystem interface to unify diverse brain dynamics models. Models at any level of resolution can be defined as DynamicalSystem classes, which further can be hierarchically composed to create higher-level models. Second, BrainPy builds an integrated platform for studying brain dynamics models, where the same BrainPy model can be used for simulation, training (e.g., offline learning, online learning, or backpropagation training), and analysis (e.g., low-dimensional bifurcation analysis or high-dimensional slow point analysis). Third, through JIT compilation and dedicated operators, BrainPy achieves impressive performance for its code execution. The same models can be deployed into different devices (such as Central Processing Unit [CPU], Graphics Processing Unit [GPU], and Tensor Processing Unit [TPU]) without additional code modification. Fourth, BrainPy is highly extensible. New extensions can be easily implemented as plug-in modules. Even the low-level primitive operators in the kernel system can be extended in the user-level Python interface. BrainPy is implemented in a robust continuous integration pipeline and is equipped with an automatic documentation building environment (Appendix 3). It is open sourced at https://github.com/brainpy/BrainPy. Rich tutorials and extensive examples are available at https://brainpy.readthedocs.io and https://brainpy-examples.readthedocs.io, respectively. Figure 1 Download asset Open asset BrainPy is an integrative framework targeting general-purpose brain dynamics programming. (A) Infrastructure: BrainPy provides infrastructure tailored for brain dynamics programming, including NumPy-like operators for computations based on dense matrices, sparse and event-based operators for event-driven computations, numerical integrators for solving diverse differential equations, the modular and composable programming interface for universal model building, and a toolbox useful for brain dynamics modeling. (B) Function: BrainPy provides an integrated platform for studying brain dynamics, including model building, simulation, training, and analysis. Models defined in BrainPy can be used for simulation, training, and analysis jointly. (C) Compilation: Based on JAX (Frostig et al., 2018) and XLA (Sabne, 2020), BrainPy provides just-in-time (JIT) compilation for Python class objects. All models defined in BrainPy can be JIT compiled into machine codes to achieve high-running performance. (D) Device: The same BrainPy model can run on different devices including Central Processing Unit (CPU), Graphics Processing Unit (GPU), or Tensor Processing Unit (TPU), without additional code modification. Method and results Infrastructure tailored for brain dynamics programming To support its goal of becoming a general-purpose programming framework, BrainPy provides the infrastructure essential for brain dynamics modeling (Figure 1A). This infrastructure is a collection of interconnected utilities designed to provide foundational services that enable users to easily, flexibly, and efficiently perform various types of modeling for brain dynamics. Specifically, BrainPy implements (1) mathematical operators for conventional computation based on dense matrices and event-driven computation based on sparse connections; (2) numerical integrators for various differential equations, the backbone of dynamical neural models; (3) a universal model-building interface for constructing multi-scale brain dynamics models and the associated JIT compilation for the efficient running of these models; and (4) a toolbox specialized for brain dynamics modeling. First, BrainPy delivers rich mathematical operators as essential elements to describe diverse brain dynamics models (Appendix 4). On the one hand, brain dynamics modeling involves conventional computation based on dense matrices. In Python scientific computing ecosystem, dense matrix operators have been standardized and popularized by NumPy (Harris et al., 2020), TensorFlow (Abadi et al., 2016), and PyTorch (Paszke et al., 2019). To reduce the cost of learning a new set of computing languages, dense matrix operators in BrainPy (including multi-dimensional arrays, mathematical operations, linear algebra routines, Fourier transforms, and random number generations) follow the syntax of those in NumPy, TensorFlow, and PyTorch so that most Python users can directly program in BrainPy with their familiar operator syntax. On the other hand, brain dynamics modeling has specific computation properties, such as sparse connections and event-driven computations, which are difficult to efficiently implement with conventional operators. To accommodate these needs, BrainPy provides dozens of dedicated operators tailored for brain dynamics modeling, including event-driven operators, sparse operators, and JIT connectivity operators. Compared to traditional dense matrix operators, these operators can reduce the running time of typical brain dynamics models by several orders of magnitude (see Efficient performance of BrainPy). Second, BrainPy offers a repertoire of numerical solvers for solving differential equations (Appendix 5). Differential equations are involved in most brain dynamics models. For ease of use, BrainPy’s numerical integration of differential equations is designed as a Python decorator. Users define differential equations as Python functions, whose numerical integration is accomplished by calling integrator functions, for example, brainpy.odeint() for ordinary differential equations (ODEs), brainpy.sdeint() for stochastic differential equations (SDEs), and brainpy.fdeint() for fractional differential equations (FDEs). These integrator functions are designed to be general, and most numerical solvers for ODEs and SDEs are provided, such as explicit Runge–Kutta methods, adaptive Runge–Kutta methods, and Exponential methods. For SDEs, BrainPy supports different stochastic integrals (Itô or Stratonovich) and different types of Wiener processes (scalar or multi-dimensional). As delays are ubiquitous in brain dynamics, BrainPy also supports the numerical integration of delayed ODEs, SDEs, and FDEs with various delay forms. Third, BrainPy supports modular and composable programming and the associated object-oriented transformations (Appendix 6). To capture the fundamental characteristics of brain dynamics, which are modular, multi-scaled, and hierarchical (Meunier et al., 2010), BrainPy follows the philosophy that ‘any dynamical model is just a Python class, and high-level models can be recursively composed by low-level ones’ (details will be illustrated in Flexible model building in BrainPy). However, such a modular and composable interface is not directly compatible with JIT compilers such as JAX and Numba, because they are designed to work with pure functions (Appendix 2). By providing object-oriented transformations, including the JIT compilation for class objects and the automatic differentiation for class variables, models defined with the above modular and composable interface can also benefit from the powerful transformations in advanced JIT compilers. Fourth, BrainPy offers a toolbox specialized for brain dynamics modeling. A typical modeling experiment involves multiple stages or processes, such as creating synaptic connectivity, initializing connection weights, presenting stimulus inputs, and analyzing simulated results. For the convenience of running these operations repeatedly, BrainPy presets a set of utility functions, including synaptic connection, weight initialization, input construction, and data analysis. However, this presetting does not prevent users from defining their utility functions in the toolbox. Flexible model building in BrainPy Brain dynamics models have the key characteristics of being modular, multi-scaled, and hierarchical, and BrainPy designs a modular, composable, and flexible programming paradigm to match these features. The paradigm is realized by the DynamicalSystem interface, which has the following appealing features. DynamicalSystem supports the definition of brain dynamics models at any organization level. Given a dynamical system, regardless of its complexity, users can implement it as a DynamicalSystem class. As an example, Figure 2A demonstrates how to define a potassium channel model with DynamicalSystem, in which the initialization function defines parameters and states, and the update function specifies how the states evolve. In this process, BrainPy toolbox can help users quickly initialize model variables, synaptic connections, weights, and delays, and BrainPy operators and integrators can support users to define model updating logic freely. In a similar fashion, other dynamical models, such as discontinuous neuron models (e.g., leaky integrate-and-fire model; Abbott, 1999), continuous neuron models (e.g., FitzHugh–Nagumo model; Fitzhugh, 1961), population models (e.g., Wilson–Cowan model; Wilson and Cowan, 1972), and network models (e.g., continuous attractor neural network; Wu et al., 2008), can be implemented by subclassing DynamicalSystem as standalone modules. Figure 2 Download asset Open asset BrainPy supports modular and composable programming for building hierarchical brain dynamics models. (A) An ion channel model is defined as a subclass of brainpy.dynz.IonChannel. The __init__() function specifies the parameters and states, while the update() function defines the updating rule for the states. (B) An Hodgkin–Huxley (HH)-typed neuron model is defined by combining multiple ion channel models as a subclass of brainpy.dyn.CondNeuGroup. (C) An E/I balanced network model is defined by combining two neuron populations and their connections as a subclass of brainpy.DynSysGroup. (D) A ventral visual system model is defined by combining several networks, including V1, V2, V4, TEo, and TEpd, as a subclass of brainpy.DynSysGroup. For detailed mathematical information about the complete model, please refer to Appendix 9. However, for complex dynamical models, such as Hodgkin–Huxley (HH)-typed neuron models or large-scale cortical networks, their model definitions can be achieved through the composition of subcomponents. All models defined with DynamicalSystem can be used as modules to form more complicated high-level models. As an example, Figure 2B demonstrates how an HH-typed neuron model is created by combining multiple ion channel models. Such composable programming is the core of DynamicalSystem, and applies to almost all BrainPy models. For example, a synapse model consists of four components: synaptic dynamics (e.g., alpha, exponential, or dual exponential dynamics), synaptic communication (e.g., dense, sparse, or convolutional connections), synaptic output (e.g., conductance-, current-, or magnesium blocking-based), and synaptic plasticity (e.g., short- or long-term plasticity). Composing different realizations of these components enables to create diverse kinds of synaptic models. Similarly, various network models can be implemented by combining different neuron groups and their synaptic projections. Remarkably, DynamicalSystem supports hierarchical composable programming, such that a model composed of lower-level components can hierarchically serve as a new component to form higher-level models. This property is highly useful for the construction of multi-scale brain models. Figure 2 demonstrates an example of recursively composing a model from channels (Figure 2A) to neurons (Figure 2B) to networks (Figure 2C) and to systems (Figure 2D, see Appendix 9 for details of the full model). It is worth pointing out that this hierarchical composition property is not shared by other brain simulators, and BrainPy allows for flexible control of composition depth according to users’ needs. Moreover, for user convenience, BrainPy provides dozens of commonly used models, including channels, neurons, synapses, populations, and networks, as building blocks to simplify the building of large-scale models. Integrated modeling in BrainPy BrainPy offers an integrated platform to comprehensively perform simulation, training, and analysis of brain dynamics models. Model simulation BrainPy designs the interface brainpy.DSRunner to simulate the dynamics of brain models. DSRunner can be used to simulate models at any level, including but not limited to channel (Figure 3A), neuron (Figure 3B), network (Figure 3C), and system (Figure 3D) levels. Figure 3 Download asset Open asset Model simulation in BrainPy. The interface DSRunner supports the simulation of brain dynamics models at various levels. (A) The simulation of the potassium channel in Figure 2A. (B) The simulation of the HH neuron model in Figure 2B. (C) The simulation of the E/I balanced network, COBAHH (Brette et al., 2007) in Figure 2C. (D) The simulation of the ventral visual system model (the code please see Figure 2D, and the model please see Appendix 9). (E) Using jax.vmap to run a batch of spiking decision-making models (Wang, 2002) with inputs of different coherence levels. The left panel shows the code used for batch simulations of different inputs, and the right panel illustrates the firing rates under different inputs. Brain dynamics models often require intensive parameter searches to fit the experimental data, which is a computationally demanding task. BrainPy facilitates this process by supporting multiple parallel simulation methods. Firstly, the brainpy.running module offers convenient routines for concurrent executions based on the python multiprocessing mechanism. This method is flexible, but may introduce additional time overhead due to the model recompilation and reinitialization in each process. Secondly, most BrainPy models inherently support the automatic vectorization of jax.vmap and automatic parallelization of jax.pmap. These methods can avoid the recompilation and reinitialization of models in the same batch, and automatically parallelize the model execution on the given machines. Figure 3E illustrates the simplicity of this batch simulation approach. By using a single line of functional calls, BrainPy models can run simultaneously with different parameter settings. Model training The use of machine-learning methods to train neural models is becoming a new trend for studying brain functions (Masse et al., 2019; Finkelstein et al., 2021; Laje and Buonomano, 2013; Sussillo et al., 2015; Saxe et al., 2021). BrainPy provides the brainpy.DSTrainer interface to support this utility. Different subclasses of DSTrainer provide different training algorithms, which can be used to train different types of models. For instance, the trainer brainpy.BPTT implements the algorithm of backpropagation through time, which is helpful for training spiking neural networks (Figure 4A) and recurrent neural networks (Figure 4B). Similarly, brainpy.OfflineTrainer implements offline learning algorithms such as ridge regression (Lukoševičius, 2012), brainpy.OnlineTrainer implements online learning algorithms such as FORCE learning (Sussillo and Abbott, 2009), which are useful for training reservoir computing models (Figure 4C). In a typical training task, one may try different algorithms that can be used to train a model. The unified syntax for defining and training models in BrainPy enables users to train the same model using multiple algorithms (see Appendix 10). Figure 4D–F demonstrates that a reservoir network model can be trained with three different algorithms (online, offline, and backpropagation) to accomplish a classical task of chaotic time series prediction (Jaeger, 2007). Figure 4 Download asset Open asset Model training in BrainPy. BrainPy supports the training of brain dynamics models from data or tasks. (A) Training a spiking neural network (Bellec et al., 2020) on an evidence accumulation task (Morcos and Harvey, 2016) using the backpropagation algorithm with brainpy.BPTT. (B) Training an artificial recurrent neural network model (Song et al., 2016) on a perceptual decision-making task (Britten et al., 1992) with brainpy.BPTT. (C) Training a reservoir computing model (Gauthier et al., 2021) to infer the Lorenz dynamics with the ridge regression algorithm implemented in brainpy.OfflineTrainer. x,y, and z are variables in the Lorenz system. (D–F) The classical echo state machine (Jaeger, 2007) has been trained using multiple algorithms to predict the chaotic Lorenz dynamics. The algorithms utilized include ridge regression (D), FORCE learning (E), and backpropagation algorithms (F) implemented in BrainPy. The mean squared errors between the predicted and actual Lorenz dynamics were 0.001057 for ridge regression, 0.171304 for FORCE learning, and 1.276112 for backpropagation. Please refer to Appendix 10 for the training details. Since the training algorithms for brain dynamics models have not been standardized in the field, BrainPy provides interfaces to support the flexible customization of training algorithms. Specifically, OfflineTrainer and OnlineTrainer provide general interfaces for offline and online learning algorithms, respectively, and users can easily select the appropriate method by specifying the fit_method parameter in OfflineTrainer or OnlineTrainer. Furthermore, the BPTT interface is designed to capture the latest advances in backpropagation algorithms. For instance, it supports eligibility propagation algorithm (Bellec et al., 2020) and surrogate gradient learning (Neftci et al., 2019) for training spiking neural networks. Model analysis Analyzing model dynamics is as essential as model simulation and training because it helps unveil the underlying mechanism of model behaviors. Given a dynamical system, BrainPy provides the interface brainpy.DSAnalyzer for automatic dynamic analysis, and different classes of DSAnalyzer implement different analytical methods. First, BrainPy supports phase plane and bifurcation analyses for low-dimensional dynamical systems. The phase plane is a classical and powerful technique for the analysis of dynamical systems and has been widely used in brain dynamics studies, including neuron models (e.g., Izhikevich model; Izhikevich, 2003) and population rate models (e.g., Wilson–Cowan model; Wilson and Cowan, 1972). Figure 5A shows an example where many features of phase plane analysis, including nullcline, vector field, fixed points, and their stability, for a complex rate-based decision-making model (Wong and Wang, 2006) are automatically evaluated by several lines of BrainPy code. Bifurcation analysis is another utility of BrainPy, which allows users to easily investigate the changing behaviors of a dynamical system when parameters are continuously varying. Figure 5B demonstrates the stability changes of the classical FitzHugh–Nagumo model (Fitzhugh, 1961) with one parameter varying can be easily inspected by the bifurcation analysis interface provided in BrainPy. Similarly, bifurcation analysis of codimension-2 (with two parameters changing simultaneously; Figure 5C) can be performed with the same interface. BrainPy also supports bifurcation analysis for three-dimensional fast–slow systems, for example, a bursting neuron model (Rinzel, 1985). This set of low-dimensional analyzers is performed numerically so that they are not restricted to equations with smooth functions, but are equally applicable to ones with strong and complex nonlinearity. Figure 5 Download asset Open asset Model analysis in BrainPy. BrainPy supports automatic dynamics analysis for low- and high-dimensional systems. (A) Phase plane analysis of a rate-based decision-making model (Wong and Wang, 2006). (B) Bifurcation analysis of codimension 1 of the FitzHugh–Nagumo model (Fitzhugh, 1961), in which the bifurcation parameter is the external input Iext. (C) Bifurcation analysis of codimension 2 of the FitzHugh–Nagumo model (Fitzhugh, 1961), in which two bifurcation parameters Iext and a are continuously varying. (D) Finding stable and unstable fixed points of a high-dimensional CANN model (Wu et al., 2008). (E) Linearization analysis of the high-dimensional CANN model (Wu et al., 2008) around one stable and one unstable fixed point. Second, BrainPy supports slow point computation and linearization analysis for high-dimensional dynamical systems. With powerful numerical optimization methods, one can find fixed or slow points of a high-dimensional nonlinear system (Sussillo and Barak, 2013). By integrating numerical methods such as gradient descent and nonlinear optimization algorithms, BrainPy provides the interface brainpy.analysis.SlowPointFinder as a fundamental tool for high-dimensional analysis. Figure 5D demonstrates that the SlowPointFinder can effectively find a line of stable and unstable attractors in a CANN network (Wu et al., 2008). Furthermore, the linearized dynamics around the found fixed points can be easily inspected and visualized with SlowPointFinder interface (Figure 5E). Efficient performance of BrainPy Simulating dynamical models efficiently in Python is notoriously challenging (Blundell et al., 2018). To resolve this problem, BrainPy leverages the JIT compilation of JAX/XLA and exploits dedicated primitive operators to accelerate the model running. JIT compilation In contrast to deep neural networks (DNNs), which mainly consist of computation-intensive operations (such as convolution and matrix multiplication), brain dynamics models are usually dominated by memory-intensive operations. Taking the classical leaky integrate-and-fire (LIF) neuron model (Abbott, 1999) as an example, its computation mainly relies on operators such as addition, multiplication, and division. As shown in Figure 6A, we measure the running times of an LIF model and a matrix multiplication with the same number of floating-point operations (FLOPs) on both CPU and GPU devices. The results indicate that the LIF model is significantly slower than the matrix multiplication on both devices, despite having the same theoretical complexity. This reveals the existence of large overheads when executing brain dynamics models in Python. Moreover, these overheads become dominant when simulating large-scale brain networks, as they grow rapidly with the number of operators in the model. Figure 6 Download asset Open asset BrainPy accelerates the running speed of brain dynamics models through just-in-time (JIT) compilation. (A) Performance comparison between an LIF neuron model (Abbott, 1999) and a matrix–vector multiplication Wv (W∈Rm×m and v∈Rm). By adjusting the number of LIF neurons in a network and the dimension m in the matrix–vector multiplication, we compare two models under the same floating-point operations (FLOPs). The top panel: On the Central Processing Unit (CPU) device, the LIF