Merge pull request #422 from brainpy/updates

upgrade Runner and Trainer for new style of ``DynamicalSystem.update()`` function

Author: chaoming0625

.github/workflows/CI-models.yml

@@ -117,7 +117,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11"]
 
     steps:
     - uses: actions/checkout@v2
@@ -128,8 +128,6 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install numpy>=1.21.0
-        python -m pip install "jaxlib==0.4.10" -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver
-        python -m pip install jax==0.4.10
         python -m pip install -r requirements-dev.txt
         python -m pip install tqdm brainpylib
         pip uninstall brainpy -y

brainpy/__init__.py

@@ -65,7 +65,7 @@
 Network = DynSysGroup
 # delays
 from brainpy._src.delay import (
-  VariDelay as VariDelay,
+  VarDelay as VarDelay,
 )
 
 # building blocks
@@ -129,12 +129,15 @@
 from brainpy._add_deprecations import deprecation_getattr2
 
 __deprecations = {
+  'Module': ('brainpy.Module', 'brainpy.DynamicalSystem', DynamicalSystem),
+  'Channel': ('brainpy.Channel', 'brainpy.dyn.IonChannel', dyn.IonChannel),
+  'SynConn': ('brainpy.SynConn', 'brainpy.dyn.SynConn', dyn.SynConn),
   'Container': ('brainpy.Container', 'brainpy.DynSysGroup', DynSysGroup),
+
   'optimizers': ('brainpy.optimizers', 'brainpy.optim', optim),
   'TensorCollector': ('brainpy.TensorCollector', 'brainpy.ArrayCollector', ArrayCollector),
   'SynSTP': ('brainpy.SynSTP', 'brainpy.synapses.SynSTP', synapses.SynSTP),
   'SynOut': ('brainpy.SynOut', 'brainpy.synapses.SynOut', synapses.SynOut),
-  'SynConn': ('brainpy.SynConn', 'brainpy.dyn.SynConn', dyn.SynConn),
   'TwoEndConn': ('brainpy.TwoEndConn', 'brainpy.synapses.TwoEndConn', synapses.TwoEndConn),
   'CondNeuGroup': ('brainpy.CondNeuGroup', 'brainpy.syn.CondNeuGroup', dyn.CondNeuGroup),
 }

brainpy/_src/analysis/highdim/slow_points.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 
+import inspect
 import math
 import time
+import warnings
 from typing import Callable, Union, Dict, Sequence, Tuple
 
 import jax.numpy as jnp
@@ -14,10 +16,12 @@
 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
+from brainpy._src.deprecations import _input_deprecate_msg
+
 
 __all__ = [
   'SlowPointFinder',
@@ -123,7 +127,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):
@@ -514,7 +518,7 @@ def exclude_outliers(self, tolerance: float = 1e0):
     # Compute pairwise distances between all fixed points.
     distances = np.asarray(utils.euclidean_distance_jax(self.fixed_points, num_fps))
 
-    # Find second smallest element in each column of the pairwise distance matrix.
+    # Find the second smallest element in each column of the pairwise distance matrix.
     # This corresponds to the closest neighbor for each fixed point.
     closest_neighbor = np.partition(distances, kth=1, axis=0)[1]
 
@@ -636,11 +640,16 @@ 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)
+      try:
+        ba = inspect.signature(self._inputs).bind(dict())
+        self._inputs(share.get_shargs())
+        warnings.warn(_input_deprecate_msg, UserWarning)
+      except TypeError:
+        self._inputs()
     else:
       for ops, values in self._inputs['fixed'].items():
         for var, data in values:
@@ -650,7 +659,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 +741,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 +757,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/analysis/lowdim/lowdim_analyzer.py

