ported over IF/quadLIF cells, minor revision to LIF cell

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/IFCell.py

@@ -2,16 +2,15 @@
 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 import deprecate_args
 from ngcsimlib.logger import info, warn
 from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \
                                             step_euler, step_rk2
-from ngclearn.utils.surrogate_fx import (secant_lif_estimator, arctan_estimator,
-                                         triangular_estimator,
-                                         straight_through_estimator)
+# from ngclearn.utils.surrogate_fx import (secant_lif_estimator, arctan_estimator,
+#                                          triangular_estimator,
+#                                          straight_through_estimator)
 
-from ngcsimlib.compilers.process import transition
-#from ngcsimlib.component import Component
+from ngcsimlib.parser import compilable
 from ngcsimlib.compartment import Compartment
 
 
@@ -35,7 +34,7 @@ class IFCell(JaxComponent): ## integrate-and-fire cell
     The specific differential equation that characterizes this cell
     is (for adjusting v, given current j, over time) is:
 
-    | tau_m * dv/dt = (v_rest - v) + j * R
+    | tau_m * dv/dt = j * R
     | where R is the membrane resistance and v_rest is the resting potential
     | also, if a spike occurs, v is set to v_reset
 
@@ -91,10 +90,10 @@ 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",
-                 surrogate_type="straight_through", lower_clamp_voltage=True,
-                 **kwargs):
+    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", surrogate_type="straight_through", lower_clamp_voltage=True, **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         ## Integration properties
@@ -118,12 +117,12 @@ 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 surrogate_type == "arctan":
-            self.spike_fx, self.d_spike_fx = arctan_estimator()
-        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()
+        # if surrogate_type == "arctan":
+        #     self.spike_fx, self.d_spike_fx = arctan_estimator()
+        # 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()
 
 
         ## Compartment setup
@@ -138,47 +137,48 @@ def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65.,
                                 units="ms") ## time-of-last-spike
         self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value")
 
-    @transition(output_compartments=["v", "s", "rfr", "tols", "key", "surrogate"])
-    @staticmethod
+    @compilable
     def advance_state(
-            t, dt, tau_m, resist_m, v_rest, v_reset, refract_T, thr, lower_clamp_voltage, intgFlag, d_spike_fx, key,
-            j, v, rfr, tols
+            self, dt, t
     ):
         ## run one integration step for neuronal dynamics
-        j = j * resist_m
+        j = self.j.get() * self.resist_m
 
         ### Runs integrator (or integrate-and-fire; IF) neuronal dynamics
         ## update voltage / membrane potential
-        v_params = (j, rfr, tau_m, refract_T)
-        if intgFlag == 1:
-            _, _v = step_rk2(0., v, _dfv, dt, v_params)
+        v_params = (j, self.rfr.get(), self.tau_m, self.refract_T)
+        if self.intgFlag == 1:
+            _, _v = step_rk2(0., self.v.get(), _dfv, dt, v_params)
         else:
-            _, _v = step_euler(0., v, _dfv, dt, v_params)
+            _, _v = step_euler(0., self.v.get(), _dfv, dt, v_params)
         ## obtain action potentials/spikes
-        s = (_v > thr) * 1.
+        s = (_v > self.thr) * 1.
         ## update refractory variables
-        rfr = (rfr + dt) * (1. - s)
+        rfr = (self.rfr.get() + dt) * (1. - s)
         ## perform hyper-polarization of neuronal cells
-        v = _v * (1. - s) + s * v_reset
+        v = _v * (1. - s) + s * self.v_reset
+
+        #surrogate = d_spike_fx(v, self.thr)
 
-        surrogate = d_spike_fx(v, thr)
         ## update tols
-        tols = (1. - s) * tols + (s * t)
-        if lower_clamp_voltage: ## ensure voltage never < v_rest
-            v = jnp.maximum(v, v_rest)
-        return v, s, rfr, tols, key, surrogate
-
-    @transition(output_compartments=["j", "v", "s", "rfr", "tols", "surrogate"])
-    @staticmethod
-    def reset(batch_size, n_units, v_rest, refract_T):
-        restVals = jnp.zeros((batch_size, n_units))
-        j = restVals #+ 0
-        v = restVals + v_rest
-        s = restVals #+ 0
-        rfr = restVals + refract_T
-        tols = restVals #+ 0
-        surrogate = restVals + 1.
-        return j, v, s, rfr, tols, surrogate
+        self.tols.set((1. - s) * self.tols.get() + (s * t))
+        if self.lower_clamp_voltage: ## ensure voltage never < v_rest
+            _v = jnp.maximum(v, self.v_rest)
+
+        self.v.set(_v)
+        self.s.set(s)
+        self.rfr.set(rfr)
+
+    @compilable
+    def reset(self):
+        restVals = jnp.zeros((self.batch_size, self.n_units))
+        if not self.j.targeted:
+            self.j.set(restVals)
+        self.v.set(restVals + self.v_rest)
+        self.s.set(restVals)
+        self.rfr.set(restVals + self.refract_T)
+        self.tols.set(restVals)
+        #surrogate = restVals + 1.
 
     def save(self, directory, **kwargs):
         ## do a protected save of constants, depending on whether they are floats or arrays

ngclearn/components/neurons/spiking/LIFCell.py

@@ -19,7 +19,7 @@ def _dfv(t, v, params): ## voltage dynamics wrapper
 
 
 #@partial(jit, static_argnums=[3, 4])
-def _update_theta(dt, v_theta, s, tau_theta, theta_plus=0.05):
+def _update_theta(dt, v_theta, s, tau_theta, theta_plus: Array=0.05):
     ### Runs homeostatic threshold update dynamics one step (via Euler integration).
     #theta_decay = 0.9999999 #0.999999762 #jnp.exp(-dt/1e7)
     #theta_plus = 0.05

ngclearn/components/neurons/spiking/quadLIFCell.py

@@ -29,7 +29,7 @@ def _dfv(t, v, params): ## voltage dynamics wrapper
     return dv_dt
 
 #@partial(jit, static_argnums=[3, 4])
-def _update_theta(dt, v_theta, s, tau_theta, theta_plus: Array | float=0.05):
+def _update_theta(dt, v_theta, s, tau_theta, theta_plus: Array=0.05):
     ### Runs homeostatic threshold update dynamics one step (via Euler integration).
     #theta_decay = 0.9999999 #0.999999762 #jnp.exp(-dt/1e7)
     #theta_plus = 0.05
@@ -138,7 +138,7 @@ def advance_state(self, dt, t):
 
         _v_thr = self.thr_theta.get() + self.thr  ## calc present voltage threshold
 
-        v_params = (j, self.rfr.get(), self.tau_m.get(), self.refract_T, self.v_rest, self.v_c, self.a0)
+        v_params = (j, self.rfr.get(), self.tau_m, self.refract_T, self.v_rest, self.v_c, self.a0)
 
         if self.intgFlag == 1:
             _, _v = step_rk2(0., self.v.get(), _dfv, dt, v_params)