nn.modules.TimedModule

class nn.modules.TimedModule(*args, **kwargs)[source]

Bases: ModuleBase

The Rockpool base class for all TimedModule modules

TimedModule provides functionality for Module s to understand time series data, and to conveniently evolve, handle and return time series data from modules.

The evolve() method provided by TimedModule can accept TimeSeries objects natively as input, or can accept clocked / rasterised input data.

See also

TimedModule provides the useful methods _prepare_input() and _gen_timeseries() to help you in rasterising data for your own TimedModule subclasses.

For more information on how to used the TimedModule API for Rockpool, see ⏱ High-level TimedModule API.

Attributes overview

class_name

Class name of self

full_name

The full name of this module (class plus module name)

input_type

The TimeSeries class accepted by this module

name

The name of this module, or an empty string if None

output_type

The TimeSeries class returned by this module

shape

The shape of this module

size

(DEPRECATED) The output size of this module

size_in

The input size of this module

size_out

The output size of this module

spiking_input

If True, this module receives spiking input.

spiking_output

If True, this module sends spiking output.

t

The current evolution time of this layer, in seconds

dt

The simulation and input rasterisation timestep for this TimedModule

Methods overview

__init__(dt[,Β spiking_input,Β ...])

Initialise this TimedModule object

as_graph()

Convert this module to a computational graph

attributes_named(name)

Search for attributes of this or submodules by time

evolve([ts_input,Β duration,Β num_timesteps,Β ...])

Evolve the state of this module over time

modules()

Return a dictionary of all sub-modules of this module

parameters([family])

Return a nested dictionary of module and submodule Parameters

reset_all()

Reset the internal state and time of this module and all sub-modules

reset_parameters()

Reset all parameters in this module

reset_state()

Reset the state of this module

reset_time()

Reset the internal time of this module and all sub-modules to zero

set_attributes(new_attributes)

Set the attributes and sub-module attributes from a dictionary

simulation_parameters([family])

Return a nested dictionary of module and submodule SimulationParameters

state([family])

Return a nested dictionary of module and submodule States

__in_TimedModule_init: bool = False

A flag indicating that this TimedModule is currently being initialised

__init__(dt: float | SimulationParameter, spiking_input: bool = False, spiking_output: bool = False, add_events: bool = True, *args, **kwargs)[source]

Initialise this TimedModule object

When initialised, the TimedModule will have a dt attribute assigned, as well as initialising the internal module _timestep, _parent_dt_factor and _is_child. The subclass evolve() method will be wrapped to update the internal timestamp clock.

Parameters:
  • dt (float) – The duration of a single time step for this module, in seconds

  • spiking_input (bool) – If True, this module accepts TSEvent event time series objects as input. If False (default), this module accepts TSContinuous continuous time series objects as input.

  • spiking_output (bool) – If True, this module sends TSEvent event time series objects as output. If False (default), this module sends TSContinuous continuous time series objects as output.

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

_abc_impl = <_abc._abc_data object>
_determine_timesteps(ts_input: TimeSeries | None = None, duration: float | None = None, num_timesteps: int | None = None) int[source]

Determine how many time steps to evolve with the given input specification

Parameters:
  • ts_input (Optional[TimeSeries]) – TxM or Tx1 time series of input signals for this layer

  • duration (Optional[float]) – Duration of the desired evolution, in seconds. If not provided, num_timesteps or the duration of ts_input will be used to determine evolution time

  • num_timesteps (Optional[int]) – Number of evolution time steps, in units of dt. If not provided, duration or the duration of ts_input will be used to determine evolution time

Return int:

num_timesteps: Number of evolution time steps

_evolve_wrapper(ts_input=None, duration=None, num_timesteps=None, kwargs_timeseries=None, record: bool = False, *args, **kwargs) Tuple[TimeSeries, Dict, Dict][source]

Wrap a call to evolve() to update the internal time-steps count

See evolve() for calling syntax.

_force_set_attributes

(bool) If True, do not sanity-check attributes when setting.

_gen_time_trace(t_start: float, num_timesteps: int) ndarray[source]

