Merge branches 'major_release_update' and 'major_release_update' of github.com:NACLab/ngc-learn into major_release_update

Author: rxng8

docs/museum/snn_dc.md

@@ -209,7 +209,7 @@ To use your saved model and examine its performance on the MNIST test-set, you
 can execute the evaluation script like so:
 
 ```console
-$ python analyze_dcsnn.py --dataX=../data/mnist/testX.npy --sample_idx=0
+$ python analyze_dcsnn.py --dataX=../../data/mnist/testX.npy --sample_idx=0
 ```
 
 while will produce a visualization of your DC-SNN's receptive fields (the

ngclearn/__init__.py

@@ -37,6 +37,9 @@
 from ngcsimlib.resolver import resolver
 from ngcsimlib import utils as sim_utils
 
+from ngclearn.utils.jaxProcess import JaxProcess
+from ngcsimlib.compilers.process import transition, Process
+
 
 from ngcsimlib import configure, preload_modules
 from ngcsimlib import logger

ngclearn/components/base_monitor.py

@@ -1,10 +1,10 @@
 import json
 
-from ngclearn import Component, Compartment
+from ngclearn import Component, Compartment, transition
 from ngclearn import numpy as np
-#from ngcsimlib.utils import add_component_resolver, add_resolver_meta, \
 from ngcsimlib.utils import get_current_path
 from ngcsimlib.logger import warn, critical
+
 import matplotlib.pyplot as plt
 
 
@@ -44,9 +44,16 @@ class Base_Monitor(Component):
     """
     auto_resolve = False
 
+    @staticmethod
+    def build_reset(component):
+        return Base_Monitor.reset(component)
+
+    @staticmethod
+    def build_advance_state(component):
+        return Base_Monitor.record(component)
 
     @staticmethod
-    def build_advance(compartments):
+    def _record_internal(compartments):
         """
         A method to build the method to advance the stored values.
 
@@ -61,8 +68,9 @@ def build_advance(compartments):
             "monitor found in ngclearn.components or "
             "ngclearn.components.lava (If using lava)")
 
+    @transition(None, True)
     @staticmethod
-    def build_reset(component):
+    def reset(component):
         """
         A method to build the method to reset the stored values.
         Args:
@@ -87,15 +95,16 @@ def _reset(**kwargs):
         # pure func, output compartments, args, params, input compartments
         return _reset, output_compartments, [], [], output_compartments
 
+    @transition(None, True)
     @staticmethod
-    def build_advance_state(component):
+    def record(component):
         output_compartments = []
         compartments = []
         for comp in component.compartments:
             output_compartments.append(comp.split("/")[-1] + "*store")
             compartments.append(comp.split("/")[-1])
 
-        _advance = component.build_advance(compartments)
+        _advance = component._record_internal(compartments)
 
         return _advance, output_compartments, [], [], compartments + output_compartments
 
@@ -124,8 +133,15 @@ def watch(self, compartment, window_length):
         """
         cs, end = self._add_path(compartment.path)
 
-        dtype = compartment.value.dtype
-        shape = compartment.value.shape
+        if hasattr(compartment.value, "dtype"):
+            dtype = compartment.value.dtype
+        else:
+            dtype = type(compartment.value)
+
+        if hasattr(compartment.value, "shape"):
+            shape = compartment.value.shape
+        else:
+            shape = (1,)
         new_comp = Compartment(np.zeros(shape, dtype=dtype))
         new_comp_store = Compartment(np.zeros((window_length, *shape), dtype=dtype))
 
@@ -268,7 +284,7 @@ def load(self, directory, **kwargs):
                 setattr(self, compartment_path, new_comp)
                 self.compartments.append(new_comp.path)
 
-            self._update_resolver()
+            # self._update_resolver()
 
     def make_plot(self, compartment, ax=None, ylabel=None, xlabel=None, title=None, n=None, plot_func=None):
         vals = self.view(compartment)

ngclearn/components/monitor.py

@@ -1,4 +1,5 @@
 from ngclearn.components.base_monitor import Base_Monitor
