revised izh-cell w/ unit test

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/izhikevichCell.py

@@ -1,27 +1,16 @@
-from jax import numpy as jnp, jit
-from ngclearn.utils import tensorstats
-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
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
-        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, w, b, tau_m): ## raw voltage dynamics
@@ -55,39 +44,6 @@ def _post_process(s, _v, _w, v, w, c, d): ## internal post-processing routine
     w_next = _w * (1. - s) + s * (w + d)
     return v_next, w_next
 
-@jit
-def _emit_spike(v, v_thr):
-    s = (v > v_thr).astype(jnp.float32)
-    return s
-
-@jit
-def _modify_current(j, R_m):
-    _j = j * R_m
-    return _j
-
-def _run_cell(dt, j, v, s, w, v_thr=30., tau_m=1., tau_w=50., b=0.2, c=-65., d=8.,
-              R_m=1., integType=0):
-    ## note: a = 0.1 --> fast spikes, a = 0.02 --> regular spikes
-    a = 1./tau_w ## we map time constant to variable "a" (a = 1/tau_w)
-    _j = _modify_current(j, R_m)
-    #_j = jnp.maximum(-30.0, _j) ## lower-bound/clip input current
-    ## check for spikes
-    s = _emit_spike(v, v_thr)
-    ## for non-spikes, evolve according to dynamics
-    if integType == 1:
-        v_params = (_j, w, b, tau_m)
-        _, _v = step_rk2(0., v, _dfv, dt, v_params) #_v = step_rk2(v, v_params, _dfv, dt)
-        w_params = (_j, v, b, tau_w)
-        _, _w = step_rk2(0., w, _dfw, dt, w_params) #_w = step_rk2(w, w_params, _dfw, dt)
-    else: # integType == 0 (default -- Euler)
-        v_params = (_j, w, b, tau_m)
-        _, _v = step_euler(0., v, _dfv, dt, v_params) #_v = step_euler(v, v_params, _dfv, dt)
-        w_params = (_j, v, b, tau_w)
-        _, _w = step_euler(0., w, _dfw, dt, w_params) #_w = step_euler(w, w_params, _dfw, dt)
-    ## for spikes, snap to particular states
-    _v, _w = _post_process(s, _v, _w, v, w, c, d)
-    return _v, _w, s
-
 class IzhikevichCell(JaxComponent): ## Izhikevich neuronal cell
     """
     A spiking cell based on Izhikevich's model of neuronal dynamics. Note that
@@ -197,24 +153,38 @@ def __init__(self, name, n_units, tau_m=1., resist_m=1., v_thr=30., v_reset=-65.
         self.s = Compartment(restVals)
         self.tols = Compartment(restVals) ## time-of-last-spike
 
+    @transition(output_compartments=["j", "v", "w", "s", "tols"])
     @staticmethod
-    def _advance_state(t, dt, tau_m, tau_w, v_thr, coupling, v_reset, w_reset, R_m,
+    def advance_state(t, dt, tau_m, tau_w, v_thr, coupling, v_reset, w_reset, R_m,
                        intgFlag, j, v, w, s, tols):
-        v, w, s = _run_cell(dt, j, v, s, w, v_thr=v_thr, tau_m=tau_m, tau_w=tau_w,
-                            b=coupling, c=v_reset, d=w_reset, R_m=R_m, integType=intgFlag)
-        tols = _update_times(t, s, tols)
+        ## note: a = 0.1 --> fast spikes, a = 0.02 --> regular spikes
+        a = 1. / tau_w  ## we map time constant to variable "a" (a = 1/tau_w)
+        _j = j * R_m
+        # _j = jnp.maximum(-30.0, _j) ## lower-bound/clip input current
+        ## check for spikes
+        s = (v > v_thr) * 1.
+        ## for non-spikes, evolve according to dynamics
+        if intgFlag == 1:
+            v_params = (_j, w, coupling, tau_m)
+            _, _v = step_rk2(0., v, _dfv, dt, v_params)  # _v = step_rk2(v, v_params, _dfv, dt)
+            w_params = (_j, v, coupling, tau_w)
+            _, _w = step_rk2(0., w, _dfw, dt, w_params)  # _w = step_rk2(w, w_params, _dfw, dt)
+        else:  # integType == 0 (default -- Euler)
+            v_params = (_j, w, coupling, tau_m)
+            _, _v = step_euler(0., v, _dfv, dt, v_params)  # _v = step_euler(v, v_params, _dfv, dt)
+            w_params = (_j, v, coupling, tau_w)
+            _, _w = step_euler(0., w, _dfw, dt, w_params)  # _w = step_euler(w, w_params, _dfw, dt)
+        ## for spikes, snap to particular states
+        _v, _w = _post_process(s, _v, _w, v, w, v_reset, w_reset)
+        v = _v
+        w = _w
+
+        tols = (1. - s) * tols + (s * t) ## update tols
         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
@@ -223,14 +193,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 = {

tests/components/neurons/spiking/test_fitzhughNagumoCell.py

@@ -14,7 +14,7 @@
 from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 
-def test_fitzhughNagumoCellCell1():
+def test_fitzhughNagumoCell1():
     name = "fh_ctx"
     ## create seeding keys
     dkey = random.PRNGKey(1234)
@@ -55,20 +55,16 @@ def clamp(x):
     y_seq = jnp.asarray([[0., 0., 0., 0., 0., 1., 0., 0., 0., 0.]], dtype=jnp.float32)
 
     outs = []
-    volts = []
-    recovs = []
     ctx.reset()
     for ts in range(x_seq.shape[1]):
         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)
-        volts.append(a.v.value)
-        recovs.append(a.w.value)
     outs = jnp.concatenate(outs, axis=1)
     #print(outs)
 
     ## output should equal input
     assert_array_equal(outs, y_seq)
 
-test_fitzhughNagumoCellCell1()
+#test_fitzhughNagumoCell1()

tests/components/neurons/spiking/test_izhikevichCell.py

@@ -0,0 +1,71 @@
+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 IzhikevichCell
+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_izhikevichCell1():
+    name = "izh_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 = IzhikevichCell(
+            name="a", n_units=1, tau_m=1., resist_m=4., v_thr=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([[0., 0., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    y_seq = jnp.asarray([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.]], dtype=jnp.float32)
+
+    outs = []
+    ctx.reset()
+    for ts in range(x_seq.shape[1]):
+        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)
+        print(a.v.value)
+    outs = jnp.concatenate(outs, axis=1)
+    print(outs)
+    #exit()
+    ## output should equal input
+    assert_array_equal(outs, y_seq)
+
+#test_izhikevichCell1()