Refactor module assignments to 'brainpy.state' for consistency across files

Author: chaoming0625

brainpy/state/__init__.py

@@ -18,14 +18,14 @@
 from ._base import __all__ as base_all
 from ._exponential import *
 from ._exponential import __all__ as exp_all
+from ._hh import *
+from ._hh import __all__ as hh_all
 from ._inputs import *
 from ._inputs import __all__ as inputs_all
-from ._lif import *
-from ._lif import __all__ as neuron_all
 from ._izhikevich import *
 from ._izhikevich import __all__ as izh_all
-from ._hh import *
-from ._hh import __all__ as hh_all
+from ._lif import *
+from ._lif import __all__ as neuron_all
 from ._projection import *
 from ._projection import __all__ as proj_all
 from ._readout import *
@@ -38,13 +38,8 @@
 from ._synaptic_projection import __all__ as synproj_all
 from ._synouts import *
 from ._synouts import __all__ as synout_all
-from .. import mixin
 
-__main__ = ['version2', 'mixin'] + inputs_all + neuron_all + izh_all + hh_all + readout_all + stp_all + synapse_all
+__main__ = inputs_all + neuron_all + izh_all + hh_all + readout_all + stp_all + synapse_all
 __main__ = __main__ + synout_all + base_all + exp_all + proj_all + synproj_all
 del inputs_all, neuron_all, izh_all, hh_all, readout_all, stp_all, synapse_all, synout_all, base_all
 del exp_all, proj_all, synproj_all
-
-if __name__ == '__main__':
-    mixin
-

brainpy/state/_base.py

@@ -40,6 +40,8 @@ def _input_label_repr(name: str, label: Optional[str] = None):
 
 
 class Dynamics(brainstate.nn.Dynamics):
+    __module__ = 'brainpy.state'
+
     def __init__(self, in_size: Size, name: Optional[str] = None):
         # initialize
         super().__init__(name=name, in_size=in_size)
@@ -595,7 +597,7 @@ class Neuron(Dynamics):
     .. [3] Gerstner, W., Kistler, W. M., Naud, R., & Paninski, L. (2014). Neuronal dynamics:
            From single neurons to networks and models of cognition. Cambridge University Press.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -849,4 +851,4 @@ class Synapse(Dynamics):
     .. [3] Gerstner, W., Kistler, W. M., Naud, R., & Paninski, L. (2014). Neuronal dynamics:
            From single neurons to networks and models of cognition. Cambridge University Press.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'

brainpy/state/_exponential.py

@@ -71,7 +71,7 @@ class Expon(Synapse, AlignPost):
     where synaptic variables are aligned with post-synaptic neurons, enabling event-driven
     computation and more efficient handling of sparse connectivity patterns.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -156,7 +156,7 @@ class DualExpon(Synapse, AlignPost):
     where synaptic variables are aligned with post-synaptic neurons, enabling event-driven
     computation and more efficient handling of sparse connectivity patterns.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,

brainpy/state/_hh.py

@@ -148,7 +148,7 @@ class HH(Neuron):
     .. [2] https://en.wikipedia.org/wiki/Hodgkin%E2%80%93Huxley_model
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -380,7 +380,7 @@ class MorrisLecar(Neuron):
     .. [3] https://en.wikipedia.org/wiki/Morris%E2%80%93Lecar_model
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -572,7 +572,7 @@ class WangBuzsakiHH(Neuron):
            neuroscience, 16(20), pp.6402-6413.
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,

brainpy/state/_inputs.py

