refactor: Reorganize import statements for improved structure and readability

Author: chaoming0625

brainpy/__init__.py

@@ -3,38 +3,31 @@
 
 __version__ = "3.0.0"
 
+#  Part: Toolbox  #
+# --------------- #
+# modules of toolbox
+from brainpy import (
+    connect,  # synaptic connection
+    initialize,  # weight initialization
+    optim,  # gradient descent optimizers
+    losses,  # loss functions
+    measure,  # methods for data analysis
+    inputs,  # methods for generating input currents
+    encoding,  # encoding schema
+    checkpoints,  # checkpoints
+    check,  # error checking
+    mixin,  # mixin classes
+    algorithms,  # online or offline training algorithms
+)
 # fundamental supporting modules
 from brainpy import errors, check, tools
 
-try:
-  import jaxlib
-  del jaxlib
-except ModuleNotFoundError:
-  raise ModuleNotFoundError(tools.jaxlib_install_info) from None
-
-#  Part: Math Foundation  #
-# ----------------------- #
-
 # math foundation
 from brainpy import math
 from .math import BrainPyObject
 
-#  Part: Toolbox  #
-# --------------- #
-# modules of toolbox
-from brainpy import (
-  connect,  # synaptic connection
-  initialize,  # weight initialization
-  optim,  # gradient descent optimizers
-  losses,  # loss functions
-  measure,  # methods for data analysis
-  inputs,  # methods for generating input currents
-  encoding,  # encoding schema
-  checkpoints,  # checkpoints
-  check,  # error checking
-  mixin,  # mixin classes
-  algorithms,  # online or offline training algorithms
-)
+#  Part: Math Foundation  #
+# ----------------------- #
 
 # convenient alias
 conn = connect
@@ -50,34 +43,35 @@
 from brainpy._src.integrators.sde.generic import (sdeint as sdeint)
 from brainpy._src.integrators.fde.generic import (fdeint as fdeint)
 
-
 #  Part: Models  #
 # -------------- #
 
 # base classes
 from brainpy._src.dynsys import (
-  DynamicalSystem as DynamicalSystem,
-  DynSysGroup as DynSysGroup,  # collectors
-  Sequential as Sequential,
-  Dynamic as Dynamic,  # category
-  Projection as Projection,
-  receive_update_input,  # decorators
-  receive_update_output,
-  not_receive_update_input,
-  not_receive_update_output,
+    DynamicalSystem as DynamicalSystem,
+    DynSysGroup as DynSysGroup,  # collectors
+    Sequential as Sequential,
+    Dynamic as Dynamic,  # category
+    Projection as Projection,
+    receive_update_input,  # decorators
+    receive_update_output,
+    not_receive_update_input,
+    not_receive_update_output,
 )
+
 DynamicalSystemNS = DynamicalSystem
 Network = DynSysGroup
 # delays
 from brainpy._src.delay import (
-  VarDelay as VarDelay,
+    VarDelay as VarDelay,
 )
 
 # building blocks
 from brainpy import (
-  dnn, layers,  # module for dnn layers
-  dyn,  # module for modeling dynamics
+    dnn, layers,  # module for dnn layers
+    dyn,  # module for modeling dynamics
 )
+
 NeuGroup = NeuGroupNS = dyn.NeuDyn
 
 # common tools
@@ -94,60 +88,58 @@
 from brainpy._src.transform import (LoopOverTime as LoopOverTime, )
 from brainpy import (running as running)
 
-
 #  Part: Training  #
 # ---------------- #
 from brainpy._src.train.base import (DSTrainer as DSTrainer, )
 from brainpy._src.train.back_propagation import (BPTT as BPTT,
-                                                 BPFF as BPFF,)
+                                                 BPFF as BPFF, )
 from brainpy._src.train.online import (OnlineTrainer as OnlineTrainer,
                                        ForceTrainer as ForceTrainer, )
 from brainpy._src.train.offline import (OfflineTrainer as OfflineTrainer,
                                         RidgeTrainer as RidgeTrainer, )
 
-
 #  Part: Analysis  #
 # ---------------- #
 from brainpy import (analysis as analysis)
 
-
 #  Part: Others    #
 # ---------------- #
 from brainpy._src.visualization import (visualize as visualize)
 
-
 #  Part: Deprecations  #
 # -------------------- #
 from brainpy._src import train
 from brainpy import (
-  channels,  # channel models
-  neurons,  # neuron groups
-  synapses,  # synapses
-  rates,  # rate models
-  experimental,
-  synouts,  # synaptic output
-  synplast,  # synaptic plasticity
+    channels,  # channel models
+    neurons,  # neuron groups
+    synapses,  # synapses
+    rates,  # rate models
+    experimental,
+    synouts,  # synaptic output
+    synplast,  # synaptic plasticity
 )
 from brainpy._src import modes
-from brainpy._src.math.object_transform.base import (Base as Base,
-                                                     ArrayCollector as ArrayCollector,
-                                                     Collector as Collector, )
+from brainpy._src.math.object_transform.base import (
+    Base as Base,
+    ArrayCollector as ArrayCollector,
+    Collector as Collector,
+)
 
 # deprecated
 from brainpy._add_deprecations import deprecation_getattr2
 
 __deprecations = {
-  'Module': ('brainpy.Module', 'brainpy.DynamicalSystem', DynamicalSystem),
-  'Channel': ('brainpy.Channel', 'brainpy.dyn.IonChannel', dyn.IonChannel),
-  'SynConn': ('brainpy.SynConn', 'brainpy.dyn.SynConn', dyn.SynConn),
-  'Container': ('brainpy.Container', 'brainpy.DynSysGroup', DynSysGroup),
-
-  'optimizers': ('brainpy.optimizers', 'brainpy.optim', optim),
-  'TensorCollector': ('brainpy.TensorCollector', 'brainpy.ArrayCollector', ArrayCollector),
-  'SynSTP': ('brainpy.SynSTP', 'brainpy.synapses.SynSTP', synapses.SynSTP),
-  'SynOut': ('brainpy.SynOut', 'brainpy.synapses.SynOut', synapses.SynOut),
-  'TwoEndConn': ('brainpy.TwoEndConn', 'brainpy.synapses.TwoEndConn', synapses.TwoEndConn),
-  'CondNeuGroup': ('brainpy.CondNeuGroup', 'brainpy.dyn.CondNeuGroup', dyn.CondNeuGroup),
+    'Module': ('brainpy.Module', 'brainpy.DynamicalSystem', DynamicalSystem),
+    'Channel': ('brainpy.Channel', 'brainpy.dyn.IonChannel', dyn.IonChannel),
+    'SynConn': ('brainpy.SynConn', 'brainpy.dyn.SynConn', dyn.SynConn),
+    'Container': ('brainpy.Container', 'brainpy.DynSysGroup', DynSysGroup),
+
+    'optimizers': ('brainpy.optimizers', 'brainpy.optim', optim),
+    'TensorCollector': ('brainpy.TensorCollector', 'brainpy.ArrayCollector', ArrayCollector),
+    'SynSTP': ('brainpy.SynSTP', 'brainpy.synapses.SynSTP', synapses.SynSTP),
+    'SynOut': ('brainpy.SynOut', 'brainpy.synapses.SynOut', synapses.SynOut),
+    'TwoEndConn': ('brainpy.TwoEndConn', 'brainpy.synapses.TwoEndConn', synapses.TwoEndConn),
+    'CondNeuGroup': ('brainpy.CondNeuGroup', 'brainpy.dyn.CondNeuGroup', dyn.CondNeuGroup),
 }
 __getattr__ = deprecation_getattr2('brainpy', __deprecations)