Optimize GaussianProb

Optimize GaussianProb with JAX and Create test_GaussianProb_opt.py to test its correction and optimization effect
Below is the test report
============================= test session starts =============================
collecting ... collected 4 items

test_GaussianProb_opt.py::test_gaussian_prob1
test_GaussianProb_opt.py::test_gaussian_prob2
test_GaussianProb_opt.py::test_gaussian_prob3
test_GaussianProb_opt.py::test_gaussian_prob4

================== 4 passed, 1 warning in 134.70s (0:02:14) ===================
PASSED                     [ 25%]
time_optimized:0.040474891662597656
time_origin:0.16180109977722168
PASSED                     [ 50%]
time_optimized:1.0694525241851807
time_origin:7.487141132354736
PASSED                     [ 75%]
time_optimized:22.59926676750183
time_origin:52.102147579193115
PASSED                     [100%]
time_optimized:2.2735044956207275
time_origin:21.08591890335083

Process finished with exit code 0

Author: Routhleck

brainpy/_src/connect/random_conn.py

@@ -1,8 +1,8 @@
 # -*- coding: utf-8 -*-
-
+from functools import partial
 from typing import Optional
 
-import jax.numpy as jnp
+from jax import vmap, jit, numpy as jnp
 import numpy as np
 
 import brainpy.math as bm
@@ -327,6 +327,21 @@ def build_csr(self):
     selected_pre_inptr = jnp.cumsum(jnp.concatenate([jnp.zeros(1), pre_nums]))
     return selected_post_ids.astype(IDX_DTYPE), selected_pre_inptr.astype(IDX_DTYPE)
 
