Welcome to Rockpool

Rockpool is a Python package for working with dynamical neural network architectures, particularly for designing event-driven networks for Neuromorphic computing hardware. Rockpool provides a convenient interface for designing, training and evaluating recurrent networks, which can operate both with continuous-time dynamics and event-driven dynamics.

Rockpool is an open-source project managed by aiCTX AG.

About Rockpool

_images/logo_Color_trimmed.png

Rockpool is an open source project released by aiCTX AG.

https://gitlab.com/aiCTX/rockpool/badges/develop/pipeline.svghttps://img.shields.io/pypi/v/rockpool.svghttps://img.shields.io/conda/v/conda-forge/rockpool https://img.shields.io/pypi/pyversions/rockpool?logo=python https://img.shields.io/badge/code%20style-black-000000.svghttps://img.shields.io/pypi/dd/rockpool

About aiCTX

_images/aictx_logo.png

aiCTX is a Neuromorphic computing hardware and solutions startup, based in Zurich Switzerland. The company specializes in developing mixed-signal neuromorphic silicon hardware for neural simulation and signal processing; it develops software for interfacing with and configuring neuromorphic hardware; and develops solutions to analyse and process bio-signals. aiCTX is a commercial spin-off from the Institute of Neuroinformatics (INI), University of Zurich (UZH) and ETH Zurich (ETHZ).

About Noodle

_images/noodle.png

Noodle is the mascot of Rockpool. Noodle is a Nudibranch, Glaucus marginatus. Nudibranches are a group of amazing sea snails that shed their shells after the larval stage, to display an incredible array of forms, patterns and colours. Glaucus marginatus is a species found in the Pacific ocean, and often seen at beaches and in rock pools of the eastern Australian coast.

Photograph of Noodle is CC BY 2.0 Taro Taylor

Installing Rockpool

Base requirements

Rockpool requires Python 3.6, numpy, scipy and numba to install. These requirements will be installed by pip when installing Rockpool. We recommend using anaconda, miniconda or another environment manager to keep your Python dependencies clean.

Installation using pip

The simplest way to install Rockpool is by using pip to download and install the latest version from PyPI.

pip install rockpool

Dependencies

Rockpool has several dependencies for various aspects of the package. However, these dependencies are compartmentalised as much as possible. For example, Jax is required to use the Jax-backed layers (e.g. RecRateEulerJax); PyTorch is required to use the Torch-backed layers (e.g. RecIAFTorch), and so on. But if these dependencies are not available, the remainder of Rockpool is still usable.

To automatically install all the extra dependencies required by Rockpool, use the command

$ pip install rockpool[all]

Contributing

If you would like to contribute to Rockpool, then you should begin by forking the public repository at https://gitlab.com/ai-ctx/rockpool to your own account. Then clone your fork to your development machine

$ git clone https://gitlab.com/your-fork-location/rockpool.git rockpool

Install the package in development mode using pip

$ cd rockpool
$ pip install -e . --user

or

$ pip install -e .[all] --user

The main branch is development. You should commit your modifications to a new feature branch.

$ git checkout -b feature/my-feature develop
...
$ git commit -m 'This is a verbose commit message.'

Then push your new branch to your repository

$ git push -u origin feature/my-feature

Use the Black code formatter on your submission during your final commit. This is required for us to merge your changes. If your modifications aren’t already covered by a unit test, please include a unit test with your merge request. Unit tests go in the tests directory.

Then when you’re ready, make a merge request on gitlab.com, from the feature branch in your fork to https://gitlab.com/ai-ctx/rockpool.

Running tests

As part of the merge review process, we’ll check that all the unit tests pass. You can check this yourself (and probably should before making your merge request), by running the unit tests locally.

To run all the unit tests for Rockpool, use pytest:

$ pytest tests

Building documentation

The Rockpool documentation requires Sphinx, NBSphinx and Sphinx-autobuild. The commands

$ cd docs
$ make livehtml

Will compile the documentation and open a web browser to the local copy of the docs.

This page was generated from docs/basics/getting_started.ipynb. Interactive online version: Binder badge

Getting started with Rockpool

Rockpool is designed to let you design, simulate, train and test dynamical neural networks – in contrast to standard ANNs, these networks include explicit temporal dynamics and simulation of time. Rockpool contains several types of neuron simulations, including continuous-time (“rate”) models as well as event-driven spiking models. Rockpool supports several simulation back-ends, and layers with varying dynamics can be combined in the same network.

Importing Rockpool modules

[1]:
# - Switch off warnings
import warnings
warnings.filterwarnings('ignore')

# - Import classes to represent time series data
from rockpool import TimeSeries, TSContinuous, TSEvent

# - Import the `Network` base class
from rockpool import Network

# - Import some `Layer` classes to use
from rockpool.layers import RecRateEuler, PassThrough
[3]:
# - Import numpy
import numpy as np

# - Import the plotting library
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['figure.figsize'] = [12, 4]

Build a small network

We will define a very small recurrent network, with one input channel, ten recurrent units, and one output channel.

[4]:
# - Define the network size
input_size = 1
rec_size = 10
output_size = 1

# - Define weights
weights_in = np.random.rand(input_size, rec_size) - .5
weights_rec = np.random.randn(rec_size, rec_size) / rec_size
weights_out = np.random.rand(rec_size, output_size) - .5

# - Construct three layers
lyrInput = PassThrough(weights_in, name = "Input")
lyrRecurrent = RecRateEuler(weights_rec, name = "Recurrent")
lyrOut = PassThrough(weights_out, name = "Output")

# - Compose these into a network
net = Network(lyrInput, lyrRecurrent, lyrOut)

# - Display the network
print(net)
Network object with 3 layers
    PassThrough object: "Input" [1 TSContinuous in -> 10 internal -> 10 TSContinuous out]
    RecRateEuler object: "Recurrent" [10 TSContinuous in -> 10 internal -> 10 TSContinuous out]
    PassThrough object: "Output" [10 TSContinuous in -> 1 internal -> 1 TSContinuous out]

Define an input signal

In order to pass data through the network, we first need to generate some input data. In Rockpool, time series data is represented by .TimeSeries subclasses. See Working with time series data for more information.

Let’s build a simple white noise signal as an input.

[5]:
# - Build a time base
timebase = np.arange(1000)

# - Create a white noise time series
ts_white = TSContinuous(timebase, np.random.rand(timebase.size))

# - Plot the time series
plt.figure()
ts_white.plot()
plt.xlabel('Time')
plt.ylabel('Value')
plt.title('White noise');
_images/basics_getting_started_9_0.png

Stimulate the network

Now we can inject the white noise series into the network, and obtain the resulting network activity.

[6]:
dResponse = net.evolve(ts_white)
Network: Evolving layer `Input` with external input as input
Network: Evolving layer `Recurrent` with Input's output as input
Network: Evolving layer `Output` with Recurrent's output as input

dResponse is a dictionary containing the time series produced by the network during evolution. The signals in dResponse are named for the layer which produced them:

[7]:
dResponse
[7]:
{'external': non-periodic TSContinuous object `unnamed` from t=0.0 to 999.0. Samples: 1000. Channels: 1,
 'Input': non-periodic TSContinuous object `unnamed` from t=0.0 to 999.0. Samples: 1000. Channels: 10,
 'Recurrent': non-periodic TSContinuous object `unnamed` from t=0.0 to 999.0. Samples: 9991. Channels: 10,
 'Output': non-periodic TSContinuous object `unnamed` from t=0.0 to 999.0. Samples: 1000. Channels: 1}
[8]:
# - Define an auxilliary function to help us with plotting
def plot_signal(k, ts):
    plt.figure()
    ts.plot()
    plt.title(k)
    plt.xlabel('Time')
    plt.ylabel('Value')

# - Plot all the signals in turn
[ plot_signal(k, ts)
    for k, ts in zip(dResponse.keys(), dResponse.values())
];
_images/basics_getting_started_14_0.png
_images/basics_getting_started_14_1.png
_images/basics_getting_started_14_2.png
_images/basics_getting_started_14_3.png

