So I have a data frame like this–
import pandas as pd import numpy as np df = pd.DataFrame([[1,22], [1,23], [1,44], [2, 33], [2, 55]], columns=['id', 'delay'])
id delay 0 1 22 1 1 23 2 1 44 3 2 33 4 2 55
What I am doing is grouping by id
and doing rolling operation on the delay
column like below–
k = [0.1, 0.5, 1] def f(d): d['new_delay'] = pd.Series([0,0]).append(d['delay']).rolling(window=3).apply(lambda x: np.sum(x*k)).iloc[2:] return d df.groupby(['id']).apply(f)
id delay new_delay 0 1 22 22.0 1 1 23 34.0 2 1 44 57.7 3 2 33 33.0 4 2 55 71.5
It is working just fine but I am curious whether .apply
on grouped data frame is vectorized or not. Since my dataset is huge, is there a better-vectorized way to do this kind of operation? Also I am curious if Python is single-threaded and I am running on CPU how pandas, numpy achieve vectorized calculation.
Advertisement
Answer
You can use strides for vectorized rolling
with GroupBy.transform
:
k = [0.1, 0.5, 1] 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) def f(d): return np.sum(rolling_window(np.append([0,0],d.to_numpy()), 3) * k, axis=1) df['new_delay'] = df.groupby('id')['delay'].transform(f) print (df) id delay new_delay 0 1 22 22.0 1 1 23 34.0 2 1 44 57.7 3 2 33 33.0 4 2 55 71.5