Generate a time trace starting at t_start, of length num_timesteps with time step dt

Parameters:
  • t_start (float) – Start time, in seconds

  • num_timesteps (int) – Number of time steps to generate, in units of dt

Return ndarray:

Generated time trace

_gen_timeseries(output: ndarray, **kwargs) TimeSeries[source]

Wrap a clocked / rasterised output array into a TimeSeries object

Output TimeSeries will be of the appropriate subclass, and will be named nicely.

Parameters:
  • output (np.ndarray) – The clocked or rasterised output data (T, N)

  • **kwargs – Additional keyword arguments to TimeSeries

Returns:

The data in output wrapped into a TimeSeries object

Return type:

TimeSeries

_gen_tscontinuous(output: ndarray, dt: float | None = None, t_start: float | None = None, name: str | None = None, periodic: bool = False, interp_kind: str = 'previous') TSContinuous[source]

Wrap a rasterised output array as a TSContinuous object to present as output for this module

Output TSContinuous s will be named nicely, with correct start times, durations, etc. Several attributes of the TSContinuous object can be set as arguments here.

Parameters:
  • output (np.ndarray) – A clocked time series data array (T, N)

  • dt (Optional[float]) – The time-step of the clocked array output. If not provided, the module dt will be used

  • t_start (Optional[float]) – The start time of the output TSContinuous object, in seconds. If not provided, the module time before evolution will be used

  • name (Optional[str]) – The desired name of the TSContinuous object. If not provided, the object will be named nicely according to the module name

  • periodic (bool) – Flag to indicate whether the returned TSContinuous should be periodic. Default: False, the TSContinuous will not be periodic

  • interp_kind (str) – The style of interpolation to apply to the returned TSContinuous object. Default: "previous"

Returns:

The wrapped output data as a TSContinuous object

Return type:

TSContinuous

_gen_tsevent(output: ndarray, dt: float | None = None, t_start: float | None = None, name: str | None = None, periodic: bool = False, num_channels: int | None = None, spikes_at_bin_start: bool = False) TSEvent[source]

Wrap a rasterised output array as a TSEvent object to present as output for this module

Output TSEvent s will be named nicely, with correct start timesm durations, etc. Several attributes of the TSEvent object can be set as arguments here.

Parameters:
  • output (np.ndarray) – A rasterised event array (T, N)

  • dt (Optional[float]) – The time-step of the rasterised array output. If not provided, the module dt will be used

  • t_start (Optional[float]) – The start time of the output series, in seconds. If not provided, the module time before evolution will be used

  • name (Optional[str]) – The desired name of the TSEvent object. If not provided, the object will be named nicely according to the module name

  • periodic (bool) – Flag to indicate whether the returned TSEvent should be periodic. Default: False, the TSEvent will not be periodic

  • num_channels (Optional[int]) – The desired number of total channels for the output TSEvent object. If not provided, the output size size_out of the current module will be used

  • spikes_at_bin_start (bool) – If False (default), spike events will be considered to fall in the middle of the time bin they fall in. If True, all spike events will be considered to occur at the start of the time bin they fall in.

Returns:

The wrapped output raster as a TSEvent object

Return type:

TSEvent

_get_attribute_family(type_name: str, family: Tuple | List | str | None = None) dict

Search for attributes of this module and submodules that match a given family

This method can be used to conveniently get all weights for a network; or all time constants; or any other family of parameters. Parameter families are defined simply by a string: "weights" for weights; "taus" for time constants, etc. These strings are arbitrary, but if you follow the conventions then future developers will thank you (that includes you in six month’s time).

Parameters:
  • type_name (str) – The class of parameters to search for. Must be one of ["Parameter", "SimulationParameter", "State"] or another future subclass of ParameterBase

  • family (Union[str, Tuple[str]]) – A string or list or tuple of strings, that define one or more attribute families to search for

Returns:

A nested dictionary of attributes that match the provided type_name and family

Return type:

dict

_get_attribute_registry() Tuple[Dict, Dict]

Return or initialise the attribute registry for this module

Returns:

registered_attributes, registered_modules

Return type:

(tuple)

