Support reset function in neuron and synapse models (#181)

Support reset functions in neuron and synapse models

Author: chaoming0625

brainpy/dyn/base.py

@@ -86,7 +86,12 @@ def update(self, _t, _dt):
     Assume any dynamical system depends on the time variable ``t`` and
     the time step ``dt``.
     """
-    raise NotImplementedError('Must implement "update" function by user self.')
+    raise NotImplementedError('Must implement "update" function by subclass self.')
+
+  def reset(self):
+    """Reset function which reset the whole variables in the model.
+    """
+    raise NotImplementedError('Must implement "reset" function by subclass self.')
 
 
 class Container(DynamicalSystem):
@@ -143,6 +148,17 @@ def __getattr__(self, item):
     else:
       return super(Container, self).__getattribute__(item)
 
+  def reset(self):
+    nodes = self.nodes(level=1, include_self=False).subset(DynamicalSystem).unique()
+    neuron_groups = nodes.subset(NeuGroup)
+    synapse_groups = nodes.subset(TwoEndConn)
+    for node in neuron_groups.values():
+      node.reset()
+    for node in synapse_groups.values():
+      node.reset()
+    for node in (nodes - neuron_groups - synapse_groups).values():
+      node.reset()
+
 
 class Network(Container):
   """Base class to model network objects, an alias of Container.
@@ -196,6 +212,7 @@ class can automatically support your batched data.
   def __init__(self, size, delay, dtype=None, dt=None, **kwargs):
     # dt
     self.dt = bm.get_dt() if dt is None else dt
+    self.dtype = dtype
 
     # data size
     if isinstance(size, int): size = (size,)
@@ -213,6 +230,7 @@ def __init__(self, size, delay, dtype=None, dt=None, **kwargs):
       self.out_idx = bm.Variable(bm.array([0], dtype=bm.uint32))
       self.in_idx = bm.Variable(bm.array([self.num_step - 1], dtype=bm.uint32))
       self.data = bm.Variable(bm.zeros((self.num_step,) + self.size, dtype=dtype))
+      self.num = 1
 
     else:  # non-uniform delay
       self.uniform_delay = False
@@ -237,6 +255,12 @@ def __init__(self, size, delay, dtype=None, dt=None, **kwargs):
 
     super(ConstantDelay, self).__init__(**kwargs)
 
+  def reset(self):
+    """Reset the variables."""
+    self.in_idx[:] = self.num_step - 1
+    self.out_idx[:] = 0
+    self.data[:] = 0
+
   @property
   def oldest(self):
     return self.pull()
@@ -265,12 +289,6 @@ def update(self, _t=None, _dt=None, **kwargs):
     self.in_idx[:] = (self.in_idx + 1) % self.num_step
     self.out_idx[:] = (self.out_idx + 1) % self.num_step
 
-  def reset(self):
-    """Reset the variables."""
-    self.in_idx[:] = self.num_step - 1
-    self.out_idx[:] = 0
-    self.data[:] = 0
-
 
 class NeuGroup(DynamicalSystem):
   """Base class to model neuronal groups.
@@ -476,15 +494,15 @@ def register_delay(
         self.local_delay_vars[name] = self.global_delay_vars[name]
       else:
         if self.global_delay_vars[name].num_delay_step - 1 < max_delay_step:
-          self.global_delay_vars[name].init(delay_target, max_delay_step, initial_delay_data)
+          self.global_delay_vars[name].reset(delay_target, max_delay_step, initial_delay_data)
       self.register_implicit_nodes(self.global_delay_vars)
     return delay_step
 
-  def get_delay(
+  def get_delay_data(
       self,
       name: str,
-      delay_step: Union[int, bm.JaxArray, bm.ndarray],
-      indices=None,
+      delay_step: Union[int, bm.JaxArray, jnp.DeviceArray],
+      indices: Union[int, bm.JaxArray, jnp.DeviceArray] = None,
   ):
     """Get delay data according to the provided delay steps.
 
@@ -494,7 +512,7 @@ def get_delay(
       The delay variable name.
     delay_step: int, JaxArray, ndarray
       The delay length.
-    indices: optional, JaxArray, ndarray
+    indices: optional, int, JaxArray, ndarray
       The indices of the delay.
 
     Returns
@@ -522,10 +540,31 @@ def get_delay(
   def update_delay(
       self,
       name: str,
-      delay_target: Union[int, bm.JaxArray, bm.ndarray]
+      delay_data: Union[float, bm.JaxArray, jnp.ndarray]
+  ):
+    """Update the delay according to the delay data.
+
+    Parameters
+    ----------
+    name: str
+      The name of the delay.
+    delay_data: float, JaxArray, ndarray
+      The delay data to update at the current time.
+    """
+    if name in self.local_delay_vars:
+      return self.local_delay_vars[name].update(delay_data)
+    else:
+      if name not in self.global_delay_vars:
+        raise ValueError(f'{name} is not defined in delay variables.')
+
+  def reset_delay(
+      self,
+      name: str,
+      delay_target: Union[bm.JaxArray, jnp.DeviceArray]
   ):
+    """Reset the delay variable."""
     if name in self.local_delay_vars:
-      return self.local_delay_vars[name].update(delay_target)
+      return self.local_delay_vars[name].reset(delay_target)
     else:
       if name not in self.global_delay_vars:
         raise ValueError(f'{name} is not defined in delay variables.')

brainpy/dyn/neurons/biological_models.py

@@ -142,37 +142,41 @@ class HH(NeuGroup):
     >>> plt.yticks([])
     >>> plt.show()
 
-
-  **Model Parameters**
-
-  ============= ============== ======== ====================================
-  **Parameter** **Init Value** **Unit** **Explanation**
-  ------------- -------------- -------- ------------------------------------
-  V_th          20.            mV       the spike threshold.
-  C             1.             ufarad   capacitance.
-  E_Na          50.            mV       reversal potential of sodium.
-  E_K           -77.           mV       reversal potential of potassium.
-  E_leak        54.387         mV       reversal potential of unspecific.
-  g_Na          120.           msiemens conductance of sodium channel.
-  g_K           36.            msiemens conductance of potassium channel.
-  g_leak        .03            msiemens conductance of unspecific channels.
-  ============= ============== ======== ====================================
-
-  **Model Variables**
-
-  ================== ================= =========================================================
-  **Variables name** **Initial Value** **Explanation**
-  ------------------ ----------------- ---------------------------------------------------------
-  V                        -65         Membrane potential.
-  m                        0.05        gating variable of the sodium ion channel.
-  n                        0.32        gating variable of the potassium ion channel.
-  h                        0.60        gating variable of the sodium ion channel.
-  input                     0          External and synaptic input current.
-  spike                    False       Flag to mark whether the neuron is spiking.
-  t_last_spike       -1e7               Last spike time stamp.
-  ================== ================= =========================================================
-
-  **References**
+  Parameters
+  ----------
+  size: sequence of int, int
+    The size of the neuron group.
+  ENa: float, JaxArray, ndarray, Initializer, callable
+    The reversal potential of sodium. Default is 50 mV.
+  gNa: float, JaxArray, ndarray, Initializer, callable
+    The maximum conductance of sodium channel. Default is 120 msiemens.
+  EK: float, JaxArray, ndarray, Initializer, callable
+    The reversal potential of potassium. Default is -77 mV.
+  gK: float, JaxArray, ndarray, Initializer, callable
+    The maximum conductance of potassium channel. Default is 36 msiemens.
+  EL: float, JaxArray, ndarray, Initializer, callable
+    The reversal potential of learky channel. Default is -54.387 mV.
+  gL: float, JaxArray, ndarray, Initializer, callable
+    The conductance of learky channel. Default is 0.03 msiemens.
+  V_th: float, JaxArray, ndarray, Initializer, callable
+    The threshold of the membrane spike. Default is 20 mV.
+  C: float, JaxArray, ndarray, Initializer, callable
+    The membrane capacitance. Default is 1 ufarad.
+  V_initializer: JaxArray, ndarray, Initializer, callable
+    The initializer of membrane potential.
+  m_initializer: JaxArray, ndarray, Initializer, callable
+    The initializer of m channel.
+  h_initializer: JaxArray, ndarray, Initializer, callable
+    The initializer of h channel.
+  n_initializer: JaxArray, ndarray, Initializer, callable
+    The initializer of n channel.
+  method: str
+    The numerical integration method.
+  name: str
+    The group name.
+
+  References
+  ----------
 
   .. [1] Hodgkin, Alan L., and Andrew F. Huxley. "A quantitative description
          of membrane current and its application to conduction and excitation
@@ -214,22 +218,37 @@ def __init__(
     self.C = init_param(C, self.num, allow_none=False)
     self.V_th = init_param(V_th, self.num, allow_none=False)
 
-    # variables
+    # initializers
     check_initializer(m_initializer, 'm_initializer', allow_none=False)
     check_initializer(h_initializer, 'h_initializer', allow_none=False)
     check_initializer(n_initializer, 'n_initializer', allow_none=False)
     check_initializer(V_initializer, 'V_initializer', allow_none=False)
-    self.m = bm.Variable(init_param(m_initializer, (self.num,)))
-    self.h = bm.Variable(init_param(h_initializer, (self.num,)))
-    self.n = bm.Variable(init_param(n_initializer, (self.num,)))
-    self.V = bm.Variable(init_param(V_initializer, (self.num,)))
+    self._m_initializer = m_initializer
+    self._h_initializer = h_initializer
+    self._n_initializer = n_initializer
+    self._V_initializer = V_initializer
+
+    # variables
+    self.m = bm.Variable(init_param(self._m_initializer, (self.num,)))
+    self.h = bm.Variable(init_param(self._h_initializer, (self.num,)))
+    self.n = bm.Variable(init_param(self._n_initializer, (self.num,)))
+    self.V = bm.Variable(init_param(self._V_initializer, (self.num,)))
     self.input = bm.Variable(bm.zeros(self.num))
     self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
     self.t_last_spike = bm.Variable(bm.ones(self.num) * -1e7)
 
     # integral
     self.integral = odeint(method=method, f=self.derivative)
 
+  def reset(self):
+    self.m.value = init_param(self._m_initializer, (self.num,))
+    self.h.value = init_param(self._h_initializer, (self.num,))
+    self.n.value = init_param(self._n_initializer, (self.num,))
+    self.V.value = init_param(self._V_initializer, (self.num,))
+    self.input[:] = 0
+    self.spike[:] = False
+    self.t_last_spike[:] = -1e7
+
   def dm(self, m, t, V):
     alpha = 0.1 * (V + 40) / (1 - bm.exp(-(V + 40) / 10))
     beta = 4.0 * bm.exp(-(V + 65) / 18)
@@ -336,20 +355,8 @@ class MorrisLecar(NeuGroup):
   V_th          10             mV       The spike threshold.
   ============= ============== ======== =======================================================
 
-  **Model Variables**
-
-  ================== ================= =========================================================
-  **Variables name** **Initial Value** **Explanation**
-  ------------------ ----------------- ---------------------------------------------------------
-  V                  -20               Membrane potential.
-  W                  0.02              Gating variable, refers to the fraction of
-                                       opened K+ channels.
-  input              0                 External and synaptic input current.
-  spike              False             Flag to mark whether the neuron is spiking.
-  t_last_spike       -1e7              Last spike time stamp.
-  ================== ================= =========================================================
-
-  **References**
+  References
+  ----------
 
   .. [1] Meier, Stephen R., Jarrett L. Lancaster, and Joseph M. Starobin.
          "Bursting regimes in a reaction-diffusion system with action
@@ -398,9 +405,13 @@ def __init__(
     self.phi = init_param(phi, self.num, allow_none=False)
     self.V_th = init_param(V_th, self.num, allow_none=False)
 
-    # vars
+    # initializers
     check_initializer(V_initializer, 'V_initializer', allow_none=False)
     check_initializer(W_initializer, 'W_initializer', allow_none=False)
+    self._W_initializer = W_initializer
+    self._V_initializer = V_initializer
+
+    # variables
     self.W = bm.Variable(init_param(W_initializer, (self.num,)))
     self.V = bm.Variable(init_param(V_initializer, (self.num,)))
     self.input = bm.Variable(bm.zeros(self.num))
@@ -410,6 +421,13 @@ def __init__(
     # integral
     self.integral = odeint(method=method, f=self.derivative)
 
+  def reset(self):
+    self.W.value = init_param(self._W_initializer, (self.num,))
+    self.V.value = init_param(self._V_initializer, (self.num,))
+    self.input.value = bm.zeros(self.num)
+    self.spike.value = bm.zeros(self.num, dtype=bool)
+    self.t_last_spike.value = bm.ones(self.num) * -1e7
+
   def dV(self, V, t, W, I_ext):
     M_inf = (1 / 2) * (1 + bm.tanh((V - self.V1) / self.V2))
     I_Ca = self.g_Ca * M_inf * (V - self.V_Ca)

brainpy/dyn/neurons/fractional_models.py

@@ -110,10 +110,15 @@ def __init__(
     self.Vth = init_param(Vth, self.num, allow_none=False)
     self.delta = init_param(delta, self.num, allow_none=False)
 
-    # variables
+    # initializers
     check_initializer(V_initializer, 'V_initializer', allow_none=False)
     check_initializer(w_initializer, 'w_initializer', allow_none=False)
     check_initializer(y_initializer, 'y_initializer', allow_none=False)
+    self._V_initializer = V_initializer
+    self._w_initializer = w_initializer
+    self._y_initializer = y_initializer
+
+    # variables
     self.V = bm.Variable(init_param(V_initializer, (self.num,)))
     self.w = bm.Variable(init_param(w_initializer, (self.num,)))
     self.y = bm.Variable(init_param(y_initializer, (self.num,)))
@@ -127,6 +132,16 @@ def __init__(
                                   num_memory=num_memory,
                                   inits=[self.V, self.w, self.y])
 
+  def reset(self):
+    self.V.value = init_param(self._V_initializer, (self.num,))
+    self.w.value = init_param(self._w_initializer, (self.num,))
+    self.y.value = init_param(self._y_initializer, (self.num,))
+    self.input[:] = 0
+    self.spike[:] = False
+    self.t_last_spike[:] = -1e7
+    # integral function reset
+    self.integral.reset([self.V, self.w, self.y])
+
   def dV(self, V, t, w, y):
     return V - V ** 3 / 3 - w + y + self.input
 
@@ -149,14 +164,6 @@ def update(self, _t, _dt):
     self.y.value = y
     self.input[:] = 0.
 
-  def set_init(self, values: dict):
-    for k, v in values.items():
-      if k not in self.integral.inits:
-        raise ValueError(f'Variable "{k}" is not defined in this model.')
-      variable = getattr(self, k)
-      variable[:] = v
-      self.integral.inits[k][:] = v
-
 
 class FractionalIzhikevich(FractionalNeuron):
   r"""Fractional-order Izhikevich model [10]_.
@@ -248,9 +255,13 @@ def __init__(
     self.R = init_param(R, self.num, allow_none=False)
     self.V_th = init_param(V_th, self.num, allow_none=False)
 
-    # variables
+    # initializers
     check_initializer(V_initializer, 'V_initializer', allow_none=False)
     check_initializer(u_initializer, 'u_initializer', allow_none=False)
+    self._V_initializer = V_initializer
+    self._u_initializer = u_initializer
+
+    # variables
     self.V = bm.Variable(init_param(V_initializer, (self.num,)))
     self.u = bm.Variable(init_param(u_initializer, (self.num,)))
     self.input = bm.Variable(bm.zeros(self.num))
@@ -264,6 +275,15 @@ def __init__(
                                    num_step=num_step,
                                    inits=[self.V, self.u])
 
+  def reset(self):
+    self.V.value = init_param(self._V_initializer, (self.num,))
+    self.u.value = init_param(self._u_initializer, (self.num,))
+    self.input[:] = 0
+    self.spike[:] = False
+    self.t_last_spike[:] = -1e7
+    # integral function reset
+    self.integral.reset([self.V, self.u])
+
   def dV(self, V, t, u, I_ext):
     dVdt = self.f * V * V + self.g * V + self.h - u + self.R * I_ext
     return dVdt / self.tau
@@ -284,11 +304,3 @@ def update(self, _t, _dt):
     self.u.value = bm.where(spikes, u + self.d, u)
     self.spike.value = spikes
     self.input[:] = 0.
-
-  def set_init(self, values: dict):
-    for k, v in values.items():
-      if k not in self.integral.inits:
-        raise ValueError(f'Variable "{k}" is not defined in this model.')
-      variable = getattr(self, k)
-      variable[:] = v
-      self.integral.inits[k][:] = v