brainunit 0.1.2__tar.gz → 0.1.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. {brainunit-0.1.2 → brainunit-0.1.4}/PKG-INFO +177 -177
  2. {brainunit-0.1.2 → brainunit-0.1.4}/README.md +133 -133
  3. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/__init__.py +65 -64
  4. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/_base.py +20 -20
  5. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/_celsius.py +20 -20
  6. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/_unit_common.py +19 -19
  7. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/_unit_constants.py +20 -20
  8. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/_unit_shortcuts.py +20 -20
  9. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/autograd/__init__.py +30 -30
  10. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/autograd/_hessian.py +19 -19
  11. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/autograd/_jacobian.py +19 -19
  12. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/autograd/_value_and_grad.py +19 -19
  13. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/autograd/_vector_grad.py +19 -19
  14. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/constants.py +20 -20
  15. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/fft/__init__.py +25 -25
  16. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/fft/_fft_change_unit.py +19 -19
  17. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/fft/_fft_keep_unit.py +20 -20
  18. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/__init__.py +45 -45
  19. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_lax_accept_unitless.py +19 -19
  20. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_lax_array_creation.py +19 -19
  21. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_lax_change_unit.py +20 -20
  22. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_lax_keep_unit.py +20 -20
  23. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_lax_linalg.py +20 -20
  24. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_lax_remove_unit.py +20 -20
  25. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/lax/_misc.py +20 -20
  26. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/linalg/__init__.py +24 -24
  27. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/linalg/_linalg_change_unit.py +19 -19
  28. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/linalg/_linalg_keep_unit.py +19 -19
  29. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/linalg/_linalg_remove_unit.py +20 -20
  30. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/__init__.py +61 -61
  31. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_activation.py +20 -20
  32. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_alias.py +20 -20
  33. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_einops.py +20 -20
  34. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_fun_accept_unitless.py +20 -20
  35. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_fun_array_creation.py +20 -20
  36. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_fun_change_unit.py +20 -20
  37. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_fun_keep_unit.py +20 -20
  38. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_fun_remove_unit.py +20 -20
  39. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/_misc.py +20 -20
  40. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/fft.py +20 -20
  41. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/math/linalg.py +20 -20
  42. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/sparse/__init__.py +22 -22
  43. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/sparse/_block_csr.py +20 -20
  44. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/sparse/_block_ell.py +20 -20
  45. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/sparse/_coo.py +20 -20
  46. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit/sparse/_csr.py +20 -20
  47. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit.egg-info/PKG-INFO +177 -177
  48. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit.egg-info/requires.txt +1 -1
  49. {brainunit-0.1.2 → brainunit-0.1.4}/pyproject.toml +68 -68
  50. {brainunit-0.1.2 → brainunit-0.1.4}/setup.cfg +4 -4
  51. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit.egg-info/SOURCES.txt +0 -0
  52. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit.egg-info/dependency_links.txt +0 -0
  53. {brainunit-0.1.2 → brainunit-0.1.4}/brainunit.egg-info/top_level.txt +0 -0
