https://github.com/chaobrain/brainunit
Physical units and unit-aware mathematical system in JAX for brain dynamics and AI4Science.
https://github.com/chaobrain/brainunit
ai4science brain-dynamics-modeling brain-modeling physical-quantities physical-units
Last synced: about 2 months ago
JSON representation
Physical units and unit-aware mathematical system in JAX for brain dynamics and AI4Science.
- Host: GitHub
- URL: https://github.com/chaobrain/brainunit
- Owner: chaobrain
- License: apache-2.0
- Created: 2024-06-07T01:28:47.000Z (10 months ago)
- Default Branch: main
- Last Pushed: 2024-10-28T10:49:18.000Z (6 months ago)
- Last Synced: 2024-10-28T12:17:51.664Z (6 months ago)
- Topics: ai4science, brain-dynamics-modeling, brain-modeling, physical-quantities, physical-units
- Language: Python
- Homepage: https://brainunit.readthedocs.io
- Size: 3.27 MB
- Stars: 2
- Watchers: 2
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- Changelog: changelog.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Security: SECURITY.md
Awesome Lists containing this project
- awesome-jax - brainunit - Physical units and unit-aware mathematical system in JAX. <img src="https://img.shields.io/github/stars/chaobrain/brainunit?style=social" align="center"> (Libraries / New Libraries)
README
BrainUnit
Physical units and unit-aware math system for general-purpose brain dynamics modeling
[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.
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**).
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
```python
import brainunit as u
```
equals to
```python
import saiunit as u
```
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.
## Documentation
The official documentation of BrainUnit is hosted on Read the Docs: [https://brainunit.readthedocs.io](https://brainunit.readthedocs.io)
## Features
`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.
A quick example for this kind of integration:
```python
import braintools
import brainevent.nn
import brainstate
import brainunit as u
class EINet(brainstate.nn.Module):
def __init__(self):
super().__init__()
self.n_exc = 3200
self.n_inh = 800
self.num = self.n_exc + self.n_inh
self.N = brainstate.nn.LIFRef(
self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
tau=20. * u.ms, tau_ref=5. * u.ms,
V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
)
self.E = brainstate.nn.AlignPostProj(
comm=brainevent.nn.FixedProb(self.n_exc, self.num, 0.02, 0.6 * u.mS),
syn=brainstate.nn.Expon.desc(self.num, tau=5. * u.ms),
out=brainstate.nn.COBA.desc(E=0. * u.mV),
post=self.N
)
self.I = brainstate.nn.AlignPostProj(
comm=brainevent.nn.FixedProb(self.n_inh, self.num, 0.02, 6.7 * u.mS),
syn=brainstate.nn.Expon.desc(self.num, tau=10. * u.ms),
out=brainstate.nn.COBA.desc(E=-80. * u.mV),
post=self.N
)
def update(self, t, inp):
with brainstate.environ.context(t=t):
spk = self.N.get_spike() != 0.
self.E(spk[:self.n_exc])
self.I(spk[self.n_exc:])
self.N(inp)
return self.N.get_spike()
def save_checkpoint(self):
braintools.file.msgpack_save('states.msgpack', self.states())
```
## Installation
You can install ``brainunit`` via pip:
```bash
pip install brainunit --upgrade
```
## See also the BDP ecosystem
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.