update examples

Author: chaoming0625

examples/dynamics_analysis/2d_fitzhugh_nagumo_model.py

@@ -33,8 +33,9 @@ def dw(w, t, V, a=0.7, b=0.8):
     self.int_V = bp.odeint(dV, method=method)
     self.int_w = bp.odeint(dw, method=method)
 
-  def update(self, tdi):
-    t, dt = tdi['t'], tdi['dt']
+  def update(self):
+    t = bp.share['t']
+    dt = bp.share['dt']
     self.V.value = self.int_V(self.V, t, self.w, self.Iext, dt)
     self.w.value = self.int_w(self.w, t, self.V, self.a, self.b, dt)
     self.Iext[:] = 0.

examples/dynamics_analysis/2d_mean_field_QIF.py

@@ -14,7 +14,7 @@ class MeanFieldQIF(bp.DynamicalSystem):
   """
 
   def __init__(self, method='exp_auto'):
-    super(MeanFieldQIF, self).__init__()
+    super().__init__()
 
     # parameters
     self.tau = 1.  # the population time constant
@@ -38,8 +38,9 @@ def dv(v, t, r, Iext=0., eta=-5.0):
     self.int_r = bp.odeint(dr, method=method)
     self.int_v = bp.odeint(dv, method=method)
 
-  def update(self, tdi):
-    t, dt = tdi['t'], tdi['dt']
+  def update(self):
+    t = bp.share['t']
+    dt = bp.share['dt']
     self.r.value = self.int_r(self.r, t, self.v, self.delta, dt)
     self.v.value = self.int_v(self.v, t, self.r, self.Iext, self.eta, dt)
     self.Iext[:] = 0.

examples/dynamics_analysis/3d_reduced_trn_model.py

@@ -7,9 +7,9 @@
 bp.math.set_platform('cpu')
 
 
-class ReducedTRNModel(bp.NeuDyn):
+class ReducedTRNModel(bp.dyn.NeuDyn):
   def __init__(self, size, name=None, T=36., method='rk4'):
-    super(ReducedTRNModel, self).__init__(size=size, name=name)
+    super().__init__(size=size, name=name)
 
     self.IT_th = -3.
     self.b = 0.5

examples/dynamics_analysis/highdim_RNN_Analysis.py

@@ -26,7 +26,7 @@ def __init__(
       w_rr=bp.init.KaimingNormal(scale=1.),
       w_ro=bp.init.KaimingNormal(scale=1.)
   ):
-    super(RNNNet, self).__init__()
+    super().__init__()
 
     self.tau = 100
     self.num_input = num_input
@@ -64,7 +64,7 @@ def cell(self, x, h):
   def readout(self, h):
     return h @ self.w_ro + self.b_ro
 
-  def update(self, sha, x):
+  def update(self, x):
     self.h.value = self.cell(x, self.h.value)
     return self.readout(self.h.value)
 

examples/dynamics_simulation/COBA-v2.py renamed to examples/dynamics_simulation/COBA.py

@@ -140,12 +140,6 @@ def __init__(self, scale=1.0, method='exp_auto'):
 # bm.set_host_device_count(num_device)
 # bm.sharding.set(mesh_axes=(bp.dyn.PNEU_AXIS,), mesh_shape=(num_device, ))
 
-def run3():
-  net = EICOBA_PreAlign(3200, 800)
-  runner = bp.DSRunner(net, monitors={'E.spike': net.E.spike})
-  print(runner.run(100., eval_time=True))
-  bp.visualize.raster_plot(runner.mon.ts, runner.mon['E.spike'], show=True)
-
 
 def run1():
   with bm.environment(mode=bm.BatchingMode(10)):
@@ -167,7 +161,23 @@ def run2():
   bp.visualize.raster_plot(runner.mon.ts, runner.mon['E.spike'], show=True)
 
 
+def run3():
+  net = EICOBA_PreAlign(3200, 800)
+  runner = bp.DSRunner(net, monitors={'E.spike': net.E.spike})
+  print(runner.run(100., eval_time=True))
+  bp.visualize.raster_plot(runner.mon.ts, runner.mon['E.spike'], show=True)
+
+
+
+def run4():
+  net = EICOBA_PostAlign(3200, 800)
+  runner = bp.DSRunner(net, monitors={'E.spike': net.E.spike})
+  print(runner.run(100., eval_time=True))
+  bp.visualize.raster_plot(runner.mon.ts, runner.mon['E.spike'], show=True)
+
+
 if __name__ == '__main__':
-  # run1()
+  run1()
   run2()
-  # run3()
+  run3()
+  run4()

examples/dynamics_simulation/hh_model.py

@@ -11,32 +11,27 @@
 
 class HH(bp.dyn.CondNeuGroup):
   def __init__(self, size):
-    super().__init__(size, keep_size=True)
+    super().__init__(size)
 
-    self.INa = bp.channels.INa_HH1952(size, keep_size=True)
-    self.IK = bp.channels.IK_HH1952(size, keep_size=True)
-    self.IL = bp.channels.IL(size, E=-54.387, g_max=0.03, keep_size=True)
+    self.INa = bp.channels.INa_HH1952(size)
+    self.IK = bp.channels.IK_HH1952(size)
+    self.IL = bp.channels.IL(size, E=-54.387, g_max=0.03)
 
 
 class HHv2(bp.dyn.CondNeuGroupLTC):
   def __init__(self, size):
-    super().__init__(size, keep_size=True)
+    super().__init__(size)
 
     self.Na = bp.dyn.SodiumFixed(size, E=50.)
-    self.Na.add(ina=bp.dyn.INa_HH1952v2(size, keep_size=True))
+    self.Na.add_elem(ina=bp.dyn.INa_HH1952v2(size))
 
     self.K = bp.dyn.PotassiumFixed(size, E=50.)
-    self.K.add(ik=bp.dyn.IK_HH1952v2(size, keep_size=True))
-
-    self.IL = bp.dyn.IL(size, E=-54.387, g_max=0.03, keep_size=True)
-
-    self.KNa = bp.dyn.mixs(self.Na, self.K)
-    self.KNa.add()
-
-
-
+    self.K.add_elem(ik=bp.dyn.IK_HH1952v2(size))
 
+    self.IL = bp.dyn.IL(size, E=-54.387, g_max=0.03)
 
+    self.KNa = bp.dyn.MixIons(self.Na, self.K)
+    self.KNa.add_elem()
 
 
 # hh = HH(1)
@@ -52,26 +47,3 @@ def __init__(self, size):
 #
 # bp.visualize.line_plot(runner.mon.ts, runner.mon.V, show=True)
 
-
-hh = HH((20, 10000))
-variables = hh.vars().unique()
-
-
-iis = np.arange(1000000000)
-
-def f(i):
-  bp.share.save(i=i, t=i * bm.get_dt(), dt=bm.get_dt())
-  hh(5.)
-
-
-@pmap
-def run(vars):
-  for v, d in vars.items():
-    variables[v]._value = d
-  bm.for_loop(f, bm.arange(1000000000))
-  print('Compiling End')
-  return hh.spike
-
-
-r = run(variables.dict())
-print(r.shape)

examples/dynamics_simulation/multi_scale_COBAHH.py

@@ -7,12 +7,9 @@
 
 import brainpy as bp
 import brainpy.math as bm
-from brainpy.channels import INa_TM1991, IL
-from brainpy.synapses import Exponential, COBA
 from brainpy.connect import FixedProb
 from jax import vmap
 
-comp_method = 'sparse'
 
 
 area_names = ['V1', 'V2', 'V4', 'TEO', 'TEpd']
@@ -47,8 +44,8 @@ class HH(bp.CondNeuGroup):
   def __init__(self, size):
     super(HH, self).__init__(size, V_initializer=bp.init.Uniform(-70, -50.))
     self.IK = IK(size, g_max=30., V_sh=-63.)
-    self.INa = INa_TM1991(size, g_max=100., V_sh=-63.)
-    self.IL = IL(size, E=-60., g_max=0.05)
+    self.INa = bp.dyn.INa_TM1991(size, g_max=100., V_sh=-63.)
+    self.IL = bp.dyn.IL(size, E=-60., g_max=0.05)
 
 
 class Network(bp.Network):

examples/dynamics_simulation/whole_brain_simulation_with_fhn.py

@@ -21,7 +21,7 @@ def bifurcation_analysis():
   pp.show_figure()
 
 
-class Network(bp.Network):
+class Network(bp.DynSysGroup):
   def __init__(self, signal_speed=20.):
     super(Network, self).__init__()
 
@@ -36,12 +36,12 @@ def __init__(self, signal_speed=20.):
     delay_mat = bm.asarray(delay_mat)
     bm.fill_diagonal(delay_mat, 0)
 
-    self.fhn = bp.rates.FHN(
+    self.fhn = bp.dyn.FHN(
       80,
       x_ou_sigma=0.01,
       y_ou_sigma=0.01,
     )
-    self.coupling = bp.synapses.DiffusiveCoupling(
+    self.coupling = bp.dyn.DiffusiveCoupling(
       self.fhn.x,
       self.fhn.x,
       var_to_output=self.fhn.input,
@@ -95,5 +95,5 @@ def net_analysis():
 
 if __name__ == '__main__':
   # bifurcation_analysis()
-  # net_simulation()
-  net_analysis()
+  net_simulation()
+  # net_analysis()

examples/dynamics_simulation/whole_brain_simulation_with_sl_oscillator.py

@@ -10,7 +10,7 @@
 
 
 def bifurcation_analysis():
-  model = bp.rates.StuartLandauOscillator(1, method='exp_auto')
+  model = bp.dyn.StuartLandauOscillator(1, method='exp_auto')
   pp = bp.analysis.Bifurcation2D(
     model,
     target_vars={'x': [-2, 2], 'y': [-2, 2]},
@@ -22,7 +22,7 @@ def bifurcation_analysis():
   pp.show_figure()
 
 
-class Network(bp.Network):
+class Network(bp.DynSysGroup):
   def __init__(self, noise=0.14):
     super(Network, self).__init__()
 
@@ -35,8 +35,8 @@ def __init__(self, noise=0.14):
     bm.fill_diagonal(conn_mat, 0)
     gc = 0.6  # global coupling strength
 
-    self.sl = bp.rates.StuartLandauOscillator(80, x_ou_sigma=noise, y_ou_sigma=noise)
-    self.coupling = bp.synapses.DiffusiveCoupling(
+    self.sl = bp.dyn.StuartLandauOscillator(80, x_ou_sigma=noise, y_ou_sigma=noise)
+    self.coupling = bp.dyn.DiffusiveCoupling(
       self.sl.x, self.sl.x,
       var_to_output=self.sl.input,
       conn_mat=conn_mat * gc
@@ -87,6 +87,6 @@ def net_analysis():
 
 
 if __name__ == '__main__':
-  bifurcation_analysis()
+  # bifurcation_analysis()
   simulation()
-  net_analysis()
+  # net_analysis()

examples/dynamics_training/Song_2016_EI_RNN.py

@@ -27,7 +27,7 @@ def __init__(
       w_rr=bp.init.KaimingUniform(scale=1.),
       w_ro=bp.init.KaimingUniform(scale=1.)
   ):
-    super(EI_RNN, self).__init__()
+    super().__init__()
 
     # parameters
     self.tau = 100