Merge pull request #38 from PKU-NIP-Lab/develop

New supports
brainpy · Mar 27, 2021 · 94b4dbe · 94b4dbe
2 parents 71af5c5 + 4f648d3
commit 94b4dbe
Show file tree

Hide file tree

Showing 48 changed files with 1,249 additions and 1,820 deletions.
diff --git a/README.md b/README.md
@@ -53,7 +53,7 @@ Install ``BrainPy`` from source:
 ``BrainPy`` is based on Python (>=3.7), and the following packages are required to be installed to use ``BrainPy``:
 
 - NumPy >= 1.13
-- Matplotlib >= 3.0
+- Matplotlib >= 3.2
 
 
 
@@ -63,12 +63,12 @@ Install ``BrainPy`` from source:
 <table border="0">
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/neurons/HH_model.html">
+            <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">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/neurons/HH_model.html">HH Neuron Model</a></h3>
+            <h3><a href="https://github.com/PKU-NIP-Lab/BrainModels/blob/main/brainmodels/tensor_backend/neurons/HodgkinHuxley_model.py">HH Neuron Model</a></h3>
             <p>The Hodgkin–Huxley model, or conductance-based model,
             is a mathematical model that describes how action potentials
             in neurons are initiated and propagated. It is a set of nonlinear
@@ -78,23 +78,23 @@ Install ``BrainPy`` from source:
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/synapses/Excitatory_models.html">
+            <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">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/synapses/Excitatory_models.html">AMPA Synapse Model</a></h3>
+            <h3><a href="https://github.com/PKU-NIP-Lab/BrainModels/blob/main/brainmodels/tensor_backend/synapses/AMPA_synapse.py">AMPA Synapse Model</a></h3>
             <p>AMPA synapse model.</p>
         </td>
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/examples/from_papers/Wang_1996_gamma_oscillation.html">
+            <a href="https://brainmodels.readthedocs.io/en/latest/examples/from_papers/Wang_1996_gamma_oscillation.html">
             <img src="docs/images/gamma_oscillation.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/examples/from_papers/Wang_1996_gamma_oscillation.html">Gamma Oscillation Model</a></h3>
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/examples/from_papers/Wang_1996_gamma_oscillation.html">Gamma Oscillation Model</a></h3>
             <p>Implementation of the paper: <i> Wang, Xiao-Jing, and György Buzsáki. “Gamma oscillation by
                   synaptic inhibition in a hippocampal interneuronal network
                   model.” Journal of neuroscience 16.20 (1996): 6402-6413. </i>
@@ -103,25 +103,25 @@ Install ``BrainPy`` from source:
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/examples/from_papers/Vreeswijk_1996_EI_net.html">
+            <a href="https://brainmodels.readthedocs.io/en/latest/examples/from_papers/Vreeswijk_1996_EI_net.html">
             <img src="docs/images/EI_balance_net.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/examples/from_papers/Vreeswijk_1996_EI_net.html">E/I Balance Network</a></h3>
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/examples/from_papers/Vreeswijk_1996_EI_net.html">E/I Balance Network</a></h3>
         <p>Implementation of the paper: <i>Van Vreeswijk, Carl, and Haim Sompolinsky. 
         “Chaos in neuronal networks with balanced excitatory and inhibitory activity.” 
         Science 274.5293 (1996): 1724-1726.</i></p>        
 </td>
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/examples/from_papers/Wu_2008_CANN.html">
+            <a href="https://brainmodels.readthedocs.io/en/latest/examples/from_papers/Wu_2008_CANN.html">
             <img src="docs/images/CANN1d.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/examples/from_papers/Wu_2008_CANN.html">Continuous-attractor Network</a></h3>
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/examples/from_papers/Wu_2008_CANN.html">Continuous-attractor Network</a></h3>
             <p>Implementation of the paper: <i> Si Wu, Kosuke Hamaguchi, and Shun-ichi Amari. "Dynamics and
                     computation of continuous attractors." Neural
                     computation 20.4 (2008): 994-1025. </i>
@@ -131,12 +131,11 @@ Install ``BrainPy`` from source:
 </table>
 
 
