implemented 'add_mean' in plot_posterior_predictive

Author: jsnyde0

pymc_marketing/mmm/base.py

@@ -360,7 +360,11 @@ def plot_prior_predictive(self, **plt_kwargs: Any) -> plt.Figure:
         return fig
 
     def plot_posterior_predictive(
-        self, original_scale: bool = False, ax: plt.Axes = None, **plt_kwargs: Any
+        self,
+        original_scale: bool = False,
+        add_mean: bool = True,
+        ax: plt.Axes = None,
+        **plt_kwargs: Any,
     ) -> plt.Figure:
         """Plot posterior distribution from the model fit.
 
@@ -425,6 +429,23 @@ def plot_posterior_predictive(
                 label=f"{hdi_prob:.0%} HDI",
             )
 
+        if add_mean:
+            mean_prediction = posterior_predictive_data[self.output_var].mean(
+                dim=["chain", "draw"]
+            )
+
+            if original_scale:
+                mean_prediction = transform_1d_array(
+                    self.get_target_transformer().inverse_transform, mean_prediction
+                )
+
+            ax.plot(
+                np.asarray(posterior_predictive_data.date),
+                mean_prediction,
+                color="C0",
+                label="Mean Prediction",
+            )
+
         ax.plot(
             np.asarray(posterior_predictive_data.date),
             target_to_plot,