Merge branch 'brainpy:master' into master

Author: AkitsuFaye

.github/workflows/docker.yml

@@ -35,13 +35,13 @@ jobs:
 
     - name: Login to DockerHub
       if: github.event_name != 'pull_request'
-      uses: docker/login-action@v2
+      uses: docker/login-action@v3
       with:
         username: ${{ secrets.DOCKERHUB_USERNAME }}
         password: ${{ secrets.DOCKERHUB_TOKEN }}
 
     - name: Docker Build & Push (version tag)
-      uses: docker/build-push-action@v4
+      uses: docker/build-push-action@v5
       with:
         context: ${{ matrix.context }}
         tags: ${{ matrix.base }}:${{ env.DOCKER_TAG_NAME }}
@@ -51,7 +51,7 @@ jobs:
     - name: Docker Build & Push (latest tag)
       if: |
           (github.event_name == 'release' && ! github.event.release.prerelease)
-      uses: docker/build-push-action@v4
+      uses: docker/build-push-action@v5
       with:
         context: ${{ matrix.context }}
         tags: ${{ matrix.base }}:latest

README.md

@@ -28,7 +28,25 @@ BrainPy is a flexible, efficient, and extensible framework for computational neu
 BrainPy is based on Python (>=3.8) and can be installed on Linux (Ubuntu 16.04 or later), macOS (10.12 or later), and Windows platforms. Install the latest version of BrainPy:
 
 ```bash
-$ pip install brainpy brainpylib -U
+$ pip install brainpy -U
+```
+
+In addition, many customized operators in BrainPy are implemented in ``brainpylib``.
+Install the latest version of `brainpylib` by:
+
+```bash
+# CPU installation for Linux, macOS and Windows
+$ pip install --upgrade brainpylib
+```
+
+```bash
+# CUDA 12 installation for Linux only
+$ pip install --upgrade brainpylib-cu12x
+```
+
+```bash
+# CUDA 11 installation for Linux only
+$ pip install --upgrade brainpylib-cu11x
 ```
 
 For detailed installation instructions, please refer to the documentation: [Quickstart/Installation](https://brainpy.readthedocs.io/en/latest/quickstart/installation.html)

brainpy/__init__.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 __version__ = "2.4.4.post4"
+_minimal_brainpylib_version = '0.1.10'
 
 # fundamental supporting modules
 from brainpy import errors, check, tools
@@ -11,6 +12,15 @@
 except ModuleNotFoundError:
   raise ModuleNotFoundError(tools.jaxlib_install_info) from None
 
+
+try:
+  import brainpylib
+  if brainpylib.__version__ < _minimal_brainpylib_version:
+    raise SystemError(f'This version of brainpy ({__version__}) needs brainpylib >= {_minimal_brainpylib_version}.')
+  del brainpylib
+except ModuleNotFoundError:
+  pass
+
 #  Part: Math Foundation  #
 # ----------------------- #
 

brainpy/_src/analysis/utils/model.py

@@ -135,11 +135,11 @@ def update(self):
       self.implicit_vars[key].update(intg(*all_vars, *self.pars, dt=share['dt']))
 
   def __getattr__(self, item):
-    child_vars = super(TrajectModel, self).__getattribute__('implicit_vars')
+    child_vars = super().__getattribute__('implicit_vars')
     if item in child_vars:
       return child_vars[item]
     else:
-      return super(TrajectModel, self).__getattribute__(item)
+      return super().__getattribute__(item)
 
   def run(self, duration):
     self.runner.run(duration)

brainpy/_src/delay.py

@@ -375,10 +375,6 @@ def reset_state(self, batch_size: int = None):
     # initialize delay data
     if self.data is not None:
       self._init_data(self.max_length, batch_size)
-    for cls in self.before_updates.values():
-      cls.reset_state(batch_size)
-    for cls in self.after_updates.values():
-      cls.reset_state(batch_size)
 
   def _init_data(self, length: int, batch_size: int = None):
     if batch_size is not None:
@@ -468,13 +464,16 @@ def __init__(
       *indices
   ):
     super().__init__(mode=delay.mode)
-    self.delay = delay
+    self.refs = {'delay': delay}
     assert isinstance(delay, Delay)
     delay.register_entry(self.name, time)
     self.indices = indices
 
   def update(self):
-    return self.delay.at(self.name, *self.indices)
+    return self.refs['delay'].at(self.name, *self.indices)
+
+  def reset_state(self, *args, **kwargs):
+    pass
 
 
 def init_delay_by_return(info: Union[bm.Variable, ReturnInfo]) -> Delay:

brainpy/_src/dnn/linear.py

@@ -44,7 +44,7 @@ class Dense(Layer, SupportOnline, SupportOffline, SupportSTDP):
     The number of the input feature. A positive integer.
   num_out: int
     The number of the output features. A positive integer.
-  weight_initializer: optional, Initializer
+  W_initializer: optional, Initializer
     The weight initialization.
   b_initializer: optional, Initializer
     The bias initialization.

brainpy/_src/dyn/projections/__init__.py

@@ -2,3 +2,4 @@
 from .aligns import *
 from .conn import *
 from .others import *
+from .inputs import *

brainpy/_src/dyn/projections/inputs.py

@@ -0,0 +1,96 @@
+from typing import Optional, Any
+
+from brainpy import math as bm
+from brainpy._src.dynsys import Dynamic
+from brainpy._src.mixin import SupportAutoDelay
+from brainpy.types import Shape
+
+__all__ = [
+    'InputVar',
+]
+
+
+class InputVar(Dynamic, SupportAutoDelay):
+    """Define an input variable.
+
+    Example::
+
+        import brainpy as bp
+
+
+        class Exponential(bp.Projection):
+            def __init__(self, pre, post, prob, g_max, tau, E=0.):
+                super().__init__()
+                self.proj = bp.dyn.ProjAlignPostMg2(
+                    pre=pre,
+                    delay=None,
+                    comm=bp.dnn.EventCSRLinear(bp.conn.FixedProb(prob, pre=pre.num, post=post.num), g_max),
+                    syn=bp.dyn.Expon.desc(post.num, tau=tau),
+                    out=bp.dyn.COBA.desc(E=E),
+                    post=post,
+                )
+
+
+        class EINet(bp.DynSysGroup):
+            def __init__(self, num_exc, num_inh, method='exp_auto'):
+                super(EINet, self).__init__()
+
+                # neurons
+                pars = dict(V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
+                            V_initializer=bp.init.Normal(-55., 2.), method=method)
+                self.E = bp.dyn.LifRef(num_exc, **pars)
+                self.I = bp.dyn.LifRef(num_inh, **pars)
+
+                # synapses
+                w_e = 0.6  # excitatory synaptic weight
+                w_i = 6.7  # inhibitory synaptic weight
+
+                # Neurons connect to each other randomly with a connection probability of 2%
+                self.E2E = Exponential(self.E, self.E, 0.02, g_max=w_e, tau=5., E=0.)
+                self.E2I = Exponential(self.E, self.I, 0.02, g_max=w_e, tau=5., E=0.)
+                self.I2E = Exponential(self.I, self.E, 0.02, g_max=w_i, tau=10., E=-80.)
+                self.I2I = Exponential(self.I, self.I, 0.02, g_max=w_i, tau=10., E=-80.)
+
+                # define input variables given to E/I populations
+                self.Ein = bp.dyn.InputVar(self.E.varshape)
+                self.Iin = bp.dyn.InputVar(self.I.varshape)
+                self.E.add_inp_fun('', self.Ein)
+                self.I.add_inp_fun('', self.Iin)
+
+
+        net = EINet(3200, 800, method='exp_auto')  # "method": the numerical integrator method
+        runner = bp.DSRunner(net, monitors=['E.spike', 'I.spike'], inputs=[('Ein.input', 20.), ('Iin.input', 20.)])
+        runner.run(100.)
+
+        # visualization
+        bp.visualize.raster_plot(runner.mon.ts, runner.mon['E.spike'],
+                                 title='Spikes of Excitatory Neurons', show=True)
+        bp.visualize.raster_plot(runner.mon.ts, runner.mon['I.spike'],
+                                 title='Spikes of Inhibitory Neurons', show=True)
+
+
+    """
+    def __init__(
+            self,
+            size: Shape,
+            keep_size: bool = False,
+            sharding: Optional[Any] = None,
+            name: Optional[str] = None,
+            mode: Optional[bm.Mode] = None,
+            method: str = 'exp_auto'
+    ):
+        super().__init__(size=size, keep_size=keep_size, sharding=sharding, name=name, mode=mode, method=method)
+
+        self.reset_state(self.mode)
+
+    def reset_state(self, batch_or_mode=None):
+        self.input = self.init_variable(bm.zeros, batch_or_mode)
+
+    def update(self, *args, **kwargs):
+        return self.input.value
+
+    def return_info(self):
+        return self.input
+
+    def clear_input(self, *args, **kwargs):
+        self.reset_state(self.mode)

brainpy/_src/dyn/projections/plasticity.py

@@ -40,7 +40,8 @@ class STDP_Song2000(Projection):
   where :math:`t_{sp}` denotes the spike time and :math:`A_1` is the increment
   of :math:`A_{pre}`, :math:`A_2` is the increment of :math:`A_{post}` produced by a spike.
 
-  Example::
+  Here is an example of the usage of this class::
+
     import brainpy as bp
     import brainpy.math as bm
 
@@ -138,10 +139,7 @@ def __init__(
       if not post.has_bef_update(self._post_repr):
         syn_cls = syn()
         out_cls = out()
-        if out_label is None:
-          out_name = self.name
-        else:
-          out_name = f'{out_label} // {self.name}'
+        out_name = self.name if out_label is None else f'{out_label} // {self.name}'
         post.add_inp_fun(out_name, out_cls)
         post.add_bef_update(self._post_repr, _AlignPost(syn_cls, out_cls))
       # references
@@ -154,35 +152,29 @@ def __init__(
       # synapse initialization
       self._syn_id = f'Delay({str(delay)}) // {syn.identifier}'
       if not delay_cls.has_bef_update(self._syn_id):
-        # delay
         delay_access = DelayAccess(delay_cls, delay)
-        # synapse
         syn_cls = syn()
-        # add to "after_updates"
         delay_cls.add_bef_update(self._syn_id, _AlignPreMg(delay_access, syn_cls))
-
       # output initialization
-      if out_label is None:
-        out_name = self.name
-      else:
-        out_name = f'{out_label} // {self.name}'
+      out_name = self.name if out_label is None else f'{out_label} // {self.name}'
       post.add_inp_fun(out_name, out)
-
       # references
       self.refs = dict(pre=pre, post=post)  # invisible to `self.nodes()`
       self.refs['delay'] = delay_cls.get_bef_update(self._syn_id)
       self.refs['syn'] = delay_cls.get_bef_update(self._syn_id).syn
       self.refs['out'] = out
 
-    self.refs['pre_trace'] = self.calculate_trace(pre, delay, Expon.desc(pre.num, tau=tau_s))
-    self.refs['post_trace'] = self.calculate_trace(post, None, Expon.desc(post.num, tau=tau_t))
-    # parameters
+    # trace initialization
+    self.refs['pre_trace'] = self._init_trace(pre, delay, Expon.desc(pre.num, tau=tau_s))
+    self.refs['post_trace'] = self._init_trace(post, None, Expon.desc(post.num, tau=tau_t))
+
+    # synapse parameters
     self.tau_s = parameter(tau_s, sizes=self.pre_num)
     self.tau_t = parameter(tau_t, sizes=self.post_num)
     self.A1 = parameter(A1, sizes=self.pre_num)
     self.A2 = parameter(A2, sizes=self.post_num)
 
-  def calculate_trace(
+  def _init_trace(
       self,
       target: DynamicalSystem,
       delay: Union[None, int, float],
@@ -234,5 +226,5 @@ def update(self):
     if issubclass(self.syn.cls, AlignPost):
       self.refs['syn'].add_current(current)  # synapse post current
     else:
-      self.refs['out'].bind_cond(current)
+      self.refs['out'].bind_cond(current)  # align pre
     return current

brainpy/_src/dyn/projections/tests/test_STDP.py

@@ -51,3 +51,4 @@ def run(i, I_pre, I_post):
 
     indices = bm.arange(0, duration, bm.dt)
     bm.for_loop(run, [indices, I_pre, I_post], jit=True)
+    bm.clear_buffer_memory()