Major release update (to 2.0.0) (#100)

* add initial patch mask features

* minor edit to bern-cell

* fixed bernoulli error cell

* example rate cell test

* made some corrections to bern err-cell and heb syn

* made some corrections to bern err-cell and heb syn

* cleaned up bern-cell, hebb-syn

* minor mod to model-utils

* attempted rewrite of bernoulli-cell

* got bernoulli-cell rewritten and unit-tested

* edit to bern-cell

* bernoulli and poisson cells revised, unit-tested

* latency-cell refactored and unit-tested

* refactored Rate Cell

* minor revisions to input-encoders, revised phasor-cell w/ unit-test

* revised and add unit-test for varTrace

* revised and added unit-test for exp-kernel

* revised and added unit-test for exp-kernel

* revised slif cell w/ unit-test; needed mod to diffeq

* revised slif-cell w/ unit-test; cleaned up ode_utils to play nicer w/ new sim-lib

* revised lif-cell w/ unit-test

* revised unit-tests to pass globally; some minor patches to phasor-cell and lif

* minor cleanup of unit-test for phasor

* revised if-cell w/ unit-test

* revised if-cell w/ unit-test

* revised quad-lif w/ unit-test

* revised adex-cell w/ unit test, minor cleanup of quad-lif

* minor edit to adex unit-test

* refactor bernoulli, laplacian, and rewarderror cells

* revised raf-cell w/ unit test; fixed typos/mistakes in all spiking cells

* revised wtas-cell w/ unit test

* revised fh-cell w/ unit test

* revised izh-cell w/ unit test

* patched ode_utils backend wrt jax, cleaned up unit-tests, added disable flag for phasor-cell

* update rate cell

* fix test rate cell

* update test for bernoulli cell

* update refactoring for gaussian error cell

* update unit testing for all graded neurons

* wrote+unit-test of hodgkin-huxley spike cell, minor tweaks/clean-up elsewhere

* added rk2 support for H-H cell

* update rate cell and fix bug of passing a tuple of (jax Array -- not hashable)  to jax jit functions. Basically, simplify the codebase by using a hashmap of functions

* update test rate cell

* refactored dense and trace-stdp syn w/ unit-test

* refactored exp-stdp syn w/ unit-test

* refactored event-stdp w/ unit-test

* cleanup of stdp-syn

* refactored bcm syn w/ unit-test

* refactored stp-syn with unit-test

* cleaned up modulated

* refactored mstdp-et syn w/ unit-test

* refactored lava components to new sim-lib

* refactored conv/hebb-conv syn w/ unit-test

* refactored/revised hebb-deconv syn w/ unit-test

* revised/refactored hebb/stdp conv/deconv syn w/ unit-tests

* updated modeling doc to point to hodgkin-huxley cell

* updated modeling docs

* fixed typo in adex-cell tutorial doc

* revised tutorials to reflect new sim-lib config/syntax

* revised tutorials to reflect new sim-lib config/syntax

* patched docs to reflect revisions/refactor

* tweaked requirements in prep for major release

* cleaned up a few unit tests to use deterministic syn init vals

* mod to requirements

* nudge toml to upcoming 2.0.0

* update to support docs in prep for 2.0.0

* update patched synapses and their test cases

* cleaned up syn modeling doc

* push hebbian synapse

* push reinforce synapse

* push np seed

* patched minor prior None arg issue in hebb-syn

* moved reinforce-syn to right spot

* update reinforce synapse and testing

* tweaked trace-stdp and mstdpet

* patched mstdpet unit-test

* update reinforce synapse and test cases

* add reinforce synapse fix

* minor mod to mstdpet

* update test code for more than 1 steps

* Updated monitors

* patched tests to use process naming

* Added wrapper for reset and advance_state

* Added a JaxProcess

Added Jax Process to allow for scanning over the process.

* update the old rate cell

* update old hebbian synapse

* minor edit to if-cell

* ported over adex tutorial to new ngclearn format

* hh-cell supports rk4 integration

* clean up and integrated hodgkin-huxley mini lesson in neurocog tutorials

* Update jaxProcess.py

Updated the jax process to allow for more configurations of inputs.

* update working reinforce synapse

* update correct reinforce and testing

* update documentation

* update features, documentation, and testing

* update testing for REINFORCE cell

* update code and test

* update code

* add clipping gradient to model utils

* update reinforce cell to the new model utils clip

* major cleanup in prep for merge over to main/prep for major release

* update test cases

* update to require file in docs

---------

Co-authored-by: Viet Dung Nguyen <vietdungnguyen233@gmail.com>
Co-authored-by: Alexander Ororbia <ago@hal3.cs.rit.edu>
Co-authored-by: Will Gebhardt <will@gebhardts.net>

Author: 4 people

README.md

@@ -2,19 +2,15 @@
 
 <img src="docs/images/ngc-learn-logo.png" width="300">
 
-<b>ngc-learn</b> is a Python library for building, simulating, and analyzing
-biomimetic systems, neurobiological agents, spiking neuronal networks,
-predictive coding circuitry, and models that learn via biologically-plausible
-forms of credit assignment. This simulation toolkit is built on top of JAX and is
-distributed under the 3-Clause BSD license.
+<b>ngc-learn</b> is a Python library for building, simulating, and analyzing biophysical / neurobiological systems, spiking neuronal networks, predictive coding circuitry, and biomimetic (NeuroAI) agents that learn in a biologically-plausible manner. This simulation toolkit, meant to support computational neuroscience and brain-inspired computing research, is built on top of JAX and is distributed under the 3-Clause BSD license.
 
 It is currently maintained by the
 <a href="https://www.cs.rit.edu/~ago/nac_lab.html">Neural Adaptive Computing (NAC) laboratory</a>.
 
 ## <b>Documentation</b>
 
 Official documentation, including tutorials, can be found
-<a href="https://ngc-learn.readthedocs.io/en/latest/#">here</a>. The model museum repo,
+<a href="https://ngc-learn.readthedocs.io/en/latest/#">here</a>. The model museum repo (ngc-museum),
 which implements several historical models, can be found
 <a href="https://github.com/NACLab/ngc-museum">here</a>.
 
@@ -36,8 +32,8 @@ ngc-learn requires:
 1) Python (>=3.10)
 2) NumPy (>=1.26.0)
 3) SciPy (>=1.7.0)
