Blogpost-Heineken5

Blogpost-Heineken5.py

class LagCreator(object):
    def __init__(self, groupby_keys, lagdict, drop_cols=True, suffix=''):
        self.groupby_keys = groupby_keys
        self.lagdict = lagdict
        self.drop_cols = drop_cols
        self.suffix = suffix
    def transform(self, X):
        cols_to_drop = list(set([item for sublist in 
                            self.lagdict.values() for item in sublist]))
        for week in self.lagdict:
            df_add = X.groupby(self.groupby_keys)[self.lagdict[week]].shift(-week)
            new_colnames = [col + str(week).replace('-', '_') + self.suffix for col in df_add.columns]
            X[new_colnames] = df_add
        if self.drop_cols:
            X = X.drop(cols_to_drop, axis=1)
        return X
static_cols = ['cpi', 'unemployment', 'fuel_price', 'markdown1', 'markdown2', 'markdown3','markdown4', 'markdown5']
lagdict = {
    -2: ['weekly_sales'],
    -1: ['weekly_sales', 'weekly_sales_sm1', 'weekly_sales_sm3'],
    0: ['weekly_sales', 'isholiday', 'temperature'] + static_cols,
    1: ['weekly_sales', 'isholiday', 'temperature', 'weekly_sales_1ly'],
    2: ['weekly_sales', 'isholiday', 'temperature', 'weekly_sales_1ly']
}
lagdict_yearly = {
    -1: ['weekly_sales']
}