具有上限/下限的 Numpy 自定义 Cumsum 函数？

Posted 2023-03-11

【中文标题】具有上限/下限的 Numpy 自定义 Cumsum 函数？

【英文标题】：Numpy custom Cumsum function with upper/lower limits?

【发布时间】：2019-04-05 21:11:03

【问题描述】：

我有一个 numpy/pandas 值列表：

a = np.random.randint(-100, 100, 10000)
b = a/100

我想应用一个自定义的 cumsum 函数，但我还没有找到没有循环的方法。自定义函数设置 cumsum 值的上限为 1，下限为 -1，如果 sum 的“add”超出这些限制，则“add”变为 0。

如果总和在 -1 和 1 的限制之间，但“添加”的值会超出限制，则“添加”的值将变为 -1 或 1 的余数。

这是循环版本：

def cumsum_with_limits(values):
    cumsum_values = []
    sum = 0
    for i in values:
        if sum+i <= 1 and sum+i >= -1: 
            sum += i
            cumsum_values.append(sum)
        elif sum+i >= 1:
            d = 1-sum # Remainder to 1
            sum += d
            cumsum_values.append(sum)
        elif sum+i <= -1:
            d = -1-sum # Remainder to -1
            sum += d
            cumsum_values.append(sum)

    return cumsum_values

有什么方法可以矢量化吗？我需要在大型数据集上运行此功能，而性能是我当前的问题。感谢任何帮助！

---

更新：稍微修正了代码，并对输出进行了一点澄清： 使用 np.random.seed(0)，前 6 个值是：

b = [0.72, -0.53, 0.17, 0.92, -0.33, 0.95]

预期输出：

o = [0.72, 0.19, 0.36, 1, 0.67, 1]

【问题讨论】：

如果我理解正确，cumsum_with_limits 会为您提供值列表，使它们的 cumsum 低于 -1 或高于 +1，对吧？那么，你想要的是数字数组，而不是 cumsum 本身，对吗？

是的，正确，输出是值列表而不是 cumsum 本身，例如[0, 0.3, 0.6, , 0.8, 1, 1, 1, 1, 1, 1, 0.6, 0.4, 0.1, -0.3, -0.6, -1, -1, -1, -1, -0.5],它不能超过 1 或 -1

您能否展示一个输入和预期输出的具体示例，展示您处理所有边界的方式？

我很确定有一些版本的 reduceat 可以为你做到这一点。只需要弄清楚如何表达它。

是的，输入 [0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]，输出 [0.2, 0.4, 0.6, 0.8, 1, 1, 1]

【参考方案1】：

从常规的 cumsum 开始：

b = ...
s = np.cumsum(b)

找到第一个剪辑点：

i = np.argmax((s[0:] > 1) | (s[0:] < -1))

调整以下所有内容：

s[i:] += (np.sign(s[i]) - s[i])

冲洗并重复。这仍然需要一个循环，但只是在调整点上进行，通常预计会比数组大小的总数小很多。

b = ...
s = np.cumsum(b)
while True:
    i = np.argmax((s[0:] > 1) | (s[0:] < -1))
    if np.abs(s[i]) <= 1:
        break
    s[i:] += (np.sign(s[i]) - s[i])

我仍然没有找到一种方法来完全预先计算调整点，所以我不得不猜测 numba 解决方案会比这更快，即使它是用 numba 编译的。

从np.seed(0) 开始，您的原始示例有 3090 个调整点，大约是 1/3。不幸的是，使用所有临时数组和额外的总和，这使得我的解决方案的算法复杂度趋于 O(n²)。这是完全不能接受的。

【参考方案2】：

循环不一定是不可取的。如果性能是一个问题，请考虑numba。在没有实质性改变您的逻辑的情况下有约 330 倍的改进：

from numba import njit

np.random.seed(0)
a = np.random.randint(-100, 100, 10000)
b = a/100

@njit
def cumsum_with_limits_nb(values):
    n = len(values)
    res = np.empty(n)
    sum_val = 0
    for i in range(n):
        x = values[i]
        if (sum_val+x <= 1) and (sum_val+x >= -1):
            res[i] = x
            sum_val += x
        elif sum_val+x >= 1:
            d = 1-sum_val # Remainder to 1
            res[i] = d
            sum_val += d
        elif sum_val+x <= -1:
            d = -1-sum_val # Remainder to -1
            res[i] = d
            sum_val += d
    return res

assert np.isclose(cumsum_with_limits(b), cumsum_with_limits_nb(b)).all()

如果你不介意牺牲一些性能，你可以更简洁地重写这个循环：

@njit
def cumsum_with_limits_nb2(values):
    n = len(values)
    res = np.empty(n)
    sum_val = 0
    for i in range(n):
        x = values[i]
        next_sum = sum_val + x
        if np.abs(next_sum) >= 1:
            x = np.sign(next_sum) - sum_val
        res[i] = x
        sum_val += x
    return res

与nb2 的性能相似，这里有一个替代方案（感谢@jdehesa）：

@njit
def cumsum_with_limits_nb3(values):
    n = len(values)
    res = np.empty(n)
    sum_val = 0
    for i in range(n):
        x = min(max(sum_val + values[i], -1) , 1) - sum_val
        res[i] = x
        sum_val += x
    return res

性能比较：

assert np.isclose(cumsum_with_limits(b), cumsum_with_limits_nb(b)).all()
assert np.isclose(cumsum_with_limits(b), cumsum_with_limits_nb2(b)).all()
assert np.isclose(cumsum_with_limits(b), cumsum_with_limits_nb3(b)).all()

%timeit cumsum_with_limits(b)      # 12.5 ms per loop
%timeit cumsum_with_limits_nb(b)   # 40.9 µs per loop
%timeit cumsum_with_limits_nb2(b)  # 54.7 µs per loop
%timeit cumsum_with_limits_nb3(b)  # 54 µs per loop

【讨论】：

@FrancWeser，没问题。这可能是因为您提到您不想要循环。但有时不一定是个坏主意。对于它的价值，一个专门的 NumPy 方法可能更可取，所以在其他人试图弄清楚之前，请不要将此标记为已接受。

我知道我确实误解了这个问题，所以投反对票对我来说很好，但为什么这个答案也会被投反对票？

@W-B 不确定，但我认为您的答案是正确的。用户在上面发布了一个示例输入输出，它实际上与您正在执行的操作相匹配：P。非常混乱

@user3483203 嗯，这就是我删除它的原因:-) 我正在远离不清楚的问题

这是一个惊人的性能改进。非常感谢！好奇是否还有办法在不循环的情况下解决它