-4) ngcsimlib (>=0.3.b4), (visit official page <a href="https://github.com/NACLab/ngc-sim-lib">here</a>)
-5) JAX (>= 0.4.28) (to enable GPU use, make sure to install one of the CUDA variants)
+4) ngcsimlib (>=1.0.0), (visit official page <a href="https://github.com/NACLab/ngc-sim-lib">here</a>)
+5) JAX (>=0.4.28) (to enable GPU use, make sure to install one of the CUDA variants)
 <!--
 5) scikit-learn (>=1.3.1) if using `ngclearn.utils.density`
 6) matplotlib (>=3.4.3) if using `ngclearn.utils.viz`
@@ -46,7 +42,7 @@ ngc-learn requires:
 -->
 
 ---
-ngc-learn 1.2.beta2 and later require Python 3.10 or newer as well as ngcsimlib >=0.3.b4.
+ngc-learn 2.0.0 and later require Python 3.10 or newer as well as ngcsimlib >=1.0.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).
@@ -66,7 +62,7 @@ running the above pip command if you want to use the GPU version.
 
 The documentation includes more detailed
 <a href="https://ngc-learn.readthedocs.io/en/latest/installation.html">installation instructions</a>.
-Note that this library was developed on Ubuntu 20.04 and tested on Ubuntu(s) 18.04 and 20.04.
+Note that this library was developed on Ubuntu 20.04/22.04 and tested on Ubuntu(s) 20.04 and 22.04.
 
 If the installation was successful, you should see the following if you test
 it against your Python interpreter, i.e., run the <code>$ python</code> command

docs/museum/snn_bfa.md

@@ -233,7 +233,7 @@ of $1000$ `SLIF` cells) similar to the one below:
 <img src="../images/museum/bfa_snn/bfasnn_codes.jpg" width="450" />
 
 Intriguingly, we see that the latent codes represented by the BFA-SNN's hidden
-layer spikes yield a rather (piecewise) linearly-separable transformation
+layer spikes yield a rather (piecewise) linearly-separable representation
 of the input digits, making the process of mapping inputs to label vectors
 much easier for the model's second layer of classification LIF units.
 Note that, in the `BFA_SNN` model exhibit class, we estimated

docs/requirements.txt

@@ -5,7 +5,7 @@ numpy>=1.26.0
 scikit-learn>=0.24.2
 scipy>=1.7.0
 matplotlib>=3.8.0
