refactored regression module to be compliant with v3

Author: Alexander Ororbia

ngclearn/modules/__init__.py

@@ -2,7 +2,3 @@
 from .regression.lasso import Iterative_Lasso
 from .regression.ridge import Iterative_Ridge
 
-
-
-
-

ngclearn/modules/regression/__init__.py

@@ -2,8 +2,3 @@
 from .lasso import Iterative_Lasso
 from .ridge import Iterative_Ridge
 
-
-
-
-
-

ngclearn/modules/regression/elastic_net.py

@@ -11,7 +11,7 @@
 class Iterative_ElasticNet():
     """
         A neural circuit implementation of the iterative Elastic Net (L1 and L2) algorithm
-        using Hebbian learning update rule.
+        using a Hebbian learning update rule.
 
         The circuit implements sparse regression through Hebbian synapses with Elastic Net regularization.
 
@@ -21,8 +21,6 @@ class Iterative_ElasticNet():
         | dW_reg = (jnp.sign(W) * l1_ratio) + (W * (1-l1_ratio)/2)
         | dW/dt = dW + lmbda * dW_reg
 
-
-
         | --- Circuit Components: ---
         | W - HebbianSynapse for learning regularized dictionary weights
         | err - GaussianErrorCell for computing prediction errors
@@ -104,14 +102,6 @@ def __init__(self, key, name, sys_dim, dict_dim, batch_size, weight_fill=0.05, l
                      >> self.W.reset)
             self.reset = reset
 
-            # advance_cmd, advance_args =self.circuit.compile_by_key(self.W,  ## execute prediction synapses
-            #                                                    self.err,  ## finally, execute error neurons
-            #                                                    compile_key="advance_state")
-            # evolve_cmd, evolve_args =self.circuit.compile_by_key(self.W, compile_key="evolve")
-            # reset_cmd, reset_args =self.circuit.compile_by_key(self.err, self.W, compile_key="reset")
-            # # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            # self.dynamic()
-
     def batch_set(self, batch_size):
         self.W.batch_size = batch_size
         self.err.batch_size = batch_size
@@ -121,33 +111,6 @@ def clamp(self, y_scaled, X):
         self.W.pre.set(X)
         self.err.target.set(y_scaled)
 
-    # def dynamic(self):  ## create dynamic commands forself.circuit
-    #     W, err = self.circuit.get_components("W", "err")
-    #     self.self = W
-    #     self.err = err
-    #
-    #     @Context.dynamicCommand
-    #     def batch_set(batch_size):
-    #         self.W.batch_size = batch_size
-    #         self.err.batch_size = batch_size
-    #
-    #     @Context.dynamicCommand
-    #     def clamps(y_scaled, X):
-    #         self.W.inputs.set(X)
-    #         self.W.pre.set(X)
-    #         self.err.target.set(y_scaled)
-    #
-    #     self.circuit.wrap_and_add_command(jit(self.circuit.evolve), name="evolve")
-    #     self.circuit.wrap_and_add_command(jit(self.circuit.advance_state), name="advance")
-    #     self.circuit.wrap_and_add_command(jit(self.circuit.reset), name="reset")
-    #
-    #     @scanner
-    #     def _process(compartment_values, args):
-    #         _t, _dt = args
-    #         compartment_values = self.circuit.advance_state(compartment_values, t=_t, dt=_dt)
-    #         return compartment_values, compartment_values[self.W.weights.path]
-
-
     def thresholding(self, scale=1.):
         coef_old = self.coef_
         new_coeff = jnp.where(jnp.abs(coef_old) >= self.threshold, coef_old, 0.)
@@ -172,18 +135,3 @@ def fit(self, y, X):
 
         return self.coef_, self.err.mu.get(), self.err.L.get()
 
-        # self.circuit.reset()
-        # self.circuit.clamps(y_scaled=y, X=X)
-        #
-        # for i in range(self.epochs):
-        #     self.circuit._process(jnp.array([[self.dt * i, self.dt] for i in range(self.T)]))
-        #     self.circuit.evolve(t=self.T, dt=self.dt)
-        #
-        # self.coef_ = np.array(self.W.weights.value)
-        #
-        # return self.coef_, self.err.mu.value, self.err.L.value
-
-
-
-
-

ngclearn/modules/regression/lasso.py

@@ -1,26 +1,19 @@
-import jax
-import pandas as pd
-from jax import random, jit
 import numpy as np
-from scipy.integrate import solve_ivp
-import matplotlib.pyplot as plt
-from ngcsimlib.utils import Get_Compartment_Batch
-from ngclearn.utils.model_utils import normalize_matrix
 from ngclearn.utils import weight_distribution as dist
-from ngclearn import Context, numpy as jnp
-from ngclearn.components import (RateCell,
-                                 HebbianSynapse,
-                                 GaussianErrorCell,
-                                 StaticSynapse)
-from ngclearn.utils.model_utils import scanner
+from ngclearn import numpy as jnp
 
+from jax import numpy as jnp, random, jit
+from ngclearn import Context, MethodProcess
+from ngclearn.components.synapses.hebbian.hebbianSynapse import HebbianSynapse
+from ngclearn.components.neurons.graded.gaussianErrorCell import GaussianErrorCell
+from ngcsimlib.global_state import stateManager
 
 class Iterative_Lasso():
     """
         A neural circuit implementation of the iterative Lasso (L1) algorithm
