我正在尝试为 N 体积分器优化双循环,我发现我的代码的问题是当我将存储的变量写入内存视图位置时会产生巨大的开销。
我最初在 numpy 中对这段代码进行了矢量化处理,但是在另一个 for 循环中调用它来更新粒子位置,而且开销非常大。我有一个 np.ndarray Nx2 位置向量(X),我想返回一个 Nx2 动量向量(XOut)——下面列出的当前代码返回一个内存视图,但这没关系,因为我想最终嵌入这个一旦我调试了这个瓶颈,就可以在其他 Cython 函数中运行。
我尝试了 cython -a "name.pyx" 命令,发现我或多或少地将所有内容都作为 C 类型。但是,我发现在循环的底部,写入 XOut[ii,0] -= valuex 的内存视图会产生大部分运行时间。如果我将其更改为常数,以使 XOut[ii,0] -= 5,则代码将快 40 倍左右。我认为这意味着我正在在那条线上进行某种复制操作,这让我放慢了速度。我的 Cython/C++ 背景是初级的,但我认为我需要更改语法,以便从指针写入内存视图。任何建议将不胜感激;谢谢!
import numpy as np
cimport numpy as np
from cython.view cimport array as cvarray
cimport cython
from libc.math cimport sinh, cosh, sin, cos, acos, exp, sqrt, fabs, M_PI
DTYPE = np.float64
ctypedef np.float64_t DTYPE_t
cdef DTYPE_t pi = 3.141592653589793
@cython.cdivision(True)
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False) # turn off negative index wrapping for entire function
def intTerms(const DTYPE_t[:,:] X, DTYPE_t epsilon, DTYPE_t[:,:] XOut):
cdef Py_ssize_t ii,jj,N
N = X.shape[0]
cdef DTYPE_t valuex,valuey,r2,xvec,yvec
for ii in range(0,N):
for jj in range(ii+1,N):
xvec = X[ii,0]-X[jj,0]
yvec = X[ii,1]-X[jj,1]
r2 = max(xvec**2+yvec**2,epsilon)
valuex = xvec/r2**2
valuey = yvec/r2**2
XOut[ii,0] -= valuex
XOut[ii,1] -= 5 #valuey
XOut[jj,0] += 5 #valuex
XOut[jj,1] += 5 #valuey
XOut[ii,0] /= 2*pi
XOut[ii,1] /= 2*pi
return XOut