Now you know how to build and simulate small networks, you can explore the other types of .Layer offered by Rockpool, especially event-driven layers. See Types of Layer available in Rockpool for more information, and :ref:’/tutorials/building_reservoir.ipynb` for a detailed view of building and training a reservoir.

Rockpool also offers facilities for training networks, and seamlessly using Neuromorphic compute hardware as simulation backends. See the tutorial on Event-based simulation on DynapSE hardware for more details.

This page was generated from docs/basics/time_series.ipynb. Interactive online version: Binder badge

Working with time series data

Concepts

In Rockpool, temporal data (“time series” data) is encapsulated in a set of classes that derive from .TimeSeries. Time series come in two basic flavours: “continuous” time series, which have been sampled at some set of time points but which represent values that can exist at any point in time; and “event” time series, which consist of discrete event times.

The .TimeSeries subclasses provide methods for extracting, resampling, shifting, trimming and manipulating time series data in a convenient fashion. Since Rockpool naturally deals with temporal dynamics and temporal data, TimeSeries objects are used to pass around time series data both as input and as output.

TimeSeries objects have an implicit shared time-base at \(t_0 = 0\) sec. However, they can easily be offset in time, concatenated, etc.

Housekeeping and import statements

[1]:
# - Import required modules and configure

# - Switch off warnings
import warnings
warnings.filterwarnings('ignore')

# - Required imports
import numpy as np

from rockpool.timeseries import (
    TimeSeries,
    TSContinuous,
    TSEvent,
    set_global_ts_plotting_backend,
)

# - Use HoloViews for plotting
import colorcet as cc
import holoviews as hv
hv.extension('bokeh')

%opts Curve [width=600]
%opts Scatter [width=600]

Continuous time series represented by TSContinuous

Continuous time series are represented by tuples \([t_{k}, a(t_{k})]\), where \(a(t_{k})\) is the amplitude of a signal, sampled at the time \(t_{k}\). A full time series is therefore the set of samples \([t_{k}, a(t_{k})]\) for \(k=1\dots K\).

Continuous time series in Rockpool are represented by the .TSContinuous class.

A time series is constructed by providing the sample times in seconds and the corresponding sample values. The full syntax for constructing a .TSContinuous object is given by

def __init__(
    self,
    times: Optional[ArrayLike] = None,
    samples: Optional[ArrayLike] = None,
    num_channels: Optional[int] = None,
    periodic: bool = False,
    t_start: Optional[float] = None,
    t_stop: Optional[float] = None,
    name: str = "unnamed",
    interp_kind: str = "linear",
)
[2]:
# - Build a time trace vector
duration = 10.
times = np.arange(0., duration, 0.1)
theta = times / duration * 2 * np.pi

# - Create a TSContinuous object containing a sin-wave time series
ts_sin = TSContinuous(
    times=times,
    samples=np.sin(theta),
    name="sine wave",
)
ts_sin
[2]:
non-periodic TSContinuous object `sine wave` from t=0.0 to 9.9. Samples: 100. Channels: 1

.TSContinuous objects provide a convenience plotting method .TSContinuous.plot for visualisation. This makes use of holoviews / bokeh or matplotlib plotting libraries, if available.

If both are available you can choose between them using the .timeseries.set_global_plotting_backend function.

[3]:
# - Set backend for Timeseries to holoviews
set_global_ts_plotting_backend("holoviews")

# # - Alternatively, it can be set for specific Timeseries instances
# ts_sin.set_plotting_backend("holoviews")

# - Plot the time series
ts_sin.plot()
Global plotting backend has been set to holoviews.
[3]:

.TSContinuous objects can represent multiple series simultaneously, as long as they share a common time base:

[4]:
# - Create a time series containing a sin and cos trace
ts_cos_sin = TSContinuous(
    times=times,
    samples=np.stack((np.sin(theta), np.cos(theta))).T,
    name="sine and cosine",
)
# - Print the representation
print(ts_cos_sin)

# - Plot the time series
ts_cos_sin.plot()
non-periodic TSContinuous object `sine and cosine` from t=0.0 to 9.9. Samples: 100. Channels: 2
[4]:

For convenience, .TimeSeries objects can be made to be periodic. This is particularly useful when simulating networks over repeated trials. To do so, use the periodic flag when constructing the .TimeSeries object:

[5]:
# - Create a periodic time series object
ts_sin_periodic = TSContinuous(
    times=times,
    samples=np.sin(theta),
    periodic=True,
    name="periodic sine wave",
)
# - Print the representation
print(ts_sin_periodic)

# - Plot the time series
plot_trace = np.arange(0, 100, .1)
ts_sin_periodic.plot(plot_trace)
periodic TSContinuous object `periodic sine wave` from t=0.0 to 9.9. Samples: 100. Channels: 1
[5]:

Continuous time series permit interpolation between sampling points, using scipy.interpolate as a back-end. By default linear interpolation is used, but any interpolation method supported by scipy.interpolate can be provided as a string when constructing the .TSContinuous object.

The interpolation interface is simple: .TSContinuous objects are callable with a list-like set of time points.; the interpolated amplitudes at those time points are returned as a numpy.ndarray.

[6]:
# - Interpolate the sine wave
print(ts_sin([1, 1.1, 1.2]))
[[0.58778525]
 [0.63742399]
 [0.68454711]]

As a convenience, .TSContinuous objects can also be indexed, which uses interpolation. A second index can be provided to choose specific channels. Indexing will return a new .TSContinuous object.

[7]:
# - Slice a time series object
ts_cos_sin[:1:.09, 0].print()

non-periodic TSContinuous object `sine and cosine` from t=0.0 to 0.99. Samples: 12. Channels: 1
0.0:     [0.]
0.09:    [0.05651147]
0.18:    [0.11282469]
0.27:    [0.16876689]
        ...
0.72:    [0.43697417]
0.8099999999999999:      [0.48716099]
0.8999999999999999:      [0.53582679]
0.99:    [0.58258941]

.TSContinuous provides a large number of methods for manipulating time series. For example, binary operations such as addition, multiplication etc. are supported between two time series as well as between time series and scalars. Most operations return a new .TSContinuous object.

See the api reference for .TSContinuous for full detail.

Attributes (TSContinuous)

Attribute name

Description

times

Vector \(T\) of sample times

samples

Matrix \(T\times N\) of samples, corresponding to sample times in times

num_channels

Scalar reporting \(N\): number of series in this object

num_traces

Synonym to num_channels

t_start, t_stop

First and last sample times, respectively

duration

Duration between t_start and t_stop

plotting_backend

Current plotting backend for this instance.

Examples of time series manipulation

[8]:
# - Perform additions, powers and subtractions of time series

(   (ts_sin + 2).plot() + \
    (ts_cos_sin ** 6).plot() + \
    (ts_sin - (ts_sin ** 3).delay(2)).plot()
).cols(1)
[8]:

Representation of event-based time series

Sequences of events (e.g. spike trains) are represented by the .TSEvent class, which inherits from .TimeSeries.

Discrete time series are represented by tuples \((t_k, c_k)\), where \(t_k\) are sample times as before and \(c_k\) is a “channel” associated with each sample (e.g. the source of an event).

Multiple samples at identical time points are explictly permitted such that (for example) multiple neurons could spike simultaneously.

.TSEvent objects are initialised with the syntax

def __init__(
    self,
    times: ArrayLike = None,
    channels: Union[int, ArrayLike] = None,
    periodic: bool = False,
    t_start: Optional[float] = None,
    t_stop: Optional[float] = None,
    name: str = None,
    num_channels: int = None,
)
[9]:
# - Build a time trace vector
times = np.sort(np.random.rand(100))
channels = np.random.randint(0, 10, (100))
ts_spikes = TSEvent(
    times = times,
    channels = channels,
)
ts_spikes
[9]:
non-periodic `TSEvent` object `None` from t=0.0003808350954009887 to 0.9896131800308735. Channels: 10. Events: 100
[10]:
# - Plot the events
ts_spikes.plot()
[10]:

If .TSEvent is called, it returns arrays of the event times and channels that fall within the defined time points and correspond to selected channels.

[11]:
# - Return events between t=0.5 and t=0.6
ts_spikes(.5, .6)
ts_spikes(.5, .6, channels=[3, 7, 8])
[11]:
(array([], dtype=float64), array([], dtype=int64))

.TSEvent also supports indexing, where indices correspond to the indices of the events in the times attribute. A new .TSEvent will be returned. For example, in order to get a new series with the first 5 events of tsSpikes one can do:

[12]:
ts_spikes[:5]
[12]:
non-periodic `TSEvent` object `None` from t=0.0003808350954009887 to 0.9896131800308735. Channels: 10. Events: 5

.TSEvent provides several methods for combining multiple .TSEvent objects and for extracting data. See the API reference for .TSEvent for full details.

Attributes (TSEvent)

Attribute name

Description

times

Vector \(T\) of sample times

channels

Vector of channels corresponding to sample times in times

num_channels

Scalar reporting \(N\): number of channels in this object

t_start, t_stop

First and last sample times, respectively

duration

Duration between t_start and t_stop

plotting_backend

Current plotting backend for this instance.

This page was generated from docs/tutorials/building_reservoir.ipynb. Interactive online version: Binder badge

Building and simulating a reservoir network

This tutorial illustrates how to use Rockpool to construct, simulate, train and visualise networks (with a focus on reservoir networks).

Housekeeping and import statements

[1]:
# - Import required modules and configure

# - Disable warning display
import warnings
warnings.filterwarnings('ignore')

# - Required imports
import numpy as np
import scipy.signal as sig

from rockpool.timeseries import (
    TimeSeries,
    TSContinuous,
    TSEvent,
    set_global_ts_plotting_backend,
)

# - Use HoloViews for plotting
import colorcet as cc
import holoviews as hv
hv.extension('bokeh')

%opts Curve [width=600]
%opts Scatter [width=600]

General concepts — Networks and Layers

Networks in Rockpool are represented as stacks of layers, currently with the restriction that layers are connected in a chain from input to output. Full recurrent connectivity is permitted within a layer.

.Layer objects combine a number of neurons of arbitrary type, along with a set of weights.

In the case of feedforward layers, the weights comprise an \(M\times N\) matrix \(W\), which describe the transformation between \(M\) input channels and the \(N\) neurons in the layer.

Most recurrent layers contain two sets of weights: An \(M\times N\) matrix for mapping the input to the neurons, as with feedforward layers, as well as an \(N\times N\) matrix for the recurrent connections. Note that some older layer classes only have the recurrent weight matrix. Here, inputs are mapped 1:1 to the neurons and need to be of the same dimension as the layer. Use feedforward layers to map between different layer dimensions.

Different layers implement different types of neurons, and define their outputs in slightly different ways. The principal difference is whether a layer expects continous-time or event-based (i.e. spiking) inputs, and similarly what output representation the layer generates.

To be connected together, two layers must match in terms of output\(\rightarrow\)input dimensionality and signal representation. There are several feedforward layers that convert between spiking and continuous signals (see below).

For a full overview of available layers, see layersdocs and the API reference.

Summary of available layers

Simple layers

Class

Structure

Input representation

Output representation

Comment

FFRateEuler

Feedforward

Continuous

Continuous

PassThrough

Feedforward

Continuous

Continuous

Simple weighting

PassThroughEvents

Feedforward

Spiking

Spiking

Route events between channels

FFIAFBrian

Feedforward

Continuous

Spiking

FFIAFSpkInBrian

Feedforward

Spiking

Spiking

FFExpSyn

Feedforward

Spiking

Continuous

Filter with exponential kernel

FFExpSynBrian

Feedforward

Spiking

Continuous

Filter with exponential kernel (Brian2-based, obsolete)

FFCLIAF

Feedforward

Spiking

Spiking

Constant leak, clocked

FFUpDown

Feedforward

Continuous

Spiking

Analog-to-spike conversion through delta modulation

SoftMaxLayer

Feedforward

Spiking

Continuous

SoftMax operation

RecRateEuler

Recurrent

Continuous

Continuous

RecIAFBrian

Recurrent

Continuous

Spiking

RecIAFSpkInBrian

Recurrent

Spiking

Spiking

RecFSSpikeEulerBT

Recurrent

Continuous

Spiking

Precise spiking timing, fast/slow synapses

RecDynapseBrian

Recurrent

Spiking

Spiking

DynapSE simulation

RecCLIAF

Recurrent

Spiking

Spiking

Constant leak, clocked

RecDIAF

Recurrent

Spiking

Spiking

Digital neuron. Event based.

RecRateEulerJax

Recurrent

Continuous

Continous

Jax-backed.

ForceRateEulerJax

Recurrent

Continuous

Continuous

Jax-backed. For reservoir transfer.

PyTorch- and Nest-accelerated versions

Base class

PyTorch class

Nest class

Purpose

FFExpSyn

FFExpSynTorch

Filter with exponential kernel

FFIAFBrian

FFIAFRefrTorch

FFIAFNest

FF Cont -> Spiking

FFIAFTorch

FF Cont -> Spiking without refractoriness but faster

FFIAFSpkInBrian

FFIAFSpkInRefrTorch

FF Spiking -> Spiking

FFIAFSpkInTorch

FF Spiking -> Spiking without refractoriness but faster

RecIAFBrian

RecIAFRefrTorch

Rec Cont -> Spiking

RecIAFTorch

Rec Cont -> Spiking without refractoriness but faster

RecIAFSpkInBrian

RecIAFSpkInRefrTorch

RecIAFSpkInNest

Rec Spiking -> Spiking

RecIAFSpkInTorch

Rec Spiking -> Spiking without refractoriness but faster

RecIAFSpkInRefrCLTorch

Rec Spiking -> Spiking with constant leak

RecAEIFSpkInNest

Rec Spiking -> Spiking with AdEx neurons

Layers for simulation of or interaction with hardware

Class

Structure

Input representation

Output representation

Purpose

VirtualDynapse

Recurrent

Spiking

Spiking

Conceptual simulation of DynapSE and related chips. RecAEIFSpkInNest as backend.

RecDynapSE

Recurrent

Spiking

Spiking

Set up reservoirs on DynapSE chip.

Importing the packages

[2]:
# - Import the network module
from rockpool.networks.network import Network

# - Import single layer classes
from rockpool.layers import RecFSSpikeEulerBT as RecSpike

# - Import the TimeSeries classes
from rockpool.timeseries import TSContinuous, TSEvent
set_global_ts_plotting_backend('holoviews')
Global plotting backend has been set to holoviews.

Building and simulating a simple reservoir network

Let’s begin by building a simple network, with a input layer (FF rate layer); a recurrent reservoir (rate); and a linear readout (passthrough). We’ll combine these into a chain of layers.

First we need to load the required classes.

[3]:
# - Import classes
from rockpool.layers import PassThrough, FFRateEuler
from rockpool.layers import RecRateEuler

# - Import weight generation functions
from rockpool.weights import unit_lambda_net

Now we need to generate layers, by providing weights to be encapsulated by each layer. We also define the network size.

[4]:
# - Define layer sizes
nInputChannels = 1
nRecurrentUnits = 100
nOutputChannels = 2

# - Define the input layer
lyrInput = PassThrough(weights = np.random.rand(nInputChannels, nRecurrentUnits)-.5,
                       name = 'Input',
                      )
print(lyrInput)

# - Define the recurrent layer, using a convenience weight generation function
# - USe a range of time constants, to improve performance
vtTimeConstants = np.random.rand(nRecurrentUnits) * 50 + 10
lyrRecurrent = RecRateEuler(weights = unit_lambda_net(nRecurrentUnits),
                            tau = vtTimeConstants,
                            dt = 1,
                            name = 'Reservoir',
                           )
print(lyrRecurrent)

# - Define the ouput layer
lyrOutput = PassThrough(weights = np.random.rand(nRecurrentUnits, nOutputChannels)-.5,
                        name = 'Readout',
                       )
print(lyrOutput)
PassThrough object: "Input" [1 TSContinuous in -> 100 TSContinuous out]
RecRateEuler object: "Reservoir" [100 TSContinuous in -> 100 TSContinuous out]
PassThrough object: "Readout" [100 TSContinuous in -> 2 TSContinuous out]

We can visualise these layers by plotting the weight matrices and eigenspectra.

[5]:
# - Display the layer weights
hv.Raster(lyrInput.weights
         ).redim(x = 'i_{res}', y = 'c_{in}', z = 'w') +\
hv.Raster(lyrRecurrent.weights
         ).redim(x = 'i_{res}', y = 'j_{res}', z = 'w') +\
hv.Raster(lyrOutput.weights
         ).redim(x = 'c_{out}', y = 'j_{res}', z = 'w')
[5]:
[6]:
# - Get the recurrent layer eigenspectrum
vfEigVals = np.linalg.eigvals(lyrRecurrent.weights)

# - Plot the eigenspectrum
vfSamples = np.linspace(0, 2*np.pi, 100)
hv.Curve((np.cos(vfSamples), np.sin(vfSamples))
        ).options(color = 'red',
                  line_dash = 'dashed',
                  width = 340) *\
hv.Scatter((np.real(vfEigVals),
            np.imag(vfEigVals)),
          ).redim(x = 'Re',
                  y = 'Im')
[6]:

Now we compose these layers into a network, using the Network class initialiser. The syntax for the initialiser is:

def __init__(self, *layers : Layer, tDt=None):
[7]:
# - Build a network of these layers
net = Network(lyrInput, lyrRecurrent, lyrOutput)
net
[7]:
Network object with 3 layers
    PassThrough object: "Input" [1 TSContinuous in -> 100 TSContinuous out]
    RecRateEuler object: "Reservoir" [100 TSContinuous in -> 100 TSContinuous out]
    PassThrough object: "Readout" [100 TSContinuous in -> 2 TSContinuous out]

Representation of time

Internally, a .Network has a discrete representation of time. During instantiation and whenever a layer is added or removed, it tries to find the smallest time step size Network.dt that is a multiple of all its layers’ Layer.dt s. For instance, if a .Network instance contains three layers with time steps 0.005, 0.003 and 0.006, respectively, the network’s Network.dt will be 0.03. This makes sure that each time the .Network.evolve method is called, the evolution duration is a multiple of all involved layers’ time step lengths and therefore all layers evolve to the same time point.

For unfortunate choices of the layer Layer.dt s, such as the combination of 0.007, 0.013 and 0.043, the smallest suitable Layer.dt for the network is rather large (3.913 in this case, which is 559 times the smallest layer time step 0.007). It may also happen for more reasonable combinations of time steps that due to numerical errors the network simply cannot find a suitable Layer.dt (however, due to improved handling of real values, this has become extremely rare). In these two cases the network will raise an AssertionError.

This can be avoided by setting the Network.dt parameter at instantiation, e.g.

net = Network(layer1, layer2, dt=0.005)

This forces the network’s time step length to the provided value. Whenever a layer is added to the network, it will make sure that Network.dt is a multiple of the new layer’s Layer.dt and raise an AssertionError otherwise. This also applies to the layers added at instantiation (layer1 and layer2 in the example above). This procedure is numerically more stable. It can also be used to set Layer.dt to rather large values, such as 3.913 in the example above. It also guarantees that Network.dt does not change over time (e.g. when new layers are added).

Let’s build a simple input time series (a ramp), and simulate the network. To do that we use the .Network.evolve method, which handles all the signal passing within the network, and ensure that all the layers are evolved appropriately.

[8]:
# - Build ramp input time series
tDt = 1
tInputDuration = 100
vtTimeTrace = np.arange(0, tInputDuration+tDt, tDt)
tsRamp = TSContinuous(
    times = vtTimeTrace,
    samples = np.repeat(vtTimeTrace, nInputChannels),
    periodic = True,
    name = 'Ramp',
)
# - Plot the input ramp
tsRamp.plot()
[8]:
[9]:
# - Evolve the network
tSimDuration = 1000
dOutput = net.evolve(ts_input = tsRamp,
                     duration = tSimDuration,
                    )
Network: Evolving layer `Input` with external input as input
Network: Evolving layer `Reservoir` with Input's output as input
Network: Evolving layer `Readout` with Reservoir's output as input
[10]:
# - Plot the signals of the network output
(   dOutput['external'].plot(dOutput['Input'].times) +\
    dOutput['Input'].clip(channels=range(0, nRecurrentUnits, 10)).plot().redim(y = 'y_{Inp}') +\
    dOutput['Reservoir'].clip(channels=range(0, nRecurrentUnits, 10)).plot().redim(y = 'y_{Res}') +\
    dOutput['Readout'].plot().redim(y = 'y_{RO}')
).cols(1)

[10]:

Training the network

Some layers support self-training, by implementing train_XXX() methods. Here we’ll use simple ridge regression to train the linear readout layer.

First we need to generate some targets for the reservoir outputs. How about a \(\sin\) and \(\cos\) curve?

[11]:
# - Generate target signals
vtTimeTrace = np.arange(0, tInputDuration+tDt, tDt)
tsTargets = TSContinuous(vtTimeTrace,
                         np.stack((np.sin(vtTimeTrace / tInputDuration * 2 * np.pi),
                                   np.cos(vtTimeTrace / tInputDuration * 2 * np.pi)
                                 )).T,
                         periodic = True,
                         name = 'Target',
                        )
# - Plot target signals
tsTargets = tsTargets.clip(channels=range(nOutputChannels))
tsTargets.plot()
[11]:

The .Network class provides a method .Network.train, which evolves the entire network over a series of input batches, then uses an auxilliary callback function fhTraining() to operate on the network layers.

The syntax for .Network.train is given by:

def train(
        self,
        training_fct: Callable,
        ts_input: TimeSeries = None,
        duration: float = None,
        batch_durs: float = None,
        verbose: bool = True,
        high_verbosity: bool = False,
        ):

The training callback function must know about the specfic network structure and composition. This is so .Network.train can operate without needing to assume anything about the network structure.

fhTraining() is called with the syntax

fhTraining(netObj, dtsSignals, bFirst, bFinal)

Here netObj is a reference to the network that is being trained, so that the trianing callback can access all layers and signals as necessary.

dtsSignals is a dictionary of signals, containing the results of evolving the network for the current batch.

bFirst is a boolean flag that is True only when fhTraining() is called on the first batch, to permit any initialisation steps to take place.

bFinal is a boolean flag that is True only when fhTraining() is called on the final batch, so it can finalise and clean up.

Pseudo-code for an example training callback would be something like

def fhTraining(netObj, dtsSignals, bFirst, bFinal, tsTarget):
    # - Perform initialisation
    if bFirst:
        # - Initialise the algorithm

    # - Perform some training, operating on dtsSignals and the network
    netObj.lyrOutput.train(dtsSignals['Output'], tsTarget)

    # - Finalise training
    if bFinal:
        # - Finalise the algorithm

.PassThrough supports ridge regression via the .PassThrough.train_rr method. This method has the syntax

def train_rr(self,
             ts_target: TimeSeries,
             ts_input: TimeSeries = None,
             regularize: float = 0,
             is_first: bool = True,
             is_final: bool = False):

So let’s define an auxilliary function for training.

[12]:
# - Define a function to set up a training problem
def train_reservoir(tsTarget):
    def training_callback(netObj, dtsSignals, bFirst, bFinal):
        # - Call layer layer training function
        netObj.output_layer.train_rr(
            ts_target = tsTarget,
            ts_input = dtsSignals['Reservoir'],
            regularize = .1,
            is_first = bFirst,
            is_last = bFinal,
        )
    return training_callback

# - Get a callback function
fhTrainingCallback = train_reservoir(tsTargets)

We will train the network over several batches of input data. Note that we also use a “burn-in” time, to allow the internal reservoir dynamics to settle away from their transient initial state. See the figures showing the internal state evolution above, to see this settling in action.

If we do not let the internal dynamics settle, then some of the training effort will go towards fitting the output to the transient state. This leads to bad performance later on.

[13]:
# - Train the network
nNumBatchesTrain = 10
tBurnInTime = 1000
net.reset_all()
net.evolve(ts_input = tsRamp, duration = tBurnInTime)
net.train(training_fct = fhTrainingCallback,
          ts_input = tsRamp,
          duration = tInputDuration,
          batch_durs = tInputDuration / nNumBatchesTrain,
          verbose = True,
         )
Network: Evolving layer `Input` with external input as input
Network: Evolving layer `Reservoir` with Input's output as input
Network: Evolving layer `Readout` with Reservoir's output as input

Network: Training successful

Now that training has been completed, let’s evaluate the performance of the system by injecting the ramp input and comparing the output to the target signal.

[14]:
dtsOutpt = net.evolve(tsRamp, duration = tInputDuration * 2)
dtsOutpt['Readout'].plot() * tsTargets.plot(dtsOutpt['Readout'].times)
Network: Evolving layer `Input` with external input as input
Network: Evolving layer `Reservoir` with Input's output as input
Network: Evolving layer `Readout` with Reservoir's output as input
[14]:

The system performs well — the signal output by the reservoir closely matches the desired target signal.

Training in batches

In the example above training was split into 10 batches of the same duration. However, there are a few other possible ways of defining batches, using arguments to net.train():

  • Setting batch_durs to a float: Training will be split into batches of the given duration, the last one might be shorter.

  • Setting batch_durs with a vector with the duration for each batch. If the times don’t add up to the full duration of the training, the last batch(es) will be shortened accordingly or an additional batch will be added.

Instead of time you can set the batch durations by network-timesteps by setting nums_ts_batch as an array of integer time steps. If both batch_durs and nums_ts_batch are provided, nums_ts_batch will be used.

Helper functions for building weight matrices

We provide a utility package .weights, containing several useful functions for constructing recurrent weight matrices. These are summarised below. See also the API reference for .weights.

Function

Description

UnitLambdaNet

Build a simple unit-circle eigenspectrum recurrent matrix, by drawing weights from a Normal distribution with std. dev. \(\sigma = \sqrt{N}\)

RandomEINet

Build a reservoir with excitatory and inhibitory populations

RndmSparseEINet

Build a network with excitatory and inhibitory populations, and sparse recurrent connectivity

WilsonCowanNet

Build a network composed of randomly-connected Wilson-Cowan units

DiscretiseWeightMatrix

Convert a fully-connected weight matrix into a discretised version by pruning and down-sampling weights

DynapseConform

Build a sparse network of Normally-distributed synapses, that conforms to neurmorphic hardware constraints

IAFSparseNet

Build a sparse network scaled nicely for IAF spiking simulations

This page was generated from docs/tutorials/jax_sgd.ipynb. Interactive online version: Binder badge

Gradient descent training of a rate-based reservoir model

This tutorial demonstrates using Rockpool and a Jax-accelerated rate-based reservoir layer to perform gradient descent training of all network parameters. The result is a trained dynamic recurrent network with long memory, optimised to perform a signal generation task.

Requirements

This example requires the Rockpool package from aiCTX, as well as jax and its dependencies.

[2]:
# - Ignore warnings
import warnings
warnings.filterwarnings('ignore')

# - Imports and boilerplate
from rockpool import TimeSeries, TSContinuous
from rockpool.layers import RecRateEulerJax, H_ReLU, H_tanh
from rockpool.layers.training import add_train_output
from tqdm import tnrange
from tqdm.autonotebook import tqdm

import numpy as np
import numpy.random as npr

import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['figure.figsize'] = [12, 4]

Triggered signal-generation task

We will use a pulse-to-chirp task as a demonstration. The reservoir receives a short pulse, and must respond by generating a chirp signal (a sinusoid increasing in frequency over time). This task is difficult, as no input is present for most of the time, and so considerable reservoir memory is required.

You can adjust the parameters of the input by changing the number of repeats num_repeats, the duration of the input pulse pulse_duration, and the maximum frequency reached by the chirp chirp_freq_factor. Shorter pulses and higher chirp frequencies make the problem more difficult. More repeats make learning more difficult, by forcing gradients to be accumulated over more time steps.

You can also adjust the time step dt, which makes learning slower (more time points evaluated per trial), but which permits shorter time constants to be used in the network. For numerical stability, time constants must be at least 10*dt.

[3]:
# - Define input and target signals
num_repeats = 1
pulse_duration = 50e-3
chirp_freq_factor = 10
padding_duration = 1
chirp_duration = 1
dt = 1e-3

# - Build chirp and trigger signals
chirp_end = int(np.round(chirp_duration / dt))
chirp_timebase = np.linspace(0, chirp_end * dt, chirp_end)
chirp = np.atleast_2d(np.sin(chirp_timebase * 2 * np.pi * (chirp_timebase * chirp_freq_factor))).T
trigger = np.atleast_2d(chirp_timebase < pulse_duration).T

# - Add padding
padding = np.zeros((int(np.round(padding_duration / dt)), 1))
chirp = np.vstack((padding, chirp, padding))
trigger = np.vstack((padding, trigger, padding))

# - Build a time base
num_samples = (chirp_end + len(padding)*2) * num_repeats
timebase = np.linspace(0, num_samples, num_samples + 1)
timebase = timebase[:-1] * dt

# - Replicate out inputs and target signals
input_t = np.tile(trigger * 1., (num_repeats, 1))
target_t = np.tile(chirp, (num_repeats, 1))

# - Generate time series objects
ts_input = TSContinuous(timebase, input_t, periodic=True)
ts_target = TSContinuous(timebase, target_t, periodic=True)

# - Plot the input and target signals
plt.figure()
plt.plot(
     timebase,
     input_t,
     timebase,
     target_t,
     # timebase, target_t + np.diff(np.vstack((target_t, target_t[0])), axis=0) / dt * tau,
 )
plt.xlabel('Time (s)')
plt.ylabel('Input / target amplitude')
plt.legend(("Input", "Target"));
_images/tutorials_jax_sgd_6_0.png

Network model

We will define a ReLU reservoir, with a single input and a single output channel, and with a chosen number of units in the recurrent layer nResSize. Larger reservoirs take longer to train, but perform the task to higher accuracy.

The dynamics of a unit \(j\) in the recurrent layer is given by

\[\tau_j\frac{\textrm{d}{x_j}}{\textrm{d}t} + {x_j} = W_r \cdot f(\textbf{x}) + b_j + i_j + \sigma\zeta_j(t)\]

where \(\tau_j\) is the time constant of the unit (tau); \(W_r\) is the \(N \times N\) weight matrix defining the recurrent connections; \(\textbf{x}\) is the vector of recurrent layer activities (w_rec); \(f(x)\) is the neuron transfer function \(\tanh(x)\); \(b_j\) is the bias input of the unit (bias); \(i_j\) is the external input to the unit; and \(\sigma\zeta_j(t)\) is a white noise process with standard deviation \(\sigma\) (noise_std).

External input is weighted such that \(\textbf{i} = W_i \cdot i(t)\), where \(W_i\) is the external input weight matrix (w_in) and \(i(t)\) is the external input function.

The output of the reservoir is also weighted such that \(z = W_o \cdot \textbf{x}\), where \(W_o\) is the output weight matrix (w_out). The goal of the training task is to match the reservoir output \(\hat{z}\) with a target signal \(z^*\).

Weight initialisation doesn’t seem to matter too much in this process; gradient descent does a good job even when weights are initially zero. Here we use a standard initialisation with unit spectral radius for the recurrent weights.

You can change the activation function to one of H_tanh or H_ReLU. You can also define your own, but must use jax.numpy to do so.

[4]:
# - Define the reservoir parameters
nResSize = 100
tau = 20e-3
bias = 0.
activation_func = H_ReLU
noise_std = 0.1

# - Build a layer
nInSize = 1
nOutSize = 1
w_in = 2*npr.rand(nInSize, nResSize)-1
w_rec = npr.randn(nResSize, nResSize) / np.sqrt(nResSize)
w_out = 2*npr.rand(nResSize, nOutSize)-1

lyrRes = RecRateEulerJax(w_in, w_rec, w_out, tau, bias,
                         dt = dt, noise_std = noise_std,
                         activation_func = activation_func)

# - Get initial output
ts_output0 = lyrRes.evolve(ts_input)
[5]:
# - Initialise training
lyrRes = add_train_output(lyrRes)

# - Force initialisation of training
lyrRes.train_adam(ts_input, ts_target, is_first=True);

Here we use the training method .train_adam() to perform stochastic descent using the Adam optimiser. You can re-run the cell below as many times as you like if you want to extend the training time.

The loss function is a combination of mean-squared-error between the reservoir output and the target signal \(\frac{1}{T} (\hat{z} - z^*)^2\); a factor that harshly penalises time constants \(\tau\) shorter than the minimum time constant \(\tau_{min}\); and a factor related to the 2-norm of the recurrent weight matrix np.mean(w_recurrent ** 2) or \(||W_r||^2\). This helps to ensure that the network remains stable.

At the end of each batch the network is testing by providing a pulse input, and plotting the reservoir response.

[6]:
num_batches = 10
trials_per_batch = 100

def two_norm(params):
    return np.sqrt(np.sum([np.sum(e ** 2) for e in params.values()]))

fig = plt.figure()
ax = plt.axes()
line_target = plt.plot(ts_target.times, ts_target.samples)
line_output = plt.plot(ts_output0.times, ts_output0.samples, '--')
plt.legend(['Target', 'Output'])
plt.xlabel('Time (s)')
plt.ylabel('Output / target');
fig.canvas.draw();

with tnrange(num_batches, desc='batches') as tqdm_batches:
    for _ in range(num_batches):
        with tnrange(trials_per_batch, desc='trials', leave=False) as tqdm_trials:
            for _ in range(trials_per_batch):
                # - Get this trial
                pass

                # - One step of Adam training
                loss, grad_loss = lyrRes.train_adam(ts_input, ts_target)

                # - Update statistics
                tqdm_trials.set_postfix(loss=loss(),
                                        #norm_g=two_norm(grad_loss()),
                                        min_tau=int(np.min(lyrRes.tau) / 1e-3),
                                        refresh=False)
                tqdm_trials.update()

            # - Update batch progress
            tqdm_batches.update()
            ts_output = lyrRes.evolve(ts_input)
            line_output[0].set_data(ts_output.times, ts_output.samples)
            line_target[0].set_data(ts_output.times, ts_target.samples)
            ax.set_xlim([np.min(ts_output.times), np.max(ts_output.times)])
            ax.set_ylim([np.min(ts_output.samples), np.max(ts_output.samples)])
            fig.canvas.draw()


_images/tutorials_jax_sgd_11_12.png

If all has gone according to plan, the output of the reservoir should closely match the target signal. We can see the effect of training by examining the distribution of network parameters below.

[9]:
# - Plot the network time constants
plt.figure()
plt.hist(lyrRes.tau / 1e-3, 21);
plt.stem([tau / 1e-3], [50], 'r-')
plt.legend(('Trained time constants', 'Initial time constants'))
plt.xlabel('Time constant (ms)');
plt.ylabel('Count');
plt.title('Distribution of time constants');
_images/tutorials_jax_sgd_13_0.png
[10]:
# - Plot the recurrent layer biases
plt.figure()
plt.stem(lyrRes.bias);
plt.title('Distribution of trained biases')
plt.xlabel('Unit')
plt.ylabel('Bias');
_images/tutorials_jax_sgd_14_0.png

We can examine something of the computational properties of the network by finding the eigenspectrum of the Jacobian of the recurrent layer. The Jacobian \(\hat{J}\) is given by

\[\hat{J} = (\hat{W_r} - I) ./ \hat{T}\]

where \(I\) is the identity matrix, \(./\) denotes element-wise division, and \(T\) is the matrix composed of all time constants \(\hat{\tau}\) of the recurrent layer.

Below we plot the eigenvalues \(\lambda\) of \(J\). In my training result, several complex eigenvalues \(\lambda\) with real parts greater than zero are present in the trained eigenspectrum. These correspond to oscillatory modes, which are obviously useful in generating the chirp output.

[14]:
# - Plot the recurrent layer eigenspectrum
J = lyrRes.w_recurrent - np.identity(nResSize)
J = J / lyrRes.tau

J0 = w_rec - np.identity(nResSize)
J0 = J0 / tau

plt.figure()
eigs = np.linalg.eigvals(J)
eigs0 = np.linalg.eigvals(J0)
plt.plot(np.real(eigs0), np.imag(eigs0), '.')
plt.plot(np.real(eigs), np.imag(eigs), 'o')
plt.plot([0, 0], [-100, 100], '--')
plt.legend(('Initial eigenspectrum', 'Trained eigenspectrum',))
plt.xlim([-350, 250])
plt.ylim([-100, 100])
plt.title('Eigenspectrum of recurrent weights')
plt.xlabel('Re($\lambda$)')
plt.ylabel('Im($\lambda$)');
_images/tutorials_jax_sgd_16_0.png

Summary

Gradient descent training does a good job of optimmising a dynamic recurrent network for a difficult task requiring significant reservoir memory. jax provides a computationally efficient back-end as well as automatic differentiation of the recurrent reservoir layer.

This page was generated from docs/tutorials/RecDynapSE.ipynb. Interactive online version: Binder badge

Event-based simulation on DynapSE hardware

This document illustrates how to use the .RecDynapSE layer, that runs directly on the DynapSE event-driven neuromorphic computing hardware from aiCTX.

Hardware basics

The layer uses the DynapSE processor, which consists of 4 chips. Each chips has 4 cores of 256 neurons. The chips, as well as each core in a chip and each neuron in a core are identified with an ID between 0 and 3 or 0 and 256, respectively. However, for this layer the neurons are given logical IDs from 0 to 4095 that range over all neurons. In other words the logical neuron ID is \(1024 \cdot \text{ChipID} + 256 \cdot \text{CoreID} + \text{NeuronID}\).

Setup

Connecting to Cortexcontrol

In order to work interface the DynapSE chip, this layer relies on cortexcontrol. It should be accessed via an RPyC connection. In order to run some examples from within this jupyter notebook, we will do the latter. For this we start cortexcontrol and run the following commands in its console (not in this notebook):

[ ]:
import rpyc.utils.classic
c = rpyc.utils.classic.SlaveService()
from rpyc.utils.server import OneShotServer
t = OneShotServer(c, port=1300)
print("RPyC: Ready to start.")
t.start()

If the cortexcontrol console prints “RPyC: Ready to start.” and nothing else, it is ready.

Using DynapseControl

The .RecDynapSE layer uses a .DynapseControl object to interact with Cortexcontrol (see …. for more details). You can either pass an existing DynapseControl object to the layer upon instantiation, or let the layer create such an object. Either way, it can be accessed via the layer’s controller attribute.

Import

Import the class with the following command:

[2]:
# - Import recurrent RecDIAF layer
from rockpool.layers import RecDynapSE

This might take a while until the dynapse_control module has prepared the hardware.

Instantiation

RecDynapSE objects have many instantiation arguments in common with other Rockpool layers, in partuclar recurrent layers such as .RecIAFTorch, .RecIAFBrian or .RecDIAF. Other arguments are related to how the hardware is used and therefore unique to this layer.

Argument

Type

Default

Meaning

weights_in

2D-ndarray

Input weights (required)

weights_rec

2D-ndarray

Recurrent weights (required)

neuron_ids

1D-ArrayLike

None

IDs of the layer neurons

virtual_neuron_ids

1D-ArrayLike

None

IDs of the virtual (input) neurons

dt

float

2e-5

Time step

max_num_trials_batch

int

None

Maximum number of trials in individual batch

max_batch_dur

float

None

Maximum duration time of individual batch

max_num_timesteps

int

None

Maximum number of time steps in individual batch

max_num_events_batch

int

None

Maximum number of input Events during individual batch

l_input_core_ids

ArrayLike

[0]

IDs of cores that receive input spikes

input_chip_id

int

0

Chip that receives input spikes

clearcores_list

ArrayLike

None

IDs of cores to be reset

controller

DynapseControl

None

DynapseControl instance

rpyc_port

int or None

None

Port for RPyC connection

name

str

“unnamed”

Layer name

weights_in and weights_rec are the input and recurrent weights and have to be provided as 2D-arrays. weights_in determines the layer’s dimensions nSIzeIn and size. weights_rec has to be of size size x size. Each weight must be an integer (positive or negative). Furthermore, the sum over the absolute values of the elements in any given column in weights_in plus the sum over the absolute values of elements in the corresponding column of weights_rec must be les or equal to 64. \(\sum_{i} |\) weights_in$ {ik}| + :nbsphinx-math:`sum`{j}|$ weights_rec$ _{jk}| \leq `64  :nbsphinx-math:forall `k$. This is due to limitations of the hardware and cannot be circumvented.

The layer neurons of .RecDynapSE objects directly correspond to physical neurons on the chip. Inputs are sent to the hardware through so-called virtual neurons. Each virtual neuron has an ID, just like a physical neuron. Each input channel of the layer corresponds to such a virtual neuron.

You can choose which physical and virtual neurons are used for the layer by passing their IDs in vnLayerNeuronIDs and vnVirtualNeuronIDs, which are 1D array-like objects with integers between 0 and 1023 or 0 and 4095, respectively. All neurons with IDs from \(j \cdot 256\) to \((j+1) \cdot 256, \ j \in {0,1,..15}\) belong to the same core (with core ID \(j \ mod \ 4\)). All neurons with IDs from \(j \cdot 1024\) to \((j+1) \cdot 1024, \ j \in {0,..,4}\) belong to same chip (with chip ID \(j\)). With this in mind you can allocate neurons to specific chips and cores.

In order for the to layer to function as expected you should stick to the following two rules: - Neurons that receive external input should be on different cores than neurons that receive input from other layer neurons. As a consequence, each neuron should not receive both types of input. The cores with neurons that receive external inputs are set with l_input_core_ids. - All neurons that receive external input should be on the same chip. This chip is set with input_chip_id and is 0 by default.

These rules are quite restrictive and it is possible to set them less strictly. Contact Felix if needed.

.dt is a positive float that on the one hand sets the discrete layer evolution timestep, as in other layers, but on the other hand also corresponds to the smallest (nonzero) interval between input events that are sent to the chip. It needs to be larger than \(1.11 \cdot 10^{-9}\) (seconds). Below, under Choosing dt you can find some more thoughts on how to choose this value.

Evolution is automatically split into batches, the size of which is determined by max_num_events_batch, max_num_timesteps, max_batch_dur and max_num_trials_batch, which control the maximum number of events, number of timesteps, duration or number of trials in a batch. All of them can be set to None, which corresponds to setting no limit, except for max_num_events_batch, where the limit will be set to 65535. If both max_num_timesteps and max_batch_dur are not None, the max_batch_dur will be ignored. max_num_trials_batch will only have an effect when the input time series to the evolve method contains a trial_start_times attribute. For more details on how evolutions are split into batches see Evolution in batches.

The list clearcores_list contains the IDs of the cores where (presynaptic) connections should be reset on instantiation. Ideally this contains all the cores that are going to be used for this RecDynapSE object. However, if you want to save time and you know what you are doing you can set this to None, so no connections are reset.

You can pass an existing DynapseControl instance to the layer that will handle the interactions with Cortexcontrol. If this argument is None, a new DynapseControl will be instantiated. In this case, you can use the rpyc_port argument to define a port at which it should try to establish the connection.

.RecDynapSE.name is a str that defines the layer’s name.

All of these values can be accessed and changed via <Layer>.value, where <Layer> is the instance of the layer.

Example

We can set up a simple layer on the chip by only passing input weights and recurrent weights. The weights are chosen so that there is a population of 3 “input neurons” that receive the input and then excite the remaining 6 neurons, which are recurrently connected. This way the constraints mentioned above are satisfied.

[3]:
import numpy as np

# - Weight matrices: 3 neurons receive external input
#   from 2 channels and stimulate the remaning
#   6 neurons, which are recurrently connected.

weights_in = np.zeros((2,9))
# Only first 3 neurons receive input
weights_in[:,:3] = np.random.randint(-2, 2, size=(2,3))

weights_rec = np.zeros((9,9))
# Excitatory connections from input neurons to rest
weights_rec[:3, 3:] = np.random.randint(3, size=(3,6))
# Recurrent connecitons between remaining 6 neurons
weights_rec[3:, 3:] = np.random.randint(-2, 2, size=(6,6))

rlRec = RecDynapSE(weights_in, weights_rec, l_input_core_ids=[0], name="example-layer")

RecDynapSE `example-layer`: Superclass initialized
dynapse_control: RPyC connection established through port 1300.
dynapse_control: RPyC namespace complete.
dynapse_control: RPyC connection has been setup successfully.
DynapseControl: Initializing DynapSE
DynapseControl: Spike generator module ready.
DynapseControl: Poisson generator module ready.
DynapseControl: Time constants of cores [] have been reset.
DynapseControl: Neurons initialized.
         0 hardware neurons and 1023 virtual neurons available.
DynapseControl: Neuron connector initialized
DynapseControl: Connectivity array initialized
DynapseControl: FPGA spike generator prepared.
DynapseControl ready.
DynapseControl: Not sufficient neurons available. Initializing  chips to make more neurons available.
dynapse_control: Chips 0 have been cleared.
DynapseControl: 1023 hardware neurons available.
Layer `example-layer`: Layer neurons allocated
Layer `example-layer`: Virtual neurons allocated
DynapseControl: Excitatory connections of type `FAST_EXC` between virtual and hardware neurons have been set.
DynapseControl: Inhibitory connections of type `FAST_INH` between virtual and hardware neurons have been set.
Layer `example-layer`: Connections to virtual neurons have been set.
DynapseControl: Excitatory connections of type `FAST_EXC` between hardware neurons have been set.
DynapseControl: Inhibitory connections of type `FAST_INH` between hardware neurons have been set.
Layer `example-layer`: Connections from input neurons to reservoir have been set.
DynapseControl: Excitatory connections of type `SLOW_EXC` between hardware neurons have been set.
DynapseControl: Inhibitory connections of type `FAST_INH` between hardware neurons have been set.
DynapseControl: Connections have been written to the chip.
Layer `example-layer`: Recurrent connections have been set.
Layer `example-layer` prepared.

Choosing dt

As with all layers, a RecDynapSE object’s evolution takes place in discrete time steps. This allows to send an num_timesteps argument to the evolve method, which is consistent with other layer classes and important for the use within a Network. Besides, event though the hardware evolves in continuous time, the input events use discrete timesteps. These timesteps are \(\frac{10^{-7}}{9} \text{ s} = 11.\bar{1} \cdot 10^{-9} \text{ s}\) and mark the smallest value that can be chosen for the layer timestep dt.

Although the effect the timestep size has on computation time is much smaller than with other layers, it is not always advisable to choose the smallest possible value. The reason is that the number of timesteps between two input spikes is currently limited to \(2^{16}-1 = 65535\). This means that with dt \(= 11.\overline{1} \cdot 10^{-9} \text{ s}\), any section in the input signal without input spikes longer than about 0.73 milliseconds will cause the layer to throw an exception. Therefore it makes sense to set dt to something between \(10^{-6}\) and \(10^{-4}\) seconds, in order to allow for sufficiently long silent parts in the input while still maintaining a good temporal resolution.

If these limitations are causing problems contact Felix so that he can implement a way around it.

Neurons

The layer neurons of RecDynapSE objects directly correspond to physical neurons on the chip. Inputs are sent to the hardware through so-called virtual neurons. Each input channel of the layer corresponds to such a virtual neuron. Every neuron on the DynapSE has a logical ID, which ranges from 0 to 4095 for the physical and from 0 to 1023 for the virtual neurons.

Neuron states

Hardware neurons’ states change constantly according to the laws of physics, even when the layer is currently not evolving, and there is no state variable that could be read out for all neurons simultaneously. Therefore the RecDynapSE has no state vector like other layer classes. The state attribute is just a 1D array of size zeros.

Synapses

There are four different synapse types on the DynapSE: fast and slow excitatory as well as fast and slow inhibitory. Each neuron can receive inputs through up to 64 synapses, each of which can be any of the given types. Via cortexcontrol the synaptic behavior can be adjusted for each type and for each core, but not for individual synapses.

There is a priori no difference between slow and fast excitatory synapses, so they can be set to have the same behavior. In fact, one could assign shorter time constants to the “slow” excitatory synapses, making them effectively the fast ones. While both excitatory and the fast inhibitory synapses work by adding or subtracting current to the neuron membrane, the slow inhibitory synapses use shunt inhibition and in practice silence a neuron very quickly.

Note that all synapses that are of the same type and that are on the same core have the same weight. Different connection strengths between neurons can only be achieved by setting the same connection multiple times. Therefore the weight matrices weights_in and weights_rec can only be positive or negative integers and thus determine the number of excitatory and inhibitory connections to a neuron.

In this layer the connections from external input to layer neurons are fast excitatory or inhibitory synapses. Outgoing connections from neurons that receive external input are fast excitatory or inhibitory. Outgoing connecitons from other neurons are slow excitatory and fast inhibitory. For different configurations the _compile_weights_and_configure method has to be modified. Felix can help you with this.

Simulation

The .evolve method takes the standard arguments ts_input, duration, num_timesteps and verbose, which is currently not being used. Evolution duration is determined by the usual rules. If ts_input is provided, it must be a .TSEvent object.

The input is sent to the hardware neurons which will evolve and spike. The neurons’ spikes are collected in a .TSEvent object with the timings and IDs of the spiking neurons.

Note that the hardware neurons continue evolving, so in contrast to software simulations, the layer will be in a different state (membrane potentials etc.) when an evolution begins than when the previous evolution ended. Because evolution happens in batches, this even happens between some of the timesteps within the evolution (see below).

Evolution in batches

As for now it is not possible to stream events (input spikes) continuously to the DynapSE. Therefore a group of events is transferred to the hardware, temporarily stored there and then translated to spikes of virtual neurons, with temporal order and inter-spike intervals matching the input signal.

The number of events that can be sent at once is limited. To allow for arbitrarily long layer evolution times, the input can be split into batches, during each of which a number of events is sent and “played back” to the hardware.

There are two ways of splitting the full input signal into batches. First, a new batch ends, as soon as it contains the maximum number of spikes or lasts the maximum allowed batch duration. These values can be set for each .RecDynapSE object and are stored in the max_num_events_batch and max_batch_dur attributes. While the former is limited by the hardware to be at most $2^{16}-1 = 65535$, the latter can be set to None, which in principle allows for arbitrarily long batch durations. However, the length of any inter-spike interval is currently limited, too, so effectively the batch duration is limited also in this case (see Choosing dt for more details).

Note that because the hardware neurons keep evolving after a batch ends, their state (membrane potentials etc.) will have changed until the next batch starts. These discontinuities could be problematic for some simulations. In particular if the input data consists of trials, no trial should be divided over two batches.

The second way of splitting batches considers this scenario and by splitting the input only at the beginnings of trials. The number of trials in a batch is determined by the layer’s nMaxTrialPerBatch attribute. If a batch with this number of trials contains more events or lasts longer than allowed, the number of trials is reduced accordingly. This method is used if the input time series has a trial_start_times attribute, that indicates the start time of each trial, and if max_num_trials_batch is not None.

Resetting

As usual the layer’s time and state can be reset by the .RecDynapSE.reset_time, .RecDynapSE.reset_state and .RecDynapSE.reset_all methods. However, since there is no state vector in this class, .RecDynapSE.reset_state has no effect and .RecDynapSE.reset_all effectively does the same as .RecDynapSE.reset_time.

Internal methods

_batch_input_data(
    self, ts_input: TSEvent, num_timesteps: int, verbose: bool = False
) -> (np.ndarray, int)

This mehtod is called by evolve, which passes it the evolution input ts_input and the number of evolution timesteps num_timesteps. It splits the input into batches according to the maximum duration, number of events and number of trials and returns a generator that for each batch yields the timesteps and channels of the input events in the batch, the time step at which the batch begins and teh duration of the batch.

_compile_weights_and_configure(self)

Configures the synaptic connections on the hardware according to the layer weights.

Class member overview

Methods

arguments:

Description

_batch_input_data

Split evolution into batches and return generator

_compile_weights_and_configure

Configure hardware synaptic connections

evolve

Evolve layer

reset_all

Reset layer time to 0

reset_state

Do nothing.

reset_time

Reset layer time to 0

Internal methods of parent class Layer are listed in corresponding documentation.

Attributes

Each argument that described in section Instantiation has a corresponding attribute that can be accessed by <Layer>.<attribute>, where <Layer> is the layer instance and <attribute> the argument name. Furthermore there are a few internal attributes:

Attribute name

Description

_vHWNeurons

1D-Array of hardware neurons used for the layer.

vVirtualNeurons

1D-Array of virtual neurons used for the layer.

This page was generated from docs/tutorials/jax_lif_tutorial.ipynb. Interactive online version: Binder badge

Working with spiking JAX layers

This tutorial illustrates how to use JAX-accelerated layers in Rockpool to simulate spiking recurrent networks. The benefit of JAX backend layers is the ability to use the automatic differentiation and CPU / GPU acceleration features of JAX to perform back-propagation through time (BPTT) to train spiking recurrent networks.

[15]:
## Housekeeping and imports

# - Numpy
import numpy as np

# - Matplotlib and plotting config
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['figure.figsize'] = [12, 4]

# - Switch off warnings
import warnings
warnings.filterwarnings('ignore')

Available layers

Rockpool provides three JAX-backed spiking recurrent layers, implementing leaky integrate-and-fire spiking neurons with no refractory periods. Layers with spiking input implement a single exponential synapse on the input of each layer neuron. Layers with direct current injection onto the neuron membrane are also supported.

Layer class

Description

RecLIFJax

Simple recurrent spiking layer, with spiking inputs and outputs. No I/O weighting.

RecLIFJax_IO

Recurrent spiking layer with spiking inputs, and a weighted surrogate output. Input / output weights are provided.

RecLIFCurrentInJax

Recurrent spiking layer, with direct current input injection. No I/O weighting.

RecLIFCUrrentInJax_IO

Recurrent spiking layer with direct current input injection, and a weighted surrogate output. Input / output weights are provided.

Layer dynamics

The membrane and synaptic dynamics evolve under the equations

\[\tau_{syn} \dot{I}_{syn} + I_{syn} = 0\]
\[\tau_{mem} \dot{V}_{mem} + V_{mem} = I_{syn} + I_{in}(t)\cdot W_{in} + b + \sigma\zeta(t)\]
\[I_{syn} = I_{syn} + S_{in}(t) \cdot W_{in}\]

where \(I_{syn}\) is the \(N\) vector of synaptic input currents; \(V_{mem}\) is the \(N\) vector of membrane potentials of the neurons; \(b\) is the \(N\) vector of bias currents; \(\sigma\zeta(t)\) is a white noise process with standard deviation \(\sigma\), injected independently into each neuron; \(\tau_{syn}\) and \(\tau_{mem}\) are the \(N\) vectors of synaptic and membrane time constants, respectively; and \(W_{in}\) is the \([N_{in} \times N]\) input weight matrix.

The input spike train \(S_{in}(t)\) is \(1\) when a spike is present on an input channel. When an input spike arrives on an input channel, the synaptic current variable is incremented by \(1\) then decays to zero.

The input currents \(I_{in}(t)\) define input channels injeted onto neuron membranes via the input weight matrix \(W_{in}\).

Spiking

When the membrane potential for neuron \(j\), \(V_{mem, j}\) exceeds the threshold voltage \(V_{thr} = 0\), then the neuron \(j\) emits a spike. These events are collected into the binary spike variable \(S_{rec}(t)\).

\[S_{rec}(t) = H(V_{mem} - V_{thr})\]
\[I_{syn} = I_{syn} + S_{rec}(t) \cdot W_{rec}\]
\[V_{mem} = V_{mem} - S_{rec}(t)\]

Where \(H(x)\) is a Heaviside function for spike production (with a pseudo gradient described below), \(H(x) = x > 0\); and \(W_{rec}\) is the matrix of recurrent synaptic weights.

All neurons therefore share a common resting potential of 0, a firing threshold of 0, and a subtractive reset of -1. The bias current for each neuron is set to -1 by default.

The spiking events emitted by the layer \(S_{rec}(t)\) are returned as the output of the

Surrogate signals for back-propagation

To assist with training, the layers describe here use surrogate signals in two ways: - To generate a non-spiking weighted output from \(V_{mem}\), which can be used to train the layer to approximate a target representation - On spike generation via \(H(x)\), used to propagate errors backwards through time

A surrogate \(U(t)\) is generated from the layer, with \(U(t) = \textrm{sig}\left[V_{mem}(t)\right]\), where \(\textrm{sig}(x)\) is the sigmoid function \(\textrm{sig}(x) = \left[1 + \exp(x)\right]^{-1}\). The surrogate weighted output of a layer is given by

\[O(t) = U(t) \cdot W_{out}\]

Where \(W_{out}\) is the \([N \times N_{out}]\) output weight matrix.

\(H(x)\) is implemented internally such that it becomes differentiable in the backwards pass.

  • Forward pass: \(H(x) = \max\left[0, \textrm{floor}(x + 1)\right]\). This form permits multiple events for each neuron per time step.

  • Backwards pass: \(\textrm{d}H(x)/\textrm{d}x = (x > -0.5)\)

[16]:
def H(x):
    return np.clip(np.floor(x + 1.0), 0.0, np.inf)

def dHdx(x):
    return x > -0.5

def H_surrogate(x):
    return np.clip(x+.5, 0.0, np.inf)

x = np.linspace(-1., 5., 500)

plt.figure()
plt.plot(x, H(x), x, dHdx(x), x, H_surrogate(x))
plt.legend(['$H(x)$', '$dH/dx(x)$', '$H_{surr}(x)$'])
plt.xlabel('$V_{mem}$')
plt.ylabel('$H$; $dH/dx$; $H_{surr}$');
_images/tutorials_jax_lif_tutorial_9_0.png

In other words, \(\textrm{d}H(x)/\textrm{d}x(x)\) acts as though each neuron is a linear-threshold unit.

Spiking versus current inputs

The classes RecLIFJax and RecLIFJax_IO receive spiking inputs; the evolve() methods for these layers accept TSEvent time series obejcts, describing spiking activity of the input channels over time.

The classes RecLIFCurrentInJax and RecLIFCurrentInJax_IO receive direct current injection inputs; the evolve() methods for these layers accept TSContinuous time series objects, describing instantaneous currents on the input channels.

Spiking versus surrogate outputs

The classes RecLIFJax and RecLIFCurrentInJax return spiking activity of the layer neurons as TSEvent objects from the evolve() method.

The classes RecLIFJax_IO and RecLIFCurrentInJax_IO return (weighted) surrogate non-spiking activity as TSContinuous objects from the evolve() method.

In both cases, several interal signals of the layers are monitored during evolution, and these are available as object attributes after evolution.

Signal

Attribute name

Spiking activity

spikes_last_evolution

Recurrent synaptic currents

i_rec_last_evolution

Surrogate signal

surrogate_last_evolution

Membrane voltage

v_mem_last_evolution

Synaptic input currents

i_syn_last_evolution

Building and simulating a spiking recurrent layer

We’ll start by building a spiking recurrent layer RecLIFJax, which accepts spiking inputs and generates spiking outputs. RecLIFJax is imported from rockpool.layers.

[17]:
## -- Import the recurrent layer
from rockpool.layers import RecLIFJax
from rockpool import TSEvent, TSContinuous
[18]:
# - Build a recurrent layer
num_neurons = 100
bias = -1
dt = 1e-3

weights_rec = 1 * np.random.randn(num_neurons, num_neurons) / np.sqrt(num_neurons)
tau_mem = np.random.rand(num_neurons) * 100e-3 + 50e-3
tau_syn = np.random.rand(num_neurons) * 100e-3 + 50e-3

lyrLIF = RecLIFJax(weights_rec, tau_mem, tau_syn, bias,
                   name = 'Recurrent LIF layer', dt = dt)
print(lyrLIF)
RecLIFJax object: "Recurrent LIF layer" [100 TSEvent in -> 100 internal -> 100 TSEvent out]

We’ll generate a Poisson input signal independently for each of the layer neurons, with a sinusoidal modulation.

[19]:
# - Build an input signal
input_duration = 1.
mod_freq = 5.
min_rate = 3.
max_rate = 10.

# - Generate a sinusoidal instantaneous rate
time_base = np.arange(0, input_duration, dt)
sinusoid = (np.sin(time_base * 2*np.pi * mod_freq)/2. + .5) * (max_rate - min_rate) + min_rate
sinusoid_ts = TSContinuous(time_base, sinusoid)

# - Generate Poisson spiking inputs
raster = np.random.rand(len(time_base), num_neurons) < np.tile(sinusoid * dt, (num_neurons, 1)).T
spikes = np.argwhere(raster)
input_sp_ts = TSEvent(time_base[spikes[:, 0]], spikes[:, 1],
                      t_start = 0., t_stop = input_duration,
                      name = "Input spikes",
                      periodic = True,
                     )

# - Plot the input events
plt.figure()
input_sp_ts.plot()
plt.ylabel('Neuron');
_images/tutorials_jax_lif_tutorial_21_0.png

As with all Rockpool layers, use the evolve() method to simulate the layer and collect the output.

[20]:
# - Evolve the layer with the input events
output_ts = lyrLIF.evolve(input_sp_ts)

# - Plot the output events
plt.figure()
output_ts.plot()
plt.ylabel('Neuron');
_images/tutorials_jax_lif_tutorial_23_0.png

As described above, several other properties are set on each evolution, to assist in monitoring the behaviour of the layers. Let’s plot those and take a look.

[21]:
# - Synaptic input currents
plt.figure()
lyrLIF.i_syn_last_evolution.plot(stagger = 1, skip = 10);
plt.xlabel('Time (s)')
plt.ylabel('$I_{syn}$')
plt.title('Synaptic input');

# - Membrane potentials
plt.figure()
lyrLIF.v_mem_last_evolution.plot(stagger = 1, skip = 10);
plt.xlabel('Time (s)')
plt.ylabel('$V_{mem}$')
plt.title('Membrane potentials');

# - Surrogate signals
plt.figure()
lyrLIF.surrogate_last_evolution.plot(stagger = 1, skip = 10);
plt.xlabel('Time (s)')
plt.ylabel('$U$')
plt.title('Surrogate signals');

_images/tutorials_jax_lif_tutorial_25_0.png
_images/tutorials_jax_lif_tutorial_25_1.png
_images/tutorials_jax_lif_tutorial_25_2.png

Building and stimulating a layer with current injection inputs

Interacting with a current injection layer is similar; using the evolve() method.

[22]:
# - Import a current-in class
from rockpool.layers import RecLIFCurrentInJax

# - Create a current-injection layer
lyrCI = RecLIFCurrentInJax(weights_rec, tau_mem, tau_syn,
                           name = 'Recurrent LIF layer - CI', dt = dt)
print(lyrCI)
RecLIFCurrentInJax object: "Recurrent LIF layer - CI" [100 TSContinuous in -> 100 internal -> 100 TSEvent out]
[23]:
# - Generate several noisy sinusoids as input currents
input_duration = 1.
mod_freq = 5.
sin_amp = 5.
noise_std = 5.

# - Generate a sinusoidal signal
time_base = np.arange(0, input_duration, dt)
sinusoid = sin_amp * np.sin(time_base * 2*np.pi * mod_freq) / 2.

inputs = np.tile(sinusoid, (num_neurons, 1)).T + noise_std * np.random.randn(len(time_base), num_neurons)
input_ts = TSContinuous(time_base, inputs, periodic = True,
                        name = "Input currents",
                        units = "$I_{in}$",
                       )

# - Plot the input currents
plt.figure()
input_ts.plot(stagger = noise_std * sin_amp, skip = 10);
_images/tutorials_jax_lif_tutorial_29_0.png

Now we can stimulate the layer using the evolve() method. The individual current traces will be injected directly onto the membrane of each layer neuron.

[24]:
# - Evolve the layer
lyrCI.evolve(input_ts)

# - Display the layer activity
plt.figure()
lyrCI.i_syn_last_evolution.plot(stagger = 1, skip = 5)
plt.ylabel('$I_{syn}$')
plt.title('Synaptic input currents')

plt.figure()
lyrCI.v_mem_last_evolution.plot(stagger = 1, skip = 5);
plt.ylabel('$V_{mem}$')
plt.title('Membrane potentials')

plt.figure()
lyrCI.spikes_last_evolution.plot();
plt.ylabel('Neuron')
plt.title('Output events')
plt.xlabel('Time(s)');
_images/tutorials_jax_lif_tutorial_31_0.png
_images/tutorials_jax_lif_tutorial_31_1.png
_images/tutorials_jax_lif_tutorial_31_2.png

Using layers with input / output weighting

The layers RecLIFJax_IO and RecLIFCurrentInJax_IO accept two extra parameters, defining input and output weight matrices. Weighted outputs from these classes are formed form the neuron surrogates in both cases.

[25]:
# - Import a weighted layer
from rockpool.layers import RecLIFJax_IO

# - Build a layer with weighted inputs and outputs
num_in = 2
num_out = 5

w_in = np.random.rand(num_in, num_neurons)
w_rec = 1.1 * np.random.randn(num_neurons, num_neurons) / np.sqrt(num_neurons)
w_out = np.random.rand(num_neurons, num_out)

lyrIO = RecLIFJax_IO(w_in, w_rec, w_out,
                     tau_mem, tau_syn,
                     name = 'LIF with I/O weighting',
                     dt = dt,
                    )
print(lyrIO)
RecLIFJax_IO object: "LIF with I/O weighting" [2 TSEvent in -> 100 internal -> 5 TSContinuous out]
[26]:
# - Build an input signal
input_duration = 1.
mod_freq = 5.
min_rate = 10.
max_rate = 20.

# - Generate a sinusoidal instantaneous rate
time_base = np.arange(0, input_duration, dt)
sinusoid = (np.sin(time_base * 2*np.pi * mod_freq)/2. + .5) * (max_rate - min_rate) + min_rate
sinusoid_ts = TSContinuous(time_base, sinusoid)

# - Generate Poisson spiking inputs
raster = np.random.rand(len(time_base), num_in) < np.tile(sinusoid * dt, (num_in, 1)).T
spikes = np.argwhere(raster)
input_sp_ts = TSEvent(time_base[spikes[:, 0]], spikes[:, 1],
                      t_start = 0., t_stop = input_duration,
                      periodic = True,
                      name = "Input spikes",
                     )

# - Plot the input events
plt.figure()
input_sp_ts.plot()
plt.ylabel('Neuron');
_images/tutorials_jax_lif_tutorial_35_0.png
[27]:
# - Evolve the layer
output_ts = lyrIO.evolve(input_sp_ts)

# - Plot layer activity
plt.figure()
lyrIO.i_syn_last_evolution.plot(stagger = 1, skip = 5);
plt.ylabel('$I_{syn}$')
plt.title('Synaptic input currents')

plt.figure()
lyrIO.v_mem_last_evolution.plot(stagger = 1, skip = 5);
plt.ylabel('$V_{mem}$')
plt.title('Membrane potentials')

plt.figure()
lyrIO.spikes_last_evolution.plot()
plt.ylabel('Layer neuron')
plt.title('Layer activity')

plt.figure()
output_ts.plot()
plt.xlabel('Time(s)')
plt.ylabel('U')
plt.title('Weighted output surrogate');
_images/tutorials_jax_lif_tutorial_36_0.png
_images/tutorials_jax_lif_tutorial_36_1.png
_images/tutorials_jax_lif_tutorial_36_2.png
_images/tutorials_jax_lif_tutorial_36_3.png

Using JAX to perform gradient-based learning

To be continued…

This page was generated from docs/advanced/extending_layers.ipynb. Interactive online version: Binder badge

Writing a new Layer subclass

Rockpool can be extended by writing additional Layer subclasses. This brief guide shows you how to get started.

Housekeeping and import statements

[1]:
# - Import required modules and configure

# - Disable warning display
import warnings
warnings.filterwarnings('ignore')

# - Required imports
import numpy as np
import scipy.signal as sig

from rockpool.timeseries import (
    TimeSeries,
    TSContinuous,
    TSEvent,
    set_global_ts_plotting_backend,
)

# - Use HoloViews for plotting
import colorcet as cc
import holoviews as hv
hv.extension('bokeh')

%opts Curve [width=600]
%opts Scatter [width=600]

Functionality provided by Layer

The Layer base class provides several attributes and methods that make implementing a new subclass easier. See also the API reference for the Layer class.

  • Initialisation of attributes weights, size_in, name, noise_std, dt, t, _time_step

  • _prepare_input() method, which discretises inputs to a clocked time base

  • _checkinput_dims() method, which ensures that an input time series is the correct shape

  • _gen_time_trace() method, which generates a simulation time trace

  • _expand_to_net_size() method, which expands scalars to the dimensions of the layer, and checks that variables have the same size as the layer

  • _expand_to_weight_size() method, which does the same as above but to the size of .weights

  • reset_state() method, which sets attribute :py:attr`~.Layer.state` to all zeros

  • reset_time() method, which sets attribute :py:attr`~.Layer._time_step` to zero

  • randomize_state() method, which sets attribute :py:attr`~.Layer.state` to uniformly distributed random variates

  • Provides setters and getters for the attributes above

  • A reset_all() method, which calls reset_state() and reset_time()

  • Provides default attributes cInput and cOutput, which define what class of TimeSeries is expected for input and output

Representation of time

Each Layer subclass has an internal discrete representation of time. The attribute _time_step indicates how many time steps the layer object has evolved since instantiation. The length of a time step dt is normally provided as argument to the __init__() method. t is not an actual attribute but a property, defined as _time_step * dt. Both dt and t are in seconds.

Writing your class

Your class must inherit from Layer or a subclass:

class MyLayer(Layer):
...

You must provide an evolve() method (see below), which manages the simulation of your layer, and a to_dict() method, which converts the layer to a dictionary for saving.

You should probably define an __init__() method, which initialises everything needed for your layer (if you need to define more things). Note that the __init__() method does not call reset_all(), so you should do that in your __init__() method after calling super().__init__().

If your layer should accept or emit time series that are not TSContinuous, you must override the attributes cInput and cOutput appropriately. You should simply use return TSEvent or any subclass of TimeSeries.

Defining the evolve() method

In your subclass you must provide a method evolve(), which simulates the activity of neurons in the layer, and generates output signals.

The signature of this method should look like this:

def evolve(
    self,
    ts_input: Optional[TimeSeries] = None,
    duration: Optional[float] = None,
    num_timesteps: Optional[int] = None,
    verbose: Optional[bool] = False,
) -> TimeSeries:
    """
    Evolve the state of this layer given an input

    :param Optional[TimeSeries] ts_input: Input time series
    :param Optional[float] duration:      Simulation/Evolution time. If not provided, the duration of `ts_input` will be used
    :param Optional[int] time_steps:      Number of evolution time steps
    :param Optional[bool] verbose:        Output information about evolution status

    :return TimeSeries: Output time series

    """

You can include further arguments but should provide default values in that case.

If time_steps is provided to the method, this is the number of time steps over which the layer has to evolve. Otherwise, this number needs to be infered from duration by rounding down duration \ dt to an integer or, if duration is None, then the duration should be inferred from ts_input. In this case duration is ts_input.duration if ts_input.periodic is True, or otherwise ts_input.t_stop - self.t. time_steps is then determined the same way as when duration is provided.

Particularly for layers with continuous-time representations, the method Layer._prepare_input() is useful:

def _prepare_input(self,
                   ts_input: TimeSeries = None,
                   duration: float = None,
                   num_timesteps: int = None,
) -> (np.ndarray, np.ndarray, float):
"""
_prepare_input - Sample input, set up time base

:param TimeSeries tsInput: TxM or Tx1 Input signals for this layer
:param float duration:     Duration of the desired evolution, in seconds
:param int num_timesteps:  Number of evolution time steps

:return: (time_base, input_steps, num_timesteps)
    time_base:      ndarray T1 Discretised time base for evolution
    input_steps:    ndarray (T1xN) Discretised input signal for layer
    num_timesteps:  int Actual number of evolution time steps
"""

This method returns a clocked time base to use for the simulation time_base; a discretised version of the input signals input_steps, clocked to the same time base; and num_timesteps, the actual number of time steps by which the layer should be evolved. If num_timesteps is provided as input argument, the returned num_timesteps is identical, otherwise it is inferred from duration or from ts_input. It is equal to time_base.size - 1 because time_base also includes the current time point of the layer, which is not counted in num_timesteps.

Noise is generated within the evolve method. How you do so and how to add this to the internal state is up to you, but the attribute fNoiseStd is defined as the expected std. dev. of the noise after 1s of integration.

You should then evolve the activity of your layer; update state to be the layer state at the last time step of evolution; update _timestep to to reflect the new layer time, and return the output of the layer.

The precise definition of the layer output is of course up to you, but you must return a TimeSeries-subclass (the same class as cOutput). Furthermore, the output time series must cover the time range from the layer time before evolution until the end of the evolution.

Layer Properties and PyTorch backends

Many layer classes use properties with getters and setters for some attribues, instead of providing direct access to the attributes. This allows restrictions on how they are set as well as processing the data before updating the actual objects. In the following example a layer has one time constant for each neuron. The values can be accessed through an property tau. They are stored in an attribute _tau, while tau is only a property. This way it can be ensured that the time constants are always stored in a vector of the size of the layer.

[2]:
from rockpool.layers import Layer

class ExampleLayer(Layer):
    def __init__(mfW, vtTau):
        super().__init__(mfW)
        self.vtTau = vtTau

    def evolve(self):
        NotImplemented

    @property
    def vtTau(self):
        return self._vtTau

    @vtTau.setter
    def vtTau(self, vtNewTau):
        # Make sure _vtTau has correct dimensions
        self._vtTau = self._expand_to_net_size(vtNewTau, "vtTau")

The method _expand_to_net_size() is inherited from the Layer base class. It makes sure that the argument gets reshaped, so that it matches the size of the layer or throws an exception if this is not possible.

For layers that use torch as a backend, it is often helpful, if the properties return a numpy array instead of a torch tensor, so that other layers can work with it. This can be accomplished the following way:

[3]:
import torch

class ExampleTorchLayer(Layer):
    def __init__(self, mfW, vtTau, device):
        super().__init__(mfW)
        # Device on which torch tensors are processed
        self.device = device
        self.vtTau = vtTau

    def evolve(self):
        NotImplemented

    def to_dict(self):
        NotImplemented

    @property
    def vtTau(self):
        return self._vtTau.cpu().numpy()

    @vtTau.setter
    def vtTau(self, vtNewTau):
        # Make sure _vtTau has correct dimensions
        _vtTau = self._expand_to_net_size(vtNewTau, "vtTau")
        # Convert to torch tensor and move to self.device
        self._vtTau = torch.from_numpy(_vtTau).float().to(self.device)

Here the property formalism takes care of the convesion between numpy arrays and float tensors. It also makes sure that _tau is on the device (e.g. cuda) that is specified in device.

The problem with this method is, that if we try to change the time constants via item assignment, it might behave in an unexpected way:

[4]:
import numpy as np
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

mfW = np.random.randn(3,4)
vtTau = np.ones(4) * 0.5
ffExample = ExampleTorchLayer(mfW, vtTau, device)

# Print the initial value
print("Initial value:", ffExample.vtTau)

# Try to assign new value to all time constants
ffExample.vtTau = 0.1
print("Changing by reassignment:", ffExample.vtTau)

# Try to assign new value to first time constant
ffExample.vtTau[0] = 1
print("Changing by item assignment:", ffExample.vtTau)
Initial value: [0.5 0.5 0.5 0.5]
Changing by reassignment: [0.1 0.1 0.1 0.1]
Changing by item assignment: [1.  0.1 0.1 0.1]

The time constant of the layer has not changed because the property returned a copy of _tau and only this copy has been modified.

This can be fixed by replacing property with RefProperty, a subclass of property that is defined in the utilities module. The object that it returns is still a copy, but with a reference to the original array which allows to change the values even with item assignment:

[6]:
from rockpool.utilities import RefProperty

class ExampleTorchLayer(Layer):
    def __init__(self, mfW, vtTau, device):
        super().__init__(mfW)
        # Device on which torch tensors are processed
        self.device = device
        self.vtTau = vtTau

    def evolve(self):
        NotImplemented

    def to_dict(self):
        NotImplemented

    @RefProperty
    def vtTau(self):
        return self._vtTau

    @vtTau.setter
    def vtTau(self, vtNewTau):
        # Make sure _vtTau has correct dimensions
        _vtTau = self._expand_to_net_size(vtNewTau, "vtTau")
        # Convert to torch tensor and move to self.device
        self._vtTau = torch.from_numpy(_vtTau).float().to(self.device)


ffExample = ExampleTorchLayer(mfW, vtTau, device)

# - Print initial value
print("Initial value:", ffExample.vtTau)

# Try to assign new value to all time constants
ffExample.vtTau = 0.1
print("Changing by reassignment:", ffExample.vtTau)

# Try to assign new value to first time constant
ffExample.vtTau[0] = 1
print("Changing by item assignment:", ffExample.vtTau)
Initial value: [0.5 0.5 0.5 0.5]
Changing by reassignment: [0.1 0.1 0.1 0.1]
Changing by item assignment: [1.  0.1 0.1 0.1]

Note that now the function of in property only returns _tau instead of Layer._tau.cpu().numpy(). This is important because the RefProperty needs a reference to the actual object that has to bee modified. It will take care of the conversion between torch tensors and other array-like objects.

Types of Layer available in Rockpool

Rate-based non-spiking layers

layers.PassThrough(weights[, dt, noise_std, …])

Feed-forward layer with neuron states directly corresponding to input with an optional delay

layers.FFRateEuler(weights[, dt, name, …])

Feedforward layer consisting of rate-based neurons

layers.RecRateEuler(weights[, bias, tau, …])

A standard recurrent non-spiking layer of dynamical neurons

JAX-based backend

layers.RecRateEulerJax(w_in, w_recurrent, …)

JAX-backed firing-rate recurrent layer

layers.ForceRateEulerJax(w_in, w_out, tau, bias)

Implements a pseudo recurrent reservoir, for use in reservoir transfer

Event-driven spiking layers

layers.PassThroughEvents(weights[, dt, …])

Pass through events by routing to different channels

layers.FFExpSyn(weights[, bias, dt, …])

Define an exponential synapse layer with spiking inputs and current outputs

layers.RecDIAF(weights_in, weights_rec[, …])

Define a spiking recurrent layer based on quantized digital IAF neurons

layers.RecFSSpikeEulerBT([weights_fast, …])

Implement a spiking reservoir with tight E/I balance.

layers.FFUpDown(weights[, repeat_output, …])

Define a spiking feedforward layer to convert analogue inputs to up and down channels

JAX-based backend

layers.RecLIFJax(w_recurrent, tau_mem, tau_syn)

Recurrent spiking neuron layer (LIF), spiking input and spiking output.

layers.RecLIFCurrentInJax(w_recurrent, …)

Recurrent spiking neuron layer (LIF), current injection input and spiking output.

layers.RecLIFJax_IO(w_in, w_recurrent, …)

Recurrent spiking neuron layer (LIF), spiking input and weighted surrogate output.

layers.RecLIFCurrentInJax_IO(w_in, …[, …])

Recurrent spiking neuron layer (LIF), weighted current input and weighted surrogate output.

Layers with constant leak

layers.CLIAF(weights_in[, bias, v_thresh, …])

Abstract layer class of integrate and fire neurons with constant leak

layers.FFCLIAF(weights[, bias, v_thresh, …])

Feedforward layer of integrate and fire neurons with constant leak

layers.RecCLIAF(weights_in, weights_rec[, …])

Recurrent layer of integrate and fire neurons with constant leak

layers.SoftMaxLayer([weights, thresh, dt, name])

A spiking SoftMax layer with spiking inputs and outputs, and constant leak

layers.FFCLIAFCNNTorch

Brian-based backend

layers.FFIAFBrian(weights[, bias, dt, …])

DEPRECATED A spiking feedforward layer with current inputs and spiking outputs

layers.FFIAFSpkInBrian(weights[, bias, dt, …])

DEPRECATED Spiking feedforward layer with spiking inputs and outputs

layers.RecIAFBrian([weights, bias, dt, …])

DEPRECATED A spiking recurrent layer with current inputs and spiking outputs, using a Brian2 backend

layers.RecIAFSpkInBrian(weights_in, weights_rec)

DEPRECATED Spiking recurrent layer with spiking in- and outputs, and a Brian2 backend

layers.FFExpSynBrian([weights, dt, …])

DEPRECATED Define an exponential synapse layer (spiking input), with a Brian2 backend

Torch-based backend

layers.FFExpSynTorch([weights, bias, dt, …])

Define an exponential synapse layer (spiking input, pytorch as backend)

layers.FFIAFTorch(weights[, bias, dt, …])

Define a spiking feedforward layer with spiking outputs, with a PyTorch backend

layers.FFIAFRefrTorch(weights[, bias, dt, …])

A spiking feedforward layer with spiking outputs and refractoriness

layers.FFIAFSpkInTorch(weights[, bias, dt, …])

Spiking feedforward layer with spiking in- and outputs

layers.FFIAFSpkInRefrTorch(weights[, bias, …])

Spiking feedforward layer with spiking in- and outputs and refractoriness, using a PyTorch backend

layers.RecIAFTorch(weights[, bias, dt, …])

A spiking recurrent layer with input currents and spiking outputs

layers.RecIAFRefrTorch(weights[, bias, dt, …])

A spiking recurrent layer with current inputs, spiking outputs and refractoriness.

layers.RecIAFSpkInTorch(weights_in, weights_rec)

A spiking recurrent layer with spiking in- and outputs

layers.RecIAFSpkInRefrTorch(weights_in, …)

A spiking recurrent layer with spiking in- and outputs and refractoriness, and a PyTorch backend

layers.RecIAFSpkInRefrCLTorch(weights_in, …)

A recurrent spiking layer with constant leak.

layers.FFCLIAFCNNTorch

Nest-based backend

layers.FFIAFNest

layers.RecIAFSpkInNest

layers.RecAEIFSpkInNest

Hardware-backed and hardware simulation

For more information on using these layers, see Event-based simulation on DynapSE hardware

layers.RecDynapSE(weights_in, weights_rec[, …])

Recurrent spiking layer implemented with a DynapSE backend.

layers.VirtualDynapse

Full API summary for Rockpool

Base classes

networks.Network(*layers[, dt])

Base class to manage networks (collections of Layer objects)

layers.Layer(weights[, dt, noise_std, name])

Base class for Layers in rockpool

Layer and Network alternative base classes

networks.NetworkDeneve()

Network base class that defines and builds networks of Deneve reservoirs to solve linear problems.

layers.training.RRTrainedLayer(weights[, …])

Base class that defines methods for training with ridge regression.

Time series classes

timeseries.TimeSeries(times[, periodic, …])

Super-class to represent a continuous or event-based time series.

timeseries.TSContinuous([times, samples, …])

Represents a continuously-sampled time series.

timeseries.TSEvent([times, channels, …])

Represents a discrete time series, composed of binary events (present or absent).

Utility modules

The weights module provides several useful functions for generating network weights.

The utilities module provides several useful utility functions.

Layer subclasses

layers.RecRateEuler(weights[, bias, tau, …])

A standard recurrent non-spiking layer of dynamical neurons

layers.FFRateEuler(weights[, dt, name, …])

Feedforward layer consisting of rate-based neurons

layers.PassThrough(weights[, dt, noise_std, …])

Feed-forward layer with neuron states directly corresponding to input with an optional delay

layers.FFIAFBrian(weights[, bias, dt, …])

DEPRECATED A spiking feedforward layer with current inputs and spiking outputs

layers.FFIAFSpkInBrian(weights[, bias, dt, …])

DEPRECATED Spiking feedforward layer with spiking inputs and outputs

layers.RecIAFBrian([weights, bias, dt, …])

DEPRECATED A spiking recurrent layer with current inputs and spiking outputs, using a Brian2 backend

layers.RecIAFSpkInBrian(weights_in, weights_rec)

DEPRECATED Spiking recurrent layer with spiking in- and outputs, and a Brian2 backend

layers.PassThroughEvents(weights[, dt, …])

Pass through events by routing to different channels

layers.FFExpSynBrian([weights, dt, …])

DEPRECATED Define an exponential synapse layer (spiking input), with a Brian2 backend

layers.FFExpSyn(weights[, bias, dt, …])

Define an exponential synapse layer with spiking inputs and current outputs

layers.RecLIFJax(w_recurrent, tau_mem, tau_syn)

Recurrent spiking neuron layer (LIF), spiking input and spiking output.

layers.RecLIFCurrentInJax(w_recurrent, …)

Recurrent spiking neuron layer (LIF), current injection input and spiking output.

layers.RecLIFJax_IO(w_in, w_recurrent, …)

Recurrent spiking neuron layer (LIF), spiking input and weighted surrogate output.

layers.RecLIFCurrentInJax_IO(w_in, …[, …])

Recurrent spiking neuron layer (LIF), weighted current input and weighted surrogate output.

layers.FFCLIAF(weights[, bias, v_thresh, …])

Feedforward layer of integrate and fire neurons with constant leak

layers.RecCLIAF(weights_in, weights_rec[, …])

Recurrent layer of integrate and fire neurons with constant leak

layers.CLIAF(weights_in[, bias, v_thresh, …])

Abstract layer class of integrate and fire neurons with constant leak

layers.SoftMaxLayer([weights, thresh, dt, name])

A spiking SoftMax layer with spiking inputs and outputs, and constant leak

layers.RecDIAF(weights_in, weights_rec[, …])

Define a spiking recurrent layer based on quantized digital IAF neurons

layers.RecFSSpikeEulerBT([weights_fast, …])

Implement a spiking reservoir with tight E/I balance.

layers.FFUpDown(weights[, repeat_output, …])

Define a spiking feedforward layer to convert analogue inputs to up and down channels

layers.FFExpSynTorch([weights, bias, dt, …])

Define an exponential synapse layer (spiking input, pytorch as backend)

layers.FFIAFTorch(weights[, bias, dt, …])

Define a spiking feedforward layer with spiking outputs, with a PyTorch backend

layers.FFIAFRefrTorch(weights[, bias, dt, …])

A spiking feedforward layer with spiking outputs and refractoriness

layers.FFIAFSpkInTorch(weights[, bias, dt, …])

Spiking feedforward layer with spiking in- and outputs

layers.FFIAFSpkInRefrTorch(weights[, bias, …])

Spiking feedforward layer with spiking in- and outputs and refractoriness, using a PyTorch backend

layers.RecIAFTorch(weights[, bias, dt, …])

A spiking recurrent layer with input currents and spiking outputs

layers.RecIAFRefrTorch(weights[, bias, dt, …])

A spiking recurrent layer with current inputs, spiking outputs and refractoriness.

layers.RecIAFSpkInTorch(weights_in, weights_rec)

A spiking recurrent layer with spiking in- and outputs

layers.RecIAFSpkInRefrTorch(weights_in, …)

A spiking recurrent layer with spiking in- and outputs and refractoriness, and a PyTorch backend

layers.RecIAFSpkInRefrCLTorch(weights_in, …)

A recurrent spiking layer with constant leak.

layers.FFIAFNest

layers.RecIAFSpkInNest

layers.RecAEIFSpkInNest

layers.RecDynapSE(weights_in, weights_rec[, …])

Recurrent spiking layer implemented with a DynapSE backend.

layers.VirtualDynapse

layers.RecRateEulerJax(w_in, w_recurrent, …)

JAX-backed firing-rate recurrent layer

layers.ForceRateEulerJax(w_in, w_out, tau, bias)

Implements a pseudo recurrent reservoir, for use in reservoir transfer