revised lif-cell w/ unit-test

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/LIFCell.py

@@ -1,32 +1,31 @@
+"""
 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 ngclearn.utils.diffeq.ode_utils import get_integrator_code, \
+                                            step_euler, step_rk2
+from ngclearn.utils.surrogate_fx import (secant_lif_estimator, arctan_estimator,
+                                         triangular_estimator,
+                                         straight_through_estimator)
+"""
+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 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 (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)
-
-        s: binary spike vector
-
-        tols: current time-of-last-spike variable
-
-    Returns:
-        updated tols variable
-    """
-    _tols = (1. - s) * tols + (s * t)
-    return _tols
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
 @jit
 def _dfv_internal(j, v, rfr, tau_m, refract_T, v_rest, v_decay=1.): ## raw voltage dynamics
@@ -41,37 +40,6 @@ def _dfv(t, v, params): ## voltage dynamics wrapper
     dv_dt = _dfv_internal(j, v, rfr, tau_m, refract_T, v_rest, v_decay)
     return dv_dt
 
-#@partial(jit, static_argnums=[7, 8, 9, 10, 11, 12])
-def _run_cell(dt, j, v, v_thr, v_theta, rfr, skey, tau_m, v_rest, v_reset,
-              v_decay, refract_T, integType=0):
-    ### Runs leaky integrator (or leaky integrate-and-fire; LIF) neuronal dynamics.
-    _v_thr = v_theta + v_thr ## calc present voltage threshold
-    #mask = (rfr >= refract_T).astype(jnp.float32) # get refractory mask
-    ## update voltage / membrane potential
-    v_params = (j, rfr, tau_m, refract_T, v_rest, v_decay)
-    if integType == 1:
-        _, _v = step_rk2(0., v, _dfv, dt, v_params)
-    else: #_v = v + (v_rest - v) * (dt/tau_m) + (j * mask)
-        _, _v = step_euler(0., v, _dfv, dt, v_params)
-    ## obtain action potentials/spikes
-    s = (_v > _v_thr).astype(jnp.float32)
-    ## update refractory variables
-    _rfr = (rfr + dt) * (1. - s)
-    ## perform hyper-polarization of neuronal cells
-    _v = _v * (1. - s) + s * v_reset
-
-    raw_s = s + 0 ## preserve un-altered spikes
-    ############################################################################
-    ## this is a spike post-processing step
-    if skey is not None:
-        m_switch = (jnp.sum(s) > 0.).astype(jnp.float32) ## TODO: not batch-able
-        rS = s * random.uniform(skey, s.shape)
-        rS = nn.one_hot(jnp.argmax(rS, axis=1), num_classes=s.shape[1],
-                        dtype=jnp.float32)
-        s = s * (1. - m_switch) + rS * m_switch
-    ############################################################################
-    return _v, s, raw_s, _rfr
-
 #@partial(jit, static_argnums=[3, 4])
 def _update_theta(dt, v_theta, s, tau_theta, theta_plus=0.05):
     ### Runs homeostatic threshold update dynamics one step (via Euler integration).
@@ -159,7 +127,7 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
 
         lower_clamp_voltage: if True, this will ensure voltage never is below
             the value of `v_rest` (default: True)
-    """
+    """ ## batch_size arg?
 
     @deprecate_args(thr_jitter=None)
     def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
@@ -220,41 +188,61 @@ 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", "s_raw", "rfr", "thr_theta", "tols", "key", "surrogate"])    
     @staticmethod
-    def _advance_state(t, dt, tau_m, resist_m, v_rest, v_reset, v_decay, refract_T,
-                       thr, tau_theta, theta_plus, one_spike, lower_clamp_voltage,
-                       intgFlag, d_spike_fx, key, j, v, rfr, thr_theta, tols):
+    def advance_state(
+            t, dt, tau_m, resist_m, v_rest, v_reset, v_decay, refract_T, thr, tau_theta, theta_plus, 
+            one_spike, lower_clamp_voltage, intgFlag, d_spike_fx, key, j, v, rfr, thr_theta, tols
+    ):
         skey = None ## this is an empty dkey if single_spike mode turned off
         if one_spike:
             key, skey = random.split(key, 2)
         ## run one integration step for neuronal dynamics
         j = j * resist_m
-        v, s, raw_spikes, rfr = _run_cell(dt, j, v, thr, thr_theta, rfr, skey,
-                                          tau_m, v_rest, v_reset, v_decay,
-                                          refract_T, intgFlag)
-        surrogate = d_spike_fx(v, thr + thr_theta)
+        ############################################################################
+        ### Runs leaky integrator (leaky integrate-and-fire; LIF) neuronal dynamics.
+        _v_thr = thr_theta + thr #v_theta + v_thr ## calc present voltage threshold
+        #mask = (rfr >= refract_T).astype(jnp.float32) # get refractory mask
+        ## update voltage / membrane potential
+        v_params = (j, rfr, tau_m, refract_T, v_rest, v_decay)
+        if intgFlag == 1:
+            _, _v = step_rk2(0., v, _dfv, dt, v_params)
+        else: #_v = v + (v_rest - v) * (dt/tau_m) + (j * mask)
+            _, _v = step_euler(0., v, _dfv, dt, v_params)
+        ## obtain action potentials/spikes
+        s = (_v > _v_thr).astype(jnp.float32)
+        ## update refractory variables
+        _rfr = (rfr + dt) * (1. - s)
+        ## perform hyper-polarization of neuronal cells
+        _v = _v * (1. - s) + s * v_reset
+
+        raw_s = s + 0 ## preserve un-altered spikes
+        ############################################################################
+        ## this is a spike post-processing step
+        if skey is not None:
+            m_switch = (jnp.sum(s) > 0.).astype(jnp.float32) ## TODO: not batch-able
+            rS = s * random.uniform(skey, s.shape)
+            rS = nn.one_hot(jnp.argmax(rS, axis=1), num_classes=s.shape[1],
+                            dtype=jnp.float32)
+            s = s * (1. - m_switch) + rS * m_switch
+        ############################################################################
+        raw_spikes = raw_s
+        v = _v
+        rfr = _rfr
+
+        surrogate = d_spike_fx(v, _v_thr) #d_spike_fx(v, thr + thr_theta)
         if tau_theta > 0.:
             ## run one integration step for threshold dynamics
             thr_theta = _update_theta(dt, thr_theta, raw_spikes, tau_theta, theta_plus)
         ## 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, raw_spikes, rfr, thr_theta, tols, key, surrogate
 
-    @resolver(_advance_state)
-    def advance_state(self, v, s, s_raw, rfr, thr_theta, tols, key, surrogate):
-        self.v.set(v)
-        self.s.set(s)
-        self.s_raw.set(s_raw)
-        self.rfr.set(rfr)
-        self.thr_theta.set(thr_theta)
-        self.tols.set(tols)
-        self.key.set(key)
-        self.surrogate.set(surrogate)
-
+    @transition(output_compartments=["j", "v", "s", "s_raw", "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
@@ -266,16 +254,6 @@ def _reset(batch_size, n_units, v_rest, refract_T):
         surrogate = restVals + 1.
         return j, v, s, s_raw, rfr, tols, surrogate
 
-    @resolver(_reset)
-    def reset(self, j, v, s, s_raw, rfr, tols, surrogate):
-        self.j.set(j)
-        self.v.set(v)
-        self.s.set(s)
-        self.s_raw.set(s_raw)
-        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_LIFCell.py

@@ -0,0 +1,68 @@
+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 LIFCell
+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_LIFCell1():
+    ## 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 = LIFCell(
+            name="a", n_units=1, tau_m=5., resist_m=30., 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., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    y_seq = jnp.asarray([[0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 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_LIFCell1()