[structure] update package structures

Author: chaoming0625

brainpy/__init__.py

@@ -55,35 +55,36 @@
                      layers,  # ANN layers
                      neurons,  # neuron groups
                      rates,  # rate models
-                     synapses,  # synaptic dynamics
-                     synouts,  # synaptic output
-                     synplast,  # synaptic plasticity
                      experimental,
                      )
-from brainpy._src.dyn.base import (DynamicalSystem as DynamicalSystem,
-                                   Container as Container,
-                                   Sequential as Sequential,
-                                   Network as Network,
-                                   NeuGroup as NeuGroup,
-                                   SynConn as SynConn,
-                                   SynOut as SynOut,
-                                   SynSTP as SynSTP,
-                                   SynLTP as SynLTP,
-                                   TwoEndConn as TwoEndConn,
-                                   CondNeuGroup as CondNeuGroup,
-                                   Channel as Channel)
-from brainpy._src.dyn.delay import Delay
+from brainpy.synapses import (synapses,  # synaptic dynamics
+                              synouts,  # synaptic output
+                              synplast,)  # synaptic plasticity
+
+from brainpy._src.dynsys import (DynamicalSystem as DynamicalSystem,
+                                 Container as Container,
+                                 Sequential as Sequential,
+                                 Network as Network,
+                                 NeuGroup as NeuGroup,
+                                 SynConn as SynConn,
+                                 SynOut as SynOut,
+                                 SynSTP as SynSTP,
+                                 SynLTP as SynLTP,
+                                 TwoEndConn as TwoEndConn,
+                                 CondNeuGroup as CondNeuGroup,
+                                 Channel as Channel)
+from brainpy._src.delay import Delay
 # shared parameters
-from brainpy._src.dyn.context import share
-from brainpy._src.dyn.base import not_pass_shared
+from brainpy._src.context import share
+from brainpy._src.dynsys import not_pass_shared
 # running
 from brainpy._src.dyn.runners import (DSRunner as DSRunner)
 from brainpy._src.dyn.transform import (LoopOverTime as LoopOverTime,)
 # DynamicalSystem base classes
-from brainpy._src.dyn.base import (DynamicalSystemNS as DynamicalSystemNS,
-                                   NeuGroupNS as NeuGroupNS,
-                                   TwoEndConnNS as TwoEndConnNS,
-                                   )
+from brainpy._src.dynsys import (DynamicalSystemNS as DynamicalSystemNS,
+                                 NeuGroupNS as NeuGroupNS,
+                                 TwoEndConnNS as TwoEndConnNS,
+                                 )
 from brainpy._src.dyn.synapses_v2.base import (SynOutNS as SynOutNS,
                                                SynSTPNS as SynSTPNS,
                                                SynConnNS as SynConnNS, )
@@ -92,7 +93,6 @@
 #  Part 4: Training  #
 # ------------------ #
 
-from . import train
 from ._src.train.base import (DSTrainer as DSTrainer)
 from ._src.train.back_propagation import (BPTT as BPTT,
                                           BPFF as BPFF,)
@@ -102,17 +102,12 @@
                                  RidgeTrainer as RidgeTrainer,)
 
 
-#  Part 5: Analysis  #
-# ------------------ #
-
-from . import analysis
-
-
 #  Part 6: Others    #
 # ------------------ #
 
-from . import running, testing, interop
+from . import running, testing, analysis
 from ._src.visualization import (visualize as visualize)
+from ._src import base, modes, train, dyn
 
 
 #  Part 7: Deprecations  #
@@ -167,7 +162,6 @@
 train.__dict__['RidgeTrainer'] = RidgeTrainer
 
 
-from ._src import base
 base.base.__dict__['BrainPyObject'] = BrainPyObject
 base.base.__dict__['Base'] = Base
 base.collector.__dict__['Collector'] = Collector
@@ -201,8 +195,6 @@
 base.__dict__['load_by_mat'] = checkpoints.io.load_by_mat
 base.__dict__['clear_name_cache'] = math.clear_name_cache
 
