consolidate the concept of OO transformation (#309)

consolidate the concept of OO transformation

Author: chaoming0625

.github/ISSUE_TEMPLATE/bug_report.md

@@ -12,7 +12,7 @@ Please:
 ```python
 import brainpy.math as bm
 print(bm.asarray([1, 2, 3]))
-# JaxArray([1, 2, 3], dtype=int32)
+# Array([1, 2, 3], dtype=int32)
 ```
 
 - [ ] If applicable, include full error messages/tracebacks.

.github/workflows/Linux_CI.yml

@@ -29,6 +29,7 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
+        pip uninstall brainpy -y
         python setup.py install
     - name: Lint with flake8
       run: |

.github/workflows/MacOS_CI.yml

@@ -29,6 +29,7 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
+        pip uninstall brainpy -y
         python setup.py install
     - name: Lint with flake8
       run: |

.github/workflows/Windows_CI.yml

@@ -29,10 +29,11 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest
         python -m pip install numpy>=1.21.0
-        python -m pip install "jaxlib==0.3.14" -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver
-        python -m pip install https://github.com/google/jax/archive/refs/tags/jax-v0.3.14.tar.gz
+        python -m pip install "jaxlib==0.3.25" -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver
+        python -m pip install https://github.com/google/jax/archive/refs/tags/jax-v0.3.25.tar.gz
         python -m pip install -r requirements-dev.txt
         python -m pip install tqdm brainpylib
+        pip uninstall brainpy -y
         python setup.py install
     - name: Lint with flake8
       run: |

LICENSE

@@ -1,17 +1,17 @@
 <p align="center">
-  	<img alt="Header image of BrainPy - brain dynamics programming in Python." src="https://github.com/PKU-NIP-Lab/BrainPy/blob/master/images/logo.png" width=80%>
+  	<img alt="Header image of BrainPy - brain dynamics programming in Python." src="https://github.com/brainpy/BrainPy/blob/master/images/logo.png" width=80%>
 </p> 
 
 
 
 <p align="center">
 	<a href="https://pypi.org/project/brainpy/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainpy"></a>
-	<a href="https://github.com/PKU-NIP-Lab/BrainPy"><img alt="LICENSE" src="https://anaconda.org/brainpy/brainpy/badges/license.svg"></a>
+	<a href="https://github.com/brainpy/BrainPy"><img alt="LICENSE" src="https://anaconda.org/brainpy/brainpy/badges/license.svg"></a>
   	<a href="https://brainpy.readthedocs.io/en/latest/?badge=latest"><img alt="Documentation" src="https://readthedocs.org/projects/brainpy/badge/?version=latest"></a>
   	<a href="https://badge.fury.io/py/brainpy"><img alt="PyPI version" src="https://badge.fury.io/py/brainpy.svg"></a>
-    <a href="https://github.com/PKU-NIP-Lab/BrainPy"><img alt="Linux CI" src="https://github.com/PKU-NIP-Lab/BrainPy/actions/workflows/Linux_CI.yml/badge.svg"></a>
-    <a href="https://github.com/PKU-NIP-Lab/BrainPy"><img alt="Windows CI" src="https://github.com/PKU-NIP-Lab/BrainPy/actions/workflows/Windows_CI.yml/badge.svg"></a>
-    <a href="https://github.com/PKU-NIP-Lab/BrainPy"><img alt="MacOS CI" src="https://github.com/PKU-NIP-Lab/BrainPy/actions/workflows/MacOS_CI.yml/badge.svg"></a>
+    <a href="https://github.com/brainpy/BrainPy"><img alt="Linux CI" src="https://github.com/brainpy/BrainPy/actions/workflows/Linux_CI.yml/badge.svg"></a>
+    <a href="https://github.com/brainpy/BrainPy"><img alt="Windows CI" src="https://github.com/brainpy/BrainPy/actions/workflows/Windows_CI.yml/badge.svg"></a>
+    <a href="https://github.com/brainpy/BrainPy"><img alt="MacOS CI" src="https://github.com/brainpy/BrainPy/actions/workflows/MacOS_CI.yml/badge.svg"></a>
 </p>
 
 
@@ -20,17 +20,17 @@
 BrainPy is a flexible, efficient, and extensible framework for computational neuroscience and brain-inspired computation based on the Just-In-Time (JIT) compilation (built on top of [JAX](https://github.com/google/jax), [Numba](https://github.com/numba/numba), and other JIT compilers). It provides an integrative ecosystem for brain dynamics programming, including brain dynamics **building**, **simulation**, **training**, **analysis**, etc. 
 
 - **Website (documentation and APIs)**: https://brainpy.readthedocs.io/en/latest
-- **Source**: https://github.com/PKU-NIP-Lab/BrainPy
-- **Bug reports**: https://github.com/PKU-NIP-Lab/BrainPy/issues
+- **Source**: https://github.com/brainpy/BrainPy
+- **Bug reports**: https://github.com/brainpy/BrainPy/issues
 - **Source on OpenI**: https://git.openi.org.cn/OpenI/BrainPy
 
 
 
 ## Ecosystem
 
-- **[BrainPy](https://github.com/PKU-NIP-Lab/BrainPy)**: The solution for the general-purpose brain dynamics programming. 
-- **[brainpylib](https://github.com/PKU-NIP-Lab/brainpylib)**: Efficient operators for the sparse and event-driven computation. 
-- **[BrainPyExamples](https://github.com/PKU-NIP-Lab/BrainPyExamples)**: Comprehensive examples of BrainPy computation. 
+- **[BrainPy](https://github.com/brainpy/BrainPy)**: The solution for the general-purpose brain dynamics programming. 
+- **[brainpylib](https://github.com/brainpy/brainpylib)**: Efficient operators for the sparse and event-driven computation. 
+- **[BrainPyExamples](https://github.com/brainpy/BrainPyExamples)**: Comprehensive examples of BrainPy computation. 
 - **[brainpy-largescale](https://github.com/NH-NCL/brainpy-largescale)**: One solution for the large-scale brain modeling. 
 
 
@@ -47,12 +47,7 @@ For detailed installation instructions, please refer to the documentation: [Quic
 
 
 
-## License
-
-[Apache License, Version 2.0](https://github.com/PKU-NIP-Lab/BrainPy/blob/master/LICENSE)
-
-
 
 ## Citing
 
-If you are using BrainPy, please consider citing [the corresponding papers](https://brainpy.readthedocs.io/en/latest/tutorial_FAQs/citing_and_publication.html). 
+If you are using ``brainpy``, please consider citing [the corresponding papers](https://brainpy.readthedocs.io/en/latest/tutorial_FAQs/citing_and_publication.html). 

README.md

@@ -1,14 +1,14 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "2.2.4.0"
+__version__ = "2.3.0"
 
 
 # fundamental modules
 from . import errors, tools, check, modes
 
 # "base" module
 from . import base
-from .base.base import Base
+from .base.base import BrainPyObject, Base
 from .base.collector import Collector, TensorCollector
 
 # math foundation
@@ -72,7 +72,6 @@
   OfflineTrainer, RidgeTrainer,
   BPFF,
   BPTT,
-  OnlineBPTT,
 )
 
 # automatic dynamics analysis

brainpy/__init__.py

@@ -6,7 +6,7 @@
 from jax.lax import while_loop
 
 import brainpy.math as bm
-from brainpy.base import Base
+from brainpy.base import BrainPyObject
 from brainpy.types import Array
 from .utils import (Sigmoid,
                     Regularization, L1Regularization, L1L2Regularization, L2Regularization,
@@ -33,7 +33,7 @@
 name2func = dict()
 
 
-class OfflineAlgorithm(Base):
+class OfflineAlgorithm(BrainPyObject):
   """Base class for offline training algorithm."""
 
   def __init__(self, name=None):
@@ -46,16 +46,16 @@ def __call__(self, identifier, target, input, output):
     ----------
     identifier: str
       The variable name.
-    target: JaxArray, ndarray
+    target: Array, ndarray
       The 2d target data with the shape of `(num_batch, num_output)`.
-    input: JaxArray, ndarray
+    input: Array, ndarray
       The 2d input data with the shape of `(num_batch, num_input)`.
-    output: JaxArray, ndarray
+    output: Array, ndarray
       The 2d output data with the shape of `(num_batch, num_output)`.
 
     Returns
     -------
-    weight: JaxArray
+    weight: Array
       The weights after fit.
     """
     return self.call(identifier, target, input, output)
@@ -68,21 +68,21 @@ def call(self, identifier, targets, inputs, outputs) -> Array:
     identifier: str
       The identifier.
 
-    inputs: JaxArray, jax.numpy.ndarray, numpy.ndarray
+    inputs: Array, jax.numpy.ndarray, numpy.ndarray
       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: JaxArray, jax.numpy.ndarray, numpy.ndarray
+    targets: Array, jax.numpy.ndarray, numpy.ndarray
       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: JaxArray, jax.numpy.ndarray, numpy.ndarray
+    outputs: Array, jax.numpy.ndarray, numpy.ndarray
       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: JaxArray
+    weight: Array
       The weights after fit.
     """
     raise NotImplementedError('Must implement the __call__ function by the subclass itself.')

brainpy/algorithms/offline.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import brainpy.math as bm
-from brainpy.base import Base
+from brainpy.base import BrainPyObject
 from jax import vmap
 import jax.numpy as jnp
 
@@ -21,7 +21,7 @@
 name2func = dict()
 
 
-class OnlineAlgorithm(Base):
+class OnlineAlgorithm(BrainPyObject):
   """Base class for online training algorithm."""
 
   def __init__(self, name=None):
@@ -34,16 +34,16 @@ def __call__(self, identifier, target, input, output):
     ----------
     identifier: str
       The variable name.
-    target: JaxArray, ndarray
+    target: Array, ndarray
       The 2d target data with the shape of `(num_batch, num_output)`.
-    input: JaxArray, ndarray
+    input: Array, ndarray
       The 2d input data with the shape of `(num_batch, num_input)`.
-    output: JaxArray, ndarray
+    output: Array, ndarray
       The 2d output data with the shape of `(num_batch, num_output)`.
 
     Returns
     -------
-    weight: JaxArray
+    weight: Array
       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: JaxArray, ndarray
+    target: Array, ndarray
       The 2d target data with the shape of `(num_batch, num_output)`.
-    input: JaxArray, ndarray
+    input: Array, ndarray
       The 2d input data with the shape of `(num_batch, num_input)`.
-    output: JaxArray, ndarray
+    output: Array, ndarray
       The 2d output data with the shape of `(num_batch, num_output)`.
 
     Returns
     -------
-    weight: JaxArray
+    weight: Array
       The weights after fit.
     """
     raise NotImplementedError('Must implement the call() function by the subclass itself.')

brainpy/algorithms/online.py

@@ -335,14 +335,14 @@ 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 JaxArray or dict of JaxArray.')
-    fixed_points = tree_map(lambda a: bm.TrainVar(a), candidates, is_leaf=lambda x: isinstance(x, bm.JaxArray))
+      raise ValueError('Candidates must be instance of Array or dict of Array.')
+    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()
 
     def f_loss():
       return f_eval_loss(tree_map(lambda a: bm.as_device_array(a),
                                   fixed_points,
-                                  is_leaf=lambda x: isinstance(x, bm.JaxArray))).mean()
+                                  is_leaf=lambda x: isinstance(x, bm.Array))).mean()
 
     grad_f = bm.grad(f_loss, grad_vars=fixed_points, return_value=True)
     optimizer.register_vars(fixed_points if isinstance(fixed_points, dict) else {'a': fixed_points})
@@ -355,7 +355,7 @@ def train(idx):
       return loss
 
     def batch_train(start_i, n_batch):
-      return bm.for_loop(train, dyn_vars, bm.arange(start_i, start_i + n_batch))
+      return bm.for_loop(train, bm.arange(start_i, start_i + n_batch), dyn_vars=dyn_vars)
 
     # Run the optimization
     if self.verbose:
@@ -387,10 +387,10 @@ def batch_train(start_i, n_batch):
     self._opt_losses = bm.concatenate(opt_losses)
     self._losses = f_eval_loss(tree_map(lambda a: bm.as_device_array(a),
                                         fixed_points,
-                                        is_leaf=lambda x: isinstance(x, bm.JaxArray)))
+                                        is_leaf=lambda x: isinstance(x, bm.Array)))
     self._fixed_points = tree_map(lambda a: bm.as_device_array(a),
                                   fixed_points,
-                                  is_leaf=lambda x: isinstance(x, bm.JaxArray))
+                                  is_leaf=lambda x: isinstance(x, bm.Array))
     self._selected_ids = jnp.arange(num_candidate)
 
     if isinstance(self.target, DynamicalSystem):
@@ -428,7 +428,7 @@ def find_fps_with_opt_solver(
     # optimizing
     res = f_opt(tree_map(lambda a: bm.as_device_array(a),
                          candidates,
-                         is_leaf=lambda a: isinstance(a, bm.JaxArray)))
+                         is_leaf=lambda a: isinstance(a, bm.Array)))
 
     # results
     valid_ids = jnp.where(res.success)[0]
@@ -546,7 +546,7 @@ def compute_jacobians(
 
     Parameters
     ----------
-    points: np.ndarray, bm.JaxArray, jax.ndarray
+    points: np.ndarray, bm.Array, jax.ndarray
       The fixed points with the shape of (num_point, num_dim).
     stack_dict_var: bool
       Stack dictionary variables to calculate Jacobian matrix?
@@ -561,7 +561,7 @@ def compute_jacobians(
     """
     # check data
     info = np.asarray([(l.ndim, l.shape[0])
-                       for l in tree_flatten(points, is_leaf=lambda a: isinstance(a, bm.JaxArray))[0]])
+                       for l in tree_flatten(points, is_leaf=lambda a: isinstance(a, bm.Array))[0]])
     ndim = np.unique(info[:, 0])
     if len(ndim) != 1: raise ValueError(f'Get multiple dimension of the evaluated points. {ndim}')
     if ndim[0] == 1:
@@ -606,7 +606,7 @@ def decompose_eigenvalues(matrices, sort_by='magnitude', do_compute_lefts=False)
 
     Parameters
     ----------
-    matrices: np.ndarray, bm.JaxArray, jax.ndarray
+    matrices: np.ndarray, bm.Array, jax.ndarray
       A 3D array with the shape of (num_matrices, dim, dim).
     sort_by: str
       The method of sorting.