revised slif cell w/ unit-test; needed mod to diffeq

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/sLIFCell.py

@@ -1,61 +1,15 @@
+from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, random, jit
 from functools import partial
-from ngclearn import resolver, Component, Compartment
-from ngclearn.components.jaxComponent import JaxComponent
+from ngclearn.utils import tensorstats
+from ngcsimlib.deprecators import deprecate_args
+from ngcsimlib.logger import info, warn
 from ngclearn.utils.diffeq.ode_utils import step_euler
 from ngclearn.utils.surrogate_fx import secant_lif_estimator
-from ngclearn.utils import tensorstats
-
-@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
-
-@partial(jit, static_argnums=[3,4])
-def _modify_current(j, spikes, inh_weights, R_m, inh_R):
-    """
-    A simple function that modifies electrical current j via application of a
-    scalar membrane resistance value and an approximate form of lateral inhibition.
-    Note that if no inhibitory resistance is set (i.e., inh_R = 0), then no
-    lateral inhibition is applied. Functionally, this routine carries out the
-    following piecewise equation:
-
-    | j * R_m - [Wi * s(t-dt)] * inh_R, if inh_R > 0
-    | j * R_m, otherwise
-
-    Args:
-        j: electrical current value
-
-        spikes: previous binary spike vector (for t-dt)
-
-        inh_weights: lateral recurrent inhibitory synapses (typically should be
-            chosen to be a scaled hollow matrix)
-
-        R_m: membrane resistance (to multiply/scale j by)
 
-        inh_R: inhibitory resistance to scale lateral inhibitory current by; if
-            inh_R = 0, NO lateral inhibitory pressure will be applied
-
-    Returns:
-        modified electrical current value
-    """
-    _j = j * R_m
-    if inh_R > 0.:
-        _j = _j - (jnp.matmul(spikes, inh_weights) * inh_R)
-    return _j
+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): ## raw voltage dynamics
@@ -97,6 +51,7 @@ def _update_refract_and_spikes(dt, rfr, s, refract_T, sticky_spikes=False):
         _s = s * mask + (1. - mask)
     return _rfr, _s
 
