Upgrade documentations(#310)

Upgrade documentations

Author: chaoming0625

brainpy/__init__.py

@@ -8,8 +8,12 @@
 
 # "base" module
 from . import base
-from .base.base import BrainPyObject, Base
-from .base.collector import Collector, TensorCollector
+from .base.base import (BrainPyObject, Base)
+from .base.collector import (
+  Collector,
+  ArrayCollector,
+  TensorCollector
+)
 
 # math foundation
 from . import math

brainpy/algorithms/offline.py

@@ -7,7 +7,7 @@
 
 import brainpy.math as bm
 from brainpy.base import BrainPyObject
-from brainpy.types import Array
+from brainpy.types import ArrayType
 from .utils import (Sigmoid,
                     Regularization, L1Regularization, L1L2Regularization, L2Regularization,
                     polynomial_features, normalize)
@@ -46,43 +46,43 @@ def __call__(self, identifier, target, input, output):
     ----------
     identifier: str
       The variable name.
-    target: Array, ndarray
+    target: ArrayType
       The 2d target data with the shape of `(num_batch, num_output)`.
-    input: Array, ndarray
+    input: ArrayType
       The 2d input data with the shape of `(num_batch, num_input)`.
-    output: Array, ndarray
+    output: ArrayType
       The 2d output data with the shape of `(num_batch, num_output)`.
 
     Returns
     -------
-    weight: Array
+    weight: ArrayType
       The weights after fit.
     """
     return self.call(identifier, target, input, output)
 
-  def call(self, identifier, targets, inputs, outputs) -> Array:
+  def call(self, identifier, targets, inputs, outputs) -> ArrayType:
     """The training procedure.
 
     Parameters
     ----------
     identifier: str
       The identifier.
 
-    inputs: Array, jax.numpy.ndarray, numpy.ndarray
+    inputs: ArrayType
       The 3d input data with the shape of `(num_batch, num_time, num_input)`,
       or, the 2d input data with the shape of `(num_time, num_input)`.
 
-    targets: Array, jax.numpy.ndarray, numpy.ndarray
+    targets: ArrayType
       The 3d target data with the shape of `(num_batch, num_time, num_output)`,
       or the 2d target data with the shape of `(num_time, num_output)`.
 
-    outputs: Array, jax.numpy.ndarray, numpy.ndarray
+    outputs: ArrayType
       The 3d output data with the shape of `(num_batch, num_time, num_output)`,
       or the 2d output data with the shape of `(num_time, num_output)`.
 
     Returns
     -------
-    weight: Array
+    weight: ArrayType
       The weights after fit.
     """
     raise NotImplementedError('Must implement the __call__ function by the subclass itself.')
@@ -355,7 +355,7 @@ def __init__(
     self.gradient_descent = gradient_descent
     self.sigmoid = Sigmoid()
 
-  def call(self, identifier, targets, inputs, outputs=None) -> Array:
+  def call(self, identifier, targets, inputs, outputs=None) -> ArrayType:
     # prepare data
     inputs = _check_data_2d_atls(bm.asarray(inputs))
     targets = _check_data_2d_atls(bm.asarray(targets))

brainpy/algorithms/online.py

@@ -34,16 +34,16 @@ def __call__(self, identifier, target, input, output):
     ----------
     identifier: str
       The variable name.
-    target: Array, ndarray
+    target: ArrayType
       The 2d target data with the shape of `(num_batch, num_output)`.
-    input: Array, ndarray
+    input: ArrayType
       The 2d input data with the shape of `(num_batch, num_input)`.
-    output: Array, ndarray
+    output: ArrayType
       The 2d output data with the shape of `(num_batch, num_output)`.
 
     Returns
     -------
-    weight: Array
+    weight: ArrayType
       The weights after fit.
     """
     return self.call(identifier, target, input, output)
@@ -58,16 +58,16 @@ def call(self, identifier, target, input, output):
     ----------
     identifier: str
       The variable name.
-    target: Array, ndarray
+    target: ArrayType
       The 2d target data with the shape of `(num_batch, num_output)`.
-    input: Array, ndarray
+    input: ArrayType
       The 2d input data with the shape of `(num_batch, num_input)`.
-    output: Array, ndarray
+    output: ArrayType
       The 2d output data with the shape of `(num_batch, num_output)`.
 
     Returns
     -------
-    weight: Array
+    weight: ArrayType
       The weights after fit.
     """
     raise NotImplementedError('Must implement the call() function by the subclass itself.')

brainpy/analysis/highdim/slow_points.py

@@ -13,12 +13,12 @@
 import brainpy.math as bm
 from brainpy import optimizers as optim, losses
 from brainpy.analysis import utils, base, constants
-from brainpy.base import TensorCollector
+from brainpy.base import ArrayCollector
 from brainpy.dyn.base import DynamicalSystem
 from brainpy.dyn.runners import build_inputs, check_and_format_inputs
 from brainpy.errors import AnalyzerError, UnsupportedError
 from brainpy.tools.others.dicts import DotDict
-from brainpy.types import Array
+from brainpy.types import ArrayType
 
 __all__ = [
   'SlowPointFinder',
@@ -136,11 +136,11 @@ def __init__(
 
     # update function
     if target_vars is None:
-      self.target_vars = TensorCollector()
+      self.target_vars = ArrayCollector()
     else:
       if not isinstance(target_vars, dict):
         raise TypeError(f'"target_vars" must be a dict but we got {type(target_vars)}')
-      self.target_vars = TensorCollector(target_vars)
+      self.target_vars = ArrayCollector(target_vars)
     excluded_vars = () if excluded_vars is None else excluded_vars
     if isinstance(excluded_vars, dict):
       excluded_vars = tuple(excluded_vars.values())
@@ -295,7 +295,7 @@ def selected_ids(self, val):
 
   def find_fps_with_gd_method(
       self,
-      candidates: Union[Array, Dict[str, Array]],
+      candidates: Union[ArrayType, Dict[str, ArrayType]],
       tolerance: Union[float, Dict[str, float]] = 1e-5,
       num_batch: int = 100,
       num_opt: int = 10000,
@@ -305,7 +305,7 @@ def find_fps_with_gd_method(
 
     Parameters
     ----------
-    candidates : Array, dict
+    candidates : ArrayType, dict
       The array with the shape of (batch size, state dim) of hidden states
       of RNN to start training for fixed points.
 
@@ -335,7 +335,7 @@ def find_fps_with_gd_method(
     # set up optimization
     num_candidate = self._check_candidates(candidates)
     if not (isinstance(candidates, (bm.ndarray, jnp.ndarray, np.ndarray)) or isinstance(candidates, dict)):
-      raise ValueError('Candidates must be instance of Array or dict of Array.')
+      raise ValueError('Candidates must be instance of ArrayType or dict of ArrayType.')
     fixed_points = tree_map(lambda a: bm.TrainVar(a), candidates, is_leaf=lambda x: isinstance(x, bm.Array))
     f_eval_loss = self._get_f_eval_loss()
 
@@ -401,14 +401,14 @@ def batch_train(start_i, n_batch):
 
   def find_fps_with_opt_solver(
       self,
-      candidates: Union[Array, Dict[str, Array]],
+      candidates: Union[ArrayType, Dict[str, ArrayType]],
       opt_solver: str = 'BFGS'
   ):
     """Optimize fixed points with nonlinear optimization solvers.
 
     Parameters
     ----------
-    candidates: Array, dict
+    candidates: ArrayType, dict
       The candidate (initial) fixed points.
     opt_solver: str
       The solver of the optimization.
@@ -535,7 +535,7 @@ def exclude_outliers(self, tolerance: float = 1e0):
 
   def compute_jacobians(
       self,
-      points: Union[Array, Dict[str, Array]],
+      points: Union[ArrayType, Dict[str, ArrayType]],
       stack_dict_var: bool = True,
       plot: bool = False,
       num_col: int = 4,
@@ -546,7 +546,7 @@ def compute_jacobians(
 
     Parameters
     ----------
-    points: np.ndarray, bm.Array, jax.ndarray
+    points: np.ndarray, bm.ArrayType, jax.ndarray
       The fixed points with the shape of (num_point, num_dim).
     stack_dict_var: bool
       Stack dictionary variables to calculate Jacobian matrix?
@@ -606,7 +606,7 @@ def decompose_eigenvalues(matrices, sort_by='magnitude', do_compute_lefts=False)
 
     Parameters
     ----------
-    matrices: np.ndarray, bm.Array, jax.ndarray
+    matrices: np.ndarray, bm.ArrayType, jax.ndarray
       A 3D array with the shape of (num_matrices, dim, dim).
     sort_by: str
       The method of sorting.

brainpy/analysis/utils/measurement.py

@@ -49,7 +49,7 @@ def euclidean_distance(points: np.ndarray, num_point=None):
 
   Parameters
   ----------
-  points: jnp.ndarray, bm.Array
+  points: ArrayType
     The points.
 
   Returns
@@ -96,13 +96,13 @@ def euclidean_distance_jax(points: Union[jnp.ndarray, bm.ndarray], num_point=Non
 
   Parameters
   ----------
-  points: jnp.ndarray, bm.Array
+  points: ArrayType
     The points.
   num_point: int
 
   Returns
   -------
-  dist_matrix: Array
+  dist_matrix: ArrayType
     The distance matrix.
   """
   if isinstance(points, dict):

brainpy/base/base.py

@@ -5,7 +5,7 @@
 
 from brainpy import errors
 from brainpy.base import io, naming
-from brainpy.base.collector import Collector, TensorCollector
+from brainpy.base.collector import Collector, ArrayCollector
 
 math = None
 
@@ -40,7 +40,7 @@ def __init__(self, name=None):
 
     # Used to wrap the implicit variables
     # which cannot be accessed by self.xxx
-    self.implicit_vars = TensorCollector()
+    self.implicit_vars = ArrayCollector()
 
     # Used to wrap the implicit children nodes
     # which cannot be accessed by self.xxx
@@ -113,14 +113,14 @@ def vars(self, method='absolute', level=-1, include_self=True):
 
     Returns
     -------
-    gather : TensorCollector
+    gather : ArrayCollector
       The collection contained (the path, the variable).
     """
     global math
     if math is None: from brainpy import math
 
     nodes = self.nodes(method=method, level=level, include_self=include_self)
-    gather = TensorCollector()
+    gather = ArrayCollector()
     for node_path, node in nodes.items():
       for k in dir(node):
         v = getattr(node, k)
@@ -144,7 +144,7 @@ def train_vars(self, method='absolute', level=-1, include_self=True):
 
     Returns
     -------
-    gather : TensorCollector
+    gather : ArrayCollector
       The collection contained (the path, the trainable variable).
     """
     global math
@@ -287,7 +287,7 @@ def save_states(self, filename, variables=None, **setting):
     ----------
     filename : str
       The file name which to store the model states.
-    variables: optional, dict, TensorCollector
+    variables: optional, dict, ArrayCollector
       The variables to save. If not provided, all variables retrieved by ``~.vars()`` will be used.
     """
     if variables is None:

brainpy/base/collector.py

@@ -7,6 +7,7 @@
 
 __all__ = [
   'Collector',
+  'ArrayCollector',
   'TensorCollector',
 ]
 
@@ -108,7 +109,7 @@ def subset(self, var_type):
     >>> # get all trainable variables
     >>> some_collector.subset(bp.math.TrainVar)
     >>>
-    >>> # get all Array
+    >>> # get all Variable
     >>> some_collector.subset(bp.math.Variable)
 
     or, it can be used to get a subset of integrators:
@@ -142,7 +143,7 @@ def unique(self):
     return gather
 
 
-class TensorCollector(Collector):
+class ArrayCollector(Collector):
   """A ArrayCollector is a dictionary (name, var)
   with some additional methods to make manipulation
   of collections of variables easy. A Collection
@@ -195,3 +196,6 @@ def from_other(cls, other: Union[Sequence, Dict]):
       return cls(other)
     else:
       raise TypeError
+
+
+TensorCollector = ArrayCollector

brainpy/base/function.py

@@ -71,7 +71,7 @@ def __init__(self, f: Optional[Callable], child_objs=None, dyn_vars=None, name=N
     # variables
     # ---
     if dyn_vars is not None:
-      self.implicit_vars = collector.TensorCollector()
+      self.implicit_vars = collector.ArrayCollector()
       global math
       if math is None: from brainpy import math
       if isinstance(dyn_vars, math.ndarray):

brainpy/base/io.py

@@ -7,7 +7,7 @@
 import numpy as np
 
 from brainpy import errors
-from brainpy.base.collector import TensorCollector
+from brainpy.base.collector import ArrayCollector
 
 logger = logging.getLogger('brainpy.base.io')
 
@@ -124,7 +124,7 @@ def _load(
 
 
 def _unique_and_duplicate(collector: dict):
-  gather = TensorCollector()
+  gather = ArrayCollector()
   id2name = dict()
   duplicates = ([], [])
   for k, v in collector.items():