fix bugs on slow point analysis

Author: chaoming0625

brainpy/analysis/highdim/slow_points.py

@@ -355,8 +355,7 @@ def train(idx):
       return loss
 
     def batch_train(start_i, n_batch):
-      f = bm.make_loop(train, dyn_vars=dyn_vars, has_return=True)
-      return f(bm.arange(start_i, start_i + n_batch))
+      return bm.for_loop(train, dyn_vars, bm.arange(start_i, start_i + n_batch))
 
     # Run the optimization
     if self.verbose:
@@ -369,7 +368,7 @@ def batch_train(start_i, n_batch):
         break
       batch_idx_start = oidx * num_batch
       start_time = time.time()
-      (_, train_losses) = batch_train(start_i=batch_idx_start, n_batch=num_batch)
+      train_losses = batch_train(start_i=batch_idx_start, n_batch=num_batch)
       batch_time = time.time() - start_time
       opt_losses.append(train_losses)
 
@@ -722,8 +721,6 @@ def _generate_ds_cell_function(
     shared = DotDict(t=t, dt=dt, i=0)
 
     def f_cell(h: Dict):
-      target.clear_input()
-
       # update target variables
       for k, v in self.target_vars.items():
         v.value = (bm.asarray(h[k], dtype=v.dtype)
@@ -735,6 +732,7 @@ def f_cell(h: Dict):
         v.value = self.excluded_data[k]
 
       # add inputs
+      target.clear_input()
       if f_input is not None:
         f_input(shared)
 
@@ -743,7 +741,7 @@ def f_cell(h: Dict):
       target.update(*args)
 
       # get new states
-      new_h = {k: (v.value if v.batch_axis is None else jnp.squeeze(v.value, axis=v.batch_axis))
+      new_h = {k: (v.value if (v.batch_axis is None) else jnp.squeeze(v.value, axis=v.batch_axis))
                for k, v in self.target_vars.items()}
       return new_h
 

brainpy/analysis/lowdim/lowdim_bifurcation.py

@@ -354,8 +354,17 @@ def plot_bifurcation(self, with_plot=True, show=False, with_return=False,
     if with_return:
       return final_fps, final_pars, jacobians
 
-  def plot_limit_cycle_by_sim(self, duration=100, with_plot=True, with_return=False,
-                              plot_style=None, tol=0.001, show=False, dt=None, offset=1.):
+  def plot_limit_cycle_by_sim(
+      self,
+      duration=100,
+      with_plot: bool = True,
+      with_return: bool = False,
+      plot_style: dict = None,
+      tol: float = 0.001,
+      show: bool = False,
+      dt: float = None,
+      offset: float = 1.
+  ):
     global pyplot
     if pyplot is None: from matplotlib import pyplot
     utils.output('I am plotting the limit cycle ...')
@@ -400,10 +409,16 @@ def plot_limit_cycle_by_sim(self, duration=100, with_plot=True, with_return=Fals
         if len(ps_limit_cycle[0]):
           for i, var in enumerate(self.target_var_names):
             pyplot.figure(var)
-            pyplot.plot(ps_limit_cycle[0], ps_limit_cycle[1], vs_limit_cycle[i]['max'],
-                        **plot_style, label='limit cycle (max)')
-            pyplot.plot(ps_limit_cycle[0], ps_limit_cycle[1], vs_limit_cycle[i]['min'],
-                        **plot_style, label='limit cycle (min)')
+            pyplot.plot(ps_limit_cycle[0],
+                        ps_limit_cycle[1],
+                        vs_limit_cycle[i]['max'],
+                        **plot_style,
+                        label='limit cycle (max)')
+            pyplot.plot(ps_limit_cycle[0],
+                        ps_limit_cycle[1],
+                        vs_limit_cycle[i]['min'],
+                        **plot_style,
+                        label='limit cycle (min)')
             pyplot.legend()
 
       elif len(self.target_par_names) == 1:
@@ -427,8 +442,16 @@ def plot_limit_cycle_by_sim(self, duration=100, with_plot=True, with_return=Fals
 
 
 class FastSlow1D(Bifurcation1D):
-  def __init__(self, model, fast_vars, slow_vars, fixed_vars=None,
-               pars_update=None, resolutions=None, options=None):
+  def __init__(
+      self,
+      model,
+      fast_vars: dict,
+      slow_vars: dict,
+      fixed_vars: dict = None,
+      pars_update: dict = None,
+      resolutions=None,
+      options: dict = None
+  ):
     super(FastSlow1D, self).__init__(model=model,
                                      target_pars=slow_vars,
                                      target_vars=fast_vars,
@@ -510,8 +533,16 @@ def plot_trajectory(self, initials, duration, plot_durations=None,
 
 
 class FastSlow2D(Bifurcation2D):
-  def __init__(self, model, fast_vars, slow_vars, fixed_vars=None,
-               pars_update=None, resolutions=0.1, options=None):
+  def __init__(
+      self,
+      model,
+      fast_vars: dict,
+      slow_vars: dict,
+      fixed_vars: dict = None,
+      pars_update: dict = None,
+      resolutions=0.1,
+      options: dict = None
+  ):
     super(FastSlow2D, self).__init__(model=model,
                                      target_pars=slow_vars,
                                      target_vars=fast_vars,

brainpy/analysis/utils/model.py

@@ -112,14 +112,14 @@ def __init__(self, integrals: dict, initial_vars: dict, pars=None, dt=None):
 
     # variables
     assert isinstance(initial_vars, dict)
-    initial_vars = {k: bm.Variable(jnp.asarray(bm.as_device_array(v), dtype=jnp.float_))
+    initial_vars = {k: bm.Variable(jnp.asarray(bm.as_device_array(v), dtype=bm.dftype()))
                     for k, v in initial_vars.items()}
     self.register_implicit_vars(initial_vars)
 
     # parameters
     pars = dict() if pars is None else pars
     assert isinstance(pars, dict)
-    self.pars = [jnp.asarray(bm.as_device_array(v), dtype=jnp.float_)
+    self.pars = [jnp.asarray(bm.as_device_array(v), dtype=bm.dftype())
                  for k, v in pars.items()]
 
     # integrals
@@ -128,7 +128,8 @@ def __init__(self, integrals: dict, initial_vars: dict, pars=None, dt=None):
     # runner
     self.runner = DSRunner(self,
                            monitors=list(initial_vars.keys()),
-                           dyn_vars=self.vars().unique(), dt=dt,
+                           dyn_vars=self.vars().unique(),
+                           dt=dt,
                            progress_bar=False)
 
   def update(self, sha):

brainpy/dyn/neurons/biological_models.py

@@ -244,20 +244,17 @@ def __init__(
 
     # variables
     self.V = variable(self._V_initializer, mode, self.varshape)
-    if self._m_initializer is None:
-      self.m = bm.Variable(self.m_inf(self.V.value))
-    else:
-      self.m = variable(self._m_initializer, mode, self.varshape)
-    if self._h_initializer is None:
-      self.h = bm.Variable(self.h_inf(self.V.value))
-    else:
-      self.h = variable(self._h_initializer, mode, self.varshape)
-    if self._n_initializer is None:
-      self.n = bm.Variable(self.n_inf(self.V.value))
-    else:
-      self.n = variable(self._n_initializer, mode, self.varshape)
-    self.input = variable(bm.zeros, mode, self.varshape)
+    self.m = (bm.Variable(self.m_inf(self.V.value))
+              if m_initializer is None else
+              variable(self._m_initializer, mode, self.varshape))
+    self.h = (bm.Variable(self.h_inf(self.V.value))
+              if h_initializer is None else
+              variable(self._h_initializer, mode, self.varshape))
+    self.n = (bm.Variable(self.n_inf(self.V.value))
+              if n_initializer is None else
+              variable(self._n_initializer, mode, self.varshape))
     self.spike = variable(lambda s: bm.zeros(s, dtype=bool), mode, self.varshape)
+    self.input = variable(bm.zeros, mode, self.varshape)
 
     # integral
     if self.noise is None:
@@ -309,7 +306,7 @@ def dV(self, V, t, m, h, n, I_ext):
 
   @property
   def derivative(self):
-    return JointEq([self.dV, self.dm, self.dh, self.dn])
+    return JointEq(self.dV, self.dm, self.dh, self.dn)
 
   def update(self, tdi, x=None):
     t, dt = tdi['t'], tdi['dt']

brainpy/inputs/currents.py

@@ -307,9 +307,9 @@ def ou_process(mean, sigma, tau, duration, dt=None, n=1, t_start=0., t_end=None,
 
   def _f(t):
     x.value = x + dt * ((mean - x) / tau) + sigma * dt_sqrt * rng.rand(n)
+    return x.value
 
-  f = bm.make_loop(_f, dyn_vars=[x, rng], out_vars=x)
-  noises = f(jnp.arange(t_start, t_end, dt))
+  noises = bm.for_loop(_f, [x, rng], jnp.arange(t_start, t_end, dt))
 
   t_end = duration if t_end is None else t_end
   i_start = int(t_start / dt)

brainpy/math/operators/utils.py

@@ -17,12 +17,14 @@ def _check_brainpylib(ops_name):
       raise PackageMissingError(
         f'"{ops_name}" operator need "brainpylib>={_BRAINPYLIB_MINIMAL_VERSION}". \n'
         f'Please install it through:\n\n'
-        f'>>> pip install brainpylib>={_BRAINPYLIB_MINIMAL_VERSION} -U'
+        f'>>> pip install brainpylib=={_BRAINPYLIB_MINIMAL_VERSION}\n'
+        f'>>> # or \n'
+        f'>>> pip install brainpylib -U'
       )
   else:
     raise PackageMissingError(
       f'"brainpylib" must be installed when the user '
       f'wants to use "{ops_name}" operator. \n'
       f'Please install "brainpylib>={_BRAINPYLIB_MINIMAL_VERSION}" through:\n\n'
-      f'>>> pip install brainpylib>={_BRAINPYLIB_MINIMAL_VERSION}'
+      f'>>> pip install brainpylib'
     )

brainpy/train/back_propagation.py

@@ -520,7 +520,7 @@ def f_train(self, shared_args=None) -> Callable:
     shared_args_str = serialize_kwargs(shared_args)
     if shared_args_str not in self._f_train_compiled:
 
-      def train_step(x):
+      def train_step(*x):
         # t, i, input_, target_ = x
         res = self.f_grad(shared_args)(*x)
         self.optimizer.update(res[0])
@@ -529,8 +529,7 @@ def train_step(x):
       if self.jit[c.FIT_PHASE]:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
-        f = bm.make_loop(train_step, dyn_vars=dyn_vars.unique(), has_return=True)
-        run_func = lambda all_inputs: f(all_inputs)[1]
+        run_func = lambda all_inputs: bm.for_loop(train_step, dyn_vars.unique(), all_inputs)
 
       else:
         def run_func(xs):
@@ -541,7 +540,7 @@ def run_func(xs):
             x = tree_map(lambda x: x[i], inputs, is_leaf=_is_jax_array)
             y = tree_map(lambda x: x[i], targets, is_leaf=_is_jax_array)
             # step at the i
-            loss = train_step((times[i], indices[i], x, y))
+            loss = train_step(times[i], indices[i], x, y)
             # append output and monitor
             losses.append(loss)
           return bm.asarray(losses)

examples/analysis/4d_HH_model.py

@@ -3,13 +3,15 @@
 import brainpy as bp
 import brainpy.math as bm
 
+
 I = 5.
 model = bp.dyn.neurons.HH(1)
-runner = bp.dyn.DSRunner(model, inputs=('input', I), monitors=['V'])
+runner = bp.dyn.DSRunner(model, inputs=(model.input, I), monitors=['V'])
 runner.run(100)
 bp.visualize.line_plot(runner.mon.ts, runner.mon.V, legend='V', show=True)
 
 # analysis
+bm.enable_x64()
 model = bp.dyn.neurons.HH(1, method='euler')
 finder = bp.analysis.SlowPointFinder(
   model,
@@ -18,13 +20,14 @@
                'm': model.m,
                'h': model.h,
                'n': model.n},
-  dt=1.
+  dt=10.
 )
-candidates = {'V': bm.random.normal(0., 5., (1000, model.num)) - 50.,
+finder.find_fps_with_opt_solver(
+  candidates={'V': bm.random.normal(0., 10., (1000, model.num)) - 50.,
               'm': bm.random.random((1000, model.num)),
               'h': bm.random.random((1000, model.num)),
               'n': bm.random.random((1000, model.num))}
-finder.find_fps_with_opt_solver(candidates=candidates)
+)
 finder.filter_loss(1e-7)
 finder.keep_unique(tolerance=1e-1)
 print('fixed_points: ', finder.fixed_points)

examples/training/Song_2016_EI_RNN.py

@@ -128,17 +128,17 @@ def __init__(self, num_input, num_hidden, num_output, num_batch,
     self.mask = bm.asarray(mask, dtype=bm.dftype())
 
     # input weight
-    self.w_ir = bm.TrainVar(w_ir(num_input, num_hidden))
+    self.w_ir = bm.TrainVar(w_ir((num_input, num_hidden)))
 
     # recurrent weight
     bound = 1 / num_hidden ** 0.5
-    self.w_rr = bm.TrainVar(w_rr(num_hidden, num_hidden))
+    self.w_rr = bm.TrainVar(w_rr((num_hidden, num_hidden)))
     self.w_rr[:, :self.e_size] /= (self.e_size / self.i_size)
     self.b_rr = bm.TrainVar(self.rng.uniform(-bound, bound, num_hidden))
 
     # readout weight
     bound = 1 / self.e_size ** 0.5
-    self.w_ro = bm.TrainVar(w_ro(self.e_size, num_output))
+    self.w_ro = bm.TrainVar(w_ro((self.e_size, num_output)))
     self.b_ro = bm.TrainVar(self.rng.uniform(-bound, bound, num_output))
 
     # variables
@@ -158,15 +158,13 @@ def make_update(self, h: bm.JaxArray, o: bm.JaxArray):
     def f(x):
       h.value = self.cell(x, h.value)
       o.value = self.readout(h.value[:, :self.e_size])
+      return h.value, o.value
 
     return f
 
   def predict(self, xs):
     self.h[:] = 0.
-    f = bm.make_loop(self.make_update(self.h, self.o),
-                     dyn_vars=self.vars(),
-                     out_vars=[self.h, self.o])
-    return f(xs)
+    return bm.for_loop(self.make_update(self.h, self.o), self.vars(), xs)
 
   def loss(self, xs, ys):
     hs, os = self.predict(xs)
@@ -247,7 +245,7 @@ def train(xs, ys):
   rnn_activity, action_pred = predict(inputs)
 
   # Compute performance
-  action_pred = action_pred.numpy()
+  action_pred = bm.as_numpy(action_pred)
   choice = np.argmax(action_pred[-1, 0, :])
   correct = choice == gt[-1]
 
@@ -257,7 +255,7 @@ def train(xs, ys):
   trial_infos[i] = trial_info
 
   # Log stimulus period activity
-  rnn_activity = rnn_activity.numpy()[:, 0, :]
+  rnn_activity = bm.as_numpy(rnn_activity)[:, 0, :]
   activity_dict[i] = rnn_activity
 
   # Compute stimulus selectivity for all units
@@ -312,7 +310,7 @@ def train(xs, ys):
 plt.show()
 
 # %%
-W = (bm.abs(net.w_rr) * net.mask).numpy()
+W = bm.as_numpy(bm.abs(net.w_rr) * net.mask)
 # Sort by selectivity
 W = W[:, ind_sort][ind_sort, :]
 wlim = np.max(np.abs(W))