-
-from . import modes
 modes.__dict__['Mode'] = math.Mode
 modes.__dict__['NormalMode'] = math.NonBatchingMode
 modes.__dict__['BatchingMode'] = math.BatchingMode
@@ -213,7 +205,6 @@
 modes.__dict__['check_mode'] = check.is_subclass
 
 
-from brainpy import dyn
 dyn.__dict__['channels'] = channels
 dyn.__dict__['neurons'] = neurons
 dyn.__dict__['rates'] = rates

brainpy/_src/analysis/highdim/slow_points.py

@@ -13,7 +13,7 @@
 import brainpy._src.math as bm
 from brainpy import optim, losses
 from brainpy._src.analysis import utils, base, constants
-from brainpy._src.dyn.base import DynamicalSystem
+from brainpy._src.dynsys import DynamicalSystem
 from brainpy._src.dyn.runners import check_and_format_inputs, _f_ops
 from brainpy._src.tools.dicts import DotDict
 from brainpy.errors import AnalyzerError, UnsupportedError

brainpy/_src/analysis/utils/model.py

@@ -4,7 +4,7 @@
 from brainpy._src.math.object_transform import Variable
 from brainpy._src.math.environment import get_float
 from brainpy._src.math.interoperability import as_jax
-from brainpy._src.dyn.base import DynamicalSystem
+from brainpy._src.dynsys import DynamicalSystem
 from brainpy._src.dyn.runners import DSRunner
 from brainpy._src.integrators.base import Integrator
 from brainpy._src.integrators.joint_eq import JointEq

brainpy/_src/dyn/context.py renamed to brainpy/_src/context.py

@@ -7,7 +7,7 @@
 from typing import Any
 from typing import Union
 
-from brainpy._src.dyn.base import DynamicalSystemNS
+from brainpy._src.dynsys import DynamicalSystemNS
 from brainpy._src.math.environment import get_dt
 from brainpy._src.tools.dicts import DotDict
 
@@ -51,7 +51,8 @@ def load(self, key, value: Any = None):
     if key in self._arguments:
       return self._arguments[key]
     if value is None:
-      raise KeyError(f'Cannot found shared data of {key}.')
+      raise KeyError(f'Cannot found shared data of {key}. '
+                     f'Please define it with "brainpy.share.save()". ')
     else:
       return value
 

brainpy/_src/dyn/delay.py renamed to brainpy/_src/delay.py

@@ -11,10 +11,10 @@
 
 from brainpy import check
 from brainpy import math as bm
-from brainpy._src.dyn.base import DynamicalSystemNS
+from brainpy._src.dynsys import DynamicalSystemNS
 from brainpy._src.math.delayvars import ROTATE_UPDATE, CONCAT_UPDATE
 from brainpy.check import is_integer, jit_error_checking
