使用numpy计算滚动加权总和

Question

我很好奇是否有更优化的方法来计算这个“滚动加权和”（不确定实际术语是什么，但我将提供一个示例以进一步澄清）。我之所以问这个问题，是因为我确信我的当前代码片段在内存使用方面不是最优的，并且可以通过使用numpy的更高级函数来提高其性能。

示例：

import numpy as np

A = np.append(np.linspace(0, 1, 10), np.linspace(1.1, 2, 30))
np.random.seed(0)

B = np.random.randint(3, size=40) + 1

# list of [(weight, (lower, upper))]
d = [(1, (-0.25, -0.20)), (0.5, (-0.20, -0.10)), (2, (-0.10, 0.15))]

在Python 3.7中：

## A
array([0.        , 0.11111111, 0.22222222, 0.33333333, 0.44444444,
       0.55555556, 0.66666667, 0.77777778, 0.88888889, 1.        ,
       1.1       , 1.13103448, 1.16206897, 1.19310345, 1.22413793,
       1.25517241, 1.2862069 , 1.31724138, 1.34827586, 1.37931034,
       1.41034483, 1.44137931, 1.47241379, 1.50344828, 1.53448276,
       1.56551724, 1.59655172, 1.62758621, 1.65862069, 1.68965517,
       1.72068966, 1.75172414, 1.78275862, 1.8137931 , 1.84482759,
       1.87586207, 1.90689655, 1.93793103, 1.96896552, 2.        ])

## B
array([1, 2, 1, 2, 2, 3, 1, 3, 1, 1, 1, 3, 2, 3, 3, 1, 2, 2, 2, 2, 1, 2,
       1, 1, 2, 3, 1, 3, 1, 2, 2, 3, 1, 2, 2, 2, 1, 3, 1, 3])

预期解决方案：

array([ 6. ,  6.5,  8. , 10.5, 12. , 11. , 11.5, 11.5,  6.5, 13.5, 25. ,
       27.5, 30.5, 34.5, 37.5, 36. , 35. , 35. , 34. , 34.5, 34. , 36.5,
       33. , 34. , 34.5, 34.5, 36. , 39. , 37. , 36. , 37. , 36.5, 37.5,
       39. , 36.5, 37.5, 34. , 31. , 27.5, 23. ])

我想要转化为代码的逻辑是：

让我们看看如何计算10.5（期望解决方案中的第四个元素）。d表示具有第一个浮点元素weight和第二个元组元素bounds（以(lower, upper)的形式）的嵌套元组集合。

我们查看A的第四个元素（0.33333333），并对d中每个元组应用bounds。对于d中的第一个元组：

0.33333333 + (-0.25) = 0.08333333
0.33333333 + (-0.20) = 0.13333333

我们回到A，看看在(0.08333333, 0.1333333)范围内是否有任何元素。因为A的第二个元素(0.11111111)落在这个范围内，所以我们提取B的第二个元素(2)，并将其乘以来自d的权重(1)，然后加上期望输出的第二个元素。

在遍历d中的所有元组之后，期望输出的第四个元素被计算为：

1 * 2 + 0.5 * 1 + 2 * (2 + 2) = 10.5

这是我的尝试代码：

D = np.zeros(len(A))
for v in d:
    weight, (_lower, _upper) = v
    lower, upper = A + _lower, A + _upper
    _A = np.tile(A, (len(A), 1))
    __A = np.bitwise_and(_A > lower.reshape(-1, 1), _A < upper.reshape(-1, 1))
    D += weight * (__A @ B)
D

希望这有意义。如有疑问，请随时提出。谢谢！

- Wilson

Answer 1

由于区间（-0.25，-0.20）、（-0.20，-0.10）和（-0.10，0.15）实际上是一个区间的分段（-0.25，0.15）的子区间，因此您可以找到索引，以便将元素插入A中以保持顺序。它们指定了在B上执行加法的切片。简而言之：

partition = np.array([-0.25, -0.20, -0.10, 0.15])
weights = np.array([1, 0.5, 2])
out = []
for n in A:
    idx = np.searchsorted(A, n + partition)
    results = np.add.reduceat(B[:idx[-1]], idx[:-1])
    out.append(np.dot(results, weights))
>>> print(out)
[7.5, 7.5, 8.0, 10.5, 12.0, 11.0, 11.5, 11.5, 6.5, 13.5, 27.5, 27.5, 31.5, 35.5, 37.5, 37.0, 36.0, 35.0, 34.0, 34.5, 34.0, 36.5, 33.0, 34.0, 34.5, 34.5, 36.0, 39.0, 37.0, 36.0, 37.0, 36.5, 37.5, 39.0, 36.5, 37.5, 34.0, 31.0, 27.5, 23.0]

请注意，如果存在空的B切片，则results将是错误的。

- mathfux

---

感谢您给出的出色第一反应！我真的很喜欢您如何定义partition和weights，因为您所做的假设适用于我的问题。这也是一种更清洁的解决方案，可以降低内存成本。只是想知道，有没有什么方法可以使用numpy矢量化来避免在A元素之间使用for循环？谢谢！ - Wilson

这里的瓶颈是np.add.reduceat方法。通常，如果您能找到一种用扁平化数组替换命令的方法，就可以将for循环向量化。for n in A: idx = np.searchsorted(A, n + partition)部分很容易，您可以用np.searchsorted(A, partition + A[:, None])替换它，并在需要时调用ravel方法。我只是不知道该怎么处理这个np.add.reduceat部分，它看起来很难在2D数组的轴上应用它。如果我有更多时间，我会再看看。 - mathfux

我找到了一个启发于你的解决方案的方法。我将标记你的为最有帮助的，因为你提供了足够的指导让我自己解决了问题。我会在这里发布我的解决方案！再次感谢！ - Wilson

感谢合作。确实，看到“np.add.reduceat”被stride技巧所取代是一件好事。虽然我不确定“np.apply_along_axis”是否有助于提高速度，因为您仍然需要逐步应用某些方法。如果您想要更快的循环，请看一下“numba”。通常它比“numpy”表现更好。 - mathfux

Answer 2

感谢 @mathfux 在这里提供了足够的指导。以下是我根据这里的对话开发的最终代码解决方案：

partition = np.array([-0.25, -0.20, -0.10, 0.15])
weights = np.array([1, 0.5, 2])
idx = np.searchsorted(A, partition + A[:, None])
_idx = np.lib.stride_tricks.sliding_window_view(idx, 2, axis = 1)
values = np.apply_along_axis(lambda x: B[slice(*(x))].sum(), 2, _idx)
values @ weights

- Wilson