Merge pull request #108 from PKU-NIP-Lab/V2.1.1

release V2.1.1

Author: chaoming0625

brainpy/dyn/neurons/fractional_models.py

@@ -1,72 +1,196 @@
 # -*- coding: utf-8 -*-
 
+from typing import Union, Sequence
+
 import brainpy.math as bm
 from brainpy.dyn.base import NeuGroup
 from brainpy.integrators.fde import CaputoL1Schema
+from brainpy.integrators.fde import GLShortMemory
 from brainpy.integrators.joint_eq import JointEq
 from brainpy.tools.checking import check_float, check_integer
+from brainpy.types import Parameter, Shape
 
 __all__ = [
-  'FractionalFHN',
+  'FractionalFHR',
   'FractionalIzhikevich',
 ]
 
 
-class FractionalFHN(NeuGroup):
-  """
+class FractionalNeuron(NeuGroup):
+  """Fractional-order neuron model."""
+  pass
+
+
+class FractionalFHR(FractionalNeuron):
+  r"""The fractional-order FH-R model [1]_.
+
+  FitzHugh and Rinzel introduced FH-R model (1976, in an unpublished article),
+  which is the modification of the classical FHN neuron model. The fractional-order
+  FH-R model is described as
+
+  .. math::
+
+     \begin{array}{rcl}
+     \frac{{d}^{\alpha }v}{d{t}^{\alpha }} & = & v-{v}^{3}/3-w+y+I={f}_{1}(v,w,y),\\
+     \frac{{d}^{\alpha }w}{d{t}^{\alpha }} & = & \delta (a+v-bw)={f}_{2}(v,w,y),\\
+     \frac{{d}^{\alpha }y}{d{t}^{\alpha }} & = & \mu (c-v-dy)={f}_{3}(v,w,y),
+     \end{array}
+
+  where :math:`v, w` and :math:`y` represent the membrane voltage, recovery variable
+  and slow modulation of the current respectively.
+  :math:`I` measures the constant magnitude of external stimulus current, and :math:`\alpha`
+  is the fractional exponent which ranges in the interval :math:`(0 < \alpha \le 1)`.
+  :math:`a, b, c, d, \delta` and :math:`\mu` are the system parameters.
+
+  The system reduces to the original classical order system when :math:`\alpha=1`.
 
