revised traces neurocog tutorial to v3

Author: Alexander Ororbia

docs/tutorials/neurocog/traces.md

@@ -1,29 +1,17 @@
 # Lecture 1B: Trace Variables and Filtering
 
-Traces represent one very important component tool in ngc-learn as these are
-often, in biophysical model simulations, used to produce real-valued
-representations of often discrete-valued patterns, e.g., spike vectors within
-a spike train, that can facilitate mechanisms such as online biological credit
-assignment. In this lesson, we will observe how one of ngc-learn's core
-trace components -- the `VarTrace` -- operates.
+Traces represent one very important component tool in ngc-learn as these are often, in biophysical model simulations, used to produce real-valued representations of often discrete-valued patterns, e.g., spike vectors within a spike train, that can facilitate mechanisms such as online biological credit assignment. In this lesson, we will observe how one of ngc-learn's core trace components -- the `VarTrace` -- operates.
 
 ## Setting Up a Variable Trace for a Poisson Spike Train
 
-To observe the value of a variable trace, we will pair it to another in-built
-ngc-component; the `PoissonCell`, which will be configured to emit spikes
-approximately at `63.75` Hertz (yielding a fairly sparse spike train). This means
-we will construct a two-component dynamical system, where the input
-compartment `outputs` of the `PoissonCell` will be wired directly into the
-`inputs` compartment of the `VarTrace`. Note that a `VarTrace` has an `inputs`
-compartment -- which is where raw signals typically go into -- and a `trace`
-output compartment -- which is where filtered signal values/by-products are emitted from.
+To observe the value of a variable trace, we will pair it to another in-built ngc-component; the `PoissonCell`, which will be configured to emit spikes approximately at `63.75` Hertz (yielding a fairly sparse spike train). This means we will construct a two-component dynamical system, where the input compartment `outputs` of the `PoissonCell` will be wired directly into the `inputs` compartment of the `VarTrace`. Note that a `VarTrace` has an `inputs` compartment -- which is where raw signals typically go into -- and a `trace` output compartment -- which is where filtered signal values/by-products are emitted from. 
 
 The code below will instantiate the paired Poisson cell and corresponding variable trace:
 
 ```python
 from jax import numpy as jnp, random, jit
-from ngclearn.utils import JaxProcess
-from ngcsimlib.context import Context
+
+from ngclearn import Context, MethodProcess
 ## import model-specific mechanisms
 from ngclearn.components.input_encoders.poissonCell import PoissonCell
 from ngclearn.components.other.varTrace import VarTrace
@@ -37,30 +25,24 @@ with Context("Model") as model:
     trace = VarTrace("tr0", n_units=1, tau_tr=30., a_delta=0.5)
 
     ## wire up cell z0 to trace tr0
-    trace.inputs << cell.outputs
+    cell.outputs >> trace.inputs 
 
-    advance_process = (JaxProcess()
+    advance_process = (MethodProcess("advance")
                        >> cell.advance_state
                        >> trace.advance_state)
-    model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (JaxProcess()
+    reset_process = (MethodProcess("reset")
                      >> cell.reset
                      >> trace.reset)
-    model.wrap_and_add_command(jit(reset_process.pure), name="reset")
-
 
-    @Context.dynamicCommand
-    def clamp(x):
-        cell.inputs.set(x)
+## set up some utility functions for the model context
+def clamp(x):
+    cell.inputs.set(x)
 ```
 
 ## Running the Paired Cell-Trace System
 
