[DOCS] Improve MMM Case Study Notebook (#2116)

* custom saturation plots

* fix scaler

* rm wrong cells

* training optimization plots

* reorder

* bug

* rerun

Author: juanitorduz

docs/source/notebooks/mmm/mmm_case_study.ipynb

@@ -288,12 +288,25 @@ def _get_group_predictive_data(
                 f"Make sure the model has been fitted and the {group} has been sampled!"
             ) from e
 
-        if original_scale:
+        has_been_scaled = (
+            hasattr(self, "_posterior_predictive_samples_original_scale")
+            and self._posterior_predictive_samples_original_scale
+        )
+
+        if original_scale and not has_been_scaled:
             group_data = apply_sklearn_transformer_across_dim(
                 data=group_data,
                 func=self.get_target_transformer().inverse_transform,
                 dim_name="date",
             )
+
+        if not original_scale and has_been_scaled:
+            group_data = apply_sklearn_transformer_across_dim(
+                data=group_data,
+                func=self.get_target_transformer().transform,
+                dim_name="date",
+            )
+
         return group_data
 
     def _get_prior_predictive_data(self, original_scale: bool = False) -> Dataset:
@@ -707,38 +720,44 @@ def get_errors(self, original_scale: bool = False) -> DataArray:
                 "Make sure the model has been fitted and the posterior_predictive has been sampled!"
             ) from e
 
-        target_array = np.asarray(
-            transform_1d_array(self.get_target_transformer().transform, self.y)
-        )
-
-        if len(target_array) != len(posterior_predictive_data.date):
-            raise ValueError(
-                "The length of the target variable doesn't match the length of the date column. "
-                "If you are computing out-of-sample errors, please overwrite `self.y` with the "
-                "corresponding (non-transformed) target variable."
+        target_array = (
+            np.asarray(self.y)
+            if original_scale
+            else np.asarray(
+                transform_1d_array(self.get_target_transformer().transform, self.y)
             )
+        )
 
         target = (
             pd.Series(target_array, index=self.posterior_predictive.date)
             .rename_axis("date")
             .to_xarray()
         )
 
-        errors = (
+        if original_scale:
+            # If posterior predictive data is not in original scale, transform it:
+            if not hasattr(self, "_posterior_predictive_samples_original_scale"):
+                posterior_predictive_data = apply_sklearn_transformer_across_dim(
+                    data=posterior_predictive_data,
+                    func=self.get_target_transformer().inverse_transform,
+                    dim_name="date",
+                )
+        else:
+            # If posterior predictive data is in original scale, transform it back
+            # to the scaled space:
+            if hasattr(self, "_posterior_predictive_samples_original_scale"):
+                posterior_predictive_data = apply_sklearn_transformer_across_dim(
+                    data=posterior_predictive_data,
+                    func=self.get_target_transformer().transform,
+                    dim_name="date",
+                )
+
+        return (
             (target - posterior_predictive_data)[self.output_var]
             .rename("errors")
             .transpose(..., "date")
         )
 
-        if original_scale:
-            return apply_sklearn_transformer_across_dim(
-                data=errors,
-                func=self.get_target_transformer().inverse_transform,
-                dim_name="date",
-            )
-
-        return errors
-
     def plot_errors(
         self, original_scale: bool = False, ax: plt.Axes = None, **plt_kwargs: Any
     ) -> plt.Figure:

docs/source/notebooks/mmm/mmm_example.ipynb

@@ -2389,21 +2389,26 @@ def legend_title_func(channel):
         fig.suptitle("Direct response curves", fontsize=16)
         return fig
 
-    def _transform_to_original_scale_new(self, samples: DataArray) -> DataArray:
+    def _transform_to_original_scale_new(
+        self, samples: DataArray, var_names: list[str] | None = None
+    ) -> DataArray:
         """Transform samples to original scale using new scaling approach.
 
         Parameters
         ----------
         samples : DataArray
             Samples in scaled space
+        var_names : list[str] | None
+            Variable names requested in sampling.
 
         Returns
         -------
         DataArray
             Samples in original scale
         """
-        if self.output_var in samples:
-            samples[self.output_var] = samples[self.output_var] * self.target_scale
+        vars_to_transform = var_names if var_names is not None else [self.output_var]
+        for var_name in (v for v in vars_to_transform if v in samples):
+            samples[var_name] *= self.target_scale
         return samples
 
     def _transform_to_original_scale_legacy(
@@ -2527,12 +2532,13 @@ def sample_posterior_predictive(
 
         # Transform to original scale if requested
         if original_scale:
+            self._posterior_predictive_samples_original_scale = True
+            var_names = sample_posterior_predictive_kwargs.get("var_names")
             if self._has_new_scaling():
                 posterior_predictive_samples = self._transform_to_original_scale_new(
-                    posterior_predictive_samples
+                    posterior_predictive_samples, var_names
                 )
             else:
-                var_names = sample_posterior_predictive_kwargs.get("var_names")
                 posterior_predictive_samples = self._transform_to_original_scale_legacy(
                     posterior_predictive_samples, var_names
                 )