Merge pull request #421 from brainpy/updates

Recent updates

Author: chaoming0625

brainpy/__init__.py

@@ -58,14 +58,14 @@
   DynamicalSystem as DynamicalSystem,
   DynSysGroup as DynSysGroup,  # collectors
   Sequential as Sequential,
-  Network as Network,
   Dynamic as Dynamic,  # category
   Projection as Projection,
 )
 DynamicalSystemNS = DynamicalSystem
+Network = DynSysGroup
 # delays
 from brainpy._src.delay import (
-  VariableDelay as VariableDelay,
+  VariDelay as VariDelay,
 )
 
 # building blocks

brainpy/_src/delay.py

@@ -21,8 +21,9 @@
 
 __all__ = [
   'Delay',
-  'VariableDelay',
+  'VariDelay',
   'DataDelay',
+  'DelayAccess',
 ]
 
 
@@ -431,8 +432,8 @@ def _check_target_sharding(sharding, ndim, mode: bm.Mode):
   return sharding
 
 
-class VariableDelay(Delay):
-  """Delay variable which has a fixed delay length.
+class VariDelay(Delay):
+  """Generate Delays for the given :py:class:`~.Variable` instance.
 
   The data in this delay variable is arranged as::
 
@@ -517,8 +518,8 @@ def __init__(
 
     # other info
     if entries is not None:
-      for entry, value in entries.items():
-        self.register_entry(entry, value)
+      for entry, delay_time in entries.items():
+        self.register_entry(entry, delay_time)
 
   def register_entry(
       self,
@@ -572,11 +573,17 @@ def at(self, entry: str, *indices) -> bm.Array:
       raise KeyError(f'Does not find delay entry "{entry}".')
     delay_step = self._registered_entries[entry]
     if delay_step is None or delay_step == 0.:
-      return self.target.value
+      if len(indices):
+        return self.target[indices]
+      else:
+        return self.target.value
     else:
       assert self.data is not None
       if delay_step == 0:
-        return self.target.value
+        if len(indices):
+          return self.target[indices]
+        else:
+          return self.target.value
       else:
         return self.retrieve(delay_step, *indices)
 
@@ -683,16 +690,15 @@ def _init_data(self, length: int, batch_size: int = None):
       self.data[:] = self._init((length,) + self.target.shape, dtype=self.target.dtype)
 
 
-class DataDelay(VariableDelay):
-  
+class DataDelay(VariDelay):
   not_desc_params = ('time', 'entries')
 
   def __init__(
       self,
 
       # delay target
-      target: bm.Variable,
-      target_init: Callable,
+      data: bm.Variable,
+      data_init: Union[Callable, bm.Array, jax.Array],
 
       # delay time
       time: Optional[Union[int, float]] = None,
@@ -710,8 +716,8 @@ def __init__(
       name: Optional[str] = None,
       mode: Optional[bm.Mode] = None,
   ):
-    self.target_init = target_init
-    super().__init__(target=target,
+    self.target_init = data_init
+    super().__init__(target=data,
                      time=time,
                      init=init,
                      entries=entries,
@@ -736,3 +742,20 @@ def update(
     super().update(latest_value)
 
 
+class DelayAccess(DynamicalSystem):
+  def __init__(
+      self,
+      delay: Delay,
+      time: Union[None, int, float],
+      *indices
+  ):
+    super().__init__(mode=delay.mode)
+    self.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)
+
+

brainpy/_src/dnn/linear.py

@@ -4,6 +4,7 @@
 from typing import Dict, Optional, Union, Callable
 
 import jax
+import numpy as np
 import jax.numpy as jnp
 
 from brainpy import math as bm
@@ -63,8 +64,8 @@ def __init__(
       num_out: int,
       W_initializer: Union[Initializer, Callable, ArrayType] = XavierNormal(),
       b_initializer: Optional[Union[Initializer, Callable, ArrayType]] = ZeroInit(),
-      mode: bm.Mode = None,
-      name: str = None,
+      mode: Optional[bm.Mode] = None,
+      name: Optional[str] = None,
   ):
     super(Dense, self).__init__(mode=mode, name=name)
 
@@ -642,7 +643,7 @@ def __init__(
       num_out: int,
       prob: float,
       weight: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -654,7 +655,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -723,7 +724,7 @@ def __init__(
       prob: float,
       w_low: float,
       w_high: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -735,7 +736,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -803,7 +804,7 @@ def __init__(
       prob: float,
       w_mu: float,
       w_sigma: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       transpose: bool = False,
       atomic: bool = False,
@@ -815,7 +816,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -881,7 +882,7 @@ def __init__(
       num_out: int,
       prob: float,
       weight: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -893,7 +894,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 1000000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -962,7 +963,7 @@ def __init__(
       prob: float,
       w_low: float,
       w_high: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -974,7 +975,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -1042,7 +1043,7 @@ def __init__(
       prob: float,
       w_mu: float,
       w_sigma: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       transpose: bool = False,
       atomic: bool = False,
@@ -1054,7 +1055,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out

brainpy/_src/dyn/channels/potassium_compatible.py

@@ -9,12 +9,11 @@
 
 import brainpy.math as bm
 from brainpy._src.context import share
-from brainpy._src.dyn.channels.leaky import LeakyChannel
+from brainpy._src.dyn.channels.base import IonChannel
 from brainpy._src.dyn.neurons.hh import HHTypedNeuron
 from brainpy._src.initialize import Initializer, parameter, variable
 from brainpy._src.integrators import odeint, JointEq
 from brainpy.types import ArrayType
-from .potassium import PotassiumChannel
 
 __all__ = [
   'IKDR_Ba2002',
@@ -29,7 +28,7 @@
 ]
 
 
-class _IK_p4_markov(PotassiumChannel):
+class _IK_p4_markov(IonChannel):
   r"""The delayed rectifier potassium channel of :math:`p^4`
   current which described with first-order Markov chain.
 
@@ -339,7 +338,7 @@ def f_p_beta(self, V):
     return 0.125 * bm.exp(-(V - self.V_sh + 20) / 80)
 
 
-class _IKA_p4q_ss(PotassiumChannel):
+class _IKA_p4q_ss(IonChannel):
   r"""The rapidly inactivating Potassium channel of :math:`p^4q`
   current which described with steady-state format.
 
@@ -634,7 +633,7 @@ def f_q_tau(self, V):
                     19.)
 
 
-class _IKK2_pq_ss(PotassiumChannel):
+class _IKK2_pq_ss(IonChannel):
   r"""The slowly inactivating Potassium channel of :math:`pq`
   current which described with steady-state format.
 
@@ -921,7 +920,7 @@ def f_q_tau(self, V):
                     8.9)
 
 
-class IKNI_Ya1989(PotassiumChannel):
+class IKNI_Ya1989(IonChannel):
   r"""A slow non-inactivating K+ current described by Yamada et al. (1989) [1]_.
 
   This slow potassium current can effectively account for spike-frequency adaptation.
@@ -1019,7 +1018,7 @@ def f_p_tau(self, V):
     return self.tau_max / (3.3 * bm.exp(temp / 20.) + bm.exp(-temp / 20.))
 
 
-class IKL(LeakyChannel):
+class IKL(IonChannel):
   """The potassium leak channel current.
 
   Parameters
@@ -1031,6 +1030,8 @@ class IKL(LeakyChannel):
     The reversal potential.
   """
 
+  master_type = HHTypedNeuron
+
   def __init__(
       self,
       size: Union[int, Sequence[int]],

brainpy/_src/dyn/channels/sodium_compatible.py

@@ -13,7 +13,7 @@
 from brainpy._src.initialize import Initializer, parameter, variable
 from brainpy._src.integrators import odeint, JointEq
 from brainpy.types import ArrayType
-from .sodium import SodiumChannel
+from .base import IonChannel
 
 __all__ = [
   'INa_Ba2002',
@@ -22,7 +22,7 @@
 ]
 
 
-class _INa_p3q_markov(SodiumChannel):
+class _INa_p3q_markov(IonChannel):
   r"""The sodium current model of :math:`p^3q` current which described with first-order Markov chain.
 
   The general model can be used to model the dynamics with:
@@ -64,7 +64,7 @@ def __init__(
       name: str = None,
       mode: bm.Mode = None,
   ):
-    super(_INa_p3q_markov, self).__init__(size=size,
+    super().__init__(size=size,
                                           keep_size=keep_size,
                                           name=name,
                                           mode=mode)
@@ -173,7 +173,7 @@ def __init__(
       name: str = None,
       mode: bm.Mode = None,
   ):
-    super(INa_Ba2002, self).__init__(size,
+    super().__init__(size,
                                      keep_size=keep_size,
                                      name=name,
                                      method=method,
@@ -260,7 +260,7 @@ def __init__(
       name: str = None,
       mode: bm.Mode = None,
   ):
-    super(INa_TM1991, self).__init__(size,
+    super().__init__(size,
                                      keep_size=keep_size,
                                      name=name,
                                      method=method,
@@ -347,7 +347,7 @@ def __init__(
       name: str = None,
       mode: bm.Mode = None,
   ):
-    super(INa_HH1952, self).__init__(size,
+    super().__init__(size,
                                      keep_size=keep_size,
                                      name=name,
                                      method=method,

brainpy/_src/dyn/ions/base.py

@@ -166,13 +166,14 @@ def update(self, V):
     for node in self.nodes(level=1, include_self=False).unique().subset(IonChaDyn).values():
       node.update(V, self.C, self.E)
 
-  def current(self, V, C=None, E=None):
+  def current(self, V, C=None, E=None, external: bool = False):
     """Generate ion channel current.
 
     Args:
       V: The membrane potential.
       C: The ion concentration.
       E: The reversal potential.
+      external: Include the external current.
 
     Returns:
       Current.
@@ -186,8 +187,9 @@ def current(self, V, C=None, E=None):
     if len(nodes) > 0:
       for node in nodes:
         current = current + node.current(V, C, E)
-    for key, node in self.external.items():
-      current = current + node(V, C, E)
+    if external:
+      for key, node in self.external.items():
+        current = current + node(V, C, E)
     return current
 
   def reset_state(self, V, batch_size=None):