update tests (#169)

update tests

Author: chaoming0625

brainpy/analysis/lowdim/tests/test_phase_plane.py

@@ -3,14 +3,15 @@
 import unittest
 
 import brainpy as bp
-import matplotlib.pyplot as plt
 
 block = False
 
 
 class TestPhasePlane(unittest.TestCase):
   def test_1d(self):
+    import matplotlib.pyplot as plt
     bp.math.enable_x64()
+
     @bp.odeint
     def int_x(x, t, Iext):
       dx = x ** 3 - x + Iext
@@ -28,6 +29,8 @@ def int_x(x, t, Iext):
     bp.math.disable_x64()
 
   def test_2d_decision_making_model(self):
+    import matplotlib.pyplot as plt
+
     bp.math.enable_x64()
     gamma = 0.641  # Saturation factor for gating variable
     tau = 0.06  # Synaptic time constant [sec]

brainpy/initialize/tests/test_decay_inits.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import unittest
 
-import matplotlib.pyplot as plt
 import numpy as np
 
 import brainpy as bp
@@ -10,7 +9,11 @@
 
 
 # visualization
-def mat_visualize(matrix, cmap=plt.cm.get_cmap('coolwarm')):
+def mat_visualize(matrix, cmap=None):
+  import matplotlib.pyplot as plt
+  if cmap is None:
+    cmap = plt.cm.get_cmap('coolwarm')
+  plt.cm.get_cmap('coolwarm')
   im = plt.matshow(matrix, cmap=cmap)
   plt.colorbar(mappable=im, shrink=0.8, aspect=15)
   plt.show()
@@ -71,6 +74,8 @@ def test_dog_decay_init2(self):
     assert isinstance(weights, bp.math.ndarray)
 
   def test_dog_decay3(self):
+    import matplotlib.pyplot as plt
+
     size = (10, 12)
     dog_init = bp.init.DOGDecay(sigmas=(1., 3.),
                                 max_ws=(10., 5.),

brainpy/inputs/tests/test_currents.py

@@ -3,15 +3,19 @@
 
 from unittest import TestCase
 
-import matplotlib.pyplot as plt
+
 import numpy as np
 
 import brainpy as bp
 
+plt = None
 block = False
 
 
 def show(current, duration, title=''):
+  global plt
+  if plt is None:
+    import matplotlib.pyplot as plt
   ts = np.arange(0, duration, bp.math.get_dt())
   plt.plot(ts, current)
   plt.title(title)

brainpy/integrators/fde/tests/test_GL.py

@@ -2,12 +2,18 @@
 
 
 import unittest
-import matplotlib.pyplot as plt
+
 import brainpy as bp
 
+plt = None
+
 
 class TestGLShortMemory(unittest.TestCase):
   def test_lorenz(self):
+    global plt
+    if plt is None:
+      import matplotlib.pyplot as plt
+
     a, b, c = 10, 28, 8 / 3
 
     def lorenz(x, y, z, t):

brainpy/integrators/ode/tests/test_ode_method_adaptive_rk.py

@@ -2,12 +2,14 @@
 
 import unittest
 
-import matplotlib.pyplot as plt
 import numpy as np
 
 import brainpy.math as bm
 from brainpy.integrators.ode import adaptive_rk
 
+plt = None
+
+
 sigma = 10
 beta = 8 / 3
 rho = 28
@@ -23,6 +25,10 @@ def f_lorenz(x, y, z, t):
 
 
 def run_integrator(method, show=False, tol=0.001, adaptive=True):
+  global plt
+  if plt is None:
+    import matplotlib.pyplot as plt
+
   f_integral = method(f_lorenz, adaptive=adaptive, tol=tol, show_code=True)
   x, y, z = bm.ones(1), bm.ones(1), bm.ones(1)
   dt = bm.ones(1) * 0.01

brainpy/integrators/ode/tests/test_ode_method_exp_euler.py

@@ -1,14 +1,14 @@
 # -*- coding: utf-8 -*-
 
 import unittest
-import numpy as np
-import matplotlib.pyplot as plt
 import pytest
 
 import brainpy as bp
 import brainpy.math as bm
 from brainpy.integrators.ode.exponential import ExponentialEuler
 
+plt = None
+
 
 class TestExpnentialEuler(unittest.TestCase):
   def test_hh_model(self):
@@ -227,6 +227,10 @@ def dev(x, t):
 
 class TestExpEulerAuto(unittest.TestCase):
   def test_hh_model(self):
+    global plt
+    if plt is None:
+      import matplotlib.pyplot as plt
+
     class HH(bp.dyn.NeuGroup):
       def __init__(self, size, ENa=55., EK=-90., EL=-65, C=1.0, gNa=35., gK=9.,
                    gL=0.1, V_th=20., phi=5.0, name=None, method='exponential_euler'):

brainpy/integrators/ode/tests/test_ode_method_rk.py

@@ -2,11 +2,11 @@
 
 import unittest
 
-import matplotlib.pyplot as plt
 import numpy as np
 
 import brainpy.math as bm
 from brainpy.integrators.ode import explicit_rk
+plt = None
 
 sigma = 10
 beta = 8 / 3
@@ -23,6 +23,10 @@ def f_lorenz(x, y, z, t):
 
 
 def run_integrator(method, show=False):
+  global plt
+  if plt is None:
+    import matplotlib.pyplot as plt
+
   f_integral = bm.jit(method(f_lorenz, dt=dt), auto_infer=False)
   x, y, z = bm.ones(1), bm.ones(1), bm.ones(1)
 

brainpy/integrators/sde/tests/test_sde_scalar.py

@@ -8,6 +8,8 @@
 import brainpy as bp
 from brainpy.integrators import sde
 
+plt = None
+
 sigma = 10
 beta = 8 / 3
 rho = 28
@@ -42,17 +44,21 @@ def lorenz_system(method, **kwargs):
     mon1.append(x)
     mon2.append(y)
     mon3.append(z)
-  mon1 = bp.math.array(mon1)
-  mon2 = bp.math.array(mon2)
-  mon3 = bp.math.array(mon3)
-
-  # fig = plt.figure()
-  # ax = fig.gca(projection='3d')
-  # plt.plot(mon1, mon2, mon3)
-  # ax.set_xlabel('x')
-  # ax.set_xlabel('y')
-  # ax.set_xlabel('z')
-  # plt.show()
+  mon1 = bp.math.array(mon1).to_numpy()
+  mon2 = bp.math.array(mon2).to_numpy()
+  mon3 = bp.math.array(mon3).to_numpy()
+
+  global plt
+  if plt is None:
+    import matplotlib.pyplot as plt
+
+  fig = plt.figure()
+  ax = fig.gca(projection='3d')
+  plt.plot(mon1, mon2, mon3)
+  ax.set_xlabel('x')
+  ax.set_xlabel('y')
+  ax.set_xlabel('z')
+  plt.show()
 
 
 class TestScalarWienerIntegral(unittest.TestCase):

brainpy/integrators/tests/test_integ_runner.py

@@ -2,13 +2,16 @@
 
 from unittest import TestCase
 
-import matplotlib.pyplot as plt
-
 import brainpy as bp
 
+plt = None
+
 
 class TestIntegratorRunnerForODEs(TestCase):
   def test_ode(self):
+    global plt
+    if plt is None:
+      import matplotlib.pyplot as plt
     sigma = 10
     beta = 8 / 3
     rho = 28

brainpy/nn/nodes/ANN/normalization.py

@@ -6,7 +6,6 @@
 import jax.numpy as jnp
 
 import brainpy.math as bm
-import brainpy
 from brainpy.initialize import ZeroInit, OneInit, Initializer
 from brainpy.nn.base import Node