-        using Hebbian learning update rule.
+        using a Hebbian learning update rule.
 
-        The circuit implements sparse coding through Hebbian synapses with L1 regularization.
+        The circuit implements sparse coding-like regression through Hebbian synapses with L1 regularization.
 
         The specific differential equation that characterizes this model is adding lmbda * sign(W)
         to the dW (the gradient of loss/energy function):
@@ -89,43 +82,32 @@ def __init__(self, key, name, sys_dim, dict_dim, batch_size, weight_fill=0.05, l
             self.W.batch_size = batch_size
             self.err.batch_size = batch_size
             # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            self.err.mu << self.W.outputs
-            self.W.post << self.err.dmu
+            self.W.outputs >> self.err.mu 
+            self.err.dmu >> self.W.post 
             # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            advance_cmd, advance_args =self.circuit.compile_by_key(self.W,  ## execute prediction synapses
-                                                               self.err,  ## finally, execute error neurons
-                                                               compile_key="advance_state")
-            evolve_cmd, evolve_args =self.circuit.compile_by_key(self.W, compile_key="evolve")
-            reset_cmd, reset_args =self.circuit.compile_by_key(self.err, self.W, compile_key="reset")
-            # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            self.dynamic()
-
-    def dynamic(self):  ## create dynamic commands for self.circuit
-        W, err = self.circuit.get_components("W", "err")
-        self.self = W
-        self.err = err
-
-        @Context.dynamicCommand
-        def batch_set(batch_size):
-            self.W.batch_size = batch_size
-            self.err.batch_size = batch_size
-
-        @Context.dynamicCommand
-        def clamps(y_scaled, X):
-            self.W.inputs.set(X)
-            self.W.pre.set(X)
-            self.err.target.set(y_scaled)
-
-        self.circuit.wrap_and_add_command(jit(self.circuit.evolve), name="evolve")
-        self.circuit.wrap_and_add_command(jit(self.circuit.advance_state), name="advance")
-        self.circuit.wrap_and_add_command(jit(self.circuit.reset), name="reset")
-
-        @scanner
-        def _process(compartment_values, args):
-            _t, _dt = args
-            compartment_values = self.circuit.advance_state(compartment_values, t=_t, dt=_dt)
-            return compartment_values, compartment_values[self.W.weights.path]
-
+           
+            advance = (MethodProcess(name="advance_state")
+                               >> self.W.advance_state
+                               >> self.err.advance_state)
+            self.advance = advance
+
+            evolve = (MethodProcess(name="evolve")
+                      >> self.W.evolve)
+            self.evolve = evolve
+
+            reset = (MethodProcess(name="reset")
+                     >> self.err.reset
+                     >> self.W.reset)
+            self.reset = reset
+
+    def batch_set(self, batch_size):
+        self.W.batch_size = batch_size
+        self.err.batch_size = batch_size
+
+    def clamp(self, y_scaled, X):
+        self.W.inputs.set(X)
+        self.W.pre.set(X)
+        self.err.target.set(y_scaled)
 
     def thresholding(self, scale=2):
         coef_old = self.coef_ 
@@ -136,23 +118,16 @@ def thresholding(self, scale=2):
 
         return self.coef_, coef_old
 
-
     def fit(self, y, X):
-
-        self.circuit.reset()
-        self.circuit.clamps(y_scaled=y, X=X)
+        self.reset.run()
+        self.clamp(y_scaled=y, X=X)
 
         for i in range(self.epochs):
-            self.circuit._process(jnp.array([[self.dt * i, self.dt] for i in range(self.T)]))
-            self.circuit.evolve(t=self.T, dt=self.dt)
-
-        self.coef_ = np.array(self.W.weights.value)
-
-        return self.coef_, self.err.mu.value, self.err.L.value
-
-
-
-
+            inputs = jnp.array(self.advance.pack_rows(self.T, t=lambda x: x, dt=self.dt))
+            stateManager.state, outputs = self.advance.scan(inputs)
+            self.evolve.run(t=self.T, dt=self.dt)
 
+        self.coef_ = np.array(self.W.weights.get())
 
+        return self.coef_, self.err.mu.get(), self.err.L.get()
 

ngclearn/modules/regression/ridge.py

@@ -1,21 +1,19 @@
-from jax import random, jit
 import numpy as np
 from ngclearn.utils import weight_distribution as dist
-from ngclearn import Context, numpy as jnp
-from ngclearn.components import (RateCell,
-                                 HebbianSynapse,
-                                 GaussianErrorCell,
-                                 StaticSynapse)
-from ngclearn.utils.model_utils import scanner
-
+from ngclearn import numpy as jnp
 
