refactored fn-cell and test passed

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/adExCell.py

@@ -32,7 +32,7 @@ def _dfw(t, w, params): ## recovery dynamics wrapper
     dv_dt = _dfw_internal(j, v, w, a, tau_m, v_rest)
     return dv_dt
 
-class AdExCell(JaxComponent):
+class AdExCell(JaxComponent): ## adaptive exponential integrate-and-fire cell
     """
     The AdEx (adaptive exponential leaky integrate-and-fire) neuronal cell
     model; a two-variable model. This cell model iteratively evolves

ngclearn/components/neurons/spiking/fitzhughNagumoCell.py

@@ -2,13 +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
 
 
@@ -34,7 +33,7 @@ def _dfw(t, w, params): ## recovery dynamics wrapper
     dv_dt = _dfw_internal(j, v, w, a, b, g, tau_m)
     return dv_dt
 
-class FitzhughNagumoCell(JaxComponent):
+class FitzhughNagumoCell(JaxComponent): ## F-H cell
     """
     The Fitzhugh-Nagumo neuronal cell model; a two-variable simplification
     of the Hodgkin-Huxley (squid axon) model. This cell model iteratively evolves
@@ -103,10 +102,10 @@ class FitzhughNagumoCell(JaxComponent):
                 at an increase in computational cost (and simulation time)
     """
 
-    # Define Functions
-    def __init__(self, name, n_units, tau_m=1., resist_m=1., tau_w=12.5, alpha=0.7,
-                 beta=0.8, gamma=3., v0=0., w0=0., v_thr=1.07, spike_reset=False,
-                 integration_type="euler", **kwargs):
+    def __init__(
+            self, name, n_units, tau_m=1., resist_m=1., tau_w=12.5, alpha=0.7, beta=0.8, gamma=3., v0=0., w0=0.,
+            v_thr=1.07, spike_reset=False, integration_type="euler", **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         ## Integration properties
@@ -115,7 +114,7 @@ def __init__(self, name, n_units, tau_m=1., resist_m=1., tau_w=12.5, alpha=0.7,
 
         ## Cell properties
         self.tau_m = tau_m
-        self.R_m = resist_m
+        self.resist_m = resist_m ## resistance R_m
         self.tau_w = tau_w
         self.alpha = alpha
         self.beta = beta
@@ -138,41 +137,44 @@ 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,
-                       beta, gamma, intgFlag, j, v, w, tols):
-        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)
+    @compilable
+    def advance_state(self, t, dt):
+        j_mod = self.j.get() * self.resist_m
+        if self.intgFlag == 1:
+            v_params = (j_mod, self.w.get(), self.alpha, self.beta, self.gamma, self.tau_m)
+            _, _v = step_rk2(0., self.v.get(), _dfv, dt, v_params)  # _v = step_rk2(v, v_params, _dfv, dt)
+            w_params = (j_mod, self.v.get(), self.alpha, self.beta, self.gamma, self.tau_w)
+            _, _w = step_rk2(0., self.w.get(), _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_params = (j_mod, self.w.get(), self.alpha, self.beta, self.gamma, self.tau_m)
+            _, _v = step_euler(0., self.v.get(), _dfv, dt, v_params)  # _v = step_euler(v, v_params, _dfv, dt)
+            w_params = (j_mod, self.v.get(), self.alpha, self.beta, self.gamma, self.tau_w)
+            _, _w = step_euler(0., self.w.get(), _dfw, dt, w_params)  # _w = step_euler(w, w_params, _dfw, dt)
+        s = (_v > self.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 = (1. - s) * tols + (s * t) ## update tols
-        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
+        if self.spike_reset: ## if spike-reset used, variables snapped back to initial conditions
+            v = v * (1. - s) + s * self.v0
+            w = w * (1. - s) + s * self.w0
+
+        ## 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
@@ -197,8 +199,7 @@ def help(cls): ## component help function
             "resist_m": "Membrane resistance value",
             "tau_w": "Recovery variable time constant",
             "v_thr": "Base voltage threshold value",
-            "spike_reset": "Should voltage/recover be snapped to initial "
-                           "condition(s) if spike emitted?",
+            "spike_reset": "Should voltage/recover be snapped to initial condition(s) if spike emitted?",
             "alpha": "Dimensionless recovery variable shift factor `a",
             "beta": "Dimensionless recovery variable scale factor `b`",
             "gamma": "Power-term divisor constant",

tests/components/neurons/spiking/test_fitzhughNagumoCell.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 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
+from ngclearn import Context, MethodProcess
+from ngclearn.components.neurons.spiking.fitzhughNagumoCell import FitzhughNagumoCell
+from numpy.testing import assert_array_equal
 
 
 def test_fitzhughNagumoCell1():
@@ -26,43 +20,39 @@ def test_fitzhughNagumoCell1():
             name="a", n_units=1, tau_m=1., resist_m=5., v_thr=2.1, 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.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 = []
-    ctx.reset()
+    reset_process.run()  # 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)
+        x_t = jnp.array([[x_seq[0, ts]]])  ## get data at time t
+        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)