revised unit-tests to pass globally; some minor patches to phasor-cell and lif

Author: Alexander Ororbia

ngclearn/components/input_encoders/phasorCell.py

@@ -44,7 +44,7 @@ def __init__(self, name, n_units, target_freq=63.75, batch_size=1,
         ## Layer Size Setup
         self.batch_size = batch_size
         self.n_units = n_units
-        _key, subkey = random.split(self.key.value, 2)
+        _key, *subkey = random.split(self.key.value, 3)
         self.key.set(_key)
         ## Compartment setup
         restVals = jnp.zeros((self.batch_size, self.n_units))
@@ -62,7 +62,7 @@ def __init__(self, name, n_units, target_freq=63.75, batch_size=1,
         # alpha = ((random.normal(subkey, self.angles.value.shape) * (jnp.sqrt(target_freq) / target_freq)) + 1)
         # beta = random.poisson(subkey, lam=target_freq, shape=self.angles.value.shape) / target_freq
 
-        self.base_scale = random.poisson(subkey, lam=target_freq, shape=self.angles.value.shape) / target_freq
+        self.base_scale = random.poisson(subkey[0], lam=target_freq, shape=self.angles.value.shape) / target_freq
 
     def validate(self, dt=None, **validation_kwargs):
         valid = super().validate(**validation_kwargs)
@@ -95,11 +95,11 @@ def advance_state(t, dt, target_freq, key, inputs, angles, tols, base_scale):
         angle_per_event = 2 * jnp.pi  # rad / e
         angle_per_timestep = angle_per_event / time_step_per_event  # rad / e
         # * e/ts -> rad / ts
-        key, subkey = random.split(key, 2)
+        key, *subkey = random.split(key, 3)
         # scatter = random.uniform(subkey, angles.shape, minval=0.5,
         #                          maxval=1.5) * base_scale
 
-        scatter = ((random.normal(subkey, angles.shape) * 0.2) + 1) * base_scale
+        scatter = ((random.normal(subkey[0], angles.shape) * 0.2) + 1) * base_scale
         scattered_update = angle_per_timestep * scatter
         scaled_scattered_update = scattered_update * inputs
 
@@ -116,7 +116,7 @@ def advance_state(t, dt, target_freq, key, inputs, angles, tols, base_scale):
     @staticmethod
     def reset(batch_size, n_units, key, target_freq):
         restVals = jnp.zeros((batch_size, n_units))
-        key, subkey = random.split(key, 2)
+        key, *subkey = random.split(key, 3)
         return restVals, restVals, restVals, restVals, key
 
     def save(self, directory, **kwargs):

ngclearn/components/neurons/spiking/LIFCell.py

@@ -1,16 +1,3 @@
-"""
-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 ngclearn import resolver, Component, Compartment
-from ngclearn.components.jaxComponent import JaxComponent
-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.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, random, jit, nn
 from functools import partial

tests/components/input_encoders/test_bernoulliCell.py

@@ -3,33 +3,34 @@
 import numpy as np
 np.random.seed(42)
 from ngclearn.components import BernoulliCell
-from ngcsimlib.compilers import compile_command, wrap_command
+#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.component import Component
+#from ngcsimlib.compartment import Compartment
 from ngcsimlib.context import Context
-from ngcsimlib.utils.compartment import Get_Compartment_Batch
+#from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 
 def test_bernoulliCell1():
+    name = "bernoulli_ctx"
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
     dt = 1.  # ms
     #T = 300  # ms
     # ---- build a simple Bernoulli cell system ----
-    with Context("Circuit") as ctx:
+    with Context(name) as ctx:
         a = BernoulliCell(name="a", n_units=1, key=subkeys[0])
 
         advance_process = (Process()
                            >> a.advance_state)
-        ctx.wrap_and_add_command(advance_process.pure, name="run")
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
         reset_process = (Process()
                          >> a.reset)
-        ctx.wrap_and_add_command(reset_process.pure, name="reset")
+        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
         ## set up non-compiled utility commands
         @Context.dynamicCommand

tests/components/input_encoders/test_latencyCell.py

@@ -14,14 +14,15 @@
 
 
 def test_latencyCell1():
+    name = "latency_ctx"
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
     T = 50  # 100 #5  ## number of simulation steps to run
     dt = 1.  # 0.1  # ms ## compute integration time constant
     tau = 1.
     # ---- build a simple Poisson cell system ----
-    with Context("Circuit") as ctx:
+    with Context(name) as ctx:
         a = LatencyCell(
             "a", n_units=4, tau=tau, threshold=0.01, linearize=True,
             normalize=True, num_steps=T, clip_spikes=False
@@ -30,13 +31,13 @@ def test_latencyCell1():
         ## create and compile core simulation commands
         advance_process = (Process()
                            >> a.advance_state)
-        ctx.wrap_and_add_command(advance_process.pure, name="advance")
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="advance")
         calc_spike_times_process = (Process()
                                     >> a.calc_spike_times)
-        ctx.wrap_and_add_command(calc_spike_times_process.pure, name="calc_spike_times")
+        ctx.wrap_and_add_command(jit(calc_spike_times_process.pure), name="calc_spike_times")
         reset_process = (Process()
                          >> a.reset)
-        ctx.wrap_and_add_command(reset_process.pure, name="reset")
+        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
         ## set up non-compiled utility commands
         @Context.dynamicCommand

tests/components/input_encoders/test_phasorCell.py

@@ -3,24 +3,25 @@
 import numpy as np
 np.random.seed(42)
 from ngclearn.components import PhasorCell
-from ngcsimlib.compilers import compile_command, wrap_command
+#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 ngcsimlib.component import Component
+#from ngcsimlib.compartment import Compartment
+#from ngcsimlib.context import Context
+#from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 
 def test_phasorCell1():
+    name = "phasor_ctx"
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
     dt = 1.  # ms
     # T = 300  # ms
     # ---- build a simple Poisson cell system ----
-    with Context("Circuit") as ctx:
+    with Context(name) as ctx:
         a = PhasorCell(name="a", n_units=1, target_freq=1000., key=subkeys[0])
 
         advance_process = (Process()
@@ -52,4 +53,4 @@ def clamp(x):
     ## output should equal input
     assert_array_equal(outs, x_seq)
 
-test_phasorCell1()
+#test_phasorCell1()

tests/components/input_encoders/test_poissonCell.py

@@ -14,22 +14,23 @@
 
 
 def test_poissonCell1():
+    name = "poisson_ctx"
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
     dt = 1.  # ms
     # T = 300  # ms
     # ---- build a simple Poisson cell system ----
-    with Context("Circuit") as ctx:
+    with Context(name) as ctx:
         a = PoissonCell(name="a", n_units=1, target_freq=1000., key=subkeys[0])
 
         advance_process = (Process()
                            >> a.advance_state)
-        ctx.wrap_and_add_command(advance_process.pure, name="run")
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
         reset_process = (Process()
                          >> a.reset)
-        ctx.wrap_and_add_command(reset_process.pure, name="reset")
+        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
         ## set up non-compiled utility commands
         @Context.dynamicCommand

tests/components/neurons/graded/test_rateCell.py

@@ -13,13 +13,14 @@
 from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 def test_LIFCell1():
+    name = "lif_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("Circuit") as ctx:
+    with Context(name) as ctx:
         a = LIFCell(
             name="a", n_units=1, tau_m=5., resist_m=30., key=subkeys[0]
         )
@@ -65,4 +66,4 @@ def clamp(x):
     ## output should equal input
     assert_array_equal(outs, y_seq)
 
-test_LIFCell1()
+#test_LIFCell1()

tests/components/neurons/spiking/test_LIFCell.py

@@ -13,13 +13,14 @@
 from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 def test_sLIFCell1():
+    name = "slif_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("Circuit") as ctx:
+    with Context(name) as ctx:
         a = SLIFCell(
             name="a", n_units=1, tau_m=50., resist_m=10., thr=0.3, key=subkeys[0]
         )

tests/components/neurons/spiking/test_sLIFCell.py

@@ -13,24 +13,25 @@
 from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 def test_expKernel1():
+    name = "expKernel_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("Circuit") as ctx:
+    with Context(name) as ctx:
         a = ExpKernel(
             name="a", n_units=1, dt=1., tau_w=500., nu=4., key=subkeys[0]
         )
 
         advance_process = (Process()
                            >> a.advance_state)
-        ctx.wrap_and_add_command(advance_process.pure, name="run")
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
 
         reset_process = (Process()
                          >> a.reset)
-        ctx.wrap_and_add_command(reset_process.pure, name="reset")
+        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
         ## set up non-compiled utility commands
         @Context.dynamicCommand