refactored mstdp-et syn w/ unit-test

Author: Alexander Ororbia

ngclearn/components/__init__.py

@@ -1,14 +1,12 @@
 from .jaxComponent import JaxComponent
 
-
 ## point to rate-coded cell component types
 from .neurons.graded.rateCell import RateCell
 from .neurons.graded.gaussianErrorCell import GaussianErrorCell
 from .neurons.graded.laplacianErrorCell import LaplacianErrorCell
 from .neurons.graded.bernoulliErrorCell import BernoulliErrorCell
 from .neurons.graded.rewardErrorCell import RewardErrorCell
 
-
 ## point to standard spiking cell component types
 from .neurons.spiking.sLIFCell import SLIFCell
 from .neurons.spiking.IFCell import IFCell
@@ -21,11 +19,9 @@
 from .neurons.spiking.hodgkinHuxleyCell import HodgkinHuxleyCell
 from .neurons.spiking.RAFCell import RAFCell
 
-## point to transformer/operater component types
+## point to transformer/operator component types
 from .other.varTrace import VarTrace
 from .other.expKernel import ExpKernel
-from ngclearn.components.synapses.modulated.eligibilityTrace import EligibilityTrace
-
 
 ## point to input encoder component types
 from .input_encoders.bernoulliCell import BernoulliCell
@@ -43,7 +39,6 @@
 from .synapses.hebbian.BCMSynapse import BCMSynapse
 from .synapses.STPDenseSynapse import STPDenseSynapse
 
-
 ## point to convolutional component types
 from .synapses.convolution.convSynapse import ConvSynapse
 from .synapses.convolution.staticConvSynapse import StaticConvSynapse
@@ -56,7 +51,6 @@
 ## point to modulated component types
 from .synapses.modulated.MSTDPETSynapse import MSTDPETSynapse
 
-
 ## point to monitors
 from .monitor import Monitor
 
@@ -65,8 +59,3 @@
 from .synapses.patched.staticPatchedSynapse import StaticPatchedSynapse
 from .synapses.patched.hebbianPatchedSynapse import HebbianPatchedSynapse
 
-
-
-
-
-

ngclearn/components/synapses/modulated/MSTDPETSynapse.py

@@ -1,13 +1,17 @@
 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.utils.weight_distribution import initialize_params
+from ngcsimlib.logger import info
 from ngclearn.components.synapses.hebbian import TraceSTDPSynapse
 from ngclearn.utils import tensorstats
 
 class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligility traces
     """
-    A synaptic cable that adjusts its efficacies via trace-based form of
-    three-factor learning, i.e., modulated spike-timing-dependent plasticity
-    (M-STDP) or modulated STDP with eligibility traces (M-STDP-ET).
+    A synaptic cable that adjusts its efficacies via trace-based form of three-factor learning, i.e., modulated
+    spike-timing-dependent plasticity (M-STDP) or modulated STDP with eligibility traces (M-STDP-ET).
 
     | --- Synapse Compartments: ---
     | inputs - input (takes in external signals)
@@ -20,11 +24,14 @@ class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligilit
     | postSpike - post-synaptic spike to drive 2nd term of STDP update (takes in external signals)
     | preTrace - pre-synaptic trace value to drive 1st term of STDP update (takes in external signals)
     | postTrace - post-synaptic trace value to drive 2nd term of STDP update (takes in external signals)
-    | dWeights - current delta matrix containing changes to be applied to synaptic efficacies
+    | dWeights - current delta matrix containing (MS-STDP/MS-STDP-ET) changes to be applied to synaptic efficacies
     | eligibility - current state of eligibility trace
-    | eta - global learning rate (multiplier beyond A_plus and A_minus)
+    | eta - global learning rate (applied to change in weights for final MS-STDP/MS-STDP-ET adjustment)
 
     | References:
+    | Florian, Răzvan V. "Reinforcement learning through modulation of spike-timing-dependent synaptic plasticity."
+    | Neural computation 19.6 (2007): 1468-1502.
+    |
     | Morrison, Abigail, Ad Aertsen, and Markus Diesmann. "Spike-timing-dependent
     | plasticity in balanced random networks." Neural computation 19.6 (2007): 1437-1467.
     |
@@ -66,29 +73,30 @@ class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligilit
     """
 
     # Define Functions
-    def __init__(self, name, shape, A_plus, A_minus, eta=1., mu=0.,
-                 pretrace_target=0., tau_elg=0., elg_decay=1.,
-                 weight_init=None, resist_scale=1., p_conn=1., w_bound=1.,
-                 batch_size=1, **kwargs):
-        super().__init__(name, shape, A_plus, A_minus, eta=eta, mu=mu,
-                         pretrace_target=pretrace_target, weight_init=weight_init,
-                         resist_scale=resist_scale, p_conn=p_conn, w_bound=w_bound,
-                         batch_size=batch_size, **kwargs)
+    def __init__(
+            self, name, shape, A_plus, A_minus, eta=1., mu=0., pretrace_target=0., tau_elg=0., elg_decay=1.,
+            weight_init=None, resist_scale=1., p_conn=1., w_bound=1., batch_size=1, **kwargs
+    ):
+        super().__init__(
+            name, shape, A_plus, A_minus, eta=eta, mu=mu, pretrace_target=pretrace_target, weight_init=weight_init,
+            resist_scale=resist_scale, p_conn=p_conn, w_bound=w_bound, batch_size=batch_size, **kwargs
+        )
         ## MSTDP/MSTDP-ET meta-parameters
         self.tau_elg = tau_elg
         self.elg_decay = elg_decay
         ## MSTDP/MSTDP-ET compartments
         self.modulator = Compartment(jnp.zeros((self.batch_size, 1)))
         self.eligibility = Compartment(jnp.zeros(shape))
 
+    @transition(output_compartments=["weights", "dWeights", "eligibility"])
     @staticmethod
-    def _evolve(dt, w_bound, preTrace_target, mu, Aplus, Aminus, tau_elg,
-                elg_decay, preSpike, postSpike, preTrace, postTrace, weights,
-                eta, modulator, eligibility):
+    def evolve(
+            dt, w_bound, preTrace_target, mu, Aplus, Aminus, tau_elg, elg_decay, preSpike, postSpike, preTrace,
+            postTrace, weights, eta, modulator, eligibility
+    ):
         ## compute local synaptic update (via STDP)
         dW_dt = 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
         ) ## produce dW/dt (ODE for synaptic change dynamics)
         if tau_elg > 0.: ## perform dynamics of M-STDP-ET
             ## update eligibility trace given current local update
@@ -107,14 +115,11 @@ def _evolve(dt, w_bound, preTrace_target, mu, Aplus, Aminus, tau_elg,
         weights = jnp.clip(weights, eps, w_bound - eps)  # jnp.abs(w_bound))
         return weights, dWeights, eligibility
 
-    @resolver(_evolve)
-    def evolve(self, weights, dWeights, eligibility):
-        self.weights.set(weights)
-        self.dWeights.set(dWeights)
-        self.eligibility.set(eligibility)
-
+    @transition(
+        output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace", "dWeights", "eligibility"]
+    )
     @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]))
         synVals = jnp.zeros(shape)
@@ -126,20 +131,7 @@ def _reset(batch_size, shape):
         postTrace = postVals
         dWeights = synVals
         eligibility = synVals
-        return (inputs, outputs, preSpike, postSpike, preTrace, postTrace,
-                dWeights, eligibility)
-
-    @resolver(_reset)
-    def reset(self, inputs, outputs, preSpike, postSpike, preTrace, postTrace,
-              dWeights, eligibility):
-        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)
-        self.eligibility.set(eligibility)
+        return inputs, outputs, preSpike, postSpike, preTrace, postTrace, dWeights, eligibility
 
     @classmethod
     def help(cls): ## component help function

ngclearn/components/synapses/modulated/__init__.py

@@ -1 +1 @@
-from .MSTDPETSynapse import MSTDPETSynapse
+from .MSTDPETSynapse import MSTDPETSynapse

tests/components/synapses/modulated/test_MSTDPETSynapse.py

@@ -0,0 +1,83 @@
+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 MSTDPETSynapse
+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
+import ngclearn.utils.weight_distribution as dist
+
+def test_MSTDPETSynapse1():
+    name = "mstdpet_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 = MSTDPETSynapse(
+            name="a", shape=(1,1), A_plus=1., A_minus=1., eta=0., key=subkeys[0]
+        )
+
+        #"""
+        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")
+
+        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")
+        
+        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")
+        """
+
+    a.weights.set(jnp.ones((1, 1)))
+
+    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
+    r_neg = -jnp.ones((1, 1))
+    r_pos = jnp.ones((1, 1))
+
+    ctx.reset()
+    a.preSpike.set(in_spike * 0)
+    a.preTrace.set(in_trace)
+    a.postSpike.set(out_spike)
+    a.postTrace.set(out_trace)
+    a.modulator.set(r_pos)
+    ctx.run(t=1. * dt, dt=dt)
+    ctx.adapt(t=1. * dt, dt=dt)
+    #print(a.dWeights.value)
+    assert_array_equal(a.dWeights.value, 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)
+    a.modulator.set(r_neg)
+    ctx.run(t=1. * dt, dt=dt)
+    ctx.adapt(t=1. * dt, dt=dt)
+    #print(a.dWeights.value)
+    assert_array_equal(a.dWeights.value, jnp.array([[-1.25]]))
+
+#test_MSTDPETSynapse1()