upgrade object_transform

Author: chaoming0625

brainpy/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "2.2.4.1"
+__version__ = "2.3.0"
 
 
 # fundamental modules

brainpy/analysis/highdim/tests/test_slow_points.py

@@ -105,6 +105,7 @@ def step(s):
 
     finder = bp.analysis.SlowPointFinder(f_cell=step, f_type=bp.analysis.CONTINUOUS)
     finder.find_fps_with_opt_solver(bm.random.random((100, 2)))
+    bm.clear_buffer_memory()
 
   def test_opt_solver_for_ds1(self):
     hh = HH(1)
@@ -117,6 +118,7 @@ def test_opt_solver_for_ds1(self):
                                      'm': bm.random.random((100, 1)),
                                      'h': bm.random.random((100, 1)),
                                      'n': bm.random.random((100, 1))})
+    bm.clear_buffer_memory()
 
   def test_gd_method_for_func1(self):
     gamma = 0.641  # Saturation factor for gating variable
@@ -149,6 +151,7 @@ def step(s):
 
     finder = bp.analysis.SlowPointFinder(f_cell=step, f_type=bp.analysis.CONTINUOUS)
     finder.find_fps_with_gd_method(bm.random.random((100, 2)), num_opt=100)
+    bm.clear_buffer_memory()
 
   def test_gd_method_for_func2(self):
     hh = HH(1)
@@ -162,4 +165,5 @@ def test_gd_method_for_func2(self):
                                     'h': bm.random.random((100, 1)),
                                     'n': bm.random.random((100, 1))},
                                    num_opt=100)
+    bm.clear_buffer_memory()
 

brainpy/base/__init__.py

@@ -7,10 +7,10 @@
   and its associated helper class ``Collector`` and ``ArrayCollector``.
 - For each instance of "BrainPyObject" class, users can retrieve all
   the variables (or trainable variables), integrators, and nodes.
-- This module also provides a ``FunAsAObject`` class to wrap user-defined
+- This module also provides a ``FunAsObject`` class to wrap user-defined
   functions. In each function, maybe several nodes are used, and
-  users can initialize a ``FunAsAObject`` by providing the nodes used
-  in the function. Unfortunately, ``FunAsAObject`` class does not have
+  users can initialize a ``FunAsObject`` by providing the nodes used
+  in the function. Unfortunately, ``FunAsObject`` class does not have
   the ability to gather nodes automatically.
 - This module provides ``io`` helper functions to help users save/load
   model states, or share user's customized model with others.

brainpy/base/base.py

@@ -24,10 +24,10 @@ class BrainPyObject(object):
 
   - ``DynamicalSystem`` in *brainpy.dyn.base.py*
   - ``Integrator`` in *brainpy.integrators.base.py*
-  - ``FunAsAObject`` in *brainpy.base.function.py*
+  - ``FunAsObject`` in *brainpy.base.function.py*
   - ``Optimizer`` in *brainpy.optimizers.py*
   - ``Scheduler`` in *brainpy.optimizers.py*
-
+  - and others.
   """
 
   _excluded_vars = ()

brainpy/base/collector.py

@@ -186,3 +186,12 @@ def dict(self):
   def data(self):
     """Get all data in each value."""
     return [x.value for x in self.values()]
+
+  @classmethod
+  def from_other(cls, other: Union[Sequence, Dict]):
+    if isinstance(other, (tuple, list)):
+      return cls({id(o): o for o in other})
+    elif isinstance(other, dict):
+      return cls(other)
+    else:
+      raise TypeError

brainpy/base/function.py

@@ -1,14 +1,15 @@
 # -*- coding: utf-8 -*-
 
 from typing import Optional, Callable
+
 from brainpy import errors
-from brainpy.base.base import BrainPyObject
 from brainpy.base import collector
+from brainpy.base.base import BrainPyObject
 
 math = None
 
 __all__ = [
-  'FunAsAObject',
+  'FunAsObject',
 ]
 
 
@@ -26,7 +27,7 @@ def _check_var(var):
                               f'{math.ndarray.__name__}, but we got {type(var)}.')
 
 
-class FunAsAObject(BrainPyObject):
+class FunAsObject(BrainPyObject):
   """The wrapper for Python functions.
 
   Parameters
