refactor: update imports to use brainpy.state_based and clean up code structure

Author: oujago

brainpy/state_based/_base.py

@@ -15,9 +15,8 @@
 
 from typing import Callable, Optional
 
-import braintools
-
 import brainstate
+import braintools
 
 __all__ = [
     'Neuron', 'Synapse',
@@ -315,7 +314,7 @@ class Synapse(brainstate.nn.Dynamics):
 
     .. code-block:: python
 
-        >>> import brainpy
+        >>> import brainpy.state_based as brainpy
         >>> import brainstate
         >>> import brainunit as u
         >>> import braintools
@@ -344,7 +343,7 @@ class Synapse(brainstate.nn.Dynamics):
 
     .. code-block:: python
 
-        >>> import brainpy
+        >>> import brainpy.state_based as brainpy
         >>> import brainstate
         >>> import brainunit as u
         >>> import jax
@@ -369,7 +368,7 @@ class Synapse(brainstate.nn.Dynamics):
 
     .. code-block:: python
 
-        >>> import brainpy
+        >>> import brainpy.state_based as brainpy
         >>> import brainstate
         >>> import brainunit as u
         >>>
@@ -410,7 +409,7 @@ class Synapse(brainstate.nn.Dynamics):
 
     .. code-block:: python
 
-        >>> import brainpy
+        >>> import brainpy.state_based as brainpy
         >>> import brainstate
         >>> import brainunit as u
         >>>

brainpy/state_based/_base_test.py

@@ -15,12 +15,12 @@
 
 import unittest
 
+import brainstate
 import braintools
 import brainunit as u
 import jax.numpy as jnp
 
-import brainstate
-from brainpy._base import Neuron, Synapse
+from brainpy.state_based import Neuron, Synapse
 
 
 class TestNeuronBaseClass(unittest.TestCase):

brainpy/state_based/_exponential.py

@@ -18,13 +18,13 @@
 
 from typing import Optional, Callable
 
+import brainstate
 import braintools
 import brainunit as u
-
-import brainstate
 from brainstate.typing import Size, ArrayLike
+
+from brainpy.mixin import AlignPost
 from ._base import Synapse
-from .mixin import AlignPost
 
 __all__ = [
     'Expon', 'DualExpon',

brainpy/state_based/_inputs.py

@@ -15,15 +15,15 @@
 
 from typing import Union, Optional, Sequence, Callable
 
+import brainstate
 import braintools
 import brainunit as u
 import jax
 import numpy as np
-
-import brainstate
-from brainpy._misc import set_module_as
 from brainstate.typing import ArrayLike, Size, DTypeLike
 
+from ._misc import set_module_as
+
 __all__ = [
     'SpikeTime',
     'PoissonSpike',

brainpy/state_based/_lif.py

@@ -17,12 +17,12 @@
 
 from typing import Callable
 
+import brainstate
 import braintools
 import brainunit as u
 import jax
-
-import brainstate
 from brainstate.typing import ArrayLike, Size
+
 from ._base import Neuron
 
 __all__ = [
@@ -74,7 +74,7 @@ class IF(Neuron):
 
     Examples
     --------
-    >>> import brainpy
+    >>> import brainpy.state_based as brainpy
     >>> import brainstate
     >>> import brainunit as u
     >>>
@@ -207,7 +207,7 @@ class LIF(Neuron):
 
     Examples
     --------
-    >>> import brainpy
+    >>> import brainpy.state_based as brainpy
     >>> import brainstate
     >>> import brainunit as u
     >>>
@@ -344,7 +344,7 @@ class LIFRef(Neuron):
 
     Examples
     --------
-    >>> import brainpy
+    >>> import brainpy.state_based as brainpy
     >>> import brainstate
     >>> import brainunit as u
     >>>
@@ -512,7 +512,7 @@ class ALIF(Neuron):
 
     Examples
     --------
-    >>> import brainpy
+    >>> import brainpy.state_based as brainpy
     >>> import brainstate
     >>> import brainunit as u
     >>>

brainpy/state_based/_lif_test.py

@@ -18,12 +18,12 @@
 
 import unittest
 
+import brainstate
 import brainunit as u
 import jax
 import jax.numpy as jnp
 
-import brainstate
-from brainpy import IF, LIF, ALIF
+from brainpy.state_based import IF, LIF, ALIF
 
 
 class TestNeuron(unittest.TestCase):

brainpy/state_based/_projection.py

@@ -17,18 +17,14 @@
 from typing import Optional
 
 import brainevent
-
 import brainstate
 from brainstate._state import State
 from brainstate.mixin import JointTypes, ParamDescriber
 from brainstate.nn._dynamics import maybe_init_prefetch
-from ._synouts import SynOut
-from .mixin import BindCondData, AlignPost
+from brainstate.util import get_unique_name
 
-if brainstate.__version__ < '0.2.0':
-    from brainstate.util.others import get_unique_name
-else:
-    from brainstate.util import get_unique_name
+from brainpy.mixin import BindCondData, AlignPost
+from ._synouts import SynOut
 
 __all__ = [
     'Projection',

brainpy/state_based/_readout.py

@@ -19,12 +19,12 @@
 import numbers
 from typing import Callable
 
+import brainstate
 import braintools
 import brainunit as u
 import jax
-
-import brainstate
 from brainstate.typing import Size, ArrayLike
+
 from ._base import Neuron
 
 __all__ = [

brainpy/state_based/_readout_test.py

@@ -16,12 +16,12 @@
 
 import unittest
 
+import brainstate
 import braintools
 import brainunit as u
 import jax.numpy as jnp
 
-import brainpy
-import brainstate
+import brainpy.state_based as brainpy
 
 
 class TestReadoutModels(unittest.TestCase):

brainpy/state_based/_stp.py

@@ -17,11 +17,11 @@
 
 from typing import Optional
 
+import brainstate
 import braintools
 import brainunit as u
-
-import brainstate
 from brainstate.typing import ArrayLike, Size
+
 from ._base import Synapse
 
 __all__ = [