Shift value by month within a group

Question

Consider the following data frame:

import pandas as pd

df = pd.DataFrame({
    "id": [0, 0, 0, 1, 1, 1, 1, 2, 2, 2],
    "date": ["2017-01-31", "2017-02-28", "2017-03-31", "2017-01-31", "2017-03-31", "2017-04-30", "2017-05-31", "2017-01-31", "2017-03-31", "2017-05-31"],
    "value": [10., 12., 15., 8., 11., 15., 17., 6., 14., 15.]
})

df["date"] = pd.to_datetime(df["date"], format="%Y-%m-%d")

I want to create monthly shifted value columns within each id group, where the monthly shifts are specified in a list and can also be negative (meaning past and future shifts should be allowed).

Desired result:

from pandas import Timestamp
from numpy import nan

data={
    'id': {0: 0, 1: 0, 2: 0, 3: 1, 4: 1, 5: 1, 6: 1, 7: 2, 8: 2, 9: 2}, 
    'date': {
         0: Timestamp('2017-01-31 00:00:00'), 
         1: Timestamp('2017-02-28 00:00:00'), 
         2: Timestamp('2017-03-31 00:00:00'), 
         3: Timestamp('2017-01-31 00:00:00'), 
         4: Timestamp('2017-03-31 00:00:00'), 
         5: Timestamp('2017-04-30 00:00:00'), 
         6: Timestamp('2017-05-31 00:00:00'), 
         7: Timestamp('2017-01-31 00:00:00'), 
         8: Timestamp('2017-03-31 00:00:00'), 
         9: Timestamp('2017-05-31 00:00:00')
    }, 
    'value': {0: 10.0, 1: 12.0, 2: 15.0, 3: 8.0, 4: 11.0, 5: 15.0, 6: 17.0, 7: 6.0, 8: 14.0, 9: 15.0}, 
    'value_1': {0: 12.0, 1: 15.0, 2: nan, 3: nan, 4: 15.0, 5: 17.0, 6: nan, 7: nan, 8: nan, 9: nan}, 
    'value_2': {0: 15.0, 1: nan, 2: nan, 3: 11.0, 4: 17.0, 5: nan, 6: nan, 7: 14.0, 8: 15.0, 9: nan}
}

df = pd.DataFrame(data=data)

    id  date        value   value_1  value_2
0   0   2017-01-31  10.0    12.0     15.0
1   0   2017-02-28  12.0    15.0     NaN
2   0   2017-03-31  15.0    NaN      NaN
3   1   2017-01-31  8.0     NaN      11.0
4   1   2017-03-31  11.0    15.0     17.0
5   1   2017-04-30  15.0    17.0     NaN
6   1   2017-05-31  17.0    NaN      NaN
7   2   2017-01-31  6.0     NaN      14.0
8   2   2017-03-31  14.0    NaN      15.0
9   2   2017-05-31  15.0    NaN      NaN

In the data frame above, the columns value_1 and value_2 shall be created.

My approach so far:

from pandas.tseries.offsets import MonthEnd

shifts = [1, 2]

tmp = df.copy()
for shift in shifts:
    tmp_shifed = tmp.rename(columns={"value": f"value_{shift}"}).assign(date=df["date"] + MonthEnd(-1 * shift))
    df = df.merge(tmp_shifed, on=["id", "date"], how="left")

It works but I am sure there is a better way to achieve this. Note that my data frame is quite large and that the shift list has the size of 7.

Any help is appreciated!

constantstranger · Answer

Here's another way to do what your question asks:

df = df.set_index('id')
for shift in [1, 2]:
    df = df.assign(shift_date=df.date + MonthEnd(shift))
    df[f'value_{shift}'] = df.set_index('shift_date', append=True).index.map(df.set_index('date', append=True).value)

df = df.drop(columns='shift_date').reset_index()

Explanation:

set id column to be the index
loop over desired shifts (in months)
set a temporary column shift_date that adds the current iteration's shift to the date column
append shift_date to the index and use map() to do a value lookup in the original dataframe (with index set to id, date) and populate a new column value_X where X has the value of the current iteration's shift in months
after the loop, drop the temporary column shift_date and use reset_index() to restore id as a column.

Output:

   id       date  value  value_1  value_2
0   0 2017-01-31   10.0     12.0     15.0
1   0 2017-02-28   12.0     15.0      NaN
2   0 2017-03-31   15.0      NaN      NaN
3   1 2017-01-31    8.0      NaN     11.0
4   1 2017-03-31   11.0     15.0     17.0
5   1 2017-04-30   15.0     17.0      NaN
6   1 2017-05-31   17.0      NaN      NaN
7   2 2017-01-31    6.0      NaN     14.0
8   2 2017-03-31   14.0      NaN     15.0
9   2 2017-05-31   15.0      NaN      NaN

UPDATE:

A different approach would be to identify the min and max dates in the dataframe and automatically create value_X columns for each month. This will create the minimum number of columns with at least one non-null entry:

dates = pd.date_range(df.date.min(), df.date.max(), freq='M')
df['sh'] = df.date.map({dt:i for i, dt in enumerate(dates)})
df = ( df
    .pivot('id','date','value')
    .reindex(columns=dates)
    .set_axis(['value' + (f'_{i}' if i else '') for i, col in enumerate(dates)], axis='columns')
    .join(df.set_index('id')[['date','sh']]) )
df = df[['date','sh'] + 
    [col for col in df.columns if col.startswith('value')]].set_index('date', append=True)
sh = df.pop('sh')
df = df.T
for col in df.columns:
    df[col] = df[col].shift(-sh[col])
df = df.T.reset_index()

Output:

   id       date  value  value_1  value_2  value_3  value_4
0   0 2017-01-31   10.0     12.0     15.0      NaN      NaN
1   0 2017-02-28   12.0     15.0      NaN      NaN      NaN
2   0 2017-03-31   15.0      NaN      NaN      NaN      NaN
3   1 2017-01-31    8.0      NaN     11.0     15.0     17.0
4   1 2017-03-31   11.0     15.0     17.0      NaN      NaN
5   1 2017-04-30   15.0     17.0      NaN      NaN      NaN
6   1 2017-05-31   17.0      NaN      NaN      NaN      NaN
7   2 2017-01-31    6.0      NaN     14.0      NaN     15.0
8   2 2017-03-31   14.0      NaN     15.0      NaN      NaN
9   2 2017-05-31   15.0      NaN      NaN      NaN      NaN

Shift value by month within a group

1 Answers