refactored lava components to new sim-lib

Author: Alexander Ororbia

ngclearn/components/lava/neurons/LIFCell.py

@@ -1,8 +1,11 @@
-from ngclearn import resolver, Component, Compartment
-from ngclearn.utils import tensorstats
 from ngclearn import numpy as jnp
-from ngclearn.utils.weight_distribution import initialize_params
 from ngcsimlib.logger import info, warn
+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.utils import tensorstats
 
 class LIFCell(Component): ## Lava-compliant leaky integrate-and-fire cell
     """
@@ -113,8 +116,9 @@ def __init__(self, name, n_units, dt, tau_m, thr_theta_init=None, resist_m=1.,
     def _init(self, thr_theta0):
         self.thr_theta.set(thr_theta0)
 
+    @transition(output_compartments=["v", "s", "rfr", "thr_theta"])
     @staticmethod
-    def _advance_state(dt, tau_m, R_m, v_rest, v_reset, v_decay, refract_T, thr, tau_theta,
+    def advance_state(dt, tau_m, R_m, v_rest, v_reset, v_decay, refract_T, thr, tau_theta,
                        theta_plus, j_exc, j_inh, v, s, rfr, thr_theta):
         #j = j * (tau_m/dt) ## scale electrical current
         j = j_exc - j_inh ## sum the excitatory and inhibitory input channels
@@ -136,16 +140,9 @@ def _advance_state(dt, tau_m, R_m, v_rest, v_reset, v_decay, refract_T, thr, tau
         #tols = (1. - s) * tols + (s * t)
         return v, s, rfr, thr_theta #, tols
 
-    @resolver(_advance_state)
-    def advance_state(self, v, s, rfr, thr_theta): #, tols):
-        self.v.set(v)
-        self.s.set(s)
-        self.rfr.set(rfr)
-        self.thr_theta.set(thr_theta)
-        #self.tols.set(tols)
-
+    @transition(output_compartments=["j_exc", "j_inh", "v", "s", "rfr"])
     @staticmethod
-    def _reset(batch_size, n_units, v_rest, refract_T):
+    def reset(batch_size, n_units, v_rest, refract_T):
         restVals = jnp.zeros((batch_size, n_units))
         j_exc = restVals #+ 0
         j_inh = restVals #+ 0
@@ -154,14 +151,6 @@ def _reset(batch_size, n_units, v_rest, refract_T):
         rfr = restVals + refract_T
         return j_exc, j_inh, v, s, rfr #, tols
 
-    @resolver(_reset)
-    def reset(self, j_exc, j_inh, v, s, rfr):#, tols):
-        self.j_exc.set(j_exc)
-        self.j_inh.set(j_inh)
-        self.v.set(v)
-        self.s.set(s)
-        self.rfr.set(rfr)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name,

ngclearn/components/lava/synapses/hebbianSynapse.py

@@ -1,14 +1,16 @@
-from ngclearn import resolver, Component, Compartment
-from ngclearn.utils import tensorstats
 from ngclearn import numpy as jnp
-from ngclearn.utils.weight_distribution import initialize_params
 from ngcsimlib.logger import info, warn
+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.utils import tensorstats
 
 class HebbianSynapse(Component): ## Lava-compliant Hebbian synapse
     """
-    A synaptic cable that adjusts its efficacies via a two-factor Hebbian
-    adjustment rule. This is a Lava-compliant synaptic cable that adjusts
-    with a hard-coded form of (stochastic) gradient ascent.
+    A synaptic cable that adjusts its efficacies via a two-factor Hebbian adjustment rule. This is a Lava-compliant
+    synaptic cable that adjusts with a hard-coded form of (stochastic) gradient ascent.
 
     | --- Synapse Input Compartments: (Takes wired-in signals) ---
     | inputs - input (pre-synaptic) stimulus
@@ -102,8 +104,9 @@ def _init(self, weights):
         self.post.set(postVals)
         self.weights.set(weights)
 
+    @transition(output_compartments=["outputs", "weights"])
     @staticmethod
-    def _advance_state(dt, Rscale, w_bounds, w_decay, inputs, weights,
+    def advance_state(dt, Rscale, w_bounds, w_decay, inputs, weights,
                        pre, post, eta):
         outputs = jnp.matmul(inputs, weights) * Rscale
         ########################################################################
@@ -119,13 +122,9 @@ def _advance_state(dt, Rscale, w_bounds, w_decay, inputs, weights,
         ########################################################################
         return outputs, weights
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, weights):
-        self.outputs.set(outputs)
-        self.weights.set(weights)
-
+    @transition(output_compartments=["inputs", "outputs", "pre", "post", "eta"])
     @staticmethod
-    def _reset(batch_size, rows, cols, eta0):
+    def reset(batch_size, rows, cols, eta0):
         preVals = jnp.zeros((batch_size, rows))
         postVals = jnp.zeros((batch_size, cols))
         return (
@@ -136,14 +135,6 @@ def _reset(batch_size, rows, cols, eta0):
             jnp.ones((1,1)) * eta0
         )
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, pre, post, eta):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.pre.set(pre)
-        self.post.set(post)
-        self.eta.set(eta)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, weights=self.weights.value)

ngclearn/components/lava/synapses/staticSynapse.py

@@ -1,14 +1,16 @@
 from ngclearn import numpy as jnp
