minor cleanup/patches to rate-cell/lif/hebb-syn/trace-stdp and dim_reduce method

Author: Alexander Ororbia

ngclearn/components/neurons/graded/rateCell.py

@@ -145,6 +145,8 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
 
         act_fx: string name of activation function/nonlinearity to use
 
+        output_scale: factor to multiply output of nonlinearity of this cell by (Default: 1.)
+
         integration_type: type of integration to use for this cell's dynamics;
             current supported forms include "euler" (Euler/RK-1 integration)
             and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler")
@@ -157,12 +159,13 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
     """
 
     # Define Functions
-    def __init__(self, name, n_units, tau_m, prior=("gaussian", 0.), act_fx="identity",
-                 threshold=("none", 0.), integration_type="euler",
-                 batch_size=1, resist_scale=1., shape=None, is_stateful=True, **kwargs):
+    def __init__(
+            self, name, n_units, tau_m, prior=("gaussian", 0.), act_fx="identity", output_scale=1., threshold=("none", 0.),
+            integration_type="euler", batch_size=1, resist_scale=1., shape=None, is_stateful=True, **kwargs):
         super().__init__(name, **kwargs)
 
         ## membrane parameter setup (affects ODE integration)
+        self.output_scale = output_scale
         self.tau_m = tau_m ## membrane time constant -- setting to 0 triggers "stateless" mode
         self.is_stateful = is_stateful
         if isinstance(tau_m, float):
@@ -211,8 +214,9 @@ def __init__(self, name, n_units, tau_m, prior=("gaussian", 0.), act_fx="identit
 
     @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, j, j_td, z):
+    def advance_state(
+            dt, fx, dfx, tau_m, priorLeakRate, intgFlag, priorType, resist_scale, thresholdType, thr_lmbda, is_stateful,
+            output_scale, j, j_td, z):
         #if tau_m > 0.:
         if is_stateful:
             ### run a step of integration over neuronal dynamics
@@ -231,12 +235,12 @@ def advance_state(dt, fx, dfx, tau_m, priorLeakRate, intgFlag, priorType,
             elif thresholdType == "cauchy_threshold":
                 tmp_z = threshold_cauchy(tmp_z, thr_lmbda)
             z = tmp_z ## pre-activation function value(s)
-            zF = fx(z) ## post-activation function value(s)
+            zF = fx(z) * 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)
+            zF = fx(z) * output_scale
         return j, j_td, z, zF
 
     @transition(output_compartments=["j", "j_td", "z", "zF"])

ngclearn/components/neurons/spiking/LIFCell.py

@@ -123,15 +123,14 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
                 "arctan" (arc-tangent estimator), and "secant_lif" (the
                 LIF-specialized secant estimator)
 
-        lower_clamp_voltage: if True, this will ensure voltage never is below
-            the value of `v_rest` (default: True)
+        v_min: minimum voltage to clamp dynamics to (Default: None)
     """ ## batch_size arg?
 
     @deprecate_args(thr_jitter=None, v_decay="conduct_leak")
     def __init__(
             self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., v_reset=-60., conduct_leak=1., tau_theta=1e7,
             theta_plus=0.05, refract_time=5., one_spike=False, integration_type="euler", surrogate_type="straight_through",
-            lower_clamp_voltage=True, **kwargs
+            v_min=None, max_one_spike=False, **kwargs
     ):
         super().__init__(name, **kwargs)
 
