Updates

Author: chaoming0625

brainpy/__init__.py

@@ -109,6 +109,8 @@
 init = initialize
 optim = optimizers
 
+from . import experimental
+
 
 # deprecated
 from . import base

brainpy/dyn/base.py

@@ -841,9 +841,9 @@ def __init__(
 
     # pre or post neuron group
     # ------------------------
-    if not isinstance(pre, NeuGroup):
+    if not isinstance(pre, (NeuGroup, DynamicalSystem)):
       raise TypeError('"pre" must be an instance of NeuGroup.')
-    if not isinstance(post, NeuGroup):
+    if not isinstance(post, (NeuGroup, DynamicalSystem)):
       raise TypeError('"post" must be an instance of NeuGroup.')
     self.pre = pre
     self.post = post

brainpy/dyn/synapses/abstract_models.py

@@ -280,7 +280,7 @@ def __init__(
       pre: NeuGroup,
       post: NeuGroup,
       conn: Union[TwoEndConnector, ArrayType, Dict[str, ArrayType]],
-      output: SynOut = CUBA(),
+      output: Optional[SynOut] = CUBA(),
       stp: Optional[SynSTP] = None,
       comp_method: str = 'sparse',
       g_max: Union[float, ArrayType, Initializer, Callable] = 1.,

brainpy/experimental/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+from .synapses import *

brainpy/experimental/synapses.py

@@ -0,0 +1,122 @@
+# -*- coding: utf-8 -*-
+
+from typing import Union, Optional
+
+import brainpylib as bl
+import jax
+
+from brainpy import (math as bm,
+                     initialize as init,
+                     connect)
+from brainpy.dyn.base import DynamicalSystem, SynSTP
+from brainpy.integrators.ode import odeint
+from brainpy.types import Initializer, ArrayType
+
+__all__ = [
+  'Exponential',
+]
+
+
+class Exponential(DynamicalSystem):
+  def __init__(
+      self,
+      conn: connect.TwoEndConnector,
+      stp: Optional[SynSTP] = None,
+      g_max: Union[float, Initializer] = 1.,
+      g_initializer: Union[float, Initializer] = init.ZeroInit(),
+      tau: Union[float, ArrayType] = 8.0,
+      method: str = 'exp_auto',
+      mode: Optional[bm.Mode] = None,
+      name: Optional[str] = None,
+  ):
+    super(Exponential, self).__init__(name=name, mode=mode)
+
+    # component
+    self.conn = conn
+    self.stp = stp
+    self.g_initializer = g_initializer
+    assert self.conn.pre_num is not None
+    assert self.conn.post_num is not None
+
+    # parameters
+    self.tau = tau
+    if bm.size(self.tau) != 1:
+      raise ValueError(f'"tau" must be a scalar or a tensor with size of 1. But we got {self.tau}')
+
+    # connections and weights
+    if isinstance(self.conn, connect.One2One):
+      self.g_max = init.parameter(g_max, (self.conn.pre_num,), allow_none=False)
+
+    elif isinstance(self.conn, connect.All2All):
+      self.g_max = init.parameter(g_max, (self.conn.pre_num, self.conn.post_num), allow_none=False)
+
+    else:
+      self.conn_mask = self.conn.require('pre2post')
+      self.g_max = init.parameter(g_max, self.conn_mask[0].shape, allow_none=False)
+
+    # variables
+    self.g = init.variable_(g_initializer, self.conn.post_num, self.mode)
+
+    # function
+    self.integral = odeint(lambda g, t: -g / self.tau, method=method)
+
+  def reset_state(self, batch_size=None):
+    self.g.value = init.variable_(bm.zeros, self.conn.post_num, batch_size)
+    if self.stp is not None:
+      self.stp.reset_state(batch_size)
+
+  def _syn2post_with_one2one(self, syn_value, syn_weight):
+    return syn_value * syn_weight
+
+  def _syn2post_with_all2all(self, syn_value, syn_weight):
+    if bm.ndim(syn_weight) == 0:
+      if isinstance(self.mode, bm.BatchingMode):
+        assert syn_value.ndim == 2
+        post_vs = bm.sum(syn_value, keepdims=True, axis=1)
+      else:
+        post_vs = bm.sum(syn_value)
+      if not self.conn.include_self:
+        post_vs = post_vs - syn_value
+      post_vs = syn_weight * post_vs
+    else:
+      assert syn_weight.ndim == 2
+      if isinstance(self.mode, bm.BatchingMode):
+        assert syn_value.ndim == 2
+        post_vs = syn_value @ syn_weight
+      else:
+        post_vs = syn_value @ syn_weight
+    return post_vs
+
+  def update(self, tdi, spike):
+    t, dt = tdi['t'], tdi.get('dt', bm.dt)
+
+    # update sub-components
+    if self.stp is not None:
+      self.stp.update(tdi, spike)
+
+    # post values
+    if isinstance(self.conn, connect.All2All):
+      syn_value = bm.asarray(spike, dtype=bm.float_)
+      if self.stp is not None:
+        syn_value = self.stp(syn_value)
+      post_vs = self._syn2post_with_all2all(syn_value, self.g_max)
+    elif isinstance(self.conn, connect.One2One):
+      syn_value = bm.asarray(spike, dtype=bm.float_)
+      if self.stp is not None:
+        syn_value = self.stp(syn_value)
+      post_vs = self._syn2post_with_one2one(syn_value, self.g_max)
+    else:
+      if isinstance(self.mode, bm.BatchingMode):
+        f = jax.vmap(bl.event_ops.event_csr_matvec, in_axes=(None, None, None, 0))
+        post_vs = f(self.g_max, self.conn_mask[0], self.conn_mask[1], spike,
+                    shape=(self.conn.pre_num, self.conn.post_num), transpose=True)
+      else:
+        post_vs = bl.event_ops.event_csr_matvec(
+          self.g_max, self.conn_mask[0], self.conn_mask[1], spike,
+          shape=(self.conn.pre_num, self.conn.post_num), transpose=True
+        )
+    # updates
+    self.g.value = self.integral(self.g.value, t, dt) + post_vs
+
+    # output
+    return self.g.value

brainpy/initialize/generic.py

@@ -94,16 +94,16 @@ def init_param(
 
 
 def variable_(
-    data: Union[Callable, ArrayType],
+    init: Union[Callable, ArrayType],
     size: Shape = None,
     batch_size_or_mode: Optional[Union[int, bool, bm.Mode]] = None,
     batch_axis: int = 0,
 ):
-  """Initialize variables. Same as `variable()`.
+  """Initialize a :math:`~.Variable` from a callable function or a data.
 
   Parameters
   ----------
-  data: callable, function, ArrayType
+  init: callable, function, ArrayType
     The data to be initialized as a ``Variable``.
   batch_size_or_mode: int, bool, Mode, optional
     The batch size, model ``Mode``, boolean state.
@@ -125,11 +125,11 @@ def variable_(
   variable, parameter, noise, delay
 
   """
-  return variable(data, batch_size_or_mode, size, batch_axis)
+  return variable(init, batch_size_or_mode, size, batch_axis)
 
 
 def variable(
-    data: Union[Callable, ArrayType],
+    init: Union[Callable, ArrayType],
     batch_size_or_mode: Optional[Union[int, bool, bm.Mode]] = None,
     size: Shape = None,
     batch_axis: int = 0,
@@ -138,7 +138,7 @@ def variable(
 
   Parameters
   ----------
-  data: callable, function, ArrayType
+  init: callable, function, ArrayType
     The data to be initialized as a ``Variable``.
   batch_size_or_mode: int, bool, Mode, optional
     The batch size, model ``Mode``, boolean state.
@@ -161,34 +161,34 @@ def variable(
 
   """
   size = to_size(size)
-  if callable(data):
+  if callable(init):
     if size is None:
       raise ValueError('"varshape" cannot be None when data is a callable function.')
     if isinstance(batch_size_or_mode, bm.NonBatchingMode):
-      return bm.Variable(data(size))
+      return bm.Variable(init(size))
     elif isinstance(batch_size_or_mode, bm.BatchingMode):
       new_shape = size[:batch_axis] + (1,) + size[batch_axis:]
-      return bm.Variable(data(new_shape), batch_axis=batch_axis)
+      return bm.Variable(init(new_shape), batch_axis=batch_axis)
     elif batch_size_or_mode in (None, False):
-      return bm.Variable(data(size))
+      return bm.Variable(init(size))
     elif isinstance(batch_size_or_mode, int):
       new_shape = size[:batch_axis] + (int(batch_size_or_mode),) + size[batch_axis:]
-      return bm.Variable(data(new_shape), batch_axis=batch_axis)
+      return bm.Variable(init(new_shape), batch_axis=batch_axis)
     else:
       raise ValueError('Unknown batch_size_or_mode.')
 
   else:
     if size is not None:
-      if bm.shape(data) != size:
-        raise ValueError(f'The shape of "data" {bm.shape(data)} does not match with "var_shape" {size}')
+      if bm.shape(init) != size:
+        raise ValueError(f'The shape of "data" {bm.shape(init)} does not match with "var_shape" {size}')
     if isinstance(batch_size_or_mode, bm.NonBatchingMode):
-      return bm.Variable(data)
+      return bm.Variable(init)
     elif isinstance(batch_size_or_mode, bm.BatchingMode):
-      return bm.Variable(bm.expand_dims(data, axis=batch_axis), batch_axis=batch_axis)
+      return bm.Variable(bm.expand_dims(init, axis=batch_axis), batch_axis=batch_axis)
     elif batch_size_or_mode in (None, False):
-      return bm.Variable(data)
+      return bm.Variable(init)
     elif isinstance(batch_size_or_mode, int):
-      return bm.Variable(bm.repeat(bm.expand_dims(data, axis=batch_axis),
+      return bm.Variable(bm.repeat(bm.expand_dims(init, axis=batch_axis),
                                    int(batch_size_or_mode),
                                    axis=batch_axis),
                          batch_axis=batch_axis)

brainpy/math/_utils.py

@@ -1,10 +1,14 @@
 # -*- coding: utf-8 -*-
 
-from typing import Callable
 import functools
-from .ndarray import Array
+from typing import Callable
+
+import jax
+import numpy as np
 from jax.tree_util import tree_map
 
+from .ndarray import Array
+
 
 def wraps(fun: Callable):
   """Specialized version of functools.wraps for wrapping numpy functions.
@@ -15,6 +19,7 @@ def wraps(fun: Callable):
   this reason, it is important that parameter names match those in the original
   numpy function.
   """
+
   def wrap(op):
     docstr = getattr(fun, "__doc__", None)
     op.__doc__ = docstr
@@ -27,20 +32,31 @@ def wrap(op):
       else:
         setattr(op, attr, value)
     return op
+
   return wrap
 
+
 def _as_jax_array(a):
   return a.value if isinstance(a, Array) else a
 
+
+def _as_brainpy_array(a):
+  return Array(a) if isinstance(a, (np.ndarray, jax.Array)) else a
+
+
 def _is_leaf(a):
   return isinstance(a, Array)
 
 
-def _compatible_with_brainpy_array(fun):
+def _compatible_with_brainpy_array(fun: Callable, return_brainpy_array: bool = False):
   @functools.wraps(fun)
   def new_fun(*args, **kwargs):
     args = tree_map(_as_jax_array, args, is_leaf=_is_leaf)
-    kwargs = tree_map(_as_jax_array, kwargs, is_leaf=_is_leaf)
-    return fun(*args, **kwargs)
-  return new_fun
+    if len(kwargs):
+      kwargs = tree_map(_as_jax_array, kwargs, is_leaf=_is_leaf)
+    r = fun(*args, **kwargs)
+    return tree_map(_as_brainpy_array, r) if return_brainpy_array else r
+
+  new_fun.__doc__ = getattr(fun, "__doc__", None)
 
+  return new_fun

brainpy/math/operators/pre_syn_post.py

@@ -8,7 +8,6 @@
 
 from brainpy.errors import MathError
 from brainpy.math.numpy_ops import as_jax
-from brainpy.types import ArrayType
 
 __all__ = [
   # pre-to-post
@@ -43,10 +42,10 @@ def _raise_pre_ids_is_none(pre_ids):
                     f'(brainpy.math.ndim(pre_values) != 0).')
 
 
-def pre2post_event_sum(events: ArrayType,
-                       pre2post: Tuple[ArrayType, ArrayType],
+def pre2post_event_sum(events,
+                       pre2post,
                        post_num: int,
-                       values: Union[float, ArrayType] = 1.):
+                       values = 1.):
   """The pre-to-post event-driven synaptic summation with `CSR` synapse structure.
 
   When ``values`` is a scalar, this function is equivalent to
@@ -103,11 +102,11 @@ def pre2post_event_sum(events: ArrayType,
                                      transpose=True)
 
 
-def pre2post_coo_event_sum(events: ArrayType,
-                           pre_ids: ArrayType,
-                           post_ids: ArrayType,
+def pre2post_coo_event_sum(events,
+                           pre_ids,
+                           post_ids,
                            post_num: int,
-                           values: Union[float, ArrayType] = 1.):
+                           values = 1.):
   """The pre-to-post synaptic computation with event-driven summation.
 
   Parameters

brainpy/math/operators/sparse_matmul.py

@@ -9,7 +9,6 @@
 
 from brainpy.math.ndarray import Array
 from brainpy.math.numpy_ops import as_jax
-from brainpy.types import ArrayType
 
 __all__ = [
   'sparse_matmul',
@@ -18,10 +17,10 @@
 ]
 
 
-def event_csr_matvec(values: ArrayType,
-                     indices: ArrayType,
-                     indptr: ArrayType,
-                     events: ArrayType,
+def event_csr_matvec(values,
+                     indices,
+                     indptr,
+                     events,
                      shape: Tuple[int, int],
                      transpose: bool = False):
   """The pre-to-post event-driven synaptic summation with `CSR` synapse structure.

brainpy/types.py

@@ -1,15 +1,17 @@
 # -*- coding: utf-8 -*-
 
-from typing import TypeVar, Tuple
+from typing import TypeVar, Tuple, Union, Callable
 
 import jax.numpy as jnp
 import numpy as np
 
 from brainpy.math.ndarray import Array, Variable, TrainVar
+from brainpy import connect as conn
+from brainpy import initialize as init
 
 __all__ = [
   'ArrayType', 'Parameter', 'PyTree',
-  'Shape',
+  'Shape', 'Initializer',
   'Output', 'Monitor'
 ]
 
@@ -26,4 +28,6 @@
 # component
 Output = TypeVar('Output') # noqa
 Monitor = TypeVar('Monitor') # noqa
+Connector = Union[conn.Connector, Array, Variable, jnp.ndarray, np.ndarray]
+Initializer = Union[init.Initializer, Callable, Array, Variable, jnp.ndarray, np.ndarray]