Source code for nengo.networks.ensemblearray

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 access to each individual neuron, typically for the purpose of inhibiting 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_per_ensemble : 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) for param in ens_kwargs: if is_iterable(ens_kwargs[param]): ens_kwargs[param] = nengo.dists.Samples(ens_kwargs[param]) self.config[nengo.Ensemble].update(ens_kwargs) label_prefix = "" if label is None else label + "_" self.n_neurons_per_ensemble = 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
[docs] @with_self def add_neuron_input(self): """Adds a node that provides input to the neurons of all ensembles. Direct neuron input is useful for inhibiting the activity of all neurons in the ensemble array. This node is accessible through the 'neuron_input' attribute of this ensemble array. """ if self.neuron_input is not None: warnings.warn("neuron_input already exists. Returning.") return self.neuron_input if isinstance(self.ea_ensembles[0].neuron_type, nengo.Direct): raise ValidationError( "Ensembles use Direct neuron type. " "Cannot give neuron input to Direct neurons.", attr='ea_ensembles[0].neuron_type', obj=self) self.neuron_input = nengo.Node( size_in=self.n_neurons_per_ensemble * self.n_ensembles, label="neuron_input") for i, ens in enumerate(self.ea_ensembles): nengo.Connection( self.neuron_input[i * self.n_neurons_per_ensemble: (i + 1) * self.n_neurons_per_ensemble], ens.neurons, synapse=None) return self.neuron_input
[docs] @with_self 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_per_ensemble * 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_per_ensemble: (i + 1) * self.n_neurons_per_ensemble], synapse=None) return self.neuron_output
[docs] @with_self 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 synaptic filtering should be performed in 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