@@ -1,177 +1,177 @@
1
- Metadata-Version: 2.4
2
- Name: brainunit
3
- Version: 0.1.2
4
- Summary: Physical Units and Unit-Aware Mathematical System for General-Purpose Brain Dynamics Modeling
5
- Author-email: BrainUnit Developers <chao.brain@qq.com>
6
- License-Expression: Apache-2.0
7
- Project-URL: homepage, https://github.com/chaobrain/brainunit
8
- Project-URL: repository, https://github.com/chaobrain/brainunit
9
- Project-URL: Bug Tracker, https://github.com/chaobrain/brainunit/issues
10
- Project-URL: Documentation, https://brainunit.readthedocs.io/
11
- Keywords: physical unit,physical quantity,brain modeling
12
- Classifier: Natural Language :: English
13
- Classifier: Operating System :: OS Independent
14
- Classifier: Development Status :: 4 - Beta
15
- Classifier: Intended Audience :: Developers
16
- Classifier: Intended Audience :: Science/Research
17
- Classifier: Programming Language :: Python
18
- Classifier: Programming Language :: Python :: 3
19
- Classifier: Programming Language :: Python :: 3.10
20
- Classifier: Programming Language :: Python :: 3.11
21
- Classifier: Programming Language :: Python :: 3.12
22
- Classifier: Programming Language :: Python :: 3.13
23
- Classifier: Programming Language :: Python :: 3.14
24
- Classifier: Programming Language :: Python
25
- Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
26
- Classifier: Topic :: Scientific/Engineering :: Mathematics
27
- Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
28
- Classifier: Topic :: Software Development :: Libraries
29
- Requires-Python: >=3.10
30
- Description-Content-Type: text/markdown
31
- Requires-Dist: numpy
32
- Requires-Dist: typing_extensions
33
- Requires-Dist: saiunit==0.1.2
34
- Provides-Extra: testing
35
- Requires-Dist: pytest; extra == "testing"
36
- Provides-Extra: cpu
37
- Requires-Dist: jax[cpu]; extra == "cpu"
38
- Provides-Extra: cuda12
39
- Requires-Dist: jax[cuda12]; extra == "cuda12"
40
- Provides-Extra: cuda13
41
- Requires-Dist: jax[cuda13]; extra == "cuda13"
42
- Provides-Extra: tpu
43
- Requires-Dist: jax[tpu]; extra == "tpu"
44
-
45
- <h1 align='center'>BrainUnit</h1>
46
- <h2 align='center'>Physical units and unit-aware math system for general-purpose brain dynamics modeling</h2>
47
-
48
- <p align="center">
49
- <img alt="Header image of brainunit." src="https://raw.githubusercontent.com/chaobrain/saiunit/main/docs/_static/brainunit.png" width=30%>
50
- </p>
51
-
52
-
53
-
54
- <p align="center">
55
- <a href="https://pypi.org/project/brainunit/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainunit"></a>
56
- <a href="https://github.com/chaobrain/brainunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
57
- <a href='https://brainunit.readthedocs.io/?badge=latest'>
58
- <img src='https://readthedocs.org/projects/brainunit/badge/?version=latest' alt='Documentation Status' />
59
- </a>
60
- <a href="https://badge.fury.io/py/brainunit"><img alt="PyPI version" src="https://badge.fury.io/py/brainunit.svg"></a>
61
- <a href="https://pepy.tech/projects/brainunit"><img src="https://static.pepy.tech/badge/brainunit" alt="PyPI Downloads"></a>
62
- </p>
63
-
64
-
65
-
66
-
67
- [BrainUnit](https://github.com/chaobrain/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics modeling. It introduces rigoirous physical units into high-performance AI-driven abstract numerical computing.
68
-
69
- BrainUnit is initially designed to enable unit-aware computations in brain dynamics modeling (see our [BDP ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/)). However, its features and capacities can be applied to general domains in scientific computing and AI4Science. Starting in 2025/02, BrainUnit has been fully integrated into [SAIUnit](https://github.com/chaobrain/saiunit) (the **Unit** system for **S**cientific **AI**).
70
-
71
- Functionalities are the same for both ``brainunit`` and ``saiunit``, and their functions and data structures are interoperable, sharing the same set of APIs, and eliminating any potential conflicts. This meas that
72
-
73
- ```python
74
- import brainunit as u
75
- ```
76
-
77
- equals to
78
-
79
- ```python
80
- import saiunit as u
81
- ```
82
-
83
- For users primarily engaged in general scientific computing, `saiunit` is likely the preferred choice. However, for those focused on brain modeling, we recommend `brainunit`, as it is more closely aligned with our specialized brain dynamics programming ecosystem.
84
-
85
-
86
-
87
- ## Documentation
88
-
89
- The official documentation of BrainUnit is hosted on Read the Docs: [https://brainunit.readthedocs.io](https://brainunit.readthedocs.io)
90
-
91
-
92
-
93
- ## Features
94
-
95
- `brainunit` can be seamlessly integrated into every aspect of our [brain dynamics programming ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/), such as, the checkpointing of [braintools](https://github.com/chaobrain/braintools), the event-driven operators in [brainevent](https://github.com/chaobrain/brainevent), the state-based JIT compilation in [brainstate](https://github.com/chaobrain/brainstate), online learning rules in [brainscale](https://github.com/chaobrain/brainscale), or event more.
96
-
97
- A quick example for this kind of integration:
98
-
99
- ```python
100
-
101
- import braintools
102
- import brainevent
103
- import brainstate
104
- import brainunit as u
105
-
106
-
107
- class EINet(brainstate.nn.Module):
108
- def __init__(self):
109
- super().__init__()
110
- self.n_exc = 3200
111
- self.n_inh = 800
112
- self.num = self.n_exc + self.n_inh
113
- self.N = brainstate.nn.LIFRef(
114
- self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
115
- tau=20. * u.ms, tau_ref=5. * u.ms,
116
- V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
117
- )
118
- self.E = brainstate.nn.AlignPostProj(
119
- comm=brainstate.nn.EventFixedProb(self.n_exc, self.num, 0.02, 0.6 * u.mS),
120
- syn=brainstate.nn.Expon.desc(self.num, tau=5. * u.ms),
121
- out=brainstate.nn.COBA.desc(E=0. * u.mV),
122
- post=self.N
123
- )
124
- self.I = brainstate.nn.AlignPostProj(
125
- comm=brainstate.nn.EventFixedProb(self.n_inh, self.num, 0.02, 6.7 * u.mS),
126
- syn=brainstate.nn.Expon.desc(self.num, tau=10. * u.ms),
127
- out=brainstate.nn.COBA.desc(E=-80. * u.mV),
128
- post=self.N
129
- )
130
-
131
- def update(self, t, inp):
132
- with brainstate.environ.context(t=t):
133
- spk = self.N.get_spike() != 0.
134
- self.E(spk[:self.n_exc])
135
- self.I(spk[self.n_exc:])
136
- self.N(inp)
137
- return self.N.get_spike()
138
-
139
- def save_checkpoint(self):
140
- braintools.file.msgpack_save('states.msgpack', self.states())
141
-
142
- ```
143
-
144
-
145
-
146
- ## Installation
147
-
148
- You can install ``brainunit`` via pip:
149
-
150
- ```bash
151
- pip install brainunit --upgrade
152
- ```
153
-
154
-
155
- ## Citation
156
-
157
- If you use `brainunit` in your research, please consider citing the following paper:
158
-
159
- ```bibtex
160
- @article{wang2025integrating,
161
- title={Integrating physical units into high-performance AI-driven scientific computing},
162
- author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},
163
- journal={Nature Communications},
164
- volume={16},
165
- number={1},
166
- pages={3609},
167
- year={2025},
168
- publisher={Nature Publishing Group UK London},
169
- url={https://doi.org/10.1038/s41467-025-58626-4}
170
- }
171
- ```
172
-
173
-
174
-
175
- ## See also the BDP ecosystem
176
-
177
- We are building the [brain dynamics programming (BDP) ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/). [brainunit](https://github.com/chaobrain/brainunit) has been deeply integrated into our BDP ecosystem.
1
+ Metadata-Version: 2.4
2
+ Name: brainunit
3
+ Version: 0.1.4
4
+ Summary: Physical Units and Unit-Aware Mathematical System for General-Purpose Brain Dynamics Modeling
5
+ Author-email: BrainUnit Developers <chao.brain@qq.com>
6
+ License-Expression: Apache-2.0
7
+ Project-URL: homepage, https://github.com/chaobrain/brainunit
8
+ Project-URL: repository, https://github.com/chaobrain/brainunit
9
+ Project-URL: Bug Tracker, https://github.com/chaobrain/brainunit/issues
10
+ Project-URL: Documentation, https://brainunit.readthedocs.io/
11
+ Keywords: physical unit,physical quantity,brain modeling
12
+ Classifier: Natural Language :: English
13
+ Classifier: Operating System :: OS Independent
14
+ Classifier: Development Status :: 4 - Beta
15
+ Classifier: Intended Audience :: Developers
16
+ Classifier: Intended Audience :: Science/Research
17
+ Classifier: Programming Language :: Python
18
+ Classifier: Programming Language :: Python :: 3
19
+ Classifier: Programming Language :: Python :: 3.10
20
+ Classifier: Programming Language :: Python :: 3.11
21
+ Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Programming Language :: Python :: 3.13
23
+ Classifier: Programming Language :: Python :: 3.14
24
+ Classifier: Programming Language :: Python
25
+ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
26
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
27
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
28
+ Classifier: Topic :: Software Development :: Libraries
29
+ Requires-Python: >=3.10
30
+ Description-Content-Type: text/markdown
31
+ Requires-Dist: numpy
32
+ Requires-Dist: typing_extensions
33
+ Requires-Dist: saiunit==0.1.4
34
+ Provides-Extra: testing
35
+ Requires-Dist: pytest; extra == "testing"
36
+ Provides-Extra: cpu
37
+ Requires-Dist: jax[cpu]; extra == "cpu"
38
+ Provides-Extra: cuda12
39
+ Requires-Dist: jax[cuda12]; extra == "cuda12"
40
+ Provides-Extra: cuda13
41
+ Requires-Dist: jax[cuda13]; extra == "cuda13"
42
+ Provides-Extra: tpu
43
+ Requires-Dist: jax[tpu]; extra == "tpu"
44
+
45
+ <h1 align='center'>BrainUnit</h1>
46
+ <h2 align='center'>Physical units and unit-aware math system for general-purpose brain dynamics modeling</h2>
47
+
48
+ <p align="center">
49
+ <img alt="Header image of brainunit." src="https://raw.githubusercontent.com/chaobrain/saiunit/main/docs/_static/brainunit.png" width=30%>
50
+ </p>
51
+
52
+
53
+
54
+ <p align="center">
55
+ <a href="https://pypi.org/project/brainunit/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainunit"></a>
56
+ <a href="https://github.com/chaobrain/brainunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
57
+ <a href='https://brainunit.readthedocs.io/?badge=latest'>
58
+ <img src='https://readthedocs.org/projects/brainunit/badge/?version=latest' alt='Documentation Status' />
59
+ </a>
60
+ <a href="https://badge.fury.io/py/brainunit"><img alt="PyPI version" src="https://badge.fury.io/py/brainunit.svg"></a>
61
+ <a href="https://pepy.tech/projects/brainunit"><img src="https://static.pepy.tech/badge/brainunit" alt="PyPI Downloads"></a>
62
+ </p>
63
+
64
+
65
+
66
+
67
+ [BrainUnit](https://github.com/chaobrain/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics modeling. It introduces rigoirous physical units into high-performance AI-driven abstract numerical computing.
68
+
69
+ BrainUnit is initially designed to enable unit-aware computations in brain dynamics modeling (see our [BDP ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/)). However, its features and capacities can be applied to general domains in scientific computing and AI4Science. Starting in 2025/02, BrainUnit has been fully integrated into [SAIUnit](https://github.com/chaobrain/saiunit) (the **Unit** system for **S**cientific **AI**).
70
+
71
+ Functionalities are the same for both ``brainunit`` and ``saiunit``, and their functions and data structures are interoperable, sharing the same set of APIs, and eliminating any potential conflicts. This meas that
72
+
73
+ ```python
74
+ import brainunit as u
75
+ ```
76
+
77
+ equals to
78
+
79
+ ```python
80
+ import saiunit as u
81
+ ```
82
+
83
+ For users primarily engaged in general scientific computing, `saiunit` is likely the preferred choice. However, for those focused on brain modeling, we recommend `brainunit`, as it is more closely aligned with our specialized brain dynamics programming ecosystem.
84
+
85
+
86
+
87
+ ## Documentation
88
+
89
+ The official documentation of BrainUnit is hosted on Read the Docs: [https://brainunit.readthedocs.io](https://brainunit.readthedocs.io)
90
+
91
+
92
+
93
+ ## Features
94
+
95
+ `brainunit` can be seamlessly integrated into every aspect of our [brain modeling ecosystem](https://brainmodeling.readthedocs.io/), such as, the checkpointing of [braintools](https://github.com/chaobrain/braintools), the event-driven operators in [brainevent](https://github.com/chaobrain/brainevent), the state-based JIT compilation in [brainstate](https://github.com/chaobrain/brainstate), online learning rules in [brainscale](https://github.com/chaobrain/brainscale), or event more.
96
+
97
+ A quick example for this kind of integration:
98
+
99
+ ```python
100
+
101
+ import braintools
102
+ import brainevent
103
+ import brainstate
104
+ import brainunit as u
105
+
106
+
107
+ class EINet(brainstate.nn.Module):
108
+ def __init__(self):
109
+ super().__init__()
110
+ self.n_exc = 3200
111
+ self.n_inh = 800
112
+ self.num = self.n_exc + self.n_inh
113
+ self.N = brainstate.nn.LIFRef(
114
+ self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
115
+ tau=20. * u.ms, tau_ref=5. * u.ms,
116
+ V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
117
+ )
118
+ self.E = brainstate.nn.AlignPostProj(
119
+ comm=brainstate.nn.EventFixedProb(self.n_exc, self.num, 0.02, 0.6 * u.mS),
120
+ syn=brainstate.nn.Expon.desc(self.num, tau=5. * u.ms),
121
+ out=brainstate.nn.COBA.desc(E=0. * u.mV),
122
+ post=self.N
123
+ )
124
+ self.I = brainstate.nn.AlignPostProj(
125
+ comm=brainstate.nn.EventFixedProb(self.n_inh, self.num, 0.02, 6.7 * u.mS),
126
+ syn=brainstate.nn.Expon.desc(self.num, tau=10. * u.ms),
127
+ out=brainstate.nn.COBA.desc(E=-80. * u.mV),
128
+ post=self.N
129
+ )
130
+
131
+ def update(self, t, inp):
132
+ with brainstate.environ.context(t=t):
133
+ spk = self.N.get_spike() != 0.
134
+ self.E(spk[:self.n_exc])
135
+ self.I(spk[self.n_exc:])
136
+ self.N(inp)
137
+ return self.N.get_spike()
138
+
139
+ def save_checkpoint(self):
140
+ braintools.file.msgpack_save('states.msgpack', self.states())
141
+
142
+ ```
143
+
144
+
145
+
146
+ ## Installation
147
+
148
+ You can install ``brainunit`` via pip:
149
+
150
+ ```bash
151
+ pip install brainunit --upgrade
152
+ ```
153
+
154
+
155
+ ## Citation
156
+
157
+ If you use `brainunit` in your research, please consider citing the following paper:
158
+
159
+ ```bibtex
160
+ @article{wang2025integrating,
161
+ title={Integrating physical units into high-performance AI-driven scientific computing},
162
+ author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},
163
+ journal={Nature Communications},
164
+ volume={16},
165
+ number={1},
166
+ pages={3609},
167
+ year={2025},
168
+ publisher={Nature Publishing Group UK London},
169
+ url={https://doi.org/10.1038/s41467-025-58626-4}
170
+ }
171
+ ```
172
+
173
+
174
+
175
+ ## See also the ecosystem
176
+
177
+ ``brainunit`` is one part of our [brain modeling ecosystem](https://brainmodeling.readthedocs.io/).
@@ -1,133 +1,133 @@
1
- <h1 align='center'>BrainUnit</h1>
2
- <h2 align='center'>Physical units and unit-aware math system for general-purpose brain dynamics modeling</h2>
3
-
4
- <p align="center">
5
- <img alt="Header image of brainunit." src="https://raw.githubusercontent.com/chaobrain/saiunit/main/docs/_static/brainunit.png" width=30%>
6
- </p>
7
-
8
-
9
-
10
- <p align="center">
11
- <a href="https://pypi.org/project/brainunit/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainunit"></a>
12
- <a href="https://github.com/chaobrain/brainunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
13
- <a href='https://brainunit.readthedocs.io/?badge=latest'>
14
- <img src='https://readthedocs.org/projects/brainunit/badge/?version=latest' alt='Documentation Status' />
15
- </a>
16
- <a href="https://badge.fury.io/py/brainunit"><img alt="PyPI version" src="https://badge.fury.io/py/brainunit.svg"></a>
17
- <a href="https://pepy.tech/projects/brainunit"><img src="https://static.pepy.tech/badge/brainunit" alt="PyPI Downloads"></a>
18
- </p>
19
-
20
-
21
-
22
-
23
- [BrainUnit](https://github.com/chaobrain/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics modeling. It introduces rigoirous physical units into high-performance AI-driven abstract numerical computing.
24
-
25
- BrainUnit is initially designed to enable unit-aware computations in brain dynamics modeling (see our [BDP ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/)). However, its features and capacities can be applied to general domains in scientific computing and AI4Science. Starting in 2025/02, BrainUnit has been fully integrated into [SAIUnit](https://github.com/chaobrain/saiunit) (the **Unit** system for **S**cientific **AI**).
26
-
27
- Functionalities are the same for both ``brainunit`` and ``saiunit``, and their functions and data structures are interoperable, sharing the same set of APIs, and eliminating any potential conflicts. This meas that
28
-
29
- ```python
30
- import brainunit as u
31
- ```
32
-
33
- equals to
34
-
35
- ```python
36
- import saiunit as u
37
- ```
38
-
39
- For users primarily engaged in general scientific computing, `saiunit` is likely the preferred choice. However, for those focused on brain modeling, we recommend `brainunit`, as it is more closely aligned with our specialized brain dynamics programming ecosystem.
40
-
41
-
42
-
43
- ## Documentation
44
-
45
- The official documentation of BrainUnit is hosted on Read the Docs: [https://brainunit.readthedocs.io](https://brainunit.readthedocs.io)
46
-
47
-
48
-
49
- ## Features
50
-
51
- `brainunit` can be seamlessly integrated into every aspect of our [brain dynamics programming ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/), such as, the checkpointing of [braintools](https://github.com/chaobrain/braintools), the event-driven operators in [brainevent](https://github.com/chaobrain/brainevent), the state-based JIT compilation in [brainstate](https://github.com/chaobrain/brainstate), online learning rules in [brainscale](https://github.com/chaobrain/brainscale), or event more.
52
-
53
- A quick example for this kind of integration:
54
-
55
- ```python
56
-
57
- import braintools
58
- import brainevent
59
- import brainstate
60
- import brainunit as u
61
-
62
-
63
- class EINet(brainstate.nn.Module):
64
- def __init__(self):
65
- super().__init__()
66
- self.n_exc = 3200
67
- self.n_inh = 800
68
- self.num = self.n_exc + self.n_inh
69
- self.N = brainstate.nn.LIFRef(
70
- self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
71
- tau=20. * u.ms, tau_ref=5. * u.ms,
72
- V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
73
- )
74
- self.E = brainstate.nn.AlignPostProj(
75
- comm=brainstate.nn.EventFixedProb(self.n_exc, self.num, 0.02, 0.6 * u.mS),
76
- syn=brainstate.nn.Expon.desc(self.num, tau=5. * u.ms),
77
- out=brainstate.nn.COBA.desc(E=0. * u.mV),
78
- post=self.N
79
- )
80
- self.I = brainstate.nn.AlignPostProj(
81
- comm=brainstate.nn.EventFixedProb(self.n_inh, self.num, 0.02, 6.7 * u.mS),
82
- syn=brainstate.nn.Expon.desc(self.num, tau=10. * u.ms),
83
- out=brainstate.nn.COBA.desc(E=-80. * u.mV),
84
- post=self.N
85
- )
86
-
87
- def update(self, t, inp):
88
- with brainstate.environ.context(t=t):
89
- spk = self.N.get_spike() != 0.
90
- self.E(spk[:self.n_exc])
91
- self.I(spk[self.n_exc:])
92
- self.N(inp)
93
- return self.N.get_spike()
94
-
95
- def save_checkpoint(self):
96
- braintools.file.msgpack_save('states.msgpack', self.states())
97
-
98
- ```
99
-
100
-
101
-
102
- ## Installation
103
-
104
- You can install ``brainunit`` via pip:
105
-
106
- ```bash
107
- pip install brainunit --upgrade
108
- ```
109
-
110
-
111
- ## Citation
112
-
113
- If you use `brainunit` in your research, please consider citing the following paper:
114
-
115
- ```bibtex
116
- @article{wang2025integrating,
117
- title={Integrating physical units into high-performance AI-driven scientific computing},
118
- author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},
119
- journal={Nature Communications},
120
- volume={16},
121
- number={1},
122
- pages={3609},
123
- year={2025},
124
- publisher={Nature Publishing Group UK London},
125
- url={https://doi.org/10.1038/s41467-025-58626-4}
126
- }
127
- ```
128
-
129
-
130
-
131
- ## See also the BDP ecosystem
132
-
133
- We are building the [brain dynamics programming (BDP) ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/). [brainunit](https://github.com/chaobrain/brainunit) has been deeply integrated into our BDP ecosystem.
1
+ <h1 align='center'>BrainUnit</h1>
2
+ <h2 align='center'>Physical units and unit-aware math system for general-purpose brain dynamics modeling</h2>
3
+
4
+ <p align="center">
5
+ <img alt="Header image of brainunit." src="https://raw.githubusercontent.com/chaobrain/saiunit/main/docs/_static/brainunit.png" width=30%>
6
+ </p>
7
+
8
+
9
+
10
+ <p align="center">
11
+ <a href="https://pypi.org/project/brainunit/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainunit"></a>
12
+ <a href="https://github.com/chaobrain/brainunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
13
+ <a href='https://brainunit.readthedocs.io/?badge=latest'>
14
+ <img src='https://readthedocs.org/projects/brainunit/badge/?version=latest' alt='Documentation Status' />
15
+ </a>
16
+ <a href="https://badge.fury.io/py/brainunit"><img alt="PyPI version" src="https://badge.fury.io/py/brainunit.svg"></a>
17
+ <a href="https://pepy.tech/projects/brainunit"><img src="https://static.pepy.tech/badge/brainunit" alt="PyPI Downloads"></a>
18
+ </p>
19
+
20
+
21
+
22
+
23
+ [BrainUnit](https://github.com/chaobrain/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics modeling. It introduces rigoirous physical units into high-performance AI-driven abstract numerical computing.
24
+
25
+ BrainUnit is initially designed to enable unit-aware computations in brain dynamics modeling (see our [BDP ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/)). However, its features and capacities can be applied to general domains in scientific computing and AI4Science. Starting in 2025/02, BrainUnit has been fully integrated into [SAIUnit](https://github.com/chaobrain/saiunit) (the **Unit** system for **S**cientific **AI**).
26
+
27
+ Functionalities are the same for both ``brainunit`` and ``saiunit``, and their functions and data structures are interoperable, sharing the same set of APIs, and eliminating any potential conflicts. This meas that
28
+
29
+ ```python
30
+ import brainunit as u
31
+ ```
32
+
33
+ equals to
34
+
35
+ ```python
36
+ import saiunit as u
37
+ ```
38
+
39
+ For users primarily engaged in general scientific computing, `saiunit` is likely the preferred choice. However, for those focused on brain modeling, we recommend `brainunit`, as it is more closely aligned with our specialized brain dynamics programming ecosystem.
40
+
41
+
42
+
43
+ ## Documentation
44
+
45
+ The official documentation of BrainUnit is hosted on Read the Docs: [https://brainunit.readthedocs.io](https://brainunit.readthedocs.io)
46
+
47
+
48
+
49
+ ## Features
50
+
51
+ `brainunit` can be seamlessly integrated into every aspect of our [brain modeling ecosystem](https://brainmodeling.readthedocs.io/), such as, the checkpointing of [braintools](https://github.com/chaobrain/braintools), the event-driven operators in [brainevent](https://github.com/chaobrain/brainevent), the state-based JIT compilation in [brainstate](https://github.com/chaobrain/brainstate), online learning rules in [brainscale](https://github.com/chaobrain/brainscale), or event more.
52
+
53
+ A quick example for this kind of integration:
54
+
55
+ ```python
56
+
57
+ import braintools
58
+ import brainevent
59
+ import brainstate
60
+ import brainunit as u
61
+
62
+
63
+ class EINet(brainstate.nn.Module):
64
+ def __init__(self):
65
+ super().__init__()
66
+ self.n_exc = 3200
67
+ self.n_inh = 800
68
+ self.num = self.n_exc + self.n_inh
69
+ self.N = brainstate.nn.LIFRef(
70
+ self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
71
+ tau=20. * u.ms, tau_ref=5. * u.ms,
72
+ V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
73
+ )
74
+ self.E = brainstate.nn.AlignPostProj(
75
+ comm=brainstate.nn.EventFixedProb(self.n_exc, self.num, 0.02, 0.6 * u.mS),
76
+ syn=brainstate.nn.Expon.desc(self.num, tau=5. * u.ms),
77
+ out=brainstate.nn.COBA.desc(E=0. * u.mV),
78
+ post=self.N
79
+ )
80
+ self.I = brainstate.nn.AlignPostProj(
81
+ comm=brainstate.nn.EventFixedProb(self.n_inh, self.num, 0.02, 6.7 * u.mS),
82
+ syn=brainstate.nn.Expon.desc(self.num, tau=10. * u.ms),
83
+ out=brainstate.nn.COBA.desc(E=-80. * u.mV),
84
+ post=self.N
85
+ )
86
+
87
+ def update(self, t, inp):
88
+ with brainstate.environ.context(t=t):
89
+ spk = self.N.get_spike() != 0.
90
+ self.E(spk[:self.n_exc])
91
+ self.I(spk[self.n_exc:])
92
+ self.N(inp)
93
+ return self.N.get_spike()
94
+
95
+ def save_checkpoint(self):
96
+ braintools.file.msgpack_save('states.msgpack', self.states())
97
+
98
+ ```
99
+
100
+
101
+
102
+ ## Installation
103
+
104
+ You can install ``brainunit`` via pip:
105
+
106
+ ```bash
107
+ pip install brainunit --upgrade
108
+ ```
109
+
110
+
111
+ ## Citation
112
+
113
+ If you use `brainunit` in your research, please consider citing the following paper:
114
+
115
+ ```bibtex
116
+ @article{wang2025integrating,
117
+ title={Integrating physical units into high-performance AI-driven scientific computing},
118
+ author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},
119
+ journal={Nature Communications},
120
+ volume={16},
121
+ number={1},
122
+ pages={3609},
123
+ year={2025},
124
+ publisher={Nature Publishing Group UK London},
125
+ url={https://doi.org/10.1038/s41467-025-58626-4}
126
+ }
127
+ ```
128
+
129
+
130
+
131
+ ## See also the ecosystem
132
+
133
+ ``brainunit`` is one part of our [brain modeling ecosystem](https://brainmodeling.readthedocs.io/).