bug fix

app.py

@@ -54,23 +54,28 @@ def prepare_dataset(parameters, df, rain, temperature, datepicker):
     upperbound = datepicker + datetime.timedelta(days = 30) 
 
     df = df.loc[(df["Date"].dt.date>lowerbound) & (df["Date"].dt.date<=upperbound)]
-    dates = df["Date"]
     
     df = TimeSeriesDataSet.from_parameters(parameters, df)
-    return df.to_dataloader(train=False, batch_size=256,num_workers = 0), dates
+    return df.to_dataloader(train=False, batch_size=256,num_workers = 0)
 
+@st.cache_data
 def predict(model, dataloader):
     out = model.predict(dataloader, mode="raw", return_x=True, return_index=True)
     preds = raw_preds_to_df(out, quantiles = None)
-    return preds[["Group", "pred"]]
+    
+    date_list = [datepicker + datetime.timedelta(days=x) for x in range(30)]
+    preds["date_imputed"] = date_list
+    
+    return preds[["date_imputed", "Group", "pred"]]
 
 
-def generate_plot(df, dates, preds):
+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') 
     # 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(dates, preds.loc[preds['Group'] == '4', 'pred'], color = 'red')
+    axs[0, 0].plot(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
     axs[0, 0].set_title('Article Group 1')
     
     axs[0, 1].scatter(df.loc[df['Group'] == '7', 'Date'], df.loc[df['Group'] == '7', 'sales'], color='blue', marker='o')
@@ -141,9 +146,9 @@ def main():
     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))
 
     if st.button("Forecast Sales", type="primary"):
-        dataloader, dates = prepare_dataset(parameters, df, rain, temperature, datepicker)
+        dataloader = prepare_dataset(parameters, df, rain, temperature, datepicker)
         preds = predict(model, dataloader)
-        generate_plot(df, dates, preds)
+        generate_plot(df, preds)
 
 if __name__ == '__main__':
     main()