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.

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

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.

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

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

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

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).

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

To overcome this limitation, we employ the JIT compilation technique to dramatically reduce these overhead costs in BrainPy. The JIT compilation transforms the dynamic Python code into the static machine code during runtime, which can significantly reduce the time cost of Python interpretation. Specifically, we utilize JAX, which implements JIT compilation based on XLA (Appendix 2). The XLA JIT engine employs specialized optimizations for memory-intensive operators, for example, operator fusion, which alleviates memory access overhead by minimizing the requirement for intermediate data storage and redundant data transfers during the sequential execution of multiple unmerged operations. This renders the JIT compilation with XLA highly suitable for handling brain dynamics models. Figure 6A demonstrates that with the JIT compilation, the LIF model achieves a running speed comparable to that of the matrix multiplication operation Dot on the CPU and outperforms or matches it on the GPU (see Figure 6A, Appendix 11—figure 6A, and Appendix 11—figure 6B). To further illustrate the benefits of the JIT compilation, we apply it to a realistic brain simulation model, namely, the E/I balanced network model COBA (Vogels and Abbott, 2005). The results show that the JIT compilation boosts the running speed by 10 times on both the CPU and GPU compared to the case without JIT compilation (for CPU acceleration, see Figure 6B and Appendix 11—figure 6C; for GPU acceleration, see Figure 6C and Appendix 11—figure 6D).

Another key feature that distinguishes brain dynamics models from DNNs is that they usually have sparse connections and perform event-driven computations. For example, neurons in a network are typically connected with each other with a probability less than 0.2 (Potjans and Diesmann, 2014), and the state of a postsynaptic neuron is updated only when a presynaptic spike event occurs. These unique features greatly impair the efficiency of brain model simulation using conventional operators, even with the help of JIT compilation. To illustrate this, Figure 7A demonstrates that when implementing a COBA network model with dense matrix-based operators, the majority of simulation time is consumed by synaptic computations on both CPU and GPU devices, and this issue becomes more pronounced as the network size increases (see ‘CPU, Dense’ and ‘GPU, Dense’ lines in Figure 7A).

Figure 7

Download asset Open asset

In order to address this challenge, BrainPy introduces specialized primitive operators designed to accelerate event-based computations within sparsely connected networks. These specialized operators encompass transformations among variables associated with presynaptic neurons, postsynaptic neurons, and synapses, as well as sparse computation operators, event-driven computation operators, and JIT connectivity operators (refer to Appendix 4 for more details). By employing these specialized operators, BrainPy significantly reduces the time required for synaptic computations. As depicted in Figure 7B, the specialized event-based operators result in a remarkable speedup of the classical COBA network model by orders of magnitude (see Appendix 11—figure 7A). Similar speed improvements are observed when utilizing GPU computations, as shown in Figure 7C and Appendix 11—figure 7B. Furthermore, an examination of the time proportion for synaptic computations indicates that the utilization of specialized operators ensures a consistent time ratio for synaptic computation, even as the network size increases (see ‘CPU, Event’ and ‘GPU, Event’ lines in Figure 7A).

Benchmarking

To conduct a formal assessment of the running efficiency of BrainPy, we conducted a comparative analysis against several widely used brain simulators, namely NEURON (Hines and Carnevale, 1997), NEST (Gewaltig and Diesmann, 2007), Brian2 (Stimberg et al., 2019), Brian2CUDA (Alevi et al., 2022), GeNN (Yavuz et al., 2016), and Brian2GeNN (Stimberg et al., 2020). Our benchmarking focused on measuring the simulation speeds of these frameworks for models with sparse and dense connectivity patterns. The tests were performed using three common computing platforms: CPU, GPU, and TPU. This comprehensive assessment provides insights into BrainPy’s efficiency relative to other mainstream simulators across different hardware configurations and network scales.

To evaluate the performance of brain simulators on sparsely connected networks, we utilized two established E/I balanced network models with LIF and HH neuron types: the COBA (Vogels and Abbott, 2005) and COBAHH (Brette et al., 2007) networks (experimental details please see Appendix 11). COBA consists of excitatory and inhibitory LIF neurons with sparse random connectivity. COBAHH uses the same network architecture but replaces the LIF neurons with biophysically detailed HH neuron models. On the CPU platform, consistent with previous benchmark experiments (Stimberg et al., 2019), we find that NEURON and NEST simulators exhibit suboptimal performance when running on a single node (see Figure 8A and Figure 8B). In contrast, BrainPy and Brian2 demonstrate comparable performance, showcasing a remarkable speed advantage of one to two orders of magnitude over NEURON and NEST. As both Brian2 and BrainPy support single-precision floating-point computation (x32), we conducted an analysis of their performances in the context of x32 computation. In order to ensure accurate simulation results with x32 computation, we examined the simulation outcomes across various simulators and platforms (refer to Appendix 11). Our evaluation demonstrated that BrainPy outperforms Brian2 in terms of speedup for numerical integration using x32 arithmetic on CPU platforms. On the GPU platform, GeNN demonstrates optimal linear scaling of execution time on both COBA and COBAHH network models as the network size increases (Figure 8D and Figure 8E). In contrast, BrainPy and Brian2CUDA exhibit a slight overhead and maintain a constant running time when dealing with small network sizes. However, when it comes to network scaling, BrainPy and Brian2CUDA outperform GeNN. Particularly as the network size grows, GeNN exhibits significantly slower performance. Additionally, the utilization of single-precision floating point in GeNN, Brian2CUDA, and BrainPy further enhances their GPU performance (excluding the COBA model in Brian2CUDA). Once again, we observed that BrainPy’s x32 mode achieves a more pronounced performance gain. Particularly, in the COBAHH model, BrainPy’s x32 computation demonstrates a substantial speedup compared to other brain simulators. BrainPy also enables model deployment on TPUs. However, since TPUs currently lack native support for sparse computations and toolchains for operator customization, we could not leverage event-driven sparse operators to simulate the sparsely connected COBA and COBAHH networks. Instead, we used dense matrix multiplication with masking to approximate the sparse connectivity. Unfortunately, this led to significantly slower performance for the two sparsely connected models compared to the results obtained on GPUs (please refer to Appendix 11—figure 8). Moreover, the use of masked matrices resulted in a quadratic increase in memory usage. Consequently, the benchmarking experiments of COBA and COBAHH networks on TPU were limited to a scale of ${\displaystyle 4e^4}$ neurons.

Figure 8

Download asset Open asset

To evaluate the performance of brain simulators on densely connected networks, we utilized the decision-making network proposed by Wang, 2002. Assessing computational efficiency for dense connectivity is important for simulating models that feature dense recurrent connections (Motta et al., 2019) and facilitating the integration with DNNs which commonly employ dense connectivity between layers (Tavanaei et al., 2019). Due to the considerably slower speeds observed and the absence of a publicly available implementation of a decision-making network model using NEURON and NEST, we have excluded them from this benchmark test. Additionally, we did not include a comparison with GeNN because Brian2GeNN does not support the translation of the advanced Brian2 feature employed in this model. Our evaluation showcases that Brian2, Brian2CUDA, and BrainPy exhibit comparable performance on networks of small sizes. However, BrainPy demonstrated substantially better scalability on larger network sizes (see Figure 8C and Figure 8F). For these types of simulation workloads with dense connectivity, TPUs significantly outperformed CPUs and GPUs. Since TPUs primarily utilize low-precision floating point (especially floating point with 16 bits) and are less optimized for double precision, we only tested the model with single-precision operations. Our evaluations clearly showcase the excellent scalability of the network as the size increases (refer to the GPU and TPU comparison in Figure 8F).

Brain science, as well as brain dynamics modeling, is progressing rapidly. Along with the gigantic projects on brain research worldwide, new data and knowledge about brain structures and functions are constantly emerging, which impose new demands on brain dynamics modeling frequently, including, for instance, the simulation and analysis of large-size neural circuits, and the training of neural models based on recorded neural data. To be a general-purpose brain dynamics programming framework, the architecture of the framework must be extensible to conveniently take up new advances in the field. Current brain simulators based on descriptive languages have difficulty achieving this goal, since the extension of a new component through the descriptive interface needs to be done in both high- and low-level programming languages (Appendix 1). Through the elaborate architecture design, BrainPy enables easy extension with new infrastructure, new utility functions, and new machine-learning methods, all performed in our convenient Python interface.

First, for infrastructure (Figure 1A), BrainPy provides a convenient way of customizing a new tool by defining a new subclass. For example, a new Runge–Kutta integrator can be created by inheriting from brainpy.ode.ExplicitRKIntegrator and specifying the Butcher tableau; a new connector can be implemented by deriving from brainpy.conn.TwoEndConnector and overriding initialization function and connection building function (see Appendix 7 for details). Since models and modeling methods have not yet been standardized in the field, the abstraction and summarization of primitive operators for brain dynamics modeling are largely lacking. Although BrainPy has provided dozens of dedicated operators, it would be too soon to establish a complete operator library for brain dynamics modeling. To simplify the process of operator customization, BrainPy provides the brainpy.math.CustomOpByNumba interface that allows users to write and register an operator directly with Python syntax. Specifically, to customize a primitive operator, users need to subclass CustomOpByNumba and implement two Python functions: the abstract evaluation function eval_shape() and concrete computation function con_compute() (see Appendix 8 for more information). Notably, this approach differs from the operator customization in most DNN frameworks, in which low-level operators must be implemented through C++ code. We confirmed that operators customized through the BrainPy interface have comparable and even better performance than those written in C++ (please refer to Figure 9 for the results and Appendix 8 for the source code for comparison).

