计算矩阵沿所有对角线的迹

Question

我需要计算一个矩阵沿着其所有对角线的迹。也就是说，对于一个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实现了最快的方式。

- shadowprice

---

我想知道这个会怎么比较：np.fromiter(map(A.trace, range(A.shape[1]-1, -A.shape[0], -1)), dtype=np.int64)。 - behzad.nouri

对于大矩阵，Cython版本可以在内存访问方面得到改进。即循环遍历行而不是对角线。 - user2379410

如果您希望跟踪信息换行而不是零填充，那么我认为在傅里叶空间中有一种很好的方法可以实现。 - eickenberg

Answer 1

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

- Jaime

---

对于一个60×10的数组，这比原来快了21倍。 - Fred Foo

不错，而且它不会让人头疼内存。 - eickenberg

对于更大的数组，执行 diag_idx = rows_arr[:, None] - (cols_arr - cols + 1) 可以明显提高速度。 - user2379410

好的，修改了答案。我猜测，但如果数组足够不规则，从最短的向量中减去可能也会有所察觉... - Jaime

Answer 2

如果您的矩阵形状与方形差别很大，也就是说它比较高或宽，则可以使用步幅技巧高效地完成此操作。您可以在任何情况下使用步幅技巧，但如果矩阵接近方形，则可能不是超级节省内存。

您需要做的是在相同的数据上创建一个新的数组视图，该数组视图构造方式使得从一行到下一行的步长也会导致列的增加。这是通过更改数组的步长来实现的。

需要注意的问题在于数组的边界处，需要进行零填充。如果数组距离方形很远，则这并不重要。如果是正方形，则需要两倍于数组大小的填充空间。

如果您不需要边缘处的较小痕迹，则无需进行零填充。

以下是示例代码（假设列数多于行数，但很容易适应其他情况）：

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]

这是一个具体的例子，有具体的数字。如果这对您的使用场景有用，我可以将其转换为通用函数。

- eickenberg

---

+1，尽管输出与原始顺序相反。 B.sum(axis=0)[::-1] 可以解决这个问题。 - Fred Foo

谢谢您的输入，我没有注意到。由于我没有看到在计算中集成这个反向顺序的方法，所以我会按照您的建议编辑输出。 - eickenberg

Answer 3

这是您的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。

- IanH

---

这个方法是所有提出的方法中最快的。虽然如此，我真的很感激大家的努力 - 我学到了很多以前不知道的关于numpy的知识（特别是bincount和strides）。感谢大家的帮助！ - shadowprice

Answer 4

如果数组很大，这将具有竞争力：

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），像我发布的第一个代码中逐行访问数组仍然可能更快，这可能是因为数组在RAM中的布局方式（获取连续块比分散位更快）。

- user2379410

---

是的，我能看出当矩阵非常宽时这会更有效，但它将关键取决于列/行比率，比例越高越好。 - eickenberg

@eickenberg；当然这取决于列/行比，但这并不是非常关键。在我的测试中，即使对于A = np.random.rand(100000, 500)，它也比@Jaime的方法更快。 - user2379410

有趣。那应该是一个不同的效果：@Jaime构建了一个大的索引数组，而你的方法避免了这种情况。因此，当数据在任何方向上增加时，他的方法应该变得更糟。而当cols * lines == constant时，你的方法应该在lines >> cols时表现更好。 - eickenberg

如果您总是迭代最小的维度，并在最大的维度上进行向量化求和，那么它应该表现得更好，不是吗？ - Jaime

@Jaime，对于较小的矩阵确实是这样。从某个大小开始，访问列中不连续的数据似乎成本太高，速度又变慢了。在哪个数组大小上发生转换取决于机器。 - user2379410

Answer 5

这可以通过（有点滥用地）使用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]))

- eickenberg

---

有趣的是，最终问题归结为“找到一个C扩展，无论该扩展的原始目的如何，都可以满足您的需求” :P 我想知道这与（几乎最优）Cython实现的性能有多接近。 - user2379410

你说得对，基本上就是这样。其中一些甚至可以很优雅（虽然我不想评判这个）。如果有人这样做，它应该非常有用，并且在速度/内存等方面具有足够的优势，以证明其晦涩性。在这种情况下，我承认我已经开始从代码高尔夫的角度看待这个线程 :) - eickenberg

Answer 6

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()。

- Dietrich

---

如果你要使用列表推导式，np.trace 已经做得更好了。关键是要避免使用 Python 级别的循环和推导式，因为它们比你想要的慢得多。 - user2357112

Answer 7

使用numpy数组的trace方法：

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

返回：

15

- Russia Must Remove Putin