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

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