bug fix

app.py

@@ -57,20 +57,13 @@ def prepare_dataset(parameters, df, rain, temperature, datepicker, mapping):
 def predict(_model, _dataloader, datepicker):
     out = _model.predict(_dataloader, mode="raw", return_x=True, return_index=True)
     preds = raw_preds_to_df(out, quantiles = None)
-
-    def add_dates(group):
-        group["date_imputed"] = [datepicker + datetime.timedelta(days=x) for x in range(30)]
-        return group
-    
-    preds["date_imputed"] = preds.groupby("Group").pred.transform(add_dates)
-    
-    return preds[["date_imputed", "Group", "pred"]]
+    return preds[["pred_idx", "Group", "pred"]
 
 
 def generate_plot(df, preds):
     fig, axs = plt.subplots(2, 2, figsize=(8, 6))
     
-    df = pd.merge(df, pred, left_on='Date', right_on='date_imputed') 
+    df = pd.merge(df, preds, left_on=["time_idx", "Group"], right_on=["pred_idx", "Group"], how = "left")
     # Plot scatter plots for each group
     axs[0, 0].scatter(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'sales'], color='grey', marker='o')
     axs[0, 0].plot(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
@@ -145,10 +138,10 @@ def main():
     
     temperature = st.slider('Change in Temperature', min_value=-10.0, max_value=10.0, value=0.0, step=0.25)
     
-    datepicker = st.date_input("Start of Forecast", datetime.date(2022, 12, 24), min_value=datetime.date(2022, 6, 26) + datetime.timedelta(days = 35), max_value=datetime.date(2023, 6, 26) - datetime.timedelta(days = 30))
+    datepicker = st.date_input("Start of Forecast", datetime.date(2022, 10, 24), min_value=datetime.date(2022, 6, 26) + datetime.timedelta(days = 35), max_value=datetime.date(2023, 6, 26) - datetime.timedelta(days = 30))
 
     if st.button("Forecast Sales", type="primary"):
-        dataloader = prepare_dataset(parameters, df, rain, temperature, datepicker, rain_mapping)
+        dataloader = prepare_dataset(parameters, df.copy(), rain, temperature, datepicker, rain_mapping)
         preds = predict(model, dataloader, datepicker)
         generate_plot(df, preds)