Merge branch 'brainpy:master' into master

.github/workflows/CI-models.yml

@@ -27,12 +27,11 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          pip install taichi-nightly -i https://pypi.taichi.graphics/simple/
           if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
           pip uninstall brainpy -y
           python setup.py install
@@ -51,7 +50,7 @@ jobs:
 #    steps:
 #    - uses: actions/checkout@v4
 #    - name: Set up Python ${{ matrix.python-version }}
-#      uses: actions/setup-python@v4
+#      uses: actions/setup-python@v5
 #      with:
 #        python-version: ${{ matrix.python-version }}
 #    - name: Install dependencies
@@ -75,12 +74,11 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install taichi-nightly -i https://pypi.taichi.graphics/simple/
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
         pip uninstall brainpy -y
         python setup.py install
@@ -99,7 +97,7 @@ jobs:
 #    steps:
 #      - uses: actions/checkout@v4
 #      - name: Set up Python ${{ matrix.python-version }}
-#        uses: actions/setup-python@v4
+#        uses: actions/setup-python@v5
 #        with:
 #          python-version: ${{ matrix.python-version }}
 #      - name: Install dependencies
@@ -124,13 +122,12 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
         python -m pip install numpy>=1.21.0
-        pip install taichi-nightly -i https://pypi.taichi.graphics/simple/
         python -m pip install -r requirements-dev.txt
         python -m pip install tqdm brainpylib
         pip uninstall brainpy -y
@@ -150,7 +147,7 @@ jobs:
 #    steps:
 #    - uses: actions/checkout@v4
 #    - name: Set up Python ${{ matrix.python-version }}
-#      uses: actions/setup-python@v4
+#      uses: actions/setup-python@v5
 #      with:
 #        python-version: ${{ matrix.python-version }}
 #    - name: Install dependencies

.github/workflows/CI.yml

@@ -29,14 +29,13 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           python -m pip install flake8 pytest
-          pip install taichi-nightly -i https://pypi.taichi.graphics/simple/
           if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
           pip uninstall brainpy -y
           python setup.py install
@@ -62,7 +61,7 @@ jobs:
 #    steps:
 #    - uses: actions/checkout@v4
 #    - name: Set up Python ${{ matrix.python-version }}
-#      uses: actions/setup-python@v4
+#      uses: actions/setup-python@v5
 #      with:
 #        python-version: ${{ matrix.python-version }}
 #    - name: Install dependencies
@@ -96,14 +95,13 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest
-        pip install taichi-nightly -i https://pypi.taichi.graphics/simple/
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
         pip uninstall brainpy -y
         python setup.py install
@@ -128,7 +126,7 @@ jobs:
 #    steps:
 #      - uses: actions/checkout@v4
 #      - name: Set up Python ${{ matrix.python-version }}
-#        uses: actions/setup-python@v4
+#        uses: actions/setup-python@v5
 #        with:
 #          python-version: ${{ matrix.python-version }}
 #      - name: Install dependencies
@@ -158,23 +156,19 @@ jobs:
 #    strategy:
 #      fail-fast: false
 #      matrix:
-#        python-version: ["3.8", "3.9", "3.10", "3.11"]
+#        python-version: ["3.9", "3.10", "3.11"]
 #
 #    steps:
 #    - uses: actions/checkout@v4
 #    - name: Set up Python ${{ matrix.python-version }}
-#      uses: actions/setup-python@v4
+#      uses: actions/setup-python@v5
 #      with:
 #        python-version: ${{ matrix.python-version }}
 #    - name: Install dependencies
 #      run: |
 #        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.4.11" -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver
-#        python -m pip install jax==0.4.11
 #        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
@@ -199,7 +193,7 @@ jobs:
 #    steps:
 #    - uses: actions/checkout@v4
 #    - name: Set up Python ${{ matrix.python-version }}
-#      uses: actions/setup-python@v4
+#      uses: actions/setup-python@v5
 #      with:
 #        python-version: ${{ matrix.python-version }}
 #    - name: Install dependencies

.github/workflows/Publish.yml

@@ -10,7 +10,7 @@ jobs:
         uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - run: python setup.py bdist_wheel
+      - run: python setup.py bdist_wheel --python-tag=py3
       - name: Publish package
         uses: pypa/gh-action-pypi-publish@release/v1
         with:

.github/workflows/docs.yml

@@ -18,7 +18,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
           python-version: "3.10"

README.md

@@ -14,7 +14,7 @@
 </p>
 
 
