Multiple functionalinaty improvements (#197)

Functionalinaty improvements

Author: chaoming0625

brainpy/dyn/__init__.py

@@ -12,4 +12,4 @@
 from .utils import *
 from .runners import *
 
-from . import neurons, synapses, channels, utils, runners
+from . import neurons, synapses, channels, rates, utils, runners

brainpy/dyn/base.py

@@ -502,14 +502,17 @@ def __init__(self,
     # size
     if isinstance(size, (list, tuple)):
       if len(size) <= 0:
-        raise ModelBuildError('size must be int, or a tuple/list of int.')
+        raise ModelBuildError(f'size must be int, or a tuple/list of int. '
+                              f'But we got {type(size)}')
       if not isinstance(size[0], int):
-        raise ModelBuildError('size must be int, or a tuple/list of int.')
+        raise ModelBuildError('size must be int, or a tuple/list of int.'
+                              f'But we got {type(size)}')
       size = tuple(size)
     elif isinstance(size, int):
       size = (size,)
     else:
-      raise ModelBuildError('size must be int, or a tuple/list of int.')
+      raise ModelBuildError('size must be int, or a tuple/list of int.'
+                            f'But we got {type(size)}')
     self.size = size
     self.num = tools.size2num(size)
 

brainpy/dyn/neurons/__init__.py

@@ -3,6 +3,4 @@
 from .biological_models import *
 from .fractional_models import *
 from .input_models import *
-from .noise_models import *
-from .rate_models import *
 from .reduced_models import *

brainpy/dyn/neurons/reduced_models.py

@@ -5,8 +5,7 @@
 import brainpy.math as bm
 from brainpy.dyn.base import NeuGroup
 from brainpy.initialize import ZeroInit, OneInit, Initializer, init_param
-from brainpy.integrators.joint_eq import JointEq
-from brainpy.integrators.ode import odeint
+from brainpy.integrators import sdeint, odeint, JointEq
 from brainpy.tools.checking import check_initializer
 from brainpy.types import Shape, Tensor
 
@@ -46,31 +45,34 @@ class LIF(NeuGroup):
 
   - `(Brette, Romain. 2004) LIF phase locking <https://brainpy-examples.readthedocs.io/en/latest/neurons/Romain_2004_LIF_phase_locking.html>`_
 
-  **Model Parameters**
-
-  ============= ============== ======== =========================================
-  **Parameter** **Init Value** **Unit** **Explanation**
-  ------------- -------------- -------- -----------------------------------------
-  V_rest         0              mV       Resting membrane potential.
-  V_reset        -5             mV       Reset potential after spike.
-  V_th           20             mV       Threshold potential of spike.
-  tau            10             ms       Membrane time constant. Compute by R * C.
-  tau_ref       5              ms       Refractory period length.(ms)
-  ============= ============== ======== =========================================
 
-  **Neuron Variables**
-
-  ================== ================= =========================================================
-  **Variables name** **Initial Value** **Explanation**
-  ------------------ ----------------- ---------------------------------------------------------
-  V                    0                Membrane potential.
-  input                0                External and synaptic input current.
-  spike                False             Flag to mark whether the neuron is spiking.
-  refractory           False             Flag to mark whether the neuron is in refractory period.
-  t_last_spike         -1e7              Last spike time stamp.
-  ================== ================= =========================================================
-
-  **References**
+  Parameters
+  ----------
+  size: sequence of int, int
+    The size of the neuron group.
+  V_rest: float, JaxArray, ndarray, Initializer, callable
+    Resting membrane potential.
+  V_reset: float, JaxArray, ndarray, Initializer, callable
+    Reset potential after spike.
+  V_th: float, JaxArray, ndarray, Initializer, callable
+    Threshold potential of spike.
+  tau: float, JaxArray, ndarray, Initializer, callable
+    Membrane time constant.
+  tau_ref: float, JaxArray, ndarray, Initializer, callable
+    Refractory period length.(ms)
+  V_initializer: JaxArray, ndarray, Initializer, callable
+    The initializer of membrane potential.
+  noise: JaxArray, ndarray, Initializer, callable
+    The noise added onto the membrane potential
+  noise_type: str
+    The type of the provided noise. Can be `value` or `func`.
+  method: str
+    The numerical integration method.
+  name: str
+    The group name.
+
+  References
+  ----------
 
   .. [1] Abbott, Larry F. "Lapicque’s introduction of the integrate-and-fire model
          neuron (1907)." Brain research bulletin 50, no. 5-6 (1999): 303-304.
@@ -85,44 +87,57 @@ def __init__(
       tau: Union[float, Tensor, Initializer, Callable] = 10.,
       tau_ref: Union[float, Tensor, Initializer, Callable] = 1.,
       V_initializer: Union[Initializer, Callable, Tensor] = ZeroInit(),
+      noise: Union[float, Tensor, Initializer, Callable] = None,
+      noise_type: str = 'value',
+      keep_size: bool=False,
       method: str = 'exp_auto',
       name: str = None
   ):
     # initialization
     super(LIF, self).__init__(size=size, name=name)
 
     # parameters
-    self.V_rest = init_param(V_rest, self.num, allow_none=False)
-    self.V_reset = init_param(V_reset, self.num, allow_none=False)
-    self.V_th = init_param(V_th, self.num, allow_none=False)
-    self.tau = init_param(tau, self.num, allow_none=False)
-    self.tau_ref = init_param(tau_ref, self.num, allow_none=False)
+    self.keep_size = keep_size
+    self.noise_type = noise_type
+    if noise_type not in ['func', 'value']:
+      raise ValueError(f'noise_type only supports `func` and `value`, but we got {noise_type}')
+    size = self.size if keep_size else self.num
+    self.V_rest = init_param(V_rest, size, allow_none=False)
+    self.V_reset = init_param(V_reset, size, allow_none=False)
+    self.V_th = init_param(V_th, size, allow_none=False)
+    self.tau = init_param(tau, size, allow_none=False)
+    self.tau_ref = init_param(tau_ref, size, allow_none=False)
+    if noise_type == 'func':
+      self.noise = noise
+    else:
+      self.noise = init_param(noise, size, allow_none=True)
 
     # initializers
     check_initializer(V_initializer, 'V_initializer')
     self._V_initializer = V_initializer
 
     # variables
-    self.V = bm.Variable(init_param(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)
-    self.refractory = bm.Variable(bm.zeros(self.num, dtype=bool))
+    self.V = bm.Variable(init_param(V_initializer, size))
+    self.input = bm.Variable(bm.zeros(size))
+    self.spike = bm.Variable(bm.zeros(size, dtype=bool))
+    self.t_last_spike = bm.Variable(bm.ones(size) * -1e7)
+    self.refractory = bm.Variable(bm.zeros(size, dtype=bool))
 
     # integral
-    self.integral = odeint(method=method, f=self.derivative)
+    f = lambda V, t, I_ext: (-V + self.V_rest + I_ext) / self.tau
+    if self.noise is not None:
+      g = noise if (noise_type == 'func') else (lambda V, t, I_ext: self.noise / bm.sqrt(self.tau))
+      self.integral = sdeint(method=method, f=f, g=g)
+    else:
+      self.integral = odeint(method=method, f=f)
 
   def reset(self):
-    self.V.value = init_param(self._V_initializer, (self.num,))
+    self.V.value = init_param(self._V_initializer, self.size if self.keep_size else self.num)
     self.input[:] = 0
     self.spike[:] = False
     self.t_last_spike[:] = -1e7
     self.refractory[:] = False
 
-  def derivative(self, V, t, I_ext):
-    dvdt = (-V + self.V_rest + I_ext) / self.tau
-    return dvdt
-
   def update(self, t, dt):
     refractory = (t - self.t_last_spike) <= self.tau_ref
     V = self.integral(self.V, t, self.input, dt=dt)

brainpy/dyn/rates/__init__.py

@@ -0,0 +1,6 @@
+# -*- coding: utf-8 -*-
+
+from .noises import *
+from .populations import *
+from .couplings import *
+

brainpy/dyn/synapses/delay_coupling.py renamed to brainpy/dyn/rates/couplings.py

@@ -25,7 +25,7 @@ class DelayCoupling(DynamicalSystem):
   ----------
   delay_var: Variable
     The delay variable.
-  output_var: Variable, sequence of Variable
+  target_var: Variable, sequence of Variable
     The target variables to output.
   conn_mat: JaxArray, ndarray
     The connection matrix.
@@ -40,7 +40,7 @@ class DelayCoupling(DynamicalSystem):
   def __init__(
       self,
       delay_var: bm.Variable,
-      output_var: Union[bm.Variable, Sequence[bm.Variable]],
+      target_var: Union[bm.Variable, Sequence[bm.Variable]],
       conn_mat: Tensor,
       required_shape: Tuple[int, ...],
       delay_steps: Optional[Union[int, Tensor, Initializer, Callable]] = None,
@@ -56,10 +56,10 @@ def __init__(
     self.delay_var = delay_var
 
     # output variables
-    if isinstance(output_var, bm.Variable):
-      output_var = [output_var]
-    check_sequence(output_var, 'output_var', elem_type=bm.Variable, allow_none=False)
-    self.output_var = output_var
+    if isinstance(target_var, bm.Variable):
+      target_var = [target_var]
+    check_sequence(target_var, 'output_var', elem_type=bm.Variable, allow_none=False)
+    self.output_var = target_var
 
     # Connection matrix
     self.conn_mat = bm.asarray(conn_mat)
@@ -122,8 +122,9 @@ class DiffusiveCoupling(DelayCoupling):
   --------
 
   >>> import brainpy as bp
-  >>> areas = bp.dyn.RateFHN(80, x_ou_sigma=0.01, y_ou_sigma=0.01, name='fhn')
-  >>> conn = bp.dyn.DiffusiveCoupling(areas.x, areas.x, areas.input,
+  >>> from brainpy.dyn import rates
+  >>> areas = rates.FHN(80, x_ou_sigma=0.01, y_ou_sigma=0.01, name='fhn')
+  >>> conn = rates.DiffusiveCoupling(areas.x, areas.x, areas.input,
   >>>                                 conn_mat=Cmat, delay_steps=Dmat,
   >>>                                 initial_delay_data=bp.init.Uniform(0, 0.05))
   >>> net = bp.dyn.Network(areas, conn)
@@ -134,7 +135,7 @@ class DiffusiveCoupling(DelayCoupling):
     The first coupling variable, used for delay.
   coupling_var2: Variable
     Another coupling variable.
-  output_var: Variable, sequence of Variable
+  target_var: Variable, sequence of Variable
     The target variables to output.
   conn_mat: JaxArray, ndarray
     The connection matrix.
@@ -150,7 +151,7 @@ def __init__(
       self,
       coupling_var1: bm.Variable,
       coupling_var2: bm.Variable,
-      output_var: Union[bm.Variable, Sequence[bm.Variable]],
+      target_var: Union[bm.Variable, Sequence[bm.Variable]],
       conn_mat: Tensor,
       delay_steps: Optional[Union[int, Tensor, Initializer, Callable]] = None,
       initial_delay_data: Union[Initializer, Callable, Tensor, float, int, bool] = None,
@@ -171,7 +172,7 @@ def __init__(
 
     super(DiffusiveCoupling, self).__init__(
       delay_var=coupling_var1,
-      output_var=output_var,
+      target_var=target_var,
       conn_mat=conn_mat,
       required_shape=(coupling_var1.size, coupling_var2.size),
       delay_steps=delay_steps,
@@ -190,15 +191,18 @@ def update(self, t, dt):
     # delays
     if self.delay_type == 'none':
       diffusive = bm.expand_dims(self.coupling_var1, axis=1) - self.coupling_var2
+      diffusive = (self.conn_mat * diffusive).sum(axis=0)
     elif self.delay_type == 'array':
       f = vmap(lambda i: delay_var(self.delay_steps[i], bm.arange(self.coupling_var1.size)))  # (pre.num,)
       delays = f(bm.arange(self.coupling_var2.size).value)
       diffusive = delays.T - self.coupling_var2  # (post.num, pre.num)
+      diffusive = (self.conn_mat * diffusive).sum(axis=0)
     elif self.delay_type == 'int':
-      diffusive = bm.expand_dims(delay_var(self.delay_steps), axis=1) - self.coupling_var2
+      delayed_var = delay_var(self.delay_steps)
+      diffusive = bm.expand_dims(delayed_var, axis=1) - self.coupling_var2
+      diffusive = (self.conn_mat * diffusive).sum(axis=0)
     else:
       raise ValueError
-    diffusive = (self.conn_mat * diffusive).sum(axis=0)
 
     # output to target variable
     for target in self.output_var:
@@ -221,7 +225,7 @@ class AdditiveCoupling(DelayCoupling):
   ----------
   coupling_var: Variable
     The coupling variable, used for delay.
-  output_var: Variable, sequence of Variable
+  target_var: Variable, sequence of Variable
     The target variables to output.
   conn_mat: JaxArray, ndarray
     The connection matrix.
@@ -236,7 +240,7 @@ class AdditiveCoupling(DelayCoupling):
   def __init__(
       self,
       coupling_var: bm.Variable,
-      output_var: Union[bm.Variable, Sequence[bm.Variable]],
+      target_var: Union[bm.Variable, Sequence[bm.Variable]],
       conn_mat: Tensor,
       delay_steps: Optional[Union[int, Tensor, Initializer, Callable]] = None,
       initial_delay_data: Union[Initializer, Callable, Tensor, float, int, bool] = None,
@@ -251,7 +255,7 @@ def __init__(
 
     super(AdditiveCoupling, self).__init__(
       delay_var=coupling_var,
-      output_var=output_var,
+      target_var=target_var,
       conn_mat=conn_mat,
       required_shape=(coupling_var.size, coupling_var.size),
       delay_steps=delay_steps,
@@ -267,7 +271,7 @@ def update(self, t, dt):
 
     # delay function
     if self.delay_steps is None:
-      additive = self.conn_mat * bm.expand_dims(self.coupling_var, axis=1)
+      additive = self.coupling_var @ self.conn_mat
     else:
       f = vmap(lambda i: delay_var(self.delay_steps[i], bm.arange(self.coupling_var.size)))  # (pre.num,)
       delays = f(bm.arange(self.coupling_var.size).value)  # (post.num, pre.num)