等于或大于 0 时对三个连续数求和 - Python

Question

I am using numpy in Python

I have an array of numbers, for example:

arr = np.array([0.1, 1, 1.2, 0.5, -0.3, -0.2, 0.1, 0.5, 1)

If i is a position in the array, I want to create a function which creates a running sum of i and the two previous numbers, but only accumulating the number if it is equal to or greater than 0.

In other words, negative numbers in the array become equal to 0 when calculating the three number running sum.

For example, the answer I would be looking for here is

2.3, 2.7, 1.7, 0.5, 0.1, 0.6, 1.6

The new array has two elements less than the original array as the calculation can't be completed for the first two number.

Thank you !

Answer 1

As Dani Mesejo answered, you can use stride tricks. You can either use clip or boolean indexing to handle the <0 elements. I have explained how stride tricks work below -

1. arr[arr<0]=0 sets all elements below 0 as 0
2. as_strided takes in the array, the expected shape of the view (7,3) and the number of strides in the respective axes, (8,8) . This is the number of bytes you have to move in axis0 and axis1 respectively to access the next element. Eg If you want to move every 2 elements, then you can set it to (16,8) . This means you would move 16 bytes each time to get the element in axis0 (which is 0.1->1.2->0->0.1->.. , till a shape of 7 ) and 8 bytes each time to get element in axis1 (which is 0.1->1->1.2 , till a shape of 3 )

Use this function with caution. Always use x.strides to define the strides parameter to avoid corrupting memory!

3. Lastly, sum this array view over axis=1 to get your rolling sum.

arr = np.array([0.1, 1, 1.2, 0.5, -0.3, -0.2, 0.1, 0.5, 1])
w = 3  #rolling window

arr[arr<0]=0

shape = arr.shape[0]-w+1, w  #Expected shape of view (7,3)
strides = arr.strides[0], arr.strides[0] #Strides (8,8) bytes
rolling = np.lib.stride_tricks.as_strided(arr, shape=shape, strides=strides)

rolling_sum = np.sum(rolling, axis=1)
rolling_sum

array([2.3, 2.7, 1.7, 0.5, 0.1, 0.6, 1.6])

Answer 2

You could clip , roll and sum :

import numpy as np


def rolling_window(a, window):
    """Recipe from https://stackoverflow.com/q/6811183/4001592"""
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)


a = np.array([0.1, 1, 1.2, 0.5, -0.3, -0.2, 0.1, 0.5, 1])

res = rolling_window(np.clip(a, 0, a.max()), 3).sum(axis=1)
print(res)

Output

[2.3 2.7 1.7 0.5 0.1 0.6 1.6]

Answer 3

You may use np.correlate to sweep an array of 3 ones over the clipped of arr to get desired output

In [20]: np.correlate(arr.clip(0), np.ones(3), mode='valid')
Out[20]: array([2.3, 2.7, 1.7, 0.5, 0.1, 0.6, 1.6])

Answer 4

arr = np.array([0.1, 1, 1.2, 0.5, -0.3, -0.2, 0.1, 0.5, 1])

def sum_3(x):
    collector = []
    
    for i in range(len(arr)-2):            
        collector.append(sum(arr[i:i+3][arr[i:i+3]>0]))
    return collector

#output

[2.3, 2.7, 1.7, 0.5, 0.1, 0.6, 1.6]

Easiest and most comprehensible way. The collector will append the sum of the 3 consecutive numbers if their indices are True otherwise, they are all turned to 0 s.

The method is not general, it is for 3 consecutives but you can adapt it.

def sum_any(x,n):
    collector = []
    for i in range(len(arr)-(n-1)):            
        collector.append(sum(arr[i:i+n][arr[i:i+n]>0]))
    return collector

Answer 5

Masked arrays and view_as_windows (which uses numpy strides under the hood) are built for this purpose:

from skimage.util import view_as_windows
arr = view_as_windows(arr, 3)
arr2 = np.ma.masked_array(arr, arr<0).sum(-1)

output:

[2.3 2.7 1.7 0.5 0.1 0.6 1.6]