refactored dense and trace-stdp syn w/ unit-test

Author: Alexander Ororbia

ngclearn/components/synapses/denseSynapse.py

@@ -1,10 +1,13 @@
 from jax import random, numpy as jnp, jit
-from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
 from ngclearn.utils import tensorstats
 from ngclearn.utils.weight_distribution import initialize_params
 from ngcsimlib.logger import info
 
+from ngcsimlib.compilers.process import transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+
 class DenseSynapse(JaxComponent): ## base dense synaptic cable
     """
     A dense synaptic cable; no form of synaptic evolution/adaptation
@@ -13,7 +16,7 @@ class DenseSynapse(JaxComponent): ## base dense synaptic cable
     | --- Synapse Compartments: ---
     | inputs - input (takes in external signals)
     | outputs - output signals
-    | weights - current value matrix of synaptic efficacies
+    | weights - current value matrix of synaptic efficacies (strength values)
     | biases - current value vector of synaptic bias values
 
     Args:
@@ -75,28 +78,21 @@ def __init__(self, name, shape, weight_init=None, bias_init=None,
                                                     (1, shape[1]))
                                   if bias_init else 0.0)
 
+    @transition(output_compartments=["outputs"])
     @staticmethod
-    def _advance_state(Rscale, inputs, weights, biases):
+    def advance_state(Rscale, inputs, weights, biases):
         outputs = (jnp.matmul(inputs, weights) * Rscale) + biases
         return outputs
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs):
-        self.outputs.set(outputs)
-
+    @transition(output_compartments=["inputs", "outputs"])
     @staticmethod
-    def _reset(batch_size, shape):
+    def reset(batch_size, shape):
         preVals = jnp.zeros((batch_size, shape[0]))
         postVals = jnp.zeros((batch_size, shape[1]))
         inputs = preVals
         outputs = postVals
         return inputs, outputs
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         if self.bias_init != None:

ngclearn/components/synapses/hebbian/traceSTDPSynapse.py

@@ -1,28 +1,11 @@
 from jax import random, numpy as jnp, jit
-from ngclearn import resolver, Component, Compartment
+from ngcsimlib.compilers.process import transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+
 from ngclearn.components.synapses import DenseSynapse
 from ngclearn.utils import tensorstats
 
-def _calc_update(dt, pre, x_pre, post, x_post, W, w_bound=1., x_tar=0.0, mu=0.,
-                 Aplus=1., Aminus=0.):
-    if mu > 0.:
-        ## equations 3, 5, & 6 from Diehl and Cook - full power-law STDP
-        post_shift = jnp.power(w_bound - W, mu)
-        pre_shift = jnp.power(W, mu)
-        dWpost = (post_shift * jnp.matmul((x_pre - x_tar).T, post)) * Aplus
-        dWpre = 0.
-        if Aminus > 0.:
-            dWpre = -(pre_shift * jnp.matmul(pre.T, x_post)) * Aminus
-    else:
-        ## calculate post-synaptic term
-        dWpost = jnp.matmul((x_pre - x_tar).T, post * Aplus)
-        dWpre = 0.
-        if Aminus > 0.:
-            ## calculate pre-synaptic term
-            dWpre = -jnp.matmul(pre.T, x_post * Aminus)
-    ## calc final weighted adjustment
-    dW = (dWpost + dWpre)
-    return dW
 
 class TraceSTDPSynapse(DenseSynapse): # power-law / trace-based STDP
     """