+from jax import numpy as jnp, random, jit
+from ngclearn import Context, MethodProcess
+from ngclearn.components.synapses.hebbian.hebbianSynapse import HebbianSynapse
+from ngclearn.components.neurons.graded.gaussianErrorCell import GaussianErrorCell
+from ngcsimlib.global_state import stateManager
 
 class Iterative_Ridge():
     """
         A neural circuit implementation of the iterative Ridge (L2) algorithm
-        using Hebbian learning update rule.
+        using a Hebbian learning update rule.
 
-        The circuit implements sparse regression through Hebbian synapses with L2 regularization.
+        This circuit implements sparse regression through Hebbian synapses with L2 regularization.
 
         The specific differential equation that characterizes this model is adding lmbda * W
         to the dW (the gradient of loss/energy function):
@@ -75,54 +73,43 @@ def __init__(self, key, name, sys_dim, dict_dim, batch_size, weight_fill=0.05, l
         feature_dim = dict_dim
 
         with Context(self.name) as self.circuit:
-            self.W = HebbianSynapse("W", shape=(feature_dim, sys_dim), eta=self.lr,
-                                   sign_value=-1, weight_init=dist.constant(weight_fill),
-                                   prior=('ridge', ridge_lmbda), w_bound=0.,
-                                   optim_type=optim_type, key=subkeys[0])
+            self.W = HebbianSynapse(
+                "W", shape=(feature_dim, sys_dim), eta=self.lr, sign_value=-1, 
+                weight_init=dist.constant(weight_fill), prior=('ridge', ridge_lmbda), w_bound=0.,
+                optim_type=optim_type, key=subkeys[0]
+            )
             self.err = GaussianErrorCell("err", n_units=sys_dim)
 
             # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
             self.W.batch_size = batch_size
             self.err.batch_size = batch_size
             # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            self.err.mu << self.W.outputs
-            self.W.post << self.err.dmu
+            self.W.outputs >> self.err.mu
+            self.err.dmu >> self.W.post
             # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            advance_cmd, advance_args =self.circuit.compile_by_key(self.W,  ## execute prediction synapses
-                                                               self.err,  ## finally, execute error neurons
-                                                               compile_key="advance_state")
-            evolve_cmd, evolve_args =self.circuit.compile_by_key(self.W, compile_key="evolve")
-            reset_cmd, reset_args =self.circuit.compile_by_key(self.err, self.W, compile_key="reset")
-            # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            self.dynamic()
-
-    def dynamic(self):  ## create dynamic commands forself.circuit
-        W, err = self.circuit.get_components("W", "err")
-        self.self = W
-        self.err = err
-
-        @Context.dynamicCommand
-        def batch_set(batch_size):
-            self.W.batch_size = batch_size
-            self.err.batch_size = batch_size
-
-        @Context.dynamicCommand
-        def clamps(y_scaled, X):
-            self.W.inputs.set(X)
-            self.W.pre.set(X)
-            self.err.target.set(y_scaled)
-
-        self.circuit.wrap_and_add_command(jit(self.circuit.evolve), name="evolve")
-        self.circuit.wrap_and_add_command(jit(self.circuit.advance_state), name="advance")
-        self.circuit.wrap_and_add_command(jit(self.circuit.reset), name="reset")
-
-
-        @scanner
-        def _process(compartment_values, args):
-            _t, _dt = args
-            compartment_values = self.circuit.advance_state(compartment_values, t=_t, dt=_dt)
-            return compartment_values, compartment_values[self.W.weights.path]
 
+            advance = (MethodProcess(name="advance_state")
+                               >> self.W.advance_state
+                               >> self.err.advance_state)
+            self.advance = advance
+
+            evolve = (MethodProcess(name="evolve")
+                      >> self.W.evolve)
+            self.evolve = evolve
+
+            reset = (MethodProcess(name="reset")
+                     >> self.err.reset
+                     >> self.W.reset)
+            self.reset = reset
+            
+    def batch_set(self, batch_size):
+        self.W.batch_size = batch_size
+        self.err.batch_size = batch_size
+
+    def clamp(self, y_scaled, X):
+        self.W.inputs.set(X)
+        self.W.pre.set(X)
+        self.err.target.set(y_scaled)
 
     def thresholding(self, scale=2):
         coef_old = self.coef_ #self.W.weights.value
@@ -135,21 +122,15 @@ def thresholding(self, scale=2):
 
 
     def fit(self, y, X):
-        self.circuit.reset()
-        self.circuit.clamps(y_scaled=y, X=X)
+        self.reset.run()
+        self.clamp(y_scaled=y, X=X)
 
         for i in range(self.epochs):
-            self.circuit._process(jnp.array([[self.dt * i, self.dt] for i in range(self.T)]))
-            self.circuit.evolve(t=self.T, dt=self.dt)
-
-        self.coef_ = np.array(self.W.weights.value)
-
-        return self.coef_, self.err.mu.value, self.err.L.value
-
-
-
-
-
+            inputs = jnp.array(self.advance.pack_rows(self.T, t=lambda x: x, dt=self.dt))
+            stateManager.state, outputs = self.advance.scan(inputs)
+            self.evolve.run(t=self.T, dt=self.dt)
 
+        self.coef_ = np.array(self.W.weights.get())
 
+        return self.coef_, self.err.mu.get(), self.err.L.get()