refactored adex/test passed; minor cleanup in lif, raf, and wtas cells

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/LIFCell.py

@@ -200,7 +200,7 @@ def advance_state(self, dt, t):
             thr_theta = _update_theta(dt, self.thr_theta.get(), raw_s, self.tau_theta, self.theta_plus) #.get())
             self.thr_theta.set(thr_theta)
 
-        ## update tols
+        ## update time-of-last spike variable(s)
         self.tols.set((1. - s) * self.tols.get() + (s * t))
 
         if self.v_min is not None: ## ensures voltage never < v_rest

ngclearn/components/neurons/spiking/RAFCell.py

@@ -139,9 +139,7 @@ def __init__(
         ) ## time-of-last-spike
 
     @compilable
-    def advance_state(
-            self, t, dt
-    ):
+    def advance_state(self, t, dt):
         ## continue with centered dynamics
         j_ = self.j.get() * self.resist_v
         if self.intgFlag == 1:  ## RK-2/midpoint

ngclearn/components/neurons/spiking/WTASCell.py

@@ -87,12 +87,8 @@ def __init__(
         self.rfr = Compartment(restVals + self.refract_T)
         self.tols = Compartment(restVals) ## time-of-last-spike
 
-    # @transition(output_compartments=["v", "s", "thr", "rfr", "tols"])
-    # @staticmethod
     @compilable
-    def advance_state(
-            self, t, dt #, tau_m, R_m, thr_gain, refract_T, j, v, thr, rfr, tols
-    ):
+    def advance_state(self, t, dt):
         mask = (self.rfr.get() >= self.refract_T) * 1.  ## check refractory period
         v = (self.j.get() * self.R_m) * mask
         vp = softmax(v)  # convert to Categorical (spike) probabilities
@@ -111,8 +107,6 @@ def advance_state(
         self.thr.set(thr)
         self.rfr.set(rfr)
 
-    # @transition(output_compartments=["j", "v", "s", "rfr", "tols"])
-    # @staticmethod
     @compilable
     def reset(self):
         restVals = jnp.zeros((self.batch_size, self.n_units))

ngclearn/components/neurons/spiking/adExCell.py

@@ -2,12 +2,12 @@
 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 ngcsimlib.compilers.process import transition
-#from ngcsimlib.component import Component
+
+from ngcsimlib.parser import compilable
 from ngcsimlib.compartment import Compartment
 
 @jit
@@ -97,7 +97,7 @@ class AdExCell(JaxComponent):
                 at an increase in computational cost (and simulation time)
     """
 
-    @deprecate_args(v_thr="thr")
+    #@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
@@ -136,39 +136,40 @@ def __init__(
         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
-    ):
-        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)
+    @compilable
+    def advance_state(self, t, dt):
+        if self.intgFlag == 1:  ## RK-2/midpoint
+            v_params = (self.j.get(), self.w.get(), self.tau_m, self.v_rest, self.sharpV, self.vT, self.R_m)
+            _, _v = step_rk2(0., self.v.get(), _dfv, dt, v_params)
+            w_params = (self.j.get(), self.v.get(), self.a, self.tau_w, self.v_rest)
+            _, _w = step_rk2(0., self.w.get(), _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
+            v_params = (self.j.get(), self.w.get(), self.tau_m, self.v_rest, self.sharpV, self.vT, self.R_m)
+            _, _v = step_euler(0., self.v.get(), _dfv, dt, v_params)
+            w_params = (self.j.get(), self.v.get(), self.a, self.tau_w, self.v_rest)
+            _, _w = step_euler(0., self.w.get(), _dfw, dt, w_params)
+        s = (_v > self.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
-
-    @transition(output_compartments=["j", "v", "w", "s", "tols"])
-    @staticmethod
-    def reset(batch_size, n_units, v0, w0):
-        restVals = jnp.zeros((batch_size, n_units))
-        j = restVals # None
-        v = restVals + v0
-        w = restVals + w0
-        s = restVals #+ 0
-        tols = restVals #+ 0
-        return j, v, w, s, tols
+        v = _v * (1. - s) + s * self.v_reset
+        w = _w * (1. - s) + s * (_w + self.b)
+
+        ## update time-of-last spike variable(s)
+        self.tols.set((1. - s) * self.tols.get() + (s * t))
+
+        #self.j.set(j) ## j is not getting modified in these dynamics
+        self.v.set(v)
+        self.w.set(w)
+        self.s.set(s)
+
+    @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.v0)
+        self.w.set(restVals + self.w0)
+        self.s.set(restVals)
+        self.tols.set(restVals)
 
     @classmethod
     def help(cls): ## component help function

tests/components/neurons/spiking/test_adExCell.py

@@ -1,17 +1,11 @@
 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
+from ngclearn import Context, MethodProcess
+from ngclearn.components.neurons.spiking.adExCell import AdExCell
+from numpy.testing import assert_array_equal
 
 
 def test_adExCell1():
@@ -26,45 +20,41 @@ def test_adExCell1():
             name="a", n_units=1, tau_m=50., resist_m=30., thr=-66., key=subkeys[0]
         )
 
-        #"""
-        advance_process = (Process("advance_proc")
+        # """
+        advance_process = (MethodProcess("advance_proc")
                            >> a.advance_state)
-        # ctx.wrap_and_add_command(advance_process.pure, name="run")
-        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
+        # ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
-        reset_process = (Process("reset_proc")
+        reset_process = (MethodProcess("reset_proc")
                          >> 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")
-        """
-
+        # ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
+        # """
         ## set up non-compiled utility commands
-        @Context.dynamicCommand
-        def clamp(x):
-            a.j.set(x)
+        # @Context.dynamicCommand
+        # def clamp(x):
+        #     a.j.set(x)
+
+    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()
+    reset_process.run()  # 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)
+        clamp(x_t)  # ctx.clamp(x_t)
+        advance_process.run(t=ts * 1., dt=dt)  # ctx.run(t=ts * 1., dt=dt)
+        outs.append(a.s.get())
+
     outs = jnp.concatenate(outs, axis=1)
-    #print(outs)
+    # print(outs)
+    # print(y_seq)
 
     ## output should equal input
     assert_array_equal(outs, y_seq)
 
-#test_adExCell1()
+test_adExCell1()