sketch of ifcell test

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/IFCell.py

@@ -1,14 +1,12 @@
 from ngclearn.components.jaxComponent import JaxComponent
-from jax import numpy as jnp, random, jit, nn
-from functools import partial
+from jax import numpy as jnp, random, nn, Array, jit
 from ngclearn.utils import tensorstats
 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.parser import compilable
 from ngcsimlib.compartment import Compartment
@@ -135,7 +133,7 @@ def __init__(
                                display_name="Refractory Time Period", units="ms")
         self.tols = Compartment(restVals, display_name="Time-of-Last-Spike",
                                 units="ms") ## time-of-last-spike
-        self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value")
+        #self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value")
 
     @compilable
     def advance_state(

tests/components/neurons/spiking/test_IFCell.py

@@ -2,68 +2,57 @@
 from ngcsimlib.context import Context
 import numpy as np
 np.random.seed(42)
-from ngclearn.components import IFCell
-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.IFCell import IFCell
+from numpy.testing import assert_array_equal
 
 def test_IFCell1():
     name = "if_ctx"
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
     dt = 1.  # ms
-    trace_increment = 0.1
     # ---- build a simple Poisson cell system ----
     with Context(name) as ctx:
         a = IFCell(
             name="a", n_units=1, tau_m=5., resist_m=10., 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([[1., 1., 1., 1., 1., 0., 0., 1., 1., 1., 1., 1., 1., 1., 0.]], dtype=jnp.float32)
     ## desired output/epsp pulses
     y_seq = jnp.asarray([[0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 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_IFCell1()
+test_IFCell1()