Figure 9

Download asset Open asset

Second, for functional modules (Figure 1B), BrainPy enables an extension of a new module with BrainPy infrastructure, as the latter can be arbitrarily fused, chained, or combined to create new functions. For example, an analysis toolkit can be customized with BrainPy operators. Moreover, all customizations in BrainPy can benefit from the acceleration of JIT compilation, and users’ attention only needs to focus on the functionalities they require.

Third, for interactions with AI, BrainPy supports the easy extension of new machine-learning methods. Machine-learning approaches are becoming important tools for brain dynamics modeling (Saxe et al., 2021). Existing brain simulators have difficulty incorporating the latest advances in machine-learning research (Appendix 1). Built on top of JAX, BrainPy has the inherent advantage of being linked to the latest developments in machine learning. We noticed that JAX has a rich ecosystem of machine learning, including DNNs, graph neural networks, reinforcement learning, and probabilistic programming. To integrate this rich ecosystem as part of the users’ program, BrainPy is designed to be compatible with other JAX libraries. First, the object-oriented transformations in BrainPy can be applied to pure functions, thus enabling most JAX libraries with a functional programming style to be directly used as a part of the BrainPy program. Second, users can transform models in other libraries as BrainPy objects. For example, using brainpy.dnn.FromFlax, users can treat any artificial neural network model in Flax (Heek, 2020) as a BrainPy module. Alternatively, users can convert a BrainPy model into a format that is compatible with other JAX libraries. For instance, brainpy.dnn.ToFlax supports interpreting a dynamical system in BrainPy as a Flax recurrent cell, so that brain models in BrainPy can also be used in a machine-learning context.

While BrainPy’s native Python-based object-oriented programming paradigm confers numerous advantages compared to existing brain simulators, this novel programming approach also imposes certain limitations that must be acknowledged.

Most existing brain simulators employ a domain-specific language to define brain dynamics models. For example, Brian2 (Stimberg et al., 2019) designs an equation-oriented specification that can describe a wide variety of neural models; NeuroML (Cannon et al., 2014) employs an XML-based specification that facilitates the sharing and reuse of neuronal models; NetPyNE (Dura-Bernal et al., 2019) utilizes a high-level JSON-compatible format composed of Python lists and dictionaries to support multi-scale neuronal modeling; BMTK (Dai et al., 2020a) similarly employs a JSON-based language built on the SONATA file format (Dai et al., 2020b) to deliver consistent multi-resolution experiences via integration with established tools like NEURON and NEST. This declarative programming approach benefits from a clear separation between the mathematical model description and its computational realization. It frees users from low-level implementation details, and enables intuitive specification of complex models in a concise and semantically clear manner. In contrast, the object-oriented programming used in BrainPy exposes the implementation details to users, and adds some complexity to the code. For example, users should be aware of the differences between dense and sparse connectivity schemes, online or offline training schemes, nonbatch or batch computing modes, etc.

The current objectives of BrainPy center on enabling an integrative platform for simulating, training, and analyzing large-scale brain network models while retaining biologically relevant details. Incorporating excessive biological details would be extremely computationally expensive and difficult for such integration. Consequently, detailed spatial modeling with complex compartmental dynamics, as facilitated by tools like NEURON (Hines and Carnevale, 1997) and Arbor (Akar et al., 2019), exceeds BrainPy’s present scope. Moreover, in order to solve the governing partial differential equations, implicit numerical methods (e.g., Crank–Nicolson, implicit Euler) are often essential for stable multi-compartment model simulation. As BrainPy does not currently support fully implicit solvers, it is ill suited to the needs and preferences of modelers focused on multi-compartment dynamics in its current form. Our emphasis remains on balancing biological fidelity and computational tractability for large-scale network modeling and training.

Based on the GSPMD mechanism of the XLA compiler (Xu et al., 2021), the current version of BrainPy supports various parallelism paradigms, such as data parallelism and model parallelism. Data parallelism involves dividing the training data across multiple devices, where each device independently computes and updates the model parameters using its assigned data subset. On the other hand, model parallelism entails partitioning the model across multiple devices, with each device responsible for computing a specific portion of the model computations. These parallelism paradigms are particularly applicable to brain dynamics models with dense connections or structured sparsity. However, the GSPMD parallelism mechanism is not straightforwardly applicable to sparse spiking neural networks, and requires non-trivial changes to support sparse computations. Therefore, another limitation of the current BrainPy framework is that it does not support the general parallelization of sparse spiking neural network models on multiple computing devices. State-of-the-art brain simulators now offer powerful parallelization capabilities for simulating large-scale SNNs. For instance, NEST (Gewaltig and Diesmann, 2007) and NEURON (Hines and Carnevale, 1997) simulators provide convenient and efficient commands through the MPI interface to define, connect, and execute large-scale networks. However, the array-based data structure in BrainPy requires a different approach to parallelize spiking neural networks.

Although BrainPy offers substantial capabilities for brain dynamics modeling, fulfilling all demands in this domain will require large efforts for further ecosystem development.

First, supporting the efficient implementation of multi-compartment neuron models is needed to enable biologically detailed modeling at the subcellular level (Poirazi and Papoutsi, 2020). Multi-compartment neurons incorporate complex dendritic morphologies and spatially distributed ion channels that more precisely capture neural information processing. A substantial number of studies have demonstrated that dendritic mechanisms convey significant advantages to simplified network models of varying levels of abstraction (Bono and Clopath, 2017; Legenstein and Maass, 2011; Wu et al., 2018). Efficiently implementing such models in BrainPy could significantly advance detailed biophysical modeling and bridge the machine-learning-oriented SNN models.

Second, developing parallel primitive operators and memory-efficient algorithms will be critical for ultra-large-scale brain simulations approaching biological realism (>billions of neurons). Massive parallelization across multiple computing devices is currently the main approach to achieve such scale. For instance, the NEST simulator uses optimized data structures and algorithms (Kunkel et al., 2011; Kunkel et al., 2014; Jordan et al., 2018) to enable large-scale simulation on supercomputers and clusters. Moving forward, a priority for BrainPy will be parallelizing its array-based data structures to simulate gigantic brain models across multiple nodes. Moreover, rather than solving large-scale networks exactly, BrainPy aims to find approximating algorithms that overcome the ${\displaystyle O(n^2)}$ complexity, permitting very large-scale modeling on much less computing devices.

Third, integrating BrainPy models with modern accelerators and neuromorphic computing systems (Schuman et al., 2017) could offer a more efficient and scalable approach for simulating large-scale brain dynamics models on cutting-edge hardware accelerators. On the one hand, the implementation of sparse and event-driven operators is necessary for TPUs. While TPUs have demonstrated promising performance and efficiency for machine-learning workloads, our experiments indicate that they are less efficient than GPUs when simulating sparse biological brain network models (see Appendix 11—figure 8). This inefficiency is primarily due to the lack of dedicated operators for sparse and event-driven brain computations in current TPUs. In the future, we plan to explore the development of TPU kernels to enable scalable and efficient brain dynamics programming on TPU hardware accelerators. On the other hand, neuromorphic systems incorporate custom analog circuits that mimic neurobiological architectures and dynamics, resulting in significantly higher power efficiency compared to conventional digital hardware. By mapping BrainPy models onto neuromorphic platforms, simulations can be accelerated, and large-scale models can be executed efficiently. However, the development of translation tools and mapping optimizations is necessary to fully harness the potential of these systems.

By addressing these limitations and enhancing BrainPy’s capabilities in these areas, we can further advance its goal of serving as a comprehensive programming framework for modeling brain dynamics. This will enable users to delve into the dynamics of brain or brain-inspired models that combine biological insights with machine learning. The BrainPy team encourages collaboration with the research community to expand this modeling ecosystem and facilitate a deeper understanding of brain dynamics.

In general, the existing tools for brain dynamics programming can be roughly divided into two categories: low-level programming and descriptive language.

The representatives of the first category include NEURON (Hines and Carnevale, 1997), NEST (Gewaltig and Diesmann, 2007), CARLsim (Beyeler et al., 2015; Chou et al., 2018), NeuronGPU (Golosio et al., 2021), Arbor (Akar et al., 2019), and others. These simulators offer a library of standard models written in C/C++ (particularly for NEST), CUDA (for CARLsim, NeuronGPU, and Arbor), or domain-specific languages (for NEURON and Arbor) to ensure efficient execution, along with a user-friendly Python interface for ease of use. Users can create neural networks in Python by utilizing the neuron and synapse models provided in the library. However, when a new model is required, users must learn to program using the low-level language. This significantly increases the learning cost and restricts the flexibility in defining new models (Tikidji-Hamburyan et al., 2017).

