update components and their related test cases

Author: rxng8

ngclearn/components/other/expKernel.py

@@ -86,7 +86,9 @@ def advance_state(self, t):
     def reset(self):
         restVals = jnp.zeros((self.batch_size, self.n_units)) ## inputs, epsp
         restTensor = jnp.zeros([self.win_len, self.batch_size, self.n_units], jnp.float32)  ## tf
-        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.epsp.set(restVals)
         self.tf.set(restTensor)
 

ngclearn/components/other/varTrace.py

@@ -126,7 +126,9 @@ def advance_state(self, dt):
     @compilable
     def reset(self):
         restVals = jnp.zeros((self.batch_size, self.n_units))
-        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.trace.set(restVals)
 

ngclearn/components/synapses/patched/hebbianPatchedSynapse.py

@@ -277,15 +277,18 @@ def evolve(self):
         self.dBiases.set(dBiases)
 
     @compilable
-    def reset(self, batch_size, shape):
-        preVals = jnp.zeros((batch_size, shape[0]))
-        postVals = jnp.zeros((batch_size, shape[1]))
-        not self.inputs.targeted and self.inputs.set(preVals) # inputs
+    def reset(self):
+        preVals = jnp.zeros((self.batch_size, self.shape[0]))
+        postVals = jnp.zeros((self.batch_size, self.shape[1]))
+        # 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(preVals) # inputs
         self.outputs.set(postVals) # outputs
         self.pre.set(preVals) # pre
         self.post.set(postVals) # post
-        self.dWeights.set(jnp.zeros(shape)) # dW
-        self.dBiases.set(jnp.zeros(shape[1])) # db
+        self.dWeights.set(jnp.zeros(self.shape)) # dW
+        self.dBiases.set(jnp.zeros(self.shape[1])) # db
+
 
     @classmethod
     def help(cls): ## component help function

ngclearn/components/synapses/patched/patchedSynapse.py

@@ -154,21 +154,23 @@ def advance_state(self):
         self.outputs.set(outputs)
 
     @compilable
-    def reset(self, batch_size, shape):
-        preVals = jnp.zeros((batch_size, shape[0]))
-        postVals = jnp.zeros((batch_size, shape[1]))
+    def reset(self):
+        preVals = jnp.zeros((self.batch_size, self.shape[0]))
+        postVals = jnp.zeros((self.batch_size, self.shape[1]))
         inputs = preVals
         outputs = postVals
-        not self.inputs.targeted and self.inputs.set(inputs)
+        # 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(inputs)
         self.outputs.set(outputs)
 
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         if self.bias_init != None:
-            jnp.savez(file_name, weights=self.weights.value,
-                      biases=self.biases.value)
+            jnp.savez(file_name, weights=self.weights.get(),
+                      biases=self.biases.get())
         else:
-            jnp.savez(file_name, weights=self.weights.value)
+            jnp.savez(file_name, weights=self.weights.get())
 
     def load(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"

tests/components/other/test_expKernel.py

@@ -1,16 +1,8 @@
 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 ExpKernel
-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_expKernel1():
     name = "expKernel_ctx"
@@ -25,16 +17,12 @@ def test_expKernel1():
             name="a", n_units=1, dt=1., tau_w=500., nu=4., 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)
 
@@ -44,16 +32,16 @@ def clamp(x):
     y_seq = jnp.asarray([[0., 1., 0.998002, 0.996008, 1.9940181]], 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.epsp.value)
+        clamp(x_t)
+        advance_process.run(t=ts * 1., dt=dt)
+        outs.append(a.epsp.get())
     outs = jnp.concatenate(outs, axis=1)
     #print(outs)
 
     ## output should equal input
     np.testing.assert_allclose(outs, y_seq, atol=1e-8)
 
-#test_expKernel1()
+test_expKernel1()

tests/components/other/test_varTrace.py

@@ -1,16 +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 VarTrace
-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_varTrace1():
     name = "trace_ctx"
@@ -26,35 +21,32 @@ def test_varTrace1():
             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)
     ## desired output pulses
-    y_seq = x_seq * trace_increment 
+    y_seq = x_seq * trace_increment
 
     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)
     #print(outs)
 
     ## output should equal input
     assert_array_equal(outs, y_seq)
 
-#test_varTrace1()
+test_varTrace1()

tests/components/synapses/patched/test_hebbianPatchedSynapse.py

@@ -1,19 +1,12 @@
 # %%
 
 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 HebbianPatchedSynapse
-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_hebbianPatchedSynapse():
   np.random.seed(42)
