Pandas: vectorize local range operations (max & sum for [i:i+2] rows)

Question

I'm looking to do calculations in a local range for each row in a dataframe while avoiding a slow for loop. For example, for each row in the data below I want to find the maximum temperature within the next 3 days (including current day) and the total amount of rain within the next 3 days:

Day Temperature Rain
0   30          4
1   31          14
2   31          0
3   30          0
4   33          5
5   34          0
6   32          0
7   33          2
8   31          5
9   29          9

The ideal output would then be the new columns as in the table below. TempMax of Day 0 shows the highest temperature between Day 0 and Day 2, RainTotal shows the sum of rain between Day 0 and Day 2:

Day  Temperature  Rain  TempMax  RainTotal
0    30           4     31       18
1    31           14    31       14
2    31           0     33       5
3    30           0     34       5
4    33           5     34       5
5    34           0     34       2
6    32           0     33       7
7    33           2     33       16
8    31           5     31       14
9    29           9     29       9

Currently I'm using a for loop:

  # Make empty arrays to store each row's max & sum values
  temp_max = np.zeros(len(df))
  rain_total = np.zeros(len(df))

  # Loop through the df and do operations in the local range [i:i+2]
  for i in range(len(df)):
    temp_max[i] = df['Temperature'].iloc[i:i+2].max()
    rain_total = df['Rain'].iloc[i:i+2].sum()

  # Insert the arrays to df
  df['TempMax'] = temp_max
  df['RainTotal'] = rain_total

The for loop gets the job done but takes 50 minutes with my dataframe. Any chance this can be vecrotized or made faster some other way?

Thanks a bunch!

jezrael · Accepted Answer · 2019-05-25T14:57:56.357

Use Series.rolling with change order by indexing and max with sum:

df['TempMax'] = df['Temperature'].iloc[::-1].rolling(3, min_periods=1).max()
df['RainTotal'] = df['Rain'].iloc[::-1].rolling(3, min_periods=1).sum()
print (df)
   Day  Temperature  Rain  TempMax  RainTotal
0    0           30     4     31.0       18.0
1    1           31    14     31.0       14.0
2    2           31     0     33.0        5.0
3    3           30     0     34.0        5.0
4    4           33     5     34.0        5.0
5    5           34     0     34.0        2.0
6    6           32     0     33.0        7.0
7    7           33     2     33.0       16.0
8    8           31     5     31.0       14.0
9    9           29     9     29.0        9.0

Another faster solution with strides in numpy for 2d array and then use numpy.nanmax with numpy.nansum:

n = 2
t = np.concatenate([df['Temperature'].values, [np.nan] * (n)])
r = np.concatenate([df['Rain'].values, [np.nan] * (n)])

def rolling_window(a, window):
    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)

df['TempMax'] = np.nanmax(rolling_window(t, n + 1), axis=1)
df['RainTotal'] = np.nansum(rolling_window(r, n + 1), axis=1)
print (df)
   Day  Temperature  Rain  TempMax  RainTotal
0    0           30     4     31.0       18.0
1    1           31    14     31.0       14.0
2    2           31     0     33.0        5.0
3    3           30     0     34.0        5.0
4    4           33     5     34.0        5.0
5    5           34     0     34.0        2.0
6    6           32     0     33.0        7.0
7    7           33     2     33.0       16.0
8    8           31     5     31.0       14.0
9    9           29     9     29.0        9.0

Performance:

#[100000 rows x 3 columns]
df = pd.concat([df] * 10000, ignore_index=True)

In [23]: %%timeit
    ...: df['TempMax'] = np.nanmax(rolling_window(t, n + 1), axis=1)
    ...: df['RainTotal'] = np.nansum(rolling_window(r, n + 1), axis=1)
    ...: 
8.36 ms ± 165 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [24]: %%timeit
    ...: df['TempMax'] = df['Temperature'].iloc[::-1].rolling(3, min_periods=1).max()
    ...: df['RainTotal'] = df['Rain'].iloc[::-1].rolling(3, min_periods=1).sum()
    ...: 
20.4 ms ± 1.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Thank you! This reduced the time from 50 minutes to 3 seconds. Amazing! — Tuppitappi, May 25 '19 at 16:04

Divakar · Answer 2 · 2019-05-26T05:55:48.057

For the case when Day has data for all successive days, we can employ fast NumPy and SciPy tools to our rescue -

from scipy.ndimage.filters import maximum_filter1d

N = 2 # window length
temp = df['Temperature'].to_numpy()
rain = df['Rain'].to_numpy()
df['TempMax'] = maximum_filter1d(temp,N+1,origin=-1,mode='nearest')
df['RainTotal'] = np.convolve(rain,np.ones(N+1,dtype=int))[N:]

Sample output -

In [27]: df
Out[27]: 
   Day  Temperature  Rain  TempMax  RainTotal
0    0           30     4       31         18
1    1           31    14       31         14
2    2           31     0       33          5
3    3           30     0       34          5
4    4           33     5       34          5
5    5           34     0       34          2
6    6           32     0       33          7
7    7           33     2       33         16
8    8           31     5       31         14
9    9           29     9       29          9

This works incredibly fast too, thanks for the solution! I do however prefer the .rolling() method of the previous answer since it is fully Pandas and does not require additional imports. — Tuppitappi, May 26 '19 at 19:59

Pandas: vectorize local range operations (max & sum for [i:i+2] rows)

2 Answers2

Linked