revised wtas-cell w/ unit test

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/WTASCell.py

@@ -1,64 +1,16 @@
 from jax import numpy as jnp, random, jit, nn
-from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
+from jax import numpy as jnp, random, jit, nn
+from functools import partial
 from ngclearn.utils import tensorstats
-from ngclearn.utils.model_utils import softmax
-
-@jit
-def _update_times(t, s, tols):
-    """
-    Updates time-of-last-spike (tols) variable.
-
-    Args:
-        t: current time (a scalar/int value)
-
-        s: binary spike vector
-
-        tols: current time-of-last-spike variable
-
-    Returns:
-        updated tols variable
-    """
-    _tols = (1. - s) * tols + (s * t)
-    return _tols
-@jit
-def _run_cell(dt, j, v, rfr, v_thr, tau_m, R_m, thr_gain=0.002, refract_T=0.):
-    """
-    Runs leaky integrator neuronal dynamics
-
-    Args:
-        dt: integration time constant (milliseconds, or ms)
-
-        j: electrical current value
-
-        v: membrane potential (voltage, in milliVolts or mV) value (at t)
-
-        rfr: refractory variable vector (one per neuronal cell)
-
-        v_thr: base voltage threshold value (in mV)
-
-        tau_m: cell membrane time constant
+from ngcsimlib.deprecators import deprecate_args
+from ngcsimlib.logger import info, warn
 
-        R_m: cell membrane resistance
-
-        thr_gain: increment to be applied to threshold upon spike occurrence
-
-        refract_T: (relative) refractory time period (in ms; Default
-            value is 1 ms)
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+from ngclearn.utils.model_utils import softmax
 
-    Returns:
-        voltage(t+dt), spikes, updated voltage thresholds, updated refactory variables
-    """
-    mask = (rfr >= refract_T).astype(jnp.float32) ## check refractory period
-    v = (j * R_m) * mask
-    vp = softmax(v) # convert to Categorical (spike) probabilities
-    #s = nn.one_hot(jnp.argmax(vp, axis=1), j.shape[1]) ## hard-max spike
-    s = (vp > v_thr).astype(jnp.float32) ## calculate action potential
-    q = 1. ## Note: thr_gain ==> "rho_b"
-    dthr = jnp.sum(s, axis=1, keepdims=True) - q
-    v_thr = jnp.maximum(v_thr + dthr * thr_gain, 0.025) ## calc new threshold
-    rfr = (rfr + dt) * (1. - s) + s * dt # set refract to dt
-    return v, s, v_thr, rfr
 
 class WTASCell(JaxComponent): ## winner-take-all spiking cell
     """
@@ -136,22 +88,26 @@ def __init__(
         self.rfr = Compartment(restVals + self.refract_T)
         self.tols = Compartment(restVals) ## time-of-last-spike
 
+    @transition(output_compartments=["v", "s", "thr", "rfr", "tols"])
     @staticmethod
-    def _advance_state(t, dt, tau_m, R_m, thr_gain, refract_T, j, v, thr, rfr, tols):
-        v, s, thr, rfr = _run_cell(dt, j, v, rfr, thr, tau_m, R_m, thr_gain, refract_T)
-        tols = _update_times(t, s, tols) ## update tols
+    def advance_state(t, dt, tau_m, R_m, thr_gain, refract_T, j, v, thr, rfr, tols):
+        mask = (rfr >= refract_T) * 1.  ## check refractory period
+        v = (j * R_m) * mask
+        vp = softmax(v)  # convert to Categorical (spike) probabilities
+        # s = nn.one_hot(jnp.argmax(vp, axis=1), j.shape[1]) ## hard-max spike
+        s = (vp > thr) * 1. ## calculate action potential
+        q = 1.  ## Note: thr_gain ==> "rho_b"
+        ## increment threshold upon spike(s) occurrence
+        dthr = jnp.sum(s, axis=1, keepdims=True) - q
+        thr = jnp.maximum(thr + dthr * thr_gain, 0.025)  ## calc new threshold
+        rfr = (rfr + dt) * (1. - s) + s * dt  # set refract to dt
+
+        tols = (1. - s) * tols + (s * t) ## update tols
         return v, s, thr, rfr, tols
 
-    @resolver(_advance_state)
-    def advance_state(self, v, s, thr, rfr, tols):
-        self.v.set(v)
-        self.s.set(s)
-        self.thr.set(thr)
-        self.rfr.set(rfr)
-        self.tols.set(tols)
-
+    @transition(output_compartments=["j", "v", "s", "rfr", "tols"])
     @staticmethod
-    def _reset(batch_size, n_units, refract_T):
+    def reset(batch_size, n_units, refract_T):
         restVals = jnp.zeros((batch_size, n_units))
         j = restVals #+ 0
         v = restVals #+ 0
@@ -160,14 +116,6 @@ def _reset(batch_size, n_units, refract_T):
         tols = restVals #+ 0
         return j, v, s, rfr, tols
 
-    @resolver(_reset)
-    def reset(self, j, v, s, rfr, tols):
-        self.j.set(j)
-        self.v.set(v)
-        self.s.set(s)
-        self.rfr.set(rfr)
-        self.tols.set(tols)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, threshold=self.thr.value)

tests/components/neurons/spiking/test_WTASCell.py

@@ -0,0 +1,70 @@
+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 WTASCell
+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_WTASCell1():
+    name = "wtas_ctx"
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    dt = 1.  # ms
+    # ---- build a simple Poisson cell system ----
+    with Context(name) as ctx:
+        a = WTASCell(
+            name="a", n_units=1, tau_m=25., resist_m=1., key=subkeys[0]
+        )
+
+        #"""
+        advance_process = (Process()
+                           >> a.advance_state)
+        # ctx.wrap_and_add_command(advance_process.pure, name="run")
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
+
+        reset_process = (Process()
+                         >> 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")
+        ctx.add_command(wrap_command(jit(ctx.reset)), name="reset")
+        advance_cmd, advance_args = ctx.compile_by_key(a, compile_key="advance_state")
+        ctx.add_command(wrap_command(jit(ctx.advance_state)), name="run")
+        """
+
+        ## set up non-compiled utility commands
+        @Context.dynamicCommand
+        def clamp(x):
+            a.j.set(x)
+
+    ## input spike train
+    x_seq = jnp.asarray([[0., 1.], [0., 1.], [1., 0.], [1., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    y_seq = x_seq
+
+    outs = []
+    ctx.reset()
+    for ts in range(x_seq.shape[0]):
+        x_t = x_seq[ts:ts+1, :]  ## get data at time t
+        ctx.clamp(x_t)
+        ctx.run(t=ts * 1., dt=dt)
+        outs.append(a.s.value)
+    outs = jnp.concatenate(outs, axis=0)
+    #print(outs)
+    #exit()
+    ## output should equal input
+    assert_array_equal(outs, y_seq)
+
+#test_WTASCell1()