rename 'FixedLenDelay' to 'TimeDelay'

Author: chaoming0625

README.md

@@ -147,16 +147,18 @@ runner.run(100.)
 
 
 
-Numerical methods for delay differential equations (SDEs). 
+Numerical methods for delay differential equations (SDEs).
 
 ```python
-xdelay = bm.FixedLenDelay(1, delay_len=1., before_t0=1., dt=0.01)
+xdelay = bm.TimeDelay(1, delay_len=1., before_t0=1., dt=0.01)
+
 
 @bp.ddeint(method='rk4', state_delays={'x': xdelay})
 def second_order_eq(x, y, t):
-    dx = y
-    dy = -y - 2*x - 0.5*xdelay(t-1)
-    return dx, dy
+  dx = y
+  dy = -y - 2 * x - 0.5 * xdelay(t - 1)
+  return dx, dy
+
 
 runner = bp.integrators.IntegratorRunner(second_order_eq, dt=0.01)
 runner.run(100.)

brainpy/datasets/chaotic_systems.py

@@ -167,7 +167,7 @@ def mackey_glass_series(duration, dt=0.1, beta=2., gamma=1., tau=2., n=9.65,
     assert isinstance(inits, (bm.ndarray, jnp.ndarray))
 
   rng = bm.random.RandomState(seed)
-  xdelay = bm.FixedLenDelay(inits.shape, tau, dt=dt)
+  xdelay = bm.TimeDelay(inits.shape, tau, dt=dt)
   xdelay.data = inits + 0.2 * (rng.random((xdelay.num_delay_step,) + inits.shape) - 0.5)
 
   @ddeint(method=method, state_delays={'x': xdelay})

brainpy/dyn/neurons/rate_models.py

@@ -211,7 +211,7 @@ def __init__(self,
     # variables
     self.w = bm.Variable(bm.zeros(self.num))
     self.V = bm.Variable(bm.zeros(self.num))
-    self.Vdelay = bm.FixedLenDelay(self.num, self.delay, interp_method='round')
+    self.Vdelay = bm.TimeDelay(self.num, self.delay, interp_method='round')
     self.input = bm.Variable(bm.zeros(self.num))
 
     # integral

brainpy/integrators/dde/base.py

@@ -25,7 +25,7 @@ def __init__(
       dt: Union[float, int] = None,
       name: str = None,
       show_code: bool = False,
-      state_delays: Dict[str, bm.FixedLenDelay] = None,
+      state_delays: Dict[str, bm.TimeDelay] = None,
       neutral_delays: Dict[str, bm.NeutralDelay] = None,
   ):
     dt = bm.get_dt() if dt is None else dt
@@ -59,7 +59,7 @@ def __init__(
     # delays
     self._state_delays = dict()
     if state_delays is not None:
-      check_dict_data(state_delays, key_type=str, val_type=bm.FixedLenDelay)
+      check_dict_data(state_delays, key_type=str, val_type=bm.TimeDelay)
       for key, delay in state_delays.items():
         if key not in self.variables:
           raise DiffEqError(f'"{key}" is not defined in the variables: {self.variables}')

brainpy/integrators/runner.py

@@ -93,7 +93,7 @@ class IntegratorRunner(Runner):
     >>> dt = 0.01; beta=2.; gamma=1.; tau=2.; n=9.65
     >>> mg_eq = lambda x, t, xdelay: (beta * xdelay(t - tau) / (1 + xdelay(t - tau) ** n)
     >>>                               - gamma * x)
-    >>> xdelay = bm.FixedLenDelay(1, delay_len=tau, dt=dt, before_t0=lambda t: 1.2)
+    >>> xdelay = bm.TimeDelay(1, delay_len=tau, dt=dt, before_t0=lambda t: 1.2)
     >>> integral = bp.ddeint(mg_eq, method='rk4', state_delays={'x': xdelay})
     >>> runner = bp.integrators.IntegratorRunner(
     >>>       integral,

brainpy/math/delay_vars.py

@@ -1,21 +1,24 @@
 # -*- coding: utf-8 -*-
 
-
+import warnings
 from typing import Union, Callable, Tuple
 
 import jax.numpy as jnp
+import numpy as np
 from jax import vmap
 from jax.experimental.host_callback import id_tap
 from jax.lax import cond
 
+from brainpy import check
 from brainpy import math as bm
 from brainpy.base.base import Base
+from brainpy.errors import UnsupportedError
 from brainpy.tools.checking import check_float
 from brainpy.tools.others import to_size
-from brainpy.errors import UnsupportedError
 
 __all__ = [
   'AbstractDelay',
+  'TimeDelay',
   'FixedLenDelay',
   'NeutralDelay',
 ]
@@ -32,35 +35,35 @@ def update(self, time, value):
 _INTERP_ROUND = 'round'
 
 
-class FixedLenDelay(AbstractDelay):
-  """Delay variable which has a fixed delay length.
+class TimeDelay(AbstractDelay):
+  """Delay variable which has a fixed delay time length.
 
   For example, we create a delay variable which has a maximum delay length of 1 ms
 
   >>> import brainpy.math as bm
-  >>> delay = bm.FixedLenDelay(bm.zeros(3), delay_len=1., dt=0.1)
+  >>> delay = bm.TimeDelay(bm.zeros(3), delay_len=1., dt=0.1)
   >>> delay(-0.5)
   [-0. -0. -0.]
 
   This function supports multiple dimensions of the tensor. For example,
 
   1. the one-dimensional delay data
 
-  >>> delay = bm.FixedLenDelay(3, delay_len=1., dt=0.1, before_t0=lambda t: t)
+  >>> delay = bm.TimeDelay(3, delay_len=1., dt=0.1, before_t0=lambda t: t)
   >>> delay(-0.2)
   [-0.2 -0.2 -0.2]
 
   2. the two-dimensional delay data
 
-  >>> delay = bm.FixedLenDelay((3, 2), delay_len=1., dt=0.1, before_t0=lambda t: t)
+  >>> delay = bm.TimeDelay((3, 2), delay_len=1., dt=0.1, before_t0=lambda t: t)
   >>> delay(-0.6)
   [[-0.6 -0.6]
    [-0.6 -0.6]
    [-0.6 -0.6]]
 
   3. the three-dimensional delay data
 
-  >>> delay = bm.FixedLenDelay((3, 2, 1), delay_len=1., dt=0.1, before_t0=lambda t: t)
+  >>> delay = bm.TimeDelay((3, 2, 1), delay_len=1., dt=0.1, before_t0=lambda t: t)
   >>> delay(-0.6)
   [[[-0.8]
     [-0.8]]
@@ -113,7 +116,7 @@ def __init__(
       dtype=None,
       interp_method='linear_interp',
   ):
-    super(FixedLenDelay, self).__init__(name=name)
+    super(TimeDelay, self).__init__(name=name)
 
     # shape
     self.shape = to_size(shape)
@@ -161,6 +164,10 @@ def __init__(
     else:
       raise ValueError(f'"before_t0" does not support {type(before_t0)}: before_t0')
 
+    self.f = jnp.interp
+    for dim in range(1, len(self.shape) + 1, 1):
+      self.f = vmap(self.f, in_axes=(None, None, dim), out_axes=dim - 1)
+
   @property
   def idx(self):
     return self._idx
@@ -191,36 +198,37 @@ def current_time(self):
 
   def _check_time(self, times, transforms):
     prev_time, current_time = times
-    current_time = bm.as_device_array(current_time)
-    prev_time = bm.as_device_array(prev_time)
+    current_time = np.asarray(current_time, dtype=bm.float_)
+    prev_time = np.asarray(prev_time, dtype=bm.float_)
     if prev_time > current_time:
       raise ValueError(f'\n'
                        f'!!! Error in {self.__class__.__name__}: \n'
                        f'The request time should be less than the '
                        f'current time {current_time}. But we '
                        f'got {prev_time} > {current_time}')
-    lower_time = jnp.asarray(current_time - self.delay_len)
+    lower_time = np.asarray(current_time - self.delay_len)
     if prev_time < lower_time:
       raise ValueError(f'\n'
                        f'!!! Error in {self.__class__.__name__}: \n'
                        f'The request time of the variable should be in '
                        f'[{lower_time}, {current_time}], but we got {prev_time}')
 
-  def __call__(self, prev_time):
+  def __call__(self, time, indices=None):
     # check
-    id_tap(self._check_time, (prev_time, self.current_time))
+    if check.is_checking():
+      id_tap(self._check_time, (time, self.current_time))
     if self._before_type == _FUNC_BEFORE:
-      return cond(prev_time < self.t0,
+      return cond(time < self.t0,
                   self._before_t0,
                   self._after_t0,
-                  prev_time)
+                  time)
     else:
-      return self._after_t0(prev_time)
+      return self._after_t0(time)
 
   def _after_t0(self, prev_time):
     diff = self.delay_len - (self.current_time - prev_time)
-    if isinstance(diff, bm.ndarray): diff = diff.value
-
+    if isinstance(diff, bm.ndarray):
+      diff = diff.value
     if self.interp_method == _INTERP_LINEAR:
       req_num_step = jnp.asarray(diff / self._dt, dtype=bm.get_dint())
       extra = diff - req_num_step * self._dt
@@ -238,31 +246,43 @@ def _true_fn(self, div_mod):
 
   def _false_fn(self, div_mod):
     req_num_step, extra = div_mod
-    f = jnp.interp
-    for dim in range(1, len(self.shape) + 1, 1):
-      f = vmap(f, in_axes=(None, None, dim), out_axes=dim - 1)
     idx = jnp.asarray([self.idx[0] + req_num_step,
                        self.idx[0] + req_num_step + 1])
     idx %= self.num_delay_step
-    return f(extra, jnp.asarray([0., self._dt]), self._data[idx])
+    return self.f(extra, jnp.asarray([0., self._dt]), self._data[idx])
 
   def update(self, time, value):
     self._data[self._idx[0]] = value
     self._current_time[0] = time
     self._idx.value = (self._idx + 1) % self.num_delay_step
 
 
-class VariedLenDelay(AbstractDelay):
-  """Delay variable which has a functional delay
-
-  """
+def FixedLenDelay(shape: Union[int, Tuple[int, ...]],
+                  delay_len: Union[float, int],
+                  before_t0: Union[Callable, bm.ndarray, jnp.ndarray, float, int] = None,
+                  t0: Union[float, int] = 0.,
+                  dt: Union[float, int] = None,
+                  name: str = None,
+                  dtype=None,
+                  interp_method='linear_interp', ):
+  warnings.warn('Please use "brainpy.math.TimeDelay" instead. '
+                '"brainpy.math.FixedLenDelay" is deprecated since version 2.1.2. ',
+                DeprecationWarning)
+  return TimeDelay(shape=shape,
+                   delay_len=delay_len,
+                   before_t0=before_t0,
+                   t0=t0,
+                   dt=dt,
+                   name=name,
+                   dtype=dtype,
+                   interp_method=interp_method)
+
+
+class NeutralDelay(TimeDelay):
+  pass
 
-  def update(self, time, value):
-    pass
 
-  def __init__(self):
-    super(VariedLenDelay, self).__init__()
+class LengthDelay(AbstractDelay):
+  pass
 
 
-class NeutralDelay(FixedLenDelay):
-  pass

brainpy/math/tests/test_delay_vars.py

@@ -12,7 +12,7 @@ def test_dim1(self):
     # linear interp
     t0 = 0.
     before_t0 = bm.repeat(bm.arange(11).reshape((-1, 1)), 10, axis=1)
-    delay = bm.FixedLenDelay(10, delay_len=1., t0=t0, dt=0.1, before_t0=before_t0)
+    delay = bm.TimeDelay(10, delay_len=1., t0=t0, dt=0.1, before_t0=before_t0)
     self.assertTrue(bm.array_equal(delay(t0 - 0.1), bm.ones(10) * 10))
     self.assertTrue(bm.array_equal(delay(t0 - 0.15), bm.ones(10) * 9.5))
     print()
@@ -21,8 +21,8 @@ def test_dim1(self):
     # self.assertTrue(bm.array_equal(delay(t0 - 0.23), bm.ones(10) * 8.7))
 
     # round interp
-    delay = bm.FixedLenDelay(10, delay_len=1., t0=t0, dt=0.1, before_t0=before_t0,
-                             interp_method='round')
+    delay = bm.TimeDelay(10, delay_len=1., t0=t0, dt=0.1, before_t0=before_t0,
+                         interp_method='round')
     self.assertTrue(bm.array_equal(delay(t0 - 0.1), bm.ones(10) * 10))
     self.assertTrue(bm.array_equal(delay(t0 - 0.15), bm.ones(10) * 10))
     self.assertTrue(bm.array_equal(delay(t0 - 0.2), bm.ones(10) * 9))
@@ -31,7 +31,7 @@ def test_dim2(self):
     t0 = 0.
     before_t0 = bm.repeat(bm.arange(11).reshape((-1, 1)), 10, axis=1)
     before_t0 = bm.repeat(before_t0.reshape((11, 10, 1)), 5, axis=2)
-    delay = bm.FixedLenDelay((10, 5), delay_len=1., t0=t0, dt=0.1, before_t0=before_t0)
+    delay = bm.TimeDelay((10, 5), delay_len=1., t0=t0, dt=0.1, before_t0=before_t0)
     self.assertTrue(bm.array_equal(delay(t0 - 0.1), bm.ones((10, 5)) * 10))
     self.assertTrue(bm.array_equal(delay(t0 - 0.15), bm.ones((10, 5)) * 9.5))
     # self.assertTrue(bm.array_equal(delay(t0 - 0.23), bm.ones((10, 5)) * 8.7))
@@ -41,27 +41,27 @@ def test_dim3(self):
       before_t0 = bm.repeat(bm.arange(11).reshape((-1, 1)), 10, axis=1)
       before_t0 = bm.repeat(before_t0.reshape((11, 10, 1)), 5, axis=2)
       before_t0 = bm.repeat(before_t0.reshape((11, 10, 5, 1)), 3, axis=3)
-      delay = bm.FixedLenDelay((10, 5, 3), delay_len=1., t0=t0, dt=0.1, before_t0=before_t0)
+      delay = bm.TimeDelay((10, 5, 3), delay_len=1., t0=t0, dt=0.1, before_t0=before_t0)
       self.assertTrue(bm.array_equal(delay(t0 - 0.1), bm.ones((10, 5, 3)) * 10))
       self.assertTrue(bm.array_equal(delay(t0 - 0.15), bm.ones((10, 5, 3)) * 9.5))
       # self.assertTrue(bm.array_equal(delay(t0 - 0.23), bm.ones((10, 5, 3)) * 8.7))
 
   def test1(self):
     print()
-    delay = bm.FixedLenDelay(3, delay_len=1., dt=0.1, before_t0=lambda t: t)
+    delay = bm.TimeDelay(3, delay_len=1., dt=0.1, before_t0=lambda t: t)
     print(delay(-0.2))
-    delay = bm.FixedLenDelay((3, 2), delay_len=1., dt=0.1, before_t0=lambda t: t)
+    delay = bm.TimeDelay((3, 2), delay_len=1., dt=0.1, before_t0=lambda t: t)
     print(delay(-0.6))
-    delay = bm.FixedLenDelay((3, 2, 1), delay_len=1., dt=0.1, before_t0=lambda t: t)
+    delay = bm.TimeDelay((3, 2, 1), delay_len=1., dt=0.1, before_t0=lambda t: t)
     print(delay(-0.8))
 
   def test_current_time2(self):
     print()
-    delay = bm.FixedLenDelay(3, delay_len=1., dt=0.1, before_t0=lambda t: t)
+    delay = bm.TimeDelay(3, delay_len=1., dt=0.1, before_t0=lambda t: t)
     print(delay(0.))
     before_t0 = bm.repeat(bm.arange(11).reshape((-1, 1)), 10, axis=1)
     before_t0 = bm.repeat(before_t0.reshape((11, 10, 1)), 5, axis=2)
-    delay = bm.FixedLenDelay((10, 5), delay_len=1., dt=0.1, before_t0=before_t0)
+    delay = bm.TimeDelay((10, 5), delay_len=1., dt=0.1, before_t0=before_t0)
     print(delay(0.))
 
   # def test_prev_time_beyond_boundary(self):