Merge branch 'brainpy:master' into master

Author: ztqakita

brainpy/__init__.py

@@ -33,6 +33,8 @@
   inputs,  # methods for generating input currents
   algorithms,  # online or offline training algorithms
   encoding,  # encoding schema
+  checkpoints,  # checkpoints
+  check,  # error checking
 )
 
 # numerical integrators

brainpy/base/base.py

@@ -2,12 +2,16 @@
 
 import logging
 import os.path
+import warnings
+from collections import namedtuple
+from typing import Dict, Any, Tuple
 
 from brainpy import errors
 from brainpy.base import io, naming
 from brainpy.base.collector import Collector, ArrayCollector
 
 math = None
+StateLoadResult = namedtuple('StateLoadResult', ['missing_keys', 'unexpected_keys'])
 
 __all__ = [
   'BrainPyObject',
@@ -57,25 +61,27 @@ def name(self, name: str = None):
     naming.check_name_uniqueness(name=self._name, obj=self)
 
   def register_implicit_vars(self, *variables, **named_variables):
-    from brainpy.math import Variable
+    global math
+    if math is None: from brainpy import math
+
     for variable in variables:
-      if isinstance(variable, Variable):
+      if isinstance(variable, math.Variable):
         self.implicit_vars[f'var{id(variable)}'] = variable
       elif isinstance(variable, (tuple, list)):
         for v in variable:
-          if not isinstance(v, Variable):
-            raise ValueError(f'Must be instance of {Variable.__name__}, but we got {type(v)}')
+          if not isinstance(v, math.Variable):
+            raise ValueError(f'Must be instance of {math.Variable.__name__}, but we got {type(v)}')
           self.implicit_vars[f'var{id(v)}'] = v
       elif isinstance(variable, dict):
         for k, v in variable.items():
-          if not isinstance(v, Variable):
-            raise ValueError(f'Must be instance of {Variable.__name__}, but we got {type(v)}')
+          if not isinstance(v, math.Variable):
+            raise ValueError(f'Must be instance of {math.Variable.__name__}, but we got {type(v)}')
           self.implicit_vars[k] = v
       else:
         raise ValueError(f'Unknown type: {type(variable)}')
     for key, variable in named_variables.items():
-      if not isinstance(variable, Variable):
-        raise ValueError(f'Must be instance of {Variable.__name__}, but we got {type(variable)}')
+      if not isinstance(variable, math.Variable):
+        raise ValueError(f'Must be instance of {math.Variable.__name__}, but we got {type(variable)}')
       self.implicit_vars[key] = variable
 
   def register_implicit_nodes(self, *nodes, node_cls: type = None, **named_nodes):
@@ -101,7 +107,11 @@ def register_implicit_nodes(self, *nodes, node_cls: type = None, **named_nodes):
         raise ValueError(f'Must be instance of {node_cls.__name__}, but we got {type(node)}')
       self.implicit_nodes[key] = node
 
-  def vars(self, method='absolute', level=-1, include_self=True):
+  def vars(self,
+           method: str = 'absolute',
+           level: int = -1,
+           include_self: bool = True,
+           exclude_types: Tuple[type, ...] = None):
     """Collect all variables in this node and the children nodes.
 
     Parameters
@@ -112,6 +122,8 @@ def vars(self, method='absolute', level=-1, include_self=True):
       The hierarchy level to find variables.
     include_self: bool
       Whether include the variables in the self.
+    exclude_types: tuple of type
+      The type to exclude.
 
     Returns
     -------
@@ -121,12 +133,19 @@ def vars(self, method='absolute', level=-1, include_self=True):
     global math
     if math is None: from brainpy import math
 
+    if exclude_types is None:
+      exclude_types = (math.VariableView, )
     nodes = self.nodes(method=method, level=level, include_self=include_self)
     gather = ArrayCollector()
     for node_path, node in nodes.items():
       for k in dir(node):
         v = getattr(node, k)
+        include = False
         if isinstance(v, math.Variable):
+          include = True
+          if len(exclude_types) and isinstance(v, exclude_types):
+              include = False
+        if include:
           if k not in node._excluded_vars:
             gather[f'{node_path}.{k}' if node_path else k] = v
       gather.update({f'{node_path}.{k}': v for k, v in node.implicit_vars.items()})
@@ -306,6 +325,49 @@ def save_states(self, filename, variables=None, **setting):
     else:
       raise errors.BrainPyError(f'Unknown file format: {filename}. We only supports {io.SUPPORTED_FORMATS}')
 
+  def state_dict(self):
+    """Returns a dictionary containing a whole state of the module.
+
+    Returns
+    -------
+    out: dict
+      A dictionary containing a whole state of the module.
+    """
+    return self.vars().unique().dict()
+
+  def load_state_dict(self, state_dict: Dict[str, Any], warn: bool = True):
+    """Copy parameters and buffers from :attr:`state_dict` into
+    this module and its descendants.
+
+    Parameters
+    ----------
+    state_dict: dict
+      A dict containing parameters and persistent buffers.
+    warn: bool
+      Warnings when there are missing keys or unexpected keys in the external ``state_dict``.
+
+    Returns
+    -------
+    out: StateLoadResult
+      ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+
+      * **missing_keys** is a list of str containing the missing keys
+      * **unexpected_keys** is a list of str containing the unexpected keys
+    """
+    variables = self.vars().unique()
+    keys1 = set(state_dict.keys())
+    keys2 = set(variables.keys())
+    unexpected_keys = list(keys1 - keys2)
+    missing_keys = list(keys2 - keys1)
+    for key in keys2.intersection(keys1):
+      variables[key].value = state_dict[key]
+    if warn:
+      if len(unexpected_keys):
+        warnings.warn(f'Unexpected keys in state_dict: {unexpected_keys}', UserWarning)
+      if len(missing_keys):
+        warnings.warn(f'Missing keys in state_dict: {missing_keys}', UserWarning)
+    return StateLoadResult(missing_keys, unexpected_keys)
+
   # def to(self, devices):
   #   global math
   #   if math is None: from brainpy import math
@@ -317,7 +379,6 @@ def save_states(self, filename, variables=None, **setting):
   #   all_vars = self.vars().unique()
   #   for data in all_vars.values():
   #     data[:] = math.asarray(data.value)
-  #     # TODO
   #
   # def cuda(self):
   #   global math

brainpy/base/collector.py

@@ -3,6 +3,9 @@
 
 from typing import Dict, Sequence, Union
 
+from jax.tree_util import register_pytree_node
+from jax.util import safe_zip
+
 math = None
 
 __all__ = [
@@ -37,11 +40,12 @@ def update(self, other, **kwargs):
     elif isinstance(other, (tuple, list)):
       num = len(self)
       for i, value in enumerate(other):
-        self[f'_var{i+num}'] = value
+        self[f'_var{i + num}'] = value
     else:
       raise ValueError(f'Only supports dict/list/tuple, but we got {type(other)}')
     for key, value in kwargs.items():
       self[key] = value
+    return self
 
   def __add__(self, other):
     """Merging two dicts.
@@ -73,8 +77,8 @@ def __sub__(self, other: Union[Dict, Sequence]):
     gather: Collector
       The new collector.
     """
-    if not isinstance(other, dict):
-      raise ValueError(f'Only support dict, but we got {type(other)}.')
+    if not isinstance(other, (dict, tuple, list)):
+      raise ValueError(f'Only support dict/tuple/list, but we got {type(other)}.')
     gather = type(self)(self)
     if isinstance(other, dict):
       for key, val in other.items():
@@ -87,7 +91,21 @@ def __sub__(self, other: Union[Dict, Sequence]):
           raise ValueError(f'Cannot remove {key}, because we do not find it '
                            f'in {self.keys()}.')
     elif isinstance(other, (list, tuple)):
+      id_to_keys = {}
+      for k, v in self.items():
+        id_ = id(v)
+        if id_ not in id_to_keys:
+          id_to_keys[id_] = []
+        id_to_keys[id_].append(k)
+
+      keys_to_remove = []
       for key in other:
+        if isinstance(key, str):
+          keys_to_remove.append(key)
+        else:
+          keys_to_remove.extend(id_to_keys[id(key)])
+
+      for key in set(keys_to_remove):
         if key in gather:
           gather.pop(key)
         else:
@@ -199,3 +217,9 @@ def from_other(cls, other: Union[Sequence, Dict]):
 
 
 TensorCollector = ArrayCollector
+
+register_pytree_node(
+  ArrayCollector,
+  lambda x: (x.values(), x.keys()),
+  lambda keys, values: ArrayCollector(safe_zip(keys, values))
+)

brainpy/base/io.py

@@ -72,7 +72,7 @@ def _check_target(target):
 
 not_found_msg = ('"{key}" is stored in {filename}. But we does '
                  'not find it is defined as variable in {target}.')
-id_dismatch_msg = ('{key1} and {key2} is the same data in {filename}. '
+id_mismatch_msg = ('{key1} and {key2} is the same data in {filename}. '
                    'But we found they are different in {target}.')
 
 DUPLICATE_KEY = 'duplicate_keys'
@@ -92,7 +92,7 @@ def _load(
   # get variables
   _check_target(target)
   variables = target.vars(method='absolute', level=-1)
-  all_names = list(variables.keys())
+  var_names_in_obj = list(variables.keys())
 
   # read data from file
   for key in load_vars.keys():
@@ -105,22 +105,22 @@ def _load(
     else:
       value = load_vars[key]
     variables[key].value = bm.asarray(value)
-    all_names.remove(key)
+    var_names_in_obj.remove(key)
 
   # check duplicate names
   duplicate_keys = duplicates[0]
   duplicate_targets = duplicates[1]
   for key1, key2 in zip(duplicate_keys, duplicate_targets):
-    if key1 not in all_names:
+    if key1 not in var_names_in_obj:
       raise KeyError(not_found_msg.format(key=key1, target=target.name, filename=filename))
     if id(variables[key1]) != id(variables[key2]):
-      raise ValueError(id_dismatch_msg.format(key1=key1, key2=target, filename=filename, target=target.name))
-    all_names.remove(key1)
+      raise ValueError(id_mismatch_msg.format(key1=key1, key2=target, filename=filename, target=target.name))
+    var_names_in_obj.remove(key1)
 
   # check missing names
-  if len(all_names):
+  if len(var_names_in_obj):
     logger.warning(f'There are variable states missed in {filename}. '
-                   f'The missed variables are: {all_names}.')
+                   f'The missed variables are: {var_names_in_obj}.')
 
 
 def _unique_and_duplicate(collector: dict):