-
 More neuron examples please see [BrainPy-Models/neurons](https://github.com/PKU-NIP-Lab/BrainPy-Models/tree/main/bpmodels/neurons);
 
 More synapse examples please see [BrainPy-Models/synapses](https://github.com/PKU-NIP-Lab/BrainPy-Models/tree/main/bpmodels/synapses);
 
-More network examples please see [BrainPy-Models/from_papers](https://brainpy-models.readthedocs.io/en/latest/examples/from_papers.html).
+More network examples please see [BrainPy-Models/from_papers](https://brainmodels.readthedocs.io/en/latest/examples/from_papers.html).
 
 
 
@@ -148,37 +147,37 @@ More network examples please see [BrainPy-Models/from_papers](https://brainpy-mo
 <table border="0">
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html">
+            <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html">
             <img src="docs/images/phase_plane_analysis1.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html">Phase Plane Analysis</a></h3>
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html">Phase Plane Analysis</a></h3>
             <p>Phase plane analysis of the I<sub>Na,p+</sub>-I<sub>K</sub> model, where
             "input" is 50., and "Vn_half" is -45..</p>
         </td>
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-1-bifurcation-analysis">
+            <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-1-bifurcation-analysis">
             <img src="docs/images/NaK_model_codimension1.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-1-bifurcation-analysis">
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-1-bifurcation-analysis">
                 Codimension 1 Bifurcation Analysis (1)</a></h3>
             <p>Codimension 1 bifurcation analysis of the I<sub>Na,p+</sub>-I<sub>K</sub> model,
                 in which "input" is varied in [0., 50.].</p>
         </td>
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-2-bifurcation-analysis">
+            <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-2-bifurcation-analysis">
             <img src="docs/images/NaK_model_codimension2.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-2-bifurcation-analysis">
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/NaK_model_analysis.html#Codimension-2-bifurcation-analysis">
                 Codimension 2 Bifurcation Analysis (1)</a></h3>
             <p>Codimension 2 bifurcation analysis of a two-variable neuron model:
                 the I<sub>Na,p+</sub>-I<sub>K</sub> model, in which "input" is varied
@@ -187,25 +186,25 @@ More network examples please see [BrainPy-Models/from_papers](https://brainpy-mo
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html">
+            <a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html">
             <img src="docs/images/FitzHugh_Nagumo_codimension1.png">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html">
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html">
                 Codimension 1 Bifurcation Analysis (2)</a></h3>
             <p>Codimension 1 bifurcation analysis of FitzHugh Nagumo model, in which
                 "a" is equal to 0.7, and "Iext" is varied in [0., 1.].</p>
         </td>
     </tr>
     <tr>
         <td border="0" width="30%">
-            <a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html#Codimension-2-bifurcation-analysis">
+            <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">
             </a>
         </td>
         <td border="0" valign="top">
-            <h3><a href="https://brainpy-models.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html#Codimension-2-bifurcation-analysis">
+            <h3><a href="https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis/FitzHugh_Nagumo_analysis.html#Codimension-2-bifurcation-analysis">
                 Codimension 2 Bifurcation Analysis (2)</a></h3>
             <p>Codimension 2 bifurcation analysis of FitzHugh Nagumo model, in which "a"
                is varied in [0.5, 1.0], and "Iext" is varied in [0., 1.].</p>
@@ -215,7 +214,7 @@ More network examples please see [BrainPy-Models/from_papers](https://brainpy-mo
 
 
 
-More examples please see [BrainPy-Models/dynamics_analysis](https://github.com/PKU-NIP-Lab/BrainPy-Models/tree/main/brainpy-examples/dynamics_analysis).
+More examples please see [BrainPy-Models/dynamics_analysis](https://brainmodels.readthedocs.io/en/latest/tutorials/dynamics_analysis.html).
 
 
 

diff --git a/brainpy/__init__.py b/brainpy/__init__.py
@@ -2,20 +2,18 @@
 
 __version__ = "1.0.0-alpha"
 
-# "analysis" module
-from . import analysis
-
 # "backend" module
 from . import backend
 
-# "connectivity" module
-from . import connectivity
-from . import connectivity as connect
+# "analysis" module
+from . import analysis
 
 # "simulation" module
 from . import simulation
+from .simulation import connectivity as connect
 from .simulation.dynamic_system import *
 from .simulation.brain_objects import *
+from .simulation.utils import size2len
 
 # "integrators" module
 from . import integrators

diff --git a/brainpy/analysis/base.py b/brainpy/analysis/base.py
@@ -4,14 +4,18 @@
 from copy import deepcopy
 
 import numpy as np
-import sympy
 
 from brainpy import errors
 from brainpy import tools
 from brainpy.analysis import solver
 from brainpy.analysis import utils
 from brainpy.integrators import sympy_analysis
 
+try:
+    import sympy
+except ModuleNotFoundError:
+    sympy = None
+
 __all__ = [
     'BaseNeuronAnalyzer',
     'Base1DNeuronAnalyzer',
@@ -79,6 +83,9 @@ def __init__(self,
                  numerical_resolution=0.1,
                  options=None):
 
+        if sympy is None:
+            raise errors.PackageMissingError('"SymPy" must be installed for dynamics analysis.')
+
         # model
         # -----
         if isinstance(model_or_integrals, utils.DynamicModel):

diff --git a/brainpy/analysis/utils.py b/brainpy/analysis/utils.py
@@ -64,7 +64,7 @@ def transform_integrals_to_model(integrals):
             variables = []
             for v in vars:
                 if len(v) > 1:
-                    raise ValueError('Cannot analyze must assignment code line.')
+                    raise ValueError('Cannot analyze multi-assignment code line.')
                 variables.append(v[0])
             expressions = expressions_for_returns[key]
             var_name = integral.variables[vi]

diff --git a/brainpy/backend/__init__.py b/brainpy/backend/__init__.py
@@ -12,13 +12,16 @@
 _dt = 0.1
 
 CLASS_KEYWORDS = ['self', 'cls']
-NEEDED_OPS = ['normal', 'exp', 'matmul', 'sum',
-              'as_tensor', 'zeros', 'ones', 'arange',
-              'eye', 'vstack', 'reshape', 'shape', ]
+
+OPS_FOR_SOLVER = ['normal', 'sum', 'exp', 'matmul', 'shape', ]
+OPS_FOR_SIMULATION = ['as_tensor', 'zeros', 'ones', 'arange',
+                      'eye', 'vstack', 'reshape', 'where', ]
+
 SUPPORTED_BACKEND = {
     'numba', 'numba-parallel', 'numba-cuda', 'jax',  # JIT framework
     'numpy', 'pytorch', 'tensorflow',
 }
+
 SYSTEM_KEYWORDS = ['_dt', '_t', '_i']
 
 
@@ -89,13 +92,15 @@ def set(backend=None, module_or_operations=None, node_runner=None, net_runner=No
         set_numba_profile(parallel=True)
 
     elif backend == 'numba-cuda':
+        raise NotImplementedError
         from .operators import bk_numba_cuda
         from .runners.numba_cuda_runner import NumbaCudaNodeRunner
 
         node_runner = NumbaCudaNodeRunner if node_runner is None else node_runner
         module_or_operations = bk_numba_cuda if module_or_operations is None else module_or_operations
 
     elif backend == 'jax':
+        raise NotImplementedError
         from .operators import bk_jax
         from .runners.jax_runner import JaxRunner
 
@@ -166,19 +171,35 @@ def get_dt():
 
 
 def set_ops_from_module(module):
+    """Set operations from a module.
+
+    Parameters
+    ----------
+    module :
+    """
     global_vars = globals()
-    for ops in NEEDED_OPS:
+    for ops in OPS_FOR_SOLVER:
         if not hasattr(module, ops):
             raise ValueError(f'Operation "{ops}" is needed, but is not '
                              f'defined in module "{module}".')
         global_vars[ops] = getattr(module, ops)
+    for ops in OPS_FOR_SIMULATION:
+        if hasattr(module, ops):
+            global_vars[ops] = getattr(module, ops)
+        else:
+            del global_vars[ops]
 
 
 def set_ops(**kwargs):
+    """Set operations.
+
+    Parameters
+    ----------
+    kwargs :
+        The key=operation setting.
+    """
     global_vars = globals()
     for key, value in kwargs.items():
-        if key not in NEEDED_OPS:
-            print(f'"{key}" is not a necessary operation.')
         global_vars[key] = value
 
 

diff --git a/brainpy/backend/operators/bk_numpy.py b/brainpy/backend/operators/bk_numpy.py
@@ -13,8 +13,12 @@
     'ones',
     'eye',
     'matmul',
-    'vstack',
     'arange',
+
+    'vstack',
+    'where',
+    'unsqueeze',
+    'squeeze',
 ]
 
 
@@ -28,6 +32,9 @@
 ones = np.ones
 eye = np.eye
 matmul = np.matmul
-vstack = np.vstack
 arange = np.arange
 
+vstack = np.vstack
+where = np.where
+unsqueeze = np.expand_dims
+squeeze = np.squeeze
diff --git a/brainpy/backend/operators/bk_pytorch.py b/brainpy/backend/operators/bk_pytorch.py
@@ -22,7 +22,19 @@
 
 def shape(x):
     if isinstance(x, (int, float)):
-        return (1,)
+        return ()
     else:
         return x.size()
 
+
+def where(tensor, x, y):
+    if isinstance(x, (int, float)):
+        x = torch.full_like(tensor, x)
+    if isinstance(y, (int, float)):
+        y = torch.full_like(tensor, y)
+    return torch.where(tensor, x, y)
+
+
+unsqueeze = torch.unsqueeze
+squeeze = torch.squeeze
+
diff --git a/brainpy/backend/operators/bk_tensorflow.py b/brainpy/backend/operators/bk_tensorflow.py
@@ -1,12 +1,12 @@
 # -*- coding: utf-8 -*-
 
-
 """
 The TensorFlow with the version of xx is needed.
 """
 
 import tensorflow as tf
 
+
 reshape = tf.reshape
 exp = tf.math.exp
 sum = tf.math.reduce_sum
@@ -22,10 +22,19 @@ def vstack(values):
 
 def shape(x):
     if isinstance(x, (int, float)):
-        return (1,)
+        return ()
     else:
         return x.shape()
 
 
 def normal(loc, scale, size):
     return tf.random.normal(size, loc, scale)
+
+
+def where(tensor, x, y):
+    return tf.where(tensor, x, y)
+
+
+unsqueeze = tf.expand_dims
+squeeze = tf.squeeze
+