revised adex-cell w/ unit test, minor cleanup of quad-lif

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/adExCell.py

@@ -1,28 +1,14 @@
-from jax import numpy as jnp, jit
-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 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
-
-@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, w, tau_m, v_rest, sharpV, vT, R_m): ## raw voltage dynamics
@@ -46,30 +32,6 @@ def _dfw(t, w, params): ## recovery dynamics wrapper
     dv_dt = _dfw_internal(j, v, w, a, tau_m, v_rest)
     return dv_dt
 
-@jit
-def _emit_spike(v, v_thr):
-    s = (v > v_thr).astype(jnp.float32)
-    return s
-
-#@partial(jit, static_argnums=[10])
-def _run_cell(dt, j, v, w, v_thr, tau_m, tau_w, a, b, sharpV, vT,
-              v_rest, v_reset, R_m, integType=0):
-    if integType == 1: ## RK-2/midpoint
-        v_params = (j, w, tau_m, v_rest, sharpV, vT, R_m)
-        _, _v = step_rk2(0., v, _dfv, dt, v_params)
-        w_params = (j, v, a, tau_w, v_rest)
-        _, _w = step_rk2(0., w, _dfw, dt, w_params)
-    else: # integType == 0 (default -- Euler)
-        v_params = (j, w, tau_m, v_rest, sharpV, vT, R_m)
-        _, _v = step_euler(0., v, _dfv, dt, v_params)
-        w_params = (j, v, a, tau_w, v_rest)
-        _, _w = step_euler(0., w, _dfw, dt, w_params)
-    s = _emit_spike(_v, v_thr)
-    ## hyperpolarize/reset/snap variables
-    _v = _v * (1. - s) + s * v_reset
-    _w = _w * (1. - s) + s * (_w + b)
-    return _v, _w, s
-
 class AdExCell(JaxComponent):
     """
     The AdEx (adaptive exponential leaky integrate-and-fire) neuronal cell
@@ -136,10 +98,10 @@ class AdExCell(JaxComponent):
     """
 
     @deprecate_args(v_thr="thr")
-    def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400.,
-                 v_sharpness=2., intrinsic_mem_thr=-55., thr=5., v_rest=-72.,
-                 v_reset=-75., a=0.1, b=0.75, v0=-70., w0=0.,
-                 integration_type="euler", batch_size=1, **kwargs):
+    def __init__(
+            self, name, n_units, tau_m=15., resist_m=1., tau_w=400., v_sharpness=2., intrinsic_mem_thr=-55., thr=5.,
+            v_rest=-72., v_reset=-75., a=0.1, b=0.75, v0=-70., w0=0., integration_type="euler", batch_size=1, **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         ## Integration properties
@@ -174,24 +136,32 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400.,
         self.tols = Compartment(restVals, display_name="Time-of-Last-Spike",
                                 units="ms") ## time-of-last-spike
 
+    @transition(output_compartments=["j", "v", "w", "s", "tols"])
     @staticmethod
-    def _advance_state(t, dt, tau_m, R_m, tau_w, thr, a, b, sharpV, vT,
-                     v_rest, v_reset, intgFlag, j, v, w, tols):
-        v, w, s = _run_cell(dt, j, v, w, thr, tau_m, tau_w, a, b, sharpV, vT,
-                            v_rest, v_reset, R_m, intgFlag)
-        tols = _update_times(t, s, tols)
+    def advance_state(
+            t, dt, tau_m, R_m, tau_w, thr, a, b, sharpV, vT, v_rest, v_reset, intgFlag, j, v, w, tols
+    ):
+        if intgFlag == 1:  ## RK-2/midpoint
+            v_params = (j, w, tau_m, v_rest, sharpV, vT, R_m)
+            _, _v = step_rk2(0., v, _dfv, dt, v_params)
+            w_params = (j, v, a, tau_w, v_rest)
+            _, _w = step_rk2(0., w, _dfw, dt, w_params)
+        else:  # intgFlag == 0 (default -- Euler)
+            v_params = (j, w, tau_m, v_rest, sharpV, vT, R_m)
+            _, _v = step_euler(0., v, _dfv, dt, v_params)
+            w_params = (j, v, a, tau_w, v_rest)
+            _, _w = step_euler(0., w, _dfw, dt, w_params)
+        s = (_v > thr) * 1. ## emit spikes/pulses
+        ## hyperpolarize/reset/snap variables
+        v = _v * (1. - s) + s * v_reset
+        w = _w * (1. - s) + s * (_w + b)
+
+        tols = (1. - s) * tols + (s * t) ## update time-of-last spike variable(s)
         return j, v, w, s, tols
 
-    @resolver(_advance_state)
-    def advance_state(self, j, v, w, s, tols):
-        self.j.set(j)
-        self.w.set(w)
-        self.v.set(v)
-        self.s.set(s)
-        self.tols.set(tols)
-
+    @transition(output_compartments=["j", "v", "w", "s", "tols"])
     @staticmethod
-    def _reset(batch_size, n_units, v0, w0):
+    def reset(batch_size, n_units, v0, w0):
         restVals = jnp.zeros((batch_size, n_units))
         j = restVals # None
         v = restVals + v0
@@ -200,14 +170,6 @@ def _reset(batch_size, n_units, v0, w0):
         tols = restVals #+ 0
         return j, v, w, s, tols
 
-    @resolver(_reset)
-    def reset(self, j, v, w, s, tols):
-        self.j.set(j)
-        self.v.set(v)
-        self.w.set(w)
-        self.s.set(s)
-        self.tols.set(tols)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/neurons/spiking/quadLIFCell.py

@@ -102,7 +102,7 @@ class QuadLIFCell(LIFCell): ## quadratic integrate-and-fire cell
             sampled from the non-zero spikes detected
     """ ## batch_size arg?
 
-    @deprecate_args(thr_jitter=None)
+    @deprecate_args(thr_jitter=None, critical_v="critical_V")
     def __init__(
             self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., v_reset=-60., v_scale=-41.6, critical_v=1.,
             tau_theta=1e7, theta_plus=0.05, refract_time=5., one_spike=False, integration_type="euler",

tests/components/neurons/spiking/test_adExCell.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 AdExCell
+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_adExCell1():
+    name = "adex_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 = AdExCell(
+            name="a", n_units=1, tau_m=50., resist_m=30., thr=-66., 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.ones((1, 10))
+    ## desired output/epsp pulses
+    y_seq = jnp.asarray([[0., 0., 0., 0., 0., 0., 0., 1., 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_adExCell1()