-from .context import share
+from brainpy._src.context import share
 
 __all__ = [
   'Delay',
@@ -86,7 +86,8 @@ def __init__(
     self.latest = latest
 
     # delay length
-    self.length = is_integer(length, allow_none=False, min_bound=0)
+    assert isinstance(length, int)
+    self.length = length
 
     # delay data
     if before_t0 is not None:

brainpy/_src/dyn/__init__.py

@@ -11,7 +11,6 @@
   runners,
   transform,
 )
-from .base import *
 from .neurons.compat import *
 from .runners import *
 from .synapses.compat import *

brainpy/_src/dyn/channels/Ca.py

@@ -8,7 +8,7 @@
 from typing import Union, Callable
 
 import brainpy.math as bm
-from brainpy._src.dyn.base import Channel
+from brainpy._src.dynsys import Channel
 from brainpy._src.initialize import OneInit, Initializer, parameter, variable
 from brainpy._src.integrators.joint_eq import JointEq
 from brainpy._src.integrators.ode.generic import odeint

brainpy/_src/dyn/channels/base.py

@@ -3,7 +3,7 @@
 from typing import Union
 
 import brainpy.math as bm
-from brainpy._src.dyn.base import Container, CondNeuGroup, Channel, check_master
+from brainpy._src.dynsys import Container, CondNeuGroup, Channel, check_master
 from brainpy.types import Shape
 
 __all__ = [

brainpy/_src/dyn/networks/cann.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 
-from brainpy._src.dyn.base import NeuGroup
+from brainpy._src.dynsys import NeuGroup
 
 __all__ = [
   'WuCANN1D',

brainpy/_src/dyn/neurons/biological_models.py

@@ -4,8 +4,8 @@
 
 import brainpy.math as bm
 from brainpy import check
-from brainpy._src.dyn.base import NeuGroupNS
-from brainpy._src.dyn.context import share
+from brainpy._src.dynsys import NeuGroupNS
+from brainpy._src.context import share
 from brainpy._src.initialize import (OneInit,
                                      Uniform,
                                      Initializer,
@@ -226,9 +226,7 @@ def __init__(
                              keep_size=keep_size,
                              name=name,
                              mode=mode)
-    check.is_subclass(self.mode,
-                      (bm.BatchingMode, bm.NonBatchingMode),
-                      self.__class__.__name__)
+    assert self.mode.is_one_of(bm.BatchingMode, bm.NonBatchingMode)
 
     # parameters
     self.ENa = parameter(ENa, self.varshape, allow_none=False)
@@ -436,7 +434,7 @@ def __init__(
                                       keep_size=keep_size,
                                       name=name,
                                       mode=mode)
-    check.is_subclass(self.mode, (bm.BatchingMode, bm.NonBatchingMode))
+    assert self.mode.is_one_of(bm.BatchingMode, bm.NonBatchingMode)
 
     # params
     self.V_Ca = parameter(V_Ca, self.varshape, allow_none=False)
@@ -700,7 +698,7 @@ def __init__(
                                             keep_size=keep_size,
                                             name=name,
                                             mode=mode)
-    check.is_subclass(self.mode, (bm.NonBatchingMode, bm.BatchingMode))
+    assert self.mode.is_one_of(bm.BatchingMode, bm.NonBatchingMode)
 
     # conductance parameters
     self.gAHP = parameter(gAHP, self.varshape, allow_none=False)
@@ -743,7 +741,6 @@ def __init__(
     self.q = bm.Variable(self.inf_q(self.Ca), batch_axis=0 if isinstance(self.mode, bm.BatchingMode) else None)
     self.Id = variable_(bm.zeros, self.varshape, self.mode)  # input to soma
     self.Is = variable_(bm.zeros, self.varshape, self.mode)  # input to dendrite
-    # self.spike = bm.Variable(bm.zeros(self.varshape, dtype=bool))
 
     # integral
     if self.noise is None:
@@ -761,10 +758,8 @@ def reset_state(self, batch_size=None):
     self.s.value = bm.Variable(self.inf_s(self.Vd), batch_axis=batch_axis)
     self.c.value = bm.Variable(self.inf_c(self.Vd), batch_axis=batch_axis)
     self.q.value = bm.Variable(self.inf_q(self.Ca), batch_axis=batch_axis)
-    if self.input_var:
-      self.Id.value = variable_(bm.zeros, self.varshape, batch_size)
-      self.Is.value = variable_(bm.zeros, self.varshape, batch_size)
-    # self.spike[:] = False
+    self.Id.value = variable_(bm.zeros, self.varshape, batch_size)
+    self.Is.value = variable_(bm.zeros, self.varshape, batch_size)
 
   def dCa(self, Ca, t, s, Vd):
     ICa = self.gCa * s * s * (Vd - self.ECa)
@@ -1011,7 +1006,7 @@ def __init__(
   ):
     # initialization
     super(WangBuzsakiModel, self).__init__(size=size, keep_size=keep_size, name=name, mode=mode)
-    check.is_subclass(self.mode, (bm.BatchingMode, bm.NonBatchingMode))
+    assert self.mode.is_one_of(bm.BatchingMode, bm.NonBatchingMode)
 
     # parameters
     self.ENa = parameter(ENa, self.varshape, allow_none=False)