编辑推荐：
在人工智能（AI）助力科研的浪潮中，主流高性能计算库缺乏对物理单位的支持，阻碍了 AI 与科研融合。研究人员开展 SAIUnit 系统的研究，该系统集成超 2000 种物理单位和 500 种函数，与 JAX 兼容。结果显示其提升计算准确性等，促进跨学科科研，意义重大。
在当今科学研究领域，人工智能（AI）的迅猛发展为各个学科带来了前所未有的机遇，它正以前所未有的方式重塑科研格局，在物理、化学、生物、气候科学、天文学和神经科学等诸多领域都发挥着关键作用。AI 能够帮助科学家提出假设、设计实验、高效处理和分析海量数据，发现那些传统方法难以触及的新见解。然而，一个不容忽视的问题是，当前主流的高性能计算库，如 PyTorch、TensorFlow 和 JAX 等，虽然在机器学习和深度学习应用中表现出色，但它们最初并非为科学研究量身定制，普遍缺乏对物理单位的原生支持。
科学研究的基石之一是对物理世界进行精确测量，这依赖于明确的单位系统。例如国际单位制（SI），它以米作为长度的标准单位，确保了不同研究间测量结果的互操作性和可比性。在科学计算中，准确应用物理单位不仅体现了科学的严谨性，更是防止维度错误、避免灾难性后果的关键，比如因单位管理不当导致的火星气候轨道器灾难。而且，在复杂计算过程中保持统一的单位系统，对于强化模型中的物理约束、确保结果的物理合理性至关重要。随着科学界越来越多地采用 AI 驱动的研究方法，在高性能计算库中提供标准化、全面的单位支持变得迫在眉睫。
为了解决这一关键问题，来自北京大学、上海交通大学和广东智能科学技术研究院等机构的研究人员开展了深入研究，他们引入了 SAIUnit 系统，旨在将物理单位无缝集成到高性能的科学 AI 库中，尤其注重与 JAX 的兼容性。这一研究成果发表在《Nature Communications》上，具有重要的科学意义。
研究人员在开展这项研究时，运用了多个关键技术方法。首先，构建了基于 SI 单位的严格单位转换系统，通过 saiunit.Dimension、saiunit.Unit 和 saiunit.Quantity 三种相互关联的数据结构来实现物理单位的处理和计算。其次，将 SAIUnit 深度集成到多个基于 JAX 的科学计算库中，如数值求解器库 Diffrax、脑动力学编程生态系统 BrainPy 和物理信息神经网络框架 PINNx 等，以验证其在不同科学计算领域的适用性和有效性。同时，对 SAIUnit 在 CPU、GPU 和 TPU 平台上的性能进行了全面评估，对比有无物理单位时模型的编译和运行时间等指标。
自动维度分析与数据结构创新
SAIUnit 创新的核心在于三种相互关联的数据结构。saiunit.Dimension 基于 SI 单位建立了严格的单位转换系统，它由七个整数组成的元组表示，每个整数对应一个基本 SI 单位的维度，通过这些基本单位的组合可以创建出各种导出单位，其运算规则严格遵循 SI 规则，确保了单位转换和计算的维度一致性。saiunit.Unit 用于表示具有不同尺度的物理单位，它由一个度量尺度和一个 Dimension 实例组成，能够处理从微观到宏观不同尺度的物理量，其运算规则保证了在不同尺度和维度下物理量运算的一致性和有意义性。saiunit.Quantity 则通过将数值与 Unit 实例集成，实现了单位感知的数组计算，它可以处理非常大或非常小的尺度数据，并且与现有的数值计算工作流程无缝集成，重载了几乎所有 NumPy 和 JAX 数组支持的数学运算符。
跨学科的单位感知计算集成
研究人员将 SAIUnit 集成到多个科学计算领域，展现了其广泛的适用性。在数值积分方面，以 Diffrax 为例，不同物理系统对时空离散化精度要求差异很大，传统抽象数值计算库难以准确表示和比较不同系统的时间尺度，而单位感知的 Diffrax 能够在一致的维度内标准化时空精度控制，防止因单位转换错误导致的计算偏差。在脑建模领域，BrainPy 集成 SAIUnit 后，能够有效表示大脑中的多尺度物理量，从神经元和突触层面的膜时间常数、膜面积，到网络和电路层面的突触连接性、网络动力学等，使多尺度脑网络的定义更直观，模拟过程更易于生物学家理解，增强了模型结果的可解释性。在物理信息神经网络（PINNs）领域，PINNx 框架集成 SAIUnit 后，对神经网络的输入、输出以及微分方程计算施加物理单位约束，自动进行单位一致性检查和尺度变换，提高了数据导入和模型拟合的效率与准确性，确保模型严格遵循物理维度规范。
单位感知计算的性能与优势
SAIUnit 在性能方面表现出色，通过将数值计算与单位处理分离，其单位感知计算的性能开销主要集中在编译阶段，在运行时的开销与无物理单位的抽象计算相比可忽略不计。在模拟洛伦兹系统的实验中，与其他 Python 单位系统（如 Quantities 和 Pint）相比，SAIUnit 在不同精度下都能保证基于单位和无单位模拟结果的一致性，表明其在处理物理单位时不会损失精度，可应用于低精度 AI 训练场景。此外，SAIUnit 的单位感知计算还显著提高了研究效率，在构建生物物理神经元模型时，它避免了手动单位转换带来的错误和复杂性。同时，其标准化的单位系统促进了跨学科研究的协作，统一了不同学科的数据单位，方便了数据共享和利用。
在研究结论和讨论部分，SAIUnit 成功解决了当前 AI 库中缺乏物理单位支持的关键问题，为单位感知的科学计算提供了强大的工具。它与 JAX 变换的完全集成，使得用户能够在保持单位一致性的同时，充分利用自动微分、即时编译、向量化和并行化等高级功能。通过在数值积分、脑建模和 PINNs 等领域的应用案例，展示了 SAIUnit 在实际场景中的实用性和有效性，能够以最小的开销实现单位处理自动化。然而，研究也指出 SAIUnit 在低精度计算中，当尾数 A 和单位缩放因子 s 都以低精度表示时，可能会出现数值不稳定的问题，未来需要开发更强大的操作框架来动态调整缩放因子 s，以确保计算稳定性和精度。总体而言，SAIUnit 为高性能、AI 驱动的科学计算带来了新的突破，为跨学科科研合作搭建了更坚实的桥梁，推动了科学研究向更精确、更高效的方向发展。
