Integrating physical units into high-performance AI-driven scientific computing

Area of Science:

• Scientific Computing
• Artificial Intelligence
• Physics-Informed Machine Learning

Background:

• Scientific research relies on physical units for accurate computation.
• Current AI libraries lack native support for physical units, hindering scientific integration.
• This gap impedes the development of reliable AI-driven scientific tools.

Purpose of the Study:

• Introduce SAIUnit, a novel system for integrating physical units into AI scientific computing.
• Ensure compatibility with JAX, a popular AI framework.
• Enhance the accuracy, reliability, and interpretability of AI in scientific research.

Main Methods:

• Developed SAIUnit with over 2000 physical units and 500 unit-aware functions.
• Ensured full compatibility with JAX transformations (automatic differentiation, JIT compilation, etc.).
• Tested SAIUnit in diverse AI-driven scientific domains like numerical integration and brain modeling.

Main Results:

• SAIUnit maintains unit consistency during JAX transformations.
• Unit checking during compilation improves accuracy and reliability.
• Demonstrated effectiveness in numerical integration, brain modeling, and physics-informed neural networks.
• No compromise in runtime performance was observed.

Conclusions:

• SAIUnit bridges the gap between abstract AI frameworks and physical units.
• It enables more robust and physically grounded AI-driven scientific computing.
• This system enhances the interpretability and collaborative potential of scientific computations.