feat: efficient checking on numerical values

brainpy/__init__.py

@@ -36,11 +36,9 @@
 from . import integrators
 from .integrators import ode
 from .integrators import sde
-from .integrators import dde
 from .integrators import fde
 from .integrators.ode import odeint
 from .integrators.sde import sdeint
-from .integrators.dde import ddeint
 from .integrators.fde import fdeint
 from .integrators.joint_eq import JointEq
 

brainpy/dyn/rates/populations.py

@@ -2,19 +2,16 @@
 
 from typing import Union, Callable
 
-import numpy as np
-from jax.experimental.host_callback import id_tap
-
 import brainpy.math as bm
 from brainpy import check
 from brainpy.dyn.base import NeuGroup
+from brainpy.dyn.others.noises import OUProcess
 from brainpy.initialize import Initializer, Uniform, init_param, ZeroInit
-from brainpy.integrators.dde import ddeint
 from brainpy.integrators.joint_eq import JointEq
 from brainpy.integrators.ode import odeint
 from brainpy.tools.checking import check_float, check_initializer
+from brainpy.tools.errors import check_error_in_jit
 from brainpy.types import Shape, Tensor
-from brainpy.dyn.others.noises import OUProcess
 
 __all__ = [
   'Population',
@@ -307,7 +304,7 @@ def __init__(
                             method=sde_method)
 
     # integral
-    self.integral = ddeint(method=method,
+    self.integral = odeint(method=method,
                            f=JointEq([self.dx, self.dy]),
                            state_delays={'V': self.x_delay})
 
@@ -327,15 +324,14 @@ def dx(self, x, t, y, x_ext):
   def dy(self, y, t, x, y_ext):
     return (x + self.a - self.b * y + y_ext) / self.tau
 
-  def _check_dt(self, dt, *args):
-    if np.absolute(dt - self.dt) > 1e-6:
-      raise ValueError(f'The "dt" {dt} used in model running is '
-                       f'not consistent with the "dt" {self.dt} '
-                       f'used in model definition.')
+  def _check_dt(self, dt):
+    raise ValueError(f'The "dt" {dt} used in model running is '
+                     f'not consistent with the "dt" {self.dt} '
+                     f'used in model definition.')
 
   def update(self, t, dt):
     if check.is_checking():
-      id_tap(self._check_dt, dt)
+      check_error_in_jit(not bm.isclose(dt, self.dt), self._check_dt, dt)
     if self.x_ou is not None:
       self.input += self.x_ou.x
       self.x_ou.update(t, dt)
@@ -882,5 +878,3 @@ def update(self, t, dt):
     self.i.value = bm.maximum(self.i + di * dt, 0.)
     self.Ie[:] = 0.
     self.Ii[:] = 0.
-
-

brainpy/integrators/__init__.py

@@ -42,10 +42,5 @@
                           set_default_fdeint,
                           register_fde_integrator)
 
-# DDE tools
-from . import dde
-from .dde import ddeint
-
-
 # PDE tools
 from . import pde