fix bugs

Author: chaoming0625

brainpy/dyn/runners.py

@@ -346,7 +346,7 @@ def __init__(
     self.i0 = bm.Variable(bm.asarray([1], dtype=bm.int_))
     self.t0 = bm.Variable(bm.asarray([t0], dtype=bm.float_))
     if data_first_axis is None:
-      data_first_axis = 'B' if isinstance(self.target, bm.BatchingMode) else 'T'
+      data_first_axis = 'B' if isinstance(self.target.mode, bm.BatchingMode) else 'T'
     assert data_first_axis in ['B', 'T']
     self.data_first_axis = data_first_axis
 
@@ -380,8 +380,8 @@ def __repr__(self):
 
   def reset_state(self):
     """Reset state of the ``DSRunner``."""
-    self.i0[0] = 0
-    self.t0[0] = self._t0
+    self.i0.value = bm.zeros_like(self.i0)
+    self.t0.value = bm.ones_like(self.t0) * self._t0
 
   def predict(
       self,
@@ -635,8 +635,10 @@ def _get_f_predict(self, shared_args: Dict = None):
 
     shared_kwargs_str = serialize_kwargs(shared_args)
     if shared_kwargs_str not in self._f_predict_compiled:
-      dyn_vars = self.vars().unique()
-      dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
+      dyn_vars = self.target.vars()
+      dyn_vars.update(self._dyn_vars)
+      dyn_vars.update(self.vars(level=0))
+      dyn_vars = dyn_vars.unique()
 
       def run_func(all_inputs):
         with jax.disable_jit(not self.jit['predict']):

brainpy/integrators/runner.py

@@ -168,8 +168,14 @@ def __init__(
       self.variables[k][:] = inits[k]
 
     # format string monitors
-    monitors = self._format_seq_monitors(monitors)
-    monitors = {k: (self.variables[k], i) for k, i in monitors}
+    if isinstance(monitors, (tuple, list)):
+      monitors = self._format_seq_monitors(monitors)
+      monitors = {k: (self.variables[k], i) for k, i in monitors}
+    elif isinstance(monitors, dict):
+      monitors = self._format_dict_monitors(monitors)
+      monitors = {k: ((self.variables[i], i) if isinstance(i, str) else i) for k, i in monitors.items()}
+    else:
+      raise ValueError
 
     # initialize super class
     super(IntegratorRunner, self).__init__(target=target,
@@ -218,12 +224,6 @@ def __init__(
     else:
       self._dyn_args = dict()
 
-    # monitors
-    for k in self.mon.var_names:
-      if k not in self.target.variables:
-        raise MonitorError(f'Variable "{k}" to monitor is not defined '
-                           f'in the integrator {self.target}.')
-
     # start simulation time and index
     self.start_t = bm.Variable(bm.zeros(1))
     self.idx = bm.Variable(bm.zeros(1, dtype=bm.int_))

brainpy/math/environment.py

@@ -6,12 +6,12 @@
 import os
 import re
 import sys
-import warnings
 from typing import Any, Callable, TypeVar, cast
 
 from jax import config, numpy as jnp, devices
 from jax.lib import xla_bridge
 
+from brainpy import errors
 from . import modes
 
 bm = None
@@ -63,8 +63,8 @@ def ditype():
   .. deprecated:: 2.3.1
      Use `brainpy.math.int_` instead.
   """
-  warnings.warn('\nGet default integer data type through `ditype()` has been deprecated. \n'
-                'Use `brainpy.math.int_` instead.')
+  # raise errors.NoLongerSupportError('\nGet default integer data type through `ditype()` has been deprecated. \n'
+  #                                   'Use `brainpy.math.int_` instead.')
   global bm
   if bm is None: from brainpy import math as bm
   return bm.int_
@@ -77,8 +77,8 @@ def dftype():
      Use `brainpy.math.float_` instead.
   """
 
-  warnings.warn('\nGet default floating data type through `dftype()` has been deprecated. \n'
-                'Use `brainpy.math.float_` instead.')
+  # raise errors.NoLongerSupportError('\nGet default floating data type through `dftype()` has been deprecated. \n'
+  #                                   'Use `brainpy.math.float_` instead.')
   global bm
   if bm is None: from brainpy import math as bm
   return bm.float_

brainpy/math/object_transform/__init__.py

@@ -16,6 +16,9 @@
 """
 
 from . import (
+  base_object,
+  base_transform,
+  collector,
   autograd,
   controls,
   jit,

brainpy/running/runner.py

@@ -99,36 +99,12 @@ def __init__(
       # format string monitors
       monitors = self._format_seq_monitors(monitors)
       # get monitor targets
-      monitors = self._find_monitor_targets(monitors)
+      monitors = self._find_seq_monitor_targets(monitors)
     elif isinstance(monitors, dict):
-      _monitors = dict()
-      for key, val in monitors.items():
-        if not isinstance(key, str):
-          raise MonitorError('Expect the key of the dict "monitors" must be a string. But got '
-                             f'{type(key)}: {key}')
-        if isinstance(val, bm.Variable):
-          val = (val, None)
-        if isinstance(val, (tuple, list)):
-          if not isinstance(val[0], bm.Variable):
-            raise MonitorError('Expect the format of (variable, index) in the monitor setting. '
-                               f'But we got {val}')
-          if len(val) == 1:
-            _monitors[key] = (val[0], None)
-          elif len(val) == 2:
-            if isinstance(val[1], (int, np.integer)):
-              idx = bm.array([val[1]])
-            else:
-              idx = None if val[1] is None else bm.asarray(val[1])
-            _monitors[key] = (val[0], idx)
-          else:
-            raise MonitorError('Expect the format of (variable, index) in the monitor setting. '
-                               f'But we got {val}')
-        elif callable(val):
-          _monitors[key] = val
-        else:
-          raise MonitorError('The value of dict monitor expect a sequence with (variable, index) '
-                             f'or a callable function. But we got {val}')
-      monitors = _monitors
+      # format string monitors
+      monitors = self._format_dict_monitors(monitors)
+      # get monitor targets
+      monitors = self._find_dict_monitor_targets(monitors)
     else:
       raise MonitorError(f'We only supports a format of list/tuple/dict of '
                          f'"vars", while we got {type(monitors)}.')
@@ -160,7 +136,7 @@ def __init__(
 
   def _format_seq_monitors(self, monitors):
     if not isinstance(monitors, (tuple, list)):
-      raise TypeError(f'Must be a sequence, but we got {type(monitors)}')
+      raise TypeError(f'Must be a tuple/list, but we got {type(monitors)}')
     _monitors = []
     for mon in monitors:
       if isinstance(mon, str):
@@ -183,7 +159,40 @@ def _format_seq_monitors(self, monitors):
         raise MonitorError(f'We do not support monitor with {type(mon)}: {mon}')
     return _monitors
 
-  def _find_monitor_targets(self, _monitors):
+  def _format_dict_monitors(self, monitors):
+    if not isinstance(monitors, dict):
+      raise TypeError(f'Must be a dict, but we got {type(monitors)}')
+    _monitors = dict()
+    for key, val in monitors.items():
+      if not isinstance(key, str):
+        raise MonitorError('Expect the key of the dict "monitors" must be a string. But got '
+                           f'{type(key)}: {key}')
+      if isinstance(val, (bm.Variable, str)):
+        val = (val, None)
+
+      if isinstance(val, (tuple, list)):
+        if not isinstance(val[0], (bm.Variable, str)):
+          raise MonitorError('Expect the format of (variable, index) in the monitor setting. '
+                             f'But we got {val}')
+        if len(val) == 1:
+          _monitors[key] = (val[0], None)
+        elif len(val) == 2:
+          if isinstance(val[1], (int, np.integer)):
+            idx = bm.array([val[1]])
+          else:
+            idx = None if val[1] is None else bm.asarray(val[1])
+          _monitors[key] = (val[0], idx)
+        else:
+          raise MonitorError('Expect the format of (variable, index) in the monitor setting. '
+                             f'But we got {val}')
+      elif callable(val):
+        _monitors[key] = val
+      else:
+        raise MonitorError('The value of dict monitor expect a sequence with (variable, index) '
+                           f'or a callable function. But we got {val}')
+    return _monitors
+
+  def _find_seq_monitor_targets(self, _monitors):
     if not isinstance(_monitors, (tuple, list)):
       raise TypeError(f'Must be a sequence, but we got {type(_monitors)}')
     # get monitor targets
@@ -214,6 +223,43 @@ def _find_monitor_targets(self, _monitors):
           monitors[key] = (getattr(master, splits[-1]), index)
     return monitors
 
+  def _find_dict_monitor_targets(self, _monitors):
+    if not isinstance(_monitors, dict):
+      raise TypeError(f'Must be a dict, but we got {type(_monitors)}')
+    # get monitor targets
+    monitors = {}
+    name2node = None
+    for _key, _mon in _monitors.items():
+      if isinstance(_mon, str):
+        if name2node is None:
+          name2node = {node.name: node for node in list(self.target.nodes(level=-1).unique().values())}
+
+        key, index = _mon[0], _mon[1]
+        splits = key.split('.')
+        if len(splits) == 1:
+          if not hasattr(self.target, splits[0]):
+            raise RunningError(f'{self.target} does not has variable {key}.')
+          monitors[key] = (getattr(self.target, splits[-1]), index)
+        else:
+          if not hasattr(self.target, splits[0]):
+            if splits[0] not in name2node:
+              raise MonitorError(f'Cannot find target {key} in monitor of {self.target}, please check.')
+            else:
+              master = name2node[splits[0]]
+              assert len(splits) == 2
+              monitors[key] = (getattr(master, splits[-1]), index)
+          else:
+            master = self.target
+            for s in splits[:-1]:
+              try:
+                master = getattr(master, s)
+              except KeyError:
+                raise MonitorError(f'Cannot find {key} in {master}, please check.')
+            monitors[key] = (getattr(master, splits[-1]), index)
+      else:
+        monitors[_key] = _mon
+    return monitors
+
   def __del__(self):
     if hasattr(self, 'mon'):
       for key in tuple(self.mon.keys()):

brainpy/train/back_propagation.py

@@ -112,7 +112,8 @@ def __init__(
       lr = optim.ExponentialDecay(lr=0.025, decay_steps=1, decay_rate=0.99975)
       optimizer = optim.Adam(lr=lr)
     self.optimizer: optim.Optimizer = optimizer
-    self.optimizer.register_vars(self.target.vars(level=-1, include_self=True).subset(bm.TrainVar).unique())
+    if len(self.optimizer.vars_to_train) == 0:
+      self.optimizer.register_vars(self.target.vars(level=-1, include_self=True).subset(bm.TrainVar).unique())
 
     # loss function
     self.loss_has_aux = loss_has_aux
@@ -146,7 +147,7 @@ def __repr__(self):
             f'{prefix}loss={self._loss_func}, \n\t'
             f'{prefix}optimizer={self.optimizer})')
 
-  def get_hist_metric(self, phase='fit', metric='loss', which='detailed'):
+  def get_hist_metric(self, phase='fit', metric='loss', which='report'):
     """Get history losses."""
     assert phase in [c.FIT_PHASE, c.TEST_PHASE, c.TRAIN_PHASE, c.PREDICT_PHASE]
     assert which in ['report', 'detailed']
@@ -332,7 +333,7 @@ def fit(
             self.target.reset_state(self._get_input_batch_size(x))
             self.reset_state()
 
-          # training
+          # testing
           res = self._get_f_loss(shared_args)(x, y)
 
           # loss
@@ -406,7 +407,7 @@ def _get_f_loss(self, shared_args=None, jit=True) -> Callable:
       if self.jit[c.LOSS_PHASE] and jit:
         dyn_vars = self.target.vars()
         dyn_vars.update(self._dyn_vars)
-        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
+        dyn_vars.update(self.vars(level=0))
         self._f_loss_compiled[shared_args_str] = bm.jit(self._f_loss_compiled[shared_args_str],
                                                         dyn_vars=dyn_vars.unique())
     return self._f_loss_compiled[shared_args_str]
@@ -437,18 +438,23 @@ def _get_f_train(self, shared_args=None) -> Callable:
 
     shared_args_str = serialize_kwargs(shared_args)
     if shared_args_str not in self._f_fit_compiled:
-      self._f_fit_compiled[shared_args_str] = partial(self._step_func_train, shared_args)
+      self._f_fit_compiled[shared_args_str] = partial(self._step_func_fit, shared_args)
       if self.jit[c.FIT_PHASE]:
-        dyn_vars = self.vars().unique()
-        dyn_vars = dyn_vars - dyn_vars.subset(bm.VariableView)
+        dyn_vars = self.target.vars()
+        dyn_vars.update(self.optimizer.vars())
+        if isinstance(self._loss_func, bm.BrainPyObject):
+          dyn_vars.update(self._loss_func)
+        dyn_vars.update(self._dyn_vars)
+        dyn_vars.update(self.vars(level=0))
+        dyn_vars = dyn_vars.unique()
         self._f_fit_compiled[shared_args_str] = bm.jit(self._f_fit_compiled[shared_args_str],
                                                        dyn_vars=dyn_vars)
     return self._f_fit_compiled[shared_args_str]
 
   def _step_func_loss(self, shared_args, inputs, targets):
     raise NotImplementedError
 
-  def _step_func_train(self, shared_args, inputs, targets):
+  def _step_func_fit(self, shared_args, inputs, targets):
     raise NotImplementedError
 
 
@@ -508,7 +514,7 @@ def _step_func_loss(self, shared_args, inputs, targets):
     predicts = (outs, mons) if len(mons) > 0 else outs
     return self._loss_func(predicts, targets)
 
-  def _step_func_train(self, shared_args, inputs, targets):
+  def _step_func_fit(self, shared_args, inputs, targets):
     res = self._get_f_grad(shared_args)(inputs, targets)
     self.optimizer.update(res[0])
     return res[1:]
@@ -529,7 +535,7 @@ def _step_func_loss(self, shared_args, inputs, targets):
     loss = self._loss_func(outs, targets)
     return loss
 
-  def _step_func_train(self, shared_args, inputs, targets):
+  def _step_func_fit(self, shared_args, inputs, targets):
     res = self._get_f_grad(shared_args)(inputs, targets)
     self.optimizer.update(res[0])
     return res[1:]