具有可变数量因子的 Numpy 矢量求和

Question

我目前正在计算一个包含对索引求和的函数。索引在0和T的整数部分之间；理想情况下，我希望能够快速计算 T 的几个值的总和。 在现实生活中，T 的大部分值都很小，但有一小部分可能比平均值大一到两个数量级。

我现在做的是：
1）我定义了向量T，例如（我的真实数据有更多的条目，只是提供一个想法）：

import numpy as np 
T = np.random.exponential(5, 10)

2) 我创建了一个矩阵，其中包含 0 和 int(T) 之间的因子，然后归零：

n = int(T.max())
j = ((np.arange(n) < T[:,np.newaxis])*np.arange(1,n+1)).astype(int).transpose()
print(j)

[[ 1 1 1 1 1 1 1 1 1 1]
[ 2 0 2 2 2 0 2 0 2 2]
[ 0 0 3 0 3 0 3 0 3 3]
[ 0 0 4 0 4 0 0 0 4 4]
[ 0 0 5 0 5 0 0 0 5 5]
[ 0 0 6 0 6 0 0 0 6 6]
[ 0 0 7 0 7 0 0 0 0 7]
[ 0 0 8 0 8 0 0 0 0 8]
[ 0 0 9 0 9 0 0 0 0 9]
[ 0 0 0 0 10 0 0 0 0 10]
[ 0 0 0 0 11 0 0 0 0 0]
[0 0 0 0 12 0 0 0 0 0]]

3) 我生成求和的单个元素，使用掩码避免将函数应用于零元素：

A =  np.log(1 + (1 + j) * 5)* (j>0) 

4) 我沿列求和：

A.sum(axis=0)

获得： 数组([ 5.170484 , 2.39789527, 29.96464821, 5.170484 , 42.29052851、2.39789527、8.21500643、2.39789527、 18.49060911, 33.9899999])

有没有最快/更好的方法来矢量化它？我觉得它很慢，因为大量的零对总和没有贡献，但由于我是 NumPy 的初学者，我想不出更好的编写方法。

编辑：在我的实际问题中，应用于 j 的函数还取决于第二个参数 tau（在与 T 大小相同的向量中）。所以每一列包含的项目都不一样。

Answer 1

查看您的j，对于每一列，它的数字从1 到N，其中N 是根据每个T 元素决定的。然后，您沿每一列求和，这与求和直到 N 相同，因为其余元素无论如何都是零。这些求和值可以用np.cumsum 计算，而那些N 值基本上是j 中每一列的限制，可以直接从T 计算。然后将这些 N 值用作索引，以索引到 cumsum-ed 值，从而为我们提供最终输出。

考虑到cumsum 是唯一完成的计算并且在一维数组上也是如此，这应该非常快且内存效率很高，与原始方法中沿每列的二维数组完成的求和相比。因此，我们有一个像这样的矢量化方法 -

n = int(T.max())
vals = (np.log(1 + (1 + np.arange(1,n+1)) * 5)).cumsum()
out = vals[(T.astype(int)).clip(max=n-1)]

在内存使用方面，我们生成了三个变量-

n    : Scalar
vals : 1D array of n elements 
out  : 1D array of T.size elements (this is the output anyway)

运行时测试和验证输出 -

In [5]: def original_app(T):
   ...:     n = int(T.max())
   ...:     j = ((np.arange(n) < T[:,None])*np.arange(1,n+1)).astype(int).transpose()
   ...:     A =  np.log(1 + (1 + j) * 5)* (j>0) 
   ...:     return A.sum(axis=0)
   ...: 
   ...: def vectorized_app(T):
   ...:     n = int(T.max())
   ...:     vals = (np.log(1 + (1 + np.arange(1,n+1)) * 5)).cumsum()
   ...:     return vals[(T.astype(int)).clip(max=n-1)]
   ...: 

In [6]: # Input array
   ...: T = np.random.exponential(5, 10000)

In [7]: %timeit original_app(T)
100 loops, best of 3: 9.62 ms per loop

In [8]: %timeit vectorized_app(T)
10000 loops, best of 3: 50.1 µs per loop

In [9]: np.allclose(original_app(T),vectorized_app(T)) # Verify outputs
Out[9]: True