Numpy пользовательская функция Cumsum с верхним / нижним пределами?

Question

У меня есть список значений numpy / pandas:

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

Я хочу применить пользовательскую функцию cumsum, но я не нашел способа сделать это без циклов. Пользовательская функция устанавливает верхний предел 1 и нижний предел -1 для значений cumsum, если «добавить» к сумме выходит за эти пределы, «добавить» становится 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]

Answer 1

Начните с обычной суммы:

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 ² ). Это совершенно недопустимо.

Answer 2

Петли не обязательно нежелательны. Если производительность является проблемой, рассмотрим numba. Улучшение ~ 330x без существенного изменения логики:

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