<p align="center">
<img alt="Header image of SAIUnit." src="https://raw.githubusercontent.com/chaobrain/saiunit/main/docs/_static/logo.png" width=80%>
</p>
<p align="center">
<a href="https://pypi.org/project/saiunit/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/saiunit"></a>
<a href="https://github.com/chaobrain/saiunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
<a href='https://saiunit.readthedocs.io/?badge=latest'>
<img src='https://readthedocs.org/projects/saiunit/badge/?version=latest' alt='Documentation Status' />
</a>
<a href="https://badge.fury.io/py/saiunit"><img alt="PyPI version" src="https://badge.fury.io/py/saiunit.svg"></a>
<a href="https://github.com/chaobrain/saiunit/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaobrain/saiunit/actions/workflows/CI.yml/badge.svg"></a>
<a href="https://pepy.tech/projects/saiunit"><img src="https://static.pepy.tech/badge/saiunit" alt="PyPI Downloads"></a>
</p>
## Motivation
[SAIUnit](https://github.com/chaobrain/saiunit) (/saɪ ˈjuːnɪt/) is designed to provide physical units and unit-aware mathematical systems tailored for **S**cientific **AI** within JAX. In this context, Scientific AI refers to the use of AI models or tools to advance scientific computations. SAIUnit evolves from our [BrainUnit](https://github.com/chaobrain/brainunit), a unit framework originally developed for brain dynamics modeling, extending its capabilities to support a broader range of scientific computing applications. SAIUnit is committed to providing rigorous and automatic physical unit conversion and analysis system for general AI-driven scientific computing.
## Features
Compared to existing unit libraries, such as [Quantities](https://github.com/python-quantities/python-quantities) and [Pint](https://github.com/hgrecco/pint), SAIUnit introduces a rigorous physical unit system specifically designed to support AI computations (e.g., automatic differentiation, just-in-time compilation, and parallelization). Its unique advantages include:
- Integration of over 2,000 commonly used physical units and constants
- Implementation of more than 500 unit-aware mathematical functions
- Deep integration with JAX, providing comprehensive support for modern AI framework features including automatic differentiation (autograd), just-in-time compilation (JIT), vectorization, and parallel computation
- Unit conversion and analysis are performed at compilation time, resulting in zero runtime overhead
- Strict physical unit type checking and dimensional inference system, detecting unit inconsistencies during compilation
```mermaid
graph TD
A[SAIUnit] --> B[Physical Units]
A --> C[Mathematical Functions]
A --> D[JAX Integration]
B --> B1[2000+ Units]
B --> B2[Physical Constants]
C --> C1[500+ Unit-aware Functions]
D --> D1[Autograd]
D --> D2[JIT Compilation]
D --> D3[Vectorization]
D --> D4[Parallelization]
```
We hope these features establish SAIUnit as a reliable physical unit handling solution for general AI-driven scientific computing scenarios.
A quick example:
```python
import saiunit as u
# Define a physical quantity
x = 3.0 * u.meter
x
# [out] 3. * meter
# autograd
f = lambda x: x ** 3
u.autograd.grad(f)(x)
# [out] 27. * meter2
# JIT
import jax
jax.jit(f)(x)
# [out] 27. * klitre
# vmap
jax.vmap(f)(u.math.arange(0. * u.mV, 10. * u.mV, 1. * u.mV))
# [out] ArrayImpl([ 0., 1., 8., 27., 64., 125., 216., 343., 512., 729.]) * mvolt3
```
## Installation
``saiunit`` has been well tested on ``python>=3.9`` + ``jax>=0.4.30`` environments, and can be installed on Windows, Linux, and MacOS.
You can install ``saiunit`` via pip:
```bash
pip install saiunit --upgrade
```
which should install in about 1 minute. If you want to install the latest version from the source, you can clone the repository and install it:
```bash
git clone https://github.com/chaobrain/saiunit.git
cd saiunit
pip install -e .
```
Alternatively, you can install `BrainX`, which bundles `saiunit` with other compatible packages for a comprehensive brain modeling ecosystem:
```bash
pip install BrainX -U
```
## Documentation
The official documentation is hosted on Read the Docs: [https://saiunit.readthedocs.io](https://saiunit.readthedocs.io)
## 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}
}
```
## Ecosystem
`saiunit` has been deeply integrated into following diverse projects, such as:
- [``brainstate``](https://github.com/chaobrain/brainstate): A State-based Transformation System for Program Compilation and Augmentation
- [``braintaichi``](https://github.com/chaobrain/braintaichi): Leveraging Taichi Lang to customize brain dynamics operators
- [``braintools``](https://github.com/chaobrain/braintools): The Common Toolbox for Brain Dynamics Programming.
- [``dendritex``](https://github.com/chaobrain/dendritex): Dendritic Modeling in JAX
- [``pinnx``](https://github.com/chaobrain/pinnx): Physics-Informed Neural Networks for Scientific Machine Learning in JAX.
Other unofficial projects include:
- [``diffrax``](https://github.com/chaoming0625/diffrax): Numerical differential equation solvers in JAX.
- [``jax-md``](https://github.com/Routhleck/jax-md): Differentiable Molecular Dynamics in JAX
- [``Catalax``](https://github.com/Routhleck/Catalax): JAX-based framework to model biological systems
- ...
## Acknowledgement
The initial version of the project benefited a lot from the following projects:
- [``astropy.units``](https://docs.astropy.org/en/stable/units/index.html): physical units in ``astropy``.
- [``brian2.units``](https://brian2.readthedocs.io/en/stable/user/units.html): physical units in ``brian2``.
Raw data
{
"_id": null,
"home_page": null,
"name": "saiunit",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.10",
"maintainer_email": null,
"keywords": "physical unit, physical quantity, brain modeling, scientific computing, AI for science",
"author": null,
"author_email": "SAIUnit Developers <chao.brain@qq.com>",
"download_url": null,
"platform": null,
"description": "<p align=\"center\">\r\n \t<img alt=\"Header image of SAIUnit.\" src=\"https://raw.githubusercontent.com/chaobrain/saiunit/main/docs/_static/logo.png\" width=80%>\r\n</p> \r\n\r\n\r\n\r\n<p align=\"center\">\r\n\t<a href=\"https://pypi.org/project/saiunit/\"><img alt=\"Supported Python Version\" src=\"https://img.shields.io/pypi/pyversions/saiunit\"></a>\r\n\t<a href=\"https://github.com/chaobrain/saiunit/blob/main/LICENSE\"><img alt=\"LICENSE\" src=\"https://img.shields.io/badge/License-Apache%202.0-blue.svg\"></a>\r\n <a href='https://saiunit.readthedocs.io/?badge=latest'>\r\n <img src='https://readthedocs.org/projects/saiunit/badge/?version=latest' alt='Documentation Status' />\r\n </a> \t\r\n <a href=\"https://badge.fury.io/py/saiunit\"><img alt=\"PyPI version\" src=\"https://badge.fury.io/py/saiunit.svg\"></a>\r\n <a href=\"https://github.com/chaobrain/saiunit/actions/workflows/CI.yml\"><img alt=\"Continuous Integration\" src=\"https://github.com/chaobrain/saiunit/actions/workflows/CI.yml/badge.svg\"></a>\r\n <a href=\"https://pepy.tech/projects/saiunit\"><img src=\"https://static.pepy.tech/badge/saiunit\" alt=\"PyPI Downloads\"></a>\r\n</p>\r\n\r\n\r\n## Motivation\r\n\r\n[SAIUnit](https://github.com/chaobrain/saiunit) (/sa\u026a \u02c8ju\u02d0n\u026at/) is designed to provide physical units and unit-aware mathematical systems tailored for **S**cientific **AI** within JAX. In this context, Scientific AI refers to the use of AI models or tools to advance scientific computations. SAIUnit evolves from our [BrainUnit](https://github.com/chaobrain/brainunit), a unit framework originally developed for brain dynamics modeling, extending its capabilities to support a broader range of scientific computing applications. SAIUnit is committed to providing rigorous and automatic physical unit conversion and analysis system for general AI-driven scientific computing. \r\n\r\n\r\n\r\n## Features\r\n\r\nCompared to existing unit libraries, such as [Quantities](https://github.com/python-quantities/python-quantities) and [Pint](https://github.com/hgrecco/pint), SAIUnit introduces a rigorous physical unit system specifically designed to support AI computations (e.g., automatic differentiation, just-in-time compilation, and parallelization). Its unique advantages include:\r\n\r\n- Integration of over 2,000 commonly used physical units and constants\r\n- Implementation of more than 500 unit-aware mathematical functions\r\n- Deep integration with JAX, providing comprehensive support for modern AI framework features including automatic differentiation (autograd), just-in-time compilation (JIT), vectorization, and parallel computation\r\n- Unit conversion and analysis are performed at compilation time, resulting in zero runtime overhead\r\n- Strict physical unit type checking and dimensional inference system, detecting unit inconsistencies during compilation\r\n\r\n\r\n```mermaid\r\ngraph TD\r\n A[SAIUnit] --> B[Physical Units]\r\n A --> C[Mathematical Functions]\r\n A --> D[JAX Integration]\r\n B --> B1[2000+ Units]\r\n B --> B2[Physical Constants]\r\n C --> C1[500+ Unit-aware Functions]\r\n D --> D1[Autograd]\r\n D --> D2[JIT Compilation]\r\n D --> D3[Vectorization]\r\n D --> D4[Parallelization]\r\n```\r\n\r\nWe hope these features establish SAIUnit as a reliable physical unit handling solution for general AI-driven scientific computing scenarios.\r\n\r\nA quick example:\r\n\r\n```python\r\n\r\nimport saiunit as u\r\n\r\n# Define a physical quantity\r\nx = 3.0 * u.meter\r\nx\r\n# [out] 3. * meter\r\n\r\n# autograd\r\nf = lambda x: x ** 3\r\nu.autograd.grad(f)(x)\r\n# [out] 27. * meter2 \r\n\r\n\r\n# JIT\r\nimport jax\r\njax.jit(f)(x)\r\n# [out] 27. * klitre\r\n\r\n# vmap\r\njax.vmap(f)(u.math.arange(0. * u.mV, 10. * u.mV, 1. * u.mV))\r\n# [out] ArrayImpl([ 0., 1., 8., 27., 64., 125., 216., 343., 512., 729.]) * mvolt3\r\n```\r\n\r\n\r\n\r\n## Installation\r\n\r\n``saiunit`` has been well tested on ``python>=3.9`` + ``jax>=0.4.30`` environments, and can be installed on Windows, Linux, and MacOS.\r\n\r\nYou can install ``saiunit`` via pip:\r\n\r\n```bash\r\npip install saiunit --upgrade\r\n```\r\n\r\nwhich should install in about 1 minute. If you want to install the latest version from the source, you can clone the repository and install it:\r\n\r\n```bash\r\ngit clone https://github.com/chaobrain/saiunit.git\r\ncd saiunit\r\npip install -e .\r\n```\r\n\r\nAlternatively, you can install `BrainX`, which bundles `saiunit` with other compatible packages for a comprehensive brain modeling ecosystem:\r\n\r\n```bash\r\npip install BrainX -U\r\n```\r\n\r\n## Documentation\r\n\r\nThe official documentation is hosted on Read the Docs: [https://saiunit.readthedocs.io](https://saiunit.readthedocs.io)\r\n\r\n\r\n\r\n## Citation\r\n\r\n```bibtex\r\n@article{wang2025integrating,\r\n title={Integrating physical units into high-performance AI-driven scientific computing},\r\n author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},\r\n journal={Nature Communications},\r\n volume={16},\r\n number={1},\r\n pages={3609},\r\n year={2025},\r\n publisher={Nature Publishing Group UK London},\r\n url={https://doi.org/10.1038/s41467-025-58626-4}\r\n}\r\n```\r\n\r\n\r\n## Ecosystem\r\n\r\n`saiunit` has been deeply integrated into following diverse projects, such as:\r\n\r\n- [``brainstate``](https://github.com/chaobrain/brainstate): A State-based Transformation System for Program Compilation and Augmentation\r\n- [``braintaichi``](https://github.com/chaobrain/braintaichi): Leveraging Taichi Lang to customize brain dynamics operators\r\n- [``braintools``](https://github.com/chaobrain/braintools): The Common Toolbox for Brain Dynamics Programming.\r\n- [``dendritex``](https://github.com/chaobrain/dendritex): Dendritic Modeling in JAX\r\n- [``pinnx``](https://github.com/chaobrain/pinnx): Physics-Informed Neural Networks for Scientific Machine Learning in JAX.\r\n\r\n\r\nOther unofficial projects include:\r\n\r\n- [``diffrax``](https://github.com/chaoming0625/diffrax): Numerical differential equation solvers in JAX.\r\n- [``jax-md``](https://github.com/Routhleck/jax-md): Differentiable Molecular Dynamics in JAX\r\n- [``Catalax``](https://github.com/Routhleck/Catalax): JAX-based framework to model biological systems\r\n- ...\r\n\r\n\r\n## Acknowledgement\r\n\r\nThe initial version of the project benefited a lot from the following projects:\r\n\r\n- [``astropy.units``](https://docs.astropy.org/en/stable/units/index.html): physical units in ``astropy``.\r\n- [``brian2.units``](https://brian2.readthedocs.io/en/stable/user/units.html): physical units in ``brian2``.\r\n\r\n",
"bugtrack_url": null,
"license": null,
"summary": "Enabling Unit-aware Computations for AI-driven Scientific Computing.",
"version": "0.0.16",
"project_urls": {
"Bug Tracker": "https://github.com/chaobrain/saiunit/issues",
"Documentation": "https://saiunit.readthedocs.io/",
"homepage": "https://github.com/chaobrain/saiunit",
"repository": "https://github.com/chaobrain/saiunit"
},
"split_keywords": [
"physical unit",
" physical quantity",
" brain modeling",
" scientific computing",
" ai for science"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "be56928dfdf6a3daa595b7f460617e9dd794a1a81ceb1c6e81989bcec6a96bc3",
"md5": "ac8121720d1297679f24c2fb46c0ec78",
"sha256": "f321090c488e699596fbe8277fbf9229150dd2653386d96e102ad634b3315a0e"
},
"downloads": -1,
"filename": "saiunit-0.0.16-py3-none-any.whl",
"has_sig": false,
"md5_digest": "ac8121720d1297679f24c2fb46c0ec78",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.10",
"size": 332121,
"upload_time": "2025-07-13T05:35:42",
"upload_time_iso_8601": "2025-07-13T05:35:42.058305Z",
"url": "https://files.pythonhosted.org/packages/be/56/928dfdf6a3daa595b7f460617e9dd794a1a81ceb1c6e81989bcec6a96bc3/saiunit-0.0.16-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-07-13 05:35:42",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "chaobrain",
"github_project": "saiunit",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [
{
"name": "numpy",
"specs": []
},
{
"name": "jax",
"specs": []
},
{
"name": "jaxlib",
"specs": []
},
{
"name": "typing_extensions",
"specs": []
}
],
"lcname": "saiunit"
}
JSON
