提高对numpy三角函数的操作性能

Question

我有一段相当庞大的代码需要进行优化。经过使用time.time()进行分析后，我发现占用了最多处理时间的是这一行代码（它被执行了数千次）：

A = np.cos(a) * np.cos(b) - np.sin(a) * np.sin(b) * np.sin(c - d)

所有变量均可随机定义为：

N = 5000
a = np.random.uniform(0., 10., N)
b = np.random.uniform(0., 50., N)
c = np.random.uniform(0., 30., N)
d = np.random.uniform(0., 25., N)

有什么方法可以提高计算 A 的性能吗？我已经使用了numpy，但是我已经没有更多的想法了。

- Gabriel

---

所有变量都可以随机定义，您是真正感兴趣的用例吗？还是只是为了在stackoverflow问题中使用np.random.uniform来获取一些数据？ - Warren Weckesser

这是在循环中发生的吗？它可以放入矩阵格式并向量化吗？ - Jamie Bull

@WarrenWeckesser 后者。实际数据来自我相当庞大的代码中缺失的部分。 - Gabriel

@njuffa 显然它没有：https://github.com/numpy/numpy/issues/2626 - Gabriel

很遗憾，大多数现代计算环境都提供了该功能。您可能需要考虑将您的意见添加到现有的 numpy 功能请求中。 - njuffa

Answer 1

通过使用和差公式，您可以减少三角函数调用的次数。在下面的示例中，func1和func2计算相同的值，但func2对三角函数的调用更少。

import numpy as np

def func1(a, b, c, d):
    A = np.cos(a) * np.cos(b) - np.sin(a) * np.sin(b) * np.sin(c - d)
    return A

def func2(a, b, c, d):
    s = np.sin(c - d)
    A = 0.5*((1 - s)*np.cos(a - b) + (1 + s)*np.cos(a + b))
    return A

这是一个与N = 5000进行的时间比较：

In [48]: %timeit func1(a, b, c, d)
1000 loops, best of 3: 374 µs per loop

In [49]: %timeit func2(a, b, c, d)
1000 loops, best of 3: 241 µs per loop

- Warren Weckesser

Answer 2

您是否尝试过使用像Numba、Cython、Pythran或其他任何Python加速器？

我进行了一些 Pythran 的测试，以下是结果：

原始代码：

• Python + numpy：1000次循环，3次中的最佳结果：每次1.43毫秒
• Pythran：1000次循环，3次中的最佳结果：每次777微秒
• Pythran + SIMD：1000次循环，3次中的最佳结果：每次488微秒

Warren提供的代码：

• Python + numpy：1000次循环，3次中的最佳结果：每次1.05毫秒
• Pythran：1000次循环，3次中的最佳结果：每次646微秒
• Pythran + SIMD：1000次循环，3次中的最佳结果：每次425微秒

这是在 N = 5000 时执行的。

• 更新*：

以下是代码：

# pythran export func1(float[], float[], float[], float[])
# pythran export func2(float[], float[], float[], float[])
import numpy as np

def func1(a, b, c, d):
    A = np.cos(a) * np.cos(b) - np.sin(a) * np.sin(b) * np.sin(c - d)
    return A

def func2(a, b, c, d):
    s = np.sin(c - d)
    A = 0.5*((1 - s)*np.cos(a - b) + (1 + s)*np.cos(a + b))
    return A

并且命令行：

$ pythran test.py  # Default compilation
$ pythran test.py -march=native -DUSE_BOOST_SIMD  # Pythran with code vectorization

- P. Brunet

---

我不熟悉Pythran，也不知道SIMD代表什么。你能提供更多细节吗？ - Gabriel

SIMD的意思是：单指令多数据。它是现代体系结构中使用的一个概念，用于同时处理多个值（即这些数组的多个单元格）。现代CPU处理AVX（或AVX2），可以一次处理4个单元格。 关于Pythran，它是一个将Python编译为C++的编译器，可优化给定代码。它仅处理Python代码的限制（静态类型代码和其他一些限制），但它适用于科学内核（如此示例）。 - P. Brunet

像 numexpr 这样简单的工具可以让你获得与 Pythran + SIMD 相当的加速。 - Daniel