fix #294: remove VariableView in dyn_vars of a runner (#295)

fix #294: remove VariableView in `dyn_vars` of a runner

Author: chaoming0625

brainpy/dyn/runners.py

@@ -585,6 +585,7 @@ def _step_func(t, i, x):
       if self.jit['predict']:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         run_func = lambda all_inputs: bm.for_loop(_step_func, dyn_vars.unique(), all_inputs)
 
       else:

brainpy/dyn/tests/test_dyn_runner.py

@@ -32,6 +32,53 @@ def update(self, tdi):
     runner = bp.dyn.DSRunner(ExampleDS(), dt=1., monitors=['i'], progress_bar=False)
     runner.run(100.)
 
+  def test_DSView(self):
+    class EINet(bp.dyn.Network):
+      def __init__(self, scale=1.0, method='exp_auto'):
+        super(EINet, self).__init__()
+
+        # network size
+        num_exc = int(800 * scale)
+        num_inh = int(200 * scale)
+
+        # neurons
+        pars = dict(V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.)
+        self.E = bp.neurons.LIF(num_exc, **pars, method=method)
+        self.I = bp.neurons.LIF(num_inh, **pars, method=method)
+        self.E.V[:] = bm.random.randn(num_exc) * 2 - 55.
+        self.I.V[:] = bm.random.randn(num_inh) * 2 - 55.
+
+        # synapses
+        we = 0.6 / scale  # excitatory synaptic weight (voltage)
+        wi = 6.7 / scale  # inhibitory synaptic weight
+        self.E2E = bp.synapses.Exponential(self.E, self.E[:100], bp.conn.FixedProb(0.02),
+                                           output=bp.synouts.COBA(E=0.), g_max=we,
+                                           tau=5., method=method)
+        self.E2I = bp.synapses.Exponential(self.E, self.I[:100], bp.conn.FixedProb(0.02),
+                                           output=bp.synouts.COBA(E=0.), g_max=we,
+                                           tau=5., method=method)
+        self.I2E = bp.synapses.Exponential(self.I, self.E[:100], bp.conn.FixedProb(0.02),
+                                           output=bp.synouts.COBA(E=-80.), g_max=wi,
+                                           tau=10., method=method)
+        self.I2I = bp.synapses.Exponential(self.I, self.I[:100], bp.conn.FixedProb(0.02),
+                                           output=bp.synouts.COBA(E=-80.), g_max=wi,
+                                           tau=10., method=method)
+
+    net = EINet(scale=1., method='exp_auto')
+    # with JIT
+    runner = bp.DSRunner(net, monitors={'E.spike': net.E.spike},
+                         inputs=[(net.E.input, 20.), (net.I.input, 20.)]).run(1.)
+
+    # without JIT
+    runner = bp.DSRunner(net, monitors={'E.spike': net.E.spike},
+                         inputs=[(net.E.input, 20.), (net.I.input, 20.)],
+                         jit=False).run(0.2)
+
+
+
+
+
+
 # class TestMonitor(TestCase):
 #   def test_1d_array(self):
 #     try1 = TryGroup(monitors=['a'])

brainpy/measure/correlation.py

@@ -166,7 +166,7 @@ def voltage_fluctuation(potentials, numpy=True, method='loop'):
   References
   ----------
   .. [1] Golomb, D. and Rinzel J. (1993) Dynamics of globally coupled
-         inhibitory neurons with heterogeneity. Phys. Rev. reversal_potential 48:4810-4814.
+         inhibitory neurons with heterogeneity. Phys. Rev. E 48:4810-4814.
   .. [2] Golomb D. and Rinzel J. (1994) Clustering in globally coupled
          inhibitory neurons. Physica D 72:259-282.
   .. [3] David Golomb (2007) Neuronal synchrony measures. Scholarpedia, 2(1):1347.

brainpy/running/multiprocess.py

@@ -21,7 +21,7 @@ def process_pool(func, all_params, num_process):
   ----------
   func : callable
       The function to run model.
-  all_params : a_list, tuple
+  all_params : list, tuple, dict
       The parameters of the function arguments.
       The parameters for each process can be a tuple, or a dictionary.
   num_process : int
@@ -47,7 +47,7 @@ def process_pool(func, all_params, num_process):
   return [r.get() for r in results]
 
 
