update reset methods

Author: rxng8

ngclearn/components/neurons/graded/bernoulliErrorCell.py

@@ -12,7 +12,7 @@
 class BernoulliErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
     A simple (non-spiking) Bernoulli error cell - this is a fixed-point solution
-    of a mismatch signal. Specifically, this cell operates as a factorized multivariate 
+    of a mismatch signal. Specifically, this cell operates as a factorized multivariate
     Bernoulli distribution.
 
     | --- Cell Input Compartments: ---

ngclearn/components/neurons/graded/gaussianErrorCell.py

@@ -108,6 +108,32 @@ def advance_state(self, dt): ## compute Gaussian error cell output
         self.L.set(jnp.squeeze(L))
         self.mask.set(mask)
 
+    # @transition(output_compartments=["dmu", "dtarget", "dSigma", "target", "mu", "modulator", "L", "mask"])
+    # @staticmethod
+    @compilable
+    def reset(self, 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)
+
+        self.dmu.set(dmu)
+        self.dtarget.set(dtarget)
+        self.dSigma.set(dSigma)
+        self.target.set(target)
+        self.mu.set(mu)
+        self.modulator.set(modulator)
+        self.L.set(L)
+        self.mask.set(mask)
+
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/neurons/graded/laplacianErrorCell.py

@@ -103,6 +103,26 @@ def advance_state(self, dt): ## compute Laplacian error cell output
         self.L.set(jnp.squeeze(L))
         self.mask.set(mask)
 
+    def reset(self, 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))
+
+        self.dshift.set(dshift)
+        self.dtarget.set(dtarget)
+        self.dScale.set(dScale)
+        self.target.set(target)
+        self.shift.set(shift)
+        self.modulator.set(modulator)
+        self.L.set(L)
+        self.mask.set(mask)
+
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/neurons/graded/rateCell.py

@@ -251,6 +251,18 @@ def advance_state(self, dt):
         self.z.set(z)
         self.zF.set(zF)
 
+    @compilable
+    def reset(self, 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)
+        self.j.set(restVals)
+        self.j_td.set(restVals)
+        self.z.set(restVals)
+        self.zF.set(restVals)
+
+
     def save(self, directory, **kwargs):
         ## do a protected save of constants, depending on whether they are floats or arrays
         tau_m = (self.tau_m if isinstance(self.tau_m, float)

ngclearn/components/neurons/graded/rewardErrorCell.py

@@ -91,6 +91,18 @@ def evolve(self, dt):
         # Update compartment
         self.mu.set(mu)
 
+    @compilable
+    def reset(self, batch_size, n_units):
+        restVals = jnp.zeros((batch_size, n_units))
+        mu = restVals
+        rpe = restVals
+        accum_reward = restVals
+        n_ep_steps = jnp.zeros((batch_size, 1))
+        self.mu.set(mu)
+        self.rpe.set(rpe)
+        self.accum_reward.set(accum_reward)
+        self.n_ep_steps.set(n_ep_steps)
+
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/other/expKernel.py

@@ -71,13 +71,13 @@ def advance_state(self, t):
         # Get the variables
         inputs = self.inputs.get()
         tf = self.tf.get()
-        
+
         s = inputs
         ## update spike time window and corresponding window volume
         tf, epsp = _apply_kernel(
             tf, s, t, self.tau_w, self.win_len, krn_start=0, krn_end=self.win_len-1
         ) #0:win_len-1)
-        
+
         # Update compartments
         self.epsp.set(epsp)
         self.tf.set(tf)