brainunit

KO | EN

gitlite — search

#hacktoberfest #react #ai #llm #golang #machine-learning #linux #docker #macos #nodejs #windows #awesome #kubernetes #vue

(★ 16)

Physical units and unit-aware mathematical system for general-purpose brain dynamics modeling.

brainunit Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1 align='center'>BrainUnit</h1>
<h2 align='center'>Physical units and unit-aware math system for general-purpose brain dynamics modeling</h2>

[BrainUnit](https://github.com/chaobrain/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics modeling. It introduces rigorous 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 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.

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=brainstate.nn.EventFixedProb(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=brainstate.nn.EventFixedProb(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
```

## Citation

```bibtex
@article{wang2025integrating,
  title={Integrating physical units into high-performance AI-driven scientific computing},
  author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},
  journal={Nature Communications},
  volume={16},
  number={1},
  pages={3609},
  year={2025},
  publisher={Nature Publishing Group UK London},
  url={https://doi.org/10.1038/s41467-025-58626-4}
}
```

## See also the ecosystem

We are building the [brain modeling ecosystem](https://brainmodeling.readthedocs.io/). [brainunit](https://github.com/chaobrain/brainunit) has been deeply integrated into our ecosystem.

brainunit

brainunit Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

brainunit

brainunit Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README