update code

Author: rxng8

tests/components/neurons/graded/test_RateCell.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 RateCell
-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
 
 
 def test_RateCell1():
@@ -26,17 +20,9 @@ def test_RateCell1():
       threshold=("none", 0.), integration_type="euler",
       batch_size=1, resist_scale=1., shape=None, is_stateful=True
     )
-    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")
+    advance_process = (MethodProcess("advance_proc") >> a.advance_state)
+    reset_process = (MethodProcess("reset_proc") >> a.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")
-
-    @Context.dynamicCommand
     def clamp(x):
       a.j.set(x)
 
@@ -46,11 +32,11 @@ def clamp(x):
   y_seq = jnp.asarray([[0.02, 0.04, 0.06, 0.08, 0.09999999999999999, 0.11999999999999998, 0.13999999999999999, 0.15999999999999998, 0.17999999999999998, 0.19999999999999998]], 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)
+      advance_process.run(t=ts * 1., dt=dt)
       outs.append(a.z.value)
   outs = jnp.concatenate(outs, axis=1)
   # print(outs)

tests/components/synapses/hebbian/test_hebbianSynapse.py

@@ -1,18 +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 HebbianSynapse
-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 numpy.testing import assert_array_equal
+from ngclearn import Context, MethodProcess
 
 
 def test_hebbianSynapse():
@@ -29,37 +24,23 @@ def test_hebbianSynapse():
 
   with Context(name) as ctx:
     a = HebbianSynapse(
-      name="a", 
-      shape=shape, 
+      name="a",
+      shape=shape,
       resist_scale=resist_scale,
       batch_size=batch_size,
       prior = ("gaussian", 0.01)
     )
 
-    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)
 
@@ -70,16 +51,18 @@ def clamp_post(x):
   in_pre = jnp.ones((1, 10)) * 1.0
   in_post = jnp.ones((1, 5)) * 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 == (10, 5), ""
-  assert a.weights.value[0, 0] == 0.5, ""
+  assert a.weights.get().shape == (10, 5), ""
+  assert a.weights.get()[0, 0] == 0.5, ""
+
+test_hebbianSynapse()
+
 
-# test_hebbianSynapse()