Refactoring neuronal and synaptic components (#123) - merge from fork to v3

* refactoring graded cells

* update refactored models

* update sLIF cell

---------

Co-authored-by: Alex Ororbia <agocse109@gmail.com>

Author: rxng8

ngclearn/components/neurons/graded/bernoulliErrorCell.py

@@ -1,9 +1,13 @@
-from ngclearn import resolver, Component, Compartment
+# %%
+
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, jit
 from ngclearn.utils import tensorstats
 from ngclearn.utils.model_utils import sigmoid, d_sigmoid
-from ngcsimlib.compilers.process import transition
+
+from ngcsimlib.logger import info
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.parser import compilable
 
 class BernoulliErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
@@ -59,14 +63,20 @@ def __init__(self, name, n_units, batch_size=1, input_logits=False, shape=None,
         self.modulator = Compartment(restVals + 1.0) # to be set/consumed
         self.mask = Compartment(restVals + 1.0)
 
-    @transition(output_compartments=["dp", "dtarget", "L", "mask"])
-    @staticmethod
-    def advance_state(dt, p, target, modulator, mask, input_logits): ## compute Bernoulli error cell output
+    # @transition(output_compartments=["dp", "dtarget", "L", "mask"])
+    @compilable
+    def advance_state(self, dt): ## compute Bernoulli error cell output
+        # Get the variables
+        p = self.p.get()
+        target = self.target.get()
+        modulator = self.modulator.get()
+        mask = self.mask.get()
+
         # Moves Bernoulli error cell dynamics one step forward. Specifically, this routine emulates the error unit
         # behavior of the local cost functional
         eps = 0.0001
         _p = p
-        if input_logits: ## convert from "logits" to probs via sigmoidal link function
+        if self.input_logits: ## convert from "logits" to probs via sigmoidal link function
             _p = sigmoid(p)
         _p = jnp.clip(_p, eps, 1. - eps) ## post-process to prevent div by 0
         x = target
@@ -78,7 +88,7 @@ def advance_state(dt, p, target, modulator, mask, input_logits): ## compute Bern
         log_p = jnp.log(_p) ## ln(p)
         log_one_min_p = jnp.log(one_min_p) ## ln(1 - p)
         L = jnp.sum(log_p * x + log_one_min_p * one_min_x) ## Bern LL
-        if input_logits:
+        if self.input_logits:
             dL_dp = x - _p ## d(Bern LL)/dp where _p = sigmoid(p)
         else:
             dL_dp = x/(_p) - one_min_x/one_min_p  ## d(Bern LL)/dp
@@ -89,14 +99,21 @@ def advance_state(dt, p, target, modulator, mask, input_logits): ## compute Bern
         dp = dp * modulator * mask ## NOTE: how does mask apply to a multivariate Bernoulli?
         dtarget = dL_dx * modulator * mask
         mask = mask * 0. + 1. ## "eat" the mask as it should only apply at time t
-        return dp, dtarget, jnp.squeeze(L), mask
-
-    @transition(output_compartments=["dp", "dtarget", "target", "p", "modulator", "L", "mask"])
-    @staticmethod
-    def reset(batch_size, shape): ## reset core components/statistics
-        _shape = (batch_size, shape[0])
-        if len(shape) > 1:
-            _shape = (batch_size, shape[0], shape[1], shape[2])
+
+        # Set state
+        # dp, dtarget, jnp.squeeze(L), mask
+        self.dp.set(dp)
+        self.dtarget.set(dtarget)
+        self.L.set(jnp.squeeze(L))
+        self.mask.set(mask)
+
+
+    # @transition(output_compartments=["dp", "dtarget", "target", "p", "modulator", "L", "mask"])
+    @compilable
+    def reset(self, batch_size): ## reset core components/statistics
+        _shape = (batch_size, self.shape[0])
+        if len(self.shape) > 1:
+            _shape = (batch_size, self.shape[0], self.shape[1], self.shape[2])
         restVals = jnp.zeros(_shape) ## "rest"/reset values
         dp = restVals
         dtarget = restVals
@@ -105,7 +122,16 @@ def reset(batch_size, shape): ## reset core components/statistics
         modulator = restVals + 1. ## reset modulator signal
         L = 0. #jnp.zeros((1, 1)) ## rest loss
         mask = jnp.ones(_shape) ## reset mask
-        return dp, dtarget, target, p, modulator, L, mask
+
+        # Set compartment
+        self.dp.set(dp)
+        self.dtarget.set(dtarget)
+        self.target.set(target)
+        self.p.set(p)
+        self.modulator.set(modulator)
+        self.L.set(L)
+        self.mask.set(mask)
+
 
     @classmethod
     def help(cls): ## component help function
@@ -136,11 +162,11 @@ def help(cls): ## component help function
         return info
 
     def __repr__(self):
-        comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
+        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
         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)
+            stats = tensorstats(getattr(self, c).get())
             if stats is not None:
                 line = [f"{k}: {v}" for k, v in stats.items()]
                 line = ", ".join(line)

ngclearn/components/neurons/graded/gaussianErrorCell.py

@@ -1,8 +1,12 @@
-from ngclearn import resolver, Component, Compartment
+# %%
+
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, jit
 from ngclearn.utils import tensorstats
-from ngcsimlib.compilers.process import transition
+
+from ngcsimlib.logger import info
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.parser import compilable
 
 class GaussianErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
@@ -71,9 +75,15 @@ def eval_log_density(target, mu, Sigma):
         log_density = -jnp.sum(jnp.square(_dmu)) * (0.5 / Sigma)
         return log_density
 
-    @transition(output_compartments=["dmu", "dtarget", "dSigma", "L", "mask"])
-    @staticmethod
-    def advance_state(dt, mu, target, Sigma, modulator, mask): ## compute Gaussian error cell output
+    @compilable
+    def advance_state(self, dt): ## compute Gaussian error cell output
+        # Get the variables
+        mu = self.mu.get()
+        target = self.target.get()
+        Sigma = self.Sigma.get()
+        modulator = self.modulator.get()
+        mask = self.mask.get()
+
         # Moves Gaussian cell dynamics one step forward. Specifically, this routine emulates the error unit
         # behavior of the local cost functional:
         # FIXME: Currently, below does: L(targ, mu) = -(1/(2*sigma)) * ||targ - mu||^2_2
@@ -90,24 +100,13 @@ def advance_state(dt, mu, target, Sigma, modulator, mask): ## compute Gaussian e
         dmu = dmu * modulator * mask ## not sure how mask will apply to a full covariance...
         dtarget = dtarget * modulator * mask
         mask = mask * 0. + 1. ## "eat" the mask as it should only apply at time t
-        return dmu, dtarget, dSigma, jnp.squeeze(L), mask
 
-    @transition(output_compartments=["dmu", "dtarget", "dSigma", "target", "mu", "modulator", "L", "mask"])
-    @staticmethod
-    def reset(batch_size, shape, sigma_shape): ## reset core components/statistics
-        _shape = (batch_size, shape[0])
-        if len(shape) > 1:
-            _shape = (batch_size, shape[0], shape[1], shape[2])
-        restVals = jnp.zeros(_shape)
-        dmu = restVals
-        dtarget = restVals
-        dSigma = jnp.zeros(sigma_shape)
-        target = restVals
-        mu = restVals
-        modulator = mu + 1.
-        L = 0. #jnp.zeros((1, 1))
-        mask = jnp.ones(_shape)
-        return dmu, dtarget, dSigma, target, mu, modulator, L, mask
+        # Update compartments
+        self.dmu.set(dmu)
+        self.dtarget.set(dtarget)
+        self.dSigma.set(dSigma)
+        self.L.set(jnp.squeeze(L))
+        self.mask.set(mask)
 
     @classmethod
     def help(cls): ## component help function
@@ -140,11 +139,11 @@ def help(cls): ## component help function
         return info
 
     def __repr__(self):
-        comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
+        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
         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)
+            stats = tensorstats(getattr(self, c).get())
             if stats is not None:
                 line = [f"{k}: {v}" for k, v in stats.items()]
                 line = ", ".join(line)

ngclearn/components/neurons/graded/laplacianErrorCell.py

@@ -1,8 +1,12 @@
-from ngclearn import resolver, Component, Compartment
+# %%
+
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, jit
 from ngclearn.utils import tensorstats
-from ngcsimlib.compilers.process import transition
+
+from ngcsimlib.logger import info
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.parser import compilable
 
 class LaplacianErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
@@ -44,7 +48,7 @@ def __init__(self, name, n_units, batch_size=1, scale=1., shape=None, **kwargs):
         else:
             _shape = (batch_size, shape[0], shape[1], shape[2])  ## shape is 4D tensor
         scale_shape = (1, 1)
-        if not isinstance(scale, float) and not isinstance(sigma, int):
+        if not isinstance(scale, float) and not isinstance(scale, int):
             scale_shape = jnp.array(scale).shape
         self.scale_shape = scale_shape
         ## Layer Size setup
@@ -67,9 +71,15 @@ def __init__(self, name, n_units, batch_size=1, scale=1., shape=None, **kwargs):
         self.modulator = Compartment(restVals + 1.0) ## to be set/consumed
         self.mask = Compartment(restVals + 1.0)
 
-    @transition(output_compartments=["dshift", "dtarget", "dScale", "L", "mask"])
-    @staticmethod
-    def advance_state(dt, shift, target, Scale, modulator, mask): ## compute Laplacian error cell output
+    @compilable
+    def advance_state(self, dt): ## compute Laplacian error cell output
+        # Get the variables
+        shift = self.shift.get()
+        target = self.target.get()
+        Scale = self.Scale.get()
+        modulator = self.modulator.get()
+        mask = self.mask.get()
+
         # Moves Laplacian cell dynamics one step forward. Specifically, this routine emulates the error unit
         # behavior of the local cost functional:
         # FIXME: Currently, below does: L(targ, shift) = -||targ - shift||_1/scale
@@ -85,21 +95,13 @@ def advance_state(dt, shift, target, Scale, modulator, mask): ## compute Laplaci
         dshift = dshift * modulator * mask
         dtarget = dtarget * modulator * mask
         mask = mask * 0. + 1.  ## "eat" the mask as it should only apply at time t
-        return dshift, dtarget, dScale, jnp.squeeze(L), mask
-
-    @transition(output_compartments=["dshift", "dtarget", "dScale", "target", "shift", "modulator", "L", "mask"])
-    @staticmethod
-    def reset(batch_size, n_units, scale_shape):
-        restVals = jnp.zeros((batch_size, n_units))
-        dshift = restVals
-        dtarget = restVals
-        dScale = jnp.zeros(scale_shape)
-        target = restVals
-        shift = restVals
-        modulator = shift + 1.
-        L = 0.
-        mask = jnp.ones((batch_size, n_units))
-        return dshift, dtarget, dScale, target, shift, modulator, L, mask
+
+        # Update compartments
+        self.dshift.set(dshift)
+        self.dtarget.set(dtarget)
+        self.dScale.set(dScale)
+        self.L.set(jnp.squeeze(L))
+        self.mask.set(mask)
 
     @classmethod
     def help(cls): ## component help function
@@ -131,11 +133,11 @@ def help(cls): ## component help function
         return info
 
     def __repr__(self):
-        comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
+        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
         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)
+            stats = tensorstats(getattr(self, c).get())
             if stats is not None:
                 line = [f"{k}: {v}" for k, v in stats.items()]
                 line = ", ".join(line)

ngclearn/components/neurons/graded/rateCell.py

@@ -11,6 +11,9 @@
 from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \
                                             step_euler, step_rk2, step_rk4
 
