revised raf-cell w/ unit test; fixed typos/mistakes in all spiking cells

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/IFCell.py

@@ -75,10 +75,10 @@ class IFCell(JaxComponent): ## integrate-and-fire cell
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
 
             :Note: setting the integration type to the midpoint method will
-                increase the accuray of the estimate of the cell's evolution
+                increase the accuracy of the estimate of the cell's evolution
                 at an increase in computational cost (and simulation time)
 
-        surrgoate_type: type of surrogate function to use for approximating a
+        surrogate_type: type of surrogate function to use for approximating a
             partial derivative of this cell's spikes w.r.t. its voltage/current
             (default: "straight_through")
 
@@ -93,7 +93,7 @@ class IFCell(JaxComponent): ## integrate-and-fire cell
     @deprecate_args(thr_jitter=None)
     def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
                  v_reset=-60., refract_time=0., integration_type="euler",
-                 surrgoate_type="straight_through", lower_clamp_voltage=True,
+                 surrogate_type="straight_through", lower_clamp_voltage=True,
                  **kwargs):
         super().__init__(name, **kwargs)
 
@@ -118,9 +118,9 @@ def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
         self.n_units = n_units
 
         ## set up surrogate function for spike emission
-        if surrgoate_type == "arctan":
+        if surrogate_type == "arctan":
             self.spike_fx, self.d_spike_fx = arctan_estimator()
-        elif surrgoate_type == "triangular":
+        elif surrogate_type == "triangular":
             self.spike_fx, self.d_spike_fx = triangular_estimator()
         else: ## default: straight_through
             self.spike_fx, self.d_spike_fx = straight_through_estimator()

ngclearn/components/neurons/spiking/LIFCell.py

@@ -100,10 +100,10 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
 
             :Note: setting the integration type to the midpoint method will
-                increase the accuray of the estimate of the cell's evolution
+                increase the accuracy of the estimate of the cell's evolution
                 at an increase in computational cost (and simulation time)
 
-        surrgoate_type: type of surrogate function to use for approximating a
+        surrogate_type: type of surrogate function to use for approximating a
             partial derivative of this cell's spikes w.r.t. its voltage/current
             (default: "straight_through")
 
@@ -120,7 +120,7 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
     def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
                  v_reset=-60., v_decay=1., tau_theta=1e7, theta_plus=0.05,
                  refract_time=5., one_spike=False, integration_type="euler",
-                 surrgoate_type="straight_through", lower_clamp_voltage=True,
+                 surrogate_type="straight_through", lower_clamp_voltage=True,
                  **kwargs):
         super().__init__(name, **kwargs)
 
@@ -150,11 +150,11 @@ def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
         self.n_units = n_units
 
         ## set up surrogate function for spike emission
-        if surrgoate_type == "secant_lif":
+        if surrogate_type == "secant_lif":
             self.spike_fx, self.d_spike_fx = secant_lif_estimator()
-        elif surrgoate_type == "arctan":
+        elif surrogate_type == "arctan":
             self.spike_fx, self.d_spike_fx = arctan_estimator()
-        elif surrgoate_type == "triangular":
+        elif surrogate_type == "triangular":
             self.spike_fx, self.d_spike_fx = triangular_estimator()
         else: ## default: straight_through
             self.spike_fx, self.d_spike_fx = straight_through_estimator()

ngclearn/components/neurons/spiking/RAFCell.py

@@ -1,28 +1,15 @@
-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, omega, b): ## "voltage" dynamics
@@ -48,11 +35,6 @@ def _dfw(t, w, params): ## angular driver dynamics wrapper
     dv_dt = _dfw_internal(j, v, w, tau_w, omega, b)
     return dv_dt
 
-@jit
-def _emit_spike(v, v_thr):
-    s = (v > v_thr).astype(jnp.float32)
-    return s
-
 class RAFCell(JaxComponent):
     """
     The resonate-and-fire (RAF) neuronal cell
@@ -112,14 +94,15 @@ class RAFCell(JaxComponent):
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
 
             :Note: setting the integration type to the midpoint method will
-                increase the accuray of the estimate of the cell's evolution
+                increase the accuracy of the estimate of the cell's evolution
                 at an increase in computational cost (and simulation time)
     """
 
     @deprecate_args(resist_m="resist_v", tau_m="tau_v")
-    def __init__(self, name, n_units, tau_v=1., tau_w=1., thr=1., omega=10.,
-                 b=-1., v_reset=1., w_reset=0., v0=0., w0=0., resist_v=1.,
-                 integration_type="euler", batch_size=1, **kwargs):
+    def __init__(
+            self, name, n_units, tau_v=1., tau_w=1., thr=1., omega=10., b=-1., v_reset=0., w_reset=0., v0=0., w0=0.,
+            resist_v=1., integration_type="euler", batch_size=1, **kwargs
+    ):
         #v_rest=-72., v_reset=-75., w_reset=0., thr=5., v0=-70., w0=0., tau_w=400., thr=5., omega=10., b=-1.
         super().__init__(name, **kwargs)
 
