revised if-cell w/ unit-test

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/IFCell.py

@@ -1,31 +1,19 @@
+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 ngcsimlib.deprecators import deprecate_args
-from ngclearn import resolver, Component, Compartment
-from ngclearn.components.jaxComponent import JaxComponent
+from ngcsimlib.logger import info, warn
 from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \
                                             step_euler, step_rk2
-from ngclearn.utils.surrogate_fx import (arctan_estimator,
+from ngclearn.utils.surrogate_fx import (secant_lif_estimator, arctan_estimator,
                                          triangular_estimator,
                                          straight_through_estimator)
 
-@jit
-def _update_times(t, s, tols):
-    """
-    Updates time-of-last-spike (tols) variable.
-
-    Args:
-        t: current time (a scalar/int value)
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
-        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 _dfv_internal(j, v, rfr, tau_m, refract_T): ## raw voltage dynamics
@@ -166,32 +154,26 @@ def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
                                 units="ms") ## time-of-last-spike
         self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value")
 
+    @transition(output_compartments=["v", "s", "rfr", "tols", "key", "surrogate"])
     @staticmethod
-    def _advance_state(t, dt, tau_m, resist_m, v_rest, v_reset, refract_T,
-                       thr, lower_clamp_voltage, intgFlag, d_spike_fx, key,
-                       j, v, rfr, tols):
+    def advance_state(
+            t, dt, tau_m, resist_m, v_rest, v_reset, refract_T, thr, lower_clamp_voltage, intgFlag, d_spike_fx, key,
+            j, v, rfr, tols
+    ):
         ## run one integration step for neuronal dynamics
         j = j * resist_m
         v, s, rfr = _run_cell(dt, j, v, thr, rfr, tau_m, v_rest, v_reset,
                               refract_T, intgFlag)
         surrogate = d_spike_fx(v, thr)
         ## update tols
-        tols = _update_times(t, s, tols)
+        tols = (1. - s) * tols + (s * t)
         if lower_clamp_voltage: ## ensure voltage never < v_rest
             v = jnp.maximum(v, v_rest)
         return v, s, rfr, tols, key, surrogate
 
-    @resolver(_advance_state)
-    def advance_state(self, v, s, rfr, tols, key, surrogate):
-        self.v.set(v)
-        self.s.set(s)
-        self.rfr.set(rfr)
-        self.tols.set(tols)
-        self.key.set(key)
-        self.surrogate.set(surrogate)
-
+    @transition(output_compartments=["j", "v", "s", "rfr", "tols", "surrogate"])
     @staticmethod
-    def _reset(batch_size, n_units, v_rest, refract_T):
+    def reset(batch_size, n_units, v_rest, refract_T):
         restVals = jnp.zeros((batch_size, n_units))
         j = restVals #+ 0
         v = restVals + v_rest
@@ -201,15 +183,6 @@ def _reset(batch_size, n_units, v_rest, refract_T):
         surrogate = restVals + 1.
         return j, v, s, rfr, tols, surrogate
 
-    @resolver(_reset)
-    def reset(self, j, v, s, rfr, tols, surrogate):
-        self.j.set(j)
-        self.v.set(v)
-        self.s.set(s)
-        self.rfr.set(rfr)
-        self.tols.set(tols)
-        self.surrogate.set(surrogate)
-
     def save(self, directory, **kwargs):
         ## do a protected save of constants, depending on whether they are floats or arrays
         tau_m = (self.tau_m if isinstance(self.tau_m, float)

tests/components/neurons/spiking/test_IFCell.py

@@ -0,0 +1,69 @@
+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 IFCell
+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_IFCell1():
+    name = "if_ctx"
+    ## 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(name) as ctx:
+        a = IFCell(
+            name="a", n_units=1, tau_m=5., resist_m=10., 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([[1., 1., 1., 1., 1., 0., 0., 1., 1., 1., 1., 1., 1., 1., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    y_seq = jnp.asarray([[0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0.]], 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.s.value)
+    outs = jnp.concatenate(outs, axis=1)
+    print(outs)
+    
+    ## output should equal input
+    assert_array_equal(outs, y_seq)
+
+#test_IFCell1()