+from ngcsimlib.logger import info
+from ngcsimlib.parser import compilable
+
 def _dfz_internal_laplace(z, j, j_td, tau_m, leak_gamma): ## raw dynamics
     z_leak = jnp.sign(z) ## d/dx of Laplace is signum
     dz_dt = (-z_leak * leak_gamma + (j + j_td)) * (1./tau_m)
@@ -198,7 +201,6 @@ def __init__(
         self.n_units = n_units
         self.batch_size = batch_size
 
-        
         omega_0 = None
         if act_fx == "sine":
             omega_0 = kwargs["omega_0"]
@@ -211,46 +213,43 @@ def __init__(
         self.j_td = Compartment(restVals, display_name="Modulatory Stimulus Current", units="mA") # top-down electrical current - pressure
         self.z = Compartment(restVals, display_name="Rate Activity", units="mA") # rate activity
 
-    @transition(output_compartments=["j", "j_td", "z", "zF"])
-    @staticmethod
-    def advance_state(
-            dt, fx, dfx, tau_m, priorLeakRate, intgFlag, priorType, resist_scale, thresholdType, thr_lmbda, is_stateful,
-            output_scale, j, j_td, z):
+    @compilable
+    def advance_state(self, dt):
+        # Get the compartment values
+        j = self.j.get()
+        j_td = self.j_td.get()
+        z = self.z.get()
+
         #if tau_m > 0.:
-        if is_stateful:
+        if self.is_stateful:
             ### run a step of integration over neuronal dynamics
             ## Notes:
             ## self.pressure <-- "top-down" expectation / contextual pressure
             ## self.current <-- "bottom-up" data-dependent signal
-            dfx_val = dfx(z)
+            dfx_val = self.dfx(z)
             j = _modulate(j, dfx_val)
-            j = j * resist_scale
+            j = j * self.resist_scale
             tmp_z = _run_cell(dt, j, j_td, z,
-                              tau_m, leak_gamma=priorLeakRate,
-                              integType=intgFlag, priorType=priorType)
+                              self.tau_m, leak_gamma=self.priorLeakRate,
+                              integType=self.intgFlag, priorType=self.priorType)
             ## apply optional thresholding sub-dynamics
-            if thresholdType == "soft_threshold":
-                tmp_z = threshold_soft(tmp_z, thr_lmbda)
-            elif thresholdType == "cauchy_threshold":
-                tmp_z = threshold_cauchy(tmp_z, thr_lmbda)
+            if self.thresholdType == "soft_threshold":
+                tmp_z = threshold_soft(tmp_z, self.thr_lmbda)
+            elif self.thresholdType == "cauchy_threshold":
+                tmp_z = threshold_cauchy(tmp_z, self.thr_lmbda)
             z = tmp_z ## pre-activation function value(s)
-            zF = fx(z) * output_scale ## post-activation function value(s)
+            zF = self.fx(z) * self.output_scale ## post-activation function value(s)
         else:
             ## run in "stateless" mode (when no membrane time constant provided)
             j_total = j + j_td
             z = _run_cell_stateless(j_total)
-            zF = fx(z) * output_scale
-        return j, j_td, z, zF
-
-    @transition(output_compartments=["j", "j_td", "z", "zF"])
-    @staticmethod
-    def reset(batch_size, shape): #n_units
-        _shape = (batch_size, shape[0])
-        if len(shape) > 1:
-            _shape = (batch_size, shape[0], shape[1], shape[2])
-        restVals = jnp.zeros(_shape)
-        return tuple([restVals for _ in range(4)])
+            zF = self.fx(z) * self.output_scale
 
+        # Update compartments
+        self.j.set(j)
+        self.j_td.set(j_td)
+        self.z.set(z)
+        self.zF.set(zF)
 
     def save(self, directory, **kwargs):
         ## do a protected save of constants, depending on whether they are floats or arrays
@@ -308,11 +307,11 @@ def help(cls): ## component help function
         return info
 
     def __repr__(self):
-        comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
+        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
         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)
+            stats = tensorstats(getattr(self, c).get())
             if stats is not None:
                 line = [f"{k}: {v}" for k, v in stats.items()]
                 line = ", ".join(line)