计算矩阵所有对角线的迹

Question

我需要计算一个矩阵所有对角线的迹。简单来说，对于一个n行m列的矩阵，这个操作应该能产生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达到了最快的方式。

Answer 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]))

Answer 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的解决方案要快（对于大数组来说）。

---

这个方法对数组的行列比例非常敏感，因为这个比例决定了在Python中循环的次数相对于Numpy的次数。正如@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），像我之前发的代码那样逐行访问数组可能还是更快，这可能是因为数组在内存中的布局方式（获取连续的块比获取分散的部分要快）。

Answer 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]

这是一个具体数字的例子。如果这个对你的使用场景有帮助，我可以把它变成一个通用的函数。

Answer 4

如果你想避免使用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代码快，但你可以自己测试一下时间。

Answer 5

这是你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的调用。