fixed minor issues in input-encoders, further revisions to docs for v3

Author: Alexander Ororbia

docs/tutorials/neurocog/hebbian.md

@@ -21,30 +21,29 @@ Specifically, we will zoom in on two particular code snippets from
 below:
 
 ```python
-Wab = HebbianSynapse(name="Wab", shape=(1, 1), eta=1.,   signVal=-1.,
-                     wInit=("constant", 1., None), w_bound=0., key=subkeys[3])
+Wab = HebbianSynapse(
+    name="Wab", shape=(1, 1), eta=1.,   signVal=-1., wInit=("constant", 1., None), w_bound=0., key=subkeys[3]
+)
 
 # wire output compartment (rate-coded output zF) of RateCell `a` to input compartment of HebbianSynapse `Wab`
-Wab.inputs << a.zF
+a.zF >> Wab.inputs
 # wire output compartment of HebbianSynapse `Wab` to input compartment (electrical current j) RateCell `b`
-b.j << Wab.outputs
+Wab.outputs >> b.j
 
 # wire output compartment (rate-coded output zF) of RateCell `a` to presynaptic compartment of HebbianSynapse `Wab`
-Wab.pre << a.zF
+a.zF >> Wab.pre
 # wire output compartment (rate-coded output zF) of RateCell `b` to postsynaptic compartment of HebbianSynapse `Wab`
-Wab.post << b.zF
+b.zF >> Wab.post
 ```
 
 as well as (a bit later in the model construction code):
 
 ```python
-evolve_process = (JaxProcess()
+evolve_process = (MethodProcess()
                   >> a.evolve)
-circuit.wrap_and_add_command(jit(evolve_process.pure), name="evolve")
 
-advance_process = (JaxProcess()
+advance_process = (MethodProcess()
                    >> a.advance_state)
-circuit.wrap_and_add_command(jit(advance_process.pure), name="advance")
 ```
 
 Notice that beyond wiring component `a`'s values into the synapse `Wab`'s input compartment
@@ -54,7 +53,7 @@ post-synaptic compartment `Wab.post`. These compartments are specifically
 used in `Wab`'s `evolve` call and are not strictly required to be exactly
 the same as its input and output compartments. Note that, if one wanted `pre`
 and `post` to be exactly identical to `inputs` and `outputs`, one would simply need
-to write `Wab.pre << Wab.inputs` and `Wab.post << Wab.outputs` in place
+to write `Wab.inputs >> Wab.pre` and `Wab.outputs >> Wab.post` in place
 of the pre- and post-synaptic compartment calls above.
 
 The above snippets highlight two key aspects of functionality that a synapse

docs/tutorials/neurocog/input_cells.md

@@ -39,8 +39,7 @@ spike train over $100$ steps in time as follows:
 
 ```python
 from jax import numpy as jnp, random, jit
-from ngcsimlib.context import Context
-from ngclearn.utils import JaxProcess
+from ngclearn import Context, MethodProcess
 
 from ngclearn.utils.viz.raster import create_raster_plot
 ## import model-specific mechanisms
@@ -56,27 +55,24 @@ T = 100  ## number time steps to simulate
 with Context("Model") as model:
     cell = BernoulliCell("z0", n_units=10, key=subkeys[0])
 
-    advance_process = (JaxProcess()
+    advance_process = (MethodProcess("advance_proc")
                        >> cell.advance_state)
-    model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (JaxProcess()
+    reset_process = (MethodProcess("reset_proc")
                      >> cell.reset)
-    model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
-
-    @Context.dynamicCommand
-    def clamp(x):
-        cell.inputs.set(x)
+def clamp(x):
+    cell.inputs.set(x)
+    
 
 probs = jnp.asarray([[0.8, 0.2, 0., 0.55, 0.9, 0, 0.15, 0., 0.6, 0.77]], dtype=jnp.float32)
 spikes = []
-model.reset()
+reset_process.run() 
 for ts in range(T):
-    model.clamp(probs)
-    model.advance(t=ts * 1., dt=dt)
+    clamp(probs)
+    advance_process.run(t=ts * 1., dt=dt)
 
-    s_t = cell.outputs.value
+    s_t = cell.outputs.get()
     spikes.append(s_t)
 spikes = jnp.concatenate(spikes, axis=0)
 create_raster_plot(spikes, plot_fname="input_cell_raster.jpg")
@@ -121,7 +117,7 @@ and by replacing the line that has the `BernoulliCell` call with the
 following line instead:
 
 ```python