@@ -150,11 +133,13 @@ def __init__(self, name, n_units, tau_v=1., tau_w=1., thr=1., omega=10.,
         self.v = Compartment(restVals + self.v0, display_name="Voltage", units="mV")
         self.w = Compartment(restVals + self.w0, display_name="Angular-Driver")
         self.s = Compartment(restVals, display_name="Spikes")
-        self.tols = Compartment(restVals, display_name="Time-of-Last-Spike",
-                                units="ms") ## time-of-last-spike
+        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_v, resist_v, tau_w, thr, omega, b,
+    def advance_state(t, dt, tau_v, resist_v, tau_w, thr, omega, b,
                        v_reset, w_reset, intgFlag, j, v, w, tols):
         ## continue with centered dynamics
         j_ = j * resist_v
@@ -170,24 +155,17 @@ def _advance_state(t, dt, tau_v, resist_v, tau_w, thr, omega, b,
             _, _w = step_euler(0., w, _dfw, dt, w_params)
             v_params = (j_, _w, tau_v, omega, b)
             _, _v = step_euler(0., v, _dfv, dt, v_params)
-        s = _emit_spike(_v, thr)
+        s = (_v > thr) * 1. ## emit spikes/pulses
         ## hyperpolarize/reset/snap variables
         w = _w * (1. - s) + s * w_reset
         v = _v * (1. - s) + s * v_reset
 
-        tols = _update_times(t, s, tols)
+        tols = (1. - s) * tols + (s * t) ## update times-of-last-spike(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
@@ -196,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 = {

ngclearn/components/neurons/spiking/WTASCell.py

@@ -102,8 +102,10 @@ class WTASCell(JaxComponent): ## winner-take-all spiking cell
     """
 
     # Define Functions
-    def __init__(self, name, n_units, tau_m, resist_m=1., thr_base=0.4, thr_gain=0.002,
-                 refract_time=0., thr_jitter=0.05, **kwargs):
+    def __init__(
+            self, name, n_units, tau_m, resist_m=1., thr_base=0.4, thr_gain=0.002, refract_time=0., thr_jitter=0.05,
+            **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         ## membrane parameter setup (affects ODE integration)

ngclearn/components/neurons/spiking/adExCell.py

@@ -93,7 +93,7 @@ class AdExCell(JaxComponent):
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
 
             :Note: setting the integration type to the midpoint method will
-                increase the accuray of the estimate of the cell's evolution
+                increase the accuracy of the estimate of the cell's evolution
                 at an increase in computational cost (and simulation time)
     """
 

ngclearn/components/neurons/spiking/fitzhughNagumoCell.py

@@ -129,7 +129,7 @@ class FitzhughNagumoCell(JaxComponent):
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
 
             :Note: setting the integration type to the midpoint method will
-                increase the accuray of the estimate of the cell's evolution
+                increase the accuracy of the estimate of the cell's evolution
                 at an increase in computational cost (and simulation time)
     """
 

ngclearn/components/neurons/spiking/izhikevichCell.py

@@ -159,7 +159,7 @@ class IzhikevichCell(JaxComponent): ## Izhikevich neuronal cell
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
 
             :Note: setting the integration type to the midpoint method will
-                increase the accuray of the estimate of the cell's evolution
+                increase the accuracy of the estimate of the cell's evolution
                 at an increase in computational cost (and simulation time)
     """
 

ngclearn/components/neurons/spiking/quadLIFCell.py

@@ -100,17 +100,34 @@ class QuadLIFCell(LIFCell): ## quadratic integrate-and-fire cell
             a single spike will be permitted to emit per step -- this means that
             if > 1 spikes emitted, a single action potential will be randomly
             sampled from the non-zero spikes detected
+            
+        integration_type: type of integration to use for this cell's dynamics;
+            current supported forms include "euler" (Euler/RK-1 integration)
+            and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
+
+            :Note: setting the integration type to the midpoint method will
+                increase the accuracy of the estimate of the cell's evolution
+                at an increase in computational cost (and simulation time)
+
+        surrogate_type: type of surrogate function to use for approximating a
+            partial derivative of this cell's spikes w.r.t. its voltage/current
+            (default: "straight_through")
+
+            :Note: surrogate options available include: "straight_through"
+                (straight-through estimator), "triangular" (triangular estimator),
+                "arctan" (arc-tangent estimator), and "secant_lif" (the
+                LIF-specialized secant estimator)
     """ ## batch_size arg?
 
     @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",
-            surrgoate_type="straight_through", lower_clamp_voltage=True, **kwargs
+            surrogate_type="straight_through", lower_clamp_voltage=True, **kwargs
     ):
         super().__init__(
             name, n_units, tau_m, resist_m, thr, v_rest, v_reset, 1., tau_theta, theta_plus, refract_time,
-            one_spike, integration_type, surrgoate_type, lower_clamp_voltage, **kwargs
+            one_spike, integration_type, surrogate_type, lower_clamp_voltage, **kwargs
         )
         ## only two distinct additional constants distinguish the Quad-LIF cell
         self.v_c = v_scale

ngclearn/components/neurons/spiking/sLIFCell.py

@@ -113,10 +113,10 @@ class SLIFCell(JaxComponent): ## leaky integrate-and-fire cell
     """
 
     # Define Functions
-    def __init__(self, name, n_units, tau_m, resist_m, thr, resist_inh=0.,
-                 thr_persist=False, thr_gain=0.0, thr_leak=0.0, rho_b=0.,
-                 refract_time=0., sticky_spikes=False, thr_jitter=0.05,
-                 batch_size=1, **kwargs):
+    def __init__(
+            self, name, n_units, tau_m, resist_m, thr, resist_inh=0., thr_persist=False, thr_gain=0.0, thr_leak=0.0,
+            rho_b=0., refract_time=0., sticky_spikes=False, thr_jitter=0.05, batch_size=1, **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         ## membrane parameter setup (affects ODE integration)

tests/components/neurons/spiking/test_RAFCell.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 RAFCell
+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_RAFCell1():
+    name = "raf_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 = RAFCell(
+            name="a", n_units=1, tau_v=20., resist_v=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., 1., 0., 0., 0., 0., 1., 0., 0.]], dtype=jnp.float32)
+    ## desired output/epsp pulses
+    y_seq = jnp.asarray([[0., 0., 0., 1., 0., 0., 0., 0., 1.]], 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_RAFCell1()