Merging over Dynamics feature branch to main (#92)

* modded bernoulli-cell to include max-frequency constraint

* added warning check to bernoulli, some cleanup

* integrated if-cell, cleaned up lif and inits

* mod to latency-cell

* updated the poissonCell to be a true poisson

* fixed minor bug in deprecation for poiss/bern

* fixed minor bug in deprecation for poiss/bern

* fixed validation fun in bern/poiss

* moved back and cleaned up bernoulli and poisson cells

* added threshold-clipping to latency cell

* updates to if/lif

* added batch-size arg to slif

* fixed minor load bug in lif-cell

* fixed a blocking jit-partial call in lif update_theta method; when loading

* minor edit to dim-reduce

* updated monitor plot code

* update to dim-reduce

* integrated phasor-cell, minor cleanup of latency

* tweak to adex thr arg

* tweak to adex thr arg

* integrated resonate-and-fire neuronal cell

* mod to raf-cell

* cleaned up raf

* cleaned up raf

* cleaned up raf-cell

* cleaned up raf-cell

* cleaned up raf-cell

* minor tweak to dim-reduce in utils

* Additions for inhibition stuff

* update to API modeling docs to reflect RAF neuronal cell

---------

Co-authored-by: Alexander Ororbia <[email protected]>
Co-authored-by: Will Gebhardt <[email protected]>

Author: ago109

docs/modeling/neurons.md

@@ -191,6 +191,23 @@ cell supports either Euler or midpoint method / RK-2 integration.)
     :noindex:
 ```
 
+### The Resonate-and-Fire (RAF) Cell
+
+This cell models dynamics over voltage `v` and a angular driver state/variable `w`; these 
+two variables result in a dampened oscillatory spiking neuronal cell). In effect, the 
+resonatoe-and-fire (RAF) model (or "resonator") evolves as a result of two coupled 
+differential equations. (Note that this cell supports either Euler or RK-2 integration.)
+
+```{eval-rst}
+.. autoclass:: ngclearn.components.RAFCell
+  :noindex:
+
+  .. automethod:: advance_state
+    :noindex:
+  .. automethod:: reset
+    :noindex:
+```
+
 ### The Izhikevich Cell
 
 This cell models dynamics over voltage `v` and a recover variable `w` (where `w`

ngclearn/components/__init__.py

@@ -10,14 +10,14 @@
 
 ## point to standard spiking cell component types
 from .neurons.spiking.sLIFCell import SLIFCell
+from .neurons.spiking.IFCell import IFCell
 from .neurons.spiking.LIFCell import LIFCell
 from .neurons.spiking.WTASCell import WTASCell
 from .neurons.spiking.quadLIFCell import QuadLIFCell
 from .neurons.spiking.adExCell import AdExCell
 from .neurons.spiking.fitzhughNagumoCell import FitzhughNagumoCell
 from .neurons.spiking.izhikevichCell import IzhikevichCell
-
-
+from .neurons.spiking.RAFCell import RAFCell
 ## point to transformer/operater component types
 from .other.varTrace import VarTrace
 from .other.expKernel import ExpKernel
@@ -28,7 +28,7 @@
 from .input_encoders.bernoulliCell import BernoulliCell
 from .input_encoders.poissonCell import PoissonCell
 from .input_encoders.latencyCell import LatencyCell
-
+from .input_encoders.phasorCell import PhasorCell
 
 ## point to synapse component types
 from .synapses.denseSynapse import DenseSynapse

ngclearn/components/base_monitor.py

@@ -124,9 +124,10 @@ def watch(self, compartment, window_length):
         """
         cs, end = self._add_path(compartment.path)
 
+        dtype = compartment.value.dtype
         shape = compartment.value.shape
-        new_comp = Compartment(np.zeros(shape))
-        new_comp_store = Compartment(np.zeros((window_length, *shape)))
+        new_comp = Compartment(np.zeros(shape, dtype=dtype))
+        new_comp_store = Compartment(np.zeros((window_length, *shape), dtype=dtype))
 
         comp_key = "*".join(compartment.path.split("/"))
         store_comp_key = comp_key + "*store"
@@ -310,4 +311,4 @@ def make_plot(self, compartment, ax=None, ylabel=None, xlabel=None, title=None,
             for k in range(n):
                 _ax.plot(vals[:, 0, k])
         else:
-            plot_func(vals, ax=_ax)
+            plot_func(vals[:, :, 0:n], ax=_ax)

ngclearn/components/input_encoders/__init__.py

@@ -1,3 +1,4 @@
 from .bernoulliCell import BernoulliCell
 from .poissonCell import PoissonCell
 from .latencyCell import LatencyCell
+from .phasorCell import PhasorCell

ngclearn/components/input_encoders/bernoulliCell.py

@@ -2,6 +2,9 @@
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, random, jit
 from ngclearn.utils import tensorstats
+from functools import partial
+from ngcsimlib.deprecators import deprecate_args
+from ngcsimlib.logger import info, warn
 
 @jit
 def _update_times(t, s, tols):
@@ -21,25 +24,10 @@ def _update_times(t, s, tols):
     _tols = (1. - s) * tols + (s * t)
     return _tols
 
-@jit
-def _sample_bernoulli(dkey, data):
-    """
-    Samples a Bernoulli spike train on-the-fly
-
-    Args:
-        dkey: JAX key to drive stochasticity/noise
-
-        data: sensory data (vector/matrix)
-
-    Returns:
-        binary spikes
-    """
-    s_t = random.bernoulli(dkey, p=data).astype(jnp.float32)
-    return s_t
-
 class BernoulliCell(JaxComponent):
     """
-    A Bernoulli cell that produces Bernoulli-distributed spikes on-the-fly.
+    A Bernoulli cell that produces spikes by sampling a Bernoulli distribution
+    on-the-fly (to produce data-scaled Bernoulli spike trains).
 
     | --- Cell Input Compartments: ---
     | inputs - input (takes in external signals)
@@ -55,7 +43,6 @@ class BernoulliCell(JaxComponent):
         n_units: number of cellular entities (neural population size)
     """
 
-    # Define Functions
     def __init__(self, name, n_units, batch_size=1, **kwargs):
         super().__init__(name, **kwargs)
 
@@ -72,9 +59,9 @@ def __init__(self, name, n_units, batch_size=1, **kwargs):
     @staticmethod
     def _advance_state(t, key, inputs, tols):
         key, *subkeys = random.split(key, 2)
-        outputs = _sample_bernoulli(subkeys[0], data=inputs)
-        timeOfLastSpike = _update_times(t, outputs, tols)
-        return outputs, timeOfLastSpike, key
+        outputs = random.bernoulli(subkeys[0], p=inputs).astype(jnp.float32)
+        tols = _update_times(t, outputs, tols)
+        return outputs, tols, key
 
     @resolver(_advance_state)
     def advance_state(self, outputs, tols, key):

ngclearn/components/input_encoders/latencyCell.py

@@ -4,6 +4,7 @@
 from ngclearn.utils.model_utils import clamp_min, clamp_max
 from jax import numpy as jnp, random, jit
 from functools import partial
+from ngcsimlib.logger import info
 
 @jit
 def _update_times(t, s, tols):
@@ -47,7 +48,7 @@ def _calc_spike_times_linear(data, tau, thr, first_spk_t, num_steps=1.,
         projected spike times
     """
     _tau = tau
-    if normalize == True:
+    if normalize:
         _tau = num_steps - 1. - first_spk_t ## linear normalization
     #torch.clamp_max((-tau * (data - 1)), -tau * (threshold - 1))
     stimes = -_tau * (data - 1.) ## calc raw latency code values
@@ -84,7 +85,7 @@ def _calc_spike_times_nonlinear(data, tau, thr, first_spk_t, eps=1e-7,
     stimes = jnp.log(_data / (_data - thr)) * tau ## calc spike times
     stimes = stimes + first_spk_t
 
-    if normalize == True:
+    if normalize:
         term1 = (stimes - first_spk_t)
         term2 = (num_steps - first_spk_t - 1.)
         term3 = jnp.max(stimes - first_spk_t)
@@ -148,13 +149,16 @@ class LatencyCell(JaxComponent):
             :Note: if this set to True, you will need to choose a useful value
                 for the "num_steps" argument (>1), depending on how many steps simulated
 
+        clip_spikes: should values under threshold be removed/suppressed?
+            (default: False)
+
         num_steps: number of discrete time steps to consider for normalized latency
             code (only useful if "normalize" is set to True) (Default: 1)
     """
 
     # Define Functions
     def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0.,
-                 linearize=False, normalize=False, num_steps=1.,
+                 linearize=False, normalize=False, clip_spikes=False, num_steps=1.,
                  batch_size=1, **kwargs):
         super().__init__(name, **kwargs)
 
@@ -163,6 +167,7 @@ def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0.,
         self.tau = tau
         self.threshold = threshold
         self.linearize = linearize
+        self.clip_spikes = clip_spikes
         ## normalize latency code s.t. final spike(s) occur w/in num_steps
         self.normalize = normalize
         self.num_steps = num_steps
@@ -175,17 +180,22 @@ def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0.,
         restVals = jnp.zeros((self.batch_size, self.n_units))
         self.inputs = Compartment(restVals, display_name="Input Stimulus") # input compartment
         self.outputs = Compartment(restVals, display_name="Spikes") # output compartment
-        self.mask = Compartment(restVals, display_name="Mask Variable")  # output compartment
+        self.mask = Compartment(restVals, display_name="Spike Time Mask")
+        self.clip_mask = Compartment(restVals, display_name="Clip Mask")
         self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike
         self.targ_sp_times = Compartment(restVals, display_name="Target Spike Time", units="ms")
         #self.reset()
 
     @staticmethod
     def _calc_spike_times(linearize, tau, threshold, first_spike_time, num_steps,
-                          normalize, inputs):
+                          normalize, clip_spikes, inputs):
         ## would call this function before processing a spike train (at start)
         data = inputs
-        if linearize == True: ## linearize spike time calculation
+        if clip_spikes:
+            clip_mask = (data < threshold) * 1. ## find values under threshold
+        else:
+            clip_mask = data * 0.  ## all values allowed to fire spikes
+        if linearize: ## linearize spike time calculation
             stimes = _calc_spike_times_linear(data, tau, threshold,
                                               first_spike_time,
                                               num_steps, normalize)
@@ -196,18 +206,20 @@ def _calc_spike_times(linearize, tau, threshold, first_spike_time, num_steps,
                                                  num_steps=num_steps,
                                                  normalize=normalize)
             targ_sp_times = stimes #* calcEvent + targ_sp_times * (1. - calcEvent)
-        return targ_sp_times
+        return targ_sp_times, clip_mask
 
     @resolver(_calc_spike_times)
-    def calc_spike_times(self, targ_sp_times):
+    def calc_spike_times(self, targ_sp_times, clip_mask):
         self.targ_sp_times.set(targ_sp_times)
+        self.clip_mask.set(clip_mask)
 
     @staticmethod
-    def _advance_state(t, dt, key, inputs, mask, targ_sp_times, tols):
+    def _advance_state(t, dt, key, inputs, mask, clip_mask, targ_sp_times, tols):
         key, *subkeys = random.split(key, 2)
         data = inputs ## get sensory pattern data / features
         spikes, spk_mask = _extract_spike(targ_sp_times, t, mask) ## get spikes at t
         tols = _update_times(t, spikes, tols)
+        spikes = spikes * (1. - clip_mask)
         return spikes, tols, spk_mask, targ_sp_times, key
 
     @resolver(_advance_state)
@@ -221,14 +233,15 @@ def advance_state(self, outputs, tols, mask, targ_sp_times, key):
     @staticmethod
     def _reset(batch_size, n_units):
         restVals = jnp.zeros((batch_size, n_units))
-        return (restVals, restVals, restVals, restVals, restVals)
+        return (restVals, restVals, restVals, restVals, restVals, restVals)
 
     @resolver(_reset)
-    def reset(self, inputs, outputs, tols, mask, targ_sp_times):
+    def reset(self, inputs, outputs, tols, mask, clip_mask, targ_sp_times):
         self.inputs.set(inputs)
         self.outputs.set(outputs)
         self.tols.set(tols)
         self.mask.set(mask)
+        self.clip_mask.set(clip_mask)
         self.targ_sp_times.set(targ_sp_times)
 
     def save(self, directory, **kwargs):