refactored bcm-syn and test passed

Author: Alexander Ororbia

ngclearn/components/synapses/hebbian/BCMSynapse.py

@@ -1,10 +1,8 @@
 from jax import random, numpy as jnp, jit
-from ngcsimlib.compilers.process import transition
-from ngcsimlib.component import Component
 from ngcsimlib.compartment import Compartment
+from ngcsimlib.parser import compilable
 
-from ngclearn.components.synapses import DenseSynapse
-from ngclearn.utils import tensorstats
+from ngclearn.components.synapses.denseSynapse import DenseSynapse
 
 class BCMSynapse(DenseSynapse): # BCM-adjusted synaptic cable
     """
@@ -71,8 +69,7 @@ def __init__(
             self, name, shape, tau_w, tau_theta, theta0=-1., w_bound=0., w_decay=0., weight_init=None, resist_scale=1.,
             p_conn=1., batch_size=1, **kwargs
     ):
-        super().__init__(name, shape, weight_init, None, resist_scale, p_conn,
-                         batch_size=batch_size, **kwargs)
+        super().__init__(name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs)
 
         ## Synapse and BCM hyper-parameters
         self.shape = shape ## shape of synaptic efficacy matrix
@@ -90,48 +87,51 @@ def __init__(
         self.post = Compartment(postVals) ## post-synaptic statistic
         self.post_term = Compartment(postVals)
         self.theta = Compartment(postVals + self.theta0) ## synaptic modification thresholds
-        self.dWeights = Compartment(self.weights.value * 0)
+        self.dWeights = Compartment(self.weights.get() * 0)
 
-    @transition(output_compartments=["weights", "theta", "dWeights", "post_term"])
-    @staticmethod
-    def evolve(t, dt, tau_w, tau_theta, w_bound, w_decay, pre, post, theta, weights):
+    @compilable
+    def evolve(self, t, dt): #t, dt, tau_w, tau_theta, w_bound, w_decay, pre, post, theta, weights):
         eps = 1e-7
-        post_term = post * (post - theta)  # post - theta
-        post_term = post_term * (1. / (theta + eps))
-        dWeights = jnp.matmul(pre.T, post_term)
-        if w_bound > 0.:
-            dWeights = dWeights * (w_bound - jnp.abs(weights))
+        post_term = self.post.get() * (self.post.get() - self.theta.get())  # post - theta
+        post_term = post_term * (1. / (self.theta.get() + eps))
+        dWeights = jnp.matmul(self.pre.get().T, post_term)
+        if self.w_bound > 0.:
+            dWeights = dWeights * (self.w_bound - jnp.abs(self.weights.get()))
         ## update synaptic efficacies according to a leaky ODE
-        dWeights = -weights * w_decay + dWeights
-        _W = weights + dWeights * dt / tau_w
+        dWeights = -self.weights.get() * self.w_decay + dWeights
+        _W = self.weights.get() + dWeights * dt / self.tau_w
         ## update synaptic modification threshold as a leaky ODE
-        dtheta = jnp.mean(jnp.square(post), axis=0, keepdims=True)  ## batch avg
-        theta = theta + (-theta + dtheta) * dt / tau_theta
-        return weights, theta, dWeights, post_term
-
-    @transition(output_compartments=["inputs", "outputs", "pre", "post", "dWeights", "post_term"])
-    @staticmethod
-    def reset(batch_size, shape):
-        preVals = jnp.zeros((batch_size, shape[0]))
-        postVals = jnp.zeros((batch_size, shape[1]))
-        inputs = preVals
-        outputs = postVals
-        pre = preVals
-        post = postVals
-        dWeights = jnp.zeros(shape)
-        post_term = postVals
-        return inputs, outputs, pre, post, dWeights, post_term
-
-    def save(self, directory, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        jnp.savez(file_name,
-                  weights=self.weights.value, theta=self.theta.value)
-
-    def load(self, directory, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        data = jnp.load(file_name)
-        self.weights.set(data['weights'])
-        self.theta.set(data['theta'])
+        dtheta = jnp.mean(jnp.square(self.post.get()), axis=0, keepdims=True)  ## batch avg
+        theta = self.theta.get() + (-self.theta.get() + dtheta) * dt / self.tau_theta
+
+        #self.weights.set(weights)
+        self.theta.set(theta)
+        self.dWeights.set(dWeights)
+        self.post_term.set(post_term)
+
+    @compilable
+    def reset(self):
+        preVals = jnp.zeros((self.batch_size.get(), self.shape.get()[0]))
+        postVals = jnp.zeros((self.batch_size.get(), self.shape.get()[1]))
+
+        if not self.inputs.targeted:
+            self.inputs.set(preVals)
+        self.outputs.set(postVals)
+        self.pre.set(preVals)
+        self.post.set(postVals)
+        self.dWeights.set(jnp.zeros(self.shape.get()))
+        self.post_term.set(postVals)
+
+    # def save(self, directory, **kwargs):
+    #     file_name = directory + "/" + self.name + ".npz"
+    #     jnp.savez(file_name,
+    #               weights=self.weights.value, theta=self.theta.value)
+    #
+    # def load(self, directory, **kwargs):
+    #     file_name = directory + "/" + self.name + ".npz"
+    #     data = jnp.load(file_name)
+    #     self.weights.set(data['weights'])
+    #     self.theta.set(data['theta'])
 
     @classmethod
     def help(cls): ## component help function
@@ -175,21 +175,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).value)
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
-
-
 if __name__ == '__main__':
     from ngcsimlib.context import Context
     with Context("Bar") as bar:

ngclearn/components/synapses/hebbian/__init__.py

@@ -1,6 +1,6 @@
-from .hebbianSynapse import HebbianSynapse
+#from .hebbianSynapse import HebbianSynapse
 from .traceSTDPSynapse import TraceSTDPSynapse
-from .expSTDPSynapse import ExpSTDPSynapse
-from .eventSTDPSynapse import EventSTDPSynapse
+#from .expSTDPSynapse import ExpSTDPSynapse
+#from .eventSTDPSynapse import EventSTDPSynapse
 from .BCMSynapse import BCMSynapse
 

tests/components/synapses/hebbian/test_BCMSynapse.py

@@ -2,14 +2,11 @@
 from ngcsimlib.context import Context
 import numpy as np
 np.random.seed(42)
-from ngclearn.components import BCMSynapse
-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
+from ngclearn import Context, MethodProcess
+import ngclearn.utils.weight_distribution as dist
+from ngclearn.components.synapses.hebbian.BCMSynapse import BCMSynapse
+from numpy.testing import assert_array_equal
 
 def test_BCMSynapse1():
     name = "bcm_stdp_ctx"
@@ -23,42 +20,26 @@ def test_BCMSynapse1():
             name="a", shape=(1,1), tau_w=40., tau_theta=20., key=subkeys[0]
         )
 
-        #"""
-        evolve_process = (Process("evolve_proc")
+        evolve_process = (MethodProcess("evolve_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("advance_proc")
+        advance_process = (MethodProcess("advance_proc")
                            >> 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("reset_proc")
+        reset_process = (MethodProcess("reset_proc")
                          >> 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")
-        """
 
     pre_value = jnp.ones((1, 1)) * 0.425
     post_value = jnp.ones((1, 1)) * 1.55
 
     truth = jnp.array([[-1.6798127]])
-    ctx.reset()
+    reset_process.run()  # ctx.reset()
     a.pre.set(pre_value)
     a.post.set(post_value)
-    ctx.run(t=1., dt=dt)
-    ctx.adapt(t=1., dt=dt)
-    #print(a.dWeights.value)
-    assert_array_equal(a.dWeights.value, truth)
-
+    advance_process.run(t=1., dt=dt)  # ctx.run(t=1., dt=dt)
+    evolve_process.run(t=1., dt=dt)  # ctx.adapt(t=1., dt=dt)
+    # print(a.dWeights.get())
+    # print(truth)
+    assert_array_equal(a.dWeights.get(), truth)
 
-#test_BCMSynapse1()
+test_BCMSynapse1()