bernoulli and poisson cells revised, unit-tested

Author: Alexander Ororbia

ngclearn/components/input_encoders/bernoulliCell.py

@@ -1,8 +1,6 @@
-#from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
-from jax import numpy as jnp, random, jit
+from jax import numpy as jnp, random
 from ngclearn.utils import tensorstats
-from functools import partial
 from ngcsimlib.deprecators import deprecate_args
 from ngcsimlib.logger import info, warn
 
@@ -49,6 +47,11 @@ def __init__(self, name, n_units, batch_size=1, **kwargs):
     @staticmethod
     def advance_state(t, key, inputs, tols):
         ## NOTE: should `inputs` be checked if bounded to [0,1]?
+        # print(key)
+        # print(t)
+        # print(inputs.shape)
+        # print(tols.shape)
+        # print("-----")
         key, *subkeys = random.split(key, 3)
         outputs = random.bernoulli(subkeys[0], p=inputs).astype(jnp.float32)
         # Updates time-of-last-spike (tols) variable:

ngclearn/components/input_encoders/poissonCell.py

@@ -1,28 +1,12 @@
-from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
-from jax import numpy as jnp, random, jit
+from jax import numpy as jnp, random
 from ngclearn.utils import tensorstats
-from functools import partial
 from ngcsimlib.deprecators import deprecate_args
 from ngcsimlib.logger import info, warn
 
-@jit
-def _update_times(t, s, tols):
-    """
-    Updates time-of-last-spike (tols) variable.
-
-    Args:
-        t: current time (a scalar/int value)
-
-        s: binary spike vector
-
-        tols: current time-of-last-spike variable
-
-    Returns:
-        updated tols variable
-    """
-    _tols = (1. - s) * tols + (s * t)
-    return _tols
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
 class PoissonCell(JaxComponent):
     """
@@ -79,33 +63,27 @@ def validate(self, dt=None, **validation_kwargs):
             )
         return valid
 
+    @transition(output_compartments=["outputs", "tols", "key"])
     @staticmethod
-    def _advance_state(t, dt, target_freq, key, inputs, tols):
+    def advance_state(t, dt, target_freq, key, inputs, tols):
         key, *subkeys = random.split(key, 2)
         pspike = inputs * (dt / 1000.) * target_freq
         eps = random.uniform(subkeys[0], inputs.shape, minval=0., maxval=1.,
                              dtype=jnp.float32)
         outputs = (eps < pspike).astype(jnp.float32)
-        tols = _update_times(t, outputs, tols)
-        return outputs, tols, key
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, tols, key):
-        self.outputs.set(outputs)
-        self.tols.set(tols)
-        self.key.set(key)
+        # Updates time-of-last-spike (tols) variable:
+        # output = s = binary spike vector
+        # tols = current time-of-last-spike variable
+        tols = (1. - outputs) * tols + (outputs * t)
+        return outputs, tols, key
 
+    @transition(output_compartments=["inputs", "outputs", "tols"])
     @staticmethod
-    def _reset(batch_size, n_units):
+    def reset(batch_size, n_units):
         restVals = jnp.zeros((batch_size, n_units))
         return restVals, restVals, restVals
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, tols):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs) #None
-        self.tols.set(tols)
-
     def save(self, directory, **kwargs):
         target_freq = (self.target_freq if isinstance(self.target_freq, float)
                        else jnp.ones([[self.target_freq]]))
@@ -121,12 +99,9 @@ def load(self, directory, **kwargs):
     @classmethod
     def help(cls):  ## component help function
         properties = {
-            "cell_type": "PoissonCell - samples input to produce spikes, "
-                         "where dimension is a probability proportional to "
-                         "the dimension's magnitude/value/intensity and "
-                         "constrained by a maximum spike frequency (spikes "
-                         "follow "
-                         "a Poisson distribution)"
+            "cell_type": "PoissonCell - samples input to produce spikes, where dimension is a probability proportional "
+                         "to the dimension's magnitude/value/intensity and constrained by a maximum spike frequency "
+                         "(spikes follow a Poisson distribution)"
         }
         compartment_props = {
             "inputs":

tests/components/input_encoders/test_bernoulliCell.py

@@ -17,11 +17,9 @@ def test_bernoulliCell1():
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
-    # in_dim = 9  # ... dimension of patch data ...
-    # hid_dim = 9  # ... number of atoms in the dictionary matrix
     dt = 1.  # ms
-    T = 300  # ms # (OR) number of E-steps to take during inference
-    # ---- build a sparse coding linear generative model with a Cauchy prior ----
+    #T = 300  # ms
+    # ---- build a simple Bernoulli cell system ----
     with Context("Circuit") as ctx:
         a = BernoulliCell(name="a", n_units=1, key=subkeys[0])
 
@@ -33,17 +31,6 @@ def test_bernoulliCell1():
                          >> a.reset)
         ctx.wrap_and_add_command(reset_process.pure, name="reset")
 
-        # ## create and compile core simulation commands
-        # 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="advance")
-
-
         ## set up non-compiled utility commands
         @Context.dynamicCommand
         def clamp(x):
@@ -53,16 +40,16 @@ def clamp(x):
     x_seq = jnp.asarray([[1., 1., 0., 0., 1.]], dtype=jnp.float32)
 
     outs = []
-    #ctx.reset()
+    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=1.)
+        ctx.run(t=ts*1., dt=dt)
         outs.append(a.outputs.value)
     outs = jnp.concatenate(outs, axis=1)
 
     ## output should equal input
     assert_array_equal(outs, x_seq)
-    #print(outs)
+    print(outs)
 
 test_bernoulliCell1()

tests/components/input_encoders/test_poissonCell.py

@@ -0,0 +1,54 @@
+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
+
+
+def test_poissonCell():
+    ## 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:
+        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")
+
+        reset_process = (Process()
+                         >> a.reset)
+        ctx.wrap_and_add_command(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()
+    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)
+    outs = jnp.concatenate(outs, axis=1)
+
+    ## output should equal input
+    assert_array_equal(outs, x_seq)
+
+test_poissonCell()