The second category tools include Brian (Goodman and Brette, 2008), Brian2 (Stimberg et al., 2014; Stimberg et al., 2019), ANNarchy (Vitay et al., 2015), GeNN (Yavuz et al., 2016), BMTK (Dai et al., 2020a), NetPyNE (Dura-Bernal et al., 2019), NeuroML (Gleeson et al., 2010), and NMODL (Krzhizhanovskaya et al., 2020), which employ a code generation approach based on descriptive languages. Descriptive simulators allow users to create new models based on convenient descriptions (such as text, Goodman and Brette, 2008; Stimberg et al., 2014; Stimberg et al., 2019; Vitay et al., 2015; JSON, Dai et al., 2020a; Dura-Bernal et al., 2019; XML files, Gleeson et al., 2010; or customized languages, Krzhizhanovskaya et al., 2020) and then translate the descriptions into low-level codes to speed up model running. In such a way, descriptive simulators enable model customization based on high-level descriptive languages and ensure efficient running by generating low-level codes.

Currently, descriptive language has become a standard approach for brain simulation (Blundell et al., 2018). Simulators employing this first approach also start to provide their descriptive language interface for code generation. For instance, the NEST simulator provided its domain-specific language NESTML (Plotnikov et al., 2016) to describe stereotypical neuron and synapse models. Similarly, the NEURON simulator recently released its modern descriptive interface NetPyNE (Dura-Bernal et al., 2019), which employs the standard JSON format to allow users to describe neural circuit models by composing the existing available NEURON building block models. However, users still need to code based on its low-level programming interface to customize new models for channels, neurons, or synapses.

Descriptive languages have been highly successful for brain simulation (Blundell et al., 2018). A major benefit of these declarative approaches is the clear separation between model specification and implementation. This frees users from low-level programming details, enabling them to intuitively specify complex models in a concise, semantically clear way. The declarative nature of descriptive languages allows modelers to focus on computational neuroscience instead of implementation specifics. This has enabled rapid prototyping and sharing of models in the field. Overall, descriptive simulation languages have greatly improved modeler productivity thanks to their high-level, implementation-agnostic nature. However, they have intrinsic limitations on transparency, extensibility, and flexibility. One prominent feature of these descriptive languages is that they completely separate the model definition from the simulation, and therefore are not directly executable (Blundell et al., 2018). This kind of programming paradigm will cause great restrictions on usability and flexibility because it disables model debugging, error correction, and direct logic controlling. Moreover, descriptive languages are usually designed for specific kinds of models or one particular modeling approach. They are written in more than two programming languages: one is based on a low-level language (e.g., C++, CUDA) to implement its core functionality, and the other is based on a high-level language (e.g., Python, Matlab) for ease of use. Once they are not tailored to users’ needs, extensions to accommodate new components must be made in both high- and low-level languages, which is hard or nearly impossible for normal users. What is more, descriptive languages greatly reduce the expressive power of a general-purpose programming language and make it hard to describe all aspects of a simulation experiment, including clipping variables out of bounds, input–output relations, model debugging, code optimization, dynamics analysis, and others.

In summary, significant challenges to transparency, flexibility, efficiency, and extensibility are still present in the existing programming paradigm for brain simulation. We can draw the conclusion that current software solutions cannot lead us to a general-purpose programming framework that allows us to freely define brain dynamics models in various application domains.

A notorious challenge in scientific computing is the trade-off between usability and efficiency. The former seeks the fast prototyping of thoughts and ideas, whereas the latter pursues efficient code execution. For a long time in the past, it was difficult to strike a balance between the two. For example, statically typed compiled programming languages such as C or C++ are incredibly efficient in code execution, but their productivity is relatively low due to their complex and heavy syntax. In contrast, dynamically typed interpreted programming languages like Python and R are easy to learn and use, but they have slow running speeds. Nowadays, with the increasing complexity of models, the demand for both usability and efficiency has increased dramatically. Fortunately, recent advancements in just-in-time (JIT) compilation technology (Lattner and Adve, 2004; Lattner et al., 2020) have provided viable answers to this two-language problem. In particular, a new generation of computational engines based on JIT compilation (Aycock, 2003) has begun to have an impact on a variety of scientific computing disciplines.

The JIT compilation can be seen as the combination of the statically typed compilation and dynamically typed interpretation. It benefits from both the convenience of dynamic high-level languages like Python and the efficiency of static low-level languages such as C++. At the start of a program execution, a JIT compiler acts like an interpreter. It runs your code step by step, and can output the intermediate results at the run-time for debugging. However, if some hot code snippets that are executed frequently, for example, certain functions or loop bodies, are detected or manually labeled, they will be submitted to the JIT compiler for compilation and storage. In this sense, it acts like a statically typed compiler. Once the compiled code snippets are entered again, the program will directly execute the compiler-generated low-level code without time-consuming interpreting again. Hot code snippets can be automatically detected by the JIT compiler, or be manually labeled by users.

JIT compilation is a mature and well-established technology. It has been adopted in modern programming languages like JAVA, Julia, and Python. JAVA language provides JIT compilation in its JAVA virtual machine (JVM) to accelerate the execution of JAVA code (Grcevski et al., 2004). Java source code is first compiled into the platform-independent Java bytecode (.class file). Then, JVM loads the bytecode, interprets it, and executes it. To increase the running speed, JVM detects code that is frequently called through the hotspot detection and submits its bytecode to the JIT compiler to compile them into machine code. For the code with lower frequency, executing it through the JVM interpretation can save the time of the JIT compilation; while for the hot code frequently called, JIT compilation can significantly improve the running speed after the code is compiled. However, compared with Python, JAVA has poor ecosystem support for numerical computing. Its JIT is not specialized in numerical computing, but in general domains.

Julia (Bezanson et al., 2017), another dynamic high-level programming language, is recently proposed for high-performance scientific computing. Julia features intuitive, productive, and general-purpose syntax inspired by the success of Python, Matlab, and C++. Moreover, it achieves attractive performance through the JIT compilation based on the LLVM compiler infrastructure (Lattner and Adve, 2004). In a remarkably short time, Julia has provided excellent routines for mathematical functions, machine-learning algorithms, data processing tools, visualization utilities, and others. However, Julia is still young. Costs, such as lack of familiarity, rough edges, correctness bugs, and continual language changes, are still imposed on normal users.

Python is a well-known and popular interactive dynamic programming language. It has a long history in numerical computing (Dubois et al., 1996; Harris et al., 2020). The ecosystem of scientific computing, including array programming (Harris et al., 2020), scientific algorithms (Virtanen et al., 2020), machine learning (Pedregosa et al., 2011), deep learning (Abadi et al., 2016; Paszke et al., 2019; Frostig et al., 2018), image processing (van der Walt et al., 2014), data analysis and statistics (McKinney, 2011), network analysis (Hagberg and Swart, 2008), visualization (Hunter, 2007), and many others, has been well established in Python. Before Julia, the JIT compilation was introduced into Python by PyPy in 2007. Later, other attempts, including Pyston, Pyjion, Psyco, JitPy, HOPE, etc., are proposed. With a long history of JIT development, Python nowadays has provided mature platforms of JIT compilation focusing on numerical computing. These numerical JIT platforms include Numba (Lam et al., 2015), JAX (Frostig et al., 2018), and XLA (Sabne, 2020). Each of them has its own characteristics.

JAX (Frostig et al., 2018) is a flourishing machine-learning library developed by Google. It aims to provide high-level numerical functions to help users fast prototype machine learning ideas. Moreover, these numerical functions can benefit from powerful functional transformations, like automatic differentiation grad, JIT compilation jit, automatic vectorization vmap, and parallelization pmap. JAX makes heavy use of XLA (Sabne, 2020) (see the following text) for code optimization. Specifically, for ease of use, high-level numerical functions in JAX are NumPy like. JAX provides many numerical functions in NumPy, including basic mathematical operators, linear algebra functions, and Fourier transform routines. However, some fundamental designs are significantly different from NumPy, for instance, the well-established syntax for in-place updating and random samplings. This is the reason why we provide another set of numerical functions consistent with NumPy. In addition to its NumPy-like API, JAX provides a wonderful set of composable functional transformations. Among them, automatic differentiation in JAX supports both forward and backward modes for arbitrary numerical functions. It can take derivatives of a function with a large subset of Python syntax, including loops, conditions, recursions, and closures. Moreover, JAX utilizes XLA to JIT compile your Python code on modern devices, like Central Processing Units (CPUs), Graphics Processing Units (GPUs), and Tensor Processing Units (TPUs). It can significantly accelerate the execution speed of your code and allows you to get maximal performance without having to leave Python. JAX also provides automatic vectorization or batching. It supports transforming loops to vector operations via a single functional call vmap. What’s more, JAX delivers pmap to express single-instruction multiple-data programs. Applying pmap to a function will JIT compile and execute the code in parallel on XLA devices, like multiple GPUs or TPU cores. Similar to vmap, pmap transformation maps a function over array axes. But what is different is that the former vectorizes functions by compiling the mapped axis as primitive operations, whereas the latter replicates the function and runs each replica on its own XLA device in parallel. Automatic differentiation and compilation in JAX can be composed arbitrarily to enable rapid experimentation of novel algorithms.