@@ -99,7 +99,8 @@ def __init__(
         raise errors.AnalyzerError(f'{key} is not a dynamical variable in {self.model}.')
       value = self.target_vars[key]
       if value[0] > value[1]:
-        raise errors.AnalyzerError(f'The range of variable {key} is reversed, which means {value[0]} should be smaller than {value[1]}.')
+        raise errors.AnalyzerError(
+          f'The range of variable {key} is reversed, which means {value[0]} should be smaller than {value[1]}.')
 
     # fixed variables
     # ----------------
@@ -246,7 +247,7 @@ class Num1DAnalyzer(LowDimAnalyzer):
   """
 
   def __init__(self, *args, **kwargs):
-    super(Num1DAnalyzer, self).__init__(*args, **kwargs)
+    super().__init__(*args, **kwargs)
     self.x_var = self.target_var_names[0]
     if len(self.target_vars) < 1:
       raise errors.AnalyzerError(f'{Num1DAnalyzer.__name__} only supports dynamical system '
@@ -407,7 +408,7 @@ class Num2DAnalyzer(Num1DAnalyzer):
   """
 
   def __init__(self, *args, **kwargs):
-    super(Num2DAnalyzer, self).__init__(*args, **kwargs)
+    super().__init__(*args, **kwargs)
     if len(self.target_vars) < 2:
       raise errors.AnalyzerError(f'{Num1DAnalyzer.__name__} only supports dynamical system '
                                  f'with >= 2 variables. But we got {len(self.target_vars)} '
@@ -1028,7 +1029,7 @@ def _get_fixed_points(self, candidates, *args, tol_aux=1e-7,
 
 class Num3DAnalyzer(Num2DAnalyzer):
   def __init__(self, *args, **kwargs):
-    super(Num3DAnalyzer, self).__init__(*args, **kwargs)
+    super().__init__(*args, **kwargs)
     if len(self.target_vars) < 3:
       raise errors.AnalyzerError(f'{Num1DAnalyzer.__name__} only supports dynamical system '
                                  f'with >= 3 variables. But we got {len(self.target_vars)} '
@@ -1045,7 +1046,3 @@ def F_fz(self):
       f = partial(f, **(self.pars_update + self.fixed_vars))
       self.analyzed_results[C.F_fz] = jax.jit(f, device=self.jit_device)
     return self.analyzed_results[C.F_fz]
-
-  def fz_signs(self, pars=(), cache=False):
-    xyz = tuple(self.resolutions.values())
-    return utils.get_sign2(self.F_fz, *xyz, args=pars)

brainpy/_src/analysis/lowdim/lowdim_bifurcation.py

@@ -31,13 +31,13 @@ class Bifurcation1D(Num1DAnalyzer):
 
   def __init__(self, model, target_pars, target_vars, fixed_vars=None,
                pars_update=None, resolutions=None, options=None):
-    super(Bifurcation1D, self).__init__(model=model,
-                                        target_pars=target_pars,
-                                        target_vars=target_vars,
-                                        fixed_vars=fixed_vars,
-                                        pars_update=pars_update,
-                                        resolutions=resolutions,
-                                        options=options)
+    super().__init__(model=model,
+                     target_pars=target_pars,
+                     target_vars=target_vars,
+                     fixed_vars=fixed_vars,
+                     pars_update=pars_update,
+                     resolutions=resolutions,
+                     options=options)
 
     if len(self.target_pars) == 0:
       raise ValueError
@@ -146,13 +146,13 @@ class Bifurcation2D(Num2DAnalyzer):
 
   def __init__(self, model, target_pars, target_vars, fixed_vars=None,
                pars_update=None, resolutions=None, options=None):
-    super(Bifurcation2D, self).__init__(model=model,
-                                        target_pars=target_pars,
-                                        target_vars=target_vars,
-                                        fixed_vars=fixed_vars,
-                                        pars_update=pars_update,
-                                        resolutions=resolutions,
-                                        options=options)
+    super().__init__(model=model,
+                     target_pars=target_pars,
+                     target_vars=target_vars,
+                     fixed_vars=fixed_vars,
+                     pars_update=pars_update,
+                     resolutions=resolutions,
+                     options=options)
 
     if len(self.target_pars) == 0:
       raise ValueError
@@ -458,13 +458,13 @@ def __init__(
       resolutions=None,
       options: dict = None
   ):
-    super(FastSlow1D, self).__init__(model=model,
-                                     target_pars=slow_vars,
-                                     target_vars=fast_vars,
-                                     fixed_vars=fixed_vars,
-                                     pars_update=pars_update,
-                                     resolutions=resolutions,
-                                     options=options)
+    super().__init__(model=model,
+                     target_pars=slow_vars,
+                     target_vars=fast_vars,
+                     fixed_vars=fixed_vars,
+                     pars_update=pars_update,
+                     resolutions=resolutions,
+                     options=options)
 
     # standard integrators
     self._std_integrators = dict()
@@ -549,13 +549,13 @@ def __init__(
       resolutions=0.1,
       options: dict = None
   ):
-    super(FastSlow2D, self).__init__(model=model,
-                                     target_pars=slow_vars,
-                                     target_vars=fast_vars,
-                                     fixed_vars=fixed_vars,
-                                     pars_update=pars_update,
-                                     resolutions=resolutions,
-                                     options=options)
+    super().__init__(model=model,
+                     target_pars=slow_vars,
+                     target_vars=fast_vars,
+                     fixed_vars=fixed_vars,
+                     pars_update=pars_update,
+                     resolutions=resolutions,
+                     options=options)
     # standard integrators
     self._std_integrators = dict()
     for key, intg in self.model.name2integral.items():

brainpy/_src/analysis/lowdim/lowdim_phase_plane.py

@@ -55,13 +55,13 @@ def __init__(self,
     if (target_pars is not None) and len(target_pars) > 0:
       raise errors.AnalyzerError(f'Phase plane analysis does not support "target_pars". '
                                  f'While we detect "target_pars={target_pars}".')
-    super(PhasePlane1D, self).__init__(model=model,
-                                       target_vars=target_vars,
-                                       fixed_vars=fixed_vars,
-                                       target_pars=target_pars,
-                                       pars_update=pars_update,
-                                       resolutions=resolutions,
-                                       **kwargs)
+    super().__init__(model=model,
+                     target_vars=target_vars,
+                     fixed_vars=fixed_vars,
+                     target_pars=target_pars,
+                     pars_update=pars_update,
+                     resolutions=resolutions,
+                     **kwargs)
     # utils.output(f'I am {PhasePlane1D.__name__}.')
 
   def plot_vector_field(self, show=False, with_plot=True, with_return=False):
@@ -150,13 +150,13 @@ def __init__(self,
     if (target_pars is not None) and len(target_pars) > 0:
       raise errors.AnalyzerError(f'Phase plane analysis does not support "target_pars". '
                                  f'While we detect "target_pars={target_pars}".')
-    super(PhasePlane2D, self).__init__(model=model,
-                                       target_vars=target_vars,
-                                       fixed_vars=fixed_vars,
-                                       target_pars=target_pars,
-                                       pars_update=pars_update,
-                                       resolutions=resolutions,
-                                       **kwargs)
+    super().__init__(model=model,
+                     target_vars=target_vars,
+                     fixed_vars=fixed_vars,
+                     target_pars=target_pars,
+                     pars_update=pars_update,
+                     resolutions=resolutions,
+                     **kwargs)
 
   @property
   def F_vmap_brentq_fy(self):
@@ -251,7 +251,7 @@ def plot_nullcline(self, with_plot=True, with_return=False,
     if with_plot:
       if x_style is None:
         x_style = dict(color='cornflowerblue', alpha=.7, fmt='.')
-      line_args = (x_style.pop('fmt'), ) if 'fmt' in x_style else tuple()
+      line_args = (x_style.pop('fmt'),) if 'fmt' in x_style else tuple()
       pyplot.plot(x_values_in_fx, y_values_in_fx, *line_args, **x_style, label=f"{self.x_var} nullcline")
 
     # Nullcline of the y variable
@@ -263,7 +263,7 @@ def plot_nullcline(self, with_plot=True, with_return=False,
     if with_plot:
       if y_style is None:
         y_style = dict(color='lightcoral', alpha=.7, fmt='.')
-      line_args = (y_style.pop('fmt'), ) if 'fmt' in y_style else tuple()
+      line_args = (y_style.pop('fmt'),) if 'fmt' in y_style else tuple()
       pyplot.plot(x_values_in_fy, y_values_in_fy, *line_args, **y_style, label=f"{self.y_var} nullcline")
 
     if with_plot:

brainpy/_src/analysis/utils/model.py

@@ -5,6 +5,7 @@
 from brainpy._src.math.environment import get_float
 from brainpy._src.math.interoperability import as_jax
 from brainpy._src.dynsys import DynamicalSystem
+from brainpy._src.context import share
 from brainpy._src.runners import DSRunner
 from brainpy._src.integrators.base import Integrator
 from brainpy._src.integrators.joint_eq import JointEq
@@ -126,16 +127,12 @@ def __init__(self, integrals: dict, initial_vars: dict, pars=None, dt=None):
     self.integrals = integrals
 
     # runner
-    self.runner = DSRunner(self,
-                           monitors=list(initial_vars.keys()),
-                           dyn_vars=self.vars().unique(),
-                           dt=dt,
-                           progress_bar=False)
+    self.runner = DSRunner(self, monitors=list(initial_vars.keys()), dt=dt, progress_bar=False)
 
-  def update(self, sha):
+  def update(self):
     all_vars = list(self.implicit_vars.values())
     for key, intg in self.integrals.items():
-      self.implicit_vars[key].update(intg(*all_vars, *self.pars, dt=sha['dt']))
+      self.implicit_vars[key].update(intg(*all_vars, *self.pars, dt=share['dt']))
 
   def __getattr__(self, item):
     child_vars = super(TrajectModel, self).__getattribute__('implicit_vars')

brainpy/_src/delay.py

@@ -21,7 +21,7 @@
 
 __all__ = [
   'Delay',
-  'VariDelay',
+  'VarDelay',
   'DataDelay',
   'DelayAccess',
 ]
@@ -432,7 +432,7 @@ def _check_target_sharding(sharding, ndim, mode: bm.Mode):
   return sharding
 
 
-class VariDelay(Delay):
+class VarDelay(Delay):
   """Generate Delays for the given :py:class:`~.Variable` instance.
 
   The data in this delay variable is arranged as::
@@ -690,7 +690,7 @@ def _init_data(self, length: int, batch_size: int = None):
       self.data[:] = self._init((length,) + self.target.shape, dtype=self.target.dtype)
 
 
-class DataDelay(VariDelay):
+class DataDelay(VarDelay):
   not_desc_params = ('time', 'entries')
 
   def __init__(