@@ -83,9 +66,10 @@ class TraceSTDPSynapse(DenseSynapse): # power-law / trace-based STDP
     """
 
     # Define Functions
-    def __init__(self, name, shape, A_plus, A_minus, eta=1., mu=0.,
-                 pretrace_target=0., weight_init=None, resist_scale=1.,
-                 p_conn=1., w_bound=1., batch_size=1, **kwargs):
+    def __init__(
+            self, name, shape, A_plus, A_minus, eta=1., mu=0., pretrace_target=0., weight_init=None, resist_scale=1.,
+            p_conn=1., w_bound=1., batch_size=1, **kwargs
+    ):
         super().__init__(name, shape, weight_init, None, resist_scale,
                          p_conn, batch_size=batch_size, **kwargs)
 
@@ -109,19 +93,41 @@ def __init__(self, name, shape, A_plus, A_minus, eta=1., mu=0.,
         self.eta = Compartment(jnp.ones((1, 1)) * eta) ## global learning rate
 
     @staticmethod
-    def _compute_update(dt, w_bound, preTrace_target, mu, Aplus, Aminus,
-                preSpike, postSpike, preTrace, postTrace, weights):
-        dW = _calc_update(dt, preSpike, preTrace, postSpike, postTrace, weights,
-                          w_bound=w_bound, x_tar=preTrace_target, mu=mu,
-                          Aplus=Aplus, Aminus=Aminus)
+    def _compute_update(
+            dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
+    ):
+        pre = preSpike
+        x_pre = preTrace
+        post = postSpike
+        x_post = postTrace
+        W = weights
+        x_tar = preTrace_target
+        if mu > 0.:
+            ## equations 3, 5, & 6 from Diehl and Cook - full power-law STDP
+            post_shift = jnp.power(w_bound - W, mu)
+            pre_shift = jnp.power(W, mu)
+            dWpost = (post_shift * jnp.matmul((x_pre - x_tar).T, post)) * Aplus
+            dWpre = 0.
+            if Aminus > 0.:
+                dWpre = -(pre_shift * jnp.matmul(pre.T, x_post)) * Aminus
+        else:
+            ## calculate post-synaptic term
+            dWpost = jnp.matmul((x_pre - x_tar).T, post * Aplus)
+            dWpre = 0.
+            if Aminus > 0.:
+                ## calculate pre-synaptic term
+                dWpre = -jnp.matmul(pre.T, x_post * Aminus)
+        ## calc final weighted adjustment
+        dW = (dWpost + dWpre)
         return dW
 
+    @transition(output_compartments=["weights", "dWeights"])
     @staticmethod
-    def _evolve(dt, w_bound, preTrace_target, mu, Aplus, Aminus,
-                preSpike, postSpike, preTrace, postTrace, weights, eta):
+    def evolve(
+            dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights, eta
+    ):
         dWeights = TraceSTDPSynapse._compute_update(
-            dt, w_bound, preTrace_target, mu, Aplus, Aminus,
-            preSpike, postSpike, preTrace, postTrace, weights
+            dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
         )
         ## do a gradient ascent update/shift
         weights = weights + dWeights * eta
@@ -130,13 +136,9 @@ def _evolve(dt, w_bound, preTrace_target, mu, Aplus, Aminus,
         weights = jnp.clip(weights, eps, w_bound - eps)  # jnp.abs(w_bound))
         return weights, dWeights
 
-    @resolver(_evolve)
-    def evolve(self, weights, dWeights):
-        self.weights.set(weights)
-        self.dWeights.set(dWeights)
-
+    @transition(output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace", "dWeights"])
     @staticmethod
-    def _reset(batch_size, shape):
+    def reset(batch_size, shape):
         preVals = jnp.zeros((batch_size, shape[0]))
         postVals = jnp.zeros((batch_size, shape[1]))
         inputs = preVals
@@ -148,16 +150,6 @@ def _reset(batch_size, shape):
         dWeights = jnp.zeros(shape)
         return inputs, outputs, preSpike, postSpike, preTrace, postTrace, dWeights
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, preSpike, postSpike, preTrace, postTrace, dWeights):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.preSpike.set(preSpike)
-        self.postSpike.set(postSpike)
-        self.preTrace.set(preTrace)
-        self.postTrace.set(postTrace)
-        self.dWeights.set(dWeights)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

tests/components/synapses/hebbian/test_traceSTDPSynapse.py

@@ -0,0 +1,85 @@
+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 TraceSTDPSynapse
+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
+
+def test_traceSTDPSynapse1():
+    name = "trace_stdp_ctx"
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    dt = 1.  # ms
+    trace_increment = 0.1
+    # ---- build a simple Poisson cell system ----
+    with Context(name) as ctx:
+        a = TraceSTDPSynapse(
+            name="a", shape=(1,1), A_plus=1., A_minus=1., key=subkeys[0]
+        )
+
+        #"""
+        evolve_process = (Process()
+                           >> a.evolve)
+        #ctx.wrap_and_add_command(evolve_process.pure, name="run")
+        ctx.wrap_and_add_command(jit(evolve_process.pure), name="adapt")
+
+        advance_process = (Process()
+                           >> a.advance_state)
+        # ctx.wrap_and_add_command(advance_process.pure, name="run")
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
+
+        reset_process = (Process()
+                         >> a.reset)
+        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
+        #"""
+
+        """
+        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="run")
+        evolve_cmd, evolve_args = ctx.compile_by_key(a, compile_key="evolve")
+        ctx.add_command(wrap_command(jit(ctx.evolve)), name="adapt")
+        """
+
+        ## set up non-compiled utility commands
+        @Context.dynamicCommand
+        def clamp(x):
+            a.j.set(x)
+
+    in_spike = jnp.ones((1, 1))
+    in_trace = jnp.ones((1, 1,)) * 1.25
+    out_spike = jnp.ones((1, 1))
+    out_trace = jnp.ones((1, 1,)) * 0.65
+
+    ## check pre-synaptic STDP only
+    truth = jnp.array([[1.25]])
+    ctx.reset()
+    a.preSpike.set(in_spike * 0)
+    a.preTrace.set(in_trace)
+    a.postSpike.set(out_spike)
+    a.postTrace.set(out_trace)
+    ctx.run(t=1., dt=dt)
+    ctx.adapt(t=1., dt=dt)
+    #print(a.dWeights.value)
+    assert_array_equal(a.dWeights.value, truth)
+
+    truth = jnp.array([[-0.65]])
+    ctx.reset()
+    a.preSpike.set(in_spike)
+    a.preTrace.set(in_trace)
+    a.postSpike.set(out_spike * 0)
+    a.postTrace.set(out_trace)
+    ctx.run(t=1., dt=dt)
+    ctx.adapt(t=1., dt=dt)
+    #print(a.dWeights.value)
+    assert_array_equal(a.dWeights.value, truth)
+
+#test_traceSTDPSynapse1()