python - 数组的 numpy 和反对角线

Question

给定一个 numpy ndarray，我想取前两个轴，并用一个新轴替换它们，这是它们的反对角线的总和。

特别是，假设我有变量 x、y、z、...，并且我的数组的条目表示概率

array[i,j,k,...] = P(x=i, y=j, z=k, ...)

我想得到

new_array[l,k,...] = P(x+y=l, z=k, ...) = sum_i P(x=i, y=l-i, z=k, ...)

即 new_array[l,k,...] 是所有 array[i,j,k,...] 的总和使得 i+j=l.

在 numpy 中执行此操作的最有效和/或最干净的方法是什么？

编辑添加:根据@hpaulj 的推荐，这里是明显的迭代解决方案:

array = numpy.arange(30).reshape((2,3,5))
array = array / float(array.sum()) # make it a probability
new_array = numpy.zeros([array.shape[0] + array.shape[1] - 1] + list(array.shape[2:]))
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        new_array[i+j,...] += array[i,j,...]
new_array.sum() # == 1

Answer 1

trace 函数可以给出对角线的总和。您可以指定偏移量和 2 个轴(0 和 1 是默认值)。为了得到反对角线，你只需要翻转一个维度。 np.flipud 可以做到这一点，尽管它只是 [::-1,...] 索引。

将它们放在一起，

np.array([np.trace(np.flipud(array),offset=k) for k in range(-1,3)])

匹配你的 new_array。

它仍然遍历 l 的可能值(在本例中为 4)。 trace 本身被编译。

在这种小情况下，它实际上比双循环(2x3 步)慢。即使我将 flipud 移出内部循环，它仍然更慢。我不知道这对于更大的阵列如何扩展。

将其进一步矢量化的部分问题在于每条对角线的长度不同。

In [331]: %%timeit
array1 = array[::-1]
np.array([np.trace(array1,offset=k) for k in range(-1,3)])
   .....: 
10000 loops, best of 3: 87.4 µs per loop

In [332]: %%timeit 
new_array = np.zeros([array.shape[0] + array.shape[1] - 1] + list(array.shape[2:]))                                                       
for i in range(2):
    for j in range(3):
        new_array[i+j] += array[i,j]
   .....: 
10000 loops, best of 3: 43.5 µs per loop

scipy.sparse 有一个 dia 格式，它存储非零对角线的值。它存储一个填充值数组以及偏移量。

array([[12,  0,  0,  0],
       [ 8, 13,  0,  0],
       [ 4,  9, 14,  0],
       [ 0,  5, 10, 15],
       [ 0,  1,  6, 11],
       [ 0,  0,  2,  7],
       [ 0,  0,  0,  3]])
array([-3, -2, -1,  0,  1,  2,  3])

虽然这是解决可变对角线长度问题的一种方法，但我认为这在您只需要它们的总和的情况下没有帮助。