Numpy 用并行Cython加速线性变换

我需要加速线性变换的计算，其大致形式如下：

import numpy as np

N=10000
input=np.random.random(N)

x=np.linspace(0,100,N)
y=np.linspace(0,30,N)
X,Y=np.meshgrid(x,y,sparse=True)

output=np.dot(np.cos(X*Y),input)

也就是说，我在一个规则的网格上计算余弦，然后将我的输入乘以得到的矩阵。实际上，核函数（这里是余弦）更复杂，特别是它不是周期性的。因此，不可能简化FFT类型

在我的多核机器上，上述转换大约需要5秒钟。现在，我确实需要加快速度。简单的第一次尝试是使用numexpr：

import numpy as np
import numexpr as ne

N=10000
input=np.random.random(N)

x=np.linspace(0,100,N)
y=np.linspace(0,30,N)
X,Y=np.meshgrid(x,y,sparse=True)

output=np.dot(ne.evaluate('cos(X*Y)'),input)

这利用了并行计算，并将执行时间减少到约0.9秒。这很好，但对我来说还不够。因此，我的下一个尝试是使用并行Cython：

import numpy as np
from cython.parallel import prange

cimport numpy as np
cimport cython
from libc.math cimport cos

DTYPE = np.float64    
ctypedef np.float64_t DTYPE_t

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def transform(double[:] x, double[:] y, double[:] input):

    cdef unsigned int N = x.shape[0]
    cdef double[:] output = np.zeros(N)
    cdef unsigned int row, col

    for row in prange(N, nogil= True):
        for col in range(N):
            output[row] += cos(x[row]*y[col])*input[col]

    return output

我通过执行

from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext

ext_modules=[
    Extension("cythontransform",
          ["cythontransform.pyx"],
          libraries=["m"],
          extra_compile_args = ["-O3", "-ffast-math", "-march=native", "-fopenmp" ],
          extra_link_args=['-fopenmp']
          ) 
]

setup(
  name = "cythontransform",
  cmdclass = {"build_ext": build_ext},
  ext_modules = ext_modules
)

从命令行。通过调用转换

import numpy as np
from cythontransform import transform

N=10000
input=np.random.random(N)

x=np.linspace(0,100,N)
y=np.linspace(0,30,N)

output=transform(x,y,input)

产生相当微弱的改善，大约0.7秒

有人知道Cython代码进一步改进的可能性吗

---

或者，是否有其他框架（PyOpenCL、Pythran、Numba等）更适合此问题？

在我的笔记本电脑上，以下版本：

编撰

pythran python -Ofast dd.py -fopenmp

---

运行速度大约是您建议的cython版本的两倍。虽然…

FYI:

np，但我没有调查为什么会发生这种情况。empty（N）

不会将新数组初始化为0，因此您可能会从cython函数中获取垃圾。使用

np.zero（N）

。行

X，Y=np.meshgrid（X，X，sparse=True）

中是否有输入错误？应该是

X，Y=np.meshgrid（X，Y，sparse=True）

？如果没有，则您定义了

y

，但从未使用过它。这可能就是计算所需的时间。Cython代码看起来相当不错。将MemoryView声明为

double[：：1]

将强制阵列对齐，并可能提供一个微小的加速。亲爱的Pierre，确实，这似乎提供了几个百分点的加速。总比没有好。谢谢

pythran python -Ofast dd.py -fopenmp