@@ -143,7 +142,8 @@ def __init__(
         self.tau_m = tau_m ## membrane time constant
         self.resist_m = resist_m ## resistance value
         self.one_spike = one_spike ## True => constrains system to simulate 1 spike per time step
-        self.lower_clamp_voltage = lower_clamp_voltage ## True ==> ensures voltage is never < v_rest
+        self.max_one_spike = max_one_spike
+        self.v_min = v_min ## ensures voltage is never < v_min
 
         self.v_rest = v_rest #-65. # mV
         self.v_reset = v_reset # -60. # -65. # mV (milli-volts)
@@ -189,11 +189,11 @@ def __init__(
     @transition(output_compartments=["v", "s", "s_raw", "rfr", "thr_theta", "tols", "key", "surrogate"])    
     @staticmethod
     def advance_state(
-            t, dt, tau_m, resist_m, v_rest, v_reset, g_L, refract_T, thr, tau_theta, theta_plus, 
-            one_spike, lower_clamp_voltage, intgFlag, d_spike_fx, key, j, v, rfr, thr_theta, tols
+            t, dt, tau_m, resist_m, v_rest, v_reset, g_L, refract_T, thr, tau_theta, theta_plus, one_spike, max_one_spike,
+            v_min, intgFlag, d_spike_fx, key, j, v, rfr, thr_theta, tols
     ):
         skey = None ## this is an empty dkey if single_spike mode turned off
-        if one_spike:
+        if one_spike and not max_one_spike:
             key, skey = random.split(key, 2)
         ## run one integration step for neuronal dynamics
         j = j * resist_m
@@ -209,6 +209,7 @@ def advance_state(
             _, _v = step_euler(0., v, _dfv, dt, v_params)
         ## obtain action potentials/spikes/pulses
         s = (_v > _v_thr) * 1.
+        v_prespike = v
         ## update refractory variables
         _rfr = (rfr + dt) * (1. - s)
         ## perform hyper-polarization of neuronal cells
@@ -223,6 +224,9 @@ def advance_state(
             rS = nn.one_hot(jnp.argmax(rS, axis=1), num_classes=s.shape[1],
                             dtype=jnp.float32)
             s = s * (1. - m_switch) + rS * m_switch
+        if max_one_spike:
+            rS = nn.one_hot(jnp.argmax(v_prespike, axis=1), num_classes=s.shape[1], dtype=jnp.float32) ## get max-volt spike
+            s = s * rS ## mask out non-max volt spikes
         ############################################################################
         raw_spikes = raw_s
         v = _v
@@ -234,8 +238,8 @@ def advance_state(
             thr_theta = _update_theta(dt, thr_theta, raw_spikes, tau_theta, theta_plus)
         ## update tols
         tols = (1. - s) * tols + (s * t)
-        if lower_clamp_voltage: ## ensure voltage never < v_rest
-            v = jnp.maximum(v, v_rest)
+        if v_min is not None: ## ensures voltage never < v_rest
+            v = jnp.maximum(v, v_min)
         return v, s, raw_spikes, rfr, thr_theta, tols, key, surrogate
 
     @transition(output_compartments=["j", "v", "s", "s_raw", "rfr", "tols", "surrogate"])

ngclearn/components/synapses/denseSynapse.py

@@ -41,8 +41,10 @@ class DenseSynapse(JaxComponent): ## base dense synaptic cable
     """
 
     # Define Functions
-    def __init__(self, name, shape, weight_init=None, bias_init=None,
-                 resist_scale=1., p_conn=1., batch_size=1, **kwargs):
+    def __init__(
+            self, name, shape, weight_init=None, bias_init=None, resist_scale=1., 
+            p_conn=1., batch_size=1, **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         self.batch_size = batch_size
@@ -63,7 +65,7 @@ def __init__(self, name, shape, weight_init=None, bias_init=None,
             mask = random.bernoulli(subkeys[1], p=p_conn, shape=shape)
             weights = weights * mask ## sparsify matrix
 
-        self.batch_size = 1
+        self.batch_size = batch_size #1
         ## Compartment setup
         preVals = jnp.zeros((self.batch_size, shape[0]))
         postVals = jnp.zeros((self.batch_size, shape[1]))

ngclearn/components/synapses/hebbian/hebbianSynapse.py

@@ -162,10 +162,11 @@ class HebbianSynapse(DenseSynapse):
 
     # Define Functions
     @deprecate_args(_rebind=False, w_decay='prior')
-    def __init__(self, name, shape, eta=0., weight_init=None, bias_init=None,
-                 w_bound=1., is_nonnegative=False, prior=("constant", 0.), w_decay=0., sign_value=1.,
-                 optim_type="sgd", pre_wght=1., post_wght=1., p_conn=1.,
-                 resist_scale=1., batch_size=1, **kwargs):
+    def __init__(
+            self, name, shape, eta=0., weight_init=None, bias_init=None, w_bound=1., is_nonnegative=False,
+            prior=("constant", 0.), w_decay=0., sign_value=1., optim_type="sgd", pre_wght=1., post_wght=1., p_conn=1.,
+            resist_scale=1., batch_size=1, **kwargs
+    ):
         super().__init__(name, shape, weight_init, bias_init, resist_scale,
                          p_conn, batch_size=batch_size, **kwargs)
 

ngclearn/components/synapses/hebbian/traceSTDPSynapse.py

@@ -68,20 +68,25 @@ class TraceSTDPSynapse(DenseSynapse): # power-law / trace-based STDP
     # Define Functions
     def __init__(
             self, name, shape, A_plus, A_minus, eta=1., mu=0., pretrace_target=0., weight_init=None, resist_scale=1.,
-            p_conn=1., w_bound=1., batch_size=1, **kwargs
+            p_conn=1., w_bound=1., tau_w=0., weight_mask=None, batch_size=1, **kwargs
     ):
         super().__init__(name, shape, weight_init, None, resist_scale,
                          p_conn, batch_size=batch_size, **kwargs)
 
         ## Synaptic hyper-parameters
         self.shape = shape ## shape of synaptic efficacy matrix
+        self.tau_w = tau_w
         self.mu = mu ## controls power-scaling of STDP rule
         self.preTrace_target = pretrace_target ## target (pre-synaptic) trace activity value # 0.7
         self.Aplus = A_plus ## LTP strength
         self.Aminus = A_minus ## LTD strength
         self.Rscale = resist_scale ## post-transformation scale factor
         self.w_bound = w_bound #1. ## soft weight constraint
         self.w_eps = 0. ## w_eps = 0.01
+        self.weight_mask = weight_mask
+        if self.weight_mask is None:
+            self.weight_mask = jnp.ones((1, 1))
+        self.weights.set(self.weights.value * self.weight_mask)
 
         ## Compartment setup
         preVals = jnp.zeros((self.batch_size, shape[0]))
@@ -93,6 +98,12 @@ def __init__(
         self.dWeights = Compartment(self.weights.value * 0)
         self.eta = Compartment(jnp.ones((1, 1)) * eta) ## global learning rate
 
+    #@transition(output_compartments=["outputs"])
+    #@staticmethod
+    #def advance_state(Rscale, inputs, weights, biases, weight_mask):
+    #    outputs = (jnp.matmul(inputs, weights * weight_mask) * Rscale) + biases
+    #    return outputs
+
     @staticmethod
     def _compute_update(
             dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
@@ -126,16 +137,23 @@ def _compute_update(
     @transition(output_compartments=["weights", "dWeights"])
     @staticmethod
     def evolve(
-            dt, w_bound, w_eps, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights, eta
+            dt, w_bound, w_eps, preTrace_target, mu, Aplus, Aminus, tau_w, preSpike, postSpike, preTrace, 
+            postTrace, weights, eta, weight_mask
     ):
+        #_wm = weight_mask #
+        _wm = (weight_mask != 0.)
         dWeights = TraceSTDPSynapse._compute_update(
             dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
         )
         ## do a gradient ascent update/shift
-        weights = weights + dWeights * eta
+        decayTerm = 0.
+        if tau_w > 0.:
+            decayTerm = weights / tau_w
+        weights = weights + (dWeights * eta) - decayTerm #weight_mask * eta)
         ## enforce non-negativity
         #w_eps = 0. # 0.01 # 0.001
         weights = jnp.clip(weights, w_eps, w_bound - w_eps)  # jnp.abs(w_bound))
+        weights = weights * _wm # weight_mask
         return weights, dWeights
 
     @transition(output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace", "dWeights"])

ngclearn/components/synapses/staticSynapse.py

@@ -27,4 +27,4 @@ class StaticSynapse(DenseSynapse):
                 this to < 1 and > 0. will result in a sparser synaptic structure
                 (lower values yield sparse structure)
         """
-    pass
+    pass

ngclearn/utils/viz/dim_reduce.py

@@ -29,7 +29,7 @@ def extract_pca_latents(vectors): ## PCA mapping routine
         z_2D = vectors
     return z_2D
 
-def extract_tsne_latents(vectors, perplexity=30, n_pca_comp=32): ## tSNE mapping routine
+def extract_tsne_latents(vectors, perplexity=30, n_pca_comp=32, batch_size=500): ## tSNE mapping routine
     """
     Projects collection of K vectors (stored in a matrix) to a two-dimensional (2D)
     visualization space via the t-distributed stochastic neighbor embedding
@@ -42,10 +42,13 @@ def extract_tsne_latents(vectors, perplexity=30, n_pca_comp=32): ## tSNE mapping
 
         perplexity: the perplexity control factor for t-SNE (Default: 30)
 
+        batch_size: number of sampled embedding vectors to use per iteration 
+            of online internal PCA
+
     Returns:
         a matrix (K x 2) of projected vectors (to 2D space)
     """
-    batch_size = 500 #50
+    #batch_size = 500 #50
     z_dim = vectors.shape[1]
     z_2D = None
     if z_dim != 2: