revised and added unit-test for exp-kernel

Author: Alexander Ororbia

ngclearn/components/other/expKernel.py

@@ -1,8 +1,13 @@
-from jax import numpy as jnp, jit
-from functools import partial
-from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
+from jax import numpy as jnp, random, jit
+from functools import partial
 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=[5,6])
 def _apply_kernel(tf_curr, s, t, tau_w, win_len, krn_start, krn_end):
@@ -40,6 +45,8 @@ class ExpKernel(JaxComponent): ## exponential kernel
         nu: (ms, spike time interval for window)
 
         tau_w: spike window time constant (in micro-secs, or nano-s)
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     # Define Functions
@@ -60,31 +67,22 @@ def __init__(self, name, n_units, dt, tau_w=500., nu=4., batch_size=1, **kwargs)
         ## window of spike times
         self.tf = Compartment(jnp.zeros((self.win_len, self.batch_size, self.n_units)))
 
+    @transition(output_compartments=["epsp", "tf"])
     @staticmethod
-    def _advance_state(t, tau_w, win_len, inputs, tf):
+    def advance_state(t, tau_w, win_len, inputs, tf):
         s = inputs
         ## update spike time window and corresponding window volume
         tf, epsp = _apply_kernel(tf, s, t, tau_w, win_len, krn_start=0,
                                  krn_end=win_len-1) #0:win_len-1)
         return epsp, tf
 
-    @resolver(_advance_state)
-    def advance_state(self, epsp, tf):
-        self.epsp.set(epsp)
-        self.tf.set(tf)
-
+    @transition(output_compartments=["inputs", "epsp", "tf"])
     @staticmethod
-    def _reset(batch_size, n_units, win_len):
+    def reset(batch_size, n_units, win_len):
         restVals = jnp.zeros((batch_size, n_units))
         restTensor = jnp.zeros([win_len, batch_size, n_units], jnp.float32) # tf
         return restVals, restVals, restTensor # inputs, epsp, tf
 
-    @resolver(_reset)
-    def reset(self, inputs, epsp, tf):
-        self.inputs.set(inputs)
-        self.epsp.set(epsp)
-        self.tf.set(tf)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

tests/components/other/test_expKernel.py

@@ -0,0 +1,58 @@
+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 ExpKernel
+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_expKernel1():
+    ## 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 = ExpKernel(
+            name="a", n_units=1, dt=1., tau_w=500., nu=4., 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/epsp pulses
+    y_seq = jnp.asarray([[0., 1., 0.998002, 0.996008, 1.9940181]], dtype=jnp.float32)
+
+    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.epsp.value)
+    outs = jnp.concatenate(outs, axis=1)
+    #print(outs)
+
+    ## output should equal input
+    assert_array_equal(outs, y_seq)
+
+#test_expKernel1()