-jax>=0.4.18
-jaxlib>=0.4.18
+jax>=0.4.28
+jaxlib>=0.4.28
 imageio>=2.31.5
-ngcsimlib>=0.3.b4
+ngcsimlib>=1.0.0

docs/tutorials/model_basics/evolving_synapses.md

@@ -19,7 +19,7 @@ We do this specifically as follows:
 ```python
 from jax import numpy as jnp, random, jit
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 from ngclearn.components import HebbianSynapse, RateCell
 import ngclearn.utils.weight_distribution as dist
 
@@ -49,16 +49,16 @@ with Context("Circuit") as circuit:
   Wab.post << b.zF
 
   ## create and compile core simulation commands  
-  evolve_process = (Process()
+  evolve_process = (JaxProcess()
                     >> a.evolve)
   circuit.wrap_and_add_command(jit(evolve_process.pure), name="evolve")
 
-  advance_process = (Process()
+  advance_process = (JaxProcess()
                      >> a.advance_state)
   circuit.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-  reset_process = (Process()
-                 >> a.reset)
+  reset_process = (JaxProcess()
+                   >> a.reset)
   circuit.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
   ## set up non-compiled utility commands
@@ -83,7 +83,6 @@ for ts in range(x_seq.shape[1]):
   circuit.advance(t=ts*1., dt=1.)
   circuit.evolve(t=ts*1., dt=1.)
   print(" {}: input = {} ~> Wab = {}".format(ts, x_t, Wab.weights.value))
-
 ```
 
 Your code should produce the same output (towards the bottom):

docs/tutorials/model_basics/model_building.md

@@ -10,6 +10,7 @@ While building our dynamical system we will set up a Context and then add the th
 ```python
 from jax import numpy as jnp, random
 from ngclearn import Context
+from ngclearn.utils import JaxProcess
 from ngcsimlib.compilers.process import Process
 from ngclearn.components import RateCell, HebbianSynapse
 import ngclearn.utils.weight_distribution as dist
@@ -71,13 +72,13 @@ This is simply done with the use of the following convenience function calls:
 
 ```python
     ## configure desired commands for simulation object
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> a.reset
                      >> Wab.reset
                      >> b.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> a.advance_state
                        >> Wab.advance_state
                        >> b.advance_state)

docs/tutorials/neurocog/adex_cell.md

@@ -24,7 +24,7 @@ import numpy as np
 
 from ngclearn.utils.model_utils import scanner
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components.neurons.spiking.adExCell import AdExCell
 
@@ -48,11 +48,11 @@ with Context("Model") as model:
     )
 
     ## create and compile core simulation commands
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/error_cell.md

@@ -54,7 +54,7 @@ The code you would write amounts to the below:
 ```python
 from jax import numpy as jnp, jit
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process, transition
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components.neurons.graded.gaussianErrorCell import GaussianErrorCell
 
@@ -64,11 +64,11 @@ T = 5  ## number time steps to simulate
 with Context("Model") as model:
     cell = GaussianErrorCell("z0", n_units=3)
 
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/fitzhugh_nagumo_cell.md

@@ -18,7 +18,7 @@ from jax import numpy as jnp, random, jit
 import numpy as np
 
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components.neurons.spiking.fitzhughNagumoCell import FitzhughNagumoCell
 
@@ -40,11 +40,11 @@ with Context("Model") as model:
                               gamma=gamma, v0=v0, w0=w0, integration_type="euler")
 
     ## create and compile core simulation commands
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/hebbian.md

@@ -38,11 +38,11 @@ Wab.post << b.zF
 as well as (a bit later in the model construction code):
 
 ```python
