revised fh-cell w/ unit test

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/fitzhughNagumoCell.py

@@ -1,27 +1,16 @@
-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
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
-    Returns:
-        updated tols variable
-    """
-    _tols = (1. - s) * tols + (s * t)
-    return _tols
 
 @jit
 def _dfv_internal(j, v, w, a, b, g, tau_m): ## raw voltage dynamics
@@ -45,25 +34,6 @@ def _dfw(t, w, params): ## recovery dynamics wrapper
     dv_dt = _dfw_internal(j, v, w, a, b, g, tau_m)
     return dv_dt
 
-@jit
-def _emit_spike(v, v_thr):
-    s = (v > v_thr).astype(jnp.float32)
-    return s
-
-def _run_cell(dt, j, v, w, v_thr, tau_m, tau_w, a, b, g=3., integType=0):
-    if integType == 1:
-        v_params = (j, w, a, b, g, tau_m)
-        _, _v = step_rk2(0., v, _dfv, dt, v_params) #_v = step_rk2(v, v_params, _dfv, dt)
-        w_params = (j, v, a, b, g, 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, a, b, g, tau_m)
-        _, _v = step_euler(0., v, _dfv, dt, v_params) #_v = step_euler(v, v_params, _dfv, dt)
-        w_params = (j, v, a, b, g, tau_w)
-        _, _w = step_euler(0., w, _dfw, dt, w_params) #_w = step_euler(w, w_params, _dfw, dt)
-    s = _emit_spike(_v, v_thr)
-    return _v, _w, s
-
 class FitzhughNagumoCell(JaxComponent):
     """
     The Fitzhugh-Nagumo neuronal cell model; a two-variable simplification
@@ -168,27 +138,34 @@ def __init__(self, name, n_units, tau_m=1., resist_m=1., tau_w=12.5, alpha=0.7,
         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, R_m, tau_w, v_thr, spike_reset, v0, w0, alpha,
+    def advance_state(t, dt, tau_m, R_m, tau_w, v_thr, spike_reset, v0, w0, alpha,
                        beta, gamma, intgFlag, j, v, w, tols):
-        v, w, s = _run_cell(dt, j * R_m, v, w, v_thr, tau_m, tau_w, alpha, beta,
-                            gamma, intgFlag)
+        j_mod = j * R_m
+        if intgFlag == 1:
+            v_params = (j_mod, w, alpha, beta, gamma, tau_m)
+            _, _v = step_rk2(0., v, _dfv, dt, v_params)  # _v = step_rk2(v, v_params, _dfv, dt)
+            w_params = (j_mod, v, alpha, beta, gamma, 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_mod, w, alpha, beta, gamma, tau_m)
+            _, _v = step_euler(0., v, _dfv, dt, v_params)  # _v = step_euler(v, v_params, _dfv, dt)
+            w_params = (j_mod, v, alpha, beta, gamma, tau_w)
+            _, _w = step_euler(0., w, _dfw, dt, w_params)  # _w = step_euler(w, w_params, _dfw, dt)
+        s = (_v > v_thr) * 1.
+        v = _v
+        w = _w
+
         if spike_reset: ## if spike-reset used, variables snapped back to initial conditions
             v = v * (1. - s) + s * v0
             w = w * (1. - s) + s * w0
-        tols = _update_times(t, s, tols)
+        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
@@ -197,14 +174,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_WTASCell.py

@@ -23,7 +23,7 @@ def test_WTASCell1():
     # ---- build a simple Poisson cell system ----
     with Context(name) as ctx:
         a = WTASCell(
-            name="a", n_units=1, tau_m=25., resist_m=1., key=subkeys[0]
+            name="a", n_units=2, tau_m=25., resist_m=1., key=subkeys[0]
         )
 
         #"""

tests/components/neurons/spiking/test_fitzhughNagumoCell.py

@@ -0,0 +1,74 @@
+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 FitzhughNagumoCell
+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_fitzhughNagumoCellCell1():
+    name = "fh_ctx"
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    dt = 0.1 #1.  # ms
+    # ---- build a simple Poisson cell system ----
+    with Context(name) as ctx:
+        a = FitzhughNagumoCell(
+            name="a", n_units=1, tau_m=1., resist_m=5., v_thr=2.1, 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., 0., 0., 0., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    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()