-cell = PoissonCell("z0", n_units=10, max_freq=63.75, key=subkeys[0])
+cell = PoissonCell("z0", n_units=10, target_freq=63.75, key=subkeys[0])
 ```
 
 Running the code with the two above small modifications will
@@ -149,12 +145,12 @@ mu = 0.
 probs = jnp.asarray([[1.]],dtype=jnp.float32)
 for _ in range(n_trials):
     spikes = []
-    model.reset()
+    reset_process.run() 
     for ts in range(T):
-        model.clamp(probs)
-        model.advance(t=ts*1., dt=dt)
+        clamp(probs)
+        advance_process.run(t=ts * 1., dt=dt)
 
-        s_t = cell.outputs.value
+        s_t = cell.outputs.get()
         spikes.append(s_t)
     count = jnp.sum(jnp.concatenate(spikes, axis=0))
     mu += count

ngclearn/components/input_encoders/bernoulliCell.py

@@ -26,7 +26,7 @@ class BernoulliCell(JaxComponent):
         batch_size: batch size dimension of this cell (Default: 1)
     """
 
-    def __init__(self, name: str, n_units: int, batch_size: int = 1, key: Union[jax.Array, None] = None):
+    def __init__(self, name: str, n_units: int, batch_size: int = 1, key: Union[jax.Array, None] = None, **kwargs):
         super().__init__(name=name, key=key)
 
         ## Layer Size Setup
@@ -80,20 +80,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-    # def __repr__(self):
-    #     comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
-    #     maxlen = max(len(c) for c in comps) + 5
-    #     lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-    #     for c in comps:
-    #         stats = tensorstats(getattr(self, c).value)
-    #         if stats is not None:
-    #             line = [f"{k}: {v}" for k, v in stats.items()]
-    #             line = ", ".join(line)
-    #         else:
-    #             line = "None"
-    #         lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-    #     return lines
-
 if __name__ == '__main__':
     from ngcsimlib.context import Context
     with Context("Bar") as bar:

ngclearn/components/input_encoders/latencyCell.py

@@ -144,11 +144,9 @@ class LatencyCell(JaxComponent):
     """
 
     def __init__(
-        self, name: str, n_units: int, tau: float = 1., threshold: float = 0.01,
-        first_spike_time: float = 0., linearize: bool = False,
-        normalize: bool = False, clip_spikes: bool = False,
-        num_steps: float = 1., batch_size: int = 1,
-        key: Union[jax.Array, None] = None
+            self, name: str, n_units: int, tau: float = 1., threshold: float = 0.01, first_spike_time: float = 0.,
+            linearize: bool = False, normalize: bool = False, clip_spikes: bool = False, num_steps: float = 1.,
+            batch_size: int = 1, key: Union[jax.Array, None] = None, **kwargs
     ):
         super().__init__(name=name, key=key)
 

ngclearn/components/input_encoders/phasorCell.py

@@ -137,16 +137,6 @@ def reset(self):
         self.angles.set(restVals)
         self.key.set(key)
 
-
-    def save(self, directory, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        jnp.savez(file_name, key=self.key.value)
-
-    def load(self, directory, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        data = jnp.load(file_name)
-        self.key.set(data['key'])
-
     @classmethod
     def help(cls):  ## component help function
         properties = {

ngclearn/components/input_encoders/poissonCell.py

@@ -3,6 +3,7 @@
 import jax
 from typing import Union
 
+from ngcsimlib import deprecate_args
 from ngcsimlib.parser import compilable
 from ngcsimlib.compartment import Compartment
 
@@ -29,8 +30,11 @@ class PoissonCell(JaxComponent):
         batch_size: batch size dimension of this cell (Default: 1)
     """
 
-    def __init__(self, name: str, n_units: int, target_freq: float = 63.75, batch_size: int = 1,
-                 key: Union[jax.Array, None] = None):
+    @deprecate_args(max_freq="target_freq")
+    def __init__(
+            self, name: str, n_units: int, target_freq: float = 63.75, batch_size: int = 1,
+            key: Union[jax.Array, None] = None, **kwargs
+    ):
         super().__init__(name=name, key=key)
 
         ## Constrained Bernoulli meta-parameters

tests/components/input_encoders/test_poissonCell.py

@@ -1,7 +1,7 @@
 from jax import numpy as jnp, random, jit
 import numpy as np
 np.random.seed(42)
-from ngclearn.components import PoissonCell
+from ngclearn.components.input_encoders.poissonCell import PoissonCell
 from numpy.testing import assert_array_equal
 
 from ngclearn import MethodProcess, Context