-evolve_process = (Process()
+evolve_process = (JaxProcess()
                   >> a.evolve)
 circuit.wrap_and_add_command(jit(evolve_process.pure), name="evolve")
 
-advance_process = (Process()
+advance_process = (JaxProcess()
                    >> a.advance_state)
 circuit.wrap_and_add_command(jit(advance_process.pure), name="advance")
 ```

docs/tutorials/neurocog/hodgkin_huxley_cell.md

@@ -24,7 +24,7 @@ import numpy as np
 
 from ngclearn.utils.model_utils import scanner
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components.neurons.spiking.hodgkinHuxleyCell import HodgkinHuxleyCell
 
@@ -52,11 +52,11 @@ with Context("Model") as model:
     )
 
     ## create and compile core simulation commands
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/input_cells.md

@@ -40,7 +40,7 @@ spike train over $100$ steps in time as follows:
 ```python
 from jax import numpy as jnp, random, jit
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 
 from ngclearn.utils.viz.raster import create_raster_plot
 ## import model-specific mechanisms
@@ -56,11 +56,11 @@ T = 100  ## number time steps to simulate
 with Context("Model") as model:
     cell = BernoulliCell("z0", n_units=10, key=subkeys[0])
 
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/izhikevich_cell.md

@@ -20,7 +20,7 @@ from jax import numpy as jnp, random, jit
 import numpy as np
 
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components.neurons.spiking.izhikevichCell import IzhikevichCell
 
@@ -44,11 +44,11 @@ with Context("Model") as model:
                           integration_type="euler", v0=v0, w0=w0, key=subkeys[0])
 
     ## create and compile core simulation commands
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/lif.md

@@ -25,7 +25,7 @@ cell, you would write code akin to the following:
 from jax import numpy as jnp, random, jit
 
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components.neurons.spiking.LIFCell import LIFCell
 from ngclearn.utils.viz.spike_plot import plot_spiking_neuron
@@ -47,11 +47,11 @@ with Context("Model") as model:
                    refract_time=2., key=subkeys[0])
 
     ## create and compile core simulation commands
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/mod_stdp.md

@@ -42,7 +42,7 @@ and the required compiled simulation and dynamic commands, can be done as follow
 ```python 
 from jax import numpy as jnp, random, jit
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components import (TraceSTDPSynapse, MSTDPETSynapse,
                                  RewardErrorCell, VarTrace)
@@ -75,13 +75,13 @@ with Context("Model") as model:
     tr1 = VarTrace("tr1", n_units=1, tau_tr=tau_post, a_delta=Aminus)
     rpe = RewardErrorCell("r", n_units=1, alpha=0.)
 
-    evolve_process = (Process()
+    evolve_process = (JaxProcess()
                       >> W_stdp.evolve
                       >> W_mstdp.evolve
                       >> W_mstdpet.evolve)
     model.wrap_and_add_command(jit(evolve_process.pure), name="evolve")
 
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> tr0.advance_state
                        >> tr1.advance_state
                        >> rpe.advance_state
@@ -90,7 +90,7 @@ with Context("Model") as model:
                        >> W_mstdpet.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> W_stdp.reset
                      >> W_mstdp.reset
                      >> W_mstdpet.reset

docs/tutorials/neurocog/rate_cell.md

@@ -16,7 +16,7 @@ specifically the rate-cell (RateCell). Let's start with the file's header
 
 ```python
 from jax import numpy as jnp, random, jit
-from ngcsimlib.compilers.process import Process, transition
+from ngclearn.utils import JaxProcess
 from ngcsimlib.context import Context
 ## import model-specific elements
 from ngclearn.components.neurons.graded.rateCell import RateCell
@@ -40,11 +40,11 @@ with Context("Model") as model:  ## model/simulation definition
                     prior=("gaussian", gamma), integration_type="euler", key=subkeys[0])
 
     ## instantiate desired core commands that drive the simulation
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> cell.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> cell.reset)
     model.wrap_and_add_command(jit(reset_process.pure), name="reset")
 

docs/tutorials/neurocog/short_term_plasticity.md

@@ -60,7 +60,7 @@ STF-dominated dynamics):
 ```python 
 from jax import numpy as jnp, random, jit
 from ngcsimlib.context import Context
-from ngcsimlib.compilers.process import Process
+from ngclearn.utils import JaxProcess
 ## import model-specific mechanisms
 from ngclearn.components import PoissonCell, STPDenseSynapse, LIFCell
 import ngclearn.utils.weight_distribution as dist
@@ -98,13 +98,13 @@ with Context("Model") as model:
     W.inputs << z0.outputs ## z0 -> W
     z1.j << W.outputs ## W -> z1
 
-    advance_process = (Process()
+    advance_process = (JaxProcess()
                        >> z0.advance_state
                        >> W.advance_state
                        >> z1.advance_state)
     model.wrap_and_add_command(jit(advance_process.pure), name="advance")
 
-    reset_process = (Process()
+    reset_process = (JaxProcess()
                      >> z0.reset
                      >> z1.reset
                      >> W.reset)