diff --git a/README.md b/README.md
index 760a4933..731a02b8 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ ngc-learn requires:
1) Python (>=3.10)
2) NumPy (>=1.22.0)
3) SciPy (>=1.7.0)
-4) ngcsimlib (>=3.0.0), (visit official page here)
+4) ngcsimlib (>=3.1.0), (visit official page here)
5) JAX (>=0.4.28) (to enable GPU use, make sure to install one of the CUDA variants)
---
-ngc-learn 3.1.0 and later require Python 3.10 or newer as well as ngcsimlib >=3.0.0.
+ngc-learn 3.2.0 and later require Python 3.10 or newer as well as ngcsimlib >=3.1.0.
ngc-learn's plotting capabilities (routines within `ngclearn.utils.viz`) require
Matplotlib (>=3.8.0) and imageio (>=2.31.5) and both plotting and density estimation
tools (routines within ``ngclearn.utils.density``) will require Scikit-learn (>=0.24.2).
@@ -72,7 +72,7 @@ Python 3.11.4 (main, MONTH DAY YEAR, TIME) [GCC XX.X.X] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import ngclearn
>>> ngclearn.__version__
-'3.1.0'
+'3.2.0'
```
Note: For access to the previous Tensorflow-2 version of ngc-learn (of
@@ -119,7 +119,7 @@ $ python install -e .
**Version:**
-3.1.1
+3.2.0
Author:
Alexander G. Ororbia II
diff --git a/docs/museum/pc_rao_ballard1999.md b/docs/museum/pc_rao_ballard1999.md
index 7bdc9952..a0f1efa9 100644
--- a/docs/museum/pc_rao_ballard1999.md
+++ b/docs/museum/pc_rao_ballard1999.md
@@ -4,7 +4,7 @@ In this exhibit, we create, simulate, and visualize the internally acquired rece
model originally proposed in (Rao & Ballard, 1999) [1].
The model code for this exhibit can be found
-[here](https://github.com/NACLab/ngc-museum/tree/main/exhibits/pc_recon).
+[here](https://github.com/NACLab/ngc-museum/tree/main/exhibits/pc_reconstruction).
## Setting Up Hierarchical Predictive Coding (HPC) with NGC-Learn
diff --git a/docs/source/ngclearn.components.input_encoders.rst b/docs/source/ngclearn.components.input_encoders.rst
index e3222214..b69b731b 100644
--- a/docs/source/ngclearn.components.input_encoders.rst
+++ b/docs/source/ngclearn.components.input_encoders.rst
@@ -4,6 +4,14 @@ ngclearn.components.input\_encoders package
Submodules
----------
+ngclearn.components.input\_encoders.TTFSCell module
+---------------------------------------------------
+
+.. automodule:: ngclearn.components.input_encoders.TTFSCell
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
ngclearn.components.input\_encoders.bernoulliCell module
--------------------------------------------------------
@@ -20,14 +28,6 @@ ngclearn.components.input\_encoders.ganglionCell module
:undoc-members:
:show-inheritance:
-ngclearn.components.input\_encoders.gridCell module
----------------------------------------------------
-
-.. automodule:: ngclearn.components.input_encoders.gridCell
- :members:
- :undoc-members:
- :show-inheritance:
-
ngclearn.components.input\_encoders.latencyCell module
------------------------------------------------------
@@ -44,14 +44,6 @@ ngclearn.components.input\_encoders.phasorCell module
:undoc-members:
:show-inheritance:
-ngclearn.components.input\_encoders.placeCell module
-----------------------------------------------------
-
-.. automodule:: ngclearn.components.input_encoders.placeCell
- :members:
- :undoc-members:
- :show-inheritance:
-
ngclearn.components.input\_encoders.poissonCell module
------------------------------------------------------
@@ -60,14 +52,6 @@ ngclearn.components.input\_encoders.poissonCell module
:undoc-members:
:show-inheritance:
-ngclearn.components.input\_encoders.populationCoderCell module
---------------------------------------------------------------
-
-.. automodule:: ngclearn.components.input_encoders.populationCoderCell
- :members:
- :undoc-members:
- :show-inheritance:
-
Module contents
---------------
diff --git a/docs/source/ngclearn.components.neurons.spiking.rst b/docs/source/ngclearn.components.neurons.spiking.rst
index 3f1e7f8e..eaf95cab 100644
--- a/docs/source/ngclearn.components.neurons.spiking.rst
+++ b/docs/source/ngclearn.components.neurons.spiking.rst
@@ -20,6 +20,14 @@ ngclearn.components.neurons.spiking.LIFCell module
:undoc-members:
:show-inheritance:
+ngclearn.components.neurons.spiking.LIFSRM module
+-------------------------------------------------
+
+.. automodule:: ngclearn.components.neurons.spiking.LIFSRM
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
ngclearn.components.neurons.spiking.RAFCell module
--------------------------------------------------
@@ -28,6 +36,14 @@ ngclearn.components.neurons.spiking.RAFCell module
:undoc-members:
:show-inheritance:
+ngclearn.components.neurons.spiking.RAFSRM module
+-------------------------------------------------
+
+.. automodule:: ngclearn.components.neurons.spiking.RAFSRM
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
ngclearn.components.neurons.spiking.WTASCell module
---------------------------------------------------
diff --git a/docs/source/ngclearn.components.synapses.competitive.rst b/docs/source/ngclearn.components.synapses.competitive.rst
index a766e3c8..1b3c0010 100644
--- a/docs/source/ngclearn.components.synapses.competitive.rst
+++ b/docs/source/ngclearn.components.synapses.competitive.rst
@@ -28,14 +28,6 @@ ngclearn.components.synapses.competitive.hopfieldSynapse module
:undoc-members:
:show-inheritance:
-ngclearn.components.synapses.competitive.simple\_ART2ASynapse module
---------------------------------------------------------------------
-
-.. automodule:: ngclearn.components.synapses.competitive.simple_ART2ASynapse
- :members:
- :undoc-members:
- :show-inheritance:
-
ngclearn.components.synapses.competitive.vectorQuantizeSynapse module
---------------------------------------------------------------------
diff --git a/docs/source/ngclearn.components.synapses.hebbian.rst b/docs/source/ngclearn.components.synapses.hebbian.rst
index 783be071..d42770a6 100644
--- a/docs/source/ngclearn.components.synapses.hebbian.rst
+++ b/docs/source/ngclearn.components.synapses.hebbian.rst
@@ -52,6 +52,14 @@ ngclearn.components.synapses.hebbian.hebbianSynapse module
:undoc-members:
:show-inheritance:
+ngclearn.components.synapses.hebbian.inhibitorySTDPSynapse module
+-----------------------------------------------------------------
+
+.. automodule:: ngclearn.components.synapses.hebbian.inhibitorySTDPSynapse
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
ngclearn.components.synapses.hebbian.traceSTDPSynapse module
------------------------------------------------------------
diff --git a/docs/source/ngclearn.components.synapses.rst b/docs/source/ngclearn.components.synapses.rst
index 087bf917..dc4eb738 100644
--- a/docs/source/ngclearn.components.synapses.rst
+++ b/docs/source/ngclearn.components.synapses.rst
@@ -64,6 +64,22 @@ ngclearn.components.synapses.mpsSynapse module
:undoc-members:
:show-inheritance:
+ngclearn.components.synapses.sparseTensorSynapse module
+-------------------------------------------------------
+
+.. automodule:: ngclearn.components.synapses.sparseTensorSynapse
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+ngclearn.components.synapses.staticSparseTensorSynapse module
+-------------------------------------------------------------
+
+.. automodule:: ngclearn.components.synapses.staticSparseTensorSynapse
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
ngclearn.components.synapses.staticSynapse module
-------------------------------------------------
diff --git a/docs/tutorials/model_basics/evolving_synapses.md b/docs/tutorials/model_basics/evolving_synapses.md
index e68ef1de..27684612 100755
--- a/docs/tutorials/model_basics/evolving_synapses.md
+++ b/docs/tutorials/model_basics/evolving_synapses.md
@@ -46,13 +46,16 @@ with Context("Circuit") as circuit:
## create and compile core simulation commands
evolve = (MethodProcess("evolve")
- >> a.evolve)
+ >> Wab.evolve)
advance = (MethodProcess("advance")
- >> a.advance_state)
+ >> a.advance_state
+ >> Wab.advance_state
+ >> b.advance_state)
reset = (MethodProcess("reset")
- >> a.reset)
+ >> a.reset
+ >> b.reset)
## set up non-compiled utility commands
def clamp(x):
@@ -67,7 +70,7 @@ and evolve the synapse every time step like so:
x_seq = jnp.asarray([[1, 1, 0, 0, 1]], dtype=jnp.float32)
reset.run()
-print("{}: Wab = {}".format(-1, Wab.weights.value))
+print("{}: Wab = {}".format(-1, Wab.weights.get()))
for ts in range(x_seq.shape[1]):
x_t = jnp.expand_dims(x_seq[0,ts], axis=0) ## get data at time t
clamp(x_t)
diff --git a/history.txt b/history.txt
index 8b42a782..13364b41 100644
--- a/history.txt
+++ b/history.txt
@@ -103,10 +103,15 @@ History
3.1.1
— — — — — — — — -
- * Several patches and fixes (ensured certain error-neurons behaved as in v3.0.0)
- * Update to docs/revisions
- * Integration of filters/filtering functionality (including simple divisive normalization) in utilities
- * Integration of several useful metrics (including for characterizing sparsity)
- * Some additional clean-up throughout
-
+ * several patches and fixes (ensured certain error-neurons behaved as in v3.0.0)
+ * ipdate to docs/revisions
+ * integration of filters/filtering functionality (including simple divisive normalization) in utilities
+ * integration of several useful metrics (including for characterizing sparsity)
+ * some additional clean-up throughout
+ 3.2.0
+ — — — — — — — — -
+ * various fixes/patches as well as clean-up/minor tweaks to utils/core components
+ * integration of additional visualization tools
+ * integration of sparse-tensor synaptic cable (locally-connected/unshared-convolutional structure)
+ * additional component integration/revisions, including updates to patched-synaptic cable components
diff --git a/ngclearn/components/__init__.py b/ngclearn/components/__init__.py
index 94a2c116..4e32c6c4 100644
--- a/ngclearn/components/__init__.py
+++ b/ngclearn/components/__init__.py
@@ -19,6 +19,9 @@
from .neurons.spiking.izhikevichCell import IzhikevichCell
from .neurons.spiking.hodgkinHuxleyCell import HodgkinHuxleyCell
from .neurons.spiking.RAFCell import RAFCell
+## point to standard spiking spike-response model component types
+from .neurons.spiking.LIFSRM import LIFSRM
+from .neurons.spiking.RAFSRM import RAFSRM
## point to transformer/operator component types
from .other.varTrace import VarTrace
@@ -30,6 +33,7 @@
from .input_encoders.ganglionCell import RetinalGanglionCell
from .input_encoders.latencyCell import LatencyCell
from .input_encoders.phasorCell import PhasorCell
+from .input_encoders.TTFSCell import TTFSCell
#from .input_encoders.populationCoderCell import PopulationCoderCell
#from .input_encoders.gridCell import GridCell
#from .input_encoders.placeCell import PlaceCell
@@ -37,12 +41,15 @@
## point to synapse component types
from .synapses.denseSynapse import DenseSynapse
from .synapses.staticSynapse import StaticSynapse
+from .synapses.sparseTensorSynapse import SparseTensorSynapse
+from .synapses.staticSparseTensorSynapse import StaticSparseTensorSynapse
from .synapses.hebbian.hebbianSynapse import HebbianSynapse
from .synapses.hebbian.STDPSynapse import STDPSynapse
from .synapses.hebbian.traceSTDPSynapse import TraceSTDPSynapse
from .synapses.hebbian.expSTDPSynapse import ExpSTDPSynapse
from .synapses.hebbian.eventSTDPSynapse import EventSTDPSynapse
from .synapses.hebbian.BCMSynapse import BCMSynapse
+from .synapses.hebbian.inhibitorySTDPSynapse import InhibitorySTDPSynapse
from .synapses.competitive.SOMSynapse import SOMSynapse
from .synapses.STPDenseSynapse import STPDenseSynapse
from .synapses.exponentialSynapse import ExponentialSynapse
diff --git a/ngclearn/components/input_encoders/TTFSCell.py b/ngclearn/components/input_encoders/TTFSCell.py
new file mode 100644
index 00000000..63265f5f
--- /dev/null
+++ b/ngclearn/components/input_encoders/TTFSCell.py
@@ -0,0 +1,161 @@
+import jax.numpy as jnp
+from ngclearn.components.jaxComponent import JaxComponent
+from ngclearn import compilable # from ngcsimlib.parser import compilable
+from ngclearn import Compartment # from ngcsimlib.compartment import Compartment
+
+
+class TTFSCell(JaxComponent): ## time-to-first-spike (en)coding cell
+ """
+ A time-to-first-spike (TTFS) iterative encoder component.
+ This input encoder converts a real-valued batch input vectors into sparse temporal spike trains, using
+ either exponential or logarithmic latency mapping schemes. Units within this encoder group will only fire once
+ within a spike train (before reset/clearing).
+
+ | --- Cell Input Compartments: ---
+ | inputs - input (takes in external signals)
+ | --- Cell State Compartments: ---
+ | key - JAX PRNG key
+ | has_fired - boolean vector denoting which units have fired thus far
+ | target_spike_times - tracks targeted spike times
+ | refractory_counter - refactory variable to ensure units only fire once
+ | --- Cell Output Compartments: ---
+ | outputs - output (binary spike train matrix output at the current time step t)
+ | tols - time-of-last-spike
+
+ Args:
+ name: the string name of this cell
+
+ n_units: number of cellular entities (neural population size)
+
+ batch_size: batch size dimension of this cell (Default: 1)
+
+ num_steps: number of total time steps of simulation to consider
+
+ tau_ref: ttfs time constant
+
+ latency_mode: which type of latency mapping function to apply (Default: "exponential");
+ "exponential" triggers exponential-latency mapping while "logarithmic" triggers logarithmic-latency mapping
+ """
+
+ def __init__(
+ self,
+ name,
+ n_units,
+ batch_size=1,
+ num_steps=100,
+ tau_ref=5,
+ latency_mode="exponential",
+ **kwargs
+ ):
+ super().__init__(name, **kwargs)
+
+ ## ttfs meta-parameters
+ self.n_units = n_units
+ self.batch_size = batch_size
+ self.num_steps = num_steps
+ self.tau_ref = tau_ref
+ self.latency_mode = latency_mode
+
+ ## verify valid latency setting selected (only two supported!)
+ if self.latency_mode not in ["exponential", "logarithmic"]:
+ raise ValueError("latency_mode must be either 'exponential' or 'logarithmic'")
+
+ ## ttfs structural compartments
+ restVals = jnp.zeros((self.batch_size, self.n_units))
+ self.inputs = Compartment(restVals)
+ self.outputs = Compartment(restVals)
+ self.has_fired = Compartment(jnp.zeros((self.batch_size, self.n_units), dtype=jnp.bool_))
+ self.target_spike_times = Compartment(jnp.zeros((self.batch_size, self.n_units), dtype=jnp.int32))
+ self.refractory_counter = Compartment(jnp.zeros((self.batch_size, self.n_units), dtype=jnp.int32))
+ self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike
+
+ @staticmethod
+ def _compute_latency_map(x, num_steps, latency_mode):
+ ## NOTE: co-routine to map input range [0, 1] to discrete target time-step w/in [0, num_steps-1]
+ x_clipped = jnp.clip(x, 1e-7, 1.0) ## clip x, to ensure inputs do not have NaNs/Inf values
+ if latency_mode == "exponential": ## exponential-latency mapping
+ ## high-intensity maps to 0 delay, low intensity maps to max delay
+ raw_latency = jnp.exp(-x_clipped)
+ ## normalize mapping range to scale from 0 to 1:
+ ### x=1.0 -> raw=exp(-1) -> norm=0.0; x=0.0 -> raw=exp(0) -> norm=1.0
+ norm_latency = (raw_latency - jnp.exp(-1.0)) / (jnp.exp(-1e-7) - jnp.exp(-1.0))
+ ## multiply by norm_latency directly s.t. x=1.0 yields 0 delay and x=0.0 yields maximal delay
+ target_times = norm_latency * (num_steps - 1)
+ else: ## logarithmic-latency mapping
+ ## find minimum expected non-zero input to scale the log dynamic range properly instead of using a
+ ## fixed 1e-7 clip floor, normalize based on the lowest value in x (or map the chosen clipping floor)
+ raw_latency = -jnp.log(x_clipped)
+ ## mapping range: x=1.0 -> raw=0.0 (Step 0); x=1e-7 -> raw=max_log (Max Step)
+ #max_log = -jnp.log(1e-7)
+ ## NOTE: if we want "provided" units/elements to spread across (spike train) window,
+ ## must scale relative to maximum log latency present in input
+ max_input_log = jnp.maximum(-jnp.log(0.01), jnp.max(raw_latency))
+ target_times = (raw_latency / max_input_log) * (num_steps - 1)
+ ## ensure calculation bounds are cleanly w/in step window
+ target_times = jnp.clip(target_times, 0, num_steps - 1)
+ return jnp.round(target_times).astype(jnp.int32) ## output target times
+
+ @compilable
+ def advance_state(self, t, dt):
+ ## NOTE: advances simulation state by resolving continuous time into a discrete step index;
+ ## computes latencies at step 0, then emits spikes based on timing targets + refractory criteria
+ x = self.inputs.get()
+ ## translate continuous time (t) + integration time-constant (dt) into discrete step indices
+ current_step = jnp.round(t / dt).astype(jnp.int32) ## cast to int32
+ ## calculate target spike schedules only on the first step of a sequence
+ is_step_zero = (current_step == 0)
+
+ computed_targets = TTFSCell._compute_latency_map(x, self.num_steps, self.latency_mode)
+ init_has_fired = jnp.zeros_like(self.has_fired.get())
+ init_refractory = jnp.zeros_like(self.refractory_counter.get())
+
+ ## conditionally update targets only if resetting/restarting at step zero
+ target_spike_times_curr = jnp.where(is_step_zero, computed_targets, self.target_spike_times.get())
+ has_fired_curr = jnp.where(is_step_zero, init_has_fired, self.has_fired.get())
+ refractory_counter_curr = jnp.where(is_step_zero, init_refractory, self.refractory_counter.get())
+
+ refractory_counter_next = jnp.maximum(0, refractory_counter_curr - 1) ## decrement refractory counters
+ ## determine units that are clear to fire (not currently in refractory window)
+ is_not_refractory = (refractory_counter_next == 0)
+ reached_target_time = (target_spike_times_curr == current_step) ## check which units hit target spike times
+ ## spike conditions: reached time, not refractory and has not fired yet
+ spikes = reached_target_time & is_not_refractory & (~has_fired_curr)
+ ## update state records
+ has_fired_next = jnp.where(spikes, True, has_fired_curr)
+ refractory_counter_next = jnp.where(spikes, self.tau_ref, refractory_counter_next)
+
+ ## update internal compartments
+ self.outputs.set(spikes.astype(jnp.float32))
+ self.target_spike_times.set(target_spike_times_curr)
+ self.has_fired.set(has_fired_next)
+ self.refractory_counter.set(refractory_counter_next)
+ self.tols.set((1. - spikes) * self.tols.get() + (spikes * t)) ## track time-of-last-spike
+
+ @compilable
+ def reset(self): ## resets internal clock, tracking spikes, refractory states
+ restVals = jnp.zeros((self.batch_size, self.n_units))
+ self.inputs.set(restVals)
+ self.outputs.set(restVals)
+ self.tols.set(restVals)
+ ## clear internal state variables
+ self.target_spike_times.set(jnp.zeros((self.batch_size, self.n_units), dtype=jnp.int32))
+ self.has_fired.set(jnp.zeros((self.batch_size, self.n_units), dtype=jnp.bool_))
+ self.refractory_counter.set(jnp.zeros((self.batch_size, self.n_units), dtype=jnp.int32))
+
+ @classmethod
+ def help(cls):
+ properties = {
+ "cell_type": "TTFSCell - Converts static vectors into time-to-first-spike encodings"
+ }
+ compartment_props = {
+ "inputs": {"inputs": "Takes in external real-valued input vectors to encode"},
+ "outputs": {"outputs": "Emits a binary spike tensor (0 or 1) at time step t"}
+ }
+ hyperparams = {
+ "n_units": "Number of neuronal units/features to model",
+ "batch_size": "Batch size dimension of this component",
+ "num_steps": "Total number of discrete time steps in the simulation window",
+ "tau_ref": "Refractory period duration (in time steps) after firing a spike",
+ "latency_mode": "Mathematical mapping method: 'exponential' or 'logarithmic'"
+ }
+ return {"properties": properties, "compartments": compartment_props, "hyperparameters": hyperparams}
diff --git a/ngclearn/components/input_encoders/__init__.py b/ngclearn/components/input_encoders/__init__.py
index bbee5280..2e4d262b 100644
--- a/ngclearn/components/input_encoders/__init__.py
+++ b/ngclearn/components/input_encoders/__init__.py
@@ -3,6 +3,7 @@
from .latencyCell import LatencyCell
from .phasorCell import PhasorCell
from .ganglionCell import RetinalGanglionCell
+from .TTFSCell import TTFSCell ## time-to-first-spike (en)coding cell
#from .populationCoderCell import PopulationCoderCell
#from .gridCell import GridCell
#from .placeCell import PlaceCell
diff --git a/ngclearn/components/neurons/__init__.py b/ngclearn/components/neurons/__init__.py
index 564577cd..b46490e1 100644
--- a/ngclearn/components/neurons/__init__.py
+++ b/ngclearn/components/neurons/__init__.py
@@ -16,4 +16,7 @@
from .spiking.izhikevichCell import IzhikevichCell
from .spiking.hodgkinHuxleyCell import HodgkinHuxleyCell
from .spiking.RAFCell import RAFCell
+## point to spike-response models (SRMs)
+from .spiking.LIFSRM import LIFSRM
+from .spiking.RAFSRM import RAFSRM
diff --git a/ngclearn/components/neurons/graded/rateCell.py b/ngclearn/components/neurons/graded/rateCell.py
index 3dcbdcc8..ea1e9dbe 100755
--- a/ngclearn/components/neurons/graded/rateCell.py
+++ b/ngclearn/components/neurons/graded/rateCell.py
@@ -2,13 +2,15 @@
from jax import numpy as jnp, random, jit
-from ngclearn import compilable #from ngcsimlib.parser import compilable
-from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
+from ngclearn import compilable
+from ngclearn import Compartment
from ngclearn.components.jaxComponent import JaxComponent
from ngclearn.utils.model_utils import create_function, threshold_soft, \
threshold_cauchy
from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \
step_euler, step_rk2, step_rk4
+from ngclearn.components.synapses.patched import HebbianPatchedSynapse, StaticPatchedSynapse
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
from ngcsimlib.logger import info
@@ -111,8 +113,9 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
The specific differential equation that characterizes this cell
is (for adjusting v, given current j, over time) is:
- | tau_m * dz/dt = lambda * prior(z) + (j + j_td)
+ | tau_m * dz/dt = lambda * prior(z) + (j + j_td + j_lat(if use_lateral==True))
| where j is the set of general incoming input signals (e.g., message-passed signals)
+ | j_lat (optional) lateral pressure via lateral synapse
| and j_td is taken to be the set of top-down pressure signals
| --- Cell Input Compartments: ---
@@ -147,6 +150,10 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
output_scale: factor to multiply output of nonlinearity of this cell by (Default: 1.)
+ threshold: the kind of iterative thresholding function to place over neuronal dynamics
+
+ batch_size: the size of each mini batch
+
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")
@@ -156,12 +163,46 @@ class RateCell(JaxComponent): ## Rate-coded/real-valued cell
at an increase in computational cost (and simulation time)
resist_scale: a scaling factor applied to incoming pressure `j` (default: 1)
+
+ use_lateral: determines wheter the RateCell uses lateral pressure during inference (Default: False)
+
+ n_patch: the number of local patches in each layer (equivalent to n_sub_models in PatchedSynapse) (Default: 1)
+
+ adaptive_lateral: if True, lateral synaptic is adaptive using HebbianPatchedSynapse component,
+ if False, lateral synaptic is a static using StaticPatchedSynapse component
+ (Default: False)
+
+ Wl_eta: learning rate for the lateral synapse update (Default: 0.001)
+
+ Wl_opt_type: optimizer used by the lateral synapse (Default: "sgd")
+
+ exc_inh: lateral synaptic initialization values: (Default: (0., 0.))
+
+ fx_lat: string name of activation function/nonlinearity to use for lateral pressure (Default: same as act_fx)
+
+ Wl_prior: a kernel to drive prior over lateral synaptic cable's values (Default: (None, 0.))
+
+ Wl_bound: maximum value to enforce over newly computed lateral synaptic values (Default: 0.)
+
"""
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):
- jax_comp_kwargs = {k: v for k, v in kwargs.items() if k not in ('omega_0',)}
+ integration_type="euler",
+ batch_size=1,
+ resist_scale=1.,
+ shape=None,
+ is_stateful=True,
+ use_lateral = False,
+ adaptive_lateral = False,
+ exc_inh=(0., 0.),
+ n_patch=1,
+ fx_lat=None,
+ Wl_eta=0.,
+ Wl_opt_type="sgd",
+ Wl_prior=(None, 0.),
+ **kwargs):
+ jax_comp_kwargs = {k: v for k, v in kwargs.items() if k in ('key',)}
this_class_kwargs = {k: v for k, v in kwargs.items() if k in ('omega_0',)}
super().__init__(name, **jax_comp_kwargs)
@@ -186,6 +227,10 @@ def __init__(
self.thr_lmbda = thr_lmbda ## scale to drive thresholding dynamics
self.resist_scale = resist_scale ## a "resistance" scaling factor
+ self.use_lateral = use_lateral
+ self.adaptive_lateral = adaptive_lateral
+ self.Wl_eta = Wl_eta
+
## integration properties
self.integrationType = integration_type
self.intgFlag = get_integrator_code(self.integrationType)
@@ -198,19 +243,65 @@ def __init__(
_shape = (batch_size, shape[0], shape[1], shape[2]) ## shape is 4D tensor
self.shape = shape
self.n_units = n_units
+ self.n_patch = n_patch
+ self.dim_patch = n_units // n_patch
self.batch_size = batch_size
omega_0 = None
if act_fx == "sine":
omega_0 = this_class_kwargs["omega_0"]
self.fx, self.dfx = create_function(fun_name=act_fx, args=omega_0)
-
- # compartments (state of the cell & parameters will be updated through stateless calls)
+ if use_lateral:
+ if fx_lat is None:
+ self.fx_lat, self.dfx_lat = self.fx, self.dfx
+ else:
+ self.fx_lat, self.dfx_lat = create_function(fun_name=fx_lat)
+
+ ## ══════════════════ lateral synapse ══════════════════
+ self.Wl = None
+ if use_lateral:
+ exc, inh = exc_inh
+ d = n_units // n_patch
+
+ if adaptive_lateral:
+ self.Wl = HebbianPatchedSynapse(f"Wl_{name}", shape=(n_units, n_units),
+ n_sub_models=n_patch,
+ weight_init=dist.diag_offdiag(on_diag=exc, off_diag=inh),
+ eta=Wl_eta,
+ prior=Wl_prior,
+ sign_value=1.,
+ optim_type=Wl_opt_type,
+ batch_size=batch_size,
+ )
+ else:
+ self.Wl = StaticPatchedSynapse(f"Wl_{name}", shape=(n_units, n_units),
+ n_sub_models=n_patch,
+ weight_init=dist.diag_offdiag(on_diag=exc, off_diag=inh),
+ batch_size=batch_size,
+ )
+
+ ## ════════════════════════════════════════════════════════════════════════════════════
+ ## compartments (state of the cell & parameters will be updated through stateless calls)
restVals = jnp.zeros(_shape)
+
self.j = Compartment(restVals, display_name="Input Stimulus Current", units="mA") # electrical current
self.zF = Compartment(restVals, display_name="Transformed Rate Activity") # rate-coded output - activity
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
+ ## lateral compartments
+ if use_lateral:
+ self.j_lat = Compartment(restVals, display_name="Lateral Stimulus Current", units="mA") # lateral current
+ self.zF_lat = Compartment(restVals, display_name="Lateral Rate Activity") # fx_lat(z)
+
+ def compile(self):
+ if self.use_lateral:
+ self.zF_lat >> self.Wl.inputs
+ self.Wl.outputs >> self.j_lat
+
+ if self.adaptive_lateral:
+ self.zF_lat >> self.Wl.pre
+ self.zF_lat >> self.Wl.post
+ super().compile()
@compilable
def advance_state(self, dt):
@@ -219,74 +310,104 @@ def advance_state(self, dt):
j_td = self.j_td.get()
z = self.z.get()
- #if tau_m > 0.:
if self.is_stateful:
+ ## ════════════════════════════════════════════════════════════════
### run a step of integration over neuronal dynamics
- ## Notes:
+ ## Note:
## self.pressure <-- "top-down" expectation / contextual pressure
+ ## ══════════════════════════
## self.current <-- "bottom-up" data-dependent signal
dfx_val = self.dfx(z)
j = _modulate(j, dfx_val)
j = j * self.resist_scale
- tmp_z = _run_cell(
- dt, j, j_td, z, self.tau_m, leak_gamma=self.priorLeakRate, integType=self.intgFlag,
- priorType=self.priorType
- )
+
+ ## ══════════════════════════
+ ## Calculate lateral pressure
+ if self.use_lateral:
+ self.Wl.advance_state()
+
+ j_lat = self.j_lat.get()
+ dfx_lat = self.dfx_lat(z)
+ j_lat = _modulate(j_lat, dfx_lat)
+ j_lat = j_lat * self.resist_scale
+ else:
+ j_lat = jnp.zeros_like(z)
+
+ ## ═════════════════════════════════════════════════
+ ## Run a single integration step for a stateful cell
+ tmp_z = _run_cell(dt,
+ j=(j + j_lat),
+ j_td=j_td,
+ z=z,
+ tau_m=self.tau_m,
+ leak_gamma=self.priorLeakRate,
+ integType=self.intgFlag,
+ priorType=self.priorType
+ )
+
## apply optional thresholding sub-dynamics
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 = self.fx(z) * self.output_scale ## post-activation function value(s)
+ z = tmp_z ## pre-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 = self.fx(z) * self.output_scale
+ z = _run_cell_stateless(j + j_td) ## (j+j_td) >> z
+
+ ## ═══════════════════════════════════════════════
+ ## fx(z) >> zF post-activation function value(s)
+ zF = self.fx(z) * self.output_scale
+ if self.use_lateral:
+ zF_lat = self.fx_lat(z) * self.output_scale
- # Update compartments
+ ## ════════════════════
+ ## Update compartments
self.j.set(j)
self.j_td.set(j_td)
self.z.set(z)
self.zF.set(zF)
+ if self.use_lateral:
+ self.j_lat.set(j_lat)
+ self.zF_lat.set(zF_lat)
+ ## ═══════════════════════════════
+ ## Update lateral synaptic weights
@compilable
- def reset(self): #, batch_size, shape): #n_units
+ def evolve(self):
+ if self.use_lateral and self.adaptive_lateral:
+ self.Wl.evolve()
+
+
+ @compilable
+ def reset(self):
_shape = (self.batch_size, self.shape[0])
if len(self.shape) > 1:
_shape = (self.batch_size, self.shape[0], self.shape[1], self.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)
- # else jnp.ones([[self.tau_m]]))
- # priorLeakRate = (self.priorLeakRate if isinstance(self.priorLeakRate, float)
- # else jnp.ones([[self.priorLeakRate]]))
- # resist_scale = (self.resist_scale if isinstance(self.resist_scale, float)
- # else jnp.ones([[self.resist_scale]]))
- #
- # file_name = directory + "/" + self.name + ".npz"
- # jnp.savez(file_name,
- # tau_m=tau_m, priorLeakRate=priorLeakRate,
- # resist_scale=resist_scale) #, key=self.key.value)
- #
- # def load(self, directory, seeded=False, **kwargs):
- # file_name = directory + "/" + self.name + ".npz"
- # data = jnp.load(file_name)
- # ## constants loaded in
- # self.tau_m = data['tau_m']
- # self.priorLeakRate = data['priorLeakRate']
- # self.resist_scale = data['resist_scale']
- # #if seeded:
- # # self.key.set(data['key'])
-
- @classmethod
+ if self.use_lateral:
+ self.Wl.reset()
+ self.zF_lat.set(restVals)
+ self.j_lat.set(restVals)
+
+
+ def save(self, directory):
+ super().save(directory)
+ if self.use_lateral:
+ self.Wl.save(directory)
+
+ def load(self, directory):
+ super().load(directory)
+ if self.use_lateral:
+ self.Wl.load(directory)
+
def help(cls): ## component help function
properties = {
"cell_type": "RateCell - evolves neurons according to rate-coded/"
@@ -300,7 +421,9 @@ def help(cls): ## component help function
"states":
{"z": "Update to rate-coded continuous dynamics; value at time t"},
"outputs":
- {"zF": "Nonlinearity/function applied to rate-coded dynamics; f(z)"},
+ {"zF": "Nonlinearity/function applied to rate-coded dynamics; f(z)",
+ "zF_lat": "(if use_lateral is True) zF for lateral fx_lat(z)",
+ "j_lat": "(if use_lateral is True) lateral pressure"},
}
hyperparams = {
"n_units": "Number of neuronal cells to model in this layer",
@@ -310,15 +433,27 @@ def help(cls): ## component help function
"act_fx": "Elementwise activation function to apply over cell state `z`",
"threshold": "What kind of iterative thresholding function to place over neuronal dynamics?",
"integration_type": "Type of numerical integration to use for the cell dynamics",
+ "use_lateral": "If lateral synaptic is active for inference",
+ "n_patch": "Number of block-diagonal local lateral synaptic patch (n_sub_models of `Wl`)",
+ "adaptive_lateral": "Adaptive lateral synapses (HebbianPatchedSynapse) or static lateral synaptic (StaticPatchedSynapse)",
+ "Wl_eta": "Learning rate for the lateral synapses",
+ "exc_inh": "(exc, inh) lateral synapses initial values. On-diagonal takes exc>0 and off-diagonal takes inh<0",
+ "fx_lat": "Activation applied to z for the lateral pressure (defaults to act_fx)",
}
info = {cls.__name__: properties,
"compartments": compartment_props,
- "dynamics": "tau_m * dz/dt = Prior(z; gamma) + (j + j_td)",
+ "dynamics": "tau_m * dz/dt = Prior(z; gamma) + (j + j_lat + j_td)",
"hyperparameters": hyperparams}
return info
+
+
if __name__ == '__main__':
from ngcsimlib.context import Context
+ from ngclearn import MethodProcess
+
with Context("Bar") as bar:
- X = RateCell("X", 9, 0.03)
- print(X)
+ X = RateCell("X", 9, 20., act_fx="relu", use_lateral=True, adaptive_lateral=True,
+ exc_inh=(+5., -5.), n_patch=3, Wl_eta=0.01)
+ advance = (MethodProcess(name="advance") >> X.advance_state)
+ evolve = (MethodProcess(name="evolve") >> X.evolve)
diff --git a/ngclearn/components/neurons/spiking/LIFCell.py b/ngclearn/components/neurons/spiking/LIFCell.py
index 6fedf559..865287e9 100644
--- a/ngclearn/components/neurons/spiking/LIFCell.py
+++ b/ngclearn/components/neurons/spiking/LIFCell.py
@@ -5,30 +5,9 @@
from ngclearn.utils.surrogate_fx import (secant_lif_estimator, arctan_estimator,
triangular_estimator,
straight_through_estimator)
-
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
-def _dfv(t, v, params): ## voltage dynamics wrapper
- j, rfr, tau_m, refract_T, v_rest, g_L = params
- mask = (rfr >= refract_T) * 1. # get refractory mask
- ## update voltage / membrane potential
- dv_dt = (v_rest - v) * g_L + (j * mask)
- dv_dt = dv_dt * (1. / tau_m)
- return dv_dt
-
-
-#@partial(jit, static_argnums=[3, 4])
-def _update_theta(dt, v_theta, s, tau_theta, theta_plus: Array=0.05):
- ### Runs homeostatic threshold update dynamics one step (via Euler integration).
- #theta_decay = 0.9999999 #0.999999762 #jnp.exp(-dt/1e7)
- #theta_plus = 0.05
- #_V_theta = V_theta * theta_decay + S * theta_plus
- theta_decay = jnp.exp(-dt/tau_theta)
- _v_theta = v_theta * theta_decay + s * theta_plus
- #_V_theta = V_theta + -V_theta * (dt/tau_theta) + S * alpha
- return _v_theta
-
class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
"""
@@ -47,6 +26,7 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
| v - membrane potential/voltage state
| rfr - (relative) refractory variable state
| thr_theta - homeostatic/adaptive threshold increment state
+ | surrogate - current estimated derivative of spikes w.r.t. j
| key - JAX PRNG key
| --- Cell Output Compartments: ---
| s - emitted binary spikes/action potentials
@@ -108,9 +88,24 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell
""" ## batch_size arg?
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",
- v_min=None, max_one_spike=False, key=None
+ 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",
+ v_min=None,
+ max_one_spike=False,
+ key=None
):
super().__init__(name, key)
@@ -141,15 +136,9 @@ def __init__(
self.batch_size = 1
self.n_units = n_units
- # ## set up surrogate function for spike emission
- # if surrogate_type == "secant_lif":
- # spike_fx, d_spike_fx = secant_lif_estimator()
- # elif surrogate_type == "arctan":
- # spike_fx, d_spike_fx = arctan_estimator()
+ self.spike_fx, self.d_spike_fx = arctan_estimator() ## arctan
# elif surrogate_type == "triangular":
# spike_fx, d_spike_fx = triangular_estimator()
- # else: ## default: straight_through
- # spike_fx, d_spike_fx = straight_through_estimator()
## Compartment setup
restVals = jnp.zeros((self.batch_size, self.n_units))
@@ -160,31 +149,58 @@ def __init__(
self.rfr = Compartment(restVals + self.refract_T, display_name="Refractory Time Period", units="ms")
self.thr_theta = Compartment(restVals, display_name="Threshold Adaptive Shift", units="mV")
self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") ## time-of-last-spike
- # self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value")
+ self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value")
+ self.v_thr = Compartment(restVals + self.thr)
+
+ @staticmethod
+ def _dfv(t, v, params): ## voltage dynamics wrapper
+ j, rfr, tau_m, refract_T, v_rest, g_L = params
+ mask = (rfr >= refract_T) * 1. ## get refractory mask
+ ## update voltage / membrane potential
+ dv_dt = (v_rest - v) * g_L + (j * mask)
+ dv_dt = dv_dt * (1. / tau_m)
+ return dv_dt
+
+ #@partial(jit, static_argnums=[3, 4])
+ @staticmethod
+ def _update_theta(dt, v_theta, s, tau_theta, theta_plus: Array=0.05):
+ ### Runs homeostatic threshold update dynamics one step (via Euler integration).
+ #theta_decay = 0.9999999 #0.999999762 #jnp.exp(-dt/1e7)
+ #theta_plus = 0.05
+ #_V_theta = V_theta * theta_decay + S * theta_plus
+ theta_decay = jnp.exp(-dt/tau_theta)
+ _v_theta = v_theta * theta_decay + s * theta_plus
+ #_V_theta = V_theta + -V_theta * (dt/tau_theta) + S * alpha
+ return _v_theta
@compilable
def advance_state(self, dt, t):
- j = self.j.get() * self.resist_m
+ j = self.j.get() * self.resist_m ## get current electrical current input
_v_thr = self.thr_theta.get() + self.thr ## calc present voltage threshold
+ ## perform step of ODE integration
v_params = (j, self.rfr.get(), self.tau_m.get(), self.refract_T, self.v_rest, self.g_L)
-
- if self.intgFlag == 1:
- _, _v = step_rk2(0., self.v.get(), _dfv, dt, v_params)
- else:
- _, _v = step_euler(0., self.v.get(), _dfv, dt, v_params)
-
+ if self.intgFlag == 1: ## midpoint method
+ _, _v = step_rk2(0., self.v.get(), LIFCell._dfv, dt, v_params)
+ else: ## take forward Euler step
+ _, _v = step_euler(0., self.v.get(), LIFCell._dfv, dt, v_params)
+ ## calculate spike emission and post-spike voltage-reset mechanism
s = (_v > _v_thr) * 1.
_rfr = (self.rfr.get() + dt) * (1. - s)
+ ## calculate a surrogate estimate of ds/dj
+ ds_dv = self.d_spike_fx(_v, _v_thr, alpha=0.2) ## deriv of spike w.r.t. v/j
+ ds_dj = (self.resist_m * dt) / self.tau_m.get()
+ self.surrogate.set(ds_dv * ds_dj)
+ ## set after-potential values
_v = _v * (1. - s) + s * self.v_reset
- raw_s = s
+ raw_s = s ## "raw" spikes
if self.one_spike and not self.max_one_spike:
key, skey = random.split(self.key.get(), 2)
-
- m_switch = (jnp.sum(s) > 0.).astype(jnp.float32) ## TODO: not batch-able
+ #m_switch = (jnp.sum(s) > 0.).astype(jnp.float32) ## TODO: this line is not batch-able
+ m_switch = (jnp.sum(s, axis=1, keepdims=True) > 0.).astype(jnp.float32)
rS = s * random.uniform(skey, s.shape)
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
@@ -196,8 +212,9 @@ def advance_state(self, dt, t):
if self.tau_theta > 0.:
## run one integration step for threshold dynamics
- thr_theta = _update_theta(dt, self.thr_theta.get(), raw_s, self.tau_theta, self.theta_plus) #.get())
+ thr_theta = LIFCell._update_theta(dt, self.thr_theta.get(), raw_s, self.tau_theta, self.theta_plus) #.get())
self.thr_theta.set(thr_theta)
+ self.v_thr.set(self.thr + thr_theta)
## update time-of-last spike variable(s)
self.tols.set((1. - s) * self.tols.get() + (s * t))
@@ -205,7 +222,7 @@ def advance_state(self, dt, t):
if self.v_min is not None: ## ensures voltage never < v_rest
_v = jnp.maximum(_v, self.v_min)
-
+ ## update internal compartment values
self.v.set(_v)
self.s.set(s)
self.s_raw.set(raw_s)
@@ -220,8 +237,10 @@ def reset(self):
self.v.set(restVals + self.v_rest)
self.s.set(restVals)
self.s_raw.set(restVals)
+ self.surrogate.set(restVals + 1.)
self.rfr.set(restVals + self.refract_T)
self.tols.set(restVals)
+ self.v_thr.set(restVals + self.thr)
@classmethod
def help(cls): ## component help function
diff --git a/ngclearn/components/neurons/spiking/LIFSRM.py b/ngclearn/components/neurons/spiking/LIFSRM.py
new file mode 100644
index 00000000..58c8e81a
--- /dev/null
+++ b/ngclearn/components/neurons/spiking/LIFSRM.py
@@ -0,0 +1,177 @@
+from ngclearn.components.jaxComponent import JaxComponent
+from jax import numpy as jnp, jit
+from ngclearn import compilable, Compartment
+
+
+class LIFSRM(JaxComponent): ## LIF spike-response model (LIF-SRM)
+ """
+ The leaky integrate-and-fire (LIF) spike-response model (SRM); this SRM computes
+ dynamics of LIF units analytically.
+
+ | --- Cell Input Compartments: ---
+ | current_j - electrical current input (takes in external signals)
+ | --- Cell State Compartments: ---
+ | v - membrane potential/voltage state
+ | j_lowpass - internal low-pass-filtered current (maintained by this SRM)
+ | key - JAX PRNG key
+ | --- Cell Output Compartments: ---
+ | s - emitted binary spikes/action potentials
+ | t_last_spike - time-of-last-spike (output)
+ | last_t_eval - time of last (SRM) evaluation
+
+ | References:
+ | Gerstner, W., 1995. Time structure of the activity in neural network
+ | models. Physical review E, 51(1), p.738.
+
+ | Pedagogical Reference:
+ | http://www.scholarpedia.org/article/Spike-response_model
+
+ Args:
+ name: the string name of this cell
+
+ n_units: number of cellular entities (neural population size)
+
+ tau_m: membrane time constant (ms)
+
+ thr: base value for adaptive thresholds that govern short-term
+ plasticity (in milliVolts, or mV; default: -52. mV)
+
+ v_rest: reversal potential or membrane resting potential (in mV; default: -65 mV)
+
+ v_reset: membrane reset potential (in mV) -- upon occurrence of a spike,
+ a neuronal cell's membrane potential will be set to this value;
+ (default: -60 mV)
+ """
+
+ def __init__(
+ self,
+ name,
+ n_units,
+ tau_m, ## membrane time constant (ms)
+ thr=-52., ## threshold (mV)
+ v_rest=-65., ## membrane resting potential (mV)
+ v_reset=-60., ## membrne reset potential (mV)
+ batch_size=1,
+ **kwargs
+ ):
+ super().__init__(name, **kwargs)
+ ## LIF-SRM meta-parameters
+ self.n_units = n_units
+ self.tau_m = tau_m ## membrane time-constant
+ self.thr = thr ## threshold (mV)
+ self.v_rest = v_rest ## resting potential (mV)
+ self.v_reset = v_reset ## reset potential (mV)
+ self.batch_size = batch_size
+
+ ## LIF-SRM key compartments
+ self.current_j = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+ self.v = Compartment(jnp.full((self.batch_size, self.n_units), self.v_rest))
+ self.s = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+
+ ## analytical SRM state compartments/variables (NOTE: designed to avoid maintaining explicit spike history tensors)
+ self.t_last_spike = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.next_spike_t = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.j_lowpass = Compartment(jnp.zeros((self.batch_size, self.n_units))) ## integrated input trace
+ self.last_t_eval = Compartment(jnp.zeros((self.batch_size, self.n_units))) ## tracks clock index (when evaluated)
+
+ @compilable
+ def advance_state(self, dt, t):
+ ## pass last evaluation clock marker into kernel co-routine (to handle time jumps analytically)
+ v_new, updated_j_trace = LIFSRM._evaluate_SRM_filter( ## apply SRM
+ t, self.t_last_spike.get(), self.j_lowpass.get(), self.last_t_eval.get(),
+ self.current_j.get(), self.tau_m, self.v_rest, self.v_reset, dt
+ )
+
+ s_new = (v_new > self.thr) * 1.0
+ updated_t_last = jnp.where(s_new == 1.0, t, self.t_last_spike.get())
+ v_output = v_new * (1.0 - s_new) + s_new * self.v_reset
+
+ ## update compartment states
+ self.v.set(v_output)
+ self.s.set(s_new)
+ self.j_lowpass.set(updated_j_trace)
+ self.t_last_spike.set(updated_t_last)
+ self.last_t_eval.set(jnp.full((self.batch_size, self.n_units), t)) ## mark this time-stamp as "evaluated"
+
+ @compilable
+ def reset(self):
+ self.current_j.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.v.set(jnp.full((self.batch_size, self.n_units), self.v_rest))
+ self.s.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.t_last_spike.set(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.next_spike_t.set(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.j_lowpass.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.last_t_eval.set(jnp.zeros((self.batch_size, self.n_units)))
+
+ @compilable
+ def predict_next_spike(self, t_start):
+ next_spike_t = LIFSRM._predict_next_spike( ## call next-spike-time predictor
+ t_start,
+ self.t_last_spike.get(),
+ self.j_lowpass.get(),
+ self.tau_m,
+ self.v_rest,
+ self.v_reset,
+ self.thr
+ )
+ self.next_spike_t.set(next_spike_t) ## store predicted next spike time(s)
+
+ @staticmethod
+ def _evaluate_SRM_filter( ## kernel co-routine
+ t,
+ t_last_spike,
+ j_lowpass,
+ last_t_eval,
+ current_j,
+ tau_m,
+ v_rest,
+ v_reset,
+ dt
+ ):
+ ## applies filter-based SRM - tracks integrated voltage contributions
+ ## calculate continuous elapsed time since this specific neuron group was last evaluated
+ delta_t_eval = t - last_t_eval
+ ## analytically decay historical input voltage trace over skipped time gap
+ decayed_j_trace = j_lowpass * jnp.exp(-delta_t_eval / tau_m)
+ ## add new incoming current pulse scaled to operate akin to single LIFCell Euler step
+ new_j_trace = decayed_j_trace + (dt / tau_m) * current_j ## epsilon-kernel
+ ## calc analytical spike-post-emission kernel values (self-reset mechanism)
+ has_spiked = (t_last_spike >= 0.0) * 1.0
+ s_post = t - t_last_spike ## kappa kernel
+ eta_val = has_spiked * (v_reset - v_rest) * jnp.exp(-s_post / tau_m) ## eta kernel
+
+ v_total = v_rest + new_j_trace + eta_val ## sum explicit kernel terms
+ return v_total, new_j_trace
+
+ @staticmethod
+ def _predict_next_spike( ## next-spike-time prediction co-routine
+ t_start, t_last_spike, j_lowpass, tau_m, v_rest, v_reset, thr
+ ):
+ ## co-routine predicts future next spike, taking particular advantage of an LIF-SRM's
+ ## closed-form setup; specifically, this function calculates a future (global)
+ ## clock time-stamp as to when this LIF model's decaying voltage would
+ ## cross a firing threshold (note, this does not require step-wise numerical integration)
+
+ ## reconstruct base self-reset kernel magnitude (eta_val)
+ ### based on what historical displacement remains from last discharge event
+ has_spiked = (t_last_spike >= 0.0) * 1.0
+ s_post = t_start - t_last_spike
+ eta_val = has_spiked * (v_reset - v_rest) * jnp.exp(-s_post / tau_m)
+ ## extract combined driving force variable
+ total_driving_trace = j_lowpass + eta_val
+
+ ## define static threshold distance displacement metric
+ thr_distance = thr - v_rest
+
+ ## calc closed-form logarithmic isolation calculation for remaining segment time
+ can_reach_thr = total_driving_trace > thr_distance
+ ## for cases where a neuronal unit does not have enough charge to cross threshold:
+ safe_ratio = jnp.where( ## handles division-by-zero / negative log errors
+ can_reach_thr, total_driving_trace / jnp.maximum(thr_distance, 1e-5), 1.0
+ )
+ s_remaining = tau_m * jnp.log(safe_ratio)
+ ## absorb into current evaluation timestamp tracking variable
+ predicted_t = t_start + s_remaining
+ ## safety check: if total driving force is insufficient to cross,
+ ## then flag output as "un-triggered" (i.e.,-1.0)
+ return jnp.where(can_reach_thr, predicted_t, -1.0) # predicted spike time(s)
diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py
index 6c2bdc5d..3843b90c 100755
--- a/ngclearn/components/neurons/spiking/RAFCell.py
+++ b/ngclearn/components/neurons/spiking/RAFCell.py
@@ -8,33 +8,6 @@
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
-########################################################################################################################
-## RAF dynamics (multi-dimensional ODEs)
-@jit
-def _dfv_internal(j, v, w, tau_m, omega, b): ## "voltage" dynamics
- # dy/dt = omega x + b y
- dv_dt = omega * w + v * b ## dv/dt
- dv_dt = dv_dt * (1./tau_m)
- return dv_dt
-
-def _dfv(t, v, params): ## voltage dynamics wrapper
- j, w, tau_m, omega, b = params
- dv_dt = _dfv_internal(j, v, w, tau_m, omega, b)
- return dv_dt
-
-@jit
-def _dfw_internal(j, v, w, tau_w, omega, b): ## raw angular driver dynamics
- # dx/dt = b x − omega y + I; I is scaled injected electrical current
- dw_dt = w * b - v * omega + j
- dw_dt = dw_dt * (1./tau_w)
- return dw_dt
-
-def _dfw(t, w, params): ## angular driver dynamics wrapper
- j, v, tau_w, omega, b = params
- dv_dt = _dfw_internal(j, v, w, tau_w, omega, b)
- return dv_dt
-########################################################################################################################
-
class RAFCell(JaxComponent):
"""
The resonate-and-fire (RAF) neuronal cell
@@ -99,34 +72,51 @@ class RAFCell(JaxComponent):
@deprecate_args(resist_m="resist_v", tau_m="tau_v", b="dampen_factor")
def __init__(
- self, name, n_units, tau_v=1., tau_w=1., thr=1., omega=10., dampen_factor=-1., v_reset=0., w_reset=0.,
- v0=0., w0=0., resist_v=1., integration_type="euler", batch_size=1, **kwargs
+ self,
+ name,
+ n_units,
+ tau_v=1.,
+ tau_w=1.,
+ thr=1.,
+ omega=10.,
+ dampen_factor=-1.,
+ v_reset=0.,
+ w_reset=0.,
+ v0=0., ## voltage/membrane potential initial conditions
+ w0=0., ## angular variable driver initial conditions
+ resist_v=1.,
+ post_spike_reset=True, ## if True, snaps states to reset values post-spike-emission
+ integration_type="euler",
+ batch_size=1,
+ **kwargs
):
#v_rest=-72., v_reset=-75., w_reset=0., thr=5., v0=-70., w0=0., tau_w=400., thr=5., omega=10., b=-1.
super().__init__(name, **kwargs)
- ## Integration properties
+ ## integration properties
self.integrationType = integration_type
self.intgFlag = get_integrator_code(self.integrationType)
- ## Cell properties
+ ## RAF cell properties
+ self.post_spike_reset = post_spike_reset
self.tau_v = tau_v
self.resist_v = resist_v
self.tau_w = tau_w
self.omega = omega ## angular frequency
self.dampen_factor = dampen_factor ## dampening factor (b)
- ## Note: the smaller that dampen_factor "b" is, the faster the oscillation dampens to resting state values
+ ## Note: the smaller that dampen_factor "b" is, faster oscillation dampens
+ ## to resting state values
self.v_reset = v_reset
self.w_reset = w_reset
self.v0 = v0
self.w0 = w0
self.thr = thr
- ## Layer Size Setup
+ ## layer size setup
self.batch_size = batch_size
self.n_units = n_units
- ## Compartment setup
+ ## RAF key compartment setup
restVals = jnp.zeros((self.batch_size, self.n_units))
self.j = Compartment(restVals, display_name="Current", units="mA")
self.v = Compartment(restVals + self.v0, display_name="Voltage", units="mV")
@@ -136,6 +126,21 @@ def __init__(
restVals, display_name="Time-of-Last-Spike", units="ms"
) ## time-of-last-spike
+ ## RAF dynamics: multi-dimensional ODE system (2 ODEs)
+ def _dfv(t, v, params): ## voltage dynamics wrapper
+ j, w, tau_m, omega, b = params
+ #dv_dt = _dfv_internal(j, v, w, tau_m, omega, b)
+ dv_dt = omega * w + v * b ## dv/dt
+ dv_dt = dv_dt * (1./tau_m)
+ return dv_dt
+
+ def _dfw(t, w, params): ## angular driver dynamics wrapper
+ j, v, tau_w, omega, b = params
+ #dv_dt = _dfw_internal(j, v, w, tau_w, omega, b)
+ dw_dt = w * b - v * omega + j
+ dw_dt = dw_dt * (1./tau_w)
+ return dw_dt
+
@compilable
def advance_state(self, t, dt):
## continue with centered dynamics
@@ -143,20 +148,23 @@ def advance_state(self, t, dt):
if self.intgFlag == 1: ## RK-2/midpoint
## Note: we integrate ODEs in order: first w, then v
w_params = (j_, self.v.get(), self.tau_w, self.omega, self.dampen_factor)
- _, _w = step_rk2(0., self.w.get(), _dfw, dt, w_params)
+ _, _w = step_rk2(0., self.w.get(), RAFCell._dfw, dt, w_params)
v_params = (j_, _w, self.tau_v, self.omega, self.dampen_factor)
- _, _v = step_rk2(0., self.v.get(), _dfv, dt, v_params)
+ _, _v = step_rk2(0., self.v.get(), RAFCell._dfv, dt, v_params)
else: # integType == 0 (default -- Euler)
## Note: we integrate ODEs in order: first w, then v
w_params = (j_, self.v.get(), self.tau_w, self.omega, self.dampen_factor)
- _, _w = step_euler(0., self.w.get(), _dfw, dt, w_params)
+ _, _w = step_euler(0., self.w.get(), RAFCell._dfw, dt, w_params)
v_params = (j_, _w, self.tau_v, self.omega, self.dampen_factor)
- _, _v = step_euler(0., self.v.get(), _dfv, dt, v_params)
+ _, _v = step_euler(0., self.v.get(), RAFCell._dfv, dt, v_params)
s = (_v > self.thr) * 1. ## emit spikes/pulses
- ## hyperpolarize/reset/snap variables
- w = _w * (1. - s) + s * self.w_reset
- v = _v * (1. - s) + s * self.v_reset
+ if self.post_spike_reset: ## hyperpolarize/reset/snap variables
+ w = _w * (1. - s) + s * self.w_reset
+ v = _v * (1. - s) + s * self.v_reset
+ else:
+ w = _w
+ v = _v
self.tols.set((1. - s) * self.tols.get() + (s * t)) ## update times-of-last-spike(s)
diff --git a/ngclearn/components/neurons/spiking/RAFSRM.py b/ngclearn/components/neurons/spiking/RAFSRM.py
new file mode 100644
index 00000000..5b59bc16
--- /dev/null
+++ b/ngclearn/components/neurons/spiking/RAFSRM.py
@@ -0,0 +1,254 @@
+from ngclearn.components.jaxComponent import JaxComponent
+from jax import numpy as jnp, jit, lax
+from ngclearn import compilable, Compartment
+
+
+class RAFSRM(JaxComponent): ## RAF spike-response model (RAF-SRM)
+ """
+ The resonate-and-fire (RAF) spike-response model (SRM); this SRM computes
+ dynamics of RAF units analytically.
+
+ | --- Cell Input Compartments: ---
+ | current_j - electrical current input (takes in external signals)
+ | --- Cell State Compartments: ---
+ | v - membrane potential/voltage state
+ | j_v - voltage variable state
+ | j_w - angular-driver variable state
+ | key - JAX PRNG key
+ | --- Cell Output Compartments: ---
+ | s - emitted binary spikes/action potentials
+ | t_last_spike - time-of-last-spike (output)
+ | last_t_eval - time of last (SRM) evaluation
+
+ | References:
+ | Richardson, M.J., Brunel, N. and Hakim, V., 2003. From subthreshold to
+ | firing-rate resonance. Journal of neurophysiology, 89(5), pp.2538-2554.
+ |
+ | Izhikevich, Eugene M. "Resonate-and-fire neurons." Neural networks 14.6-7 (2001): 883-894.
+
+ | Pedagogical References:
+ | http://www.scholarpedia.org/article/Spike-response_model
+
+ Args:
+ name: the string name of this cell
+
+ n_units: number of cellular entities (neural population size)
+
+ tau_v: membrane/voltage time constant (Default: 1 ms)
+
+ tau_w: angular driver variable time constant (Default: 1 ms)
+
+ thr: voltage/membrane threshold (to obtain action potentials in terms
+ of binary spikes) (Default: 1 mV)
+
+ omega: angular frequency (Default: 10)
+
+ dampen_factor: oscillation dampening factor (Default: -1) ("b" as in RAF-Cell; Izhikevich 2001)
+
+ v_reset: reset condition for membrane potential (Default: 1 mV)
+
+ w_reset: reset condition for angular current driver (Default: 0)
+ """
+
+ def __init__(
+ self,
+ name,
+ n_units,
+ tau_v=1.,
+ tau_w=1.,
+ thr=1., ## threshold
+ omega=10.,
+ dampen_factor=-1.,
+ v_reset=0., ## membrane reset potential
+ w_reset=0., ## angular-driver variable reset potential
+ batch_size=1,
+ **kwargs
+ ):
+ super().__init__(name, **kwargs)
+ ## RAF-SRM meta-parameters
+ self.n_units = n_units
+ self.tau_v = tau_v
+ self.tau_w = tau_w
+ self.omega = omega
+ self.dampen_factor = dampen_factor
+ self.thr = thr
+ self.v_reset = v_reset
+ self.w_reset = w_reset
+ self.batch_size = batch_size
+
+ ## set up SRM's key compartments
+ self.current_j = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+ self.v = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+ self.s = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+ ## set up analytical state filters (to match RAFCell construction)
+ self.t_last_spike = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.next_spike_t = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.j_v = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+ self.j_w = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+ self.last_t_eval = Compartment(jnp.zeros((self.batch_size, self.n_units)))
+
+ @compilable
+ def advance_state(self, dt, t):
+ v_new, updated_j_v, updated_j_w = RAFSRM._evaluate_SRM_filter( ## apply SRM to get new states
+ t, self.t_last_spike.get(), self.j_v.get(), self.j_w.get(),
+ self.last_t_eval.get(), self.current_j.get(), self.tau_v, self.tau_w,
+ self.omega, self.dampen_factor, self.v_reset, self.w_reset, dt
+ )
+ ## calculate spike/pulse emission
+ s_new = (v_new > self.thr) * 1.0
+ updated_t_last = jnp.where(s_new == 1.0, t, self.t_last_spike.get())
+ v_output = v_new * (1.0 - s_new) + s_new * self.v_reset
+ ## update internal compartments with updated SRM states
+ self.v.set(v_output)
+ self.s.set(s_new)
+ self.j_v.set(updated_j_v)
+ self.j_w.set(updated_j_w)
+ self.t_last_spike.set(updated_t_last)
+ self.last_t_eval.set(jnp.full((self.batch_size, self.n_units), t))
+
+ @compilable
+ def reset(self): #
+ self.current_j.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.v.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.s.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.t_last_spike.set(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.next_spike_t.set(jnp.full((self.batch_size, self.n_units), -1.0))
+ self.j_v.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.j_w.set(jnp.zeros((self.batch_size, self.n_units)))
+ self.last_t_eval.set(jnp.zeros((self.batch_size, self.n_units)))
+
+ @compilable
+ def predict_next_spike(self, t_start):
+ next_spike_t = RAFSRM._predict_next_spike( ## next-spike-time predictor
+ t_start,
+ self.j_v.get(),
+ self.j_w.get(),
+ self.tau_v,
+ self.tau_w,
+ self.omega,
+ self.dampen_factor,
+ self.thr,
+ max_iters=10
+ )
+ self.next_spike_t.set(next_spike_t) ## store predicted next spike time(s)
+
+ @staticmethod
+ def _evaluate_SRM_filter( ## kernel co-routine
+ t,
+ t_last_spike,
+ j_v,
+ j_w,
+ last_t_eval,
+ current_j,
+ tau_v,
+ tau_w,
+ omega,
+ dampen_factor,
+ v_reset,
+ w_reset,
+ dt
+ ):
+ ## applies filter-based cumulative SRM for resonate-and-fire (RAF) units
+ ## NOTE: this analytical SRM is designed to match exact system matrix of RAFCell
+ delta_t_eval = t - last_t_eval ## calc (Dirac) delta time since last spike evaluation
+ ## reconstruct continuous system (of matrix parameters) of RAF dampened oscilator
+ ### dv/dt = a*v + b*w, and, dw/dt = c*v + d*w + j/tau_w
+ a = dampen_factor / tau_v
+ b = omega / tau_v
+ c = -omega / tau_w
+ d = dampen_factor / tau_w
+ ## calculate trace dampening (gamma) & determinant natural frequency
+ gamma = -0.5 * (a + d)
+ omega_0_sq = (a * d) - (b * c)
+ omega_d = jnp.sqrt(jnp.maximum(omega_0_sq - gamma**2, 1e-6))
+ ## analytically propagate ongoing sub-threshold state across time jump
+ decay = jnp.exp(-gamma * delta_t_eval)
+ cos_wd = jnp.cos(omega_d * delta_t_eval)
+ sin_wd = jnp.sin(omega_d * delta_t_eval)
+
+ v_old = j_v ## store current v
+ w_old = j_w ## store current w
+ ## calculate decoupled matrix exponential transformation equations
+ v_prop = decay * (
+ v_old * (cos_wd + ((a + gamma) / omega_d) * sin_wd) + w_old * (b / omega_d) * sin_wd
+ )
+ w_prop = decay * (
+ v_old * (c / omega_d) * sin_wd + w_old * (cos_wd + ((d + gamma) / omega_d) * sin_wd)
+ )
+ ## apply input current as "Euler-step" velocity shift to match exact step time
+ new_v = v_prop
+ new_w = w_prop + (current_j / tau_w) * dt
+ ## calcuate spike post-emission / self-reset kernel contribution
+ has_spiked = (t_last_spike >= 0.0) * 1.0
+ s_post = t - t_last_spike ## kappa-kernel
+ eta_decay = jnp.exp(-gamma * s_post)
+ eta_cos = jnp.cos(omega_d * s_post)
+ eta_sin = jnp.sin(omega_d * s_post)
+ ## set up eta kernel
+ eta_v = has_spiked * (
+ v_reset * (eta_cos + ((a + gamma) / omega_d) * eta_sin) + w_reset * (b / omega_d) * eta_sin
+ )
+ ## compute total combined membrane voltage output
+ v_total = new_v + eta_v
+ return v_total, new_v, new_w
+
+ @staticmethod
+ def _predict_next_spike( ## co-routine used to compute when a future spike will occur
+ t_start,
+ j_v,
+ j_w,
+ tau_v,
+ tau_w,
+ omega,
+ dampen_factor,
+ thr,
+ max_iters=10
+ ):
+ ## NOTE: this co-routine is based on a Newton-Raphson root finding process to predict when exactly
+ ## continuous sub-threshold wave equation will cross threshold line
+
+ ## standardize SRM system parameters
+ a = dampen_factor / tau_v
+ b = omega / tau_v
+ c = -omega / tau_w
+ d = dampen_factor / tau_w
+
+ gamma = -0.5 * (a + d)
+ omega_0_sq = (a * d) - (b * c)
+ omega_d = jnp.sqrt(jnp.maximum(omega_0_sq - gamma**2, 1e-6))
+
+ ## define analytical function V(s) and its temporal derivative dV(s)/ds
+ def evaluate_v_and_dv(s):
+ decay = jnp.exp(-gamma * s)
+ cos_wd = jnp.cos(omega_d * s)
+ sin_wd = jnp.sin(omega_d * s)
+ ## calc continuous voltage coordinate position equation
+ v_prop = decay * (j_v * (cos_wd + ((a + gamma) / omega_d) * sin_wd) + j_w * (b / omega_d) * sin_wd)
+ ## calc continuous velocity coordinate position equation (dV/ds)
+ ## based on underlying coupled system matrix row: dv/dt = a*v + b*w
+ w_prop = decay * (j_v * (c / omega_d) * sin_wd + j_w * (cos_wd + ((d + gamma) / omega_d) * sin_wd))
+ dv_ds = a * v_prop + b * w_prop
+ return v_prop, dv_ds
+
+ ## initialize Newton-Raphson search sequence loops; initial guess: s ~ 0.1ms
+ s_guess = jnp.full_like(j_v, 0.1)
+
+ def scan_body(carry, _): ## set up scanner for this search process
+ s_curr = carry
+ v_curr, dv_curr = evaluate_v_and_dv(s_curr)
+ f_s = v_curr - thr
+ ## update root estimation step while avoiding divide-by-zero bounds
+ s_next = s_curr - f_s / jnp.where(jnp.abs(dv_curr) < 1e-5, 1e-5, dv_curr)
+ ## keep estimated elapsed intervals bounded strictly positive
+ s_next = jnp.maximum(s_next, 0.0)
+ return s_next, None
+
+ ## execute unrolled iteration passes completely w/in JAX/JIT compiler context
+ final_s, _ = lax.scan(scan_body, s_guess, None, length=max_iters)
+ predicted_t = t_start + final_s ## compute absolute (executed) global clock projection coordinate
+ ## employ a safety check:
+ ### if initial slope is negative (or diverging), flag it as un-triggered (i.e., -1.0)
+ v_final, _ = evaluate_v_and_dv(final_s)
+ is_valid = (final_s > 0.0) & (jnp.abs(v_final - thr) < 1e-2)
+ return jnp.where(is_valid, predicted_t, -1.0) ## spit out guessed future spike time
+
diff --git a/ngclearn/components/neurons/spiking/__init__.py b/ngclearn/components/neurons/spiking/__init__.py
index b4c0b3db..45fac156 100644
--- a/ngclearn/components/neurons/spiking/__init__.py
+++ b/ngclearn/components/neurons/spiking/__init__.py
@@ -10,3 +10,7 @@
from .RAFCell import RAFCell
from .hodgkinHuxleyCell import HodgkinHuxleyCell
+## spike-response models (SRMs)
+from .LIFSRM import LIFSRM
+from .RAFSRM import RAFSRM
+
diff --git a/ngclearn/components/neurons/spiking/quadLIFCell.py b/ngclearn/components/neurons/spiking/quadLIFCell.py
index 6d7c95b6..4a132500 100755
--- a/ngclearn/components/neurons/spiking/quadLIFCell.py
+++ b/ngclearn/components/neurons/spiking/quadLIFCell.py
@@ -120,9 +120,23 @@ class QuadLIFCell(LIFCell): ## quadratic integrate-and-fire cell
@deprecate_args(thr_jitter=None, critical_V="critical_v")
def __init__(
- self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., v_reset=-60., v_scale=-41.6, critical_v=1.,
- tau_theta=1e7, theta_plus=0.05, refract_time=5., one_spike=False, integration_type="euler",
- surrogate_type="straight_through", v_min=None, **kwargs
+ self,
+ name,
+ n_units,
+ tau_m, resist_m=1.,
+ thr=-52.,
+ v_rest=-65.,
+ v_reset=-60.,
+ v_scale=-41.6,
+ critical_v=1.,
+ tau_theta=1e7,
+ theta_plus=0.05,
+ refract_time=5.,
+ one_spike=False,
+ integration_type="euler",
+ surrogate_type="straight_through",
+ v_min=None,
+ **kwargs
):
super().__init__(
name, n_units, tau_m, resist_m, thr, v_rest, v_reset, 1., tau_theta, theta_plus, refract_time,
diff --git a/ngclearn/components/synapses/STPDenseSynapse.py b/ngclearn/components/synapses/STPDenseSynapse.py
index 31cf7c67..3ce8c1ab 100755
--- a/ngclearn/components/synapses/STPDenseSynapse.py
+++ b/ngclearn/components/synapses/STPDenseSynapse.py
@@ -1,5 +1,6 @@
from jax import random, numpy as jnp, jit
from ngcsimlib.logger import info
+from ngcsimlib import deprecate_args
from ngclearn.utils.distribution_generator import DistributionGenerator
from ngclearn import compilable #from ngcsimlib.parser import compilable
@@ -38,8 +39,8 @@ class STPDenseSynapse(DenseSynapse): ## short-term plastic synaptic cable
bias_init: a kernel to drive initialization of biases for this synaptic cable
(Default: None, which turns off/disables biases)
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1 and > 0. will result in a sparser synaptic structure
@@ -54,11 +55,23 @@ class STPDenseSynapse(DenseSynapse): ## short-term plastic synaptic cable
resources_int: initialization kernel for synaptic resources matrix
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, weight_init=None, bias_init=None, resist_scale=1., p_conn=1., tau_f=750., tau_d=50.,
- resources_init=None, **kwargs
+ self,
+ name,
+ shape,
+ weight_init=None,
+ bias_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ tau_f=750.,
+ tau_d=50.,
+ resources_init=None,
+ **kwargs
):
- super().__init__(name, shape, weight_init, bias_init, resist_scale, p_conn, **kwargs)
+ super().__init__(
+ name, shape, weight_init, bias_init, g_conduct_factor, p_conn, **kwargs
+ )
## STP meta-parameters
self.resources_init = resources_init
self.tau_f = tau_f
@@ -94,7 +107,7 @@ def advance_state(self, t, dt):
if self.tau_d > 0.:
x = x + (1. - x) * (1./self.tau_d) - u * x * s
## else, do nothing with x (keep it pointing to current x compartment)
- outputs = jnp.matmul(self.inputs.get(), Wdyn * self.resist_scale) + self.biases.get()
+ outputs = jnp.matmul(self.inputs.get(), Wdyn * self.g_conduct_factor) + self.biases.get()
self.outputs.set(outputs)
self.u.set(u)
@@ -156,7 +169,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance/average level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"tau_f": "Short-term facilitation time constant",
"tau_d": "Short-term depression time constant"
diff --git a/ngclearn/components/synapses/__init__.py b/ngclearn/components/synapses/__init__.py
index 28bd001b..5bb09fbd 100644
--- a/ngclearn/components/synapses/__init__.py
+++ b/ngclearn/components/synapses/__init__.py
@@ -1,5 +1,7 @@
from .denseSynapse import DenseSynapse
from .staticSynapse import StaticSynapse
+from .sparseTensorSynapse import SparseTensorSynapse
+from .staticSparseTensorSynapse import StaticSparseTensorSynapse
## short-term plasticity components
from .STPDenseSynapse import STPDenseSynapse
@@ -14,6 +16,7 @@
from .hebbian.expSTDPSynapse import ExpSTDPSynapse
from .hebbian.eventSTDPSynapse import EventSTDPSynapse
from .hebbian.BCMSynapse import BCMSynapse
+from .hebbian.inhibitorySTDPSynapse import InhibitorySTDPSynapse
from .mpsSynapse import MPSSynapse
### dense competitive synaptic components/elements
from .competitive.SOMSynapse import SOMSynapse
diff --git a/ngclearn/components/synapses/alphaSynapse.py b/ngclearn/components/synapses/alphaSynapse.py
index cbdbb8c8..f77fdff6 100644
--- a/ngclearn/components/synapses/alphaSynapse.py
+++ b/ngclearn/components/synapses/alphaSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
@@ -47,7 +48,7 @@ class AlphaSynapse(DenseSynapse): ## dynamic alpha synapse cable
bias_init: a kernel to drive initialization of biases for this synaptic cable
(Default: None, which turns off/disables biases)
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
transform (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
p_conn: probability of a connection existing (default: 1.); setting
@@ -58,11 +59,22 @@ class AlphaSynapse(DenseSynapse): ## dynamic alpha synapse cable
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, tau_decay, g_syn_bar, syn_rest, weight_init=None, bias_init=None, resist_scale=1.,
- p_conn=1., is_nonplastic=True, **kwargs
+ self,
+ name,
+ shape,
+ tau_decay,
+ g_syn_bar,
+ syn_rest,
+ weight_init=None,
+ bias_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ is_nonplastic=True,
+ **kwargs
):
- super().__init__(name, shape, weight_init, bias_init, resist_scale, p_conn, **kwargs)
+ super().__init__(name, shape, weight_init, bias_init, g_conduct_factor, p_conn, **kwargs)
## dynamic synapse meta-parameters
self.tau_decay = tau_decay
self.g_syn_bar = g_syn_bar
@@ -138,7 +150,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"tau_decay": "Conductance decay time constant (ms)",
"g_bar_syn": "Maximum conductance value",
diff --git a/ngclearn/components/synapses/competitive/ART2ASynapse.py b/ngclearn/components/synapses/competitive/ART2ASynapse.py
index de0a9f4c..d204e553 100644
--- a/ngclearn/components/synapses/competitive/ART2ASynapse.py
+++ b/ngclearn/components/synapses/competitive/ART2ASynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit, nn
+from ngcsimlib import deprecate_args
from functools import partial
from ngclearn import compilable
from ngclearn import Compartment
@@ -6,15 +7,6 @@
from ngclearn.components.synapses.denseSynapse import DenseSynapse
-@partial(jit, static_argnums=[1])
-def _normalize(x_in, norm_fx=0):
- if norm_fx == 1:
- xmin = jnp.min(x, axis=1, keepdims=True)
- xmax = jnp.max(x, axis=1, keepdims=True)
- x = (x_in - xmin)/(xmax - xmin)
- else:
- x = x_in / jnp.linalg.norm(x_in, ord=2, axis=1, keepdims=True)
- return x
class ART2ASynapse(DenseSynapse): # Adaptive resonance theory (ART) 2A synaptic cable
"""
@@ -58,12 +50,13 @@ class ART2ASynapse(DenseSynapse): # Adaptive resonance theory (ART) 2A synaptic
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform (Default: 1.)
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform (Default: 1.)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1. will result in a sparser synaptic structure
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
self,
name,
@@ -72,13 +65,13 @@ def __init__(
eta_decrement=0., ## linear scheduled decrement over eta
vigilance=0.3, ## vigilance parameter (rho)
weight_init=None,
- resist_scale=1.,
+ g_conduct_factor=1.,
p_conn=1.,
batch_size=1,
**kwargs
):
super().__init__(
- name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
)
### Synapse and ART-2A hyper-parameters
@@ -142,13 +135,24 @@ def consolidate(self): ## memory consolition co-routine (for non-resonant signal
## can add a function that "grows" out block matrix by a chunk (to control growth)
## TODO: add pruning mechanism for low-usage slots
+ #@partial(jit, static_argnums=[1])
+ @staticmethod
+ def _normalize(x_in, norm_fx=0):
+ if norm_fx == 1:
+ xmin = jnp.min(x, axis=1, keepdims=True)
+ xmax = jnp.max(x, axis=1, keepdims=True)
+ x = (x_in - xmin) / (xmax - xmin)
+ else:
+ x = x_in / jnp.linalg.norm(x_in, ord=2, axis=1, keepdims=True)
+ return x
+
@compilable
def advance_state(self): ## forward-inference step of ART2A
x_in = self.inputs.get()
W = self.weights.get() ## get (transposed) memory matrix
used = self.used.get()
- x = _normalize(x_in, norm_fx=self.norm_fx)
+ x = ART2ASynapse._normalize(x_in, norm_fx=self.norm_fx)
self.xprobe.set(x)
sims = jnp.matmul(x, W) ## compute similarities (parallel dot products)
## we correct activities by masking out unused slots
@@ -227,7 +231,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"eta": "Global learning rate",
"eta_decrement": "Constant amount to decrease global learning by each call to `evolve`"
diff --git a/ngclearn/components/synapses/competitive/SOMSynapse.py b/ngclearn/components/synapses/competitive/SOMSynapse.py
index c95a48d4..4e340f30 100644
--- a/ngclearn/components/synapses/competitive/SOMSynapse.py
+++ b/ngclearn/components/synapses/competitive/SOMSynapse.py
@@ -136,7 +136,7 @@ class SOMSynapse(DenseSynapse): # Self-organizing map (SOM) synaptic cable
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
(Default: 1.), i.e., yields: out = ((W * Rscale) * in)
p_conn: probability of a connection existing (default: 1.); setting
@@ -153,14 +153,14 @@ def __init__(
distance_function="euclidean",
neighbor_function="gaussian",
weight_init=None,
- resist_scale=1.,
+ g_conduct_factor=1.,
p_conn=1.,
batch_size=1,
**kwargs
):
shape = (n_inputs, n_units_x * n_units_y)
super().__init__(
- name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
)
### build (rectangular) topology coordinates
@@ -265,9 +265,9 @@ def advance_state(self): ## forward-inference step of SOM
self.neighbor_weights.set(neighbor_weights) ## store neighborhood weightings
## compute an approximate weighted activity output for input pattern
- #activity = jnp.sum(self.weights * self.resist_scale * neighbor_weights, axis=1, keepdims=True)
+ #activity = jnp.sum(self.weights * self.g_conduct_factor * neighbor_weights, axis=1, keepdims=True)
### obtain weighted competitive activations (via softmax probs)
- activity = softmax(neighbor_weights * self.resist_scale)
+ activity = softmax(neighbor_weights * self.g_conduct_factor)
self.outputs.set(activity)
@compilable
@@ -357,7 +357,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"eta": "Global learning rate",
"radius": "Radius parameter to control influence of neighborhood function",
diff --git a/ngclearn/components/synapses/competitive/hopfieldSynapse.py b/ngclearn/components/synapses/competitive/hopfieldSynapse.py
index cf22a8f9..f72d4635 100644
--- a/ngclearn/components/synapses/competitive/hopfieldSynapse.py
+++ b/ngclearn/components/synapses/competitive/hopfieldSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit, vmap
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.utils.model_utils import softmax, bkwta
@@ -58,13 +59,14 @@ class HopfieldSynapse(DenseSynapse): # (Modern) Hopfield synaptic cable
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
- (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
+ (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1. will result in a sparser synaptic structure
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
self,
name,
@@ -75,13 +77,13 @@ def __init__(
initial_patterns=None,
update_rule = "delta", ## memory plasticity rule
weight_init=None,
- resist_scale=1.,
+ g_conduct_factor=1.,
p_conn=1.,
batch_size=1,
**kwargs
):
super().__init__(
- name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
)
### Synapse and Hopfield hyper-parameters
@@ -208,7 +210,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"eta": "Global learning rate (to control update to memory matrix)",
"beta": "Inverse temperature (controls softmax sharpness",
diff --git a/ngclearn/components/synapses/competitive/vectorQuantizeSynapse.py b/ngclearn/components/synapses/competitive/vectorQuantizeSynapse.py
index 37bea309..06f4e7b4 100644
--- a/ngclearn/components/synapses/competitive/vectorQuantizeSynapse.py
+++ b/ngclearn/components/synapses/competitive/vectorQuantizeSynapse.py
@@ -1,13 +1,11 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.utils.model_utils import softmax, bkwta #, chebyshev_norm
from ngclearn.components.synapses.denseSynapse import DenseSynapse
-def _gaussian_kernel(dist, sigma): ## Gaussian weighting function
- density = jnp.exp(-jnp.power(dist, 2) / (2 * (sigma ** 2))) # n_units x 1
- return density
class VectorQuantizeSynapse(DenseSynapse): # Vector quantization (VQ) synaptic cable
"""
@@ -74,13 +72,14 @@ class VectorQuantizeSynapse(DenseSynapse): # Vector quantization (VQ) synaptic c
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
- (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
+ (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1. will result in a sparser synaptic structure
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
self,
name,
@@ -94,13 +93,13 @@ def __init__(
initial_patterns=None, ## possible class-based prototypes to init by
langevin_noise_scale=0., ## scale of Langevin noise to apply to updates
weight_init=None,
- resist_scale=1.,
+ g_conduct_factor=1.,
p_conn=1.,
batch_size=1,
**kwargs
):
super().__init__(
- name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
)
### Synapse / VQ hyper-parameters
@@ -153,6 +152,12 @@ def __init__(
self.weights.set(initX)
if self.label_dim > 0: ## do we preload label matrix?
self.label_weights.set(initY)
+
+ @staticmethod
+ def _gaussian_kernel(dist, sigma): ## Gaussian weighting function
+ density = jnp.exp(-jnp.power(dist, 2) / (2 * (sigma ** 2))) # n_units x 1
+ return density
+
@compilable
def advance_state(self): ## forward-inference step of VQ
x_in = self.inputs.get()
@@ -267,7 +272,7 @@ def help(cls): ## component help function
"label_dim": "Dimensionality of labels (if this VQ is supervised)",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"eta": "Global learning rate",
"eta_decrement": "Constant to decrement `eta` by per update/call to `evolve()`",
diff --git a/ngclearn/components/synapses/convolution/convSynapse.py b/ngclearn/components/synapses/convolution/convSynapse.py
index 5af6810f..17977b2a 100755
--- a/ngclearn/components/synapses/convolution/convSynapse.py
+++ b/ngclearn/components/synapses/convolution/convSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngcsimlib.logger import info
@@ -37,16 +38,26 @@ class ConvSynapse(JaxComponent): ## base-level convolutional cable
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((K @ in) * resist_scale) + b
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((K @ in) * g_conduct_factor) + b
where `@` denotes convolution
batch_size: batch size dimension of this component
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, x_shape, filter_init=None, bias_init=None, stride=1, padding=None, resist_scale=1.,
- batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ x_shape,
+ filter_init=None,
+ bias_init=None,
+ stride=1,
+ padding=None,
+ g_conduct_factor=1.,
+ batch_size=1,
+ **kwargs
):
super().__init__(name, **kwargs)
@@ -57,7 +68,7 @@ def __init__(
self.shape = shape ## shape of synaptic filter tensor
x_size, x_size = x_shape
self.x_size = x_size
- self.resist_scale = resist_scale ## post-transformation scale factor
+ self.g_conduct_factor = g_conduct_factor ## post-transformation scale factor
self.padding = padding
self.stride = stride
@@ -105,10 +116,10 @@ def __init__(
@compilable
def advance_state(self): #Rscale, padding, stride, weights, biases, inputs):
_x = self.inputs.get()
- ## FIXME: does resist_scale affect update rules?
+ ## FIXME: does g_conduct_factor affect update rules?
outputs = conv2d(
_x, self.weights.get(), stride_size=self.stride, padding=self.padding
- ) * self.resist_scale + self.biases.get()
+ ) * self.g_conduct_factor + self.biases.get()
self.outputs.set(outputs)
@compilable
@@ -155,7 +166,7 @@ def help(cls): ## component help function
"x_shape": "Shape of any single incoming/input feature map",
"filter_init": "Initialization conditions for synaptic filter (K) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level output scaling factor (R)",
+ "g_conduct_factor": "Conductance level output scaling factor (R)",
"stride": "length / size of stride",
"padding": "pre-operator padding to use, i.e., `VALID` `SAME`"
}
diff --git a/ngclearn/components/synapses/convolution/deconvSynapse.py b/ngclearn/components/synapses/convolution/deconvSynapse.py
index cf52d9d7..3e5a3f74 100755
--- a/ngclearn/components/synapses/convolution/deconvSynapse.py
+++ b/ngclearn/components/synapses/convolution/deconvSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngcsimlib.logger import info
@@ -38,15 +39,16 @@ class DeconvSynapse(JaxComponent): ## base-level deconvolutional cable
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W @.T Rscale) * in) + b
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W @.T g_conduct_factor) * in) + b
where `@.T` denotes deconvolution
batch_size: batch size dimension of this component
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, x_shape, filter_init=None, bias_init=None, stride=1, padding=None, resist_scale=1.,
+ self, name, shape, x_shape, filter_init=None, bias_init=None, stride=1, padding=None, g_conduct_factor=1.,
batch_size=1, **kwargs
):
super().__init__(name, **kwargs)
@@ -58,7 +60,7 @@ def __init__(
self.shape = shape ## shape of synaptic filter tensor
x_size, x_size = x_shape
self.x_size = x_size
- self.resist_scale = resist_scale ## post-transformation scale factor
+ self.g_conduct_factor = g_conduct_factor ## post-transformation scale factor
self.padding = padding
self.stride = stride
@@ -84,8 +86,7 @@ def __init__(
self.outputs = Compartment(jnp.zeros(self.out_shape))
self.weights = Compartment(weights)
if self.bias_init is None:
- info(self.name, "is using default bias value of zero (no bias "
- "kernel provided)!")
+ info(self.name, "is using default bias value of zero (no bias kernel provided)!")
self.biases = Compartment(
# dist.initialize_params(subkeys[2], bias_init, (1, shape[1])) if bias_init else 0.0
self.bias_init((1, shape[1]), subkeys[2]) if bias_init else 0.0
@@ -96,7 +97,7 @@ def advance_state(self):
_x = self.inputs.get()
out = deconv2d(
_x, self.weights.get(), stride_size=self.stride, padding=self.padding
- ) * self.resist_scale + self.biases.get()
+ ) * self.g_conduct_factor + self.biases.get()
self.outputs.set(out)
@compilable
@@ -143,7 +144,7 @@ def help(cls): ## component help function
"x_shape": "Shape of any single incoming/input feature map",
"filter_init": "Initialization conditions for synaptic filter (K) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level output scaling factor (R)",
+ "g_conduct_factor": "Conductance level output scaling factor (R)",
"stride": "length / size of stride",
"padding": "pre-operator padding to use, i.e., `VALID` `SAME`"
}
diff --git a/ngclearn/components/synapses/convolution/hebbianConvSynapse.py b/ngclearn/components/synapses/convolution/hebbianConvSynapse.py
index a66242a4..2d2b7851 100755
--- a/ngclearn/components/synapses/convolution/hebbianConvSynapse.py
+++ b/ngclearn/components/synapses/convolution/hebbianConvSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.convolution.convSynapse import ConvSynapse
@@ -50,8 +51,8 @@ class HebbianConvSynapse(ConvSynapse): ## Hebbian-evolved convolutional cable
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((K @ in) * resist_scale) + b
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((K @ in) * g_conduct_factor) + b
where `@` denotes convolution
w_bound: maximum weight to softly bound this cable's value matrix to; if
@@ -81,13 +82,28 @@ class HebbianConvSynapse(ConvSynapse): ## Hebbian-evolved convolutional cable
batch_size: batch size dimension of this component
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, x_shape, eta=0., filter_init=None, bias_init=None, stride=1, padding=None,
- resist_scale=1., w_bound=0., is_nonnegative=False, w_decay=0., sign_value=1., optim_type="sgd",
- batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ x_shape,
+ eta=0.,
+ filter_init=None,
+ bias_init=None,
+ stride=1,
+ padding=None,
+ g_conduct_factor=1.,
+ w_bound=0.,
+ is_nonnegative=False,
+ w_decay=0.,
+ sign_value=1.,
+ optim_type="sgd",
+ batch_size=1,
+ **kwargs
):
super().__init__(
- name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=bias_init, resist_scale=resist_scale,
+ name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=bias_init, g_conduct_factor=g_conduct_factor,
stride=stride, padding=padding, batch_size=batch_size, **kwargs
)
@@ -235,7 +251,7 @@ def help(cls): ## component help function
"x_shape": "Shape of any single incoming/input feature map",
"filter_init": "Initialization conditions for synaptic filter (K) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level output scaling factor (R)",
+ "g_conduct_factor": "Conductance level output scaling factor (R)",
"stride": "length / size of stride",
"padding": "pre-operator padding to use, i.e., `VALID` `SAME`",
"is_nonnegative": "Should filters be constrained to be non-negative post-updates?",
diff --git a/ngclearn/components/synapses/convolution/hebbianDeconvSynapse.py b/ngclearn/components/synapses/convolution/hebbianDeconvSynapse.py
index d3317728..6d6d6005 100755
--- a/ngclearn/components/synapses/convolution/hebbianDeconvSynapse.py
+++ b/ngclearn/components/synapses/convolution/hebbianDeconvSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.convolution.deconvSynapse import DeconvSynapse
@@ -49,8 +50,8 @@ class HebbianDeconvSynapse(DeconvSynapse): ## Hebbian-evolved deconvolutional ca
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W @.T Rscale) * in) + b
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W @.T g_conduct_factor) * in) + b
where `@.T` denotes deconvolution
w_bound: maximum weight to softly bound this cable's value matrix to; if
@@ -80,13 +81,28 @@ class HebbianDeconvSynapse(DeconvSynapse): ## Hebbian-evolved deconvolutional ca
batch_size: batch size dimension of this component
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, x_shape, eta=0., filter_init=None, bias_init=None, stride=1, padding=None,
- resist_scale=1., w_bound=0., is_nonnegative=False, w_decay=0., sign_value=1., optim_type="sgd",
- batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ x_shape,
+ eta=0.,
+ filter_init=None,
+ bias_init=None,
+ stride=1,
+ padding=None,
+ g_conduct_factor=1.,
+ w_bound=0.,
+ is_nonnegative=False,
+ w_decay=0.,
+ sign_value=1.,
+ optim_type="sgd",
+ batch_size=1,
+ **kwargs
):
super().__init__(
- name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=bias_init, resist_scale=resist_scale,
+ name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=bias_init, g_conduct_factor=g_conduct_factor,
stride=stride, padding=padding, batch_size=batch_size, **kwargs
)
@@ -221,7 +237,7 @@ def help(cls): ## component help function
"x_shape": "Shape of any single incoming/input feature map",
"filter_init": "Initialization conditions for synaptic filter (K) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level output scaling factor (R)",
+ "g_conduct_factor": "Conductance level output scaling factor (R)",
"stride": "length / size of stride",
"padding": "pre-operator padding to use, i.e., `VALID` `SAME`",
"is_nonnegative": "Should filters be constrained to be non-negative post-updates?",
diff --git a/ngclearn/components/synapses/convolution/staticConvSynapse.py b/ngclearn/components/synapses/convolution/staticConvSynapse.py
index 4ffda128..a38c2a09 100755
--- a/ngclearn/components/synapses/convolution/staticConvSynapse.py
+++ b/ngclearn/components/synapses/convolution/staticConvSynapse.py
@@ -31,7 +31,7 @@ class StaticConvSynapse(ConvSynapse):
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
+ g_conduct_factor: a fixed (conductance) scaling factor to apply to synaptic
transform (Default: 1.), i.e., yields: out = ((K @ in) * resist_scale) + b
where `@` denotes convolution
diff --git a/ngclearn/components/synapses/convolution/staticDeconvSynapse.py b/ngclearn/components/synapses/convolution/staticDeconvSynapse.py
index 3fd9aeb4..69770150 100755
--- a/ngclearn/components/synapses/convolution/staticDeconvSynapse.py
+++ b/ngclearn/components/synapses/convolution/staticDeconvSynapse.py
@@ -31,7 +31,7 @@ class StaticDeconvSynapse(DeconvSynapse):
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
+ g_conduct_factor: a fixed (conductance) scaling factor to apply to synaptic
transform (Default: 1.), i.e., yields: out = ((K @ in) * resist_scale) + b
where `@` denotes convolution
diff --git a/ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py b/ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py
index 86aa33c4..6c878763 100755
--- a/ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py
+++ b/ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.convolution.convSynapse import ConvSynapse
@@ -55,8 +56,8 @@ class TraceSTDPConvSynapse(ConvSynapse): ## trace-based STDP convolutional cable
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((K @ in) * resist_scale) + b
+ g_conduct_factor: a fixed (conductance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((K @ in) * g_conduct_factor) + b
where `@` denotes convolution
w_bound: maximum weight to softly bound this cable's value matrix to; if
@@ -68,12 +69,28 @@ class TraceSTDPConvSynapse(ConvSynapse): ## trace-based STDP convolutional cable
batch_size: batch size dimension of this component
"""
+ from ngcsimlib import deprecate_args
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, x_shape, A_plus, A_minus, eta=0., pretrace_target=0., filter_init=None, stride=1,
- padding=None, resist_scale=1., w_bound=0., w_decay=0., batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ x_shape,
+ A_plus,
+ A_minus,
+ eta=0.,
+ pretrace_target=0.,
+ filter_init=None,
+ stride=1,
+ padding=None,
+ g_conduct_factor=1.,
+ w_bound=0.,
+ w_decay=0.,
+ batch_size=1,
+ **kwargs
):
super().__init__(
- name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=None, resist_scale=resist_scale,
+ name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=None, g_conduct_factor=g_conduct_factor,
stride=stride, padding=padding, batch_size=batch_size, **kwargs
)
@@ -215,7 +232,7 @@ def help(cls): ## component help function
"x `number input channels` x `number output channels`",
"x_shape": "Shape of any single incoming/input feature map",
"filter_init": "Initialization conditions for synaptic filter (K) values",
- "resist_scale": "Resistance level output scaling factor (R)",
+ "g_conduct_factor": "Conductance level output scaling factor (R)",
"stride": "length / size of stride",
"padding": "pre-operator padding to use, i.e., `VALID` `SAME`",
"A_plus": "Strength of long-term potentiation (LTP)",
diff --git a/ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py b/ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py
index a894213e..97951957 100755
--- a/ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py
+++ b/ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.convolution.deconvSynapse import DeconvSynapse
@@ -53,8 +54,8 @@ class TraceSTDPDeconvSynapse(DeconvSynapse): ## trace-based STDP deconvolutional
padding: pre-operator padding to use -- "VALID" (none), "SAME"
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((K @ in) * resist_scale) + b
+ g_conduct_factor: a fixed (conductance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((K @ in) * g_conduct_factor) + b
where `@` denotes convolution
w_bound: maximum weight to softly bound this cable's value matrix to; if
@@ -66,12 +67,27 @@ class TraceSTDPDeconvSynapse(DeconvSynapse): ## trace-based STDP deconvolutional
batch_size: batch size dimension of this component
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, x_shape, A_plus, A_minus, eta=0., pretrace_target=0., filter_init=None, stride=1,
- padding=None, resist_scale=1., w_bound=0., w_decay=0., batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ x_shape,
+ A_plus,
+ A_minus,
+ eta=0.,
+ pretrace_target=0.,
+ filter_init=None,
+ stride=1,
+ padding=None,
+ g_conduct_factor=1.,
+ w_bound=0.,
+ w_decay=0.,
+ batch_size=1,
+ **kwargs
):
super().__init__(
- name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=None, resist_scale=resist_scale,
+ name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=None, g_conduct_factor=g_conduct_factor,
stride=stride, padding=padding, batch_size=batch_size, **kwargs
)
@@ -201,7 +217,7 @@ def help(cls): ## component help function
"x_shape": "Shape of any single incoming/input feature map",
"batch_size": "Batch size dimension of this component",
"filter_init": "Initialization conditions for synaptic filter (K) values",
- "resist_scale": "Resistance level output scaling factor (R)",
+ "g_conduct_factor": "Conductance level output scaling factor (R)",
"stride": "length / size of stride",
"padding": "pre-operator padding to use, i.e., `VALID` `SAME`",
"A_plus": "Strength of long-term potentiation (LTP)",
diff --git a/ngclearn/components/synapses/denseSynapse.py b/ngclearn/components/synapses/denseSynapse.py
index 92c9c5e9..01b998b8 100755
--- a/ngclearn/components/synapses/denseSynapse.py
+++ b/ngclearn/components/synapses/denseSynapse.py
@@ -2,6 +2,7 @@
from ngclearn.components.jaxComponent import JaxComponent
from ngclearn.utils.distribution_generator import DistributionGenerator
from ngcsimlib.logger import info
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
@@ -11,11 +12,14 @@ class DenseSynapse(JaxComponent): ## base dense synaptic cable
A dense synaptic cable; no form of synaptic evolution/adaptation
is in-built to this component.
- | --- Synapse Compartments: ---
+ | --- Synapse Input Compartments: ---
| inputs - input (takes in external signals)
- | outputs - output signals
+ | --- Synapse State Compartments: ---
| weights - current value matrix of synaptic efficacies (strength values)
| biases - current value vector of synaptic bias values
+ | gate - current values of multiplicative (output) gate/modulator (Default: 1)
+ | --- Synapse Output Compartments: ---
+ | outputs - output signals
Args:
name: the string name of this cell
@@ -30,24 +34,36 @@ class DenseSynapse(JaxComponent): ## base dense synaptic cable
bias_init: a kernel to drive initialization of biases for this synaptic cable
(Default: None, which turns off/disables biases)
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W * in) * resist_scale) + bias
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * in) * g_conduct_factor) + bias
+
+ p_release_mean: probability of pre-synaptic transmission; only if this value is > 0 and < 1,
+ this synapse will enforce stochastic synaptic tranmission on pre-synaptic signals,
+ meaning that each pre-synaptic signal will make it across the synaptic cable with
+ a probability of `p(transmit) = p_release_mean +/- 0.1` (Default: 1)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1 and > 0. will result in a sparser synaptic structure
(lower values yield sparse structure)
+ max_delay_steps: maximum delay length (in terms of discrete simulation time-steps) to
+ delay transmission of pre-synaptic signals across this synaptic cable; note
+ that setting this to 0 disables the use of synaptic delay (Default: 0)
+
mask: if non-None, a (multiplicative) mask is applied to this synaptic weight matrix
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
self,
name,
shape,
weight_init=None,
bias_init=None,
- resist_scale=1.,
+ g_conduct_factor=1.,
+ p_release_mean=1.,
p_conn=1.,
+ max_delay_steps=0, ## "Tk"
mask=None,
batch_size=1,
**kwargs
@@ -61,18 +77,18 @@ def __init__(
## Synapse meta-parameters
self.shape = shape
- self.resist_scale = resist_scale
+ self.g_conduct_factor = g_conduct_factor
+ self.p_release_mean = p_release_mean
## Set up synaptic weight values
tmp_key, *subkeys = random.split(self.key.get(), 4)
if weight_init is None:
info(self.name, "is using default weight initializer!")
- # self.weight_init = {"dist": "uniform", "amin": 0.025, "amax": 0.8}
weight_init = DistributionGenerator.uniform(0.025, 0.8)
weights = weight_init(shape, subkeys[0])
- if 0. < p_conn < 1.: ## Modifier/constraint: only non-zero and <1 probs allowed
+ if 0. < p_conn < 1.: ## modifier/constraint: only non-zero and <1 probs allowed
p_mask = random.bernoulli(subkeys[1], p=p_conn, shape=shape)
weights = weights * p_mask ## sparsify matrix
@@ -91,22 +107,84 @@ def __init__(
if bias_init is None:
info(self.name, "is using default bias value of zero (no bias kernel provided)!")
self.biases = Compartment(bias_init((1, shape[1]), subkeys[2]) if bias_init else 0.0)
+ self.gate = Compartment(postVals + 1.)
## pin weight/bias initializers to component
self.weight_init = weight_init
self.bias_init = bias_init
+ ## Stochastic synaptic transmission - create static vector of heterogeneous release probabilities
+ key, *skey = random.split(self.key.get(), 4)
+ pre_units = shape[0]
+ self.p_release_mean = p_release_mean
+ p_jitter = 0.1 ## NOTE: this is hard-coded jitter
+ self.use_one_spike = True #False
+ self.p_release = jnp.ones((1, pre_units))
+ if 0. < self.p_release_mean < 1.: ## if proper p(transmit) mean given
+ self.p_release = random.uniform(
+ skey[0], shape=(1, pre_units), minval=self.p_release_mean - p_jitter, maxval=p_release_mean + p_jitter
+ ) ## probability of spike release
+
+ ## Implement staggered (pre-synaptic/axonal) cable delays
+ self.max_delay_steps = max_delay_steps ## extends up to a `Tk` ms temporal jitter window
+ pre_units = shape[0]
+ ## jitter axonal delays from 0 up to Tk
+ self.syn_delay_indices = random.randint(skey[1], shape=(pre_units,), minval=0, maxval=self.max_delay_steps)
+ ## create fixed memory grid to store a rolling history of incoming pre-synaptic signals
+ initial_buffer_state = jnp.zeros((self.max_delay_steps, self.batch_size, pre_units))
+ self.input_delay_buffer = Compartment(initial_buffer_state, display_name="Synaptic Input Queue")
+ self.delayed_inputs = Compartment(preVals) ## record of actual emitted delayed inputs (if delay > 0)
+
@compilable
def advance_state(self):
+ gate = self.gate.get()
weights = self.weights.get()
weights = weights * self.mask.get()
- self.outputs.set((jnp.matmul(self.inputs.get(), weights) * self.resist_scale) + self.biases.get())
+ raw_inputs = self.inputs.get()
+
+ inputs = raw_inputs
+ if self.max_delay_steps > 0: ## implements synaptic jitter via axonal delay
+ buffer = self.input_delay_buffer.get()
+ ## gather historical timestep slice, independently for each input axon line
+ time_indices = self.syn_delay_indices ## shape: (D_pre,)
+ pre_indices = jnp.arange(raw_inputs.shape[1]) ## shape: (D_pre,)
+ ## advanced gather outputs parallelized, jittered spike matrix (shape: (batch_size, D_pre))
+ inputs = buffer[time_indices, :, pre_indices].T ## get delay pre-synaptic signal(s)
+ ## roll input conveyor belt forward - shift historical slots down (& drop oldest timestep)
+ rolled_buffer = jnp.roll(buffer, shift=-1, axis=0)
+ ## update buffer - current input spikes 'raw_inputs' go to back of queue
+ updated_buffer = rolled_buffer.at[-1, :, :].set(raw_inputs)
+ self.input_delay_buffer.set(updated_buffer)
+ ## else, leave inputs = raw_inputs untouched
+
+ if 0. < self.p_release_mean < 1.: ## engage in stochastic synaptic transmission (in probability form)
+ ## Reference:
+ ## Del Castillo, J. and Katz, B., 1954. Quantal components of the end-plate potential.
+ ## The Journal of physiology, 124(3), p.560.
+ p_matrix = self.p_release ## get per-neuron release probs
+ key, skey = random.split(self.key.get(), 2) ## generate random Bernoulli mask
+ if not self.use_one_spike: ## does per-neuron sampled firing
+ release_mask = (random.uniform(skey, shape=raw_inputs.shape) < p_matrix).astype(jnp.float32)
+ else: ## does blind "guarantee-one-signal-fires" sampling
+ rP = raw_inputs * random.uniform(skey, raw_inputs.shape)
+ release_mask = nn.one_hot(jnp.argmax(rP, axis=1), num_classes=raw_inputs.shape[1], dtype=jnp.float32)
+ ## apply stochastic transmission: fully sparse, event-driven signals
+ inputs = raw_inputs * release_mask
+ self.key.set(key) ## update noise key of this component
+ ## else, leave inputs un-corrupted/untouched
+
+ self.delayed_inputs.set(inputs) ## store emitted delayed inputs
+ ## carry signals across synaptic cable
+ self.outputs.set((jnp.matmul(inputs, weights) * gate * self.g_conduct_factor) + self.biases.get())
@compilable
def reset(self):
if not self.inputs.targeted:
self.inputs.set(jnp.zeros((self.batch_size, self.shape[0])))
-
+ if not self.gate.targeted:
+ self.gate.set(jnp.ones((self.batch_size, self.shape[1])))
+ self.delayed_inputs.set(self.delayed_inputs.get() * 0)
self.outputs.set(jnp.zeros((self.batch_size, self.shape[1])))
+ self.input_delay_buffer.set(self.input_delay_buffer.get() * 0)
@classmethod
def help(cls): ## component help function
@@ -130,8 +208,10 @@ def help(cls): ## component help function
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level scaling factor (Rscale); applied to output of transformation",
- "p_conn": "Probability of a connection existing (otherwise, it is masked to zero)"
+ "g_conduct_factor": "Conductance/average level scaling factor; applied to output of transformation",
+ "p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
+ "p_release_mean": "Probability of pre-synaptic signal firing across axon and into synaptic cable line",
+ "max_delay_steps": "Maximum number of simulation steps to delay signal making it over axon & into cable line"
}
info = {cls.__name__: properties,
"compartments": compartment_props,
diff --git a/ngclearn/components/synapses/doubleExpSynapse.py b/ngclearn/components/synapses/doubleExpSynapse.py
index 91a05d60..c8e3bfce 100644
--- a/ngclearn/components/synapses/doubleExpSynapse.py
+++ b/ngclearn/components/synapses/doubleExpSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
@@ -49,8 +50,8 @@ class DoubleExpSynapse(DenseSynapse): ## dynamic double-exponential synapse cabl
bias_init: a kernel to drive initialization of biases for this synaptic cable
(Default: None, which turns off/disables biases)
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1 and > 0. will result in a sparser synaptic structure
@@ -60,11 +61,22 @@ class DoubleExpSynapse(DenseSynapse): ## dynamic double-exponential synapse cabl
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, tau_decay, tau_rise, g_syn_bar, syn_rest, weight_init=None, bias_init=None,
- resist_scale=1., p_conn=1., is_nonplastic=True, **kwargs
+ self, name,
+ shape,
+ tau_decay,
+ tau_rise,
+ g_syn_bar,
+ syn_rest,
+ weight_init=None,
+ bias_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ is_nonplastic=True,
+ **kwargs
):
- super().__init__(name, shape, weight_init, bias_init, resist_scale, p_conn, **kwargs)
+ super().__init__(name, shape, weight_init, bias_init, g_conduct_factor, p_conn, **kwargs)
## dynamic synapse meta-parameters
self.tau_decay = tau_decay
self.tau_rise = tau_rise
@@ -144,7 +156,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"tau_decay": "Conductance decay time constant (ms)",
"tau_rise": "Conductance rise/increase time constant (ms)",
diff --git a/ngclearn/components/synapses/exponentialSynapse.py b/ngclearn/components/synapses/exponentialSynapse.py
index dc20c362..ff0ef249 100644
--- a/ngclearn/components/synapses/exponentialSynapse.py
+++ b/ngclearn/components/synapses/exponentialSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
@@ -46,8 +47,8 @@ class ExponentialSynapse(DenseSynapse): ## dynamic exponential synapse cable
bias_init: a kernel to drive initialization of biases for this synaptic cable
(Default: None, which turns off/disables biases)
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1 and > 0. will result in a sparser synaptic structure
@@ -57,11 +58,21 @@ class ExponentialSynapse(DenseSynapse): ## dynamic exponential synapse cable
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, tau_decay, g_syn_bar, syn_rest, weight_init=None, bias_init=None, resist_scale=1.,
- p_conn=1., is_nonplastic=True, **kwargs
+ self, name,
+ shape,
+ tau_decay,
+ g_syn_bar,
+ syn_rest,
+ weight_init=None,
+ bias_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ is_nonplastic=True,
+ **kwargs
):
- super().__init__(name, shape, weight_init, bias_init, resist_scale, p_conn, **kwargs)
+ super().__init__(name, shape, weight_init, bias_init, g_conduct_factor, p_conn, **kwargs)
## dynamic synapse meta-parameters
self.tau_decay = tau_decay
self.g_syn_bar = g_syn_bar
@@ -85,10 +96,10 @@ def advance_state(self, t, dt):
dgsyn_dt = -self.g_syn.get()/self.tau_decay + (_out * self.g_syn_bar) * (1./dt)
g_syn = self.g_syn.get() + dgsyn_dt * dt ## run Euler step to move conductance
## compute derive electrical current variable
- i_syn = -g_syn * self.resist_scale
+ i_syn = -g_syn * self.g_conduct_factor
if self.syn_rest is not None:
- i_syn = -(g_syn * self.resist_scale) * (self.v.get() - self.syn_rest)
- outputs = i_syn #jnp.matmul(inputs, Wdyn * self.resist_scale) + biases
+ i_syn = -(g_syn * self.g_conduct_factor) * (self.v.get() - self.syn_rest)
+ outputs = i_syn #jnp.matmul(inputs, Wdyn * self.g_conduct_factor) + biases
self.outputs.set(outputs)
self.i_syn.set(i_syn)
@@ -130,7 +141,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"tau_decay": "Conductance decay time constant (ms)",
"g_bar_syn": "Maximum conductance value",
diff --git a/ngclearn/components/synapses/hebbian/BCMSynapse.py b/ngclearn/components/synapses/hebbian/BCMSynapse.py
index 088418a9..df32c24c 100755
--- a/ngclearn/components/synapses/hebbian/BCMSynapse.py
+++ b/ngclearn/components/synapses/hebbian/BCMSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
@@ -57,18 +58,32 @@ class BCMSynapse(DenseSynapse): # BCM-adjusted synaptic cable
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
- (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
+ (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1. will result in a sparser synaptic structure
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, tau_w, tau_theta, theta0=-1., w_bound=0., w_decay=0., weight_init=None, resist_scale=1.,
- p_conn=1., batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ tau_w,
+ tau_theta,
+ theta0=-1.,
+ w_bound=0.,
+ w_decay=0.,
+ weight_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ batch_size=1,
+ **kwargs
):
- super().__init__(name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs)
+ super().__init__(
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
+ )
## Synapse and BCM hyper-parameters
self.shape = shape ## shape of synaptic efficacy matrix
@@ -76,7 +91,7 @@ def __init__(
self.tau_theta = tau_theta ## time constant of threshold delta variables
self.w_decay = w_decay ## synaptic decay factor
self.w_bound = w_bound ## soft weight constraint
- self.Rscale = resist_scale ## post-transformation scale factor
+ self.Rscale = g_conduct_factor ## post-transformation scale factor
self.theta0 = theta0 #-1. ## initial condition for theta/threshold variables
## Compartment setup
@@ -157,7 +172,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"tau_theta": "Time constant for synaptic threshold variable `theta`",
"tau_w": "Time constant for BCM synaptic adjustment",
diff --git a/ngclearn/components/synapses/hebbian/STDPSynapse.py b/ngclearn/components/synapses/hebbian/STDPSynapse.py
index b16df8db..f5fa8403 100755
--- a/ngclearn/components/synapses/hebbian/STDPSynapse.py
+++ b/ngclearn/components/synapses/hebbian/STDPSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses import DenseSynapse
@@ -53,7 +54,7 @@ class STDPSynapse(DenseSynapse): # power-law / trace-based STDP
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
(Default: 1.), i.e., yields: out = ((W * Rscale) * in)
p_conn: probability of a connection existing (default: 1); setting
@@ -62,11 +63,26 @@ class STDPSynapse(DenseSynapse): # power-law / trace-based STDP
w_bound: maximum value/magnitude any synaptic efficacy can be (default: 1)
"""
- # Define Functions
- def __init__(self, name, shape, A_plus, A_minus, tau_plus=10., tau_minus=10., w_decay=0.,
- eta=1., tau_w=0., weight_init=None, resist_scale=1., p_conn=1., w_bound=1.,
- batch_size=1, **kwargs):
- super().__init__(name, shape, weight_init, None, resist_scale,
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
+ def __init__(
+ self,
+ name,
+ shape,
+ A_plus,
+ A_minus,
+ au_plus=10.,
+ tau_minus=10.,
+ w_decay=0.,
+ eta=1.,
+ tau_w=0.,
+ weight_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ w_bound=1.,
+ batch_size=1,
+ **kwargs
+ ):
+ super().__init__(name, shape, weight_init, None, g_conduct_factor,
p_conn, batch_size=batch_size, **kwargs)
assert self.batch_size == 1 ## note: STDP only supports online learning in this implementation
## Synaptic hyper-parameters
@@ -75,7 +91,7 @@ def __init__(self, name, shape, A_plus, A_minus, tau_plus=10., tau_minus=10., w_
self.Aminus = A_minus ## LTD strength
self.tau_plus = tau_plus ## LTP time constant
self.tau_minus = tau_minus ## LTD time constant
- self.Rscale = resist_scale ## post-transformation scale factor
+ self.Rscale = g_conduct_factor ## post-transformation scale factor
self.w_bound = w_bound #1. ## soft weight constraint
self.tau_w = tau_w ## synaptic update time constant
self.w_decay = w_decay
@@ -192,7 +208,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"A_plus": "Strength of long-term potentiation (LTP)",
"A_minus": "Strength of long-term depression (LTD)",
diff --git a/ngclearn/components/synapses/hebbian/__init__.py b/ngclearn/components/synapses/hebbian/__init__.py
index 0a1630c3..1a27c09a 100644
--- a/ngclearn/components/synapses/hebbian/__init__.py
+++ b/ngclearn/components/synapses/hebbian/__init__.py
@@ -5,4 +5,5 @@
from .eventSTDPSynapse import EventSTDPSynapse
from .BCMSynapse import BCMSynapse
from .gerstnerHebbianSynapse import GerstnerHebbianSynapse ## Taylor-expansion Hebbian model
+from .inhibitorySTDPSynapse import InhibitorySTDPSynapse ## iSTDP rule (trace-based)
diff --git a/ngclearn/components/synapses/hebbian/eventSTDPSynapse.py b/ngclearn/components/synapses/hebbian/eventSTDPSynapse.py
index 826b9ff9..cc3f2e09 100755
--- a/ngclearn/components/synapses/hebbian/eventSTDPSynapse.py
+++ b/ngclearn/components/synapses/hebbian/eventSTDPSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.denseSynapse import DenseSynapse
@@ -46,18 +47,33 @@ class EventSTDPSynapse(DenseSynapse): # event-driven, post-synaptic STDP
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
(Default: 1), i.e., yields: out = ((W * Rscale) * in) + b
p_conn: probability of a connection existing (default: 1.); setting
this to < 1. will result in a sparser synaptic structure
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, eta, lmbda=0.01, A_plus=1., A_minus=1., presyn_win_len=2., w_bound=1.,
- weight_init=None, resist_scale=1., p_conn=1., batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ eta,
+ lmbda=0.01,
+ A_plus=1.,
+ A_minus=1.,
+ presyn_win_len=2.,
+ w_bound=1.,
+ weight_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ batch_size=1,
+ **kwargs
):
- super().__init__(name, shape, weight_init, None, resist_scale, p_conn, batch_size=batch_size, **kwargs)
+ super().__init__(
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
+ )
## Synaptic hyper-parameters
self.eta = eta ## global learning rate governing plasticity
@@ -66,7 +82,7 @@ def __init__(
assert self.presyn_win_len >= 0. ## pre-synaptic window must be non-negative
self.Aplus = A_plus
self.Aminus = A_minus
- self.Rscale = resist_scale ## post-transformation scale factor
+ self.Rscale = g_conduct_factor ## post-transformation scale factor
self.w_bound = w_bound ## soft weight constraint
## Compartment setup
@@ -143,7 +159,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"lmbda": "Degree of synaptic disconnect",
"eta": "Global learning rate (multiplier beyond A_plus and A_minus)",
diff --git a/ngclearn/components/synapses/hebbian/expSTDPSynapse.py b/ngclearn/components/synapses/hebbian/expSTDPSynapse.py
index bb481512..69bf4db6 100644
--- a/ngclearn/components/synapses/hebbian/expSTDPSynapse.py
+++ b/ngclearn/components/synapses/hebbian/expSTDPSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.denseSynapse import DenseSynapse
@@ -51,7 +52,7 @@ class ExpSTDPSynapse(DenseSynapse):
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling (resistance) factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling (resistance) factor to apply to synaptic transform
(Default: 1.), i.e., yields: out = ((W * Rscale) * in) + b
p_conn: probability of a connection existing (default: 1.); setting
@@ -64,11 +65,27 @@ class ExpSTDPSynapse(DenseSynapse):
weight_mask: synaptic binary masking matrix to apply (to enforce a constant sparse structure; default: None)
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, A_plus, A_minus, exp_beta, eta=1., pretrace_target=0., weight_init=None, resist_scale=1.,
- p_conn=1., w_bound=1., tau_w=0., weight_mask=None, batch_size=1, **kwargs
+ self, name,
+ shape,
+ A_plus,
+ A_minus,
+ exp_beta,
+ eta=1.,
+ pretrace_target=0.,
+ weight_init=None,
+ g_conduct_factor=1.,
+ 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)
+ super().__init__(
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
+ )
self.tau_w = tau_w
## Exp-STDP meta-parameters
@@ -78,7 +95,7 @@ def __init__(
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.Rscale = g_conduct_factor ## post-transformation scale factor
self.w_bound = w_bound #1. ## soft weight constraint
if weight_mask is None:
@@ -179,7 +196,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"A_plus": "Strength of long-term potentiation (LTP)",
"A_minus": "Strength of long-term depression (LTD)",
diff --git a/ngclearn/components/synapses/hebbian/gerstnerHebbianSynapse.py b/ngclearn/components/synapses/hebbian/gerstnerHebbianSynapse.py
index d1efca5a..a9a66c96 100644
--- a/ngclearn/components/synapses/hebbian/gerstnerHebbianSynapse.py
+++ b/ngclearn/components/synapses/hebbian/gerstnerHebbianSynapse.py
@@ -4,25 +4,52 @@
from ngclearn import compilable
from ngclearn import Compartment
from ngclearn.components.synapses import DenseSynapse
-from ngclearn.utils import tensorstats
-from ngcsimlib import deprecate_args
-#from ngclearn.utils.io_utils import save_pkl, load_pkl
class GerstnerHebbianSynapse(DenseSynapse):
"""
- A synapse component that implements Gerstner's general Hebbian
- learning (Taylor) expansion (Equation 3 from Gerstner & Kistler, 2002).
+ A synapse component that implements Gerstner's general Hebbian learning (Taylor) expansion (Equation 3 from
+ Gerstner & Kistler, 2002).
- Note that this synpatic update model can recover several classical forms
+ Note that this synaptic update model can recover several classical forms
of Hebbian-like update rules, including the covariance rule.
- There are other higher-order terms possible, i.e., \Theta(xy), such as
+ There are other higher-order terms possible, i.e., Theta(xy), such as
x * y2 and y x^2, etc.
| c2_corr > 0 and c0 = c1_pre = c1_post = 0 => Hebbian update
| c2_corr < 0 and c0 = c1_pre = c1_post = 0 => anti-Hebbian update
| c2_corr = 1 and c1_pre = -x_theta < 0
+ | References:
+ | Gerstner, W. and Kistler, W.M., 2002. Mathematical formulations of Hebbian
+ | learning. Biological cybernetics, 87(5), pp.404-415.
+
+ Args:
+ name: the string name of this cell
+
+ shape: tuple specifying shape of this synaptic cable (usually a 2-tuple
+ with number of inputs by number of outputs)
+
+ eta: global learning rate
+
+ coeffs: dictionary containing relevant (meta-)parameter coefficients for this hebbian-synapse;
+ this dictionary must contain/specificy `c0`, `c1_pre`, `c1_post`, `c2_corr`
+
+ weight_init: a kernel to drive initialization of this synaptic cable's values;
+ typically a tuple with 1st element as a string calling the name of
+ initialization to use
+
+ sign_value: multiplicative factor to apply to final synaptic update before
+ it is applied to synapses; this is useful if gradient descent style
+ optimization is required (as Hebbian rules typically yield
+ adjustments for ascent)
+
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
+ (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in) + b
+
+ p_conn: probability of a connection existing (default: 1.); setting
+ this to < 1. will result in a sparser synaptic structure
+
"""
def __init__(
self,
@@ -32,27 +59,27 @@ def __init__(
coeffs=None, ## these configure which kind of Hebb learning is done
weight_init=None,
p_conn=1.,
- resist_scale=1.,
+ g_conduct_factor=1.,
sign_value=1.,
batch_size=1,
**kwargs
):
- bias_init = None ## no biases are included in Gerster's formulation
+ bias_init = None ## NOTE: no biases are included in Gerster's formulation
super().__init__(
name,
shape=shape,
weight_init=weight_init,
bias_init=bias_init,
- resist_scale=resist_scale,
+ g_conduct_factor=g_conduct_factor,
p_conn=p_conn,
batch_size=batch_size,
**kwargs
)
- ## General Hebbian meta-parameters
+ ## general Hebbian meta-parameters
self.eta = eta
self.sign_value = sign_value
- ## Expansion coefficients (c0, c1_pre, c1_post, c2_corr)
+ ## Gerstner and Kisler's expansion coefficients (c0, c1_pre, c1_post, c2_corr)
if coeffs is None: ## Default to standard bilinear Hebb
self.coeffs = {
'c0': 0., 'c1_pre': 0., 'c1_post': 0., 'c2_corr': 1.0
@@ -64,50 +91,43 @@ def __init__(
self.c1_post = self.coeffs['c1_post']
self.c2_corr = self.coeffs['c2_corr']
- # Initialize Weights (using JAX PRNG)
- #init_key, _ = random.split(self.key)
- #w_init = random.normal(init_key, shape) * 0.05
-
- # Compartments (ngc-learn state management)
- #self.weights = Compartment(w_init)
+ ## set up relevant compartments
self.pre = Compartment(jnp.zeros((1, shape[1])))
self.post = Compartment(jnp.zeros((1, shape[0])))
+ self.dWeights = Compartment(jnp.zeros(shape))
@compilable
- def evolve(self, **kwargs):
- """
- Updates weights using the Gerstner general expansion.
- Assumes pre_act and post_act compartments have been populated.
- """
- # Retrieve current states
+ def evolve(self, **kwargs): ## perform update via Gerstner's general expansion
+ ## retrieve current compartment state values
W = self.weights.get()
- x = self.pre.get() # pre-synaptic activity (batch, pre_dim)
- y = self.post.get() # post-synaptic activity (batch, post_dim)
+ x = self.pre.get() ## pre-synaptic activity (batch, pre_dim)
+ y = self.post.get() ## post-synaptic activity (batch, post_dim)
batch_size = self.batch_size
- ## Bilinear Term (c2): correlation matrix
- ### (post_dim, batch) @ (batch, pre_dim) -> (post_dim, pre_dim)
+ ## calculate bilinear Term (c2), i.e., correlation matrix
+ ### (pre_dim, batch) @ (batch, post_dim) -> (pre_dim, post_dim)
dW_corr = jnp.matmul(x.T, y) * (1./batch_size)
- ## Linear pre-synaptic term (c1_pre)
- ### Average over batch then broadcast to match weight matrix
+ ## linear pre-synaptic term (c1_pre)
+ ### get mean over batch then broadcast to match weight matrix
dW_pre = jnp.sum(x, axis=0, keepdims=True).T * (1./batch_size)
- ## Linear post-synaptic term (c1_post)
- dW_post = jnp.sum(y, axis=0, keepdims=True) * (1./batch_size)
+ ## linear post-synaptic term (c1_post), mean over post-syn values
+ dW_post = jnp.sum(y, axis=0, keepdims=True) * (1./batch_size)
- ## Apply Equation 3 Taylor expansion
+ ## apply Taylor expansion from Equation 3 (Gerstner and Kistler)
dW = (self.c0 * W + ## synaptic decay
self.c1_pre * dW_pre + ## bilinear term
self.c1_post * dW_post + ## pre-synaptic gating term
self.c2_corr * dW_corr ## post-synpatic gating term
)
+ self.dWeights.set(dW)
+
## perform a step of Hebbian ascent
- W = W + self.eta * dW
- ## Update weights
+ W = W + self.eta * dW ## update synaptic efficacies
self.weights.set(W)
@compilable
- def reset(self, **kwargs):
- """Clears activity compartments"""
+ def reset(self, **kwargs): ## clear compartment values
self.pre.set( jnp.zeros((self.batch_size, self.shape[1])) )
self.post.set( jnp.zeros((self.batch_size, self.shape[0])) )
+ self.dWeights.set(self.dWeights.get() * 0)
diff --git a/ngclearn/components/synapses/hebbian/hebbianSynapse.py b/ngclearn/components/synapses/hebbian/hebbianSynapse.py
index d558e46c..c161a238 100644
--- a/ngclearn/components/synapses/hebbian/hebbianSynapse.py
+++ b/ngclearn/components/synapses/hebbian/hebbianSynapse.py
@@ -160,22 +160,37 @@ class HebbianSynapse(DenseSynapse):
post_wght: post-synaptic weighting factor (Default: 1.)
- resist_scale: a fixed scaling factor to apply to synaptic transform
- (Default: 1.), i.e., yields: out = ((W * Rscale) * in) + b
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
+ (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in) + b
p_conn: probability of a connection existing (default: 1.); setting
this to < 1. will result in a sparser synaptic structure
"""
- @deprecate_args(_rebind=False, w_decay='prior')
+ @deprecate_args(_rebind=True, w_decay='prior', resist_scale='g_conduct_factor')
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
+ 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.,
+ g_conduct_factor=1.,
+ batch_size=1,
+ **kwargs
):
super().__init__(
- name, shape=shape, weight_init=weight_init, bias_init=bias_init, resist_scale=resist_scale, p_conn=p_conn,
- batch_size=batch_size, **kwargs
+ name, shape=shape, weight_init=weight_init, bias_init=bias_init, g_conduct_factor=g_conduct_factor,
+ p_conn=p_conn, batch_size=batch_size, **kwargs
)
if w_decay > 0.:
@@ -186,7 +201,7 @@ def __init__(
prior_type = "constant"
## synaptic plasticity properties and characteristics
self.shape = shape
- self.Rscale = resist_scale
+ self.Rscale = g_conduct_factor
self.prior_type = prior_type
if self.prior_type.lower() == "gaussian":
self.prior_type = "ridge"
@@ -338,7 +353,7 @@ def help(cls): ## component help function
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"bias_init": "Initialization conditions for bias/base-rate (b) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"is_nonnegative": "Should synapses be constrained to be non-negative post-updates?",
"sign_value": "Scalar `flipping` constant -- changes direction to Hebbian descent if < 0",
diff --git a/ngclearn/components/synapses/hebbian/inhibitorySTDPSynapse.py b/ngclearn/components/synapses/hebbian/inhibitorySTDPSynapse.py
new file mode 100644
index 00000000..25d49468
--- /dev/null
+++ b/ngclearn/components/synapses/hebbian/inhibitorySTDPSynapse.py
@@ -0,0 +1,255 @@
+from jax import random, numpy as jnp
+from ngcsimlib.logger import info
+from ngcsimlib import deprecate_args
+from ngclearn.utils.distribution_generator import DistributionGenerator
+from ngclearn import compilable
+from ngclearn import Compartment
+from ngclearn.components.synapses import DenseSynapse
+
+class InhibitorySTDPSynapse(DenseSynapse): ## inhibitory-STDP synaptic cable
+ """
+ A synaptic cable that adjusts its efficacies via the trace-based, inhibitory
+ spike-timing-dependent plasticity (iSTDP); this rule can also be configured to
+ utilize a voltage-dependent format of Foldiak anti-Hebbian plasticity.
+
+ | --- Synapse Compartments: ---
+ | inputs - input (takes in external signals)
+ | outputs - output signals (transformation induced by synapses)
+ | weights - current value matrix of synaptic efficacies
+ | key - JAX PRNG key
+ | --- Synaptic Plasticity Compartments: ---
+ | s_pre - pre-synaptic spike to drive 1st term of STDP update (takes in external signals)
+ | s_post - post-synaptic spike to drive 2nd term of STDP update (takes in external signals)
+ | x_pre - pre-synaptic trace value to drive 1st term of STDP update (takes in external signals)
+ | x_post - post-synaptic trace value to drive 2nd term of STDP update (takes in external signals)
+ | v_post - post-synaptic voltage/membrane potential state to drive voltage-dependent Foldiak rule (optional)
+ | dWeights - current delta matrix containing changes to be applied to synaptic efficacies
+ | eta - global learning rate (multiplier beyond A_plus and A_minus)
+
+ | References:
+ | Vogels, T.P., Sprekeler, H., Zenke, F., Clopath, C. and Gerstner, W., 2011. Inhibitory
+ | plasticity balances excitation and inhibition in sensory pathways and memory
+ | networks. Science, 334(6062), pp.1569-1573.
+ |
+ | Földiak, P., 1990. Forming sparse representations by local anti-Hebbian learning. Biological
+ | cybernetics, 64(2), pp.165-170.
+
+ Args:
+ name: the string name of this cell
+
+ shape: tuple specifying shape of this synaptic cable (usually a 2-tuple with number of
+ inputs by number of outputs)
+
+ A_plus: strength of long-term potentiation (LTP)
+
+ A_minus: strength of long-term depression (LTD)
+
+ eta: global learning rate initial value/condition (Default: 1)
+
+ rho: controls target firing rate value (Default: 0.2)
+
+ tau_x_pre: pre-synaptic trace time constant (ms) (Default: 30 ms)
+
+ tau_x_post: post-synaptic trace time constant (ms) (Default: 30 ms)
+
+ use_soft_bounds: trigggers weight-dependent soft-bounding variant of iSTDP/Foldiak
+ plasticity (Default: False)
+
+ w_max: maximal synaptic efficacy allowed (hard upper bound; Default: 1)
+
+ w_min: minimal synaptic efficacy allowed (hard lower bound; Default: 0)
+
+ is_voltage_dependent: if True, this synaptic cables adapts via voltage-dependent Foldiak
+ plasticity (Default: False)
+
+ weight_init: a kernel to drive initialization of this synaptic cable's values;
+ typically a tuple with 1st element as a string calling the name of initialization to use
+
+ bias_init: a kernel to drive initialization to this cable's fixed bias/shift values
+
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
+ (Default: 1.), i.e., yields: out = ((W * mask) * in) * g_conduct_factor
+
+ p_conn: probability of a connection existing (default: 1); setting this to < 1. will
+ result in a sparser synaptic structure
+
+ weight_mask: synaptic binary masking matrix to apply (to enforce a constant sparse
+ structure; default: None)
+
+ """
+
+ @deprecate_args(_rebind=True, w_bound='w_max')
+ def __init__(
+ self,
+ name,
+ shape,
+ eta=1.,
+ rho=0.2, ## target/desired baseline activity level
+ Aplus=1., ## weight on LTP term
+ Aminus=1., ## weight on LTD term
+ tau_x_pre=30.,
+ tau_x_post=30.,
+ use_soft_bounds=True,
+ w_max=1.,
+ w_min=0.,
+ is_volt_dependent=False, ## triggers voltage-dependent format
+ weight_init=None,
+ bias_init=None,
+ g_conduct_factor=1.,
+ p_release_mean=1.0,
+ p_conn=1.,
+ mask=None, ## None input -> triggers default mask=1
+ batch_size=1,
+ **kwargs
+ ):
+ super().__init__(
+ name,
+ shape=shape,
+ weight_init=weight_init,
+ bias_init=bias_init,
+ g_conduct_factor=g_conduct_factor,
+ p_release_mean=p_release_mean,
+ p_conn=p_conn,
+ batch_size=batch_size,
+ mask=mask,
+ **kwargs
+ )
+ ## key anti-Hebbian (iSTDP) synapse meta-parameters
+ self.eta = eta
+ self.Aplus = Aplus # 5-10 -> scales up rare coincidences in LTP
+ self.Aminus = Aminus
+ self.use_soft_bounds = use_soft_bounds
+ self.tau_x_pre = tau_x_pre
+ self.tau_x_post = tau_x_post
+ self.rho = rho
+ self.w_max = w_max
+ self.w_min = w_min
+ self.is_volt_dependent = is_volt_dependent
+
+ ## set up anti-Hebbian synapse key (extra) compartments
+ _pre_reset = jnp.zeros((batch_size, shape[0]))
+ _post_reset = jnp.zeros((batch_size, shape[1]))
+ self.s_pre = Compartment(_pre_reset) ## input compartment
+ self.x_pre = Compartment(_pre_reset) ## internal
+ self.s_post = Compartment(_post_reset) ## input compartment
+ self.x_post = Compartment(_post_reset) ## internal
+ self.v_post = Compartment(_post_reset) ## input compartment (optional)
+
+ self.dWeights = Compartment(jnp.zeros(shape))
+ self.weights.set(self.weights.get() * self.mask.get()) ## make sure mask is enforced
+
+ @compilable
+ def evolve(self, t, dt): ## NOTE: spike-based anti-Hebbian rule
+ W = self.weights.get()
+ s_pre = self.s_pre.get() ## pre-synaptic inhibitory spikes ## NOTE: use delayed inputs?
+ s_post = self.s_post.get()
+ x_pre = self.x_pre.get() ## filtered inhibitory spikes
+ x_post = self.x_post.get() ## post-synaptic target's spikes
+ v_post = self.v_post.get() ## post-synaptic target's voltage
+
+ ## (low-pass) filter synaptic spikes over time
+ x_pre = x_pre + (-x_pre + s_pre) * dt/self.tau_x_pre ## pre-trace
+ self.x_pre.set(x_pre)
+ x_post = x_post + (-x_post + s_post) * dt/self.tau_x_post ## post-trace
+ self.x_post.set(x_post)
+
+ batch_size = x_pre.shape[0]
+ if self.is_volt_dependent: ## trigger voltage-dependent foldiak anti-Hebbian rule
+ ## this simple NAC lab rule assumes:
+ ### W - excitatory-to-inhibitory synaptic efficacies
+ ### x_pre - inhibitory spike trace (shape = batch_size x num_inh)
+ ### v_post - excitatory voltage (shape = batch_size x num_exc)
+ ### rho - target voltage rate
+ ## compute (scaled) synaptic adjustment
+ dW = (x_pre.T @ (v_post - self.rho)) / batch_size
+ bound_scale = jnp.where(dW > 0, self.w_max - W, W)
+ dW = dW * bound_scale
+ else:
+ ## Vogels-Sprekeler rule iSTDP assumes:
+ ### W - excitatory-to-inhibitory synaptic efficacies
+ ### x_pre - inhibitory spike trace (shape = batch_size x num_inh)
+ ### s_pre - raw inhibitory spikes (shape = batch_size x num_inh)
+ ### x_post - excitatory spike trace (shape = batch_size x num_exc)
+ ### s_post - raw excitatory spikes (shape = batch_size x num_exc)
+ ### rho - target firing rate fraction
+
+ ## calculate Vogels-Sprekeler coincidence + homeostatic matrices
+ potentiation = ( (x_pre.T @ s_post) + (s_pre.T @ x_post) ) ## coincidence matrix shape = num_inh x num_exc
+ ## target-rate suppression matrix (broadcasted across excitatory dimension)
+ ### this depends on x_pre (scales depression based on recent inhibitory activity)
+ ### 2 * rho * x_pre(t) (where ltd_bias = 2)
+ ltd_bias = 2.0 ## NOTE: ltd bias value could be increased if needed
+ depression = (x_pre.T @ jnp.ones_like(s_post)) * (ltd_bias * self.rho)
+ ## compute raw weight change averaged over batch-length
+ ltp = potentiation
+ ltd = -depression
+ if self.use_soft_bounds: ## apply weight-dependency (NAC-lab extension)
+ ltp = potentiation * (self.w_max - W)
+ ltd = -depression * W
+ dW = (ltp * self.Aplus + ltd * self.Aminus) / batch_size
+
+ ## apply update rule to adjust synaptic efficacies
+ W = W + dW * self.eta
+ W = jnp.clip(W, self.w_min, self.w_max) ## constrains W to stay w/in bounds
+ self.weights.set(W * self.mask.get())
+ self.dWeights.set(dW)
+
+ @compilable
+ def reset(self):
+ in_reset_vals = jnp.zeros((self.batch_size, self.shape[0]))
+ out_reset_vals = jnp.zeros((self.batch_size, self.shape[1]))
+ if not self.inputs.targeted:
+ self.inputs.set(in_reset_vals)
+ self.x_pre.set(in_reset_vals)
+ self.s_pre.set(in_reset_vals)
+ self.outputs.set(out_reset_vals)
+ self.s_post.set(out_reset_vals)
+ self.x_post.set(out_reset_vals)
+ if not self.v_post.targeted:
+ self.v_post.set(out_reset_vals)
+ self.dWeights.set(self.dWeights.get() * 0)
+
+ @classmethod
+ def help(cls): ## component help function
+ properties = {
+ "synapse_type": "InhibitorySTDPSynapse - performs an adaptable synaptic "
+ "transformation of inputs to produce output signals; "
+ "synapses are adjusted with trace-based "
+ "inhibitory spike-timing-dependent plasticity (iSTDP)"
+ }
+ compartment_props = {
+ "inputs":
+ {"inputs": "Takes in external input signal values",
+ "s_pre": "Pre-synaptic spike compartment value/term for iSTDP (s_j)",
+ "s_post": "Post-synaptic spike compartment value/term for iSTDP (s_i)",
+ "v_post": "Post-synaptic voltage value term for Foldiak anti-Hebbian format (v_i)"},
+ "states":
+ {"weights": "Synapse efficacy/strength parameter values",
+ "biases": "Base-rate/bias parameter values",
+ "eta": "Global learning rate (multiplier beyond A_plus and A_minus)",
+ "x_pre": "Pre-synaptic trace value term for iSTDP (z_j)",
+ "x_post": "Post-synaptic trace value term for iSTDP (z_i)",
+ "key": "JAX PRNG key"},
+ "analytics":
+ {"dWeights": "Synaptic weight value adjustment matrix produced at time t"},
+ "outputs":
+ {"outputs": "Output of synaptic transformation"},
+ }
+ hyperparams = {
+ "shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
+ "batch_size": "Batch size dimension of this component",
+ "weight_init": "Initialization conditions for synaptic weight (W) values",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
+ "p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
+ "A_plus": "Strength of long-term potentiation (LTP)",
+ "A_minus": "Strength of long-term depression (LTD)",
+ "eta": "Global learning rate initial condition",
+ "weight_mask" : "Binary synaptic weight mask to apply to enforce a sparsity structure"
+ }
+ info = {cls.__name__: properties,
+ "compartments": compartment_props,
+ "dynamics": "outputs = [(W * g_conduct_factor) * inputs] ;"
+ "dW_{ij}/dt = A_plus * (z_j * s_i + s_j * z_i) - A_minus * (2 * rho * z_j)",
+ "hyperparameters": hyperparams}
+ return info
+
diff --git a/ngclearn/components/synapses/hebbian/traceSTDPSynapse.py b/ngclearn/components/synapses/hebbian/traceSTDPSynapse.py
index 1c7ac3ab..2cdae169 100755
--- a/ngclearn/components/synapses/hebbian/traceSTDPSynapse.py
+++ b/ngclearn/components/synapses/hebbian/traceSTDPSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
from ngclearn.components.synapses.denseSynapse import DenseSynapse
@@ -53,7 +54,7 @@ class TraceSTDPSynapse(DenseSynapse): # power-law / trace-based STDP
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
(Default: 1.), i.e., yields: out = ((W * resistance) * in)
p_conn: probability of a connection existing (default: 1); setting
@@ -66,11 +67,27 @@ class TraceSTDPSynapse(DenseSynapse): # power-law / trace-based STDP
weight_mask: synaptic binary masking matrix to apply (to enforce a constant sparse structure; default: None)
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
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., tau_w=0., weight_mask=None, batch_size=1, **kwargs
+ self, name,
+ shape,
+ A_plus,
+ A_minus,
+ eta=1.,
+ mu=0.,
+ pretrace_target=0.,
+ weight_init=None,
+ g_conduct_factor=1.,
+ 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)
+ super().__init__(
+ name, shape, weight_init, None, g_conduct_factor, p_conn, batch_size=batch_size, **kwargs
+ )
self.tau_w = tau_w
self.mu = mu ## controls power-scaling of STDP rule
@@ -176,7 +193,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"batch_size": "Batch size dimension of this component",
"weight_init": "Initialization conditions for synaptic weight (W) values",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"A_plus": "Strength of long-term potentiation (LTP)",
"A_minus": "Strength of long-term depression (LTD)",
diff --git a/ngclearn/components/synapses/modulated/MSTDPETSynapse.py b/ngclearn/components/synapses/modulated/MSTDPETSynapse.py
index 150ebc9b..b979a5c8 100755
--- a/ngclearn/components/synapses/modulated/MSTDPETSynapse.py
+++ b/ngclearn/components/synapses/modulated/MSTDPETSynapse.py
@@ -1,4 +1,5 @@
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable #from ngcsimlib.parser import compilable
from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
@@ -63,7 +64,7 @@ class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligilit
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
(Default: 1.), i.e., yields: out = ((W * Rscale) * in)
p_conn: probability of a connection existing (default: 1.); setting
@@ -72,13 +73,29 @@ class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligilit
w_bound: maximum value/magnitude any synaptic efficacy can be (default: 1)
"""
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, A_plus, A_minus, eta=1., mu=0., pretrace_target=0., tau_elg=0., elg_decay=1.,
- tau_w=0., weight_init=None, resist_scale=1., p_conn=1., w_bound=1., batch_size=1, **kwargs
+ self,
+ name,
+ shape,
+ A_plus,
+ A_minus,
+ eta=1.,
+ mu=0.,
+ pretrace_target=0.,
+ tau_elg=0.,
+ elg_decay=1.,
+ tau_w=0.,
+ weight_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ w_bound=1.,
+ batch_size=1,
+ **kwargs
):
super().__init__( # call to parent trace-stdp component
name, shape, A_plus, A_minus, eta=eta, mu=mu, pretrace_target=pretrace_target, weight_init=weight_init,
- resist_scale=resist_scale, p_conn=p_conn, w_bound=w_bound, batch_size=batch_size, **kwargs
+ g_conduct_factor=g_conduct_factor, p_conn=p_conn, w_bound=w_bound, batch_size=batch_size, **kwargs
)
self.w_eps = 0.
self.tau_w = tau_w
@@ -164,7 +181,7 @@ def help(cls): ## component help function
"shape": "Shape of synaptic weight value matrix; number inputs x number outputs",
"weight_init": "Initialization conditions for synaptic weight (W) values",
"batch_size": "Batch size dimension of this component",
- "resist_scale": "Resistance level scaling factor (applied to output of transformation)",
+ "g_conduct_factor": "Conductance level scaling factor (applied to output of transformation)",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"A_plus": "Strength of long-term potentiation (LTP)",
"A_minus": "Strength of long-term depression (LTD)",
diff --git a/ngclearn/components/synapses/modulated/REINFORCESynapse.py b/ngclearn/components/synapses/modulated/REINFORCESynapse.py
index 9219e930..f9e9209e 100644
--- a/ngclearn/components/synapses/modulated/REINFORCESynapse.py
+++ b/ngclearn/components/synapses/modulated/REINFORCESynapse.py
@@ -1,6 +1,7 @@
# %%
from jax import random, numpy as jnp, jit
+from ngcsimlib import deprecate_args
from ngclearn import compilable, Compartment
from ngclearn.utils.model_utils import clip, d_clip
@@ -11,60 +12,6 @@
from ngclearn.utils import tensorstats
from ngclearn.utils.model_utils import create_function
-def _gaussian_logpdf(event, mean, stddev):
- scale_sqrd = stddev ** 2
- log_normalizer = jnp.log(2 * jnp.pi * scale_sqrd)
- quadratic = (event - mean)**2 / scale_sqrd
- return - 0.5 * (log_normalizer + quadratic)
-
-
-def _compute_update(
- dt, inputs, rewards, act_fx, weights, seed, mu_act_fx, dmu_act_fx, mu_out_min, mu_out_max, scalar_stddev
-):
- learning_stddev_mask = jnp.asarray(scalar_stddev <= 0.0, dtype=jnp.float32)
- # (input_dim, output_dim * 2) => (input_dim, output_dim), (input_dim, output_dim)
- W_mu, W_logstd = jnp.split(weights, 2, axis=-1)
- # Forward pass
- activation = act_fx(inputs)
- mean = activation @ W_mu
- fx_mean = mu_act_fx(mean)
- logstd = activation @ W_logstd
- clip_logstd = clip(logstd, -10.0, 2.0)
- std = jnp.exp(clip_logstd)
- std = learning_stddev_mask * std + (1.0 - learning_stddev_mask) * scalar_stddev # masking trick
- # Sample using reparameterization trick
- epsilon = jax.random.normal(seed, fx_mean.shape)
- sample = epsilon * std + fx_mean
- sample = jnp.clip(sample, mu_out_min, mu_out_max)
- outputs = sample # the actual action that we take
- # Compute log probability density of the Gaussian
- log_prob = _gaussian_logpdf(sample, fx_mean, std).sum(-1)
- # Compute objective (negative REINFORCE objective)
- objective = (-log_prob * rewards).mean() * 1e-2
-
- # Backward pass
- batch_size = inputs.shape[0] # B
- dL_dlogp = -rewards[:, None] * 1e-2 / batch_size # (B, 1)
-
- # Compute gradients manually based on the derivation
- # dL/dmu = -(r-r_hat) * dlog_prob/dmu = -(r-r_hat) * -(sample-mu)/sigma^2
- dlog_prob_dfxmean = (sample - fx_mean) / (std ** 2)
- dL_dmean = dL_dlogp * dlog_prob_dfxmean * dmu_act_fx(mean) # (B, A)
- dL_dWmu = activation.T @ dL_dmean
-
- # dL/dlog(sigma) = -(r-r_hat) * dlog_prob/dlog(sigma) = -(r-r_hat) * (((sample-mu)/sigma)^2 - 1)
- dlog_prob_dlogstd = - 1.0 / std + (sample - fx_mean)**2 / std**3
- dL_dstd = dL_dlogp * dlog_prob_dlogstd
- # Apply gradient clipping for logstd
- dL_dlogstd = d_clip(logstd, -10.0, 2.0) * dL_dstd * std
- dL_dWlogstd = activation.T @ dL_dlogstd # (I, B) @ (B, A) = (I, A)
- dL_dWlogstd = dL_dWlogstd * learning_stddev_mask # there is no learning for the scalar stddev
-
- # Update weights, negate the gradient because gradient ascent in ngc-learn
- dW = jnp.concatenate([-dL_dWmu, -dL_dWlogstd], axis=-1)
- # Finally, return metrics if needed
- return dW, objective, outputs
-
class REINFORCESynapse(DenseSynapse):
"""
@@ -97,7 +44,7 @@ class REINFORCESynapse(DenseSynapse):
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed scaling factor to apply to synaptic transform
+ g_conduct_factor: a fixed scaling factor to apply to synaptic transform
(Default: 1.)
act_fx: activation function to apply to inputs (Default: "identity")
@@ -114,22 +61,36 @@ class REINFORCESynapse(DenseSynapse):
mu_act_fx: activation function to apply to the mean of the Gaussian distribution (Default: "identity")
"""
- # Define Functions
+ @deprecate_args(_rebind=True, resist_scale='g_conduct_factor')
def __init__(
- self, name, shape, eta=1e-4, decay=0.99, weight_init=None, resist_scale=1., act_fx=None,
- p_conn=1., w_bound=1., batch_size=1, seed=None, mu_act_fx=None, mu_out_min=-jnp.inf, mu_out_max=jnp.inf,
- scalar_stddev=-1.0, **kwargs
+ self,
+ name,
+ shape,
+ eta=1e-4,
+ decay=0.99,
+ weight_init=None,
+ g_conduct_factor=1.,
+ act_fx=None,
+ p_conn=1.,
+ w_bound=1.,
+ batch_size=1,
+ seed=None,
+ mu_act_fx=None,
+ mu_out_min=-jnp.inf,
+ mu_out_max=jnp.inf,
+ scalar_stddev=-1.0,
+ **kwargs
) -> None:
# This is because we have weights mu and weight log sigma
input_dim, output_dim = shape
super().__init__(
- name, (input_dim, output_dim * 2), weight_init, None, resist_scale, p_conn,
+ name, (input_dim, output_dim * 2), weight_init, None, g_conduct_factor, p_conn,
batch_size=batch_size, **kwargs
)
## Synaptic hyper-parameters
self.shape = shape ## shape of synaptic efficacy matrix
- self.Rscale = resist_scale ## post-transformation scale factor
+ self.Rscale = g_conduct_factor ## post-transformation scale factor
self.w_bound = w_bound #1. ## soft weight constraint
self.eta = eta ## learning rate
# self.out_min = out_min
@@ -152,6 +113,72 @@ def __init__(
self.learning_mask = Compartment(jnp.zeros(()))
self.seed = Compartment(jax.random.PRNGKey(seed if seed is not None else 42))
+ @staticmethod
+ def _gaussian_logpdf(event, mean, stddev): ## internal Gauss log pdf function
+ ## NOTE: might want to replace this later with external ngclearn.utils.model_utils (or "kernels/pdfs") import
+ scale_sqrd = stddev ** 2
+ log_normalizer = jnp.log(2 * jnp.pi * scale_sqrd)
+ quadratic = (event - mean) ** 2 / scale_sqrd
+ return - 0.5 * (log_normalizer + quadratic)
+
+ @staticmethod
+ def _compute_update( ## internal update calculation routine for REINFORCE
+ dt,
+ inputs,
+ rewards,
+ act_fx,
+ weights,
+ seed,
+ mu_act_fx,
+ dmu_act_fx,
+ mu_out_min,
+ mu_out_max,
+ scalar_stddev
+ ):
+ learning_stddev_mask = jnp.asarray(scalar_stddev <= 0.0, dtype=jnp.float32)
+ # (input_dim, output_dim * 2) => (input_dim, output_dim), (input_dim, output_dim)
+ W_mu, W_logstd = jnp.split(weights, 2, axis=-1)
+ # Forward pass
+ activation = act_fx(inputs)
+ mean = activation @ W_mu
+ fx_mean = mu_act_fx(mean)
+ logstd = activation @ W_logstd
+ clip_logstd = clip(logstd, -10.0, 2.0)
+ std = jnp.exp(clip_logstd)
+ std = learning_stddev_mask * std + (1.0 - learning_stddev_mask) * scalar_stddev # masking trick
+ # Sample using reparameterization trick
+ epsilon = jax.random.normal(seed, fx_mean.shape)
+ sample = epsilon * std + fx_mean
+ sample = jnp.clip(sample, mu_out_min, mu_out_max)
+ outputs = sample # the actual action that we take
+ # Compute log probability density of the Gaussian
+ log_prob = REINFORCESynapse._gaussian_logpdf(sample, fx_mean, std).sum(-1)
+ # Compute objective (negative REINFORCE objective)
+ objective = (-log_prob * rewards).mean() * 1e-2
+
+ # Backward pass
+ batch_size = inputs.shape[0] # B
+ dL_dlogp = -rewards[:, None] * 1e-2 / batch_size # (B, 1)
+
+ # Compute gradients manually based on the derivation
+ # dL/dmu = -(r-r_hat) * dlog_prob/dmu = -(r-r_hat) * -(sample-mu)/sigma^2
+ dlog_prob_dfxmean = (sample - fx_mean) / (std ** 2)
+ dL_dmean = dL_dlogp * dlog_prob_dfxmean * dmu_act_fx(mean) # (B, A)
+ dL_dWmu = activation.T @ dL_dmean
+
+ # dL/dlog(sigma) = -(r-r_hat) * dlog_prob/dlog(sigma) = -(r-r_hat) * (((sample-mu)/sigma)^2 - 1)
+ dlog_prob_dlogstd = - 1.0 / std + (sample - fx_mean) ** 2 / std ** 3
+ dL_dstd = dL_dlogp * dlog_prob_dlogstd
+ # Apply gradient clipping for logstd
+ dL_dlogstd = d_clip(logstd, -10.0, 2.0) * dL_dstd * std
+ dL_dWlogstd = activation.T @ dL_dlogstd # (I, B) @ (B, A) = (I, A)
+ dL_dWlogstd = dL_dWlogstd * learning_stddev_mask # there is no learning for the scalar stddev
+
+ # Update weights, negate the gradient because gradient ascent in ngc-learn
+ dW = jnp.concatenate([-dL_dWmu, -dL_dWlogstd], axis=-1)
+ # Finally, return metrics if needed
+ return dW, objective, outputs
+
@compilable
def evolve(self, dt):
# Get compartment values
@@ -167,7 +194,7 @@ def evolve(self, dt):
# Main logic
main_seed, sub_seed = jax.random.split(seed)
- dWeights, objective, outputs = _compute_update(
+ dWeights, objective, outputs = REINFORCESynapse._compute_update(
dt, inputs, rewards, self.act_fx, weights, sub_seed, self.mu_act_fx, self.dmu_act_fx, self.mu_out_min, self.mu_out_max, self.scalar_stddev
)
## do a gradient ascent update/shift
@@ -236,7 +263,7 @@ def help(cls): ## component help function
"eta": "Learning rate for weight updates",
"decay": "Decay factor for EMA of gradients",
"weight_init": "Initialization conditions for synaptic weight values",
- "resist_scale": "Resistance level scaling factor applied to output",
+ "g_conduct_factor": "Conductance level scaling factor applied to output",
"act_fx": "Activation function to apply to inputs",
"p_conn": "Probability of a connection existing (otherwise, it is masked to zero)",
"w_bound": "Upper bound for weight clipping",
diff --git a/ngclearn/components/synapses/mpsSynapse.py b/ngclearn/components/synapses/mpsSynapse.py
index 0b7efc7d..17ad09c9 100644
--- a/ngclearn/components/synapses/mpsSynapse.py
+++ b/ngclearn/components/synapses/mpsSynapse.py
@@ -11,9 +11,9 @@ class MPSSynapse(JaxComponent):
A Matrix Product State (MPS) compressed synaptic cable.
This component represents a synaptic weight matrix decomposed into a
- contracted chain of low-rank tensor cores (also known as a Tensor Train).
+ contracted chain of low-rank tensor cores (also known as a "tensor train").
This architecture drastically reduces parameter counts for high-dimensional
- layers—from O(N*M) to O(N*K + M*K)—while maintaining high expressive power
+ layers -- from O(N*M) to O(N*K + M*K) -- while maintaining high expressive power
and biological plausibility through local error-driven updates.
| References:
diff --git a/ngclearn/components/synapses/patched/hebbianPatchedSynapse.py b/ngclearn/components/synapses/patched/hebbianPatchedSynapse.py
index 9ccaa8db..d9555d8b 100644
--- a/ngclearn/components/synapses/patched/hebbianPatchedSynapse.py
+++ b/ngclearn/components/synapses/patched/hebbianPatchedSynapse.py
@@ -1,18 +1,15 @@
# %%
-import matplotlib.pyplot as plt
from jax import random, numpy as jnp, jit
-from functools import partial
from ngclearn.utils.optim import get_opt_init_fn, get_opt_step_fn
-
-from ngcsimlib.logger import info
-from ngclearn import compilable #from ngcsimlib.parser import compilable
-from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
+from ngclearn.utils.io_utils import save_pkl, load_pkl
+from ngclearn import compilable
+from ngclearn import Compartment
from ngclearn.components.synapses.patched import PatchedSynapse
from ngclearn.utils import tensorstats
-# @partial(jit, static_argnums=[3, 4, 5, 6, 7, 8, 9])
+
def _calc_update(
pre, post, W, mask, w_bound, is_nonnegative=True, signVal=1., prior_type=None, prior_lmbda=0., pre_wght=1.,
post_wght=1.
@@ -68,42 +65,13 @@ def _calc_update(
prior_lmbda = prior_lmbda[0]
dW_reg = jnp.sign(W) * l1_ratio + W * (1-l1_ratio)/2
- dW = dW + prior_lmbda * dW_reg
+ dW = dW - prior_lmbda * dW_reg
- if mask != None:
+ if mask is not None:
dW = dW * mask
return dW * signVal, db * signVal
-# @partial(jit, static_argnums=[1,2, 3])
-def _enforce_constraints(W, block_mask, w_bound, is_nonnegative=True):
- """
- Enforces constraints that the (synaptic) efficacies/values within matrix
- `W` must adhere to.
-
- Args:
- W: synaptic weight values (at time t)
-
- block_mask: weight mask matrix
-
- w_bound: maximum value to enforce over newly computed efficacies
-
- is_nonnegative: ensure updated value matrix is strictly non-negative
-
- Returns:
- the newly evolved synaptic weight value matrix
- """
- _W = W
- if w_bound > 0.:
- if is_nonnegative:
- _W = jnp.clip(_W, 0., w_bound)
- else:
- _W = jnp.clip(_W, -w_bound, w_bound)
-
- if block_mask != None:
- _W = _W * block_mask
-
- return _W
class HebbianPatchedSynapse(PatchedSynapse):
@@ -234,10 +202,10 @@ def __init__(
self.dWeights = Compartment(jnp.zeros(self.shape))
self.dBiases = Compartment(jnp.zeros(self.shape[1]))
- #key, subkey = random.split(self.key.get())
self.opt_params = Compartment(get_opt_init_fn(optim_type)(
[self.weights.get(), self.biases.get()]
- if bias_init else [self.weights.get()]))
+ if bias_init else [self.weights.get()]),
+ auto_save=False)
@staticmethod
def _compute_update(block_mask, w_bound, is_nonnegative, sign_value, prior_type, prior_lmbda, pre_wght,
@@ -250,6 +218,30 @@ def _compute_update(block_mask, w_bound, is_nonnegative, sign_value, prior_type,
return dW * jnp.where(0 != jnp.abs(weights), 1, 0) , db
+ @staticmethod
+ def _enforce_constraints(W, block_mask, w_bound, is_nonnegative=True):
+ ## Enforces constraints that the (synaptic) efficacies/values within matrix `W` must adhere to
+ _W = W
+ if w_bound > 0.:
+ if is_nonnegative:
+ _W = jnp.clip(_W, 0., w_bound)
+ else:
+ _W = jnp.clip(_W, -w_bound, w_bound)
+
+ if block_mask is not None:
+ _W = _W * block_mask
+ return _W
+
+ def save(self, directory: str):
+ super().save(directory)
+ # Also save the optimizer parameters
+ save_pkl(directory, self.name + "_opt_params", self.opt_params.get())
+
+ def load(self, directory: str):
+ super().load(directory)
+ # load the optimizer parameters in a custom way
+ self.opt_params.set(load_pkl(directory, self.name + "_opt_params"))
+
@compilable
def evolve(self):
# Get the variables
@@ -271,7 +263,9 @@ def evolve(self):
# ignore db since no biases configured
opt_params, [weights] = self.opt(opt_params, [weights], [dWeights])
## ensure synaptic efficacies adhere to constraints
- weights = _enforce_constraints(weights, self.block_mask, self.w_bound, is_nonnegative=self.is_nonnegative)
+ weights = HebbianPatchedSynapse._enforce_constraints(
+ weights, self.block_mask, self.w_bound, is_nonnegative=self.is_nonnegative
+ )
# Update compartments
self.opt_params.set(opt_params)
@@ -352,10 +346,13 @@ def help(cls): ## component help function
if __name__ == '__main__':
from ngcsimlib.context import Context
+ import matplotlib.pyplot as plt
+
with Context("Bar") as bar:
Wab = HebbianPatchedSynapse("Wab", (9, 30), 3, (0, 0), optim_type='adam',
sign_value=-1.0, prior=("l1l2", 0.001))
print(Wab)
+
plt.imshow(Wab.weights.get(), cmap='gray')
plt.show()
@@ -366,3 +363,6 @@ def help(cls): ## component help function
+
+
+
diff --git a/ngclearn/components/synapses/sparseTensorSynapse.py b/ngclearn/components/synapses/sparseTensorSynapse.py
new file mode 100644
index 00000000..aea25558
--- /dev/null
+++ b/ngclearn/components/synapses/sparseTensorSynapse.py
@@ -0,0 +1,624 @@
+from jax import random, numpy as jnp, jit
+from functools import partial
+import numpy as np
+
+from ngclearn import compilable # from ngcsimlib.parser import compilable
+from ngclearn import Compartment # from ngcsimlib.compartment import Compartment
+from ngclearn.components.synapses import DenseSynapse
+#from ngclearn.utils.model_utils import normalize_block_matrix
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
+
+
+########################################################################################################################
+## helper-functions for sparse-synaptic tensors
+def _make_connectivity_map(
+ P,
+ S,
+ stride=1,
+ padding=0,
+ dilation=1,
+ convergent_factor=0.0,
+ rewire_prob=0.0,
+ seed=None,
+ total_inputs=None,
+):
+ """
+ This generator produces an (S, P) connectivity map with options for dilated, convergent, and small-world (randomized)
+ biological synaptic (tensor/locally-connected) structures.
+
+ Args:
+ P: maximum connection window size (defines the tensor shape)
+ S: total number of output streams
+ stride: stride step between successive output receptive fields
+ padding: boundary padding to offset the starting index
+ dilation: gap between input indices inside a block (1 = contiguous)
+ convergent_factor: if > 0, windows contract symmetrically towards the center of the layer
+ (simulating foveation / variable resolution)
+ rewire_prob: probability [0, 1] of rewiring any connection to a random shortcut
+ seed: random seed for reproducible small-world rewiring
+ total_inputs: maximum valid input stream index. Required if using rewire_prob or padding boundaries to
+ ensure validity
+
+ Returns:
+ a connectivity map for dictating information flow in synaptic locally-connected tensor structure
+ """
+ ## initialize a map filled with -1 (dead-zone padding marker)
+ conn_map = np.full((S, P), -1, dtype=np.int32)
+ rng = np.random.default_rng(seed)
+ ## if total_inputs isn't specified, estimate it from the basic sliding layout
+ if total_inputs is None:
+ total_inputs = (S - 1) * stride + (P - 1) * dilation + 1
+
+ for s in range(S):
+ ## base starting point for sliding window
+ start_idx = s * stride - padding
+
+ ## calculate localized window size if convergent_factor is active; this shrinks window as we move away from layer's center
+ if convergent_factor > 0.0:
+ center = S / 2.0
+ distance_from_center = abs(s - center) / center
+ ## shrink window proportionally, ensuring it uses at least 1 connection
+ local_P = max(1, int(P * (1.0 - convergent_factor * distance_from_center)))
+ else:
+ local_P = P
+
+ ## generate indices incorporating dilation gaps; example: start=0, local_P=3, dilation=2 -> [0, 2, 4]
+ indices = start_idx + np.arange(local_P) * dilation
+
+ ## save to map (unfilled slots at tail end remain -1)
+ conn_map[s, :local_P] = indices
+
+ ## apply small-world rewiring
+ if rewire_prob > 0.0:
+ for s in range(S):
+ for p in range(P):
+ ## only rewire valid existing connections, skipping padding slots (-1)
+ if conn_map[s, p] != -1 and rng.random() < rewire_prob:
+ ## swap connection with completely random global input stream
+ conn_map[s, p] = rng.integers(0, total_inputs)
+ return jnp.array(conn_map)
+
+# def make_sliding_connectivity_map(P, S, stride=1, padding=0):
+# conn_map = np.zeros((S, P), dtype=np.int32)
+# for s in range(S):
+# start_idx = s * stride - padding
+# conn_map[s, :] = np.arange(start_idx, start_idx + P)
+# return jnp.array(conn_map)
+
+@partial(jit, static_argnums=(2, 3, 4, 5))
+def _reconstruct_global_2d_matrix(
+ weights: jnp.ndarray,
+ conn_map: jnp.ndarray,
+ total_input_streams: int,
+ K_local: int,
+ O_local: int,
+ S: int,
+) -> jnp.ndarray:
+ """
+ Reconstructs the full 2D block matrix (with zeros) from the dense 4D weight tensor.
+ Operates entirely in parallel via JAX JIT, making it highly efficient.
+
+ Args:
+ weights: 4D Tensor of shape (P, S, K_local, O_local)
+ conn_map: 2D array of shape (S, P) mapping output streams to input streams
+ total_input_streams: total unique global input streams (including any padded ones)
+ K_local:
+ O_local:
+ S:
+
+ Returns:
+ reconstructed/re-created global block matrix
+ """
+ P, _, _, _ = weights.shape
+ ## Initialize the massive, zero-filled global matrix structure
+ ### shape mirrors classical block format: (K_local * total_in, O_local * S)
+ global_matrix_h = total_input_streams * K_local
+ global_matrix_w = S * O_local
+ global_2d_matrix = jnp.zeros((global_matrix_h, global_matrix_w))
+ ## to avoid loops, construct coordinate indices for every element in 'weights'
+ ### Note: below creates mesh-grids matching dimensions of an individual unshared block
+ k_indices, o_indices = jnp.meshgrid(
+ jnp.arange(K_local), jnp.arange(O_local), indexing="ij"
+ )
+ ## broadcast local block dimensions across all S + P positions => expanded shape for coordinate grids: (P, S, K_local, O_local)
+ k_grid = jnp.broadcast_to(k_indices[None, None, :, :], weights.shape)
+ o_grid = jnp.broadcast_to(o_indices[None, None, :, :], weights.shape)
+ ## resolve starting stream-level offsets for every block position;
+ ### s_offsets calculates where each output stream block starts horizontally
+ s_offsets = jnp.arange(S)[None, :, None, None] * O_local
+ ## p_offsets reads conn_map to find where each input block starts vertically
+ ### conn_map has shape (S, P) -> transpose to (P, S) to align with weight axes
+ p_offsets = jnp.transpose(conn_map, (1, 0))[:, :, None, None] * K_local
+ ## compute final absolute 2D global destination coordinates
+ global_y_indices = p_offsets + k_grid
+ global_x_indices = s_offsets + o_grid
+ ## scatter all dense 4D weights into giant 2D matrix in one shot
+ global_2d_matrix = global_2d_matrix.at[
+ global_y_indices.ravel(), global_x_indices.ravel()
+ ].set(weights.ravel())
+ return global_2d_matrix
+
+def _make_backwards_connectivity_map(
+ forward_conn_map, total_input_streams
+):
+ ## generates connectivity map + local window size for the backwards pass, safely ignoring -1 boundary padding markers
+ S, P = forward_conn_map.shape
+ ## initialize tracking dictionary for all valid input stream slots
+ in_to_out_links = {i: [] for i in range(total_input_streams)}
+
+ for s in range(S):
+ for p in range(P):
+ global_in_idx = int(forward_conn_map[s, p])
+
+ ## only map connections that fall inside valid bounds;
+ ### this ignores -1 padding placeholders + any indices exceeding total_input_streams
+ if 0 <= global_in_idx < total_input_streams:
+ in_to_out_links[global_in_idx].append(s)
+
+ ## determine backwards window size (P_back);
+ ### if an input has zero connections (e.g. dead boundary edge), default max to 1
+ P_back = max(len(links) for links in in_to_out_links.values())
+ P_back = max(1, P_back)
+
+ ## build backwards connectivity map of shape (S_back, P_back)
+ back_conn_map = np.zeros((total_input_streams, P_back), dtype=np.int32)
+ for in_idx, out_indices in in_to_out_links.items():
+ ## fill available slots, pad remainder with -1 for safety
+ for slot_idx, out_idx in enumerate(out_indices):
+ back_conn_map[in_idx, slot_idx] = out_idx
+ ## if this input stream connects to fewer outputs than P_back, pad out with -1
+ for slot_idx in range(len(out_indices), P_back):
+ back_conn_map[in_idx, slot_idx] = -1
+
+ return np.array(back_conn_map), P_back
+########################################################################################################################
+
+class SparseTensorSynapse(DenseSynapse):
+ """
+ A sparse tensor-synaptic cable. Note this component cable implements a full, locally-connected structure or an
+ unshared convolutional synaptic tensor structural component (cable).
+
+ | --- Synapse Input Compartments: ---
+ | inputs - input (takes in external signals)
+ | --- Synapse State Compartments: ---
+ | weights - current value matrix of synaptic efficacies (strength values)
+ | biases - current value vector of synaptic bias values
+ | --- Synapse Output Compartments: ---
+ | outputs - output signals
+
+ Args:
+ name: the string name of this cell
+
+ n_in_streams: total number of incoming streams
+
+ K_local: input local block feature size
+
+ O_local: output local block feature size (note: total number of output streams is computed internally)
+
+ P_l: local window size (how many input streams to grab), i.e., maximum connection window size
+ (defines tensor shape) (Default: 1)
+
+ stride: stride factor for locally-connected structure to "skip over" (in terms of incoming streams);
+ stride step between successive output receptive field streams (Default: 1)
+
+ convergent_factor: if > 0, windows contract symmetrically towards center of layer
+ (simulates foveation or variable resolution) (Default: 0)
+
+ dilation: dilation factor; gap between input indices inside a block (1 = contiguous) (Default: 1)
+
+ invert_conn: if True, this tensor shape will internally "transpose" itself to formulate the appropriate
+ locally-connected inverted/transposed structure (note that this means the shape will no longer be what
+ the constructor's argument dictate - it will be the effective transpose of these arguments) (Default: False)
+
+ normalize: if True, this synaptic tensor will normalize its internal blocks (Default: False)
+
+ norm_axis: axis upon which block norms are computed (Default: 1)
+
+ weight_init: a kernel to drive initialization of this synaptic cable's values;
+ typically a tuple with 1st element as a string calling the name of
+ initialization to use
+
+ bias_init: a kernel to drive initialization of biases for this synaptic cable
+ (Default: None, which turns off/disables biases)
+
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * in) * g_conduct_factor) + bias
+
+ p_conn: probability of a connection existing (default: 1.); setting this to < 1 and > 0. will result in a
+ sparser synaptic structure (lower values yield sparse structure)
+
+ use_block_matrix_format: if True, this tensor synapse resorts to a memory-intensive block-matrix
+ computational format (Default: False)
+
+ """
+
+ def __init__(
+ self,
+ name,
+ n_in_streams, ## total incoming streams (nlm1_streams)
+ K_local, ## input local block feature size
+ O_local, ## output local block feature size
+ P_l=1, ## local window size (how many input streams to grab)
+ stride=1,
+ convergent_factor=0.,
+ dilation=1,
+ invert_conn=False,
+ normalize=False,
+ norm_axis=1,
+ weight_init=None,
+ bias_init=None,
+ g_conduct_factor=1.,
+ p_conn=1.,
+ use_block_matrix_format=False, ## if True, triggers block-matrix implementation
+ batch_size=1,
+ **kwargs
+ ):
+ ################################################################################################################
+ ## set up internal structure/topology
+ ### run geometry math internally first; if invert_conn is True, 'incoming' streams to this component are actually
+ ### errors from upper layer, meaning 'in_streams' is upper layer's output streams (S_forward)
+ if invert_conn:
+ dims = SparseTensorSynapse.initialize_layer_geometry(
+ n_in_streams,
+ K_local,
+ O_local,
+ P_l,
+ stride,
+ auto_pad=False,
+ dilation_l=dilation,
+ convergent_factor_l=convergent_factor
+ )
+
+ total_forward_inputs = n_in_streams
+ total_forward_outputs = dims["total_num_output_streams"]
+
+ ## build base forward connectivity map so this can later be "inverted"
+ f_map = _make_connectivity_map(
+ P_l, S=n_in_streams, stride=stride, convergent_factor=convergent_factor, dilation=dilation
+ )
+ self.connectivity_map, P_back = _make_backwards_connectivity_map(f_map, total_forward_inputs) ## invert map
+
+ ## explicitly compile exact shapes that ngc-learn expects
+ self.n_in_streams = total_forward_outputs ## e.g., 1 (number of incoming error maps)
+ self.n_out_streams = total_forward_inputs ## e.g., 3 (number of feedback target maps)
+ self.K_local = O_local # Error size (9)
+ self.O_local = K_local # Hidden feature target size (16)
+
+ # Final 4D tensor shape for backwards weights
+ calculated_shape = (P_back, total_forward_inputs, O_local, K_local)
+ self.io_shape = calculated_io_shape = (self.n_in_streams * self.K_local, self.n_out_streams * self.O_local)
+
+ # Clear window strides for pure backward routing execution
+ self.stride = 0
+ self.dilation = 1
+ else:
+ ## standard forward-projection synaptic cable
+ dims = SparseTensorSynapse.initialize_layer_geometry(
+ n_in_streams,
+ K_local,
+ O_local,
+ P_l,
+ stride,
+ auto_pad=False,
+ dilation_l=dilation,
+ convergent_factor_l=convergent_factor
+ )
+
+ self.n_in_streams = n_in_streams # dims["total_input_streams"] # nlm1_streams
+ self.n_out_streams = dims["total_num_output_streams"] # Computed S
+ self.K_local = K_local
+ self.O_local = O_local
+
+ self.connectivity_map = _make_connectivity_map(
+ P_l, dims["total_num_output_streams"], stride, convergent_factor=convergent_factor, dilation=dilation
+ )
+
+ calculated_shape = dims["forward_weight_shape"] # (P_l, S, K_local, O_local)
+ self.io_shape = calculated_io_shape = (dims["global_input_dim"], dims["global_output_dim"])
+ self.stride = stride
+ self.dilation = dilation
+ ################################################################################################################
+ super().__init__( ## now, call parent constructor using internally calculated shapes
+ name,
+ shape=calculated_shape,
+ weight_init=weight_init,
+ bias_init=bias_init,
+ g_conduct_factor=g_conduct_factor,
+ p_conn=p_conn,
+ batch_size=batch_size,
+ **kwargs
+ )
+
+ ## store framework properties
+ self.convergent_factor = convergent_factor
+ self.invert_conn = invert_conn
+ self.use_block_matrix_format = use_block_matrix_format
+ self.Rscale = g_conduct_factor
+ self.normalize = normalize
+ self.norm_axis = norm_axis
+
+ #if self.use_effective_dim_prior:
+ self.norm_trigger = Compartment(jnp.zeros((1, 1)))
+ if self.normalize:
+ self.norm_trigger.set(self.norm_trigger.get() + 1)
+ P, S, Klocal, Olocal = self.shape
+ norm_axis = 1 #0 #1
+ Weffdim = Klocal * P
+ if norm_axis == 1:
+ Weffdim = Olocal * self.n_out_streams #Weffdim = Olocal
+ weight_init = dist.gaussian(mean=0., std=float(1.0 / jnp.sqrt(Weffdim)))
+ weights = weight_init(self.shape, self.key.get())
+ self.weights.set(weights)
+
+ if self.use_block_matrix_format: ## set up block-diagonal algorithm backend if flagged
+ block_weights = _reconstruct_global_2d_matrix(
+ self.weights.get(),
+ self.connectivity_map,
+ self.n_in_streams, ## total num in-streams
+ K_local=self.shape[2],
+ O_local=self.shape[3],
+ S=self.n_out_streams #self.shape[1] ## total num out-streams
+ )
+ self.weights.set(block_weights)
+
+ if self.normalize: ## pre-normalize synapses if norm-constraints are to be used
+ if not self.use_block_matrix_format:
+ weights = SparseTensorSynapse._normalize_outgoing_weights(self.weights.get())
+ else:
+ weights = SparseTensorSynapse._normalize_block_matrix_outgoing_weights(
+ self.weights.get(),
+ self.connectivity_map,
+ self.n_in_streams,
+ K_local,
+ O_local,
+ self.n_out_streams,
+ axis=1, #0, #1,
+ order=2,
+ norm_targ=1.
+ )
+ self.weights.set(weights)
+
+
+ self.initial_weights = Compartment(self.weights.get()) ## store synaptic initial conditions (never updated)
+ preVals = jnp.zeros((self.batch_size, self.io_shape[0]))
+ postVals = jnp.zeros((self.batch_size, self.io_shape[1]))
+ self.inputs.set(preVals)
+ self.outputs.set(postVals)
+
+ @staticmethod
+ def initialize_layer_geometry( ## internal co-routine for tensor-synaptic projection shaping
+ nlm1_streams,
+ Nlm1,
+ Nl,
+ P_l,
+ stride_l,
+ auto_pad=False,
+ dilation_l=1,
+ convergent_factor_l=0.0,
+ ):
+ ## this sets up geometrical constraints for this synaptic tensor component, notably this:
+ ### * calculates output stream count (S), maximum structural input reach, and
+ ### * allocates flat sizes for ngc-learn buffers in one unified sweep
+
+ ## check 1: dilation stride minimum bounds
+ if stride_l < 1:
+ raise ValueError(
+ "Geometry Error: stride_l must be greater than or equal to 1."
+ )
+ ## check 2: check physical window limits if padding is disabled
+ max_physical_span = (P_l - 1) * dilation_l + 1
+ if max_physical_span > nlm1_streams and not auto_pad:
+ raise ValueError(
+ f"Geometry Error: Receptive field physical span ({max_physical_span} inputs) "
+ f"exceeds incoming streams ({nlm1_streams}). Enable auto_pad."
+ )
+
+ ## sweep over potential output streams to find exact boundary
+ s = 0
+ max_seen_input_idx = -1
+ while True:
+ start_idx = s * stride_l
+ ## under foveation/convergence, calculate local window footprint dynamically
+ ### we estimate center distance dynamically as we look ahead
+ if convergent_factor_l > 0.0:
+ ## we use a running estimation of center scaling
+ ## for a self-contained lookup, we approximate shrinking footprint relative to stride
+ estimated_S_guess = max(1, nlm1_streams // stride_l)
+ center = estimated_S_guess / 2.0
+ distance_from_center = abs(s - center) / center
+ local_P = max(
+ 1, int(P_l * (1.0 - convergent_factor_l * distance_from_center))
+ )
+ else:
+ local_P = P_l
+ ## compute absolute furthest index this specific stream row will touch
+ last_index_in_row = start_idx + (local_P - 1) * dilation_l
+
+ if auto_pad: ## boundary evaluation gates
+ ## stop if starting position completely falls off available canvas
+ if start_idx >= nlm1_streams:
+ break
+ else:
+ ## without padding, entire synaptic projection footprint must fit w/in source bounds
+ if last_index_in_row >= nlm1_streams:
+ break
+ ## if stream is valid, track its absolute furthest input coordinate reach
+ if last_index_in_row > max_seen_input_idx:
+ max_seen_input_idx = last_index_in_row
+ s += 1
+
+ ## finalize verified number of output streams (S)
+ nl_streams = max(1, s)
+ ## calculate unique input streams required to fully span generated matrix/tensor
+ total_input_streams = max(1, max_seen_input_idx + 1)
+ ## calculate how many extra virtual padding streams are absorbed by "dead zone"
+ pad_streams = max(0, total_input_streams - nlm1_streams)
+ ## build final (structural) metadata values
+ forward_weight_shape = (P_l, nl_streams, Nlm1, Nl)
+ return {
+ "total_num_output_streams": nl_streams, ## computed num output streams (S)
+ "pad_streams": pad_streams, ## num streams absorbed by padding filters
+ "total_num_input_streams": total_input_streams, ## total span across input map
+ "global_input_dim": nlm1_streams * Nlm1, ## raw size of incoming layer nodes
+ "global_output_dim": nl_streams * Nl, ## raw size of output nodes
+ "forward_weight_shape": forward_weight_shape, ## locally-connected tensor shape
+ }
+
+ @compilable
+ def advance_state(self):
+ weights = self.weights.get() ## get synaptic tensor
+ biases = self.biases.get()
+ inputs = self.inputs.get() ## get inputs
+ conn_map = self.connectivity_map ## get connectivity structure
+
+ B = inputs.shape[0] ## get batch size for subsequent averaging
+ if not self.use_block_matrix_format:
+ P, S, K_local, O_local = weights.shape
+ ## reshape global inputs into standard patches
+ patched_inputs = inputs.reshape(B, -1, K_local)
+ total_input_streams = patched_inputs.shape[1]
+
+ ## append an explicit row of absolute zeros to the end of our inputs;
+ ### this acts as dead-zone buffer for all invalid/padded indices
+ zero_padding_block = jnp.zeros((B, 1, K_local))
+ padded_inputs = jnp.concatenate(
+ [patched_inputs, zero_padding_block], axis=1
+ )
+
+ ## create a safe connectivity map by replacing all -1 flags or
+ ### out-of-bounds indices with a pointer to our dead-zone block;
+ ### (dead-zone block sits at very last index: total_input_streams)
+ is_invalid = (conn_map == -1) | (conn_map < 0) | (conn_map >= total_input_streams)
+ safe_conn_map = jnp.where(is_invalid, total_input_streams, conn_map)
+
+ ## gather overlapping streams using our safe map => (B, S, P, K_local)
+ ### invalid indices safely pull from dead-zone block of zeros
+ gathered_inputs = padded_inputs[:, safe_conn_map, :]
+ ## transpose to move P to front: (P, B, S, K_local)
+ gathered_inputs = jnp.transpose(gathered_inputs, (2, 0, 1, 3))
+
+ ## use None axes to broadcast (P, B, S, K_local, 1) * (P, 1, S, K_local, O_local)
+ ### output of this elementwise multiply is shape: (P, B, S, K_local, O_local)
+ elementwise_prod = (
+ gathered_inputs[..., None] * weights[:, None, :, :, :]
+ )
+
+ ## sum out the input stream dimensions: axis 0 (P) and axis 3 (K_local)
+ ### resulting shape: (B, S, O_local)
+ outputs_3d = jnp.sum(elementwise_prod, axis=(0, 3))
+ outputs = outputs_3d.reshape(B, -1) ## flatten to 2D formatted output
+ else:
+ outputs = inputs @ weights ## block-diagonal/block-matrix multiply
+ outputs + biases
+ self.outputs.set(outputs)
+
+ @compilable
+ def reset(self): ## reset compartments/statistics
+ if not self.inputs.targeted:
+ self.inputs.set(self.inputs.get() * 0)
+ self.outputs.set(self.outputs.get() * 0) # outputs
+
+ @staticmethod
+ def _normalize_outgoing_weights( ## internal tensor normalizer co-routine
+ weights: jnp.ndarray,
+ axis: int = 1, # 0,
+ order: int = 2,
+ epsilon: float = 1e-8
+ ) -> jnp.ndarray:
+ ## normalizes outgoing weights of each input neuron to a unit L2 norm;
+ ### weight tensor shape: (P, S, K_local, O_local)
+ ### - axis 0 (P) and Axis 2 (K_local) isolate an individual input neuron.
+ ### - axis 1 (S) and Axis 3 (O_local) represent all the places its synapses land.
+
+ if axis == 0:
+ ## input-wise normalization: reduce across patch connections (P) and local features (K_local)
+ reduce_axes = (0, 2)
+ elif axis == 1:
+ ## output-wise normalization: reduce across output streams (S) and local output neurons (O_local)
+ reduce_axes = (1, 3)
+ else:
+ raise ValueError("Norm.axis must be 0 or 1.")
+
+ if order == 1:
+ ## compute L1 norm across all target output streams (S) and output units (O_local);
+ outgoing_norms = jnp.sum(
+ jnp.abs(weights),
+ axis=reduce_axes,
+ # axis=(0, 2), ## equiv to axis=0 for 2D synapse projection
+ # axis=(1, 3), ## equiv to axis=1 for 2D synapse projection
+ keepdims=True ## keepdims=True ensures shape matches for broadcasted division
+ )
+ else: # order==2
+ ## compute L2 norm across all target output streams (S) and output units (O_local);
+ outgoing_norms = jnp.sqrt(
+ jnp.sum(
+ jnp.square(weights),
+ axis=reduce_axes,
+ # axis=(0, 2), ## equiv to axis=0 for 2D synapse projection
+ #axis=(1, 3), ## equiv to axis=1 for 2D synapse projection
+ keepdims=True ## keepdims=True ensures shape matches for broadcasted division
+ )
+ )
+ ## divide elementwise to ensure sum of squared outgoing weights for any input neuron equals exactly 1.0.
+ return weights / (outgoing_norms + epsilon)
+
+
+ @staticmethod #@partial(jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9))
+ def _normalize_block_matrix_outgoing_weights( ## internal block-matrix normalizer co-routine
+ global_matrix: jnp.ndarray,
+ conn_map: jnp.ndarray,
+ total_input_streams: int,
+ K_local: int,
+ O_local: int,
+ S: int,
+ axis: int = 1, #0,
+ order: int = 2,
+ norm_targ: float = 1.0,
+ eps: float = 1e-8,
+ ) -> jnp.ndarray:
+ ## normalizes only active blocks within an arbitrary overlapping 2D block matrix;
+ ### NOTE: leaves all structural zero blocks completely untouched
+ P = conn_map.shape[1]
+ ## generate the exact coordinate grids for the active weights
+ ### this matches the coordinate mapping logic from our reconstruction code
+ k_indices, o_indices = jnp.meshgrid(
+ jnp.arange(K_local), jnp.arange(O_local), indexing="ij"
+ )
+
+ k_grid = jnp.broadcast_to(k_indices[None, None, :, :], (P, S, K_local, O_local))
+ o_grid = jnp.broadcast_to(o_indices[None, None, :, :], (P, S, K_local, O_local))
+ s_offsets = jnp.arange(S)[None, :, None, None] * O_local
+ p_offsets = jnp.transpose(conn_map, (1, 0))[:, :, None, None] * K_local
+
+ global_y_indices = p_offsets + k_grid
+ global_x_indices = s_offsets + o_grid
+ ## pull out ONLY the valid weights from the 2D matrix into a 4D tensor
+ ### gathered_weights shape: (P, S, K_local, O_local)
+ gathered_weights = global_matrix[global_y_indices, global_x_indices]
+
+ ## compute norms on the isolated valid blocks based on the desired axis
+ if axis == 0:
+ ## input-wise normalization: reduce across patch connections (P) and local features (K_local)
+ reduce_axes = (0, 2)
+ elif axis == 1:
+ ## output-wise normalization: reduce across output streams (S) and local output neurons (O_local)
+ reduce_axes = (1, 3)
+ else:
+ raise ValueError("Norm.axis must be 0 or 1.")
+
+ norms = jnp.linalg.norm(
+ gathered_weights, ord=order, axis=reduce_axes, keepdims=True
+ )
+ normalized_gathered = gathered_weights * (norm_targ / (norms + eps)) ## normalize the isolated valid blocks
+
+ ## initialize a pristine zero matrix and scatter the normalized weights back into it
+ ### this guarantees that all structural zeros remain mathematically clean 0.0s
+ global_matrix_h = total_input_streams * K_local
+ global_matrix_w = S * O_local
+ new_global_matrix = jnp.zeros((global_matrix_h, global_matrix_w))
+
+ new_global_matrix = new_global_matrix.at[
+ global_y_indices.ravel(), global_x_indices.ravel()
+ ].set(normalized_gathered.ravel())
+ return new_global_matrix
+
diff --git a/ngclearn/components/synapses/staticSparseTensorSynapse.py b/ngclearn/components/synapses/staticSparseTensorSynapse.py
new file mode 100644
index 00000000..282b9933
--- /dev/null
+++ b/ngclearn/components/synapses/staticSparseTensorSynapse.py
@@ -0,0 +1,62 @@
+from .sparseTensorSynapse import SparseTensorSynapse
+
+class StaticSparseTensorSynapse(SparseTensorSynapse):
+ """"
+ A static sparse tensor-synaptic cable; no form of synaptic evolution/adaptation is in-built to this component. Note
+ this component cable implements a full, locally-connected structure or an unshared convolutional synaptic
+ tensor structural component (cable).
+
+ | --- Synapse Input Compartments: ---
+ | inputs - input (takes in external signals)
+ | --- Synapse State Compartments: ---
+ | weights - current value matrix of synaptic efficacies (strength values)
+ | biases - current value vector of synaptic bias values
+ | --- Synapse Output Compartments: ---
+ | outputs - output signals
+
+ Args:
+ name: the string name of this cell
+
+ n_in_streams: total number of incoming streams
+
+ K_local: input local block feature size
+
+ O_local: output local block feature size (note: total number of output streams is computed internally)
+
+ P_l: local window size (how many input streams to grab), i.e., maximum connection window size
+ (defines tensor shape) (Default: 1)
+
+ stride: stride factor for locally-connected structure to "skip over" (in terms of incoming streams);
+ stride step between successive output receptive field streams (Default: 1)
+
+ convergent_factor: if > 0, windows contract symmetrically towards center of layer
+ (simulates foveation or variable resolution) (Default: 0)
+
+ dilation: dilation factor; gap between input indices inside a block (1 = contiguous) (Default: 1)
+
+ invert_conn: if True, this tensor shape will internally "transpose" itself to formulate the appropriate
+ locally-connected inverted/transposed structure (note that this means the shape will no longer be what
+ the constructor's argument dictate - it will be the effective transpose of these arguments) (Default: False)
+
+ normalize: if True, this synaptic tensor will normalize its internal blocks (Default: False)
+
+ norm_axis: axis upon which block norms are computed (Default: 1)
+
+ weight_init: a kernel to drive initialization of this synaptic cable's values;
+ typically a tuple with 1st element as a string calling the name of
+ initialization to use
+
+ bias_init: a kernel to drive initialization of biases for this synaptic cable
+ (Default: None, which turns off/disables biases)
+
+ g_conduct_factor: a fixed (resistance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * in) * g_conduct_factor) + bias
+
+ p_conn: probability of a connection existing (default: 1.); setting this to < 1 and > 0. will result in a
+ sparser synaptic structure (lower values yield sparse structure)
+
+ use_block_matrix_format: if True, this tensor synapse resorts to a memory-intensive block-matrix
+ computational format (Default: False)
+
+ """
+ pass
\ No newline at end of file
diff --git a/ngclearn/components/synapses/staticSynapse.py b/ngclearn/components/synapses/staticSynapse.py
index 0cf06757..bca7349e 100755
--- a/ngclearn/components/synapses/staticSynapse.py
+++ b/ngclearn/components/synapses/staticSynapse.py
@@ -20,8 +20,8 @@ class StaticSynapse(DenseSynapse):
typically a tuple with 1st element as a string calling the name of
initialization to use
- resist_scale: a fixed (resistance) scaling factor to apply to synaptic
- transform (Default: 1.), i.e., yields: out = ((W * Rscale) * in)
+ g_conduct_factor: a fixed (conductance) scaling factor to apply to synaptic
+ transform (Default: 1.), i.e., yields: out = ((W * g_conduct_factor) * in)
p_conn: probability of a connection existing (default: 1.); setting
this to < 1 and > 0. will result in a sparser synaptic structure
diff --git a/ngclearn/utils/analysis/effective_dim.py b/ngclearn/utils/analysis/effective_dim.py
index 7b46f318..d4327f4e 100644
--- a/ngclearn/utils/analysis/effective_dim.py
+++ b/ngclearn/utils/analysis/effective_dim.py
@@ -1,12 +1,19 @@
-from jax import numpy as jnp
+from functools import partial
+from jax import numpy as jnp, jit
-def participation_ratio(latent_codes):
+@partial(jit, static_argnums=[1])
+def participation_ratio(
+ latent_codes, use_NaN_fallback=False
+):
"""
Calculates the participation ratio coefficient for a set of latent codes
Args:
latent_codes: a set of (N x D) latent code vectors (one row per vector code)
+ use_NaN_fallback: if True, this function returns NaN for a squared covariance
+ trace of zero; else, it returns an eff-dim of 1 (Default: False)
+
Returns:
scalar measurement of the effective dimension
"""
@@ -18,29 +25,41 @@ def participation_ratio(latent_codes):
tr2_cov = tr * tr
cov2_tr = jnp.trace(cov @ cov)
+ ## this algorithm supports one of two fallback cases
+ if not use_NaN_fallback: ## use fallback-to-1 eff-dim check
+ ## use JAX-friendly conditional / direct switch to fallback to 1.0.
+ ### if squared trace of covariance is 0 then effective dimension is 1.0
+ return jnp.where(cov2_tr > 0.0, tr2_cov / cov2_tr, 1.0)
+ ##else, use ML-oriented NaN return value fallback
return tr2_cov / cov2_tr if cov2_tr > 0 else float("nan")
-
-def rankme(Z, eps=1e-7):
+@partial(jit, static_argnums=[1])
+def rankme(latent_codes, eps=1e-7):
"""
- Calculates the effective rank of for a code matrix Z
- effective rank = exp(Shannon entropy), from Garrido, Balestriero,
- Najman & LeCun, "RankMe: Assessing the Downstream Performance of Pretrained
- Self-Supervised Representations by Their Rank" (ICML 2023, arXiv:2210.02885).
+ Calculates the effective rank of for a code matrix latent_codes
+
+ effective rank = exp(Shannon entropy), adapted from:
+ | Garrido, Balestriero, Najman & LeCun, "RankMe: Assessing the Downstream Performance of Pretrained
+ | Self-Supervised Representations by Their Rank" (ICML 2023, arXiv:2210.02885).
Args:
latent_codes: a set of (N x D) latent code vectors (one row per vector code)
+ eps: (regularization) constant to prevent division by zero
+
Returns:
scalar measurement of the effective dimension
"""
- singular_values = jnp.linalg.svd(Z, compute_uv=False) ## singular values of Z
- sum_singular_vals = jnp.sum(singular_values) ## L1
- if sum_singular_vals <= 0:
- return float("nan")
- p = singular_values / sum_singular_vals + eps ## L1-normalized singular value
- shannon_entropy = -jnp.sum(p * jnp.log(p)) ## Shannon entropy
- return jnp.exp(shannon_entropy) ## exp(Shannon entropy) = effective rank
+ singular_values = jnp.linalg.svd(latent_codes, compute_uv=False) ## singular values of latent_codes
+ sum_singular_values = jnp.sum(singular_values) ## L1
+ sum_S_vals = jnp.where(sum_singular_values > 0.0, sum_singular_values, 1.0)
+ p = singular_values / sum_S_vals + eps ## L1-normalized singular value
+ shannon_entropy = -jnp.sum(p * jnp.log(p)) ## calc Shannon entropy
+
+ return jnp.exp(jnp.where(sum_singular_values > 0.0,
+ shannon_entropy, ## compute final exp(Shannon entropy) = effective rank
+ jnp.nan))
+
diff --git a/ngclearn/utils/data_loader.py b/ngclearn/utils/data_loader.py
index e90df8d9..ea211d97 100644
--- a/ngclearn/utils/data_loader.py
+++ b/ngclearn/utils/data_loader.py
@@ -25,8 +25,14 @@ class DataLoader(object):
key: PRNG key to control determinism of any underlying random values
associated with this synaptic cable
"""
- def __init__(self, design_matrices, batch_size, disable_shuffle=False,
- ensure_equal_batches=True, key=None):
+ def __init__(
+ self,
+ design_matrices,
+ batch_size,
+ disable_shuffle=False,
+ ensure_equal_batches=True,
+ key=None
+ ):
self.key = key
if self.key is None:
self.key = random.PRNGKey(time.time_ns())
@@ -47,23 +53,25 @@ def __init__(self, design_matrices, batch_size, disable_shuffle=False,
def __iter__(self):
"""
- Yields a mini-batch of the form: [("name", batch),("name",batch),...]
+ Yields a mini-batch of the form:
+
+ | batch = [("name", batchx), ("name", batchy),...("name", batchz)]
"""
- if self.disable_shuffle == False:
+ if not self.disable_shuffle: #self.disable_shuffle == False:
self.key, *subkeys = random.split(self.key, 2)
self.ptrs = random.permutation(subkeys[0], self.data_len)
idx = 0
- while idx < len(self.ptrs): # go through each sample via the sampling pointer
+ while idx < len(self.ptrs): ## go through each sample via the sampling pointer
e_idx = idx + self.batch_size
- if e_idx > len(self.ptrs): # prevents reaching beyond length of dataset
+ if e_idx > len(self.ptrs): ## prevents reaching beyond length of dataset
e_idx = len(self.ptrs)
- # extract sampling integer pointers
+ ## extract sampling integer pointers
indices = self.ptrs[idx:e_idx]
- if self.ensure_equal_batches == True:
+ if self.ensure_equal_batches: # == True:
if indices.shape[0] < self.batch_size:
diff = self.batch_size - indices.shape[0]
indices = jnp.concatenate((indices, self.ptrs[0:diff]))
- # create the actual pattern vector batch block matrices
+ ## create the actual pattern vector batch block matrices
data_batch = []
for dname, dmatrix in self.design_matrices:
x_batch = dmatrix[indices]
diff --git a/ngclearn/utils/distribution_generator.py b/ngclearn/utils/distribution_generator.py
index d7877991..60374ade 100644
--- a/ngclearn/utils/distribution_generator.py
+++ b/ngclearn/utils/distribution_generator.py
@@ -73,6 +73,26 @@ def constant_generator(shape: Sequence[int], dKey: jax.dtypes.prng_key | None =
return matrix
return constant_generator
+ @staticmethod
+ def diag_offdiag(on_diag: float, off_diag: float, **params: Unpack[DistributionParams]) -> DistributionInitializer:
+ using_np = params.get("use_numpy", False)
+ if using_np:
+ def diag_offdiag_generator(shape: Sequence[int], seed: int | None = None) -> numpy.ndarray:
+ eye = numpy.eye(*shape).astype(params.get("dtype", numpy.float32))
+ ones = numpy.ones(shape).astype(params.get("dtype", numpy.float32))
+ matrix = on_diag * eye + off_diag * (ones - eye)
+ matrix = DistributionGenerator._process_params_numpy(matrix, params, seed)
+ return matrix
+ else:
+ def diag_offdiag_generator(shape: Sequence[int], dKey: jax.dtypes.prng_key | None = None) -> jax.Array:
+ eye = jax.numpy.eye(*shape).astype(params.get("dtype", jax.numpy.float32))
+ ones = jax.numpy.ones(shape).astype(params.get("dtype", jax.numpy.float32))
+ matrix = on_diag * eye + off_diag * (ones - eye)
+ matrix = DistributionGenerator._process_params_jax(matrix, params, dKey)
+ return matrix
+ return diag_offdiag_generator
+
+
@staticmethod
def uniform(low: float = 0.0, high: float = 1.0, **params: Unpack[DistributionParams]) -> DistributionInitializer:
"""
diff --git a/ngclearn/utils/metric_utils.py b/ngclearn/utils/metric_utils.py
index e30f3d9c..409a2f44 100755
--- a/ngclearn/utils/metric_utils.py
+++ b/ngclearn/utils/metric_utils.py
@@ -4,6 +4,7 @@
"""
from jax import numpy as jnp, jit
from functools import partial
+from ngcsimlib import deprecated
from sklearn.metrics import confusion_matrix, precision_score, recall_score
@partial(jit, static_argnums=[1])
@@ -53,35 +54,135 @@ def measure_firingRate(spikes, preserve_batch=False):
fireRates = jnp.mean(fireRates)
return fireRates
+#@partial(jit, static_argnums=[1])
+@jit
+def measure_workload_entropy(spike_count_matrix):
+ """
+ Calculates the workload entropy (or balance) across a D-length population of neuronal units across
+ N samples of stimuli. This routine also returns the number of "dead neurons" (or those that never activate
+ in the context of all N stimuli samples). This measurement function is useful for characterizing the
+ "lifetime sparsity" of a population of D spiking/pulsing neuronal units.
+
+ Args:
+ spike_count_matrix: N x D tensor containing N rows of D-length spike count vectors (over D neuronal units)
+
+ Returns:
+ (min, max, mean) pulse count/rate, workload entropy, number of "dead"/silent neurons
+ """
+ global_neuron_activity = jnp.sum(spike_count_matrix, axis=0)
+ min_act = jnp.min(global_neuron_activity)
+ max_act = jnp.max(global_neuron_activity)
+ mean_act = jnp.mean(global_neuron_activity)
+ dead_neurons = jnp.sum(global_neuron_activity == 0)
+ ### compute workload entropy across neuron pool
+ total_epoch_spikes = jnp.sum(global_neuron_activity)
+
+ _flag = (total_epoch_spikes > 0) ## mask
+ probs = global_neuron_activity / total_epoch_spikes
+ full_entropy = -jnp.sum(jnp.where(probs > 0, probs * jnp.log2(probs), 0.0)) ## prevents log2(0)
+ entropy = full_entropy * _flag
+ # if total_epoch_spikes > 0:
+ # probs = global_neuron_activity / total_epoch_spikes
+ # entropy = -jnp.sum(jnp.where(probs > 0, probs * jnp.log2(probs), 0.0)) ## prevents log2(0)
+ # else:
+ # entropy = 0.0
+ return (min_act, max_act, mean_act), entropy, dead_neurons
+
@partial(jit, static_argnums=[1])
-def measure_breadth_TC(spikes, preserve_batch=False):
+def measure_breadthOfTuningCurve(responses, preserve_batch=False):
"""
- Calculates the breath tuning curve (BTC) of a group of neurons given full
- spike train.(s). BTC measures the neural selectivity such that the
- sparse code distribution concentrates near zero with a heavy tail. For a
- neural layer where most of the neurons fire, the activity distribution is
+ Calculates the breadth of tuning curve (BTC) of a population of neurons given full
+ response window(s), across multiple trials. BTC measures the neural selectivity such that the
+ sparse code distribution concentrates near zero with a heavy tail.
+ Note that, within `responses`, neural response values will include measurements of
+ characteristics such as firing rates, spike counts (burstiness measures), etc.
+ at a particular point in time/stimulus condition (each row represents a condition).
+
+ For a neural layer where most of the neurons fire, the activity distribution is
more uniformly spread and BTC > 0.5. When most of the neurons do not fire,
the firing distribution is peaked at zero and BTC < 0.5.
Args:
- spikes: full spike train matrix; shape is (T x D) where D is number of
- neurons in a group/cluster
-
- preserve_batch: if True, will return one score per neuron in train/window
- (Default: False), otherwise, returns scalar average score
+ responses: tensor of shape (B x T x D), containing "evoked activity", where:
+ B = number of trials (batch length; trial axis=0);
+ T = number of unique stimulus conditions / orientations (stimulus axis=1); and,
+ D = number of neurons in the population (neural index axis=2)
+ preserve_batch: if True, returns (1 x D) vector of tuning widths (BTC) per neuron;
+ if False, returns scalar representing population-wide average tuning width (BTC)
- Returns:
- a 1 x D BTC vector (one factor per neuron) OR a single
- average BTC across the neuronal group
+ Returns:
+ a 1 x D BTC vector (one factor per neuron) OR a single average BTC across
+ the neuronal group / population
"""
- mu = jnp.mean(spikes, axis=0, keepdims=True)
- sigSqr = jnp.square(jnp.std(spikes, axis=0, keepdims=True))
- C = sigSqr/mu
- BTC = 1./(1 + jnp.square(C))
- if not preserve_batch:
+ ## clean out trial noise: mean across"trials" axis (axis=0)
+ clean_responses = jnp.mean(responses, axis=0, keepdims=True) ## (1xTxD)
+
+ ## calculate mean & variance across stimulus axis (i.e., `axis=1`)
+ mu = jnp.mean(clean_responses, axis=1, keepdims=True) ## (1x1xD)
+ sigSqr = jnp.square(jnp.std(clean_responses, axis=1, keepdims=True)) ## (1x1xD)
+
+ safe_mu = jnp.where(mu == 0.0, 1.0, mu) ## check for 0/0 division for non-responsive neurons
+ C = sigSqr / safe_mu
+ raw_BTC = 1. / (1 + jnp.square(C)) ## raw original BTC metric
+
+ ## NOTE: if neuron never fires for any stimulus (mu == 0), its breadth of
+ ## tuning is exactly 0.0 (not 1.0); thus assign it maximal sparsity
+ BTC = jnp.where(mu == 0.0, 0.0, raw_BTC) ## masking check
+ BTC = jnp.squeeze(BTC) ## obtain flat vector of length D (number of neurons)
+ if not preserve_batch: ## calc population average
BTC = jnp.mean(BTC)
return BTC
+@deprecated(replaced_by=measure_breadthOfTuningCurve)
+def measure_breadth_TC(responses, preserve_batch=False): ## old BTC function
+ """
+ WARNING: this function is deprecated (renamed to `measure_breadthOfTuningCurve(.)`).
+ """
+ return measure_breadthOfTuningCurve(responses, preserve_batch)
+
+@partial(jit, static_argnums=[2])
+def measure_Pearson_correlation(a, b, preserve_batch=False):
+ """
+ Computes general Pearson correlation `r` between two arbitrary signal sets.
+ Tracks raw positive/negative directional correlations and assumes time steps are along `axis=1`.
+
+ Args:
+ a: response set/signals a (shape: N x T x D)
+
+ b: response set/signals b (shape: N x T x D)
+
+ preserve_batch: if True, returns one score per sample across batch.
+ (Default: False), otherwise, returns scalar average score.
+
+ Returns:
+ A (N x 1) Pearson correlation score vector (one score per sample) OR a single
+ global scalar average Pearson score for the entire signal set.
+ """
+ t_axis = 1 ## temporal axis
+
+ ## calculate means of raw signals
+ mean_a = jnp.mean(a, axis=t_axis, keepdims=True)
+ mean_b = jnp.mean(b, axis=t_axis, keepdims=True)
+ ## calculate deviation & variance/covariance terms
+ dev_a = a - mean_a
+ dev_b = b - mean_b
+ covariance = jnp.sum(dev_a * dev_b, axis=t_axis)
+ var_a = jnp.sum(jnp.square(dev_a), axis=t_axis)
+ var_b = jnp.sum(jnp.square(dev_b), axis=t_axis)
+
+ ## now compute Pearson correlation
+ denominator = jnp.sqrt(var_a * var_b)
+ safe_denominator = jnp.where(denominator == 0.0, 1e-7, denominator)
+ r_per_neuron = covariance / safe_denominator
+ r_per_neuron = jnp.where(denominator == 0.0, 0.0, r_per_neuron) ## handle silent/dead channels
+
+ if preserve_batch:
+ ## collapse neuron dimension to return one score per batch sample
+ r_mu = jnp.mean(r_per_neuron, axis=1, keepdims=True)
+ return r_mu
+ return jnp.mean(r_per_neuron) ## output global scalar score (averaged over neurons and batch)
+
+
@partial(jit, static_argnums=[1])
def measure_gini_index(codes, preserve_batch=True):
"""
@@ -164,8 +265,7 @@ def measure_sparsity(codes, tolerance=0., preserve_batch=True, flip_measure=Fals
rho = jnp.mean(rho)
return rho
-#@partial(jit, static_argnums=[2])
-def analyze_scores(mu, y, extract_label_indx=True): ## examines classifcation statistics
+def analyze_categorization_performance(mu, y, extract_label_indx=True): ## examines classifcation statistics
"""
Analyzes a set of prediction matrix and target/ground-truth matrix or vector.
@@ -204,6 +304,13 @@ def analyze_scores(mu, y, extract_label_indx=True): ## examines classifcation st
## output analysis statistics
return conf_matrix, precision, recall, misses, acc, adj_acc
+@deprecated(replaced_by=analyze_categorization_performance) ## deprecated name of funct was "analyze_scores(.)"
+def analyze_scores(*args, **kwargs):
+ """
+ WARNING: this function is deprecated (renamed to `analyze_categorization_performance(.)`).
+ """
+ return analyze_categorization_performance(*args, **kwargs)
+
@partial(jit, static_argnums=[2])
def measure_ACC(mu, y, extract_label_indx=True): ## measures/calculates accuracy
"""
diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py
index aa8ef9cb..0e2306ee 100755
--- a/ngclearn/utils/model_utils.py
+++ b/ngclearn/utils/model_utils.py
@@ -256,6 +256,20 @@ def chebyshev_norm(d, axis=-1, keepdims=False):
dist_vals = jnp.max(abs_diff, axis=axis, keepdims=keepdims)
return dist_vals
+@partial(jit, static_argnums=[2])
+def sample_values_by_group(skey, values, n_groups): ## multi-group value sampling co-routine
+ s = values
+ K = n_groups
+ batch_size, num_neurons = s.shape
+ L = num_neurons // K
+ s_reshaped = s.reshape((batch_size, K, L))
+ m_switch = (jnp.sum(s_reshaped, axis=2, keepdims=True) > 0.).astype(jnp.float32)
+ rS = s_reshaped * random.uniform(skey, s_reshaped.shape)
+ winning_indices = jnp.argmax(rS, axis=2)
+ rS_one_hot = nn.one_hot(winning_indices, num_classes=L, dtype=jnp.float32)
+ s_partitioned = s_reshaped * (1. - m_switch) + rS_one_hot * m_switch
+ return s_partitioned.reshape((batch_size, num_neurons))
+
@jit
def binarize(data, threshold=0.5):
"""
@@ -657,9 +671,11 @@ def swish(x, beta):
def d_swish(x, beta):
# df/dx = beta * [ 1/(exp(-x) + 1) + (exp(-x) * x) / (exp(-x) + 1)^2]
# df/dx = beta * sigmoid(x * beta) * (1 - sigmoid(x) * beta)
- exp_neg_x = jnp.exp(-x)
- _x = (1./(exp_neg_x + 1.)) + (exp_neg_x * x)/jnp.square(exp_neg_x+1)
- return _x * beta
+ # exp_neg_x = jnp.exp(-x)
+ # _x = (1./(exp_neg_x + 1.)) + (exp_neg_x * x)/jnp.square(exp_neg_x+1)
+ # return _x * beta
+ return sigmoid(x * beta) + beta * x * d_sigmoid(x * beta)
+
@jit
def silu(x):
@@ -913,20 +929,22 @@ def eye_wrapped(N, k, values):
return matrix.at[row_indices, col_indices].set(values) ## Fill diagonal using efficient indexing
@partial(jit, static_argnums=[1, 2, 3, 4])
-def normalize_block_matrix(matrix, block_size, order=2, axis=0, norm_targ=1.):
+def normalize_block_matrix(
+ matrix, block_size, order=2, axis=0, norm_targ=1.
+):
"""
- Normalizes columns of blocks within a matrix.
+ Normalizes blocks within a block-matrix.
Args:
matrix: 2D JAX Array (M, N)
- block_size: Tuple (block_rows, block_cols)
+ block_size: block shape tuple (n_block_rows, n_block_cols)
- order:
+ order: order (p) of p-norm
axis: (relative) axis for normalization within block; 0 -> by rows, 1 -> by cols
- norm_targ:
+ norm_targ: target p-norm value to enforce
Returns:
block-normalized (M, N) matrix
@@ -939,18 +957,18 @@ def normalize_block_matrix(matrix, block_size, order=2, axis=0, norm_targ=1.):
## else, we leave row-axis as target for normalization
M, N = matrix.shape
r_blk, c_blk = block_size
- # Reshape to 4D to isolate blocks: (num_blocks_row, block_rows, num_blocks_col, block_cols)
+ ## reshape to 4D to isolate blocks: (num_blocks_row, block_rows, num_blocks_col, block_cols)
reshaped = matrix.reshape(M // r_blk, r_blk, N // c_blk, c_blk)
- # Transpose to group block data: (num_blocks_row, num_blocks_col, block_rows, block_cols)
+ ## transpose to group block data: (num_blocks_row, num_blocks_col, block_rows, block_cols)
transposed = jnp.transpose(reshaped, (0, 2, 1, 3))
# Calculate norm for each column "w/in" each block
## (over axis 2 -> block_rows); (over axis 3 -> block_cols)
norms = jnp.linalg.norm(transposed, ord=order, axis=_tensor_axis, keepdims=True)
#normalized_blocks = jnp.divide(transposed, norms + 1e-8) ## normalize (w/ safe-division)
normalized_blocks = transposed * (norm_targ/(norms + 1e-8))
- # Reverse transpose: (num_blocks_row, block_rows, num_blocks_col, block_cols)
+ ## reverse-transpose: (num_blocks_row, block_rows, num_blocks_col, block_cols)
reverted = jnp.transpose(normalized_blocks, (0, 2, 1, 3))
- # Reshape back to original 2D shape (M, N)
+ ## reshape back to original 2D shape (M, N)
return reverted.reshape(M, N)
@partial(jit, static_argnums=[2, 3])
diff --git a/ngclearn/utils/surrogate_fx.py b/ngclearn/utils/surrogate_fx.py
index 3371829f..9b4d77ed 100644
--- a/ngclearn/utils/surrogate_fx.py
+++ b/ngclearn/utils/surrogate_fx.py
@@ -32,7 +32,7 @@ def spike_fx(v, thr):
@jit
def d_spike_fx(v, thr):
return v * 0 + 1.
- if get_surr_fx == True:
+ if get_surr_fx:
return spike_fx, spike_fx, d_spike_fx
else:
return spike_fx, d_spike_fx
@@ -48,12 +48,44 @@ def triangular_estimator(get_surr_fx=False):
@jit
def spike_fx(v, thr):
return (v > thr).astype(jnp.float32)
+ # @jit
+ # def d_spike_fx1(v, thr=0., alpha=0.2):
+ # mask = (v < thr).astype(jnp.float32)
+ # dfx = mask * thr - (1. - mask) * thr
+ # return dfx
@jit
- def d_spike_fx(v, thr):
- mask = (v < thr).astype(jnp.float32)
- dfx = mask * thr - (1. - mask) * thr
- return dfx
- if get_surr_fx == True:
+ def d_spike_fx(v, thr=0., alpha=2.0):
+ ## alpha: controls the width of triangle base (larger alpha = narrower window).
+ v_centered = v - thr
+ # Peak value is bounded, slope drops off proportional to alpha
+ gradient = 1.0 - alpha * jnp.abs(v_centered)
+ return jnp.maximum(0.0, gradient)
+ if get_surr_fx:
+ return spike_fx, spike_fx, d_spike_fx
+ else:
+ return spike_fx, d_spike_fx
+
+def explicit_triangular_estimator(get_surr_fx=False):
+ """
+ The explicit triangular surrogate gradient estimator for binary spike emission.
+ This is controlled by a "window" argument, instead of the alpha-scaled variant
+ "triangular_estimator()".
+
+ Returns:
+ ( spike_fx(x, thr), d_spike_fx(x, thr) ) OR
+ ( spike_fx(x, thr), surr_fx(x, thr, args), d_spike_fx(x, thr, args) )
+ """
+ @jit
+ def spike_fx(v, thr):
+ return (v > thr).astype(jnp.float32)
+ @jit
+ def d_spike_fx(v, thr=0., window=5.0):
+ ## sloped surrogate gate with an explicit window size
+ distance = jnp.abs(v - thr)
+ ## linear ramp down to 0.0 at edge of window
+ gradient = 1.0 - (distance / window)
+ return jnp.maximum(0.0, gradient) ## clip negative values to 0.0
+ if get_surr_fx:
return spike_fx, spike_fx, d_spike_fx
else:
return spike_fx, d_spike_fx
@@ -79,15 +111,38 @@ def surr_fx(v, alpha=2.):
return jnp.arctan(v * (alpha/2.) * pi) * (1. / pi)
@jit
def d_spike_fx(v, thr=0., alpha=2.):
+ v_centered = v - thr
pi = jnp.pi
- divTerm = jnp.square(v * pi * (alpha/2.))
+ divTerm = jnp.square(v_centered * pi * (alpha/2.))
dfx = (1./(1. + divTerm)) * (1./pi)
return dfx
- if get_surr_fx == True:
+ if get_surr_fx:
return spike_fx, surr_fx, d_spike_fx
else:
return spike_fx, d_spike_fx
+def boxcar_estimator(get_surr_fx=False):
+ """
+ The box-car surrogate gradient estimator for binary spike emission.
+
+ Returns:
+ ( spike_fx(x, thr), d_spike_fx(x, thr) ) OR
+ ( spike_fx(x, thr), surr_fx(x, thr, args), d_spike_fx(x, thr, args) )
+ """
+ @jit
+ def spike_fx(v, thr):
+ return (v > thr).astype(jnp.float32)
+ @jit
+ def d_spike_fx(v, thr=0., window=5.0):
+ ## returns 1.0 if voltage is w/in a "window" millivolts of threshold; else 0
+ distance = jnp.abs(v - thr)
+ dfx = jnp.where(distance <= window, 1.0, 0.0) ## box-car gate
+ return dfx
+ if get_surr_fx:
+ return spike_fx, spike_fx, d_spike_fx
+ else:
+ return spike_fx, d_spike_fx
+
def secant_lif_estimator(get_surr_fx=False):
"""
Surrogate function for computing derivative of (binary) spike function
diff --git a/ngclearn/utils/viz/dim_reduce.py b/ngclearn/utils/viz/dim_reduce.py
index 3300feef..864c7da7 100755
--- a/ngclearn/utils/viz/dim_reduce.py
+++ b/ngclearn/utils/viz/dim_reduce.py
@@ -32,7 +32,8 @@ def extract_tsne_latents(
vectors,
perplexity=30,
n_pca_comp=32,
- batch_size=500
+ batch_size=500,
+ verbose=1
): ## tSNE mapping routine
"""
Projects collection of K vectors (stored in a matrix) to a two-dimensional (2D) visualization space via the
@@ -67,7 +68,7 @@ def extract_tsne_latents(
z_2D = ipca.transform(vectors)
print(" PCA.lat.shape = ",z_2D.shape)
print(" > Finishing projection via t-SNE...")
- z_2D = TSNE(n_components=2,perplexity=perplexity, verbose=1).fit_transform(z_2D)
+ z_2D = TSNE(n_components=2,perplexity=perplexity, verbose=verbose).fit_transform(z_2D)
else:
z_2D = vectors
return z_2D
diff --git a/ngclearn/utils/viz/synapse_plot.py b/ngclearn/utils/viz/synapse_plot.py
index f62a47da..cae28e5c 100644
--- a/ngclearn/utils/viz/synapse_plot.py
+++ b/ngclearn/utils/viz/synapse_plot.py
@@ -9,6 +9,93 @@
import jax.numpy as jnp
+def visualize_macro_grid( ## more complex filter visualization co-routine
+ thetas,
+ sizes,
+ macro_grid_shape,
+ prefix,
+ order=None,
+ suffix='.jpg',
+ contrast_by_data=True
+):
+ """
+ Stitches filter sets into a 2D Grid-of-Grids layout with bright white padding.
+
+ Args:
+ thetas:
+ sizes:
+ macro_grid_shape:
+ prefix:
+ order:
+ suffix:
+ contrast_by_data:
+
+ Returns:
+
+ """
+ if order is None:
+ order = ['C' for _ in range(len(thetas))]
+
+ H_macro, W_macro = macro_grid_shape
+ num_banks = len(thetas)
+
+ T_sample = thetas[0].T
+ filters_per_bank = T_sample.shape[0]
+ f_cols = int(math.ceil(math.sqrt(filters_per_bank)))
+ f_rows = int(math.ceil(filters_per_bank / f_cols))
+
+ p_h, p_w = sizes[0]
+ ## total pixel height/width of an individual composite filter bank block
+ bank_px_h = f_rows * p_h
+ bank_px_w = f_cols * p_w
+
+ ## Use stark white background; initialize the entire master sheet canvas with 1.0 (White in grayscale)
+ pad = 2
+ canvas_h = H_macro * bank_px_h + (H_macro + 1) * pad
+ canvas_w = W_macro * bank_px_w + (W_macro + 1) * pad
+ master_canvas = np.ones((canvas_h, canvas_w))
+
+ for b_idx in range(num_banks):
+ m_row = b_idx // W_macro
+ m_col = b_idx % W_macro
+ if m_row >= H_macro:
+ break
+
+ T = thetas[b_idx].T
+ b_start_y = m_row * bank_px_h + (m_row + 1) * pad
+ b_start_x = m_col * bank_px_w + (m_col + 1) * pad
+ for f_idx in range(filters_per_bank):
+ if f_idx >= T.shape[0]:
+ break
+ i_row = f_idx // f_cols
+ i_col = f_idx % f_cols
+ single_filter = np.reshape(T[f_idx, :], (p_h, p_w), order=order[b_idx])
+
+ ## shift values to map correctly to the bone color scheme; max absolute value normalization
+ max_val = float(np.max(np.abs(single_filter)))
+ if max_val > 0:
+ single_filter = single_filter / max_val
+
+ y_loc = b_start_y + (i_row * p_h)
+ x_loc = b_start_x + (i_col * p_w)
+ master_canvas[y_loc:y_loc + p_h, x_loc:x_loc + p_w] = single_filter
+
+ ## render out the crisp grid sheet
+ plt.figure(figsize=(10, 10), dpi=300)
+
+ ## use vmin=-1.0 and vmax=1.0 so that the 1.0 canvas background registers as absolute white
+ max_val = 1.
+ min_val = -1.
+ if contrast_by_data:
+ max_val = float(jnp.max(jnp.abs(master_canvas)))
+ min_val = float(jnp.min(jnp.abs(master_canvas)))
+ plt.imshow(master_canvas, cmap=plt.cm.bone, interpolation='nearest', vmin=min_val, vmax=max_val)
+ plt.axis("off")
+ plt.savefig(prefix + suffix, bbox_inches='tight', pad_inches=0.0)
+ plt.clf()
+ plt.close()
+
+
def visualize(
thetas,
sizes,
@@ -56,9 +143,11 @@ def visualize(
point = start + 1 + i + (r * extra)
plt.subplot(n_rows_total, n_cols_total, point)
_filter = T[i, :]
+ max_val = float(jnp.max(jnp.abs(_filter)))
+ min_val = float(jnp.min(jnp.abs(_filter)))
plt.imshow(
np.reshape(_filter, (sizes[idx][0], sizes[idx][1]), order=order[idx]),
- cmap=plt.cm.bone, interpolation='nearest'
+ cmap=plt.cm.bone, interpolation='nearest', vmin=min_val, vmax=max_val
)
plt.axis("off")
diff --git a/pyproject.toml b/pyproject.toml
index 3458b230..1c14055a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" # using setuptool building engine
[project]
name = "ngclearn"
-version = "3.1.1"
+version = "3.2.0"
description = "Simulation software for building and analyzing computational neuroscience models, brain-inspired computing systems, and NeuroAI agents."
authors = [
{name = "Alexander Ororbia", email = "ago@cs.rit.edu"},
diff --git a/tests/components/synapses/modulated/test_REINFORCESynapse.py b/tests/components/synapses/modulated/test_REINFORCESynapse.py
index f0235789..493dd4bc 100644
--- a/tests/components/synapses/modulated/test_REINFORCESynapse.py
+++ b/tests/components/synapses/modulated/test_REINFORCESynapse.py
@@ -4,7 +4,7 @@
from jax import numpy as jnp, random, jit
import numpy as np
np.random.seed(42)
-from ngclearn.components.synapses.modulated.REINFORCESynapse import REINFORCESynapse, _gaussian_logpdf
+from ngclearn.components.synapses.modulated.REINFORCESynapse import REINFORCESynapse #, _gaussian_logpdf
from numpy.testing import assert_array_equal
from ngclearn import Context, MethodProcess
@@ -58,7 +58,7 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
std = jnp.exp(logstd.clip(-10.0, 2.0))
sample = jax.random.normal(seed, mean.shape) * std + mean
sample = jnp.clip(sample, mu_out_min, mu_out_max)
- logp = _gaussian_logpdf(jax.lax.stop_gradient(sample), mean, std).sum(-1)
+ logp = REINFORCESynapse._gaussian_logpdf(jax.lax.stop_gradient(sample), mean, std).sum(-1)
return (-logp * outputs).mean() * 1e-2
grad_fn = jax.value_and_grad(fn)
@@ -177,7 +177,7 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
std = scalar_stddev
sample = jax.random.normal(seed, mean.shape) * std + mean
sample = jnp.clip(sample, mu_out_min, mu_out_max)
- logp = _gaussian_logpdf(jax.lax.stop_gradient(sample), mean, std).sum(-1)
+ logp = REINFORCESynapse._gaussian_logpdf(jax.lax.stop_gradient(sample), mean, std).sum(-1)
return (-logp * outputs).mean() * 1e-2
grad_fn = jax.value_and_grad(fn)