_has_registered_attribute(name: str) bool

Check if the module has a registered attribute

Parameters:

name (str) – The name of the attribute to check

Returns:

True if the attribute name is in the attribute registry, False otherwise.

Return type:

bool

_in_Module_init

(bool) If exists and True, indicates that the module is in the __init__ chain.

_is_child: bool

Flag indicating that this is a child module

_name: str | None

Name of this module, if assigned

_parent_dt_factor: float

The factor between the parent’s dt and this module’s dt. Given by self.dt / parent.dt

_prepare_input(ts_input: TimeSeries | None = None, duration: float | None = None, num_timesteps: int | None = None) Tuple[ndarray, ndarray, int][source]

Sample input, set up time base

This function checks an input signal, and prepares a discretised time base according to the time step of the current module

Parameters:
  • ts_input (Optional[TimeSeries]) – TimeSeries of TxM or Tx1 Input signals for this layer

  • duration (Optional[float]) – Duration of the desired evolution, in seconds. If not provided, then either num_timesteps or the duration of ts_input will define the evolution time

  • num_timesteps (Optional[int]) – Integer number of evolution time steps, in units of dt. If not provided, then duration or the duration of ts_input will define the evolution time

Return (ndarray, ndarray, int):

(time_base, input_steps, num_timesteps) time_base: T1 Discretised time base for evolution input_raster (T1xN) Discretised input signal for layer num_timesteps: Actual number of evolution time steps, in units of dt

_prepare_input_continuous(ts_input: TSContinuous | None = None, duration: float | None = None, num_timesteps: int | None = None) Tuple[ndarray, ndarray, int][source]

Sample input, set up time base

This function checks an input signal, and prepares a discretised time base according to the time step of the current module

Parameters:
  • ts_input (Optional[TSContinuous]) – TSContinuous of TxM or Tx1 Input signals for this layer

  • duration (Optional[float]) – Duration of the desired evolution, in seconds. If not provided, then either num_timesteps or the duration of ts_input will define the evolution time

  • num_timesteps (Optional[int]) – Integer number of evolution time steps, in units of dt. If not provided, then duration or the duration of ts_input will define the evolution time

Return (ndarray, ndarray, int):

(time_base, input_raster, num_timesteps) time_base: T1 Discretised time base for evolution input_raster: (T1xN) Discretised input signal for layer num_timesteps: Actual number of evolution time steps, in units of dt

_prepare_input_events(ts_input: TSEvent | None = None, duration: float | None = None, num_timesteps: int | None = None, add_events: bool = False) Tuple[ndarray, ndarray, int][source]

Sample input from a TSEvent time series, set up evolution time base

This function checks an input signal, and prepares a discretised time base according to the time step of the current module

Parameters:
  • ts_input (Optional[TSEvent]) – TimeSeries of TxM or Tx1 Input signals for this layer

  • duration (Optional[float]) – Duration of the desired evolution, in seconds. If not provided, then either num_timesteps or the duration of ts_input will determine evolution itme

  • num_timesteps (Optional[int]) – Number of evolution time steps, in units of dt. If not provided, then either duration or the duration of ts_input will determine evolution time

Return (ndarray, ndarray, int):

time_base: T1X1 vector of time points – time base for the rasterisation spike_raster: Boolean or integer raster containing spike information. T1xM array num_timesteps: Actual number of evolution time steps, in units of dt

_register_attribute(name: str, val: ParameterBase)

Record an attribute in the attribute registry

Parameters:
  • name (str) – The name of the attribute to register

  • val (ParameterBase) – The ParameterBase subclass object to register. e.g. Parameter, SimulationParameter or State.

_register_module(name: str, mod: ModuleBase)

Register a sub-module in the module registry

Parameters:
  • name (str) – The name of the module to register

  • mod (ModuleBase) – The ModuleBase object to register

_reset_attribute(name: str) ModuleBase

Reset an attribute to its initialisation value

Parameters:

name (str) – The name of the attribute to reset

Returns:

For compatibility with the functional API

Return type:

self (Module)

_set_dt(max_factor: float = 100) None[source]

