python - 使用 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,上)).

我们查看 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

希望这是有道理的。请随时提出澄清问题。谢谢!

Answer 1

因为区间 (-0.25, -0.20)、(-0.20, -0.10) 和 (-0.10, 0.15) 实际上是 partition of an interval 的子区间(-0.25, 0.15) 你可以 find indices应在 A 中插入元素以维持顺序。他们指定 slices of B to perform添加上。简而言之:

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 是错误的