计算所有对角线上矩阵的轨迹

时间:2014-06-28 23:07:28

标签: python numpy cython

我需要在其所有对角线上计算矩阵的轨迹。也就是说,对于nxm矩阵,操作应该产生n + m-1'跟踪。这是一个示例程序:

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的最快方式。

7 个答案:

答案 0 :(得分:7)

如果你想远离Cython,建立一个对角线索引数组并使用np.bincount可能会有所帮助:

>>> import numpy as np
>>> a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> rows, cols = a.shape
>>> rows_arr = np.arange(rows)
>>> cols_arr = np.arange(cols)
>>> diag_idx = rows_arr[:, None] - (cols_arr - (cols - 1))
>>> diag_idx
array([[3, 2, 1, 0],
       [4, 3, 2, 1],
       [5, 4, 3, 2]])
>>> np.bincount(diag_idx.ravel(), weights=a.ravel())
array([  3.,   9.,  18.,  15.,  13.,   8.])

根据我的时间,对于您的示例输入,它比原始的纯Python方法快4倍。所以我不认为它会比你的Cython代码更快,但你可能想要计时。

答案 1 :(得分:3)

如果您的矩阵形状远离方形,即如果它是高或宽,那么您可以有效地使用步幅技巧来做到这一点。你可以在任何情况下使用步幅技巧,但如果矩阵接近正方形,它可能不是超级内存效率。

您需要做的是在相同的数据上创建一个新的数组视图,该数据的构建方式是从一行到另一行的步骤也会导致列中的增量。这是通过改变数组的步幅来实现的。

需要处理的问题在于阵列的边界,其中需要零填充。如果阵列远非正方形,这无关紧要。如果它是方形的,那么我们需要两倍大小的数组来填充。

如果边缘处不需要较小的迹线,则无需进行零焊接。

这里(假设列多于行,但很容易适应):

import numpy as np
from numpy.lib.stride_tricks import as_strided

A = np.arange(30).reshape(3, 10)
A_embedded = np.hstack([np.zeros([3, 2]), A, np.zeros([3, 2])])
A = A_embedded[:, 2:-2]  # We are now sure that the memory around A is padded with 0, but actually we never really need A again

new_strides = (A.strides[0] + A.strides[1], A.strides[1])
B = as_strided(A_embedded, shape=A_embedded[:, :-2].shape, strides=new_strides)

traces = B.sum(0)

print A
print B
print traces

为了符合您在示例中显示的输出,您需要将其反转(请参阅@larsmans评论)

traces = traces[::-1]

这是具体数字的具体示例。如果这对你的用例有用,我可以把它变成一般功能。

答案 2 :(得分:2)

如果阵列很大,这是有竞争力的:

def f5(A):
    rows, cols = A.shape
    N = rows + cols -1
    out = np.zeros(N, A.dtype)
    for idx in range(rows):
        out[N-idx-cols:N-idx] += A[idx]
    return out[::-1]

虽然它使用Python循环,但它比bincount解决方案更快(对于我的系统上的大型数组..)


此方法对阵列列/行比率具有高灵敏度,因为此比率决定了相对于Numpy在Python中完成的循环次数。 由于@Jaime指出迭代最小维度是有效的,例如:

def f6(A):
    rows, cols = A.shape
    N = rows + cols -1
    out = np.zeros(N, A.dtype)

    if rows > cols:
        for idx in range(cols):
            out[N-idx-rows:N-idx] += A[:, idx]
    else:
        for idx in range(rows):
            out[N-idx-cols:N-idx] += A[idx]
        out = out[::-1]
    return out

但是应该注意的是,对于更大的数组大小(例如我的系统上的100000 x 500),我在第一个代码中逐行访问数组仍然可能更快,可能是因为数组是如何放置的在RAM中 (获取连续块比分散位更快)。

答案 3 :(得分:2)

这是Cython功能的改进版本。 老实说,如果Cython是一个选项,我就是这样做的。

import numpy as np
from libc.stdint cimport int64_t as i64
from cython cimport boundscheck, wraparound

@boundscheck(False)
@wraparound(False)
def all_trace_int64(i64[:,::1] A):
    cdef:
        int i,j
        i64[:] t = np.zeros(A.shape[0] + A.shape[1] - 1, dtype=np.int64)
    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            t[A.shape[0]-i+j-1] += A[i,j]
    return np.array(t)

这将比您在问题中提供的版本快得多,因为它按照存储在内存中的顺序迭代数组。 对于小型阵列,这两种方法几乎相同,不过这个方法在我的机器上略快一些。

我编写了这个函数,因此它需要一个C连续的数组。 如果你有一个Fortran连续数组,转置它,然后反转输出的顺序。

这确实以与示例中显示的函数相反的顺序返回答案,因此如果顺序特别重要,则需要反转数组的顺序。

您还可以通过使用更重的优化进行编译来提高性能。 例如,您可以通过替换

在IPython笔记本中使用其他编译器标志构建您的Cython代码
%%cython

类似

%%cython -c=-O3 -c=-march=native -c=-funroll-loops -f

编辑: 执行此操作时,您还需要确保您的值不是由外部产品生成的。如果您的值来自外部产品,则此操作可以与外部产品合并为np.convolve的单个调用。

答案 4 :(得分:1)

这可以通过(略微滥用)以两种方式使用scipy.sparse.dia_matrix来完成,一种比另一种稀疏。

第一个产生确切结果,使用dia_matrix存储的数据向量

import numpy as np
from scipy.sparse import dia_matrix
A = np.arange(30).reshape(3, 10)
traces = dia_matrix(A).data.sum(1)[::-1]

内存密集程度较低的方法是以相反的方式工作:

import numpy as np
from scipy.sparse import dia_matrix
A = np.arange(30).reshape(3, 10)
A_dia = dia_matrix((A, range(len(A))), shape=(A.shape[1],) * 2)
traces = np.array(A_dia.sum(1)).ravel()[::-1]

但请注意,此解决方案中缺少两个条目。这可以通过聪明的方式进行纠正,但我还不确定。


@moarningsun找到了解决方案:

rows, cols = A.shape

A_dia = dia_matrix((A, np.arange(rows)), shape=(cols,)*2)
traces1 = A_dia.sum(1).A.ravel()

A_dia = dia_matrix((A, np.arange(-rows+1, 1)), shape=(rows,)*2)
traces2 = A_dia.sum(1).A.ravel()

traces = np.concatenate((traces1[::-1], traces2[-2::-1]))

答案 5 :(得分:-1)

np.trace做你想做的事:

import numpy as np

A = array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]])

n = A.shape[0]
[np.trace(A, i) for i in range(-n+1, n+1)]

修改:根据@ user2357112的建议将np.sum(np.diag())更改为np.trace()

答案 6 :(得分:-2)

使用numpy数组trace方法:

import numpy as np
A = np.array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]])
A.trace()

返回:

15