Set a time step size for the network which is the lowest common multiple of all sub-module’s dt s.

Parameters:

max_factor (float) – Factor by which the module dt may exceed the largest sub-module dt before an error is raised. Default: 100.

Raises:

ValueError – If a sensible dt cannot be found

_shape

The shape of this module

_spiking_input: bool

Whether this module receives spiking input

_spiking_output: bool

Whether this module produces spiking output

_submodulenames: List[str]

Registry of sub-module names

_timestep: int

The current time-step count in units of dt

_wrap_recorded_state(recorded_dict: dict, t_start: float) Dict[str, TimeSeries]

Convert a recorded dictionary to a TimeSeries representation

This method is optional, and is provided to make the timed() conversion to a TimedModule work better. You should override this method in your custom Module, to wrap each element of your recorded state dictionary as a TimeSeries

Parameters:
  • state_dict (dict) – A recorded state dictionary as returned by evolve()

  • t_start (float) – The initial time of the recorded state, to use as the starting point of the time series

Returns:

The mapped recorded state dictionary, wrapped as TimeSeries objects

Return type:

Dict[str, TimeSeries]

as_graph() GraphModuleBase

Convert this module to a computational graph

Returns:

The computational graph corresponding to this module

Return type:

GraphModuleBase

Raises:

NotImplementedError – If as_graph() is not implemented for this subclass

attributes_named(name: Tuple[str] | List[str] | str) dict

Search for attributes of this or submodules by time

Parameters:

