Numpy：将函数应用于每个（或子集）对角线

Question

我想对数组的每个对角线应用一个函数（即 np.var，但一般方法会更好）。我的阵列是方形的。我可以这样做：

offset_list = np.arange(-1 * len(arr) + 1, len(arr))
diag_var_list = [np.var(np.diagonal(arr, k)) for k in offset_list]

如果我想要对角线的子集，我可以更改 offset_list。

但是使用列表理解似乎效率低下，因为我将在许多大型数组上执行此操作。有没有更好的方法？

Answer 1

线性代数中使用的稀疏矩阵通常具有对角线结构（尽管并非所有对角线）。 scipy.sparse 有两种指定对角线的方法 - 作为数组列表，每个数组的长度不同，以及作为二维填充数组。

https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.sparse.diags.html

https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.sparse.spdiags.html#scipy.sparse.spdiags

对于 numpy（密集）数组，对角线没有那么特殊。所有函数一次只处理一个对角线（可以偏移）。 np.diag*

Answer 2

你可以使用下面的思路。如果 major, minor = A.strides 然后将 A 的步幅设置为 major + minor, minor （这应该小心完成以避免跨出数组边界）你得到的数组的每一列都是对角线。这样用 A.sum(axis=0) 之类的东西你可以计算对角线的总和。对于意味着您可以使用相同但乘以某些值，如 A.shape[0] / [1, 2, ... A.shape[0], ... 2, 1] 来修复对角线长度的变化。对于方差，您可以使用 variance = <(x - <x>)**2> = <x**2> - <x>**2.

import numpy as np
def rot45(A):
    """
    >>> A = np.triu(np.arange(25).reshape(5, 5), 1)
    >>> print(A)
    [[ 0  1  2  3  4]
     [ 0  0  7  8  9]
     [ 0  0  0 13 14]
     [ 0  0  0  0 19]
     [ 0  0  0  0  0]]
    >>> print(rot45(A))
    [[ 1  2  3  4]
     [ 7  8  9  0]
     [13 14  0  0]
     [19  0  0  0]]
    """
    major, minor = A.strides
    strides = major + minor, minor
    shape = A.shape[0] - 1, A.shape[1]
    return np.lib.stride_tricks.as_strided(A, shape, strides)[:, 1:]


def apply_diag(A, func):
    """
    >>> A = np.arange(25).reshape(5, 5)
    >>> print(A)
    [[ 0  1  2  3  4]
     [ 5  6  7  8  9]
     [10 11 12 13 14]
     [15 16 17 18 19]
     [20 21 22 23 24]]
    >>> offset_list = np.arange(-1 * len(A) + 1, len(A))
    >>> diag_var_list = [np.sum(np.diagonal(A, k)) for k in offset_list]
    >>> diag_var_list
    [20, 36, 48, 56, 60, 40, 24, 12, 4]
    >>> print(apply_diag(A, np.sum))
    [20, 36, 48, 56, 60, 40, 24, 12, 4]
    """
    U = np.triu(A, 1)
    U = rot45(U)
    D = np.tril(A, -1).T.copy()
    D = rot45(D)
    return func(D, axis=0)[::-1].tolist() + [func(np.diag(A))] + func(U, axis=0)[::-1].tolist()[::-1]


def using_numpy(A):
    """
    >>> A = np.arange(25).reshape(5, 5)
    >>> print(A)
    [[ 0  1  2  3  4]
     [ 5  6  7  8  9]
     [10 11 12 13 14]
     [15 16 17 18 19]
     [20 21 22 23 24]]
    >>> offset_list = np.arange(-1 * len(A) + 1, len(A))
    >>> diag_var_list = [np.var(np.diagonal(A, k)) for k in offset_list]
    >>> diag_var_list
    [0.0, 9.0, 24.0, 45.0, 72.0, 45.0, 24.0, 9.0, 0.0]
    >>> print(using_numpy(A))
    [ 0.  9. 24. 45. 72. 45. 24.  9.  0.]
    """
    multiply = (A.shape[0] - 1) / np.r_[1:A.shape[0], A.shape[0] - 1, A.shape[0] - 1:0:-1]
    return multiply * apply_diag(A ** 2, np.mean) - (multiply * apply_diag(A, np.mean))**2


def list_comp(A, func):
    """
    >>> A = np.arange(25).reshape(5, 5)
    >>> list_comp(A, np.sum)
    [20, 36, 48, 56, 60, 40, 24, 12, 4]
    """
    offset_list = np.arange(-1 * len(A) + 1, len(A))
    return [func(np.diagonal(A, k)) for k in offset_list]

(100, 100) 大小的矩阵似乎有 10 倍的加速，但对于更大的矩阵，速度差异下降得更小。

In [9]: A = np.random.randn(100, 100)

In [10]: %timeit using_numpy(A)
761 µs ± 3.35 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [11]: %timeit list_comp(A, np.var)
9.57 ms ± 19.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [12]: A = np.random.randn(1000, 1000)

In [13]: %timeit using_numpy(A)
37.4 ms ± 125 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [14]: %timeit list_comp(A, np.var)
112 ms ± 927 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Numpy：将函数应用于每个（或子集）对角线

Numpy: apply function to every (or subset of) diagonal

python

arrays

numpy

vectorization