refactored bcm syn w/ unit-test

Author: Alexander Ororbia

docs/modeling/synapses.md

@@ -63,6 +63,23 @@ used for fixed value deconvolution/transposed convolution synaptic filters.
     :noindex:
 ```
 
+## Dynamic Synapse Types
+
+### Short-Term Plasticity(Dense) Synapse
+
+This synapse performs a linear transform of its input signals. Note that this 
+synapse is "dynamic" in the sense that it engages in short-term plasticity (STP), meaning that its efficacy values change as a function of its inputs (and simulated consumed resources), but it does not provide any long-term form of plasticity/adjustment.
+
+```{eval-rst}
+.. autoclass:: ngclearn.components.STPDenseSynapse
+  :noindex:
+
+  .. automethod:: advance_state
+    :noindex:
+  .. automethod:: reset
+    :noindex:
+```
+
 ## Multi-Factor Learning Synapse Types
 
 Hebbian rules operate in a local manner -- they generally use information more

ngclearn/components/synapses/hebbian/BCMSynapse.py

@@ -1,5 +1,8 @@
 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
 
@@ -64,8 +67,10 @@ class BCMSynapse(DenseSynapse): # BCM-adjusted synaptic cable
     """
 
     # Define Functions
-    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):
+    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)
 
@@ -87,8 +92,9 @@ def __init__(self, name, shape, tau_w, tau_theta, theta0=-1., w_bound=0., w_deca
         self.theta = Compartment(postVals + self.theta0) ## synaptic modification thresholds
         self.dWeights = Compartment(self.weights.value * 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):
+    def evolve(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))
@@ -101,18 +107,11 @@ def _evolve(t, dt, tau_w, tau_theta, w_bound, w_decay, pre, post, theta, weights
         ## 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
 
-    @resolver(_evolve)
-    def evolve(self, weights, theta, dWeights, post_term):
-        self.weights.set(weights)
-        self.theta.set(theta)
-        self.dWeights.set(dWeights)
-        self.post_term.set(post_term)
-
+    @transition(output_compartments=["inputs", "outputs", "pre", "post", "dWeights", "post_term"])
     @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
@@ -123,15 +122,6 @@ def _reset(batch_size, shape):
         post_term = postVals
         return inputs, outputs, pre, post, dWeights, post_term
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, pre, post, dWeights, post_term):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.pre.set(pre)
-        self.post.set(post)
-        self.dWeights.set(dWeights)
-        self.post_term.set(post_term)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name,

tests/components/synapses/hebbian/test_BCMSynapse.py

@@ -0,0 +1,64 @@
+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 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
+
+def test_BCMSynapse1():
+    name = "bcm_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 = BCMSynapse(
+            name="a", shape=(1,1), tau_w=40., tau_theta=20., 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")
+        """
+
+    pre_value = jnp.ones((1, 1)) * 0.425
+    post_value = jnp.ones((1, 1)) * 1.55
+
+    truth = jnp.array([[-1.6798127]])
+    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)
+
+
+#test_BCMSynapse1()