我需要计算一个矩阵沿着其所有对角线的迹。也就是说,对于一个nxm的矩阵,该操作应该生成n+m-1个“迹”。 下面是一个示例程序:
Cython程序的性能优于通过np.trace循环的程序:
import numpy as np
A=np.arange(12).reshape(3,4)
def function_1(A):
output=np.zeros(A.shape[0]+A.shape[1]-1)
for i in range(A.shape[0]+A.shape[1]-1):
output[i]=np.trace(A,A.shape[1]-1-i)
return output
A
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
function_1(A)
array([ 3., 9., 18., 15., 13., 8.])
我的希望是找到一种替换程序中循环的方法,因为我需要在非常大的矩阵上进行多次这种计算。看起来有前途的一个方向是使用numpy.einsum,但我还不太会如何实现。另外,我已经考虑使用Cython在循环中完全重写问题:
%load_ext cythonmagic
%%cython
import numpy as np
cimport numpy as np
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
def function_2(long [:,:] A):
cdef int n=A.shape[0]
cdef int m=A.shape[1]
cdef long [::1] output = np.empty(n+m-1,dtype=np.int64)
cdef size_t l1
cdef int i,j, k1
cdef long out
it_list1=range(m)
it_list2=range(m,m+n-1)
for l1 in range(len(it_list1)):
k1=it_list1[l1]
i=0
j=m-1-k1
out=0
while (i<n)&(j<m):
out+=A[i,j]
i+=1
j+=1
output[k1]=out
for l1 in range(len(it_list2)):
k1=it_list2[l1]
i=k1-m+1
j=0
out=0
while (i<n)&(j<m):
out+=A[i,j]
i+=1
j+=1
output[k1]=out
return np.array(output)
Cython程序的性能优于通过np.trace循环的程序:
%timeit function_1(A)
10000 loops, best of 3: 62.7 µs per loop
%timeit function_2(A)
100000 loops, best of 3: 9.66 µs per loop
所以,基本上我希望得到反馈,关于是否有更有效的使用numpy/scipy例程的方式,或者我是否已经通过使用cython实现了最快的方式。
np.fromiter(map(A.trace, range(A.shape[1]-1, -A.shape[0], -1)), dtype=np.int64)
。 - behzad.nouri