bug fix

app.py

@@ -49,27 +49,23 @@ 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)
+    return df.to_dataloader(train=False, batch_size=256,num_workers = 0), dates
 
-def predict(model, dataloader): 
-    return model.predict(dataloader, mode="raw", return_x=True, return_index=True)
-    
- #on_click=None,
-
-# %%
-#preds = raw_preds_to_df(out, quantiles = None)
+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"]]
 
-#preds = preds.merge(data_selected[['time_idx','Group','Branch','sales','weight','Date','MTXWTH_Day_precip','MTXWTH_Temp_max','MTXWTH_Temp_min']],how='left',left_on=['pred_idx','Group','Branch'],right_on=['time_idx','Group','Branch'])
-#preds.rename(columns={'time_idx_x':'time_idx'},inplace=True)
-#preds.drop(columns=['time_idx_y'],inplace=True)
 
-def generate_plot(df): #, predictions):
+def generate_plot(df, dates, preds)
     fig, axs = plt.subplots(2, 2, figsize=(8, 6))
     
     # Plot scatter plots for each group
-    axs[0, 0].scatter(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'sales'], color='red', marker='o')
+    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].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')
@@ -137,8 +133,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"):
-        converted_data = prepare_dataset(parameters, df, rain, temperature, datepicker)
-        generate_plot(df)
+        dataloader, dates = prepare_dataset(parameters, df, rain, temperature, datepicker)
+        preds = predict(model, dataloader)
+        generate_plot(df, dates, preds)
 
 if __name__ == '__main__':
     main()