From 71e0874ab54e8289c6990014246946423658ed19 Mon Sep 17 00:00:00 2001
From: chaoming <adaduo@outlook.com>
Date: Mon, 18 Apr 2022 22:16:30 +0800
Subject: [PATCH] update tests

---
 .../analysis/lowdim/tests/test_phase_plane.py |  5 +++-
 brainpy/initialize/tests/test_decay_inits.py  |  9 ++++--
 brainpy/inputs/tests/test_currents.py         |  6 +++-
 brainpy/integrators/fde/tests/test_GL.py      |  8 +++++-
 .../ode/tests/test_ode_method_adaptive_rk.py  |  8 +++++-
 .../ode/tests/test_ode_method_exp_euler.py    |  8 ++++--
 .../ode/tests/test_ode_method_rk.py           |  6 +++-
 .../integrators/sde/tests/test_sde_scalar.py  | 28 +++++++++++--------
 .../integrators/tests/test_integ_runner.py    |  7 +++--
 brainpy/nn/nodes/ANN/normalization.py         |  1 -
 brainpy/nn/nodes/ANN/pooling.py               |  1 -
 brainpy/nn/nodes/ANN/tests/test_conv.py       |  1 -
 brainpy/nn/nodes/ANN/tests/test_pooling.py    |  1 -
 brainpy/nn/tests/test_vis.py                  |  1 -
 14 files changed, 63 insertions(+), 27 deletions(-)

diff --git a/brainpy/analysis/lowdim/tests/test_phase_plane.py b/brainpy/analysis/lowdim/tests/test_phase_plane.py
index 53f06b1c4..735029623 100644
--- a/brainpy/analysis/lowdim/tests/test_phase_plane.py
+++ b/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]
diff --git a/brainpy/initialize/tests/test_decay_inits.py b/brainpy/initialize/tests/test_decay_inits.py
index 50d39739f..d03bd4c4d 100644
--- a/brainpy/initialize/tests/test_decay_inits.py
+++ b/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.),
diff --git a/brainpy/inputs/tests/test_currents.py b/brainpy/inputs/tests/test_currents.py
index 8aa0664de..d7b7ddbb2 100644
--- a/brainpy/inputs/tests/test_currents.py
+++ b/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)
diff --git a/brainpy/integrators/fde/tests/test_GL.py b/brainpy/integrators/fde/tests/test_GL.py
index fed98ac62..75cbc3118 100644
--- a/brainpy/integrators/fde/tests/test_GL.py
+++ b/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):
diff --git a/brainpy/integrators/ode/tests/test_ode_method_adaptive_rk.py b/brainpy/integrators/ode/tests/test_ode_method_adaptive_rk.py
index ae5c3a926..29cf849a8 100644
--- a/brainpy/integrators/ode/tests/test_ode_method_adaptive_rk.py
+++ b/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
diff --git a/brainpy/integrators/ode/tests/test_ode_method_exp_euler.py b/brainpy/integrators/ode/tests/test_ode_method_exp_euler.py
index cafec77db..4aa10bdbd 100644
--- a/brainpy/integrators/ode/tests/test_ode_method_exp_euler.py
+++ b/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'):
diff --git a/brainpy/integrators/ode/tests/test_ode_method_rk.py b/brainpy/integrators/ode/tests/test_ode_method_rk.py
index f8f1aec7c..62f2ba61f 100644
--- a/brainpy/integrators/ode/tests/test_ode_method_rk.py
+++ b/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)
 
diff --git a/brainpy/integrators/sde/tests/test_sde_scalar.py b/brainpy/integrators/sde/tests/test_sde_scalar.py
index 4e2e9000b..9f319b737 100644
--- a/brainpy/integrators/sde/tests/test_sde_scalar.py
+++ b/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):
diff --git a/brainpy/integrators/tests/test_integ_runner.py b/brainpy/integrators/tests/test_integ_runner.py
index a20ca1920..bcda4be1e 100644
--- a/brainpy/integrators/tests/test_integ_runner.py
+++ b/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
diff --git a/brainpy/nn/nodes/ANN/normalization.py b/brainpy/nn/nodes/ANN/normalization.py
index b289ba022..86db5dc0c 100644
--- a/brainpy/nn/nodes/ANN/normalization.py
+++ b/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
 
diff --git a/brainpy/nn/nodes/ANN/pooling.py b/brainpy/nn/nodes/ANN/pooling.py
index 7ae25b50e..e4c674369 100644
--- a/brainpy/nn/nodes/ANN/pooling.py
+++ b/brainpy/nn/nodes/ANN/pooling.py
@@ -3,7 +3,6 @@
 
 import jax.lax
 import brainpy.math as bm
-from brainpy.initialize import XavierNormal, ZeroInit, init_param
 from brainpy.nn.base import Node
 
 __all__ = [
diff --git a/brainpy/nn/nodes/ANN/tests/test_conv.py b/brainpy/nn/nodes/ANN/tests/test_conv.py
index 7fd245a0d..b2d57a782 100644
--- a/brainpy/nn/nodes/ANN/tests/test_conv.py
+++ b/brainpy/nn/nodes/ANN/tests/test_conv.py
@@ -6,7 +6,6 @@
 import brainpy as bp
 import jax.numpy as jnp
 import numpy as np
-import matplotlib.pyplot as plt
 
 class TestConv(TestCase):
   def test_Conv2D_img(self):
diff --git a/brainpy/nn/nodes/ANN/tests/test_pooling.py b/brainpy/nn/nodes/ANN/tests/test_pooling.py
index 046d5552b..8ca5720a9 100644
--- a/brainpy/nn/nodes/ANN/tests/test_pooling.py
+++ b/brainpy/nn/nodes/ANN/tests/test_pooling.py
@@ -7,7 +7,6 @@
 import jax.numpy as jnp
 import jax
 import numpy as np
-import matplotlib.pyplot as plt
 
 
 class TestPool(TestCase):
diff --git a/brainpy/nn/tests/test_vis.py b/brainpy/nn/tests/test_vis.py
index 4476a0127..09d0d503a 100644
--- a/brainpy/nn/tests/test_vis.py
+++ b/brainpy/nn/tests/test_vis.py
@@ -3,7 +3,6 @@
 
 import unittest
 import brainpy as bp
-import brainpy.math as bm
 
 
 class TestVisualize(unittest.TestCase):