-from ngclearn import resolver, Component, Compartment
-from ngclearn.utils import tensorstats
+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, warn
+from ngclearn.components.synapses.hebbian import TraceSTDPSynapse
+from ngclearn.utils import tensorstats
 
 class StaticSynapse(Component): ## Lava-compliant fixed/non-evolvable synapse
     """
-    A static (dense) synaptic cable; no form of synaptic evolution/adaptation
-    is in-built to this component. This a Lava-compliant version of the
-    static synapse component from the synapses sub-package of components.
+    A static (dense) synaptic cable; no form of synaptic evolution/adaptation is in-built to this component. This a
+    Lava-compliant version of the static synapse component from the synapses sub-package of components.
 
     | --- Synapse Input Compartments: (Takes wired-in signals) ---
     | inputs - input (pre-synaptic) stimulus
@@ -79,29 +81,22 @@ def _init(self, weights):
         self.outputs.set(postVals)
         self.weights.set(weights)
 
+    @transition(output_compartments=["outputs"])
     @staticmethod
-    def _advance_state(dt, Rscale, inputs, weights):
+    def advance_state(dt, Rscale, inputs, weights):
         outputs = jnp.matmul(inputs, weights) * Rscale
         return outputs
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs):
-        self.outputs.set(outputs)
-
+    @transition(output_compartments=["inputs", "outputs"])
     @staticmethod
-    def _reset(batch_size, rows, cols):
+    def reset(batch_size, rows, cols):
         preVals = jnp.zeros((batch_size, rows))
         postVals = jnp.zeros((batch_size, cols))
         return (
             preVals, # inputs
             postVals, # outputs
         )
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, weights=self.weights.value)

ngclearn/components/lava/synapses/traceSTDPSynapse.py

@@ -1,14 +1,16 @@
-from ngclearn import resolver, Component, Compartment
-from ngclearn.utils import tensorstats
 from ngclearn import numpy as jnp
-from ngclearn.utils.weight_distribution import initialize_params
 from ngcsimlib.logger import info, warn
+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.utils import tensorstats
 
 class TraceSTDPSynapse(Component): ## Lava-compliant trace-STDP synapse
     """
-    A synaptic cable that adjusts its efficacies via trace-based form of
-    spike-timing-dependent plasticity (STDP). This is a Lava-compliant synaptic
-    cable that adjusts with a hard-coded form of (stochastic) gradient ascent.
+    A synaptic cable that adjusts its efficacies via trace-based form of spike-timing-dependent plasticity (STDP).
+    This is a Lava-compliant synaptic cable that adjusts with a hard-coded form of (stochastic) gradient ascent.
 
     | --- Synapse Input Compartments: (Takes wired-in signals) ---
     | inputs - input (pre-synaptic) stimulus
@@ -120,8 +122,9 @@ def _init(self, weights):
         self.x_post.set(postVals) ## post-synaptic trace
         self.weights.set(weights)
 
+    @transition(output_compartments=["outputs", "weights"])
     @staticmethod
-    def _advance_state(dt, Rscale, Aplus, Aminus, w_bounds, w_decay, x_tar,
+    def advance_state(dt, Rscale, Aplus, Aminus, w_bounds, w_decay, x_tar,
                        inputs, weights, pre, x_pre, post, x_post, eta):
         outputs = jnp.matmul(inputs, weights) * Rscale
         ########################################################################
@@ -139,13 +142,9 @@ def _advance_state(dt, Rscale, Aplus, Aminus, w_bounds, w_decay, x_tar,
         ########################################################################
         return outputs, weights
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs, weights):
-        self.outputs.set(outputs)
-        self.weights.set(weights)
-
+    @transition(output_compartments=["inputs", "outputs", "pre", "post", "x_pre", "x_post", "eta"])
     @staticmethod
-    def _reset(batch_size, rows, cols, eta0):
+    def reset(batch_size, rows, cols, eta0):
         preVals = jnp.zeros((batch_size, rows))
         postVals = jnp.zeros((batch_size, cols))
         return (
@@ -158,16 +157,6 @@ def _reset(batch_size, rows, cols, eta0):
             jnp.ones((1, 1)) * eta0
         )
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, pre, post, x_pre, x_post, eta):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.pre.set(pre)
-        self.post.set(post)
-        self.x_pre.set(x_pre)
-        self.x_post.set(x_post)
-        self.eta.set(eta)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         jnp.savez(file_name, weights=self.weights.value)

ngclearn/components/lava/traces/gatedTrace.py

@@ -1,12 +1,15 @@
-from ngclearn import resolver, Component, Compartment
-from ngclearn.utils import tensorstats
-
 from ngclearn import numpy as jnp
-import time, sys
+from ngcsimlib.logger import info, warn
+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.utils import tensorstats
 
 class GatedTrace(Component): ## gated/piecewise low-pass filter
     """
-    A gated/piecewise variable trace (filter).
+    A gated/piecewise variable trace (filter). This is a Lava-compliant trace component.
 
     | --- Cell Input Compartments: (Takes wired-in signals) ---
     | inputs - input (takes wired-in external signals)
@@ -39,25 +42,18 @@ def __init__(self, name, n_units, dt, tau_tr, **kwargs):
         self.inputs = Compartment(restVals) # input compartment
         self.trace = Compartment(restVals)
 
+    @transition(output_compartments=["trace"])
     @staticmethod
-    def _advance_state(dt, tau_tr, inputs, trace):
+    def advance_state(dt, tau_tr, inputs, trace):
         trace = (trace * (1. - dt/tau_tr)) * (1. - inputs) + inputs
         return trace
 
-    @resolver(_advance_state)
-    def advance_state(self, trace):
-        self.trace.set(trace)
-
+    @transition(output_compartments=["inputs", "trace"])
     @staticmethod
-    def _reset(batch_size, n_units):
+    def reset(batch_size, n_units):
         restVals = jnp.zeros((batch_size, n_units))
         return restVals, restVals
 
-    @resolver(_reset)
-    def reset(self, inputs, trace):
-        self.inputs.set(inputs)
-        self.trace.set(trace)
-
     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