XLA (Sabne, 2020) is a domain-specific linear algebra compiler developed by Google that aims to improve the execution speed, memory usage, portability, and mobile footprint reduction of machine-learning algorithms. XLA compiler provides support for JIT compilation based on LLVM (Lattner and Adve, 2004). The front-end program (e.g., JAX) which wants to take advantage of the JIT compilation of XLA should first define the computation graph as ‘High-Level Optimizer IR’ (HLO IR). Then, XLA takes this graph defined in HLO IR and compiles it into machine instructions for different backend architectures. Currently, XLA supports JIT compilation on backend devices of x86-64 CPUs, NVIDIA GPUs, and Google TPUs. XLA is designed for easy portability on new hardware. It provides an abstract interface that a new hardware device can implement to create a backend to run existing computation graphs. Instead of implementing every existing operator for new hardware, XLA provides a simple and scalable mechanism that can retarget different backends. This advantage may be valuable for neuromorphic computing (Schuman et al., 2017), because new neuromorphic hardware can be interfaced as a new backend of XLA computation.

Numba (Lam et al., 2015) is a JIT compiler for numerical functions in Python. Similar to JAX, Numba supports the JIT compilation of a subset of Python and NumPy code based on the LLVM compiler. However, different from JAX which accelerates the computation flow composed of high-level operators, Numba pays more attention to the acceleration of loop-based functions. Numba achieves excellent optimizations on Python functions with a lot of loops. It allows users to write fine-grained code with native Python control flows and meanwhile obtains the running speed approaching C. This is a huge advantage compared to JAX because JAX does not support the automatic vectorization of a for-loop function.

To ensure any code changes do not introduce unintended bugs or side effects, we standardized the development process and enabled the automatic continuous integration (CI) of BrainPy with GitHub Actions. Moreover, to update the tutorial and documentation with the latest code changes, we automated the documentation building of BrainPy with Read the Docs. Appendix 3—figure 1 illustrates the whole workflow of BrainPy development. First, any code change should be proposed through GitHub Pull Request. Once a Pull Request is opened, CI pipelines are triggered to test BrainPy codes on Windows, Linux, and macOS platforms (Appendix 3—figure 1 (1)). After all test suites are passed, the code reviewer should manually inspect the significance and correctness of the proposed code changes again. If all things are fine, the reviewer can merge the code into the master branch (Appendix 3—figure 1 (2)). After merging, a new set of test cases is triggered automatically to test that the latest BrainPy codebase does not have bugs (Appendix 3—figure 1 (3)). Besides, the merging operation also triggers the automatic documentation generation through the documentation hosting platform Read the Docs (Appendix 3—figure 1 (4)), in which the latest documentation, including the code annotation, user manual, and tutorials, are automatically built with Sphinx and hosted online at https://brainpy.readthedocs.io/.

Appendix 3—figure 1

Download asset Open asset

Brain dynamics modeling involves conventional computation based on dense matrices and event-driven computation based on sparse connections. BrainPy provides operators for these two kinds of computations. For the list of the number of currently implemented operators please see Appendix 4—table 1.

JAX (Frostig et al., 2018) has provided most numerical functions in NumPy (Harris et al., 2020). However, there are several significant differences between JAX and NumPy. First, the array structure in JAX does not support in-place updating. Second, numerical functions that need in-place updating are missing in JAX. Third, random sampling functions are significantly different from NumPy.

The dense matrix operators in BrainPy are based on JAX’s implementations but are designed to be seamlessly consistent with NumPy. First, we provide brainpy.math.Array which is consistent with NumPy’s ndarray structure, and a series of mathematical operations for brainpy.math.Array which is similar to those for ndarray. Second, mathematical operators for brainpy.math.Array, such as indexing, slicing, sorting, rounding, arithmetic operations, linear algebraic functions, and Fourier transform routines, are all supported. Many of these operators (nearly 85%) are directly implemented through the NumPy-like functions in JAX, while BrainPy provides dozens of APIs missing or inconsistent in JAX. Third, to unify random number generators, BrainPy implements most of the random sampling functions in NumPy, including its univariate distributions, multivariate distributions, standard distributions, and utility functions.

Moreover, BrainPy is working on dense operators provided in PyTorch (Paszke et al., 2019) and TensorFlow (Abadi et al., 2016) libraries. In the future, BrainPy will continue to cover dense array operators in TensorFlow and PyTorch, since these implementation syntaxes have been widely accepted in the community.

Brain dynamics models differ from deep neural network (DNN) models in the way they perform computation. Brain dynamics models typically have sparse connections (less than 20% of neurons are connected to each other) and perform event-driven computations (synaptic currents are only transmitted when a presynaptic neuron spikes). These unique features make brain dynamics models less efficient when conventional dense array operators are used. To tackle this efficiency issue, traditional brain simulators heavily rely on event-driven synaptic operations. Previous works have explored event-driven synaptic operations on both CPU platforms (see Vitay et al., 2015; Plesser et al., 2007; Stimberg et al., 2019) and GPU platforms (see Fidjeland et al., 2009; Brette and Goodman, 2011; Yavuz et al., 2016; Alevi et al., 2022).

Despite the effectiveness of these simulators, one limitation is the lack of abstraction of event-driven synaptic operations as primitive low-level operators. In other words, these operations are not treated as fundamental building blocks that can be easily manipulated and optimized. This absence of abstraction hampers the development of more efficient algorithms and restricts the flexibility and extensibility of the simulators. Therefore, it is crucial to bridge this gap and provide a higher level of abstraction for event-driven synaptic operations, which we refer to as ‘event-driven operators’. Note here in BrainPy, event-driven operators are employed within a clock-driven simulation schema, where the simulation advances in a synchronized manner, updating all neurons and synapses at each time step. These event-driven operators process information based on the detection of spatial spiking events. They execute computations when the presynaptic neuron fires spikes at each time step. This contrasts with an event-driven simulation approach (Ros et al., 2006), where neuronal or synaptic state updates are triggered by temporal spiking events, rather than updating all elements at each time step. Therefore, the key distinction between our event-driven operators and an event-driven simulation scheme lies in their scope and application: event-driven operators serve as fundamental building blocks that define how individual components of an SNN respond to spatial events at the current time step, while an event-driven simulation scheme serves as a methodology for simulating the collective behavior of the entire network based on the occurrence of temporal spikes.

Moreover, by abstracting these operations into primitive low-level operators, BrainPy offers automatic differentiation support compatible with jax.grad. It also enables vectorization and parallelization support compatible with jax.vmap and jax.pmap. This compatibility further enhances the applicability of event-driven operators across a wider range of contexts.

Specifically, BrainPy provides these sparse and event-driven operators in the following modules: (1) The brainpy.math.sparse module. This module provides a set of sparse operators that can store synaptic connectivity compactly and compute synaptic currents efficiently. (2) The brainpy.math.event module. This module implements event-driven operators that only perform computations when spikes arrive. This can lead to significant improvements in efficiency, as the state of the system does not need to be constantly updated when no spikes arrive.

In addition, efficient modeling of brain dynamics encounters scalability issues. The computational resources and device memory requirements for brain dynamics models increase quadratically with the number of neurons, as the synaptic connectivity grows almost quadratically in relation to the number of neurons. This characteristic severely restricts the size of the network that can be simulated on a single device using traditional array programming solutions.

One established approach to address this challenge is the utilization of JIT connectivity (Lytton et al., 2008; Azevedo Carvalho et al., 2020; Knight and Nowotny, 2021). This method involves regenerating the synaptic connectivity during computation by controlling the same random seed. Specifically, the brainpy.math.jitconn module provides JIT connectivity as primitive operators, specifically designed for cases where synaptic connections follow a fixed connectivity rule and do not require modifications after initialization. These operators eliminate the memory overhead required for connectivity storage, as synaptic connectivity can be generated JIT during the execution of the operators. Notably, when compared to conventional operators, they enable the execution of large-scale networks that are two to three orders of magnitude larger on a single device.

Moreover, BrainPy also provides operators that are involved in transformations among presynaptic neuronal data, synaptic data, and postsynaptic neuronal data in the brainpy.math module. Specifically, this module provides operators for mapping presynaptic neuronal data to its connected synaptic dimension (pre-to-syn), arithmetic operators including summation, product, max, min, mean, and softmax for transforming synaptic data to postsynaptic neuronal dimension (syn-to-post), and arithmetic operators such as summation, product, max, min, and mean for directly computing postsynaptic data from presynaptic events (pre-to-post).

To meet the large demand for solving differential equations in brain dynamics modeling, BrainPy implements numerical solvers for ODEs, SDEs, fractional differential equations (FDEs), and delay differential equations (DDEs). In general, numerical integration in BrainPy defines the system evolution of ${\displaystyle x(t)\to x(t+dt)}$, where ${\displaystyle x}$ is the state, ${\displaystyle t}$ is the current time, and ${\displaystyle dt}$ is the integration step.

