BrainUnit

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

Header image of brainunit.

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). 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 (the Unit system for Scientific 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

import brainunit as u

equals to

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://brainx.chaobrain.com/brainunit

Features

brainunit can be seamlessly integrated into every aspect of our brain modeling ecosystem, such as, the checkpointing of braintools, the event-driven operators in brainevent, the state-based JIT compilation in brainstate, online learning rules in brainscale, or event more.

A quick example for this kind of integration:

import braintools
import brainevent
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:

pip install brainunit --upgrade

Citation

If you use brainunit in your research, please consider citing the following paper:

@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}
}

Metadata

Release files for brainunit 0.5.2

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for brainunit 0.5.2
File	Size	Uploaded
brainunit-0.5.2.tar.gz	11.9 kB	Jul 22, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for brainunit 0.5.2
File	Interpreter	ABI	Platform	Reset
brainunit-0.5.2-py3-none-any.whl 37.0 kB Jul 22, 2026	Python 3	none	any	Details

Total release size: 49.0 kB

Release files / brainunit-0.5.2.tar.gz

Download URL	brainunit-0.5.2.tar.gz
Size	11.9 kB
Tags	Source
SHA-256 checksum How to use checksums	`fca089d0ec6badbd4ca0e1c3321ad7ef3b667e54b37295a832b0f48b498fec13`
BLAKE2b-256 checksum How to use checksums	`10f4f427c41a5f730131c994ca7d4c04e09c87235777541362b4eef14aa53953`
Upload date	Jul 22, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.13.14`

Release files / brainunit-0.5.2-py3-none-any.whl

Download URL	brainunit-0.5.2-py3-none-any.whl
Size	37.0 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`a863df4cc78d9f4f9b4b5ec92f204e7d78903864a7492554cac384cdbb000ad6`
BLAKE2b-256 checksum How to use checksums	`95962c715c3ab2a9dfbb722b9031c146ddebab6ba8a1b1e5c17a6ee889914c25`
Upload date	Jul 22, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.13.14`