Ctrl+K

NumPy-like Operations

NumPy-like Operations#

Colab Open in Kaggle

@Xiaoyu Chen @Chaoming Wang

brainpy.math.array#

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#

For 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 numpy.xxx to brainpy.math.xxx.

Specifically, the NumPy-like operations for brainpy.math.array include:

  1. Creating an array: users can use brainpy.math.array()， brainpy.math.ones()， brainpy.math.zeros()， brainpy.math.arange()， and so on to create an array.

  2. Element-wise operations: +, -, *, /, …

  3. Aggregation functions: .max(), .min(), .sum(), .mean(), .prod(), …

  4. Broadcasting: operating two arrays that have different dimensions (usually a smaller array is broadcast to a larger array with compatible shapes to complete operation).

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

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

  7. Other mathematical and logical functions: brainpy.math.sin(), brainpy.math.sort(), brainpy.math.argmax(), brainpy.math.where(), …

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[2] = 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([[1],
                   [2],
                   [3],
                   [4],
                   [5],
                   [6]]),
      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)

True

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

0.3499875
True

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 brianpy.math module.

When to use brainpy.math.array?#

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 Variable and BrainPbject.

Contents