@@ -46,8 +47,8 @@ def __init__(self, f: Optional[Callable], child_objs=None, dyn_vars=None, name=N
     # ---
     self._f = f
     if name is None:
-      name = self.unique_name(type_=f.__name__ if hasattr(f, '__name__') else 'FunAsAObject')
-    super(FunAsAObject, self).__init__(name=name)
+      name = self.unique_name(type_=f.__name__ if hasattr(f, '__name__') else 'FunAsObject')
+    super(FunAsObject, self).__init__(name=name)
 
     # nodes 
     # ---
@@ -89,3 +90,19 @@ def __init__(self, f: Optional[Callable], child_objs=None, dyn_vars=None, name=N
 
   def __call__(self, *args, **kwargs):
     return self._f(*args, **kwargs)
+
+  def __repr__(self):
+    name = self.__class__.__name__
+    # indent = ' ' * (len(name) + 10)
+    # child_nodes = ['\n'.join([('' if i == 0 else indent) + l for i, l in enumerate(repr(node).split('\n'))])
+    #                for node in self.implicit_nodes.values()]
+    # first_line = f'{name}(objects=['
+    # format_res = (
+    #     first_line +
+    #     (',\n' + ' ' * (len(name) + 10)).join(child_nodes) +
+    #     '],\n'
+    #     + (" " * (len(name) + 1)) + f'number of variables = {len(self.implicit_vars)})'
+    # )
+    format_ref = (f'{name}(nodes=[{", ".join([n.name for n in tuple(self.implicit_nodes.values())])}],\n' +
+                  " " * (len(name) + 1) + f'num_of_vars={len(self.implicit_vars)})')
+    return format_ref

brainpy/dyn/base.py

@@ -59,17 +59,16 @@ class DynamicalSystem(BrainPyObject):
     The model computation mode. It should be instance of :py:class:`~.Mode`.
   """
 
-  '''Online fitting method.'''
   online_fit_by: Optional[OnlineAlgorithm]
+  '''Online fitting method.'''
 
-  '''Offline fitting method.'''
   offline_fit_by: Optional[OfflineAlgorithm]
+  '''Offline fitting method.'''
 
+  global_delay_data: Dict[str, Tuple[Union[bm.LengthDelay, None], bm.Variable]] = dict()
   '''Global delay data, which stores the delay variables and corresponding delay targets. 
-   
   This variable is useful when the same target variable is used in multiple mappings, 
   as it can reduce the duplicate delay variable registration.'''
-  global_delay_data: Dict[str, Tuple[Union[bm.LengthDelay, None], bm.Variable]] = dict()
 
   def __init__(
       self,
@@ -435,15 +434,45 @@ def clear_input(self):
       node.clear_input()
 
 
-class Sequential(Container):
+class Sequential(DynamicalSystem):
   def __init__(
       self,
       *modules,
       name: str = None,
       mode: Mode = normal,
       **kw_modules
   ):
-    super(Sequential, self).__init__(*modules, name=name, mode=mode, **kw_modules)
+    super().__init__(name=name, mode=mode)
+    self._modules = tuple(modules) + tuple(kw_modules.values())
+
+    seq_modules = [m for m in modules if isinstance(m, BrainPyObject)]
+    dict_modules = {k: m for k, m in kw_modules.items() if isinstance(m, BrainPyObject)}
+
+    # add tuple-typed components
+    for module in seq_modules:
+      if isinstance(module, BrainPyObject):
+        self.implicit_nodes[module.name] = module
+      elif isinstance(module, (list, tuple)):
+        for m in module:
+          if not isinstance(m, BrainPyObject):
+            raise ValueError(f'Should be instance of {BrainPyObject.__name__}. '
+                             f'But we got {type(m)}')
+          self.implicit_nodes[m.name] = module
+      elif isinstance(module, dict):
+        for k, v in module.items():
+          if not isinstance(v, BrainPyObject):
+            raise ValueError(f'Should be instance of {BrainPyObject.__name__}. '
+                             f'But we got {type(v)}')
+          self.implicit_nodes[k] = v
+      else:
+        raise ValueError(f'Cannot parse sub-systems. They should be {BrainPyObject.__name__} '
+                         f'or a list/tuple/dict of  {BrainPyObject.__name__}.')
+    # add dict-typed components
+    for k, v in dict_modules.items():
+      if not isinstance(v, BrainPyObject):
+        raise ValueError(f'Should be instance of {BrainPyObject.__name__}. '
+                         f'But we got {type(v)}')
+      self.implicit_nodes[k] = v
 
   def __getattr__(self, item):
     """Wrap the dot access ('self.'). """
@@ -463,7 +492,7 @@ def __getitem__(self, key: Union[int, slice]):
       components = tuple(self.implicit_nodes.values())[key]
       return Sequential(dict(zip(keys, components)))
     elif isinstance(key, int):
-      return self.implicit_nodes.values()[key]
+      return tuple(self.implicit_nodes.values())[key]
     elif isinstance(key, (tuple, list)):
       all_keys = tuple(self.implicit_nodes.keys())
       all_vals = tuple(self.implicit_nodes.values())
@@ -478,27 +507,7 @@ def __getitem__(self, key: Union[int, slice]):
       raise KeyError(f'Unknown type of key: {type(key)}')
 
   def __repr__(self):
-    def f(x):
-      if not isinstance(x, DynamicalSystem) and callable(x):
-        signature = inspect.signature(x)
-        args = [f'{k}={v.default}' for k, v in signature.parameters.items()
-                if v.default is not inspect.Parameter.empty]
-        args = ', '.join(args)
-        while not hasattr(x, '__name__'):
-          if not hasattr(x, 'func'):
-            break
-          x = x.func  # Handle functools.partial
-        if not hasattr(x, '__name__') and hasattr(x, '__class__'):
-          return x.__class__.__name__
-        if args:
-          return f'{x.__name__}(*, {args})'
-        return x.__name__
-      else:
-        x = repr(x).split('\n')
-        x = [x[0]] + ['  ' + y for y in x[1:]]
-        return '\n'.join(x)
-
-    entries = '\n'.join(f'  [{i}] {f(x)}' for i, x in enumerate(self))
+    entries = '\n'.join(f'  [{i}] {tools.repr_object(x)}' for i, x in enumerate(self._modules))
     return f'{self.__class__.__name__}(\n{entries}\n)'
 
   def update(self, sha: dict, x: Any) -> Array:
@@ -516,8 +525,11 @@ def update(self, sha: dict, x: Any) -> Array:
     y: Array
       The output tensor.
     """
-    for node in self.implicit_nodes.values():
-      x = node(sha, x)
+    for m in self._modules:
+      if isinstance(m, DynamicalSystem):
+        x = m(sha, x)
+      else:
+        x = m(x)
     return x
 
 

brainpy/dyn/synapses/abstract_models.py

@@ -54,7 +54,7 @@ class Delta(TwoEndConn):
     >>>
     >>> neu1 = neurons.LIF(1)
     >>> neu2 = neurons.LIF(1)
-    >>> syn1 = synapses.Alpha(neu1, neu2, bp.connect.All2All(), weights=5.)
+    >>> syn1 = synapses.Alpha(neu1, neu2, bp.connect.All2All(), g_max=5.)
     >>> net = bp.dyn.Network(pre=neu1, syn=syn1, post=neu2)
     >>>
     >>> runner = bp.dyn.DSRunner(net, inputs=[('pre.input', 25.), ('post.input', 10.)], monitors=['pre.V', 'post.V', 'pre.spike'])

brainpy/dyn/synapses/biological_models.py

@@ -412,7 +412,7 @@ class BioNMDA(TwoEndConn):
     >>>
     >>> neu1 = neurons.HH(1)
     >>> neu2 = neurons.HH(1)
-    >>> syn1 = synapses.BioNMDA(neu1, neu2, bp.connect.All2All(), E=0.)
+    >>> syn1 = synapses.BioNMDA(neu1, neu2, bp.connect.All2All())
     >>> net = bp.dyn.Network(pre=neu1, syn=syn1, post=neu2)
     >>>
     >>> runner = bp.dyn.DSRunner(net, inputs=[('pre.input', 5.)], monitors=['pre.V', 'post.V', 'syn.g', 'syn.x'])

brainpy/dyn/synapses/learning_rules.py

@@ -76,9 +76,9 @@ class STP(TwoEndConn):
     >>> import brainpy as bp
     >>> import matplotlib.pyplot as plt
     >>>
-    >>> neu1 = bp.dyn.LIF(1)
-    >>> neu2 = bp.dyn.LIF(1)
-    >>> syn1 = bp.dyn.STP(neu1, neu2, bp.connect.All2All(), U=0.2, tau_d=150., tau_f=2.)
+    >>> neu1 = bp.neurons.LIF(1)
+    >>> neu2 = bp.neurons.LIF(1)
+    >>> syn1 = bp.synapses.STP(neu1, neu2, bp.connect.All2All(), U=0.2, tau_d=150., tau_f=2.)
     >>> net = bp.dyn.Network(pre=neu1, syn=syn1, post=neu2)
     >>>
     >>> runner = bp.dyn.DSRunner(net, inputs=[('pre.input', 28.)], monitors=['syn.I', 'syn.u', 'syn.x'])