name (Union[str, Tuple[str]) – The name of the attribute to search for

Returns:

A nested dictionary of attributes that match name

Return type:

dict

property class_name: str

Class name of self

Type:

str

dt: float | SimulationParameter

The simulation and input rasterisation timestep for this TimedModule

Type:

float

abstract evolve(ts_input: TimeSeries | ndarray | None = None, duration: float | None = None, num_timesteps: int | None = None, kwargs_timeseries: dict | None = None, record: bool = False, *args, **kwargs) Tuple[TimeSeries, Dict, Dict][source]

Evolve the state of this module over time

Warning

If you are seeing this message in documentation for a TimedModule subclass, then THIS CLASS HAS NOT PROVIDED DOCUMENTATION FOR ITS EVOLVE METHOD. PLEASE UPDATE THE DOCUMENTATION TO INCLUDE SPECIFIC DETAILS FOR THIS CLASS.

You need to implement an evolve() method for each class which inherits from TimedModule.

Here is an example evolve() method that rasterises a time series and uses the rasterised version for further processing. The output data is re-wrapped as a time series and returned.

def evolve(...):
    # - Rasterise input and prepare input time steps
    time_base, input_raster, num_timesteps = self._prepare_input(
        ts_input, duration, num_timesteps
    )

    # - Call sub-modules, do your evolution, etc.

    # - Return and wrap outputs if necessary
    return (
        self._gen_timeseries(output, **kwargs_timeseries),
        new_state,
        record_dict,
    )

Here is an example evolve() method that uses TimeSeries objects natively. Any rasterisation would be taken care of by submodules, if and when required.

def evolve(...):
    new_state = {}
    record = {}

    x1, new_state1, record1 = self.submodule(input_ts)
    new_state.update({'submodule': new_state1})
    record.update({'submodule': record1})

    x2, new_state2, record2 = self.submodule2(x1)
    new_state.update({'submodule2': new_state2})
    record.update({'submodule2': record2})

    return x2, new_state, record

You can of course use a mixture of these approaches.

Parameters:
  • ts_input (Union[TimeSeries, np.ndarray]) – The input time series over which to evolve

  • duration (float) – The duration over which to evolve, in seconds

  • num_timesteps (int) – The number of time steps (in terms of the dt attribute of this module) to evolve over

  • kwargs_timeseries (Optional[dict]) – Any additional arguments to pass when generating output time series

  • record (bool) – If True, this module and sub-modules must record their state during evolution and return it in the record_state dict. If False (default), no recording is requested

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

Returns:

(output_ts, new_state, record_state)

output_ts TimeSeries: A time series containing the output time series produces by this module. new_state dict: A dictionary containing the updated state of this module and sub-modules, after evolution record_state dict: If the argument record is True, record_state must contain a dictionary of the recorded states o this and all sub-modules during evolution. Otherwise it may be an empty dict.

Return type:

tuple

property full_name: str

The full name of this module (class plus module name)

Type:

str

property input_type: type

The TimeSeries class accepted by this module

Type:

type

modules() Dict

Return a dictionary of all sub-modules of this module

Returns:

A dictionary containing all sub-modules. Each item will be named with the sub-module name.

Return type:

dict

property name: str

The name of this module, or an empty string if None

Type:

str

property output_type: type

The TimeSeries class returned by this module

Type:

type

parameters(family: Tuple | List | str | None = None) Dict

Return a nested dictionary of module and submodule Parameters

Use this method to inspect the Parameters from this and all submodules. The optional argument family allows you to search for Parameters in a particular family β€” for example "weights" for all weights of this module and nested submodules.

Although the family argument is an arbitrary string, reasonable choises are "weights", "taus" for time constants, "biases" for biases…

Examples

Obtain a dictionary of all Parameters for this module (including submodules):

>>> mod.parameters()
dict{ ... }

Obtain a dictionary of Parameters from a particular family:

>>> mod.parameters("weights")
dict{ ... }
Parameters:

family (str) – The family of Parameters to search for. Default: None; return all parameters.

Returns:

A nested dictionary of Parameters of this module and all submodules

Return type:

dict

reset_all() None[source]

Reset the internal state and time of this module and all sub-modules

reset_parameters()

Reset all parameters in this module

Returns:

The updated module is returned for compatibility with the functional API

Return type:

Module

reset_state() ModuleBase

Reset the state of this module

Returns:

The updated module is returned for compatibility with the functional API

Return type:

Module

reset_time() None[source]

Reset the internal time of this module and all sub-modules to zero

set_attributes(new_attributes: dict) ModuleBase

Set the attributes and sub-module attributes from a dictionary

This method can be used with the dictionary returned from module evolution to set the new state of the module. It can also be used to set multiple parameters of a module and submodules.

Examples

Use the functional API to evolve, obtain new states, and set those states:

>>> _, new_state, _ = mod(input)
>>> mod = mod.set_attributes(new_state)

Obtain a parameter dictionary, modify it, then set the parameters back:

>>> params = mod.parameters()
>>> params['w_input'] *= 0.
>>> mod.set_attributes(params)
Parameters:

new_attributes (dict) – A nested dictionary containing parameters of this module and sub-modules.

property shape: tuple

The shape of this module

Type:

tuple

simulation_parameters(family: Tuple | List | str | None = None) Dict

Return a nested dictionary of module and submodule SimulationParameters

Use this method to inspect the SimulationParameters from this and all submodules. The optional argument family allows you to search for SimulationParameters in a particular family.

Examples

Obtain a dictionary of all SimulationParameters for this module (including submodules):

>>> mod.simulation_parameters()
dict{ ... }
Parameters:

family (str) – The family of SimulationParameters to search for. Default: None; return all SimulationParameter attributes.

Returns:

A nested dictionary of SimulationParameters of this module and all submodules

Return type:

dict

property size: int

(DEPRECATED) The output size of this module

Type:

int

property size_in: int

The input size of this module

Type:

int

property size_out: int

The output size of this module

Type:

int

property spiking_input: bool

If True, this module receives spiking input. If False, this module expects continuous input.

Type:

bool

property spiking_output

If True, this module sends spiking output. If False, this module sends continuous output.

Type:

bool

state(family: Tuple | List | str | None = None) Dict

Return a nested dictionary of module and submodule States

Use this method to inspect the States from this and all submodules. The optional argument family allows you to search for States in a particular family.

Examples

Obtain a dictionary of all States for this module (including submodules):

>>> mod.state()
dict{ ... }
Parameters:

family (str) – The family of States to search for. Default: None; return all State attributes.

Returns:

A nested dictionary of States of this module and all submodules

Return type:

dict

property t: float

The current evolution time of this layer, in seconds

Type:

float