Metadata-Version: 2.4
Name: saiunit
Version: 0.5.1
Summary: Enabling Unit-aware Computations for AI-driven Scientific Computing.
Author-email: SAIUnit Developers <chao.brain@qq.com>
License-Expression: Apache-2.0
Project-URL: homepage, https://github.com/chaobrain/saiunit
Project-URL: repository, https://github.com/chaobrain/saiunit
Project-URL: Bug Tracker, https://github.com/chaobrain/saiunit/issues
Project-URL: Documentation, https://saiunit.readthedocs.io/
Keywords: physical unit,physical quantity,scientific computing,AI for science
Classifier: Natural Language :: English
Classifier: Operating System :: OS Independent
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: typing_extensions
Requires-Dist: array_api_compat>=1.9
Requires-Dist: opt_einsum
Provides-Extra: testing
Requires-Dist: pytest; extra == "testing"
Requires-Dist: brainstate; extra == "testing"
Provides-Extra: jax
Requires-Dist: jax; extra == "jax"
Provides-Extra: cpu
Requires-Dist: jax[cpu]; extra == "cpu"
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]; extra == "cuda12"
Provides-Extra: cuda13
Requires-Dist: jax[cuda13]; extra == "cuda13"
Provides-Extra: tpu
Requires-Dist: jax[tpu]; extra == "tpu"
Provides-Extra: cupy
Requires-Dist: cupy-cuda12x>=13.0; extra == "cupy"
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == "torch"
Provides-Extra: dask
Requires-Dist: dask[array]>=2024.1; extra == "dask"
Provides-Extra: ndonnx
Requires-Dist: ndonnx>=0.9; extra == "ndonnx"
Provides-Extra: all
Requires-Dist: saiunit[cupy,dask,jax,ndonnx,torch]; extra == "all"
Dynamic: license-file

<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
```



## Multiple-backend support

`saiunit` is **backend-agnostic**: a `Quantity` pairs a unit with an array
mantissa, and that mantissa can live on any of the supported array libraries.
Every unit-aware operation dispatches to the matching backend, so you can stay
in one library end-to-end or convert with a single method call.

| Backend  | Mantissa            | Install                  | When to use                                 |
|----------|---------------------|--------------------------|---------------------------------------------|
| `numpy`  | `numpy.ndarray`     | core (always installed)  | eager CPU, scipy/pandas/sklearn interop     |
| `jax`    | `jax.Array`         | `saiunit[jax]` (or `[cpu]`/`[cuda12]`/`[cuda13]`/`[tpu]`) | autograd, JIT, vmap, accelerators           |
| `cupy`   | `cupy.ndarray`      | `saiunit[cupy]`          | NVIDIA GPU arrays                           |
| `torch`  | `torch.Tensor`      | `saiunit[torch]`         | PyTorch models, torch autograd              |
| `dask`   | `dask.array.Array`  | `saiunit[dask]`          | out-of-core / parallel, lazy compute        |
| `ndonnx` | `ndonnx.Array`      | `saiunit[ndonnx]`        | symbolic graph for ONNX export              |

Select or override the backend explicitly with `u.using_backend(...)` or
`u.set_default_backend(...)`; convert with `q.to_jax()` / `q.to_numpy()` /
`q.to_cupy()` / `q.to_torch()` / `q.to_dask()` / `q.to_ndonnx()`. Requesting an
uninstalled backend raises `saiunit.BackendError` with the install command —
not a bare `ImportError`. See the
[Backends documentation](https://saiunit.readthedocs.io/backends/overview.html)
for the full story.


## Installation

``saiunit`` has been well tested on ``python>=3.10`` and can be installed on
Windows, Linux, and MacOS. The core package depends only on NumPy. JAX is
optional — install it to enable the ``saiunit.autograd``, ``saiunit.lax``,
and ``saiunit.sparse`` submodules, the custom ``exprel`` primitive, and the
``"jax"`` backend.

Install the NumPy-only core:

```bash
pip install saiunit --upgrade
```

Or pull in JAX with the accelerator build that matches your hardware:

```bash
pip install -U saiunit[jax]      # plain JAX
pip install -U saiunit[cpu]      # pinned JAX CPU wheels
pip install -U saiunit[cuda12]   # JAX on CUDA 12
pip install -U saiunit[cuda13]   # JAX on CUDA 13
pip install -U saiunit[tpu]      # JAX on TPU
```

Opt into additional array backends with the matching extra:

```bash
pip install -U saiunit[cupy]     # CuPy (NVIDIA GPU)
pip install -U saiunit[torch]    # PyTorch
pip install -U saiunit[dask]     # Dask
pip install -U saiunit[ndonnx]   # ndonnx
pip install -U saiunit[all]      # jax + cupy + torch + dask + ndonnx
```

Without JAX, the NumPy backend is auto-selected and any access to a
JAX-only submodule (``saiunit.autograd``, ``saiunit.lax``, ``saiunit.sparse``)
raises ``saiunit.BackendError`` with an install hint. The optional extras are
independent and can be combined freely.

To install the latest version from source:

```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``.

