Implementing a custom module

This example demonstrates how custom SPA modules can be created that can take advantage of all the features of the SPA syntax. We will adapt the InputGatedMemory from nengo.

import nengo
import nengo_spa as spa

Implementing a SPA module requires a few steps:

1. Implement a class inheriting from spa.Network.

2. Use VocabularyOrDimParam to declare class variables for storing vocublary parameters. This will allow the usage of integer dimensions instead of vocabularies without any further additions.

3. Declare inputs and outputs with their respective vocabularies.

Not that parameters in SPA modules should usually be defined as readonly because changing them will usually not update the network accordingly.

class GatedMemory(spa.Network):

    # The vocabulary parameter.
    vocab = spa.vocabulary.VocabularyOrDimParam("vocab", default=None, readonly=True)
    # The number of neurons per dimensions.
    neurons_per_dimension = nengo.params.IntParam(
        "neurons_per_dimension", default=200, low=1, readonly=True
    )

    # Arguments assigned to parameters should be assigned
    # nengo.params.Default as default value. This makes sure they work
    # properly with nengo.Config. It is a good idea to pass on the keyword
    # arguments **kwargs to the spa.Network constructor to allow the user to
    # set the network label etc.
    def __init__(
        self,
        vocab=nengo.params.Default,
        neurons_per_dimension=nengo.params.Default,
        **kwargs
    ):
        super(GatedMemory, self).__init__(**kwargs)

        # Assign parameter values
        # If vocab is an integer dimension, the appropriate Vocabulary
        # instance will assigned to self.vocab.
        self.vocab = vocab
        self.neurons_per_dimension = neurons_per_dimension

        # Construct the network
        with self:
            self.mem = nengo.networks.InputGatedMemory(
                self.neurons_per_dimension, self.vocab.dimensions
            )

        # Assign inputs to root object for easier referencing
        self.input = self.mem.input
        self.input_gate = self.mem.gate
        self.input_reset = self.mem.reset
        self.output = self.mem.output

        # Declare inputs and outputs
        # Use None as vocabulary for scalar inputs/outputs
        self.declare_input(self.input, self.vocab)
        self.declare_input(self.input_gate, None)
        self.declare_input(self.input_reset, None)
        self.declare_output(self.output, self.vocab)

We can then use our new module as we would any other module.

dimensions = 32

with spa.Network() as model:
    # The module can be configured
    model.config[GatedMemory].neurons_per_dimension = 150

    spa_in = spa.Transcode("OKAY", output_vocab=dimensions)
    gate_in = nengo.Node(lambda t: 1 if t < 0.1 else 0)

    g_mem = GatedMemory(dimensions)

    # It can be in routing rules
    spa_in >> g_mem
    gate_in >> g_mem.input_gate