+  :math:`\mu` indicates a small parameter that determines the pace of the slow system
+  variable :math:`y`. The fast subsystem (:math:`v-w`) presents a relaxation oscillator
+  in the phase plane where :math:`\delta`  is a small parameter.
+  :math:`v` is expressed in mV (millivolt) scale. Time :math:`t` is in ms (millisecond) scale.
+  It exhibits tonic spiking or quiescent state depending on the parameter sets for a fixed
+  value of :math:`I`. The parameter :math:`a` in the 2D FHN model corresponds to the
+  parameter :math:`c` of the FH-R neuron model. If we decrease the value of :math:`a`,
+  it causes longer intervals between two burstings, however there exists :math:`a`
+  relatively fixed time of bursting duration. With the increasing of :math:`a`, the
+  interburst intervals become shorter and periodic bursting changes to tonic spiking.
+
+  Examples
+  --------
+
+  - [(Mondal, et, al., 2019): Fractional-order FitzHugh-Rinzel bursting neuron model](https://brainpy-examples.readthedocs.io/en/latest/neurons/2019_Fractional_order_FHR_model.html)
+
+
+  Parameters
+  ----------
+  size: int, sequence of int
+    The size of the neuron group.
+  alpha: float, tensor
+    The fractional order.
+  num_memory: int
+    The total number of the short memory.
 
   References
   ----------
   .. [1] Mondal, A., Sharma, S.K., Upadhyay, R.K. *et al.* Firing activities of a fractional-order FitzHugh-Rinzel bursting neuron model and its coupled dynamics. *Sci Rep* **9,** 15721 (2019). https://doi.org/10.1038/s41598-019-52061-4
   """
-  def __init__(self, size, alpha, num_step,
-               a=0.7, b=0.8, c=-0.775, d=1., delta=0.08, mu=0.0001):
-    super(FractionalFHN, self).__init__(size)
 
+  def __init__(self,
+               size: Shape,
+               alpha: Union[float, Sequence[float]],
+               num_memory: int = 1000,
+               a: Parameter = 0.7,
+               b: Parameter = 0.8,
+               c: Parameter = -0.775,
+               d: Parameter = 1.,
+               delta: Parameter = 0.08,
+               mu: Parameter = 0.0001,
+               Vth: Parameter = 1.8,
+               name: str = None):
+    super(FractionalFHR, self).__init__(size, name=name)
+
+    # fractional order
     self.alpha = alpha
-    self.num_step = num_step
+    check_integer(num_memory, 'num_memory', allow_none=False)
 
+    # parameters
     self.a = a
     self.b = b
     self.c = c
     self.d = d
     self.delta = delta
     self.mu = mu
+    self.Vth = Vth
 
+    # variables
     self.input = bm.Variable(bm.zeros(self.num))
-    self.v = bm.Variable(bm.ones(self.num) * 2.5)
+    self.V = bm.Variable(bm.ones(self.num) * 2.5)
     self.w = bm.Variable(bm.zeros(self.num))
     self.y = 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.integral = CaputoL1Schema(self.derivative,
-                                   alpha=alpha,
-                                   num_step=num_step,
-                                   inits=[self.v, self.w, self.y])
+    # integral function
+    self.integral = GLShortMemory(self.derivative,
+                                  alpha=alpha,
+                                  num_memory=num_memory,
+                                  inits=[self.V, self.w, self.y])
 
-  def dv(self, v, t, w, y):
-    return v - v ** 3 / 3 - w + y + self.input
+  def dV(self, V, t, w, y):
+    return V - V ** 3 / 3 - w + y + self.input
 
-  def dw(self, w, t, v):
-    return self.delta * (self.a + v - self.b * w)
+  def dw(self, w, t, V):
+    return self.delta * (self.a + V - self.b * w)
 
-  def dy(self, y, t, v):
-    return self.mu * (self.c - v - self.d * y)
+  def dy(self, y, t, V):
+    return self.mu * (self.c - V - self.d * y)
 
   @property
   def derivative(self):
-    return JointEq([self.dv, self.dw, self.dy])
+    return JointEq([self.dV, self.dw, self.dy])
 
   def update(self, _t, _dt):
-    v, w, y = self.integral(self.v, self.w, self.y, _t, _dt)
-    self.v.value = v
+    V, w, y = self.integral(self.V, self.w, self.y, _t, _dt)
+    self.spike.value = bm.logical_and(V >= self.Vth, self.V < self.Vth)
+    self.t_last_spike.value = bm.where(self.spike, _t, self.t_last_spike)
+    self.V.value = V
     self.w.value = w
     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]_.
+
+  The fractional-order Izhikevich model is given by
+
+  .. math::
+
+     \begin{aligned}
+      &\tau \frac{d^{\alpha} v}{d t^{\alpha}}=\mathrm{f} v^{2}+g v+h-u+R I \\
+      &\tau \frac{d^{\alpha} u}{d t^{\alpha}}=a(b v-u)
+      \end{aligned}
+
+  where :math:`\alpha` is the fractional order (exponent) such that :math:`0<\alpha\le1`.
+  It is a commensurate system that reduces to classical Izhikevich model at :math:`\alpha=1`.
 
-class FractionalIzhikevich(NeuGroup):
-  """Fractional-order Izhikevich model [10]_.
+  The time :math:`t` is in ms; and the system variable :math:`v` expressed in mV
+  corresponds to membrane voltage. Moreover, :math:`u` expressed in mV is the
+  recovery variable that corresponds to the activation of K+ ionic current and
+  inactivation of Na+ ionic current.
+
+  The parameters :math:`f, g, h` are fixed constants (should not be changed) such
+  that :math:`f=0.04` (mV)−1, :math:`g=5, h=140` mV; and :math:`a` and :math:`b` are
+  dimensionless parameters. The time constant :math:`\tau=1` ms; the resistance
+  :math:`R=1` Ω; and :math:`I` expressed in mA measures the injected (applied)
+  dc stimulus current to the system.
+
+  When the membrane voltage reaches the spike peak :math:`v_{peak}`, the two variables
+  are rest as follow:
+
+  .. math::
+
+     \text { if } v \geq v_{\text {peak }} \text { then }\left\{\begin{array}{l}
+      v \leftarrow c \\
+      u \leftarrow u+d
+      \end{array}\right.
+
+  we used :math:`v_{peak}=30` mV, and :math:`c` and :math:`d` are parameters expressed
+  in mV. When the spike reaches its peak value, the membrane voltage :math:`v` and the
+  recovery variable :math:`u` are reset according to the above condition.
+
+  Examples
+  --------
+
+  - [(Teka, et. al, 2018): Fractional-order Izhikevich neuron model](https://brainpy-examples.readthedocs.io/en/latest/neurons/2018_Fractional_Izhikevich_model.html)
 
 
   References
@@ -77,9 +201,21 @@ class FractionalIzhikevich(NeuGroup):
 
   """
 
-  def __init__(self, size, num_step, alpha=0.9,
-               a=0.02, b=0.20, c=-65., d=8., f=0.04,
-               g=5., h=140., tau=1., R=1., V_th=30., name=None):
+  def __init__(self,
+               size: Shape,
+               alpha: Union[float, Sequence[float]],
+               num_step: int,
+               a: Parameter = 0.02,
+               b: Parameter = 0.20,
+               c: Parameter = -65.,
+               d: Parameter = 8.,
+               f: Parameter = 0.04,
+               g: Parameter = 5.,
+               h: Parameter = 140.,
+               tau: Parameter = 1.,
+               R: Parameter = 1.,
+               V_th: Parameter = 30.,
+               name: str = None):
     # initialization
     super(FractionalIzhikevich, self).__init__(size=size, name=name)
 
@@ -131,3 +267,11 @@ 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

brainpy/integrators/fde/GL.py

@@ -123,7 +123,7 @@ def __init__(self, f, alpha, num_memory, inits, dt=None, name=None):
     super(GLShortMemory, self).__init__(f=f, alpha=alpha, dt=dt, name=name)
 
     # fractional order
-    if not jnp.all(jnp.logical_and(self.alpha < 1, self.alpha > 0)):
+    if not jnp.all(jnp.logical_and(self.alpha <= 1, self.alpha > 0)):
       raise UnsupportedError(f'Only support the fractional order in (0, 1), '
                              f'but we got {self.alpha}.')
 
@@ -132,11 +132,11 @@ def __init__(self, f, alpha, num_memory, inits, dt=None, name=None):
     self.num_memory = num_memory
 
     # initial values
-    inits = check_inits(inits, self.variables)
+    self.inits = check_inits(inits, self.variables)
 
     # delays
     self.delays = {}
-    for key, val in inits.items():
+    for key, val in self.inits.items():
       delay = bm.Variable(bm.zeros((self.num_memory,) + val.shape, dtype=bm.float_))
       delay[0] = val
       self.delays[key] = delay

changelog.rst

@@ -6,6 +6,28 @@ brainpy 2.x (LTS)
 *****************
 
 
+Version 2.1.1 (2022.03.18)
+==========================
+
+This release continues to update the functionality of BrainPy. Core changes include
+
+- numerical solvers for fractional differential equations
+- more standard ``brainpy.nn`` interfaces
+
+
+New Features
+~~~~~~~~~~~~
+
+- Numerical solvers for fractional differential equations
+    - ``brainpy.fde.CaputoEuler``
+    - ``brainpy.fde.CaputoL1Schema``
+    - ``brainpy.fde.GLShortMemory``
+- Fractional neuron models
+    - ``brainpy.dyn.FractionalFHR``
+    - ``brainpy.dyn.FractionalIzhikevich``
+- support ``shared_kwargs`` in `RNNTrainer` and `RNNRunner`
+
+
 Version 2.1.0 (2022.03.14)
 ==========================
 

docs/index.rst

@@ -7,7 +7,8 @@ high-performance Brain Dynamics Programming (BDP). Among its key ingredients, Br
 - **JIT compilation** and **automatic differentiation** for class objects.
 - **Numerical methods** for ordinary differential equations (ODEs),
   stochastic differential equations (SDEs),
-  delay differential equations (DDEs), etc.
+  delay differential equations (DDEs),
+  fractional differential equations (FDEs), etc.
 - **Dynamics simulation** tools for various brain objects, like
   neurons, synapses, networks, soma, dendrites, channels, and even more.
 - **Dynamics training** tools with various machine learning algorithms,