1 integral=brainpy.odeint(f=<function>, 
2                              method=<str>, 
3                              dt=<float>)

1 integral = brainpy.sdeint(f=<function > , 
2                                 g=<function > , 
3                                 method =<str > , 
4                                 dt =<float > , 
5                                 wiener_type=<’scalar’ or ’vector’>, 
6                                 intg_type=<’Ito’ or ’Stratonovich’>)

FDE numerical solvers

The numerical integration of FDEs is very similar to that of ODEs, except that the initial value, memory length, and fractional order should be provided. Given the fractional-order differential equation

$\frac{{\mathrm{d}}^{\alpha }x}{\mathrm{d}{t}^{\alpha }}=F(x,t),$

where the fractional order ${\displaystyle 0<\alpha \le 1}$. BrainPy supports its numerical integration with the following format of

1 brainpy.fdeint(f=<function> , 
2                  alpha=<float> , 
3                  num_memory=<int> , 
4                  inits=<dict> , 
5                  method=<str> )

BrainPy supports two types of FDE, that is, the Caputo derivative and Grünwald–Letnikov derivative. Caputo derivatives are widely used in neuroscience modeling (Kaslik and Sivasundaram, 2012). However, the numerical integration of Caputo derivatives has a high memory consumption, because it requires integration over all past activities. This implies that FDEs with the Caputo derivative definition cannot be used to simulate realistic neural systems. However, the numerical method for Grünwald–Letnikov FDEs, brainpy.fde.GLShortMemory, is highly efficient because it does not require an infinite memory length for numerical solutions. With the increasing length of short memory, the accuracy of the numerical approximation will increase. Therefore, brainpy.fde.GLShortMemory can be applied to real-world neural modeling. See Appendix 5—table 3 for the full list of FDE solvers currently implemented in BrainPy.

Appendix 5—table 3

Derivative type	Solver name	Keyword
Caputo derivative	L1 schema	l1
	Euler method	euler
Grünwald–Letnikov derivative	Short Memory Principle	short-memory

DDE numerical solvers

Delays occur in any type of differential equation. In realistic neural modeling, delays are often inevitable. BrainPy supports equations with variables of constant delays, like

$y^{\prime }(t)=f(t,y(t),y(t-{\tau }_1),y(t-{\tau }_2),\mathrm{\dots },y(t-{\tau }_k)),$

where the time lags ${\displaystyle {\tau }_j}$ are the positive constants. It also supports systems with state-dependent delays, for example,

$y^{\prime }(t)=f(t,y(t),y(t-f_1(y(t))),\mathrm{\dots },y(t-f_k(y(t)))),$

where ${\displaystyle f_k}$ is the function that computes the delay length by the system state ${\displaystyle y(t)}$. For neutral-typed equations in which delays appear in derivative terms,

$y^{\prime }(t)=f(t,y(t),y^{\prime }(t-{\tau }_1),y^{\prime }(t-{\tau }_2),\mathrm{\dots },y^{\prime }(t-{\tau }_k))$

BrainPy also supports its numerical integration.

BrainPy, in particular, implements interfaces to define these various delay variables. brainpy.math.TimeDelay and brainpy.math.LengthDelay are provided to support state-dependent variables. Both brainpy.math.TimeDelay and brainpy.math.LengthDelay are used to delay neuronal signals in BrainPy, for example, Boolean spikes, real-valued synaptic conductance, or membrane potentials.

The TimeDelay is the class representing a connection delay measured in time units. For example, one might specify a ${\displaystyle T}$ ms delay. The TimeDelay class stores history values at times ${\displaystyle [T_0-T,T_0-T+\mathrm{\Delta }t,\cdots ,T_0-2\mathrm{\Delta }t,T_0-\mathrm{\Delta }t,T_0]}$, where ${\displaystyle T_0}$ is the current time and ${\displaystyle \mathrm{\Delta }t}$ is the time step. Users can retrieve the history values for any ${\displaystyle t}$ in the interval ${\displaystyle [T_0-T,T_0]}$. For a ${\displaystyle t}$ that exactly matches one of the stored time points, TimeDelay directly returns the stored history values; otherwise, TimeDelay uses linear interpolation to estimate values between the nearest stored data points, or rounding to return the value from the nearest data point in time.

The LengthDelay class represents a connection delay based on the number or length of running iterations. For instance, one can specify a delay of ${\displaystyle L}$ iterations. It stores historical values at previous time steps, such as ${\displaystyle [L,L-1,\cdots ,1,0]}$. Unlike TimeDelay, users can only retrieve the historical values at specific discrete time steps that have been stored.

Despite the distinct usage characteristics of TimeDelay and LengthDelay, both of them employ a circular array for their delay updates. This implementation involves utilizing an array of fixed length along with a cursor that indicates the current updating position within the array (see Appendix 5—figure 1).

Appendix 5—figure 1

Download asset Open asset

The classes brainpy.math.NeuTimeDelay and brainpy.math.NeuLenDelay are identical to brainpy.math.TimeDelay and brainpy.math.LengthDelay, respectively. However, they are specifically designed to model neutral delay variables.

With these delay supports, the differential equations with delayed variables are intrinsically supported in each integral function.

For delayed ODEs, users can use

1 brainpy.odeint(f=<function> , 
2                  method=<str> , 
3                  dt=<float> , 
4                  state_delays=<dict> , 
5                  neutral_delays=<dict>)

Similarly, the numerical integration of delayed SDEs should utilize

1 brainpy.sdeint(f=<function> , 
2                  g=<function> , 
3                  method=<str> , 
4                  dt=<float> , 
5                  state_delays=<dict>)

For delayed FDEs, one can use

1 brainpy.fdeint(f=<function> , 
2                  method=<str> , 
3                  dt=<float> , 
4                  num_memory=<int> , 
5                  inits=<dict>, 
6                  state_delays=<dict>)

Note that we currently do not support neutral delays for delayed SDEs and FDEs. However, users can easily customize their supports for equations with neutral delays.

Under minimal constraints and assumptions, BrainPy enables the JIT compilation for class objects. These assumptions include the following:

• The class for JIT must be a subclass of brainpy.BrainPyObject.

• Dynamically changed variables in the class must be labeled as brainpy.math.Variable. Otherwise, they will be compiled as constants and cannot be updated during the program execution.

To take advantage of the JIT compilation, we can directly apply brainpy.math.jit() onto the instantiated class objects, or functions of a class object.

brainpy.math.grad() takes a function/object (${\displaystyle f:{\mathbb{R}}^n\to \mathbb{R}}$, which returns a scalar value) as the input and returns a new function (${\displaystyle \mathrm{\partial }f(x)\to {\mathbb{R}}^n}$) which computes the gradient of the original function/object.

1 grad_f=brainpy.math.grad(f, grad_vars=<variables to take gradients>)

brainpy.math.vector_grad() takes vector-valued gradients for a function/object (${\displaystyle f:{\mathbb{R}}^n\to {\mathbb{R}}^m}$). It returns a new function (${\displaystyle \mathrm{\partial }f(x):{\mathbb{R}}^m\to {\mathbb{R}}^n}$) which evaluates the vector-Jacobian products.

1 grad_f=brainpy.math.vector_grad(f, grad_vars<variables to take gradients>)

Another way to take gradients of a vector-output value is using brainpy.math.jacobian(). It aims to automatically compute the Jacobian matrices ${\displaystyle \mathrm{\partial }f(x)\in {\mathbb{R}}^{m\times n}}$ by the given function/object ${\displaystyle f:{\mathbb{R}}^n\to {\mathbb{R}}^m}$ at the given point of ${\displaystyle x\in {\mathbb{R}}^n}$.

1 grad_f=brainpy.math.jacobian(f, grad_vars<variables to take gradients>)

BrainPy features extensible architecture. New extensions can be easily made by using BrainPy infrastructure. Even new tools at the infrastructure level can be customized by using BrainPy operators.

At the toolbox level, BrainPy provides a mechanism in which the extension of a new tool can be made by adding a new subclass. For instance, a new synaptic connection method can be extended by subclassing brainpy.connect.TwoEndConnector:

1 class NewConnector(bp.conn.TwoEndConnector): 
2   def __init__(self): 
3    # initializing connector 
4 
5   def build_csr(self): 
6    # build the CSR data structure 
7 
8   def build_coo(self): 
9    # build the COO data structure

As another example, customizing a new weight initialization method can be added by inheriting from brainpy.initialize.Initializer base class:

1 class NewInitializer(bp.init.Initializer): 
2   def __init__(self): 
3    # initializing connector 
4 
5   def __call__(self, shape, dtype =None): 
6    # building weights

At the model building level, BrainPy enables to flexibly customize the user’s own dynamical systems by simply subclassing brainpy.DynamicalSystem base class.

At the numerical integrator level, BrainPy provides an easy way to write new integrator methods. For example, adding a new Runge–Kutta integrator can be done by subclassing brainpy.ode.ExplicitRKIntegrator and providing the corresponding Butcher tableau:

1 class NewExplicitRK(bp.ode.ExplicitRKIntegrator): 
2   A = [] # The A matrix in the Butcher tableau. 
3   B = [] # The B vector in the Butcher tableau. 
4   C = [] # The C vector in the Butcher tableau.

In a similar way, a customized adaptive Runge–Kutta integrator can be extended by subclassing brainpy.ode.AdaptiveRKIntegrator with the specification of the corresponding Butcher tableau:

1 class NewAdaptiveRK(bp.ode.ExplicitRKIntegrator): 
2   A = [] # The A matrix in the Butcher tableau. 
3   B1 = [] # The B1 vector in the Butcher tableau. 
4   B2 = [] # The B2 vector in the Butcher tableau. 
5   C = [] # The C vector in the Butcher tableau.

At the operator level, BrainPy is trying to provide a way to write fine-granularity low-level operators in the same Python interface by adopting the same JIT compilation technology (please refer to Appendix 8).

By bridging Numba (Lam et al., 2015), JAX (Frostig et al., 2018), and XLA (Sabne, 2020), BrainPy enables the customization of primitive operators through the native Python syntax. Exposing a custom operator to JAX requires registering an XLA ‘custom call’, and providing its C callback for Python. Based on the following two properties of Numba, we are aware of the possibility of using Numba as a convenient method for writing low-level kernels that support JAX’s JIT compilation. First, unlike JAX, which only supports the JIT compilation of high-level functions, Numba is a JIT compiler that allows users to write a function with low-level fine-grained operations, for instance, looping over an array, or conditional branching over array values. This implies that numba.jit() can be used as a means to write low-level kernel functions. The second property of Numba is that it provides a mechanism to create a compiled function that is callable from the foreign C code, such as XLA. Specifically, numba.cfunc() can be used to create a C callback for Numba’s JIT function to interface with XLA. Therefore, by integrating Numba with JAX and XLA, BrainPy provides an interface where users write primitive operators directly with Python syntax. Note that Numba supports various native Python features and a large subset of NumPy functions. Therefore, there is a large flexibility in coding low-level primitives with Numba.

Below, we illustrate how to write a primitive operator in BrainPy. Particularly, to customize a primitive operator we need to provide two functions. The first is an abstract evaluation function that tells JAX what shapes and types of outputs are produced according to the input shapes and types:

1 def abstract_eval(*ins): 
2   return outs

in which ins specifies the information of input shapes and types, outs denotes the array information of shapes and types of outputs. The other function is the concrete computation function, in which the output data is calculated according to the input data:

1 def concrete_comp(outs, ins): 
2   # calculate outputs according to inputs

where outs and ins are output and input data owned by XLA, respectively. Note this function should not return anything. All computations must be made in place. Finally, by using

1 # "name" is the operator name 
2 # "concrete_comp" is the concrete computation function 
3 # "abstract_eval" is the abstract evaluation function 
4 op =brainpy.math.CustomOpByNumba(abstract_eval, concrete_comp, name)

we register a primitive operator op. This operator op can be used anywhere the user wants. The Numba-based approach for operator customization demonstrates comparable performance to the native C++-based XLA operator registration. To illustrate this, we conducted a performance comparison of the COBA network model (Vogels and Abbott, 2005) using both an event-driven operator implementation based on Numba and C++.

This event-driven operator implements the matrix–vector multiplication ${\displaystyle \mathbf{y}=\mathbf{v}\cdot \mathbf{M}}$ for synaptic computation, where ${\displaystyle \mathbf{v}}$ is the presynaptic spikes, ${\displaystyle \mathbf{M}}$ the synaptic connection matrix, and ${\displaystyle \mathbf{y}}$ the postsynaptic current. Specifically, it performs matrix–vector multiplication in a sparse and efficient way by exploiting the event property of the input vector ${\displaystyle \mathbf{v}}$. Particularly, we performs event-driven matrix–vector multiplication only for the non-zero elements of the vector, which are called events. This can reduce the number of operations and memory accesses, and improve the running performance of matrix–vector multiplication.

Based on our Python-level registration syntax, this operator can be implemented as:

1 from jax.core import ShapedArray 
2 import brainpy.math as bm 
3 
4 # the abstract evaluation function 
5 def abs_eval(events, indices, indptr, *, weight, post_num): 
6    return [ShapedArray((post_num,), bm.float32), ] 
7 
8 # the concrete function implementing the event-driven computation 
9 def con_compute(outs, ins): 
10   post_val, =outs 
11   post_val.fill(0) 
12   events, indices, indptr, weight, _=ins 
13   weight=weight[()] 
14   for i in range(events.size): 
15     if events[i]: 
16       for j in range(indptr[i], indptr[i+1]): 
17         index =indices[j] 
18         post_val[index]+=weight 
19 
20 # operator registration 
21 event_sum =bm.CustomOpByNumba(eval_shape =abs_eval, cpu_func =con_compute)

Listing 1: The event-driven operator implemented with our Python syntax for the computation of ${\displaystyle \mathbf{y}=\mathbf{v}\cdot \mathbf{M}}$.

This operator can also be implemented with the pure C++ syntax:

1 #include <cstdint> 
2 #include <cstring> 
3 #include <cmath> 
4 
5 template <typename F, typename I>
6 void cpu_csr_event_sum_homo(void *out, const void **in) { 
7    const std::uint32_t pre_size = *reinterpret_cast<const std::uint32_t *>(in[0]); 
8    const std::uint32_t post_size = *reinterpret_cast<const std::uint32_t *>(in(1)); 
9    const bool *events=reinterpret_cast<const bool *>(in(2)); 
10   const I *indices=reinterpret_cast <const I *>(in(3)); 
11   const I *indptr=reinterpret_cast <const I *>(in(4)); 
12   const F weight=reinterpret_cast<const F *>(in(5))[0]; 
13   F *result=reinterpret_cast<F *>(out); 
14 
15   // algorithm 
16   memset(&result[0], 0, sizeof(F) * post_size); 
17   for (std::uint32_t i=0; i<pre_size; ++i) { 
18     if (events[i]){ 
19       for (I j=indptr[i]; j<indptr[i+1]; ++j) { 
20         result[indices[j]]+=weight; 
21       } 
22     } 
23   } 
24 }

Listing 2: The event-driven operator utilizes C++ for the computation of ${\displaystyle \mathbf{y}=\mathbf{v}\cdot \mathbf{M}}$. Here, we present the main code snippet that implements the event-driven matrix–vector multiplication. Please note that we have omitted the code for Python binding, compilation, registration, and other related aspects. For detailed instructions and a comprehensive tutorial on XLA operator customization, we encourage users to refer to the appropriate resource.

The speed comparison between the two approaches has been depicted in Figure 9. Although the approach shows promising results on CPU hardware, it is not directly compatible with other computing platforms like GPUs. This restricts the applicability and scalability of the proposed method, as GPUs are increasingly used to accelerate brain dynamics models. To overcome this limitation, currently, we are working on alternative approach that can be used in both CPU and GPU devices, allowing for broader utilization of available hardware resources and unlocking new possibilities for customizing any complex brain dynamics operators.

We build a spiking network model examined in Figure 2 and Figure 3 for modeling the visual system (V1, V2, V4, TEO, and TEpd). Simulations are performed using a network of Hodgkin–Huxley neurons, with the local circuit and long-range connectivity structure derived from Markov et al., 2014. Each of the five areas consists of 2000 neurons, with 1600 excitatory and 400 inhibitory neurons.

For each neuron, the membrane potential dynamics is modified from Traub and Miles, 1991 and is described by the following equations:

where ${\displaystyle V}$ is the membrane potential, ${\displaystyle G(t)}$ stands for synaptic interactions, ${\displaystyle C_m}$ is the membrane capacitance per unit area, ${\displaystyle E_K=-90\mathrm{m}\mathrm{V}}$ and ${\displaystyle E_{Na}=50\text{mV}}$ are the potassium and sodium reversal potentials, respectively, ${\displaystyle E_l=-60\text{mV}}$ is the leak reversal potential, ${\displaystyle {\overline{g}}_K=30{\text{mS/cm}}^2}$ and ${\displaystyle {\overline{g}}_{Na}=100{\text{mS/cm}}^2}$ are the potassium and sodium conductance per unit area, respectively, and ${\displaystyle {\overline{g}}_l=0.05{\text{mS/cm}}^2}$ is the leak conductance per unit area.

Each neuron is composed of two active ion channels, including the potassium and sodium channels. Because the potassium and sodium channels are voltage sensitive, according to the biological experiments, ${\displaystyle m}$, ${\displaystyle n}$, and ${\displaystyle h}$ are used to simulate the activation of the channels.

Specifically, ${\displaystyle n}$ measures the activation of potassium channels, and ${\displaystyle m}$ and ${\displaystyle h}$ measures the activation and inactivation of sodium channels, respectively. ${\displaystyle {\alpha }_x}$ and ${\displaystyle {\beta }_x}$ are rate constants for the ion channel ${\displaystyle x}$ and depend exclusively on the membrane potential. The voltage-dependent expressions of the rate constants were modified from the model described by Traub and Miles, 1991:

