import warnings
import numpy as np
import nengo
from nengo.exceptions import ValidationError
from nengo.utils.compat import is_iterable, range
from nengo.utils.network import with_self
[docs]class EnsembleArray(nengo.Network):
"""An array of ensembles.
This acts, in some ways, like a single high-dimensional ensemble,
but actually consists of many sub-ensembles, each one representing
a separate dimension. This tends to be much faster to create
and can be more accurate than having one huge high-dimensional ensemble.
However, since the neurons represent different dimensions separately,
we cannot compute nonlinear interactions between those dimensions.
Note that in addition to the parameters below, parameters affecting
all of the sub-ensembles can be passed to the ensemble array.
For example::
ea = nengo.networks.EnsembleArray(20, 2, radius=1.5)
creates an ensemble array with 2 sub-ensembles, each with 20 neurons,
and a radius of 1.5.
Parameters
----------
n_neurons : int
The number of neurons in each sub-ensemble.
n_ensembles : int
The number of sub-ensembles to create.
ens_dimensions : int, optional (Default: 1)
The dimensionality of each sub-ensemble.
neuron_nodes : bool, optional (Default: False)
Whether to create a node that provides each access to each individual
neuron, typically for the purpose of inibiting the entire EnsembleArray.
.. note:: Deprecated in Nengo 2.1.0.
Call `~.EnsembleArray.add_neuron_input` or
`~.EnsembleArray.add_neuron_output` instead.
label : str, optional (Default: None)
A name to assign this EnsembleArray.
Used for visualization and debugging.
seed : int, optional (Default: None)
Random number seed that will be used in the build step.
add_to_container : bool, optional (Default: None)
Determines if this network will be added to the current container.
If None, this network will be added to the network at the top of the
``Network.context`` stack unless the stack is empty.
Attributes
----------
dimensions_per_ensemble : int
The dimensionality of each sub-ensemble.
ea_ensembles : list
The sub-ensembles in the ensemble array.
input : Node
A node that provides input to all of the ensembles in the array.
n_ensembles : int
The number of sub-ensembles to create.
n_neurons : int
The number of neurons in each sub-ensemble.
neuron_input : Node or None
A node that provides input to all the neurons in the ensemble array.
None unless created in `~.EnsembleArray.add_neuron_input`.
neuron_output : Node or None
A node that gathers neural output from all the neurons in the ensemble
array. None unless created in `~.EnsembleArray.add_neuron_output`.
output : Node
A node that gathers decoded output from all of the ensembles
in the array.
"""
def __init__(self, n_neurons, n_ensembles, ens_dimensions=1,
neuron_nodes=False, label=None, seed=None,
add_to_container=None, **ens_kwargs):
if "dimensions" in ens_kwargs:
raise ValidationError(
"'dimensions' is not a valid argument to EnsembleArray. "
"To set the number of ensembles, use 'n_ensembles'. To set "
"the number of dimensions per ensemble, use 'ens_dimensions'.",
attr='dimensions', obj=self)
super(EnsembleArray, self).__init__(label, seed, add_to_container)
self.config[nengo.Ensemble].update(ens_kwargs)
label_prefix = "" if label is None else label + "_"
self.n_neurons = n_neurons
self.n_ensembles = n_ensembles
self.dimensions_per_ensemble = ens_dimensions
# These may be set in add_neuron_input and add_neuron_output
self.neuron_input, self.neuron_output = None, None
self.ea_ensembles = []
with self:
self.input = nengo.Node(size_in=self.dimensions, label="input")
for i in range(n_ensembles):
e = nengo.Ensemble(n_neurons, self.dimensions_per_ensemble,
label="%s%d" % (label_prefix, i))
nengo.Connection(self.input[i * ens_dimensions:
(i + 1) * ens_dimensions],
e, synapse=None)
self.ea_ensembles.append(e)
if neuron_nodes:
self.add_neuron_input()
self.add_neuron_output()
warnings.warn(
"'neuron_nodes' argument will be removed in Nengo 2.2. Use "
"'add_neuron_input' and 'add_neuron_output' methods instead.",
DeprecationWarning)
self.add_output('output', function=None)
@property
def dimensions(self):
"""(int) Dimensionality of the ensemble array."""
return self.n_ensembles * self.dimensions_per_ensemble
@with_self
@with_self
[docs] def add_neuron_output(self):
"""Adds a node that collects the neural output of all ensembles.
Direct neuron output is useful for plotting the spike raster of
all neurons in the ensemble array.
This node is accessible through the 'neuron_output' attribute
of this ensemble array.
"""
if self.neuron_output is not None:
warnings.warn("neuron_output already exists. Returning.")
return self.neuron_output
if isinstance(self.ea_ensembles[0].neuron_type, nengo.Direct):
raise ValidationError(
"Ensembles use Direct neuron type. "
"Cannot get neuron output from Direct neurons.",
attr='ea_ensembles[0].neuron_type', obj=self)
self.neuron_output = nengo.Node(
size_in=self.n_neurons * self.n_ensembles, label="neuron_output")
for i, ens in enumerate(self.ea_ensembles):
nengo.Connection(ens.neurons,
self.neuron_output[i * self.n_neurons:
(i + 1) * self.n_neurons],
synapse=None)
return self.neuron_output
@with_self
[docs] def add_output(self, name, function, synapse=None, **conn_kwargs):
"""Adds a node that collects the decoded output of all ensembles.
By default, this is called once in ``__init__`` with ``function=None``.
However, this can be called multiple times with different functions,
similar to the way in which an ensemble can be connected to many
downstream ensembles with different functions.
Note that in addition to the parameters below, parameters affecting
all of the connections from the sub-ensembles to the new node
can be passed to this function. For example::
ea.add_output('output', None, solver=nengo.solers.Lstsq())
creates a new output with the decoders of each connection solved for
with the `.Lstsq` solver.
Parameters
----------
name : str
The name of the output. This will also be the name of the attribute
set on the ensemble array.
function : callable or iterable of callables
The function to compute across the connection from sub-ensembles
to the new output node. If function is an iterable, it must be
an iterable consisting of one function for each sub-ensemble.
synapse : Synapse, optional (Default: None)
The synapse model with which to filter the connections from
sub-ensembles to the new output node. This is kept separate from
the other ``conn_kwargs`` because this defaults to None rather
than the default synapse model. In almost all cases, the synapse
should stay as None, and instead applied to the connection from
the output node.
"""
dims_per_ens = self.dimensions_per_ensemble
# get output size for each ensemble
sizes = np.zeros(self.n_ensembles, dtype=int)
if is_iterable(function) and all(callable(f) for f in function):
if len(list(function)) != self.n_ensembles:
raise ValidationError(
"Must have one function per ensemble", attr='function')
for i, func in enumerate(function):
sizes[i] = np.asarray(func(np.zeros(dims_per_ens))).size
elif callable(function):
sizes[:] = np.asarray(function(np.zeros(dims_per_ens))).size
function = [function] * self.n_ensembles
elif function is None:
sizes[:] = dims_per_ens
function = [None] * self.n_ensembles
else:
raise ValidationError("'function' must be a callable, list of "
"callables, or None", attr='function')
output = nengo.Node(output=None, size_in=sizes.sum(), label=name)
setattr(self, name, output)
indices = np.zeros(len(sizes) + 1, dtype=int)
indices[1:] = np.cumsum(sizes)
for i, e in enumerate(self.ea_ensembles):
nengo.Connection(
e, output[indices[i]:indices[i+1]], function=function[i],
synapse=synapse, **conn_kwargs)
return output