[refactor] Improve JAX configuration handling and clean up setup.py

Author: chaoming0625

README.md

@@ -19,6 +19,7 @@ BrainPy is a flexible, efficient, and extensible framework for computational neu
 - **Website (documentation and APIs)**: https://brainpy.readthedocs.io/en/latest
 - **Source**: https://github.com/brainpy/BrainPy
 - **Bug reports**: https://github.com/brainpy/BrainPy/issues
+- **Ecosystem**: https://brainmodeling.readthedocs.io/
 - **Source on OpenI**: https://git.openi.org.cn/OpenI/BrainPy
 
 

brainpy/__init__.py

@@ -153,8 +153,11 @@
 
 del deprecation_getattr2
 
-# jax config
-import os
-os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false'
-import jax
-jax.config.update('jax_cpu_enable_async_dispatch', False)
+try:
+    # jax config
+    import os
+    os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false'
+    import jax
+    jax.config.update('jax_cpu_enable_async_dispatch', False)
+except:
+    pass

examples/dynamics_simulation/hh_model.py

@@ -1,9 +1,6 @@
 # -*- coding: utf-8 -*-
 
-import numpy as np
-
 import brainpy as bp
-from jax import pmap
 import brainpy.math as bm
 
 bm.set_host_device_count(20)

examples/dynamics_simulation/stdp.py

@@ -13,52 +13,52 @@
 
 
 class STDPNet(bp.DynSysGroup):
-  def __init__(self, num_poisson, num_lif=1, g_max=0.01):
-    super().__init__()
+    def __init__(self, num_poisson, num_lif=1, g_max=0.01):
+        super().__init__()
 
-    self.g_max = g_max
+        self.g_max = g_max
 
-    # neuron groups
-    self.noise = bp.dyn.PoissonGroup(num_poisson, freqs=15.)
-    self.group = bp.dyn.Lif(num_lif, V_reset=-60., V_rest=-74, V_th=-54, tau=10.,
-                            V_initializer=bp.init.Normal(-60., 1.))
+        # neuron groups
+        self.noise = bp.dyn.PoissonGroup(num_poisson, freqs=15.)
+        self.group = bp.dyn.Lif(num_lif, V_reset=-60., V_rest=-74, V_th=-54, tau=10.,
+                                V_initializer=bp.init.Normal(-60., 1.))
 
-    # synapses
-    syn = bp.dyn.Expon.desc(num_lif, tau=5.)
-    out = bp.dyn.COBA.desc(E=0.)
-    comm = bp.dnn.AllToAll(num_poisson, num_lif, bp.init.Uniform(0., g_max))
-    self.syn = bp.dyn.STDP_Song2000(self.noise, None, syn, comm, out, self.group,
-                                    tau_s=20, tau_t=20, W_max=g_max, W_min=0.,
-                                    A1=0.01 * g_max, A2=0.0105 * g_max)
+        # synapses
+        syn = bp.dyn.Expon.desc(num_lif, tau=5.)
+        out = bp.dyn.COBA.desc(E=0.)
+        comm = bp.dnn.AllToAll(num_poisson, num_lif, bp.init.Uniform(0., g_max))
+        self.syn = bp.dyn.STDP_Song2000(self.noise, None, syn, comm, out, self.group,
+                                        tau_s=20, tau_t=20, W_max=g_max, W_min=0.,
+                                        A1=0.01 * g_max, A2=0.0105 * g_max)
 
-  def update(self, *args, **kwargs):
-    self.noise()
-    self.syn()
-    self.group()
-    return self.syn.comm.weight.flatten()[:10]
+    def update(self, *args, **kwargs):
+        self.noise()
+        self.syn()
+        self.group()
+        return self.syn.comm.weight.flatten()[:10]
 
 
 def run_model():
-  net = STDPNet(1000, 1)
-  indices = np.arange(int(100.0e3 / bm.dt))  # 100 s
-  ws = bm.for_loop(net.step_run, indices, progress_bar=True)
-  weight = bm.as_numpy(net.syn.comm.weight.flatten())
+    net = STDPNet(1000, 1)
+    indices = np.arange(int(100.0e3 / bm.dt))  # 100 s
+    ws = bm.for_loop(net.step_run, indices, progress_bar=True)
+    weight = bm.as_numpy(net.syn.comm.weight.flatten())
 
-  fig, gs = bp.visualize.get_figure(3, 1, 3, 10)
-  fig.add_subplot(gs[0, 0])
-  plt.plot(weight / net.g_max, '.k')
-  plt.xlabel('Weight / gmax')
+    fig, gs = bp.visualize.get_figure(3, 1, 3, 10)
+    fig.add_subplot(gs[0, 0])
+    plt.plot(weight / net.g_max, '.k')
+    plt.xlabel('Weight / gmax')
 
-  fig.add_subplot(gs[1, 0])
-  plt.hist(weight / net.g_max, 20)
-  plt.xlabel('Weight / gmax')
+    fig.add_subplot(gs[1, 0])
+    plt.hist(weight / net.g_max, 20)
+    plt.xlabel('Weight / gmax')
 
-  fig.add_subplot(gs[2, 0])
-  plt.plot(indices * bm.dt, bm.as_numpy(ws) / net.g_max)
-  plt.xlabel('Time (s)')
-  plt.ylabel('Weight / gmax')
-  plt.show()
+    fig.add_subplot(gs[2, 0])
+    plt.plot(indices * bm.dt, bm.as_numpy(ws) / net.g_max)
+    plt.xlabel('Time (s)')
+    plt.ylabel('Weight / gmax')
+    plt.show()
 
 
 if __name__ == '__main__':
