pandas dataframe groupby categorical column but keep original index after explode

Question

Dummy Data:

df = pd.DataFrame({'id': [1, 1, 1, 2, 2, 2],
                   'num': [12, 14, 18, 10, 10 ,11]})

Now I want to calculate the gradient of num for each id . Hence:

df.groupby(['id'])['num'].apply(np.gradient).explode()

However, I am not sure how to preserve the original index. The fix I came up with is just ugly, I was wondering if there was a better way.

df['gradient'] = df.groupby(['id'])['num'].apply(np.gradient).explode()\
                                          .reset_index().set_index(df.index)['num']

Which yields the desired result:

    id  num gradient
0   1   12  2.0
1   1   14  3.0
2   1   18  4.0
3   2   10  0.0
4   2   10  0.5
5   2   11  1.0

Answer 1

groupby().transform seems to do the job:

df['gradient'] = df.groupby(['id'])['num'].transform(np.gradient)

Output:

   id  num  gradient
0   1   12       2.0
1   1   14       3.0
2   1   18       4.0
3   2   10       0.0
4   2   10       0.5
5   2   11       1.0