-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. 
+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), [Taichi](https://github.com/taichi-dev/taichi), [Numba](https://github.com/numba/numba), and others). 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/brainpy/BrainPy
@@ -77,7 +77,9 @@ We provide a Binder environment for BrainPy. You can use the following button to
 - **[BrainPy](https://github.com/brainpy/BrainPy)**: The solution for the general-purpose brain dynamics programming. 
 - **[brainpy-examples](https://github.com/brainpy/examples)**: Comprehensive examples of BrainPy computation. 
 - **[brainpy-datasets](https://github.com/brainpy/datasets)**: Neuromorphic and Cognitive Datasets for Brain Dynamics Modeling.
+- [《神经计算建模实战》 (Neural Modeling in Action)](https://github.com/c-xy17/NeuralModeling)
 - [第一届神经计算建模与编程培训班 (First Training Course on Neural Modeling and Programming)](https://github.com/brainpy/1st-neural-modeling-and-programming-course)
+- [第二届神经计算建模与编程培训班 (Second Training Course on Neural Modeling and Programming)](https://github.com/brainpy/2nd-neural-modeling-and-programming-course)
 
 
 ## Citing 
@@ -102,4 +104,6 @@ We also welcome your contributions
 - [ ] pipeline parallelization on multiple devices for sparse spiking network models
 - [ ] multi-compartment modeling
 - [ ] measurements, analysis, and visualization methods for large-scale spiking data
+- [ ] Online learning methods for large-scale spiking network models
+- [ ] Classical plasticity rules for large-scale spiking network models
 

brainpy/__init__.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "2.4.6"
+
+__version__ = "2.5.0"
 
 # fundamental supporting modules
 from brainpy import errors, check, tools
@@ -75,9 +76,10 @@
 )
 NeuGroup = NeuGroupNS = dyn.NeuDyn
 
-# shared parameters
+# common tools
 from brainpy._src.context import (share as share)
-from brainpy._src.helpers import (reset_state as reset_state,
+from brainpy._src.helpers import (reset_level as reset_level,
+                                  reset_state as reset_state,
                                   save_state as save_state,
                                   load_state as load_state,
                                   clear_input as clear_input)

brainpy/_add_deprecations.py

@@ -88,6 +88,16 @@
   # neurons
   'NeuGroup': ('brainpy.dyn.NeuGroup', 'brainpy.dyn.NeuDyn', NeuDyn),
 
+  # projections
+  'ProjAlignPostMg1': ('brainpy.dyn.ProjAlignPostMg1', 'brainpy.dyn.HalfProjAlignPostMg', dyn.HalfProjAlignPostMg),
+  'ProjAlignPostMg2': ('brainpy.dyn.ProjAlignPostMg2', 'brainpy.dyn.FullProjAlignPostMg', dyn.FullProjAlignPostMg),
+  'ProjAlignPost1': ('brainpy.dyn.ProjAlignPost1', 'brainpy.dyn.HalfProjAlignPost', dyn.HalfProjAlignPost),
+  'ProjAlignPost2': ('brainpy.dyn.ProjAlignPost2', 'brainpy.dyn.FullProjAlignPost', dyn.FullProjAlignPost),
+  'ProjAlignPreMg1': ('brainpy.dyn.ProjAlignPreMg1', 'brainpy.dyn.FullProjAlignPreSDMg', dyn.FullProjAlignPreSDMg),
+  'ProjAlignPreMg2': ('brainpy.dyn.ProjAlignPreMg2', 'brainpy.dyn.FullProjAlignPreDSMg', dyn.FullProjAlignPreDSMg),
+  'ProjAlignPre1': ('brainpy.dyn.ProjAlignPre1', 'brainpy.dyn.FullProjAlignPreSD', dyn.FullProjAlignPreSD),
+  'ProjAlignPre2': ('brainpy.dyn.ProjAlignPre2', 'brainpy.dyn.FullProjAlignPreDS', dyn.FullProjAlignPreDS),
+
   # synapses
   'TwoEndConn': ('brainpy.dyn.TwoEndConn', 'brainpy.synapses.TwoEndConn', synapses.TwoEndConn),
   'SynSTP': ('brainpy.dyn.SynSTP', 'brainpy.synapses.SynSTP', synapses.SynSTP),

brainpy/_src/checkpoints/serialization.py

@@ -19,21 +19,19 @@
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 import jax
-import msgpack
 import numpy as np
+from jax import monitoring
 from jax import process_index
 from jax.experimental.multihost_utils import sync_global_devices
-
 try:
-  from jax import monitoring
-except (ModuleNotFoundError, ImportError):
-  monitoring = None
+  from jax.experimental.array_serialization import get_tensorstore_spec, GlobalAsyncCheckpointManager  # noqa
+except:
+  get_tensorstore_spec = GlobalAsyncCheckpointManager = None
 
 try:
-  from jax.experimental.array_serialization import get_tensorstore_spec, GlobalAsyncCheckpointManager  # noqa
-except (ModuleNotFoundError, ImportError):
-  get_tensorstore_spec = None
-  GlobalAsyncCheckpointManager = None
+  import msgpack
+except ModuleNotFoundError:
+  msgpack = None
 
 from brainpy._src.math.ndarray import Array
 from brainpy.errors import (AlreadyExistsError,
@@ -116,6 +114,12 @@ def _record_path(name):
     _error_context.path.pop()
 
 
+def check_msgpack():
+  if msgpack is None:
+    raise ModuleNotFoundError('\nbrainpy.checkpoints needs "msgpack" package. Please install msgpack via:\n'
+                              '> pip install msgpack')
+
+
 def current_path():
   """Current state_dict path during deserialization for error messages."""
   return '/'.join(_error_context.path)
@@ -1126,6 +1130,7 @@ def save(
   out: str
     Filename of saved checkpoint.
   """
+  check_msgpack()
   start_time = time.time()
   # Make sure all saves are finished before the logic of checking and removing
   # outdated checkpoints happens.
@@ -1257,6 +1262,7 @@ def save_pytree(
   out: str
     Filename of saved checkpoint.
   """
+  check_msgpack()
   if verbose:
     print(f'Saving checkpoint into {filename}')
   start_time = time.time()
@@ -1344,6 +1350,7 @@ def multiprocess_save(
   out: str
     Filename of saved checkpoint.
   """
+  check_msgpack()
   start_time = time.time()
   # Make sure all saves are finished before the logic of checking and removing
   # outdated checkpoints happens.
@@ -1488,6 +1495,7 @@ def load(
     returned. This is to match the behavior of the case where a directory path
     is specified but the directory has not yet been created.
   """
+  check_msgpack()
   start_time = time.time()
 
   ckpt_dir = os.fspath(ckpt_dir)  # Pathlib -> str
@@ -1582,6 +1590,7 @@ def load_pytree(
     returned. This is to match the behavior of the case where a directory path
     is specified but the directory has not yet been created.
   """
+  check_msgpack()
   start_time = time.time()
   if not os.path.exists(filename):
     raise ValueError(f'Checkpoint not found: {filename}')

brainpy/_src/dependency_check.py

@@ -1,38 +1,41 @@
+import os
+import sys
 from jax.lib import xla_client
 
-
 __all__ = [
   'import_taichi',
   'import_brainpylib_cpu_ops',
   'import_brainpylib_gpu_ops',
 ]
 
-
-_minimal_brainpylib_version = '0.1.10'
+_minimal_brainpylib_version = '0.2.6'
 _minimal_taichi_version = (1, 7, 0)
 
 taichi = None
 brainpylib_cpu_ops = None
 brainpylib_gpu_ops = None
 
+taichi_install_info = (f'We need taichi=={_minimal_taichi_version}. '
+                       f'Currently you can install taichi=={_minimal_taichi_version} through:\n\n'
+                       '> pip install taichi==1.7.0')
+os.environ["TI_LOG_LEVEL"] = "error"
+
 
 def import_taichi():
   global taichi
   if taichi is None:
-    try:
-      import taichi as taichi  # noqa
-    except ModuleNotFoundError:
-      raise ModuleNotFoundError(
-        'Taichi is needed. Please install taichi through:\n\n'
-        '> pip install -i https://pypi.taichi.graphics/simple/ taichi-nightly'
-      )
-
-  if taichi.__version__ < _minimal_taichi_version:
-    raise RuntimeError(
-      f'We need taichi>={_minimal_taichi_version}. '
-      f'Currently you can install taichi>={_minimal_taichi_version} through taichi-nightly:\n\n'
-      '> pip install -i https://pypi.taichi.graphics/simple/ taichi-nightly'
-    )
+    with open(os.devnull, 'w') as devnull:
+      old_stdout = sys.stdout
+      sys.stdout = devnull
+      try:
+        import taichi as taichi  # noqa
+      except ModuleNotFoundError:
+        raise ModuleNotFoundError(taichi_install_info)
+      finally:
+        sys.stdout = old_stdout
+
+  if taichi.__version__ != _minimal_taichi_version:
+    raise RuntimeError(taichi_install_info)
   return taichi
 
 
@@ -82,6 +85,3 @@ def import_brainpylib_gpu_ops():
                         'See https://brainpy.readthedocs.io for installation instructions.')
 
   return brainpylib_gpu_ops
-
-
-