Merge pull request #406 from Routhleck/master

Optimize GaussianProb

Author: chaoming0625

brainpy/_src/connect/base.py

@@ -9,8 +9,6 @@
 from brainpy import tools, math as bm
 from brainpy.errors import ConnectorError
 
-import matplotlib.pyplot as plt
-import seaborn as sns
 import textwrap
 
 __all__ = [
@@ -729,6 +727,12 @@ def coo2csc(coo, post_num, data=None):
 
 
 def visualizeMat(mat, description):
+  try:
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+  except (ModuleNotFoundError, ImportError):
+    print('Please install seaborn and matplotlib for this function')
+    return
   sns.heatmap(mat, cmap='viridis')
   warpped_title = textwrap.fill(description, width=60)
   plt.title(warpped_title)

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=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_prob2():
+    conn = bp.connect.GaussianProb(sigma=4, seed=123)(pre_size=(10, 10))
+    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=(10, 10))
+    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=(10, 10, 10))
+    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}')

requirements-dev.txt

@@ -8,7 +8,6 @@ scipy>=1.1.0
 brainpylib
 h5py
 pathos
-seaborn
 
 # test requirements
 pytest