brainpy.base module

The base module for whole BrainPy ecosystem.

This module provides the most fundamental class Base, and its associated helper class Collector.

For each instance of “Base” class, users can retrieve all the variables (or trainable variables), integrators, and nodes.

This module also provides a Function class to wrap user-defined functions. In each function, maybe several nodes are used, and users can initialize a Function by providing the nodes used in the function. Unfortunately, Function class does not have the ability to gather nodes automatically.

Base Class

Base([name])

The Base class for whole BrainPy ecosystem.

class brainpy.base.Base(name=None)[source]

The Base class for whole BrainPy ecosystem.

The subclass of Base includes:

  • Module in brainpy.dnn.base.py

  • DynamicalSystem in brainpy.simulation.brainobjects.base.py

ints(method='absolute')[source]

Collect all integrators in this node and the children nodes.

Parameters

method (str) – The method to access the integrators.

Returns

collector – The collection contained (the path, the integrator).

Return type

collector.Collector

load_states(filename)[source]

Load the model states.

Parameters

filename (str) – The filename which stores the model states.

nodes(method='absolute', _paths=None)[source]

Collect all children nodes.

Parameters
  • method (str) – The method to access the nodes.

  • _paths (set, Optional) – The data structure to solve the circular reference.

Returns

gather – The collection contained (the path, the node).

Return type

collector.Collector

save_states(filename, **setting)[source]

Save the model states.

Parameters

filename (str) – The file name which to store the model states.

train_vars(method='absolute')[source]

The shortcut for retrieving all trainable variables.

Parameters

method (str) – The method to access the variables. Support ‘absolute’ and ‘relative’.

Returns

gather – The collection contained (the path, the trainable variable).

Return type

collector.ArrayCollector

unique_name(name=None, type=None)[source]

Get the unique name for this object.

Parameters
  • name (str, optional) – The expected name. If None, the default unique name will be returned. Otherwise, the provided name will be checked to guarantee its uniqueness.

  • type (str, optional) – The type of this class, used for object naming.

Returns

name – The unique name for this object.

Return type

str

vars(method='absolute')[source]

Collect all variables in this node and the children nodes.

Parameters

method (str) – The method to access the variables.

Returns

gather – The collection contained (the path, the variable).

Return type

collector.ArrayCollector

Function Wrapper

Function(f, nodes[, name])

The wrapper for Python functions.

class brainpy.base.Function(f, nodes, name=None)[source]

The wrapper for Python functions.

Parameters
  • f (function) – The function to wrap.

  • nodes (list, tuple, dict) – The nodes in the defined function f.

  • name (optional, str) – The function name.

nodes(method='absolute', _paths=None)[source]

Collect all children nodes.

Parameters
  • method (str) – The method to access the nodes.

  • _paths (set, Optional) – The data structure to solve the circular reference.

Returns

gather – The collection contained (the path, the node).

Return type

collector.Collector

Collectors

Collector

A Collector is a dictionary (name, var) with some additional methods to make manipulation of collections of variables easy.

ArrayCollector

A ArrayCollector is a dictionary (name, var) with some additional methods to make manipulation of collections of variables easy.

class brainpy.base.Collector[source]

A Collector is a dictionary (name, var) with some additional methods to make manipulation of collections of variables easy. A Collector is ordered by insertion order. It is the object returned by Base.vars() and used as input in many Collector instance: optimizers, jit, etc…

dict()[source]

Get a dict with the key and the value data.

replace(key, new_value)[source]

Replace the original key with the new value.

subset(var_type, judge_func=None)[source]

Get the subset of the (key, value) pair.

subset() can be used to get a subset of some class:

>>> import brainpy as bp
>>>
>>> # get all trainable variables
>>> some_collector.subset(bp.math.TrainVar)
>>>
>>> # get all JaxArray
>>> some_collector.subset(bp.math.Variable)

or, it can be used to get a subset of integrators:

>>> # get all ODE integrators
>>> some_collector.subset(bp.integrators.ODE_INT)
Parameters
  • var_type (Any) – The type/class to match.

  • judge_func (optional, callable) –

unique()[source]

Get a new type of collector with unique values.

If one value is assigned to two or more keys, then only one pair of (key, value) will be returned.

update([E, ]**F) None.  Update D from dict/iterable E and F.[source]

If E is present and has a .keys() method, then does: for k in E: D[k] = E[k] If E is present and lacks a .keys() method, then does: for k, v in E: D[k] = v In either case, this is followed by: for k in F: D[k] = F[k]

class brainpy.base.ArrayCollector[source]

A ArrayCollector is a dictionary (name, var) with some additional methods to make manipulation of collections of variables easy. A Collection is ordered by insertion order. It is the object returned by DynamicalSystem.vars() and used as input in many DynamicalSystem instance: optimizers, Jit, etc…

assign(inputs)[source]

Assign data to all values.

Parameters

inputs (dict) – The data for each value in this collector.

data()[source]

Get all data in each value.

replicate()[source]

A context manager to use in a with statement that replicates the variables in this collection to multiple devices.

Important: replicating also updates the random state in order to have a new one per device.

Exporting and Loading

save_h5(filename, all_vars)

save_npz(filename, all_vars[, compressed])

save_pkl(filename, all_vars)

save_mat(filename, all_vars)

load_h5(filename, target)

load_npz(filename, target)

load_pkl(filename, target)

load_mat(filename, target)