pandas interpolate date series

Viewed 40

I'm trying to optimize my code a bit (run-time as well as readability). Essentially I want to spread the days between stay_from to stay_until across calendar months, creating a new record for each calendar month.

Current code:

import pandas as pd
import numpy as np
from datetime import date
stays = {
    1: {'customer_id':1, 'stay_from': date(2021, 1, 1), 'stay_until': date(2021, 12, 31)},
    2: {'customer_id':2, 'stay_from': date(2021, 1, 1), 'stay_until': date(2021, 12, 31)},
    3: {'customer_id':1, 'stay_from': date(2021, 1, 1), 'stay_until': date(2021, 12, 31)}
}
stays = pd.DataFrame.from_dict(stays, orient='index')
stays = stays.astype({'stay_from': 'datetime64[ns]', 'stay_until': 'datetime64[ns]'})

df = stays.copy()

df['earnings_bop'] = - 1.
df['earnings_eop'] = 0.

# This creates an index for the 1st of each month following the start
# Want the 1st of the month of every month:
df['upr_idx_bom'] = df.apply(
    lambda x: pd.date_range(x.stay_from, x.stay_until + pd.DateOffset(months=1), freq='MS', closed='right'),
    axis=1)
# This creates an end-of-month index equivalent to start of the 1st day of the next month:
df['upr_idx_eom'] = df.apply(
    lambda x: pd.date_range(x.stay_from, x.stay_until + pd.DateOffset(months=1), freq='M', closed='left'), axis=1)

# This is a heavy step when 100'000 records
# The month-start index is joined to the from and until dates
# The earnings are interpolated along the index
# The index is reset and samples the month end figures only, and drops the start and end
df['upr_s'] = df.apply(
    lambda x:
    pd.Series([x.earnings_bop, x.earnings_eop], index=[x.stay_from, x.stay_until])
        .reindex(np.unique([x.stay_from, x.stay_until] + x.upr_idx_bom.tolist()))
        .interpolate('index')
        .reindex(x.upr_idx_bom).to_list(),
    axis=1)

df['stay_x'] = df.apply(lambda x: list(zip(x.upr_idx_bom, x.upr_s)), axis=1)
df = df.drop(columns=['stay_until', 'earnings_bop', 'earnings_eop', 'upr_idx_bom', 'upr_idx_eom', 'upr_s'])

df = df.explode('stay_x')
df[['earnings_month', 'earnings_remaining']] = pd.DataFrame(df.stay_x.to_list(), index=df.index)
df = df.drop(columns=['stay_x'])

This works fine so far, albeit it the interpolation step is quite slow. Now I also want to add a column 'earnings_current':

df['earnings_current'] = df.groupby(by=['customer_id', 'stay_from'])['earnings_remaining'].transform(lambda x: x.diff())

But this gives unexpected results where there is a duplicate record in the groupby. There are good reasons why there may be duplicate records.

Appreciate there are similar questions out there but looking for help on neat syntax in this specific case.

0 Answers
Related