add update() deprecation warning

Author: chaoming0625

brainpy/_src/analysis/highdim/slow_points.py

@@ -14,8 +14,8 @@
 from brainpy import optim, losses
 from brainpy._src.analysis import utils, base, constants
 from brainpy._src.dynsys import DynamicalSystem
+from brainpy._src.context import share
 from brainpy._src.runners import check_and_format_inputs, _f_ops
-from brainpy._src.tools.dicts import DotDict
 from brainpy.errors import AnalyzerError, UnsupportedError
 from brainpy.types import ArrayType
 
@@ -123,7 +123,7 @@ def __init__(
       f_loss_batch: Callable = None,
       fun_inputs: Callable = None,
   ):
-    super(SlowPointFinder, self).__init__()
+    super().__init__()
 
     # static arguments
     if not isinstance(args, tuple):
@@ -636,11 +636,11 @@ def decompose_eigenvalues(matrices, sort_by='magnitude', do_compute_lefts=False)
                              'L': L})
     return decompositions
 
-  def _step_func_input(self, shared):
+  def _step_func_input(self):
     if self._inputs is None:
       return
     elif callable(self._inputs):
-      self._inputs(shared)
+      self._inputs(share.get_shargs())
     else:
       for ops, values in self._inputs['fixed'].items():
         for var, data in values:
@@ -650,7 +650,7 @@ def _step_func_input(self, shared):
           raise UnsupportedError
       for ops, values in self._inputs['functional'].items():
         for var, data in values:
-          _f_ops(ops, var, data(shared))
+          _f_ops(ops, var, data(share.get_shargs()))
       for ops, values in self._inputs['iterated'].items():
         if len(values) > 0:
           raise UnsupportedError
@@ -732,9 +732,10 @@ def _generate_ds_cell_function(
   ):
     if dt is None: dt = bm.get_dt()
     if t is None: t = 0.
-    shared = DotDict(t=t, dt=dt, i=0)
 
     def f_cell(h: Dict):
+      share.save(t=t, i=0, dt=dt)
+
       # update target variables
       for k, v in self.target_vars.items():
         v.value = (bm.asarray(h[k], dtype=v.dtype)
@@ -747,11 +748,10 @@ def f_cell(h: Dict):
 
       # add inputs
       target.clear_input()
-      self._step_func_input(shared)
+      self._step_func_input()
 
       # call update functions
-      args = (shared,) + self.args
-      target(*args)
+      target(*self.args)
 
       # get new states
       new_h = {k: (v.value if (v.batch_axis is None) else jnp.squeeze(v.value, axis=v.batch_axis))

brainpy/_src/dynsys.py

@@ -3,6 +3,7 @@
 import collections
 import gc
 import inspect
+import warnings
 from typing import Union, Dict, Callable, Sequence, Optional, Any
 
 import numpy as np
@@ -28,6 +29,21 @@
 
 SLICE_VARS = 'slice_vars'
 
+_update_deprecate_msg = '''
+From brainpy>=2.4.3, update() function no longer needs to receive a global shared argument.
+
+Instead of using:
+          
+  def update(self, tdi, *args, **kwagrs):
+     ...
+          
+Please use:
+          
+  def update(self, *args, **kwagrs):
+     t = bp.share['t']
+     ...
+'''
+
 
 def not_pass_shared(func: Callable):
   """Label the update function as the one without passing shared arguments.
@@ -160,13 +176,38 @@ def clear_input(self):
     pass
 
   def step_run(self, i, *args, **kwargs):
+    """The step run function.
+
+    This function can be directly applied to run the dynamical system.
+    Particularly, ``i`` denotes the running index.
+
+    Args:
+      i: The current running index.
+      *args: The arguments of ``update()`` function.
+      **kwargs: The arguments of ``update()`` function.
+
+    Returns:
+      out: The update function returns.
+    """
     global share
     if share is None:
       from brainpy._src.context import share
     share.save(i=i, t=i * bm.dt)
     return self.update(*args, **kwargs)
 
-  jit_step_run = bm.cls_jit(step_run, inline=True)
+  @bm.cls_jit(inline=True)
+  def jit_step_run(self, i, *args, **kwargs):
+    """The jitted step function for running.
+
+    Args:
+      i: The current running index.
+      *args: The arguments of ``update()`` function.
+      **kwargs: The arguments of ``update()`` function.
+
+    Returns:
+      out: The update function returns.
+    """
+    return self.step_run(i, *args, **kwargs)
 
   @property
   def mode(self) -> bm.Mode:
@@ -189,32 +230,35 @@ def _compatible_update(self, *args, **kwargs):
 
     if len(update_args) and update_args[0].name in ['tdi', 'sh', 'sha']:
       if len(args) > 0:
-        if isinstance(args[0], dict):
+        if isinstance(args[0], dict) and all([bm.isscalar(v) for v in args[0].values()]):
           # define:
           #    update(tdi, *args, **kwargs)
           # call:
           #    update(tdi, *args, **kwargs)
           ret = update_fun(*args, **kwargs)
-          # TODO: deprecation
+          warnings.warn(_update_deprecate_msg, UserWarning)
         else:
           # define:
           #    update(tdi, *args, **kwargs)
           # call:
           #    update(*args, **kwargs)
           ret = update_fun(share.get_shargs(), *args, **kwargs)
+          warnings.warn(_update_deprecate_msg, UserWarning)
       else:
         if update_args[0].name in kwargs:
           # define:
           #    update(tdi, *args, **kwargs)
           # call:
           #    update(tdi=??, **kwargs)
           ret = update_fun(**kwargs)
+          warnings.warn(_update_deprecate_msg, UserWarning)
         else:
           # define:
           #    update(tdi, *args, **kwargs)
           # call:
           #    update(**kwargs)
           ret = update_fun(share.get_shargs(), *args, **kwargs)
+          warnings.warn(_update_deprecate_msg, UserWarning)
       return ret
 
     try:
@@ -230,6 +274,7 @@ def _compatible_update(self, *args, **kwargs):
         #    update(*args, **kwargs)
         share.save(**args[0])
         ret = update_fun(*args[1:], **kwargs)
+        warnings.warn(_update_deprecate_msg, UserWarning)
         return ret
       else:
         # user define ``update()`` function which receives the shared argument,
@@ -240,6 +285,7 @@ def _compatible_update(self, *args, **kwargs):
         # as
         #    update(tdi, *args, **kwargs)
         ret = update_fun(share.get_shargs(), *args, **kwargs)
+        warnings.warn(_update_deprecate_msg, UserWarning)
         return ret
     else:
       return update_fun(*args, **kwargs)

brainpy/dyn/__init__.py

@@ -7,4 +7,5 @@
 from .projections import *
 from .others import *
 from .outs import *
+from .rates import *
 from .compat import NeuGroup