+from ngclearn import transition
 
 class Monitor(Base_Monitor):
     """
@@ -8,9 +9,9 @@ class Monitor(Base_Monitor):
     auto_resolve = False
 
     @staticmethod
-    def build_advance(compartments):
+    def _record_internal(compartments):
         @staticmethod
-        def _advance(**kwargs):
+        def _record(**kwargs):
             return_vals = []
             for comp in compartments:
                 new_val = kwargs[comp]
@@ -19,7 +20,7 @@ def _advance(**kwargs):
                 current_store = current_store.at[-1].set(new_val)
                 return_vals.append(current_store)
             return return_vals if len(compartments) > 1 else return_vals[0]
-        return _advance
+        return _record
 
     @staticmethod
     def build_advance_state(component):

ngclearn/utils/jaxProcess.py

@@ -0,0 +1,20 @@
+from ngcsimlib.compilers.process import Process
+from jax.lax import scan as _scan
+
+
+class JaxProcess(Process):
+    def scan(self, xs, arg_order=None, compartments_to_monitor=None, save_state=True):
+        if compartments_to_monitor is None:
+            compartments_to_monitor = []
+        if arg_order is None:
+            arg_order = list(self.get_required_args())
+
+        def _pure(current_state, x):
+            v = self.pure(current_state, **{key: value for key, value in zip(arg_order, x)})
+            return v, [v[c.path] for c in compartments_to_monitor]
+
+        vals, stacked = _scan(_pure, init=self.get_required_state(include_special_compartments=True), xs=xs)
+        if save_state:
+            self.updated_modified_state(vals)
+        return stacked
+    

tests/components/input_encoders/test_bernoulliCell.py

@@ -24,11 +24,11 @@ def test_bernoulliCell1():
     with Context(name) as ctx:
         a = BernoulliCell(name="a", n_units=1, key=subkeys[0])
 
-        advance_process = (Process()
+        advance_process = (Process("advance_proc")
                            >> a.advance_state)
         ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
-        reset_process = (Process()
+        reset_process = (Process("reset_proc")
                          >> a.reset)
         ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

tests/components/input_encoders/test_latencyCell.py

@@ -29,13 +29,13 @@ def test_latencyCell1():
         )
 
         ## create and compile core simulation commands
-        advance_process = (Process()
+        advance_process = (Process("advance_proc")
                            >> a.advance_state)
         ctx.wrap_and_add_command(jit(advance_process.pure), name="advance")
-        calc_spike_times_process = (Process()
+        calc_spike_times_process = (Process("calc_sptimes_proc")
                                     >> a.calc_spike_times)
         ctx.wrap_and_add_command(jit(calc_spike_times_process.pure), name="calc_spike_times")
-        reset_process = (Process()
+        reset_process = (Process("reset_proc")
                          >> a.reset)
         ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

tests/components/input_encoders/test_phasorCell.py

@@ -24,11 +24,11 @@ def test_phasorCell1():
     with Context(name) as ctx:
         a = PhasorCell(name="a", n_units=1, target_freq=1000., disable_phasor=True, key=subkeys[0])
 
-        advance_process = (Process()
+        advance_process = (Process("advance_proc")
                            >> a.advance_state)
         ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
-        reset_process = (Process()
+        reset_process = (Process("reset_proc")
                          >> a.reset)
         ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

tests/components/input_encoders/test_poissonCell.py

@@ -24,11 +24,11 @@ def test_poissonCell1():
     with Context(name) as ctx:
         a = PoissonCell(name="a", n_units=1, target_freq=1000., key=subkeys[0])
 
-        advance_process = (Process()
+        advance_process = (Process("advance_proc")
                            >> a.advance_state)
         ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
-        reset_process = (Process()
+        reset_process = (Process("reset_proc")
                          >> a.reset)
         ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

tests/components/neurons/graded/test_RateCell.py

@@ -26,9 +26,9 @@ def test_RateCell1():
       threshold=("none", 0.), integration_type="euler",
       batch_size=1, resist_scale=1., shape=None, is_stateful=True
     )
-    advance_process = (Process() >> a.advance_state)
+    advance_process = (Process("advance_proc") >> a.advance_state)
     ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
-    reset_process = (Process() >> a.reset)
+    reset_process = (Process("reset_proc") >> a.reset)
     ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
     # reset_cmd, reset_args = ctx.compile_by_key(a, compile_key="reset")