minor revisions to input-encoders, revised phasor-cell w/ unit-test

Author: Alexander Ororbia

ngclearn/components/input_encoders/latencyCell.py

@@ -142,6 +142,8 @@ class LatencyCell(JaxComponent):
 
         num_steps: number of discrete time steps to consider for normalized latency
             code (only useful if "normalize" is set to True) (Default: 1)
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     # Define Functions

ngclearn/components/input_encoders/phasorCell.py

@@ -1,8 +1,13 @@
-from ngclearn import resolver, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
-from ngclearn.utils import tensorstats
 from jax import numpy as jnp, random
-from ngcsimlib.logger import warn
+from ngclearn.utils import tensorstats
+from ngcsimlib.deprecators import deprecate_args
+from ngcsimlib.logger import info, warn
+
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+
 
 class PhasorCell(JaxComponent):
     """
@@ -12,8 +17,9 @@ class PhasorCell(JaxComponent):
     | inputs - input (takes in external signals)
     | --- Cell State Compartments: ---
     | key - JAX PRNG key
+    | angles - current angle of phasor 
     | --- Cell Output Compartments: ---
-    | outputs - output
+    | outputs - output of phasor cell
     | tols - time-of-last-spike
 
     Args:
@@ -23,6 +29,8 @@ class PhasorCell(JaxComponent):
 
         target_freq: maximum frequency (in Hertz) of this spike train
             (must be > 0.)
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     # Define Functions
@@ -63,8 +71,7 @@ def validate(self, dt=None, **validation_kwargs):
             return False
         ## check for unstable combinations of dt and target-frequency
         # meta-params
-        events_per_timestep = (
-                                  dt / 1000.) * self.target_freq  ##
+        events_per_timestep = (dt / 1000.) * self.target_freq  ##
         # compute scaled probability
         if events_per_timestep > 1.:
             valid = False
@@ -78,9 +85,9 @@ def validate(self, dt=None, **validation_kwargs):
             )
         return valid
 
+    @transition(output_compartments=["outputs", "tols", "key", "angles"])
     @staticmethod
-    def _advance_state(t, dt, target_freq, key,
-                       inputs, angles, tols, base_scale):
+    def advance_state(t, dt, target_freq, key, inputs, angles, tols, base_scale):
         ms_per_second = 1000  # ms/s
         events_per_ms = target_freq / ms_per_second  # e/s s/ms -> e/ms
         ms_per_event = 1 / events_per_ms  # ms/e
@@ -105,27 +112,13 @@ def _advance_state(t, dt, target_freq, key,
 
         return outputs, tols, key, updated_angles
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, tols, key, angles):
-        self.outputs.set(outputs)
-        self.tols.set(tols)
-        self.key.set(key)
-        self.angles.set(angles)
-
+    @transition(output_compartments=["inputs", "outputs", "tols", "angles", "key"])
     @staticmethod
-    def _reset(batch_size, n_units, key, target_freq):
+    def reset(batch_size, n_units, key, target_freq):
         restVals = jnp.zeros((batch_size, n_units))
         key, subkey = random.split(key, 2)
         return restVals, restVals, restVals, restVals, key
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, tols, angles, key):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.tols.set(tols)
-        self.key.set(key)
-        self.angles.set(angles)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, key=self.key.value)
@@ -154,7 +147,7 @@ def help(cls):  ## component help function
         hyperparams = {
             "n_units": "Number of neuronal cells to model in this layer",
             "batch_size": "Batch size dimension of this component",
-            "target_freq": "Maximum spike frequency of the train produced",
+            "target_freq": "Maximum spike frequency of the (spike) train produced",
         }
         info = {cls.__name__: properties,
                 "compartments": compartment_props,

ngclearn/components/input_encoders/poissonCell.py

@@ -27,6 +27,8 @@ class PoissonCell(JaxComponent):
         n_units: number of cellular entities (neural population size)
 
         target_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.)
+
+        batch_size: batch size dimension of this cell (Default: 1)
     """
 
     @deprecate_args(max_freq="target_freq")

ngclearn/utils/patch_utils.py

@@ -118,7 +118,7 @@ def create_patches(self, add_frame=False, center=True):
 
 
 
-def generate_patch_set(x_batch, patch_size=(8, 8), max_patches=50, center=True): ## scikit
+def generate_patch_set(x_batch, patch_size=(8, 8), max_patches=50, center=True, seed=1234): ## scikit
     """
     Generates a set of patches from an array/list of image arrays (via
     random sampling with replacement). This uses scikit-learn's patch creation
@@ -134,6 +134,8 @@ def generate_patch_set(x_batch, patch_size=(8, 8), max_patches=50, center=True):
 
         center: centers each patch by subtracting the patch mean (per-patch)
 
+        seed: seed to control the random state of internal patch sampling
+
     Returns:
         an array (D x (pH * pW)), where each row is a flattened patch sample
     """
@@ -143,7 +145,7 @@ def generate_patch_set(x_batch, patch_size=(8, 8), max_patches=50, center=True):
     for s in range(_x_batch.shape[0]):
         xs = _x_batch[s, :]
         xs = xs.reshape(px, py)
-        patches = extract_patches_2d(xs, patch_size, max_patches=max_patches)#, random_state=69)
+        patches = extract_patches_2d(xs, patch_size, max_patches=max_patches, random_state=seed)#, random_state=69)
         patches = np.reshape(patches, (len(patches), -1)) # flatten each patch in set
         if s > 0:
             p_batch = np.concatenate((p_batch,patches),axis=0)

tests/components/input_encoders/test_bernoulliCell.py

@@ -52,4 +52,4 @@ def clamp(x):
     assert_array_equal(outs, x_seq)
     print(outs)
 
-test_bernoulliCell1()
+#test_bernoulliCell1()

tests/components/input_encoders/test_latencyCell.py

@@ -13,7 +13,7 @@
 from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 
-def test_latencyCell():
+def test_latencyCell1():
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
@@ -69,4 +69,4 @@ def clamp(x):
     ## output should equal input
     assert_array_equal(outs, targets)
 
-test_latencyCell()
+#test_latencyCell1()

tests/components/input_encoders/test_phasorCell.py

@@ -0,0 +1,55 @@
+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 PhasorCell
+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_phasorCell1():
+    ## 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 = PhasorCell(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)
+        #print(a.outputs.value)
+    outs = jnp.concatenate(outs, axis=1)
+
+    ## output should equal input
+    assert_array_equal(outs, x_seq)
+
+test_phasorCell1()

tests/components/input_encoders/test_poissonCell.py

@@ -13,7 +13,7 @@
 from ngcsimlib.utils.compartment import Get_Compartment_Batch
 
 
-def test_poissonCell():
+def test_poissonCell1():
     ## create seeding keys
     dkey = random.PRNGKey(1234)
     dkey, *subkeys = random.split(dkey, 6)
@@ -51,4 +51,4 @@ def clamp(x):
     ## output should equal input
     assert_array_equal(outs, x_seq)
 
-test_poissonCell()
+#test_poissonCell1()