feat: support "func_monitors" in DSRunner

brainpy · May 1, 2022 · e8394df · e8394df
1 parent 1604fde
commit e8394df
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/brainpy/dyn/runners.py b/brainpy/dyn/runners.py
@@ -163,6 +163,7 @@ def func(_t, _dt):
              for k, v in return_without_idx.items()}
       res.update({k: (v.flatten()[idx] if bm.ndim(v) > 1 else v[idx])
                   for k, (v, idx) in return_with_idx.items()})
+      res.update({k: f(_t, _dt) for k, f in self.fun_monitors.items()})
       return res
 
     return func

diff --git a/brainpy/running/runner.py b/brainpy/running/runner.py
@@ -54,6 +54,9 @@ def __init__(self, target, monitors=None, fun_monitors=None,
       fun_monitors = dict()
     check_dict_data(fun_monitors, key_type=str, val_type=types.FunctionType)
     self.fun_monitors = fun_monitors
+    for key in self.fun_monitors.keys():
+      self.mon.item_names.append(key)
+      self.mon.item_contents[key] = []
 
     # progress bar
     assert isinstance(progress_bar, bool), 'Must be a boolean variable.'