-  run_model()
+    run_model()

setup.py

@@ -3,96 +3,68 @@
 import io
 import os
 import re
-import time
-import sys
-
-from setuptools import find_packages
-from setuptools import setup
-
-try:
-  # require users to uninstall previous brainpy releases.
-  import pkg_resources
-
-  installed_packages = pkg_resources.working_set
-  for i in installed_packages:
-    if i.key == 'brainpy-simulator':
-      raise SystemError('Please uninstall the older version of brainpy '
-                        f'package "brainpy-simulator={i.version}" '
-                        f'(located in {i.location}) first. \n'
-                        '>>> pip uninstall brainpy-simulator')
-    if i.key == 'brain-py':
-      raise SystemError('Please uninstall the older version of brainpy '
-                        f'package "brain-py={i.version}" '
-                        f'(located in {i.location}) first. \n'
-                        '>>> pip uninstall brain-py')
-except ModuleNotFoundError:
-  pass
 
+from setuptools import find_packages, setup
 
 # version
 here = os.path.abspath(os.path.dirname(__file__))
 with open(os.path.join(here, 'brainpy', '__init__.py'), 'r') as f:
-  init_py = f.read()
+    init_py = f.read()
 version = re.search('__version__ = "(.*)"', init_py).groups()[0]
-if len(sys.argv) > 2 and sys.argv[2] == '--python-tag=py3':
-  version = version
-else:
-  version += '.post{}'.format(time.strftime("%Y%m%d", time.localtime()))
 
 # obtain long description from README
 with io.open(os.path.join(here, 'README.md'), 'r', encoding='utf-8') as f:
-  README = f.read()
+    README = f.read()
 
 # installation packages
 packages = find_packages(exclude=['lib*', 'docs', 'tests'])
 
 # setup
 setup(
-  name='brainpy',
-  version=version,
-  description='BrainPy: Brain Dynamics Programming in Python',
-  long_description=README,
-  long_description_content_type="text/markdown",
-  author='BrainPy Team',
-  author_email='chao.brain@qq.com',
-  packages=packages,
-  python_requires='>=3.10',
-  install_requires=['numpy>=1.15', 'jax>=0.4.13,<0.6.0', 'tqdm'],
-  url='https://github.com/brainpy/BrainPy',
-  project_urls={
-    "Bug Tracker": "https://github.com/brainpy/BrainPy/issues",
-    "Documentation": "https://brainpy.readthedocs.io/",
-    "Source Code": "https://github.com/brainpy/BrainPy",
-  },
-  dependency_links=[
-    'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html',
-  ],
-  extras_require={
-    'cpu': ['jaxlib>=0.4.13', 'numba', 'braintaichi'],
-    'cuda12': ['jaxlib[cuda12_pip]', 'numba', 'braintaichi'],
-    'tpu': ['jaxlib[tpu]', 'numba',],
-    'cpu_mini': ['jaxlib>=0.4.13'],
-    'cuda12_mini': ['jaxlib[cuda12_pip]'],
-  },
-  keywords=('computational neuroscience, '
-            'brain-inspired computation, '
-            'brain modeling, '
-            'brain dynamics modeling, '
-            'brain dynamics programming'),
-  classifiers=[
-    'Natural Language :: English',
-    'Operating System :: OS Independent',
-    'Programming Language :: Python',
-    'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.10',
-    'Programming Language :: Python :: 3.11',
-    'Programming Language :: Python :: 3.12',
-    'Intended Audience :: Science/Research',
-    'License :: OSI Approved :: Apache Software License',
-    'Topic :: Scientific/Engineering :: Bio-Informatics',
-    'Topic :: Scientific/Engineering :: Mathematics',
-    'Topic :: Scientific/Engineering :: Artificial Intelligence',
-    'Topic :: Software Development :: Libraries',
-  ],
-  license='GPL-3.0 license',
+    name='brainpy',
+    version=version,
+    description='BrainPy: Brain Dynamics Programming in Python',
+    long_description=README,
+    long_description_content_type="text/markdown",
+    author='BrainPy Team',
+    author_email='chao.brain@qq.com',
+    packages=packages,
+    python_requires='>=3.10',
+    install_requires=['numpy>=1.15', 'jax', 'tqdm'],
+    url='https://github.com/brainpy/BrainPy',
+    project_urls={
+        "Bug Tracker": "https://github.com/brainpy/BrainPy/issues",
+        "Documentation": "https://brainpy.readthedocs.io/",
+        "Source Code": "https://github.com/brainpy/BrainPy",
+    },
+    dependency_links=[
+        'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html',
+    ],
+    extras_require={
+        'cpu': ['jax[cpu]<0.5.0', 'numba', 'braintaichi'],
+        'cuda12': ['jax[cuda12]<0.5.0', 'numba', 'braintaichi'],
+        'tpu': ['jax[tpu]<0.5.0', 'numba', ],
+    },
+    keywords=('computational neuroscience, '
+              'brain-inspired computation, '
+              'brain modeling, '
+              'brain dynamics modeling, '
+              'brain dynamics programming'),
+    classifiers=[
+        'Natural Language :: English',
+        'Operating System :: OS Independent',
+        'Programming Language :: Python',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
+        'Programming Language :: Python :: 3.12',
+        'Programming Language :: Python :: 3.13',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Topic :: Scientific/Engineering :: Bio-Informatics',
+        'Topic :: Scientific/Engineering :: Mathematics',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development :: Libraries',
+    ],
+    license='GPL-3.0 license',
 )