upgrade brainpy.analysis for new version of DynamicalSystem

Author: chaoming0625

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
@@ -18,6 +20,8 @@
 from brainpy._src.runners import check_and_format_inputs, _f_ops
 from brainpy.errors import AnalyzerError, UnsupportedError
 from brainpy.types import ArrayType
+from brainpy._src.deprecations import _input_deprecate_msg
+
 
 __all__ = [
   'SlowPointFinder',
@@ -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]
 
@@ -640,7 +644,12 @@ def _step_func_input(self):
     if self._inputs is None:
       return
     elif callable(self._inputs):
-      self._inputs(share.get_shargs())
+      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:

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')