where ${\displaystyle V_T}$ = −63 mV adjusts the threshold.

For the synapse, we use conductance-based synaptic interactions. Particularly, ${\displaystyle G(t)}$ is given by:

where ${\displaystyle V_i}$ is the membrane potential of neuron ${\displaystyle i}$. The synaptic conductance from neuron ${\displaystyle j}$ to neuron ${\displaystyle i}$ is represented by ${\displaystyle g_{ji}(t)}$, while ${\displaystyle E_j}$ signifies the reversal potential of that synapse. For excitatory synapses, ${\displaystyle E_j}$ was set to 0 mV, whereas for inhibitory synapses, it was −80 mV. The dynamics of the synaptic conductance is given by:

where ${\displaystyle t_j^k}$ is the spiking time of the presynaptic spike. Whenever a spike occurred in neuron ${\displaystyle j}$, the synaptic conductance ${\displaystyle g_{ji}}$ experienced an immediate increase by a fixed amount ${\displaystyle g_{\mathrm{m}\mathrm{a}\mathrm{x}}}$. Subsequently, the conductance ${\displaystyle g_{ji}}$ decayed exponentially with a time constant of ${\displaystyle {\tau }_{\mathrm{d}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{y}}=5}$ ms for excitation and ${\displaystyle {\tau }_{\mathrm{d}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{y}}=10}$ ms for inhibition.

The connection density is set according to the experimental connectivity data (Markov et al., 2014). The inter-areal connectivity is measured as a weight index (see Appendix 9—table 1), called the extrinsic fraction of labeled neurons (Markov et al., 2014). The intra-area connectivity is set according to the setting in a standard E/I balanced network (Brette et al., 2007).

Moreover, we introduce distance-dependent inter-areal synaptic delays by assuming a conduction velocity of 3.5 m/s (Swadlow, 1990) and using a distance matrix based on experimentally measured wiring distances across areas (Markov et al., 2014).

Reservoir computing is a type of recurrent neural network that is often used for processing temporal data. Unlike traditional recurrent neural networks, reservoir computing fixes the weights of the recurrent layer (known as the ‘reservoir’) and only trains the weights of the output layer. This makes training much more efficient.

A reservoir computing model can be trained using various algorithms, such as online learning, offline learning, and backpropagation learning, to optimize its performance. Online learning refers to the process of updating the model in real time as new data becomes available. This algorithm allows the model to adapt and adjust its predictions continuously. Offline learning, on the other hand, involves training the model using a fixed dataset, where the entire dataset is used to update the model’s parameters. This method is particularly useful when a large amount of data is available upfront. Lastly, backpropagation learning utilizes a technique called backpropagation to train the model by computing the gradients of the model’s parameters and adjusting them accordingly.

The unified building and training interface of BrainPy enables the training of the same reservoir model with multiple training algorithms. By employing BrainPy, we can compare the performance of different training algorithms and determine which approach yields the best results for the reservoir computing model.

The following lists the details of such training.

The dynamics of the reservoir model used here are given by:

where ${\displaystyle \mathbf{x}(t)}$ is the reservoir state, ${\displaystyle \mathbf{y}(t)}$ the readout value, ${\displaystyle {\mathbf{W}}_{in}}$ and ${\displaystyle {\mathbf{W}}_{rec}}$ are input and recurrent connection matrices, respectively, ${\displaystyle {\mathbf{W}}_{out}}$ the readout weight matrix which can be trained by either offline learning or online learning algorithms, ${\displaystyle \alpha \in (0,1]}$ the leaky rate, ${\displaystyle \mathbf{u}(t)}$ the input at time step ${\displaystyle t}$, and ${\displaystyle f}$ the nonlinear activation function.

In BrainPy, the reservoir model can be easily instantiated as the following code:

1 reservoir =brainpy.dyn.Reservoir(input_shape, output_shape, leaky_rate)

The reservoir model is utilized for inference tasks in this work. To generate training and testing data, we numerically integrate a simplified model of a weather system originally developed by Lorenz, 1963. This model comprises three coupled nonlinear differential equations:

The state ${\displaystyle X(t)=[x(t),y(t),z(t){]}^T}$ represents a vector of Rayleigh–Bénard convection observables. The system exhibits deterministic chaos, displaying sensitivity to initial conditions and forming a strange attractor in the phase space trajectory (Figure 4D).

In this task, we train a reservoir model to predict the ${\displaystyle T}$-step-ahead value of all Lorenz variables, ${\displaystyle x}$, ${\displaystyle y}$, and ${\displaystyle z}$. During training, we provide the model with all three variables. During testing, we only provide the model with ${\displaystyle x}$, ${\displaystyle y}$, and ${\displaystyle z}$ at ${\displaystyle T}$-step-ahead, and it must infer ${\displaystyle x,y}$, and ${\displaystyle z}$ after ${\displaystyle T}$ steps. Here, we choose ${\displaystyle T=5}$.

In this task, the input size was set to 3, the recurrent layer of the reservoir model contained 100 units, and the output size was 3, as used in Figure 4D-F.

The training objective of reservoir models is to find the optimal ${\displaystyle {\mathbf{W}}_{out}}$ that minimizes the square error between ${\displaystyle \mathbf{y}(t)}$ and ${\displaystyle {\mathbf{y}}^{target}(t)}$. The common way to learn the linear output weight ${\displaystyle {\mathbf{W}}_{out}}$ is using the ridge regression algorithm. The ridge regression, also known as regression with Tikhonov regularization, is given by:

where ${\displaystyle \beta }$ is a regularization coefficient, ${\displaystyle I}$ is the identity matrix, and ${\displaystyle X}$ is the concatenated hidden states of all samples.

In BrainPy, the reservoir model trained with ridge regression can be implemented as:

1 model =ESN(num_in, num_rec, num_out) 
2 trainer =bp.RidgeTrainer(model) 
3 trainer.fit([X, Y]) 
4 outputs =trainer.predict(X)

Ridge regression is an offline learning method, meaning that it needs all of the data to be present before it can learn the model parameters. This can be a problem when training reservoir models with a lot of data, because the memory requirements can be too high for some devices. FORCE learning (Sussillo and Abbott, 2009), on the other hand, is an online learning algorithm. This means that it can learn the model parameters one sample of data at a time. This makes it a much more efficient way to train reservoir models with large datasets. Therefore, FORCE learning is a good choice for training reservoir models when the amount of data is large or the memory resources are limited.

The FORCE learning is done using the recursive least squares (RLS) algorithm. It is a supervised error-driven learning rule, that is the weight changes depending on the error made by each neuron: the difference between the output of a neuron ${\displaystyle {\mathbf{y}}_i(t)}$ and a desired value ${\displaystyle {\mathbf{y}}_i^{\mathrm{t}\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{e}\mathrm{t}}(t)}$.

Contrary to the delta learning rule which modifies weights proportionally to the error and to the direct input to a synapse (${\displaystyle \mathrm{\Delta }{w}_{ij}=-\eta \cdot e_i\cdot {\mathbf{x}}_j}$), the RLS learning uses a running estimate of the inverse correlation matrix of the inputs to each neuron:

Each neuron ${\displaystyle i}$ therefore stores a square matrix ${\displaystyle P^i}$, whose size depends on the number of weights arriving at the neuron. Readout neurons receive synapses from all ${\displaystyle N}$ recurrent neurons, so the ${\displaystyle P}$ matrix is ${\displaystyle N\ast N}$.

The inverse correlation matrix ${\displaystyle P}$ is updated at each time step with the following rule:

Each matrix ${\displaystyle P^i}$ is initialized to the diagonal matrix and scaled by a factor ${\displaystyle 1/\delta }$, where ${\displaystyle \delta }$ is 1 in the current implementation and can be used to modify implicitly the learning rate (Sussillo and Abbott, 2009).

In BrainPy, the reservoir model trained with FORCE learning can be implemented as:

1 model =ESN(num_in, num_rec, num_out) 
2 trainer =bp.ForceTrainer(model, alpha =0.1) 
3 trainer.fit([X, Y]) 
4 trainer.fit([X, Y]) 
5 outputs =trainer.predict(X)

The readout layer in the reservoir model is just a single-layer Perceptron. To train its weights, we can use the backpropagation algorithm. Backpropagation is a method used in artificial neural networks to calculate a gradient that is needed in the calculation of the weights to be used in the network. The loss function used here is the mean square error between the reservoir output and the target output:

The updated weight between neuron ${\displaystyle i}$ and ${\displaystyle j}$ is calculated by:

In BrainPy, the reservoir model trained with backpropagation algorithms can be implemented as:

1 reservoir = bp.dyn.Reservoir(num_in, num_rec) 
2 readout = bp.dnn.Dense(num_rec, num_out, mode =bm.training_mode) 
3 trainer = bp.BPFF(target =readout, 
4                     loss_fun =bp.losses.mean_squared_error, 
5                     optimizer =bp.optim.Adam(1e-3)) 
6 # batch_data: the data generated by "reservoir" 
7 trainer.fit(batch_data, num_epoch =100)