@@ -31,58 +24,45 @@ def test_hebbianPatchedSynapse():
 
   with Context(name) as ctx:
     a = HebbianPatchedSynapse(
-      name="a", 
-      shape=shape, 
-      n_sub_models=n_sub_models, 
+      name="a",
+      shape=shape,
+      n_sub_models=n_sub_models,
       stride_shape=stride_shape,
       resist_scale=resist_scale,
       batch_size=batch_size
     )
 
-    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")
-    evolve_process = (Process("evolve_proc") >> a.evolve)
-    ctx.wrap_and_add_command(jit(evolve_process.pure), name="evolve")
-
-    # Compile and add commands
-    # reset_cmd, reset_args = ctx.compile_by_key(a, compile_key="reset")
-    # ctx.add_command(wrap_command(jit(reset_cmd)), name="reset")
-    # advance_cmd, advance_args = ctx.compile_by_key(a, compile_key="advance_state")
-    # ctx.add_command(wrap_command(jit(advance_cmd)), name="run")
-    # evolve_cmd, evolve_args = ctx.compile_by_key(a, compile_key="evolve")
-    # ctx.add_command(wrap_command(jit(evolve_cmd)), name="evolve")
-
-    @Context.dynamicCommand
+    advance_process = (MethodProcess("advance_proc") >> a.advance_state)
+    reset_process = (MethodProcess("reset_proc") >> a.reset)
+    evolve_process = (MethodProcess("evolve_proc") >> a.evolve)
+
     def clamp_inputs(x):
       a.inputs.set(x)
 
-    @Context.dynamicCommand
     def clamp_pre(x):
       a.pre.set(x)
 
-    @Context.dynamicCommand
     def clamp_post(x):
       a.post.set(x)
 
-  a.weights.set(jnp.ones((12, 12)) * 0.5) 
+  a.weights.set(jnp.ones((12, 12)) * 0.5)
 
   in_pre = jnp.ones((10, 12)) * 1.0
   in_post = jnp.ones((10, 12)) * 0.75
 
-  ctx.reset()
+  reset_process.run()
   clamp_pre(in_pre)
   clamp_post(in_post)
-  ctx.run(t=1. * dt, dt=dt)
-  ctx.evolve(t=1. * dt, dt=dt)
+  advance_process.run(t=1. * dt, dt=dt)
+  evolve_process.run(t=1. * dt, dt=dt)
 
-  print(a.weights.value)
+  print(a.weights.get())
 
   # Basic assertions to check learning dynamics
-  assert a.weights.value.shape == (12, 12), ""
-  assert a.weights.value[0, 0] == 0.5, ""
+  assert a.weights.get().shape == (12, 12), ""
+  assert a.weights.get()[0, 0] == 0.5, ""
+
 
+test_hebbianPatchedSynapse()
 
 
-# test_hebbianPatchedSynapse()

tests/components/synapses/patched/test_patchedSynapse.py

@@ -1,18 +1,12 @@
 # %%
 
 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 PatchedSynapse
-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_patchedSynapse():
@@ -39,31 +33,25 @@ def test_patchedSynapse():
       bias_init={"dist": "constant", "value": 0.0}
     )
 
-    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")
-
-    # Compile and add commands
-    # reset_cmd, reset_args = ctx.compile_by_key(a, compile_key="reset")
-    # ctx.add_command(wrap_command(jit(reset_cmd)), name="reset")
-    # advance_cmd, advance_args = ctx.compile_by_key(a, compile_key="advance_state")
-    # ctx.add_command(wrap_command(jit(advance_cmd)), name="run")
+    advance_process = (MethodProcess("advance_proc") >> a.advance_state)
+    reset_process = (MethodProcess("reset_proc") >> a.reset)
 
-    @Context.dynamicCommand
     def clamp_inputs(x):
       a.inputs.set(x)
 
   inputs_seq = jnp.asarray(np.random.randn(1, 12))
-  weights = a.weights.value
-  biases = a.biases.value
+  weights = a.weights.get()
+  biases = a.biases.get()
   expected_outputs = (jnp.matmul(inputs_seq, weights) * resist_scale) + biases
   outputs_outs = []
-  ctx.reset()
-  ctx.clamp_inputs(inputs_seq)
-  ctx.run(t=0., dt=dt)
-  outputs_outs.append(a.outputs.value)
+  reset_process.run()
+  clamp_inputs(inputs_seq)
+  advance_process.run(t=0., dt=dt)
+  outputs_outs.append(a.outputs.get())
   outputs_outs = jnp.concatenate(outputs_outs, axis=1)
   # Verify outputs match expected values
   np.testing.assert_allclose(outputs_outs, expected_outputs, atol=1e-5)
 
+
+test_patchedSynapse()
+