cleaned up stdp syn

Author: Alexander Ororbia

ngclearn/components/synapses/hebbian/STDPSynapse.py

@@ -3,7 +3,7 @@
 from ngclearn.components.synapses import DenseSynapse
 from ngclearn.utils import tensorstats
 
-class STDPSynapse(DenseSynapse): # power-law / trace-based STDP
+class STDPSynapse(DenseSynapse): # classical STDP
     """
     A synaptic cable that adjusts its efficacies via raw
     spike-timing-dependent plasticity (STDP).
@@ -92,17 +92,20 @@ def __init__(self, name, shape, A_plus, A_minus, tau_plus=10., tau_minus=10.,
     def _compute_update(Aplus, Aminus, tau_plus, tau_minus, preSpike, postSpike,
                         pre_tols, post_tols, weights):
         ## calculate time deltas matrix block --> (t_post - t_pre)
-        post_m = (post_tols > 0.) #* 1.
-        pre_m = (pre_tols > 0.).T # * 1.
+        post_m = (post_tols > 0.) ## zero post-tols mask
+        pre_m = (pre_tols > 0.).T ## zero pre-tols mask
         t_delta = ((weights * 0 + 1.) * post_tols) - pre_tols.T ## t_delta.shape = weights.shape
-        t_delta = t_delta * post_m * pre_m ## mask out zero tols
+        t_delta = t_delta * post_m * pre_m  ## mask out zero tols and same-time spikes
+        pos_t_delta_m = (t_delta > 0.) ## positive t-delta mask
+        neg_t_delta_m = (t_delta < 0.) ## negative t-delta mask
+        #t_delta = t_delta * pos_t_delta_m + t_delta * neg_t_delta_m ## mask out same time spikes
         ## calculate post-synaptic term
-        postTerm = jnp.exp(-t_delta/tau_plus) * (t_delta > 0.) #* post_m * pre_m
+        postTerm = jnp.exp(-t_delta/tau_plus) * pos_t_delta_m
         dWpost = postTerm * (postSpike * Aplus)
         dWpre = 0.
         if Aminus > 0.:
             ## calculate pre-synaptic term
-            preTerm = jnp.exp(-t_delta / tau_minus) * (t_delta < 0.) #* post_m * pre_m
+            preTerm = jnp.exp(-t_delta / tau_minus) * neg_t_delta_m
             dWpre = -preTerm * (preSpike.T * Aminus)
         ## calc final weighted adjustment
         dW = (dWpost + dWpre)
@@ -121,7 +124,7 @@ def _evolve(dt, w_bound, tau_w, Aplus, Aminus, tau_plus, tau_minus, preSpike,
         else: ## raw simple ascent-style update
             weights = weights + dWeights * eta
         ## enforce non-negativity
-        eps = 0.001 # 0.01
+        eps = 0.01
         weights = jnp.clip(weights, eps, w_bound - eps)  # jnp.abs(w_bound))
         return weights, dWeights