brainpy.base module

The base module for whole BrainPy ecosystem.

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

  • 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.

  • This module provides io helper functions to help users save/load model states, or share user’s customized model with others.

Details please see the following.

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:

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

  • Function in brainpy.base.function.py

  • AutoGrad in brainpy.math.jax.autograd.py

  • Optimizer in brainpy.math.jax.optimizers.py

  • Scheduler in brainpy.math.jax.optimizers.py

implicit_nodes = None

Used to wrap the implicit children nodes which cannot be accessed by self.xxx

implicit_vars = None

Used to wrap the implicit variables which cannot be accessed by self.xxx

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

load_states(filename, verbose=False, check=False)[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

save_states(filename, all_vars=None, **setting)[source]

Save the model states.

Parameters

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

target_backend = None

Used to specify the target backend which the model to run.

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

TensorCollector

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

TensorCollector

Function Wrapper

Function(f[, nodes, dyn_vars, name])

The wrapper for Python functions.

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

The wrapper for Python functions.

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

  • nodes (optional, Base, sequence of Base, dict) – The nodes in the defined function f.

  • dyn_vars (optional, ndarray, sequence of ndarray, dict) – The dynamically changed variables.

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

Collectors

Collector

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

TensorCollector

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…

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
>>>
>>> some_collector = Collector()
>>>
>>> # 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.ode.ODEIntegrator)
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.TensorCollector[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.

dict()[source]

Get a dict with the key and the value data.

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[, verbose, check])

load_npz(filename, target[, verbose, check])

load_pkl(filename, target[, verbose, check])

load_mat(filename, target[, verbose, check])