+@partial(jit, static_argnums=[6, 7, 8, 9, 10, 11])
 def _run_cell(dt, j, v, v_thr, tau_m, rfr, spike_fx, refract_T=1., thrGain=0.002,
               thrLeak=0.0005, rho_b = 0., sticky_spikes=False, v_min=None):
     """
@@ -206,6 +161,8 @@ class SLIFCell(JaxComponent): ## leaky integrate-and-fire cell
             a key setting used by Samadi et al., 2017
 
         thr_jitter: scale of uniform jitter to add to initialization of thresholds
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     # Define Functions
@@ -258,36 +215,44 @@ def __init__(self, name, n_units, tau_m, resist_m, thr, resist_inh=0.,
         self.rfr = Compartment(restVals + self.refract_T) ## refractory variable(s)
         self.surrogate = Compartment(restVals + 1.) ## surrogate signal
 
+    @transition(output_compartments=["j", "s", "tols", "v", "thr", "rfr", "surrogate"])
     @staticmethod
-    def _advance_state(t, dt, inh_weights, R_m, inh_R, d_spike_fx, tau_m,
-                       spike_fx, refract_T, thrGain, thrLeak, rho_b,
-                       sticky_spikes, v_min, j, s, v, thr, rfr, tols):
-        ## run one step of Euler integration over neuronal dynamics
-        j_curr = j
-        ## apply simplified inhibitory pressure
-        j_curr = _modify_current(j_curr, s, inh_weights, R_m, inh_R)
-        j = j_curr # None ## store electrical current
-        surrogate = d_spike_fx(j_curr, c1=0.82, c2=0.08)
+    def advance_state(
+            t, dt, inh_weights, R_m, inh_R, d_spike_fx, tau_m, spike_fx, refract_T, thrGain, 
+            thrLeak, rho_b, sticky_spikes, v_min, j, s, v, thr, rfr, tols
+    ):
+        #####################################################################################
+        #The following 3 lines of code modify electrical current j via application of a
+        #scalar membrane resistance value and an approximate form of lateral inhibition.
+        #Functionally, this routine carries out the following piecewise equation:
+        #| j * R_m - [Wi * s(t-dt)] * inh_R, if inh_R > 0
+        #| j * R_m, otherwise
+        #| where j: electrical current value, spikes: previous binary spike vector (for t-dt), 
+        #    inh_weights: lateral recurrent inhibitory synapses (typically should be chosen 
+        #    to be a scaled hollow matrix), 
+        #| R_m: membrane resistance (to multiply/scale j by), 
+        #| inh_R: inhibitory resistance to scale lateral inhibitory current by; if inh_R = 0, 
+        #    NO lateral inhibitory pressure will be applied
+        j = j * R_m
+        if inh_R > 0.: ## if inh_R > 0, then lateral inhibition is applied
+            j = j - (jnp.matmul(spikes, inh_weights) * inh_R)
+        #####################################################################################
+
+        surrogate = d_spike_fx(j, c1=0.82, c2=0.08)
+        #surrogate = d_spike_fx(j_curr, c1=0.82, c2=0.08)
+    
         v, s, thr, rfr = \
-            _run_cell(dt, j_curr, v, thr, tau_m,
+            _run_cell(dt, j, v, thr, tau_m, 
                       rfr, spike_fx, refract_T, thrGain, thrLeak,
                       rho_b, sticky_spikes=sticky_spikes, v_min=v_min)
+
         ## update tols
-        tols = _update_times(t, s, tols)
+        tols = (1. - s) * tols + (s * t)
         return j, s, tols, v, thr, rfr, surrogate
 
-    @resolver(_advance_state)
-    def advance_state(self, j, s, tols, v, thr, rfr, surrogate):
-        self.j.set(j)
-        self.s.set(s)
-        self.tols.set(tols)
-        self.thr.set(thr)
-        self.rfr.set(rfr)
-        self.surrogate.set(surrogate)
-        self.v.set(v)
-
+    @transition(output_compartments=["j", "s", "tols", "v", "thr", "rfr", "surrogate"])
     @staticmethod
-    def _reset(refract_T, thr_persist, threshold0, batch_size, n_units, thr):
+    def reset(refract_T, thr_persist, threshold0, batch_size, n_units, thr):
         restVals = jnp.zeros((batch_size, n_units))
         voltage = restVals
         refract = restVals + refract_T
@@ -299,16 +264,6 @@ def _reset(refract_T, thr_persist, threshold0, batch_size, n_units, thr):
             thr = threshold0 + 0
         return current, spikes, timeOfLastSpike, voltage, thr, refract, surrogate
 
-    @resolver(_reset)
-    def reset(self, j, s, tols, v, thr, rfr, surrogate):
-        self.j.set(j)
-        self.s.set(s)
-        self.tols.set(tols)
-        self.thr.set(thr)
-        self.rfr.set(rfr)
-        self.surrogate.set(surrogate)
-        self.v.set(v)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         if self.thr_persist == False:

ngclearn/utils/diffeq/ode_utils.py

@@ -57,7 +57,7 @@ def _step_forward(t, x, dx_dt, dt, x_scale): ## internal step co-routine
     _x = x * x_scale + dx_dt * dt
     return _t, _x
 
-
+@partial(jit, static_argnums=(2, 3, 4, 5, ))
 def step_euler(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the Euler method, i.e., a

tests/components/neurons/spiking/test_sLIFCell.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 SLIFCell
+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_sLIFCell1():
+    ## 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 = SLIFCell(
+            name="a", n_units=1, tau_m=50., resist_m=10., thr=0.3, key=subkeys[0]
+        )
+
+        #"""
+        advance_process = (Process()
+                           >> a.advance_state)
+        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., 0., 0., 1., 1., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    y_seq = jnp.asarray([[0., 1., 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_sLIFCell1()