revised and add unit-test for varTrace

Author: Alexander Ororbia

ngclearn/components/other/varTrace.py

@@ -1,8 +1,13 @@
+from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, random, jit
 from functools import partial
-from ngclearn import resolver, Component, Compartment
-from ngclearn.components.jaxComponent import JaxComponent
 from ngclearn.utils import tensorstats
+from ngcsimlib.deprecators import deprecate_args
+from ngcsimlib.logger import info, warn
+
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
 @partial(jit, static_argnums=[4])
 def _run_varfilter(dt, x, x_tr, decayFactor, gamma_tr, a_delta=0.):
@@ -54,13 +59,17 @@ class VarTrace(JaxComponent): ## low-pass filter
         a_delta: value to increment a trace by in presence of a spike; note if set
             to a value <= 0, then a piecewise gated trace will be used instead
 
+        gamma_tr: an extra multiplier in front of the leak of the trace (Default: 1)
+
         decay_type: string indicating the decay type to be applied to ODE
             integration; low-pass filter configuration
 
             :Note: string values that this can be (Default: "exp") are:
                 1) `'lin'` = linear trace filter, i.e., decay = x_tr + (-x_tr) * (dt/tau_tr);
                 2) `'exp'` = exponential trace filter, i.e., decay = exp(-dt/tau_tr) * x_tr;
                 3) `'step'` = step trace, i.e., decay = 0 (a pulse applied upon input value)
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     # Define Functions
@@ -83,38 +92,28 @@ def __init__(self, name, n_units, tau_tr, a_delta, gamma_tr=1, decay_type="exp",
         self.outputs = Compartment(restVals) # output compartment
         self.trace = Compartment(restVals)
 
+    @transition(output_compartments=["outputs", "trace"])
     @staticmethod
-    def _advance_state(dt, decay_type, tau_tr, a_delta, gamma_tr, inputs, trace):
+    def advance_state(dt, decay_type, tau_tr, a_delta, gamma_tr, inputs, trace):
         decayFactor = 0.
         if "exp" in decay_type:
             decayFactor = jnp.exp(-dt/tau_tr)
         elif "lin" in decay_type:
             decayFactor = (1. - dt/tau_tr)
-
         _x_tr = gamma_tr * trace * decayFactor
         if a_delta > 0.:
             _x_tr = _x_tr + inputs * a_delta
         else:
             _x_tr = _x_tr * (1. - inputs) + inputs
-
+        trace = _x_tr
         return trace, trace
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, trace):
-        self.outputs.set(outputs)
-        self.trace.set(trace)
-
+    @transition(output_compartments=["inputs", "outputs", "trace"])
     @staticmethod
-    def _reset(batch_size, n_units):
+    def reset(batch_size, n_units):
         restVals = jnp.zeros((batch_size, n_units))
         return restVals, restVals, restVals
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, trace):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.trace.set(trace)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

tests/components/other/test_varTrace.py

@@ -0,0 +1,60 @@
+from jax import numpy as jnp, random, jit
+from ngcsimlib.context import Context
+import numpy as np
+np.random.seed(42)
+from ngclearn.components import VarTrace
+from ngcsimlib.compilers import compile_command, wrap_command
+from numpy.testing import assert_array_equal
+
+from ngcsimlib.compilers.process import Process, transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.context import Context
+from ngcsimlib.utils.compartment import Get_Compartment_Batch
+
+
+def test_varTrace1():
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    dt = 1.  # ms
+    trace_increment = 0.1
+    # ---- build a simple Poisson cell system ----
+    with Context("Circuit") as ctx:
+        a = VarTrace(
+            name="a", n_units=1, a_delta=trace_increment, decay_type="step", tau_tr=1., 
+            key=subkeys[0]
+        )
+
+        advance_process = (Process()
+                           >> a.advance_state)
+        ctx.wrap_and_add_command(advance_process.pure, name="run")
+
+        reset_process = (Process()
+                         >> a.reset)
+        ctx.wrap_and_add_command(reset_process.pure, name="reset")
+
+        ## set up non-compiled utility commands
+        @Context.dynamicCommand
+        def clamp(x):
+            a.inputs.set(x)
+
+    ## input spike train
+    x_seq = jnp.asarray([[1., 1., 0., 0., 1.]], dtype=jnp.float32)
+    ## desired output pulses
+    y_seq = x_seq * trace_increment 
+
+    outs = []
+    ctx.reset()
+    for ts in range(x_seq.shape[1]):
+        x_t = jnp.array([[x_seq[0, ts]]])  ## get data at time t
+        ctx.clamp(x_t)
+        ctx.run(t=ts * 1., dt=dt)
+        outs.append(a.outputs.value)
+    outs = jnp.concatenate(outs, axis=1)
+    #print(outs)
+
+    ## output should equal input
+    assert_array_equal(outs, y_seq)
+
+#test_varTrace1()