update testing for graded neurons and input encoders

Author: rxng8

ngclearn/components/input_encoders/bernoulliCell.py

@@ -48,7 +48,9 @@ def advance_state(self, t):
     @compilable
     def reset(self):
         restVals = jnp.zeros((self.batch_size.get(), self.n_units.get()))
-        not self.inputs.targeted and self.inputs.set(restVals)
+        # BUG: the self.inputs here does not have the targeted field
+        # NOTE: Quick workaround is to check if targeted is in the input or not
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(restVals)
         self.outputs.set(restVals)
         self.tols.set(restVals)
 

ngclearn/components/input_encoders/latencyCell.py

@@ -211,7 +211,9 @@ def advance_state(self, t):
     @compilable
     def reset(self):
         restVals = jnp.zeros((self.batch_size.get(), self.n_units.get()))
-        not self.inputs.targeted and self.inputs.set(restVals)
+        # BUG: the self.inputs here does not have the targeted field
+        # NOTE: Quick workaround is to check if targeted is in the input or not
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(restVals)
         self.outputs.set(restVals)
         self.tols.set(restVals)
         self.mask.set(restVals)

ngclearn/components/input_encoders/phasorCell.py

@@ -88,7 +88,9 @@ def advance_state(self, t, dt):
     @compilable
     def reset(self):
         restVals = jnp.zeros((self.batch_size.get(), self.n_units.get()))
-        not self.inputs.targeted and self.inputs.set(restVals)
+        # BUG: the self.inputs here does not have the targeted field
+        # NOTE: Quick workaround is to check if targeted is in the input or not
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(restVals)
         self.outputs.set(restVals)
         self.tols.set(restVals)
         self.angles.set(restVals)

ngclearn/components/input_encoders/poissonCell.py

@@ -60,8 +60,9 @@ def advance_state(self, t, dt):
     @compilable
     def reset(self):
         restVals = jnp.zeros((self.batch_size, self.n_units))
-        if not self.inputs.targeted:
-            self.inputs.set(restVals)
+        # BUG: the self.inputs here does not have the targeted field
+        # NOTE: Quick workaround is to check if targeted is in the input or not
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(restVals)
         self.outputs.set(restVals)
         self.tols.set(restVals)
 

tests/components/input_encoders/test_bernoulliCell.py

@@ -1,15 +1,13 @@
+# %%
+
 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 BernoulliCell
 #from ngcsimlib.compilers import compile_command, wrap_command
 from numpy.testing import assert_array_equal
 
-from ngcsimlib.compilers.process import Process, transition
-from ngclearn.utils import JaxProcess
-from ngcsimlib.context import Context
-#from ngcsimlib.utils.compartment import Get_Compartment_Batch
+from ngclearn import MethodProcess, Context
 
 
 def test_bernoulliCell1():
@@ -23,28 +21,25 @@ def test_bernoulliCell1():
     with Context(name) as ctx:
         a = BernoulliCell(name="a", n_units=1, key=subkeys[0])
 
-        advance_process = (JaxProcess("advance_proc")
+        advance_process = (MethodProcess("advance_proc")
                            >> a.advance_state)
-        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")
 
         ## set up non-compiled utility commands
-        @Context.dynamicCommand
         def clamp(x):
             a.inputs.set(x)
 
     ## input spike train
     x_seq = jnp.asarray([[1., 1., 0., 0., 1.]], dtype=jnp.float32)
 
     outs = []
-    ctx.reset()
+    reset_process.run()
     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)
+        clamp(x_t)
+        advance_process.run(t=ts*1., dt=dt)
         outs.append(a.outputs.value)
     outs = jnp.concatenate(outs, axis=1)
 

tests/components/input_encoders/test_latencyCell.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 LatencyCell
-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 MethodProcess, Context
 
 def test_latencyCell1():
     name = "latency_ctx"
@@ -29,23 +23,19 @@ def test_latencyCell1():
         )
 
         ## create and compile core simulation commands
