refactored exp-stdp syn w/ unit-test

Alexander Ororbia · Alexander Ororbia · commit 6a4889a4d201 · 2025-04-01T19:12:22.000-04:00
diff --git a/ngclearn/components/synapses/hebbian/expSTDPSynapse.py b/ngclearn/components/synapses/hebbian/expSTDPSynapse.py
@@ -1,25 +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.7,
-                 exp_beta=1., Aplus=1., Aminus=0.): ## internal dynamics method
-    ## equations 4 from Diehl and Cook - full exponential weight-dependent STDP
-    ## calculate post-synaptic term
-    post_term1 = jnp.exp(-exp_beta * W) * jnp.matmul(x_pre.T, post)
-    x_tar_vec = x_pre * 0 + x_tar  # need to broadcast scalar x_tar to mat/vec form
-    post_term2 = jnp.exp(-exp_beta * (w_bound - W)) * jnp.matmul(x_tar_vec.T,
-                                                                 post)
-    dWpost = (post_term1 - post_term2) * Aplus
-    ## calculate pre-synaptic term
-    dWpre = 0.
-    if Aminus > 0.:
-        dWpre = -jnp.exp(-exp_beta * W) * jnp.matmul(pre.T, x_post) * Aminus
-    ## calc final weighted adjustment
-    dW = (dWpost + dWpre)
-    return dW
-
 class ExpSTDPSynapse(DenseSynapse):
     """
     A synaptic cable that adjusts its efficacies via trace-based form of
@@ -78,9 +64,10 @@ class ExpSTDPSynapse(DenseSynapse):
     """
 
     # Define Functions
-    def __init__(self, name, shape, A_plus, A_minus, exp_beta, eta=1.,
-                 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, exp_beta, eta=1., 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)
 
@@ -105,16 +92,36 @@ def __init__(self, name, shape, A_plus, A_minus, exp_beta, eta=1.,
         self.eta = Compartment(jnp.ones((1, 1)) * eta) ## global learning rate governing plasticity
 
     @staticmethod
-    def _compute_update(dt, w_bound, preTrace_target, exp_beta, Aplus, Aminus,
-                        preSpike, postSpike, preTrace, postTrace, weights):
-        dW = _calc_update(dt, preSpike, preTrace, postSpike, postTrace, weights,
-                          w_bound=w_bound, x_tar=preTrace_target, exp_beta=exp_beta,
-                          Aplus=Aplus, Aminus=Aminus)
+    def _compute_update(
+            dt, w_bound, preTrace_target, exp_beta, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
+    ):
+        pre = preSpike
+        x_pre = preTrace
+        post = postSpike
+        x_post = postTrace
+        W = weights
+        x_tar = preTrace_target
+        ## equations 4 from Diehl and Cook - full exponential weight-dependent STDP
+        ## calculate post-synaptic term
+        post_term1 = jnp.exp(-exp_beta * W) * jnp.matmul(x_pre.T, post)
+        x_tar_vec = x_pre * 0 + x_tar  # need to broadcast scalar x_tar to mat/vec form
+        post_term2 = jnp.exp(-exp_beta * (w_bound - W)) * jnp.matmul(x_tar_vec.T,
+                                                                     post)
+        dWpost = (post_term1 - post_term2) * Aplus
+        ## calculate pre-synaptic term
+        dWpre = 0.
+        if Aminus > 0.:
+            dWpre = -jnp.exp(-exp_beta * W) * 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, exp_beta, Aplus, Aminus,
-                preSpike, postSpike, preTrace, postTrace, weights, eta):
+    def evolve(
+            dt, w_bound, preTrace_target, exp_beta, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace,
+            weights, eta
+    ):
         dW = ExpSTDPSynapse._compute_update(
             dt, w_bound, preTrace_target, exp_beta, Aplus, Aminus,
             preSpike, postSpike, preTrace, postTrace, weights
@@ -126,13 +133,9 @@ def _evolve(dt, w_bound, preTrace_target, exp_beta, Aplus, Aminus,
         _W = jnp.clip(_W, eps, w_bound - eps)
         return weights, dW
 
-    @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
@@ -144,16 +147,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 = {
diff --git a/tests/components/synapses/hebbian/test_expSTDPSynapse.py b/tests/components/synapses/hebbian/test_expSTDPSynapse.py
@@ -0,0 +1,79 @@
+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 ExpSTDPSynapse
+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_expSTDPSynapse1():
+    name = "exp_stdp_ctx"
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    dt = 1.  # ms
+    # ---- build a simple Poisson cell system ----
+    with Context(name) as ctx:
+        a = ExpSTDPSynapse(
+            name="a", shape=(1,1), A_plus=1., A_minus=1., exp_beta=1.25, 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")
+        """
+
+    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([[0.57342285]])
+    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.29817986]])
+    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_expSTDPSynapse1()
diff --git a/tests/components/synapses/hebbian/test_traceSTDPSynapse.py b/tests/components/synapses/hebbian/test_traceSTDPSynapse.py
@@ -17,7 +17,6 @@ def test_traceSTDPSynapse1():
     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(
@@ -49,11 +48,6 @@ def test_traceSTDPSynapse1():
         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))