-We can then run the above two-component dynamical system by injecting a fixed
-(valid) probability value into the Poisson input encoder and then record the
-resulting spikes and trace values. We will do this for `T = 200` milliseconds (ms)
-with the code below:
+We can then run the above two-component dynamical system by injecting a fixed (valid) probability value into the Poisson input encoder and then record the resulting spikes and trace values. We will do this for `T = 200` milliseconds (ms) with the code below:
 
 ```python
 dt = 1. # ms # integration time constant
@@ -70,22 +52,21 @@ probs = jnp.asarray([[0.35]],dtype=jnp.float32)
 time_span = []
 spikes = []
 traceVals = []
-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)
 
-    print("{}  {}".format(cell.outputs.value, trace.trace.value), end="")
-    spikes.append( cell.outputs.value )
-    traceVals.append( trace.trace.value )
+    print(f"\r{cell.outputs.get()}  {trace.trace.get()}", end="")
+    spikes.append( cell.outputs.get() )
+    traceVals.append( trace.trace.get() )
     time_span.append(ts * dt)
 print()
 spikes = jnp.concatenate(spikes,axis=0)
 traceVals = jnp.concatenate(traceVals,axis=0)
 ```
 
-We can plot the above simulation's trace outputs with the discrete spikes
-super-imposed at their times of occurrence with the code below:
+We can plot the above simulation's trace outputs with the discrete spikes super-imposed at their times of occurrence with the code below:
 
 ```python
 import matplotlib #.pyplot as plt
@@ -100,8 +81,7 @@ stat = jnp.where(spikes > 0.)
 indx = (stat[0] * 1. - 1.).tolist()
 spk = ax.vlines(x=indx, ymin=0.985, ymax=1.05, colors='black', ls='-', lw=5)
 
-ax.set(xlabel='Time (ms)', ylabel='Trace Output',
-      title='Variable Trace of Poisson Spikes')
+ax.set(xlabel='Time (ms)', ylabel='Trace Output', title='Variable Trace of Poisson Spikes')
 #ax.legend([zTr[0],spk[0]],['z','phi(z)'])
 ax.grid()
 fig.savefig("poisson_trace.jpg")
@@ -111,29 +91,16 @@ to get the following output saved to disk:
 
 <img src="../../images/tutorials/neurocog/poisson_trace.jpg" width="400" />
 
-Notice that every time a spike is produced by the Poisson encoding cell, the trace
-increments by `0.5` -- the result of the `a_delta` hyper-parameter we set when
-crafting the model and simulation object -- and then exponentially decays in
-the absence of a spike (with the time constant of `tau_tr = 30` milliseconds).
+Notice that every time a spike is produced by the Poisson encoding cell, the trace increments by `0.5` -- the result of the `a_delta` hyper-parameter we set when crafting the model and simulation object -- and then exponentially decays in the absence of a spike (with the time constant of `tau_tr = 30` milliseconds). 
 
-The variable trace can be further configured to filter signals in different ways
-if desired; specifically by manipulating its `decay_type` and `a_delta` arguments.
-Notably, if a piecewise-gated variable trace is desired (a very common choice
-in some neuronal circuit models), then all one would have to do is set `a_delta = 0`,
-yielding the following line in the model creation code earlier in this tutorial:
+The variable trace can be further configured to filter signals in different ways if desired; specifically by manipulating its `decay_type` and `a_delta` arguments. Notably, if a piecewise-gated variable trace is desired (a very common choice in some neuronal circuit models), then all one would have to do is set `a_delta = 0`, yielding the following line in the model creation code earlier in this tutorial: 
 
 ```python
 trace = VarTrace("tr0", n_units=1, tau_tr=30., a_delta=0., decay_type="exp")
 ```
 
-Running the same code from before but with the above alteration would yield the
-plot below:
+Running the same code from before but with the above alteration would yield the plot below:
 
 <img src="../../images/tutorials/neurocog/poisson_trace_gate.jpg" width="400" />
 
-Notice that, this time, when a spike is emitted from the Poisson cell, the trace
-is "clamped" to the value of one and then exponentially decays. Such a trace
-configuration is useful if one requires the maintained trace to never increase
-beyond a value of one, preventing divergence or run-away values if a spike train
-is particularly dense and yielding friendlier values for biological learning
-rules.
+Notice that, this time, when a spike is emitted from the Poisson cell, the trace is "clamped" to the value of one and then exponentially decays. Such a trace configuration is useful if one requires the maintained trace to never increase beyond a value of one, preventing divergence or run-away values if a spike train is particularly dense and yielding friendlier values for biological learning rules.