bug fix

app.py

@@ -40,7 +40,7 @@ def raw_preds_to_df(raw,quantiles = None):
 
 def prepare_dataset(parameters, df, rain, temperature, datepicker, mapping):
     if rain != "Default":
-        df["MTXWTH_Day_precip"] = rain_mapping[rain]
+        df["MTXWTH_Day_precip"] = mapping[rain]
     
     df["MTXWTH_Temp_min"] = df["MTXWTH_Temp_min"] + temperature
     df["MTXWTH_Temp_max"] = df["MTXWTH_Temp_max"] + temperature
@@ -59,37 +59,34 @@ def predict(_model, _dataloader, datepicker):
     preds = raw_preds_to_df(out, quantiles = None)
     return preds[["pred_idx", "Group", "pred"]]
 
-
-def generate_plot(df, preds):
-    fig, axs = plt.subplots(2, 2, figsize=(8, 6))
-    
+def update_plot(df, preds):
     df = pd.merge(df, preds, left_on=["time_idx", "Group"], right_on=["pred_idx", "Group"], how = "left")
     df = df[~df["pred"].isna()]
     df[["sales", "pred"]] = df[["sales", "pred"]].replace(0.0, np.nan)
+
+    axs[0, 0].plot(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
+    axs[0, 1].plot(df.loc[df['Group'] == '7', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
+    axs[1, 0].plot(df.loc[df['Group'] == '1', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
+    axs[1, 1].plot(df.loc[df['Group'] == '6', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
+    return st.pyplot(fig)
+
+@st.cache_resource
+def generate_plot():
+    fig, axs = plt.subplots(2, 2, figsize=(8, 6))
+
     # Plot scatter plots for each group
     axs[0, 0].plot(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'sales'], color='grey')
-    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, 0].set_title('Article Group 1')
     axs[0, 1].plot(df.loc[df['Group'] == '7', 'Date'], df.loc[df['Group'] == '7', 'sales'], color='grey')
-    axs[0, 1].plot(df.loc[df['Group'] == '7', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
     axs[0, 1].set_title('Article Group 2')
     
     axs[1, 0].plot(df.loc[df['Group'] == '1', 'Date'], df.loc[df['Group'] == '1', 'sales'], color='grey')
-    axs[1, 0].plot(df.loc[df['Group'] == '1', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
     axs[1, 0].set_title('Article Group 3')
     
     axs[1, 1].plot(df.loc[df['Group'] == '6', 'Date'], df.loc[df['Group'] == '6', 'sales'], color='grey')
-    axs[1, 1].plot(df.loc[df['Group'] == '6', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
     axs[1, 1].set_title('Article Group 4')
-    
-    # Adjust spacing between subplots
     plt.tight_layout()
-    
-    #for ax in axs.flat:
-    #    ax.set_xlim(df['Date'].min(), df['Date'].max())
-    #    ax.set_ylim(df['sales'].min(), df['sales'].max())
-    
     st.pyplot(fig)
 
 @st.cache_data
@@ -145,10 +142,12 @@ def main():
     
     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))
 
+    generate_plot()
+
     if st.button("Forecast Sales", type="primary"):
         dataloader = prepare_dataset(parameters, df.copy(), rain, temperature, datepicker, rain_mapping)
         preds = predict(model, dataloader, datepicker)
-        generate_plot(df, preds)
+        update_plot(df, preds)
 
 if __name__ == '__main__':
     main()