attempted rewrite of bernoulli-cell

Author: Alexander Ororbia

ngclearn/components/base_monitor.py

@@ -2,8 +2,8 @@
 
 from ngclearn import Component, Compartment
 from ngclearn import numpy as np
-from ngcsimlib.utils import add_component_resolver, add_resolver_meta, \
-    get_current_path
+#from ngcsimlib.utils import add_component_resolver, add_resolver_meta, \
+from ngcsimlib.utils import get_current_path
 from ngcsimlib.logger import warn, critical
 import matplotlib.pyplot as plt
 

ngclearn/components/input_encoders/bernoulliCell.py

@@ -1,36 +1,22 @@
-from ngclearn import resolver, Component, Compartment
+#from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, random, jit
 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 Process, transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
 class BernoulliCell(JaxComponent):
     """
     A Bernoulli cell that produces spikes by sampling a Bernoulli distribution
     on-the-fly (to produce data-scaled Bernoulli spike trains).
 
     | --- Cell Input Compartments: ---
-    | inputs - input (takes in external signals)
+    | inputs - input (takes in external signals -- should be probabilities w/ values in [0,1])
     | --- Cell State Compartments: ---
     | key - JAX PRNG key
     | --- Cell Output Compartments: ---
@@ -41,10 +27,13 @@ class BernoulliCell(JaxComponent):
         name: the string name of this cell
 
         n_units: number of cellular entities (neural population size)
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     def __init__(self, name, n_units, batch_size=1, **kwargs):
-        super().__init__(name, **kwargs)
+        #super().__init__(name, **kwargs)
+        super(JaxComponent, self).__init__(name, **kwargs)
 
         ## Layer Size Setup
         self.batch_size = batch_size
@@ -56,30 +45,24 @@ def __init__(self, name, n_units, batch_size=1, **kwargs):
         self.outputs = Compartment(restVals, display_name="Spikes") # output compartment
         self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike
 
+    @transition(output_compartments=["outputs", "tols", "key"])
     @staticmethod
-    def _advance_state(t, key, inputs, tols):
+    def advance_state(t, key, inputs, tols):
+        ## NOTE: should `inputs` be checked if bounded to [0,1]?
         key, *subkeys = random.split(key, 2)
         outputs = random.bernoulli(subkeys[0], p=inputs).astype(jnp.float32)
-        tols = _update_times(t, outputs, tols)
+        # 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
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, tols, key):
-        self.outputs.set(outputs)
-        self.tols.set(tols)
-        self.key.set(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):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, key=self.key.value)

ngclearn/components/jaxComponent.py

@@ -1,6 +1,8 @@
 import time
 from jax import random
-from ngclearn import resolver, Component, Compartment
+#from ngclearn import resolver, Component, Compartment
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
 
 class JaxComponent(Component):
     """

tests/components/input_encoders/test_bernoulliCell.py

@@ -0,0 +1,65 @@
+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 ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.context import Context
+from ngcsimlib.utils.compartment import Get_Compartment_Batch
+
+
+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 ----
+    with Context("Circuit") as ctx:
+        a = BernoulliCell(name="a", n_units=1, key=subkeys[0])
+
+        myProcess = (Process()
+                     >> a.advance_state)
+
+        ctx.wrap_and_add_command(myProcess.pure, name="run")
+
+        ## 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):
+            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.advance(t=ts*1., dt=1.)
+        outs.append(a.outputs.value)
+    outs = jnp.concatenate(outs, axis=1)
+
+    ## output should equal input
+    assert_array_equal(outs, x_seq)
+    #print(outs)
+
+test_bernoulliCell1()

tests/components/neurons/graded/test_rateCell.py

@@ -48,4 +48,4 @@ def clamp(x):
           circuit.advance(t=ts*1., dt=1.)
 
   print(a.zF.value)
-  # assertion here if needed!
+  # assertion here if needed!