CODEWITHSUNDEEP
pythonpandasrolling-computation
Daven
Daven

Reputation: 131

Creating a rolling sum column that resets once it reaches a threshold

This question is unlike other similar ones that I could find because I am trying to combine a lookback window and a threshold into one rolling sum. I'm not actually sure what I'm trying to do is achievable in one step:

I have a pandas dataframe with a datetime column and a value column. I have created a column that sums the value column (V) over a rolling time window. However I would like this rolling sum to reset to 0 once it reaches a certain threshold.

I don't know if it's possible to do this in one column manipulation step since there are two conditions at play at each step in the sum- the lookback window and the threshold. If anyone has any ideas about if this is possible and how I might be able to achieve it please let me know. I know how to do this iteratively however it is very very slow (my dataframe has >1 million entries).

Example:

Lookback time: 3 minutes

Threshold: 3

+---+-----------------------+-------+--------------------------+
|   |           myDate      |   V   | rolling | desired_column |
+---+-----------------------+-------+---------+----------------+
| 1 | 2020-04-01 10:00:00   | 0     |  0      |       0        |   
| 2 | 2020-04-01 10:01:00   | 1     |  1      |       1        | 
| 3 | 2020-04-01 10:02:00   | 2     |  3      |       3        | 
| 4 | 2020-04-01 10:03:00   | 1     |  4      |       1        | 
| 5 | 2020-04-01 10:04:00   | 0     |  4      |       1        | 
| 6 | 2020-04-01 10:05:00   | 4     |  7      |       5        | 
| 7 | 2020-04-01 10:06:00   | 1     |  6      |       1        | 
| 8 | 2020-04-01 10:07:00   | 1     |  6      |       2        | 
| 9 | 2020-04-01 10:08:00   | 0     |  6      |       0        |       
| 10| 2020-04-01 10:09:00   | 3     |  5      |       5        | 
+---+-----------------------+-------+---------+----------------+

In this example the sum rulling sum will not take into account any values on or before a row that breaches (or is equal to) the threshold of 3.

Upvotes: 3

Views: 1016

Answers (2)

RichieV
RichieV

Reputation: 5183

The following approach is not memory efficient by any means, but it should be faster than looping. It assumes time is continuous in order to delegate to numpy methods, otherwise you can include the missing times before calling.

def rolling_window(a, window):
    b = np.concatenate((np.zeros(window-1), a)) # only for 1d
    return np.array([b[..., i:i+window] for i in range(a.size)])


def dynamic_window(w: np.array, reset):
    regions = np.hstack([
        np.zeros((w.shape[0], 1)),
        np.cumsum(w, axis=-1)[:, :-1]
    ]) // reset
    return w * (regions == regions[:, -1][:, np.newaxis])

Use it as

# sample df
# please always provide a callable line of code
# you could get it with `df.head(10).to_dict('split')`
df = pd.DataFrame({
    'myDate': pd.date_range('2020-04-01 10:00', periods=10, freq='T'),
    'V': [0, 1, 2, 1, 0, 4, 1, 1, 0, 3]
})
# include all time increments
df = pd.concat([
    df,
    pd.DataFrame(pd.date_range(df['myDate'].min(),
        df['myDate'].max(), freq='T'), columns=['myDate'])
]).drop_duplicates(subset=['myDate']).fillna(0).sort_values('myDate')

df['4min_sum'] = df.rolling('4min', on='myDate')['V'].sum()

# use the functions
df['desired_column'] = dynamic_window(
    rolling_window(df['V'].to_numpy(), 4),
    3).sum(axis=-1)

Output

               myDate    V  4min_sum  desired_column
0 2020-04-01 10:00:00  0.0       0.0             0.0
1 2020-04-01 10:01:00  1.0       1.0             1.0
2 2020-04-01 10:02:00  2.0       3.0             3.0
3 2020-04-01 10:03:00  1.0       4.0             1.0
4 2020-04-01 10:04:00  0.0       4.0             1.0
5 2020-04-01 10:05:00  4.0       7.0             4.0
6 2020-04-01 10:06:00  1.0       6.0             1.0
7 2020-04-01 10:07:00  1.0       6.0             2.0
8 2020-04-01 10:08:00  0.0       6.0             0.0
9 2020-04-01 10:09:00  3.0       5.0             5.0

Notice how at 10:05 it outputs 4 instead of the 5 you have in your expected output. According to your logic it should be 4; that window contains [2, 1, 0, 4] and, since the two first numbers sum 3, the window should reset and return 0 + 4.

Upvotes: 1

Serge Ballesta
Serge Ballesta

Reputation: 149185

I could not find a vectorized way to do a reset to 0 every time a threshold value is reached.

But the underlying container of a Pandas column is a numpy array, and iterating a numpy array takes an acceptable time. So I would to:

arr = np.zeros(len(df), dtype='int')
cum = 0
src = df['V'].values
dt = df['myDate'].values
start = 0
for i in range(len(df)):
    cum += src[i]
    while dt[start] < dt[i] - np.timedelta64(4, 'm'):
        cum -= src[start]
        start +=1
    arr[i] = cum
    if cum >=3:
        cum = 0
        start = i

df['desired_column'] = arr

It gives :

                myDate  V  rolling  desired_column
1  2020-04-01 10:00:00  0        0               0
2  2020-04-01 10:01:00  1        1               1
3  2020-04-01 10:02:00  2        3               3
4  2020-04-01 10:03:00  1        4               1
5  2020-04-01 10:04:00  0        4               1
6  2020-04-01 10:05:00  4        7               5
7  2020-04-01 10:06:00  1        6               1
8  2020-04-01 10:07:00  1        6               2
9  2020-04-01 10:08:00  0        6               2
10 2020-04-01 10:09:00  3        5               5

It only takes a few seconds for an array of length 1000000 on my i5 machine (~90s for 10 000 000)

Upvotes: 1

Related Questions