Neural modeling often involves the computation of thousands of elements and requires large-scale computation. To support this, BrainPy provides a basic data structure
brainpy.math.array with abundant operations. In fact, the concept of
brainpy.math.array is the same as
ndarray in NumPy, and the operations for
brainpy.math.array closely resemble NumPy operations.
Users who are not familiar with the NumPy package may refer to https://numpy.org/ for more information.
NumPy-like operations for arrays#
brainpy.math.array, BrainPy provides array-related operations that have the same form as NumPy while capable of performing JIT transformations with the help of JAX. Almost all NumPy operations can be used in BrainPy as long as users transform
Specifically, the NumPy-like operations for
Creating an array: users can use
brainpy.math.arange()， and so on to create an array.
Broadcasting: operating two arrays that have different dimensions (usually a smaller array is broadcast to a larger array with compatible shapes to complete operation).
Indexing, slicing, and iterating: Though based on JAX, BrainPy provides NumPy-like indexing, slicing, and iterating operations that are more convenient than those in JAX. Arrays can be treated much the same as other Python sequence data structures such as lists.
Random functions: BrainPy also offers NumPy-like random functions that are different from JAX. To enable random functions in JIT compilation, JAX requires users to pass an explicit key to define a random state, whereas random functions in BrainPy can be used in the same way as in NumPy where the seed does not need to be explicitly defined.
Other mathematical and logical functions:
import brainpy.math as bm bm.set_platform('cpu')
Here are some examples of array-related operations in BrainPy:
# to create an array a1 = bm.array([[0, 1, 2], [3, 4, 5]]) a1
Array(value=Array([[0, 1, 2], [3, 4, 5]]), dtype=int32)
# element-wise operation a1 += 1 a1
Array(value=Array([[1, 2, 3], [4, 5, 6]]), dtype=int32)
a2 = bm.ones(3) a2
Array(value=Array([1., 1., 1.]), dtype=float32)
# slicing and indexing a2[0:2] = 10. a2 = 0. a2
Array(value=Array([10., 10., 0.]), dtype=float32)
# broadcasting a3 = a1 + a2 a3
Array(value=Array([[11., 12., 3.], [14., 15., 6.]]), dtype=float32)
# sorting bm.sort(a3, axis=-1)
Array(value=Array([[ 3., 11., 12.], [ 6., 14., 15.]]), dtype=float32)
# reshaping a1.reshape(6, -1)
Array(value=Array([, , , , , ]), dtype=int32)
Here is an example for random number generation in JAX and in BrainPy:
# random number generation in JAX from jax import random key = random.PRNGKey(0) # define a key explicitly r1 = random.uniform(key) # a random number generated from a uniform distribution r2 = random.uniform(key) # reusing the same key leads to the same result print(r1 == r2)
# random number generation in BrainPy r3 = bm.random.uniform() # the seed does not need to be passed explicitly r4 = bm.random.uniform(key=1) # users can also assign an explicit key to generate a constant result r5 = bm.random.uniform(key=1) print(r3) print(r4 == r5)
Users can also generate a random array using random functions in BrainPy:
bm.random.random((5, 3)) # to generate random floats number in the interval [0.0, 1.0)
Array(value=Array([[0.9790665 , 0.68106973, 0.63952565], [0.27893245, 0.00639188, 0.53392684], [0.7553079 , 0.219872 , 0.8078612 ], [0.21869218, 0.10861421, 0.6631783 ], [0.11138558, 0.9173161 , 0.5229598 ]]), dtype=float32)
For all array-related operations, please refer to the
When to use
Simply speaking, users can use
brainpy.math.array whenever array computation is involved in JIT compilation. While JIT compilation speeds up code execution, it also brings some limitations. Once an array is given to the JIT compiler, the values inside the array cannot be changed. This means
brainpy.math.array can only be used to store values that do not change over time. Static neural properties such as the membrane capacitance, the voltage threshold, and the time constant of a neuron population, and static synaptic connections between neurons are good candidates to be defined as arrays (or a unified single value).
To extend arrays to store changing values, BrainPy offers a new data structure named
brainpy.math.Variable. More details about Variables can be found in