-        advance_process = (Process("advance_proc")
+        advance_process = (MethodProcess("advance_proc")
                            >> a.advance_state)
-        ctx.wrap_and_add_command(jit(advance_process.pure), name="advance")
-        calc_spike_times_process = (Process("calc_sptimes_proc")
+        calc_spike_times_process = (MethodProcess("calc_sptimes_proc")
                                     >> a.calc_spike_times)
-        ctx.wrap_and_add_command(jit(calc_spike_times_process.pure), name="calc_spike_times")
-        reset_process = (Process("reset_proc")
+        reset_process = (MethodProcess("reset_proc")
                          >> a.reset)
-        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
         ## set up non-compiled utility commands
-        @Context.dynamicCommand
         def clamp(x):
             a.inputs.set(x)
 
     ## input spike train
-    inputs = jnp.asarray([[0.02, 0.5, 1., 0.0]])
+    x_t = jnp.asarray([[0.02, 0.5, 1., 0.0]])
 
     targets = np.zeros((T, 4))
     targets[0, 2] = 1.
@@ -55,19 +45,19 @@ def clamp(x):
     targets = jnp.array(targets) ## gold-standard solution to check against
 
     outs = []
-    ctx.reset()
-    ctx.clamp(inputs)
-    ctx.calc_spike_times()
+    reset_process.run()
+    clamp(x_t)
+    calc_spike_times_process.run()
     for ts in range(T):
-        ctx.clamp(inputs)
-        ctx.advance(t=ts * dt, dt=dt)
+        clamp(x_t)
+        advance_process.run(t=ts * dt, dt=dt)
         ## naively extract simple statistics at time ts and print them to I/O
-        s = a.outputs.value
+        s = a.outputs.get()
         outs.append(s)
         #print(" {}: s {} ".format(ts, jnp.squeeze(s)))
     outs = jnp.concatenate(outs, axis=0)
 
     ## output should equal input
     assert_array_equal(outs, targets)
 
-#test_latencyCell1()
+test_latencyCell1()

tests/components/input_encoders/test_phasorCell.py

@@ -1,16 +1,9 @@
 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 PhasorCell
-#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 MethodProcess, Context
 
 
 def test_phasorCell1():
@@ -24,29 +17,26 @@ def test_phasorCell1():
     with Context(name) as ctx:
         a = PhasorCell(name="a", n_units=1, target_freq=1000., disable_phasor=True, key=subkeys[0])
 
-        advance_process = (Process("advance_proc")
+        advance_process = (MethodProcess("advance_proc")
                            >> a.advance_state)
-        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")
 
         ## set up non-compiled utility commands
-        @Context.dynamicCommand
         def clamp(x):
             a.inputs.set(x)
 
     ## input spike train
     x_seq = jnp.asarray([[1., 1., 0., 0., 1.]], dtype=jnp.float32)
 
     outs = []
-    ctx.reset()
+    reset_process.run()
     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.outputs.value)
+        clamp(x_t)
+        advance_process.run(t=ts * 1., dt=dt)
+        outs.append(a.outputs.get())
         #print(a.outputs.value)
     outs = jnp.concatenate(outs, axis=1)
     #print(outs)

tests/components/input_encoders/test_poissonCell.py

@@ -1,16 +1,10 @@
 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 PoissonCell
-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 MethodProcess, Context
 
 
 def test_poissonCell1():
@@ -24,29 +18,26 @@ def test_poissonCell1():
     with Context(name) as ctx:
         a = PoissonCell(name="a", n_units=1, target_freq=1000., key=subkeys[0])
 
-        advance_process = (Process("advance_proc")
+        advance_process = (MethodProcess("advance_proc")
                            >> a.advance_state)
-        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")
 
         ## set up non-compiled utility commands
-        @Context.dynamicCommand
         def clamp(x):
             a.inputs.set(x)
 
     ## input spike train
     x_seq = jnp.asarray([[1., 1., 0., 0., 1.]], dtype=jnp.float32)
 
     outs = []
-    ctx.reset()
+    reset_process.run()
     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.outputs.value)
+        clamp(x_t)
+        advance_process.run(t=ts * 1., dt=dt)
+        outs.append(a.outputs.get())
     outs = jnp.concatenate(outs, axis=1)
 
     ## output should equal input

tests/components/neurons/graded/test_bernoulliErrorCell.py

@@ -1,19 +1,10 @@
 # %%
 
 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 BernoulliErrorCell
-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 MethodProcess, Context
 
 def test_bernoulliErrorCell():
   np.random.seed(42)
@@ -25,21 +16,12 @@ def test_bernoulliErrorCell():
     a = BernoulliErrorCell(
       name="a", n_units=1, batch_size=1, input_logits=False, shape=None
     )
-    advance_process = (Process("advance_proc") >> a.advance_state)
-    ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
-    reset_process = (Process("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")
+    advance_process = (MethodProcess("advance_proc") >> a.advance_state)
+    reset_process = (MethodProcess("reset_proc") >> a.reset)
 
-    @Context.dynamicCommand
     def clamp(x):
       a.p.set(x)
 
-    @Context.dynamicCommand
     def clamp_target(x):
       a.target.set(x)
 
@@ -50,13 +32,13 @@ def clamp_target(x):
   y_seq = jnp.asarray([[-2.8193381, -4976.9263, -2.1224928, -2939.0425, -1233.3916, -0.24662945, -708.30042, 0.28213939, 3550.8477, 1.3651246]], dtype=jnp.float32)
 
   outs = []
-  ctx.reset()
+  reset_process.run()
   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)
+      clamp(x_t)
       target_xt = jnp.array([[target_seq[0, ts]]])
-      ctx.clamp_target(target_xt)
-      ctx.run(t=ts * 1., dt=dt)
+      clamp_target(target_xt)
+      advance_process.run(t=ts * 1., dt=dt)
       outs.append(a.dp.value)
   outs = jnp.concatenate(outs, axis=1)
   # print(outs)