RSTSR v0.10:Rust 原生、全面兼容 Python Array API 标准的科学计算高维张量库

💡 原文中文,约3500字,阅读约需9分钟。
📝

内容提要

RSTSR v0.10 是一个原生 Rust 的 n 维张量科学计算库,接口高度兼容 NumPy,并全面实现 Python Array API 2025 标准(FFT 除外)。它支持多种 BLAS 后端、行/列主序布局、并行管理、零开销 reshape 和 Einstein Summation,性能与 NumPy 相当。该库已用于 REST 电子结构程序,实现高效杂化泛函二阶解析梯度。

🔎

延伸解读

Array API 兼容性的实际含义

RSTSR 宣称全面兼容 Python Array API 2025 标准(FFT 除外),测试套件结果为 1336 通过、17 失败、33 跳过。其中 14 个失败和 33 个跳过源于测试套件自身或 Faer 后端未实现 FFT,剩余 3 个失败是实际应用中罕见的 SVD 浮点数溢出问题。这意味着除 FFT 外,RSTSR 在接口层面已基本对齐标准,但用户需注意 SVD 在极端数值下的稳定性风险。

Rust 语言特性带来的接口取舍

RSTSR 并非机械照搬所有 Array API 接口,而是受 Rust 语言特性限制做了取舍。例如 Python 允许的 __ipow__、__imatmul__ 运算符重载在 Rust 中不存在,动态类型列表也难以静态实现,导致 concat、stack 等函数在通用场景下受限。不过,对于相同类型的张量,RSTSR 仍提供 rt::concat 和 rt::stack 作为替代。读者应理解这些排除项是语言层面的必然,而非实现遗漏。

并行管理与后端选择的实用考量

RSTSR 通过 rayon 控制线程数,确保在并行线程内外调用时都不会超出用户指定的线程上限,这对嵌套并行场景尤为重要。它支持 Faer、OpenBLAS、Intel MKL、BLIS、KML 等多种 BLAS 后端,且性能与 NumPy 相当。用户可根据现有环境选择后端,而泛型后端设计允许自定义数据承载方式(如 GPU、硬盘、网络),但需自行实现 trait 接口。

在电子结构计算中的验证价值

RSTSR 已用于 REST 电子结构程序,实现了杂化泛函二阶解析梯度(Hessian),可计算分子振动频率、热力学能,并衍生红外与拉曼光谱。在 RI-JK 近似下,Skeleton 导数性能达到甚至超过同类程序,CP-HF 求解过程 FLOPS 利用率超过 40%。这表明 RSTSR 的高维张量操作和 Vecdot 优化在真实科学计算负载中得到了检验,而非仅停留在基准测试。

❓

Q&A

RSTSR 是什么?它和 NumPy 有什么关系?

RSTSR 是一个原生 Rust 的 n 维张量科学计算库,接口高度兼容 NumPy,并全面实现 Python Array API 2025 标准(FFT 除外)。它旨在让 Rust 科学计算写起来像 NumPy,同时更高效、更安全、更适合并行。

RSTSR 对 Python Array API 2025 标准的兼容性如何?

RSTSR 除 FFT 功能外,正确实现了所有 Array API 2025 接口。测试套件结果为 1336 Pass / 17 Fail / 33 Skip,其中 14 Fail / 33 Skip 是测试套件自身或 Faer 后端未实现 FFT 导致,其余 3 Fail 是罕见的 SVD 浮点数溢出问题。

RSTSR 支持哪些 BLAS 后端?性能如何?

RSTSR 支持 Faer 默认多线程后端,以及 OpenBLAS、Intel MKL、BLIS、KML 等多种 BLAS 后端。在 CPU 下性能表现与 NumPy 相当,支持 solve、svd、eigh 等矩阵线性代数运算。

RSTSR 有哪些主要特性?

主要特性包括:使用 % 作为矩阵乘法算符;优秀的并行管理(通过 rayon 控制线程数);同时支持 row-major 与 column-major 内存布局;支持多种 BLAS 后端;可扩展后端(泛型支持自定义后端);支持多种数据类型(包括字符串等);零开销 reshape;Einstein Summation 支持;高效的 Vecdot 支持。

RSTSR 在实际项目中有哪些应用?

RSTSR 基于 REST 电子结构程序平台开发,已用于在 REST 中实现高效的杂化泛函二阶解析梯度(Hessian),在 RI-JK 近似下 Skeleton 导数性能达到甚至超过其他类似电子结构程序,CP-HF 求解过程 FLOPS 利用率达 40% 以上。

RSTSR 有哪些功能限制或未实现的接口?

RSTSR 未实现 FFT 功能;由于 Rust 语言特性,无法实现 Python 中的 __ipow__、__imatmul__ 等运算符重载,以及动态类型列表相关的 concat、stack 等函数(但相同类型张量仍可使用 rt::concat 和 rt::stack)。

🏷️

标签

➡️

继续阅读