-def process_pool_lock(func, all_net_params, nb_process):
+def process_pool_lock(func, all_params, nb_process):
   """Run multiple models in multi-processes with lock.
 
   Sometimes, you want to synchronize the processes. For example,
@@ -73,7 +73,7 @@ def some_func(..., lock, ...):
   ----------
   func : callable
       The function to run model.
-  all_net_params : a_list, tuple
+  all_params : list, tuple, dict
       The parameters of the function arguments.
   nb_process : int
       The number of the processes.
@@ -83,12 +83,12 @@ def some_func(..., lock, ...):
   results : list
       Process results.
   """
-  print('{} jobs total.'.format(len(all_net_params)))
+  print('{} jobs total.'.format(len(all_params)))
   pool = multiprocessing.Pool(processes=nb_process)
   m = multiprocessing.Manager()
   lock = m.Lock()
   results = []
-  for net_params in all_net_params:
+  for net_params in all_params:
     if isinstance(net_params, (list, tuple)):
       results.append(pool.apply_async(func, args=tuple(net_params) + (lock,)))
     elif isinstance(net_params, dict):

brainpy/train/back_propagation.py

@@ -300,6 +300,7 @@ def loss_fun(inputs, targets):
       if self.jit[c.LOSS_PHASE] and jit:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         self._f_loss_compiled[shared_args_str] = bm.jit(self._f_loss_compiled[shared_args_str],
                                                         dyn_vars=dyn_vars)
     return self._f_loss_compiled[shared_args_str]
@@ -311,6 +312,7 @@ def f_grad(self, shared_args=None) -> Callable:
       _f_loss_internal = self.f_loss(shared_args, jit=False)
       dyn_vars = self.target.vars()
       dyn_vars.update(self.dyn_vars)
+      dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
       tran_vars = dyn_vars.subset(bm.TrainVar)
       grad_f = bm.grad(_f_loss_internal,
                        dyn_vars=dyn_vars.unique(),
@@ -339,6 +341,7 @@ def train_func(inputs, targets):
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
         dyn_vars.update(self.optimizer.vars())
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         self._f_train_compiled[shared_args_str] = bm.jit(train_func, dyn_vars=dyn_vars.unique())
       else:
         self._f_train_compiled[shared_args_str] = train_func
@@ -453,6 +456,7 @@ def loss_fun(inputs, targets):
       if self.jit[c.LOSS_PHASE] and jit:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         self._f_loss_compiled[shared_args_str] = bm.jit(self._f_loss_compiled[shared_args_str],
                                                         dyn_vars=dyn_vars)
       else:
@@ -480,6 +484,7 @@ def run_func(xs):
       if self.jit[c.PREDICT_PHASE] and jit:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         self._f_predict_compiled[shared_args_str] = bm.jit(run_func, dyn_vars=dyn_vars.unique())
       else:
         self._f_predict_compiled[shared_args_str] = run_func
@@ -505,6 +510,7 @@ def loss_fun(t, i, input_, target_):
       if self.jit[c.LOSS_PHASE] and jit:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         self._f_loss_compiled[shared_args_str] = bm.jit(self._f_loss_compiled[shared_args_str],
                                                         dyn_vars=dyn_vars)
       else:
@@ -529,6 +535,7 @@ def train_step(*x):
       if self.jit[c.FIT_PHASE]:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         run_func = lambda all_inputs: bm.for_loop(train_step, dyn_vars.unique(), all_inputs)
 
       else:
@@ -582,6 +589,7 @@ def run_func(t, i, x):
       if self.jit[c.FIT_PHASE] and jit:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
+        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
         self._f_predict_compiled[shared_args_str] = bm.jit(run_func, dyn_vars=dyn_vars.unique())
       else:
         self._f_predict_compiled[shared_args_str] = run_func

brainpy/train/offline.py

@@ -231,6 +231,7 @@ def train_func(monitor_data: Dict[str, Array], target_data: Dict[str, Array]):
     if self.jit['fit']:
       dyn_vars = self.target.vars()
       dyn_vars.update(self.dyn_vars)
+      dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
       train_func = bm.jit(train_func, dyn_vars=dyn_vars.unique())
     return train_func
 

brainpy/train/online.py

@@ -261,6 +261,7 @@ def _step_func(t, i, x, ys):
     if self.jit['fit']:
       dyn_vars = self.target.vars()
       dyn_vars.update(self.dyn_vars)
+      dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
       return lambda all_inputs: bm.for_loop(_step_func, dyn_vars.unique(), all_inputs)
 
     else: