Merge pull request brainpy#43 from PKU-NIP-Lab/develop

Add detailed document for V1.0.0

README.md

@@ -1,5 +1,5 @@
 
-![Logo](docs/images/logo.png)
+![Logo](docs/_static/logo.png)
 
 [![LICENSE](https://anaconda.org/brainpy/brainpy/badges/license.svg)](https://github.com/PKU-NIP-Lab/BrainPy)    [![Documentation](https://readthedocs.org/projects/brainpy/badge/?version=latest)](https://brainpy.readthedocs.io/en/latest/?badge=latest)     [![Conda](https://anaconda.org/brainpy/brainpy-simulator/badges/version.svg)](https://anaconda.org/brainpy/brainpy-simulator)  [![PyPI version](https://badge.fury.io/py/brainpy-simulator.svg)](https://badge.fury.io/py/brainpy-simulator)
 
@@ -21,7 +21,7 @@ Moreover, `BrainPy` is designed to effectively satisfy your basic requirements:
 - *Efficient running speed*, because BrainPy is compatible with the latest JIT compilers or any other accelerating framework you prefer (below we list the speed comparison based on Numba JIT).
 
 
-![Speed Comparison](docs/images/speed.png)
+![Speed Comparison](docs/_static/speed.png)
 
 `BrainPy` is a backend-independent neural simulator. Users can define models with any backend they prefer. Intrinsically, BrainPy supports the array/tensor-oriented backends such like [NumPy](https://numpy.org/), [PyTorch](https://pytorch.org/), and [TensorFlow](https://www.tensorflow.org/), it also supports the JIT compilers such as [Numba](https://numba.pydata.org/) on CPU or CUDA devices. Extending BrainPy to support other backend frameworks you prefer is very easy. The details please see documents coming soon. 
 
@@ -71,7 +71,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://github.com/PKU-NIP-Lab/BrainModels/blob/main/brainmodels/tensor_backend/neurons/HodgkinHuxley_model.py">
-            <img src="docs/images/HH_neuron.png">
+            <img src="docs/_static/HH_neuron.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -82,7 +82,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://github.com/PKU-NIP-Lab/BrainModels/blob/main/brainmodels/tensor_backend/synapses/AMPA_synapse.py">
-            <img src="docs/images/AMPA_model.png">
+            <img src="docs/_static/AMPA_model.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -93,7 +93,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://brainmodels.readthedocs.io/en/latest/from_papers/Wang_1996_gamma_oscillation.html">
-            <img src="docs/images/gamma_oscillation.png">
+            <img src="docs/_static/gamma_oscillation.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -107,7 +107,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://brainmodels.readthedocs.io/en/latest/from_papers/Vreeswijk_1996_EI_net.html">
-            <img src="docs/images/EI_balance_net.png">
+            <img src="docs/_static/EI_balance_net.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -120,7 +120,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://brainmodels.readthedocs.io/en/latest/from_papers/Wu_2008_CANN.html">
-            <img src="docs/images/CANN1d.png">
+            <img src="docs/_static/CANN1d.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -134,7 +134,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html">
-            <img src="docs/images/phase_plane_analysis1.png">
+            <img src="docs/_static/phase_plane_analysis1.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -146,7 +146,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html">
-            <img src="docs/images/FitzHugh_Nagumo_codimension1.png">
+            <img src="docs/_static/FitzHugh_Nagumo_codimension1.png">
             </a>
         </td>
         <td border="0" valign="top">
@@ -159,7 +159,7 @@ Here list several simple examples for neurodynamics simulation and analysis. Com
     <tr>
         <td border="0" width="30%">
             <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html#Codimension-2-bifurcation-analysis">
-            <img src="docs/images/FitzHugh_Nagumo_codimension2.png">
+            <img src="docs/_static/FitzHugh_Nagumo_codimension2.png">
             </a>
         </td>
         <td border="0" valign="top">

brainpy/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "1.0.0rc2"
+__version__ = "1.0.0"
 
 # "backend" module
 from . import backend

brainpy/backend/drivers/numba_cuda.py

@@ -18,6 +18,7 @@
 from brainpy.simulation import drivers
 from brainpy.simulation.brainobjects import SynConn, NeuGroup
 from brainpy.simulation.delays import ConstantDelay
+from brainpy.simulation.monitors import Monitor
 from . import utils
 from .general import GeneralNetDriver
 from .numba_cpu import NumbaCPUNodeDriver
@@ -47,10 +48,11 @@
 ]
 
 _num_thread_gpu = 1024
+
 # Monitor can be done in :
 # 1. 'cpu'
 # 2. 'cuda'
-_monitor_done_in = 'cpu'
+_monitor_done_in = 'cuda'
 
 
 def set_monitor_done_in(place):
@@ -82,6 +84,20 @@ def get_cuda_size(num):
     return num_block, num_thread
 
 
+def get_categories(category):
+    if category is None:
+        category = [NeuGroup, SynConn, Monitor, ConstantDelay]
+    else:
+        str2target = {'mon': Monitor, 'neu': NeuGroup,
+                      'syn': SynConn, 'delay': ConstantDelay}
+        if isinstance(category, str):
+            category = [str2target[category]]
+        elif isinstance(category, (tuple, list)):
+            category = [str2target[c] for c in category]
+        else:
+            raise ValueError
+    return category
+
 def data_shape(data):
     if isinstance(data, DeviceNDArray):
         return data.shape
@@ -549,6 +565,7 @@ def _reprocess_steps(self, f, func_name=None, show_code=False):
                 code_scope[utils.attr_replace(data)] = obj
             else:
                 if callable(obj):
+                    code_scope[utils.attr_replace(data)] = obj
                     continue
                 if isinstance(obj, np.ndarray):  # 2. transform the cpu data to cuda data
                     splits[-1] = self.transfer_cpu_data_to_gpu(host, cpu_key=splits[-1], cpu_data=obj)
@@ -832,12 +849,16 @@ def reshape_mon_items(self, run_length):
             shape = ops.shape(data)
             if run_length < shape[0]:
                 data = data[:run_length]
+                setattr(self.host.mon, var, data)
+                if _monitor_done_in == 'cuda':
+                    setattr(self.host.mon, cuda_name_of(var), self.cpu2gpu(self.host, cpu_data=data))
+
             elif run_length > shape[0]:
                 append = ops.zeros((run_length - shape[0],) + shape[1:])
                 data = ops.vstack([data, append])
-            setattr(self.host.mon, var, data)
-            if _monitor_done_in == 'cuda':
-                setattr(self.host.mon, cuda_name_of(var), self.cpu2gpu(self.host, cpu_data=data))
+                setattr(self.host.mon, var, data)
+                if _monitor_done_in == 'cuda':
+                    setattr(self.host.mon, cuda_name_of(var), self.cpu2gpu(self.host, cpu_data=data))
 
     def get_monitor_func(self, mon_length, show_code=False):
         """Get monitor function.
@@ -979,38 +1000,49 @@ def get_monitor_func(self, mon_length, show_code=False):
                 raise ValueError(f'Monitor is set to an unknown place by "_monitor_done_in = '
                                  f'{_monitor_done_in}".')
 
-    def to_host(self):
+    def to_host(self, category=None):
+        categories = get_categories(category)
         for host, keys in self.host_cpukey_gpukey.items():
-            for cpukey, gpukey in keys.items():
-                setattr(host, cpukey, getattr(host, gpukey).copy_to_host(stream=host.stream))
+            if type(host) in categories:
+                for cpukey, gpukey in keys.items():
+                    setattr(host, cpukey, getattr(host, gpukey).copy_to_host(stream=host.stream))
 
-    def to_device(self):
+    def to_device(self, category=None):
+        categories = get_categories(category)
         for host, keys in self.host_cpukey_gpukey.items():
-            for cpukey, gpukey in keys.items():
-                setattr(host, gpukey, cuda.to_device(getattr(host, cpukey), stream=host.stream))
+            if type(host) in categories:
+                for cpukey, gpukey in keys.items():
+                    setattr(host, gpukey, cuda.to_device(getattr(host, cpukey),
+                                                         stream=host.stream))
 
 
 class NumbaCUDANetDriver(GeneralNetDriver):
-    def to_host(self):
+    def to_host(self, category=None):
+        categories = get_categories(category)
+
         host_cpukey_gpukey = {}
         for node in self.host.all_nodes.values():
             for host, keys in node.driver.host_cpukey_gpukey.items():
-                if host not in host_cpukey_gpukey:
-                    host_cpukey_gpukey[host] = {}
-                for cpukey, gpukey in keys.items():
-                    if cpukey in host_cpukey_gpukey[host]:
-                        continue
-                    host_cpukey_gpukey[host][cpukey] = gpukey
-                    setattr(host, cpukey, getattr(host, gpukey).copy_to_host(stream=host.stream))
+                if type(host) in categories:
+                    if host not in host_cpukey_gpukey:
+                        host_cpukey_gpukey[host] = {}
+                    for cpukey, gpukey in keys.items():
+                        if cpukey in host_cpukey_gpukey[host]:
+                            continue
+                        host_cpukey_gpukey[host][cpukey] = gpukey
+                        setattr(host, cpukey, getattr(host, gpukey).copy_to_host(stream=host.stream))
+
+    def to_device(self, category=None):
+        categories = get_categories(category)
 
-    def to_device(self):
         host_cpukey_gpukey = {}
         for node in self.host.all_nodes.values():
             for host, keys in node.driver.host_cpukey_gpukey.items():
-                if host not in host_cpukey_gpukey:
-                    host_cpukey_gpukey[host] = {}
-                for cpukey, gpukey in keys.items():
-                    if cpukey in host_cpukey_gpukey[host]:
-                        continue
-                    host_cpukey_gpukey[host][cpukey] = gpukey
-                    setattr(host, gpukey, cuda.to_device(getattr(host, cpukey), stream=host.stream))
+                if type(host) in categories:
+                    if host not in host_cpukey_gpukey:
+                        host_cpukey_gpukey[host] = {}
+                    for cpukey, gpukey in keys.items():
+                        if cpukey in host_cpukey_gpukey[host]:
+                            continue
+                        host_cpukey_gpukey[host][cpukey] = gpukey
+                        setattr(host, gpukey, cuda.to_device(getattr(host, cpukey), stream=host.stream))

brainpy/backend/ops/__init__.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import types
 
 from .numpy_ import *
 
@@ -16,9 +17,12 @@
 
 _backend = 'numpy'
 BUFFER = {}
-OPS_FOR_SOLVER = ['normal', 'sum', 'exp', 'matmul', 'shape', ]
+OPS_FOR_SOLVER = ['normal', 'sum', 'exp', 'shape', ]
 OPS_FOR_SIMULATION = ['as_tensor', 'zeros', 'ones', 'arange',
                       'vstack', 'where', 'unsqueeze', 'squeeze']
+OPS_OF_DTYPE = ['bool',
+                'int', 'int32', 'int64',
+                'float', 'float32', 'float64']
 
 
 def switch_to(backend):
@@ -54,7 +58,7 @@ def switch_to(backend):
         ops_in_buffer = get_buffer(backend)
         for ops in OPS_FOR_SOLVER:
             if ops not in ops_in_buffer:
-                raise ValueError(f'Operation "{ops}" is needed, but is not '
+                raise ValueError(f'Necessary operation "{ops}" is not '
                                  f'defined in "{backend}" backend\'s buffers.')
 
     # set operations from BUFFER
@@ -71,17 +75,30 @@ def set_ops_from_module(module):
     ----------
     module :
     """
+
+    ops_in_module = {p: getattr(module, p) for p in dir(module)
+                     if (not p.startswith('__')) and
+                     (not isinstance(getattr(module, p), types.ModuleType))}
     global_vars = globals()
+
     for ops in OPS_FOR_SOLVER:
-        if not hasattr(module, ops):
+        if ops not in ops_in_module:
             raise ValueError(f'Operation "{ops}" is needed, but is not '
                              f'defined in module "{module}".')
-        global_vars[ops] = getattr(module, ops)
+        global_vars[ops] = ops_in_module.pop(ops)
     for ops in OPS_FOR_SIMULATION:
-        if hasattr(module, ops):
-            global_vars[ops] = getattr(module, ops)
+        if ops in ops_in_module:
+            global_vars[ops] = ops_in_module.pop(ops)
         else:
             del global_vars[ops]
+    for ops in OPS_OF_DTYPE:
+        if ops in ops_in_module:
+            global_vars[ops] = ops_in_module.pop(ops)
+        else:
+            del global_vars[ops]
+
+    for ops, val in ops_in_module.items():
+        global_vars[ops] = val
 
 
 def set_ops(**kwargs):

brainpy/backend/ops/jax_.py

@@ -26,7 +26,6 @@ def normal(loc, scale, size):
 
 exp = numpy.exp
 sum = numpy.sum
-matmul = numpy.matmul
 shape = numpy.shape
 
 
@@ -39,3 +38,13 @@ def normal(loc, scale, size):
 where = numpy.where
 unsqueeze = numpy.expand_dims
 squeeze = numpy.squeeze
+
+# necessary ops for dtypes
+
+bool = numpy.bool_
+int = numpy.int_
+int32 = numpy.int32
+int64 = numpy.int64
+float = numpy.float_
+float32 = numpy.float32
+float64 = numpy.float64

brainpy/backend/ops/numba_cpu.py

@@ -10,7 +10,6 @@
 sum = np.sum
 shape = np.shape
 exp = np.exp
-matmul = np.matmul
 
 # necessary ops for dynamics simulation
 as_tensor = np.asarray
@@ -23,5 +22,16 @@
 squeeze = np.squeeze
 
 
+# necessary ops for dtypes
+
+bool = np.bool_
+int = np.int_
+int32 = np.int32
+int64 = np.int64
+float = np.float_
+float32 = np.float32
+float64 = np.float64
+
+
 if __name__ == '__main__':
     numba_overload

brainpy/backend/ops/numba_cuda.py

@@ -12,7 +12,6 @@
 sum = np.sum
 shape = np.shape
 exp = math.exp
-matmul = np.matmul
 
 # necessary ops for dynamics simulation
 as_tensor = np.asarray
@@ -25,5 +24,16 @@
 squeeze = np.squeeze
 
 
+# necessary ops for dtypes
+
+bool = np.bool_
+int = np.int_
+int32 = np.int32
+int64 = np.int64
+float = np.float_
+float32 = np.float32
+float64 = np.float64
+
+
 if __name__ == '__main__':
     numba_overload