+@jit
+@partial(vmap, in_axes=(0, None, None))
+def gaussian_prob_dist_cal1(i_value, post_values, sigma):
+  dists = jnp.abs(i_value - post_values)
+  exp_dists = jnp.exp(-(jnp.sqrt(jnp.sum(dists ** 2, axis=0)) / sigma) ** 2 / 2)
+  return bm.asarray(exp_dists)
+
+@jit
+@partial(vmap, in_axes=(0, None, None, None))
+def gaussian_prob_dist_cal2(i_value, post_values, value_sizes, sigma):
+  dists = jnp.abs(i_value - post_values)
+  dists = jnp.where(dists > (value_sizes / 2), value_sizes - dists, dists)
+  exp_dists = jnp.exp(-(jnp.sqrt(jnp.sum(dists ** 2, axis=0)) / sigma) ** 2 / 2)
+  return bm.asarray(exp_dists)
+
 
 class GaussianProb(OneEndConnector):
   r"""Builds a Gaussian connectivity pattern within a population of neurons,
@@ -392,7 +407,8 @@ def __repr__(self):
             f'include_self={self.include_self}, '
             f'seed={self.seed})')
 
-  def build_mat(self, pre_size=None, post_size=None):
+  def build_mat(self, isOptimized=True):
+    self.rng = np.random.RandomState(self.seed)
     # value range to encode
     if self.encoding_values is None:
       value_ranges = tuple([(0, s) for s in self.pre_size])
@@ -426,24 +442,45 @@ def build_mat(self, pre_size=None, post_size=None):
       value_sizes = np.expand_dims(value_sizes, axis=tuple([i + 1 for i in range(post_values.ndim - 1)]))
 
     # probability of connections
-    prob_mat = []
-    for i in range(self.pre_num):
-      # values for node i
-      i_coordinate = tuple()
-      for s in self.pre_size[:-1]:
-        i, pos = divmod(i, s)
-        i_coordinate += (pos,)
-      i_coordinate += (i,)
-      i_value = np.array([values[i][c] for i, c in enumerate(i_coordinate)])
-      if i_value.ndim < post_values.ndim:
-        i_value = np.expand_dims(i_value, axis=tuple([i + 1 for i in range(post_values.ndim - 1)]))
-      # distances
-      dists = np.abs(i_value - post_values)
+    if isOptimized:
+      i_value_list = np.zeros(shape=(self.pre_num, len(self.pre_size), 1))
+      for i in range(self.pre_num):
+        list_index = i
+        # values for node i
+        i_coordinate = tuple()
+        for s in self.pre_size[:-1]:
+          i, pos = divmod(i, s)
+          i_coordinate += (pos,)
+        i_coordinate += (i,)
+        i_value = np.array([values[i][c] for i, c in enumerate(i_coordinate)])
+        if i_value.ndim < post_values.ndim:
+          i_value = np.expand_dims(i_value, axis=tuple([i + 1 for i in range(post_values.ndim - 1)]))
+        i_value_list[list_index] = i_value
+
       if self.periodic_boundary:
-        dists = np.where(dists > value_sizes / 2, value_sizes - dists, dists)
-      exp_dists = np.exp(-(np.linalg.norm(dists, axis=0) / self.sigma) ** 2 / 2)
-      prob_mat.append(exp_dists)
-    prob_mat = np.stack(prob_mat)
+        prob_mat = gaussian_prob_dist_cal2(i_value_list, post_values, value_sizes, self.sigma)
+      else:
+        prob_mat = gaussian_prob_dist_cal1(i_value_list, post_values, self.sigma)
+    else:
+      prob_mat = []
+      for i in range(self.pre_num):
+        # values for node i
+        i_coordinate = tuple()
+        for s in self.pre_size[:-1]:
+          i, pos = divmod(i, s)
+          i_coordinate += (pos,)
+        i_coordinate += (i,)
+        i_value = np.array([values[i][c] for i, c in enumerate(i_coordinate)])
+        if i_value.ndim < post_values.ndim:
+          i_value = np.expand_dims(i_value, axis=tuple([i + 1 for i in range(post_values.ndim - 1)]))
+        # distances
+        dists = np.abs(i_value - post_values)
+        if self.periodic_boundary:
+          dists = np.where(dists > value_sizes / 2, value_sizes - dists, dists)
+        exp_dists = np.exp(-(np.linalg.norm(dists, axis=0) / self.sigma) ** 2 / 2)
+        prob_mat.append(exp_dists)
+      prob_mat = np.stack(prob_mat)
+
     if self.normalize:
       prob_mat /= prob_mat.max()
 

brainpy/_src/connect/tests/test_GaussianProb_opt.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+
+import unittest
+
+import brainpy as bp
+
+from time import time
+
+
+def test_gaussian_prob1():
+    conn = bp.connect.GaussianProb(sigma=1., include_self=False, seed=123)(pre_size=2000)
+
+    mat = conn.build_mat(isOptimized=True)
+    time0 = time()
+    mat1 = conn.build_mat(isOptimized=True)
+    time_optimized = time() - time0
+
+    time0 = time()
+    mat2 = conn.build_mat(isOptimized=False)
+    time_origin = time() - time0
+
+    assert bp.math.array_equal(mat1, mat2)
+    print()
+    print(f'time_optimized:{time_optimized}\ntime_origin:{time_origin}')
+
+
+def test_gaussian_prob2():
+    conn = bp.connect.GaussianProb(sigma=4, seed=123)(pre_size=(100, 100))
+    mat = conn.build_mat(isOptimized=True)
+    time0 = time()
+    mat1 = conn.build_mat(isOptimized=True)
+    time_optimized = time() - time0
+
+    time0 = time()
+    mat2 = conn.build_mat(isOptimized=False)
+    time_origin = time() - time0
+
+    assert bp.math.array_equal(mat1, mat2)
+    print()
+    print(f'time_optimized:{time_optimized}\ntime_origin:{time_origin}')
+
+
+def test_gaussian_prob3():
+    conn = bp.connect.GaussianProb(sigma=4, periodic_boundary=True, seed=123)(pre_size=(200, 200))
+    mat = conn.build_mat(isOptimized=True)
+    time0 = time()
+    mat1 = conn.build_mat(isOptimized=True)
+    time_optimized = time() - time0
+
+    time0 = time()
+    mat2 = conn.build_mat(isOptimized=False)
+    time_origin = time() - time0
+
+    assert bp.math.array_equal(mat1, mat2)
+    print()
+    print(f'time_optimized:{time_optimized}\ntime_origin:{time_origin}')
+
+
+def test_gaussian_prob4():
+    conn = bp.connect.GaussianProb(sigma=4, periodic_boundary=True, seed=123)(pre_size=(25, 25, 25))
+    mat = conn.build_mat(isOptimized=True)
+    time0 = time()
+    mat1 = conn.build_mat(isOptimized=True)
+    time_optimized = time() - time0
+
+    time0 = time()
+    mat2 = conn.build_mat(isOptimized=False)
+    time_origin = time() - time0
+
+    assert bp.math.array_equal(mat1, mat2)
+    print()
+    print(f'time_optimized:{time_optimized}\ntime_origin:{time_origin}')