@@ -60,7 +60,7 @@ class SpikeTime(brainstate.nn.Dynamics):
     name : str, optional
         The name of the dynamic system.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -116,7 +116,7 @@ class PoissonSpike(brainstate.nn.Dynamics):
     """
     Poisson Neuron Group.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -203,7 +203,7 @@ class PoissonEncoder(brainstate.nn.Dynamics):
     - The independence of spike generation between time steps results in renewal process
       statistics without memory of previous spiking history.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -316,7 +316,7 @@ class PoissonInput(brainstate.nn.Module):
     - The update method internally calls the poisson_input function which handles the
       spike generation and target state updates.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -348,7 +348,7 @@ def update(self):
         )
 
 
-@set_module_as('brainpy')
+@set_module_as('brainpy.state')
 def poisson_input(
     freq: u.Quantity[u.Hz],
     num_input: int,

brainpy/state/_izhikevich.py

@@ -119,7 +119,7 @@ class Izhikevich(Neuron):
            IEEE transactions on neural networks, 15(5), 1063-1070.
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -302,7 +302,7 @@ class IzhikevichRef(Neuron):
            IEEE transactions on neural networks, 15(5), 1063-1070.
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,

brainpy/state/_lif.py

@@ -113,7 +113,7 @@ class IF(Neuron):
            I. Homogeneous synaptic input. Biological cybernetics, 95(1), 1-19.
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -237,7 +237,7 @@ class LIF(Neuron):
     .. [2] Burkitt, A. N. (2006). A review of the integrate-and-fire neuron model:
            I. Homogeneous synaptic input. Biological cybernetics, 95(1), 1-19.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -380,7 +380,7 @@ class ExpIF(Neuron):
            conductance-based synaptic drive." Physical Review E 76, no. 2 (2007): 021919.
     .. [5] https://en.wikipedia.org/wiki/Exponential_integrate-and-fire
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -506,7 +506,7 @@ class ExpIFRef(Neuron):
     >>>         with brainstate.environ.context(t=t*0.1*u.ms):
     >>>             spikes = expif.update(x=inputs[t])
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -680,7 +680,7 @@ class AdExIF(Neuron):
 
        :class:`brainpy.dyn.AdExIF` for the dynamical-system counterpart.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -863,7 +863,7 @@ class AdExIFRef(Neuron):
 
        :class:`brainpy.dyn.AdExIFRef` for the dynamical-system counterpart.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -1072,7 +1072,7 @@ class LIFRef(Neuron):
     .. [3] Izhikevich, E. M. (2003). Simple model of spiking neurons. IEEE Transactions on
            neural networks, 14(6), 1569-1572.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -1244,7 +1244,7 @@ class ALIF(Neuron):
            the adaptive exponential integrate-and-fire model. Biological cybernetics,
            99(4), 335-347.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -1390,7 +1390,7 @@ class QuaIF(Neuron):
            (2000) Intrinsic dynamics in neuronal networks. I. Theory.
            J. Neurophysiology 83, pp. 808–827.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -1545,7 +1545,7 @@ class AdQuaIF(Neuron):
            nonlinear integrate-and-fire neurons." SIAM Journal on Applied
            Mathematics 68, no. 4 (2008): 1045-1079.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -1706,7 +1706,7 @@ class AdQuaIFRef(Neuron):
     - Refractory period prevents unrealistically high firing rates.
     - More biologically realistic than AdQuaIF without refractory period.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -1939,7 +1939,7 @@ class Gif(Neuron):
            leaky integrate-and-fire models classify multiple neuron types."
            Nature communications 9, no. 1 (2018): 1-15.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -2151,7 +2151,7 @@ class GifRef(Neuron):
     - Can still exhibit diverse firing patterns: regular, bursting, adaptation.
     - Refractory period prevents unrealistically high firing rates.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,

brainpy/state/_projection.py

@@ -71,7 +71,7 @@ class Projection(brainstate.nn.Module):
     Derived classes should implement specific projection behaviors, such as
     dense connectivity, sparse connectivity, or specific weight update rules.
     """
-    __module__ = 'brainstate.nn'
+    __module__ = 'brainpy.state'
 
     def update(self, *args, **kwargs):
         sub_nodes = tuple(self.nodes(allowed_hierarchy=(1, 1)).values())
@@ -181,7 +181,7 @@ class AlignPostProj(Projection):
     >>> exe_current = E(pop.get_spike())
 
     """
-    __module__ = 'brainstate.nn'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -303,7 +303,7 @@ class DeltaProj(Projection):
     ... )
     >>> delta_input(1.0)  # Apply voltage increment directly
     """
-    __module__ = 'brainstate.nn'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -380,7 +380,7 @@ class CurrentProj(Projection):
     ... )
     >>> current_input(0.2)  # Apply external current
     """
-    __module__ = 'brainstate.nn'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -439,6 +439,7 @@ class align_pre_projection(Projection):
     short-term plasticity.
 
     """
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -486,6 +487,7 @@ class align_post_projection(Projection):
     properties, synaptic outputs, post-synaptic dynamics, and short-term plasticity.
 
     """
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,

brainpy/state/_readout.py

@@ -74,7 +74,7 @@ class LeakyRateReadout(brainstate.nn.Module):
     r : HiddenState
         Hidden state representing the output values
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -160,7 +160,7 @@ class LeakySpikeReadout(Neuron):
         Synaptic weight matrix
     """
 
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,

brainpy/state/_stp.py

@@ -97,7 +97,7 @@ class STP(Synapse):
     .. [2] Tsodyks, M., Pawelzik, K., & Markram, H. (1998). Neural networks with dynamic
            synapses. Neural computation, 10(4), 821-835.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,
@@ -203,7 +203,7 @@ class STD(Synapse):
            pyramidal neurons depends on neurotransmitter release probability.
            Proceedings of the National Academy of Sciences, 94(2), 719-723.
     """
-    __module__ = 'brainpy'
+    __module__ = 'brainpy.state'
 
     def __init__(
         self,