The performance of a reservoir computing model is usually measured with the mean squared error, that is, the average squared difference between the predicted and actual values:

where ${\displaystyle N_y}$ is the number of training samples and ${\displaystyle T}$ is the number of time steps.

In this section, we provide a comprehensive overview of the experimental benchmark details used in Figure 6, Figure 7, and Figure 8. The purpose is to present a complete picture of the experimental setup and methodology employed in our study. The details encompass the following aspects: hardware specifications, software versions, algorithm hyperparameters, and performance measurements.

By providing these experimental benchmark details, we aim to ensure transparency and reproducibility, allowing readers and researchers to understand and replicate our experiments accurately.

We conducted a series of experiments on various computing devices, namely the CPU, GPU, and TPU, in order to validate and compare the simulation speeds of several widely utilized brain simulators. The brain simulators under investigation included Brian2 (Stimberg et al., 2019), NEURON (Hines and Carnevale, 1997), NEST (Gewaltig and Diesmann, 2007), GeNN (Yavuz et al., 2016), and Brian2CUDA (Alevi et al., 2022).

Particularly, the computing devices we used here are:

• CPU: The CPU used was the Intel Xeon W-2255 processor. This is a 10-core, 20-thread CPU based on Intel’s Cascade Lake microarchitecture. It has a base clock frequency of 3.7 GHz and a max turbo boost up to 4.5 GHz. The Xeon W-2255 utilizes the LGA2066 socket and supports up to 512 GB of ECC-registered DDR4-2933 RAM across 6 channels. It has 24.75 MB of L3 cache. As a server-grade CPU with a high core count, the Xeon W-2255 is well suited for heavily parallelized simulations.

• GPU: The GPU used was an NVIDIA RTX A6000. This is a high-end Ampere architecture GPU aimed at professional visualization, AI, and compute workloads. The RTX A6000 has 10,752 CUDA cores, 336 tensor cores, and 84 RT cores. It comes equipped with 48 GB of GDDR6 memory clocked at 16 Gbps, providing up to 1 TB/s of memory bandwidth. The RTX A6000 has a maximum power consumption of 300 W and requires auxiliary power connectors. It uses a PCIe 4.0x16 interface. With its large number of CUDA cores and abundant memory, the RTX A6000 is well suited for massively parallel simulations and neural network training.

• TPU: The TPU simulations leveraged the Kaggle free TPU v3-8 cloud instance. This provides access to one of Google’s TPU v3 chips via their cloud platform. Specifically, the v3-8 instance gives 8 TPU v3 cores, each providing 128 GB/s of bandwidth to high-performance HBM memory. The TPU v3 is optimized for both training and inferencing of DNNs, providing up to 420 TFLOPS of mixed precision computing. By utilizing Google’s cloud TPUs, researchers can run models with minimal coding changes while leveraging Google’s optimized deep learning frameworks.

We carried out benchmark experiments using Python 3.10.12 running on Ubuntu 20.04.6 LTS. For any experiments utilizing a GPU, we used version 11.6 of the NVIDIA CUDA Toolkit to take advantage of the parallel processing capabilities of the GPU hardware. Other dependent software versions are: NumPy 1.24.3, Numba 0.57.1, jax 0.4.16, and jaxlib 0.4.16.

The comparison brain simulators for benchmarking are:

• NEURON at version 8.2.0

• NEST at version 3.6

• Brian2 at version 2.5.4

• GeNN at version 4.8.1

• Brian2GeNN at version of 1.7.0: Brian2GeNN is a bridge between Brian2 and GeNN that allows Brian2 users to run their simulations on GPUs using GeNN.

• Brian2CUDA at version of 1.0a3: Brian2CUDA is a native CUDA implementation of Brian2 that allows Brian2 users to run their simulations on GPUs without the need for a bridge.

By testing against these simulators on benchmark tasks, we aimed to thoroughly evaluate the performance in terms of simulation speed and accuracy across different model sizes and paradigms.

When evaluating simulation speed, we focused solely on measuring the execution time of the simulation itself, excluding any time spent on compilation, initialization, or other preparatory steps. This allowed us to make a direct comparison of the raw simulation performance between our simulator and others. We simulated networks of varying sizes. This range of model scales allowed us to determine how well our simulator performs as network size and complexity increase. The final experimental results can be obtained in Figure 6, Figure 7, and Figure 8.

To evaluate the accuracy, we configured all simulators to use identical model parameters for a fair comparison.

First, we verified that all simulators generated consistent average firing rates for a given network model (see Appendix 11—figure 1 and Appendix 11—figure 2).

Appendix 11—figure 1

Download asset Open asset

Appendix 11—figure 2

Download asset Open asset

Second, we qualitatively compared the overall network activity patterns produced by each simulator to ensure they exhibited the same dynamics. While exact spike-to-spike reproducibility was not guaranteed between different simulator implementations, we confirmed that our simulator produced activity consistent with the reference simulators for both firing rates and network-level dynamics (see Appendix 11—figure 3 and Appendix 11—figure 4).

Appendix 11—figure 3

Download asset Open asset

Appendix 11—figure 4

Download asset Open asset

Third, we recognize that the precision of numerical computation plays a crucial role in accurately simulating biologically detailed neural networks, such as the COBAHH network model used in our benchmarks. To assess the numerical integration accuracy of each neuron in the COBAHH network model, we conducted a 5-s simulation and examined the presence of NaN membrane potentials at the end of the simulation. The occurrence of NaN membrane potentials indicates that the corresponding neurons are no longer active and signifies a loss of simulation accuracy. This analysis was performed on the GPU backend of Brian2GeNN, Brian2CUDA, and BrainPy. The results of this analysis can be found in Appendix 11—figure 5. We specifically focused on the analysis results of single-precision floating-point simulations, as we did not encounter any NaN results when simulating with double precision. Our observations reveal that the single-precision computation with XLA in BrainPy exhibits a higher proportion of NaN results compared to Brian2CUDA and GeNN. Notably, GeNN demonstrated very few occurrences of NaN membrane potentials after simulation, which may be attributed to the special handling of NaN within the GeNN computation.

Appendix 11—figure 5

Download asset Open asset

To address the issue of a larger number of NaN membrane potentials leading to reduced neuronal spiking and communication spikes, we took steps to resolve this problem in BrainPy. Our aim was to ensure a fair comparison when benchmarking the simulation speed against other brain simulators. Specifically, the occurrence of NaN membrane potentials was attributed to the use of the ${\displaystyle x/(\exp (x)-1)}$ operation during the integration of the Hodgkin–Huxley neuron model with low-precision floating-point calculations. In order to mitigate the catastrophic loss of precision when ${\displaystyle x}$ is close to zero, we replaced this operation with the relative error exponential function, represented by ${\displaystyle 1/\mathrm{e}\mathrm{x}\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{l}(x)}$. This modification ensures that the numerical calculations do not result in NaN values even during significantly long simulations. As a result of this fix, we did not encounter any instances of NaN membrane potentials, even after conducting extensive and prolonged simulations. We believe that these revised operations contribute to a more equitable benchmarking process, eliminating any potential bias caused by the presence of NaN values. The speed benchmarking can be obtained in Figure 8.

In the below, we describe the experimental setting details used in Figure 6, Figure 7, and Figure 8.

Appendix 11—figure 6

Download asset Open asset

We evaluate the speedup of event-driven operators over traditional dense matrix operators on both CPU and GPU devices. Evaluation results are listed in Appendix 11—figure 7. Both CPU and GPU devices demonstrate a significant speed advantage of several orders of magnitude when utilizing event-driven operators, in comparison to dense ones. Notably, as the number of neurons increases, event-driven operators showcase an even greater speedup on both platforms.

Appendix 11—figure 7

Download asset Open asset

We also conducted COBA and COBAHH experiments on TPU devices. The experimental results are shown in Appendix 11—figure 8. Although TPUs can perform double precision operations, they are not as efficient at it as they are at lower precisions such as float16 or bfloat16. Therefore, we here only tested models with single-precision operations.

Appendix 11—figure 8

Download asset Open asset

The funders had no role in study design, data collection, and interpretation, or the decision to submit the work for publication.

This work was supported by Science and Technology Innovation 2030-Brain Science and Brain-inspired Intelligence Project (No. 2021ZD0200204) and Beijing Academy of Artificial Intelligence. We would like to express our sincere gratitude to Marcel Stimberg for his valuable insights and assistance with benchmarking brain simulators, which greatly contributed to this paper. We would like to acknowledge Xiaohan Lin, Yifeng Gong, Hongyaoxing Gu, Linfei Lu, Xiaolong Zou, Zhiyu Zhao, Yingqian Jiang, Xinyu Liu, and all other members of the Wu laboratory for their helpful discussions. We thank all GitHub users who contributed codes to BrainPy.

© 2023, Wang et al.

This article is distributed under the terms of the Creative Commons Attribution License, which permits unrestricted use and redistribution provided that the original author and source are credited.