latency-cell refactored and unit-tested

Author: Alexander Ororbia

ngclearn/components/input_encoders/latencyCell.py

@@ -1,28 +1,16 @@
-from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
-from ngclearn.utils import tensorstats
-from ngclearn.utils.model_utils import clamp_min, clamp_max
 from jax import numpy as jnp, random, jit
 from functools import partial
-from ngcsimlib.logger import info
-
-@jit
-def _update_times(t, s, tols):
-    """
-    Updates time-of-last-spike (tols) variable.
-
-    Args:
-        t: current time (a scalar/int value)
+from ngclearn.utils import tensorstats
+from ngcsimlib.deprecators import deprecate_args
+from ngcsimlib.logger import info, warn
 
-        s: binary spike vector
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
-        tols: current time-of-last-spike variable
+from ngclearn.utils.model_utils import clamp_min, clamp_max
 
-    Returns:
-        updated tols variable
-    """
-    _tols = (1. - s) * tols + (s * t)
-    return _tols
 
 @partial(jit, static_argnums=[5])
 def _calc_spike_times_linear(data, tau, thr, first_spk_t, num_steps=1.,
@@ -157,9 +145,10 @@ class LatencyCell(JaxComponent):
     """
 
     # Define Functions
-    def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0.,
-                 linearize=False, normalize=False, clip_spikes=False, num_steps=1.,
-                 batch_size=1, **kwargs):
+    def __init__(
+            self, name, n_units, tau=1., threshold=0.01, first_spike_time=0., linearize=False, normalize=False,
+            clip_spikes=False, num_steps=1., batch_size=1, **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         ## latency meta-parameters
@@ -186,9 +175,11 @@ def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0.,
         self.targ_sp_times = Compartment(restVals, display_name="Target Spike Time", units="ms")
         #self.reset()
 
+    @transition(output_compartments=["targ_sp_times", "clip_mask"])
     @staticmethod
-    def _calc_spike_times(linearize, tau, threshold, first_spike_time, num_steps,
-                          normalize, clip_spikes, inputs):
+    def calc_spike_times(
+            linearize, tau, threshold, first_spike_time, num_steps, normalize, clip_spikes, inputs
+    ):
         ## would call this function before processing a spike train (at start)
         data = inputs
         if clip_spikes:
@@ -208,42 +199,27 @@ def _calc_spike_times(linearize, tau, threshold, first_spike_time, num_steps,
             targ_sp_times = stimes #* calcEvent + targ_sp_times * (1. - calcEvent)
         return targ_sp_times, clip_mask
 
-    @resolver(_calc_spike_times)
-    def calc_spike_times(self, targ_sp_times, clip_mask):
-        self.targ_sp_times.set(targ_sp_times)
-        self.clip_mask.set(clip_mask)
-
+    @transition(output_compartments=["outputs", "tols", "mask", "targ_sp_times", "key"])
     @staticmethod
-    def _advance_state(t, dt, key, inputs, mask, clip_mask, targ_sp_times, tols):
+    def advance_state(t, dt, key, inputs, mask, clip_mask, targ_sp_times, tols):
         key, *subkeys = random.split(key, 2)
-        data = inputs ## get sensory pattern data / features
-        spikes, spk_mask = _extract_spike(targ_sp_times, t, mask) ## get spikes at t
-        tols = _update_times(t, spikes, tols)
+        data = inputs  ## get sensory pattern data / features
+        spikes, spk_mask = _extract_spike(targ_sp_times, t, mask)  ## get spikes at t
+
+        # Updates time-of-last-spike (tols) variable:
+        # output = s = binary spike vector
+        # tols = current time-of-last-spike variable
+        tols = (1. - spikes) * tols + (spikes * t)
+
         spikes = spikes * (1. - clip_mask)
         return spikes, tols, spk_mask, targ_sp_times, key
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, tols, mask, targ_sp_times, key):
-        self.outputs.set(outputs)
-        self.tols.set(tols)
-        self.mask.set(mask)
-        self.targ_sp_times.set(targ_sp_times)
-        self.key.set(key)
-
+    @transition(output_compartments=["inputs", "outputs", "tols", "mask", "clip_mask", "targ_sp_times"])
     @staticmethod
-    def _reset(batch_size, n_units):
+    def reset(batch_size, n_units):
         restVals = jnp.zeros((batch_size, n_units))
         return (restVals, restVals, restVals, restVals, restVals, restVals)
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, tols, mask, clip_mask, targ_sp_times):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.tols.set(tols)
-        self.mask.set(mask)
-        self.clip_mask.set(clip_mask)
-        self.targ_sp_times.set(targ_sp_times)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, key=self.key.value)

tests/components/input_encoders/test_latencyCell.py

@@ -0,0 +1,72 @@
+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 LatencyCell
+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_latencyCell():
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    T = 50  # 100 #5  ## number of simulation steps to run
+    dt = 1.  # 0.1  # ms ## compute integration time constant
+    tau = 1.
+    # ---- build a simple Poisson cell system ----
+    with Context("Circuit") as ctx:
+        a = LatencyCell(
+            "a", n_units=4, tau=tau, threshold=0.01, linearize=True,
+            normalize=True, num_steps=T, clip_spikes=False
+        )
+
+        ## create and compile core simulation commands
+        advance_process = (Process()
+                           >> a.advance_state)
+        ctx.wrap_and_add_command(advance_process.pure, name="advance")
+        calc_spike_times_process = (Process()
+                                    >> a.calc_spike_times)
+        ctx.wrap_and_add_command(calc_spike_times_process.pure, name="calc_spike_times")
+        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
+    inputs = jnp.asarray([[0.02, 0.5, 1., 0.0]])
+
+    targets = np.zeros((T, 4))
+    targets[0, 2] = 1.
+    targets[24, 1] = 1.
+    targets[48, 0] = 1.
+    targets[49, 3] = 1.
+    targets = jnp.array(targets) ## gold-standard solution to check against
+
+    outs = []
+    ctx.reset()
+    ctx.clamp(inputs)
+    ctx.calc_spike_times()
+    for ts in range(T):
+        ctx.clamp(inputs)
+        ctx.advance(t=ts * dt, dt=dt)
+        ## naively extract simple statistics at time ts and print them to I/O
+        s = a.outputs.value
+        outs.append(s)
+        #print(" {}: s {} ".format(ts, jnp.squeeze(s)))
+    outs = jnp.concatenate(outs, axis=0)
+
+    ## output should equal input
+    assert_array_equal(outs, targets)
+
+test_latencyCell()