refactored event-stdp w/ unit-test

Author: Alexander Ororbia

ngclearn/components/synapses/hebbian/eventSTDPSynapse.py

@@ -1,8 +1,10 @@
 from jax import numpy as jnp, jit
-from ngclearn import resolver, Component, Compartment
-from ngclearn.utils import tensorstats
-## import parent synapse class/component
+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
 
 class EventSTDPSynapse(DenseSynapse): # event-driven, post-synaptic STDP
     """
@@ -80,8 +82,9 @@ def __init__(self, name, shape, eta, lmbda=0.01, A_plus=1., A_minus=1.,
         self.eta = Compartment(jnp.ones((1, 1)) * eta)  ## global learning rate governing plasticity
 
     @staticmethod
-    def _compute_update(t, lmbda, presyn_win_len, Aminus, Aplus, w_bound, pre_tols,
-                        postSpike, weights):
+    def _compute_update(
+            t, lmbda, presyn_win_len, Aminus, Aplus, w_bound, pre_tols, postSpike, weights
+    ): ## synaptic adjustment calculation co-routine
         ## check if a spike occurred in window of (t - presyn_win_len, t]
         m = (pre_tols > 0.) * 1.  ## ignore default value of tols = 0 ms
         if presyn_win_len > 0.:
@@ -99,40 +102,30 @@ def _compute_update(t, lmbda, presyn_win_len, Aminus, Aplus, w_bound, pre_tols,
         dW = (dW * postSpike) ## gate to make sure only post-spikes trigger updates
         return dW
 
+    @transition(output_compartments=["weights", "dWeights"])
     @staticmethod
-    def _evolve(t, lmbda, presyn_win_len, Aminus, Aplus, w_bound, pre_tols,
-                postSpike, weights, eta):
+    def evolve(
+            t, lmbda, presyn_win_len, Aminus, Aplus, w_bound, pre_tols, postSpike, weights, eta
+    ):
         dWeights = EventSTDPSynapse._compute_update(
             t, lmbda, presyn_win_len, Aminus, Aplus, w_bound, pre_tols, postSpike, weights
         )
         weights = weights + dWeights * eta  # * (1. - w) * eta
-        weights = jnp.clip(weights, 0.01, w_bound)  # not in source paper
+        weights = jnp.clip(weights, 0.01, w_bound)  ## Note: this step not in source paper
         return weights, dWeights
 
-    @resolver(_evolve)
-    def evolve(self, weights, dWeights):
-        self.weights.set(weights)
-        self.dWeights.set(dWeights)
-
+    @transition(output_compartments=["inputs", "outputs", "pre_tols", "postSpike", "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
         outputs = postVals
-        pre_tols = preVals ## pre-synaptic time-of-last-spike record
+        pre_tols = preVals ## pre-synaptic time-of-last-spike(s) record
         postSpike = postVals
         dWeights = jnp.zeros(shape)
         return inputs, outputs, pre_tols, postSpike, dWeights
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, pre_tols, postSpike, dWeights):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.pre_tols.set(pre_tols)
-        self.postSpike.set(postSpike)
-        self.dWeights.set(dWeights)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

tests/components/synapses/hebbian/test_eventSTDPSynapse.py

@@ -0,0 +1,81 @@
+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 EventSTDPSynapse
+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_eventSTDPSynapse1():
+    name = "event_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 = EventSTDPSynapse(
+            name="a", shape=(1,1), eta=0., presyn_win_len=2., 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")
+        """
+
+    t = 12. ## fake out current time
+    ## Case 1: outside of pre-syn time window
+    input_tols = jnp.ones((1, 1,)) * 9.
+    out_spike = jnp.ones((1, 1))
+
+    ## check pre-synaptic STDP only
+    truth = jnp.array([[-0.6296545]])
+    ctx.reset()
+    a.pre_tols.set(input_tols)
+    a.postSpike.set(out_spike)
+    ctx.run(t=t, dt=dt)
+    ctx.adapt(t=t, dt=dt)
+    #print(a.dWeights.value)
+    assert_array_equal(a.dWeights.value, truth)
+
+    ## Case 2: within pre-syn time window
+    input_tols = jnp.ones((1, 1,)) * 11.
+    out_spike = jnp.ones((1, 1))
+
+    ## check pre-synaptic STDP only
+    truth = jnp.array([[0.37034547]])
+    ctx.reset()
+    a.pre_tols.set(input_tols)
+    a.postSpike.set(out_spike)
+    ctx.run(t=t, dt=dt)
+    ctx.adapt(t=t, dt=dt)
+    #print(a.dWeights.value)
+    assert_array_equal(a.dWeights.value, truth)
+
+#test_eventSTDPSynapse1()