Эффективное применение пользовательских функций к группам в Pandas

Question

Этот вопрос касается эффективного применения пользовательской функции к логическим группам строк в кадре данных Pandas, которые совместно используют значение в некотором столбце.

Рассмотрим следующий пример кадра данных:

sID = [1,1,1,2,4,4,5,5,5]
data = np.random.randn(len(sID))
dates = pd.date_range(start='1/1/2018', periods=len(sID))

mydf = pd.DataFrame({"subject_id":sID, "data":data, "date":dates})
mydf['date'][5] += pd.Timedelta('2 days')

, который выглядит следующим образом:

       data       date  subject_id
0  0.168150 2018-01-01           1
1 -0.484301 2018-01-02           1
2 -0.522980 2018-01-03           1
3 -0.724524 2018-01-04           2
4  0.563453 2018-01-05           4
5  0.439059 2018-01-08           4
6 -1.902182 2018-01-07           5
7 -1.433561 2018-01-08           5
8  0.586191 2018-01-09           5

Представьте себе, что для каждого subject_id мы хотим вычесть из каждой даты первую дату, встреченную для этого subject_id.Сохранение результата в новом столбце days_elapsed, результат будет выглядеть следующим образом:

       data       date  subject_id  days_elapsed
0  0.168150 2018-01-01           1             0
1 -0.484301 2018-01-02           1             1
2 -0.522980 2018-01-03           1             2
3 -0.724524 2018-01-04           2             0
4  0.563453 2018-01-05           4             0
5  0.439059 2018-01-08           4             3
6 -1.902182 2018-01-07           5             0
7 -1.433561 2018-01-08           5             1
8  0.586191 2018-01-09           5             2

Один из естественных способов сделать это - использовать groupby и apply:

g_df = mydf.groupby('subject_id')
mydf.loc[:, "days_elapsed"] = g_df["date"].apply(lambda x: x - x.iloc[0]).astype('timedelta64[D]').astype(int)

Однако, если количество групп (идентификаторов субъектов) велико (например, 10 ^ 4) и, скажем, только в 10 раз меньше длины кадра данных, эта очень простая операция действительно медленная.

Есть ли более быстрый метод?

---

PS: я также попытался установить индекс на subject_id и затем использовать следующее понимание списка:

def get_first(series, ind):
    "Return the first row in a group within a series which (group) potentially can span multiple rows and corresponds to a given index"
    group = series.loc[ind]

    if hasattr(group, 'iloc'):
        return group.iloc[0]
    else: # this is for indices with a single element
        return group

hind_df = mydf.set_index('subject_id')
A = pd.concat([hind_df["date"].loc[ind] - get_first(hind_df["date"], ind) for ind in np.unique(hind_df.index)])

Однако, это еще медленнее.

Answer 1

mydf['days_elapsed'] = (mydf['date'] - mydf.groupby(['subject_id'])['date'].transform('min')).dt.days

Answer 2

Вы можете использовать GroupBy + transform с first.Это должно быть более эффективным, поскольку позволяет избежать дорогостоящих lambda вызовов функций.

Вы также можете увидеть улучшение производительности, работая с массивом NumPy через pd.Series.values:

first = df.groupby('subject_id')['date'].transform('first').values

df['days_elapsed'] = (df['date'].values - first).astype('timedelta64[D]').astype(int)

print(df)

   subject_id      data       date  days_elapsed
0           1  1.079472 2018-01-01             0
1           1 -0.197255 2018-01-02             1
2           1 -0.687764 2018-01-03             2
3           2  0.023771 2018-01-04             0
4           4 -0.538191 2018-01-05             0
5           4  1.479294 2018-01-08             3
6           5 -1.993196 2018-01-07             0
7           5 -2.111831 2018-01-08             1
8           5 -0.934775 2018-01-09             2